diff --git a/README.md b/README.md index f3e0fe2..19b4a92 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ Runs on Bun >= 1.2 or Node >= 24. The built `dist/` entry is the default import. bun add @corbits/reranking ``` -`rerankDocuments(query, docs, config, options)` validates config with `RerankConfigSchema`, short-circuits empty input with no request, and returns ranked results. Options carry `deps`, with optional `retryPolicy`, `registry`, and `signal`. Config takes `baseURL` and `apiStyle`, with optional `model` (used by Cohere and Voyage), `apiKey`, and `timeoutMs`. +`rerankDocuments(query, docs, config, options?)` validates config with `RerankConfigSchema`, short-circuits empty input with no request, and returns ranked results. Options are optional: `deps` (defaults to global `fetch` and `createDefaultScheduler()`), `retryPolicy`, `registry`, and `signal`. Config takes `baseURL` and `apiStyle`, with optional `model` (used by Cohere and Voyage), `apiKey`, and `timeoutMs`. A retrieval pipeline reads its reranker endpoint from host config at boot — a half-configured reranker should fail boot, not degrade silently later — then reranks each query's fused candidates, falling back to the caller's own ordering on failure so a reranker outage never fails search: @@ -77,7 +77,7 @@ Reranking has no single OpenAI-style standard, so this package draws the adapter Jina's `/rerank` follows the Cohere shape and is served by the `cohere` adapter. -`RerankAdapter` mirrors inference's `ProviderAdapter` with `buildRequest` / `parseResponse` plus optional `extractRetryAfterMs`. Styles resolve through `createRerankAdapterRegistry`, which keeps a private `Map` copy so config-supplied style names resolve safely. Pass a custom `registry` to add a house format alongside the built-ins (`rerankAdapters`, `rerankAdapterRegistry`). +`RerankAdapter` mirrors inference's `ProviderAdapter` with `buildRequest` / `parseResponse` plus optional `extractRetryAfterMs`. Styles resolve through `createRerankAdapterRegistry`, which keeps a private `Map` copy so config-supplied style names resolve safely. Pass a custom `registry` to add a house format alongside the built-ins in `rerankAdapterRegistry`. Every protocol addresses documents by position in the request array, and TEI replies unordered. Results map back through `docs[index]`, with an out-of-range index raising so a score always lands on the right document. `RerankDoc.id` carries your stable identifier across that index-based wire, and the final list sorts by score descending. diff --git a/src/adapters.ts b/src/adapters.ts index 3b982d3..06600e7 100644 --- a/src/adapters.ts +++ b/src/adapters.ts @@ -193,7 +193,7 @@ export function createRerankAdapterRegistry( }; } -export const rerankAdapters: Readonly> = { +const rerankAdapters: Readonly> = { tei, cohere, voyage, diff --git a/src/index.ts b/src/index.ts index ed01d42..ab11a52 100644 --- a/src/index.ts +++ b/src/index.ts @@ -5,7 +5,6 @@ export { type RerankOptions, } from "./rerank.js"; export { - rerankAdapters, rerankAdapterRegistry, createRerankAdapterRegistry, type RerankAdapter, diff --git a/src/request.ts b/src/request.ts index ba9e30d..70713eb 100644 --- a/src/request.ts +++ b/src/request.ts @@ -3,7 +3,6 @@ import { classifyHTTPError, classifyNetworkError, classifyProtocolMismatch, - createDefaultRetryPolicy, type BuiltRequest, type Dependencies, } from "@intx/inference"; @@ -59,29 +58,25 @@ export type RetryAfterExtractor = (headers: Headers) => number | undefined; export type RunRequestOptions = { deps: RequestDependencies; - /** Defaults to Interchange's policy: back off retryables, abort the rest. */ - retryPolicy?: RetryPolicy; + retryPolicy: RetryPolicy; /** * Per-attempt ceiling, enforced alongside any caller `signal` rather than * instead of it. A cold local model can take a while to page in. */ - timeoutMs?: number; + timeoutMs: number; /** * Reads `Retry-After` off a failed response. Mirrors `ProviderAdapter`'s * member of the same name, so a provider that signals pacing its own way can - * override without touching the transport. Defaults to - * {@link extractRetryAfterMs}. + * override without touching the transport. */ - extractRetryAfterMs?: RetryAfterExtractor; - signal?: AbortSignal; + extractRetryAfterMs: RetryAfterExtractor; + signal: AbortSignal | undefined; }; type Attempt = | { ok: true; body: unknown } | { ok: false; error: InferenceError }; -const DEFAULT_TIMEOUT_MS = 30_000; - /** * `Retry-After` in seconds or as an HTTP date; undefined when absent. * @@ -199,10 +194,13 @@ export async function runJSONRequest( request: BuiltRequest, options: RunRequestOptions, ): Promise { - const { deps, signal } = options; - const policy = options.retryPolicy ?? createDefaultRetryPolicy(); - const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS; - const extractRetryAfter = options.extractRetryAfterMs ?? extractRetryAfterMs; + const { + deps, + signal, + retryPolicy: policy, + timeoutMs, + extractRetryAfterMs: extractRetryAfter, + } = options; // Time comes from the harness scheduler, not globals, so a virtual-clock // test scheduler drives the retry loop deterministically. diff --git a/src/rerank.test.ts b/src/rerank.test.ts index 43a0eb1..09ec363 100644 --- a/src/rerank.test.ts +++ b/src/rerank.test.ts @@ -5,7 +5,7 @@ import { rerankDocuments, type RerankConfig } from "./rerank"; import { createRerankAdapterRegistry, rerankAdapterRegistry, - rerankAdapters, + type RerankAdapter, type RerankAPIStyle, } from "./adapters"; import { extractRetryAfterMs, RerankRequestError } from "./request"; @@ -232,7 +232,7 @@ describe("retry behaviour inherited from the inference policy", () => { const seen: Array = []; const registry = createRerankAdapterRegistry({ tei: { - ...rerankAdapters.tei, + ...rerankAdapterRegistry.resolve("tei"), extractRetryAfterMs: (headers) => { seen.push(headers.get("x-ratelimit-reset")); return 0; @@ -301,11 +301,11 @@ describe("adapter registry", () => { }); it("is not mutable through the map it was built from", () => { - const source: Record = { - tei: rerankAdapters.tei, + const source: Record = { + tei: rerankAdapterRegistry.resolve("tei"), }; const registry = createRerankAdapterRegistry(source); - source.voyage = rerankAdapters.voyage; + source.voyage = rerankAdapterRegistry.resolve("voyage"); expect(registry.has("voyage")).toBe(false); }); }); diff --git a/src/rerank.ts b/src/rerank.ts index b455e99..ca12192 100644 --- a/src/rerank.ts +++ b/src/rerank.ts @@ -1,6 +1,8 @@ import { type } from "arktype"; import { classifyProtocolMismatch, + createDefaultRetryPolicy, + createDefaultScheduler, ProtocolMismatchError, } from "@intx/inference"; import type { RetryPolicy } from "@intx/types/runtime"; @@ -8,22 +10,24 @@ import type { RetryPolicy } from "@intx/types/runtime"; import { rerankAdapterRegistry, type RerankAdapterRegistry, - type RerankAPIStyle, type RerankDoc, type RerankResult, } from "./adapters.js"; import { + extractRetryAfterMs, RerankRequestError, runJSONRequest, type RequestDependencies, } from "./request.js"; +const DEFAULT_TIMEOUT_MS = 30_000; + export const RerankConfigSchema = type({ - /** Provider root, e.g. `http://localhost:8085` for a TEI server. */ + /** Provider root, e.g. `http://localhost:8080` for a TEI server. */ baseURL: "string", /** * A key into the registry. Left open rather than pinned to - * {@link RerankAPIStyle} so a custom registry can name a house format; the + * `RerankAPIStyle` so a custom registry can name a house format; the * registry rejects an unknown style by name. The built-ins are `tei`, * `cohere` (also Jina) and `voyage`. */ @@ -36,7 +40,8 @@ export const RerankConfigSchema = type({ export type RerankConfig = typeof RerankConfigSchema.infer; export type RerankOptions = { - deps: RequestDependencies; + /** Defaults to global `fetch` and `createDefaultScheduler()`. */ + deps?: RequestDependencies; retryPolicy?: RetryPolicy; /** Defaults to the built-in registry: TEI, Cohere/Jina and Voyage. */ registry?: RerankAdapterRegistry; @@ -55,7 +60,7 @@ export async function rerankDocuments( query: string, docs: readonly RerankDoc[], config: RerankConfig, - options: RerankOptions, + options: RerankOptions = {}, ): Promise { const parsed = RerankConfigSchema.assert(config); @@ -70,15 +75,14 @@ export async function rerankDocuments( }); const body = await runJSONRequest(request, { - deps: options.deps, - ...(options.retryPolicy !== undefined - ? { retryPolicy: options.retryPolicy } - : {}), - ...(parsed.timeoutMs !== undefined ? { timeoutMs: parsed.timeoutMs } : {}), - ...(adapter.extractRetryAfterMs !== undefined - ? { extractRetryAfterMs: adapter.extractRetryAfterMs } - : {}), - ...(options.signal !== undefined ? { signal: options.signal } : {}), + deps: options.deps ?? { + fetch: globalThis.fetch, + scheduler: createDefaultScheduler(), + }, + retryPolicy: options.retryPolicy ?? createDefaultRetryPolicy(), + timeoutMs: parsed.timeoutMs ?? DEFAULT_TIMEOUT_MS, + extractRetryAfterMs: adapter.extractRetryAfterMs ?? extractRetryAfterMs, + signal: options.signal, }); let scored: Array<{ index: number; score: number }>; @@ -108,5 +112,3 @@ export async function rerankDocuments( }) .sort((a, b) => b.score - a.score); } - -export type { RerankAPIStyle, RerankDoc, RerankResult };