From ec3f96dab170fb2333780438bdec3628cfc58116 Mon Sep 17 00:00:00 2001 From: ginnoir Date: Wed, 8 Jul 2026 19:50:46 -0500 Subject: [PATCH] feat(agent): let users choose model route --- drizzle/0026_assistant_model_route.sql | 5 + drizzle/meta/_journal.json | 9 +- src/app/api/agent/chat/route.ts | 2 +- src/app/api/agent/models/route.ts | 13 +- src/app/layout.tsx | 4 + src/app/settings/assistant-actions.ts | 24 +++- src/lib/assistant-preference.ts | 7 ++ src/lib/llm/models.ts | 37 ++++-- src/modules/_core/schema.ts | 1 + .../agent/components/assistant-bubble.tsx | 10 +- .../agent/components/assistant-panel.tsx | 116 ++++++++++++++---- tests/e2e/assistant.spec.ts | 10 +- tests/unit/llm-models.test.ts | 42 ++++++- 13 files changed, 235 insertions(+), 45 deletions(-) create mode 100644 drizzle/0026_assistant_model_route.sql diff --git a/drizzle/0026_assistant_model_route.sql b/drizzle/0026_assistant_model_route.sql new file mode 100644 index 0000000..152d1f3 --- /dev/null +++ b/drizzle/0026_assistant_model_route.sql @@ -0,0 +1,5 @@ +ALTER TABLE "users" ADD COLUMN "assistant_model_route" text;--> statement-breakpoint +UPDATE "users" +SET "assistant_model_route" = "assistant_model", + "assistant_model" = NULL +WHERE "assistant_model" IN ('auto', 'uncensored'); diff --git a/drizzle/meta/_journal.json b/drizzle/meta/_journal.json index a0c11fb..0eb35e2 100644 --- a/drizzle/meta/_journal.json +++ b/drizzle/meta/_journal.json @@ -183,6 +183,13 @@ "when": 1783560000000, "tag": "0025_assistant_model", "breakpoints": true + }, + { + "idx": 26, + "version": "7", + "when": 1783561000000, + "tag": "0026_assistant_model_route", + "breakpoints": true } ] -} \ No newline at end of file +} diff --git a/src/app/api/agent/chat/route.ts b/src/app/api/agent/chat/route.ts index 596947c..471b939 100644 --- a/src/app/api/agent/chat/route.ts +++ b/src/app/api/agent/chat/route.ts @@ -32,7 +32,7 @@ export async function POST(request: Request) { return apiError(parsed.error.issues[0]?.message ?? "Validation error", 400); } - const modelList = await listLlmModels(); + const modelList = await listLlmModels({ route: assistant.modelRoute }); const modelResolution = resolveAssistantModel({ requestedModel: parsed.data.model, savedModel: assistant.model, diff --git a/src/app/api/agent/models/route.ts b/src/app/api/agent/models/route.ts index bfe5670..bdef4a9 100644 --- a/src/app/api/agent/models/route.ts +++ b/src/app/api/agent/models/route.ts @@ -1,7 +1,7 @@ import { apiError, apiJson } from "@/lib/api-handler"; import { resolveApiAuth } from "@/lib/api-auth"; import { getAssistantPreferences } from "@/lib/assistant-preference"; -import { listLlmModels, resolveAssistantModel } from "@/lib/llm/models"; +import { isValidAssistantModelRoute, listLlmModels, resolveAssistantModel } from "@/lib/llm/models"; export const dynamic = "force-dynamic"; @@ -21,10 +21,16 @@ export async function GET(request: Request) { return noStore(apiError("Assistant not enabled", 403)); } - const modelList = await listLlmModels(); + const requestedRoute = new URL(request.url).searchParams.get("route"); + if (requestedRoute !== null && !isValidAssistantModelRoute(requestedRoute)) { + return noStore(apiError("Invalid assistant model route", 400)); + } + + const modelRoute = requestedRoute ?? assistant.modelRoute; + const modelList = await listLlmModels({ route: modelRoute }); const resolved = resolveAssistantModel({ requestedModel: null, - savedModel: assistant.model, + savedModel: requestedRoute === null ? assistant.model : null, fallbackModel: modelList.fallbackModel, models: modelList.models, }); @@ -34,6 +40,7 @@ export async function GET(request: Request) { models: modelList.models, selectedModel: resolved.model, fallbackModel: modelList.fallbackModel, + route: modelList.route, degraded: modelList.degraded, }), ); diff --git a/src/app/layout.tsx b/src/app/layout.tsx index 7da8be5..9c88140 100644 --- a/src/app/layout.tsx +++ b/src/app/layout.tsx @@ -104,6 +104,7 @@ export default async function RootLayout({ children }: { children: React.ReactNo let signedIn = false; let assistantEnabled = false; let assistantName = DEFAULT_ASSISTANT_NAME; + let assistantModelRoute: string | null = null; let assistantModel: string | null = null; const session = await auth(); @@ -118,6 +119,7 @@ export default async function RootLayout({ children }: { children: React.ReactNo themeNavStyle: users.themeNavStyle, assistantEnabled: users.assistantEnabled, assistantName: users.assistantName, + assistantModelRoute: users.assistantModelRoute, assistantModel: users.assistantModel, }) .from(users) @@ -131,6 +133,7 @@ export default async function RootLayout({ children }: { children: React.ReactNo navStyle = row.themeNavStyle as NavStyle; assistantEnabled = row.assistantEnabled; assistantName = row.assistantName?.trim() || DEFAULT_ASSISTANT_NAME; + assistantModelRoute = row.assistantModelRoute?.trim() || null; assistantModel = row.assistantModel?.trim() || null; } userDashboards = await db @@ -189,6 +192,7 @@ export default async function RootLayout({ children }: { children: React.ReactNo configured={isLlmConfigured()} userId={session.user.id} assistantName={assistantName} + assistantModelRoute={assistantModelRoute} assistantModel={assistantModel} /> ) : null} diff --git a/src/app/settings/assistant-actions.ts b/src/app/settings/assistant-actions.ts index 385fd56..2dcd73d 100644 --- a/src/app/settings/assistant-actions.ts +++ b/src/app/settings/assistant-actions.ts @@ -9,9 +9,15 @@ import { MAX_ASSISTANT_NAME_LENGTH, MAX_ASSISTANT_SYSTEM_PROMPT_LENGTH, } from "@/lib/assistant-config"; -import { isValidLlmModelId, listLlmModels } from "@/lib/llm/models"; +import { + isValidAssistantModelRoute, + isValidLlmModelId, + listLlmModels, + type AssistantModelRoute, +} from "@/lib/llm/models"; import { users } from "@/modules/_core/schema"; import { getCurrentSession } from "@/lib/session"; +import { getAssistantPreferences } from "@/lib/assistant-preference"; const assistantNameSchema = z .string() @@ -51,6 +57,19 @@ export async function setAssistantSystemPrompt(prompt: string | null): Promise { + if (!isValidAssistantModelRoute(route)) { + throw new Error("Invalid assistant model route"); + } + + const { user } = await getCurrentSession(); + await db + .update(users) + .set({ assistantModelRoute: route, assistantModel: null }) + .where(eq(users.id, user.id)); + revalidateAssistantSurfaces(); +} + export async function setAssistantModel(model: string | null): Promise { const { user } = await getCurrentSession(); const normalized = model?.trim() || null; @@ -59,7 +78,8 @@ export async function setAssistantModel(model: string | null): Promise { throw new Error("Invalid assistant model"); } - const available = await listLlmModels(); + const assistant = await getAssistantPreferences(user.id); + const available = await listLlmModels({ route: assistant.modelRoute }); const requested = normalized === available.fallbackModel ? null : normalized; if (requested !== null && !available.models.some((option) => option.id === requested)) { diff --git a/src/lib/assistant-preference.ts b/src/lib/assistant-preference.ts index babe897..09f80c9 100644 --- a/src/lib/assistant-preference.ts +++ b/src/lib/assistant-preference.ts @@ -1,6 +1,7 @@ import { eq } from "drizzle-orm"; import { db } from "@/lib/db"; import { DEFAULT_ASSISTANT_NAME } from "@/lib/assistant-config"; +import { isValidAssistantModelRoute, type AssistantModelRoute } from "@/lib/llm/models"; import { users } from "@/modules/_core/schema"; export { @@ -14,6 +15,7 @@ export type AssistantPreferences = { enabled: boolean; name: string; systemPrompt: string | null; + modelRoute: AssistantModelRoute | null; model: string | null; }; @@ -23,6 +25,7 @@ export async function getAssistantPreferences(userId: string): Promise route === value); +} + +function resolveModelRoute( + route: AssistantModelRoute | null | undefined, + config: LlmConfig, +): { route: AssistantModelRoute; fallbackModel: string } { + if (route) { + return { route, fallbackModel: route }; + } + + const defaultRoute: AssistantModelRoute = config.model === "uncensored" ? "uncensored" : "auto"; + return { route: defaultRoute, fallbackModel: config.model }; +} + export function normalizeLlmModelsPayload(payload: unknown): LlmModelOption[] { const data = typeof payload === "object" && payload !== null && "data" in payload @@ -51,14 +72,15 @@ export function normalizeLlmModelsPayload(payload: unknown): LlmModelOption[] { export async function listLlmModels(options?: { config?: LlmConfig; fetchImpl?: typeof fetch; + route?: AssistantModelRoute | null; }): Promise { const config = options?.config ?? getLlmConfig(); const fetchImpl = options?.fetchImpl ?? fetch; - const fallbackModel = config.model; + const { route, fallbackModel } = resolveModelRoute(options?.route, config); const fallbackOption = { id: fallbackModel, label: fallbackModel }; if (config.provider === "mock" || !config.baseUrl) { - return { models: [fallbackOption], fallbackModel, degraded: false }; + return { models: [fallbackOption], fallbackModel, route, degraded: false }; } try { @@ -66,7 +88,7 @@ export async function listLlmModels(options?: { if (config.apiKey) headers.Authorization = `Bearer ${config.apiKey}`; const modelsUrl = new URL(`${config.baseUrl.replace(/\/$/, "")}/models`); - if (fallbackModel === "uncensored") { + if (route === "uncensored") { modelsUrl.searchParams.set("type", "uncensored"); } @@ -76,26 +98,27 @@ export async function listLlmModels(options?: { }); if (!response.ok) { - return { models: [fallbackOption], fallbackModel, degraded: true }; + return { models: [fallbackOption], fallbackModel, route, degraded: true }; } const models = normalizeLlmModelsPayload(await response.json()); if (models.length === 0) { - return { models: [fallbackOption], fallbackModel, degraded: true }; + return { models: [fallbackOption], fallbackModel, route, degraded: true }; } if (!models.some((model) => model.id === fallbackModel)) { - return { models: [fallbackOption], fallbackModel, degraded: false }; + return { models: [fallbackOption], fallbackModel, route, degraded: false }; } return { models, fallbackModel, + route, degraded: false, }; } catch { - return { models: [fallbackOption], fallbackModel, degraded: true }; + return { models: [fallbackOption], fallbackModel, route, degraded: true }; } } diff --git a/src/modules/_core/schema.ts b/src/modules/_core/schema.ts index d404740..285014b 100644 --- a/src/modules/_core/schema.ts +++ b/src/modules/_core/schema.ts @@ -38,6 +38,7 @@ export const users = pgTable("users", { assistantEnabled: boolean("assistant_enabled").notNull().default(false), assistantName: text("assistant_name").notNull().default("Assistant"), assistantSystemPrompt: text("assistant_system_prompt"), + assistantModelRoute: text("assistant_model_route"), assistantModel: text("assistant_model"), defaultEventReminderOffsets: jsonb("default_event_reminder_offsets") .notNull() diff --git a/src/modules/agent/components/assistant-bubble.tsx b/src/modules/agent/components/assistant-bubble.tsx index 904722b..f9cc4ec 100644 --- a/src/modules/agent/components/assistant-bubble.tsx +++ b/src/modules/agent/components/assistant-bubble.tsx @@ -8,10 +8,17 @@ type Props = { configured: boolean; userId: string; assistantName: string; + assistantModelRoute: string | null; assistantModel: string | null; }; -export function AssistantBubble({ configured, userId, assistantName, assistantModel }: Props) { +export function AssistantBubble({ + configured, + userId, + assistantName, + assistantModelRoute, + assistantModel, +}: Props) { const [open, setOpen] = useState(false); return ( @@ -42,6 +49,7 @@ export function AssistantBubble({ configured, userId, assistantName, assistantMo configured={configured} userId={userId} assistantName={assistantName} + assistantModelRoute={assistantModelRoute} assistantModel={assistantModel} /> diff --git a/src/modules/agent/components/assistant-panel.tsx b/src/modules/agent/components/assistant-panel.tsx index 6c06688..9cc59d2 100644 --- a/src/modules/agent/components/assistant-panel.tsx +++ b/src/modules/agent/components/assistant-panel.tsx @@ -2,9 +2,10 @@ import { useCallback, useEffect, useRef, useState, useTransition } from "react"; import { ChevronDown, ImagePlus, Loader2, Mic, RefreshCw, Send, Square } from "lucide-react"; -import { setAssistantModel } from "@/app/settings/assistant-actions"; +import { setAssistantModel, setAssistantModelRoute } from "@/app/settings/assistant-actions"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; +import type { AssistantModelRoute } from "@/lib/llm/models"; import { consumeAgentChatStream } from "../assistant-chat-stream"; import { clearAssistantChat, @@ -19,6 +20,7 @@ type Props = { configured: boolean; userId: string; assistantName: string; + assistantModelRoute: string | null; assistantModel: string | null; }; @@ -35,6 +37,7 @@ type ModelsResponse = { models: LlmModelOption[]; selectedModel: string; fallbackModel: string; + route: AssistantModelRoute; degraded: boolean; }; @@ -50,7 +53,17 @@ async function uploadAssistantImage(file: File): Promise { return payload.url; } -export function AssistantPanel({ configured, userId, assistantName, assistantModel }: Props) { +function normalizeAssistantModelRoute(value: string | null): AssistantModelRoute { + return value === "uncensored" ? "uncensored" : "auto"; +} + +export function AssistantPanel({ + configured, + userId, + assistantName, + assistantModelRoute, + assistantModel, +}: Props) { const [messages, setMessages] = useState(() => loadAssistantChat(userId)); const [input, setInput] = useState(""); const [pendingImage, setPendingImage] = useState(null); @@ -59,6 +72,9 @@ export function AssistantPanel({ configured, userId, assistantName, assistantMod const [isPending, setIsPending] = useState(false); const [activityLabel, setActivityLabel] = useState(null); const [modelOptions, setModelOptions] = useState([]); + const [selectedRoute, setSelectedRoute] = useState(() => + normalizeAssistantModelRoute(assistantModelRoute), + ); const [selectedModel, setSelectedModel] = useState(assistantModel ?? ""); const [fallbackModel, setFallbackModel] = useState(""); const [modelsDegraded, setModelsDegraded] = useState(false); @@ -87,31 +103,40 @@ export function AssistantPanel({ configured, userId, assistantName, assistantMod }; }, []); - const loadModels = useCallback(async (options?: { refresh?: boolean; signal?: AbortSignal }) => { - if (options?.signal?.aborted) return; - - setModelsLoading(true); - try { - const response = await fetch(`/api/agent/models${options?.refresh ? "?refresh=1" : ""}`, { - cache: "no-store", - signal: options?.signal, - }); - if (!response.ok) throw new Error("Model discovery unavailable"); - const payload = (await response.json()) as ModelsResponse; + const loadModels = useCallback( + async (options?: { refresh?: boolean; route?: AssistantModelRoute; signal?: AbortSignal }) => { if (options?.signal?.aborted) return; - setModelOptions(payload.models); - setSelectedModel(payload.selectedModel); - setFallbackModel(payload.fallbackModel); - setModelsDegraded(payload.degraded); - setError(null); - } catch (err) { - if (err instanceof Error && err.name === "AbortError") return; - setModelsDegraded(true); - setError("Model discovery unavailable"); - } finally { - if (!options?.signal?.aborted) setModelsLoading(false); - } - }, []); + + setModelsLoading(true); + try { + const params = new URLSearchParams(); + if (options?.refresh) params.set("refresh", "1"); + if (options?.route) params.set("route", options.route); + const query = params.size > 0 ? `?${params.toString()}` : ""; + + const response = await fetch(`/api/agent/models${query}`, { + cache: "no-store", + signal: options?.signal, + }); + if (!response.ok) throw new Error("Model discovery unavailable"); + const payload = (await response.json()) as ModelsResponse; + if (options?.signal?.aborted) return; + setModelOptions(payload.models); + setSelectedRoute(payload.route); + setSelectedModel(payload.selectedModel); + setFallbackModel(payload.fallbackModel); + setModelsDegraded(payload.degraded); + setError(null); + } catch (err) { + if (err instanceof Error && err.name === "AbortError") return; + setModelsDegraded(true); + setError("Model discovery unavailable"); + } finally { + if (!options?.signal?.aborted) setModelsLoading(false); + } + }, + [], + ); useEffect(() => { const controller = new AbortController(); @@ -175,6 +200,26 @@ export function AssistantPanel({ configured, userId, assistantName, assistantMod }); } + function changeModelRoute(nextRoute: string) { + const route = normalizeAssistantModelRoute(nextRoute); + + setSelectedRoute(route); + setSelectedModel(route); + setFallbackModel(route); + setModelOptions([]); + setModelsDegraded(false); + setError(null); + + startTransition(async () => { + try { + await setAssistantModelRoute(route); + await loadModels({ route, refresh: true }); + } catch (err) { + setError(err instanceof Error ? err.message : "Could not save assistant route"); + } + }); + } + async function sendMessage() { const text = input.trim(); const hasImage = pendingImage !== null; @@ -242,6 +287,7 @@ export function AssistantPanel({ configured, userId, assistantName, assistantMod const inputDisabled = isPending || voiceState === "transcribing" || uploadingImage; const canSend = !isPending && + !savingModel && voiceState === "idle" && !uploadingImage && (input.trim().length > 0 || pendingImage !== null); @@ -260,6 +306,22 @@ export function AssistantPanel({ configured, userId, assistantName, assistantMod ) : null}
+
+ +
{modelOptions.length > 0 ? (