diff --git a/apps/web/src/lib/ai-gateway/auto-model/resolution.ts b/apps/web/src/lib/ai-gateway/auto-model/resolution.ts index c82a69ccb7..aec50da220 100644 --- a/apps/web/src/lib/ai-gateway/auto-model/resolution.ts +++ b/apps/web/src/lib/ai-gateway/auto-model/resolution.ts @@ -29,11 +29,11 @@ import { ORG_AUTO_MODEL, } from '@/lib/ai-gateway/auto-model'; import { userIsWithinFirstKiloClawInstanceWindow } from '@/lib/kiloclaw/setup-promo'; -import { getRandomNumber } from '@/lib/ai-gateway/getRandomNumber'; import { autoFreeModels, findKiloExclusiveModel, isKiloExclusiveFreeModel, + selectAutoFreeModel, } from '@/lib/ai-gateway/models'; import { getOpenRouterModelsFromRedis } from '@/lib/ai-gateway/providers/gateway-models-cache'; import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions'; @@ -84,7 +84,7 @@ export async function getAutoFreeCandidates( ): Promise> { const openRouterModels = await getOpenRouterModelsFromRedis(); const candidates = new Set(); - for (const model of autoFreeModels) { + for (const { model } of autoFreeModels) { if (isKiloExclusiveFreeModel(model)) { const kiloModel = findKiloExclusiveModel(model); if (kiloModel && gatewaySupportsApiKind(kiloModel.gateway, apiKind)) { @@ -286,11 +286,16 @@ export async function resolveAutoModel( if (candidates.length === 0) { return { kind: 'no_free_models_available' }; } - const randomNumber = getRandomNumber( - 'free_routing_' + (sessionId ?? (await userPromise)?.id ?? clientIp), - candidates.length + const candidateIds = new Set(candidates); + const selectedModel = selectAutoFreeModel( + autoFreeModels + .filter(candidate => candidateIds.has(candidate.model)) + .toSorted((a, b) => a.model.localeCompare(b.model)), + 'free_routing_' + (sessionId ?? (await userPromise)?.id ?? clientIp) ); - return { kind: 'ok', resolved: { model: candidates[randomNumber] } }; + return selectedModel + ? { kind: 'ok', resolved: { model: selectedModel } } + : { kind: 'no_free_models_available' }; } if (model === KILO_AUTO_SMALL_MODEL.id) { return { diff --git a/apps/web/src/lib/ai-gateway/models.test.ts b/apps/web/src/lib/ai-gateway/models.test.ts index 389094c002..c5cbcca4c7 100644 --- a/apps/web/src/lib/ai-gateway/models.test.ts +++ b/apps/web/src/lib/ai-gateway/models.test.ts @@ -4,6 +4,7 @@ import { findKiloExclusiveModel, isKiloExclusiveRateLimitedModel, kiloExclusiveModels, + selectAutoFreeModel, } from './models'; import { hasBestEffortGuessDataCollectionRequirement, isFreeModel } from './is-free-model'; import { getInferenceProvider } from './providers/kilo-exclusive-model'; @@ -16,6 +17,7 @@ import { import { gpt_5_6_sol_stealth_model } from './providers/openai-exclusive'; import { tencent_hy3_free_model } from './providers/tencent'; import { gemma_4_26b_a4b_it_free_model } from './providers/google'; +import { getRandomNumber } from './getRandomNumber'; describe('rate-limited Kilo-exclusive models', () => { test('only includes free Gemma', () => { @@ -77,7 +79,9 @@ describe('isFreeModel', () => { expect(findKiloExclusiveModel('tencent/hy3:free')).toBe(tencent_hy3_free_model); expect(tencent_hy3_free_model.internal_id).toBe('tencent/hy3'); expect(tencent_hy3_free_model.inference_provider_restriction).toEqual(['tencent']); - expect(autoFreeModels).not.toContain(tencent_hy3_free_model.public_id); + expect(autoFreeModels.map(({ model }) => model)).not.toContain( + tencent_hy3_free_model.public_id + ); }); test('routes the discounted Claude Opus offering through the stealth provider identity', () => { @@ -144,14 +148,36 @@ describe('isFreeModel', () => { test('all autoFreeModels should pass isFreeModel', async () => { expect(autoFreeModels.length).toBeGreaterThan(0); - for (const model of autoFreeModels) { + for (const { model } of autoFreeModels) { expect(await isFreeModel(model)).toBe(true); } }); + test('all autoFreeModels should have positive integer weights', () => { + for (const { weight } of autoFreeModels) { + expect(Number.isInteger(weight)).toBe(true); + expect(weight).toBeGreaterThan(0); + } + }); + + test('uses autoFreeModels weights when selecting a model', () => { + const candidates = [ + { model: 'preferred/model', weight: 3 }, + { model: 'other/model', weight: 1 }, + ]; + const randomSeed = Array.from({ length: 100 }, (_, index) => `weight-test-${index}`).find( + seed => getRandomNumber(seed, 4) === 1 + ); + expect(randomSeed).toBeDefined(); + if (!randomSeed) return; + + expect(getRandomNumber(randomSeed, 4)).toBe(1); + expect(selectAutoFreeModel(candidates, randomSeed)).toBe('preferred/model'); + }); + test('all autoFreeModels should use the same AI SDK provider', () => { expect(autoFreeModels.length).toBeGreaterThan(0); - const providers = new Set(autoFreeModels.map(model => getAiSdkProvider(model, null))); + const providers = new Set(autoFreeModels.map(({ model }) => getAiSdkProvider(model, null))); expect(providers.size).toBe(1); }); diff --git a/apps/web/src/lib/ai-gateway/models.ts b/apps/web/src/lib/ai-gateway/models.ts index b4381e9f5f..c3eed1c935 100644 --- a/apps/web/src/lib/ai-gateway/models.ts +++ b/apps/web/src/lib/ai-gateway/models.ts @@ -35,14 +35,32 @@ import { deepseekDiscountedModels, } from '@/lib/ai-gateway/providers/deepseek'; import { type ProviderId } from '@/lib/ai-gateway/providers/types'; +import { getRandomNumber } from '@/lib/ai-gateway/getRandomNumber'; export const PRIMARY_DEFAULT_MODEL = CLAUDE_SONNET_CURRENT_MODEL_ID; +export type AutoFreeModel = { model: string; weight: number }; + export const autoFreeModels = [ - stepfun_37_flash_free_model.status === 'public' ? stepfun_37_flash_free_model.public_id : null, - 'inclusionai/ling-3.0-flash:free', - 'poolside/laguna-s-2.1:free', -].filter(m => m !== null); + ...(stepfun_37_flash_free_model.status === 'public' + ? [{ model: stepfun_37_flash_free_model.public_id, weight: 1 }] + : []), + { model: 'inclusionai/ling-3.0-flash:free', weight: 1 }, + { model: 'poolside/laguna-s-2.1:free', weight: 1 }, +] satisfies ReadonlyArray; + +export function selectAutoFreeModel(candidates: ReadonlyArray, randomSeed: string) { + const totalWeight = candidates.reduce((total, candidate) => total + candidate.weight, 0); + if (totalWeight === 0) return null; + + const bucket = getRandomNumber(randomSeed, totalWeight); + let cumulativeWeight = 0; + for (const candidate of candidates) { + cumulativeWeight += candidate.weight; + if (bucket < cumulativeWeight) return candidate.model; + } + return null; +} export const preferredModels = [ KILO_AUTO_FRONTIER_MODEL.id, @@ -50,7 +68,7 @@ export const preferredModels = [ KILO_AUTO_EFFICIENT_MODEL.id, KILO_AUTO_FREE_MODEL.id, - ...autoFreeModels, + ...autoFreeModels.map(({ model }) => model), ...(tencent_hy3_free_model.status === 'public' ? [tencent_hy3_free_model.public_id] : []), CLAUDE_SONNET_CURRENT_MODEL_ID,