93 lines
3.5 KiB
TypeScript
93 lines
3.5 KiB
TypeScript
import { MODELS } from "./models.generated.ts";
|
|
import type { Api, KnownProvider, Model, ModelThinkingLevel, Usage } from "./types.ts";
|
|
|
|
const modelRegistry: Map<string, Map<string, Model<Api>>> = new Map();
|
|
|
|
// Initialize registry from MODELS on module load
|
|
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 calculateCost<TApi extends Api>(model: Model<TApi>, usage: Usage): 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 = (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"];
|
|
|
|
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") 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";
|
|
}
|
|
|
|
/**
|
|
* Check if two models are equal by comparing both their id and provider.
|
|
* Returns false if either model is null or undefined.
|
|
*/
|
|
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;
|
|
}
|