From 9b7a04431cb12609cbcd766e05fbc411599e0b34 Mon Sep 17 00:00:00 2001 From: ginnoir Date: Wed, 8 Jul 2026 16:34:22 -0500 Subject: [PATCH] feat(agent): route chat through selected model --- src/app/api/agent/chat/route.ts | 15 ++++++++++++++ src/lib/llm/index.ts | 6 +++--- src/modules/agent/messages.ts | 7 +++++++ src/modules/agent/server/run.ts | 3 ++- tests/unit/agent-chat.test.ts | 33 +++++++++++++++++++++++++++++++ tests/unit/agent-messages.test.ts | 18 +++++++++++++++++ 6 files changed, 78 insertions(+), 4 deletions(-) diff --git a/src/app/api/agent/chat/route.ts b/src/app/api/agent/chat/route.ts index 1c8c335..596947c 100644 --- a/src/app/api/agent/chat/route.ts +++ b/src/app/api/agent/chat/route.ts @@ -2,6 +2,7 @@ import { apiError, apiJson } from "@/lib/api-handler"; import { resolveApiAuth } from "@/lib/api-auth"; import { getAssistantPreferences, resolveAssistantSystemPrompt } from "@/lib/assistant-preference"; import { isLlmConfigured } from "@/lib/llm"; +import { listLlmModels, resolveAssistantModel } from "@/lib/llm/models"; import { clientChatInputSchema } from "@/modules/agent/messages"; import { encodeSseEvent } from "@/modules/agent/server/progress"; import { runAgentChat } from "@/modules/agent/server/run"; @@ -31,6 +32,18 @@ export async function POST(request: Request) { return apiError(parsed.error.issues[0]?.message ?? "Validation error", 400); } + const modelList = await listLlmModels(); + const modelResolution = resolveAssistantModel({ + requestedModel: parsed.data.model, + savedModel: assistant.model, + fallbackModel: modelList.fallbackModel, + models: modelList.models, + }); + + if (!modelResolution.ok) { + return apiError(modelResolution.error, 400); + } + if (parsed.data.stream) { const stream = new ReadableStream({ async start(controller) { @@ -44,6 +57,7 @@ export async function POST(request: Request) { messages: parsed.data.messages, request, systemPrompt, + model: modelResolution.model, onProgress: send, }); @@ -78,6 +92,7 @@ export async function POST(request: Request) { messages: parsed.data.messages, request, systemPrompt, + model: modelResolution.model, }); return apiJson({ diff --git a/src/lib/llm/index.ts b/src/lib/llm/index.ts index 64d7506..5e3af58 100644 --- a/src/lib/llm/index.ts +++ b/src/lib/llm/index.ts @@ -13,8 +13,8 @@ export type { export { getLlmConfig, isLlmConfigured } from "./config"; export { createMockLlmClient } from "./mock"; -export function createLlmClient(override?: LlmClient): LlmClient { - if (override) return override; +export function createLlmClient(options?: { model?: string; override?: LlmClient }): LlmClient { + if (options?.override) return options.override; const config = getLlmConfig(); if (config.provider === "mock" || !config.baseUrl) { @@ -24,6 +24,6 @@ export function createLlmClient(override?: LlmClient): LlmClient { return createOpenAiCompatibleClient({ baseUrl: config.baseUrl, apiKey: config.apiKey, - model: config.model, + model: options?.model ?? config.model, }); } diff --git a/src/modules/agent/messages.ts b/src/modules/agent/messages.ts index e148acd..68828a1 100644 --- a/src/modules/agent/messages.ts +++ b/src/modules/agent/messages.ts @@ -1,4 +1,5 @@ import { z } from "zod"; +import { isValidLlmModelId } from "@/lib/llm/models"; export const clientChatAttachmentSchema = z.object({ type: z.literal("image"), @@ -11,8 +12,14 @@ export const clientChatMessageSchema = z.object({ attachments: z.array(clientChatAttachmentSchema).max(3).optional(), }); +export const clientChatModelSchema = z + .string() + .trim() + .refine((value) => isValidLlmModelId(value), "Invalid assistant model"); + export const clientChatInputSchema = z.object({ stream: z.boolean().optional(), + model: clientChatModelSchema.optional(), messages: z.array(clientChatMessageSchema).min(1).max(40), }); diff --git a/src/modules/agent/server/run.ts b/src/modules/agent/server/run.ts index 341e267..46d54f8 100644 --- a/src/modules/agent/server/run.ts +++ b/src/modules/agent/server/run.ts @@ -45,11 +45,12 @@ export async function runAgentChat(options: { messages: ClientChatMessage[]; request: Request; systemPrompt?: string; + model?: string; llm?: LlmClient; executeTool?: ToolExecutor; onProgress?: AgentProgressHandler; }): Promise { - const llm = options.llm ?? createLlmClient(); + const llm = options.llm ?? createLlmClient({ model: options.model }); const executeTool = options.executeTool ?? createApiToolExecutor(options.request); const onProgress = options.onProgress; const systemPrompt = options.systemPrompt ?? AGENT_SYSTEM_PROMPT; diff --git a/tests/unit/agent-chat.test.ts b/tests/unit/agent-chat.test.ts index 18e399a..5e06efe 100644 --- a/tests/unit/agent-chat.test.ts +++ b/tests/unit/agent-chat.test.ts @@ -84,3 +84,36 @@ describe("runAgentChat", () => { assert.ok(result.message.content.length > 0); }); }); + +it("passes a model override to the OpenAI-compatible client", async () => { + const originalBaseUrl = process.env.LLM_BASE_URL; + const originalModel = process.env.LLM_MODEL; + const originalProvider = process.env.LLM_PROVIDER; + const originalFetch = globalThis.fetch; + let requestBody: unknown = null; + + process.env.LLM_BASE_URL = "https://llm.example.test/v1"; + process.env.LLM_MODEL = "llama3.2"; + delete process.env.LLM_PROVIDER; + + globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { + requestBody = JSON.parse(String(init?.body)); + return Response.json({ + choices: [{ message: { role: "assistant", content: "done" }, finish_reason: "stop" }], + }); + }) as typeof fetch; + + const { createLlmClient } = await import("../../src/lib/llm/index"); + const client = createLlmClient({ model: "qwen2.5-coder" }); + await client.chatCompletion({ messages: [{ role: "user", content: "hello" }] }); + + assert.equal((requestBody as { model?: string }).model, "qwen2.5-coder"); + + globalThis.fetch = originalFetch; + if (originalBaseUrl === undefined) delete process.env.LLM_BASE_URL; + else process.env.LLM_BASE_URL = originalBaseUrl; + if (originalModel === undefined) delete process.env.LLM_MODEL; + else process.env.LLM_MODEL = originalModel; + if (originalProvider === undefined) delete process.env.LLM_PROVIDER; + else process.env.LLM_PROVIDER = originalProvider; +}); diff --git a/tests/unit/agent-messages.test.ts b/tests/unit/agent-messages.test.ts index 9815d11..c038599 100644 --- a/tests/unit/agent-messages.test.ts +++ b/tests/unit/agent-messages.test.ts @@ -25,4 +25,22 @@ describe("clientChatInputSchema", () => { assert.equal(parsed.success, false); }); + + it("accepts an optional model ID", () => { + const parsed = clientChatInputSchema.safeParse({ + model: "qwen2.5-coder", + messages: [{ role: "user", content: "hello" }], + }); + + assert.equal(parsed.success, true); + }); + + it("rejects invalid model IDs", () => { + const parsed = clientChatInputSchema.safeParse({ + model: "bad model", + messages: [{ role: "user", content: "hello" }], + }); + + assert.equal(parsed.success, false); + }); });