Merge pull request #33 from mervyn-teo/model-seperation

Seperated model used for various tasks

Author: mervyn-teo

cmd/api/main.go

@@ -16,7 +16,7 @@ import (
 )
 
 var (
-	client   *openai.Client
+	client   *router.Clients
 	settings storage.Settings
 	messages map[string][]openai.ChatCompletionMessage
 )
@@ -37,7 +37,11 @@ func init() {
 	messages = storage.ReadChatHistory(settings.ChatHistoryFilePath)
 
 	// Create the OpenAI client
-	client, err = router.CreateClient(settings.Model, settings.ApiKey)
+	client = &router.Clients{}
+	client.BaseClient, err = router.CreateClient(settings.Model, settings.ApiKey)
+	client.CompressionClient, err = router.CreateClient(settings.CompressionModel, settings.ApiKey)
+	client.ImageClient, err = router.CreateClient(settings.ImageModel, settings.ApiKey)
+
 	if err != nil {
 		log.Fatalf("Failed to create OpenAI client: %v", err)
 	}

internal/router/chat.go

@@ -215,7 +215,7 @@ func SplitString(s string, chunkSize int) []string {
 // MessageLoop listens for messages from the bot and processes them
 // It handles user messages, tool calls, and manages the conversation history
 // It also handles reminders and music-related commands
-func MessageLoop(ctx context.Context, Mybot *bot.Bot, client *openai.Client, messageChannel chan *bot.MessageForCompletion, instructions string, messages map[string][]ChatCompletionMessage, chatFilepath string, initSystemMessage string) {
+func MessageLoop(ctx context.Context, Mybot *bot.Bot, client *Clients, messageChannel chan *bot.MessageForCompletion, instructions string, messages map[string][]ChatCompletionMessage, chatFilepath string, initSystemMessage string) {
 	initialSystemMessage = initSystemMessage
 
 	// Load reminders and song map from files
@@ -274,7 +274,7 @@ func MessageLoop(ctx context.Context, Mybot *bot.Bot, client *openai.Client, mes
 			if userInput.Message.Attachments != nil {
 				for i, attachment := range userInput.Message.Attachments {
 					if strings.HasPrefix(attachment.ContentType, "image/") {
-						imageDescriptions.WriteString("{\n\"index\" : " + strconv.Itoa(i) + ",\n\"description\": " + getImageDescription(client, attachment, Mybot) + "}, \n")
+						imageDescriptions.WriteString("{\n\"index\" : " + strconv.Itoa(i) + ",\n\"description\": " + getImageDescription(client.ImageClient, attachment, Mybot) + "}, \n")
 					}
 				}
 			}
@@ -317,7 +317,7 @@ func MessageLoop(ctx context.Context, Mybot *bot.Bot, client *openai.Client, mes
 
 			storage.SaveChatHistory(messages, chatFilepath)
 
-			aiResponseContent, err := SendMessage(client, &updatedMessages, Mybot)
+			aiResponseContent, err := SendMessage(client.BaseClient, &updatedMessages, Mybot)
 
 			if err != nil {
 				log.Printf("Error getting response from OpenRouter: %v", err)
@@ -335,7 +335,7 @@ func MessageLoop(ctx context.Context, Mybot *bot.Bot, client *openai.Client, mes
 			// Compress messages to reduce length
 			if len(messages[userID]) > MaxMessagesToKeep {
 				log.Printf("Compressing messages for user %s to reduce length", userID)
-				messages[userID] = compressMsg(client, messages[userID], Mybot)
+				messages[userID] = compressMsg(client.CompressionClient, messages[userID], Mybot)
 			}
 
 			storage.SaveChatHistory(messages, chatFilepath)

internal/router/chat_test.go

@@ -485,7 +485,9 @@ func TestMessageLoop(t *testing.T) {
 	messages := make(map[string][]ChatCompletionMessage)
 
 	// Start message loop in goroutine
-	go MessageLoop(ctx, &bot.Bot{}, client, messageChannel, "Test instructions", messages, tempChatFile.Name(), "test init message")
+	clients := &Clients{}
+	clients.BaseClient = client
+	go MessageLoop(ctx, &bot.Bot{}, clients, messageChannel, "Test instructions", messages, tempChatFile.Name(), "test init message")
 
 	// Test normal message
 	testMessage := &bot.MessageForCompletion{

internal/router/client.go

@@ -7,6 +7,12 @@ import (
 	openai "github.com/sashabaranov/go-openai"
 )
 
+type Clients struct {
+	BaseClient        *openai.Client
+	ImageClient       *openai.Client
+	CompressionClient *openai.Client
+}
+
 const (
 	OpenRouterBaseURL = "https://openrouter.ai/api/v1"
 	RefererURL        = "http://localhost"

internal/storage/settings.go

@@ -8,11 +8,17 @@ import (
 	"os"
 )
 
+const (
+	DefaultModel = "google/gemini-2.5-flash"
+)
+
 type Settings struct {
 	ApiKey              string `json:"api_key"`
 	DiscordToken        string `json:"discord_bot_token"`
 	Instructions        string `json:"instructions"`
 	Model               string `json:"model"`
+	CompressionModel    string `json:"compression_model"`
+	ImageModel          string `json:"image_model"`
 	NewsAPIToken        string `json:"news_api_key"`
 	YoutubeToken        string `json:"youtube_api_key"`
 	ChatHistoryFilePath string `json:"chat_history_file_path"`
@@ -49,6 +55,10 @@ func LoadSettings(filePath string) (Settings, error) {
 		return Setting, fmt.Errorf("error decoding settings JSON: %w", err)
 	}
 
+	if Setting.Model == "" {
+		Setting.Model = DefaultModel
+		log.Println("WARNING: no model found in settings file, setting default model: " + DefaultModel)
+	}
 	if Setting.ApiKey == "" {
 		return Setting, fmt.Errorf("API key is missing in settings file")
 	}
@@ -61,6 +71,18 @@ func LoadSettings(filePath string) (Settings, error) {
 	if Setting.YoutubeToken == "" {
 		return Setting, fmt.Errorf("Youtube API key is missing in settings file")
 	}
+	if Setting.CompressionModel == "" {
+		Setting.CompressionModel = Setting.Model
+		log.Println("Compression Model not set, setting to: ", Setting.CompressionModel)
+	}
+	if Setting.ImageModel == "" {
+		Setting.ImageModel = Setting.Model
+		log.Println("Image Model not set, setting to: ", Setting.CompressionModel)
+	}
+
+	log.Println("Base model: " + Setting.Model)
+	log.Println("Compression model: " + Setting.CompressionModel)
+	log.Println("Image model: " + Setting.ImageModel)
 
 	return Setting, nil
 }

settings.example.json

@@ -2,7 +2,9 @@
   "api_key": "[Your OpenAI API Key]",
   "discord_bot_token": "[Your Discord Bot Token]",
   "instructions": "You are a Discord bot with an anime cat girl personality that helps the users with what they want to know. You like to use the UWU speak. Reply users in the language you've been spoken to. Describe yourself as a \"uwu cute machine\". Do not mention that you are a AI. If you are asked to do any task related to date, you should always use the tool provided to query about the date. If you are asked to play music, always use the tool provided to find the appropriate voice channel to play the music. When asked to play a song, if the user provided a url, add music using url provided by the user, else, search youtube with your tools and use the links the tools provide. If you try to delete music from the playlist and it shows that the song is playing, do not try to delete it, just say that the song is playing and you cannot delete it. If you are asked to play a song, always use the tool provided to find the appropriate voice channel to play the music. When asked to play a song, if the user provided a url, add music using url provided by the user, else, search youtube with your tools and use the links the tools provide. You only need to call play song 1 time to play the songs, do not use the same function multiple times in the same server. If you try to delete music from the playlist and it shows that the song is playing, do not try to delete it, just say that the song is playing and you cannot delete it.",
-  "model": "openai/gpt-4.1-mini",
+  "model": "deepseek/deepseek-chat-v3-0324",
+  "compression_model": "deepseek/deepseek-chat-v3-0324",
+  "image_model": "google/gemini-2.5-pro",
   "youtube_api_key": "[Your YouTube API Key]",
   "news_api_key": "[Your News API Key]",
   "chat_history_file_path" : "chat_history.json"