Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/klaus/__tests__/models.test.ts

Raw
import type { Api, Model } from "@earendil-works/pi-ai";
import { getBuiltinModels } from "@earendil-works/pi-ai/providers/all";
import { describe, expect, it } from "vitest";
import {
	isProjectableClaudeModel,
	klausModels,
	projectBuiltinModels,
	projectModel,
	projectModels,
} from "../src/models";

function anthropic(modelId: string): Model<Api> {
	const model = getBuiltinModels("anthropic").find((m) => m.id === modelId);
	if (!model) throw new Error(`missing builtin anthropic model ${modelId}`);
	return model;
}

describe("Klaus model projection", () => {
	it("projects every non-dated builtin Claude catalog model", () => {
		const builtin = getBuiltinModels("anthropic");
		const expected = builtin
			.filter((model) => isProjectableClaudeModel(model))
			.map((model) => model.id);
		expect(projectBuiltinModels().map((model) => model.id)).toEqual(
			expect.arrayContaining(expected),
		);
		expect(projectBuiltinModels()).toHaveLength(expected.length);
		for (const model of builtin) {
			if (/-\d{8}$/.test(model.id) || model.id.endsWith("-latest"))
				expect(projectBuiltinModels().some((m) => m.id === model.id)).toBe(
					false,
				);
		}
	});

	it("keeps the curated alias table inside the projected catalog", () => {
		const ids = projectBuiltinModels().map((model) => model.id);
		for (const id of [
			"claude-fable-5-1",
			"claude-haiku-4-5",
			"claude-opus-5",
			"claude-sonnet-5",
			"claude-sonnet-4-6",
		])
			expect(ids).toContain(id);
	});

	it("filters foreign providers from a registry enumeration", () => {
		const openai = { ...anthropic("claude-sonnet-5"), provider: "openrouter" };
		expect(projectModels([openai])).toEqual([]);
	});

	it("copies Pi catalog metadata without transport fields", () => {
		const source = anthropic("claude-sonnet-5");
		const projected = projectModel(source);
		expect(projected).toEqual({
			id: source.id,
			name: source.name,
			reasoning: source.reasoning,
			thinkingLevelMap: source.thinkingLevelMap,
			input: [...source.input],
			cost: source.cost,
			compat: source.compat,
			contextWindow: source.contextWindow,
			maxTokens: source.maxTokens,
		});
		expect(projected).not.toHaveProperty("api");
		expect(projected).not.toHaveProperty("baseUrl");
		expect(projected).not.toHaveProperty("headers");
		expect(projected).not.toHaveProperty("samplingParams");
	});

	it("inherits user overrides that stay within the builtin baseline", () => {
		const baseline = anthropic("claude-opus-5");
		const projected = projectModel({
			...baseline,
			name: "Opus (work account)",
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
			contextWindow: baseline.contextWindow - 1000,
			maxTokens: baseline.maxTokens - 1000,
		});
		expect(projected.name).toBe("Opus (work account)");
		expect(projected.cost).toEqual({
			input: 0,
			output: 0,
			cacheRead: 0,
			cacheWrite: 0,
		});
		expect(projected.contextWindow).toBe(baseline.contextWindow - 1000);
		expect(projected.maxTokens).toBe(baseline.maxTokens - 1000);
	});

	it("clamps overrides that exceed what the Agent SDK path can deliver", () => {
		const baseline = anthropic("claude-haiku-4-5");
		const projected = projectModel({
			...baseline,
			contextWindow: 1_000_000,
			maxTokens: baseline.maxTokens * 4,
		});
		expect(projected.contextWindow).toBe(baseline.contextWindow);
		expect(projected.maxTokens).toBe(baseline.maxTokens);
	});

	it("keeps Fable fallback models priced from the builtin catalog", () => {
		const projected = projectModel(anthropic("claude-fable-5-1"));
		expect(projected.compat).toMatchObject({
			allowedFallbackModels: [
				{
					provider: "anthropic",
					model: "claude-opus-4-8",
					cost: anthropic("claude-opus-4-8").cost,
				},
				{
					provider: "anthropic",
					model: "claude-opus-5",
					cost: anthropic("claude-opus-5").cost,
				},
			],
		});
	});

	it("does not mutate the Pi catalog entry it projects", () => {
		const source = anthropic("claude-fable-5-1");
		const before = structuredClone(source);
		const projected = projectModel(source);
		projected.cost.input = 999;
		projected.input.push("image");
		expect(source).toEqual(before);
	});

	it("filters the enumeration it is given without touching other models", () => {
		const catalog = getBuiltinModels("anthropic");
		const models = projectModels(
			catalog.filter((model) => model.id !== "claude-haiku-4-5"),
		);
		expect(models.map((model) => model.id)).not.toContain("claude-haiku-4-5");
		expect(models).toHaveLength(
			catalog.filter(isProjectableClaudeModel).length - 1,
		);
	});

	it("falls back to the builtin catalog without a registry or Claude models", () => {
		const expected = projectBuiltinModels().map((model) => model.id);
		const empty = {
			getProvider: () => ({ getModels: () => [] }),
		} as unknown as Parameters<typeof klausModels>[0];
		const throwing = {
			getProvider: () => {
				throw new Error("registry unavailable");
			},
		} as unknown as Parameters<typeof klausModels>[0];
		for (const registry of [undefined, empty, throwing]) {
			const models = klausModels(registry);
			expect(models.map((model) => model.id)).toEqual(expected);
			for (const model of models)
				expect(model).toMatchObject({
					provider: "klaus",
					api: "klaus-agent-sdk",
					baseUrl: "https://klaus.invalid",
				});
		}
	});

	it("projects the registry's Anthropic catalog when it has Claude models", () => {
		const registry = {
			getProvider: (id: string) =>
				id === "anthropic"
					? {
							getModels: () => [
								{ ...anthropic("claude-opus-5"), name: "renamed" },
							],
						}
					: undefined,
		} as unknown as Parameters<typeof klausModels>[0];
		expect(klausModels(registry)).toMatchObject([
			{ id: "claude-opus-5", name: "renamed", provider: "klaus" },
		]);
	});
});