Update llmProvider

Author: kah-seng

internal/api/chat/create_conversation_message_stream_v2.go

@@ -277,6 +277,7 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 	}
 
 	// Check if user has an API key for requested model
+	var llmProvider *models.LLMProviderConfig
 	var customModel *models.CustomModel
 	customModel = nil
 	for _, m := range settings.CustomModels {
@@ -287,20 +288,22 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 	}
 
 	// Usage is the same as ChatCompletion, just passing the stream parameter
-	var llmProvider *models.LLMProviderConfig
+
 	if customModel == nil {
 		// User did not specify API key for this model
 		llmProvider = &models.LLMProviderConfig{
-			APIKey: "",
+			APIKey:        "",
+			IsCustomModel: false,
 		}
 	} else {
 		llmProvider = &models.LLMProviderConfig{
-			APIKey:   customModel.APIKey,
-			Endpoint: customModel.BaseUrl,
+			APIKey:        customModel.APIKey,
+			Endpoint:      customModel.BaseUrl,
+			IsCustomModel: true,
 		}
 	}
 
-	openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.ID.Hex(), modelSlug, customModel != nil, conversation.OpenaiChatHistoryCompletion, llmProvider)
+	openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.ID.Hex(), modelSlug, conversation.OpenaiChatHistoryCompletion, llmProvider)
 	if err != nil {
 		return s.sendStreamError(stream, err)
 	}
@@ -326,7 +329,7 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 			for i, bsonMsg := range conversation.InappChatHistory {
 				protoMessages[i] = mapper.BSONToChatMessageV2(bsonMsg)
 			}
-			title, err := s.aiClientV2.GetConversationTitleV2(ctx, protoMessages, llmProvider, modelSlug, customModel != nil)
+			title, err := s.aiClientV2.GetConversationTitleV2(ctx, protoMessages, llmProvider, modelSlug)
 			if err != nil {
 				s.logger.Error("Failed to get conversation title", "error", err, "conversationID", conversation.ID.Hex())
 				return

internal/models/llm_provider.go

@@ -3,9 +3,10 @@ package models
 // LLMProviderConfig holds the configuration for LLM API calls.
 // If both Endpoint and APIKey are empty, the system default will be used.
 type LLMProviderConfig struct {
-	Endpoint  string
-	APIKey    string
-	ModelName string
+	Endpoint      string
+	APIKey        string
+	ModelName     string
+	IsCustomModel bool
 }
 
 // IsCustom returns true if the user has configured custom LLM provider settings.

internal/services/toolkit/client/client_v2.go

@@ -32,18 +32,20 @@ func (a *AIClientV2) GetOpenAIClient(llmConfig *models.LLMProviderConfig) *opena
 	var Endpoint string = llmConfig.Endpoint
 	var APIKey string = llmConfig.APIKey
 
-	if Endpoint == "" {
-		if APIKey != "" {
-			// User provided their own API key, use the OpenAI-compatible endpoint
-			Endpoint = a.cfg.OpenAIBaseURL // standard openai base url
-		} else {
-			// suffix needed for cloudflare gateway
-			Endpoint = a.cfg.InferenceBaseURL + "/openrouter"
+	if !llmConfig.IsCustomModel {
+		if Endpoint == "" {
+			if APIKey != "" {
+				// User provided their own API key, use the OpenAI-compatible endpoint
+				Endpoint = a.cfg.OpenAIBaseURL // standard openai base url
+			} else {
+				// suffix needed for cloudflare gateway
+				Endpoint = a.cfg.InferenceBaseURL + "/openrouter"
+			}
 		}
-	}
 
-	if APIKey == "" {
-		APIKey = a.cfg.InferenceAPIKey
+		if APIKey == "" {
+			APIKey = a.cfg.InferenceAPIKey
+		}
 	}
 
 	opts := []option.RequestOption{

internal/services/toolkit/client/completion_v2.go

@@ -25,8 +25,8 @@ import (
 //  1. The full chat history sent to the language model (including any tool call results).
 //  2. The incremental chat history visible to the user (including tool call results and assistant responses).
 //  3. An error, if any occurred during the process.
-func (a *AIClientV2) ChatCompletionV2(ctx context.Context, modelSlug string, isCustomModel bool, messages OpenAIChatHistory, llmProvider *models.LLMProviderConfig) (OpenAIChatHistory, AppChatHistory, error) {
-	openaiChatHistory, inappChatHistory, err := a.ChatCompletionStreamV2(ctx, nil, "", modelSlug, isCustomModel, messages, llmProvider)
+func (a *AIClientV2) ChatCompletionV2(ctx context.Context, modelSlug string, messages OpenAIChatHistory, llmProvider *models.LLMProviderConfig) (OpenAIChatHistory, AppChatHistory, error) {
+	openaiChatHistory, inappChatHistory, err := a.ChatCompletionStreamV2(ctx, nil, "", modelSlug, messages, llmProvider)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -54,7 +54,7 @@ func (a *AIClientV2) ChatCompletionV2(ctx context.Context, modelSlug string, isC
 //   - If tool calls are required, it handles them and appends the results to the chat history, then continues the loop.
 //   - If no tool calls are needed, it appends the assistant's response and exits the loop.
 //   - Finally, it returns the updated chat histories and any error encountered.
-func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream chatv2.ChatService_CreateConversationMessageStreamServer, conversationId string, modelSlug string, isCustomModel bool, messages OpenAIChatHistory, llmProvider *models.LLMProviderConfig) (OpenAIChatHistory, AppChatHistory, error) {
+func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream chatv2.ChatService_CreateConversationMessageStreamServer, conversationId string, modelSlug string, messages OpenAIChatHistory, llmProvider *models.LLMProviderConfig) (OpenAIChatHistory, AppChatHistory, error) {
 	openaiChatHistory := messages
 	inappChatHistory := AppChatHistory{}
 
@@ -65,7 +65,7 @@ func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream
 		streamHandler.SendFinalization()
 	}()
 
-	if isCustomModel {
+	if llmProvider.IsCustomModel {
 		// e.g., Strip "google/" from "google/gemini-2.5-flash"
 		modelSlug = modelSlug[strings.Index(modelSlug, "/")+1:]
 	}

internal/services/toolkit/client/get_conversation_title_v2.go

@@ -13,7 +13,7 @@ import (
 	"github.com/samber/lo"
 )
 
-func (a *AIClientV2) GetConversationTitleV2(ctx context.Context, inappChatHistory []*chatv2.Message, llmProvider *models.LLMProviderConfig, modelSlug string, isCustomModel bool) (string, error) {
+func (a *AIClientV2) GetConversationTitleV2(ctx context.Context, inappChatHistory []*chatv2.Message, llmProvider *models.LLMProviderConfig, modelSlug string) (string, error) {
 	messages := lo.Map(inappChatHistory, func(message *chatv2.Message, _ int) string {
 		if _, ok := message.Payload.MessageType.(*chatv2.MessagePayload_Assistant); ok {
 			return fmt.Sprintf("Assistant: %s", message.Payload.GetAssistant().GetContent())
@@ -31,11 +31,11 @@ func (a *AIClientV2) GetConversationTitleV2(ctx context.Context, inappChatHistor
 
 	// Default model if user is not using their own
 	modelToUse := "gpt-5-nano"
-	if isCustomModel {
+	if llmProvider.IsCustomModel {
 		modelToUse = modelSlug
 	}
 
-	_, resp, err := a.ChatCompletionV2(ctx, modelToUse, isCustomModel, OpenAIChatHistory{
+	_, resp, err := a.ChatCompletionV2(ctx, modelToUse, OpenAIChatHistory{
 		openai.SystemMessage("You are a helpful assistant that generates a title for a conversation."),
 		openai.UserMessage(message),
 	}, llmProvider)