Files
2026-07-13 12:28:55 +08:00

236 lines
7.2 KiB
TypeScript

#!/usr/bin/env bun
import { z } from "zod";
import path from "node:path";
import { mkdir, rm, readdir, stat } from "node:fs/promises";
// Helicone public model registry endpoint
const DEFAULT_ENDPOINT =
"https://jawn.helicone.ai/v1/public/model-registry/models";
// Zod schemas to validate the Helicone response
const Pricing = z
.object({
prompt: z.number().optional(),
completion: z.number().optional(),
cacheRead: z.number().optional(),
cacheWrite: z.number().optional(),
reasoning: z.number().optional(),
})
.passthrough();
const Endpoint = z
.object({
provider: z.string(),
providerSlug: z.string().optional(),
supportsPtb: z.boolean().optional(),
pricing: Pricing.optional(),
})
.passthrough();
const ModelItem = z
.object({
id: z.string(),
name: z.string(),
author: z.string().optional(),
contextLength: z.number().optional(),
maxOutput: z.number().optional(),
trainingDate: z.string().optional(),
description: z.string().optional(),
inputModalities: z.array(z.string()).optional(),
outputModalities: z.array(z.string()).optional(),
supportedParameters: z.array(z.string()).optional(),
endpoints: z.array(Endpoint).optional(),
})
.passthrough();
const HeliconeResponse = z
.object({
data: z.object({
models: z.array(ModelItem),
total: z.number().optional(),
filters: z.any().optional(),
}),
})
.passthrough();
interface ExistingModel {
base_model?: string;
base_model_omit?: string[];
}
async function loadExistingModel(filePath: string): Promise<ExistingModel | undefined> {
const file = Bun.file(filePath);
if (!(await file.exists())) return undefined;
return await import(filePath, { with: { type: "toml" } }).then(
(mod) => mod.default as ExistingModel,
);
}
function pickEndpoint(m: z.infer<typeof ModelItem>) {
if (!m.endpoints || m.endpoints.length === 0) return undefined;
// Prefer endpoint that matches author if available
if (m.author) {
const match = m.endpoints.find((e) => e.provider === m.author);
if (match) return match;
}
return m.endpoints[0];
}
function boolFromParams(params: string[] | undefined, keys: string[]): boolean {
if (!params) return false;
const set = new Set(params.map((p) => p.toLowerCase()));
return keys.some((k) => set.has(k.toLowerCase()));
}
function sanitizeModalities(values: string[] | undefined): string[] {
if (!values) return ["text"]; // default to text
const allowed = new Set(["text", "audio", "image", "video", "pdf"]);
const out = values.map((v) => v.toLowerCase()).filter((v) => allowed.has(v));
return out.length > 0 ? out : ["text"];
}
function formatToml(model: z.infer<typeof ModelItem>, existing: ExistingModel | undefined) {
const ep = pickEndpoint(model);
const pricing = ep?.pricing;
const supported = model.supportedParameters ?? [];
const nowISO = new Date().toISOString().slice(0, 10);
const rdRaw = model.trainingDate ? String(model.trainingDate) : nowISO;
const releaseDate = rdRaw.slice(0, 10);
const lastUpdated = releaseDate;
const knowledge = model.trainingDate
? String(model.trainingDate).slice(0, 7)
: undefined;
const attachment = false; // Not exposed by Helicone registry
const temperature = boolFromParams(supported, ["temperature"]);
const toolCall = boolFromParams(supported, ["tools", "tool_choice"]);
const reasoning = boolFromParams(supported, [
"reasoning",
"include_reasoning",
]);
const inputMods = sanitizeModalities(model.inputModalities);
const outputMods = sanitizeModalities(model.outputModalities);
const lines: string[] = [];
if (existing?.base_model !== undefined) {
lines.push(`base_model = "${existing.base_model}"`);
}
if (existing?.base_model_omit !== undefined) {
lines.push(
`base_model_omit = [${existing.base_model_omit.map((item) => `"${item}"`).join(", ")}]`,
);
}
lines.push(`name = "${model.name.replaceAll('"', '\\"')}"`);
lines.push(`release_date = "${releaseDate}"`);
lines.push(`last_updated = "${lastUpdated}"`);
lines.push(`attachment = ${attachment}`);
lines.push(`reasoning = ${reasoning}`);
lines.push(`temperature = ${temperature}`);
lines.push(`tool_call = ${toolCall}`);
if (knowledge) lines.push(`knowledge = "${knowledge}"`);
lines.push(`open_weights = false`);
lines.push("");
if (
pricing &&
(pricing.prompt ??
pricing.completion ??
pricing.cacheRead ??
pricing.cacheWrite ??
(reasoning && pricing.reasoning)) !== undefined
) {
lines.push(`[cost]`);
if (pricing.prompt !== undefined) lines.push(`input = ${pricing.prompt}`);
if (pricing.completion !== undefined)
lines.push(`output = ${pricing.completion}`);
if (reasoning && pricing.reasoning !== undefined)
lines.push(`reasoning = ${pricing.reasoning}`);
if (pricing.cacheRead !== undefined)
lines.push(`cache_read = ${pricing.cacheRead}`);
if (pricing.cacheWrite !== undefined)
lines.push(`cache_write = ${pricing.cacheWrite}`);
lines.push("");
}
const context = model.contextLength ?? 0;
const output = model.maxOutput ?? 4096;
lines.push(`[limit]`);
lines.push(`context = ${context}`);
lines.push(`output = ${output}`);
lines.push("");
lines.push(`[modalities]`);
lines.push(`input = [${inputMods.map((m) => `"${m}"`).join(", ")}]`);
lines.push(`output = [${outputMods.map((m) => `"${m}"`).join(", ")}]`);
return lines.join("\n") + "\n";
}
async function main() {
const endpoint = DEFAULT_ENDPOINT;
const outDir = path.join(
import.meta.dirname,
"..",
"..",
"..",
"providers",
"helicone",
"models",
);
const res = await fetch(endpoint);
if (!res.ok) {
console.error(`Failed to fetch registry: ${res.status} ${res.statusText}`);
process.exit(1);
}
const json = await res.json();
const parsed = HeliconeResponse.safeParse(json);
if (!parsed.success) {
parsed.error.cause = json;
console.error("Invalid Helicone response:", parsed.error.errors);
console.error("When parsing:", parsed.error.cause);
process.exit(1);
}
const models = parsed.data.data.models;
const existing = new Map<string, ExistingModel>();
await mkdir(outDir, { recursive: true });
for await (const file of new Bun.Glob("**/*.toml").scan({ cwd: outDir })) {
const filePath = path.join(outDir, file);
const model = await loadExistingModel(filePath);
if (model !== undefined) existing.set(file, model);
}
// Clean output directory: remove subfolders and existing TOML files
for (const entry of await readdir(outDir)) {
const p = path.join(outDir, entry);
const st = await stat(p);
if (st.isDirectory()) {
await rm(p, { recursive: true, force: true });
} else if (st.isFile() && entry.endsWith(".toml")) {
await rm(p, { force: true });
}
}
let created = 0;
for (const m of models) {
const fileSafeId = m.id.replaceAll("/", "-");
const filePath = path.join(outDir, `${fileSafeId}.toml`);
const toml = formatToml(m, existing.get(`${fileSafeId}.toml`));
await Bun.write(filePath, toml);
created++;
}
console.log(
`Generated ${created} model file(s) under providers/helicone/models/*.toml`,
);
}
await main();