Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/ai-cloudflare-embed-batch.md
Original file line number Diff line number Diff line change
@@ -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.
43 changes: 32 additions & 11 deletions packages/ai-cloudflare/src/adapters/embedding.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, unknown>

Expand Down Expand Up @@ -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<Array<number>> }
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<Array<number>> = []
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<Array<number>> }
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,
})),
Expand Down
39 changes: 39 additions & 0 deletions packages/ai-cloudflare/tests/cloudflare-adapter.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string> }).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<string> }).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<string> }).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', {
Expand Down
Loading