From d890faf4612ecefa87457077a39f792a0a43b012 Mon Sep 17 00:00:00 2001 From: Alexey Skobkin Date: Wed, 13 Mar 2024 00:32:18 +0300 Subject: [PATCH 1/3] Fix #20 disallowing any URL except http:// and https://. Extracting helper methods to separate file. --- bot/bot.go | 32 ++-------------------------- bot/helpers.go | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 30 deletions(-) create mode 100644 bot/helpers.go diff --git a/bot/bot.go b/bot/bot.go index c12c070..f741a21 100644 --- a/bot/bot.go +++ b/bot/bot.go @@ -6,7 +6,6 @@ import ( th "github.com/mymmrac/telego/telegohandler" tu "github.com/mymmrac/telego/telegoutil" "log/slog" - "net/url" "strings" "telegram-ollama-reply-bot/extractor" "telegram-ollama-reply-bot/llm" @@ -153,9 +152,8 @@ func (b *Bot) summarizeHandler(bot *telego.Bot, update telego.Update) { return } - _, err := url.ParseRequestURI(args[1]) - if err != nil { - slog.Error("Provided URL is not valid", "url", args[1]) + 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, @@ -295,29 +293,3 @@ func (b *Bot) createLlmRequestContext(update telego.Update) llm.RequestContext { func (b *Bot) escapeMarkdownV1Symbols(input string) string { return b.markdownV1Replacer.Replace(input) } - -func (b *Bot) reply(originalMessage *telego.Message, newMessage *telego.SendMessageParams) *telego.SendMessageParams { - return newMessage.WithReplyParameters(&telego.ReplyParameters{ - MessageID: originalMessage.MessageID, - }) -} - -func (b *Bot) sendTyping(chatId telego.ChatID) { - slog.Debug("Setting 'typing' chat action") - - err := b.api.SendChatAction(tu.ChatAction(chatId, "typing")) - if err != nil { - slog.Error("Cannot set chat action", "error", err) - } -} - -func (b *Bot) trySendReplyError(message *telego.Message) { - if message == nil { - return - } - - _, _ = b.api.SendMessage(b.reply(message, tu.Message( - tu.ID(message.Chat.ID), - "Error occurred while trying to send reply.", - ))) -} diff --git a/bot/helpers.go b/bot/helpers.go new file mode 100644 index 0000000..3c20598 --- /dev/null +++ b/bot/helpers.go @@ -0,0 +1,57 @@ +package bot + +import ( + "github.com/mymmrac/telego" + "github.com/mymmrac/telego/telegoutil" + "log/slog" + "net/url" + "slices" + "strings" +) + +var ( + allowedUrlSchemes = []string{"http", "https"} +) + +func (b *Bot) reply(originalMessage *telego.Message, newMessage *telego.SendMessageParams) *telego.SendMessageParams { + return newMessage.WithReplyParameters(&telego.ReplyParameters{ + MessageID: originalMessage.MessageID, + }) +} + +func (b *Bot) sendTyping(chatId telego.ChatID) { + slog.Debug("Setting 'typing' chat action") + + err := b.api.SendChatAction(telegoutil.ChatAction(chatId, "typing")) + if err != nil { + slog.Error("Cannot set chat action", "error", err) + } +} + +func (b *Bot) trySendReplyError(message *telego.Message) { + if message == nil { + return + } + + _, _ = b.api.SendMessage(b.reply(message, telegoutil.Message( + telegoutil.ID(message.Chat.ID), + "Error occurred while trying to send reply.", + ))) +} + +func isValidAndAllowedUrl(text string) bool { + u, err := url.ParseRequestURI(text) + if err != nil { + slog.Debug("Provided text is not an URL", "text", text) + + return false + } + + if !slices.Contains(allowedUrlSchemes, strings.ToLower(u.Scheme)) { + slog.Debug("Provided URL has disallowed scheme", "scheme", u.Scheme, "allowed-schemes", allowedUrlSchemes) + + return false + } + + return true +} -- 2.43.5 From ca005a93707f7345d9d23455a9168f5bbf92babd Mon Sep 17 00:00:00 2001 From: Alexey Skobkin Date: Wed, 13 Mar 2024 00:32:52 +0300 Subject: [PATCH 2/3] Extracting request context creation to separate file. --- bot/bot.go | 33 ------------------------ bot/request_context.go | 58 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 33 deletions(-) create mode 100644 bot/request_context.go diff --git a/bot/bot.go b/bot/bot.go index f741a21..eaac998 100644 --- a/bot/bot.go +++ b/bot/bot.go @@ -257,39 +257,6 @@ func (b *Bot) statsHandler(bot *telego.Bot, update telego.Update) { } } -func (b *Bot) createLlmRequestContext(update telego.Update) llm.RequestContext { - message := update.Message - - rc := llm.RequestContext{} - - if message == nil { - slog.Debug("request context creation problem: no message provided. returning empty context.", "request-context", rc) - - return rc - } - - user := message.From - if user != nil { - rc.User = llm.UserContext{ - Username: user.Username, - FirstName: user.FirstName, - LastName: user.LastName, - IsPremium: user.IsPremium, - } - } - - chat := message.Chat - rc.Chat = llm.ChatContext{ - Title: chat.Title, - Description: chat.Description, - Type: chat.Type, - } - - slog.Debug("request context created", "request-context", rc) - - return rc -} - func (b *Bot) escapeMarkdownV1Symbols(input string) string { return b.markdownV1Replacer.Replace(input) } diff --git a/bot/request_context.go b/bot/request_context.go new file mode 100644 index 0000000..6c155ae --- /dev/null +++ b/bot/request_context.go @@ -0,0 +1,58 @@ +package bot + +import ( + "github.com/mymmrac/telego" + "log/slog" + "telegram-ollama-reply-bot/llm" +) + +func (b *Bot) createLlmRequestContext(update telego.Update) llm.RequestContext { + message := update.Message + iq := update.InlineQuery + + rc := llm.RequestContext{ + Empty: true, + Inline: false, + } + + switch { + case message == nil && iq == nil: + slog.Debug("request context creation problem: no message provided. returning empty context.", "request-context", rc) + + return rc + case iq != nil: + rc.Inline = true + } + + rc.Empty = false + + var user *telego.User + + if rc.Inline { + user = &iq.From + } else { + user = message.From + } + + if user != nil { + rc.User = llm.UserContext{ + Username: user.Username, + FirstName: user.FirstName, + LastName: user.LastName, + IsPremium: user.IsPremium, + } + } + + if !rc.Inline { + chat := message.Chat + rc.Chat = llm.ChatContext{ + Title: chat.Title, + Description: chat.Description, + Type: chat.Type, + } + } + + slog.Debug("request context created", "request-context", rc) + + return rc +} -- 2.43.5 From 7bb5c65d5939cc0164b1cb5cf6d220ca9620a790 Mon Sep 17 00:00:00 2001 From: Alexey Skobkin Date: Wed, 13 Mar 2024 01:18:01 +0300 Subject: [PATCH 3/3] Closes #14. Adding inline queries. Also small refactoring of context prompt based on RequestContext. --- bot/bot.go | 101 +++++++++++++++++++++++++++++++++++++++-- bot/helpers.go | 23 ++++++++-- bot/request_context.go | 2 +- llm/llm.go | 15 ++++-- llm/request_context.go | 27 +++++++---- stats/stats.go | 10 ++++ 6 files changed, 155 insertions(+), 23 deletions(-) diff --git a/bot/bot.go b/bot/bot.go index eaac998..0d6601d 100644 --- a/bot/bot.go +++ b/bot/bot.go @@ -74,18 +74,109 @@ func (b *Bot) Run() error { // Middlewares bh.Use(b.chatTypeStatsCounter) - // Handlers + // Command handlers bh.Handle(b.startHandler, th.CommandEqual("start")) bh.Handle(b.heyHandler, th.CommandEqual("hey")) bh.Handle(b.summarizeHandler, th.CommandEqual("summarize")) bh.Handle(b.statsHandler, th.CommandEqual("stats")) bh.Handle(b.helpHandler, th.CommandEqual("help")) + // Inline query handlers + bh.Handle(b.inlineHandler, th.AnyInlineQuery()) + bh.Start() return nil } +func (b *Bot) inlineHandler(bot *telego.Bot, update telego.Update) { + iq := update.InlineQuery + slog.Info("inline query received", "query", iq.Query) + + slog.Debug("query", "query", iq) + + if len(iq.Query) < 3 { + return + } + + b.stats.InlineQuery() + + queryParts := strings.SplitN(iq.Query, " ", 2) + + if len(queryParts) < 1 { + slog.Debug("Empty query. Skipping.") + + return + } + + var response *telego.AnswerInlineQueryParams + + switch isValidAndAllowedUrl(queryParts[0]) { + case true: + slog.Info("Inline /summarize request", "url", queryParts[0]) + + b.stats.SummarizeRequest() + + article, err := b.extractor.GetArticleFromUrl(queryParts[0]) + if err != nil { + slog.Error("Cannot retrieve an article using extractor", "error", err) + } + + llmReply, err := b.llm.Summarize(article.Text, llm.ModelMistralUncensored) + if err != nil { + slog.Error("Cannot get reply from LLM connector") + + b.trySendInlineQueryError(iq, "LLM request error. Try again later.") + + return + } + + slog.Debug("Got completion. Going to send.", "llm-completion", llmReply) + + response = tu.InlineQuery( + iq.ID, + tu.ResultArticle( + "reply_"+iq.ID, + "Summary for "+queryParts[0], + tu.TextMessage(b.escapeMarkdownV1Symbols(llmReply)).WithParseMode("Markdown"), + ), + ) + case false: + b.stats.HeyRequest() + + slog.Info("Inline /hey request", "text", iq.Query) + + requestContext := createLlmRequestContextFromUpdate(update) + + llmReply, err := b.llm.HandleSingleRequest(iq.Query, llm.ModelMistralUncensored, requestContext) + if err != nil { + slog.Error("Cannot get reply from LLM connector") + + b.trySendInlineQueryError(iq, "LLM request error. Try again later.") + + return + } + + slog.Debug("Got completion. Going to send.", "llm-completion", llmReply) + + response = tu.InlineQuery( + iq.ID, + tu.ResultArticle( + "reply_"+iq.ID, + "LLM reply to\""+iq.Query+"\"", + tu.TextMessage(b.escapeMarkdownV1Symbols(llmReply)).WithParseMode("Markdown"), + ), + ) + } + + err := bot.AnswerInlineQuery(response) + if err != nil { + slog.Error("Can't answer to inline query", "error", err) + + b.trySendInlineQueryError(iq, "Couldn't send intended reply, sorry") + } +} + func (b *Bot) heyHandler(bot *telego.Bot, update telego.Update) { slog.Info("/hey", "message-text", update.Message.Text) @@ -101,7 +192,7 @@ func (b *Bot) heyHandler(bot *telego.Bot, update telego.Update) { b.sendTyping(chatID) - requestContext := b.createLlmRequestContext(update) + requestContext := createLlmRequestContextFromUpdate(update) llmReply, err := b.llm.HandleSingleRequest(userMessage, llm.ModelMistralUncensored, requestContext) if err != nil { @@ -115,7 +206,7 @@ func (b *Bot) heyHandler(bot *telego.Bot, update telego.Update) { return } - slog.Debug("Got completion. Going to send.", "llm-reply", llmReply) + slog.Debug("Got completion. Going to send.", "llm-completion", llmReply) message := tu.Message( chatID, @@ -139,7 +230,7 @@ func (b *Bot) summarizeHandler(bot *telego.Bot, update telego.Update) { b.sendTyping(chatID) - args := strings.Split(update.Message.Text, " ") + args := strings.SplitN(update.Message.Text, " ", 2) if len(args) < 2 { _, _ = bot.SendMessage(tu.Message( @@ -180,7 +271,7 @@ func (b *Bot) summarizeHandler(bot *telego.Bot, update telego.Update) { return } - slog.Debug("Got completion. Going to send.", "llm-reply", llmReply) + slog.Debug("Got completion. Going to send.", "llm-completion", llmReply) message := tu.Message( chatID, diff --git a/bot/helpers.go b/bot/helpers.go index 3c20598..2fff50d 100644 --- a/bot/helpers.go +++ b/bot/helpers.go @@ -2,7 +2,7 @@ package bot import ( "github.com/mymmrac/telego" - "github.com/mymmrac/telego/telegoutil" + tu "github.com/mymmrac/telego/telegoutil" "log/slog" "net/url" "slices" @@ -22,7 +22,7 @@ func (b *Bot) reply(originalMessage *telego.Message, newMessage *telego.SendMess func (b *Bot) sendTyping(chatId telego.ChatID) { slog.Debug("Setting 'typing' chat action") - err := b.api.SendChatAction(telegoutil.ChatAction(chatId, "typing")) + err := b.api.SendChatAction(tu.ChatAction(chatId, "typing")) if err != nil { slog.Error("Cannot set chat action", "error", err) } @@ -33,12 +33,27 @@ func (b *Bot) trySendReplyError(message *telego.Message) { return } - _, _ = b.api.SendMessage(b.reply(message, telegoutil.Message( - telegoutil.ID(message.Chat.ID), + _, _ = b.api.SendMessage(b.reply(message, tu.Message( + tu.ID(message.Chat.ID), "Error occurred while trying to send reply.", ))) } +func (b *Bot) trySendInlineQueryError(iq *telego.InlineQuery, text string) { + if iq == nil { + return + } + + _ = b.api.AnswerInlineQuery(tu.InlineQuery( + iq.ID, + tu.ResultArticle( + string("error_"+iq.ID), + "Error: "+text, + tu.TextMessage(text), + ), + )) +} + func isValidAndAllowedUrl(text string) bool { u, err := url.ParseRequestURI(text) if err != nil { diff --git a/bot/request_context.go b/bot/request_context.go index 6c155ae..969abf6 100644 --- a/bot/request_context.go +++ b/bot/request_context.go @@ -6,7 +6,7 @@ import ( "telegram-ollama-reply-bot/llm" ) -func (b *Bot) createLlmRequestContext(update telego.Update) llm.RequestContext { +func createLlmRequestContextFromUpdate(update telego.Update) llm.RequestContext { message := update.Message iq := update.InlineQuery diff --git a/llm/llm.go b/llm/llm.go index df52d28..6b2d2c5 100644 --- a/llm/llm.go +++ b/llm/llm.go @@ -30,15 +30,20 @@ func NewConnector(baseUrl string, token string) *LlmConnector { } func (l *LlmConnector) HandleSingleRequest(text string, model string, requestContext RequestContext) (string, error) { + systemPrompt := "You're a bot in the Telegram chat. " + + "You're using a free model called \"" + model + "\". " + + "You see only messages addressed to you using commands due to privacy settings." + + if !requestContext.Empty { + systemPrompt += " " + requestContext.Prompt() + } + req := openai.ChatCompletionRequest{ Model: model, Messages: []openai.ChatCompletionMessage{ { - Role: openai.ChatMessageRoleSystem, - Content: "You're a bot in the Telegram chat. " + - "You're using a free model called \"" + model + "\". " + - "You see only messages addressed to you using commands due to privacy settings. " + - requestContext.Prompt(), + Role: openai.ChatMessageRoleSystem, + Content: systemPrompt, }, }, } diff --git a/llm/request_context.go b/llm/request_context.go index 12e72d4..fff4456 100644 --- a/llm/request_context.go +++ b/llm/request_context.go @@ -1,8 +1,10 @@ package llm type RequestContext struct { - User UserContext - Chat ChatContext + Empty bool + Inline bool + User UserContext + Chat ChatContext } type UserContext struct { @@ -19,13 +21,22 @@ type ChatContext struct { } func (c RequestContext) Prompt() string { - prompt := "The type of chat you're in is \"" + c.Chat.Type + "\". " - - if c.Chat.Title != "" { - prompt += "Chat is called \"" + c.Chat.Title + "\". " + if c.Empty { + return "" } - if c.Chat.Description != "" { - prompt += "Chat description is \"" + c.Chat.Description + "\". " + + prompt := "" + if !c.Inline { + prompt += "The type of chat you're in is \"" + c.Chat.Type + "\". " + + if c.Chat.Title != "" { + prompt += "Chat is called \"" + c.Chat.Title + "\". " + } + if c.Chat.Description != "" { + prompt += "Chat description is \"" + c.Chat.Description + "\". " + } + } else { + prompt += "You're responding to inline query, so you're not in the chat right now. " } prompt += "According to their profile, first name of the user who wrote you is \"" + c.User.FirstName + "\". " diff --git a/stats/stats.go b/stats/stats.go index 44f8eda..733d5c6 100644 --- a/stats/stats.go +++ b/stats/stats.go @@ -13,6 +13,7 @@ type Stats struct { GroupRequests uint64 PrivateRequests uint64 + InlineQueries uint64 HeyRequests uint64 SummarizeRequests uint64 @@ -24,6 +25,7 @@ func NewStats() *Stats { GroupRequests: 0, PrivateRequests: 0, + InlineQueries: 0, HeyRequests: 0, SummarizeRequests: 0, @@ -36,6 +38,7 @@ func (s *Stats) MarshalJSON() ([]byte, error) { GroupRequests uint64 `json:"group_requests"` PrivateRequests uint64 `json:"private_requests"` + InlineQueries uint64 `json:"inline_queries"` HeyRequests uint64 `json:"hey_requests"` SummarizeRequests uint64 `json:"summarize_requests"` @@ -44,6 +47,7 @@ func (s *Stats) MarshalJSON() ([]byte, error) { GroupRequests: s.GroupRequests, PrivateRequests: s.PrivateRequests, + InlineQueries: s.InlineQueries, HeyRequests: s.HeyRequests, SummarizeRequests: s.SummarizeRequests, @@ -59,6 +63,12 @@ func (s *Stats) String() string { return string(data) } +func (s *Stats) InlineQuery() { + s.mu.Lock() + defer s.mu.Unlock() + s.InlineQueries++ +} + func (s *Stats) GroupRequest() { s.mu.Lock() defer s.mu.Unlock() -- 2.43.5