spb/polyllm

1import { describe, it, expect } from "vitest";2import type { PolyModel } from "@/lib/ai/core/types";3import { parseSearchQuery, searchModels, isEmptyIntent } from "@/lib/models/search";4import { buildBadgeContext } from "@/lib/models/badges";56const caps = (over: Partial<PolyModel["capabilities"]> = {}): PolyModel["capabilities"] => ({ text: true, vision: false, audioInput: false, audioOutput: false, imageGeneration: false, video: false, reasoning: false, tools: true, structuredOutput: false, streaming: true, files: false, webSearch: false, ...over });78const NOW = Date.parse("2026-09-11T12:00:00Z");9const daysAgo = (d: number) => new Date(NOW - d * 86_400_000).toISOString();1011function model(p: Partial<PolyModel> & { id: string; provider: PolyModel["provider"] }): PolyModel {12  return { key: `${p.provider}/${p.id}`, displayName: p.id, capabilities: caps(), parameters: {}, status: "active", pricing: null, metadata: { sortWeight: 0, firstSeenAt: daysAgo(200) }, ...p };13}1415const MODELS: PolyModel[] = [16  model({ id: "gpt-5.5", provider: "openai", displayName: "GPT-5.5", family: "GPT-5.5", capabilities: caps({ vision: true, reasoning: true, structuredOutput: true, files: true, webSearch: true }), limits: { contextTokens: 1_050_000, maxOutputTokens: 128_000 }, pricing: { inputPerMillion: 5, outputPerMillion: 30 }, metadata: { sortWeight: 90, firstSeenAt: daysAgo(200) } }),17  model({ id: "gpt-5.4-mini", provider: "openai", displayName: "GPT-5.4 mini", family: "GPT-5.4", capabilities: caps({ vision: true, reasoning: true, structuredOutput: true }), limits: { contextTokens: 400_000 }, pricing: { inputPerMillion: 0.4, outputPerMillion: 1.6 }, metadata: { sortWeight: 82, firstSeenAt: daysAgo(200) } }),18  model({ id: "claude-sonnet-5", provider: "anthropic", displayName: "Claude Sonnet 5", family: "Sonnet", capabilities: caps({ vision: true, reasoning: true, files: true, webSearch: true }), limits: { contextTokens: 1_000_000, maxOutputTokens: 64_000 }, pricing: { inputPerMillion: 2, outputPerMillion: 10 }, metadata: { sortWeight: 90, firstSeenAt: daysAgo(10) } }),19  model({ id: "claude-haiku-4-5", provider: "anthropic", displayName: "Claude Haiku 4.5", family: "Haiku", capabilities: caps({ vision: true }), limits: { contextTokens: 200_000 }, pricing: { inputPerMillion: 1, outputPerMillion: 5 }, metadata: { sortWeight: 60, firstSeenAt: daysAgo(200) } }),20  model({ id: "gemini-3.8-flash", provider: "gemini", displayName: "Gemini 3.8 Flash", family: "Gemini", capabilities: caps({ vision: true, reasoning: true, structuredOutput: true, webSearch: true }), limits: { contextTokens: 1_048_576 }, pricing: { inputPerMillion: 0.3, outputPerMillion: 2.5 }, metadata: { sortWeight: 100, firstSeenAt: daysAgo(200) } }),21  model({ id: "gemini-3.1-pro", provider: "gemini", displayName: "Gemini 3.1 Pro", family: "Gemini", capabilities: caps({ vision: true, reasoning: true, structuredOutput: true }), limits: { contextTokens: 2_000_000 }, pricing: { inputPerMillion: 2, outputPerMillion: 12 }, metadata: { sortWeight: 96, firstSeenAt: daysAgo(200) } }),22  model({ id: "codestral-latest", provider: "mistral", displayName: "Codestral", family: "Codestral", capabilities: caps({ structuredOutput: true }), limits: { contextTokens: 256_000 }, pricing: { inputPerMillion: 0.3, outputPerMillion: 0.9 }, metadata: { sortWeight: 50, firstSeenAt: daysAgo(200) } }),23  model({ id: "meta-llama/llama-3.1-8b-instruct", provider: "openrouter", displayName: "Meta: Llama 3.1 8B Instruct", family: "Meta", capabilities: caps(), limits: { contextTokens: 131_072 }, pricing: { inputPerMillion: 0.02, outputPerMillion: 0.05 }, metadata: { sortWeight: 20, vendor: "meta-llama", firstSeenAt: daysAgo(200) } }),24  model({ id: "deepseek-v4-flash", provider: "deepseek", displayName: "DeepSeek V4 Flash", family: "DeepSeek", capabilities: caps({ reasoning: true }), limits: { contextTokens: 128_000 }, pricing: { inputPerMillion: 0.14, outputPerMillion: 0.28 }, metadata: { sortWeight: 70, firstSeenAt: daysAgo(200) } }),25  model({ id: "old-model", provider: "openai", displayName: "Old Model", status: "deprecated", pricing: { inputPerMillion: 1, outputPerMillion: 2 }, capabilities: caps({ vision: true }) }),26];2728const ctx = buildBadgeContext(MODELS, NOW);29const ids = (r: { model: PolyModel }[]) => r.map((x) => x.model.id);3031describe("parseSearchQuery", () => {32  it("understands capability + cost intents", () => {33    const i = parseSearchQuery("cheap vision model");34    expect(i.capabilities).toEqual(["vision"]);35    expect(i.sort).toBe("cheapest");36    expect(i.text).toEqual([]);37    expect(i.chips).toContain("Vision");38  });39  it("parses context sizes", () => {40    expect(parseSearchQuery("1M context").minContext).toBe(1_000_000);41    expect(parseSearchQuery("200k").minContext).toBe(200_000);42    expect(parseSearchQuery("128k ctx").minContext).toBe(128_000);43    expect(parseSearchQuery("long context").minContext).toBe(400_000);44  });45  it("parses price caps", () => {46    expect(parseSearchQuery("under $1/M").maxInputPrice).toBe(1);47    expect(parseSearchQuery("< $2 per million output").maxOutputPrice).toBe(2);48    expect(parseSearchQuery("0.5$").maxInputPrice).toBe(0.5);49  });50  it("recognises brands, json schema, open source and reasoning", () => {51    const g = parseSearchQuery("fastest gemini");52    expect(g.brands.map((b) => b.key)).toEqual(["gemini"]);53    expect(g.sort).toBe("fastest");54    expect(parseSearchQuery("json schema").capabilities).toEqual(["structuredOutput"]);55    expect(parseSearchQuery("open source").openWeights).toBe(true);56    expect(parseSearchQuery("open-weights coder").coding).toBe(true);57    expect(parseSearchQuery("reasoning").capabilities).toEqual(["reasoning"]);58    expect(parseSearchQuery("OpenAI").brands[0].provider).toBe("openai");59  });60  it("keeps free text and detects empty intents", () => {61    expect(parseSearchQuery("sonnet").text).toEqual(["sonnet"]);62    expect(isEmptyIntent(parseSearchQuery("   "))).toBe(true);63    expect(isEmptyIntent(parseSearchQuery("the model"))).toBe(true);64  });65});6667describe("searchModels", () => {68  it("filters by capability and ranks cheapest first", () => {69    const r = searchModels(MODELS, "cheap vision model", { ctx });70    expect(r.every((x) => x.model.capabilities.vision)).toBe(true);71    expect(r[0].model.id).toBe("gemini-3.8-flash");72    expect(ids(r)).not.toContain("codestral-latest");73    // deprecated models sink to the bottom74    expect(r[r.length - 1].model.id).toBe("old-model");75  });76  it("applies context and price constraints", () => {77    expect(ids(searchModels(MODELS, "1M context", { ctx }))).toEqual(expect.arrayContaining(["gpt-5.5", "claude-sonnet-5", "gemini-3.8-flash", "gemini-3.1-pro"]));78    expect(ids(searchModels(MODELS, "1M context", { ctx }))).not.toContain("gpt-5.4-mini");79    const cheap = searchModels(MODELS, "under $1/M", { ctx });80    expect(cheap.every((x) => (x.model.pricing?.inputPerMillion ?? 99) <= 1)).toBe(true);81  });82  it("matches brands across native providers and gateways", () => {83    const g = searchModels(MODELS, "fastest gemini", { ctx });84    expect(ids(g)).toEqual(["gemini-3.8-flash", "gemini-3.1-pro"]);85    const llama = searchModels(MODELS, "llama", { ctx });86    expect(ids(llama)).toEqual(["meta-llama/llama-3.1-8b-instruct"]);87  });88  it("ranks exact and prefix name matches first and tolerates a typo", () => {89    expect(searchModels(MODELS, "GPT-5.5", { ctx })[0].model.id).toBe("gpt-5.5");90    expect(searchModels(MODELS, "sonnet", { ctx })[0].model.id).toBe("claude-sonnet-5");91    expect(searchModels(MODELS, "sonet", { ctx })[0].model.id).toBe("claude-sonnet-5");92    expect(searchModels(MODELS, "zzzz", { ctx })).toHaveLength(0);93  });94  it("uses user labels and favorites", () => {95    const r = searchModels(MODELS, "research", { ctx, labels: { "anthropic/claude-sonnet-5": "My research model" } });96    expect(ids(r)).toEqual(["claude-sonnet-5"]);97    const fav = searchModels(MODELS, "gemini", { ctx, favorites: new Set(["gemini/gemini-3.1-pro"]) });98    expect(fav[0].model.id).toBe("gemini-3.1-pro");99  });100  it("handles open source, coding, json schema and new", () => {101    expect(ids(searchModels(MODELS, "open source", { ctx }))).toEqual(expect.arrayContaining(["meta-llama/llama-3.1-8b-instruct", "deepseek-v4-flash"]));102    expect(searchModels(MODELS, "coding", { ctx })[0].model.id).toBe("codestral-latest");103    expect(searchModels(MODELS, "json schema", { ctx }).every((x) => x.model.capabilities.structuredOutput)).toBe(true);104    expect(ids(searchModels(MODELS, "new", { ctx }))).toEqual(["claude-sonnet-5"]);105  });106});107