fix(ai-openai): split embed() inputs into requests of at most 2048

.changeset/ai-openai-embed-batch.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/ai-openai': patch
+---
+
+`embed()` with the OpenAI adapter now sends inputs in sequential requests of at most 2048 inputs each, the endpoint's limit, and sums the token usage. Before, more than 2048 texts failed with a 400 error. OpenAI also limits one request to 300,000 tokens across all inputs, so a large array of long texts can still fail. Split those calls yourself.

packages/ai-openai/src/adapters/embedding.ts

@@ -24,14 +24,23 @@ import type { OpenAIClientConfig } from '../utils/client'
  */
 export interface OpenAIEmbeddingConfig extends OpenAIClientConfig {}
 
+/**
+ * The embeddings endpoint rejects requests with more than 2048 inputs, so
+ * larger arrays are sent as several sequential requests.
+ */
+// ponytail: splits by input count only. OpenAI also caps a request at 300k
+// tokens summed across inputs, so 2048 long texts can still 400. Add a
+// token-budget split if that bites (needs a tokeniser or char estimate).
+const OPENAI_MAX_BATCH_SIZE = 2048
+
 /**
  * OpenAI Embedding Adapter
  *
  * Tree-shakeable adapter for OpenAI text embeddings.
  * Supports text-embedding-3-small and text-embedding-3-large.
  *
  * Features:
- * - Batch embedding (one request for the whole input array)
+ * - Batch embedding (sequential requests of at most 2048 inputs each; usage is summed)
  * - Matryoshka dimension reduction via the top-level `dimensions` option
  */
 export class OpenAIEmbeddingAdapter<
@@ -62,40 +71,59 @@ export class OpenAIEmbeddingAdapter<
       // fields (server routes often pass modelOptions through from untyped
       // client input). encoding_format is pinned to float so vectors are
       // always number[].
-      const request: OpenAI_SDK.EmbeddingCreateParams = {
+      const baseRequest = {
         ...modelOptions,
         model,
-        input: texts,
         encoding_format: 'float',
-      }
-      if (options.dimensions !== undefined) {
-        request.dimensions = options.dimensions
-      }
+      } as const
+      const dimensions =
+        options.dimensions !== undefined
+          ? { dimensions: options.dimensions }
+          : {}
 
       logger.request(
         `activity=embed provider=${this.name} model=${model} inputs=${texts.length}`,
         { provider: this.name, model },
       )
 
-      const response = options.abortSignal
-        ? await this.client.embeddings.create(request, {
-            signal: options.abortSignal,
-          })
-        : await this.client.embeddings.create(request)
+      const embeddings: EmbeddingResult['embeddings'] = []
+      let promptTokens = 0
+      let totalTokens = 0
+
+      for (
+        let start = 0;
+        start < texts.length;
+        start += OPENAI_MAX_BATCH_SIZE
+      ) {
+        const request: OpenAI_SDK.EmbeddingCreateParams = {
+          ...baseRequest,
+          ...dimensions,
+          input: texts.slice(start, start + OPENAI_MAX_BATCH_SIZE),
+        }
+
+        const response = options.abortSignal
+          ? await this.client.embeddings.create(request, {
+              signal: options.abortSignal,
+            })
+          : await this.client.embeddings.create(request)
+
+        for (const item of response.data) {
+          embeddings.push({ vector: item.embedding, index: start + item.index })
+        }
+        promptTokens += response.usage.prompt_tokens
+        totalTokens += response.usage.total_tokens
+      }
 
       const usage: TokenUsage = {
-        promptTokens: response.usage.prompt_tokens,
+        promptTokens,
         completionTokens: 0,
-        totalTokens: response.usage.total_tokens,
+        totalTokens,
       }
 
       return {
         id: generateId(this.name),
         model,
-        embeddings: response.data.map((item) => ({
-          vector: item.embedding,
-          index: item.index,
-        })),
+        embeddings,
         usage,
       }
     } catch (error: unknown) {

packages/ai-openai/tests/embedding-adapter.test.ts

@@ -87,6 +87,82 @@ describe('OpenAI Embedding Adapter', () => {
       })
     })
 
+    it('splits inputs over 2048 into batches and keeps order, indexes and usage', async () => {
+      const adapter = new TestOpenAIEmbeddingAdapter(
+        { apiKey: 'test' },
+        'text-embedding-3-small',
+      )
+      const spy = adapter.spyOnEmbeddingsCreate()
+      // Like the real API: reject any request with more than 2048 inputs.
+      spy.mockImplementation((params) => {
+        const input = params.input as Array<string>
+        if (input.length > 2048) {
+          return Promise.reject(
+            new Error('400 max 2048 embedding inputs per request'),
+          ) as never
+        }
+        // Distinct prompt/total counts per batch prove both fields are summed.
+        return Promise.resolve({
+          ...mockResponse(input.map((text) => [Number(text)])),
+          usage: {
+            prompt_tokens: input.length,
+            total_tokens: input.length + 1,
+          },
+        }) as never
+      })
+
+      const texts = Array.from({ length: 5000 }, (_, i) => String(i))
+      const result = await adapter.createEmbeddings({
+        model: 'text-embedding-3-small',
+        input: texts,
+        dimensions: 256,
+        modelOptions: { user: 'u' },
+        logger: testLogger,
+      })
+
+      expect(spy.mock.calls.map(([params]) => params.input.length)).toEqual([
+        2048, 2048, 904,
+      ])
+      for (const [params] of spy.mock.calls) {
+        expect(params).toMatchObject({
+          model: 'text-embedding-3-small',
+          encoding_format: 'float',
+          dimensions: 256,
+          user: 'u',
+        })
+      }
+      expect(result.embeddings).toHaveLength(5000)
+      expect(result.embeddings[2048]).toEqual({ vector: [2048], index: 2048 })
+      expect(result.embeddings[4999]).toEqual({ vector: [4999], index: 4999 })
+      expect(result.usage).toEqual({
+        promptTokens: 5000,
+        completionTokens: 0,
+        totalTokens: 5003,
+      })
+    })
+
+    it('sends no request for an empty input array', async () => {
+      const adapter = new TestOpenAIEmbeddingAdapter(
+        { apiKey: 'test' },
+        'text-embedding-3-small',
+      )
+      const spy = adapter.spyOnEmbeddingsCreate()
+
+      const result = await adapter.createEmbeddings({
+        model: 'text-embedding-3-small',
+        input: [],
+        logger: testLogger,
+      })
+
+      expect(spy).not.toHaveBeenCalled()
+      expect(result.embeddings).toEqual([])
+      expect(result.usage).toEqual({
+        promptTokens: 0,
+        completionTokens: 0,
+        totalTokens: 0,
+      })
+    })
+
     it('passes the top-level dimensions option through', async () => {
       const adapter = new TestOpenAIEmbeddingAdapter(
         { apiKey: 'test' },