139 lines
3.6 KiB
TypeScript
139 lines
3.6 KiB
TypeScript
import { prisma } from '@/lib/prisma';
|
|
import { decryptApiKey, encryptApiKey, hashApiKey } from '@/lib/crypto';
|
|
import {
|
|
VALID_PROVIDERS,
|
|
type AiGatewayProvider,
|
|
} from '@/lib/ai/router';
|
|
import { getProviderConfigKeys } from '@/lib/ai/factory';
|
|
import type { SubscriptionTier } from '@/lib/entitlements';
|
|
|
|
const PRO_BYOK_PROVIDERS: readonly AiGatewayProvider[] = [
|
|
'openai',
|
|
'anthropic',
|
|
'deepseek',
|
|
'openrouter',
|
|
'minimax',
|
|
'zai',
|
|
];
|
|
|
|
const BUSINESS_BYOK_PROVIDERS: readonly AiGatewayProvider[] = [
|
|
...VALID_PROVIDERS,
|
|
].filter((p) => p !== 'ollama' && p !== 'lmstudio') as AiGatewayProvider[];
|
|
|
|
export function getAllowedByokProviders(
|
|
tier: SubscriptionTier,
|
|
): readonly AiGatewayProvider[] {
|
|
if (tier === 'BASIC') return [];
|
|
if (tier === 'PRO') return PRO_BYOK_PROVIDERS;
|
|
return BUSINESS_BYOK_PROVIDERS;
|
|
}
|
|
|
|
export function isByokProviderAllowed(
|
|
tier: SubscriptionTier,
|
|
provider: string,
|
|
): boolean {
|
|
return getAllowedByokProviders(tier).includes(provider as AiGatewayProvider);
|
|
}
|
|
|
|
export async function hasAnyActiveByok(userId: string): Promise<boolean> {
|
|
const count = await prisma.userAPIKey.count({
|
|
where: { userId, isActive: true },
|
|
});
|
|
return count > 0;
|
|
}
|
|
|
|
export async function getActiveByokKey(userId: string, provider: string) {
|
|
return prisma.userAPIKey.findFirst({
|
|
where: { userId, provider, isActive: true },
|
|
});
|
|
}
|
|
|
|
export async function resolveByokApiKey(
|
|
userId: string,
|
|
providerType: string,
|
|
): Promise<{ plaintext: string; provider: string; model: string | null } | null> {
|
|
const row = await getActiveByokKey(userId, providerType);
|
|
if (!row) return null;
|
|
try {
|
|
const plaintext = await decryptApiKey(row.encryptedKey);
|
|
return { plaintext, provider: row.provider, model: row.model };
|
|
} catch (err) {
|
|
console.error('[byok] Failed to decrypt key for provider', providerType, err);
|
|
return null;
|
|
}
|
|
}
|
|
|
|
export async function applyByokToConfig(
|
|
billingUserId: string,
|
|
providerType: string,
|
|
config: Record<string, string>,
|
|
): Promise<{ config: Record<string, string>; usedByok: boolean; model: string | null }> {
|
|
const byok = await resolveByokApiKey(billingUserId, providerType);
|
|
if (!byok) return { config, usedByok: false, model: null };
|
|
|
|
const { apiKeyConfigKey } = getProviderConfigKeys(providerType);
|
|
if (!apiKeyConfigKey) return { config, usedByok: false, model: null };
|
|
|
|
return {
|
|
config: { ...config, [apiKeyConfigKey]: byok.plaintext },
|
|
usedByok: true,
|
|
model: byok.model,
|
|
};
|
|
}
|
|
|
|
export async function upsertUserApiKey(params: {
|
|
userId: string;
|
|
provider: AiGatewayProvider;
|
|
plaintext: string;
|
|
alias?: string;
|
|
model?: string;
|
|
}) {
|
|
const encryptedKey = await encryptApiKey(params.plaintext);
|
|
const keyHash = hashApiKey(params.plaintext);
|
|
|
|
return prisma.userAPIKey.upsert({
|
|
where: {
|
|
userId_provider: {
|
|
userId: params.userId,
|
|
provider: params.provider,
|
|
},
|
|
},
|
|
create: {
|
|
userId: params.userId,
|
|
provider: params.provider,
|
|
alias: params.alias ?? '',
|
|
encryptedKey,
|
|
keyHash,
|
|
model: params.model ?? null,
|
|
isActive: true,
|
|
},
|
|
update: {
|
|
alias: params.alias ?? '',
|
|
encryptedKey,
|
|
keyHash,
|
|
model: params.model ?? null,
|
|
isActive: true,
|
|
},
|
|
});
|
|
}
|
|
|
|
export function toPublicApiKey(row: {
|
|
provider: string;
|
|
alias: string;
|
|
model: string | null;
|
|
isActive: boolean;
|
|
lastUsedAt: Date | null;
|
|
createdAt: Date;
|
|
updatedAt: Date;
|
|
}) {
|
|
return {
|
|
provider: row.provider,
|
|
alias: row.alias,
|
|
model: row.model,
|
|
isActive: row.isActive,
|
|
lastUsedAt: row.lastUsedAt,
|
|
createdAt: row.createdAt,
|
|
updatedAt: row.updatedAt,
|
|
};
|
|
}
|