feat(agent): add llm model discovery helper

2026-07-08 16:11:22 -05:00
parent eb8e562565
commit 0bf63cf8a8
2 changed files with 254 additions and 0 deletions
@@ -0,0 +1,116 @@
+import { getLlmConfig, type LlmConfig } from "./config";
+
+export type LlmModelOption = {
+  id: string;
+  label: string;
+};
+
+export type LlmModelsResult = {
+  models: LlmModelOption[];
+  fallbackModel: string;
+  degraded: boolean;
+};
+
+export type AssistantModelResolution =
+  | { ok: true; model: string }
+  | { ok: false; model: string; error: string };
+
+const MODEL_ID_PATTERN = /^[A-Za-z0-9._:/-]+$/;
+const MAX_MODEL_ID_LENGTH = 128;
+
+export function isValidLlmModelId(value: string): boolean {
+  const trimmed = value.trim();
+  return (
+    trimmed.length > 0 &&
+    trimmed.length <= MAX_MODEL_ID_LENGTH &&
+    trimmed === value &&
+    MODEL_ID_PATTERN.test(trimmed)
+  );
+}
+
+export function normalizeLlmModelsPayload(payload: unknown): LlmModelOption[] {
+  const data =
+    typeof payload === "object" && payload !== null && "data" in payload
+      ? (payload as { data?: unknown }).data
+      : null;
+
+  if (!Array.isArray(data)) return [];
+
+  const ids = new Set<string>();
+  for (const row of data) {
+    if (typeof row !== "object" || row === null || !("id" in row)) continue;
+    const id = (row as { id?: unknown }).id;
+    if (typeof id !== "string") continue;
+    const trimmed = id.trim();
+    if (!isValidLlmModelId(trimmed)) continue;
+    ids.add(trimmed);
+  }
+
+  return [...ids].sort((a, b) => a.localeCompare(b)).map((id) => ({ id, label: id }));
+}
+
+export async function listLlmModels(options?: {
+  config?: LlmConfig;
+  fetchImpl?: typeof fetch;
+}): Promise<LlmModelsResult> {
+  const config = options?.config ?? getLlmConfig();
+  const fetchImpl = options?.fetchImpl ?? fetch;
+  const fallbackModel = config.model;
+  const fallbackOption = { id: fallbackModel, label: fallbackModel };
+
+  if (config.provider === "mock" || !config.baseUrl) {
+    return { models: [fallbackOption], fallbackModel, degraded: false };
+  }
+
+  try {
+    const headers: Record<string, string> = {};
+    if (config.apiKey) headers.Authorization = `Bearer ${config.apiKey}`;
+
+    const response = await fetchImpl(`${config.baseUrl.replace(/\/$/, "")}/models`, {
+      method: "GET",
+      headers,
+    });
+
+    if (!response.ok) {
+      return { models: [fallbackOption], fallbackModel, degraded: true };
+    }
+
+    const models = normalizeLlmModelsPayload(await response.json());
+    const merged = new Map<string, LlmModelOption>();
+    merged.set(fallbackModel, fallbackOption);
+    for (const model of models) merged.set(model.id, model);
+
+    return {
+      models: [...merged.values()].sort((a, b) => a.id.localeCompare(b.id)),
+      fallbackModel,
+      degraded: models.length === 0,
+    };
+  } catch {
+    return { models: [fallbackOption], fallbackModel, degraded: true };
+  }
+}
+
+export function resolveAssistantModel(options: {
+  requestedModel: string | null | undefined;
+  savedModel: string | null | undefined;
+  fallbackModel: string;
+  models: LlmModelOption[];
+}): AssistantModelResolution {
+  const available = new Set(options.models.map((model) => model.id));
+  const fallback = available.has(options.fallbackModel)
+    ? options.fallbackModel
+    : (options.models[0]?.id ?? options.fallbackModel);
+
+  if (options.requestedModel) {
+    if (!available.has(options.requestedModel)) {
+      return { ok: false, model: fallback, error: "Invalid assistant model" };
+    }
+    return { ok: true, model: options.requestedModel };
+  }
+
+  if (options.savedModel && available.has(options.savedModel)) {
+    return { ok: true, model: options.savedModel };
+  }
+
+  return { ok: true, model: fallback };
+}
@@ -0,0 +1,138 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+import {
+  isValidLlmModelId,
+  listLlmModels,
+  normalizeLlmModelsPayload,
+  resolveAssistantModel,
+} from "../../src/lib/llm/models";
+import type { LlmConfig } from "../../src/lib/llm/config";
+
+const openAiConfig: LlmConfig = {
+  provider: "openai",
+  baseUrl: "https://llm.example.test/v1",
+  apiKey: "secret",
+  model: "llama3.2",
+};
+
+describe("normalizeLlmModelsPayload", () => {
+  it("normalizes OpenAI-compatible data arrays", () => {
+    const models = normalizeLlmModelsPayload({
+      data: [{ id: "qwen2.5-coder" }, { id: "llama3.2" }, { id: "qwen2.5-coder" }],
+    });
+
+    assert.deepEqual(models, [
+      { id: "llama3.2", label: "llama3.2" },
+      { id: "qwen2.5-coder", label: "qwen2.5-coder" },
+    ]);
+  });
+
+  it("ignores invalid or empty model rows", () => {
+    const models = normalizeLlmModelsPayload({
+      data: [{ id: "" }, { id: "   " }, { id: "bad model" }, { object: "model" }],
+    });
+
+    assert.deepEqual(models, []);
+  });
+});
+
+describe("isValidLlmModelId", () => {
+  it("accepts common provider model IDs", () => {
+    assert.equal(isValidLlmModelId("llama3.2"), true);
+    assert.equal(isValidLlmModelId("qwen2.5-coder:latest"), true);
+    assert.equal(isValidLlmModelId("hf.co/ginnoir/model-v1"), true);
+  });
+
+  it("rejects empty, whitespace, and overlong model IDs", () => {
+    assert.equal(isValidLlmModelId(""), false);
+    assert.equal(isValidLlmModelId("bad model"), false);
+    assert.equal(isValidLlmModelId("x".repeat(129)), false);
+  });
+});
+
+describe("listLlmModels", () => {
+  it("fetches provider models with API key auth and includes the fallback model", async () => {
+    const requests: Request[] = [];
+    const result = await listLlmModels({
+      config: openAiConfig,
+      fetchImpl: async (input, init) => {
+        requests.push(new Request(input, init));
+        return Response.json({ data: [{ id: "qwen2.5-coder" }] });
+      },
+    });
+
+    assert.equal(requests[0]?.url, "https://llm.example.test/v1/models");
+    assert.equal(requests[0]?.headers.get("authorization"), "Bearer secret");
+    assert.deepEqual(result.models, [
+      { id: "llama3.2", label: "llama3.2" },
+      { id: "qwen2.5-coder", label: "qwen2.5-coder" },
+    ]);
+    assert.equal(result.fallbackModel, "llama3.2");
+    assert.equal(result.degraded, false);
+  });
+
+  it("falls back to LLM_MODEL when provider discovery fails", async () => {
+    const result = await listLlmModels({
+      config: openAiConfig,
+      fetchImpl: async () => new Response("nope", { status: 500 }),
+    });
+
+    assert.deepEqual(result.models, [{ id: "llama3.2", label: "llama3.2" }]);
+    assert.equal(result.fallbackModel, "llama3.2");
+    assert.equal(result.degraded, true);
+  });
+
+  it("uses fallback only for mock provider config", async () => {
+    const result = await listLlmModels({
+      config: { provider: "mock", baseUrl: null, apiKey: null, model: "llama3.2" },
+      fetchImpl: async () => {
+        throw new Error("fetch should not run for mock config");
+      },
+    });
+
+    assert.deepEqual(result.models, [{ id: "llama3.2", label: "llama3.2" }]);
+    assert.equal(result.degraded, false);
+  });
+});
+
+describe("resolveAssistantModel", () => {
+  it("uses a valid requested model before saved and fallback values", () => {
+    const resolved = resolveAssistantModel({
+      requestedModel: "qwen2.5-coder",
+      savedModel: "llama3.2",
+      fallbackModel: "llama3.2",
+      models: [
+        { id: "llama3.2", label: "llama3.2" },
+        { id: "qwen2.5-coder", label: "qwen2.5-coder" },
+      ],
+    });
+
+    assert.deepEqual(resolved, { ok: true, model: "qwen2.5-coder" });
+  });
+
+  it("rejects invalid requested models", () => {
+    const resolved = resolveAssistantModel({
+      requestedModel: "missing",
+      savedModel: null,
+      fallbackModel: "llama3.2",
+      models: [{ id: "llama3.2", label: "llama3.2" }],
+    });
+
+    assert.deepEqual(resolved, {
+      ok: false,
+      model: "llama3.2",
+      error: "Invalid assistant model",
+    });
+  });
+
+  it("silently falls back when a saved model is gone", () => {
+    const resolved = resolveAssistantModel({
+      requestedModel: null,
+      savedModel: "old-model",
+      fallbackModel: "llama3.2",
+      models: [{ id: "llama3.2", label: "llama3.2" }],
+    });
+
+    assert.deepEqual(resolved, { ok: true, model: "llama3.2" });
+  });
+});