spb/polyllm

1import { describe, it, expect } from "vitest";2import { normalizeOpenRouterModel } from "@/lib/ai/providers/openrouter";3import { normalizeMistralModel } from "@/lib/ai/providers/mistral";4import { normalizeDeepSeekModel } from "@/lib/ai/providers/deepseek";5import { normalizeKimiModel } from "@/lib/ai/providers/kimi";6import { normalizeCerebrasModel } from "@/lib/ai/providers/cerebras";7import { normalizeChatCompletionStream, toChatCompletionMessages } from "@/lib/ai/providers/shared/openai-compat/chat-completions";8import { stripJsonFence } from "@/lib/chat/service";9import type { ChatCompletionChunk } from "openai/resources/chat/completions";1011describe("OpenRouter model mapping", () => {12  it("converts per-token USD strings and supported_parameters into a PolyModel", () => {13    const m = normalizeOpenRouterModel({14      id: "openai/gpt-6-astra",15      name: "OpenAI: GPT-6 Astra",16      created: 1788552838,17      context_length: 1050000,18      architecture: { input_modalities: ["file", "image", "text"], output_modalities: ["text"] },19      pricing: { prompt: "0.00001", completion: "0.00005", input_cache_read: "0.000001", overrides: [{ min_prompt_tokens: 272000, prompt: "0.00002", completion: "0.000075" }] },20      top_provider: { context_length: 1050000, max_completion_tokens: 128000, is_moderated: true },21      supported_parameters: ["include_reasoning", "max_completion_tokens", "max_tokens", "reasoning", "reasoning_effort", "response_format", "seed", "structured_outputs", "tool_choice", "tools"],22      reasoning: { mandatory: true, default_enabled: true, supported_efforts: ["max", "xhigh", "high", "medium", "low"], default_effort: "medium" },23    })!;24    expect(m.key).toBe("openrouter/openai/gpt-6-astra");25    expect(m.pricing?.inputPerMillion).toBe(10);26    expect(m.pricing?.outputPerMillion).toBe(50);27    expect(m.pricing?.cachedInputPerMillion).toBe(1);28    expect(m.pricing?.longContext?.thresholdTokens).toBe(272000);29    expect(m.capabilities.vision).toBe(true);30    expect(m.capabilities.files).toBe(true);31    expect(m.capabilities.reasoning).toBe(true);32    expect(m.parameters.temperature).toBe(false); // not in supported_parameters33    expect(m.parameters.seed).toBe(true);34    expect(m.parameters.reasoningEffortLevels).toEqual(["low", "medium", "high", "xhigh", "max"]); // mandatory → no "none"35    expect(m.limits?.maxOutputTokens).toBe(128000);36    expect(m.family).toBe("OpenAI");37  });38  it("hides non-text and expired models", () => {39    expect(normalizeOpenRouterModel({ id: "x/img", name: "img", architecture: { output_modalities: ["image"] } })).toBeNull();40    expect(normalizeOpenRouterModel({ id: "x/old", name: "old", expiration_date: "2020-01-01", architecture: { output_modalities: ["text"] } })).toBeNull();41  });42});4344describe("Mistral model mapping", () => {45  it("uses the listing capabilities and skips non-chat models", () => {46    const m = normalizeMistralModel({ id: "mistral-small-latest", name: "mistral-small-2603", max_context_length: 262144, aliases: ["mistral-small-2603"], capabilities: { completion_chat: true, function_calling: true, vision: true, reasoning: true } })!;47    expect(m.capabilities.vision).toBe(true);48    expect(m.capabilities.reasoning).toBe(true);49    expect(m.limits?.contextTokens).toBe(262144);50    expect(m.metadata?.resolvesTo).toBe("mistral-small-2603");51    expect(normalizeMistralModel({ id: "mistral-embed", capabilities: { completion_chat: false } })).toBeNull();52    expect(normalizeMistralModel({ id: "mistral-ocr-latest", capabilities: { completion_chat: false, ocr: true } })).toBeNull();53  });54});5556describe("id-only listings", () => {57  it("normalizes DeepSeek / Kimi / Cerebras ids with sane defaults", () => {58    expect(normalizeDeepSeekModel({ id: "deepseek-v4-flash-vision-exp" })?.capabilities.vision).toBe(true);59    expect(normalizeDeepSeekModel({ id: "text-embedding" })).toBeNull();60    expect(normalizeKimiModel({ id: "kimi-k3" })?.family).toBe("Kimi K3");61    expect(normalizeKimiModel({ id: "kimi-k2.6" })?.parameters.temperature).toBe(false); // only temperature=1 is accepted62    expect(normalizeKimiModel({ id: "kimi-k3", supports_thinking_type: "only", reasoning_efforts: { support: true, valid_efforts: ["low", "high", "max"] } })?.parameters.reasoningEffortLevels).toEqual(["low", "high", "max"]);63    expect(normalizeCerebrasModel({ id: "gpt-oss-120b" })?.parameters.reasoningEffortLevels).toEqual(["low", "medium", "high"]);64  });65});6667describe("chat-completions normalizer extras", () => {68  async function* fake(chunks: Partial<ChatCompletionChunk>[]): AsyncIterable<ChatCompletionChunk> {69    for (const c of chunks) yield { id: "c", object: "chat.completion.chunk", created: 0, model: "m", choices: [], ...c } as ChatCompletionChunk;70  }71  it("handles Mistral thinking chunks, DeepSeek cache hits and OpenRouter cost", async () => {72    const events = [];73    const stream = fake([74      { choices: [{ index: 0, delta: { content: [{ type: "thinking", thinking: [{ type: "text", text: "hmm " }] }, { type: "text", text: "Hi" }] as never }, finish_reason: null }] },75      { choices: [{ index: 0, delta: { reasoning_content: "more" } as never, finish_reason: "stop" }] },76      { choices: [], usage: { prompt_tokens: 100, completion_tokens: 10, total_tokens: 110, prompt_cache_hit_tokens: 60, cost: 0.0042 } as never },77    ]);78    for await (const ev of normalizeChatCompletionStream(stream, "openrouter")) events.push(ev);79    const reasoning = events.filter((e) => e.type === "reasoning-delta").map((e) => (e as { text: string }).text).join("");80    const text = events.filter((e) => e.type === "text-delta").map((e) => (e as { text: string }).text).join("");81    expect(reasoning).toBe("hmm more");82    expect(text).toBe("Hi");83    const usage = events.find((e) => e.type === "usage") as { usage: { cachedInputTokens?: number } };84    expect(usage.usage.cachedInputTokens).toBe(60);85    const cost = events.find((e) => e.type === "provider-data") as unknown as { data: { exactCostUsd: number } };86    expect(cost.data.exactCostUsd).toBe(0.0042);87    expect(events.at(-1)).toEqual({ type: "finish", reason: "stop" });88  });89  it("replays reasoning_content on assistant turns when asked", () => {90    const msgs = toChatCompletionMessages(undefined, [{ role: "assistant", content: [{ type: "reasoning", text: "think" }, { type: "text", text: "answer" }] }], { replayReasoningContent: true });91    expect((msgs[0] as unknown as { reasoning_content?: string }).reasoning_content).toBe("think");92    const plain = toChatCompletionMessages(undefined, [{ role: "assistant", content: [{ type: "reasoning", text: "think" }, { type: "text", text: "answer" }] }]);93    expect((plain[0] as unknown as { reasoning_content?: string }).reasoning_content).toBeUndefined();94  });95  it("strips markdown fences around JSON-mode answers", () => {96    expect(stripJsonFence('```json\n{"a":1}\n```')).toBe('{"a":1}');97    expect(stripJsonFence('{"a":1}')).toBe('{"a":1}');98    expect(stripJsonFence("text ```json x``` more")).toBe("text ```json x``` more");99  });100});101102describe("strict schema helper", () => {103  it("adds additionalProperties:false recursively", async () => {104    const { withNoAdditionalProps } = await import("@/lib/ai/providers/shared/openai-compat/factory");105    const out = withNoAdditionalProps({ type: "object", properties: { a: { type: "object", properties: { b: { type: "string" } } }, list: { type: "array", items: { type: "object", properties: { c: { type: "number" } } } } }, required: ["a"] }) as { additionalProperties: boolean; properties: { a: { additionalProperties: boolean }; list: { items: { additionalProperties: boolean } } } };106    expect(out.additionalProperties).toBe(false);107    expect(out.properties.a.additionalProperties).toBe(false);108    expect(out.properties.list.items.additionalProperties).toBe(false);109  });110});111