Building a Chatbot using the Go SDK for Google Generative AI

With the recent AI buzz that companies have been pushing as of late, I wanted to join in and try out Google AI Studio to build a simple chatbot that you can interface via a cli. Hopefully with this guide, you can understand how simple it really is to interact with these large language models.

Here are some helpful resources that I used to get me going:

1. Go SDK for Google Generative AI
2. Obtaining an API Key
3. Google AI Documentation

You do need an API key, but thankfully their is a free tier to play around with!

Chatbot CLI

Here is a quick demo of what we will be building:

All the code is located here: https://github.com/kavinaravind/go-genai

Code Walkthrough

We will be focusing on two packages:

1. client: Comprising of the genai.Client with methods to start a chat session.
2. main: Starting the chat session and providing I/O functionality.

Client

package client

import (
	"bufio"
	"context"
	"fmt"
	"time"

	"github.com/google/generative-ai-go/genai"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

// The Gemini 1.5 models are versatile and work with both text-only and multimodal prompts.
const model = "gemini-1.5-flash"

// GenAIClient is a client for the Generative AI API.
type GenAIClient struct {
	client *genai.Client
}

// NewGenAIClient creates a new Generative AI client.
func NewGenAIClient(ctx context.Context, apiKey string) (*GenAIClient, error) {
	client, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
	if err != nil {
		return nil, err
	}
	return &GenAIClient{client: client}, nil
}

// StartNewChatSession starts a new chat session.
func (c *GenAIClient) StartNewChatSession(ctx context.Context, writer *bufio.Writer) func(parts ...genai.Part) error {
	// Creates a new instance of the named generative model
	model := c.client.GenerativeModel(model)

	// Start a chat session
	cs := model.StartChat()

	// Return a closure that can send messages and handle responses
	return func(parts ...genai.Part) error {
		// Retrieve a streaming request
		iter := cs.SendMessageStream(ctx, parts...)

		// Iterate over the responses
		for {
			resp, err := iter.Next()
			if err != nil {
				if err == iterator.Done {
					break
				}
				return err
			}
			printResponse(writer, resp)
		}

		return nil
	}
}

// Close closes the client
func (c *GenAIClient) Close() error {
	return c.client.Close()
}

// printResponse prints the response to the writer
func printResponse(writer *bufio.Writer, resp *genai.GenerateContentResponse) {
	for _, cand := range resp.Candidates {
		if cand.Content != nil {
			for _, part := range cand.Content.Parts {
				response := fmt.Sprintf("%s", part)
				for _, char := range response {
					fmt.Fprintf(writer, "%c", char)
					writer.Flush()
					time.Sleep(15 * time.Millisecond)
				}
			}
		}
	}
}

Lets break this package down from the top:

1. The model we are using is gemini-1.5-flash and currently there are two other models to choose from.
2. The GenAIClient struct holds a pointer to a genai.Client that will be used to interface with the API.
3. We need to pass in the api key to NewGenAIClient which will be pulled in from main via a flag.
4. StartNewChatSession is the main focus of this package. As the name implies, it is a method that starts a new chat session
   • GenerativeModel creates a new instance of the named generative model.
   • StartChat will start a chat and preserve history.
   • We are leveraging a closure that can send messages and handle responses using the *genai.ChatSession and context!
   • SendMessageStream returns an iterator that we can iterate over via Next().
5. printResponse is a utility function that parses *genai.GenerateContentResponse and writes to a buffer.
   • time.Sleep(15 * time.Millisecond) is used to mimic typing to os.Stdout.

Main

Now let's walk through the main package which instantiates the genai client and sets up I/O.

package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/google/generative-ai-go/genai"
	"github.com/kavinaravind/go-genai/client"
)

// welcomeMessage is the welcome message displayed to the user when the chatbot starts.
const welcomeMessage = `Welcome!
You can start chatting by typing your questions or statements and pressing Enter.
Type 'fresh' to start a new chat session.
Type 'exit' to quit the chatbot.

ChatBot: Hello! Ask me anything.`

func main() {
	apiKey := flag.String("api-key", "", "API key for the Generative AI API")
	flag.Parse()

	if *apiKey == "" {
		log.Fatal("Please provide an API key with the -api-key flag")
	}

	ctx := context.Background()

	// Create a new Generative AI client
	client, err := client.NewGenAIClient(ctx, *apiKey)
	if err != nil {
		log.Fatal(err)
	}
	defer func() {
		err := client.Close()
		if err != nil {
			log.Fatal(err)
		}
	}()

	// Create a new reader and writer
	reader := bufio.NewReader(os.Stdin)
	writer := bufio.NewWriter(os.Stdout)

	// Start a new chat session
	chat := client.StartNewChatSession(ctx, writer)

	// Start chatting
	fmt.Println(welcomeMessage)
	for {
		fmt.Print("You: ")

		input, err := reader.ReadString('\n')
		if err != nil {
			log.Fatal(err)
		}
		input = strings.TrimSpace(input)

		switch input {
		case "exit":
			fmt.Println("ChatBot: Goodbye!")
			os.Exit(0)
		case "fresh":
			chat = client.StartNewChatSession(ctx, writer)
			fmt.Println("ChatBot: Chat History Cleared.")
			continue
		default:
			fmt.Print("ChatBot: ")
			err = chat(genai.Text(input))
			if err != nil {
				log.Fatal(err)
			}
		}
	}
}

Lets break this package down from the top:

1. welcomeMessage defines a string that will output from the jump.
2. We use flag from the standard library to retrieve the API to pass into the client.
3. NewReader and NewWriter is used to read from std in and write to std out.
4. Lastly we create a chat loop to read input, process that input, and write the output.
   • The switch case is used to handle commands:
     • exit: To exit the application
     • fresh: To clear chat history

---

As you can see, with very little code we can create a functioning chatbot!

Goals

1. Dabble with the Go SDK for Google Generative AI to learn a bit more about the various Go Types:
   • Client: A Client is a Google generative AI client.
   • Candidate: Candidate is a response candidate generated from the model.
   • Part: A Part is either a Text, a Blob or a FunctionResponse.
   • Text: A Text is a piece of text, like a question or phrase.
2. Make use of an effective closure that references variables from outside its body via the chat session shown in the client package.
3. See how easy it can be to interface with Google Generative AI.

I hope this helps! Thanks for reading. :)