pi-mono/packages/ai/scripts/generate-image-models.ts

#!/usr/bin/env tsx

import { writeFileSync } from "fs";
import { dirname, join } from "path";
import { fileURLToPath } from "url";
import type { ImagesModel } from "../src/types.js";

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);
const packageRoot = join(__dirname, "..");
const OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1";

interface OpenRouterModelRecord {
	id: string;
	name: string;
	context_length?: number;
	architecture?: {
		input_modalities?: string[];
		output_modalities?: string[];
	};
	pricing?: {
		prompt?: string;
		completion?: string;
		input_cache_read?: string;
		input_cache_write?: string;
	};
}

async function fetchOpenRouterImageModels(): Promise<ImagesModel<"openrouter-images">[]> {
	try {
		console.log("Fetching image models from OpenRouter API...");
		const response = await fetch(`${OPENROUTER_BASE_URL}/models?output_modalities=image`);
		const data = (await response.json()) as { data?: OpenRouterModelRecord[] };
		const models: ImagesModel<"openrouter-images">[] = [];

		for (const model of data.data ?? []) {
			const input = Array.from(
				new Set(
					(model.architecture?.input_modalities ?? [])
						.filter((modality): modality is "text" | "image" => modality === "text" || modality === "image"),
				),
			);
			const output = Array.from(
				new Set(
					(model.architecture?.output_modalities ?? []).filter(
						(modality): modality is "text" | "image" => modality === "text" || modality === "image",
					),
				),
			);

			if (!output.includes("image")) continue;
			if (input.length === 0) input.push("text");

			models.push({
				id: model.id,
				name: model.name,
				api: "openrouter-images",
				provider: "openrouter",
				baseUrl: OPENROUTER_BASE_URL,
				input,
				output,
				cost: {
					input: parseFloat(model.pricing?.prompt || "0") * 1_000_000,
					output: parseFloat(model.pricing?.completion || "0") * 1_000_000,
					cacheRead: parseFloat(model.pricing?.input_cache_read || "0") * 1_000_000,
					cacheWrite: parseFloat(model.pricing?.input_cache_write || "0") * 1_000_000,
				},
			});
		}

		console.log(`Fetched ${models.length} image models from OpenRouter`);
		return models;
	} catch (error) {
		console.error("Failed to fetch OpenRouter image models:", error);
		return [];
	}
}

function generateImageModelsFile(models: ImagesModel<"openrouter-images">[]): string {
	const imageModelsByProvider = {
		openrouter: Object.fromEntries(
			models
				.sort((a, b) => a.id.localeCompare(b.id))
				.map((model) => [
					model.id,
					`{
			id: ${JSON.stringify(model.id)},
			name: ${JSON.stringify(model.name)},
			api: ${JSON.stringify(model.api)},
			provider: ${JSON.stringify(model.provider)},
			baseUrl: ${JSON.stringify(model.baseUrl)},
			input: ${JSON.stringify(model.input)},
			output: ${JSON.stringify(model.output)},
			cost: ${JSON.stringify(model.cost, null, 2).replace(/^/gm, "\t")}
		} satisfies ImagesModel<${JSON.stringify(model.api)}>`,
				]),
		),
	};

	const providerEntries = Object.entries(imageModelsByProvider)
		.map(([provider, providerModels]) => {
			const modelEntries = Object.entries(providerModels)
				.map(([id, serialized]) => `\t\t${JSON.stringify(id)}: ${serialized},`)
				.join("\n");
			return `\t${JSON.stringify(provider)}: {\n${modelEntries}\n\t},`;
		})
		.join("\n");

	return `// This file is auto-generated by scripts/generate-image-models.ts
// Do not edit manually - run 'npm run generate-image-models' to update

import type { ImagesApi, ImagesModel } from "./types.ts";

export const IMAGE_MODELS = {
${providerEntries}
} as const satisfies Record<string, Record<string, ImagesModel<ImagesApi>>>;
`;
}

async function main(): Promise<void> {
	const models = await fetchOpenRouterImageModels();
	const output = generateImageModelsFile(models);
	const outputPath = join(packageRoot, "src", "image-models.generated.ts");
	writeFileSync(outputPath, output, "utf-8");
	console.log(`Generated ${outputPath}`);
}

main().catch((error) => {
	console.error(error);
	process.exit(1);
});