From edfd67a83cd128662faa715d5d5f0bc2e1cdf01c Mon Sep 17 00:00:00 2001 From: Neal Shah <30693865+ProgrammerIn-wonderland@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:11:04 -0400 Subject: [PATCH] fix PUT-1444 (#3515) --- .../ChatCompletionDriver.routing.test.ts | 193 ++++++++++++++++++ .../drivers/ai-chat/ChatCompletionDriver.ts | 37 +--- .../ai-chat/utils/modelRouting.test.ts | 135 ++++++++++++ .../drivers/ai-chat/utils/modelRouting.ts | 67 ++++++ 4 files changed, 406 insertions(+), 26 deletions(-) create mode 100644 src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts create mode 100644 src/backend/drivers/ai-chat/utils/modelRouting.test.ts create mode 100644 src/backend/drivers/ai-chat/utils/modelRouting.ts diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts new file mode 100644 index 000000000..db07d7161 --- /dev/null +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts @@ -0,0 +1,193 @@ +/** + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify it under the + * terms of the GNU Affero General Public License as published by the Free + * Software Foundation, either version 3 of the License, or (at your option) any + * later version. + * + * This program is distributed in the hope that it will be useful, but WITHOUT + * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS + * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more + * details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see + * [https://www.gnu.org/licenses/](https://www.gnu.org/licenses/). + */ + +/** + * Which provider actually serves a model that several providers advertise. + * + * Lives apart from ChatCompletionDriver.test.ts because `vi.mock` is + * file-scoped and hoisted — mocking the OpenAI SDK and axios here would + * otherwise leak into every test in that file. Both mocks sit at the real + * network egress points, so the driver's registration, model-map build and + * resolution all run for real. + */ + +import { + afterAll, + beforeAll, + describe, + expect, + it, + vi, + type MockInstance, +} from 'vitest'; + +import { HttpError } from '../../core/http/HttpError.js'; +import { PuterServer } from '../../server.js'; +import { setupTestServer } from '../../testUtil.js'; +import { kv } from '../../util/kvSingleton.js'; +import { withTestActor } from '../integrationTestUtil.js'; +import { ChatCompletionDriver } from './ChatCompletionDriver.js'; + +// -- OpenAI SDK mock ------------------------------------------------ +// Gemini reaches Google through `new openai.OpenAI()` (default export) and +// the gateway through the named one, so both must resolve to the same ctor. + +const { createMock } = vi.hoisted(() => ({ createMock: vi.fn() })); + +vi.mock('openai', () => { + const OpenAICtor = vi.fn().mockImplementation(function ( + this: Record, + ) { + this.chat = { completions: { create: createMock } }; + }); + return { OpenAI: OpenAICtor, default: { OpenAI: OpenAICtor } }; +}); + +// -- axios mock (gateway model catalog) ----------------------------- + +const { axiosRequestMock } = vi.hoisted(() => ({ axiosRequestMock: vi.fn() })); + +vi.mock('axios', () => ({ + default: { request: axiosRequestMock }, + request: axiosRequestMock, +})); + +// -- Harness -------------------------------------------------------- + +let server: PuterServer; +let driver: ChatCompletionDriver; + +const INFRON_KV_KEY = 'infronChat:models'; + +// Google lists gemini-2.5-flash input at $0.30/MTok. The gateway quotes a +// floor price across its upstream routes, so it undercuts — which is exactly +// the condition that used to hand it the traffic. +const GATEWAY_CATALOG = [ + { + id: 'google/gemini-2.5-flash', + display_name: 'Google: Gemini 2.5 Flash', + category_type: 'LLM', + supported_endpoint_types: ['openai'], + context_length: 1_048_576, + max_output_tokens: 65_536, + min_prompt_price: 0.15, + min_completion_price: 1.0, + }, + { + // Only the gateway carries this one — no first-party counterpart. + id: 'google/gemini-2.5-flash-image-preview', + display_name: 'Google: Gemini 2.5 Flash Image Preview', + category_type: 'LLM', + supported_endpoint_types: ['openai'], + context_length: 32_768, + max_output_tokens: 8_192, + min_prompt_price: 0.3, + min_completion_price: 2.5, + }, +]; + +beforeAll(async () => { + server = await setupTestServer(); + kv.del?.(INFRON_KV_KEY); + axiosRequestMock.mockResolvedValue({ data: { data: GATEWAY_CATALOG } }); + + // Built once, not per-test: `#buildModelMap` mutates the catalogs + // providers hand back (lowercasing ids, pushing `puterId` onto the + // shared `aliases` array), and GeminiChatProvider returns its + // module-level GEMINI_MODELS by reference. + driver = new ChatCompletionDriver( + { + providers: { + gemini: { apiKey: 'test-key' }, + infron: { apiKey: 'test-key' }, + ollama: { enabled: false }, + }, + } as never, + server.clients, + server.stores, + server.services, + ); + driver.onServerStart(); + // `onServerStart` doesn't await `#buildModelMap`, and the gateway + // catalog resolves on a microtask — poll until both providers land. + for (let i = 0; i < 200; i++) { + const ids = await driver.list(); + if (ids.some((id) => id.startsWith('infron:'))) break; + await new Promise((r) => setTimeout(r, 5)); + } +}); + +afterAll(async () => { + await server?.shutdown(); +}); + +/** + * Route a request and report who was tried, in order. Forcing the upstream to + * reject is what makes the whole chain observable: the driver records every + * attempt on the thrown error, and `attempts[0]` is who it chose first. + */ +const attemptsFor = async (model: string) => { + createMock.mockRejectedValue(new Error('upstream down')); + let caught: HttpError | undefined; + try { + await withTestActor(() => + driver.complete({ + model, + messages: [{ role: 'user', content: 'hi' }], + }), + ); + } catch (e) { + caught = e as HttpError; + } + expect(caught).toBeInstanceOf(HttpError); + return (caught as unknown as { fields: { attempts: { model: string; provider: string }[] } }) + .fields.attempts; +}; + +describe('ChatCompletionDriver gemini routing', () => { + it('serves gemini models from Google, with the gateway only as fallback', async () => { + const attempts = await attemptsFor('gemini-2.5-flash'); + + expect(attempts[0]).toMatchObject({ + provider: 'gemini', + model: 'gemini-2.5-flash', + }); + expect(attempts[1]).toMatchObject({ + provider: 'infron', + model: 'infron:google/gemini-2.5-flash', + }); + }); + + it('routes the prefixed and puterId forms to Google too', async () => { + for (const alias of [ + 'google/gemini-2.5-flash', + 'google:google/gemini-2.5-flash', + ]) { + const attempts = await attemptsFor(alias); + expect(attempts[0].provider).toBe('gemini'); + } + }); + + it('still routes models only the gateway carries to the gateway', async () => { + const attempts = await attemptsFor('google/gemini-2.5-flash-image-preview'); + + expect(attempts[0]).toMatchObject({ provider: 'infron' }); + }); +}); diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts index 5d6dca9c3..9c2c49c57 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts @@ -61,6 +61,10 @@ import { normalize_messages, normalize_single_message, } from './utils/Messages.js'; +import { + AGGREGATOR_PROVIDERS, + compareModelPreference, +} from './utils/modelRouting.js'; import { AIChatStream } from './utils/Streaming.js'; const MAX_FALLBACKS = 4; // includes first attempt @@ -1019,16 +1023,9 @@ export class ChatCompletionDriver extends PuterDriver { // -- Model map --------------------------------------------------- async #buildModelMap() { - const AGGREGATORS = new Set([ - 'together-ai', - 'openrouter', - 'infron', - 'neuralwatt', - ]); - for (const providerName in this.#providers) { const provider = this.#providers[providerName]; - const isAggregator = AGGREGATORS.has(providerName); + const isAggregator = AGGREGATOR_PROVIDERS.has(providerName); for (const model of await provider.models()) { model.id = model.id.trim().toLowerCase(); @@ -1058,8 +1055,11 @@ export class ChatCompletionDriver extends PuterDriver { existing !== this.#modelIdMap[model.id] ) { if (existing.some((m) => m.provider === 'gemini')) { - // Gemini exception — let the aggregator - // entry through. + // Gemini is the one vendor whose resold + // duplicates we keep, so a Google outage has + // somewhere to fall back to. Ranking (see + // `compareModelPreference`) keeps the direct + // provider ahead of them. continue; } skip = true; @@ -1098,22 +1098,7 @@ export class ChatCompletionDriver extends PuterDriver { } } - // Sort: together-ai always last; then cheapest input-cost - // first; ties break by shorter id (usually the official - // name over a long aggregator-qualified one). - this.#modelIdMap[model.id].sort((a, b) => { - const aAgg = a.provider === 'together-ai'; - const bAgg = b.provider === 'together-ai'; - if (aAgg !== bAgg) return aAgg ? 1 : -1; - const aCost = a.costs[ - (a.input_cost_key as string) || 'input_tokens' - ] as number; - const bCost = b.costs[ - (b.input_cost_key as string) || 'input_tokens' - ] as number; - if (aCost === bCost) return a.id.length - b.id.length; - return aCost - bCost; - }); + this.#modelIdMap[model.id].sort(compareModelPreference); } } } diff --git a/src/backend/drivers/ai-chat/utils/modelRouting.test.ts b/src/backend/drivers/ai-chat/utils/modelRouting.test.ts new file mode 100644 index 000000000..2073364d4 --- /dev/null +++ b/src/backend/drivers/ai-chat/utils/modelRouting.test.ts @@ -0,0 +1,135 @@ +/** + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { describe, expect, it } from 'vitest'; +import { GEMINI_MODELS } from '../providers/gemini/models.js'; +import type { IChatModel } from '../types.js'; +import { compareModelPreference } from './modelRouting.js'; + +// `#buildModelMap` mutates the catalogs providers hand back, and +// `GeminiChatProvider.models()` returns the module-level `GEMINI_MODELS` by +// reference — clone so these fixtures can't be perturbed by another suite. +const geminiModel = (id: string, provider = 'gemini'): IChatModel => { + const found = GEMINI_MODELS.find((m) => m.id === id); + if (!found) throw new Error(`no such gemini model: ${id}`); + return { ...structuredClone(found), provider }; +}; + +// Mirrors how the reseller providers coerce a catalog entry: `:` on +// the id, `input_cost_key: 'prompt'`, and prices as microcents per token. +const resoldModel = ( + provider: string, + catalogId: string, + promptCost: number, +): IChatModel => + ({ + id: `${provider}:${catalogId}`, + name: `${catalogId} (${provider})`, + aliases: [catalogId, catalogId.split('/').slice(1).join('/')], + costs_currency: 'usd-cents', + input_cost_key: 'prompt', + output_cost_key: 'completion', + costs: { tokens: 1_000_000, prompt: promptCost, completion: 100 }, + provider, + }) as IChatModel; + +const winner = (...models: IChatModel[]) => + [...models].sort(compareModelPreference)[0]; + +describe('compareModelPreference', () => { + it('serves the vendor directly even when a reseller quotes a lower price', () => { + // Google lists gemini-2.5-flash input at 30 microcents/token; the + // gateway advertises a floor price across its upstream routes and + // undercuts it. Price must not decide who serves the request. + const direct = geminiModel('gemini-2.5-flash'); + const resold = resoldModel('infron', 'google/gemini-2.5-flash', 15); + + expect(direct.costs.prompt_tokens).toBeGreaterThan( + resold.costs.prompt as number, + ); + expect(winner(resold, direct).provider).toBe('gemini'); + expect(winner(direct, resold).provider).toBe('gemini'); + }); + + it('ranks every reseller behind the vendor, not just one of them', () => { + const direct = geminiModel('gemini-3.1-pro-preview'); + const bucket = [ + resoldModel('openrouter', 'google/gemini-3.1-pro-preview', 90), + resoldModel('infron', 'google/gemini-3.1-pro-preview', 80), + resoldModel('together-ai', 'google/gemini-3.1-pro-preview', 70), + resoldModel('neuralwatt', 'google/gemini-3.1-pro-preview', 60), + direct, + ]; + + expect( + bucket.sort(compareModelPreference).map((m) => m.provider), + ).toEqual([ + 'gemini', + 'neuralwatt', + 'infron', + 'openrouter', + 'together-ai', + ]); + }); + + it('keeps together-ai strictly behind the other resellers', () => { + // A flat direct/reseller split would let these tie and re-order by + // price; together-ai stays last regardless of how cheap it quotes. + const together = resoldModel('together-ai', 'meta/llama-4', 1); + const openrouter = resoldModel('openrouter', 'meta/llama-4', 500); + + expect(winner(together, openrouter).provider).toBe('openrouter'); + }); + + it('still orders two direct providers by cheapest input cost', () => { + const cheap = geminiModel('gemini-2.5-flash-lite'); + const pricey = geminiModel('gemini-2.5-pro'); + + expect(cheap.costs.prompt_tokens).toBeLessThan( + pricey.costs.prompt_tokens as number, + ); + expect(winner(pricey, cheap).id).toBe('gemini-2.5-flash-lite'); + }); + + it('breaks price ties on the shorter id', () => { + const short = { ...geminiModel('gemini-2.5-flash'), id: 'gemini-x' }; + const long = { + ...geminiModel('gemini-2.5-flash'), + id: 'some-vendor/gemini-x-2025-preview', + provider: 'azure-openai', + }; + + expect(winner(long, short).id).toBe('gemini-x'); + }); + + it('leaves a reseller serving models no vendor provider carries', () => { + // The image-preview models are absent from GEMINI_MODELS, so the + // gateway is the only route and must stay the winner. + expect( + GEMINI_MODELS.some((m) => m.id === 'gemini-2.5-flash-image-preview'), + ).toBe(false); + + const onlyRoute = resoldModel( + 'infron', + 'google/gemini-2.5-flash-image-preview', + 20, + ); + expect(winner(onlyRoute).provider).toBe('infron'); + }); +}); diff --git a/src/backend/drivers/ai-chat/utils/modelRouting.ts b/src/backend/drivers/ai-chat/utils/modelRouting.ts new file mode 100644 index 000000000..3c1635ddb --- /dev/null +++ b/src/backend/drivers/ai-chat/utils/modelRouting.ts @@ -0,0 +1,67 @@ +/** + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify it under the + * terms of the GNU Affero General Public License as published by the Free + * Software Foundation, either version 3 of the License, or (at your option) any + * later version. + * + * This program is distributed in the hope that it will be useful, but WITHOUT + * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS + * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more + * details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see + * [https://www.gnu.org/licenses/](https://www.gnu.org/licenses/). + */ + +import type { IChatModel } from '../types.js'; + +/** + * Providers that resell other vendors' models rather than serving their own. + * Their catalogs duplicate models we already reach directly, so their entries + * are kept only as fallback routes. + */ +export const AGGREGATOR_PROVIDERS = new Set([ + 'together-ai', + 'openrouter', + 'infron', + 'neuralwatt', +]); + +// Lower rank is served first. `together-ai` sits behind the other resellers — +// a pre-existing guarantee this ranking extends rather than replaces. +const providerRank = (provider?: string): number => { + if (provider === 'together-ai') return 2; + if (provider && AGGREGATOR_PROVIDERS.has(provider)) return 1; + return 0; +}; + +/** + * Orders the candidates that share a model bucket; the first one gets served. + * + * Direct vendors outrank resellers regardless of quoted price. Resellers + * advertise a floor price across their upstream routes, so on price alone they + * undercut the vendor's list price and capture traffic for models we hold a + * direct integration for. Within a rank, cheapest input cost wins and ties + * break by shorter id — usually the official name over a qualified one. + */ +export const compareModelPreference = ( + a: IChatModel, + b: IChatModel, +): number => { + const rankDiff = providerRank(a.provider) - providerRank(b.provider); + if (rankDiff !== 0) return rankDiff; + + const aCost = a.costs[ + (a.input_cost_key as string) || 'input_tokens' + ] as number; + const bCost = b.costs[ + (b.input_cost_key as string) || 'input_tokens' + ] as number; + if (aCost === bCost) return a.id.length - b.id.length; + return aCost - bCost; +};