diff --git a/src/core/ai/gateway.ts b/src/core/ai/gateway.ts index 8a2674484..970d6b606 100644 --- a/src/core/ai/gateway.ts +++ b/src/core/ai/gateway.ts @@ -599,6 +599,8 @@ function warnRecipesMissingBatchTokens(): void { // LiteLLM proxy, llama-server) — they ship without a static cap because // the cap depends on a user-launched server. Warning is noise for them. if (embedding.no_batch_cap === true) continue; + // A declared item-count cap is a real batch cap — no warning needed. + if (embedding.max_batch_items !== undefined) continue; if (_warnedRecipes.has(recipe.id)) continue; _warnedRecipes.add(recipe.id); // eslint-disable-next-line no-console @@ -1517,10 +1519,17 @@ export async function embed(texts: string[], opts?: EmbedOpts): Promise capBatchItems(b, maxBatchItems)) + : tokenBatches; + const allEmbeddings: Float32Array[] = []; let _embedThrew = false; try { @@ -1596,6 +1605,23 @@ export function splitByTokenBudget( return batches; } +/** + * Split a batch into sub-batches of at most `maxItems` inputs. Enforces a + * hard COUNT cap that the token-budget split can't (many tiny inputs fit + * under any token budget). Used for endpoints like llama.cpp's llama-server + * that reject requests exceeding their launch batch size. + * + * @internal exported for tests; not part of the public gateway API. + */ +export function capBatchItems(texts: string[], maxItems: number): string[][] { + if (maxItems <= 0 || texts.length <= maxItems) return [texts]; + const batches: string[][] = []; + for (let i = 0; i < texts.length; i += maxItems) { + batches.push(texts.slice(i, i + maxItems)); + } + return batches; +} + /** * Returns true if the error looks like a provider batch-token-limit error. * diff --git a/src/core/ai/recipes/llama-server.ts b/src/core/ai/recipes/llama-server.ts index 212bd5478..a51650b27 100644 --- a/src/core/ai/recipes/llama-server.ts +++ b/src/core/ai/recipes/llama-server.ts @@ -35,9 +35,12 @@ export const llamaServer: Recipe = { trust_custom_dims: true, // #2271: user knows the launched model's native dim cost_per_1m_tokens_usd: 0, price_last_verified: '2026-05-10', - // llama-server's batch capacity is set by `--ctx-size` at launch - // time; no static cap to declare. v0.32 (#779). - no_batch_cap: true, + // llama-server enforces a hard request-COUNT cap equal to its launch + // batch size (`--batch-size`, default 32): it rejects requests with + // more inputs with `batch size N > maximum allowed batch size 32`. + // The token-budget split can't bound item count, so cap it here. A + // server launched with a larger `-b` can raise this. v0.32 (#779). + max_batch_items: 32, }, }, /** diff --git a/src/core/ai/types.ts b/src/core/ai/types.ts index 798551eec..8fc785e3d 100644 --- a/src/core/ai/types.ts +++ b/src/core/ai/types.ts @@ -54,6 +54,16 @@ export interface EmbeddingTouchpoint { * `max_batch_tokens` is also set. */ safety_factor?: number; + /** + * Maximum number of inputs per embedding request. Some endpoints enforce a + * hard COUNT cap independent of token budget — notably llama.cpp's + * `llama-server`, which rejects requests with more inputs than its launch + * batch size (e.g. `batch size 100 > maximum allowed batch size 32`). The + * token-budget pre-split cannot bound item count (many tiny chunks fit under + * any token budget), so this is enforced as a separate hard re-split after + * the token split. When unset, no count cap is applied. + */ + max_batch_items?: number; /** * v0.27.1: when true, at least one model in this recipe accepts image * inputs via a multimodal embedding endpoint (e.g. Voyage's diff --git a/test/ai/adaptive-embed-batch.test.ts b/test/ai/adaptive-embed-batch.test.ts index 505c7c82f..144668eca 100644 --- a/test/ai/adaptive-embed-batch.test.ts +++ b/test/ai/adaptive-embed-batch.test.ts @@ -34,6 +34,7 @@ import { resetGateway, embed, splitByTokenBudget, + capBatchItems, isTokenLimitError, __setEmbedTransportForTests, __getShrinkStateForTests, @@ -151,6 +152,41 @@ describe('splitByTokenBudget (pure helper)', () => { }); }); +describe('capBatchItems (hard COUNT cap helper)', () => { + test('batch at or under the cap is returned as a single batch (no copy of contents)', () => { + const texts = ['a', 'b', 'c']; + expect(capBatchItems(texts, 3)).toEqual([texts]); + expect(capBatchItems(texts, 10)).toEqual([texts]); + }); + + test('oversized batch splits into chunks of at most maxItems', () => { + const texts = Array.from({ length: 100 }, (_, i) => `t${i}`); + const result = capBatchItems(texts, 32); + expect(result.map(b => b.length)).toEqual([32, 32, 32, 4]); + expect(result.every(b => b.length <= 32)).toBe(true); + }); + + test('exact multiple splits evenly with no trailing empty batch', () => { + const texts = Array.from({ length: 64 }, (_, i) => `t${i}`); + expect(capBatchItems(texts, 32).map(b => b.length)).toEqual([32, 32]); + }); + + test('order is preserved across the split (concatenation round-trips)', () => { + const texts = Array.from({ length: 70 }, (_, i) => `t${i}`); + expect(capBatchItems(texts, 32).flat()).toEqual(texts); + }); + + test('maxItems <= 0 is a no-op (single batch) — never produces empty/infinite batches', () => { + const texts = ['a', 'b', 'c']; + expect(capBatchItems(texts, 0)).toEqual([texts]); + expect(capBatchItems(texts, -5)).toEqual([texts]); + }); + + test('empty input returns a single empty batch', () => { + expect(capBatchItems([], 32)).toEqual([[]]); + }); +}); + describe('isTokenLimitError (pure helper)', () => { test('matches Voyage error format', () => { expect(isTokenLimitError(VOYAGE_TOKEN_LIMIT_ERROR)).toBe(true); diff --git a/test/ai/no-batch-cap-suppression.serial.test.ts b/test/ai/no-batch-cap-suppression.serial.test.ts index 1241b3b5a..5a5d00f60 100644 --- a/test/ai/no-batch-cap-suppression.serial.test.ts +++ b/test/ai/no-batch-cap-suppression.serial.test.ts @@ -28,8 +28,8 @@ describe('v0.32 #779: no_batch_cap suppresses the missing-max_batch_tokens warni resetGateway(); }); - test('Ollama, LiteLLM, llama-server all declare no_batch_cap: true', () => { - for (const id of ['ollama', 'litellm', 'llama-server']) { + test('Ollama, LiteLLM declare no_batch_cap: true', () => { + for (const id of ['ollama', 'litellm']) { const r = getRecipe(id); expect(r, `${id} not registered`).toBeDefined(); expect( @@ -39,6 +39,16 @@ describe('v0.32 #779: no_batch_cap suppresses the missing-max_batch_tokens warni } }); + test('llama-server declares a hard item-count cap (max_batch_items: 32)', () => { + // llama.cpp enforces a request-COUNT cap equal to its launch --batch-size + // (default 32); declaring max_batch_items both bounds batches AND suppresses + // the missing-max_batch_tokens warning. Replaces the prior no_batch_cap flag. + const r = getRecipe('llama-server'); + expect(r, 'llama-server not registered').toBeDefined(); + expect(r!.touchpoints.embedding?.max_batch_items).toBe(32); + expect(r!.touchpoints.embedding?.no_batch_cap).toBeUndefined(); + }); + test('configureGateway does NOT warn for ollama/litellm/llama-server', () => { warnSpy.mockClear(); resetGateway();