diff --git a/.changeset/ai-cloudflare-embed-batch.md b/.changeset/ai-cloudflare-embed-batch.md new file mode 100644 index 0000000000..39ff066744 --- /dev/null +++ b/.changeset/ai-cloudflare-embed-batch.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-cloudflare': patch +--- + +`embed()` with more than 100 inputs no longer fails on Workers AI embedding models that accept 100 texts per request, such as the BGE models. The adapter now sends the inputs in sequential requests of at most 100 texts and returns the vectors in input order. An aborted signal stops the remaining requests. Models with a lower limit, such as `@cf/qwen/qwen3-embedding-0.6b`, can still reject large inputs. diff --git a/packages/ai-cloudflare/src/adapters/embedding.ts b/packages/ai-cloudflare/src/adapters/embedding.ts index cd66449d60..c9059e31bb 100644 --- a/packages/ai-cloudflare/src/adapters/embedding.ts +++ b/packages/ai-cloudflare/src/adapters/embedding.ts @@ -8,6 +8,13 @@ import type { EmbeddingOptions, EmbeddingResult } from '@tanstack/ai' import type { CloudflareConfig, CloudflareConfigInput } from '../utils/config' import type { CloudflareEmbeddingModel } from '../utils/models' +/** + * Texts per Workers AI request. 100 is the cap of the BGE and EmbeddingGemma + * models. Models with a lower cap (such as `@cf/qwen/qwen3-embedding-0.6b`) + * can still reject a full batch. + */ +const CLOUDFLARE_MAX_BATCH_SIZE = 100 + /** Extra inputs forwarded to the embedding model (model specific). */ export type CloudflareEmbeddingProviderOptions = Record @@ -43,21 +50,35 @@ export class CloudflareEmbeddingAdapter< `activity=embed provider=${this.name} model=${model} inputs=${texts.length}`, { provider: this.name, model }, ) - const output = (await runModel( - this.cfConfig, - model, - { ...options.modelOptions, text: texts }, - { signal: options.abortSignal }, - )) as { data?: Array> } - if (!Array.isArray(output.data) || output.data.length !== texts.length) { - throw new Error( - `Workers AI ${model} returned ${output.data?.length ?? 0} embeddings for ${texts.length} inputs`, - ) + const vectors: Array> = [] + for ( + let start = 0; + start < texts.length; + start += CLOUDFLARE_MAX_BATCH_SIZE + ) { + // The binding path ignores the signal, so check it between batches. + options.abortSignal?.throwIfAborted() + const batch = texts.slice(start, start + CLOUDFLARE_MAX_BATCH_SIZE) + const output = (await runModel( + this.cfConfig, + model, + { ...options.modelOptions, text: batch }, + { signal: options.abortSignal }, + )) as { data?: Array> } + if ( + !Array.isArray(output.data) || + output.data.length !== batch.length + ) { + throw new Error( + `Workers AI ${model} returned ${output.data?.length ?? 0} embeddings for ${batch.length} inputs`, + ) + } + vectors.push(...output.data) } return { id: generateId(this.name), model, - embeddings: output.data.map((vector, index) => ({ + embeddings: vectors.map((vector, index) => ({ vector, index, })), diff --git a/packages/ai-cloudflare/tests/cloudflare-adapter.test.ts b/packages/ai-cloudflare/tests/cloudflare-adapter.test.ts index 2e21b48cba..a17b5f51e6 100644 --- a/packages/ai-cloudflare/tests/cloudflare-adapter.test.ts +++ b/packages/ai-cloudflare/tests/cloudflare-adapter.test.ts @@ -438,6 +438,45 @@ describe('embedding adapter', () => { ]) }) + it('splits more than 100 inputs into requests of at most 100', async () => { + const { binding, run } = fakeBinding((_model, inputs) => { + const text = (inputs as { text: Array }).text + if (text.length > 100) throw new Error('too many: maxItems is 100') + return { data: text.map((t) => [Number(t)]) } + }) + const adapter = createCloudflareEmbedding('@cf/baai/bge-m3', { binding }) + const input = Array.from({ length: 250 }, (_, i) => String(i)) + const result = await adapter.createEmbeddings({ + model: '@cf/baai/bge-m3', + input, + logger, + }) + expect( + run.mock.calls.map((c) => (c[1] as { text: Array }).text.length), + ).toEqual([100, 100, 50]) + expect(result.embeddings).toHaveLength(250) + expect(result.embeddings[0]).toEqual({ vector: [0], index: 0 }) + expect(result.embeddings[249]).toEqual({ vector: [249], index: 249 }) + }) + + it('stops sending batches once the signal aborts', async () => { + const controller = new AbortController() + const { binding, run } = fakeBinding((_model, inputs) => { + controller.abort(new Error('stop')) + return { data: (inputs as { text: Array }).text.map(() => [0]) } + }) + const adapter = createCloudflareEmbedding('@cf/baai/bge-m3', { binding }) + await expect( + adapter.createEmbeddings({ + model: '@cf/baai/bge-m3', + input: Array.from({ length: 250 }, (_, i) => String(i)), + abortSignal: controller.signal, + logger, + }), + ).rejects.toThrow('stop') + expect(run).toHaveBeenCalledTimes(1) + }) + it('sends defaultHeaders as extraHeaders', async () => { const { binding, run } = fakeBinding(() => ({ data: [[0.1]] })) const adapter = createCloudflareEmbedding('@cf/baai/bge-m3', {