mirror of
https://github.com/fosrl/pangolin.git
synced 2026-08-12 15:30:53 +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();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user