From 77be151533433900f40016d52aabf1fd2287732f Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Wed, 4 Nov 2026 18:56:13 -0400 Subject: [PATCH] fix(ai-gemini): split embed() input into requests of at most 100 Co-Authored-By: Claude Sonnet 5.5 --- .changeset/gemini-embed-batch-limit.md | 5 ++ packages/ai-gemini/src/adapters/embedding.ts | 46 +++++++++++++------ .../ai-gemini/tests/embedding-adapter.test.ts | 37 +++++++++++++++ 3 files changed, 75 insertions(+), 13 deletions(-) create mode 100644 .changeset/gemini-embed-batch-limit.md diff --git a/.changeset/gemini-embed-batch-limit.md b/.changeset/gemini-embed-batch-limit.md new file mode 100644 index 0000000000..59dbb18e88 --- /dev/null +++ b/.changeset/gemini-embed-batch-limit.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-gemini': patch +--- + +`embed()` with the Gemini adapter now sends inputs in requests of at most 100 texts, one after the other, instead of one request for the whole array. Vectors keep input order and indexes, and Vertex billable characters are summed across requests. Calls with 100 or fewer inputs are unchanged. diff --git a/packages/ai-gemini/src/adapters/embedding.ts b/packages/ai-gemini/src/adapters/embedding.ts index 62500fa1c1..a12ad8c340 100644 --- a/packages/ai-gemini/src/adapters/embedding.ts +++ b/packages/ai-gemini/src/adapters/embedding.ts @@ -7,7 +7,7 @@ import { getGeminiApiKeyFromEnv, } from '../utils' import type { EmbeddingOptions, EmbeddingResult } from '@tanstack/ai' -import type { EmbedContentConfig } from '@google/genai' +import type { ContentEmbedding, EmbedContentConfig } from '@google/genai' import type { GeminiClient } from '../client/index' import type { GeminiEmbeddingModel, @@ -17,6 +17,12 @@ import type { import type { GeminiEmbeddingProviderOptions } from '../embedding/embedding-provider-options' import type { GeminiClientConfig } from '../utils/client' +/** + * Largest number of inputs sent in one `embedContent` call. The Gemini API + * rejects larger batches, so longer inputs are split into sequential requests. + */ +const GEMINI_MAX_BATCH_SIZE = 100 + /** * Configuration for Gemini Embedding adapter */ @@ -29,7 +35,7 @@ export interface GeminiEmbeddingConfig extends GeminiClientConfig {} * Supports gemini-embedding-001. * * Features: - * - Batch embedding (one request for the whole input array) + * - Batch embedding (inputs are sent in requests of at most 100 texts) * - Matryoshka dimension reduction via the top-level `dimensions` option * (mapped to the SDK's `outputDimensionality`) * - Task-type hints (`taskType`, `title`) via provider options @@ -79,17 +85,32 @@ export class GeminiEmbeddingAdapter< { provider: this.name, model }, ) - const response = await this.client.models.embedContent({ - model, - contents: texts, - config, - }) + const embeddings: Array = [] + let billableCharacters: number | undefined + for ( + let start = 0; + start < texts.length; + start += GEMINI_MAX_BATCH_SIZE + ) { + const batch = texts.slice(start, start + GEMINI_MAX_BATCH_SIZE) + const response = await this.client.models.embedContent({ + model, + contents: batch, + config, + }) + + const batchEmbeddings = response.embeddings + if (!batchEmbeddings || batchEmbeddings.length !== batch.length) { + throw new Error( + `Gemini embedContent returned ${batchEmbeddings ? batchEmbeddings.length : 'no'} embeddings for ${batch.length} inputs (model: ${model}).`, + ) + } + embeddings.push(...batchEmbeddings) - const embeddings = response.embeddings - if (!embeddings || embeddings.length !== texts.length) { - throw new Error( - `Gemini embedContent returned ${embeddings ? embeddings.length : 'no'} embeddings for ${texts.length} inputs (model: ${model}).`, - ) + const batchCharacters = response.metadata?.billableCharacterCount + if (batchCharacters !== undefined) { + billableCharacters = (billableCharacters ?? 0) + batchCharacters + } } // Only Vertex reports usage (per-input token counts and billable @@ -97,7 +118,6 @@ export class GeminiEmbeddingAdapter< const tokenCounts = embeddings.flatMap( (embedding) => embedding.statistics?.tokenCount ?? [], ) - const billableCharacters = response.metadata?.billableCharacterCount const tokens = tokenCounts.reduce((sum, count) => sum + count, 0) const hasUsage = tokenCounts.length > 0 || billableCharacters !== undefined diff --git a/packages/ai-gemini/tests/embedding-adapter.test.ts b/packages/ai-gemini/tests/embedding-adapter.test.ts index f955556f2d..e0f0c1b3f8 100644 --- a/packages/ai-gemini/tests/embedding-adapter.test.ts +++ b/packages/ai-gemini/tests/embedding-adapter.test.ts @@ -76,6 +76,43 @@ describe('Gemini Embedding Adapter', () => { expect(result.usage).toBeUndefined() }) + it('splits inputs over the 100-per-request limit into sequential requests', async () => { + const adapter = new TestGeminiEmbeddingAdapter( + { apiKey: 'test' }, + 'gemini-embedding-001', + ) + const spy = adapter.spyOnEmbedContent() + // Like the Gemini API, reject any request with more than 100 contents. + spy.mockImplementation(({ contents }) => { + const items = Array.isArray(contents) ? contents : [contents] + if (items.length > 100) { + return Promise.reject(new Error('at most 100 requests per batch')) + } + return Promise.resolve( + mockResponse(items.map((text) => [Number(String(text).slice(1))])), + ) + }) + + const input = Array.from({ length: 201 }, (_, i) => `t${i}`) + const result = await adapter.createEmbeddings({ + model: 'gemini-embedding-001', + input, + logger: testLogger, + }) + + expect(spy).toHaveBeenCalledTimes(3) + expect( + spy.mock.calls.map(([args]) => (args.contents as Array).length), + ).toEqual([100, 100, 1]) + expect(result.embeddings).toHaveLength(201) + expect(result.embeddings.map((e) => e.index)).toEqual( + input.map((_, i) => i), + ) + expect(result.embeddings.map((e) => e.vector[0])).toEqual( + input.map((_, i) => i), + ) + }) + it('maps the top-level dimensions option to outputDimensionality', async () => { const adapter = new TestGeminiEmbeddingAdapter( { apiKey: 'test' },