feat: add rerank() activity with Cohere and OpenRouter adapters#845
Open
AlemTuzlak wants to merge 7 commits into
Open
feat: add rerank() activity with Cohere and OpenRouter adapters#845AlemTuzlak wants to merge 7 commits into
AlemTuzlak wants to merge 7 commits into
Conversation
Add a provider-agnostic rerank() activity to @tanstack/ai for reordering documents by relevance to a query, plus a new rerank-only @tanstack/ai-cohere provider package. - rerank() is generic over the document type: strings or JSON-serializable objects. Object documents are serialized for the provider and the original element is returned in the result, fully typed. - Supports topN, per-request cancellation via abortSignal, observe-only GenerationMiddleware (start/usage/finish/abort/error), and rerank:* devtools events. Rerank bills in provider search units, surfaced on usage.unitsBilled. - @tanstack/ai-cohere ships cohereRerank / createCohereRerank, talking to Cohere /v2/rerank over fetch (no SDK) with per-model provider options. - Unit tests for the activity and the Cohere adapter; docs for the rerank guide and the Cohere adapter page.
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds rerank support to the AI core, Cohere and OpenRouter adapters, rerank devtools events, and matching documentation, navigation, and tests. ChangesRerank feature
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
🚀 Changeset Version Preview4 package(s) bumped directly, 29 bumped as dependents. 🟥 Major bumps
🟨 Minor bumps
🟩 Patch bumps
|
|
View your CI Pipeline Execution ↗ for commit f0f2fbb
☁️ Nx Cloud last updated this comment at |
@tanstack/ai
@tanstack/ai-angular
@tanstack/ai-anthropic
@tanstack/ai-client
@tanstack/ai-code-mode
@tanstack/ai-code-mode-skills
@tanstack/ai-cohere
@tanstack/ai-devtools-core
@tanstack/ai-elevenlabs
@tanstack/ai-event-client
@tanstack/ai-fal
@tanstack/ai-gemini
@tanstack/ai-grok
@tanstack/ai-groq
@tanstack/ai-isolate-cloudflare
@tanstack/ai-isolate-node
@tanstack/ai-isolate-quickjs
@tanstack/ai-mcp
@tanstack/ai-ollama
@tanstack/ai-openai
@tanstack/ai-openrouter
@tanstack/ai-preact
@tanstack/ai-react
@tanstack/ai-react-ui
@tanstack/ai-solid
@tanstack/ai-solid-ui
@tanstack/ai-svelte
@tanstack/ai-utils
@tanstack/ai-vue
@tanstack/ai-vue-ui
@tanstack/openai-base
@tanstack/preact-ai-devtools
@tanstack/react-ai-devtools
@tanstack/solid-ai-devtools
commit: |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
packages/ai/src/activities/rerank/adapter.ts (1)
78-81: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: parameter order makes the default unusable.
confighas a default but precedes the requiredmodel, so callers can never omitconfigto rely on the default. Consider reordering (modelfirst) or making both explicit.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/ai/src/activities/rerank/adapter.ts` around lines 78 - 81, The RerankAdapter constructor currently puts the optional config parameter before the required model, so the default config cannot actually be omitted by callers. Update the constructor signature in RerankAdapter to either place model before config or make both parameters explicit, and make sure any call sites or related initialization in RerankAdapterConfig usage still pass arguments in the correct order.packages/ai-cohere/src/adapters/rerank.ts (1)
117-120: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winJSON parse / shape-validation failures bypass
logger.errors().
await response.json()can reject on invalid JSON, and the shape-mismatch path throws too — neither routes throughlogger.errors(). The upstreamRerankOptionscontract states adapters must calllogger.errors()in catch blocks, so these failures will be silent in adapter logging unlike the fetch and!response.okpaths.🛠️ Proposed logging on parse/validation failure
- const json: unknown = await response.json() - if (!isCohereRerankResponse(json)) { - throw new Error('Cohere rerank response had an unexpected shape') - } + let json: unknown + try { + json = await response.json() + } catch (error) { + logger.errors(`${this.name}.rerank fatal`, { + error, + source: `${this.name}.rerank`, + }) + throw error + } + if (!isCohereRerankResponse(json)) { + const error = new Error('Cohere rerank response had an unexpected shape') + logger.errors(`${this.name}.rerank fatal`, { + error, + source: `${this.name}.rerank`, + }) + throw error + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/ai-cohere/src/adapters/rerank.ts` around lines 117 - 120, The JSON parse and response-shape validation path in rerank.ts bypasses adapter error logging, so update the rerank adapter flow around `response.json()` and `isCohereRerankResponse` to ensure failures are handled in a catch block that calls `logger.errors()`. Keep the existing fetch and `!response.ok` logging behavior, but wrap the parse/validation steps so any invalid JSON or unexpected shape is logged before rethrowing, using the same `RerankOptions` logger contract.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/rerank/rerank.md`:
- Around line 161-165: The server example in rerank.md uses an `as` type
assertion on `body`, which violates the docs guideline and should be removed.
Update the example to type-check without assertions by validating `body` first
with a runtime guard or schema, then access `query`, `documents`, and `topN`
through a narrowed type. Refer to the server example destructuring of `body` in
the rerank docs and keep the sample aligned with the existing request-handling
flow.
In `@packages/ai-cohere/src/utils/client.ts`:
- Around line 23-43: The env lookup in getCohereApiKeyFromEnv currently prefers
window.env and skips process.env whenever window exists, which breaks cases
where window.env is unset but COHERE_API_KEY is still available via process.env.
Update the fallback logic in getCohereApiKeyFromEnv so it first reads window.env
when present, but if that is missing or undefined it then checks process.env
before throwing. Keep the existing error behavior and use the same
COHERE_API_KEY symbol path so the function works in browser-like and
Electron/bundler setups.
In `@packages/ai/src/activities/rerank/adapter.ts`:
- Around line 87-89: The fallback ID generation in generateId() can return an
empty or very short random suffix because substring(7) is applied to a
non-fixed-length Math.random().toString(36) value. Update the generateId()
implementation in the rerank adapter to use a more robust random segment, and
align it with the createId() approach from index.ts by slicing a stable range
such as slice(2, 9) so the tracing/devtools correlation IDs remain sufficiently
unique.
In `@packages/ai/src/activities/rerank/index.ts`:
- Around line 210-215: The rerank mapping in the `ranking` construction is
trusting provider-supplied `r.index` values and masking invalid lookups with a
type cast, so add a validation step in the `rerank` flow before remapping
`documents[r.index]`. If any index is out of range or resolves to `undefined`,
throw a clear error instead of building `ranking` and `rerankedDocuments` with
corrupted entries, and keep the guard close to the `result.ranking.map(...)`
logic so the failure is caught where the provider contract is consumed.
---
Nitpick comments:
In `@packages/ai-cohere/src/adapters/rerank.ts`:
- Around line 117-120: The JSON parse and response-shape validation path in
rerank.ts bypasses adapter error logging, so update the rerank adapter flow
around `response.json()` and `isCohereRerankResponse` to ensure failures are
handled in a catch block that calls `logger.errors()`. Keep the existing fetch
and `!response.ok` logging behavior, but wrap the parse/validation steps so any
invalid JSON or unexpected shape is logged before rethrowing, using the same
`RerankOptions` logger contract.
In `@packages/ai/src/activities/rerank/adapter.ts`:
- Around line 78-81: The RerankAdapter constructor currently puts the optional
config parameter before the required model, so the default config cannot
actually be omitted by callers. Update the constructor signature in
RerankAdapter to either place model before config or make both parameters
explicit, and make sure any call sites or related initialization in
RerankAdapterConfig usage still pass arguments in the correct order.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 24968d99-0670-48fd-ac64-4f4ae084631a
⛔ Files ignored due to path filters (1)
pnpm-lock.yamlis excluded by!**/pnpm-lock.yaml
📒 Files selected for processing (23)
.changeset/rerank-cohere.mddocs/adapters/cohere.mddocs/config.jsondocs/rerank/rerank.mdpackages/ai-cohere/LICENSEpackages/ai-cohere/README.mdpackages/ai-cohere/package.jsonpackages/ai-cohere/src/adapters/rerank.tspackages/ai-cohere/src/index.tspackages/ai-cohere/src/model-meta.tspackages/ai-cohere/src/utils/client.tspackages/ai-cohere/tests/rerank-adapter.test.tspackages/ai-cohere/tsconfig.jsonpackages/ai-cohere/vite.config.tspackages/ai-event-client/src/index.tspackages/ai/src/activities/index.tspackages/ai/src/activities/middleware/types.tspackages/ai/src/activities/rerank/adapter.tspackages/ai/src/activities/rerank/index.tspackages/ai/src/index.tspackages/ai/src/middlewares/otel.tspackages/ai/src/types.tspackages/ai/tests/rerank.test.ts
Add openRouterRerank / createOpenRouterRerank for reordering documents by relevance through OpenRouter's unified /v1/rerank endpoint, using the @openrouter/sdk rerank method. - Model-agnostic: supports every rerank model OpenRouter offers (Cohere rerank-v3.5 / 4-fast / 4-pro, NVIDIA llama-nemotron-rerank, and any future slug) via an open model type with known-model autocomplete. - Maps the SDK response to the shared RerankAdapterResult, including usage (searchUnits -> unitsBilled, cost, totalTokens). - Optional provider-routing preferences via modelOptions.provider; abortSignal forwarded through the SDK fetchOptions. - Unit tests (SDK mocked) and docs (rerank guide + OpenRouter adapter page).
AlemTuzlak
changed the title
…Title) The SDK-based OpenRouter adapters use SDKOptions, whose attribution field is appTitle, not xTitle (xTitle only exists on the legacy raw-fetch client helper). Fix the rerank section of the OpenRouter adapter docs and the changeset.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
docs/rerank/rerank.md (1)
189-193: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRemove the remaining
ascast from the server example.The server sample still uses a type assertion, which breaks the docs rule and makes the example stop type-checking cleanly without casts. Narrow
bodywith a guard/schema instead.As per coding guidelines,
docs/**/*.md: Noastype-assertion casts in code samples underdocs/.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/rerank/rerank.md` around lines 189 - 193, The server example in the rerank docs still uses a type assertion on body, which should be removed to match the docs rule. Update the code sample around the request handling in the rerank server example to narrow body with a guard or schema instead of using an as cast, and keep the destructuring of query, documents, and topN inside the validated path.Source: Coding guidelines
🧹 Nitpick comments (1)
docs/adapters/openrouter.md (1)
263-275: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the current Cohere rerank model in the example.
The example still anchors on
cohere/rerank-v3.5, even though this section already calls out newer Cohere slugs. Please switch the sample (and the explicit-key variant in the paragraph below) to whatevermodel-meta.tsmarks as the latest Cohere rerank model. As per coding guidelines, "docs/**/*.md: Use the latest model per provider in documentation examples, sourced from each adapter's model-meta.ts`".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/adapters/openrouter.md` around lines 263 - 275, The OpenRouter rerank example is still using an outdated Cohere model slug. Update the sample in the openRouterRerank/createOpenRouterRerank snippet to the latest Cohere rerank model listed in model-meta.ts, and make sure the explicit-key example in the explanatory paragraph uses the same current slug.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/ai-openrouter/tests/rerank-adapter.test.ts`:
- Around line 1-134: This test suite is in the wrong location: the
`OpenRouterRerankAdapter` tests should be colocated with the adapter source
rather than kept under a package-level `tests` folder. Move the
`rerank-adapter.test.ts` suite next to `createOpenRouterRerank` in
`src/adapters/rerank.ts`, keeping the same Vitest coverage and imports, so the
`*.test.ts` convention is followed.
---
Duplicate comments:
In `@docs/rerank/rerank.md`:
- Around line 189-193: The server example in the rerank docs still uses a type
assertion on body, which should be removed to match the docs rule. Update the
code sample around the request handling in the rerank server example to narrow
body with a guard or schema instead of using an as cast, and keep the
destructuring of query, documents, and topN inside the validated path.
---
Nitpick comments:
In `@docs/adapters/openrouter.md`:
- Around line 263-275: The OpenRouter rerank example is still using an outdated
Cohere model slug. Update the sample in the
openRouterRerank/createOpenRouterRerank snippet to the latest Cohere rerank
model listed in model-meta.ts, and make sure the explicit-key example in the
explanatory paragraph uses the same current slug.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b3cac01f-540f-4e5a-9dbe-198bb4005e74
📒 Files selected for processing (8)
.changeset/openrouter-rerank.mddocs/adapters/openrouter.mddocs/config.jsondocs/rerank/rerank.mdpackages/ai-openrouter/src/adapters/rerank.tspackages/ai-openrouter/src/index.tspackages/ai-openrouter/src/rerank/rerank-provider-options.tspackages/ai-openrouter/tests/rerank-adapter.test.ts
✅ Files skipped from review due to trivial changes (2)
- .changeset/openrouter-rerank.md
- packages/ai-openrouter/src/rerank/rerank-provider-options.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/config.json
Comment on lines
+1
to
+134
| import { beforeEach, describe, expect, it, vi } from 'vitest' | ||
| import { rerank } from '@tanstack/ai' | ||
| import { createOpenRouterRerank } from '../src/adapters/rerank' | ||
|
|
||
| // Mock the OpenRouter SDK so the adapter's `new OpenRouter().rerank.rerank()` | ||
| // call resolves to a controlled response. `vi.hoisted` lets the (hoisted) | ||
| // vi.mock factory reference the spy. | ||
| const { rerankFn } = vi.hoisted(() => ({ rerankFn: vi.fn() })) | ||
|
|
||
| vi.mock('@openrouter/sdk', () => ({ | ||
| // A class so `new OpenRouter(config)` is constructable; each instance exposes | ||
| // the spied `rerank.rerank`. | ||
| OpenRouter: class { | ||
| rerank = { rerank: rerankFn } | ||
| }, | ||
| })) | ||
|
|
||
| beforeEach(() => { | ||
| rerankFn.mockReset() | ||
| }) | ||
|
|
||
| const documents = ['sunny day at the beach', 'rainy afternoon in the city'] | ||
|
|
||
| /** A parsed SDK rerank response (camelCase, as the SDK returns it). */ | ||
| function sdkResponse() { | ||
| return { | ||
| id: 'or-1', | ||
| model: 'cohere/rerank-v3.5', | ||
| results: [ | ||
| { document: { text: documents[1] }, index: 1, relevanceScore: 0.97 }, | ||
| { document: { text: documents[0] }, index: 0, relevanceScore: 0.1 }, | ||
| ], | ||
| usage: { searchUnits: 1, cost: 0.002, totalTokens: 20 }, | ||
| } | ||
| } | ||
|
|
||
| const adapter = () => createOpenRouterRerank('cohere/rerank-v3.5', 'sk-or-test') | ||
|
|
||
| describe('OpenRouterRerankAdapter', () => { | ||
| it('calls the SDK rerank with the request body and maps the response', async () => { | ||
| rerankFn.mockResolvedValue(sdkResponse()) | ||
|
|
||
| const result = await rerank({ | ||
| adapter: adapter(), | ||
| query: 'talk about rain', | ||
| documents, | ||
| topN: 2, | ||
| }) | ||
|
|
||
| expect(rerankFn).toHaveBeenCalledTimes(1) | ||
| expect(rerankFn.mock.calls[0]![0]).toEqual({ | ||
| requestBody: { | ||
| model: 'cohere/rerank-v3.5', | ||
| query: 'talk about rain', | ||
| documents, | ||
| topN: 2, | ||
| }, | ||
| }) | ||
| expect(result.id).toBe('or-1') | ||
| expect(result.ranking).toEqual([ | ||
| { index: 1, score: 0.97, document: documents[1] }, | ||
| { index: 0, score: 0.1, document: documents[0] }, | ||
| ]) | ||
| }) | ||
|
|
||
| it('maps SDK usage (searchUnits/cost/totalTokens)', async () => { | ||
| rerankFn.mockResolvedValue(sdkResponse()) | ||
|
|
||
| const result = await rerank({ adapter: adapter(), query: 'q', documents }) | ||
|
|
||
| expect(result.usage.unitsBilled).toBe(1) | ||
| expect(result.usage.cost).toBe(0.002) | ||
| expect(result.usage.totalTokens).toBe(20) | ||
| }) | ||
|
|
||
| it('works with a non-Cohere model slug', async () => { | ||
| rerankFn.mockResolvedValue({ | ||
| ...sdkResponse(), | ||
| model: 'nvidia/llama-nemotron-rerank-vl-1b-v2', | ||
| }) | ||
|
|
||
| await rerank({ | ||
| adapter: createOpenRouterRerank( | ||
| 'nvidia/llama-nemotron-rerank-vl-1b-v2', | ||
| 'sk-or-test', | ||
| ), | ||
| query: 'q', | ||
| documents, | ||
| }) | ||
|
|
||
| expect(rerankFn.mock.calls[0]![0].requestBody.model).toBe( | ||
| 'nvidia/llama-nemotron-rerank-vl-1b-v2', | ||
| ) | ||
| }) | ||
|
|
||
| it('forwards provider routing preferences into the request body', async () => { | ||
| rerankFn.mockResolvedValue(sdkResponse()) | ||
|
|
||
| await rerank({ | ||
| adapter: adapter(), | ||
| query: 'q', | ||
| documents, | ||
| modelOptions: { provider: { order: ['cohere'] } }, | ||
| }) | ||
|
|
||
| expect(rerankFn.mock.calls[0]![0].requestBody.provider).toEqual({ | ||
| order: ['cohere'], | ||
| }) | ||
| }) | ||
|
|
||
| it('forwards the abort signal via fetchOptions', async () => { | ||
| rerankFn.mockResolvedValue(sdkResponse()) | ||
| const controller = new AbortController() | ||
|
|
||
| await rerank({ | ||
| adapter: adapter(), | ||
| query: 'q', | ||
| documents, | ||
| abortSignal: controller.signal, | ||
| }) | ||
|
|
||
| expect(rerankFn.mock.calls[0]![1]).toEqual({ | ||
| fetchOptions: { signal: controller.signal }, | ||
| }) | ||
| }) | ||
|
|
||
| it('throws when the SDK returns a bare string response', async () => { | ||
| rerankFn.mockResolvedValue('error: bad request') | ||
|
|
||
| await expect( | ||
| rerank({ adapter: adapter(), query: 'q', documents, debug: false }), | ||
| ).rejects.toThrow('unexpected response') | ||
| }) | ||
| }) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Move this test beside the adapter it exercises.
This new suite lives under packages/ai-openrouter/tests/, but the repo rule for *.test.ts files is to colocate unit tests with their source. Please move it next to src/adapters/rerank.ts instead of keeping a separate package-level test folder. As per coding guidelines, "**/*.test.ts: Place unit tests alongside source code in *.test.ts files`".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/ai-openrouter/tests/rerank-adapter.test.ts` around lines 1 - 134,
This test suite is in the wrong location: the `OpenRouterRerankAdapter` tests
should be colocated with the adapter source rather than kept under a
package-level `tests` folder. Move the `rerank-adapter.test.ts` suite next to
`createOpenRouterRerank` in `src/adapters/rerank.ts`, keeping the same Vitest
coverage and imports, so the `*.test.ts` convention is followed.
Source: Coding guidelines
Replace the @openrouter/sdk module mock with a network-layer fetch intercept driving the real SDK. Module mocking was sensitive to the CI runner's test isolation and failed there while passing locally; intercepting fetch exercises the SDK's real request building and zod response parsing and is environment- independent.
- rerank() activity: guard provider-returned indices and remove the as-cast, throwing a clear error on an out-of-range index instead of yielding undefined. - BaseRerankAdapter.generateId: use slice(2, 9) so the fallback id is never empty. - ai-cohere env detection: fall back to process.env when window.env is absent (bundler/Electron browser builds). - docs: replace the `as` cast in the rerank server example with a runtime guard.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/ai/tests/rerank.test.ts`:
- Around line 231-237: The rerank bounds test currently uses an obviously
invalid ranking index, which does not protect the exact upper-bound edge. Update
the test in rerank.test.ts near the rerank() rejection assertion to use the
first invalid index for the two-document input, so it specifically validates
that the rerank validation rejects index equal to documents.length and catches
off-by-one regressions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a90eada4-643a-4f65-817d-15cb5afea3ad
📒 Files selected for processing (5)
docs/rerank/rerank.mdpackages/ai-cohere/src/utils/client.tspackages/ai/src/activities/rerank/adapter.tspackages/ai/src/activities/rerank/index.tspackages/ai/tests/rerank.test.ts
✅ Files skipped from review due to trivial changes (1)
- docs/rerank/rerank.md
🚧 Files skipped from review as they are similar to previous changes (3)
- packages/ai-cohere/src/utils/client.ts
- packages/ai/src/activities/rerank/adapter.ts
- packages/ai/src/activities/rerank/index.ts
Comment on lines
+231
to
+237
| ranking: [{ index: 5, score: 0.9 }], | ||
| usage: { ...zeroUsage }, | ||
| })) | ||
|
|
||
| await expect( | ||
| rerank({ adapter, query: 'q', documents: ['a', 'b'] }), | ||
| ).rejects.toThrow('out-of-range') |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Use the first invalid index here to catch off-by-one regressions.
index: 5 proves the guard rejects obviously bad values, but it would still pass if the implementation accidentally allowed index === documents.length. Using index: 2 for ['a', 'b'] exercises the exact upper-bound edge this test is meant to protect.
Suggested change
- ranking: [{ index: 5, score: 0.9 }],
+ ranking: [{ index: 2, score: 0.9 }],📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| ranking: [{ index: 5, score: 0.9 }], | |
| usage: { ...zeroUsage }, | |
| })) | |
| await expect( | |
| rerank({ adapter, query: 'q', documents: ['a', 'b'] }), | |
| ).rejects.toThrow('out-of-range') | |
| ranking: [{ index: 2, score: 0.9 }], | |
| usage: { ...zeroUsage }, | |
| })) | |
| await expect( | |
| rerank({ adapter, query: 'q', documents: ['a', 'b'] }), | |
| ).rejects.toThrow('out-of-range') |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/ai/tests/rerank.test.ts` around lines 231 - 237, The rerank bounds
test currently uses an obviously invalid ranking index, which does not protect
the exact upper-bound edge. Update the test in rerank.test.ts near the rerank()
rejection assertion to use the first invalid index for the two-document input,
so it specifically validates that the rerank validation rejects index equal to
documents.length and catches off-by-one regressions.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
A provider-agnostic
rerank()activity in@tanstack/aifor reordering documents by relevance to a query, plus rerank adapters in two provider packages:@tanstack/ai-cohere(new package) —cohereRerank/createCohereRerank, Cohere/v2/rerankover rawfetch.@tanstack/ai-openrouter—openRouterRerank/createOpenRouterRerank, via the@openrouter/sdkrerank method. Supports every OpenRouter rerank model (Coherererank-v3.5/4-fast/4-pro, NVIDIAllama-nemotron-rerank, and any future slug).Highlights
topN,abortSignalcancellation, observe-onlyGenerationMiddleware(start/usage/finish/abort/error),rerank:*devtools events. Bills in search units viausage.unitsBilled(OpenRouter also surfacescost).fetch; OpenRouter uses the SDK so non-Cohere models route correctly, with optionalproviderrouting viamodelOptions.Tests & docs
Unit tests for the activity, Cohere adapter (mocked fetch), and OpenRouter adapter (mocked SDK). New Reranking guide + Cohere adapter page; OpenRouter adapter page gains a Reranking section. E2E intentionally skipped in favor of unit tests.
Changesets: minor for
@tanstack/ai,@tanstack/ai-event-client,@tanstack/ai-cohere,@tanstack/ai-openrouter.Summary by CodeRabbit
rerank()activity withtopNandabortSignalcancellation.rerank:*devtools events (start/complete/usage).