feat: BYOK

internal/api/chat/create_conversation_message_stream_v2.go

@@ -2,12 +2,14 @@ package chat
 
 import (
 	"context"
+	"fmt"
 	"paperdebugger/internal/api/mapper"
 	"paperdebugger/internal/libs/contextutil"
 	"paperdebugger/internal/libs/shared"
 	"paperdebugger/internal/models"
 	"paperdebugger/internal/services"
 	chatv2 "paperdebugger/pkg/gen/api/chat/v2"
+	"strings"
 
 	"github.com/google/uuid"
 	"github.com/openai/openai-go/v3"
@@ -276,12 +278,50 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 		return s.sendStreamError(stream, err)
 	}
 
-	// Usage is the same as ChatCompletion, just passing the stream parameter
-	llmProvider := &models.LLMProviderConfig{
-		APIKey: settings.OpenAIAPIKey,
+	// Check if user has an API key for requested model
+	var llmProvider *models.LLMProviderConfig
+	var customModel *models.CustomModel
+	customModel = nil
+
+	customModelID := req.GetCustomModelId()
+	if customModelID != "" {
+		for i := range settings.CustomModels {
+			if settings.CustomModels[i].Id.Hex() == customModelID {
+				customModel = &settings.CustomModels[i]
+				break
+			}
+		}
+		if customModel == nil {
+			return s.sendStreamError(stream, fmt.Errorf("custom model not found: %q", customModelID))
+		}
+		modelSlug = customModel.Slug
+	}
+
+	if customModel == nil {
+		// User did not specify API key for this model
+		llmProvider = &models.LLMProviderConfig{
+			APIKey:        "",
+			IsCustomModel: false,
+		}
+	} else {
+		customModel.BaseUrl = strings.ToLower(customModel.BaseUrl)
+
+		if strings.Contains(customModel.BaseUrl, "paperdebugger.com") {
+			customModel.BaseUrl = ""
+		}
+		if !strings.HasPrefix(customModel.BaseUrl, "https://") {
+			customModel.BaseUrl = strings.Replace(customModel.BaseUrl, "http://", "", 1)
+			customModel.BaseUrl = "https://" + customModel.BaseUrl
+		}
+
+		llmProvider = &models.LLMProviderConfig{
+			APIKey:        customModel.APIKey,
+			Endpoint:      customModel.BaseUrl,
+			IsCustomModel: true,
+		}
 	}
 
-	openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.ID.Hex(), modelSlug, conversation.OpenaiChatHistoryCompletion, llmProvider)
+	openaiChatHistory, inappChatHistory, err := s.aiClientV2.ChatCompletionStreamV2(ctx, stream, conversation.ID.Hex(), modelSlug, conversation.OpenaiChatHistoryCompletion, llmProvider, customModel)
 	if err != nil {
 		return s.sendStreamError(stream, err)
 	}
@@ -307,7 +347,7 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 			for i, bsonMsg := range conversation.InappChatHistory {
 				protoMessages[i] = mapper.BSONToChatMessageV2(bsonMsg)
 			}
-			title, err := s.aiClientV2.GetConversationTitleV2(ctx, protoMessages, llmProvider)
+			title, err := s.aiClientV2.GetConversationTitleV2(ctx, protoMessages, llmProvider, modelSlug, customModel)
 			if err != nil {
 				s.logger.Error("Failed to get conversation title", "error", err, "conversationID", conversation.ID.Hex())
 				return

internal/api/chat/list_supported_models_v2.go

@@ -2,7 +2,6 @@ package chat
 
 import (
 	"context"
-	"strings"
 
 	"paperdebugger/internal/libs/contextutil"
 	chatv2 "paperdebugger/pkg/gen/api/chat/v2"
@@ -220,32 +219,35 @@ func (s *ChatServerV2) ListSupportedModels(
 		return nil, err
 	}
 
-	hasOwnAPIKey := strings.TrimSpace(settings.OpenAIAPIKey) != ""
-
 	var models []*chatv2.SupportedModel
-	for _, config := range allModels {
-		// Choose the appropriate slug based on whether user has their own API key.
-		//
-		// Some models are only available via OpenRouter; for those, slugOpenAI may be empty.
-		// In that case, keep using the OpenRouter slug to avoid returning an empty model slug.
-		slug := config.slugOpenRouter
-		if hasOwnAPIKey && strings.TrimSpace(config.slugOpenAI) != "" {
-			slug = config.slugOpenAI
-		}
 
+	for _, model := range settings.CustomModels {
+		modelID := model.Id.Hex()
+		models = append(models, &chatv2.SupportedModel{
+			Id:           &modelID,
+			Name:         model.Name,
+			Slug:         model.Slug,
+			TotalContext: int64(model.ContextWindow),
+			MaxOutput:    int64(model.MaxOutput),
+			InputPrice:   int64(model.InputPrice),
+			OutputPrice:  int64(model.OutputPrice),
+			IsCustom:     true,
+		})
+	}
+
+	for _, config := range allModels {
 		model := &chatv2.SupportedModel{
 			Name:         config.name,
-			Slug:         slug,
+			Slug:         config.slugOpenRouter,
 			TotalContext: config.totalContext,
 			MaxOutput:    config.maxOutput,
 			InputPrice:   config.inputPrice,
 			OutputPrice:  config.outputPrice,
 		}
 
 		// If model requires own key but user hasn't provided one, mark as disabled
-		if config.requireOwnKey && !hasOwnAPIKey {
-			model.Disabled = true
-			model.DisabledReason = stringPtr("Requires your own OpenAI API key. Configure it in Settings.")
+		if config.requireOwnKey {
+			continue
 		}
 
 		models = append(models, model)

internal/api/mapper/user.go

@@ -3,26 +3,75 @@ package mapper
 import (
 	"paperdebugger/internal/models"
 	userv1 "paperdebugger/pkg/gen/api/user/v1"
+
+	"go.mongodb.org/mongo-driver/v2/bson"
 )
 
 func MapProtoSettingsToModel(settings *userv1.Settings) *models.Settings {
+	// Map the slice of custom models
+	customModels := make([]models.CustomModel, len(settings.CustomModels))
+	for i, m := range settings.CustomModels {
+		var id bson.ObjectID
+
+		id, err := bson.ObjectIDFromHex(m.Id)
+		if err != nil {
+			id = bson.NewObjectID()
+		}
+
+		customModels[i] = models.CustomModel{
+			Id:                id,
+			Slug:              m.Slug,
+			Name:              m.Name,
+			BaseUrl:           m.BaseUrl,
+			APIKey:            m.ApiKey,
+			ContextWindow:     m.ContextWindow,
+			MaxOutput:         m.MaxOutput,
+			InputPrice:        m.InputPrice,
+			OutputPrice:       m.OutputPrice,
+			Temperature:       m.Temperature,
+			ParallelToolCalls: m.ParallelToolCalls,
+			Store:             m.Store,
+		}
+	}
+
 	return &models.Settings{
 		ShowShortcutsAfterSelection:  settings.ShowShortcutsAfterSelection,
 		FullWidthPaperDebuggerButton: settings.FullWidthPaperDebuggerButton,
-		EnableCitationSuggestion:             settings.EnableCitationSuggestion,
+		EnableCitationSuggestion:     settings.EnableCitationSuggestion,
 		FullDocumentRag:              settings.FullDocumentRag,
 		ShowedOnboarding:             settings.ShowedOnboarding,
 		OpenAIAPIKey:                 settings.OpenaiApiKey,
+		CustomModels:                 customModels,
 	}
 }
 
 func MapModelSettingsToProto(settings *models.Settings) *userv1.Settings {
+	// Map the slice back to Proto
+	customModels := make([]*userv1.CustomModel, len(settings.CustomModels))
+	for i, m := range settings.CustomModels {
+		customModels[i] = &userv1.CustomModel{
+			Id:                m.Id.Hex(),
+			Slug:              m.Slug,
+			Name:              m.Name,
+			BaseUrl:           m.BaseUrl,
+			ApiKey:            m.APIKey,
+			ContextWindow:     m.ContextWindow,
+			MaxOutput:         m.MaxOutput,
+			InputPrice:        m.InputPrice,
+			OutputPrice:       m.OutputPrice,
+			Temperature:       m.Temperature,
+			ParallelToolCalls: m.ParallelToolCalls,
+			Store:             m.Store,
+		}
+	}
+
 	return &userv1.Settings{
 		ShowShortcutsAfterSelection:  settings.ShowShortcutsAfterSelection,
 		FullWidthPaperDebuggerButton: settings.FullWidthPaperDebuggerButton,
-		EnableCitationSuggestion:             settings.EnableCitationSuggestion,
+		EnableCitationSuggestion:     settings.EnableCitationSuggestion,
 		FullDocumentRag:              settings.FullDocumentRag,
 		ShowedOnboarding:             settings.ShowedOnboarding,
 		OpenaiApiKey:                 settings.OpenAIAPIKey,
+		CustomModels:                 customModels,
 	}
 }

internal/models/llm_provider.go

@@ -2,10 +2,13 @@ package models
 
 // LLMProviderConfig holds the configuration for LLM API calls.
 // If both Endpoint and APIKey are empty, the system default will be used.
+// If IsCustomModel is true, the user-requested slug with corresponding
+// API keys and endpoint should 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/models/user.go

@@ -2,13 +2,29 @@ package models
 
 import "go.mongodb.org/mongo-driver/v2/bson"
 
+type CustomModel struct {
+	Id                bson.ObjectID `bson:"_id"`
+	Slug              string        `bson:"slug"`
+	Name              string        `bson:"name"`
+	BaseUrl           string        `bson:"base_url"`
+	APIKey            string        `bson:"api_key"`
+	ContextWindow     int32         `bson:"context_window"`
+	MaxOutput         int32         `bson:"max_output"`
+	InputPrice        int32         `bson:"input_price"`
+	OutputPrice       int32         `bson:"output_price"`
+	Temperature       float32       `bson:"temperature"`
+	ParallelToolCalls bool          `bson:"parallel_tool_calls"`
+	Store             bool          `bson:"store"`
+}
+
 type Settings struct {
-	ShowShortcutsAfterSelection  bool   `bson:"show_shortcuts_after_selection"`
-	FullWidthPaperDebuggerButton bool   `bson:"full_width_paper_debugger_button"`
-	EnableCitationSuggestion     bool   `bson:"enable_citation_suggestion"`
-	FullDocumentRag              bool   `bson:"full_document_rag"`
-	ShowedOnboarding             bool   `bson:"showed_onboarding"`
-	OpenAIAPIKey                 string `bson:"openai_api_key"`
+	ShowShortcutsAfterSelection  bool          `bson:"show_shortcuts_after_selection"`
+	FullWidthPaperDebuggerButton bool          `bson:"full_width_paper_debugger_button"`
+	EnableCitationSuggestion     bool          `bson:"enable_citation_suggestion"`
+	FullDocumentRag              bool          `bson:"full_document_rag"`
+	ShowedOnboarding             bool          `bson:"showed_onboarding"`
+	OpenAIAPIKey                 string        `bson:"openai_api_key"`
+	CustomModels                 []CustomModel `bson:"custom_models"`
 }
 
 type User struct {

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, messages OpenAIChatHistory, llmProvider *models.LLMProviderConfig) (OpenAIChatHistory, AppChatHistory, error) {
-	openaiChatHistory, inappChatHistory, err := a.ChatCompletionStreamV2(ctx, nil, "", modelSlug, messages, llmProvider)
+func (a *AIClientV2) ChatCompletionV2(ctx context.Context, modelSlug string, messages OpenAIChatHistory, llmProvider *models.LLMProviderConfig, customModel *models.CustomModel) (OpenAIChatHistory, AppChatHistory, error) {
+	openaiChatHistory, inappChatHistory, err := a.ChatCompletionStreamV2(ctx, nil, "", modelSlug, messages, llmProvider, customModel)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -54,7 +54,7 @@ func (a *AIClientV2) ChatCompletionV2(ctx context.Context, modelSlug string, mes
 //   - 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, 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, customModel *models.CustomModel) (OpenAIChatHistory, AppChatHistory, error) {
 	openaiChatHistory := messages
 	inappChatHistory := AppChatHistory{}
 
@@ -66,7 +66,7 @@ func (a *AIClientV2) ChatCompletionStreamV2(ctx context.Context, callbackStream
 	}()
 
 	oaiClient := a.GetOpenAIClient(llmProvider)
-	params := getDefaultParamsV2(modelSlug, a.toolCallHandler.Registry)
+	params := getDefaultParamsV2(modelSlug, a.toolCallHandler.Registry, customModel)
 
 	for {
 		params.Messages = openaiChatHistory

internal/services/toolkit/client/get_citation_keys.go

@@ -244,7 +244,7 @@ func (a *AIClientV2) GetCitationKeys(ctx context.Context, sentence string, userI
 	_, resp, err := a.ChatCompletionV2(ctx, "gpt-5.2", OpenAIChatHistory{
 		openai.SystemMessage("You are a helpful assistant that suggests relevant citation keys."),
 		openai.UserMessage(message),
-	}, llmProvider)
+	}, llmProvider, nil)
 
 	if err != nil {
 		return nil, err

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) (string, error) {
+func (a *AIClientV2) GetConversationTitleV2(ctx context.Context, inappChatHistory []*chatv2.Message, llmProvider *models.LLMProviderConfig, modelSlug string, customModel *models.CustomModel) (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())
@@ -29,10 +29,16 @@ func (a *AIClientV2) GetConversationTitleV2(ctx context.Context, inappChatHistor
 	message := strings.Join(messages, "\n")
 	message = fmt.Sprintf("%s\nBased on above conversation, generate a short, clear, and descriptive title that summarizes the main topic or purpose of the discussion. The title should be concise, specific, and use natural language. Avoid vague or generic titles. Use abbreviation and short words if possible. Use 3-5 words if possible. Give me the title only, no other text including any other words.", message)
 
-	_, resp, err := a.ChatCompletionV2(ctx, "gpt-5-nano", OpenAIChatHistory{
+	// Default model if user is not using their own
+	modelToUse := "gpt-5-nano"
+	if llmProvider.IsCustomModel {
+		modelToUse = modelSlug
+	}
+
+	_, 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)
+	}, llmProvider, customModel)
 	if err != nil {
 		return "", err
 	}