Momento/memento-note/lib/ai/provider-for-user.ts

import {
  getProviderInstance,
  type ProviderType,
} from '@/lib/ai/factory';
import { applyByokToConfig } from '@/lib/byok';
import {
  resolveAiRoute,
  type AiFeatureLane,
  type ResolvedAiRoute,
} from '@/lib/ai/router';
import { withAiProviderFallback } from '@/lib/ai/fallback';
import type { AIProvider } from '@/lib/ai/types';

export interface ProviderForUserResult {
  provider: AIProvider;
  usedByok: boolean;
  route: ResolvedAiRoute;
}

async function resolveProviderForLane(
  lane: AiFeatureLane,
  config: Record<string, string>,
  billingUserId?: string,
): Promise<ProviderForUserResult> {
  const cfg = { ...config };
  const route = resolveAiRoute(lane, cfg);
  let usedByok = false;

  if (billingUserId) {
    const overlay = await applyByokToConfig(
      billingUserId,
      route.providerType,
      cfg,
    );
    Object.assign(cfg, overlay.config);
    usedByok = overlay.usedByok;
  }

  const provider = getProviderInstance(
    route.providerType as ProviderType,
    cfg,
    route.modelName,
    route.embeddingModelName,
    route.ollamaBaseUrl,
  );

  return { provider, usedByok, route };
}

async function getChatProviderForBillingUser(
  config: Record<string, string>,
  billingUserId?: string,
): Promise<ProviderForUserResult> {
  return resolveProviderForLane('chat', config, billingUserId);
}

async function getTagsProviderForBillingUser(
  config: Record<string, string>,
  billingUserId?: string,
): Promise<ProviderForUserResult> {
  return resolveProviderForLane('tags', config, billingUserId);
}

async function getEmbeddingsProviderForBillingUser(
  config: Record<string, string>,
  billingUserId?: string,
): Promise<ProviderForUserResult> {
  return resolveProviderForLane('embedding', config, billingUserId);
}

/** Run a lane with BYOK overlay; skips system fallback when user key is active. */
export async function willUseByokForLane(
  lane: AiFeatureLane,
  config: Record<string, string>,
  billingUserId?: string,
): Promise<{ providerType: string; usedByok: boolean }> {
  if (!billingUserId) {
    const route = resolveAiRoute(lane, config)
    return { providerType: route.providerType, usedByok: false }
  }
  const route = resolveAiRoute(lane, config)
  const overlay = await applyByokToConfig(billingUserId, route.providerType, config)
  return { providerType: route.providerType, usedByok: overlay.usedByok }
}

export async function runLaneWithBillingUser<T>(
  lane: AiFeatureLane,
  config: Record<string, string>,
  billingUserId: string | undefined,
  run: (provider: AIProvider) => Promise<T>,
): Promise<{ result: T; usedByok: boolean }> {
  if (billingUserId) {
    const resolved =
      lane === 'chat'
        ? await getChatProviderForBillingUser(config, billingUserId)
        : lane === 'tags'
          ? await getTagsProviderForBillingUser(config, billingUserId)
          : await getEmbeddingsProviderForBillingUser(config, billingUserId);

    if (resolved.usedByok) {
      const result = await run(resolved.provider);
      return { result, usedByok: true };
    }
  }

  const result = await withAiProviderFallback(lane, config, run);
  return { result, usedByok: false };
}