Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
4.7 kB · 132 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133import { MODELS } from "./models.generated.js";import type { Api, KnownProvider, Model, ModelThinkingLevel, Usage } from "./types.js";
const modelRegistry: Map<string, Map<string, Model<Api>>> = new Map();
for (const [provider, models] of Object.entries(MODELS)) { const providerModels = new Map<string, Model<Api>>(); for (const [id, model] of Object.entries(models)) { providerModels.set(id, model as Model<Api>); } modelRegistry.set(provider, providerModels);}
type ModelApi< TProvider extends KnownProvider, TModelId extends keyof (typeof MODELS)[TProvider],> = (typeof MODELS)[TProvider][TModelId] extends { api: infer TApi } ? (TApi extends Api ? TApi : never) : never;
export function getModel<TProvider extends KnownProvider, TModelId extends keyof (typeof MODELS)[TProvider]>( provider: TProvider, modelId: TModelId,): Model<ModelApi<TProvider, TModelId>> { const providerModels = modelRegistry.get(provider); return providerModels?.get(modelId as string) as Model<ModelApi<TProvider, TModelId>>;}
export function getProviders(): KnownProvider[] { return Array.from(modelRegistry.keys()) as KnownProvider[];}
export function getModels<TProvider extends KnownProvider>( provider: TProvider,): Model<ModelApi<TProvider, keyof (typeof MODELS)[TProvider]>>[] { const models = modelRegistry.get(provider); return models ? (Array.from(models.values()) as Model<ModelApi<TProvider, keyof (typeof MODELS)[TProvider]>>[]) : [];}
export function supportsFastMode<TApi extends Api>(model: Model<TApi>): boolean { const eligibleId = model.id === "gpt-5.4" || model.id === "gpt-5.5" || model.id === "gpt-5.6" || model.id.startsWith("gpt-5.6-") || model.id === "gpt-6-astra" || model.id.startsWith("gpt-6-"); return ( eligibleId && ((model.provider === "openai-codex" && model.api === "openai-codex-responses") || (model.provider === "openai" && model.api === "openai-responses")) );}
export interface CostOverrides { cacheWrite?: number;}
export function calculateCost<TApi extends Api>( model: Model<TApi>, usage: Usage, overrides?: CostOverrides,): Usage["cost"] { usage.cost.input = (model.cost.input / 1000000) * usage.input; usage.cost.output = (model.cost.output / 1000000) * usage.output; usage.cost.cacheRead = (model.cost.cacheRead / 1000000) * usage.cacheRead; usage.cost.cacheWrite = ((overrides?.cacheWrite ?? model.cost.cacheWrite) / 1000000) * usage.cacheWrite; usage.cost.total = usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite; return usage.cost;}
const EXTENDED_THINKING_LEVELS: ModelThinkingLevel[] = ["off", "minimal", "low", "medium", "high", "xhigh", "max"];
export function getSupportedThinkingLevels<TApi extends Api>(model: Model<TApi>): ModelThinkingLevel[] { if (!model.reasoning) return ["off"];
return EXTENDED_THINKING_LEVELS.filter((level) => { const mapped = model.thinkingLevelMap?.[level]; if (mapped === null) return false; if (level === "xhigh" || level === "max") return mapped !== undefined; return true; });}
export function clampThinkingLevel<TApi extends Api>( model: Model<TApi>, level: ModelThinkingLevel,): ModelThinkingLevel { const availableLevels = getSupportedThinkingLevels(model); if (availableLevels.includes(level)) return level;
const requestedIndex = EXTENDED_THINKING_LEVELS.indexOf(level); if (requestedIndex === -1) return availableLevels[0] ?? "off";
for (let i = requestedIndex; i < EXTENDED_THINKING_LEVELS.length; i++) { const candidate = EXTENDED_THINKING_LEVELS[i]; if (availableLevels.includes(candidate)) return candidate; } for (let i = requestedIndex - 1; i >= 0; i--) { const candidate = EXTENDED_THINKING_LEVELS[i]; if (availableLevels.includes(candidate)) return candidate; } return availableLevels[0] ?? "off";}
export function modelsAreEqual<TApi extends Api>( a: Model<TApi> | null | undefined, b: Model<TApi> | null | undefined,): boolean { if (!a || !b) return false; return a.id === b.id && a.provider === b.provider;}
export function registerLiveModels(models: Model<Api>[]): void { for (const model of models) { let providerModels = modelRegistry.get(model.provider); if (!providerModels) { providerModels = new Map<string, Model<Api>>(); modelRegistry.set(model.provider, providerModels); } providerModels.set(model.id, model); }}
export function resetLiveModels(): void { modelRegistry.clear(); for (const [provider, models] of Object.entries(MODELS)) { const providerModels = new Map<string, Model<Api>>(); for (const [id, model] of Object.entries(models)) { providerModels.set(id, model as Model<Api>); } modelRegistry.set(provider, providerModels); }}