122 lines
4.0 KiB
TypeScript
122 lines
4.0 KiB
TypeScript
import {
|
|
getProviderInstance,
|
|
getProviderConfigKeys,
|
|
type ProviderType,
|
|
} from '@/lib/ai/factory';
|
|
import { getAnyActiveByokForUser, hasAnyActiveByok, ByokUnavailableError } 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;
|
|
}
|
|
|
|
/** Resolve the best AI provider for a user's lane — BYOK first, then admin config. */
|
|
async function resolveProviderForLane(
|
|
lane: AiFeatureLane,
|
|
config: Record<string, string>,
|
|
billingUserId?: string,
|
|
): Promise<ProviderForUserResult> {
|
|
const cfg = { ...config };
|
|
const adminRoute = resolveAiRoute(lane, cfg);
|
|
|
|
if (billingUserId) {
|
|
// Prefer admin's provider, fallback to any active BYOK key
|
|
const byok = await getAnyActiveByokForUser(billingUserId, adminRoute.providerType, lane);
|
|
|
|
if (byok) {
|
|
const { apiKeyConfigKey, baseUrlConfigKey } = getProviderConfigKeys(byok.provider);
|
|
const byokCfg: Record<string, string> = { ...cfg };
|
|
if (apiKeyConfigKey) byokCfg[apiKeyConfigKey] = byok.plaintext;
|
|
if (baseUrlConfigKey && byok.baseUrl) byokCfg[baseUrlConfigKey] = byok.baseUrl;
|
|
|
|
const resolvedModel = (byok.model && byok.model.trim()) ? byok.model : adminRoute.modelName;
|
|
console.log(`[byok] Using BYOK key: provider=${byok.provider} model=${resolvedModel} user=${billingUserId}`);
|
|
|
|
const provider = getProviderInstance(
|
|
byok.provider as ProviderType,
|
|
byokCfg,
|
|
resolvedModel,
|
|
adminRoute.embeddingModelName,
|
|
adminRoute.ollamaBaseUrl,
|
|
);
|
|
|
|
return {
|
|
provider,
|
|
usedByok: true,
|
|
route: {
|
|
...adminRoute,
|
|
providerType: byok.provider as ResolvedAiRoute['providerType'],
|
|
modelName: resolvedModel,
|
|
},
|
|
};
|
|
}
|
|
|
|
// No key resolved — if user HAS active BYOK rows, decryption failed → throw, no silent fallback
|
|
const hasByok = await hasAnyActiveByok(billingUserId);
|
|
if (hasByok) {
|
|
throw new ByokUnavailableError();
|
|
}
|
|
}
|
|
|
|
// No BYOK configured → use admin config
|
|
const provider = getProviderInstance(
|
|
adminRoute.providerType as ProviderType,
|
|
cfg,
|
|
adminRoute.modelName,
|
|
adminRoute.embeddingModelName,
|
|
adminRoute.ollamaBaseUrl,
|
|
);
|
|
|
|
return { provider, usedByok: false, route: adminRoute };
|
|
}
|
|
|
|
/** Check if a lane will use BYOK for a given user. */
|
|
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 byok = await getAnyActiveByokForUser(billingUserId, route.providerType);
|
|
return { providerType: byok?.provider ?? route.providerType, usedByok: !!byok };
|
|
}
|
|
|
|
/**
|
|
* Run an AI lane with BYOK priority.
|
|
* - If user has active BYOK → uses it, no quota counted.
|
|
* - If user has BYOK configured but it can't be loaded → throws ByokUnavailableError (no fallback).
|
|
* - If user has no BYOK → uses admin config with system fallback.
|
|
*/
|
|
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) {
|
|
// May throw ByokUnavailableError — let it propagate, callers should handle it
|
|
const resolved = await resolveProviderForLane(lane, config, billingUserId);
|
|
|
|
if (resolved.usedByok) {
|
|
const result = await run(resolved.provider);
|
|
return { result, usedByok: true };
|
|
}
|
|
}
|
|
|
|
// No BYOK configured → use admin config with system fallback
|
|
const result = await withAiProviderFallback(lane, config, run);
|
|
return { result, usedByok: false };
|
|
}
|