diff --git a/apps/web/src/lib/ai-gateway/providers/kilo-exclusive-model.test.ts b/apps/web/src/lib/ai-gateway/providers/kilo-exclusive-model.test.ts index acffdf51ac..d7906d7d2b 100644 --- a/apps/web/src/lib/ai-gateway/providers/kilo-exclusive-model.test.ts +++ b/apps/web/src/lib/ai-gateway/providers/kilo-exclusive-model.test.ts @@ -3,6 +3,7 @@ import { applyKiloExclusiveModelSettings, calculateCost_mUsd, convertFromKiloExclusiveModel, + getInferenceProvider, type KiloExclusiveModel, type PricingTiers, } from '@/lib/ai-gateway/providers/kilo-exclusive-model'; @@ -117,6 +118,38 @@ describe('convertFromKiloExclusiveModel', () => { }); }); +describe('getInferenceProvider', () => { + it('uses a single inference provider restriction before the gateway', () => { + const model = makeModel({ + internal_id: 'vendor/x', + gateway: 'vercel', + inference_provider_restriction: ['openai'], + }); + + expect(getInferenceProvider(model)).toEqual({ + slug: 'openai', + name: 'OPENAI', + training: false, + retainsPrompts: true, + }); + }); + + it('reports data collection for a concrete gateway provider', () => { + const model = makeModel({ + internal_id: 'vendor/x', + gateway: 'alibaba', + flags: ['requires-data-collection'], + }); + + expect(getInferenceProvider(model)).toEqual({ + slug: 'alibaba', + name: 'ALIBABA', + training: true, + retainsPrompts: true, + }); + }); +}); + describe('applyKiloExclusiveModelSettings', () => { it('rewrites the public model id to the internal id', () => { const req = makeRequest(undefined, 'kilo/test-model'); diff --git a/apps/web/src/lib/ai-gateway/providers/kilo-exclusive-model.ts b/apps/web/src/lib/ai-gateway/providers/kilo-exclusive-model.ts index ed498298f1..0b51d5b0d9 100644 --- a/apps/web/src/lib/ai-gateway/providers/kilo-exclusive-model.ts +++ b/apps/web/src/lib/ai-gateway/providers/kilo-exclusive-model.ts @@ -194,14 +194,19 @@ export function getInferenceProvider(model: KiloExclusiveModel): InferenceProvid if (model.flags.includes('stealth')) { return { slug: 'stealth', name: 'Stealth', training: true, retainsPrompts: true }; } - if (model.gateway === 'openrouter' || model.gateway === 'vercel') { - return null; - } - const slug = OpenRouterInferenceProviderIdSchema.parse(model.gateway); + + const slug: OpenRouterInferenceProviderId | null = + model.inference_provider_restriction.length === 1 + ? model.inference_provider_restriction[0] + : model.gateway === 'openrouter' || model.gateway === 'vercel' + ? null + : OpenRouterInferenceProviderIdSchema.parse(model.gateway); + if (!slug) return null; + return { slug, name: slug.toUpperCase(), - training: false, + training: model.flags.includes('requires-data-collection'), retainsPrompts: true, }; }