fix(ai-cohere): split embed() input into batches of at most 96

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

@@ -0,0 +1,5 @@
+---
+'@tanstack/ai-cohere': patch
+---
+
+`embed()` with the Cohere adapter no longer fails with a 400 when given more than 96 inputs. The adapter now splits the input into requests of at most 96 (the `v2/embed` per-request limit) and returns the vectors in input order, with token usage summed across requests.

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

@@ -28,6 +28,9 @@ export interface CohereEmbeddingConfig extends CohereClientConfig {}
 
 const DEFAULT_TIMEOUT_MS = 30_000
 
+/** Cohere v2/embed accepts at most 96 inputs per request. */
+const COHERE_MAX_BATCH_SIZE = 96
+
 /**
  * Returns true when `url` is malformed, non-http(s), or targets a private /
  * loopback / link-local host. Used to block SSRF via `allowUrlFetch`.
@@ -116,7 +119,8 @@ interface CohereEmbedResponse {
  * SDK dependency.
  *
  * Features:
- * - Batch embedding (one request for the whole input array)
+ * - Batch embedding (the input array is sent in requests of at most 96
+ *   inputs, the v2/embed per-request limit, and the results are merged)
  * - Multimodal inputs: text, images, and fused text+image items (one vector
  *   per input item)
  * - Matryoshka dimension reduction via the top-level `dimensions` option
@@ -173,75 +177,94 @@ export class CohereEmbeddingAdapter<
         }),
       )
 
-      // embedding_types is pinned to ['float'] (overriding any disagreeing
-      // modelOptions.embeddingTypes) so vectors are always number[].
-      const body: CohereEmbedRequestBody = {
-        model,
-        inputs,
-        input_type: inputType,
-        embedding_types: ['float'],
-      }
-      const truncate = modelOptions?.truncate
-      if (truncate !== undefined) {
-        body.truncate = truncate
-      }
-      if (options.dimensions !== undefined) {
-        body.output_dimension = options.dimensions
-      }
-
-      logger.request(
-        `activity=embed provider=${this.name} model=${model} inputs=${inputs.length}`,
-        { provider: this.name, model },
-      )
-
       const timeoutMs = this.clientConfig.timeout ?? DEFAULT_TIMEOUT_MS
       const transport = resolveCohereTransport(this.clientConfig)
-      const response = await fetchWithTimeout(
-        `${transport.baseUrl}/v2/embed`,
-        {
-          method: 'POST',
-          headers: {
-            Authorization: `Bearer ${this.clientConfig.apiKey}`,
-            'Content-Type': 'application/json',
-            ...transport.headers,
+
+      const vectors: Array<Array<number>> = []
+      let inputTokens: number | undefined
+      for (
+        let start = 0;
+        start < inputs.length;
+        start += COHERE_MAX_BATCH_SIZE
+      ) {
+        const batch = inputs.slice(start, start + COHERE_MAX_BATCH_SIZE)
+
+        // embedding_types is pinned to ['float'] (overriding any disagreeing
+        // modelOptions.embeddingTypes) so vectors are always number[].
+        const body: CohereEmbedRequestBody = {
+          model,
+          inputs: batch,
+          input_type: inputType,
+          embedding_types: ['float'],
+        }
+        const truncate = modelOptions?.truncate
+        if (truncate !== undefined) {
+          body.truncate = truncate
+        }
+        if (options.dimensions !== undefined) {
+          body.output_dimension = options.dimensions
+        }
+
+        logger.request(
+          `activity=embed provider=${this.name} model=${model} inputs=${batch.length}`,
+          { provider: this.name, model },
+        )
+
+        const response = await fetchWithTimeout(
+          `${transport.baseUrl}/v2/embed`,
+          {
+            method: 'POST',
+            headers: {
+              Authorization: `Bearer ${this.clientConfig.apiKey}`,
+              'Content-Type': 'application/json',
+              ...transport.headers,
+            },
+            body: JSON.stringify(body),
+            ...(options.abortSignal && { signal: options.abortSignal }),
           },
-          body: JSON.stringify(body),
-          ...(options.abortSignal && { signal: options.abortSignal }),
-        },
-        timeoutMs,
-      )
+          timeoutMs,
+        )
 
-      if (!response.ok) {
-        const bodyText = await response.text()
-        let message = bodyText
-        try {
-          const parsed: unknown = JSON.parse(bodyText)
-          if (
-            typeof parsed === 'object' &&
-            parsed !== null &&
-            'message' in parsed &&
-            typeof parsed.message === 'string'
-          ) {
-            message = parsed.message
+        if (!response.ok) {
+          const bodyText = await response.text()
+          let message = bodyText
+          try {
+            const parsed: unknown = JSON.parse(bodyText)
+            if (
+              typeof parsed === 'object' &&
+              parsed !== null &&
+              'message' in parsed &&
+              typeof parsed.message === 'string'
+            ) {
+              message = parsed.message
+            }
+          } catch {
+            // Not JSON — fall back to the raw body text.
           }
-        } catch {
-          // Not JSON — fall back to the raw body text.
+          throw new Error(
+            `Cohere embed failed (${response.status}): ${message}`,
+          )
         }
-        throw new Error(`Cohere embed failed (${response.status}): ${message}`)
-      }
 
-      const data = (await response.json()) as CohereEmbedResponse
+        const data = (await response.json()) as CohereEmbedResponse
 
-      const vectors = data.embeddings?.float
-      if (!vectors) {
-        throw new Error(
-          'Cohere embed response did not include float embeddings',
-        )
-      }
-      if (vectors.length !== inputs.length) {
-        throw new Error(
-          `Cohere embed returned ${vectors.length} embeddings for ${inputs.length} inputs`,
-        )
+        const batchVectors = data.embeddings?.float
+        if (!batchVectors) {
+          throw new Error(
+            'Cohere embed response did not include float embeddings',
+          )
+        }
+        if (batchVectors.length !== batch.length) {
+          throw new Error(
+            `Cohere embed returned ${batchVectors.length} embeddings for ${batch.length} inputs`,
+          )
+        }
+        vectors.push(...batchVectors)
+
+        const batchTokens = data.meta?.billed_units?.input_tokens
+        if (batchTokens !== undefined) {
+          inputTokens = (inputTokens ?? 0) + batchTokens
+        }
       }
 
       const result: EmbeddingResult = {
@@ -250,7 +273,6 @@ export class CohereEmbeddingAdapter<
         embeddings: vectors.map((vector, index) => ({ vector, index })),
       }
 
-      const inputTokens = data.meta?.billed_units?.input_tokens
       if (inputTokens !== undefined) {
         const usage: TokenUsage = {
           promptTokens: inputTokens,

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

@@ -391,6 +391,54 @@ describe('Cohere Embedding Adapter', () => {
       expect(init.signal).toBeInstanceOf(AbortSignal)
     })
 
+    it('splits more than 96 inputs into batches of at most 96 and merges the results', async () => {
+      // Cohere v2/embed rejects requests with more than 96 inputs, so 200
+      // texts must go out as 96 + 96 + 8 and come back as one ordered result.
+      const texts = Array.from({ length: 200 }, (_, i) => `text ${i}`)
+      fetchMock.mockImplementation(async (_url: string, init: any) => {
+        const body = JSON.parse(init.body)
+        const inputs = body.inputs as Array<{
+          content: Array<{ text: string }>
+        }>
+        if (inputs.length > 96) {
+          return new Response(
+            JSON.stringify({
+              message: 'invalid request: too many texts, max 96 per call',
+            }),
+            { status: 400 },
+          )
+        }
+        return embedResponse({
+          vectors: inputs.map((input) => [
+            Number(input.content[0]!.text.replace('text ', '')),
+          ]),
+          inputTokens: inputs.length,
+        })
+      })
+      const adapter = createCohereEmbedding('embed-v4.0', 'test-key')
+
+      const result = await adapter.createEmbeddings({
+        model: 'embed-v4.0',
+        input: texts,
+        modelOptions: { inputType: 'search_document' },
+        logger: testLogger,
+      })
+
+      expect(fetchMock).toHaveBeenCalledTimes(3)
+      expect(
+        fetchMock.mock.calls.map((c) => JSON.parse(c[1].body).inputs.length),
+      ).toEqual([96, 96, 8])
+      expect(result.embeddings).toHaveLength(200)
+      expect(result.embeddings[0]).toEqual({ vector: [0], index: 0 })
+      expect(result.embeddings[96]).toEqual({ vector: [96], index: 96 })
+      expect(result.embeddings[199]).toEqual({ vector: [199], index: 199 })
+      expect(result.usage).toEqual({
+        promptTokens: 200,
+        completionTokens: 0,
+        totalTokens: 200,
+      })
+    })
+
     it('omits usage when billed_units.input_tokens is absent', async () => {
       fetchMock.mockResolvedValue(embedResponse({ vectors: [[0.1]] }))
       const adapter = createCohereEmbedding('embed-v4.0', 'test-key')