From 43e1f2dd2662540045bc632e05091b8637e60a9e Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Tue, 3 Nov 2026 09:59:40 -0500 Subject: [PATCH] fix(ai-cohere): split embed() input into batches of at most 96 Cohere's v2/embed rejects requests with more than 96 inputs. Send the input in sequential requests of at most 96, return the vectors in input order and sum the billed input tokens. --- .changeset/ai-cohere-embed-batch.md | 5 + packages/ai-cohere/src/adapters/embedding.ts | 148 ++++++++++-------- .../ai-cohere/tests/embedding-adapter.test.ts | 48 ++++++ 3 files changed, 138 insertions(+), 63 deletions(-) create mode 100644 .changeset/ai-cohere-embed-batch.md diff --git a/.changeset/ai-cohere-embed-batch.md b/.changeset/ai-cohere-embed-batch.md new file mode 100644 index 0000000000..fba4ec67e0 --- /dev/null +++ b/.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. diff --git a/packages/ai-cohere/src/adapters/embedding.ts b/packages/ai-cohere/src/adapters/embedding.ts index 9d457de1db..f73ee02222 100644 --- a/packages/ai-cohere/src/adapters/embedding.ts +++ b/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> = [] + 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, diff --git a/packages/ai-cohere/tests/embedding-adapter.test.ts b/packages/ai-cohere/tests/embedding-adapter.test.ts index ac9e92ab81..4ae3667369 100644 --- a/packages/ai-cohere/tests/embedding-adapter.test.ts +++ b/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')