spb/rareindex

1import { AnthropicProvider } from './providers/anthropic.js';2import { OpenAICompatibleProvider } from './providers/openai-compatible.js';3import { installDbCostSink } from './costs.js';4import { AiNotConfiguredError, type CompletionRequest, type Completion, type EmbedRequest, type EmbedResult, type ExtractRequest, type ExtractResult, type ModelProvider, type Role, type StreamDelta, type ToolAgentRequest } from './types.js';56export interface RouterConfig {7  anthropicApiKey?: string;8  openaiApiKey?: string;9  localBaseUrl?: string;10  localApiKey?: string;11  /** role → "provider:model" or "model" overrides, e.g. { research: 'anthropic:claude-opus-5' } */12  overrides?: Partial<Record<Role, string>>;13  persistCosts?: boolean;14}1516export const ROLES: Role[] = ['normalize', 'classify', 'summarize', 'resolve', 'research', 'vision', 'embed'];1718/** Read router config from process.env (server-side only). AI_MODEL_<ROLE> overrides the default model. */19export function configFromEnv(env: NodeJS.ProcessEnv = process.env): RouterConfig {20  const overrides: Partial<Record<Role, string>> = {};21  for (const r of ROLES) {22    const v = env[`AI_MODEL_${r.toUpperCase()}`];23    if (v) overrides[r] = v;24  }25  return {26    anthropicApiKey: env.ANTHROPIC_API_KEY,27    openaiApiKey: env.OPENAI_API_KEY,28    localBaseUrl: env.LOCAL_LLM_BASE_URL,29    localApiKey: env.LOCAL_LLM_API_KEY,30    overrides,31    persistCosts: env.NODE_ENV !== 'test',32  };33}3435/**36 * Picks a provider per role (§168): Anthropic for everything but embeddings when configured,37 * OpenAI(-compatible) otherwise, embeddings only via OpenAI-compatible endpoints. Callers never38 * see vendor SDKs; a missing provider surfaces as AiNotConfiguredError which the UI renders as39 * "AI provider not configured".40 */41export class ModelRouter implements ModelProvider {42  readonly id = 'router';43  readonly providers: ModelProvider[] = [];44  private readonly overrides: Partial<Record<Role, string>>;4546  constructor(cfg: RouterConfig) {47    this.overrides = cfg.overrides ?? {};48    const perProviderModels = (pid: string): Partial<Record<Role, string>> => {49      const out: Partial<Record<Role, string>> = {};50      for (const [role, spec] of Object.entries(this.overrides) as Array<[Role, string]>) {51        const [p, m] = spec.includes(':') ? (spec.split(':', 2) as [string, string]) : [null, spec];52        if (!p || p === pid) out[role] = m;53      }54      return out;55    };56    if (cfg.anthropicApiKey) this.providers.push(new AnthropicProvider({ apiKey: cfg.anthropicApiKey, models: perProviderModels('anthropic') }));57    if (cfg.openaiApiKey) this.providers.push(new OpenAICompatibleProvider({ id: 'openai', apiKey: cfg.openaiApiKey, models: perProviderModels('openai') }));58    if (cfg.localBaseUrl) this.providers.push(new OpenAICompatibleProvider({ id: 'local', apiKey: cfg.localApiKey ?? 'local', baseURL: cfg.localBaseUrl, models: perProviderModels('local') }));59    if (cfg.persistCosts) void installDbCostSink();60  }6162  get configured(): boolean {63    return this.providers.length > 0;64  }6566  /** Provider for a role: explicit "provider:" override first, else first provider supporting the role. */67  providerFor(role: Role): ModelProvider {68    const spec = this.overrides[role];69    if (spec?.includes(':')) {70      const pid = spec.split(':', 1)[0];71      const p = this.providers.find((x) => x.id === pid);72      if (p) return p;73    }74    const p = this.providers.find((x) => x.supports(role));75    if (!p) throw new AiNotConfiguredError(role === 'embed' ? 'No embeddings provider configured (set OPENAI_API_KEY)' : 'AI provider not configured (set ANTHROPIC_API_KEY or OPENAI_API_KEY)');76    return p;77  }7879  supports(role: Role): boolean {80    return this.providers.some((p) => p.supports(role));81  }82  modelFor(role: Role): string {83    return this.providerFor(role).modelFor(role);84  }85  complete(role: Role, req: CompletionRequest): Promise<Completion> {86    return this.providerFor(role).complete(role, req);87  }88  stream(role: Role, req: CompletionRequest): AsyncIterable<StreamDelta> {89    return this.providerFor(role).stream(role, req);90  }91  runTools(role: Role, req: ToolAgentRequest): AsyncIterable<StreamDelta> {92    return this.providerFor(role).runTools(role, req);93  }94  extract<T>(role: Role, req: ExtractRequest<T>): Promise<ExtractResult<T>> {95    return this.providerFor(role).extract(role, req);96  }97  vision<T>(req: ExtractRequest<T>): Promise<ExtractResult<T>> {98    return this.providerFor('vision').vision(req);99  }100  embed(req: EmbedRequest): Promise<EmbedResult> {101    return this.providerFor('embed').embed(req);102  }103104  /** Table of role → provider/model for diagnostics and /admin. */105  table(): Array<{ role: Role; provider: string | null; model: string | null }> {106    return ROLES.map((role) => {107      try {108        const p = this.providerFor(role);109        return { role, provider: p.id, model: p.modelFor(role) };110      } catch {111        return { role, provider: null, model: null };112      }113    });114  }115}116117let _router: ModelRouter | null = null;118/** Process-wide router built from env. */119export function getRouter(): ModelRouter {120  if (!_router) _router = new ModelRouter(configFromEnv());121  return _router;122}123export function resetRouter(): void {124  _router = null;125}126