diff --git a/README.md b/README.md index 5258dac..728e0ef 100644 --- a/README.md +++ b/README.md @@ -53,15 +53,15 @@ console.log(vectors.length, vectors[0]?.length); // 2 768 ### Config -| Field | Type | Description | -| ---------------- | ---------------------- | ------------------------------------------------------------------------ | -| `baseURL` | `string` | Server root with its version path, such as `http://host:11434/v1`. | -| `model` | `string` | Model name. | -| `apiKey` | `string?` | Sent as a bearer token. | -| `dimensions` | `number?` | Requested output dimension. Sent only when set; not all models honor it. | -| `encodingFormat` | `"float" \| "base64"?` | Wire format. Defaults to `float`. Results are always `number[]`. | -| `batchSize` | `number?` | Texts per request. Defaults to 32. | -| `timeoutMs` | `number?` | Per-request timeout. Defaults to 30000. | +| Field | Type | Description | +| ---------------- | ---------------------- | ------------------------------------------------------------------------------------------- | +| `baseURL` | `string` | Server root with its version path, such as `http://host:11434/v1`. | +| `model` | `string` | Model name. | +| `apiKey` | `string?` | Sent as a bearer token. | +| `dimensions` | `number?` | Requested output dimension. Sent only when set; a reply of any other dimension is rejected. | +| `encodingFormat` | `"float" \| "base64"?` | Wire format. Defaults to `float`. Results are always `number[]`. | +| `batchSize` | `number?` | Texts per request. Defaults to 32. | +| `timeoutMs` | `number?` | Per-attempt timeout; a timed-out attempt is retried. Defaults to 30000. | ### Options @@ -78,7 +78,7 @@ A failed request throws `EmbeddingRequestError` with a classified `reason` and t ### Dimensions -Each model has a fixed output dimension: 768 for `nomic-embed-text`, 1536 for OpenAI `text-embedding-3-small`. If you store vectors, call `probeEmbedDims` at startup. Changing models changes the dimension, so treat it as a schema migration. +Each model has a fixed output dimension: 768 for `nomic-embed-text`, 1536 for OpenAI `text-embedding-3-small`. If you store vectors, call `probeEmbedDims` at startup. Changing models changes the dimension, so treat it as a schema migration. Every vector from one call has the same non-zero dimension, or the call throws `EmbeddingRequestError` with `protocol_mismatch`. ## Using with Interchange diff --git a/e2e/retry.test.ts b/e2e/retry.test.ts index d378e64..b917ad8 100644 --- a/e2e/retry.test.ts +++ b/e2e/retry.test.ts @@ -1,4 +1,5 @@ import { afterEach, expect, test } from "bun:test"; +import { createDefaultScheduler } from "@intx/inference"; import { setupHarness, type Harness } from "@intx/inference-testing"; import { EmbeddingRequestError, embedTexts } from "../src/index"; @@ -95,17 +96,72 @@ test("honors a caller-supplied extractRetryAfterMs", async () => { expect(harness.clock.now()).toBeGreaterThanOrEqual(45_000); }); -test("applies timeoutMs even when the caller passes a signal", async () => { - harness = setupHarness(); - const stream = harness.scenario.createStream(); - harness.scenario.whenRequestMatches(() => true, stream); +/** Hangs until aborted for the first `hangs` calls, then replies. */ +function hangingFetch(hangs: number): { + fetch: typeof fetch; + calls: () => number; +} { + let calls = 0; + const hangThenReply = (async (_url: unknown, init?: RequestInit) => { + calls++; + if (calls > hangs) { + return new Response( + JSON.stringify({ data: [{ index: 0, embedding: [1] }] }), + ); + } + const signal = init?.signal ?? undefined; + return await new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => reject(signal.reason)); + }); + }) as typeof fetch; + return { fetch: hangThenReply, calls: () => calls }; +} + +test("retries an attempt that hits timeoutMs", async () => { + const stub = hangingFetch(1); + const vectors = await embedTexts( + ["a"], + { ...CONFIG, timeoutMs: 50 }, + { deps: { fetch: stub.fetch, scheduler: createDefaultScheduler() } }, + ); + + expect(vectors).toEqual([[1]]); + expect(stub.calls()).toBe(2); +}); +test("applies timeoutMs even when the caller passes a signal", async () => { + const stub = hangingFetch(Infinity); const pending = embedTexts( ["a"], { ...CONFIG, timeoutMs: 50 }, - { deps: harness.deps, signal: new AbortController().signal }, + { + deps: { fetch: stub.fetch, scheduler: createDefaultScheduler() }, + signal: new AbortController().signal, + }, + ); + + await expect(pending).rejects.toMatchObject({ + reason: { category: "timeout" }, + }); + expect(stub.calls()).toBe(3); +}); + +test("reports a caller abort as aborted, without retrying", async () => { + harness = setupHarness({ enableInferenceTimers: true }); + harness.scenario.whenRequestMatches( + () => true, + harness.scenario.createStream(), ); + const controller = new AbortController(); + + const pending = embedTexts(["a"], CONFIG, { + deps: harness.deps, + signal: controller.signal, + }); + setTimeout(() => controller.abort(), 20); await harness.run(); - await expect(pending).rejects.toBeInstanceOf(EmbeddingRequestError); + await expect(pending).rejects.toMatchObject({ + reason: { category: "aborted" }, + }); }); diff --git a/src/embed.test.ts b/src/embed.test.ts index 236a4a3..48bdd9d 100644 --- a/src/embed.test.ts +++ b/src/embed.test.ts @@ -115,3 +115,26 @@ test("classifies a 200 that is not JSON instead of leaking a SyntaxError", async reason: { category: "protocol_mismatch" }, }); }); + +test("rejects empty, uneven, or off-dimension vectors", async () => { + for (const [config, bodies] of [ + [CONFIG, [{ data: [{ index: 0, embedding: [] }] }]], + [ + { ...CONFIG, batchSize: 1 }, + [ + { data: [{ index: 0, embedding: [1, 2] }] }, + { data: [{ index: 0, embedding: [1] }] }, + ], + ], + [{ ...CONFIG, dimensions: 3 }, [{ data: [{ index: 0, embedding: [1] }] }]], + ] as const) { + harness?.dispose(); + harness = replies(...bodies); + const texts = bodies.map((_, i) => String(i)); + const pending = embedTexts(texts, config, { deps: harness.deps }); + await harness.run(); + await expect(pending).rejects.toMatchObject({ + reason: { category: "protocol_mismatch" }, + }); + } +}); diff --git a/src/embed.ts b/src/embed.ts index 2a12416..b796b42 100644 --- a/src/embed.ts +++ b/src/embed.ts @@ -205,7 +205,20 @@ export async function embedTexts( extractRetryAfterMs: extractRetryAfter, signal: options.signal, }); - vectors.push(...parseResponse(body, request.url, batch.length)); + for (const vector of parseResponse(body, request.url, batch.length)) { + const width = valid.dimensions ?? vectors[0]?.length ?? vector.length; + if (vector.length === 0 || vector.length !== width) { + throw new EmbeddingRequestError( + classifyProtocolMismatch( + vector.length === 0 + ? "empty embedding" + : `expected ${String(width)}-dimensional embeddings, got ${String(vector.length)}`, + ), + request.url, + ); + } + vectors.push(vector); + } } return vectors; } diff --git a/src/request.ts b/src/request.ts index 89dd74c..14f9456 100644 --- a/src/request.ts +++ b/src/request.ts @@ -85,15 +85,6 @@ export const extractRetryAfterMs: RetryAfterExtractor = (headers) => { return Number.isNaN(at) ? undefined : Math.max(0, at - Date.now()); }; -/** Aborts on the caller's signal or the per-attempt timeout, whichever first. */ -function attemptSignal( - timeoutMs: number, - signal: AbortSignal | undefined, -): AbortSignal { - const timeout = AbortSignal.timeout(timeoutMs); - return signal === undefined ? timeout : AbortSignal.any([signal, timeout]); -} - async function attemptOnce( request: BuiltRequest, deps: RequestDependencies, @@ -101,6 +92,7 @@ async function attemptOnce( extractRetryAfter: RetryAfterExtractor, signal: AbortSignal | undefined, ): Promise { + const timeout = AbortSignal.timeout(timeoutMs); let res: Response; let raw: string; try { @@ -109,7 +101,8 @@ async function attemptOnce( headers: request.headers, body: request.body, redirect: "manual", - signal: attemptSignal(timeoutMs, signal), + signal: + signal === undefined ? timeout : AbortSignal.any([signal, timeout]), }); if (!res.ok) { @@ -131,9 +124,20 @@ async function attemptOnce( // as a bare SyntaxError. raw = await res.text(); } catch (cause) { - const aborted = - cause instanceof Error && - (cause.name === "AbortError" || cause.name === "TimeoutError"); + // A caller abort is final; the per-attempt timeout is retryable. + if (signal?.aborted === true) { + return { ok: false, error: classifyAbortError() }; + } + if (timeout.aborted) { + return { + ok: false, + error: { + category: "timeout", + message: `embedding request exceeded ${String(timeoutMs)} ms`, + }, + }; + } + const aborted = cause instanceof Error && cause.name === "AbortError"; return { ok: false, error: aborted ? classifyAbortError() : classifyNetworkError(cause),