Skip to content

Commit 56d815e

Browse files
authored
fix: BYOK (#150)
### Key changes: - Select models by ID instead of slug so users can have multiple API keys for the same slug - Add loading spinner for save/edit actions
1 parent 56cbe48 commit 56d815e

File tree

13 files changed

+178
-61
lines changed

13 files changed

+178
-61
lines changed

internal/api/chat/create_conversation_message_stream_v2.go

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package chat
22

33
import (
44
"context"
5+
"fmt"
56
"paperdebugger/internal/api/mapper"
67
"paperdebugger/internal/libs/contextutil"
78
"paperdebugger/internal/libs/shared"
@@ -281,14 +282,21 @@ func (s *ChatServerV2) CreateConversationMessageStream(
281282
var llmProvider *models.LLMProviderConfig
282283
var customModel *models.CustomModel
283284
customModel = nil
284-
for i := range settings.CustomModels {
285-
if settings.CustomModels[i].Slug == modelSlug {
286-
customModel = &settings.CustomModels[i]
285+
286+
customModelID := req.GetCustomModelId()
287+
if customModelID != "" {
288+
for i := range settings.CustomModels {
289+
if settings.CustomModels[i].Id.Hex() == customModelID {
290+
customModel = &settings.CustomModels[i]
291+
break
292+
}
287293
}
294+
if customModel == nil {
295+
return s.sendStreamError(stream, fmt.Errorf("custom model not found: %q", customModelID))
296+
}
297+
modelSlug = customModel.Slug
288298
}
289299

290-
// Usage is the same as ChatCompletion, just passing the stream parameter
291-
292300
if customModel == nil {
293301
// User did not specify API key for this model
294302
llmProvider = &models.LLMProviderConfig{

internal/api/chat/list_supported_models_v2.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,9 @@ func (s *ChatServerV2) ListSupportedModels(
222222
var models []*chatv2.SupportedModel
223223

224224
for _, model := range settings.CustomModels {
225+
modelID := model.Id.Hex()
225226
models = append(models, &chatv2.SupportedModel{
227+
Id: &modelID,
226228
Name: model.Name,
227229
Slug: model.Slug,
228230
TotalContext: int64(model.ContextWindow),

pkg/gen/api/chat/v2/chat.pb.go

Lines changed: 27 additions & 6 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

proto/chat/v2/chat.proto

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ message SupportedModel {
137137
bool disabled = 7; // If true, the model is disabled and cannot be used
138138
optional string disabled_reason = 8; // The reason why the model is disabled
139139
bool is_custom = 9;
140+
optional string id = 10; // Custom model unique ID (empty for built-in models)
140141
}
141142

142143
message ListSupportedModelsRequest {
@@ -223,6 +224,7 @@ message CreateConversationMessageStreamRequest {
223224
optional string user_selected_text = 5;
224225
optional ConversationType conversation_type = 6;
225226
optional string surrounding = 8;
227+
optional string custom_model_id = 9; // Selected custom model ID
226228
}
227229

228230
// Response for streaming a message within an existing conversation

webapp/_webapp/src/hooks/useLanguageModels.ts

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { useListSupportedModelsQuery } from "../query";
55
import { useConversationUiStore } from "../stores/conversation/conversation-ui-store";
66

77
export type Model = {
8+
id?: string;
89
name: string;
910
slug: string;
1011
provider: string;
@@ -39,6 +40,7 @@ const fallbackModels: Model[] = [
3940
];
4041

4142
const mapSupportedModelToModel = (supportedModel: SupportedModel): Model => ({
43+
id: supportedModel.id || undefined,
4244
name: supportedModel.name,
4345
slug: supportedModel.slug,
4446
provider: extractProvider(supportedModel.slug),
@@ -53,7 +55,7 @@ const mapSupportedModelToModel = (supportedModel: SupportedModel): Model => ({
5355

5456
export const useLanguageModels = () => {
5557
const { currentConversation, setCurrentConversation } = useConversationStore();
56-
const { setLastUsedModelSlug } = useConversationUiStore();
58+
const { lastUsedCustomModelId, setLastUsedModelSlug, setLastUsedCustomModelId } = useConversationUiStore();
5759
const { data: supportedModelsResponse } = useListSupportedModelsQuery();
5860

5961
const models: Model[] = useMemo(() => {
@@ -64,19 +66,25 @@ export const useLanguageModels = () => {
6466
}, [supportedModelsResponse]);
6567

6668
const currentModel = useMemo(() => {
67-
const model = models.find((m) => m.slug === currentConversation.modelSlug);
69+
if (lastUsedCustomModelId) {
70+
const customModel = models.find((m) => m.isCustom && m.id === lastUsedCustomModelId);
71+
if (customModel) return customModel;
72+
}
73+
74+
const model = models.find((m) => !m.isCustom && m.slug === currentConversation.modelSlug);
6875
return model || models[0];
69-
}, [models, currentConversation.modelSlug]);
76+
}, [models, currentConversation.modelSlug, lastUsedCustomModelId]);
7077

7178
const setModel = useCallback(
7279
(model: Model) => {
7380
setLastUsedModelSlug(model.slug);
81+
setLastUsedCustomModelId(model.isCustom ? (model.id ?? "") : "");
7482
setCurrentConversation({
7583
...currentConversation,
7684
modelSlug: model.slug,
7785
});
7886
},
79-
[setCurrentConversation, currentConversation, setLastUsedModelSlug],
87+
[setCurrentConversation, currentConversation, setLastUsedModelSlug, setLastUsedCustomModelId],
8088
);
8189

8290
return { models, currentModel, setModel };

webapp/_webapp/src/hooks/useSendMessageStream.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ import { useAuthStore } from "../stores/auth-store";
4141
import { useDevtoolStore } from "../stores/devtool-store";
4242
import { useSelectionStore } from "../stores/selection-store";
4343
import { useSettingStore } from "../stores/setting-store";
44+
import { useConversationUiStore } from "../stores/conversation/conversation-ui-store";
4445
import { useSync } from "./useSync";
4546
import { useAdapter } from "../adapters";
4647
import { getProjectId } from "../libs/helpers";
@@ -86,6 +87,7 @@ export function useSendMessageStream(): UseSendMessageStreamResult {
8687
const surroundingText = useSelectionStore((s) => s.surroundingText);
8788
const alwaysSyncProject = useDevtoolStore((s) => s.alwaysSyncProject);
8889
const conversationMode = useSettingStore((s) => s.conversationMode);
90+
const lastUsedCustomModelId = useConversationUiStore((s) => s.lastUsedCustomModelId);
8991

9092
/**
9193
* Add the user message to the streaming state.
@@ -165,6 +167,7 @@ export function useSendMessageStream(): UseSendMessageStreamResult {
165167
projectId,
166168
conversationId: currentConversation.id,
167169
modelSlug: currentConversation.modelSlug,
170+
customModelId: lastUsedCustomModelId || undefined,
168171
surroundingText: surroundingText ?? undefined,
169172
conversationMode: conversationMode === "debug" ? "debug" : "default",
170173
};
@@ -251,6 +254,7 @@ export function useSendMessageStream(): UseSendMessageStreamResult {
251254
alwaysSyncProject,
252255
conversationMode,
253256
surroundingText,
257+
lastUsedCustomModelId,
254258
addUserMessageToStream,
255259
truncateConversationIfEditing,
256260
],

0 commit comments

Comments
 (0)