fix(ai-gemini): split embed() input into requests of at most 100

.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.

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,25 +85,39 @@ export class GeminiEmbeddingAdapter<
         { provider: this.name, model },
       )
 
-      const response = await this.client.models.embedContent({
-        model,
-        contents: texts,
-        config,
-      })
+      const embeddings: Array<ContentEmbedding> = []
+      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
       // characters). The Developer API reports none, so `usage` is omitted.
       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

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<string>).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' },