mirror of
https://github.com/fosrl/pangolin.git
synced 2026-08-11 23:18:24 +02:00
use class with helper methods for model catelog
This commit is contained in:
miloschwartz
+216
-79
@@ -3,13 +3,18 @@ import axios from "axios";
|
||||
import config from "@server/lib/config";
|
||||
import logger from "@server/logger";
|
||||
|
||||
export type CatalogProvider =
|
||||
| "openai"
|
||||
| "anthropic"
|
||||
| "gemini"
|
||||
| "vertex"
|
||||
| "azure"
|
||||
| "bedrock";
|
||||
export const CATALOG_PROVIDERS = [
|
||||
"openai",
|
||||
"anthropic",
|
||||
"gemini",
|
||||
"vertex",
|
||||
"azure",
|
||||
"bedrock"
|
||||
] as const;
|
||||
|
||||
export type CatalogProvider = (typeof CATALOG_PROVIDERS)[number];
|
||||
|
||||
const CATALOG_PROVIDER_SET = new Set<string>(CATALOG_PROVIDERS);
|
||||
|
||||
export type AiModelCatalogEntry = {
|
||||
provider: CatalogProvider;
|
||||
@@ -22,90 +27,222 @@ export type AiModelCatalogEntry = {
|
||||
};
|
||||
};
|
||||
|
||||
let catalog: AiModelCatalogEntry[] = [];
|
||||
let refreshTimer: NodeJS.Timeout | null = null;
|
||||
type RawCatalogEntry = {
|
||||
id?: string;
|
||||
name?: string;
|
||||
model?: string;
|
||||
provider: string;
|
||||
input_cost_per_token?: number | null;
|
||||
output_cost_per_token?: number | null;
|
||||
cache_read_input_token_cost?: number | null;
|
||||
output_cost_per_reasoning_token?: number | null;
|
||||
pricing?: {
|
||||
input?: number | null;
|
||||
output?: number | null;
|
||||
cacheRead?: number | null;
|
||||
reasoningOutput?: number | null;
|
||||
};
|
||||
};
|
||||
|
||||
async function fetchFromFile(filePath: string): Promise<AiModelCatalogEntry[] | null> {
|
||||
try {
|
||||
if (!fs.existsSync(filePath)) {
|
||||
function normalizeCatalogProvider(raw: string): CatalogProvider | null {
|
||||
if (CATALOG_PROVIDER_SET.has(raw)) {
|
||||
return raw as CatalogProvider;
|
||||
}
|
||||
if (raw.startsWith("bedrock")) {
|
||||
return "bedrock";
|
||||
}
|
||||
if (raw.startsWith("vertex")) {
|
||||
return "vertex";
|
||||
}
|
||||
if (raw.startsWith("azure")) {
|
||||
return "azure";
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
function normalizeEntry(raw: RawCatalogEntry): AiModelCatalogEntry | null {
|
||||
const provider = normalizeCatalogProvider(raw.provider);
|
||||
if (!provider) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const model = raw.model ?? raw.name ?? raw.id;
|
||||
if (!model) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
provider,
|
||||
model,
|
||||
pricing: {
|
||||
input: raw.pricing?.input ?? raw.input_cost_per_token ?? null,
|
||||
output: raw.pricing?.output ?? raw.output_cost_per_token ?? null,
|
||||
cacheRead:
|
||||
raw.pricing?.cacheRead ??
|
||||
raw.cache_read_input_token_cost ??
|
||||
null,
|
||||
reasoningOutput:
|
||||
raw.pricing?.reasoningOutput ??
|
||||
raw.output_cost_per_reasoning_token ??
|
||||
null
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
function providerKey(provider: CatalogProvider, key: string): string {
|
||||
return `${provider}\0${key}`;
|
||||
}
|
||||
|
||||
export class AiModelCatalog {
|
||||
private entries: AiModelCatalogEntry[] = [];
|
||||
private byProvider = new Map<CatalogProvider, AiModelCatalogEntry[]>();
|
||||
private byProviderAndKey = new Map<string, AiModelCatalogEntry>();
|
||||
private byKey = new Map<string, AiModelCatalogEntry[]>();
|
||||
private refreshTimer: NodeJS.Timeout | null = null;
|
||||
|
||||
/**
|
||||
* Loads the catalog into memory and schedules periodic background refreshes.
|
||||
* Call once at server startup.
|
||||
*/
|
||||
async init(): Promise<void> {
|
||||
await this.refresh();
|
||||
this.scheduleNextRefresh();
|
||||
}
|
||||
|
||||
/** Exact lookup by catalog provider and model key. */
|
||||
get(
|
||||
provider: CatalogProvider,
|
||||
key: string
|
||||
): AiModelCatalogEntry | undefined {
|
||||
return this.byProviderAndKey.get(providerKey(provider, key));
|
||||
}
|
||||
|
||||
/** All models for a catalog provider. */
|
||||
list(provider: CatalogProvider): AiModelCatalogEntry[] {
|
||||
return this.byProvider.get(provider) ?? [];
|
||||
}
|
||||
|
||||
/** All catalog entries that share a model key, across providers. */
|
||||
listByKey(key: string): AiModelCatalogEntry[] {
|
||||
return this.byKey.get(key) ?? [];
|
||||
}
|
||||
|
||||
/** Full in-memory catalog. */
|
||||
getAll(): AiModelCatalogEntry[] {
|
||||
return this.entries;
|
||||
}
|
||||
|
||||
private setEntries(entries: AiModelCatalogEntry[]): void {
|
||||
const byProvider = new Map<CatalogProvider, AiModelCatalogEntry[]>();
|
||||
const byProviderAndKey = new Map<string, AiModelCatalogEntry>();
|
||||
const byKey = new Map<string, AiModelCatalogEntry[]>();
|
||||
|
||||
for (const entry of entries) {
|
||||
const list = byProvider.get(entry.provider) ?? [];
|
||||
list.push(entry);
|
||||
byProvider.set(entry.provider, list);
|
||||
|
||||
const mapKey = providerKey(entry.provider, entry.model);
|
||||
if (!byProviderAndKey.has(mapKey)) {
|
||||
byProviderAndKey.set(mapKey, entry);
|
||||
}
|
||||
|
||||
const keyList = byKey.get(entry.model) ?? [];
|
||||
keyList.push(entry);
|
||||
byKey.set(entry.model, keyList);
|
||||
}
|
||||
|
||||
this.entries = entries;
|
||||
this.byProvider = byProvider;
|
||||
this.byProviderAndKey = byProviderAndKey;
|
||||
this.byKey = byKey;
|
||||
}
|
||||
|
||||
private async fetchFromFile(
|
||||
filePath: string
|
||||
): Promise<AiModelCatalogEntry[] | null> {
|
||||
try {
|
||||
if (!fs.existsSync(filePath)) {
|
||||
logger.warn(
|
||||
`AI model catalog file not found at ${filePath}; cost calculation will fall back to unknown pricing`
|
||||
);
|
||||
return null;
|
||||
}
|
||||
const raw = fs.readFileSync(filePath, "utf-8");
|
||||
const parsed = JSON.parse(raw) as { data: RawCatalogEntry[] };
|
||||
return (parsed.data ?? [])
|
||||
.map(normalizeEntry)
|
||||
.filter((e): e is AiModelCatalogEntry => e != null);
|
||||
} catch (error) {
|
||||
logger.warn("Failed to read AI model catalog file", { error });
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private async fetchFromUpstream(
|
||||
upstreamUrl: string
|
||||
): Promise<AiModelCatalogEntry[] | null> {
|
||||
try {
|
||||
const res = await axios.get<{ data: RawCatalogEntry[] }>(
|
||||
upstreamUrl,
|
||||
{ timeout: 15_000 }
|
||||
);
|
||||
return (res.data?.data ?? [])
|
||||
.map(normalizeEntry)
|
||||
.filter((e): e is AiModelCatalogEntry => e != null);
|
||||
} catch (error: any) {
|
||||
logger.warn(
|
||||
`AI model catalog file not found at ${filePath}; cost calculation will fall back to unknown pricing`
|
||||
`Failed to fetch AI model catalog from ${upstreamUrl}: ${error.message || error}`
|
||||
);
|
||||
return null;
|
||||
}
|
||||
const raw = fs.readFileSync(filePath, "utf-8");
|
||||
const parsed = JSON.parse(raw) as { data: AiModelCatalogEntry[] };
|
||||
return parsed.data ?? [];
|
||||
} catch (error) {
|
||||
logger.warn("Failed to read AI model catalog file", { error });
|
||||
return null;
|
||||
}
|
||||
|
||||
private async refresh(): Promise<void> {
|
||||
const { file, upstream_url } = config.getRawConfig().ai.model_catalog;
|
||||
|
||||
const fetched = file
|
||||
? await this.fetchFromFile(file)
|
||||
: await this.fetchFromUpstream(upstream_url);
|
||||
|
||||
if (fetched) {
|
||||
this.setEntries(fetched);
|
||||
logger.debug(
|
||||
`AI model catalog refreshed: ${this.entries.length} models loaded`
|
||||
);
|
||||
} else {
|
||||
logger.debug(
|
||||
"AI model catalog refresh failed; keeping previously loaded catalog in memory"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private scheduleNextRefresh(): void {
|
||||
const { refresh_interval_min_hours, refresh_interval_max_hours } =
|
||||
config.getRawConfig().ai.model_catalog;
|
||||
|
||||
// Jittered rather than fixed so that many self-hosted instances don't
|
||||
// all hit the upstream catalog endpoint at the same moment.
|
||||
const minMs = refresh_interval_min_hours * 60 * 60 * 1000;
|
||||
const maxMs = refresh_interval_max_hours * 60 * 60 * 1000;
|
||||
const delayMs = minMs + Math.random() * Math.max(0, maxMs - minMs);
|
||||
|
||||
if (this.refreshTimer) {
|
||||
clearTimeout(this.refreshTimer);
|
||||
}
|
||||
this.refreshTimer = setTimeout(async () => {
|
||||
await this.refresh();
|
||||
this.scheduleNextRefresh();
|
||||
}, delayMs);
|
||||
}
|
||||
}
|
||||
|
||||
async function fetchFromUpstream(
|
||||
upstreamUrl: string
|
||||
): Promise<AiModelCatalogEntry[] | null> {
|
||||
try {
|
||||
const res = await axios.get<{ data: AiModelCatalogEntry[] }>(
|
||||
upstreamUrl,
|
||||
{ timeout: 15_000 }
|
||||
);
|
||||
return res.data?.data ?? [];
|
||||
} catch (error: any) {
|
||||
logger.warn(
|
||||
`Failed to fetch AI model catalog from ${upstreamUrl}: ${error.message || error}`
|
||||
);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
async function refreshCatalog(): Promise<void> {
|
||||
const { file, upstream_url } = config.getRawConfig().ai.model_catalog;
|
||||
|
||||
const fetched = file
|
||||
? await fetchFromFile(file)
|
||||
: await fetchFromUpstream(upstream_url);
|
||||
|
||||
if (fetched) {
|
||||
catalog = fetched;
|
||||
logger.debug(
|
||||
`AI model catalog refreshed: ${catalog.length} models loaded`
|
||||
);
|
||||
} else {
|
||||
logger.debug(
|
||||
"AI model catalog refresh failed; keeping previously loaded catalog in memory"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
function scheduleNextRefresh(): void {
|
||||
const { refresh_interval_min_hours, refresh_interval_max_hours } =
|
||||
config.getRawConfig().ai.model_catalog;
|
||||
|
||||
// Jittered rather than fixed so that many self-hosted instances don't
|
||||
// all hit the upstream catalog endpoint at the same moment.
|
||||
const minMs = refresh_interval_min_hours * 60 * 60 * 1000;
|
||||
const maxMs = refresh_interval_max_hours * 60 * 60 * 1000;
|
||||
const delayMs = minMs + Math.random() * Math.max(0, maxMs - minMs);
|
||||
|
||||
if (refreshTimer) {
|
||||
clearTimeout(refreshTimer);
|
||||
}
|
||||
refreshTimer = setTimeout(async () => {
|
||||
await refreshCatalog();
|
||||
scheduleNextRefresh();
|
||||
}, delayMs);
|
||||
}
|
||||
export const aiModelCatalog = new AiModelCatalog();
|
||||
|
||||
/**
|
||||
* Loads the AI model pricing catalog into memory and schedules periodic
|
||||
* background refreshes. Call once at server startup.
|
||||
*/
|
||||
export async function initAiModelCatalog(): Promise<void> {
|
||||
await refreshCatalog();
|
||||
scheduleNextRefresh();
|
||||
}
|
||||
|
||||
export function getAiModelCatalog(): AiModelCatalogEntry[] {
|
||||
return catalog;
|
||||
await aiModelCatalog.init();
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import type { AiProviderType } from "@server/lib/aiProviderDefaults";
|
||||
import type { AiUsage } from "@server/lib/aiUsageExtraction";
|
||||
import {
|
||||
getAiModelCatalog,
|
||||
aiModelCatalog,
|
||||
type AiModelCatalogEntry,
|
||||
type CatalogProvider
|
||||
} from "@server/lib/aiModelCatalog";
|
||||
@@ -36,31 +36,6 @@ const PROVIDER_CATALOG_MAP: Record<
|
||||
vercelAiGateway: null
|
||||
};
|
||||
|
||||
// Indexed view over the in-memory catalog, rebuilt only when
|
||||
// getAiModelCatalog() returns a different array instance (i.e. after a
|
||||
// background refresh swaps it out), not on every lookup.
|
||||
let indexedCatalog: AiModelCatalogEntry[] | null = null;
|
||||
let indexedByName: Map<string, AiModelCatalogEntry[]> = new Map();
|
||||
|
||||
function getIndexedCatalog(): Map<string, AiModelCatalogEntry[]> {
|
||||
const catalog = getAiModelCatalog();
|
||||
if (catalog === indexedCatalog) {
|
||||
return indexedByName;
|
||||
}
|
||||
|
||||
const byName = new Map<string, AiModelCatalogEntry[]>();
|
||||
for (const entry of catalog) {
|
||||
if (!entry.model) continue;
|
||||
const list = byName.get(entry.model) ?? [];
|
||||
list.push(entry);
|
||||
byName.set(entry.model, list);
|
||||
}
|
||||
|
||||
indexedCatalog = catalog;
|
||||
indexedByName = byName;
|
||||
return byName;
|
||||
}
|
||||
|
||||
function stripVendorPrefix(modelId: string): string | null {
|
||||
const idx = modelId.indexOf("/");
|
||||
if (idx === -1 || idx === modelId.length - 1) {
|
||||
@@ -83,7 +58,6 @@ function toPricing(
|
||||
}
|
||||
|
||||
function findEntry(
|
||||
byName: Map<string, AiModelCatalogEntry[]>,
|
||||
modelId: string,
|
||||
provider: CatalogProvider | null
|
||||
): AiModelCatalogEntry | null {
|
||||
@@ -92,11 +66,15 @@ function findEntry(
|
||||
);
|
||||
|
||||
for (const key of candidates) {
|
||||
const entries = byName.get(key);
|
||||
if (!entries) continue;
|
||||
const match = provider
|
||||
? entries.find((e) => e.provider === provider)
|
||||
: entries[0];
|
||||
if (provider) {
|
||||
const match = aiModelCatalog.get(provider, key);
|
||||
if (match) {
|
||||
return match;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
const match = aiModelCatalog.listByKey(key)[0];
|
||||
if (match) {
|
||||
return match;
|
||||
}
|
||||
@@ -118,18 +96,17 @@ export function getModelPricing(
|
||||
return null;
|
||||
}
|
||||
|
||||
const byName = getIndexedCatalog();
|
||||
const catalogProvider =
|
||||
providerType === "custom" ? null : PROVIDER_CATALOG_MAP[providerType];
|
||||
|
||||
if (catalogProvider) {
|
||||
const scoped = findEntry(byName, modelId, catalogProvider);
|
||||
const scoped = findEntry(modelId, catalogProvider);
|
||||
if (scoped) {
|
||||
return toPricing(scoped, false);
|
||||
}
|
||||
}
|
||||
|
||||
const fallback = findEntry(byName, modelId, null);
|
||||
const fallback = findEntry(modelId, null);
|
||||
if (fallback) {
|
||||
return toPricing(fallback, true);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user