chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,291 @@
|
||||
#!/usr/bin/env bun
|
||||
|
||||
/**
|
||||
* Generates Databricks model TOML files from the Foundation Model API endpoint.
|
||||
*
|
||||
* Each Databricks endpoint exposes a model from another provider (Anthropic,
|
||||
* OpenAI, Google, etc.), so the generated TOML uses base_model to inherit
|
||||
* provider-agnostic metadata from models.dev.
|
||||
*
|
||||
* Usage:
|
||||
* DATABRICKS_HOST=<host> DATABRICKS_TOKEN=<pat> bun run databricks:generate
|
||||
* bun run databricks:generate --workspace <host> --token <pat>
|
||||
*
|
||||
* Flags:
|
||||
* --dry-run: Preview changes without writing files
|
||||
* --new-only: Only create new models, skip updating existing ones
|
||||
*/
|
||||
|
||||
import { z } from "zod";
|
||||
import path from "node:path";
|
||||
import { mkdir, readFile } from "node:fs/promises";
|
||||
import { existsSync } from "node:fs";
|
||||
|
||||
const args = process.argv.slice(2);
|
||||
const flag = (name: string) => {
|
||||
const i = args.indexOf(`--${name}`);
|
||||
return i !== -1 ? args[i + 1] : undefined;
|
||||
};
|
||||
const dryRun = args.includes("--dry-run");
|
||||
const newOnly = args.includes("--new-only");
|
||||
|
||||
const host = flag("workspace") ?? process.env.DATABRICKS_HOST;
|
||||
const token = flag("token") ?? process.env.DATABRICKS_TOKEN;
|
||||
|
||||
if (!host || !token) {
|
||||
console.error(
|
||||
"Usage: DATABRICKS_HOST=<host> DATABRICKS_TOKEN=<pat> bun run databricks:generate",
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const workspace = host.replace(/^https?:\/\//, "").replace(/\/$/, "");
|
||||
const PROVIDERS_DIR = path.join(import.meta.dirname, "..", "..", "..", "providers");
|
||||
const MODEL_METADATA_DIR = path.join(import.meta.dirname, "..", "..", "..", "models");
|
||||
const MODELS_DIR = path.join(PROVIDERS_DIR, "databricks", "models");
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// API schemas
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const FoundationModel = z
|
||||
.object({
|
||||
ai_gateway_v2_supported: z.boolean().optional(),
|
||||
api_types: z.array(z.string()).optional(),
|
||||
})
|
||||
.passthrough();
|
||||
|
||||
const ServedEntity = z
|
||||
.object({
|
||||
foundation_model: FoundationModel.optional(),
|
||||
})
|
||||
.passthrough();
|
||||
|
||||
const Endpoint = z
|
||||
.object({
|
||||
name: z.string(),
|
||||
config: z
|
||||
.object({
|
||||
served_entities: z.array(ServedEntity).optional(),
|
||||
})
|
||||
.passthrough()
|
||||
.optional(),
|
||||
})
|
||||
.passthrough();
|
||||
|
||||
const FoundationModelsResponse = z
|
||||
.object({
|
||||
endpoints: z.array(Endpoint),
|
||||
})
|
||||
.passthrough();
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Canonical resolution: map a Databricks endpoint name to a models.dev entry
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const PREFIX_TO_PROVIDER: [string, string][] = [
|
||||
["claude-", "anthropic"],
|
||||
["gpt-", "openai"],
|
||||
["gemini-", "google"],
|
||||
["mistral-", "mistral"],
|
||||
["mixtral-", "mistral"],
|
||||
];
|
||||
|
||||
type Resolution =
|
||||
| { type: "base_model"; from: string }
|
||||
| { type: "inline"; content: string }
|
||||
| null;
|
||||
|
||||
async function resolveCanonical(endpointName: string): Promise<Resolution> {
|
||||
const bare = endpointName.replace(/^databricks-/, "");
|
||||
|
||||
// Models in provider subdirectories may not have provider-agnostic metadata
|
||||
// yet, so inline when no model-only entry exists.
|
||||
if (bare.startsWith("gpt-oss-")) {
|
||||
const p = path.join(PROVIDERS_DIR, "openrouter", "models", "openai", `${bare}.toml`);
|
||||
if (existsSync(p)) {
|
||||
return { type: "inline", content: await readFile(p, "utf8") };
|
||||
}
|
||||
}
|
||||
|
||||
// Meta Llama: "meta-llama-3-3-70b-instruct" → "llama-3.3-70b-instruct"
|
||||
if (bare.startsWith("meta-llama-") || bare.startsWith("llama-")) {
|
||||
const llamaId = bare
|
||||
.replace(/^meta-llama-/, "llama-")
|
||||
.replace(/^(llama-\d+)-(\d+)-/, "$1.$2-");
|
||||
const p = path.join(PROVIDERS_DIR, "llama", "models", `${llamaId}.toml`);
|
||||
const metadata = path.join(MODEL_METADATA_DIR, "meta", `${llamaId}.toml`);
|
||||
if (existsSync(p) && existsSync(metadata)) {
|
||||
return { type: "base_model", from: `meta/${llamaId}` };
|
||||
}
|
||||
}
|
||||
|
||||
for (const [prefix, provider] of PREFIX_TO_PROVIDER) {
|
||||
if (!bare.startsWith(prefix)) continue;
|
||||
|
||||
const exact = path.join(PROVIDERS_DIR, provider, "models", `${bare}.toml`);
|
||||
if (existsSync(exact)) return { type: "base_model", from: `${provider}/${bare}` };
|
||||
|
||||
// Try with hyphens-as-dots in version (e.g. gpt-5-4 → gpt-5.4)
|
||||
const dotted = bare.replace(/^((?:[a-z]+-)+\d+)-(\d)/, "$1.$2");
|
||||
if (dotted !== bare) {
|
||||
const dottedExact = path.join(PROVIDERS_DIR, provider, "models", `${dotted}.toml`);
|
||||
if (existsSync(dottedExact)) return { type: "base_model", from: `${provider}/${dotted}` };
|
||||
}
|
||||
|
||||
// Fuzzy: longest filename that shares a prefix with bare or its dotted form
|
||||
const candidates = [bare, ...(dotted !== bare ? [dotted] : [])];
|
||||
const files: string[] = [];
|
||||
try {
|
||||
for await (const f of new Bun.Glob("*.toml").scan({
|
||||
cwd: path.join(PROVIDERS_DIR, provider, "models"),
|
||||
})) {
|
||||
files.push(f);
|
||||
}
|
||||
} catch {
|
||||
// provider directory may not exist
|
||||
}
|
||||
const match = files
|
||||
.map((f) => f.replace(/\.toml$/, ""))
|
||||
.filter((id) => candidates.some((c) => id.startsWith(c) || c.startsWith(id)))
|
||||
.sort((a, b) => b.length - a.length)[0];
|
||||
if (match) return { type: "base_model", from: `${provider}/${match}` };
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
function formatToml(resolution: Resolution, endpointName: string): string {
|
||||
if (resolution?.type === "base_model") {
|
||||
return `base_model = "${resolution.from}"\n`;
|
||||
}
|
||||
if (resolution?.type === "inline") {
|
||||
return resolution.content;
|
||||
}
|
||||
return `# TODO: fill in details for ${endpointName}\nname = "${endpointName}"\n`;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Main
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const IGNORE_PREFIXES = [
|
||||
"databricks-llama-",
|
||||
"databricks-meta-llama-",
|
||||
"databricks-qwen",
|
||||
"databricks-gemma-",
|
||||
];
|
||||
|
||||
async function main() {
|
||||
console.log(
|
||||
`${dryRun ? "[DRY RUN] " : ""}${newOnly ? "[NEW ONLY] " : ""}Fetching Databricks foundation-models...`,
|
||||
);
|
||||
|
||||
const url = `https://${workspace}/api/2.0/serving-endpoints:foundation-models`;
|
||||
const res = await fetch(url, {
|
||||
headers: { Authorization: `Bearer ${token}` },
|
||||
});
|
||||
if (!res.ok) {
|
||||
console.error(`Failed to fetch API: ${res.status} ${res.statusText}`);
|
||||
console.error(await res.text().catch(() => ""));
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const json = await res.json();
|
||||
const parsed = FoundationModelsResponse.safeParse(json);
|
||||
if (!parsed.success) {
|
||||
console.error("Invalid API response:", parsed.error.errors);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const endpoints = parsed.data.endpoints.filter(
|
||||
(e) =>
|
||||
!IGNORE_PREFIXES.some((p) => e.name.startsWith(p)) &&
|
||||
e.config?.served_entities?.some(
|
||||
(se) =>
|
||||
se.foundation_model?.ai_gateway_v2_supported === true &&
|
||||
se.foundation_model?.api_types?.includes("mlflow/v1/chat/completions"),
|
||||
),
|
||||
);
|
||||
|
||||
const existingFiles = new Set<string>();
|
||||
try {
|
||||
for await (const f of new Bun.Glob("*.toml").scan({ cwd: MODELS_DIR })) {
|
||||
existingFiles.add(f);
|
||||
}
|
||||
} catch {
|
||||
// directory may not exist yet
|
||||
}
|
||||
|
||||
console.log(
|
||||
`Found ${endpoints.length} models in API, ${existingFiles.size} existing files\n`,
|
||||
);
|
||||
|
||||
const apiModelIds = new Set<string>();
|
||||
let created = 0;
|
||||
let updated = 0;
|
||||
let unchanged = 0;
|
||||
|
||||
for (const ep of endpoints) {
|
||||
const filename = `${ep.name}.toml`;
|
||||
apiModelIds.add(filename);
|
||||
const filePath = path.join(MODELS_DIR, filename);
|
||||
|
||||
const resolution = await resolveCanonical(ep.name);
|
||||
const newContent = formatToml(resolution, ep.name);
|
||||
const tag = resolution?.type === "base_model" ? `base_model ${resolution.from}` : resolution?.type ?? "stub";
|
||||
|
||||
const existed = existsSync(filePath);
|
||||
if (!existed) {
|
||||
created++;
|
||||
if (dryRun) {
|
||||
console.log(`[DRY RUN] Would create: ${filename} → ${tag}`);
|
||||
} else {
|
||||
await mkdir(MODELS_DIR, { recursive: true });
|
||||
await Bun.write(filePath, newContent);
|
||||
console.log(`Created: ${filename} → ${tag}`);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (newOnly) {
|
||||
unchanged++;
|
||||
continue;
|
||||
}
|
||||
|
||||
const existingContent = await readFile(filePath, "utf8");
|
||||
if (existingContent === newContent) {
|
||||
unchanged++;
|
||||
continue;
|
||||
}
|
||||
|
||||
updated++;
|
||||
if (dryRun) {
|
||||
console.log(`[DRY RUN] Would update: ${filename} → ${tag}`);
|
||||
} else {
|
||||
await Bun.write(filePath, newContent);
|
||||
console.log(`Updated: ${filename} → ${tag}`);
|
||||
}
|
||||
}
|
||||
|
||||
const orphaned: string[] = [];
|
||||
for (const file of existingFiles) {
|
||||
if (!apiModelIds.has(file)) {
|
||||
orphaned.push(file);
|
||||
console.log(`Warning: Orphaned file (not in API): ${file}`);
|
||||
}
|
||||
}
|
||||
|
||||
console.log("");
|
||||
if (dryRun) {
|
||||
console.log(
|
||||
`Summary: ${created} would be created, ${updated} would be updated, ${unchanged} unchanged, ${orphaned.length} orphaned`,
|
||||
);
|
||||
} else {
|
||||
console.log(
|
||||
`Summary: ${created} created, ${updated} updated, ${unchanged} unchanged, ${orphaned.length} orphaned`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
await main();
|
||||
Reference in New Issue
Block a user