diff --git a/server/lib/aiModelCatalog.ts b/server/lib/aiModelCatalog.ts index 39b3a70c9..faee7a345 100644 --- a/server/lib/aiModelCatalog.ts +++ b/server/lib/aiModelCatalog.ts @@ -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(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 { - 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(); + private byProviderAndKey = new Map(); + private byKey = new Map(); + private refreshTimer: NodeJS.Timeout | null = null; + + /** + * Loads the catalog into memory and schedules periodic background refreshes. + * Call once at server startup. + */ + async init(): Promise { + 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(); + const byProviderAndKey = new Map(); + const byKey = new Map(); + + 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 { + 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 { + 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 { + 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 { - 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 { - 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 { - await refreshCatalog(); - scheduleNextRefresh(); -} - -export function getAiModelCatalog(): AiModelCatalogEntry[] { - return catalog; + await aiModelCatalog.init(); } diff --git a/server/lib/aiModelPricing.ts b/server/lib/aiModelPricing.ts index 873e4213a..4cc398068 100644 --- a/server/lib/aiModelPricing.ts +++ b/server/lib/aiModelPricing.ts @@ -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 = new Map(); - -function getIndexedCatalog(): Map { - const catalog = getAiModelCatalog(); - if (catalog === indexedCatalog) { - return indexedByName; - } - - const byName = new Map(); - 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, 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); } diff --git a/server/routers/site/socketIntegration.ts b/server/routers/site/socketIntegration.ts index 85d5bc9dd..bf63eb686 100644 --- a/server/routers/site/socketIntegration.ts +++ b/server/routers/site/socketIntegration.ts @@ -201,7 +201,7 @@ async function checkSocket( ): Promise<{ siteId: number; newtId: string }> { const { newt } = await getSiteAndNewt(siteId); - logger.info( + logger.debug( `Checking Docker socket for site ${siteId} with Newt ${newt.newtId}` );