From a0f8a3fd23d10b03d752c3ff2ddb09759bb4dedb Mon Sep 17 00:00:00 2001 From: Sawyer Cutler Date: Thu, 24 Sep 2026 23:55:59 -0700 Subject: [PATCH] test: embed and rerank clients against a local HTTP server The clients now send real requests to a Bun.serve stub instead of a mocked fetch, so each test checks the wire shape, headers and body the endpoint actually receives. The rerank config validation tests move with them. --- package.json | 2 +- src/core/embed-client.test.ts | 193 ---------------- src/core/rerank-client.test.ts | 403 --------------------------------- tests/embed-client.test.ts | 106 +++++++++ tests/lib/http-stub.ts | 46 ++++ tests/rerank-client.test.ts | 267 ++++++++++++++++++++++ 6 files changed, 420 insertions(+), 597 deletions(-) delete mode 100644 src/core/embed-client.test.ts delete mode 100644 src/core/rerank-client.test.ts create mode 100644 tests/embed-client.test.ts create mode 100644 tests/lib/http-stub.ts create mode 100644 tests/rerank-client.test.ts diff --git a/package.json b/package.json index 89a3312..f5b828b 100644 --- a/package.json +++ b/package.json @@ -62,7 +62,7 @@ "build": "tsc -p tsconfig.build.json", "prepack": "npm run build", "test": "bun test ./src ./tests", - "test:coverage": "bun test --coverage --coverage-reporter=lcov --coverage-reporter=text ./src" + "test:coverage": "bun test --coverage --coverage-reporter=lcov --coverage-reporter=text ./src ./tests" }, "dependencies": { "arktype": "^2.1.29" diff --git a/src/core/embed-client.test.ts b/src/core/embed-client.test.ts deleted file mode 100644 index b31dfb7..0000000 --- a/src/core/embed-client.test.ts +++ /dev/null @@ -1,193 +0,0 @@ -import { describe, expect, it, mock } from "bun:test"; - -import { - EmbedHttpError, - EmbedTimeoutError, - embedTexts, - probeEmbedDims, -} from "./embed-client.js"; -import type { EmbedClientConfig } from "./embed-client.js"; - -function jsonResponse(body: unknown, status = 200): Response { - return new Response(JSON.stringify(body), { - status, - headers: { "content-type": "application/json" }, - }); -} - -const openAiConfig: EmbedClientConfig = { - baseUrl: "https://api.example.com", - modelId: "text-embed-3", - apiStyle: "openai", -}; - -const teiConfig: EmbedClientConfig = { - baseUrl: "https://tei.example.com", - modelId: "bge-m3", - apiStyle: "tei", -}; - -const ollamaConfig: EmbedClientConfig = { - baseUrl: "http://localhost:11434", - modelId: "nomic-embed-text", - apiStyle: "ollama", -}; - -describe("embedTexts", () => { - it("returns [] on empty input without calling fetch (T4)", async () => { - const fetchImpl = mock(() => Promise.resolve(jsonResponse({}))); - const result = await embedTexts([], openAiConfig, fetchImpl as unknown as typeof fetch); - expect(result).toEqual([]); - expect(fetchImpl).not.toHaveBeenCalled(); - }); - - it("embeds via the openai-compat shape, preserving order (T1)", async () => { - const fetchImpl = mock((url: string, init: RequestInit) => { - expect(url).toBe("https://api.example.com/v1/embeddings"); - const body = JSON.parse(init.body as string) as { model: string; input: string[] }; - expect(body.model).toBe("text-embed-3"); - return Promise.resolve( - jsonResponse({ - data: body.input.map((_, i) => ({ embedding: [i, i + 1, i + 2] })), - }), - ); - }); - const result = await embedTexts( - ["a", "b"], - openAiConfig, - fetchImpl as unknown as typeof fetch, - ); - expect(result).toEqual([ - [0, 1, 2], - [1, 2, 3], - ]); - }); - - it("embeds via the TEI shape", async () => { - const fetchImpl = mock((url: string, init: RequestInit) => { - expect(url).toBe("https://tei.example.com/embed"); - const body = JSON.parse(init.body as string) as { inputs: string[] }; - expect(body.inputs).toEqual(["a", "b"]); - return Promise.resolve(jsonResponse([[1, 2], [3, 4]])); - }); - const result = await embedTexts( - ["a", "b"], - teiConfig, - fetchImpl as unknown as typeof fetch, - ); - expect(result).toEqual([ - [1, 2], - [3, 4], - ]); - }); - - it("embeds via the Ollama /api/embed batch shape with truncate", async () => { - const calls: unknown[] = []; - const fetchImpl = mock((url: string, init: RequestInit) => { - const body = JSON.parse(init.body as string) as { - model: string; - input: string[]; - truncate: boolean; - }; - calls.push({ url, body }); - expect(url).toBe("http://localhost:11434/api/embed"); - expect(body.model).toBe("nomic-embed-text"); - expect(body.truncate).toBe(true); - return Promise.resolve( - jsonResponse({ embeddings: body.input.map((t) => [t.length, 0, 0]) }), - ); - }); - const result = await embedTexts( - ["hi", "hello"], - ollamaConfig, - fetchImpl as unknown as typeof fetch, - ); - expect((calls[0] as { body: { input: string[] } }).body.input).toEqual([ - "hi", - "hello", - ]); - expect(result).toEqual([ - [2, 0, 0], - [5, 0, 0], - ]); - }); - - it("sends a bearer token only when apiKey is set", async () => { - let seenAuth: string | undefined; - const fetchImpl = mock((_url: string, init: RequestInit) => { - seenAuth = (init.headers as Record).authorization; - return Promise.resolve(jsonResponse({ data: [{ embedding: [1] }] })); - }); - await embedTexts( - ["a"], - { ...openAiConfig, apiKey: "secret-key" }, - fetchImpl as unknown as typeof fetch, - ); - expect(seenAuth).toBe("Bearer secret-key"); - - const fetchImplNoKey = mock((_url: string, init: RequestInit) => { - expect((init.headers as Record).authorization).toBeUndefined(); - return Promise.resolve(jsonResponse({ data: [{ embedding: [1] }] })); - }); - await embedTexts(["a"], openAiConfig, fetchImplNoKey as unknown as typeof fetch); - }); - - it("batches sequentially per config.batchSize", async () => { - const seenBatches: string[][] = []; - const fetchImpl = mock((_url: string, init: RequestInit) => { - const body = JSON.parse(init.body as string) as { input: string[] }; - seenBatches.push(body.input); - return Promise.resolve( - jsonResponse({ data: body.input.map(() => ({ embedding: [1] })) }), - ); - }); - await embedTexts( - ["a", "b", "c"], - { ...openAiConfig, batchSize: 2 }, - fetchImpl as unknown as typeof fetch, - ); - expect(seenBatches).toEqual([["a", "b"], ["c"]]); - }); - - it("rejects with EmbedHttpError on a non-2xx response, capturing status (T2)", async () => { - const fetchImpl = mock(() => - Promise.resolve(new Response("server exploded", { status: 500 })), - ); - await expect( - embedTexts(["a"], openAiConfig, fetchImpl as unknown as typeof fetch), - ).rejects.toThrow(EmbedHttpError); - try { - await embedTexts(["a"], openAiConfig, fetchImpl as unknown as typeof fetch); - throw new Error("expected embedTexts to reject"); - } catch (err) { - expect(err).toBeInstanceOf(EmbedHttpError); - expect((err as EmbedHttpError).status).toBe(500); - } - }); - - it("rejects with EmbedTimeoutError before the test itself times out (T3)", async () => { - const neverResponds = (_input: string | URL | Request, init?: RequestInit) => - new Promise((_resolve, reject) => { - init?.signal?.addEventListener("abort", () => { - reject(new DOMException("The operation timed out.", "TimeoutError")); - }); - }); - await expect( - embedTexts( - ["a"], - { ...openAiConfig, timeoutMs: 15 }, - neverResponds as unknown as typeof fetch, - ), - ).rejects.toThrow(EmbedTimeoutError); - }); -}); - -describe("probeEmbedDims", () => { - it("returns the discovered vector length (T9)", async () => { - const fetchImpl = mock(() => - Promise.resolve(jsonResponse({ data: [{ embedding: new Array(768).fill(0) }] })), - ); - const dims = await probeEmbedDims(openAiConfig, fetchImpl as unknown as typeof fetch); - expect(dims).toBe(768); - }); -}); diff --git a/src/core/rerank-client.test.ts b/src/core/rerank-client.test.ts deleted file mode 100644 index 2d179d4..0000000 --- a/src/core/rerank-client.test.ts +++ /dev/null @@ -1,403 +0,0 @@ -import { describe, expect, it, mock } from "bun:test"; - -import { - DEFAULT_MAX_DOC_CHARS, - DEFAULT_RERANK_MODEL, - KNOWN_TEI_RERANK_MODEL_TOKEN_LIMITS, - RerankConfigError, - RerankHttpError, - RerankQueryTooLongError, - RerankTimeoutError, - defaultMaxDocCharsForModel, - rerankDocuments, - validateRerankConfig, -} from "./rerank-client.js"; -import type { RerankClientConfig } from "./rerank-client.js"; - -// The default env in .env.example ships RERANK_MODEL=bge-reranker-base with -// RERANK_MAX_DOC_CHARS unset (-> defaultMaxDocCharsForModel's per-model -// value). That exact combination must not throw — an earlier version of -// this budget failed its own validation on the shipped defaults, which -// would have taken down every host that never touched the rerank env vars. -const DEFAULT_SHIPPED_MODEL = "bge-reranker-base"; - -function jsonResponse(body: unknown, status = 200): Response { - return new Response(JSON.stringify(body), { - status, - headers: { "content-type": "application/json" }, - }); -} - -const teiConfig: RerankClientConfig = { - baseUrl: "https://tei.example.com", - apiStyle: "tei", -}; - -const cohereConfig: RerankClientConfig = { - baseUrl: "https://api.cohere.example.com", - apiStyle: "cohere", - apiKey: "secret-key", -}; - -const voyageConfig: RerankClientConfig = { - baseUrl: "https://api.voyage.example.com", - apiStyle: "voyage", -}; - -const docs = [ - { id: "chunk-a", text: "alpha content" }, - { id: "chunk-b", text: "beta content" }, -]; - -describe("rerankDocuments", () => { - it("returns [] on empty docs without calling fetch", async () => { - const fetchImpl = mock(() => Promise.resolve(jsonResponse([]))); - const result = await rerankDocuments( - "query", - [], - teiConfig, - fetchImpl as unknown as typeof fetch, - ); - expect(result).toEqual([]); - expect(fetchImpl).not.toHaveBeenCalled(); - }); - - it("issues the TEI shape and maps index -> requested doc id, sorted desc", async () => { - const fetchImpl = mock((url: string, init: RequestInit) => { - expect(url).toBe("https://tei.example.com/rerank"); - const body = JSON.parse(init.body as string) as { - query: string; - texts: string[]; - }; - expect(body.query).toBe("my query"); - expect(body.texts).toEqual(["alpha content", "beta content"]); - return Promise.resolve( - jsonResponse([ - { index: 1, score: 0.9 }, - { index: 0, score: 0.2 }, - ]), - ); - }); - - const result = await rerankDocuments( - "my query", - docs, - teiConfig, - fetchImpl as unknown as typeof fetch, - ); - - expect(result).toEqual([ - { id: "chunk-b", score: 0.9 }, - { id: "chunk-a", score: 0.2 }, - ]); - }); - - it("issues the Cohere v2 shape with the bearer auth header and default model", async () => { - const fetchImpl = mock((url: string, init: RequestInit) => { - expect(url).toBe("https://api.cohere.example.com/v2/rerank"); - expect((init.headers as Record).authorization).toBe( - "Bearer secret-key", - ); - const body = JSON.parse(init.body as string) as { - model: string; - query: string; - documents: string[]; - }; - expect(body.model).toBe("bge-reranker-v2-m3"); - expect(body.documents).toEqual(["alpha content", "beta content"]); - return Promise.resolve( - jsonResponse({ - results: [{ index: 0, relevance_score: 0.75 }], - }), - ); - }); - - const result = await rerankDocuments( - "my query", - docs, - cohereConfig, - fetchImpl as unknown as typeof fetch, - ); - - expect(result).toEqual([{ id: "chunk-a", score: 0.75 }]); - }); - - it("issues the Voyage shape", async () => { - const fetchImpl = mock((url: string, init: RequestInit) => { - expect(url).toBe("https://api.voyage.example.com/v1/rerank"); - const body = JSON.parse(init.body as string) as { - model: string; - query: string; - documents: string[]; - }; - expect(body.model).toBe("bge-reranker-v2-m3"); - return Promise.resolve( - jsonResponse({ - data: [ - { index: 0, relevance_score: 0.4 }, - { index: 1, relevance_score: 0.6 }, - ], - }), - ); - }); - - const result = await rerankDocuments( - "my query", - docs, - voyageConfig, - fetchImpl as unknown as typeof fetch, - ); - - expect(result).toEqual([ - { id: "chunk-b", score: 0.6 }, - { id: "chunk-a", score: 0.4 }, - ]); - }); - - it("throws RerankHttpError on a non-ok response", async () => { - const fetchImpl = mock(() => - Promise.resolve(jsonResponse({ error: "boom" }, 500)), - ); - await expect( - rerankDocuments( - "q", - docs, - teiConfig, - fetchImpl as unknown as typeof fetch, - ), - ).rejects.toBeInstanceOf(RerankHttpError); - }); - - it("throws RerankTimeoutError when the request aborts", async () => { - const fetchImpl = mock(() => { - const err = new Error("timed out"); - err.name = "TimeoutError"; - return Promise.reject(err); - }); - await expect( - rerankDocuments( - "q", - docs, - teiConfig, - fetchImpl as unknown as typeof fetch, - ), - ).rejects.toBeInstanceOf(RerankTimeoutError); - }); - - const baseBudget = defaultMaxDocCharsForModel("bge-reranker-base"); - it.each([ - // The default-model row is the Finding A case: with no `model` set the - // unconfigured default must use DEFAULT_RERANK_MODEL's own budget — NOT - // the smaller bge-reranker-base-calibrated value. - ["default model's own budget", "", DEFAULT_MAX_DOC_CHARS + 500, {}, DEFAULT_MAX_DOC_CHARS], - ["explicit bge-reranker-base budget", "", baseBudget + 500, { model: "bge-reranker-base" }, baseBudget], - ["configured maxDocChars", "q", 1_200, { maxDocChars: 1_000 }, 999], - ["at-budget text passes through untruncated", "", DEFAULT_MAX_DOC_CHARS, {}, DEFAULT_MAX_DOC_CHARS], - ])("resolves the document budget — %s", async (_label, query, docChars, override, expected) => { - const fetchImpl = mock((_url: string, init: RequestInit) => { - const body = JSON.parse(init.body as string) as { texts: string[] }; - expect(body.texts[0]?.length).toBe(expected); - return Promise.resolve(jsonResponse([{ index: 0, score: 0.5 }])); - }); - - await rerankDocuments( - query, - [{ id: "chunk-c", text: "x".repeat(docChars) }], - { ...teiConfig, ...override }, - fetchImpl as unknown as typeof fetch, - ); - expect(fetchImpl).toHaveBeenCalledTimes(1); - }); - - it.each([ - ["reserves the query's length out of the pair cap", "q".repeat(150), 1_200, 1_000, 850], - // Query 800 chars against a 1000 budget leaves exactly MIN_DOC_CHARS - // (200): must still run, not skip. - ["exact MIN_DOC_CHARS boundary still runs", "q".repeat(800), 500, 1_000, 200], - // One char longer leaves 199, one under MIN_DOC_CHARS: must throw and - // never touch the network rather than forcing the budget back up. - ["one char under MIN_DOC_CHARS skips without fetch", "q".repeat(801), 500, 1_000, "throws"], - ["query dwarfing maxDocChars skips outright", "q".repeat(10_000), 500, 300, "throws"], - ])("query-length reserve — %s", async (_label, query, docChars, maxDocChars, expected) => { - const expectedChars: number | null = - expected === "throws" ? null : (expected as number); - const fetchImpl = mock((_url: string, init?: RequestInit) => { - const body = JSON.parse(init?.body as string) as { texts: string[] }; - if (expectedChars !== null) { - expect(body.texts[0]?.length).toBe(expectedChars); - } - return Promise.resolve(jsonResponse([{ index: 0, score: 0.5 }])); - }); - - const run = rerankDocuments( - query, - [{ id: "chunk-c", text: "x".repeat(docChars) }], - { ...teiConfig, maxDocChars }, - fetchImpl as unknown as typeof fetch, - ); - if (expected === "throws") { - await expect(run).rejects.toBeInstanceOf(RerankQueryTooLongError); - expect(fetchImpl).not.toHaveBeenCalled(); - } else { - await run; - expect(fetchImpl).toHaveBeenCalledTimes(1); - } - }); - - it("trims before truncating so leading padding doesn't yield an all-whitespace document", async () => { - const paddedDoc = { - id: "chunk-c", - text: " ".repeat(220) + "real content that matters", - }; - const fetchImpl = mock((_url: string, init: RequestInit) => { - const body = JSON.parse(init.body as string) as { texts: string[] }; - expect(body.texts[0]?.trim().length).toBeGreaterThan(0); - expect(body.texts[0]).toContain("real content"); - return Promise.resolve(jsonResponse([{ index: 0, score: 0.5 }])); - }); - - await rerankDocuments( - "q", - [paddedDoc], - { ...teiConfig, maxDocChars: 220 }, - fetchImpl as unknown as typeof fetch, - ); - }); - - it("does not truncate documents at or under the budget", async () => { - const shortDoc = { id: "chunk-c", text: "x".repeat(DEFAULT_MAX_DOC_CHARS) }; - const fetchImpl = mock((_url: string, init: RequestInit) => { - const body = JSON.parse(init.body as string) as { texts: string[] }; - expect(body.texts[0]?.length).toBe(DEFAULT_MAX_DOC_CHARS); - return Promise.resolve(jsonResponse([{ index: 0, score: 0.5 }])); - }); - - await rerankDocuments( - "", // empty query: isolate document-only truncation from the query reserve - [shortDoc], - teiConfig, - fetchImpl as unknown as typeof fetch, - ); - }); -}); - -describe("validateRerankConfig", () => { - it("passes for the REAL shipped default (bge-reranker-base, its own default budget) — the exact combination .env.example ships unmodified", () => { - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - model: DEFAULT_SHIPPED_MODEL, - maxDocChars: defaultMaxDocCharsForModel(DEFAULT_SHIPPED_MODEL), - }), - ).not.toThrow(); - // Also exercise the config-omitted path, since that's what - // toRerankClientConfig actually produces when RERANK_MAX_DOC_CHARS is unset. - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - model: DEFAULT_SHIPPED_MODEL, - }), - ).not.toThrow(); - }); - - // Finding A: `model` is optional and most deployments never set it, which - // resolves to DEFAULT_RERANK_MODEL (bge-reranker-v2-m3, 8,192 tokens) — - // not bge-reranker-base. Validation must run against THAT resolution, not - // early-return because `config.model` is undefined. - it("validates the engine's true default (no model set at all) against DEFAULT_RERANK_MODEL, not bge-reranker-base", () => { - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - }), - ).not.toThrow(); - - // A maxDocChars sized for bge-reranker-base's smaller limit is nowhere - // near enough to trip DEFAULT_RERANK_MODEL's much larger one. - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - maxDocChars: defaultMaxDocCharsForModel("bge-reranker-base"), - }), - ).not.toThrow(); - - // But a maxDocChars that overflows DEFAULT_RERANK_MODEL's real 8,192-token - // limit must still be caught, not skipped because model was omitted. - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - maxDocChars: (KNOWN_TEI_RERANK_MODEL_TOKEN_LIMITS[DEFAULT_RERANK_MODEL] ?? 0) * 100, - }), - ).toThrow(RerankConfigError); - }); - - it("passes for a smaller custom budget against a known model with a large enough limit", () => { - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - model: "bge-reranker-base", - maxDocChars: 1_200, - }), - ).not.toThrow(); - }); - - it("throws RerankConfigError when maxDocChars can overflow a known model's token limit", () => { - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - model: "bge-reranker-base", - maxDocChars: 5_000, - }), - ).toThrow(RerankConfigError); - }); - - it("passes for bge-reranker-v2-m3 with its own much larger default budget", () => { - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - model: "bge-reranker-v2-m3", - }), - ).not.toThrow(); - }); - - // An unrecognized model no longer skips validation outright — it resolves - // to the conservative fallback limit (512 tokens, the smallest known TEI - // cross-encoder limit) instead. - it("validates an unlisted model against the conservative fallback limit rather than skipping", () => { - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - model: "some-custom-reranker", - maxDocChars: 100_000, - }), - ).toThrow(RerankConfigError); - - expect(() => - validateRerankConfig({ - baseUrl: "https://tei.example.com", - apiStyle: "tei", - model: "some-custom-reranker", - // Within the conservative (512-token) fallback's own default budget. - }), - ).not.toThrow(); - }); - - it("skips validation for non-TEI api styles", () => { - expect(() => - validateRerankConfig({ - baseUrl: "https://api.cohere.example.com", - apiStyle: "cohere", - model: "bge-reranker-base", - maxDocChars: 100_000, - }), - ).not.toThrow(); - }); -}); diff --git a/tests/embed-client.test.ts b/tests/embed-client.test.ts new file mode 100644 index 0000000..6c9abd5 --- /dev/null +++ b/tests/embed-client.test.ts @@ -0,0 +1,106 @@ +import { afterAll, beforeEach, describe, expect, test } from "bun:test"; + +import { + EmbedHttpError, + EmbedTimeoutError, + embedTexts, + probeEmbedDims, +} from "../src/core/embed-client.ts"; +import { startHttpStub } from "./lib/http-stub.ts"; + +const stub = startHttpStub(); +afterAll(() => stub.stop()); +beforeEach(() => stub.reset()); + +const openai = { baseUrl: stub.url, modelId: "text-embed-3", apiStyle: "openai" } as const; + +describe("embedTexts", () => { + test("sends nothing for empty input", async () => { + expect(await embedTexts([], openai)).toEqual([]); + expect(stub.requests).toEqual([]); + }); + + test("speaks the openai, tei and ollama wire shapes", async () => { + stub.reply = (req) => { + const body = req.body as { input?: string[]; inputs?: string[] }; + const texts = body.input ?? body.inputs ?? []; + const vectors = texts.map((t) => [t.length, 0]); + if (req.path === "/v1/embeddings") { + return Response.json({ data: vectors.map((embedding) => ({ embedding })) }); + } + if (req.path === "/embed") return Response.json(vectors); + return Response.json({ embeddings: vectors }); + }; + + expect(await embedTexts(["a", "bb"], openai)).toEqual([[1, 0], [2, 0]]); + expect( + await embedTexts(["a", "bb"], { baseUrl: stub.url, modelId: "bge-m3", apiStyle: "tei" }), + ).toEqual([[1, 0], [2, 0]]); + expect( + await embedTexts(["a", "bb"], { + baseUrl: stub.url, + modelId: "nomic-embed-text", + apiStyle: "ollama", + }), + ).toEqual([[1, 0], [2, 0]]); + + expect(stub.requests.map((r) => [r.path, r.body])).toEqual([ + ["/v1/embeddings", { model: "text-embed-3", input: ["a", "bb"] }], + ["/embed", { inputs: ["a", "bb"] }], + ["/api/embed", { model: "nomic-embed-text", input: ["a", "bb"], truncate: true }], + ]); + }); + + test("sends the bearer token only when apiKey is set, and batches per batchSize", async () => { + stub.reply = (req) => + Response.json({ + data: (req.body as { input: string[] }).input.map(() => ({ embedding: [1] })), + }); + + await embedTexts(["a", "b", "c"], { + baseUrl: stub.url, + modelId: "text-embed-3", + apiStyle: "openai", + apiKey: "secret", + batchSize: 2, + }); + await embedTexts(["d"], openai); + + expect( + stub.requests.map((r) => [r.authorization, (r.body as { input: string[] }).input]), + ).toEqual([ + ["Bearer secret", ["a", "b"]], + ["Bearer secret", ["c"]], + [null, ["d"]], + ]); + }); + + test("rejects a non-2xx reply with EmbedHttpError carrying the status", async () => { + stub.reply = () => new Response("server exploded", { status: 500 }); + const err = await embedTexts(["a"], openai).catch((e: unknown) => e); + expect(err).toBeInstanceOf(EmbedHttpError); + expect((err as EmbedHttpError).status).toBe(500); + }); + + test("rejects a reply slower than timeoutMs with EmbedTimeoutError", async () => { + stub.reply = async () => { + await Bun.sleep(200); + return Response.json({ data: [{ embedding: [1] }] }); + }; + await expect( + embedTexts(["a"], { + baseUrl: stub.url, + modelId: "text-embed-3", + apiStyle: "openai", + timeoutMs: 20, + }), + ).rejects.toThrow( + EmbedTimeoutError, + ); + }); +}); + +test("probeEmbedDims returns the served vector length", async () => { + stub.reply = () => Response.json({ data: [{ embedding: new Array(768).fill(0) }] }); + expect(await probeEmbedDims(openai)).toBe(768); +}); diff --git a/tests/lib/http-stub.ts b/tests/lib/http-stub.ts new file mode 100644 index 0000000..ce6ad85 --- /dev/null +++ b/tests/lib/http-stub.ts @@ -0,0 +1,46 @@ +// A local HTTP server standing in for an embed or rerank endpoint. Each test +// sets `reply`; every request is recorded with its path, auth header and +// JSON body (both clients only send JSON). + +export type StubRequest = { + path: string; + authorization: string | null; + body: unknown; +}; + +export type HttpStub = { + url: string; + requests: StubRequest[]; + reply: (req: StubRequest) => Response | Promise; + reset: () => void; + stop: () => void; +}; + +const unconfigured = () => new Response("no reply configured", { status: 500 }); + +export function startHttpStub(): HttpStub { + const stub: HttpStub = { + url: "", + requests: [], + reply: unconfigured, + reset: () => { + stub.requests.length = 0; + stub.reply = unconfigured; + }, + stop: () => server.stop(true), + }; + const server = Bun.serve({ + port: 0, + async fetch(req) { + const recorded: StubRequest = { + path: new URL(req.url).pathname, + authorization: req.headers.get("authorization"), + body: await req.json(), + }; + stub.requests.push(recorded); + return stub.reply(recorded); + }, + }); + stub.url = server.url.href.replace(/\/$/, ""); + return stub; +} diff --git a/tests/rerank-client.test.ts b/tests/rerank-client.test.ts new file mode 100644 index 0000000..ab65075 --- /dev/null +++ b/tests/rerank-client.test.ts @@ -0,0 +1,267 @@ +import { afterAll, beforeEach, describe, expect, test } from "bun:test"; + +import { + DEFAULT_MAX_DOC_CHARS, + DEFAULT_RERANK_MODEL, + KNOWN_TEI_RERANK_MODEL_TOKEN_LIMITS, + RerankConfigError, + RerankHttpError, + RerankQueryTooLongError, + RerankTimeoutError, + defaultMaxDocCharsForModel, + rerankDocuments, + validateRerankConfig, +} from "../src/core/rerank-client.ts"; +import { startHttpStub } from "./lib/http-stub.ts"; + +// .env.example ships RERANK_MODEL=bge-reranker-base with RERANK_MAX_DOC_CHARS +// unset; that combination must validate. +const DEFAULT_SHIPPED_MODEL = "bge-reranker-base"; + +const stub = startHttpStub(); +afterAll(() => stub.stop()); +beforeEach(() => stub.reset()); + +const tei = { baseUrl: stub.url, apiStyle: "tei" } as const; +const teiBaseModel = { + baseUrl: stub.url, + apiStyle: "tei", + model: "bge-reranker-base", +} as const; +const teiBudget = (maxDocChars: number) => + ({ baseUrl: stub.url, apiStyle: "tei", maxDocChars }) as const; +const docs = [ + { id: "chunk-a", text: "alpha content" }, + { id: "chunk-b", text: "beta content" }, +]; + +function teiTexts(): string[] { + return stub.requests.map((r) => (r.body as { texts: string[] }).texts[0] ?? ""); +} + +describe("rerankDocuments", () => { + test("sends nothing for no documents", async () => { + expect(await rerankDocuments("q", [], tei)).toEqual([]); + expect(stub.requests).toEqual([]); + }); + + test("speaks the tei, cohere and voyage wire shapes and maps scores to ids", async () => { + stub.reply = (req) => { + if (req.path === "/rerank") { + return Response.json([ + { index: 1, score: 0.9 }, + { index: 0, score: 0.2 }, + ]); + } + if (req.path === "/v2/rerank") { + return Response.json({ results: [{ index: 0, relevance_score: 0.75 }] }); + } + return Response.json({ + data: [ + { index: 0, relevance_score: 0.4 }, + { index: 1, relevance_score: 0.6 }, + ], + }); + }; + + expect(await rerankDocuments("my query", docs, tei)).toEqual([ + { id: "chunk-b", score: 0.9 }, + { id: "chunk-a", score: 0.2 }, + ]); + expect( + await rerankDocuments("my query", docs, { + baseUrl: stub.url, + apiStyle: "cohere", + apiKey: "secret-key", + }), + ).toEqual([{ id: "chunk-a", score: 0.75 }]); + expect( + await rerankDocuments("my query", docs, { baseUrl: stub.url, apiStyle: "voyage" }), + ).toEqual([ + { id: "chunk-b", score: 0.6 }, + { id: "chunk-a", score: 0.4 }, + ]); + + const texts = ["alpha content", "beta content"]; + expect(stub.requests.map((r) => [r.path, r.authorization, r.body])).toEqual([ + ["/rerank", null, { query: "my query", texts }], + [ + "/v2/rerank", + "Bearer secret-key", + { model: DEFAULT_RERANK_MODEL, query: "my query", documents: texts }, + ], + ["/v1/rerank", null, { model: DEFAULT_RERANK_MODEL, query: "my query", documents: texts }], + ]); + }); + + test("rejects a non-2xx reply with RerankHttpError", async () => { + stub.reply = () => Response.json({ error: "boom" }, { status: 500 }); + await expect(rerankDocuments("q", docs, tei)).rejects.toBeInstanceOf(RerankHttpError); + }); + + test("rejects a reply slower than timeoutMs with RerankTimeoutError", async () => { + stub.reply = async () => { + await Bun.sleep(200); + return Response.json([]); + }; + await expect( + rerankDocuments("q", docs, { baseUrl: stub.url, apiStyle: "tei", timeoutMs: 20 }), + ).rejects.toBeInstanceOf(RerankTimeoutError); + }); + + const baseBudget = defaultMaxDocCharsForModel("bge-reranker-base"); + test.each([ + ["the default model's own budget", "", DEFAULT_MAX_DOC_CHARS + 500, tei, DEFAULT_MAX_DOC_CHARS], + ["an explicit bge-reranker-base budget", "", baseBudget + 500, teiBaseModel, baseBudget], + ["a configured maxDocChars, less a one-character query", "q", 1_200, teiBudget(1_000), 999], + ["an at-budget document, untouched", "", DEFAULT_MAX_DOC_CHARS, tei, DEFAULT_MAX_DOC_CHARS], + ["a configured maxDocChars, less a long query", "q".repeat(150), 1_200, teiBudget(1_000), 850], + ["exactly the minimum document budget", "q".repeat(800), 500, teiBudget(1_000), 200], + ] as const)("truncates to %s", async (_label, query, docChars, config, expected) => { + stub.reply = () => Response.json([{ index: 0, score: 0.5 }]); + await rerankDocuments(query, [{ id: "c", text: "x".repeat(docChars) }], config); + expect(teiTexts()[0]?.length).toBe(expected); + }); + + test("refuses without a request when the query leaves too little document budget", async () => { + const one = [{ id: "c", text: "x".repeat(500) }]; + await expect( + rerankDocuments("q".repeat(801), one, teiBudget(1_000)), + ).rejects.toBeInstanceOf(RerankQueryTooLongError); + await expect( + rerankDocuments("q".repeat(10_000), one, teiBudget(300)), + ).rejects.toBeInstanceOf(RerankQueryTooLongError); + expect(stub.requests).toEqual([]); + }); + + test("trims before truncating so leading padding is not all that is sent", async () => { + stub.reply = () => Response.json([{ index: 0, score: 0.5 }]); + await rerankDocuments( + "q", + [{ id: "c", text: " ".repeat(220) + "real content that matters" }], + teiBudget(220), + ); + expect(teiTexts()[0]?.trim().length).toBeGreaterThan(0); + expect(teiTexts()[0]).toContain("real content"); + }); +}); + +describe("validateRerankConfig", () => { + test("passes for the REAL shipped default (bge-reranker-base, its own default budget) — the exact combination .env.example ships unmodified", () => { + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + model: DEFAULT_SHIPPED_MODEL, + maxDocChars: defaultMaxDocCharsForModel(DEFAULT_SHIPPED_MODEL), + }), + ).not.toThrow(); + // Also exercise the config-omitted path, since that's what + // toRerankClientConfig actually produces when RERANK_MAX_DOC_CHARS is unset. + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + model: DEFAULT_SHIPPED_MODEL, + }), + ).not.toThrow(); + }); + + // `model` is optional and most deployments never set it, which + // resolves to DEFAULT_RERANK_MODEL (bge-reranker-v2-m3, 8,192 tokens) — + // not bge-reranker-base. Validation must run against THAT resolution, not + // early-return because `config.model` is undefined. + test("validates the engine's true default (no model set at all) against DEFAULT_RERANK_MODEL, not bge-reranker-base", () => { + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + }), + ).not.toThrow(); + + // A maxDocChars sized for bge-reranker-base's smaller limit is nowhere + // near enough to trip DEFAULT_RERANK_MODEL's much larger one. + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + maxDocChars: defaultMaxDocCharsForModel("bge-reranker-base"), + }), + ).not.toThrow(); + + // But a maxDocChars that overflows DEFAULT_RERANK_MODEL's real 8,192-token + // limit must still be caught, not skipped because model was omitted. + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + maxDocChars: (KNOWN_TEI_RERANK_MODEL_TOKEN_LIMITS[DEFAULT_RERANK_MODEL] ?? 0) * 100, + }), + ).toThrow(RerankConfigError); + }); + + test("passes for a smaller custom budget against a known model with a large enough limit", () => { + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + model: "bge-reranker-base", + maxDocChars: 1_200, + }), + ).not.toThrow(); + }); + + test("throws RerankConfigError when maxDocChars can overflow a known model's token limit", () => { + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + model: "bge-reranker-base", + maxDocChars: 5_000, + }), + ).toThrow(RerankConfigError); + }); + + test("passes for bge-reranker-v2-m3 with its own much larger default budget", () => { + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + model: "bge-reranker-v2-m3", + }), + ).not.toThrow(); + }); + + // An unrecognized model is validated against the conservative fallback + // limit (512 tokens, the smallest known TEI cross-encoder limit). + test("validates an unlisted model against the conservative fallback limit rather than skipping", () => { + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + model: "some-custom-reranker", + maxDocChars: 100_000, + }), + ).toThrow(RerankConfigError); + + expect(() => + validateRerankConfig({ + baseUrl: "https://tei.example.com", + apiStyle: "tei", + model: "some-custom-reranker", + // Within the conservative (512-token) fallback's own default budget. + }), + ).not.toThrow(); + }); + + test("skips validation for non-TEI api styles", () => { + expect(() => + validateRerankConfig({ + baseUrl: "https://api.cohere.example.com", + apiStyle: "cohere", + model: "bge-reranker-base", + maxDocChars: 100_000, + }), + ).not.toThrow(); + }); +});