Merge pull request 'URL scheme whitelist and Inline Queries.' (#21) from fix_chat_type_middleware_nil_pointer into main

Reviewed-on: #21

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"
@@ -75,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)
 
@@ -102,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 {
@@ -116,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,
@@ -140,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(
@@ -153,9 +243,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,
@@ -182,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,
@@ -259,65 +348,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)
 }
-
-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.",
-	)))
-}

bot/helpers.go

@@ -0,0 +1,72 @@
+package bot
+
+import (
+	"github.com/mymmrac/telego"
+	tu "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(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.",
+	)))
+}
+
+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 {
+		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
+}

bot/request_context.go

@@ -0,0 +1,58 @@
+package bot
+
+import (
+	"github.com/mymmrac/telego"
+	"log/slog"
+	"telegram-ollama-reply-bot/llm"
+)
+
+func createLlmRequestContextFromUpdate(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
+}

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,
 			},
 		},
 	}

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 + "\". "

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()