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]*MessageHistory 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]*MessageHistory), 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.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.Handle(b.textMessageHandler, th.AnyMessageWithText()) 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", "MessageData 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) userMessageData := tgUserMessageToMessageData(message) llmReply, err := b.llm.HandleChatMessage( messageDataToLlmMessage(userMessageData), 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 \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) replyMarkdown := b.escapeMarkdownV1Symbols(llmReply) message := tu.Message( chatID, replyMarkdown, ).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) } b.saveBotReplyToHistory(update.Message, replyMarkdown) } 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 - Ask something from LLM\r\n"+ "/summarize - Summarize text from the provided link\r\n"+ "/s - 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) }