115 lines
3.3 KiB
TypeScript
115 lines
3.3 KiB
TypeScript
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;
|
|
let byokModel: string | null = null;
|
|
|
|
if (billingUserId) {
|
|
const overlay = await applyByokToConfig(
|
|
billingUserId,
|
|
route.providerType,
|
|
cfg,
|
|
);
|
|
Object.assign(cfg, overlay.config);
|
|
usedByok = overlay.usedByok;
|
|
byokModel = overlay.model;
|
|
}
|
|
|
|
const resolvedModel = byokModel && byokModel.trim() !== '' ? byokModel : route.modelName;
|
|
|
|
const provider = getProviderInstance(
|
|
route.providerType as ProviderType,
|
|
cfg,
|
|
resolvedModel,
|
|
route.embeddingModelName,
|
|
route.ollamaBaseUrl,
|
|
);
|
|
|
|
const updatedRoute = { ...route, modelName: resolvedModel };
|
|
|
|
return { provider, usedByok, route: updatedRoute };
|
|
}
|
|
|
|
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 };
|
|
}
|