From 7243cd532c76ec33870b9e11b265ecfa7ed41204 Mon Sep 17 00:00:00 2001 From: chrarnoldus <12196001+chrarnoldus@users.noreply.github.com> Date: Fri, 10 Jul 2026 12:56:04 +0000 Subject: [PATCH] feat(ai-gateway): route embeddings and transcription via Vercel Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com> --- .../src/app/admin/gateway/RoutingContent.tsx | 140 ++++++++++++------ .../audio/transcriptions/route.test.ts | 44 ++++++ .../openrouter/audio/transcriptions/route.ts | 34 ++++- .../app/api/openrouter/embeddings/route.ts | 28 +++- .../embeddings/embedding-request.ts | 4 +- .../src/lib/ai-gateway/gateway-config.test.ts | 88 +++++++++-- apps/web/src/lib/ai-gateway/gateway-config.ts | 53 ++++++- .../lib/ai-gateway/llm-proxy-helpers.test.ts | 34 +++++ .../src/lib/ai-gateway/llm-proxy-helpers.ts | 34 +++-- .../lib/ai-gateway/providers/get-provider.ts | 37 ++++- .../providers/getEmbeddingProvider.test.ts | 81 +++++++++- .../lib/ai-gateway/providers/vercel/index.ts | 66 ++++++--- .../transcription-request.test.ts | 23 +++ .../transcriptions/transcription-request.ts | 40 ++++- .../routers/admin/gateway-config-router.ts | 6 +- packages/db/src/schema-types.ts | 2 +- 16 files changed, 591 insertions(+), 123 deletions(-) diff --git a/apps/web/src/app/admin/gateway/RoutingContent.tsx b/apps/web/src/app/admin/gateway/RoutingContent.tsx index 6688c759e7..431f17e9d2 100644 --- a/apps/web/src/app/admin/gateway/RoutingContent.tsx +++ b/apps/web/src/app/admin/gateway/RoutingContent.tsx @@ -9,7 +9,23 @@ import { Input } from '@/components/ui/input'; import { Textarea } from '@/components/ui/textarea'; import { Label } from '@/components/ui/label'; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/components/ui/card'; -import { DEFAULT_VERCEL_PERCENTAGE, NOTE_MAX_LENGTH } from '@/lib/ai-gateway/gateway-config'; +import { + DEFAULT_VERCEL_PERCENTAGE, + NOTE_MAX_LENGTH, + type VercelRoutingApiType, +} from '@/lib/ai-gateway/gateway-config'; + +const ROUTING_API_LABELS: Record = { + chat: 'Chat', + embeddings: 'Embeddings', + transcription: 'Transcription', +}; + +const EMPTY_PERCENTAGES: Record = { + chat: '', + embeddings: '', + transcription: '', +}; export function RoutingContent() { const trpc = useTRPC(); @@ -17,13 +33,17 @@ export function RoutingContent() { const { data, isLoading } = useQuery(trpc.admin.gatewayConfig.get.queryOptions()); - const [inputValue, setInputValue] = useState(''); + const [percentages, setPercentages] = useState(EMPTY_PERCENTAGES); const [noteValue, setNoteValue] = useState(''); const [hasChanges, setHasChanges] = useState(false); useEffect(() => { if (data) { - setInputValue(data.vercel_routing_percentage?.toString() ?? ''); + setPercentages({ + chat: data.vercel_chat_routing_percentage?.toString() ?? '', + embeddings: data.vercel_embeddings_routing_percentage?.toString() ?? '', + transcription: data.vercel_transcription_routing_percentage?.toString() ?? '', + }); setNoteValue(''); setHasChanges(false); } @@ -35,7 +55,7 @@ export function RoutingContent() { void queryClient.invalidateQueries({ queryKey: trpc.admin.gatewayConfig.get.queryKey(), }); - toast.success('Vercel routing percentage updated'); + toast.success('Vercel routing percentages updated'); }, onError: error => { toast.error(error.message || 'Failed to update'); @@ -49,61 +69,85 @@ export function RoutingContent() { } function handleSave() { - const trimmed = inputValue.trim(); - const note = noteInput(); - if (trimmed === '') { - mutation.mutate({ vercel_routing_percentage: null, note }); - } else { - const num = Number(trimmed); - if (!Number.isInteger(num) || num < 0 || num > 100) { - toast.error('Please enter a whole number between 0 and 100, or leave empty for default'); - return; - } - mutation.mutate({ vercel_routing_percentage: num, note }); + function parsePercentage(input: string): number | null | undefined { + const trimmed = input.trim(); + if (trimmed === '') return null; + const value = Number(trimmed); + return Number.isInteger(value) && value >= 0 && value <= 100 ? value : undefined; + } + + const chat = parsePercentage(percentages.chat); + const embeddings = parsePercentage(percentages.embeddings); + const transcription = parsePercentage(percentages.transcription); + if (chat === undefined || embeddings === undefined || transcription === undefined) { + toast.error('Please enter a whole number between 0 and 100, or leave empty for default'); + return; } + mutation.mutate({ + vercel_chat_routing_percentage: chat, + vercel_embeddings_routing_percentage: embeddings, + vercel_transcription_routing_percentage: transcription, + note: noteInput(), + }); } function handleClear() { - mutation.mutate({ vercel_routing_percentage: null, note: noteInput() }); + mutation.mutate({ + vercel_chat_routing_percentage: null, + vercel_embeddings_routing_percentage: null, + vercel_transcription_routing_percentage: null, + note: noteInput(), + }); } if (isLoading) { return
Loading...
; } - const currentOverride = data?.vercel_routing_percentage; - const isOverrideActive = currentOverride !== null && currentOverride !== undefined; + const currentPercentages = { + chat: data?.vercel_chat_routing_percentage, + embeddings: data?.vercel_embeddings_routing_percentage, + transcription: data?.vercel_transcription_routing_percentage, + }; + const isOverrideActive = Object.values(currentPercentages).some(value => value != null); return (
- Vercel Routing Percentage + Vercel Routing Percentages - For models available on the Vercel AI Gateway, controls the percentage of traffic routed - to Vercel (vs OpenRouter). Models not available on Vercel always go to OpenRouter, so - overall traffic may still be skewed towards OpenRouter. Leave empty to use the default ( - {DEFAULT_VERCEL_PERCENTAGE}%). + Controls the Vercel traffic split independently for each API. Unsupported models always + use OpenRouter. Leave a field empty to use the {DEFAULT_VERCEL_PERCENTAGE}% default. +
+ {Object.entries(ROUTING_API_LABELS).map(([apiType, label]) => ( +
+ +
+ { + setPercentages(current => ({ ...current, [apiType]: e.target.value })); + setHasChanges(true); + }} + onKeyDown={e => { + if (e.key === 'Enter') handleSave(); + }} + /> + % +
+
+ ))} +
- { - setInputValue(e.target.value); - setHasChanges(true); - }} - onKeyDown={e => { - if (e.key === 'Enter') handleSave(); - }} - className="w-48" - /> - % @@ -114,7 +158,7 @@ export function RoutingContent() { variant="outline" size="sm" > - Clear override + Clear overrides )}
@@ -136,17 +180,19 @@ export function RoutingContent() {
{isOverrideActive ? ( -

- Override active:{' '} - {currentOverride}% of traffic - goes to Vercel. +

+

+ Active routing: chat {currentPercentages.chat ?? DEFAULT_VERCEL_PERCENTAGE}%, + embeddings {currentPercentages.embeddings ?? DEFAULT_VERCEL_PERCENTAGE}%, + transcription {currentPercentages.transcription ?? DEFAULT_VERCEL_PERCENTAGE}%. +

{data?.updated_by_email && ( - +

Set by {data.updated_by_email} {data.updated_at && <> at {new Date(data.updated_at).toLocaleString()}}. - +

)} -

+
) : (

No override set. Using default routing ({DEFAULT_VERCEL_PERCENTAGE}% to Vercel). diff --git a/apps/web/src/app/api/openrouter/audio/transcriptions/route.test.ts b/apps/web/src/app/api/openrouter/audio/transcriptions/route.test.ts index 4edbc85c37..787b7768b9 100644 --- a/apps/web/src/app/api/openrouter/audio/transcriptions/route.test.ts +++ b/apps/web/src/app/api/openrouter/audio/transcriptions/route.test.ts @@ -4,6 +4,8 @@ import { getBalanceAndOrgSettings } from '@/lib/organizations/organization-usage import type { User } from '@kilocode/db/schema'; import { emitApiMetricsForResponse } from '@/lib/ai-gateway/o11y/api-metrics.server'; import type { OrganizationSettings } from '@/lib/organizations/organization-types'; +import { getTranscriptionProvider } from '@/lib/ai-gateway/providers/get-provider'; +import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions'; jest.mock('next/server', () => { return { @@ -14,6 +16,7 @@ jest.mock('next/server', () => { jest.mock('@/lib/user/server'); jest.mock('@/lib/organizations/organization-usage'); +jest.mock('@/lib/ai-gateway/providers/get-provider'); jest.mock('@/lib/ai-gateway/o11y/api-metrics.server', () => ({ emitApiMetricsForResponse: jest.fn(), })); @@ -28,6 +31,7 @@ jest.mock('@/lib/ai-gateway/llm-proxy-helpers', () => { const mockedGetUserFromAuth = jest.mocked(getUserFromAuth); const mockedGetBalanceAndOrgSettings = jest.mocked(getBalanceAndOrgSettings); const mockedEmitApiMetricsForResponse = jest.mocked(emitApiMetricsForResponse); +const mockedGetTranscriptionProvider = jest.mocked(getTranscriptionProvider); const mockedFetch = jest.fn() as jest.MockedFunction; const originalFetch = globalThis.fetch; @@ -71,6 +75,10 @@ describe('POST /api/gateway/v1/audio/transcriptions', () => { beforeEach(() => { jest.resetAllMocks(); globalThis.fetch = mockedFetch; + mockedGetTranscriptionProvider.mockResolvedValue({ + provider: PROVIDERS.OPENROUTER, + userByok: null, + }); }); afterAll(() => { @@ -149,6 +157,42 @@ describe('POST /api/gateway/v1/audio/transcriptions', () => { expect(upstream.provider).toEqual({ only: ['openai'], data_collection: 'deny' }); }); + it('adapts transcription requests to the Vercel REST API', async () => { + setUserAuth(); + mockedGetTranscriptionProvider.mockResolvedValue({ + provider: PROVIDERS.VERCEL_AI_GATEWAY, + userByok: null, + }); + mockedFetch.mockResolvedValue( + makeUpstreamResponse({ + text: 'hello world', + durationInSeconds: 1.5, + providerMetadata: { gateway: { cost: '0.00002' } }, + }) + ); + + const { POST } = await import('./route'); + const response = await POST( + makeRequest({ + model: 'openai/gpt-4o-mini-transcribe', + input_audio: { data: 'UklGRiQA', format: 'mp3' }, + provider: { only: ['openai'] }, + }) as never + ); + + expect(response.status).toBe(200); + const [url, init] = mockedFetch.mock.calls[0]; + expect(url).toBe('https://ai-gateway.vercel.sh/v4/ai/transcription-model'); + const headers = init?.headers as Headers; + expect(headers.get('ai-model-id')).toBe('openai/gpt-4o-mini-transcribe'); + expect(headers.get('ai-transcription-model-specification-version')).toBe('4'); + expect(JSON.parse(init?.body as string)).toMatchObject({ + audio: 'UklGRiQA', + mediaType: 'audio/mpeg', + providerOptions: { gateway: { only: ['openai'] } }, + }); + }); + it('rejects malformed transcription bodies before proxying', async () => { setUserAuth(); diff --git a/apps/web/src/app/api/openrouter/audio/transcriptions/route.ts b/apps/web/src/app/api/openrouter/audio/transcriptions/route.ts index 6eb767b78e..bc16c38795 100644 --- a/apps/web/src/app/api/openrouter/audio/transcriptions/route.ts +++ b/apps/web/src/app/api/openrouter/audio/transcriptions/route.ts @@ -26,10 +26,12 @@ import { emitApiMetricsForResponse } from '@/lib/ai-gateway/o11y/api-metrics.ser import { normalizeModelId } from '@/lib/ai-gateway/model-utils'; import { buildUpstreamBody, + buildVercelTranscriptionBody, extractTranscriptionPromptInfo, TranscriptionRequestSchema, } from '@/lib/ai-gateway/transcriptions/transcription-request'; import type { Provider } from '@/lib/ai-gateway/providers/types'; +import { mapModelIdToVercel } from '@/lib/ai-gateway/providers/vercel/mapModelIdToVercel'; export const maxDuration = 300; @@ -38,9 +40,10 @@ const PAID_MODEL_AUTH_REQUIRED = 'PAID_MODEL_AUTH_REQUIRED'; async function transcriptionProxyRequest(params: { body: Record; provider: Provider; + model: string; signal?: AbortSignal; }) { - const { body, provider, signal } = params; + const { body, provider, model, signal } = params; const headers = new Headers(); headers.set('Content-Type', 'application/json'); headers.set('Authorization', `Bearer ${provider.apiKey}`); @@ -49,10 +52,20 @@ async function transcriptionProxyRequest(params: { headers.set(key, value); } + if (provider.id === 'vercel') { + headers.set('ai-model-id', mapModelIdToVercel(model, false)); + headers.set('ai-transcription-model-specification-version', '4'); + } + const timeout = AbortSignal.timeout(10 * 60 * 1000); const combined = signal ? AbortSignal.any([signal, timeout]) : timeout; - return await fetch(`${provider.apiUrl}/audio/transcriptions`, { + const targetUrl = + provider.id === 'vercel' + ? `${provider.apiUrl.replace(/\/v1$/, '')}/v4/ai/transcription-model` + : `${provider.apiUrl}/audio/transcriptions`; + + return await fetch(targetUrl, { method: 'POST', headers, body: JSON.stringify(body), @@ -133,7 +146,11 @@ export async function POST(request: NextRequest): Promise 0 && provider.id === 'vercel') { + if (provider.id === 'vercel') { requestBodyParsed.model = mapModelIdToVercel(requestBodyParsed.model, false); } @@ -265,8 +281,12 @@ export async function POST(request: NextRequest): Promise; + provider?: OpenRouterProviderConfig; + providerOptions?: Record; input_type?: string; // Mistral-specific output_dtype?: string; diff --git a/apps/web/src/lib/ai-gateway/gateway-config.test.ts b/apps/web/src/lib/ai-gateway/gateway-config.test.ts index 67a6d5c382..3f26b756b9 100644 --- a/apps/web/src/lib/ai-gateway/gateway-config.test.ts +++ b/apps/web/src/lib/ai-gateway/gateway-config.test.ts @@ -8,20 +8,41 @@ import { describe('GatewayPercentageSchema', () => { test('accepts a numeric percentage', () => { - expect(GatewayPercentageSchema.parse({ vercel_routing_percentage: 25 })).toEqual({ - vercel_routing_percentage: 25, + expect( + GatewayPercentageSchema.parse({ + vercel_chat_routing_percentage: 25, + vercel_embeddings_routing_percentage: 30, + vercel_transcription_routing_percentage: 35, + }) + ).toEqual({ + vercel_chat_routing_percentage: 25, + vercel_embeddings_routing_percentage: 30, + vercel_transcription_routing_percentage: 35, }); }); test('accepts null (written when an admin clears the override)', () => { - expect(GatewayPercentageSchema.parse({ vercel_routing_percentage: null })).toEqual({ - vercel_routing_percentage: null, + expect( + GatewayPercentageSchema.parse({ + vercel_chat_routing_percentage: null, + vercel_embeddings_routing_percentage: null, + vercel_transcription_routing_percentage: null, + }) + ).toEqual({ + vercel_chat_routing_percentage: null, + vercel_embeddings_routing_percentage: null, + vercel_transcription_routing_percentage: null, }); }); test('rejects out-of-range values', () => { - expect(() => GatewayPercentageSchema.parse({ vercel_routing_percentage: 101 })).toThrow(); - expect(() => GatewayPercentageSchema.parse({ vercel_routing_percentage: -1 })).toThrow(); + expect(() => + GatewayPercentageSchema.parse({ + vercel_chat_routing_percentage: 101, + vercel_embeddings_routing_percentage: 0, + vercel_transcription_routing_percentage: 0, + }) + ).toThrow(); }); }); @@ -34,11 +55,15 @@ describe('GatewayConfigSchema', () => { updated_by_email: 'a@example.com', }); expect(parsed.note).toBeNull(); + expect(parsed.vercel_chat_routing_percentage).toBe(25); + expect(parsed.vercel_embeddings_routing_percentage).toBeNull(); }); test('round-trips a note', () => { const parsed = GatewayConfigSchema.parse({ - vercel_routing_percentage: 25, + vercel_chat_routing_percentage: 25, + vercel_embeddings_routing_percentage: 30, + vercel_transcription_routing_percentage: 35, updated_at: '2026-01-01T00:00:00.000Z', updated_by: 'u1', updated_by_email: 'a@example.com', @@ -46,25 +71,62 @@ describe('GatewayConfigSchema', () => { }); expect(parsed.note).toBe('Ramping down Vercel due to incident.'); }); + + test('keeps the legacy chat percentage for rolling deployments', () => { + const parsed = GatewayConfigSchema.parse({ + vercel_routing_percentage: 25, + vercel_chat_routing_percentage: 25, + vercel_embeddings_routing_percentage: 30, + vercel_transcription_routing_percentage: 35, + updated_at: '2026-01-01T00:00:00.000Z', + updated_by: 'u1', + updated_by_email: 'a@example.com', + note: null, + }); + + expect(parsed.vercel_routing_percentage).toBe(25); + }); }); describe('GatewayConfigInputSchema', () => { test('accepts a note alongside a percentage', () => { expect( - GatewayConfigInputSchema.parse({ vercel_routing_percentage: 75, note: 'Rollout stable' }) - ).toEqual({ vercel_routing_percentage: 75, note: 'Rollout stable' }); + GatewayConfigInputSchema.parse({ + vercel_chat_routing_percentage: 75, + vercel_embeddings_routing_percentage: 25, + vercel_transcription_routing_percentage: 10, + note: 'Rollout stable', + }) + ).toEqual({ + vercel_chat_routing_percentage: 75, + vercel_embeddings_routing_percentage: 25, + vercel_transcription_routing_percentage: 10, + note: 'Rollout stable', + }); }); test('accepts a null note', () => { - expect(GatewayConfigInputSchema.parse({ vercel_routing_percentage: null, note: null })).toEqual( - { vercel_routing_percentage: null, note: null } - ); + expect( + GatewayConfigInputSchema.parse({ + vercel_chat_routing_percentage: null, + vercel_embeddings_routing_percentage: null, + vercel_transcription_routing_percentage: null, + note: null, + }) + ).toEqual({ + vercel_chat_routing_percentage: null, + vercel_embeddings_routing_percentage: null, + vercel_transcription_routing_percentage: null, + note: null, + }); }); test('rejects notes longer than the maximum', () => { expect(() => GatewayConfigInputSchema.parse({ - vercel_routing_percentage: 50, + vercel_chat_routing_percentage: 50, + vercel_embeddings_routing_percentage: 50, + vercel_transcription_routing_percentage: 50, note: 'x'.repeat(NOTE_MAX_LENGTH + 1), }) ).toThrow(); diff --git a/apps/web/src/lib/ai-gateway/gateway-config.ts b/apps/web/src/lib/ai-gateway/gateway-config.ts index a5ef6d7db3..8e889d6bc3 100644 --- a/apps/web/src/lib/ai-gateway/gateway-config.ts +++ b/apps/web/src/lib/ai-gateway/gateway-config.ts @@ -2,24 +2,60 @@ import * as z from 'zod'; export const DEFAULT_VERCEL_PERCENTAGE = 50; +export const VERCEL_ROUTING_API_TYPES = ['chat', 'embeddings', 'transcription'] as const; +export type VercelRoutingApiType = (typeof VERCEL_ROUTING_API_TYPES)[number]; + +export const VERCEL_ROUTING_PERCENTAGE_FIELDS = { + chat: 'vercel_chat_routing_percentage', + embeddings: 'vercel_embeddings_routing_percentage', + transcription: 'vercel_transcription_routing_percentage', +} as const satisfies Record; + const vercelRoutingPercentage = z.number().int().min(0).max(100); export const NOTE_MAX_LENGTH = 500; const note = z.string().max(NOTE_MAX_LENGTH); -export const GatewayConfigSchema = z.object({ - vercel_routing_percentage: vercelRoutingPercentage.nullable(), +const gatewayConfigMetadata = { updated_at: z.string().nullable(), updated_by: z.string().nullable(), updated_by_email: z.string().nullable(), note: note.nullable().default(null), +}; + +const CurrentGatewayConfigSchema = z.object({ + vercel_routing_percentage: vercelRoutingPercentage.nullable().optional(), + vercel_chat_routing_percentage: vercelRoutingPercentage.nullable(), + vercel_embeddings_routing_percentage: vercelRoutingPercentage.nullable(), + vercel_transcription_routing_percentage: vercelRoutingPercentage.nullable(), + ...gatewayConfigMetadata, }); +const LegacyGatewayConfigSchema = z + .object({ + vercel_routing_percentage: vercelRoutingPercentage.nullable(), + ...gatewayConfigMetadata, + }) + .transform(config => ({ + vercel_chat_routing_percentage: config.vercel_routing_percentage, + vercel_embeddings_routing_percentage: null, + vercel_transcription_routing_percentage: null, + updated_at: config.updated_at, + updated_by: config.updated_by, + updated_by_email: config.updated_by_email, + note: config.note, + })); + +export const GatewayConfigSchema = z.union([CurrentGatewayConfigSchema, LegacyGatewayConfigSchema]); + export type GatewayConfig = z.infer; export const DEFAULT_GATEWAY_CONFIG: GatewayConfig = { vercel_routing_percentage: null, + vercel_chat_routing_percentage: null, + vercel_embeddings_routing_percentage: null, + vercel_transcription_routing_percentage: null, updated_at: null, updated_by: null, updated_by_email: null, @@ -29,16 +65,19 @@ export const DEFAULT_GATEWAY_CONFIG: GatewayConfig = { /** * Schema for parsing just the percentage from Redis (used on the hot path). * - * `vercel_routing_percentage` is nullable because clearing the override in - * the admin UI persists an explicit `null`. Callers should treat `null` as - * "no override, use DEFAULT_VERCEL_PERCENTAGE". + * Percentages are nullable because clearing an override in the admin UI persists + * an explicit `null`. Callers treat `null` as the default percentage. */ export const GatewayPercentageSchema = z.object({ - vercel_routing_percentage: vercelRoutingPercentage.nullable(), + vercel_chat_routing_percentage: vercelRoutingPercentage.nullable(), + vercel_embeddings_routing_percentage: vercelRoutingPercentage.nullable(), + vercel_transcription_routing_percentage: vercelRoutingPercentage.nullable(), }); /** Schema for the admin set-mutation input. */ export const GatewayConfigInputSchema = z.object({ - vercel_routing_percentage: vercelRoutingPercentage.nullable(), + vercel_chat_routing_percentage: vercelRoutingPercentage.nullable(), + vercel_embeddings_routing_percentage: vercelRoutingPercentage.nullable(), + vercel_transcription_routing_percentage: vercelRoutingPercentage.nullable(), note: note.nullable(), }); diff --git a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.test.ts b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.test.ts index 7537d93c68..33d7ae550b 100644 --- a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.test.ts +++ b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.test.ts @@ -691,6 +691,20 @@ describe('parseEmbeddingUsageFromResponse', () => { expect(result.cost_mUsd).toBe(50); }); + it('uses Vercel provider metadata cost when available', () => { + const response = JSON.stringify({ + object: 'list', + data: [], + model: 'openai/text-embedding-3-small', + usage: { prompt_tokens: 42, total_tokens: 42 }, + providerMetadata: { gateway: { cost: '0.00002', marketCost: '0.00003' } }, + }); + + const result = parseEmbeddingUsageFromResponse(response, 200); + + expect(result.cost_mUsd).toBe(30); + }); + it('should default to 0 cost when upstream cost field is absent', () => { const response = makeResponse({ usage: { prompt_tokens: 1000, total_tokens: 1000 }, @@ -817,6 +831,26 @@ describe('parseTranscriptionUsageFromResponse', () => { expect(result.generation_time).toBe(2.5); }); + it('uses Vercel provider metadata cost and duration', () => { + const result = parseTranscriptionUsageFromResponse( + JSON.stringify({ + text: 'hello world', + durationInSeconds: 3.5, + providerMetadata: { + gateway: { + cost: '0.00002', + marketCost: '0.00003', + }, + }, + }), + 200, + 'openai/whisper-1' + ); + + expect(result.cost_mUsd).toBe(30); + expect(result.generation_time).toBe(3.5); + }); + it('falls back to requested model when response model is absent', () => { const parsed = JSON.parse(makeResponse()); delete parsed.model; diff --git a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts index b14420d2c2..cb91284028 100644 --- a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts +++ b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts @@ -39,11 +39,16 @@ import type { MicrodollarUsageContext, MicrodollarUsageStats, PromptInfo, + VercelProviderMetaData, } from '@/lib/ai-gateway/processUsage.types'; import { detectContextOverflow } from '@/lib/ai-gateway/context-overflow'; import { KILO_AUTO_BALANCED_MODEL, KILO_AUTO_FREE_MODEL } from '@/lib/ai-gateway/auto-model'; import type { GatewayChatApiKind, ProviderId } from '@/lib/ai-gateway/providers/types'; -import { computeOpenRouterCostFields } from '@/lib/ai-gateway/processUsage.shared'; +import { + computeOpenRouterCostFields, + computeVercelCostMicrodollars, + extractVercelIsByok, +} from '@/lib/ai-gateway/processUsage.shared'; import { persistExperimentAttribution } from '@/lib/ai-gateway/experiments/persist'; export { proxyErrorTypeSchema, ProxyErrorType } from '@/lib/proxy-error-types'; import { ProxyErrorType } from '@/lib/proxy-error-types'; @@ -910,6 +915,7 @@ type EmbeddingResponse = { object: 'list'; model: string; usage: EmbeddingUsage; + providerMetadata?: VercelProviderMetaData; }; type TranscriptionUsage = { @@ -927,6 +933,8 @@ type TranscriptionResponse = { model?: string; text?: string; usage?: TranscriptionUsage; + durationInSeconds?: number; + providerMetadata?: VercelProviderMetaData; }; export function parseEmbeddingUsageFromResponse( @@ -935,8 +943,12 @@ export function parseEmbeddingUsageFromResponse( ): MicrodollarUsageStats { const json: EmbeddingResponse = JSON.parse(responseText); - // Upstream providers (OpenRouter, Vercel) include cost in USD → convert to microdollars. - const cost_mUsd = json.usage.cost != null ? toMicrodollars(json.usage.cost) : 0; + const vercelGateway = json.providerMetadata?.gateway; + const cost_mUsd = vercelGateway + ? computeVercelCostMicrodollars(vercelGateway) + : json.usage.cost != null + ? toMicrodollars(json.usage.cost) + : 0; return { messageId: json.id ?? null, @@ -949,7 +961,7 @@ export function parseEmbeddingUsageFromResponse( cacheHitTokens: 0, cacheWriteTokens: 0, cost_mUsd, - is_byok: null, + is_byok: extractVercelIsByok(vercelGateway), upstream_id: null, finish_reason: null, latency: null, @@ -977,16 +989,18 @@ export function parseTranscriptionUsageFromResponse( finish_reason: null, latency: null, moderation_latency: null, - generation_time: json.usage?.seconds ?? null, + generation_time: json.usage?.seconds ?? json.durationInSeconds ?? null, streamed: false, cancelled: false, status_code: statusCode, }; - const cost = computeOpenRouterCostFields( - json.usage ?? {}, - base, - 'transcription_usage_processing' - ); + const vercelGateway = json.providerMetadata?.gateway; + const cost = vercelGateway + ? { + cost_mUsd: computeVercelCostMicrodollars(vercelGateway), + is_byok: extractVercelIsByok(vercelGateway), + } + : computeOpenRouterCostFields(json.usage ?? {}, base, 'transcription_usage_processing'); return { ...base, diff --git a/apps/web/src/lib/ai-gateway/providers/get-provider.ts b/apps/web/src/lib/ai-gateway/providers/get-provider.ts index c892134d44..82f30912b3 100644 --- a/apps/web/src/lib/ai-gateway/providers/get-provider.ts +++ b/apps/web/src/lib/ai-gateway/providers/get-provider.ts @@ -1,5 +1,6 @@ import type { GatewayRequest } from '@/lib/ai-gateway/providers/openrouter/types'; import { shouldRouteToVercel } from '@/lib/ai-gateway/providers/vercel'; +import { isReasoningExplicitlyDisabled } from '@/lib/ai-gateway/providers/openrouter/request-helpers'; import { findKiloExclusiveModel, isKiloExclusiveModel } from '@/lib/ai-gateway/models'; import { CUSTOM_LLM_PREFIX } from '@/lib/ai-gateway/model-utils'; import { @@ -231,7 +232,14 @@ export async function getProvider(input: GetProviderInput): Promise { // 1. BYOK check — route through Vercel AI Gateway when user has their own key const userByok = await checkVercelBYOK(user, requestedModel, organizationId); @@ -262,13 +272,30 @@ export async function getEmbeddingProvider( return { provider: PROVIDERS.VERCEL_AI_GATEWAY, userByok }; } - // 2. All non-BYOK embedding requests go through OpenRouter + if (await shouldRouteToVercel(requestedModel, null, provider, randomSeed, 'embeddings')) { + return { provider: PROVIDERS.VERCEL_AI_GATEWAY, userByok: null }; + } + return { provider: PROVIDERS.OPENROUTER, userByok: null }; } -export async function getTranscriptionProvider(): Promise<{ +export async function getTranscriptionProvider( + requestedModel: string, + provider: GatewayRequest['body']['provider'] | undefined, + randomSeed: string +): Promise<{ provider: Provider; userByok: BYOKResult[] | null; }> { - return { provider: PROVIDERS.OPENROUTER, userByok: null }; + const routeToVercel = await shouldRouteToVercel( + requestedModel, + null, + provider, + randomSeed, + 'transcription' + ); + return { + provider: routeToVercel ? PROVIDERS.VERCEL_AI_GATEWAY : PROVIDERS.OPENROUTER, + userByok: null, + }; } diff --git a/apps/web/src/lib/ai-gateway/providers/getEmbeddingProvider.test.ts b/apps/web/src/lib/ai-gateway/providers/getEmbeddingProvider.test.ts index 9da04edb89..eb61b722f7 100644 --- a/apps/web/src/lib/ai-gateway/providers/getEmbeddingProvider.test.ts +++ b/apps/web/src/lib/ai-gateway/providers/getEmbeddingProvider.test.ts @@ -11,12 +11,17 @@ import { getBYOKforOrganization, } from '@/lib/ai-gateway/byok'; import type { User } from '@kilocode/db/schema'; +import { shouldRouteToVercel } from '@/lib/ai-gateway/providers/vercel'; jest.mock('@/lib/ai-gateway/byok'); +jest.mock('@/lib/ai-gateway/providers/vercel', () => ({ + shouldRouteToVercel: jest.fn(), +})); const mockedGetModelUserByokProviders = getModelUserByokProviders as jest.Mock; const mockedGetBYOKforUser = getBYOKforUser as jest.Mock; const mockedGetBYOKforOrganization = getBYOKforOrganization as jest.Mock; +const mockedShouldRouteToVercel = jest.mocked(shouldRouteToVercel); function createTestUser(overrides: Partial = {}): User { return { @@ -34,6 +39,7 @@ describe('getEmbeddingProvider', () => { mockedGetModelUserByokProviders.mockClear().mockResolvedValue([]); mockedGetBYOKforUser.mockClear().mockResolvedValue(null); mockedGetBYOKforOrganization.mockClear().mockResolvedValue(null); + mockedShouldRouteToVercel.mockReset().mockResolvedValue(false); }); it('should route all non-BYOK models to PROVIDERS.OPENROUTER', async () => { @@ -44,7 +50,7 @@ describe('getEmbeddingProvider', () => { 'openai/text-embedding-3-small', 'google/text-embedding-004', ]) { - const result = await getEmbeddingProvider(model, user, undefined); + const result = await getEmbeddingProvider(model, user, undefined, undefined, user.id); expect(result.provider.id).toBe('openrouter'); expect(result.provider).toBe(PROVIDERS.OPENROUTER); expect(result.userByok).toBeNull(); @@ -58,7 +64,13 @@ describe('getEmbeddingProvider', () => { mockedGetModelUserByokProviders.mockResolvedValue(['openai']); mockedGetBYOKforUser.mockResolvedValue(mockByokResult); - const result = await getEmbeddingProvider('openai/text-embedding-3-small', user, undefined); + const result = await getEmbeddingProvider( + 'openai/text-embedding-3-small', + user, + undefined, + undefined, + user.id + ); expect(result.provider.id).toBe('vercel'); expect(result.provider).toBe(PROVIDERS.VERCEL_AI_GATEWAY); @@ -72,7 +84,13 @@ describe('getEmbeddingProvider', () => { mockedGetModelUserByokProviders.mockResolvedValue(['mistral']); mockedGetBYOKforOrganization.mockResolvedValue(mockByokResult); - const result = await getEmbeddingProvider('mistralai/mistral-embed-2312', user, 'org-123'); + const result = await getEmbeddingProvider( + 'mistralai/mistral-embed-2312', + user, + 'org-123', + undefined, + user.id + ); expect(result.provider.id).toBe('vercel'); expect(result.userByok).toBe(mockByokResult); @@ -83,7 +101,13 @@ describe('getEmbeddingProvider', () => { it('should skip BYOK check for anonymous users', async () => { const anonUser = createAnonymousContext('127.0.0.1'); - const result = await getEmbeddingProvider('openai/text-embedding-3-small', anonUser, undefined); + const result = await getEmbeddingProvider( + 'openai/text-embedding-3-small', + anonUser, + undefined, + undefined, + anonUser.id + ); expect(result.provider.id).toBe('openrouter'); expect(result.userByok).toBeNull(); @@ -95,18 +119,63 @@ describe('getEmbeddingProvider', () => { mockedGetModelUserByokProviders.mockResolvedValue(['openai']); mockedGetBYOKforUser.mockResolvedValue(null); - const result = await getEmbeddingProvider('mistralai/codestral-embed-2505', user, undefined); + const result = await getEmbeddingProvider( + 'mistralai/codestral-embed-2505', + user, + undefined, + undefined, + user.id + ); expect(result.provider.id).toBe('openrouter'); expect(result.userByok).toBeNull(); }); + + it('routes eligible non-BYOK embeddings to Vercel', async () => { + const user = createTestUser(); + mockedShouldRouteToVercel.mockResolvedValue(true); + + const result = await getEmbeddingProvider( + 'openai/text-embedding-3-small', + user, + undefined, + { only: ['openai'] }, + user.id + ); + + expect(result.provider).toBe(PROVIDERS.VERCEL_AI_GATEWAY); + expect(result.userByok).toBeNull(); + expect(mockedShouldRouteToVercel).toHaveBeenCalledWith( + 'openai/text-embedding-3-small', + null, + { only: ['openai'] }, + user.id, + 'embeddings' + ); + }); }); describe('getTranscriptionProvider', () => { + beforeEach(() => { + mockedShouldRouteToVercel.mockReset().mockResolvedValue(false); + }); + it('routes transcription requests to OpenRouter', async () => { - const result = await getTranscriptionProvider(); + const result = await getTranscriptionProvider('openai/whisper-1', undefined, 'user-123'); expect(result.provider).toBe(PROVIDERS.OPENROUTER); expect(result.userByok).toBeNull(); }); + + it('routes eligible transcription requests to Vercel', async () => { + mockedShouldRouteToVercel.mockResolvedValue(true); + + const result = await getTranscriptionProvider( + 'openai/whisper-1', + { only: ['openai'] }, + 'user-123' + ); + + expect(result.provider).toBe(PROVIDERS.VERCEL_AI_GATEWAY); + }); }); diff --git a/apps/web/src/lib/ai-gateway/providers/vercel/index.ts b/apps/web/src/lib/ai-gateway/providers/vercel/index.ts index 4a8064624f..a2267b44d2 100644 --- a/apps/web/src/lib/ai-gateway/providers/vercel/index.ts +++ b/apps/web/src/lib/ai-gateway/providers/vercel/index.ts @@ -17,43 +17,48 @@ import { redisClient } from '@/lib/redis'; import { createCachedFetch } from '@/lib/cached-fetch'; import { GatewayPercentageSchema, + GatewayConfigSchema, DEFAULT_VERCEL_PERCENTAGE, + VERCEL_ROUTING_PERCENTAGE_FIELDS, + type VercelRoutingApiType, } from '@/lib/ai-gateway/gateway-config'; import { VERCEL_ROUTING_REDIS_KEY } from '@/lib/redis-keys'; import { getRandomNumber } from '@/lib/ai-gateway/getRandomNumber'; import { getCachedVercelInferenceProviderIdsForModel, getVercelModels, + getVercelModelsMetadata, } from '@/lib/ai-gateway/providers/gateway-models-cache'; import type { AnthropicProviderOptions } from '@ai-sdk/anthropic'; import { isDeepseekModel } from '@/lib/ai-gateway/providers/deepseek'; import type { KiloExclusiveModel } from '@/lib/ai-gateway/providers/kilo-exclusive-model'; -const getVercelRoutingPercentage = createCachedFetch( +const getVercelRoutingPercentages = createCachedFetch( async () => { const raw = await redisClient.get(VERCEL_ROUTING_REDIS_KEY); - if (!raw) return DEFAULT_VERCEL_PERCENTAGE; - const { vercel_routing_percentage } = GatewayPercentageSchema.parse(JSON.parse(raw)); - return vercel_routing_percentage ?? DEFAULT_VERCEL_PERCENTAGE; + if (!raw) return null; + return GatewayPercentageSchema.parse(GatewayConfigSchema.parse(JSON.parse(raw))); }, 600_000, - DEFAULT_VERCEL_PERCENTAGE + null ); -function hasOpenRouterExclusiveProviderOptions(request: GatewayRequest) { +function hasOpenRouterExclusiveProviderOptions( + provider: GatewayRequest['body']['provider'] | undefined +) { // Vercel has options for disallowPromptTraining and zeroDataRetention // but they are not enforced for BYOK requests, which is most requests, // so we don't use them for now. // https://vercel.com/docs/ai-gateway/security-and-compliance/disallow-prompt-training - if (request.body.provider?.data_collection === 'deny') { + if (provider?.data_collection === 'deny') { console.debug('[hasOpenRouterExclusiveProviderOptions] has data_collection==deny'); return true; } - if (request.body.provider?.zdr) { + if (provider?.zdr) { console.debug('[hasOpenRouterExclusiveProviderOptions] has zdr'); return true; } - if ((request.body.provider?.ignore?.length ?? 0) > 0) { + if ((provider?.ignore?.length ?? 0) > 0) { console.debug('[hasOpenRouterExclusiveProviderOptions] has ignore'); return true; } @@ -76,10 +81,12 @@ export function hasCompatibleVercelInferenceProvider( export async function shouldRouteToVercel( requestedModel: string, kiloExclusiveModel: KiloExclusiveModel | null, - request: GatewayRequest, - randomSeed: string + provider: GatewayRequest['body']['provider'] | undefined, + randomSeed: string, + apiType: VercelRoutingApiType, + reasoningExplicitlyDisabled = false ) { - if (hasOpenRouterExclusiveProviderOptions(request)) { + if (hasOpenRouterExclusiveProviderOptions(provider)) { console.debug( '[shouldRouteToVercel] not routing to Vercel because of unsupported provider options' ); @@ -95,23 +102,31 @@ export async function shouldRouteToVercel( } console.debug('[shouldRouteToVercel] randomizing user to either OpenRouter or Vercel'); - const routingPercentage = await getVercelRoutingPercentage(); + const percentages = await getVercelRoutingPercentages(); + const routingPercentage = + percentages?.[VERCEL_ROUTING_PERCENTAGE_FIELDS[apiType]] ?? DEFAULT_VERCEL_PERCENTAGE; - const passedRandomization = - getRandomNumber('vercel_routing_' + randomSeed, 100) < routingPercentage; + const randomizationKey = + apiType === 'chat' ? `vercel_routing_${randomSeed}` : `vercel_${apiType}_routing_${randomSeed}`; + const passedRandomization = getRandomNumber(randomizationKey, 100) < routingPercentage; if (!passedRandomization) { return false; } - const vercelModels = await getVercelModels(); - const vercelModelId = mapModelIdToVercel(requestedModel, isReasoningExplicitlyDisabled(request)); - if (!vercelModels.has(vercelModelId)) { + const vercelModelId = mapModelIdToVercel(requestedModel, reasoningExplicitlyDisabled); + const modelIsAvailable = + apiType === 'embeddings' + ? (await getVercelModelsMetadata())[vercelModelId]?.type === 'embedding' + : apiType === 'transcription' + ? (await getVercelModelsMetadata())[vercelModelId]?.type === 'transcription' + : (await getVercelModels()).has(vercelModelId); + if (!modelIsAvailable) { console.debug(`[shouldRouteToVercel] model not found in Vercel model list`); return false; } - const only = request.body.provider?.only; + const only = provider?.only; if (only) { const vercelInferenceProviders = await getCachedVercelInferenceProviderIdsForModel(vercelModelId); @@ -126,15 +141,17 @@ export async function shouldRouteToVercel( return true; } -function convertProviderOptions(requestToMutate: GatewayRequest): VercelProviderConfig { - const provider = requestToMutate.body.provider; +export function convertProviderOptions( + provider: GatewayRequest['body']['provider'] | undefined, + models?: string[] +): VercelProviderConfig { return { gateway: { only: provider?.only?.map(p => openRouterToVercelInferenceProviderId(p)), order: provider?.order?.map(p => openRouterToVercelInferenceProviderId(p)), zeroDataRetention: provider?.zdr, disallowPromptTraining: provider?.data_collection === 'deny', - models: requestToMutate.body.models, + models, }, }; } @@ -224,7 +241,10 @@ export function applyVercelSettings( }, }; } else { - requestToMutate.body.providerOptions = convertProviderOptions(requestToMutate); + requestToMutate.body.providerOptions = convertProviderOptions( + requestToMutate.body.provider, + requestToMutate.body.models + ); } if (requestToMutate.body.providerOptions) { diff --git a/apps/web/src/lib/ai-gateway/transcriptions/transcription-request.test.ts b/apps/web/src/lib/ai-gateway/transcriptions/transcription-request.test.ts index 93ee9b66eb..2c014c0321 100644 --- a/apps/web/src/lib/ai-gateway/transcriptions/transcription-request.test.ts +++ b/apps/web/src/lib/ai-gateway/transcriptions/transcription-request.test.ts @@ -1,6 +1,7 @@ import { describe, it, expect } from '@jest/globals'; import { buildUpstreamBody, + buildVercelTranscriptionBody, extractTranscriptionPromptInfo, TranscriptionRequestSchema, } from './transcription-request'; @@ -47,6 +48,28 @@ describe('buildUpstreamBody', () => { }); }); +describe('buildVercelTranscriptionBody', () => { + it('converts OpenRouter audio and routing fields to the Vercel contract', () => { + const body = TranscriptionRequestSchema.parse({ + model: 'openai/gpt-4o-mini-transcribe', + input_audio: { data: 'UklGRiQA', format: 'mp3' }, + language: 'en', + temperature: 0, + provider: { only: ['openai'] }, + safety_identifier: 'hash-abc', + }); + + expect(buildVercelTranscriptionBody(body)).toMatchObject({ + audio: 'UklGRiQA', + mediaType: 'audio/mpeg', + providerOptions: { + gateway: { only: ['openai'] }, + openai: { language: 'en', temperature: 0 }, + }, + }); + }); +}); + describe('extractTranscriptionPromptInfo', () => { it('describes audio format without including encoded audio', () => { const body = TranscriptionRequestSchema.parse({ diff --git a/apps/web/src/lib/ai-gateway/transcriptions/transcription-request.ts b/apps/web/src/lib/ai-gateway/transcriptions/transcription-request.ts index a1c85f71b1..a08e4eecc3 100644 --- a/apps/web/src/lib/ai-gateway/transcriptions/transcription-request.ts +++ b/apps/web/src/lib/ai-gateway/transcriptions/transcription-request.ts @@ -1,4 +1,5 @@ import { z } from 'zod'; +import { convertProviderOptions } from '@/lib/ai-gateway/providers/vercel'; export const TranscriptionInputAudioSchema = z.object({ data: z.string().min(1), @@ -11,7 +12,17 @@ export const TranscriptionRequestSchema = z input_audio: TranscriptionInputAudioSchema, language: z.string().min(1).optional(), temperature: z.number().optional(), - provider: z.record(z.string(), z.unknown()).optional(), + provider: z + .object({ + order: z.array(z.string()).optional(), + only: z.array(z.string()).optional(), + ignore: z.array(z.string()).optional(), + data_collection: z.enum(['allow', 'deny']).optional(), + zdr: z.boolean().optional(), + require_parameters: z.boolean().optional(), + }) + .passthrough() + .optional(), safety_identifier: z.string().optional(), user: z.string().optional(), }) @@ -23,6 +34,33 @@ export function buildUpstreamBody(body: TranscriptionRequest): Record = { + mp3: 'audio/mpeg', + m4a: 'audio/mp4', + ogg: 'audio/ogg', + wav: 'audio/wav', + webm: 'audio/webm', +}; + +export function buildVercelTranscriptionBody(body: TranscriptionRequest): Record { + const mediaType = body.input_audio.format.includes('/') + ? body.input_audio.format + : (AUDIO_MEDIA_TYPES[body.input_audio.format.toLowerCase()] ?? + `audio/${body.input_audio.format.toLowerCase()}`); + const providerOptions = { + ...convertProviderOptions(body.provider), + ...(body.language !== undefined || body.temperature !== undefined + ? { openai: { language: body.language, temperature: body.temperature } } + : {}), + }; + + return { + audio: body.input_audio.data, + mediaType, + ...(providerOptions.gateway && { providerOptions }), + }; +} + export function extractTranscriptionPromptInfo(body: TranscriptionRequest) { const format = body.input_audio.format.slice(0, 100); const language = body.language ? ` language=${body.language.slice(0, 32)}` : ''; diff --git a/apps/web/src/routers/admin/gateway-config-router.ts b/apps/web/src/routers/admin/gateway-config-router.ts index 53990e427b..05446b0d97 100644 --- a/apps/web/src/routers/admin/gateway-config-router.ts +++ b/apps/web/src/routers/admin/gateway-config-router.ts @@ -26,7 +26,11 @@ export const adminGatewayConfigRouter = createTRPCRouter({ set: adminProcedure.input(GatewayConfigInputSchema).mutation(async ({ input, ctx }) => { const config: GatewayConfig = { - vercel_routing_percentage: input.vercel_routing_percentage, + // Keep the legacy field while old instances may still read this Redis value. + vercel_routing_percentage: input.vercel_chat_routing_percentage, + vercel_chat_routing_percentage: input.vercel_chat_routing_percentage, + vercel_embeddings_routing_percentage: input.vercel_embeddings_routing_percentage, + vercel_transcription_routing_percentage: input.vercel_transcription_routing_percentage, updated_at: new Date().toISOString(), updated_by: ctx.user.id, updated_by_email: ctx.user.google_user_email, diff --git a/packages/db/src/schema-types.ts b/packages/db/src/schema-types.ts index c145d03033..3c8a2f4586 100644 --- a/packages/db/src/schema-types.ts +++ b/packages/db/src/schema-types.ts @@ -1663,7 +1663,7 @@ export type CustomLlmDefinition = z.infer; export const ModelSchema = z.object({ id: z.string(), name: z.string(), - type: z.enum(['language', 'embedding', 'image']).optional().catch(undefined), + type: z.enum(['language', 'embedding', 'image', 'transcription']).optional().catch(undefined), }); export const ModelsSchema = z.object({ data: z.array(ModelSchema) });