236 lines
7.2 KiB
TypeScript
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();
|