From b03143ba6194880a837519c54b4e6b1cd8c44c6a Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 25 Jun 2026 17:55:46 +0200 Subject: [PATCH 1/9] feat: add rerank() activity and @tanstack/ai-cohere adapter Add a provider-agnostic rerank() activity to @tanstack/ai for reordering documents by relevance to a query, plus a new rerank-only @tanstack/ai-cohere provider package. - rerank() is generic over the document type: strings or JSON-serializable objects. Object documents are serialized for the provider and the original element is returned in the result, fully typed. - Supports topN, per-request cancellation via abortSignal, observe-only GenerationMiddleware (start/usage/finish/abort/error), and rerank:* devtools events. Rerank bills in provider search units, surfaced on usage.unitsBilled. - @tanstack/ai-cohere ships cohereRerank / createCohereRerank, talking to Cohere /v2/rerank over fetch (no SDK) with per-model provider options. - Unit tests for the activity and the Cohere adapter; docs for the rerank guide and the Cohere adapter page. --- .changeset/rerank-cohere.md | 19 ++ docs/adapters/cohere.md | 132 ++++++++ docs/config.json | 15 + docs/rerank/rerank.md | 303 ++++++++++++++++++ packages/ai-cohere/LICENSE | 21 ++ packages/ai-cohere/README.md | 55 ++++ packages/ai-cohere/package.json | 63 ++++ packages/ai-cohere/src/adapters/rerank.ts | 177 ++++++++++ packages/ai-cohere/src/index.ts | 24 ++ packages/ai-cohere/src/model-meta.ts | 50 +++ packages/ai-cohere/src/utils/client.ts | 43 +++ .../ai-cohere/tests/rerank-adapter.test.ts | 133 ++++++++ packages/ai-cohere/tsconfig.json | 8 + packages/ai-cohere/vite.config.ts | 36 +++ packages/ai-event-client/src/index.ts | 37 +++ packages/ai/src/activities/index.ts | 22 ++ .../ai/src/activities/middleware/types.ts | 1 + packages/ai/src/activities/rerank/adapter.ts | 91 ++++++ packages/ai/src/activities/rerank/index.ts | 290 +++++++++++++++++ packages/ai/src/index.ts | 7 + packages/ai/src/middlewares/otel.ts | 1 + packages/ai/src/types.ts | 63 ++++ packages/ai/tests/rerank.test.ts | 219 +++++++++++++ pnpm-lock.yaml | 34 +- 24 files changed, 1833 insertions(+), 11 deletions(-) create mode 100644 .changeset/rerank-cohere.md create mode 100644 docs/adapters/cohere.md create mode 100644 docs/rerank/rerank.md create mode 100644 packages/ai-cohere/LICENSE create mode 100644 packages/ai-cohere/README.md create mode 100644 packages/ai-cohere/package.json create mode 100644 packages/ai-cohere/src/adapters/rerank.ts create mode 100644 packages/ai-cohere/src/index.ts create mode 100644 packages/ai-cohere/src/model-meta.ts create mode 100644 packages/ai-cohere/src/utils/client.ts create mode 100644 packages/ai-cohere/tests/rerank-adapter.test.ts create mode 100644 packages/ai-cohere/tsconfig.json create mode 100644 packages/ai-cohere/vite.config.ts create mode 100644 packages/ai/src/activities/rerank/adapter.ts create mode 100644 packages/ai/src/activities/rerank/index.ts create mode 100644 packages/ai/tests/rerank.test.ts diff --git a/.changeset/rerank-cohere.md b/.changeset/rerank-cohere.md new file mode 100644 index 000000000..27b4e8f8b --- /dev/null +++ b/.changeset/rerank-cohere.md @@ -0,0 +1,19 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-event-client': minor +'@tanstack/ai-cohere': minor +--- + +feat: add `rerank()` activity for reordering documents by relevance to a query + +Adds a provider-agnostic `rerank()` activity (with `createRerankOptions`, the +`RerankAdapter` interface, and `BaseRerankAdapter`). Documents may be strings +or JSON-serializable objects — object documents are serialized for the +provider and the original element is returned in the result, fully typed. +Supports `topN`, per-request cancellation via `abortSignal`, and the standard +observe-only `GenerationMiddleware` (`onStart`/`onUsage`/`onFinish`/`onAbort`/ +`onError`) plus `rerank:*` devtools events. Rerank bills in provider-defined +search units, surfaced on `usage.unitsBilled`. + +The first adapter ships in the new `@tanstack/ai-cohere` package as +`cohereRerank` / `createCohereRerank`. diff --git a/docs/adapters/cohere.md b/docs/adapters/cohere.md new file mode 100644 index 000000000..8093ceec7 --- /dev/null +++ b/docs/adapters/cohere.md @@ -0,0 +1,132 @@ +--- +title: Cohere +id: cohere-adapter +order: 11 +description: "Rerank documents by relevance to a query with Cohere's rerank models in TanStack AI via the @tanstack/ai-cohere adapter." +keywords: + - tanstack ai + - cohere + - rerank + - reranking + - relevance + - retrieval + - adapter +--- + +The Cohere adapter is **rerank-focused**. It exposes one capability: + +- **Reranking** (`cohereRerank`) — reorder documents by relevance to a query via `rerank()`. + +It does not support text `chat()`, `summarize()`, embeddings, or media — use +OpenAI, Anthropic, or Gemini for those. The adapter talks to Cohere's +`/v2/rerank` endpoint directly over `fetch` (no SDK dependency). + +## Installation + +```bash +npm install @tanstack/ai-cohere +``` + +Peer dependency: + +```bash +npm install @tanstack/ai +``` + +## Basic Usage + +```typescript +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' + +const { rerankedDocuments } = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query: 'talk about rain', + documents: ['sunny day at the beach', 'rainy afternoon in the city'], +}) + +console.log(rerankedDocuments[0]) // 'rainy afternoon in the city' +``` + +For the full reranking guide — object documents, RAG pipelines, options, and +the result shape — see [Reranking](../rerank/rerank). + +## Models + +| Model | Description | +| -------------------------- | ---------------------------------------- | +| `rerank-v3.5` | Latest multilingual reranker (recommended) | +| `rerank-english-v3.0` | English-optimized reranker | +| `rerank-multilingual-v3.0` | Multilingual reranker | + +## Configuration + +`cohereRerank(model, config?)` reads `COHERE_API_KEY` from the environment. +`config` accepts: + +| Option | Type | Default | Description | +| --------- | -------------------------- | -------------------------- | ------------------------------------ | +| `baseUrl` | `string` | `https://api.cohere.com` | Override the API base URL | +| `headers` | `Record` | — | Extra headers merged into requests | + +### Provider Options + +Per-request options are passed via `modelOptions` on `rerank()`: + +```typescript +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' + +const { ranking } = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query: 'refund policy', + documents: ['Returns accepted within 30 days.', 'Free shipping over $50.'], + modelOptions: { + maxTokensPerDoc: 512, // Cap tokens kept per document (Cohere default: 4096) + }, +}) + +console.log(ranking) +``` + +## Explicit API Keys + +To pass an API key directly instead of reading the environment: + +```typescript +import { createCohereRerank } from '@tanstack/ai-cohere' + +const adapter = createCohereRerank('rerank-v3.5', 'your-cohere-api-key') +``` + +## Environment Variables + +```bash +COHERE_API_KEY=your-cohere-api-key +``` + +| Variable | Required | Description | +| ---------------- | -------- | ------------------- | +| `COHERE_API_KEY` | Yes | Your Cohere API key | + +Get your API key from the [Cohere dashboard](https://dashboard.cohere.com/). + +## API Reference + +### `cohereRerank(model, config?)` + +Creates a Cohere rerank adapter for use with `rerank()`, reading +`COHERE_API_KEY` from the environment. + +### `createCohereRerank(model, apiKey, config?)` + +Same as `cohereRerank`, but takes an explicit API key. + +## Limitations + +- **Rerank only** — Use OpenAI, Anthropic, or Gemini for `chat()`, `summarize()`, embeddings, or media generation. + +## Next Steps + +- [Reranking Guide](../rerank/rerank) — full walkthrough including RAG pipelines +- [OpenAI Adapter](./openai) — text, embeddings, and media diff --git a/docs/config.json b/docs/config.json index 86021af3f..a07271efe 100644 --- a/docs/config.json +++ b/docs/config.json @@ -277,6 +277,16 @@ } ] }, + { + "label": "Reranking", + "children": [ + { + "label": "Reranking", + "to": "rerank/rerank", + "addedAt": "2026-06-25" + } + ] + }, { "label": "Middleware", "children": [ @@ -468,6 +478,11 @@ "to": "adapters/openai-compatible", "addedAt": "2026-06-01", "updatedAt": "2026-06-20" + }, + { + "label": "Cohere", + "to": "adapters/cohere", + "addedAt": "2026-06-25" } ] }, diff --git a/docs/rerank/rerank.md b/docs/rerank/rerank.md new file mode 100644 index 000000000..77f1f31e4 --- /dev/null +++ b/docs/rerank/rerank.md @@ -0,0 +1,303 @@ +--- +title: Reranking +id: rerank +order: 1 +description: "Reorder candidate documents by relevance to a query with TanStack AI's rerank() API and the Cohere adapter — the precision step for RAG and search." +keywords: + - tanstack ai + - rerank + - reranking + - relevance + - rag + - retrieval + - semantic search + - cohere +--- + +# Reranking + +You have a query and a list of candidate documents — chunks from a vector +search, rows from a keyword query, FAQ entries — and you need them ordered by +how well they actually answer the query. Vector similarity gets you close, but +a dedicated reranking model is far more precise. By the end of this guide +you'll have that list reordered, with a relevance score per document. + +`rerank()` is the precision step in a retrieval pipeline: retrieve a broad set +of candidates cheaply, then rerank to surface the few that matter. + +## Installation + +Reranking is provided by the Cohere adapter: + +```bash +npm install @tanstack/ai-cohere +``` + +Peer dependency: + +```bash +npm install @tanstack/ai +``` + +## Basic Usage + +Pass a `query` and an array of `documents`. The result's `rerankedDocuments` +are ordered most-relevant first, and `ranking` carries the relevance score and +the original index of each. + +```typescript +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' + +const { ranking, rerankedDocuments } = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query: 'talk about rain', + documents: ['sunny day at the beach', 'rainy afternoon in the city'], + topN: 2, +}) + +console.log(rerankedDocuments[0]) // 'rainy afternoon in the city' +console.log(ranking[0]) // { index: 1, score: 0.98, document: 'rainy afternoon in the city' } +``` + +The adapter reads `COHERE_API_KEY` from the environment. To pass a key +explicitly, use `createCohereRerank('rerank-v3.5', 'co-...')`. + +## Reranking Object Documents + +Documents don't have to be strings. Pass JSON-serializable objects and the +original object is returned in the result — fully typed — so you can carry an +id or metadata through the rerank and read it back off the ranked results. + +```typescript +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' + +const chunks = [ + { id: 'doc-1', text: 'A heavy gaming desktop with an RTX card.' }, + { id: 'doc-2', text: 'A lightweight ultrabook with all-day battery.' }, +] + +const { ranking } = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query: 'best laptop for travel', + documents: chunks, +}) + +// `document` is the original object — `id` is available and type-safe. +console.log(ranking[0]?.document.id) // 'doc-2' +``` + +Object documents are serialized to JSON before being sent to the provider; the +ranking is mapped back to your original elements by index. + +## Options + +| Option | Type | Description | +| ------------- | ----------------------------- | ------------------------------------------------------------------------ | +| `adapter` | `RerankAdapter` | A rerank adapter created with a model (e.g. `cohereRerank('rerank-v3.5')`) | +| `query` | `string` | The search query documents are scored against — required | +| `documents` | `Array` | Candidate documents to rerank — required | +| `topN` | `number` | Return only the top N results | +| `abortSignal` | `AbortSignal` | Cancel the in-flight request | +| `modelOptions`| provider options | Provider-specific options (see below) | +| `middleware` | `Array` | Observe-only lifecycle hooks (usage, finish, error, abort) | + +### Provider Options + +Cohere rerank accepts: + +```typescript +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' + +const { ranking } = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query: 'refund policy', + documents: ['Returns accepted within 30 days.', 'Free shipping over $50.'], + modelOptions: { + // Cap tokens kept per document when chunking long inputs (Cohere default: 4096). + maxTokensPerDoc: 512, + }, +}) + +console.log(ranking) +``` + +## Result Shape + +```typescript +import type { TokenUsage } from '@tanstack/ai' + +interface RerankResult { + id: string + model: string + // Scored results, most relevant first. + ranking: Array<{ index: number; score: number; document: TDocument }> + // The documents reordered by relevance (ranking.map(r => r.document)). + rerankedDocuments: Array + // Rerank bills in provider "search units", surfaced on usage.unitsBilled; + // token counts are 0. + usage: TokenUsage +} +``` + +## Server Endpoint + +Reranking runs on the server (it needs your API key). Wrap it in an API route +and call it from the client over `fetch`: + +```typescript ignore +// routes/api/rerank.ts +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' +import { createFileRoute } from '@tanstack/react-router' + +export const Route = createFileRoute('/api/rerank')({ + server: { + handlers: { + POST: async ({ request }) => { + const body = await request.json() + const { query, documents, topN } = body as { + query: string + documents: Array + topN?: number + } + + const result = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query, + documents, + topN, + }) + + return Response.json(result) + }, + }, + }, +}) +``` + +```typescript ignore +// client.ts — call the endpoint and use the reordered documents +async function rerankDocuments(query: string, documents: Array) { + const res = await fetch('/api/rerank', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ query, documents, topN: 3 }), + }) + const result = await res.json() + return result.rerankedDocuments +} +``` + +## In a RAG Pipeline + +Reranking shines as the second stage after a cheap, broad retrieval. Over-fetch +candidates with vector search, then rerank to keep only the most relevant few +for the prompt: + +```typescript ignore +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' +import { vectorSearch } from './my-vector-store' + +async function retrieveContext(query: string) { + // 1. Over-fetch candidates cheaply. + const candidates = await vectorSearch(query, { limit: 50 }) + + // 2. Rerank and keep the most relevant handful for the prompt. + const { rerankedDocuments } = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query, + documents: candidates.map((c) => c.text), + topN: 5, + }) + + return rerankedDocuments +} +``` + +## Cancellation + +Pass an `abortSignal` to cancel an in-flight request: + +```typescript +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' + +const controller = new AbortController() +setTimeout(() => controller.abort(), 5000) + +const result = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query: 'q', + documents: ['a', 'b'], + abortSignal: controller.signal, +}) + +console.log(result.rerankedDocuments) +``` + +## Observability + +Attach observe-only middleware to track usage, completion, errors, and +cancellation — the same `GenerationMiddleware` contract the media activities +use: + +```typescript +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' + +const result = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query: 'q', + documents: ['a', 'b'], + middleware: [ + { + name: 'usage-logger', + onUsage: (_ctx, usage) => { + console.log('search units billed:', usage.unitsBilled) + }, + }, + ], +}) + +console.log(result.rerankedDocuments) +``` + +> **Tip:** Pass `otelMiddleware()` to emit OpenTelemetry spans for rerank +> calls. See [OpenTelemetry](../advanced/otel). + +## Environment Variables + +The Cohere rerank adapter uses: + +- `COHERE_API_KEY`: Your Cohere API key + +## Error Handling + +```typescript +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' + +try { + const result = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query: 'q', + documents: ['a', 'b'], + }) + console.log(result.rerankedDocuments) +} catch (error) { + if (error instanceof Error) { + console.error('Rerank failed:', error.message) + } +} +``` + +> Passing an empty `documents` array throws before any request is made. + +## Next Steps + +- [Cohere Adapter](../adapters/cohere) — models, configuration, and explicit API keys +- [Middleware](../advanced/middleware) — lifecycle hooks for usage and errors diff --git a/packages/ai-cohere/LICENSE b/packages/ai-cohere/LICENSE new file mode 100644 index 000000000..308cb68dc --- /dev/null +++ b/packages/ai-cohere/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Tanner Linsley + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/packages/ai-cohere/README.md b/packages/ai-cohere/README.md new file mode 100644 index 000000000..1fe07f546 --- /dev/null +++ b/packages/ai-cohere/README.md @@ -0,0 +1,55 @@ +
+ TanStack AI +
+ +
+ + + +# @tanstack/ai-cohere + +Cohere adapter for [TanStack AI](https://tanstack.com/ai). Reorder candidate +documents by relevance to a query with Cohere's rerank models — the precision +step for RAG and search pipelines. + +This adapter is **rerank-only**. For chat, summarization, embeddings, or media, +use OpenAI, Anthropic, or Gemini. + +## Install + +```bash +pnpm add @tanstack/ai @tanstack/ai-cohere +``` + +## Usage + +```typescript +import { rerank } from '@tanstack/ai' +import { cohereRerank } from '@tanstack/ai-cohere' + +const { ranking, rerankedDocuments } = await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query: 'talk about rain', + documents: ['sunny day at the beach', 'rainy afternoon in the city'], + topN: 2, +}) + +console.log(rerankedDocuments[0]) // 'rainy afternoon in the city' +``` + +The adapter reads `COHERE_API_KEY` from the environment. To pass a key +explicitly, use `createCohereRerank('rerank-v3.5', 'co-...')`. + +## Read the docs -> + +- [Reranking Guide](https://tanstack.com/ai/latest/docs/rerank/rerank) — object + documents, RAG pipelines, options, and the result shape. +- [Cohere Adapter](https://tanstack.com/ai/latest/docs/adapters/cohere) — + models, configuration, and explicit API keys. diff --git a/packages/ai-cohere/package.json b/packages/ai-cohere/package.json new file mode 100644 index 000000000..909799cb0 --- /dev/null +++ b/packages/ai-cohere/package.json @@ -0,0 +1,63 @@ +{ + "name": "@tanstack/ai-cohere", + "version": "0.0.0", + "description": "Cohere adapter for TanStack AI — document reranking.", + "author": "Tanner Linsley", + "license": "MIT", + "homepage": "https://tanstack.com/ai", + "repository": { + "type": "git", + "url": "git+https://github.com/TanStack/ai.git", + "directory": "packages/ai-cohere" + }, + "bugs": { + "url": "https://github.com/TanStack/ai/issues" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + }, + "type": "module", + "module": "./dist/esm/index.js", + "types": "./dist/esm/index.d.ts", + "exports": { + ".": { + "types": "./dist/esm/index.d.ts", + "import": "./dist/esm/index.js" + } + }, + "files": [ + "dist", + "src" + ], + "scripts": { + "build": "vite build", + "clean": "premove ./build ./dist", + "lint:fix": "eslint ./src --fix", + "test:build": "publint --strict", + "test:eslint": "eslint ./src", + "test:lib": "vitest --passWithNoTests", + "test:lib:dev": "pnpm test:lib --watch", + "test:types": "tsc" + }, + "keywords": [ + "ai", + "ai-sdk", + "typescript", + "tanstack", + "cohere", + "rerank", + "reranking", + "search", + "retrieval", + "adapter" + ], + "peerDependencies": { + "@tanstack/ai": "workspace:^" + }, + "devDependencies": { + "@tanstack/ai": "workspace:*", + "@vitest/coverage-v8": "4.0.14", + "vite": "^7.3.3" + } +} diff --git a/packages/ai-cohere/src/adapters/rerank.ts b/packages/ai-cohere/src/adapters/rerank.ts new file mode 100644 index 000000000..5be678f49 --- /dev/null +++ b/packages/ai-cohere/src/adapters/rerank.ts @@ -0,0 +1,177 @@ +import { BaseRerankAdapter } from '@tanstack/ai/adapters' +import { + COHERE_DEFAULT_BASE_URL, + getCohereApiKeyFromEnv, +} from '../utils/client' +import type { CohereClientConfig } from '../utils/client' +import type { + CohereRerankModel, + InferCohereRerankProviderOptions, +} from '../model-meta' +import type { RerankAdapterResult, RerankOptions, TokenUsage } from '@tanstack/ai' + +/** Shape of the Cohere `/v2/rerank` response we depend on. */ +interface CohereRerankResponse { + id?: string + results: Array<{ index: number; relevance_score: number }> + meta?: { billed_units?: { search_units?: number } } +} + +function isCohereRerankResponse(value: unknown): value is CohereRerankResponse { + if (typeof value !== 'object' || value === null) return false + const results = (value as { results?: unknown }).results + return ( + Array.isArray(results) && + results.every( + (r) => + typeof r === 'object' && + r !== null && + typeof (r as { index?: unknown }).index === 'number' && + typeof (r as { relevance_score?: unknown }).relevance_score === + 'number', + ) + ) +} + +/** + * Cohere rerank adapter. + * + * Talks to Cohere's `/v2/rerank` endpoint over raw `fetch` — no SDK. Returns + * scored indices into the submitted documents; the `rerank()` activity maps + * those back to the caller's original documents. + */ +export class CohereRerankAdapter< + TModel extends CohereRerankModel, +> extends BaseRerankAdapter> { + readonly name = 'cohere' as const + + private readonly apiKey: string + private readonly baseUrl: string + private readonly headers: Record + + constructor(config: CohereClientConfig, model: TModel) { + super({}, model) + this.apiKey = config.apiKey + this.baseUrl = (config.baseUrl ?? COHERE_DEFAULT_BASE_URL).replace( + /\/+$/, + '', + ) + this.headers = config.headers ?? {} + } + + async rerank( + options: RerankOptions>, + ): Promise { + const { model, query, documents, topN, modelOptions, abortSignal, logger } = + options + + const body: Record = { model, query, documents } + if (topN !== undefined) body['top_n'] = topN + if (modelOptions?.maxTokensPerDoc !== undefined) { + body['max_tokens_per_doc'] = modelOptions.maxTokensPerDoc + } + + logger.request( + `activity=rerank provider=${this.name} model=${model} documents=${documents.length}`, + { provider: this.name, model }, + ) + + let response: Response + try { + response = await fetch(`${this.baseUrl}/v2/rerank`, { + method: 'POST', + headers: { + Authorization: `Bearer ${this.apiKey}`, + 'Content-Type': 'application/json', + ...this.headers, + }, + body: JSON.stringify(body), + ...(abortSignal ? { signal: abortSignal } : {}), + }) + } catch (error) { + logger.errors(`${this.name}.rerank fatal`, { + error, + source: `${this.name}.rerank`, + }) + throw error + } + + if (!response.ok) { + const detail = await response.text().catch(() => '') + const error = new Error( + `Cohere rerank request failed: ${response.status} ${response.statusText}${ + detail ? ` — ${detail}` : '' + }`, + ) + logger.errors(`${this.name}.rerank fatal`, { + error, + source: `${this.name}.rerank`, + }) + throw error + } + + const json: unknown = await response.json() + if (!isCohereRerankResponse(json)) { + throw new Error('Cohere rerank response had an unexpected shape') + } + + const searchUnits = json.meta?.billed_units?.search_units + const usage: TokenUsage = { + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + ...(searchUnits !== undefined ? { unitsBilled: searchUnits } : {}), + } + + return { + id: json.id ?? this.generateId(), + ranking: json.results.map((r) => ({ + index: r.index, + score: r.relevance_score, + })), + usage, + } + } +} + +/** + * Creates a Cohere rerank adapter with an explicit API key. Type resolution + * (per-model provider options) happens here at the call site. + * + * @example + * ```typescript + * const adapter = createCohereRerank('rerank-v3.5', 'co-...') + * ``` + */ +export function createCohereRerank( + model: TModel, + apiKey: string, + config?: Omit, +): CohereRerankAdapter { + return new CohereRerankAdapter({ apiKey, ...config }, model) +} + +/** + * Creates a Cohere rerank adapter, reading `COHERE_API_KEY` from the + * environment. + * + * @throws Error if `COHERE_API_KEY` is not found. + * + * @example + * ```typescript + * import { rerank } from '@tanstack/ai' + * import { cohereRerank } from '@tanstack/ai-cohere' + * + * const { rerankedDocuments } = await rerank({ + * adapter: cohereRerank('rerank-v3.5'), + * query: 'talk about rain', + * documents: ['sunny day', 'rainy afternoon'], + * }) + * ``` + */ +export function cohereRerank( + model: TModel, + config?: Omit, +): CohereRerankAdapter { + return createCohereRerank(model, getCohereApiKeyFromEnv(), config) +} diff --git a/packages/ai-cohere/src/index.ts b/packages/ai-cohere/src/index.ts new file mode 100644 index 000000000..81fcece02 --- /dev/null +++ b/packages/ai-cohere/src/index.ts @@ -0,0 +1,24 @@ +// ============================================================================ +// Cohere Adapters (tree-shakeable) +// ============================================================================ + +// Rerank adapter - document reranking via Cohere's /v2/rerank endpoint +export { + CohereRerankAdapter, + createCohereRerank, + cohereRerank, +} from './adapters/rerank' + +// ============================================================================ +// Type Exports +// ============================================================================ + +export { + COHERE_RERANK_MODELS, + type CohereRerankModel, + type CohereRerankProviderOptions, + type CohereRerankModelProviderOptionsByName, + type InferCohereRerankProviderOptions, +} from './model-meta' + +export type { CohereClientConfig } from './utils/client' diff --git a/packages/ai-cohere/src/model-meta.ts b/packages/ai-cohere/src/model-meta.ts new file mode 100644 index 000000000..144086a16 --- /dev/null +++ b/packages/ai-cohere/src/model-meta.ts @@ -0,0 +1,50 @@ +/** + * Cohere rerank model metadata. + * + * Provider options are resolved per model at the `cohereRerank('model')` call + * site via {@link CohereRerankModelProviderOptionsByName}. Cohere's rerank + * models currently share the same options, but the per-model map keeps the + * surface symmetric with the other adapters and lets divergent options be + * expressed later without changing the adapter contract. + */ + +/** Available Cohere rerank models. */ +export const COHERE_RERANK_MODELS = [ + 'rerank-v3.5', + 'rerank-english-v3.0', + 'rerank-multilingual-v3.0', +] as const + +/** Union of supported Cohere rerank model names. */ +export type CohereRerankModel = (typeof COHERE_RERANK_MODELS)[number] + +/** + * Provider-specific options for a Cohere rerank request. Forwarded on the + * `modelOptions` field of `rerank()`. + */ +export interface CohereRerankProviderOptions { + /** + * Long documents are chunked to fit the model's context. This caps the + * number of tokens kept per document. Cohere defaults to 4096. + */ + maxTokensPerDoc?: number +} + +/** + * Per-model provider-options map. Each model resolves to its own options type + * at the factory call site (see {@link InferCohereRerankProviderOptions}). + */ +export interface CohereRerankModelProviderOptionsByName { + 'rerank-v3.5': CohereRerankProviderOptions + 'rerank-english-v3.0': CohereRerankProviderOptions + 'rerank-multilingual-v3.0': CohereRerankProviderOptions +} + +/** + * Resolve the provider options for a given rerank model. Falls back to the + * base options for any model not in the map. + */ +export type InferCohereRerankProviderOptions = + TModel extends keyof CohereRerankModelProviderOptionsByName + ? CohereRerankModelProviderOptionsByName[TModel] + : CohereRerankProviderOptions diff --git a/packages/ai-cohere/src/utils/client.ts b/packages/ai-cohere/src/utils/client.ts new file mode 100644 index 000000000..9c19a668b --- /dev/null +++ b/packages/ai-cohere/src/utils/client.ts @@ -0,0 +1,43 @@ +/** + * Cohere client configuration shared by the rerank adapter. + */ +export interface CohereClientConfig { + /** Cohere API key. Required by the adapter factories. */ + apiKey: string + /** Override the API base URL. Defaults to `https://api.cohere.com`. */ + baseUrl?: string + /** Extra headers merged into every request. */ + headers?: Record +} + +export const COHERE_DEFAULT_BASE_URL = 'https://api.cohere.com' + +/** + * Reads the Cohere API key from the environment. + * + * Looks for `COHERE_API_KEY` in `process.env` (Node) or `window.env` + * (browser with injected env). + * + * @throws Error if `COHERE_API_KEY` is not found. + */ +export function getCohereApiKeyFromEnv(): string { + const env = + typeof globalThis !== 'undefined' && + (globalThis as Record).window + ? (( + (globalThis as Record).window as Record< + string, + unknown + > + ).env as Record | undefined) + : typeof process !== 'undefined' + ? process.env + : undefined + const key = env?.['COHERE_API_KEY'] + if (!key) { + throw new Error( + 'COHERE_API_KEY not found in environment. Pass an API key explicitly via createCohereRerank(model, apiKey).', + ) + } + return key +} diff --git a/packages/ai-cohere/tests/rerank-adapter.test.ts b/packages/ai-cohere/tests/rerank-adapter.test.ts new file mode 100644 index 000000000..c3b0c8a9e --- /dev/null +++ b/packages/ai-cohere/tests/rerank-adapter.test.ts @@ -0,0 +1,133 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { rerank } from '@tanstack/ai' +import { createCohereRerank } from '../src/adapters/rerank' + +const fetchMock = vi.fn() + +beforeEach(() => { + vi.stubGlobal('fetch', fetchMock) + fetchMock.mockReset() +}) + +afterEach(() => { + vi.unstubAllGlobals() +}) + +/** A 200 response carrying a Cohere-shaped rerank body. */ +function cohereResponse(body: unknown): Response { + return new Response(JSON.stringify(body), { + status: 200, + headers: { 'content-type': 'application/json' }, + }) +} + +/** Default well-formed Cohere rerank payload reordering [1, 0]. */ +function defaultBody() { + return { + id: 'cohere-1', + results: [ + { index: 1, relevance_score: 0.98 }, + { index: 0, relevance_score: 0.12 }, + ], + meta: { billed_units: { search_units: 1 } }, + } +} + +/** The parsed request body of the most recent fetch call. */ +function lastRequestBody() { + const init = fetchMock.mock.calls[0]![1] + return JSON.parse(String(init?.body)) +} + +const adapter = () => createCohereRerank('rerank-v3.5', 'test-key') +const documents = ['sunny day at the beach', 'rainy afternoon in the city'] + +describe('CohereRerankAdapter', () => { + it('POSTs to /v2/rerank with auth and the expected request body', async () => { + fetchMock.mockResolvedValue(cohereResponse(defaultBody())) + + await rerank({ + adapter: adapter(), + query: 'talk about rain', + documents, + topN: 2, + modelOptions: { maxTokensPerDoc: 512 }, + }) + + const [url, init] = fetchMock.mock.calls[0]! + expect(url).toBe('https://api.cohere.com/v2/rerank') + expect(init?.method).toBe('POST') + expect(new Headers(init?.headers).get('Authorization')).toBe( + 'Bearer test-key', + ) + expect(lastRequestBody()).toEqual({ + model: 'rerank-v3.5', + query: 'talk about rain', + documents, + top_n: 2, + max_tokens_per_doc: 512, + }) + }) + + it('maps results to ranking and search_units to usage.unitsBilled', async () => { + fetchMock.mockResolvedValue(cohereResponse(defaultBody())) + + const result = await rerank({ + adapter: adapter(), + query: 'talk about rain', + documents, + }) + + expect(result.id).toBe('cohere-1') + expect(result.ranking).toEqual([ + { index: 1, score: 0.98, document: documents[1] }, + { index: 0, score: 0.12, document: documents[0] }, + ]) + expect(result.usage.unitsBilled).toBe(1) + expect(result.usage.totalTokens).toBe(0) + }) + + it('omits top_n and max_tokens_per_doc when not provided', async () => { + fetchMock.mockResolvedValue(cohereResponse(defaultBody())) + + await rerank({ adapter: adapter(), query: 'q', documents }) + + expect(lastRequestBody()).toEqual({ + model: 'rerank-v3.5', + query: 'q', + documents, + }) + }) + + it('throws with status detail on a non-200 response', async () => { + fetchMock.mockResolvedValue( + new Response('rate limited', { status: 429, statusText: 'Too Many Requests' }), + ) + + await expect( + rerank({ adapter: adapter(), query: 'q', documents, debug: false }), + ).rejects.toThrow('429') + }) + + it('throws when the response shape is unexpected', async () => { + fetchMock.mockResolvedValue(cohereResponse({ nope: true })) + + await expect( + rerank({ adapter: adapter(), query: 'q', documents, debug: false }), + ).rejects.toThrow('unexpected shape') + }) + + it('forwards the abort signal to fetch', async () => { + fetchMock.mockResolvedValue(cohereResponse(defaultBody())) + const controller = new AbortController() + + await rerank({ + adapter: adapter(), + query: 'q', + documents, + abortSignal: controller.signal, + }) + + expect(fetchMock.mock.calls[0]![1]?.signal).toBe(controller.signal) + }) +}) diff --git a/packages/ai-cohere/tsconfig.json b/packages/ai-cohere/tsconfig.json new file mode 100644 index 000000000..c38689f4e --- /dev/null +++ b/packages/ai-cohere/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "dist" + }, + "include": ["src", "tests"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/ai-cohere/vite.config.ts b/packages/ai-cohere/vite.config.ts new file mode 100644 index 000000000..77bcc2e60 --- /dev/null +++ b/packages/ai-cohere/vite.config.ts @@ -0,0 +1,36 @@ +import { defineConfig, mergeConfig } from 'vitest/config' +import { tanstackViteConfig } from '@tanstack/vite-config' +import packageJson from './package.json' + +const config = defineConfig({ + test: { + name: packageJson.name, + dir: './', + watch: false, + globals: true, + environment: 'node', + include: ['tests/**/*.test.ts'], + coverage: { + provider: 'v8', + reporter: ['text', 'json', 'html', 'lcov'], + exclude: [ + 'node_modules/', + 'dist/', + 'tests/', + '**/*.test.ts', + '**/*.config.ts', + '**/types.ts', + ], + include: ['src/**/*.ts'], + }, + }, +}) + +export default mergeConfig( + config, + tanstackViteConfig({ + entry: ['./src/index.ts'], + srcDir: './src', + cjs: false, + }), +) diff --git a/packages/ai-event-client/src/index.ts b/packages/ai-event-client/src/index.ts index 0bd1fdc12..fa739b4bf 100644 --- a/packages/ai-event-client/src/index.ts +++ b/packages/ai-event-client/src/index.ts @@ -607,6 +607,38 @@ export interface SummarizeUsageEvent extends BaseEventContext { usage: TokenUsage } +// =========================== +// Rerank Events +// =========================== + +/** Emitted when a rerank request starts. */ +export interface RerankRequestStartedEvent extends BaseEventContext { + requestId: string + provider: string + model: string + /** Number of documents submitted for reranking. */ + documentCount: number +} + +/** Emitted when rerank completes. */ +export interface RerankRequestCompletedEvent extends BaseEventContext { + requestId: string + provider: string + model: string + /** Number of documents submitted for reranking. */ + documentCount: number + /** Number of ranked results returned (after any `topN`). */ + resultCount: number + duration: number +} + +/** Emitted when rerank usage metrics are available. */ +export interface RerankUsageEvent extends BaseEventContext { + requestId: string + model: string + usage: TokenUsage +} + // =========================== // Image Events // =========================== @@ -1015,6 +1047,11 @@ export interface AIDevtoolsEventMap { 'summarize:request:completed': SummarizeRequestCompletedEvent 'summarize:usage': SummarizeUsageEvent + // Rerank events + 'rerank:request:started': RerankRequestStartedEvent + 'rerank:request:completed': RerankRequestCompletedEvent + 'rerank:usage': RerankUsageEvent + // Image events 'image:request:started': ImageRequestStartedEvent 'image:request:completed': ImageRequestCompletedEvent diff --git a/packages/ai/src/activities/index.ts b/packages/ai/src/activities/index.ts index 07fdbe73a..eef1efc2f 100644 --- a/packages/ai/src/activities/index.ts +++ b/packages/ai/src/activities/index.ts @@ -21,6 +21,7 @@ import type { AnyAudioAdapter } from './generateAudio/adapter' import type { AnyVideoAdapter } from './generateVideo/adapter' import type { AnyTTSAdapter } from './generateSpeech/adapter' import type { AnyTranscriptionAdapter } from './generateTranscription/adapter' +import type { AnyRerankAdapter } from './rerank/adapter' // =========================== // Chat Activity @@ -66,6 +67,25 @@ export { type InferTextProviderOptions, } from './summarize/chat-stream-summarize' +// =========================== +// Rerank Activity +// =========================== + +export { + kind as rerankKind, + rerank, + createRerankOptions, + type RerankActivityOptions, + type RerankProviderOptions, +} from './rerank/index' + +export { + BaseRerankAdapter, + type RerankAdapter, + type RerankAdapterConfig, + type AnyRerankAdapter, +} from './rerank/adapter' + // =========================== // Image Activity // =========================== @@ -183,6 +203,7 @@ export type AIAdapter = | AnyVideoAdapter | AnyTTSAdapter | AnyTranscriptionAdapter + | AnyRerankAdapter /** Union type of all adapter kinds */ export type AdapterKind = @@ -193,3 +214,4 @@ export type AdapterKind = | 'video' | 'tts' | 'transcription' + | 'rerank' diff --git a/packages/ai/src/activities/middleware/types.ts b/packages/ai/src/activities/middleware/types.ts index 99ae3e44e..c3fd7b1f3 100644 --- a/packages/ai/src/activities/middleware/types.ts +++ b/packages/ai/src/activities/middleware/types.ts @@ -38,6 +38,7 @@ export type GenerationActivity = | 'audio' | 'tts' | 'transcription' + | 'rerank' /** * Stable context passed to every {@link GenerationMiddleware} hook. Created diff --git a/packages/ai/src/activities/rerank/adapter.ts b/packages/ai/src/activities/rerank/adapter.ts new file mode 100644 index 000000000..7e5f54bf4 --- /dev/null +++ b/packages/ai/src/activities/rerank/adapter.ts @@ -0,0 +1,91 @@ +import type { RerankAdapterResult, RerankOptions } from '../../types' + +/** + * Configuration for rerank adapter instances + */ +export interface RerankAdapterConfig { + apiKey?: string + baseUrl?: string + timeout?: number + headers?: Record +} + +/** + * Rerank adapter interface with pre-resolved generics. + * + * An adapter is created by a provider function: `provider('model')` → `adapter` + * All type resolution happens at the provider call site, not in this interface. + * + * Generic parameters: + * - TModel: The specific model name (e.g. 'rerank-v3.5') + * - TProviderOptions: Provider-specific options (already resolved) + */ +export interface RerankAdapter< + TModel extends string = string, + TProviderOptions extends object = Record, +> { + /** Discriminator for adapter kind */ + readonly kind: 'rerank' + /** Adapter name identifier */ + readonly name: string + /** The model this adapter is configured for */ + readonly model: TModel + + /** + * @internal Type-only properties for inference. Not assigned at runtime. + */ + '~types': { + providerOptions: TProviderOptions + } + + /** + * Rerank the given (pre-serialized) documents against the query, returning + * scored indices into `options.documents`. The activity layer maps these + * back to the caller's original documents. + */ + rerank: ( + options: RerankOptions, + ) => Promise +} + +/** + * A RerankAdapter with any/unknown type parameters. + * Useful as a constraint in generic functions and interfaces. + */ +export type AnyRerankAdapter = RerankAdapter + +/** + * Abstract base class for rerank adapters. + * Extend this class to implement a rerank adapter for a specific provider. + * + * Generic parameters match RerankAdapter - all pre-resolved by the provider function. + */ +export abstract class BaseRerankAdapter< + TModel extends string = string, + TProviderOptions extends object = Record, +> implements RerankAdapter +{ + readonly kind = 'rerank' as const + abstract readonly name: string + readonly model: TModel + + // Type-only property - never assigned at runtime + declare '~types': { + providerOptions: TProviderOptions + } + + protected config: RerankAdapterConfig + + constructor(config: RerankAdapterConfig = {}, model: TModel) { + this.config = config + this.model = model + } + + abstract rerank( + options: RerankOptions, + ): Promise + + protected generateId(): string { + return `${this.name}-${Date.now()}-${Math.random().toString(36).substring(7)}` + } +} diff --git a/packages/ai/src/activities/rerank/index.ts b/packages/ai/src/activities/rerank/index.ts new file mode 100644 index 000000000..b0cf8c9da --- /dev/null +++ b/packages/ai/src/activities/rerank/index.ts @@ -0,0 +1,290 @@ +/** + * Rerank Activity + * + * Reorders a set of documents by semantic relevance to a query. + * This is a self-contained module with implementation, types, and JSDoc. + */ + +import { aiEventClient } from '@tanstack/ai-event-client' +import { resolveDebugOption } from '../../logger/resolve' +import { + createGenerationContext, + runGenerationAbort, + runGenerationError, + runGenerationFinish, + runGenerationStart, + runGenerationUsage, +} from '../middleware' +import type { InternalLogger } from '../../logger/internal-logger' +import type { DebugOption } from '../../logger/types' +import type { GenerationMiddleware } from '../middleware' +import type { RerankAdapter } from './adapter' +import type { RerankResult } from '../../types' + +// =========================== +// Activity Kind +// =========================== + +/** The adapter kind this activity handles */ +export const kind = 'rerank' as const + +// =========================== +// Type Extraction Helpers +// =========================== + +/** Extract provider options from a RerankAdapter via ~types */ +export type RerankProviderOptions = + TAdapter extends RerankAdapter + ? TAdapter['~types']['providerOptions'] + : object + +// =========================== +// Activity Options Type +// =========================== + +/** + * Options for the rerank activity. The model is extracted from the adapter's + * model property. + * + * @template TAdapter - The rerank adapter type + * @template TDocument - The document element type (string or object) + */ +export interface RerankActivityOptions< + TAdapter extends RerankAdapter, + TDocument extends string | object = string, +> { + /** The rerank adapter to use (must be created with a model) */ + adapter: TAdapter & { kind: typeof kind } + /** The query documents are scored against. */ + query: string + /** + * Documents to rerank. Either strings or JSON-serializable objects — object + * documents are serialized with `JSON.stringify` before being sent to the + * provider, and the original element (string or object) is returned in the + * result, preserving its type. + */ + documents: Array + /** Return only the top N results. */ + topN?: number + /** Provider-specific options */ + modelOptions?: RerankProviderOptions + /** Forwarded to the provider request for cancellation. */ + abortSignal?: AbortSignal + /** + * Observe-only middleware notified on start, usage, success, abort, and + * error. Pass `otelMiddleware()` to emit OpenTelemetry spans, or implement + * the `GenerationMiddleware` contract for a custom backend. + */ + middleware?: Array + /** + * Enable debug logging. Pass `true` to enable all categories, `false` to + * silence everything including errors, or a `DebugConfig` object for granular + * control and/or a custom `Logger`. + */ + debug?: DebugOption +} + +// =========================== +// Helper Functions +// =========================== + +function createId(prefix: string): string { + return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}` +} + +/** Serialize a document for the provider. Strings pass through untouched. */ +function serializeDocument(document: string | object): string { + return typeof document === 'string' ? document : JSON.stringify(document) +} + +function isAbortError(error: unknown, signal?: AbortSignal): boolean { + return ( + signal?.aborted === true || + (error instanceof Error && error.name === 'AbortError') + ) +} + +// =========================== +// Activity Implementation +// =========================== + +/** + * Rerank activity - reorders documents by relevance to a query. + * + * @example Basic reranking + * ```ts + * import { rerank } from '@tanstack/ai' + * import { cohereRerank } from '@tanstack/ai-cohere' + * + * const { ranking, rerankedDocuments } = await rerank({ + * adapter: cohereRerank('rerank-v3.5'), + * query: 'talk about rain', + * documents: ['sunny day at the beach', 'rainy afternoon in the city'], + * topN: 2, + * }) + * + * console.log(rerankedDocuments[0]) // 'rainy afternoon in the city' + * ``` + * + * @example Reranking object documents + * ```ts + * const { ranking } = await rerank({ + * adapter: cohereRerank('rerank-v3.5'), + * query: 'best laptop for travel', + * documents: [ + * { id: 1, text: 'A heavy gaming desktop' }, + * { id: 2, text: 'A lightweight ultrabook with all-day battery' }, + * ], + * }) + * + * // ranking[0].document is the original object, fully typed. + * console.log(ranking[0].document.id) + * ``` + */ +export async function rerank< + TAdapter extends RerankAdapter, + TDocument extends string | object = string, +>( + options: RerankActivityOptions, +): Promise> { + const { + adapter, + query, + documents, + topN, + modelOptions, + abortSignal, + middleware, + } = options + const model = adapter.model + const requestId = createId('rerank') + const startTime = Date.now() + const logger: InternalLogger = resolveDebugOption(options.debug) + + if (documents.length === 0) { + throw new Error('rerank() requires at least one document') + } + + const mwCtx = createGenerationContext({ + requestId, + // `rerank` joins the GenerationActivity union; otel maps it to its own + // gen_ai.operation.name. + activity: 'rerank', + provider: adapter.name, + model, + modelOptions, + createId, + }) + + await runGenerationStart(middleware, mwCtx) + + aiEventClient.emit('rerank:request:started', { + requestId, + provider: adapter.name, + model, + documentCount: documents.length, + timestamp: startTime, + }) + + logger.request(`activity=rerank provider=${adapter.name}`, { + provider: adapter.name, + model, + documentCount: documents.length, + }) + + // Serialize once; reuse for the request only. Original documents are mapped + // back by index below so the caller's element type is preserved. + const serialized = documents.map(serializeDocument) + + try { + const result = await adapter.rerank({ + model, + query, + documents: serialized, + topN, + modelOptions, + abortSignal, + logger, + }) + + const ranking = result.ranking.map((r) => ({ + index: r.index, + score: r.score, + document: documents[r.index] as TDocument, + })) + const rerankedDocuments = ranking.map((r) => r.document) + + const duration = Date.now() - startTime + + aiEventClient.emit('rerank:request:completed', { + requestId, + provider: adapter.name, + model, + documentCount: documents.length, + resultCount: ranking.length, + duration, + timestamp: Date.now(), + }) + + aiEventClient.emit('rerank:usage', { + requestId, + model, + usage: result.usage, + timestamp: Date.now(), + }) + + logger.output(`activity=rerank results=${ranking.length}`, { + resultCount: ranking.length, + }) + + await runGenerationUsage(middleware, mwCtx, result.usage) + await runGenerationFinish(middleware, mwCtx, { + duration, + usage: result.usage, + }) + + return { + id: result.id, + model, + ranking, + rerankedDocuments, + usage: result.usage, + } + } catch (error) { + const duration = Date.now() - startTime + if (isAbortError(error, abortSignal)) { + await runGenerationAbort(middleware, mwCtx, { + reason: error instanceof Error ? error.message : undefined, + duration, + }) + } else { + await runGenerationError(middleware, mwCtx, { error, duration }) + } + logger.errors('rerank activity failed', { error, source: 'rerank' }) + throw error + } +} + +// =========================== +// Options Factory +// =========================== + +/** + * Create typed options for the rerank() function without executing. + */ +export function createRerankOptions< + TAdapter extends RerankAdapter, + TDocument extends string | object = string, +>( + options: RerankActivityOptions, +): RerankActivityOptions { + return options +} + +// Re-export adapter types +export type { + RerankAdapter, + RerankAdapterConfig, + AnyRerankAdapter, +} from './adapter' +export { BaseRerankAdapter } from './adapter' diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 4a38e2e99..9e9645abe 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -2,6 +2,7 @@ export { chat, summarize, + rerank, generateImage, generateAudio, generateVideo, @@ -13,6 +14,7 @@ export { // Create options functions - for pre-defining typed configurations export { createChatOptions } from './activities/chat/index' export { createSummarizeOptions } from './activities/summarize/index' +export { createRerankOptions } from './activities/rerank/index' export { createImageOptions } from './activities/generateImage/index' export { createAudioOptions } from './activities/generateAudio/index' export { createVideoOptions } from './activities/generateVideo/index' @@ -36,8 +38,13 @@ export type { TranscriptionAdapter, AnyVideoAdapter, VideoAdapter, + AnyRerankAdapter, + RerankAdapter, } from './activities/index' +// Rerank adapter base + types +export { BaseRerankAdapter } from './activities/rerank/adapter' + // Tool definition export { toolDefinition, diff --git a/packages/ai/src/middlewares/otel.ts b/packages/ai/src/middlewares/otel.ts index f4ddba35e..63f0e8fd4 100644 --- a/packages/ai/src/middlewares/otel.ts +++ b/packages/ai/src/middlewares/otel.ts @@ -83,6 +83,7 @@ const OPERATION_NAME: Record = { audio: 'audio_generation', tts: 'text_to_speech', transcription: 'transcription', + rerank: 'rerank', } export interface OtelMiddlewareOptions { diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 3a7fe0807..95cb8b7a1 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -1488,6 +1488,69 @@ export interface SummarizationResult { usage: TokenUsage } +// ============================================================================ +// Rerank Types +// ============================================================================ + +/** + * Options passed to a {@link RerankAdapter}. Documents reach the adapter + * already serialized to strings — the `rerank()` activity stringifies object + * documents and maps results back to the original elements, so adapters never + * deal with the caller's document type. + */ +export interface RerankOptions< + TProviderOptions extends object = Record, +> { + model: string + /** The search query documents are scored against. */ + query: string + /** Documents to rerank, pre-serialized to strings by the activity. */ + documents: Array + /** Return only the top N results. Passed through to the provider. */ + topN?: number + /** Provider-specific options forwarded by the rerank() activity. */ + modelOptions?: TProviderOptions + /** Forwarded to the provider request for cancellation. */ + abortSignal?: AbortSignal + /** + * Internal logger threaded from the rerank() entry point. Adapters must call + * logger.request() before the provider call and logger.errors() in catch + * blocks. + */ + logger: InternalLogger +} + +/** + * Provider-level rerank result. Adapters return scored indices into the + * (serialized) `documents` array plus usage — never the documents themselves. + * The activity attaches the original documents. + */ +export interface RerankAdapterResult { + id: string + /** Scored results, highest relevance first, as indices into `documents`. */ + ranking: Array<{ index: number; score: number }> + usage: TokenUsage +} + +/** + * Public result of the `rerank()` activity, generic over the caller's document + * element type so `document` / `rerankedDocuments` carry the original values + * (strings or objects), not their serialized form. + */ +export interface RerankResult { + id: string + model: string + /** Scored results, highest relevance first. */ + ranking: Array<{ index: number; score: number; document: TDocument }> + /** The documents reordered by relevance — `ranking.map(r => r.document)`. */ + rerankedDocuments: Array + /** + * Usage for the request. Rerank bills in provider-defined "search units" + * rather than tokens, surfaced via `usage.unitsBilled`; token counts are 0. + */ + usage: TokenUsage +} + // ============================================================================ // Image Generation Types // ============================================================================ diff --git a/packages/ai/tests/rerank.test.ts b/packages/ai/tests/rerank.test.ts new file mode 100644 index 000000000..2862068a1 --- /dev/null +++ b/packages/ai/tests/rerank.test.ts @@ -0,0 +1,219 @@ +import { describe, expect, it } from 'vitest' +import { rerank } from '../src/index' +import type { RerankAdapter } from '../src/activities/rerank/adapter' +import type { + GenerationAbortInfo, + GenerationErrorInfo, + GenerationFinishInfo, + GenerationMiddleware, + GenerationMiddlewareContext, + GenerationUsageInfo, +} from '../src/activities/middleware' +import type { RerankAdapterResult, RerankOptions, TokenUsage } from '../src/types' + +// ============================================================================ +// Helpers +// ============================================================================ + +const zeroUsage: TokenUsage = { + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, +} + +/** + * Build a fully-typed mock rerank adapter. `rerankFn` receives the options the + * activity hands the adapter (documents already serialized to strings) and + * returns the provider-level result. + */ +function mockRerankAdapter( + rerankFn: (opts: RerankOptions) => Promise, +): RerankAdapter & { calls: Array> } { + const calls: Array> = [] + return { + kind: 'rerank', + name: 'mock', + model: 'mock-model', + '~types': { providerOptions: {} }, + calls, + rerank: (opts) => { + calls.push(opts) + return rerankFn(opts) + }, + } +} + +/** A scored result built from a descending list of indices. */ +function ranked(...indices: Array): RerankAdapterResult { + return { + id: 'rr-1', + ranking: indices.map((index, i) => ({ index, score: 1 - i * 0.1 })), + usage: { ...zeroUsage, unitsBilled: 1 }, + } +} + +/** Recording middleware capturing each lifecycle hook's context + payload. */ +function recordingMiddleware() { + const events = { + start: [] as Array, + usage: [] as Array, + finish: [] as Array, + abort: [] as Array, + error: [] as Array, + } + const middleware: GenerationMiddleware = { + name: 'rec', + onStart: (ctx) => { + events.start.push(ctx) + }, + onUsage: (_ctx, info) => { + events.usage.push(info) + }, + onFinish: (_ctx, info) => { + events.finish.push(info) + }, + onAbort: (_ctx, info) => { + events.abort.push(info) + }, + onError: (_ctx, info) => { + events.error.push(info) + }, + } + return { middleware, events } +} + +// ============================================================================ +// Tests +// ============================================================================ + +describe('rerank() activity', () => { + it('maps scored indices back to the original documents in ranked order', async () => { + const adapter = mockRerankAdapter(async () => ranked(1, 0)) + const documents = ['sunny day at the beach', 'rainy afternoon in the city'] + + const result = await rerank({ adapter, query: 'talk about rain', documents }) + + expect(result.ranking.map((r) => r.index)).toEqual([1, 0]) + expect(result.ranking[0]!.document).toBe('rainy afternoon in the city') + expect(result.rerankedDocuments).toEqual([ + 'rainy afternoon in the city', + 'sunny day at the beach', + ]) + expect(result.usage.unitsBilled).toBe(1) + }) + + it('serializes object documents with JSON.stringify before the adapter call', async () => { + const adapter = mockRerankAdapter(async () => ranked(0)) + const documents = [{ id: 1, text: 'a lightweight ultrabook' }] + + await rerank({ adapter, query: 'best travel laptop', documents }) + + expect(adapter.calls[0]!.documents).toEqual([JSON.stringify(documents[0])]) + }) + + it('returns the original object (not its serialized form) in the result', async () => { + const documents = [ + { id: 1, text: 'heavy gaming desktop' }, + { id: 2, text: 'lightweight ultrabook' }, + ] + const adapter = mockRerankAdapter(async () => ranked(1, 0)) + + const result = await rerank({ adapter, query: 'travel laptop', documents }) + + // document is the original object, fully typed — id is accessible. + expect(result.ranking[0]!.document.id).toBe(2) + expect(result.ranking[0]!.document).toBe(documents[1]) + }) + + it('throws on empty documents before calling the adapter', async () => { + const adapter = mockRerankAdapter(async () => ranked()) + + await expect( + rerank({ adapter, query: 'x', documents: [] }), + ).rejects.toThrow('at least one document') + expect(adapter.calls).toHaveLength(0) + }) + + it('fires middleware start, usage, then finish on success', async () => { + const { middleware, events } = recordingMiddleware() + const adapter = mockRerankAdapter(async () => ranked(0, 1)) + + await rerank({ + adapter, + query: 'q', + documents: ['a', 'b'], + middleware: [middleware], + }) + + expect(events.start).toHaveLength(1) + expect(events.start[0]!.activity).toBe('rerank') + expect(events.start[0]!.provider).toBe('mock') + expect(events.usage[0]!.unitsBilled).toBe(1) + expect(events.finish).toHaveLength(1) + expect(events.error).toHaveLength(0) + expect(events.abort).toHaveLength(0) + }) + + it('fires onError (not onAbort) and rethrows when the adapter throws', async () => { + const { middleware, events } = recordingMiddleware() + const adapter = mockRerankAdapter(async () => { + throw new Error('rerank boom') + }) + + await expect( + rerank({ + adapter, + query: 'q', + documents: ['a'], + middleware: [middleware], + debug: false, + }), + ).rejects.toThrow('rerank boom') + + expect(events.error).toHaveLength(1) + expect(events.abort).toHaveLength(0) + expect(events.finish).toHaveLength(0) + }) + + it('fires onAbort (not onError) when the request is cancelled', async () => { + const { middleware, events } = recordingMiddleware() + const controller = new AbortController() + const adapter = mockRerankAdapter(async () => { + controller.abort() + const err = new Error('aborted') + err.name = 'AbortError' + throw err + }) + + await expect( + rerank({ + adapter, + query: 'q', + documents: ['a'], + abortSignal: controller.signal, + middleware: [middleware], + debug: false, + }), + ).rejects.toThrow('aborted') + + expect(events.abort).toHaveLength(1) + expect(events.error).toHaveLength(0) + expect(events.finish).toHaveLength(0) + }) + + it('forwards topN and abortSignal to the adapter', async () => { + const controller = new AbortController() + const adapter = mockRerankAdapter(async () => ranked(0)) + + await rerank({ + adapter, + query: 'q', + documents: ['a', 'b', 'c'], + topN: 1, + abortSignal: controller.signal, + }) + + expect(adapter.calls[0]!.topN).toBe(1) + expect(adapter.calls[0]!.abortSignal).toBe(controller.signal) + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 40f5181e4..95d1d9894 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -49,7 +49,7 @@ importers: version: 0.5.0 knip: specifier: ^5.70.2 - version: 5.73.4(@types/node@24.10.3)(typescript@5.9.3) + version: 5.73.4(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)(@types/node@24.10.3)(typescript@5.9.3) markdown-link-extractor: specifier: ^4.0.3 version: 4.0.3 @@ -1275,6 +1275,18 @@ importers: specifier: ^4.21.0 version: 4.21.0 + packages/ai-cohere: + devDependencies: + '@tanstack/ai': + specifier: workspace:* + version: link:../ai + '@vitest/coverage-v8': + specifier: 4.0.14 + version: 4.0.14(vitest@4.1.4) + vite: + specifier: ^7.3.3 + version: 7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2) + packages/ai-devtools: dependencies: '@tanstack/ai': @@ -1699,7 +1711,7 @@ importers: version: 1.9.10 tsdown: specifier: ^0.17.0-beta.6 - version: 0.17.3(oxc-resolver@11.15.0)(publint@0.3.16)(typescript@5.9.3) + version: 0.17.3(oxc-resolver@11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1))(publint@0.3.16)(typescript@5.9.3) typescript: specifier: 5.9.3 version: 5.9.3 @@ -1819,7 +1831,7 @@ importers: version: 27.3.0(postcss@8.5.15) tsdown: specifier: ^0.17.0-beta.6 - version: 0.17.3(oxc-resolver@11.15.0)(publint@0.3.16)(typescript@5.9.3) + version: 0.17.3(oxc-resolver@11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1))(publint@0.3.16)(typescript@5.9.3) typescript: specifier: 5.9.3 version: 5.9.3 @@ -17244,7 +17256,7 @@ snapshots: '@img/sharp-wasm32@0.34.5': dependencies: - '@emnapi/runtime': 1.10.0 + '@emnapi/runtime': 1.11.1 optional: true '@img/sharp-win32-arm64@0.34.5': @@ -17838,7 +17850,7 @@ snapshots: '@oxc-resolver/binding-wasm32-wasi@11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)': dependencies: - '@napi-rs/wasm-runtime': 1.1.4(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1) + '@napi-rs/wasm-runtime': 1.1.5(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1) transitivePeerDependencies: - '@emnapi/core' - '@emnapi/runtime' @@ -22957,7 +22969,7 @@ snapshots: dotenv@8.6.0: {} - dts-resolver@2.1.3(oxc-resolver@11.15.0): + dts-resolver@2.1.3(oxc-resolver@11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)): optionalDependencies: oxc-resolver: 11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1) @@ -24829,7 +24841,7 @@ snapshots: klona@2.0.6: {} - knip@5.73.4(@types/node@24.10.3)(typescript@5.9.3): + knip@5.73.4(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)(@types/node@24.10.3)(typescript@5.9.3): dependencies: '@nodelib/fs.walk': 1.2.8 '@types/node': 24.10.3 @@ -27469,14 +27481,14 @@ snapshots: robot3@0.4.1: {} - rolldown-plugin-dts@0.18.3(oxc-resolver@11.15.0)(rolldown@1.0.0-beta.53)(typescript@5.9.3): + rolldown-plugin-dts@0.18.3(oxc-resolver@11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1))(rolldown@1.0.0-beta.53)(typescript@5.9.3): dependencies: '@babel/generator': 7.28.5 '@babel/parser': 7.29.0 '@babel/types': 7.29.0 ast-kit: 2.2.0 birpc: 3.0.0 - dts-resolver: 2.1.3(oxc-resolver@11.15.0) + dts-resolver: 2.1.3(oxc-resolver@11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)) get-tsconfig: 4.13.0 magic-string: 0.30.21 obug: 2.1.1 @@ -28433,7 +28445,7 @@ snapshots: minimist: 1.2.8 strip-bom: 3.0.0 - tsdown@0.17.3(oxc-resolver@11.15.0)(publint@0.3.16)(typescript@5.9.3): + tsdown@0.17.3(oxc-resolver@11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1))(publint@0.3.16)(typescript@5.9.3): dependencies: ansis: 4.2.0 cac: 6.7.14 @@ -28443,7 +28455,7 @@ snapshots: import-without-cache: 0.2.3 obug: 2.1.1 rolldown: 1.0.0-beta.53 - rolldown-plugin-dts: 0.18.3(oxc-resolver@11.15.0)(rolldown@1.0.0-beta.53)(typescript@5.9.3) + rolldown-plugin-dts: 0.18.3(oxc-resolver@11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1))(rolldown@1.0.0-beta.53)(typescript@5.9.3) semver: 7.7.3 tinyexec: 1.0.2 tinyglobby: 0.2.15 From fe6a81cefb37905bcdb5c1f857e4db21d535fa3b Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 25 Jun 2026 15:57:26 +0000 Subject: [PATCH 2/9] ci: apply automated fixes --- packages/ai-cohere/src/adapters/rerank.ts | 6 +++++- packages/ai-cohere/tests/rerank-adapter.test.ts | 5 ++++- packages/ai/src/activities/rerank/adapter.ts | 3 +-- packages/ai/tests/rerank.test.ts | 12 ++++++++++-- 4 files changed, 20 insertions(+), 6 deletions(-) diff --git a/packages/ai-cohere/src/adapters/rerank.ts b/packages/ai-cohere/src/adapters/rerank.ts index 5be678f49..5298175cf 100644 --- a/packages/ai-cohere/src/adapters/rerank.ts +++ b/packages/ai-cohere/src/adapters/rerank.ts @@ -8,7 +8,11 @@ import type { CohereRerankModel, InferCohereRerankProviderOptions, } from '../model-meta' -import type { RerankAdapterResult, RerankOptions, TokenUsage } from '@tanstack/ai' +import type { + RerankAdapterResult, + RerankOptions, + TokenUsage, +} from '@tanstack/ai' /** Shape of the Cohere `/v2/rerank` response we depend on. */ interface CohereRerankResponse { diff --git a/packages/ai-cohere/tests/rerank-adapter.test.ts b/packages/ai-cohere/tests/rerank-adapter.test.ts index c3b0c8a9e..de6ce7688 100644 --- a/packages/ai-cohere/tests/rerank-adapter.test.ts +++ b/packages/ai-cohere/tests/rerank-adapter.test.ts @@ -101,7 +101,10 @@ describe('CohereRerankAdapter', () => { it('throws with status detail on a non-200 response', async () => { fetchMock.mockResolvedValue( - new Response('rate limited', { status: 429, statusText: 'Too Many Requests' }), + new Response('rate limited', { + status: 429, + statusText: 'Too Many Requests', + }), ) await expect( diff --git a/packages/ai/src/activities/rerank/adapter.ts b/packages/ai/src/activities/rerank/adapter.ts index 7e5f54bf4..ed8258023 100644 --- a/packages/ai/src/activities/rerank/adapter.ts +++ b/packages/ai/src/activities/rerank/adapter.ts @@ -63,8 +63,7 @@ export type AnyRerankAdapter = RerankAdapter export abstract class BaseRerankAdapter< TModel extends string = string, TProviderOptions extends object = Record, -> implements RerankAdapter -{ +> implements RerankAdapter { readonly kind = 'rerank' as const abstract readonly name: string readonly model: TModel diff --git a/packages/ai/tests/rerank.test.ts b/packages/ai/tests/rerank.test.ts index 2862068a1..b28e7ca54 100644 --- a/packages/ai/tests/rerank.test.ts +++ b/packages/ai/tests/rerank.test.ts @@ -9,7 +9,11 @@ import type { GenerationMiddlewareContext, GenerationUsageInfo, } from '../src/activities/middleware' -import type { RerankAdapterResult, RerankOptions, TokenUsage } from '../src/types' +import type { + RerankAdapterResult, + RerankOptions, + TokenUsage, +} from '../src/types' // ============================================================================ // Helpers @@ -91,7 +95,11 @@ describe('rerank() activity', () => { const adapter = mockRerankAdapter(async () => ranked(1, 0)) const documents = ['sunny day at the beach', 'rainy afternoon in the city'] - const result = await rerank({ adapter, query: 'talk about rain', documents }) + const result = await rerank({ + adapter, + query: 'talk about rain', + documents, + }) expect(result.ranking.map((r) => r.index)).toEqual([1, 0]) expect(result.ranking[0]!.document).toBe('rainy afternoon in the city') From a7fe8ce43430ed7a4ba2fdda5af1e10b10fc0dea Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 25 Jun 2026 18:28:43 +0200 Subject: [PATCH 3/9] feat(ai-openrouter): add rerank adapter via the OpenRouter SDK Add openRouterRerank / createOpenRouterRerank for reordering documents by relevance through OpenRouter's unified /v1/rerank endpoint, using the @openrouter/sdk rerank method. - Model-agnostic: supports every rerank model OpenRouter offers (Cohere rerank-v3.5 / 4-fast / 4-pro, NVIDIA llama-nemotron-rerank, and any future slug) via an open model type with known-model autocomplete. - Maps the SDK response to the shared RerankAdapterResult, including usage (searchUnits -> unitsBilled, cost, totalTokens). - Optional provider-routing preferences via modelOptions.provider; abortSignal forwarded through the SDK fetchOptions. - Unit tests (SDK mocked) and docs (rerank guide + OpenRouter adapter page). --- .changeset/openrouter-rerank.md | 11 ++ docs/adapters/openrouter.md | 32 ++++ docs/config.json | 3 +- docs/rerank/rerank.md | 33 ++++- packages/ai-openrouter/src/adapters/rerank.ts | 140 ++++++++++++++++++ packages/ai-openrouter/src/index.ts | 14 ++ .../src/rerank/rerank-provider-options.ts | 43 ++++++ .../tests/rerank-adapter.test.ts | 131 ++++++++++++++++ 8 files changed, 404 insertions(+), 3 deletions(-) create mode 100644 .changeset/openrouter-rerank.md create mode 100644 packages/ai-openrouter/src/adapters/rerank.ts create mode 100644 packages/ai-openrouter/src/rerank/rerank-provider-options.ts create mode 100644 packages/ai-openrouter/tests/rerank-adapter.test.ts diff --git a/.changeset/openrouter-rerank.md b/.changeset/openrouter-rerank.md new file mode 100644 index 000000000..00af0da15 --- /dev/null +++ b/.changeset/openrouter-rerank.md @@ -0,0 +1,11 @@ +--- +'@tanstack/ai-openrouter': minor +--- + +feat: add `openRouterRerank` / `createOpenRouterRerank` rerank adapters + +Rerank documents by relevance to a query through OpenRouter's unified +`/v1/rerank` endpoint (e.g. `cohere/rerank-v3.5`) with the `rerank()` activity. +Reads `OPENROUTER_API_KEY` from the environment and forwards the optional +`httpReferer` / `xTitle` attribution headers, consistent with the other +OpenRouter adapters. diff --git a/docs/adapters/openrouter.md b/docs/adapters/openrouter.md index 6dfca45d1..cbbbe9cd2 100644 --- a/docs/adapters/openrouter.md +++ b/docs/adapters/openrouter.md @@ -247,10 +247,42 @@ fields are simply absent and the stream completes normally. Both `openRouterText` and `openRouterResponsesText` populate cost when OpenRouter returns it. +## Reranking + +OpenRouter exposes rerank models through its unified `/v1/rerank` endpoint +(served via the `@openrouter/sdk` SDK). Any rerank model OpenRouter offers works +by passing its slug — for example `cohere/rerank-v3.5`, `cohere/rerank-4-fast`, +`cohere/rerank-4-pro`, or `nvidia/llama-nemotron-rerank-vl-1b-v2`. Use +`openRouterRerank` with the `rerank()` activity to reorder candidate documents +by relevance to a query: + +```typescript +import { rerank } from "@tanstack/ai"; +import { openRouterRerank } from "@tanstack/ai-openrouter"; + +const { rerankedDocuments } = await rerank({ + adapter: openRouterRerank("cohere/rerank-v3.5"), + query: "talk about rain", + documents: ["sunny day at the beach", "rainy afternoon in the city"], + topN: 2, +}); + +console.log(rerankedDocuments[0]); // 'rainy afternoon in the city' +``` + +`openRouterRerank` reads `OPENROUTER_API_KEY` from the environment; pass a key +explicitly with `createOpenRouterRerank("cohere/rerank-v3.5", "sk-or-...")`. The +optional `httpReferer` / `xTitle` config fields are forwarded as OpenRouter +attribution headers, just like the chat adapter. + +See the [Reranking guide](../rerank/rerank) for object documents, RAG +pipelines, options, and the result shape. + ## Next Steps - [Getting Started](../getting-started/quick-start) - Learn the basics - [Tools Guide](../tools/tools) - Learn about tools +- [Reranking](../rerank/rerank) - Reorder documents by relevance ## Provider Tools diff --git a/docs/config.json b/docs/config.json index a07271efe..c9f822385 100644 --- a/docs/config.json +++ b/docs/config.json @@ -471,7 +471,8 @@ { "label": "OpenRouter Adapter", "to": "adapters/openrouter", - "addedAt": "2026-04-15" + "addedAt": "2026-04-15", + "updatedAt": "2026-06-25" }, { "label": "OpenAI-Compatible", diff --git a/docs/rerank/rerank.md b/docs/rerank/rerank.md index 77f1f31e4..3ce0b265c 100644 --- a/docs/rerank/rerank.md +++ b/docs/rerank/rerank.md @@ -25,12 +25,21 @@ you'll have that list reordered, with a relevance score per document. `rerank()` is the precision step in a retrieval pipeline: retrieve a broad set of candidates cheaply, then rerank to surface the few that matter. -## Installation +## Providers + +Reranking is available from two adapters today: + +- **Cohere** (`@tanstack/ai-cohere`) — `cohereRerank('rerank-v3.5')`, talking to Cohere directly. +- **OpenRouter** (`@tanstack/ai-openrouter`) — `openRouterRerank('cohere/rerank-v3.5')`, routing rerank through your existing OpenRouter key. -Reranking is provided by the Cohere adapter: +Both implement the same `rerank()` activity — swap the adapter, keep the call. + +## Installation ```bash npm install @tanstack/ai-cohere +# or, to rerank through OpenRouter: +npm install @tanstack/ai-openrouter ``` Peer dependency: @@ -63,6 +72,25 @@ console.log(ranking[0]) // { index: 1, score: 0.98, document: 'rainy afternoon i The adapter reads `COHERE_API_KEY` from the environment. To pass a key explicitly, use `createCohereRerank('rerank-v3.5', 'co-...')`. +To rerank through OpenRouter instead, swap the adapter — everything else stays +the same: + +```typescript +import { rerank } from '@tanstack/ai' +import { openRouterRerank } from '@tanstack/ai-openrouter' + +const { rerankedDocuments } = await rerank({ + adapter: openRouterRerank('cohere/rerank-v3.5'), + query: 'talk about rain', + documents: ['sunny day at the beach', 'rainy afternoon in the city'], + topN: 2, +}) + +console.log(rerankedDocuments[0]) // 'rainy afternoon in the city' +``` + +`openRouterRerank` reads `OPENROUTER_API_KEY` from the environment. + ## Reranking Object Documents Documents don't have to be strings. Pass JSON-serializable objects and the @@ -300,4 +328,5 @@ try { ## Next Steps - [Cohere Adapter](../adapters/cohere) — models, configuration, and explicit API keys +- [OpenRouter Adapter](../adapters/openrouter) — rerank through your OpenRouter key - [Middleware](../advanced/middleware) — lifecycle hooks for usage and errors diff --git a/packages/ai-openrouter/src/adapters/rerank.ts b/packages/ai-openrouter/src/adapters/rerank.ts new file mode 100644 index 000000000..d959d9b83 --- /dev/null +++ b/packages/ai-openrouter/src/adapters/rerank.ts @@ -0,0 +1,140 @@ +import { OpenRouter } from '@openrouter/sdk' +import { BaseRerankAdapter } from '@tanstack/ai/adapters' +import { toRunErrorPayload } from '@tanstack/ai/adapter-internals' +import { getOpenRouterApiKeyFromEnv } from '../utils' +import type { SDKOptions } from '@openrouter/sdk' +import type { + OpenRouterRerankModel, + OpenRouterRerankProviderOptions, +} from '../rerank/rerank-provider-options' +import type { + RerankAdapterResult, + RerankOptions, + TokenUsage, +} from '@tanstack/ai' + +export interface OpenRouterRerankConfig extends SDKOptions {} + +/** + * OpenRouter rerank adapter. + * + * Reorders documents by relevance to a query through OpenRouter's unified + * `/v1/rerank` endpoint via the `@openrouter/sdk` SDK. The endpoint is + * model-agnostic, so any rerank model OpenRouter offers works by passing its + * slug (Cohere, NVIDIA, …). Returns scored indices into the submitted + * documents; the `rerank()` activity maps those back to the caller's original + * documents. + */ +export class OpenRouterRerankAdapter< + TModel extends OpenRouterRerankModel, +> extends BaseRerankAdapter { + readonly name = 'openrouter' as const + + private readonly client: OpenRouter + + constructor(config: OpenRouterRerankConfig, model: TModel) { + super({}, model) + this.client = new OpenRouter(config) + } + + async rerank( + options: RerankOptions, + ): Promise { + const { model, query, documents, topN, modelOptions, abortSignal, logger } = + options + + logger.request( + `activity=rerank provider=${this.name} model=${model} documents=${documents.length}`, + { provider: this.name, model }, + ) + + try { + const response = await this.client.rerank.rerank( + { + requestBody: { + model, + query, + documents, + ...(topN !== undefined ? { topN } : {}), + ...(modelOptions?.provider + ? { provider: modelOptions.provider } + : {}), + }, + }, + abortSignal ? { fetchOptions: { signal: abortSignal } } : undefined, + ) + + // The SDK types the response as `CreateRerankResponseBody | string`; the + // bare-string form is an error/non-JSON payload, not a valid result. + if (typeof response === 'string') { + throw new Error('OpenRouter rerank returned an unexpected response') + } + + const usage: TokenUsage = { + promptTokens: 0, + completionTokens: 0, + totalTokens: response.usage?.totalTokens ?? 0, + ...(response.usage?.searchUnits !== undefined + ? { unitsBilled: response.usage.searchUnits } + : {}), + ...(response.usage?.cost !== undefined + ? { cost: response.usage.cost } + : {}), + } + + return { + id: response.id ?? this.generateId(), + ranking: response.results.map((r) => ({ + index: r.index, + score: r.relevanceScore, + })), + usage, + } + } catch (error) { + logger.errors(`${this.name}.rerank fatal`, { + error: toRunErrorPayload(error, `${this.name}.rerank failed`), + source: `${this.name}.rerank`, + }) + throw error + } + } +} + +/** + * Creates an OpenRouter rerank adapter with an explicit API key. + * + * @example + * ```typescript + * const adapter = createOpenRouterRerank('cohere/rerank-v3.5', 'sk-or-...') + * ``` + */ +export function createOpenRouterRerank( + model: TModel, + apiKey: string, + config?: Omit, +): OpenRouterRerankAdapter { + return new OpenRouterRerankAdapter({ apiKey, ...config }, model) +} + +/** + * Creates an OpenRouter rerank adapter, reading `OPENROUTER_API_KEY` from the + * environment. + * + * @example + * ```typescript + * import { rerank } from '@tanstack/ai' + * import { openRouterRerank } from '@tanstack/ai-openrouter' + * + * const { rerankedDocuments } = await rerank({ + * adapter: openRouterRerank('cohere/rerank-v3.5'), + * query: 'talk about rain', + * documents: ['sunny day', 'rainy afternoon'], + * }) + * ``` + */ +export function openRouterRerank( + model: TModel, + config?: Omit, +): OpenRouterRerankAdapter { + return createOpenRouterRerank(model, getOpenRouterApiKeyFromEnv(), config) +} diff --git a/packages/ai-openrouter/src/index.ts b/packages/ai-openrouter/src/index.ts index e55f9a9c7..097828fe2 100644 --- a/packages/ai-openrouter/src/index.ts +++ b/packages/ai-openrouter/src/index.ts @@ -41,6 +41,20 @@ export type { OpenRouterImageModelSizeByName, } from './image/image-provider-options' +// Rerank adapter - document reranking via OpenRouter's /v1/rerank endpoint +export { + OpenRouterRerankAdapter, + createOpenRouterRerank, + openRouterRerank, + type OpenRouterRerankConfig, +} from './adapters/rerank' +export { + OPENROUTER_RERANK_MODELS, + type OpenRouterRerankModel, + type KnownOpenRouterRerankModel, + type OpenRouterRerankProviderOptions, +} from './rerank/rerank-provider-options' + // ============================================================================ // Type Exports // ============================================================================ diff --git a/packages/ai-openrouter/src/rerank/rerank-provider-options.ts b/packages/ai-openrouter/src/rerank/rerank-provider-options.ts new file mode 100644 index 000000000..21c099352 --- /dev/null +++ b/packages/ai-openrouter/src/rerank/rerank-provider-options.ts @@ -0,0 +1,43 @@ +import type { ProviderPreferences } from '@openrouter/sdk/models' + +/** + * OpenRouter rerank model metadata and provider options. + * + * OpenRouter exposes rerank models through its unified `/v1/rerank` endpoint. + * The endpoint is model-agnostic — any rerank model OpenRouter offers works by + * passing its slug as the model, so the model type is open (a known-model + * union for autocomplete, widened with `string`). + */ + +/** + * A non-exhaustive list of known OpenRouter rerank model slugs, surfaced for + * editor autocomplete. Any other rerank model OpenRouter offers also works — + * see {@link OpenRouterRerankModel}. + */ +export const OPENROUTER_RERANK_MODELS = [ + 'cohere/rerank-v3.5', + 'cohere/rerank-4-fast', + 'cohere/rerank-4-pro', + 'nvidia/llama-nemotron-rerank-vl-1b-v2', +] as const + +/** A rerank model slug known to OpenRouter (for autocomplete). */ +export type KnownOpenRouterRerankModel = (typeof OPENROUTER_RERANK_MODELS)[number] + +/** + * Any OpenRouter rerank model. Known slugs autocomplete; any other rerank + * model OpenRouter offers is also accepted. + */ +export type OpenRouterRerankModel = KnownOpenRouterRerankModel | (string & {}) + +/** + * Provider-specific options for an OpenRouter rerank request, forwarded on the + * `modelOptions` field of `rerank()`. + */ +export interface OpenRouterRerankProviderOptions { + /** + * OpenRouter provider routing preferences — pin, order, or allow fallback + * across the providers that serve the chosen rerank model. + */ + provider?: ProviderPreferences +} diff --git a/packages/ai-openrouter/tests/rerank-adapter.test.ts b/packages/ai-openrouter/tests/rerank-adapter.test.ts new file mode 100644 index 000000000..77b401223 --- /dev/null +++ b/packages/ai-openrouter/tests/rerank-adapter.test.ts @@ -0,0 +1,131 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { rerank } from '@tanstack/ai' +import { createOpenRouterRerank } from '../src/adapters/rerank' + +// Mock the OpenRouter SDK so the adapter's `new OpenRouter().rerank.rerank()` +// call resolves to a controlled response. `vi.hoisted` lets the (hoisted) +// vi.mock factory reference the spy. +const { rerankFn } = vi.hoisted(() => ({ rerankFn: vi.fn() })) + +vi.mock('@openrouter/sdk', () => ({ + // A class so `new OpenRouter(config)` is constructable; each instance exposes + // the spied `rerank.rerank`. + OpenRouter: class { + rerank = { rerank: rerankFn } + }, +})) + +beforeEach(() => { + rerankFn.mockReset() +}) + +const documents = ['sunny day at the beach', 'rainy afternoon in the city'] + +/** A parsed SDK rerank response (camelCase, as the SDK returns it). */ +function sdkResponse() { + return { + id: 'or-1', + model: 'cohere/rerank-v3.5', + results: [ + { document: { text: documents[1] }, index: 1, relevanceScore: 0.97 }, + { document: { text: documents[0] }, index: 0, relevanceScore: 0.1 }, + ], + usage: { searchUnits: 1, cost: 0.002, totalTokens: 20 }, + } +} + +const adapter = () => createOpenRouterRerank('cohere/rerank-v3.5', 'sk-or-test') + +describe('OpenRouterRerankAdapter', () => { + it('calls the SDK rerank with the request body and maps the response', async () => { + rerankFn.mockResolvedValue(sdkResponse()) + + const result = await rerank({ + adapter: adapter(), + query: 'talk about rain', + documents, + topN: 2, + }) + + expect(rerankFn).toHaveBeenCalledTimes(1) + expect(rerankFn.mock.calls[0]![0]).toEqual({ + requestBody: { + model: 'cohere/rerank-v3.5', + query: 'talk about rain', + documents, + topN: 2, + }, + }) + expect(result.id).toBe('or-1') + expect(result.ranking).toEqual([ + { index: 1, score: 0.97, document: documents[1] }, + { index: 0, score: 0.1, document: documents[0] }, + ]) + }) + + it('maps SDK usage (searchUnits/cost/totalTokens)', async () => { + rerankFn.mockResolvedValue(sdkResponse()) + + const result = await rerank({ adapter: adapter(), query: 'q', documents }) + + expect(result.usage.unitsBilled).toBe(1) + expect(result.usage.cost).toBe(0.002) + expect(result.usage.totalTokens).toBe(20) + }) + + it('works with a non-Cohere model slug', async () => { + rerankFn.mockResolvedValue({ ...sdkResponse(), model: 'nvidia/llama-nemotron-rerank-vl-1b-v2' }) + + await rerank({ + adapter: createOpenRouterRerank( + 'nvidia/llama-nemotron-rerank-vl-1b-v2', + 'sk-or-test', + ), + query: 'q', + documents, + }) + + expect(rerankFn.mock.calls[0]![0].requestBody.model).toBe( + 'nvidia/llama-nemotron-rerank-vl-1b-v2', + ) + }) + + it('forwards provider routing preferences into the request body', async () => { + rerankFn.mockResolvedValue(sdkResponse()) + + await rerank({ + adapter: adapter(), + query: 'q', + documents, + modelOptions: { provider: { order: ['cohere'] } }, + }) + + expect(rerankFn.mock.calls[0]![0].requestBody.provider).toEqual({ + order: ['cohere'], + }) + }) + + it('forwards the abort signal via fetchOptions', async () => { + rerankFn.mockResolvedValue(sdkResponse()) + const controller = new AbortController() + + await rerank({ + adapter: adapter(), + query: 'q', + documents, + abortSignal: controller.signal, + }) + + expect(rerankFn.mock.calls[0]![1]).toEqual({ + fetchOptions: { signal: controller.signal }, + }) + }) + + it('throws when the SDK returns a bare string response', async () => { + rerankFn.mockResolvedValue('error: bad request') + + await expect( + rerank({ adapter: adapter(), query: 'q', documents, debug: false }), + ).rejects.toThrow('unexpected response') + }) +}) From a5f1578f83e274f9c88d387e664535e8d1483692 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 25 Jun 2026 16:30:31 +0000 Subject: [PATCH 4/9] ci: apply automated fixes --- packages/ai-openrouter/src/rerank/rerank-provider-options.ts | 3 ++- packages/ai-openrouter/tests/rerank-adapter.test.ts | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/packages/ai-openrouter/src/rerank/rerank-provider-options.ts b/packages/ai-openrouter/src/rerank/rerank-provider-options.ts index 21c099352..f35b21bea 100644 --- a/packages/ai-openrouter/src/rerank/rerank-provider-options.ts +++ b/packages/ai-openrouter/src/rerank/rerank-provider-options.ts @@ -22,7 +22,8 @@ export const OPENROUTER_RERANK_MODELS = [ ] as const /** A rerank model slug known to OpenRouter (for autocomplete). */ -export type KnownOpenRouterRerankModel = (typeof OPENROUTER_RERANK_MODELS)[number] +export type KnownOpenRouterRerankModel = + (typeof OPENROUTER_RERANK_MODELS)[number] /** * Any OpenRouter rerank model. Known slugs autocomplete; any other rerank diff --git a/packages/ai-openrouter/tests/rerank-adapter.test.ts b/packages/ai-openrouter/tests/rerank-adapter.test.ts index 77b401223..8ade9a8df 100644 --- a/packages/ai-openrouter/tests/rerank-adapter.test.ts +++ b/packages/ai-openrouter/tests/rerank-adapter.test.ts @@ -74,7 +74,10 @@ describe('OpenRouterRerankAdapter', () => { }) it('works with a non-Cohere model slug', async () => { - rerankFn.mockResolvedValue({ ...sdkResponse(), model: 'nvidia/llama-nemotron-rerank-vl-1b-v2' }) + rerankFn.mockResolvedValue({ + ...sdkResponse(), + model: 'nvidia/llama-nemotron-rerank-vl-1b-v2', + }) await rerank({ adapter: createOpenRouterRerank( From b6b969e666855fd913e499d9e5a06d4932f040fb Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 25 Jun 2026 18:33:24 +0200 Subject: [PATCH 5/9] docs: correct OpenRouter rerank attribution field name (xTitle -> appTitle) The SDK-based OpenRouter adapters use SDKOptions, whose attribution field is appTitle, not xTitle (xTitle only exists on the legacy raw-fetch client helper). Fix the rerank section of the OpenRouter adapter docs and the changeset. --- .changeset/openrouter-rerank.md | 2 +- docs/adapters/openrouter.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.changeset/openrouter-rerank.md b/.changeset/openrouter-rerank.md index 00af0da15..c1013fd74 100644 --- a/.changeset/openrouter-rerank.md +++ b/.changeset/openrouter-rerank.md @@ -7,5 +7,5 @@ feat: add `openRouterRerank` / `createOpenRouterRerank` rerank adapters Rerank documents by relevance to a query through OpenRouter's unified `/v1/rerank` endpoint (e.g. `cohere/rerank-v3.5`) with the `rerank()` activity. Reads `OPENROUTER_API_KEY` from the environment and forwards the optional -`httpReferer` / `xTitle` attribution headers, consistent with the other +`httpReferer` / `appTitle` attribution headers, consistent with the other OpenRouter adapters. diff --git a/docs/adapters/openrouter.md b/docs/adapters/openrouter.md index cbbbe9cd2..4983a457b 100644 --- a/docs/adapters/openrouter.md +++ b/docs/adapters/openrouter.md @@ -272,7 +272,7 @@ console.log(rerankedDocuments[0]); // 'rainy afternoon in the city' `openRouterRerank` reads `OPENROUTER_API_KEY` from the environment; pass a key explicitly with `createOpenRouterRerank("cohere/rerank-v3.5", "sk-or-...")`. The -optional `httpReferer` / `xTitle` config fields are forwarded as OpenRouter +optional `httpReferer` / `appTitle` config fields are forwarded as OpenRouter attribution headers, just like the chat adapter. See the [Reranking guide](../rerank/rerank) for object documents, RAG From 224c407213c21ed37a43a4d8e1c30914867616ae Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 25 Jun 2026 18:56:10 +0200 Subject: [PATCH 6/9] test(ai-openrouter): drive real SDK via fetch mock in rerank test Replace the @openrouter/sdk module mock with a network-layer fetch intercept driving the real SDK. Module mocking was sensitive to the CI runner's test isolation and failed there while passing locally; intercepting fetch exercises the SDK's real request building and zod response parsing and is environment- independent. --- .../tests/rerank-adapter.test.ts | 128 +++++++++--------- 1 file changed, 63 insertions(+), 65 deletions(-) diff --git a/packages/ai-openrouter/tests/rerank-adapter.test.ts b/packages/ai-openrouter/tests/rerank-adapter.test.ts index 8ade9a8df..805516033 100644 --- a/packages/ai-openrouter/tests/rerank-adapter.test.ts +++ b/packages/ai-openrouter/tests/rerank-adapter.test.ts @@ -1,44 +1,63 @@ -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { rerank } from '@tanstack/ai' import { createOpenRouterRerank } from '../src/adapters/rerank' -// Mock the OpenRouter SDK so the adapter's `new OpenRouter().rerank.rerank()` -// call resolves to a controlled response. `vi.hoisted` lets the (hoisted) -// vi.mock factory reference the spy. -const { rerankFn } = vi.hoisted(() => ({ rerankFn: vi.fn() })) - -vi.mock('@openrouter/sdk', () => ({ - // A class so `new OpenRouter(config)` is constructable; each instance exposes - // the spied `rerank.rerank`. - OpenRouter: class { - rerank = { rerank: rerankFn } - }, -})) +// Intercept at the network layer and drive the REAL @openrouter/sdk. This keeps +// the test free of module mocking (which is sensitive to runner/isolation +// differences) and exercises the SDK's real request building and response +// parsing. +const fetchMock = vi.fn() beforeEach(() => { - rerankFn.mockReset() + vi.stubGlobal('fetch', fetchMock) + fetchMock.mockReset() +}) + +afterEach(() => { + vi.unstubAllGlobals() }) const documents = ['sunny day at the beach', 'rainy afternoon in the city'] -/** A parsed SDK rerank response (camelCase, as the SDK returns it). */ -function sdkResponse() { +/** A 200 response in the wire shape the SDK's zod schema parses. */ +function rerankResponse(body: unknown): Response { + return new Response(JSON.stringify(body), { + status: 200, + headers: { 'content-type': 'application/json' }, + }) +} + +/** Default wire-format payload reordering documents [1, 0]. */ +function wireBody() { return { id: 'or-1', model: 'cohere/rerank-v3.5', results: [ - { document: { text: documents[1] }, index: 1, relevanceScore: 0.97 }, - { document: { text: documents[0] }, index: 0, relevanceScore: 0.1 }, + { document: { text: documents[1] }, index: 1, relevance_score: 0.97 }, + { document: { text: documents[0] }, index: 0, relevance_score: 0.1 }, ], - usage: { searchUnits: 1, cost: 0.002, totalTokens: 20 }, + usage: { search_units: 1, cost: 0.002, total_tokens: 20 }, } } const adapter = () => createOpenRouterRerank('cohere/rerank-v3.5', 'sk-or-test') +/** The request the SDK handed to fetch, as { url, body }. */ +async function capturedRequest() { + const [input, init] = fetchMock.mock.calls[0]! + // The SDK may call fetch with a Request object or (url, init). + if (input instanceof Request) { + return { url: input.url, body: await input.clone().json() } + } + return { + url: String(input), + body: init?.body ? JSON.parse(String(init.body)) : undefined, + } +} + describe('OpenRouterRerankAdapter', () => { - it('calls the SDK rerank with the request body and maps the response', async () => { - rerankFn.mockResolvedValue(sdkResponse()) + it('hits the /rerank endpoint and maps the response', async () => { + fetchMock.mockResolvedValue(rerankResponse(wireBody())) const result = await rerank({ adapter: adapter(), @@ -47,15 +66,15 @@ describe('OpenRouterRerankAdapter', () => { topN: 2, }) - expect(rerankFn).toHaveBeenCalledTimes(1) - expect(rerankFn.mock.calls[0]![0]).toEqual({ - requestBody: { - model: 'cohere/rerank-v3.5', - query: 'talk about rain', - documents, - topN: 2, - }, + const { url, body } = await capturedRequest() + expect(url).toContain('/rerank') + expect(body).toMatchObject({ + model: 'cohere/rerank-v3.5', + query: 'talk about rain', + documents, + top_n: 2, }) + expect(result.id).toBe('or-1') expect(result.ranking).toEqual([ { index: 1, score: 0.97, document: documents[1] }, @@ -63,8 +82,8 @@ describe('OpenRouterRerankAdapter', () => { ]) }) - it('maps SDK usage (searchUnits/cost/totalTokens)', async () => { - rerankFn.mockResolvedValue(sdkResponse()) + it('maps usage (search_units/cost/total_tokens)', async () => { + fetchMock.mockResolvedValue(rerankResponse(wireBody())) const result = await rerank({ adapter: adapter(), query: 'q', documents }) @@ -74,27 +93,21 @@ describe('OpenRouterRerankAdapter', () => { }) it('works with a non-Cohere model slug', async () => { - rerankFn.mockResolvedValue({ - ...sdkResponse(), - model: 'nvidia/llama-nemotron-rerank-vl-1b-v2', - }) + const model = 'nvidia/llama-nemotron-rerank-vl-1b-v2' + fetchMock.mockResolvedValue(rerankResponse({ ...wireBody(), model })) await rerank({ - adapter: createOpenRouterRerank( - 'nvidia/llama-nemotron-rerank-vl-1b-v2', - 'sk-or-test', - ), + adapter: createOpenRouterRerank(model, 'sk-or-test'), query: 'q', documents, }) - expect(rerankFn.mock.calls[0]![0].requestBody.model).toBe( - 'nvidia/llama-nemotron-rerank-vl-1b-v2', - ) + const { body } = await capturedRequest() + expect(body.model).toBe(model) }) it('forwards provider routing preferences into the request body', async () => { - rerankFn.mockResolvedValue(sdkResponse()) + fetchMock.mockResolvedValue(rerankResponse(wireBody())) await rerank({ adapter: adapter(), @@ -103,32 +116,17 @@ describe('OpenRouterRerankAdapter', () => { modelOptions: { provider: { order: ['cohere'] } }, }) - expect(rerankFn.mock.calls[0]![0].requestBody.provider).toEqual({ - order: ['cohere'], - }) - }) - - it('forwards the abort signal via fetchOptions', async () => { - rerankFn.mockResolvedValue(sdkResponse()) - const controller = new AbortController() - - await rerank({ - adapter: adapter(), - query: 'q', - documents, - abortSignal: controller.signal, - }) - - expect(rerankFn.mock.calls[0]![1]).toEqual({ - fetchOptions: { signal: controller.signal }, - }) + const { body } = await capturedRequest() + expect(body.provider).toEqual({ order: ['cohere'] }) }) - it('throws when the SDK returns a bare string response', async () => { - rerankFn.mockResolvedValue('error: bad request') + it('throws on a non-200 response', async () => { + fetchMock.mockResolvedValue( + new Response('bad request', { status: 400, statusText: 'Bad Request' }), + ) await expect( rerank({ adapter: adapter(), query: 'q', documents, debug: false }), - ).rejects.toThrow('unexpected response') + ).rejects.toThrow() }) }) From f0f2fbb87a70bfc5d9d5fb25e994799ef44204e7 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 25 Jun 2026 19:03:54 +0200 Subject: [PATCH 7/9] fix: address review feedback on rerank - rerank() activity: guard provider-returned indices and remove the as-cast, throwing a clear error on an out-of-range index instead of yielding undefined. - BaseRerankAdapter.generateId: use slice(2, 9) so the fallback id is never empty. - ai-cohere env detection: fall back to process.env when window.env is absent (bundler/Electron browser builds). - docs: replace the `as` cast in the rerank server example with a runtime guard. --- docs/rerank/rerank.md | 19 ++++++++++++++----- packages/ai-cohere/src/utils/client.ts | 12 +++++++----- packages/ai/src/activities/rerank/adapter.ts | 2 +- packages/ai/src/activities/rerank/index.ts | 14 +++++++++----- packages/ai/tests/rerank.test.ts | 12 ++++++++++++ 5 files changed, 43 insertions(+), 16 deletions(-) diff --git a/docs/rerank/rerank.md b/docs/rerank/rerank.md index 3ce0b265c..71780ff31 100644 --- a/docs/rerank/rerank.md +++ b/docs/rerank/rerank.md @@ -185,12 +185,21 @@ export const Route = createFileRoute('/api/rerank')({ server: { handlers: { POST: async ({ request }) => { - const body = await request.json() - const { query, documents, topN } = body as { - query: string - documents: Array - topN?: number + const body: unknown = await request.json() + if ( + typeof body !== 'object' || + body === null || + !('query' in body) || + typeof body.query !== 'string' || + !('documents' in body) || + !Array.isArray(body.documents) + ) { + return new Response('Invalid request body', { status: 400 }) } + const { query, documents } = body + const topN = 'topN' in body && typeof body.topN === 'number' + ? body.topN + : undefined const result = await rerank({ adapter: cohereRerank('rerank-v3.5'), diff --git a/packages/ai-cohere/src/utils/client.ts b/packages/ai-cohere/src/utils/client.ts index 9c19a668b..2cf2b8ec0 100644 --- a/packages/ai-cohere/src/utils/client.ts +++ b/packages/ai-cohere/src/utils/client.ts @@ -21,7 +21,7 @@ export const COHERE_DEFAULT_BASE_URL = 'https://api.cohere.com' * @throws Error if `COHERE_API_KEY` is not found. */ export function getCohereApiKeyFromEnv(): string { - const env = + const windowEnv = typeof globalThis !== 'undefined' && (globalThis as Record).window ? (( @@ -30,10 +30,12 @@ export function getCohereApiKeyFromEnv(): string { unknown > ).env as Record | undefined) - : typeof process !== 'undefined' - ? process.env - : undefined - const key = env?.['COHERE_API_KEY'] + : undefined + const processEnv = typeof process !== 'undefined' ? process.env : undefined + // Prefer an injected `window.env` (browser builds) but fall back to + // `process.env` — bundlers and Electron can populate it even when `window` + // exists. + const key = windowEnv?.['COHERE_API_KEY'] ?? processEnv?.['COHERE_API_KEY'] if (!key) { throw new Error( 'COHERE_API_KEY not found in environment. Pass an API key explicitly via createCohereRerank(model, apiKey).', diff --git a/packages/ai/src/activities/rerank/adapter.ts b/packages/ai/src/activities/rerank/adapter.ts index ed8258023..015e05631 100644 --- a/packages/ai/src/activities/rerank/adapter.ts +++ b/packages/ai/src/activities/rerank/adapter.ts @@ -85,6 +85,6 @@ export abstract class BaseRerankAdapter< ): Promise protected generateId(): string { - return `${this.name}-${Date.now()}-${Math.random().toString(36).substring(7)}` + return `${this.name}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}` } } diff --git a/packages/ai/src/activities/rerank/index.ts b/packages/ai/src/activities/rerank/index.ts index b0cf8c9da..9dfdb6a82 100644 --- a/packages/ai/src/activities/rerank/index.ts +++ b/packages/ai/src/activities/rerank/index.ts @@ -207,11 +207,15 @@ export async function rerank< logger, }) - const ranking = result.ranking.map((r) => ({ - index: r.index, - score: r.score, - document: documents[r.index] as TDocument, - })) + const ranking = result.ranking.map((r) => { + const document = documents[r.index] + if (document === undefined) { + throw new Error( + `rerank(): provider ${adapter.name} returned out-of-range index ${r.index}`, + ) + } + return { index: r.index, score: r.score, document } + }) const rerankedDocuments = ranking.map((r) => r.document) const duration = Date.now() - startTime diff --git a/packages/ai/tests/rerank.test.ts b/packages/ai/tests/rerank.test.ts index b28e7ca54..3430ee258 100644 --- a/packages/ai/tests/rerank.test.ts +++ b/packages/ai/tests/rerank.test.ts @@ -224,4 +224,16 @@ describe('rerank() activity', () => { expect(adapter.calls[0]!.topN).toBe(1) expect(adapter.calls[0]!.abortSignal).toBe(controller.signal) }) + + it('throws when the provider returns an out-of-range index', async () => { + const adapter = mockRerankAdapter(async () => ({ + id: 'rr-1', + ranking: [{ index: 5, score: 0.9 }], + usage: { ...zeroUsage }, + })) + + await expect( + rerank({ adapter, query: 'q', documents: ['a', 'b'] }), + ).rejects.toThrow('out-of-range') + }) }) From dbb015e1d3f52a3fa853c16a25b20316b1cb2628 Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Sun, 28 Jun 2026 16:29:34 +1000 Subject: [PATCH 8/9] refactor(ai): remove any from rerank provider-option inference Adopt the audio activity's clean type pattern: extract provider options by structurally matching the adapter's `~types` shape with `infer P` instead of routing through `RerankAdapter`, and tighten the adapter constraints from the loose `object` to the self-referential `RerankProviderOptions`. Co-Authored-By: Claude Opus 4.8 (1M context) --- packages/ai/src/activities/rerank/index.ts | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/packages/ai/src/activities/rerank/index.ts b/packages/ai/src/activities/rerank/index.ts index 9dfdb6a82..5b4b4f175 100644 --- a/packages/ai/src/activities/rerank/index.ts +++ b/packages/ai/src/activities/rerank/index.ts @@ -33,10 +33,11 @@ export const kind = 'rerank' as const // =========================== /** Extract provider options from a RerankAdapter via ~types */ -export type RerankProviderOptions = - TAdapter extends RerankAdapter - ? TAdapter['~types']['providerOptions'] - : object +export type RerankProviderOptions = TAdapter extends { + '~types': { providerOptions: infer P extends object } +} + ? P + : object // =========================== // Activity Options Type @@ -50,7 +51,7 @@ export type RerankProviderOptions = * @template TDocument - The document element type (string or object) */ export interface RerankActivityOptions< - TAdapter extends RerankAdapter, + TAdapter extends RerankAdapter>, TDocument extends string | object = string, > { /** The rerank adapter to use (must be created with a model) */ @@ -142,7 +143,7 @@ function isAbortError(error: unknown, signal?: AbortSignal): boolean { * ``` */ export async function rerank< - TAdapter extends RerankAdapter, + TAdapter extends RerankAdapter>, TDocument extends string | object = string, >( options: RerankActivityOptions, @@ -277,7 +278,7 @@ export async function rerank< * Create typed options for the rerank() function without executing. */ export function createRerankOptions< - TAdapter extends RerankAdapter, + TAdapter extends RerankAdapter>, TDocument extends string | object = string, >( options: RerankActivityOptions, From 28d895f7bd2ed2a2b39531a8a86cf66088ea2aef Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Sun, 28 Jun 2026 16:36:20 +1000 Subject: [PATCH 9/9] fix(ai): harden rerank abort classification and clarify usage docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review feedback: - Prefer error identity over signal state in rerank's abort detection. Classifying on `signal.aborted` alone misrouted a genuine failure (e.g. the out-of-range-index throw) to the abort hook whenever a shared/ long-lived signal happened to already be aborted, hiding it from `onError` observers. Reuse the shared abort-name set (now exported as `isAbortShapedError`), which also covers the OpenRouter SDK's `RequestAbortedError`. - Correct the "token counts are 0" claim in `RerankResult.usage` JSDoc and the reranking guide: OpenRouter also reports `totalTokens`/`cost`; only Cohere leaves token counts at 0. - Note in the OpenRouter adapter that the SDK zod-validates responses, so a malformed object body throws "Response validation failed" before mapping — the feared cryptic `undefined.map` TypeError can't occur, so no redundant runtime guard is added. Add a regression test proving a malformed body rejects loudly. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/rerank/rerank.md | 5 ++-- packages/ai-openrouter/src/adapters/rerank.ts | 3 +++ .../tests/rerank-adapter.test.ts | 13 ++++++++++ packages/ai/src/activities/error-payload.ts | 22 ++++++++++++---- packages/ai/src/activities/rerank/index.ts | 15 ++++++++--- packages/ai/src/types.ts | 6 +++-- packages/ai/tests/rerank.test.ts | 25 +++++++++++++++++++ 7 files changed, 76 insertions(+), 13 deletions(-) diff --git a/docs/rerank/rerank.md b/docs/rerank/rerank.md index 71780ff31..45be4bba6 100644 --- a/docs/rerank/rerank.md +++ b/docs/rerank/rerank.md @@ -164,8 +164,9 @@ interface RerankResult { ranking: Array<{ index: number; score: number; document: TDocument }> // The documents reordered by relevance (ranking.map(r => r.document)). rerankedDocuments: Array - // Rerank bills in provider "search units", surfaced on usage.unitsBilled; - // token counts are 0. + // Rerank typically bills in provider "search units" (usage.unitsBilled). + // Some providers (e.g. OpenRouter) also report totalTokens and cost; Cohere + // reports only search units and leaves token counts at 0. usage: TokenUsage } ``` diff --git a/packages/ai-openrouter/src/adapters/rerank.ts b/packages/ai-openrouter/src/adapters/rerank.ts index d959d9b83..c6b92055a 100644 --- a/packages/ai-openrouter/src/adapters/rerank.ts +++ b/packages/ai-openrouter/src/adapters/rerank.ts @@ -66,6 +66,9 @@ export class OpenRouterRerankAdapter< // The SDK types the response as `CreateRerankResponseBody | string`; the // bare-string form is an error/non-JSON payload, not a valid result. + // (A malformed object body never reaches here — the SDK zod-validates the + // response and throws "Response validation failed" first, so `results` is + // a runtime-guaranteed array below.) if (typeof response === 'string') { throw new Error('OpenRouter rerank returned an unexpected response') } diff --git a/packages/ai-openrouter/tests/rerank-adapter.test.ts b/packages/ai-openrouter/tests/rerank-adapter.test.ts index 805516033..5d1fa2764 100644 --- a/packages/ai-openrouter/tests/rerank-adapter.test.ts +++ b/packages/ai-openrouter/tests/rerank-adapter.test.ts @@ -120,6 +120,19 @@ describe('OpenRouterRerankAdapter', () => { expect(body.provider).toEqual({ order: ['cohere'] }) }) + it('rejects loudly on a malformed response body (SDK runtime validation)', async () => { + // The SDK zod-validates the response, so a body missing `results` throws + // before the adapter maps it — the caller never sees a silent bad result + // or a cryptic `undefined.map` TypeError. + fetchMock.mockResolvedValue( + rerankResponse({ id: 'or-1', model: 'cohere/rerank-v3.5', usage: {} }), + ) + + await expect( + rerank({ adapter: adapter(), query: 'q', documents, debug: false }), + ).rejects.toThrow() + }) + it('throws on a non-200 response', async () => { fetchMock.mockResolvedValue( new Response('bad request', { status: 400, statusText: 'Bad Request' }), diff --git a/packages/ai/src/activities/error-payload.ts b/packages/ai/src/activities/error-payload.ts index 42b322b01..6ce3ed547 100644 --- a/packages/ai/src/activities/error-payload.ts +++ b/packages/ai/src/activities/error-payload.ts @@ -18,6 +18,21 @@ const ABORT_ERROR_NAMES = new Set([ 'RequestAbortedError', ]) +/** + * True when a thrown value is an abort-shaped error (DOM `AbortError`, OpenAI + * `APIUserAbortError`, OpenRouter `RequestAbortedError`) — i.e. user-initiated + * cancellation rather than a genuine failure. Matches on the error `name` so + * callers can discriminate aborts without depending on a signal's state or on + * provider-specific message strings. + */ +export function isAbortShapedError(error: unknown): boolean { + if (error && typeof error === 'object') { + const name = (error as { name?: unknown }).name + return typeof name === 'string' && ABORT_ERROR_NAMES.has(name) + } + return false +} + // HTTP status codes carried as numbers (e.g. `error.status = 429`) are a // common variant on SDK error classes; coerce so the resulting `code` field // is stable as a string for downstream consumers. @@ -33,11 +48,8 @@ export function toRunErrorPayload( error: unknown, fallbackMessage = 'Unknown error occurred', ): { message: string; code: string | undefined } { - if (error && typeof error === 'object') { - const name = (error as { name?: unknown }).name - if (typeof name === 'string' && ABORT_ERROR_NAMES.has(name)) { - return { message: 'Request aborted', code: 'aborted' } - } + if (isAbortShapedError(error)) { + return { message: 'Request aborted', code: 'aborted' } } if (error instanceof Error) { const codeField = (error as Error & { code?: unknown }).code diff --git a/packages/ai/src/activities/rerank/index.ts b/packages/ai/src/activities/rerank/index.ts index 5b4b4f175..cf2e23fb0 100644 --- a/packages/ai/src/activities/rerank/index.ts +++ b/packages/ai/src/activities/rerank/index.ts @@ -7,6 +7,7 @@ import { aiEventClient } from '@tanstack/ai-event-client' import { resolveDebugOption } from '../../logger/resolve' +import { isAbortShapedError } from '../error-payload' import { createGenerationContext, runGenerationAbort, @@ -99,10 +100,16 @@ function serializeDocument(document: string | object): string { } function isAbortError(error: unknown, signal?: AbortSignal): boolean { - return ( - signal?.aborted === true || - (error instanceof Error && error.name === 'AbortError') - ) + // Prefer the error's own identity over the signal state. A genuine + // cancellation throws an abort-shaped error (DOM `AbortError`, the OpenRouter + // SDK's `RequestAbortedError`, …). Classifying on `signal.aborted` alone would + // misroute a real failure — e.g. the out-of-range-index throw below — to the + // abort hook whenever a shared/long-lived signal happens to already be + // aborted, hiding it from `onError` observers. + if (isAbortShapedError(error)) return true + // Fall back to signal state only for non-Error throws we can't otherwise + // identify; a real Error with a non-abort name is never an abort. + return error instanceof Error ? false : signal?.aborted === true } // =========================== diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 95cb8b7a1..8897d967c 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -1545,8 +1545,10 @@ export interface RerankResult { /** The documents reordered by relevance — `ranking.map(r => r.document)`. */ rerankedDocuments: Array /** - * Usage for the request. Rerank bills in provider-defined "search units" - * rather than tokens, surfaced via `usage.unitsBilled`; token counts are 0. + * Usage for the request. Rerank typically bills in provider-defined "search + * units" (`usage.unitsBilled`) rather than tokens. Some providers (e.g. + * OpenRouter) may also report `totalTokens` and `cost`; Cohere reports only + * search units and leaves the token counts at 0. */ usage: TokenUsage } diff --git a/packages/ai/tests/rerank.test.ts b/packages/ai/tests/rerank.test.ts index 3430ee258..2015fe625 100644 --- a/packages/ai/tests/rerank.test.ts +++ b/packages/ai/tests/rerank.test.ts @@ -209,6 +209,31 @@ describe('rerank() activity', () => { expect(events.finish).toHaveLength(0) }) + it('classifies a real error as onError even when the signal is already aborted', async () => { + // A shared/long-lived signal can be aborted while a genuine (non-abort) + // error is thrown. The error's identity — not the signal state — decides. + const { middleware, events } = recordingMiddleware() + const controller = new AbortController() + const adapter = mockRerankAdapter(async () => { + controller.abort() + throw new Error('genuine provider failure') + }) + + await expect( + rerank({ + adapter, + query: 'q', + documents: ['a'], + abortSignal: controller.signal, + middleware: [middleware], + debug: false, + }), + ).rejects.toThrow('genuine provider failure') + + expect(events.error).toHaveLength(1) + expect(events.abort).toHaveLength(0) + }) + it('forwards topN and abortSignal to the adapter', async () => { const controller = new AbortController() const adapter = mockRerankAdapter(async () => ranked(0))