Skip to content

Commit 0a02b3a

Browse files
chrarnolduskilo-code-bot[bot]
chrarnoldus
and
kilo-code-bot[bot]
authored
feat(models): expose user BYOK availability (#4033)
* feat(models): expose user BYOK availability Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com> * style(models): format BYOK availability changes Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com> * fix(models): omit BYOK metadata where irrelevant Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com> * refactor(models): simplify BYOK availability assignment Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com> * Update logging --------- Co-authored-by: chrarnoldus <12196001+chrarnoldus@users.noreply.github.com> Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com>
1 parent 058b72f commit 0a02b3a

6 files changed

Lines changed: 185 additions & 4 deletions

File tree

‎apps/web/src/app/api/openrouter/models/route.test.ts‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
import { beforeEach, describe, expect, test } from '@jest/globals';
2+
import { NextRequest } from 'next/server';
3+
import type { OpenRouterModel } from '@/lib/organizations/organization-types';
4+
import { getEnhancedOpenRouterModels } from '@/lib/ai-gateway/providers/openrouter';
5+
import { getUserFromAuth } from '@/lib/user/server';
6+
import { getDirectByokModelsForUser } from '@/lib/ai-gateway/providers/direct-byok';
7+
import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models';
8+
import { addUserByokAvailability, getUserByokProviderIds } from '@/lib/ai-gateway/byok';
9+
import { getAvailableModelsForOrganization } from '@/lib/organizations/organization-models';
10+
import { GET } from './route';
11+
12+
jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn() }));
13+
jest.mock('@/lib/user/server', () => ({ getUserFromAuth: jest.fn() }));
14+
jest.mock('@/lib/ai-gateway/providers/openrouter', () => ({
15+
getEnhancedOpenRouterModels: jest.fn(),
16+
}));
17+
jest.mock('@/lib/ai-gateway/providers/direct-byok', () => ({
18+
getDirectByokModelsForUser: jest.fn(),
19+
}));
20+
jest.mock('@/lib/ai-gateway/experiments/list-available-experiment-models', () => ({
21+
listAvailableExperimentModels: jest.fn(),
22+
}));
23+
jest.mock('@/lib/ai-gateway/byok', () => ({
24+
addUserByokAvailability: jest.fn(),
25+
getUserByokProviderIds: jest.fn(),
26+
}));
27+
jest.mock('@/lib/organizations/organization-models', () => ({
28+
getAvailableModelsForOrganization: jest.fn(),
29+
}));
30+
jest.mock('@/lib/drizzle', () => ({ readDb: {} }));
31+
32+
const mockedGetUserFromAuth = jest.mocked(getUserFromAuth);
33+
const mockedGetEnhancedOpenRouterModels = jest.mocked(getEnhancedOpenRouterModels);
34+
const mockedGetDirectByokModelsForUser = jest.mocked(getDirectByokModelsForUser);
35+
const mockedListAvailableExperimentModels = jest.mocked(listAvailableExperimentModels);
36+
const mockedAddUserByokAvailability = jest.mocked(addUserByokAvailability);
37+
const mockedGetUserByokProviderIds = jest.mocked(getUserByokProviderIds);
38+
const mockedGetAvailableModelsForOrganization = jest.mocked(getAvailableModelsForOrganization);
39+
40+
function makeModel(id: string): OpenRouterModel {
41+
return {
42+
id,
43+
name: id,
44+
created: 0,
45+
description: '',
46+
architecture: {
47+
input_modalities: ['text'],
48+
output_modalities: ['text'],
49+
tokenizer: 'test',
50+
},
51+
top_provider: { is_moderated: false },
52+
pricing: { prompt: '0', completion: '0' },
53+
context_length: 0,
54+
};
55+
}
56+
57+
function request() {
58+
return new NextRequest('http://localhost:3000/api/openrouter/models');
59+
}
60+
61+
describe('GET /api/openrouter/models', () => {
62+
beforeEach(() => {
63+
jest.resetAllMocks();
64+
mockedGetUserFromAuth.mockResolvedValue({
65+
user: null,
66+
organizationId: null,
67+
authFailedResponse: null,
68+
} as never);
69+
mockedGetEnhancedOpenRouterModels.mockResolvedValue({ data: [makeModel('public/model')] });
70+
mockedGetDirectByokModelsForUser.mockResolvedValue([]);
71+
mockedListAvailableExperimentModels.mockResolvedValue([]);
72+
mockedGetUserByokProviderIds.mockResolvedValue([]);
73+
mockedGetAvailableModelsForOrganization.mockResolvedValue(null);
74+
});
75+
76+
test('leaves BYOK availability undefined for unauthenticated requests', async () => {
77+
const response = await GET(request());
78+
79+
expect(response.status).toBe(200);
80+
await expect(response.json()).resolves.toEqual({ data: [makeModel('public/model')] });
81+
expect(mockedGetUserByokProviderIds).not.toHaveBeenCalled();
82+
expect(mockedAddUserByokAvailability).not.toHaveBeenCalled();
83+
});
84+
85+
test('returns BYOK availability for regular and direct authenticated models', async () => {
86+
const publicModel = makeModel('public/model');
87+
const directModel = { ...makeModel('direct/model'), hasUserByokAvailable: true };
88+
const experimentModel = makeModel('experiment/model');
89+
mockedGetUserFromAuth.mockResolvedValue({
90+
user: { id: 'user-id' },
91+
organizationId: null,
92+
authFailedResponse: null,
93+
} as never);
94+
mockedGetDirectByokModelsForUser.mockResolvedValue([directModel] as never);
95+
mockedListAvailableExperimentModels.mockResolvedValue([experimentModel]);
96+
mockedGetUserByokProviderIds.mockResolvedValue(['anthropic']);
97+
mockedAddUserByokAvailability.mockResolvedValue([
98+
{ ...publicModel, hasUserByokAvailable: true },
99+
]);
100+
101+
const response = await GET(request());
102+
103+
expect(response.status).toBe(200);
104+
await expect(response.json()).resolves.toEqual({
105+
data: [{ ...publicModel, hasUserByokAvailable: true }, directModel, experimentModel],
106+
});
107+
});
108+
});

‎apps/web/src/app/api/openrouter/models/route.ts‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import { getAvailableModelsForOrganization } from '@/lib/organizations/organizat
99
import { FEATURE_HEADER, validateFeatureHeader } from '@/lib/feature-detection';
1010
import { filterByFeature } from '@/lib/ai-gateway/models';
1111
import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models';
12+
import { addUserByokAvailability, getUserByokProviderIds } from '@/lib/ai-gateway/byok';
13+
import { readDb } from '@/lib/drizzle';
1214

1315
async function tryGetUserFromAuth() {
1416
try {
@@ -43,10 +45,27 @@ export async function GET(
4345
if (!Array.isArray(data.data)) {
4446
return NextResponse.json(data);
4547
}
46-
const byokModels = auth?.user ? await getDirectByokModelsForUser(auth.user.id) : [];
47-
const experimentModels = await listAvailableExperimentModels();
48+
if (!auth?.user) {
49+
const experimentModels = await listAvailableExperimentModels();
50+
return NextResponse.json({
51+
data: filterByFeature(data.data.concat(experimentModels), feature),
52+
});
53+
}
54+
55+
const [byokModels, experimentModels, enabledByokProviderIds] = await Promise.all([
56+
getDirectByokModelsForUser(auth.user.id),
57+
listAvailableExperimentModels(),
58+
getUserByokProviderIds(readDb, auth.user.id),
59+
]);
60+
const modelsWithByokAvailability = await addUserByokAvailability(
61+
data.data,
62+
enabledByokProviderIds
63+
);
4864
return NextResponse.json({
49-
data: filterByFeature(data.data.concat(byokModels, experimentModels), feature),
65+
data: filterByFeature(
66+
modelsWithByokAvailability.concat(byokModels, experimentModels),
67+
feature
68+
),
5069
});
5170
} catch (error) {
5271
captureException(error, {

‎apps/web/src/lib/ai-gateway/byok/index.ts‎

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,13 @@ import { isCodestralModel } from '@/lib/ai-gateway/providers/mistral';
1313
import { mapModelIdToVercel } from '@/lib/ai-gateway/providers/vercel/mapModelIdToVercel';
1414
import type { BYOKResult } from '@/lib/ai-gateway/providers/types';
1515
import { getVercelModelsMetadata } from '@/lib/ai-gateway/providers/gateway-models-cache';
16+
import type { OpenRouterModel } from '@/lib/organizations/organization-types';
17+
import { isKiloExclusiveModel } from '@/lib/ai-gateway/models';
1618

1719
export async function getModelUserByokProviders(modelId: string): Promise<UserByokProviderId[]> {
1820
const vercelModelMetadata = await getVercelModelsMetadata();
1921
if (Object.keys(vercelModelMetadata).length === 0) {
20-
console.error('[getModelUserByokProviders] no Vercel model metadata in the database');
22+
console.error('[getModelUserByokProviders] no Vercel model metadata for model %s', modelId);
2123
return [];
2224
}
2325
const providers: UserByokProviderId[] =
@@ -35,6 +37,49 @@ export async function getModelUserByokProviders(modelId: string): Promise<UserBy
3537
return providers;
3638
}
3739

40+
export async function getUserByokProviderIds(
41+
fromDb: typeof db,
42+
userId: string
43+
): Promise<UserByokProviderId[]> {
44+
const rows = await fromDb
45+
.select({ provider_id: byok_api_keys.provider_id })
46+
.from(byok_api_keys)
47+
.where(and(eq(byok_api_keys.kilo_user_id, userId), eq(byok_api_keys.is_enabled, true)));
48+
49+
return rows.map(row => UserByokProviderIdSchema.parse(row.provider_id));
50+
}
51+
52+
export async function getOrganizationByokProviderIds(
53+
fromDb: typeof db,
54+
organizationId: string
55+
): Promise<UserByokProviderId[]> {
56+
const rows = await fromDb
57+
.select({ provider_id: byok_api_keys.provider_id })
58+
.from(byok_api_keys)
59+
.where(
60+
and(eq(byok_api_keys.organization_id, organizationId), eq(byok_api_keys.is_enabled, true))
61+
);
62+
63+
return rows.map(row => UserByokProviderIdSchema.parse(row.provider_id));
64+
}
65+
66+
export async function addUserByokAvailability(
67+
models: OpenRouterModel[],
68+
enabledProviderIds: UserByokProviderId[]
69+
): Promise<OpenRouterModel[]> {
70+
const enabledProviders = new Set(enabledProviderIds);
71+
return Promise.all(
72+
models.map(async model => {
73+
const hasUserByokAvailable =
74+
!isKiloExclusiveModel(model.id) &&
75+
(await getModelUserByokProviders(model.id)).some(provider =>
76+
enabledProviders.has(provider)
77+
);
78+
return { ...model, hasUserByokAvailable };
79+
})
80+
);
81+
}
82+
3883
export function decryptByokRow({
3984
encrypted_api_key,
4085
provider_id,

‎apps/web/src/lib/ai-gateway/providers/direct-byok/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ function convertModel(
5757
supported_parameters: ['max_tokens', 'temperature', 'tools', 'reasoning', 'include_reasoning'],
5858
default_parameters: {},
5959
preferredIndex: model.flags?.includes('recommended') ? preferredIndex : undefined,
60+
hasUserByokAvailable: true,
6061
opencode: {
6162
ai_sdk_provider: getAiSdkProvider(id) ?? provider.default_ai_sdk_provider,
6263
variants: getModelVariants(id),

‎apps/web/src/lib/organizations/organization-models.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ import { getDirectByokModelsForOrganization } from '@/lib/ai-gateway/providers/d
1010
import { getOrganizationById } from '@/lib/organizations/organizations';
1111
import { getEffectiveModelRestrictions } from '@/lib/organizations/model-restrictions';
1212
import { listAvailableExperimentModels } from '@/lib/ai-gateway/experiments/list-available-experiment-models';
13+
import { addUserByokAvailability, getOrganizationByokProviderIds } from '@/lib/ai-gateway/byok';
14+
import { readDb } from '@/lib/drizzle';
1315

1416
export async function getAvailableModelsForOrganization(
1517
organizationId: string
@@ -40,6 +42,11 @@ export async function getAvailableModelsForOrganization(
4042
filteredModels = models;
4143
}
4244

45+
filteredModels = await addUserByokAvailability(
46+
filteredModels,
47+
await getOrganizationByokProviderIds(readDb, organizationId)
48+
);
49+
4350
if (organization.plan === 'teams' && organization.settings.data_collection === 'deny') {
4451
filteredModels = filteredModels.filter(model => model.mayTrainOnYourPrompts !== true);
4552
}

‎apps/web/src/lib/organizations/organization-types.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,7 @@ const OpenRouterModelSchema = z.object({
200200
preferredIndex: z.number().optional(),
201201
isFree: z.boolean().optional(),
202202
mayTrainOnYourPrompts: z.boolean().optional(),
203+
hasUserByokAvailable: z.boolean().optional(),
203204
terminalBench: z
204205
.object({
205206
overallScore: z.number(),

0 commit comments

Comments
 (0)