Model configuration and small prompt improvements (#24)

- Model can be configured using ENV
- Summarization prompt improved
- Single text request prompt improved
- Chat context prompt improved

Co-authored-by: Alexey Skobkin <skobkin-ru@ya.ru>
Reviewed-on: #24

Dockerfile

@@ -13,4 +13,11 @@ WORKDIR /app
 
 COPY --from=builder /build/app .
 
+# Do not forget "/v1" in the end
+ENV OPENAI_API_BASE_URL="" \
+    OPENAI_API_TOKEN="" \
+    TELEGRAM_TOKEN="" \
+    MODEL_TEXT_REQUEST="llama3.1:8b-instruct-q6_K" \
+    MODEL_SUMMARIZE_REQUEST="llama3.1:8b-instruct-q6_K"
+
 CMD ["/app/app"]

README.md

@@ -8,8 +8,10 @@
 
 ```shell
 docker run \
-  -e OLLAMA_TOKEN=123 \
-  -e OLLAMA_BASE_URL=http://ollama.localhost:11434/v1 \
+  -e OPENAI_API_TOKEN=123 \
+  -e OPENAI_API_BASE_URL=http://ollama.localhost:11434/v1 \
   -e TELEGRAM_TOKEN=12345 \
+  -e MODEL_TEXT_REQUEST=llama3.1:8b-instruct-q6_K
+  -e MODEL_TEXT_REQUEST=mistral-nemo:12b-instruct-2407-q4_K_M
   skobkin/telegram-llm-bot
 ```

bot/bot.go

@@ -23,16 +23,23 @@ type Bot struct {
 	llm       *llm.LlmConnector
 	extractor *extractor.Extractor
 	stats     *stats.Stats
+	models    ModelSelection
 
 	markdownV1Replacer *strings.Replacer
 }
 
-func NewBot(api *telego.Bot, llm *llm.LlmConnector, extractor *extractor.Extractor) *Bot {
+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,
 
 		markdownV1Replacer: strings.NewReplacer(
 			// https://core.telegram.org/bots/api#markdown-style
@@ -122,7 +129,7 @@ func (b *Bot) inlineHandler(bot *telego.Bot, update telego.Update) {
 			slog.Error("Cannot retrieve an article using extractor", "error", err)
 		}
 
-		llmReply, err := b.llm.Summarize(article.Text, llm.ModelLlama3Uncensored )
+		llmReply, err := b.llm.Summarize(article.Text, b.models.TextRequestModel)
 		if err != nil {
 			slog.Error("Cannot get reply from LLM connector")
 
@@ -148,7 +155,7 @@ func (b *Bot) inlineHandler(bot *telego.Bot, update telego.Update) {
 
 		requestContext := createLlmRequestContextFromUpdate(update)
 
-		llmReply, err := b.llm.HandleSingleRequest(iq.Query, llm.ModelLlama3Uncensored, requestContext)
+		llmReply, err := b.llm.HandleSingleRequest(iq.Query, b.models.TextRequestModel, requestContext)
 		if err != nil {
 			slog.Error("Cannot get reply from LLM connector")
 
@@ -194,7 +201,7 @@ func (b *Bot) heyHandler(bot *telego.Bot, update telego.Update) {
 
 	requestContext := createLlmRequestContextFromUpdate(update)
 
-	llmReply, err := b.llm.HandleSingleRequest(userMessage, llm.ModelLlama3Uncensored, requestContext)
+	llmReply, err := b.llm.HandleSingleRequest(userMessage, b.models.TextRequestModel, requestContext)
 	if err != nil {
 		slog.Error("Cannot get reply from LLM connector")
 
@@ -259,7 +266,7 @@ func (b *Bot) summarizeHandler(bot *telego.Bot, update telego.Update) {
 		slog.Error("Cannot retrieve an article using extractor", "error", err)
 	}
 
-	llmReply, err := b.llm.Summarize(article.Text, llm.ModelMistralUncensored)
+	llmReply, err := b.llm.Summarize(article.Text, b.models.SummarizeModel)
 	if err != nil {
 		slog.Error("Cannot get reply from LLM connector")
 

bot/models.go

@@ -0,0 +1,6 @@
+package bot
+
+type ModelSelection struct {
+	TextRequestModel string
+	SummarizeModel   string
+}

bot/request_context.go

@@ -44,10 +44,12 @@ func createLlmRequestContextFromUpdate(update telego.Update) llm.RequestContext
 	}
 
 	if !rc.Inline {
+		// TODO: implement retrieval of chat description
 		chat := message.Chat
 		rc.Chat = llm.ChatContext{
 			Title: chat.Title,
-			Description: chat.Description,
+			// TODO: fill when ChatFullInfo retrieved
+			//Description: chat.Description,
 			Type: chat.Type,
 		}
 	}

llm/llm.go

@@ -10,9 +10,6 @@ import (
 var (
 	ErrLlmBackendRequestFailed = errors.New("llm back-end request failed")
 	ErrNoChoices               = errors.New("no choices in LLM response")
-
-	ModelMistralUncensored = "dolphin-mistral:7b-v2.8-q4_K_M"
-	ModelLlama3Uncensored = "dolphin-llama3:8b-v2.9-q4_K_M"
 )
 
 type LlmConnector struct {
@@ -31,12 +28,12 @@ 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."
+	systemPrompt := "You're a bot in the Telegram chat.\n" +
+		"You're using a free model called \"" + model + "\".\n" +
+		"Currently you're not able to access chat history, so each message will be replied from a clean slate."
 
 	if !requestContext.Empty {
-		systemPrompt += " " + requestContext.Prompt()
+		systemPrompt += "\n" + requestContext.Prompt()
 	}
 
 	req := openai.ChatCompletionRequest{
@@ -79,9 +76,10 @@ func (l *LlmConnector) Summarize(text string, model string) (string, error) {
 			{
 				Role: openai.ChatMessageRoleSystem,
 				Content: "You're a text shortener. Give a very brief summary of the main facts " +
-					"point by point.  Format them as a list of bullet points. " +
+					"point by point.  Format them as a list of bullet points each starting with \"-\". " +
 					"Avoid any commentaries and value judgement on the matter. " +
-					"If possible, use the same language as the original text.",
+					"If possible, respond in the same language as the original text." +
+					"Do not use any non-ASCII characters.",
 			},
 		},
 	}
@@ -108,3 +106,37 @@ func (l *LlmConnector) Summarize(text string, model string) (string, error) {
 
 	return resp.Choices[0].Message.Content, nil
 }
+
+func (l *LlmConnector) GetModels() []string {
+	var result []string
+
+	models, err := l.client.ListModels(context.Background())
+	if err != nil {
+		slog.Error("llm: Model list request failed", "error", err)
+
+		return result
+	}
+
+	slog.Info("Model list retrieved", "models", models)
+
+	for _, model := range models.Models {
+		result = append(result, model.ID)
+	}
+
+	return result
+}
+
+func (l *LlmConnector) HasModel(id string) bool {
+	model, err := l.client.GetModel(context.Background(), id)
+	if err != nil {
+		slog.Error("llm: Model request failed", "error", err)
+	}
+
+	slog.Debug("llm: Returned model", "model", model)
+
+	if model.ID != "" {
+		return true
+	}
+
+	return false
+}

llm/request_context.go

@@ -39,15 +39,16 @@ func (c RequestContext) Prompt() string {
 		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 + "\". "
+	prompt += "User profile:" +
+		"First name: \"" + c.User.FirstName + "\"\n"
 	if c.User.Username != "" {
-		prompt += "Their username is @" + c.User.Username + ". "
+		prompt += "Username: @" + c.User.Username + ".\n"
 	}
 	if c.User.LastName != "" {
-		prompt += "Their last name is \"" + c.User.LastName + "\". "
+		prompt += "Last name: \"" + c.User.LastName + "\"\n"
 	}
 	if c.User.IsPremium {
-		prompt += "They have Telegram Premium subscription. "
+		prompt += "Telegram Premium subscription: active."
 	}
 
 	return prompt

main.go

@@ -12,12 +12,31 @@ import (
 )
 
 func main() {
-	ollamaToken := os.Getenv("OLLAMA_TOKEN")
-	ollamaBaseUrl := os.Getenv("OLLAMA_BASE_URL")
+	apiToken := os.Getenv("OPENAI_API_TOKEN")
+	apiBaseUrl := os.Getenv("OPENAI_API_BASE_URL")
+
+	models := bot.ModelSelection{
+		TextRequestModel: os.Getenv("MODEL_TEXT_REQUEST"),
+		SummarizeModel:   os.Getenv("MODEL_SUMMARIZE_REQUEST"),
+	}
+
+	slog.Info("Selected", "models", models)
 
 	telegramToken := os.Getenv("TELEGRAM_TOKEN")
 
-	llmc := llm.NewConnector(ollamaBaseUrl, ollamaToken)
+	llmc := llm.NewConnector(apiBaseUrl, apiToken)
+
+	slog.Info("Checking models availability")
+
+	for _, model := range []string{models.TextRequestModel, models.SummarizeModel} {
+		if !llmc.HasModel(model) {
+			slog.Error("Model not unavailable", "model", model)
+			os.Exit(1)
+		}
+	}
+
+	slog.Info("All needed models are available")
+
 	ext := extractor.NewExtractor()
 
 	telegramApi, err := tg.NewBot(telegramToken, tg.WithLogger(bot.NewLogger("telego: ")))
@@ -26,7 +45,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	botService := bot.NewBot(telegramApi, llmc, ext)
+	botService := bot.NewBot(telegramApi, llmc, ext, models)
 
 	err = botService.Run()
 	if err != nil {