telegram-ollama-reply-bot/bot/bot.go

package bot

import (
	"errors"
	"github.com/mymmrac/telego"
	th "github.com/mymmrac/telego/telegohandler"
	tu "github.com/mymmrac/telego/telegoutil"
	"log/slog"
	"strings"
	"telegram-ollama-reply-bot/extractor"
	"telegram-ollama-reply-bot/llm"
	"telegram-ollama-reply-bot/stats"
)

var (
	ErrGetMe          = errors.New("cannot retrieve api user")
	ErrUpdatesChannel = errors.New("cannot get updates channel")
	ErrHandlerInit    = errors.New("cannot initialize handler")
)

type BotInfo struct {
	Id       int64
	Username string
	Name     string
}

type Bot struct {
	api       *telego.Bot
	llm       *llm.LlmConnector
	extractor *extractor.Extractor
	stats     *stats.Stats
	models    ModelSelection
	history   map[int64]*MessageRingBuffer
	profile   BotInfo

	markdownV1Replacer *strings.Replacer
}

func NewBot(
	api *telego.Bot,
	llm *llm.LlmConnector,
	extractor *extractor.Extractor,
	models ModelSelection,
) *Bot {
	return &Bot{
		api:       api,
		llm:       llm,
		extractor: extractor,
		stats:     stats.NewStats(),
		models:    models,
		history:   make(map[int64]*MessageRingBuffer),
		profile:   BotInfo{0, "", ""},

		markdownV1Replacer: strings.NewReplacer(
			// https://core.telegram.org/bots/api#markdown-style
			"_", "\\_",
			//"*", "\\*",
			//"`", "\\`",
			//"[", "\\[",
		),
	}
}

func (b *Bot) Run() error {
	botUser, err := b.api.GetMe()
	if err != nil {
		slog.Error("Cannot retrieve api user", "error", err)

		return ErrGetMe
	}

	slog.Info("Running api as", "id", botUser.ID, "username", botUser.Username, "name", botUser.FirstName, "is_bot", botUser.IsBot)

	b.profile = BotInfo{
		Id:       botUser.ID,
		Username: botUser.Username,
		Name:     botUser.FirstName,
	}

	updates, err := b.api.UpdatesViaLongPolling(nil)
	if err != nil {
		slog.Error("Cannot get update channel", "error", err)

		return ErrUpdatesChannel
	}

	bh, err := th.NewBotHandler(b.api, updates)
	if err != nil {
		slog.Error("Cannot initialize bot handler", "error", err)

		return ErrHandlerInit
	}

	defer bh.Stop()
	defer b.api.StopLongPolling()

	// Middlewares
	bh.Use(b.chatHistory)
	bh.Use(b.chatTypeStatsCounter)

	// Command handlers
	bh.Handle(b.textMessageHandler, th.AnyMessageWithText())
	bh.Handle(b.startHandler, th.CommandEqual("start"))
	bh.Handle(b.summarizeHandler, th.Or(th.CommandEqual("summarize"), th.CommandEqual("s")))
	bh.Handle(b.statsHandler, th.CommandEqual("stats"))
	bh.Handle(b.helpHandler, th.CommandEqual("help"))

	bh.Start()

	return nil
}

func (b *Bot) textMessageHandler(bot *telego.Bot, update telego.Update) {
	slog.Debug("/any-message")

	message := update.Message

	switch {
	// Mentions
	case b.isMentionOfMe(update):
		slog.Info("/any-message", "type", "mention")
		b.processMention(message)
	// Replies
	case b.isReplyToMe(update):
		slog.Info("/any-message", "type", "reply")
		b.processMention(message)
	// Private chat
	case b.isPrivateWithMe(update):
		slog.Info("/any-message", "type", "private")
		b.processMention(message)
	default:
		slog.Debug("/any-message", "info", "Message is not mention, reply or private chat. Skipping.")
	}
}

func (b *Bot) processMention(message *telego.Message) {
	b.stats.Mention()

	slog.Info("/mention", "chat", message.Chat.ID)

	chatID := tu.ID(message.Chat.ID)

	b.sendTyping(chatID)

	requestContext := b.createLlmRequestContextFromMessage(message)

	llmReply, err := b.llm.HandleChatMessage(message.Text, b.models.TextRequestModel, requestContext)
	if err != nil {
		slog.Error("Cannot get reply from LLM connector")

		_, _ = b.api.SendMessage(b.reply(message, tu.Message(
			chatID,
			"LLM request error. Try again later.",
		)))

		return
	}

	slog.Debug("Got completion. Going to send.", "llm-completion", llmReply)

	reply := tu.Message(
		chatID,
		b.escapeMarkdownV1Symbols(llmReply),
	).WithParseMode("Markdown")

	_, err = b.api.SendMessage(b.reply(message, reply))
	if err != nil {
		slog.Error("Can't send reply message", "error", err)

		b.trySendReplyError(message)

		return
	}

	b.saveBotReplyToHistory(message, llmReply)
}

func (b *Bot) summarizeHandler(bot *telego.Bot, update telego.Update) {
	slog.Info("/summarize", "message-text", update.Message.Text)

	b.stats.SummarizeRequest()

	chatID := tu.ID(update.Message.Chat.ID)

	b.sendTyping(chatID)

	args := strings.SplitN(update.Message.Text, " ", 2)

	if len(args) < 2 {
		_, _ = bot.SendMessage(tu.Message(
			tu.ID(update.Message.Chat.ID),
			"Usage: /summarize <link>\r\n\r\n"+
				"Example:\r\n"+
				"/summarize https://kernel.org/get-notifications-for-your-patches.html",
		))

		return
	}

	if !isValidAndAllowedUrl(args[1]) {
		slog.Error("Provided text is not a valid URL", "text", args[1])

		_, _ = b.api.SendMessage(b.reply(update.Message, tu.Message(
			chatID,
			"URL is not valid.",
		)))

		return
	}

	article, err := b.extractor.GetArticleFromUrl(args[1])
	if err != nil {
		slog.Error("Cannot retrieve an article using extractor", "error", err)
	}

	llmReply, err := b.llm.Summarize(article.Text, b.models.SummarizeModel)
	if err != nil {
		slog.Error("Cannot get reply from LLM connector")

		_, _ = b.api.SendMessage(b.reply(update.Message, tu.Message(
			chatID,
			"LLM request error. Try again later.",
		)))

		return
	}

	slog.Debug("Got completion. Going to send.", "llm-completion", llmReply)

	message := tu.Message(
		chatID,
		b.escapeMarkdownV1Symbols(llmReply),
	).WithParseMode("Markdown")

	_, err = bot.SendMessage(b.reply(update.Message, message))

	if err != nil {
		slog.Error("Can't send reply message", "error", err)

		b.trySendReplyError(update.Message)
	}
}

func (b *Bot) helpHandler(bot *telego.Bot, update telego.Update) {
	slog.Info("/help")

	chatID := tu.ID(update.Message.Chat.ID)

	b.sendTyping(chatID)

	_, err := bot.SendMessage(b.reply(update.Message, tu.Messagef(
		chatID,
		"Instructions:\r\n"+
			"/hey <text> - Ask something from LLM\r\n"+
			"/summarize <link> - Summarize text from the provided link\r\n"+
			"/s <link> - Shorter version\r\n"+
			"/help - Show this help\r\n\r\n"+
			"Mention bot or reply to it's message to communicate with it",
	)))
	if err != nil {
		slog.Error("Cannot send a message", "error", err)

		b.trySendReplyError(update.Message)
	}
}

func (b *Bot) startHandler(bot *telego.Bot, update telego.Update) {
	slog.Info("/start")

	chatID := tu.ID(update.Message.Chat.ID)

	b.sendTyping(chatID)

	_, err := bot.SendMessage(b.reply(update.Message, tu.Message(
		chatID,
		"Hey!\r\n"+
			"Check out /help to learn how to use this bot.",
	)))
	if err != nil {
		slog.Error("Cannot send a message", "error", err)

		b.trySendReplyError(update.Message)
	}
}

func (b *Bot) statsHandler(bot *telego.Bot, update telego.Update) {
	slog.Info("/stats")

	chatID := tu.ID(update.Message.Chat.ID)

	b.sendTyping(chatID)

	_, err := bot.SendMessage(b.reply(update.Message, tu.Message(
		chatID,
		"Current bot stats:\r\n"+
			"```json\r\n"+
			b.stats.String()+"\r\n"+
			"```",
	)).WithParseMode("Markdown"))
	if err != nil {
		slog.Error("Cannot send a message", "error", err)

		b.trySendReplyError(update.Message)
	}
}

func (b *Bot) escapeMarkdownV1Symbols(input string) string {
	return b.markdownV1Replacer.Replace(input)
}