pi-ext

import type { Api, Model } from "@earendil-works/pi-ai";
import { describe, expect, it } from "vitest";
import {
	fromPiCatalog,
	modelsDevZaiEntry,
	parseZaiPricing,
	unknownZaiIds,
	zaiFreeModels,
} from "../src/catalog.ts";

/** Shape of Z.ai's pricing.md as served in September 2026. */
const ZAI_PRICING = `# Pricing

### Latest Models

| Model | Input | Cached Input | Cached Input Storage | Output |
| :- | :- | :- | :- | :- |
| GLM-5.3-Flash | \\$0.15 | \\$0.03 | Limited-time Free | \\$0.50 |

### Text Models

| Model | Input | Cached Input | Cached Input Storage | Output |
| :- | :- | :- | :- | :- |
| GLM-4.7 | \\$0.6 | \\$0.11 | Limited-time Free | \\$2.2 |
| GLM-4.7-Flash | Free | Free | Free | Free |
| GLM-4.5-Flash | Free | Free | Free | Free |

### Vision Models

| Model | Input | Cached Input | Cached Input Storage | Output |
| :- | :- | :- | :- | :- |
| GLM-4.6V-Flash | Free | Free | Free | Free |

### Image Generation Models

| Model | Price |
| :- | :- |
| CogView-4 | \\$0.01 |
`;

const piModel = (id: string, extra: Partial<Model<Api>> = {}): Model<Api> => ({
	id,
	name: id,
	api: "openai-completions",
	provider: "x",
	baseUrl: "https://example.invalid/v1",
	reasoning: false,
	input: ["text"],
	cost: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 },
	contextWindow: 131_072,
	maxTokens: 8_192,
	...extra,
});

describe("Z.ai pricing page", () => {
	it("finds models whose input and output are Free, and which are vision", () => {
		expect(parseZaiPricing(ZAI_PRICING)).toEqual({
			free: ["glm-4.7-flash", "glm-4.5-flash", "glm-4.6v-flash"],
			vision: new Set(["glm-4.6v-flash"]),
		});
	});

	it("never treats limited-time free storage as a free model", () => {
		expect(parseZaiPricing(ZAI_PRICING)?.free).not.toContain("glm-5.3-flash");
	});

	it("reports a page without model tables as unparseable", () => {
		expect(parseZaiPricing("# Pricing\n\nMoved.")).toBeUndefined();
		expect(
			parseZaiPricing(
				"| Model | Input | Cached Input | Cached Input Storage | Output |\n| GLM-4.7 | $1 | $1 | Free | $2 |",
			),
		).toEqual({ free: [], vision: new Set() });
	});
});

describe("models.dev confirmation", () => {
	const payload = (api: string, cost: number) => ({
		zai: {
			api,
			models: {
				"glm-9-flash": {
					name: "GLM 9 Flash",
					cost: { input: cost, output: cost },
					limit: { context: 300_000, output: 64_000 },
					modalities: { input: ["text", "image"] },
				},
			},
		},
	});

	it("confirms a cost-0 model on the general endpoint with its limits", () => {
		expect(
			modelsDevZaiEntry(
				payload("https://api.z.ai/api/paas/v4", 0),
				"glm-9-flash",
			),
		).toMatchObject({
			id: "glm-9-flash",
			baseUrl: "https://api.z.ai/api/paas/v4",
			contextWindow: 300_000,
			maxTokens: 64_000,
			input: ["text", "image"],
		});
	});

	it("rejects priced models and other endpoints such as the coding plan", () => {
		expect(
			modelsDevZaiEntry(
				payload("https://api.z.ai/api/paas/v4", 0.1),
				"glm-9-flash",
			),
		).toBeUndefined();
		expect(
			modelsDevZaiEntry(
				payload("https://api.z.ai/api/coding/paas/v4", 0),
				"glm-9-flash",
			),
		).toBeUndefined();
	});

	it("lists known free ids at once and unknown ones only when confirmed", () => {
		const pricing = {
			free: ["glm-4.7-flash", "glm-9-flash"],
			vision: new Set<string>(),
		};
		const entry = modelsDevZaiEntry(
			payload("https://api.z.ai/api/paas/v4", 0),
			"glm-9-flash",
		);
		if (!entry) throw new Error("models.dev entry not confirmed");
		const confirmed = new Map([["glm-9-flash", entry]]);

		expect(unknownZaiIds(pricing)).toEqual(["glm-9-flash"]);
		expect(zaiFreeModels(pricing, new Map()).map((m) => m.id)).toEqual([
			"zai/glm-4.7-flash",
		]);
		expect(zaiFreeModels(pricing, confirmed).map((m) => m.id)).toEqual([
			"zai/glm-4.7-flash",
			"zai/glm-9-flash",
		]);
	});
});

describe("Pi catalog free rules", () => {
	it("keeps concrete Gemini Flash ids, newest first, flash before lite", () => {
		const catalog = [
			"gemini-2.5-flash",
			"gemini-3.5-flash-lite",
			"gemini-flash-latest",
			"gemini-3.9-flash",
			"gemini-3-flash-preview",
			"gemini-3.5-flash",
			"gemini-3.5-pro",
		].map((id) => piModel(id, { api: "google-generative-ai" }));

		expect(fromPiCatalog("google", catalog).map((m) => m.id)).toEqual([
			"google/gemini-3.9-flash",
			"google/gemini-3.5-flash",
			"google/gemini-3.5-flash-lite",
			"google/gemini-2.5-flash",
		]);
	});

	it("carries Pi's metadata and zeroes cost", () => {
		const map = { off: null, minimal: null, low: "low", high: "high" } as const;
		const [model] = fromPiCatalog("google", [
			piModel("gemini-3.9-flash", {
				api: "google-generative-ai",
				thinkingLevelMap: map,
				contextWindow: 2_000_000,
				input: ["text", "image"],
			}),
		]);

		expect(model).toMatchObject({
			api: "google-generative-ai",
			thinkingLevelMap: map,
			contextWindow: 2_000_000,
			input: ["text", "image"],
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
		});
	});

	it("drops non-chat Groq models and ranks known preferences first", () => {
		const catalog = [
			"qwen/qwen3.8-27b",
			"openai/gpt-oss-safeguard-20b",
			"whisper-large-v3",
			"openai/gpt-oss-120b",
		].map((id) => piModel(id));

		expect(fromPiCatalog("groq", catalog).map((m) => m.id)).toEqual([
			"groq/openai/gpt-oss-120b",
			"groq/qwen/qwen3.8-27b",
		]);
	});

	it("keeps only Mistral -latest text models", () => {
		const catalog = [
			"devstral-2512",
			"devstral-latest",
			"voxtral-small-latest",
			"magistral-medium-latest",
		].map((id) => piModel(id, { api: "mistral-conversations" }));

		expect(fromPiCatalog("mistral", catalog).map((m) => m.id)).toEqual([
			"mistral/devstral-latest",
			"mistral/magistral-medium-latest",
		]);
	});

	it("keeps only NVIDIA models Pi prices at 0 that NVIDIA serves", () => {
		const free = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
		const catalog = [
			piModel("z-ai/glm-5.3", { cost: free }),
			piModel("nvidia/nemotron-3-ultra-550b-a55b"),
			piModel("vendor/gone", { cost: free }),
		];

		expect(
			fromPiCatalog(
				"nvidia",
				catalog,
				new Set(["z-ai/glm-5.3", "nvidia/nemotron-3-ultra-550b-a55b"]),
			).map((m) => m.id),
		).toEqual(["nvidia/z-ai/glm-5.3"]);
	});
});