|
| 1 | +// forge gateway model map — remap complexity tiers onto a CUSTOM gateway's real model IDs. |
| 2 | +// |
| 3 | +// The problem: model_tiers.json pins public Anthropic IDs (claude-haiku-4-5-20251001, …). |
| 4 | +// A self-hosted LiteLLM/proxy gateway rarely exposes those exact names — it advertises its |
| 5 | +// OWN ids (e.g. "bedrock-claude-haiku", "prod-sonnet", "claude-3-5-sonnet-v2"). Sending a |
| 6 | +// stock id straight to such a gateway 404s. So we ask the gateway what it actually serves |
| 7 | +// (GET /v1/models, once per process) and SCORE each advertised id against every tier's family |
| 8 | +// — the same DATA-is-a-table / DECISION-is-a-formula rule the rest of forge follows: the tier |
| 9 | +// families are data, the pick is a graded overlap score (src/math.js setOverlap), inspectable |
| 10 | +// and testable. |
| 11 | +// |
| 12 | +// Contract (zero breaking change): |
| 13 | +// - Only engages for a NON-default gateway base URL. Direct api.anthropic.com → no-op, no net. |
| 14 | +// - FAIL-SAFE. No gateway, unreachable, unparseable, or no family match → returns the stock |
| 15 | +// id unchanged. Callers are byte-identical to before when there is nothing to remap. |
| 16 | +// - The MODELS export shape is untouched; nothing here mutates model_tiers. |
| 17 | +import { spawnSync } from "node:child_process"; |
| 18 | +import { setOverlap } from "./math.js"; |
| 19 | +import { MODELS, TIER_ORDER } from "./model_tiers.js"; |
| 20 | + |
| 21 | +const ANTHROPIC_DEFAULT = "https://api.anthropic.com"; |
| 22 | + |
| 23 | +// GET {base}/v1/models in a spawned node child so this module stays synchronous like every |
| 24 | +// other forge faculty (embed.js / llm.js pattern). The auth key travels via the child's env |
| 25 | +// (_FORGE_LLM_KEY) — never in argv, never logged. Accepts both OpenAI-shaped ({data:[{id}]}) |
| 26 | +// and Anthropic-shaped ({data:[{id}]}) catalogs; both key the list under data[].id. |
| 27 | +const FETCH_CHILD = `let raw="";process.stdin.on("data",(d)=>{raw+=d;});process.stdin.on("end",async()=>{try{const{url,timeoutMs}=JSON.parse(raw);const key=process.env._FORGE_LLM_KEY||"";const headers={"anthropic-version":"2023-06-01"};if(key.startsWith("Bearer ")){headers.authorization=key;}else if(key){headers["x-api-key"]=key;headers.authorization="Bearer "+key;}const ac=new AbortController();const timer=setTimeout(()=>ac.abort(),timeoutMs||5000);let res;try{res=await fetch(url,{headers,signal:ac.signal});}finally{clearTimeout(timer);}if(!res.ok){process.stderr.write("http "+res.status);process.exit(1);}const data=await res.json();const rows=Array.isArray(data)?data:Array.isArray(data&&data.data)?data.data:[];const ids=rows.map((m)=>(typeof m==="string"?m:m&&m.id)).filter((x)=>typeof x==="string"&&x);process.stdout.write(JSON.stringify(ids));}catch(e){process.stderr.write(String((e&&e.message)||e));process.exit(1);}});`; |
| 28 | + |
| 29 | +// Process-lifetime cache: base URL -> string[] (advertised ids) | null (fetched, none usable). |
| 30 | +// "Once per process" is the whole point — the ambient LLM path must not re-probe on every call. |
| 31 | +const _catalogCache = new Map(); |
| 32 | + |
| 33 | +/** Clear the per-process /v1/models cache (tests only). */ |
| 34 | +export function _resetGatewayCache() { |
| 35 | + _catalogCache.clear(); |
| 36 | +} |
| 37 | + |
| 38 | +/** |
| 39 | + * The active gateway base URL to remap against, or null when there is nothing to remap. |
| 40 | + * Mirrors llm.js resolution (LITELLM_BASE_URL wins, then ANTHROPIC_BASE_URL). The default |
| 41 | + * Anthropic endpoint returns null so direct-API users never trigger a probe or a remap. |
| 42 | + * @returns {string|null} |
| 43 | + */ |
| 44 | +export function gatewayBase() { |
| 45 | + const url = (process.env.LITELLM_BASE_URL || process.env.ANTHROPIC_BASE_URL || "").replace( |
| 46 | + /\/+$/, |
| 47 | + "", |
| 48 | + ); |
| 49 | + if (!url) return null; |
| 50 | + if (url.toLowerCase() === ANTHROPIC_DEFAULT) return null; // direct Anthropic — stock ids are correct |
| 51 | + return url; |
| 52 | +} |
| 53 | + |
| 54 | +function apiKey() { |
| 55 | + return ( |
| 56 | + process.env.ANTHROPIC_API_KEY || |
| 57 | + process.env.ANTHROPIC_AUTH_TOKEN || |
| 58 | + process.env.LITELLM_API_KEY || |
| 59 | + "" |
| 60 | + ); |
| 61 | +} |
| 62 | + |
| 63 | +function spawnFetch(base, timeoutMs) { |
| 64 | + const r = spawnSync(process.execPath, ["-e", FETCH_CHILD], { |
| 65 | + input: JSON.stringify({ url: `${base}/v1/models`, timeoutMs }), |
| 66 | + encoding: "utf8", |
| 67 | + timeout: timeoutMs + 1000, |
| 68 | + maxBuffer: 4 * 1024 * 1024, |
| 69 | + env: { ...process.env, _FORGE_LLM_KEY: apiKey() }, |
| 70 | + stdio: ["pipe", "pipe", "pipe"], |
| 71 | + }); |
| 72 | + if (r.error || r.status !== 0 || !r.stdout) return null; |
| 73 | + const ids = JSON.parse(r.stdout); |
| 74 | + return Array.isArray(ids) ? ids : null; |
| 75 | +} |
| 76 | + |
| 77 | +/** |
| 78 | + * Fetch (and cache once per process) the model ids a gateway advertises at /v1/models. |
| 79 | + * @param {string} base gateway base URL (no trailing slash) |
| 80 | + * @param {{timeoutMs?: number, fetchImpl?: (base:string)=>string[]}} [opts] fetchImpl is injectable for tests |
| 81 | + * @returns {string[]|null} advertised ids, or null on any failure |
| 82 | + */ |
| 83 | +export function fetchModelIds(base, { timeoutMs = 5000, fetchImpl } = {}) { |
| 84 | + if (!base) return null; |
| 85 | + if (_catalogCache.has(base)) return _catalogCache.get(base); |
| 86 | + let ids = null; |
| 87 | + try { |
| 88 | + ids = fetchImpl ? fetchImpl(base) : spawnFetch(base, timeoutMs); |
| 89 | + } catch { |
| 90 | + ids = null; |
| 91 | + } |
| 92 | + const clean = Array.isArray(ids) |
| 93 | + ? [...new Set(ids.filter((x) => typeof x === "string" && x))] |
| 94 | + : null; |
| 95 | + const result = clean && clean.length ? clean : null; |
| 96 | + _catalogCache.set(base, result); |
| 97 | + return result; |
| 98 | +} |
| 99 | + |
| 100 | +const tokenize = (s) => |
| 101 | + new Set( |
| 102 | + String(s) |
| 103 | + .toLowerCase() |
| 104 | + .split(/[^a-z0-9]+/) |
| 105 | + .filter(Boolean), |
| 106 | + ); |
| 107 | + |
| 108 | +/** Reference token set for a tier: the family key plus its marketing-name tokens (e.g. haiku → {haiku,4,5}). */ |
| 109 | +export function familyTokens(tier) { |
| 110 | + return tokenize(`${tier} ${MODELS[tier]?.name ?? ""}`); |
| 111 | +} |
| 112 | + |
| 113 | +/** |
| 114 | + * Score how well a gateway model id belongs to a tier family, in [0,1]. |
| 115 | + * The family word itself (haiku/sonnet/opus/fable) is a HARD gate — absent it, the id is not a |
| 116 | + * candidate for that tier (score 0), so an unrelated model can never be mis-assigned. Present it, |
| 117 | + * the score is the overlap coefficient of the tier's reference tokens with the id's tokens, which |
| 118 | + * rewards a version match ("claude-sonnet-5" scores 1.0 for sonnet; "prod-sonnet" scores lower). |
| 119 | + * @param {string} modelId |
| 120 | + * @param {string} tier |
| 121 | + * @returns {number} |
| 122 | + */ |
| 123 | +export function familyScore(modelId, tier) { |
| 124 | + const toks = tokenize(modelId); |
| 125 | + if (!toks.has(tier)) return 0; // family word MUST be present |
| 126 | + return setOverlap(familyTokens(tier), toks); |
| 127 | +} |
| 128 | + |
| 129 | +// Deterministic tie-break among equal-scoring candidates: prefer the id closest to the canonical |
| 130 | +// name (fewest tokens — less vendor/deployment noise), then lexicographic for stability. |
| 131 | +function tieBreak(a, b) { |
| 132 | + const na = tokenize(a).size; |
| 133 | + const nb = tokenize(b).size; |
| 134 | + if (na !== nb) return na - nb; |
| 135 | + return a < b ? -1 : a > b ? 1 : 0; |
| 136 | +} |
| 137 | + |
| 138 | +/** |
| 139 | + * Pure: given a gateway's advertised ids, pick the best id per tier by family score. |
| 140 | + * @param {string[]} ids |
| 141 | + * @returns {Record<string,{id:string, score:number}>} only tiers with a family match appear |
| 142 | + */ |
| 143 | +export function buildGatewayMap(ids = []) { |
| 144 | + const list = [...new Set((ids || []).filter((x) => typeof x === "string" && x))]; |
| 145 | + /** @type {Record<string,{id:string, score:number}>} */ |
| 146 | + const map = {}; |
| 147 | + for (const tier of TIER_ORDER) { |
| 148 | + let best = null; |
| 149 | + for (const id of list) { |
| 150 | + const score = familyScore(id, tier); |
| 151 | + if (score <= 0) continue; |
| 152 | + if (!best || score > best.score || (score === best.score && tieBreak(id, best.id) < 0)) { |
| 153 | + best = { id, score }; |
| 154 | + } |
| 155 | + } |
| 156 | + if (best) map[tier] = best; |
| 157 | + } |
| 158 | + return map; |
| 159 | +} |
| 160 | + |
| 161 | +/** |
| 162 | + * The tier→gateway-model mapping for the active gateway. Fetches /v1/models (cached) and scores. |
| 163 | + * @param {{base?: string, fetchImpl?: (base:string)=>string[], timeoutMs?: number}} [opts] |
| 164 | + * @returns {{active:boolean, base:(string|null), reachable?:boolean, catalog?:string[], models:Record<string,{id:string,score:number}>}} |
| 165 | + */ |
| 166 | +export function gatewayModelMap({ base, fetchImpl, timeoutMs } = {}) { |
| 167 | + const b = base ?? gatewayBase(); |
| 168 | + if (!b) return { active: false, base: null, models: {} }; |
| 169 | + const ids = fetchModelIds(b, { fetchImpl, timeoutMs }); |
| 170 | + if (!ids) return { active: true, base: b, reachable: false, models: {} }; |
| 171 | + return { |
| 172 | + active: true, |
| 173 | + base: b, |
| 174 | + reachable: true, |
| 175 | + catalog: ids, |
| 176 | + models: buildGatewayMap(ids), |
| 177 | + }; |
| 178 | +} |
| 179 | + |
| 180 | +/** |
| 181 | + * Resolve a tier to a gateway model id, or return `fallbackId` unchanged (silent fallback). |
| 182 | + * This is the one function callers reach for: it never throws and never blocks a direct-API user. |
| 183 | + * @param {string} tier |
| 184 | + * @param {string} fallbackId the stock id to use when there is nothing to remap |
| 185 | + * @param {{base?: string, fetchImpl?: (base:string)=>string[], timeoutMs?: number}} [opts] |
| 186 | + * @returns {string} |
| 187 | + */ |
| 188 | +export function gatewayModelId(tier, fallbackId, opts = {}) { |
| 189 | + try { |
| 190 | + const m = gatewayModelMap(opts); |
| 191 | + return m.models?.[tier]?.id ?? fallbackId; |
| 192 | + } catch { |
| 193 | + return fallbackId; |
| 194 | + } |
| 195 | +} |
0 commit comments