diff --git a/apps/web/src/lib/ai-gateway/providers/direct-byok/index.test.ts b/apps/web/src/lib/ai-gateway/providers/direct-byok/index.test.ts index 64536da4e3..ac8a0c193c 100644 --- a/apps/web/src/lib/ai-gateway/providers/direct-byok/index.test.ts +++ b/apps/web/src/lib/ai-gateway/providers/direct-byok/index.test.ts @@ -28,6 +28,13 @@ jest.mock('./direct-byok-definitions', () => ({ { id: 'supported-model', name: 'Supported Model', + flags: ['reasoning'], + context_length: 4096, + max_completion_tokens: 1024, + }, + { + id: 'non-reasoning-model', + name: 'Non-reasoning Model', context_length: 4096, max_completion_tokens: 1024, }, @@ -56,9 +63,9 @@ jest.mock('./direct-byok-definitions', () => ({ async function loadDirectByokModule() { const directByokProviders = (await import('./direct-byok-definitions')).default; - const { getDirectByokModel } = await import('.'); + const { getDirectByokModel, getDirectByokModelsForUser } = await import('.'); - return { directByokProviders, getDirectByokModel }; + return { directByokProviders, getDirectByokModel, getDirectByokModelsForUser }; } describe('getDirectByokModel', () => { @@ -88,4 +95,23 @@ describe('getDirectByokModel', () => { expect(directByokProviders[0].models).toHaveBeenCalledTimes(1); expect(directByokProviders[1].models).not.toHaveBeenCalled(); }); + + test('advertises only the reasoning parameter for reasoning models', async () => { + const { getDirectByokModelsForUser } = await loadDirectByokModule(); + const { getBYOKforUser } = await import('@/lib/ai-gateway/byok'); + jest + .mocked(getBYOKforUser) + .mockResolvedValueOnce([{ providerId: 'chutes-byok', decryptedAPIKey: 'test-key' }]); + + const models = await getDirectByokModelsForUser('user-id'); + + expect(models[0].supported_parameters).toEqual([ + 'max_tokens', + 'temperature', + 'tools', + 'reasoning', + ]); + expect(models[1].supported_parameters).toEqual(['max_tokens', 'temperature', 'tools']); + expect(models.flatMap(model => model.supported_parameters)).not.toContain('include_reasoning'); + }); }); diff --git a/apps/web/src/lib/ai-gateway/providers/direct-byok/index.ts b/apps/web/src/lib/ai-gateway/providers/direct-byok/index.ts index 4ed2d9b567..1c400427c0 100644 --- a/apps/web/src/lib/ai-gateway/providers/direct-byok/index.ts +++ b/apps/web/src/lib/ai-gateway/providers/direct-byok/index.ts @@ -54,7 +54,9 @@ function convertModel( is_moderated: false, }, per_request_limits: null, - supported_parameters: ['max_tokens', 'temperature', 'tools', 'reasoning', 'include_reasoning'], + supported_parameters: ['max_tokens', 'temperature', 'tools'].concat( + model.flags?.includes('reasoning') ? ['reasoning'] : [] + ), default_parameters: {}, preferredIndex: model.flags?.includes('recommended') ? preferredIndex : undefined, hasUserByokAvailable: true,