Building an AI Text Generation Bot in Go with Telegram and Cohere

Photo by bady abbas on Unsplash

Building an AI Text Generation Bot in Go with Telegram and Cohere

This article will explore how to build a Telegram chatbot using Golang and integrate it with the Cohere API for natural language generation.

This article will explore how to build a Telegram chatbot using Golang and integrate it with the Cohere API for natural language generation.

So, let's dive in and build our bot.

Get your Cohere API key🔑

  • Sign up for an account on Cohere's website (cohere.ai).

  • Once you have an account, follow these steps to obtain the API key:

    • Log in to your Cohere account.

    • Navigate to the API Key section in your account settings.

    • Generate an API key if you don't have one already.

    • Copy the generated API key to use in your code.

Set up your Telegram bot 🤖

  • Open the Telegram app and search for BotFather.

  • Start a chat with BotFather and follow the instructions.

  • Use the /newbot command to create a new bot.

  • Provide a name for your bot.

  • Choose a username for your bot (must end with "bot").

  • Once your bot is created, BotFather will provide you with a bot token. We'll keep that for later.

Before we begin, make sure you have Go installed on your machine. You can download and install it from the official Go website (golang.org).

Or you can use GitHub Codespaces like I did.

Setting Up Your Project 💻

  1. Create a new directory for your project and navigate to it.

  2. Create a main.go file for all your fancy code to go into.

    PS: This only works for Unix/Linux OS

     touch main.go
    

    On Windows, you can use the echo command to create a new file. Here's an example:

     echo. > main.go
    
  3. Initialize a new Go module using the following command:

     go mod init <module-name>
    
  4. Install the required dependencies by running the following commands in Terminal:

     go mod tidy
    

    This creates the go.mod and go.sum files.

    Or alternatively (not necessary if you do the above):

     go get github.com/NicoNex/echotron/v3
     go get github.com/cohere-ai/cohere-go
    

Let's take a look at the code 👀

Importing Dependencies:

  • Import the necessary packages and dependencies for our chatbot.

    • log: for logging messages

    • strings: for string manipulation

    • time: for time-related operations

    • echotron: library for the Telegram bot API in Go

    • cohere-go: Cohere, but in Go

package main

import (
    _ "embed"
    "log"
    "strings"
    "time"

    "github.com/NicoNex/echotron/v3"
    "github.com/cohere-ai/cohere-go"
)

In Terminal run:

go build

Type Definitions:

type stateFn func(*echotron.Update) stateFn

This is a recursive type definition. It allows functions to have the ability to return other functions of the same type, creating a recursive-like behaviour.

Here, we define a custom type stateFn that represents a function taking an *echotron.Update parameter and returning another stateFn function.

Bot Struct:

type bot struct {
    chatID int64
    echotron.API
    cohereClient *cohere.Client
    state        stateFn
}

We define a bot struct that encapsulates the necessary fields for our chatbot, such as:

  • chatID to store the ID of the chat; ensures that the bot replies to the correct chat when messages are sent

  • echotron.API to interact with the Telegram Bot API

  • cohereClient to interact with the Cohere API

  • state to keep track of the current state function.

API Keys and Commands:

var (
    //go:embed tgtoken
    token         string
    //go:embed chtoken
    cohereAPIKey  string

    commands      = []echotron.BotCommand{
        {Command: "/start", Description: "Activate the bot."},
        {Command: "/generate", Description: "Generate an answer."},
    }
)

Here, we declare variables for our API keys (token for the Telegram Bot API and cohereAPIKey for the Cohere API).

Additionally, we define an array of echotron.BotCommand structs to represent the bot commands and their descriptions.

Creating a New Bot:

func newBot(chatID int64) echotron.Bot {
    cohereClient, err := cohere.CreateClient(cohereAPIKey)
    if err != nil {
        log.Fatalln(err)
    }

    b := &bot{
        chatID:       chatID,
        API:          echotron.NewAPI(token),
        cohereClient: cohereClient,
    }
    b.state = b.handleMessage
    return b
}

This function is called by the Echotron dispatcher to create a new instance of the bot for each chat that interacts with it.

It initializes a new Cohere client, creates a bot instance with the provided chatID, sets up the Echotron API with the Telegram Bot API token, and sets the initial state to handleMessage.

Handling the Prompt:

func (b *bot) handlePrompt(update *echotron.Update) stateFn {
    b.SendChatAction(echotron.Typing, b.chatID, nil)
    response, err := b.generateText(message(update))
    if err != nil {
        log.Println("handlePrompt", err)
        b.SendMessage("An error occurred!", b.chatID, nil)
        return b.handleMessage
    }
    b.SendMessage(response, b.chatID, nil)
    return b.handleMessage
}

This method handles the user prompt and generates a response using the Cohere API.

It sends a "typing" chat action to indicate the bot is processing the request, generates a response based on the user prompt using the generateText method, and sends the response back to the chat.

If an error occurs during text generation, it logs the error, sends an error message, and transitions back to the handleMessage state.

Handling Messages:

func (b *bot) handleMessage(update *echotron.Update) stateFn {
    switch m := message(update); {
    case strings.HasPrefix(m, "/start"):
        b.SendMessage("Hi, I'm textjen. I use AI to answer your questions!", b.chatID, nil)
    case strings.HasPrefix(m, "/generate"):
        b.SendMessage("Please enter a prompt:", b.chatID, nil)
        return b.handlePrompt
    }
    return b.handleMessage
}

This method handles different message commands received by the bot.

  • It checks the message content using a switch statement.

  • If the message starts with "/start": it sends a "Hi, I'm textjen. I use AI to answer your questions!" message to the chat.

  • If the message starts with "/generate": it prompts the user to enter a prompt and transitions to the handlePrompt state.

  • If none of the conditions are met, it stays in the handleMessage state.

Updating the Bot State:

func (b *bot) Update(update *echotron.Update) {
    b.state = b.state(update)
}

This method is needed to implement the echotron.Bot interface. It updates the current state of the bot based on the received update.

Generating Text using Cohere API:

func (b *bot) generateText(prompt string) (string, error) {
    options := cohere.GenerateOptions{
        Model:             "command",
        Prompt:            prompt,
        MaxTokens:         300,
        Temperature:       0.9,
        K:                 0,
        StopSequences:     []string{},
        ReturnLikelihoods: "NONE",
    }

    response, err := b.cohereClient.Generate(options)
    if err != nil {
        return "", err
    }

    return response.Generations[0].Text, nil
}

This method generates text using Cohere's API.

  • It takes a prompt as input and returns the generated text or an error.

  • It creates a GenerateOptions struct with various options for text generation, such as the model, prompt, maximum tokens, temperature, and more.

  • It then calls the Generate method of the cohereClient to generate text based on the provided options.

  • If an error occurs during generation, it returns the error. Otherwise, it returns the generated text.

Extracting the Message:

func message(u *echotron.Update) string {
    if u.Message != nil {
        return u.Message.Text
    } else if u.EditedMessage != nil {
        return u.EditedMessage.Text
    } else if u.CallbackQuery != nil {
        return u.CallbackQuery.Data
    }
    return ""
}

This function extracts the message from the given update.

  • It checks the update for different message types (normal message, edited message, and callback query) and returns the corresponding text.

  • If none of the conditions are met, it returns an empty string.

Main Function:

func main() {
    echotron.NewAPI(token).SetMyCommands(nil, commands...)

    dsp := echotron.NewDispatcher(token, newBot)
    for {
        err := dsp.Poll()
        if err != nil {
            log.Println("Error polling updates:", err)
        }

        time.Sleep(5 * time.Second)
    }
}

The main function serves as the entry point for our chatbot.

  • It initializes Echotron with the Telegram Bot API token and sets the bot commands.

  • It creates a new dispatcher with the token and newBot function as parameters.

  • The dispatcher continuously polls for updates and handles them accordingly.

  • If an error occurs during polling, it logs the error.

  • The program sleeps for 5 seconds before retrying the polling process.

Protecting API Keys with .gitignore:

To protect your API keys and prevent them from being committed and pushed to a remote repository, you can use a .gitignore file. Here's how you can do it:

  • Create a new file named .gitignore in the root directory of your project (where the main.go file is located).

  • Open the .gitignore file and add the following lines:

      tgtoken
      chtoken
    

    These lines specify the files or patterns that Git should ignore.

  • Save the .gitignore file.

By adding tgtoken and chtoken to the .gitignore file, Git will ignore these binary files, and they won't be included in your repository.

Make sure to store your API keys in separate files (e.g., tgtoken and chtoken) and add them to the .gitignore file to keep them private.

Run your code:

go run main.go

I think I'm missing something here but go run main.go works so... 🤔

Test your bot:

Here's the main.go code:

package main

import (
    _ "embed"
    "log"
    "strings"
    "time"

    "github.com/NicoNex/echotron/v3"
    "github.com/cohere-ai/cohere-go"
)


type stateFn func(*echotron.Update) stateFn


type bot struct {
    chatID int64
    echotron.API
    cohereClient *cohere.Client
    state        stateFn
}

var (
    //go:embed tgtoken
    token string
    //go:embed chtoken
    cohereAPIKey string

    commands = []echotron.BotCommand{
        {Command: "/start", Description: "Activate the bot."},
        {Command: "/generate", Description: "Generate an answer."},
    }
)

func newBot(chatID int64) echotron.Bot {
    cohereClient, err := cohere.CreateClient(cohereAPIKey)
    if err != nil {
        log.Fatalln(err)
    }

    b := &bot{
        chatID:       chatID,
        API:          echotron.NewAPI(token),
        cohereClient: cohereClient,
    }
    b.state = b.handleMessage
    return b
}

func (b *bot) handlePrompt(update *echotron.Update) stateFn {
    b.SendChatAction(echotron.Typing, b.chatID, nil)
    response, err := b.generateText(message(update))
    if err != nil {
        log.Println("handlePrompt", err)
        b.SendMessage("An error occurred!", b.chatID, nil)
        return b.handleMessage
    }
    b.SendMessage(response, b.chatID, nil)
    return b.handleMessage
}

func (b *bot) handleMessage(update *echotron.Update) stateFn {
    switch m := message(update); {

    case strings.HasPrefix(m, "/start"):
        b.SendMessage("Hi, I'm textjen. I use AI to answer your questions!", b.chatID, nil)

    case strings.HasPrefix(m, "/generate"):
        b.SendMessage("Please enter a prompt:", b.chatID, nil)
        return b.handlePrompt
    }
    return b.handleMessage
}

// This method is needed to implement the echotron.Bot interface.
func (b *bot) Update(update *echotron.Update) {
    b.state = b.state(update)
}

// Generate text using the Cohere API
func (b *bot) generateText(prompt string) (string, error) {
    options := cohere.GenerateOptions{
        Model:             "command",
        Prompt:            prompt,
        MaxTokens:         300,
        Temperature:       0.9,
        K:                 0,
        StopSequences:     []string{},
        ReturnLikelihoods: "NONE",
    }

    response, err := b.cohereClient.Generate(options)
    if err != nil {
        return "", err
    }

    return response.Generations[0].Text, nil
}

// Returns the message from the given update.
func message(u *echotron.Update) string {
    if u.Message != nil {
        return u.Message.Text
    } else if u.EditedMessage != nil {
        return u.EditedMessage.Text
    } else if u.CallbackQuery != nil {
        return u.CallbackQuery.Data
    }
    return ""
}

func main() {
    echotron.NewAPI(token).SetMyCommands(nil, commands...)

    // This is the entry point of echotron library.
    dsp := echotron.NewDispatcher(token, newBot)
    for {
        err := dsp.Poll()
        if err != nil {
            log.Println("Error polling updates:", err)
        }

        // In case of connection issues, wait 5 seconds before trying to reconnect.
        time.Sleep(5 * time.Second)
    }
}

Conclusion 🎉

Awesome! We've created a text generation bot using Cohere's API and deployed it on Telegram.

You can now take this and extend the functionality of your chatbot.

  • You can find the code here.

  • You can learn more about Echotron here. (Super grateful that the developer behind the library took a look at the code, helped me make it better and showed me a couple of cool stuff!💜)

  • You can check out how to build a text summarization Telegram bot in Python using Cohere here.