diff --git a/.changeset/openrouter-rerank.md b/.changeset/openrouter-rerank.md new file mode 100644 index 000000000..c1013fd74 --- /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` / `appTitle` attribution headers, consistent with the other +OpenRouter adapters. 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/adapters/openrouter.md b/docs/adapters/openrouter.md index 8bc02dba9..abf393d74 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` / `appTitle` 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 0176bb66c..bd28aa9ca 100644 --- a/docs/config.json +++ b/docs/config.json @@ -452,6 +452,16 @@ } ] }, + { + "label": "Reranking", + "children": [ + { + "label": "Reranking", + "to": "rerank/rerank", + "addedAt": "2026-06-25" + } + ] + }, { "label": "Middleware", "children": [ @@ -809,7 +819,8 @@ { "label": "OpenRouter Adapter", "to": "adapters/openrouter", - "addedAt": "2026-04-15" + "addedAt": "2026-04-15", + "updatedAt": "2026-06-25" }, { "label": "OpenAI-Compatible", @@ -817,6 +828,11 @@ "addedAt": "2026-06-01", "updatedAt": "2026-07-20" }, + { + "label": "Cohere", + "to": "adapters/cohere", + "addedAt": "2026-06-25" + }, { "label": "Claude Code", "to": "adapters/claude-code", diff --git a/docs/rerank/rerank.md b/docs/rerank/rerank.md new file mode 100644 index 000000000..f3b5b299e --- /dev/null +++ b/docs/rerank/rerank.md @@ -0,0 +1,357 @@ +--- +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. + +## 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. + +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: + +```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-...')`. + +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 +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 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 +} +``` + +## 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: 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'), + 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. + +## Runnable Example + +`examples/ts-react-rerank` is a small TanStack Start app that runs everything +on this page: a fixed corpus of support articles listed newest-first, a query +box, and a side-by-side view of the original order against the reranked order +with scores. The provider dropdown switches between the Cohere and OpenRouter +adapters over the same `rerank()` call. + +```bash +cd examples/ts-react-rerank +pnpm install +cp .env.example .env # add COHERE_API_KEY and/or OPENROUTER_API_KEY +pnpm dev +``` + +## 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/examples/README.md b/examples/README.md index 9217850fb..d23370be3 100644 --- a/examples/README.md +++ b/examples/README.md @@ -94,6 +94,40 @@ Open `http://localhost:4000` in multiple browser tabs to test multi-user functio --- +### Reranking (ts-react-rerank) + +A single-page app that reorders a fixed set of support articles by relevance to +a query, showing the original order and the reranked order side by side. + +**Tech Stack:** + +- TanStack Start (full-stack React framework) +- `@tanstack/ai` (the `rerank()` activity) +- `@tanstack/ai-cohere` (Cohere rerank adapter) +- `@tanstack/ai-openrouter` (OpenRouter rerank adapter) + +**Features:** + +- ✅ One `rerank()` call, two providers behind a dropdown +- ✅ Object documents — the original typed object comes back on each result +- ✅ Relevance scores and original position per document +- ✅ `topN` control and search-unit usage reporting +- ✅ API keys stay on the server + +**Getting Started:** + +```bash +cd examples/ts-react-rerank +pnpm install +cp .env.example .env +# Add COHERE_API_KEY and/or OPENROUTER_API_KEY — you only need the one you pick +pnpm dev +``` + +📖 [Full Documentation](ts-react-rerank/README.md) + +--- + ### Vanilla Chat A framework-free chat application using pure JavaScript and `@tanstack/ai-client`. diff --git a/examples/ts-react-rerank/.env.example b/examples/ts-react-rerank/.env.example new file mode 100644 index 000000000..1f3dd3dfd --- /dev/null +++ b/examples/ts-react-rerank/.env.example @@ -0,0 +1,11 @@ +# Duplicate this file, rename it to .env, and fill in the key(s) you want to use. +# Only the provider you pick in the UI needs a key. + +# Cohere — https://dashboard.cohere.com/api-keys +# Used by the `cohereRerank(...)` adapter (@tanstack/ai-cohere). +COHERE_API_KEY= + +# OpenRouter — https://openrouter.ai/keys +# Used by the `openRouterRerank(...)` adapter (@tanstack/ai-openrouter). +# Routes to Cohere and NVIDIA rerank models through one key. +OPENROUTER_API_KEY= diff --git a/examples/ts-react-rerank/README.md b/examples/ts-react-rerank/README.md new file mode 100644 index 000000000..0a84fc3f5 --- /dev/null +++ b/examples/ts-react-rerank/README.md @@ -0,0 +1,85 @@ +# Reranking (ts-react-rerank) + +A small TanStack Start app that shows the `rerank()` activity reordering a fixed +set of support articles by relevance to a query. + +The corpus is listed newest-first, which is a bad answer to every query in the +demo. The right-hand column shows what the rerank model does with it, with the +relevance score and the document's original position. + +## Tech stack + +- TanStack Start (full-stack React) +- `@tanstack/ai` — the `rerank()` activity +- `@tanstack/ai-cohere` — `cohereRerank(...)` +- `@tanstack/ai-openrouter` — `openRouterRerank(...)` + +## Getting started + +```bash +cd examples/ts-react-rerank +pnpm install +cp .env.example .env +# Add COHERE_API_KEY, OPENROUTER_API_KEY, or both +pnpm dev +``` + +Open http://localhost:3000. You only need a key for the provider you select. + +- Cohere key: https://dashboard.cohere.com/api-keys +- OpenRouter key: https://openrouter.ai/keys + +## What the example shows + +**One activity, two providers.** `src/lib/server-functions.ts` calls the same +`rerank()` for both providers. Only the adapter changes: + +```ts +// Cohere +await rerank({ + adapter: cohereRerank('rerank-v3.5'), + query, + documents: SUPPORT_DOCS, + topN, + modelOptions: { maxTokensPerDoc: 4096 }, +}) + +// OpenRouter — also reaches NVIDIA rerank models through one key +await rerank({ + adapter: openRouterRerank('cohere/rerank-4-fast'), + query, + documents: SUPPORT_DOCS, + topN, +}) +``` + +**Object documents stay typed.** `SUPPORT_DOCS` is `Array`, not +`Array`. `rerank()` serializes each object with `JSON.stringify` for the +provider, then hands the original object back: + +```tsx +result.ranking[0].document.title // string — no cast, no id lookup +result.ranking[0].score // number — relevance +result.ranking[0].index // number — position in the input array +``` + +**Keys stay on the server.** Both adapters read their key from the environment +inside the server function, so nothing is exposed to the browser. + +**Usage is reported.** Rerank bills in search units rather than tokens, so the +result panel reads `usage.unitsBilled`. OpenRouter also reports `usage.cost`. + +## Files worth reading + +| File | What's in it | +| -------------------------------- | -------------------------------------------------- | +| `src/lib/server-functions.ts` | The `rerank()` calls and adapter selection | +| `src/lib/documents.ts` | The corpus and why its order is deliberately bad | +| `src/lib/models.ts` | Model lists, re-exported from the adapter packages | +| `src/components/RerankPanel.tsx` | The before/after UI | + +## Learn more + +- [Reranking guide](../../docs/rerank/rerank.md) +- [Cohere adapter](../../docs/adapters/cohere.md) +- [OpenRouter adapter](../../docs/adapters/openrouter.md) diff --git a/examples/ts-react-rerank/package.json b/examples/ts-react-rerank/package.json new file mode 100644 index 000000000..7b0dd191e --- /dev/null +++ b/examples/ts-react-rerank/package.json @@ -0,0 +1,34 @@ +{ + "name": "ts-react-rerank", + "private": true, + "type": "module", + "scripts": { + "dev": "vite dev --port 3000", + "build": "vite build", + "serve": "vite preview", + "test": "exit 0", + "test:types": "tsc" + }, + "dependencies": { + "@tailwindcss/vite": "^4.1.18", + "@tanstack/ai": "workspace:*", + "@tanstack/ai-cohere": "workspace:*", + "@tanstack/ai-openrouter": "workspace:*", + "@tanstack/react-router": "^1.158.4", + "@tanstack/react-start": "^1.159.0", + "@tanstack/router-plugin": "^1.158.4", + "lucide-react": "^0.561.0", + "nitro": "3.0.260610-beta", + "react": "^19.2.3", + "react-dom": "^19.2.3", + "tailwindcss": "^4.1.18" + }, + "devDependencies": { + "@types/node": "^24.10.1", + "@types/react": "^19.2.7", + "@types/react-dom": "^19.2.3", + "@vitejs/plugin-react": "^5.1.2", + "typescript": "5.9.3", + "vite": "^8.1.4" + } +} diff --git a/examples/ts-react-rerank/src/components/RerankPanel.tsx b/examples/ts-react-rerank/src/components/RerankPanel.tsx new file mode 100644 index 000000000..d223bbae9 --- /dev/null +++ b/examples/ts-react-rerank/src/components/RerankPanel.tsx @@ -0,0 +1,285 @@ +import { useState } from 'react' +import { ArrowRight, Loader2, TriangleAlert } from 'lucide-react' +import { + MODELS_BY_PROVIDER, + PROVIDERS, + PROVIDER_ENV_VARS, + PROVIDER_LABELS, + defaultModelFor, + isProvider, +} from '@/lib/models' +import { EXAMPLE_QUERIES, SUPPORT_DOCS } from '@/lib/documents' +import { rerankDocumentsFn } from '@/lib/server-functions' +import type { SupportDoc } from '@/lib/documents' +import type { Provider } from '@/lib/models' +import type { RerankResult } from '@tanstack/ai' + +function scoreColor(score: number): string { + if (score >= 0.5) return 'bg-emerald-500' + if (score >= 0.1) return 'bg-amber-500' + return 'bg-gray-600' +} + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error) +} + +export default function RerankPanel() { + const [query, setQuery] = useState(EXAMPLE_QUERIES[0] ?? '') + const [provider, setProvider] = useState('cohere') + const [model, setModel] = useState(() => defaultModelFor('cohere')) + const [topN, setTopN] = useState(5) + const [isRunning, setIsRunning] = useState(false) + const [error, setError] = useState(null) + const [result, setResult] = useState | null>(null) + + function changeProvider(next: Provider) { + setProvider(next) + setModel(defaultModelFor(next)) + } + + async function run() { + setIsRunning(true) + setError(null) + try { + const data = await rerankDocumentsFn({ + data: { query, provider, model, topN }, + }) + setResult(data) + } catch (caught) { + setError(errorMessage(caught)) + setResult(null) + } finally { + setIsRunning(false) + } + } + + const usage = result?.usage + + return ( +
+ {/* ---------- Controls ---------- */} +
+
+ {EXAMPLE_QUERIES.map((example) => ( + + ))} +
+ +
{ + event.preventDefault() + void run() + }} + > + + +
+ + + + + +
+ + +
+
+ + {/* ---------- Error ---------- */} + {error !== null && ( +
+ +
+

Rerank failed

+

{error}

+

+ Set{' '} + {PROVIDER_ENV_VARS[provider]}{' '} + in .env and restart the dev + server. +

+
+
+ )} + + {/* ---------- Results ---------- */} +
+
+

+ Original order +

+
    + {SUPPORT_DOCS.map((doc, index) => ( +
  1. + + {index + 1} + + {doc.title} +
    2. + ))} +
+
+ +
+

+ Reranked{result ? ` — top ${result.ranking.length}` : ''} +

+ + {result === null ? ( +

+ Run a query to see the model reorder these documents by relevance. +

+ ) : ( + <> +
    + {result.ranking.map((entry, position) => ( +
  1. +
    + + {position + 1} + + {/* `entry.document` is the original SupportDoc object, + typed — no id lookup needed. */} + + {entry.document.title} + + + {entry.score.toFixed(3)} + +
    +
    +
    +
    +

    + was #{entry.index + 1} in the original order +

    +
    2. + ))} +
+ +
+
+
model
+
{result.model}
+
+ {usage?.unitsBilled !== undefined && ( +
+
search units billed
+
+ {usage.unitsBilled} +
+
+ )} + {usage !== undefined && usage.totalTokens > 0 && ( +
+
total tokens
+
+ {usage.totalTokens} +
+
+ )} + {usage?.cost !== undefined && ( +
+
cost
+
+ ${usage.cost.toFixed(6)} +
+
+ )} +
+ + )} +
+
+
+ ) +} diff --git a/examples/ts-react-rerank/src/lib/documents.ts b/examples/ts-react-rerank/src/lib/documents.ts new file mode 100644 index 000000000..7307ff1ee --- /dev/null +++ b/examples/ts-react-rerank/src/lib/documents.ts @@ -0,0 +1,81 @@ +/** + * The corpus this example reranks. + * + * These are *objects*, not strings, on purpose: `rerank()` is generic over the + * document element type. Object documents are serialized with `JSON.stringify` + * before they go to the provider, and the original object comes back on + * `ranking[n].document` — fully typed, so the UI can read `.title` off a result + * without a cast or an id lookup. + * + * The array order below is deliberately unhelpful. It is roughly "newest + * article first", which is what a plain CMS listing gives you, and it is a poor + * answer to every query in `EXAMPLE_QUERIES`. That contrast is the whole point + * of the demo: the left column is this order, the right column is what the + * rerank model does with it. + */ +export interface SupportDoc { + id: string + title: string + body: string +} + +export const SUPPORT_DOCS: Array = [ + { + id: 'shipping-zones', + title: 'Shipping zones and delivery estimates', + body: 'Orders ship from the closest fulfilment centre. Domestic delivery takes two to five business days; international delivery takes seven to twenty-one business days and may be held by customs.', + }, + { + id: 'password-reset', + title: 'Resetting your password', + body: 'Use "Forgot password" on the sign-in screen. The reset link is valid for one hour. If it expires, request a new one — old links cannot be reused.', + }, + { + id: 'gift-cards', + title: 'Buying and redeeming gift cards', + body: 'Gift cards are delivered by email and never expire. Redeem one by entering its code at checkout. Gift card balances cannot be transferred back to a bank account.', + }, + { + id: 'cancel-subscription', + title: 'Cancelling your subscription', + body: 'Open Settings → Billing → Cancel plan. Cancellation takes effect at the end of the current billing period, so you keep access until then. You are not charged again after cancelling.', + }, + { + id: 'two-factor', + title: 'Turning on two-factor authentication', + body: 'Settings → Security → Two-factor. We support authenticator apps and hardware keys. Save your recovery codes somewhere safe — support cannot regenerate them for you.', + }, + { + id: 'refund-window', + title: 'Refund window and how refunds are paid', + body: 'Ask for a refund within thirty days of a charge. Approved refunds go back to the original payment method and usually clear within five to ten business days.', + }, + { + id: 'seat-management', + title: 'Adding and removing team seats', + body: 'Workspace owners can add or remove seats at any time. Adding a seat is billed pro rata immediately; removing a seat credits the unused time to your next invoice.', + }, + { + id: 'downgrade-plan', + title: 'Downgrading instead of cancelling', + body: 'If you want to stop paying but keep your data, downgrade to the free tier rather than cancelling. Downgrades apply at the next renewal and your projects stay read-only rather than being deleted.', + }, + { + id: 'export-data', + title: 'Exporting your data', + body: 'Settings → Data → Export produces a ZIP archive of your projects as JSON. Exports are generated in the background; you get an email with a download link valid for 24 hours.', + }, + { + id: 'api-rate-limits', + title: 'API rate limits', + body: 'The API allows 600 requests per minute per key. Exceeding it returns HTTP 429 with a Retry-After header. Rate limits are per key, not per workspace.', + }, +] + +/** Prompts shown as one-click chips above the query box. */ +export const EXAMPLE_QUERIES = [ + 'how do I stop being billed every month?', + 'I want my money back for a charge', + 'I lost access to my account', + 'can I keep my projects without paying?', +] diff --git a/examples/ts-react-rerank/src/lib/models.ts b/examples/ts-react-rerank/src/lib/models.ts new file mode 100644 index 000000000..66628f155 --- /dev/null +++ b/examples/ts-react-rerank/src/lib/models.ts @@ -0,0 +1,42 @@ +/** + * The provider / model registry the UI drives. + * + * The model lists are the ones each adapter package exports, not copies — so + * this example stays correct when a package gains a model. + */ +import { COHERE_RERANK_MODELS } from '@tanstack/ai-cohere' +import { OPENROUTER_RERANK_MODELS } from '@tanstack/ai-openrouter' + +export const PROVIDERS = ['cohere', 'openrouter'] as const + +export type Provider = (typeof PROVIDERS)[number] + +export const PROVIDER_LABELS: Record = { + cohere: 'Cohere', + openrouter: 'OpenRouter', +} + +/** The env var each provider's adapter reads, shown in the "no key" hint. */ +export const PROVIDER_ENV_VARS: Record = { + cohere: 'COHERE_API_KEY', + openrouter: 'OPENROUTER_API_KEY', +} + +export const MODELS_BY_PROVIDER: Record> = { + cohere: COHERE_RERANK_MODELS, + // OpenRouter's rerank model type is open — any rerank slug it serves works, + // so this list is autocomplete sugar rather than an exhaustive set. + openrouter: OPENROUTER_RERANK_MODELS, +} + +export function defaultModelFor(provider: Provider): string { + const [first] = MODELS_BY_PROVIDER[provider] + if (first === undefined) { + throw new Error(`No rerank models registered for provider ${provider}`) + } + return first +} + +export function isProvider(value: string): value is Provider { + return PROVIDERS.some((provider) => provider === value) +} diff --git a/examples/ts-react-rerank/src/lib/server-functions.ts b/examples/ts-react-rerank/src/lib/server-functions.ts new file mode 100644 index 000000000..3df9035e9 --- /dev/null +++ b/examples/ts-react-rerank/src/lib/server-functions.ts @@ -0,0 +1,75 @@ +import { createServerFn } from '@tanstack/react-start' +import { rerank } from '@tanstack/ai' +import { COHERE_RERANK_MODELS, cohereRerank } from '@tanstack/ai-cohere' +import { openRouterRerank } from '@tanstack/ai-openrouter' +import { SUPPORT_DOCS } from './documents' +import { isProvider } from './models' +import type { CohereRerankModel } from '@tanstack/ai-cohere' +import type { Provider } from './models' + +interface RerankInput { + query: string + provider: Provider + model: string + topN: number +} + +/** + * Narrows a wire string to a Cohere rerank model. The `cohereRerank` factory is + * generic over the model literal, so the model has to be a known slug before it + * reaches the adapter — a plain `string` would not type-check. + */ +function isCohereRerankModel(model: string): model is CohereRerankModel { + return COHERE_RERANK_MODELS.some((known) => known === model) +} + +/** + * Reranks the support corpus against a query. + * + * The API key never leaves the server: both adapters read their key from the + * environment (`COHERE_API_KEY` / `OPENROUTER_API_KEY`) inside this handler. + * + * Note the two branches call the *same* `rerank()` with a different adapter — + * that is the provider-agnostic contract. They are written out separately + * rather than sharing an `adapter` variable so each call site keeps the + * adapter's literal model type, and with it the per-model `modelOptions` + * inference. + */ +export const rerankDocumentsFn = createServerFn({ method: 'POST' }) + .inputValidator((data: RerankInput) => { + if (!data.query.trim()) throw new Error('Query is required') + if (!isProvider(data.provider)) { + throw new Error(`Unknown provider: ${data.provider}`) + } + if (!data.model) throw new Error('Model is required') + if (!Number.isInteger(data.topN) || data.topN < 1) { + throw new Error('topN must be a positive integer') + } + return data + }) + .handler(async ({ data }) => { + const { query, model, topN } = data + + if (data.provider === 'cohere') { + if (!isCohereRerankModel(model)) { + throw new Error(`Unknown Cohere rerank model: ${model}`) + } + return await rerank({ + adapter: cohereRerank(model), + query, + // Object documents: JSON-serialized on the way out, and the original + // `SupportDoc` comes back on `ranking[n].document`. + documents: SUPPORT_DOCS, + topN, + // Per-model provider option, typed by the `cohereRerank(model)` literal. + modelOptions: { maxTokensPerDoc: 4096 }, + }) + } + + return await rerank({ + adapter: openRouterRerank(model), + query, + documents: SUPPORT_DOCS, + topN, + }) + }) diff --git a/examples/ts-react-rerank/src/routeTree.gen.ts b/examples/ts-react-rerank/src/routeTree.gen.ts new file mode 100644 index 000000000..dceedffdc --- /dev/null +++ b/examples/ts-react-rerank/src/routeTree.gen.ts @@ -0,0 +1,68 @@ +/* eslint-disable */ + +// @ts-nocheck + +// noinspection JSUnusedGlobalSymbols + +// This file was automatically generated by TanStack Router. +// You should NOT make any changes in this file as it will be overwritten. +// Additionally, you should also exclude this file from your linter and/or formatter to prevent it from being checked or modified. + +import { Route as rootRouteImport } from './routes/__root' +import { Route as IndexRouteImport } from './routes/index' + +const IndexRoute = IndexRouteImport.update({ + id: '/', + path: '/', + getParentRoute: () => rootRouteImport, +} as any) + +export interface FileRoutesByFullPath { + '/': typeof IndexRoute +} +export interface FileRoutesByTo { + '/': typeof IndexRoute +} +export interface FileRoutesById { + __root__: typeof rootRouteImport + '/': typeof IndexRoute +} +export interface FileRouteTypes { + fileRoutesByFullPath: FileRoutesByFullPath + fullPaths: '/' + fileRoutesByTo: FileRoutesByTo + to: '/' + id: '__root__' | '/' + fileRoutesById: FileRoutesById +} +export interface RootRouteChildren { + IndexRoute: typeof IndexRoute +} + +declare module '@tanstack/react-router' { + interface FileRoutesByPath { + '/': { + id: '/' + path: '/' + fullPath: '/' + preLoaderRoute: typeof IndexRouteImport + parentRoute: typeof rootRouteImport + } + } +} + +const rootRouteChildren: RootRouteChildren = { + IndexRoute: IndexRoute, +} +export const routeTree = rootRouteImport + ._addFileChildren(rootRouteChildren) + ._addFileTypes() + +import type { getRouter } from './router.tsx' +import type { createStart } from '@tanstack/react-start' +declare module '@tanstack/react-start' { + interface Register { + ssr: true + router: Awaited> + } +} diff --git a/examples/ts-react-rerank/src/router.tsx b/examples/ts-react-rerank/src/router.tsx new file mode 100644 index 000000000..a59544464 --- /dev/null +++ b/examples/ts-react-rerank/src/router.tsx @@ -0,0 +1,10 @@ +import { createRouter } from '@tanstack/react-router' +import { routeTree } from './routeTree.gen' + +export const getRouter = () => { + return createRouter({ + routeTree, + scrollRestoration: true, + defaultPreloadStaleTime: 0, + }) +} diff --git a/examples/ts-react-rerank/src/routes/__root.tsx b/examples/ts-react-rerank/src/routes/__root.tsx new file mode 100644 index 000000000..450a2bd82 --- /dev/null +++ b/examples/ts-react-rerank/src/routes/__root.tsx @@ -0,0 +1,41 @@ +import { HeadContent, Scripts, createRootRoute } from '@tanstack/react-router' +import appCss from '../styles.css?url' + +export const Route = createRootRoute({ + head: () => ({ + meta: [ + { + charSet: 'utf-8', + }, + { + name: 'viewport', + content: 'width=device-width, initial-scale=1', + }, + { + title: 'TanStack AI — Reranking', + }, + ], + links: [ + { + rel: 'stylesheet', + href: appCss, + }, + ], + }), + + shellComponent: RootDocument, +}) + +function RootDocument({ children }: { children: React.ReactNode }) { + return ( + + + + + + {children} + + + + ) +} diff --git a/examples/ts-react-rerank/src/routes/index.tsx b/examples/ts-react-rerank/src/routes/index.tsx new file mode 100644 index 000000000..9f2e991b0 --- /dev/null +++ b/examples/ts-react-rerank/src/routes/index.tsx @@ -0,0 +1,28 @@ +import { createFileRoute } from '@tanstack/react-router' +import RerankPanel from '@/components/RerankPanel' + +function RerankPage() { + return ( +
+
+
+

+ Document Reranking +

+

+ A fixed set of support articles, listed newest-first. The rerank + model reorders them by how well each one answers the query — the + same rerank() call, + with a Cohere or an OpenRouter adapter. +

+
+ + +
+
+ ) +} + +export const Route = createFileRoute('/')({ + component: RerankPage, +}) diff --git a/examples/ts-react-rerank/src/styles.css b/examples/ts-react-rerank/src/styles.css new file mode 100644 index 000000000..2cd2c65a0 --- /dev/null +++ b/examples/ts-react-rerank/src/styles.css @@ -0,0 +1,10 @@ +@import 'tailwindcss'; + +body { + @apply m-0; + font-family: + -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Roboto', 'Oxygen', 'Ubuntu', + 'Cantarell', 'Fira Sans', 'Droid Sans', 'Helvetica Neue', sans-serif; + -webkit-font-smoothing: antialiased; + -moz-osx-font-smoothing: grayscale; +} diff --git a/examples/ts-react-rerank/tsconfig.json b/examples/ts-react-rerank/tsconfig.json new file mode 100644 index 000000000..afc73b0a0 --- /dev/null +++ b/examples/ts-react-rerank/tsconfig.json @@ -0,0 +1,19 @@ +{ + "extends": "../../tsconfig.json", + "include": ["**/*.ts", "**/*.tsx"], + "compilerOptions": { + "target": "ES2022", + "jsx": "react-jsx", + "module": "ESNext", + "types": ["vite/client"], + "declarationMap": true, + "allowImportingTsExtensions": true, + "verbatimModuleSyntax": false, + "noFallthroughCasesInSwitch": true, + "noUncheckedSideEffectImports": true, + "baseUrl": ".", + "paths": { + "@/*": ["./src/*"] + } + } +} diff --git a/examples/ts-react-rerank/vite.config.ts b/examples/ts-react-rerank/vite.config.ts new file mode 100644 index 000000000..a8459f2e5 --- /dev/null +++ b/examples/ts-react-rerank/vite.config.ts @@ -0,0 +1,11 @@ +import { defineConfig } from 'vite' +import { tanstackStart } from '@tanstack/react-start/plugin/vite' +import { nitro } from 'nitro/vite' +import viteReact from '@vitejs/plugin-react' +import tailwindcss from '@tailwindcss/vite' + +export default defineConfig({ + resolve: { tsconfigPaths: true }, + plugins: [tailwindcss(), tanstackStart(), nitro(), viteReact()], + nitro: {}, +}) 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..915d92a79 --- /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": "oxlint src --type-aware --fix", + "test:build": "publint --strict", + "test:oxlint": "oxlint src --type-aware", + "test:lib": "vitest run", + "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": "^8.1.4" + } +} diff --git a/packages/ai-cohere/src/adapters/rerank.ts b/packages/ai-cohere/src/adapters/rerank.ts new file mode 100644 index 000000000..5298175cf --- /dev/null +++ b/packages/ai-cohere/src/adapters/rerank.ts @@ -0,0 +1,181 @@ +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..2cf2b8ec0 --- /dev/null +++ b/packages/ai-cohere/src/utils/client.ts @@ -0,0 +1,45 @@ +/** + * 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 windowEnv = + typeof globalThis !== 'undefined' && + (globalThis as Record).window + ? (( + (globalThis as Record).window as Record< + string, + unknown + > + ).env as Record | undefined) + : 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).', + ) + } + 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..de6ce7688 --- /dev/null +++ b/packages/ai-cohere/tests/rerank-adapter.test.ts @@ -0,0 +1,136 @@ +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 7465168be..74ee00e60 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 // =========================== @@ -1152,6 +1184,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-openrouter/src/adapters/rerank.ts b/packages/ai-openrouter/src/adapters/rerank.ts new file mode 100644 index 000000000..c6b92055a --- /dev/null +++ b/packages/ai-openrouter/src/adapters/rerank.ts @@ -0,0 +1,143 @@ +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. + // (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') + } + + 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 63ceffaed..c16dbbfa7 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..f35b21bea --- /dev/null +++ b/packages/ai-openrouter/src/rerank/rerank-provider-options.ts @@ -0,0 +1,44 @@ +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..5d1fa2764 --- /dev/null +++ b/packages/ai-openrouter/tests/rerank-adapter.test.ts @@ -0,0 +1,145 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { rerank } from '@tanstack/ai' +import { createOpenRouterRerank } from '../src/adapters/rerank' + +// 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(() => { + vi.stubGlobal('fetch', fetchMock) + fetchMock.mockReset() +}) + +afterEach(() => { + vi.unstubAllGlobals() +}) + +const documents = ['sunny day at the beach', 'rainy afternoon in the city'] + +/** 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, relevance_score: 0.97 }, + { document: { text: documents[0] }, index: 0, relevance_score: 0.1 }, + ], + 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('hits the /rerank endpoint and maps the response', async () => { + fetchMock.mockResolvedValue(rerankResponse(wireBody())) + + const result = await rerank({ + adapter: adapter(), + 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] }, + { index: 0, score: 0.1, document: documents[0] }, + ]) + }) + + it('maps usage (search_units/cost/total_tokens)', async () => { + fetchMock.mockResolvedValue(rerankResponse(wireBody())) + + 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 () => { + const model = 'nvidia/llama-nemotron-rerank-vl-1b-v2' + fetchMock.mockResolvedValue(rerankResponse({ ...wireBody(), model })) + + await rerank({ + adapter: createOpenRouterRerank(model, 'sk-or-test'), + query: 'q', + documents, + }) + + const { body } = await capturedRequest() + expect(body.model).toBe(model) + }) + + it('forwards provider routing preferences into the request body', async () => { + fetchMock.mockResolvedValue(rerankResponse(wireBody())) + + await rerank({ + adapter: adapter(), + query: 'q', + documents, + modelOptions: { provider: { order: ['cohere'] } }, + }) + + const { body } = await capturedRequest() + 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' }), + ) + + await expect( + rerank({ adapter: adapter(), query: 'q', documents, debug: false }), + ).rejects.toThrow() + }) +}) 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/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 8ac41db0a..a3b792451 100644 --- a/packages/ai/src/activities/middleware/types.ts +++ b/packages/ai/src/activities/middleware/types.ts @@ -41,6 +41,7 @@ export type GenerationActivity = | 'audio' | 'tts' | 'transcription' + | 'rerank' | 'summarize' /** diff --git a/packages/ai/src/activities/rerank/adapter.ts b/packages/ai/src/activities/rerank/adapter.ts new file mode 100644 index 000000000..015e05631 --- /dev/null +++ b/packages/ai/src/activities/rerank/adapter.ts @@ -0,0 +1,90 @@ +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).slice(2, 9)}` + } +} diff --git a/packages/ai/src/activities/rerank/index.ts b/packages/ai/src/activities/rerank/index.ts new file mode 100644 index 000000000..15ad1d16c --- /dev/null +++ b/packages/ai/src/activities/rerank/index.ts @@ -0,0 +1,302 @@ +/** + * 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 { isAbortShapedError } from '../error-payload' +import { + createGenerationContext, + runGenerationAbort, + runGenerationError, + runGenerationFinish, + runGenerationStart, + runGenerationUsage, +} from '../middleware/run' +import type { InternalLogger } from '../../logger/internal-logger' +import type { DebugOption } from '../../logger/types' +import type { GenerationMiddleware } from '../middleware/types' +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 { + '~types': { providerOptions: infer P extends object } +} + ? P + : 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 { + // 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 +} + +// =========================== +// 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) => { + 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 + + 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 20e5e7cb7..3388c86cf 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 e5938bd8c..82858dd6f 100644 --- a/packages/ai/src/middlewares/otel.ts +++ b/packages/ai/src/middlewares/otel.ts @@ -85,6 +85,7 @@ const OPERATION_NAME: Record = { audio: 'audio_generation', tts: 'text_to_speech', transcription: 'transcription', + rerank: 'rerank', summarize: 'summarize', } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 71ea2a39f..6893dd710 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -2007,6 +2007,71 @@ 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 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 +} + // ============================================================================ // 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..2015fe625 --- /dev/null +++ b/packages/ai/tests/rerank.test.ts @@ -0,0 +1,264 @@ +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('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)) + + 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) + }) + + 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') + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index d5ff75a6a..a1c43ede1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -977,6 +977,64 @@ importers: specifier: 5.9.3 version: 5.9.3 + examples/ts-react-rerank: + dependencies: + '@tailwindcss/vite': + specifier: ^4.1.18 + version: 4.1.18(vite@8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0)) + '@tanstack/ai': + specifier: workspace:* + version: link:../../packages/ai + '@tanstack/ai-cohere': + specifier: workspace:* + version: link:../../packages/ai-cohere + '@tanstack/ai-openrouter': + specifier: workspace:* + version: link:../../packages/ai-openrouter + '@tanstack/react-router': + specifier: ^1.158.4 + version: 1.159.5(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + '@tanstack/react-start': + specifier: ^1.159.0 + version: 1.159.5(crossws@0.4.6(srvx@0.11.17))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(vite-plugin-solid@2.11.10(solid-js@1.9.10)(vite@8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0)))(vite@8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0)) + '@tanstack/router-plugin': + specifier: ^1.158.4 + version: 1.159.5(@tanstack/react-router@1.159.5(react-dom@19.2.3(react@19.2.3))(react@19.2.3))(vite-plugin-solid@2.11.10(solid-js@1.9.10)(vite@8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0)))(vite@8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0)) + lucide-react: + specifier: ^0.561.0 + version: 0.561.0(react@19.2.3) + nitro: + specifier: 3.0.260610-beta + version: 3.0.260610-beta(@electric-sql/pglite@0.4.3)(aws4fetch@1.0.20)(better-sqlite3@12.11.1)(chokidar@5.0.0)(dotenv@17.4.2)(drizzle-orm@0.45.2(@electric-sql/pglite@0.4.3)(@opentelemetry/api@1.9.1)(@prisma/client@7.9.0(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3))(typescript@5.9.3))(better-sqlite3@12.11.1)(bun-types@1.3.14)(mysql2@3.15.3)(postgres@3.4.7)(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3)))(giget@3.3.0)(jiti@2.7.0)(miniflare@4.20260617.1)(mysql2@3.15.3)(rollup@4.60.1)(vite@8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0))(wrangler@4.103.0) + react: + specifier: ^19.2.3 + version: 19.2.3 + react-dom: + specifier: ^19.2.3 + version: 19.2.3(react@19.2.3) + tailwindcss: + specifier: ^4.1.18 + version: 4.1.18 + devDependencies: + '@types/node': + specifier: ^24.10.1 + version: 24.10.3 + '@types/react': + specifier: ^19.2.7 + version: 19.2.7 + '@types/react-dom': + specifier: ^19.2.3 + version: 19.2.3(@types/react@19.2.7) + '@vitejs/plugin-react': + specifier: ^5.1.2 + version: 5.1.2(vite@8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0)) + typescript: + specifier: 5.9.3 + version: 5.9.3 + vite: + specifier: ^8.1.4 + version: 8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0) + examples/ts-react-search: dependencies: '@radix-ui/react-slot': @@ -1706,6 +1764,18 @@ importers: specifier: 4.0.14 version: 4.0.14(supports-color@7.2.0)(vitest@4.1.10) + packages/ai-cohere: + devDependencies: + '@tanstack/ai': + specifier: workspace:* + version: link:../ai + '@vitest/coverage-v8': + specifier: 4.0.14 + version: 4.0.14(supports-color@7.2.0)(vitest@4.1.10) + vite: + specifier: ^8.1.4 + version: 8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0) + packages/ai-devtools: dependencies: '@tanstack/ai': @@ -29698,6 +29768,60 @@ snapshots: - uploadthing - wrangler + nitro@3.0.260610-beta(@electric-sql/pglite@0.4.3)(aws4fetch@1.0.20)(better-sqlite3@12.11.1)(chokidar@5.0.0)(dotenv@17.4.2)(drizzle-orm@0.45.2(@electric-sql/pglite@0.4.3)(@opentelemetry/api@1.9.1)(@prisma/client@7.9.0(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3))(typescript@5.9.3))(better-sqlite3@12.11.1)(bun-types@1.3.14)(mysql2@3.15.3)(postgres@3.4.7)(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3)))(giget@3.3.0)(jiti@2.7.0)(miniflare@4.20260617.1)(mysql2@3.15.3)(rollup@4.60.1)(vite@8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0))(wrangler@4.103.0): + dependencies: + consola: 3.4.2 + crossws: 0.4.6(srvx@0.11.17) + db0: 0.3.4(@electric-sql/pglite@0.4.3)(better-sqlite3@12.11.1)(drizzle-orm@0.45.2(@electric-sql/pglite@0.4.3)(@opentelemetry/api@1.9.1)(@prisma/client@7.9.0(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3))(typescript@5.9.3))(better-sqlite3@12.11.1)(bun-types@1.3.14)(mysql2@3.15.3)(postgres@3.4.7)(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3)))(mysql2@3.15.3) + env-runner: 0.1.14(miniflare@4.20260617.1)(wrangler@4.103.0) + h3: 2.0.1-rc.22(crossws@0.4.6(srvx@0.11.17)) + hookable: 6.1.1 + nf3: 0.3.17 + ocache: 0.1.5 + ofetch: 2.0.0-alpha.3 + ohash: 2.0.11 + rolldown: 1.1.5 + srvx: 0.11.17 + unenv: 2.0.0-rc.24 + unstorage: 2.0.0-alpha.7(aws4fetch@1.0.20)(chokidar@5.0.0)(db0@0.3.4(@electric-sql/pglite@0.4.3)(better-sqlite3@12.11.1)(drizzle-orm@0.45.2(@electric-sql/pglite@0.4.3)(@opentelemetry/api@1.9.1)(@prisma/client@7.9.0(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3))(typescript@5.9.3))(better-sqlite3@12.11.1)(bun-types@1.3.14)(mysql2@3.15.3)(postgres@3.4.7)(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3)))(mysql2@3.15.3))(ofetch@2.0.0-alpha.3) + optionalDependencies: + dotenv: 17.4.2 + giget: 3.3.0 + jiti: 2.7.0 + rollup: 4.60.1 + vite: 8.1.4(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.9.0) + transitivePeerDependencies: + - '@azure/app-configuration' + - '@azure/cosmos' + - '@azure/data-tables' + - '@azure/identity' + - '@azure/keyvault-secrets' + - '@azure/storage-blob' + - '@capacitor/preferences' + - '@deno/kv' + - '@electric-sql/pglite' + - '@libsql/client' + - '@netlify/blobs' + - '@netlify/runtime' + - '@planetscale/database' + - '@upstash/redis' + - '@vercel/blob' + - '@vercel/functions' + - '@vercel/kv' + - aws4fetch + - better-sqlite3 + - chokidar + - drizzle-orm + - idb-keyval + - ioredis + - lru-cache + - miniflare + - mongodb + - mysql2 + - sqlite3 + - uploadthing + - wrangler + nitropack@2.13.1(@electric-sql/pglite@0.4.3)(aws4fetch@1.0.20)(better-sqlite3@12.11.1)(drizzle-orm@0.45.2(@cloudflare/workers-types@4.20260317.1)(@electric-sql/pglite@0.4.3)(@opentelemetry/api@1.9.1)(@prisma/client@7.9.0(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3))(typescript@5.9.3))(better-sqlite3@12.11.1)(bun-types@1.3.14)(mysql2@3.15.3)(postgres@3.4.7)(prisma@7.9.0(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(better-sqlite3@12.11.1)(magicast@0.5.2)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(typescript@5.9.3)))(mysql2@3.15.3)(rolldown@1.1.5): dependencies: '@cloudflare/kv-asset-handler': 0.4.2