spb/polyllm

1import "server-only";2import { and, desc, eq, inArray, sql } from "drizzle-orm";3import { getDb, models, modelSyncRuns, userModelFavorites, userModelRecents, type ModelRow } from "@/db";4import { ids } from "@/lib/ids";5import { log } from "@/lib/log";6import { PROVIDER_IDS, type PolyModel, type ProviderId } from "@/lib/ai/core/types";7import { getAdapter } from "@/lib/ai/providers";8import { getCatalog } from "./catalog";9import { ownerKey } from "@/lib/ai/providers/env-keys";1011export function rowToModel(r: ModelRow): PolyModel {12  return {13    key: r.key,14    id: r.modelId,15    provider: r.provider,16    displayName: r.displayName,17    family: r.family ?? undefined,18    capabilities: r.capabilities,19    limits: r.limits,20    parameters: r.parameters,21    status: r.status,22    pricing: r.pricing ?? null,23    metadata: { ...r.metadata, sortWeight: r.sortWeight, source: r.source, firstSeenAt: r.firstSeenAt.toISOString(), lastSeenAt: r.lastSeenAt.toISOString() },24  };25}2627/** Merge discovered models (from the provider API) with the documented catalog (pricing, limits, quirks). */28export function mergeWithCatalog(discovered: PolyModel[], provider: ProviderId): PolyModel[] {29  const catalog = getCatalog(provider);30  const byId = new Map(discovered.map((m) => [m.id, m]));31  const merged: PolyModel[] = [];32  for (const m of discovered) {33    const c = catalog.get(m.id) ?? catalog.get(m.id.replace(/-\d{4}-\d{2}-\d{2}$|-\d{8}$/, ""));34    if (!c) {35      merged.push({ ...m, metadata: { ...m.metadata, source: "discovery" } });36      continue;37    }38    merged.push({39      ...m,40      displayName: c.displayName ?? m.displayName,41      family: m.family ?? c.family,42      capabilities: { ...c.capabilities, ...m.capabilities, ...(c.capabilityOverrides ?? {}) },43      limits: { ...(c.limits ?? {}), ...(m.limits ?? {}) },44      parameters: { ...(c.parameters ?? {}), ...m.parameters, ...(c.parameterOverrides ?? {}) },45      pricing: m.pricing ?? c.pricing ?? null,46      status: c.status ?? m.status,47      metadata: { ...(c.metadata ?? {}), ...m.metadata, source: "merged", sortWeight: c.sortWeight ?? (m.metadata?.sortWeight as number | undefined) ?? 0 },48    });49  }50  // Catalog-only entries (documented but not returned by this key) are not added: the key can't use them.51  void byId;52  return merged;53}5455export interface SyncResult {56  provider: ProviderId;57  ok: boolean;58  found: number;59  added: number;60  removed: number;61  latencyMs: number;62  error?: string;63}6465/**66 * Sync one provider's model list into the registry. Models not seen anymore are kept but67 * marked `hidden` (so historical usage keeps its metadata) unless they reappear.68 */69export async function syncProvider(provider: ProviderId, apiKey: string, triggeredBy = "schedule"): Promise<SyncResult> {70  const db = getDb();71  const runId = ids.sync();72  const t0 = Date.now();73  await db.insert(modelSyncRuns).values({ id: runId, provider, triggeredBy });74  try {75    const discovered = await getAdapter(provider).listModels(apiKey);76    const merged = mergeWithCatalog(discovered, provider);77    const existing = await db.select({ key: models.key, hidden: models.hidden }).from(models).where(eq(models.provider, provider));78    const existingKeys = new Set(existing.map((e) => e.key));79    const seen = new Set<string>();80    let added = 0;81    const now = new Date();82    for (const m of merged) {83      seen.add(m.key);84      if (!existingKeys.has(m.key)) added++;85      const { sortWeight, source, ...meta } = (m.metadata ?? {}) as Record<string, unknown>;86      await db87        .insert(models)88        .values({89          key: m.key,90          provider,91          modelId: m.id,92          displayName: m.displayName,93          family: m.family ?? null,94          capabilities: m.capabilities,95          limits: m.limits ?? {},96          parameters: m.parameters,97          pricing: m.pricing ?? null,98          status: m.status,99          source: (source as string) ?? "discovery",100          hidden: false,101          sortWeight: Number(sortWeight ?? 0),102          metadata: meta,103          firstSeenAt: now,104          lastSeenAt: now,105          updatedAt: now,106        })107        .onConflictDoUpdate({108          target: models.key,109          set: {110            displayName: m.displayName,111            family: m.family ?? null,112            capabilities: m.capabilities,113            limits: m.limits ?? {},114            parameters: m.parameters,115            pricing: m.pricing ?? null,116            status: m.status,117            source: (source as string) ?? "discovery",118            hidden: false,119            sortWeight: Number(sortWeight ?? 0),120            metadata: meta,121            lastSeenAt: now,122            updatedAt: now,123          },124        });125    }126    const gone = existing.filter((e) => !seen.has(e.key) && !e.hidden).map((e) => e.key);127    if (gone.length) await db.update(models).set({ hidden: true, updatedAt: now }).where(inArray(models.key, gone));128    const latencyMs = Date.now() - t0;129    await db.update(modelSyncRuns).set({ finishedAt: new Date(), ok: true, modelsFound: merged.length, modelsAdded: added, modelsRemoved: gone.length, latencyMs }).where(eq(modelSyncRuns.id, runId));130    log.info("model sync ok", { provider, found: merged.length, added, removed: gone.length, latencyMs, triggeredBy });131    return { provider, ok: true, found: merged.length, added, removed: gone.length, latencyMs };132  } catch (e) {133    const err = e as { code?: string; message?: string };134    const latencyMs = Date.now() - t0;135    await db.update(modelSyncRuns).set({ finishedAt: new Date(), ok: false, latencyMs, errorCode: err.code ?? "UNKNOWN", errorMessage: (err.message ?? String(e)).slice(0, 500) }).where(eq(modelSyncRuns.id, runId));136    log.warn("model sync failed", { provider, code: err.code, error: err.message, triggeredBy });137    return { provider, ok: false, found: 0, added: 0, removed: 0, latencyMs, error: err.message };138  }139}140141/** Sync every provider for which an owner/dev key exists in the environment. */142export async function syncAllWithEnvKeys(triggeredBy = "schedule"): Promise<SyncResult[]> {143  const results: SyncResult[] = [];144  await Promise.all(145    PROVIDER_IDS.map(async (p) => {146      const k = ownerKey(p);147      if (!k) return;148      results.push(await syncProvider(p, k, triggeredBy));149    }),150  );151  return results;152}153154export async function listRegistryModels(opts: { provider?: ProviderId; includeHidden?: boolean; includeDeprecated?: boolean } = {}): Promise<PolyModel[]> {155  const db = getDb();156  const conds = [];157  if (opts.provider) conds.push(eq(models.provider, opts.provider));158  if (!opts.includeHidden) conds.push(eq(models.hidden, false));159  const rows = await db160    .select()161    .from(models)162    .where(conds.length ? and(...conds) : undefined)163    .orderBy(desc(models.sortWeight), models.displayName);164  return rows.filter((r) => opts.includeDeprecated || r.status !== "deprecated").map(rowToModel);165}166167export async function getModel(key: string): Promise<PolyModel | null> {168  const [row] = await getDb().select().from(models).where(eq(models.key, key)).limit(1);169  return row ? rowToModel(row) : null;170}171172export async function registryStats() {173  const db = getDb();174  const counts = await db175    .select({ provider: models.provider, total: sql<number>`count(*)::int`, hidden: sql<number>`sum(case when ${models.hidden} then 1 else 0 end)::int` })176    .from(models)177    .groupBy(models.provider);178  const lastRuns = await db.select().from(modelSyncRuns).orderBy(desc(modelSyncRuns.startedAt)).limit(40);179  return { counts, lastRuns };180}181182export async function favoriteModels(userId: string): Promise<string[]> {183  const rows = await getDb().select({ key: userModelFavorites.modelKey }).from(userModelFavorites).where(eq(userModelFavorites.userId, userId));184  return rows.map((r) => r.key);185}186187export async function toggleFavorite(userId: string, key: string): Promise<boolean> {188  const db = getDb();189  const existing = await db190    .select()191    .from(userModelFavorites)192    .where(and(eq(userModelFavorites.userId, userId), eq(userModelFavorites.modelKey, key)))193    .limit(1);194  if (existing.length) {195    await db.delete(userModelFavorites).where(and(eq(userModelFavorites.userId, userId), eq(userModelFavorites.modelKey, key)));196    return false;197  }198  await db.insert(userModelFavorites).values({ userId, modelKey: key });199  return true;200}201202export async function recentModels(userId: string, limit = 8): Promise<string[]> {203  const rows = await getDb().select({ key: userModelRecents.modelKey }).from(userModelRecents).where(eq(userModelRecents.userId, userId)).orderBy(desc(userModelRecents.usedAt)).limit(limit);204  return rows.map((r) => r.key);205}206207export async function touchRecent(userId: string, key: string) {208  await getDb()209    .insert(userModelRecents)210    .values({ userId, modelKey: key })211    .onConflictDoUpdate({ target: [userModelRecents.userId, userModelRecents.modelKey], set: { usedAt: new Date(), uses: sql`${userModelRecents.uses} + 1` } });212}213