From aaf00baec35aa94073582301a9e1412cef35f109 Mon Sep 17 00:00:00 2001 From: phrog Date: Mon, 9 Nov 2026 22:29:16 -0400 Subject: [PATCH 1/2] feat(ai-bedrock): add support for Bedrock reranking --- .changeset/bedrock-rerank.md | 5 + docs/adapters/bedrock.md | 82 ++- docs/config.json | 4 +- docs/rerank/rerank.md | 18 + packages/ai-bedrock/README.md | 23 +- packages/ai-bedrock/package.json | 2 +- packages/ai-bedrock/src/adapters/rerank.ts | 188 +++++++ packages/ai-bedrock/src/index.ts | 15 +- packages/ai-bedrock/src/model-meta.ts | 9 + packages/ai-bedrock/src/utils/auth.ts | 8 + .../ai-bedrock/tests/rerank-adapter.test.ts | 474 ++++++++++++++++++ .../ai-bedrock/tests/rerank-types.test-d.ts | 60 +++ testing/e2e/global-setup.ts | 86 ++++ testing/e2e/playwright.config.ts | 9 + testing/e2e/src/routeTree.gen.ts | 21 + testing/e2e/src/routes/api.bedrock-rerank.ts | 43 ++ testing/e2e/tests/bedrock-rerank.spec.ts | 39 ++ 17 files changed, 1074 insertions(+), 12 deletions(-) create mode 100644 .changeset/bedrock-rerank.md create mode 100644 packages/ai-bedrock/src/adapters/rerank.ts create mode 100644 packages/ai-bedrock/tests/rerank-adapter.test.ts create mode 100644 packages/ai-bedrock/tests/rerank-types.test-d.ts create mode 100644 testing/e2e/src/routes/api.bedrock-rerank.ts create mode 100644 testing/e2e/tests/bedrock-rerank.spec.ts diff --git a/.changeset/bedrock-rerank.md b/.changeset/bedrock-rerank.md new file mode 100644 index 0000000000..125d12dcf0 --- /dev/null +++ b/.changeset/bedrock-rerank.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-bedrock': minor +--- + +Add `bedrockRerank()` for Amazon Rerank and Cohere Rerank models through the Bedrock Rerank API. Supports the existing `rerank()` activity with pagination and cancellation. Uses AWS credentials (SigV4), not Bedrock API keys. diff --git a/docs/adapters/bedrock.md b/docs/adapters/bedrock.md index a86f4164ae..74a6952ddd 100644 --- a/docs/adapters/bedrock.md +++ b/docs/adapters/bedrock.md @@ -19,7 +19,7 @@ keywords: - adapter --- -The Bedrock adapter calls Claude, Nova, Llama, Mistral, DeepSeek, and gpt-oss in your AWS account. It supports streaming, tools, reasoning, and embeddings. +The Bedrock adapter calls Claude, Nova, Llama, Mistral, DeepSeek, and gpt-oss in your AWS account. It supports streaming, tools, reasoning, embeddings, and reranking. ## 1. Install @@ -65,6 +65,8 @@ If a package is missing, the first request fails with an error that names the pa Bedrock takes an API key or your AWS credentials. Pick one. +For reranking, use AWS credentials (SigV4). Bedrock API keys are not supported. + **API key.** Create a key in the AWS Console. See the [Bedrock API keys guide](https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys.html). Then set one of these variables: ```bash @@ -81,7 +83,7 @@ AWS_SECRET_ACCESS_KEY=... AWS_SESSION_TOKEN=... # optional, for temporary credentials ``` -With `auth: 'auto'` (the default), the adapter uses the first of these that it finds: +For chat and embeddings, `auth: 'auto'` (the default) uses the first of these that it finds: 1. An explicit `apiKey` passed to the factory. 2. The `BEDROCK_API_KEY` environment variable. @@ -220,6 +222,78 @@ const result = await embed({ See the [Embeddings guide](../embeddings.md). +## Reranking + +Call `rerank()` with `bedrockRerank()` to rank documents in your AWS account. + +1. Install the rerank packages: + + + +react: @tanstack/ai @tanstack/ai-bedrock @aws-sdk/credential-providers +vue: @tanstack/ai @tanstack/ai-bedrock @aws-sdk/credential-providers +solid: @tanstack/ai @tanstack/ai-bedrock @aws-sdk/credential-providers +svelte: @tanstack/ai @tanstack/ai-bedrock @aws-sdk/credential-providers +preact: @tanstack/ai @tanstack/ai-bedrock @aws-sdk/credential-providers +angular: @tanstack/ai @tanstack/ai-bedrock @aws-sdk/credential-providers +vanilla: @tanstack/ai @tanstack/ai-bedrock @aws-sdk/credential-providers +octane: @tanstack/ai @tanstack/ai-bedrock @aws-sdk/credential-providers + + + +2. Configure [AWS credentials](#2-set-the-credentials) on the server. The adapter uses the AWS credential chain, including profiles and IAM roles. Bedrock API keys cannot authenticate Rerank requests. + +Your role needs model access and these [permissions](https://docs.aws.amazon.com/bedrock/latest/userguide/rerank-prereq.html): + +- `bedrock:Rerank` +- `bedrock:InvokeModel` for the selected rerank model + +3. Call `rerank()` on the server: + +```typescript +import { rerank } from '@tanstack/ai' +import { bedrockRerank } from '@tanstack/ai-bedrock' + +const { ranking, rerankedDocuments } = await rerank({ + adapter: bedrockRerank('cohere.rerank-v3-5:0'), + 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 +``` + +| Model | ID | Example region | +| --- | --- | --- | +| Cohere Rerank 3.5 | `cohere.rerank-v3-5:0` | `us-east-1` (default) | +| Amazon Rerank 1.0 | `amazon.rerank-v1:0` | `us-west-2` | + +Amazon Rerank is not available in `us-east-1`. Select a [supported region](https://docs.aws.amazon.com/bedrock/latest/userguide/rerank-supported.html): + +```typescript +import { bedrockRerank } from '@tanstack/ai-bedrock' + +const adapter = bedrockRerank('amazon.rerank-v1:0', { region: 'us-west-2' }) +``` + +Pass these options in `bedrockRerank(model, config)` (`BedrockRerankConfig`): + +| Option | Description | +| --- | --- | +| `region` | AWS region for both the model and request signing. Defaults to `us-east-1`. | +| `baseURL` | Overrides the agent-runtime URL, without the `/rerank` suffix. | +| `defaultHeaders` | Headers to add before SigV4 signing. | +| `fetch` | Custom fetch transport. The adapter still signs each request. | + +- Pass 1 to 1,000 documents. `topN` must be an integer from 1 to 1,000. +- Without `topN`, the adapter requests all documents and collects any result pages. +- Pass model-specific inference fields through `modelOptions.additionalModelRequestFields`. +- AWS reports no usage for this API. Token counts are zero, and `usage.billed` and `usage.cost` are absent. Zero counts do not mean the request is free. + +The result keeps your original documents in relevance order. For object documents and cancellation, see [Reranking](../rerank/rerank). + ## Models and options ### Choose the API @@ -277,6 +351,7 @@ The package ships a catalog of confirmed model ids (`src/model-catalog.generated | Meta Llama | `us.meta.llama4-maverick-17b-instruct-v1:0`, `us.meta.llama3-3-70b-instruct-v1:0` | | Other chat | `openai.gpt-oss-120b-1:0`, `openai.gpt-oss-20b-1:0`, `us.mistral.pixtral-large-2502-v1:0`, `us.deepseek.r1-v1:0`, `google.gemma-4-31b`, `google.gemma-4-26b-a4b`, `google.gemma-4-e2b` | | Embeddings | `amazon.titan-embed-text-v2:0`, `amazon.titan-embed-image-v1`, `cohere.embed-english-v3`, `cohere.embed-multilingual-v3` | +| Reranking | `amazon.rerank-v1:0`, `cohere.rerank-v3-5:0` | `adapter.inputModalities` lists the input kinds of a model in the catalog. If this package does not know the model id, `bedrockText` and `createBedrockText` throw. `createBedrockChat`, `createBedrockResponsesText`, and `createBedrockConverse` return `undefined` for that id. See [Check what a model accepts](../advanced/extend-adapter#check-what-a-model-accepts). @@ -457,8 +532,9 @@ for await (const chunk of chat({ | `bedrockText(model, config?)` | Chat adapter with auth from the environment. Returns the adapter of the `api` that you pick | | `createBedrockText(model, apiKey, config?)` | Chat adapter with an explicit API key. Skips the environment lookup | | `bedrockEmbedding(model, config?)` / `createBedrockEmbedding(model, apiKey, config?)` | Embedding adapter for Titan and Cohere models | +| `bedrockRerank(model, config?)` | Rerank adapter for Amazon and Cohere models. Uses AWS credentials only | -`config` takes `api`, `region`, `auth`, `apiKey`, `baseURL`, `defaultHeaders`, and `endpoint`. See [Configuration](#configuration). +For chat, `config` takes `api`, `region`, `auth`, `apiKey`, `baseURL`, `defaultHeaders`, and `endpoint`. See [Configuration](#configuration). More from AWS: diff --git a/docs/config.json b/docs/config.json index cab1bb46ce..daa5bf87be 100644 --- a/docs/config.json +++ b/docs/config.json @@ -994,7 +994,7 @@ "label": "Reranking", "to": "rerank/rerank", "addedAt": "2026-06-25", - "updatedAt": "2026-11-02" + "updatedAt": "2026-11-09" } ], "tab": "guides" @@ -1908,7 +1908,7 @@ "label": "Amazon Bedrock", "to": "adapters/bedrock", "addedAt": "2026-06-25", - "updatedAt": "2026-11-02" + "updatedAt": "2026-11-09" }, { "label": "BytePlus", diff --git a/docs/rerank/rerank.md b/docs/rerank/rerank.md index eafd9b2912..808a356a4d 100644 --- a/docs/rerank/rerank.md +++ b/docs/rerank/rerank.md @@ -185,6 +185,24 @@ console.log(rerankedDocuments[0]) // 'rainy afternoon in the city' `openRouterRerank` reads `OPENROUTER_API_KEY` from the environment. +## Use Amazon Bedrock + +Use `bedrockRerank()` with AWS credentials, not a Bedrock API key. See the [Bedrock setup](../adapters/bedrock#reranking) for installation and permissions. + +```typescript +import { rerank } from '@tanstack/ai' +import { bedrockRerank } from '@tanstack/ai-bedrock' + +const { rerankedDocuments } = await rerank({ + adapter: bedrockRerank('cohere.rerank-v3-5:0'), + 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' +``` + ## Try the example `examples/ts-react-rerank` is a small TanStack Start app. It shows a list of support articles, a query box, and the original order next to the reranked order with scores. A dropdown switches between the Cohere and OpenRouter adapters. diff --git a/packages/ai-bedrock/README.md b/packages/ai-bedrock/README.md index 43ed33e9ec..75a05c34f7 100644 --- a/packages/ai-bedrock/README.md +++ b/packages/ai-bedrock/README.md @@ -85,16 +85,35 @@ const adapter = createBedrockText( ) ``` +## Reranking + +Call `rerank()` with `bedrockRerank()` on the server to sort documents by relevance. + +First, install `@aws-sdk/credential-providers` and configure AWS credentials. Bedrock API keys do not work for reranking. See [Bedrock reranking](https://tanstack.com/ai/latest/docs/adapters/bedrock#reranking) for permissions and model configuration. + +```typescript +import { rerank } from '@tanstack/ai' +import { bedrockRerank } from '@tanstack/ai-bedrock' + +const { rerankedDocuments } = await rerank({ + adapter: bedrockRerank('cohere.rerank-v3-5:0'), + query: 'talk about rain', + documents: ['sunny day at the beach', 'rainy afternoon in the city'], + topN: 2, +}) +``` + ## Authentication -Auth is resolved in this order: +For chat and embeddings, auth is resolved in this order: 1. Explicit `apiKey` passed to the factory 2. `BEDROCK_API_KEY` environment variable 3. `AWS_BEARER_TOKEN_BEDROCK` environment variable 4. SigV4 via the AWS credential chain -SigV4 signing is built in — no additional packages required. On the Converse +SigV4 signing is built in. The AWS credential chain requires +`@aws-sdk/credential-providers`. On the Converse (default) path the `@aws-sdk/client-bedrock-runtime` client signs requests itself; the OpenAI-compatible chat/responses paths sign via the bundled `@smithy/signature-v4`. Set `auth: 'sigv4'` and provide AWS credentials through diff --git a/packages/ai-bedrock/package.json b/packages/ai-bedrock/package.json index 04e28723b2..6924b80660 100644 --- a/packages/ai-bedrock/package.json +++ b/packages/ai-bedrock/package.json @@ -2,7 +2,7 @@ "name": "@tanstack/ai-bedrock", "version": "1.2.1", "type": "module", - "description": "Amazon Bedrock adapter for TanStack AI — OpenAI-compatible chat, responses, tools, and reasoning.", + "description": "Amazon Bedrock adapter for TanStack AI with chat, tools, reasoning, embeddings, and reranking.", "author": "", "license": "MIT", "repository": { diff --git a/packages/ai-bedrock/src/adapters/rerank.ts b/packages/ai-bedrock/src/adapters/rerank.ts new file mode 100644 index 0000000000..87bc72bc20 --- /dev/null +++ b/packages/ai-bedrock/src/adapters/rerank.ts @@ -0,0 +1,188 @@ +import { BaseRerankAdapter } from '@tanstack/ai/adapters' +import { toRunErrorPayload } from '@tanstack/ai/adapter-internals' +import { resolveBedrockAuth } from '../utils/auth' +import { createSigV4Fetch } from '../utils/openai-sigv4-fetch' +import type { RerankAdapterResult, RerankOptions } from '@tanstack/ai' +import type { BedrockRerankModel } from '../model-meta' + +/** Bedrock rerank uses AWS credentials, not Bedrock API keys. */ +export interface BedrockRerankConfig { + /** AWS region; defaults to us-east-1. Amazon Rerank requires another region. */ + region?: string + /** Overrides the agent-runtime base URL, without the /rerank suffix. */ + baseURL?: string + /** Headers added before SigV4 signing. */ + defaultHeaders?: Record + /** Custom transport; requests are still SigV4-signed. */ + fetch?: typeof fetch +} + +/** Additional inference fields for the selected Bedrock rerank model. */ +export interface BedrockRerankProviderOptions { + additionalModelRequestFields?: Record +} + +interface BedrockRerankResponse { + results: Array<{ index: number; relevanceScore: number }> + nextToken?: string +} + +function isBedrockRerankResponse( + value: unknown, +): value is BedrockRerankResponse { + return ( + typeof value === 'object' && + value !== null && + 'results' in value && + Array.isArray(value.results) && + value.results.every( + (result) => + typeof result === 'object' && + result !== null && + Number.isInteger(result.index) && + Number.isFinite(result.relevanceScore), + ) && + (!('nextToken' in value) || + (typeof value.nextToken === 'string' && + value.nextToken.length > 0 && + value.nextToken.length <= 2048 && + !/\s/.test(value.nextToken))) + ) +} + +/** Bedrock rerank adapter; core maps scored indices to original documents. */ +export class BedrockRerankAdapter< + TModel extends BedrockRerankModel, +> extends BaseRerankAdapter { + readonly name = 'bedrock' as const + private readonly region: string + private readonly baseURL: string + private readonly clientConfig: BedrockRerankConfig + + constructor(config: BedrockRerankConfig, model: TModel) { + super({}, model) + this.clientConfig = config + this.region = config.region ?? 'us-east-1' + this.baseURL = ( + config.baseURL ?? + `https://bedrock-agent-runtime.${this.region}.amazonaws.com` + ).replace(/\/+$/, '') + } + + async rerank( + options: RerankOptions, + ): Promise { + const { model, query, documents, topN, modelOptions, abortSignal, logger } = + options + + try { + if (documents.length < 1 || documents.length > 1000) { + throw new Error('Bedrock rerank requires between 1 and 1000 documents') + } + const numberOfResults = topN ?? documents.length + if ( + !Number.isInteger(numberOfResults) || + numberOfResults < 1 || + numberOfResults > 1000 + ) { + throw new Error( + 'Bedrock rerank topN must be an integer between 1 and 1000', + ) + } + + const body = { + queries: [{ type: 'TEXT', textQuery: { text: query } }], + sources: documents.map((text) => ({ + type: 'INLINE', + inlineDocumentSource: { type: 'TEXT', textDocument: { text } }, + })), + rerankingConfiguration: { + type: 'BEDROCK_RERANKING_MODEL', + bedrockRerankingConfiguration: { + modelConfiguration: { + modelArn: `arn:aws:bedrock:${this.region}::foundation-model/${model}`, + additionalModelRequestFields: + modelOptions?.additionalModelRequestFields, + }, + numberOfResults, + }, + }, + } + const signedFetch = createSigV4Fetch( + resolveBedrockAuth({ region: this.region, auth: 'sigv4' }, 'runtime'), + this.clientConfig.fetch, + ) + const ranking: RerankAdapterResult['ranking'] = [] + const seenTokens = new Set() + let nextToken: string | undefined + let id: string | undefined + + do { + abortSignal?.throwIfAborted() + logger.request( + `activity=rerank provider=${this.name} model=${model} documents=${documents.length}`, + { provider: this.name, model }, + ) + const response = await signedFetch(`${this.baseURL}/rerank`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + ...this.clientConfig.defaultHeaders, + }, + body: JSON.stringify({ ...body, nextToken }), + signal: abortSignal, + }) + if (!response.ok) { + const detail = await response.text() + throw new Error( + `Bedrock rerank request failed: ${response.status} ${response.statusText}${detail ? `: ${detail}` : ''}`, + ) + } + const page: unknown = await response.json() + abortSignal?.throwIfAborted() + if (!isBedrockRerankResponse(page)) { + throw new Error('Bedrock rerank response had an unexpected shape') + } + id ??= response.headers.get('x-amzn-requestid') ?? undefined + ranking.push( + ...page.results.map((result) => ({ + index: result.index, + score: result.relevanceScore, + })), + ) + nextToken = page.nextToken + if (nextToken) { + if (seenTokens.has(nextToken) || page.results.length === 0) { + throw new Error('Bedrock rerank pagination did not advance') + } + seenTokens.add(nextToken) + } + } while (nextToken && ranking.length < numberOfResults) + + return { + id: id ?? this.generateId(), + ranking: ranking.slice(0, numberOfResults), + // The Bedrock Rerank API reports neither token counts nor billed units. + usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 }, + } + } catch (error) { + logger.errors(`${this.name}.rerank fatal`, { + error: toRunErrorPayload(error, 'Bedrock rerank failed'), + source: `${this.name}.rerank`, + }) + throw error + } + } +} + +/** + * Creates a Bedrock rerank adapter using the AWS credential chain (SigV4). + * Install `@aws-sdk/credential-providers`; Bedrock API keys are not supported. + * For `amazon.rerank-v1:0`, select a supported region such as `us-west-2`. + */ +export function bedrockRerank( + model: TModel, + config?: BedrockRerankConfig, +): BedrockRerankAdapter { + return new BedrockRerankAdapter(config ?? {}, model) +} diff --git a/packages/ai-bedrock/src/index.ts b/packages/ai-bedrock/src/index.ts index bd6a9aedbe..f7e5076ee5 100644 --- a/packages/ai-bedrock/src/index.ts +++ b/packages/ai-bedrock/src/index.ts @@ -1,10 +1,9 @@ /** * @module @tanstack/ai-bedrock * - * Amazon Bedrock adapter for TanStack AI via Bedrock's OpenAI-compatible APIs - * and the native Converse API. The public `bedrockText` / `createBedrockText` - * factory branches between the Converse adapter (DEFAULT), the Chat Completions - * adapter (`api: 'chat'`), and the Responses adapter (`api: 'responses'`). + * Bedrock chat uses Converse by default, or the OpenAI-compatible APIs with + * `api: 'chat'` or `api: 'responses'`. Embeddings and reranking use + * `bedrockEmbedding` and `bedrockRerank`. */ import { BedrockTextAdapter } from './adapters/text' @@ -188,6 +187,12 @@ export { createBedrockEmbedding, type BedrockEmbeddingConfig, } from './adapters/embedding' +export { + BedrockRerankAdapter, + bedrockRerank, + type BedrockRerankConfig, + type BedrockRerankProviderOptions, +} from './adapters/rerank' export type { BedrockCohereEmbeddingInputType, BedrockCohereEmbeddingProviderOptions, @@ -208,6 +213,8 @@ export { BEDROCK_RESPONSES_MODELS, BEDROCK_CONVERSE_MODELS, BEDROCK_EMBEDDING_MODELS, + BEDROCK_RERANK_MODELS, + type BedrockRerankModel, type BedrockEmbeddingModel, type BedrockEmbeddingModelProviderOptionsByName, type BedrockEmbeddingModelInputModalitiesByName, diff --git a/packages/ai-bedrock/src/model-meta.ts b/packages/ai-bedrock/src/model-meta.ts index 8fffaed654..12c6fb512a 100644 --- a/packages/ai-bedrock/src/model-meta.ts +++ b/packages/ai-bedrock/src/model-meta.ts @@ -150,6 +150,15 @@ export const BEDROCK_CLAUDE_NO_FORCED_TOOL_FAMILIES: ReadonlyArray = [ 'claude-sonnet-5-5', ] +/** Models supported by the Bedrock Rerank API, separate from the chat catalog. */ +export const BEDROCK_RERANK_MODELS = [ + 'amazon.rerank-v1:0', + 'cohere.rerank-v3-5:0', +] as const + +/** Short model IDs accepted by the Bedrock rerank adapter. */ +export type BedrockRerankModel = (typeof BEDROCK_RERANK_MODELS)[number] + // ============================================================================ // Embedding models // ============================================================================ diff --git a/packages/ai-bedrock/src/utils/auth.ts b/packages/ai-bedrock/src/utils/auth.ts index c6f1b6cdd8..269a898497 100644 --- a/packages/ai-bedrock/src/utils/auth.ts +++ b/packages/ai-bedrock/src/utils/auth.ts @@ -57,6 +57,14 @@ export interface BedrockAuthConfig { } /** apiKey -> BEDROCK_API_KEY -> AWS_BEARER_TOKEN_BEDROCK -> SigV4 (credential chain). */ +export function resolveBedrockAuth( + config: BedrockAuthConfig & { auth: 'sigv4' }, + endpoint: BedrockEndpoint, +): Extract +export function resolveBedrockAuth( + config: BedrockAuthConfig, + endpoint: BedrockEndpoint, +): ResolvedBedrockAuth export function resolveBedrockAuth( config: BedrockAuthConfig, endpoint: BedrockEndpoint, diff --git a/packages/ai-bedrock/tests/rerank-adapter.test.ts b/packages/ai-bedrock/tests/rerank-adapter.test.ts new file mode 100644 index 0000000000..866f88e03c --- /dev/null +++ b/packages/ai-bedrock/tests/rerank-adapter.test.ts @@ -0,0 +1,474 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { rerank } from '@tanstack/ai' +import { BEDROCK_RERANK_MODELS, bedrockRerank } from '../src' + +const { credentials, providerChain } = vi.hoisted(() => { + const resolveCredentials = vi.fn() + return { + credentials: resolveCredentials, + providerChain: vi.fn(() => resolveCredentials), + } +}) +vi.mock('@aws-sdk/credential-providers', () => ({ + fromNodeProviderChain: providerChain, +})) + +const fetchMock = vi.fn() +const documents = ['sunny day at the beach', 'rainy afternoon in the city'] +const results = [ + { index: 1, relevanceScore: 0.98 }, + { index: 0, relevanceScore: 0.12 }, +] + +function bedrockResponse( + body: unknown = { results }, + requestId = 'aws-request-1', +) { + return new Response(JSON.stringify(body), { + headers: { + 'Content-Type': 'application/json', + ...(requestId ? { 'x-amzn-requestid': requestId } : {}), + }, + }) +} + +function requestBody(index = 0): unknown { + const call = fetchMock.mock.calls[index] + if (!call) throw new Error(`Missing rerank request ${index}`) + return JSON.parse(String(call[1]?.body)) +} + +function rerankOptions() { + return { + adapter: bedrockRerank('cohere.rerank-v3-5:0'), + query: 'talk about rain', + documents, + debug: false, + } +} + +beforeEach(() => { + vi.clearAllMocks() + fetchMock.mockReset().mockImplementation(async () => bedrockResponse()) + credentials.mockReset().mockResolvedValue({ + accessKeyId: 'AKIA_TEST', + secretAccessKey: 'test-secret', + sessionToken: 'test-session', + }) + vi.stubGlobal('fetch', fetchMock) +}) + +afterEach(() => { + vi.unstubAllGlobals() + vi.unstubAllEnvs() + vi.restoreAllMocks() +}) + +describe('BedrockRerankAdapter', () => { + it('constructs without resolving AWS credentials', () => { + const adapter = bedrockRerank('cohere.rerank-v3-5:0') + expect(adapter).toMatchObject({ + name: 'bedrock', + kind: 'rerank', + model: 'cohere.rerank-v3-5:0', + }) + expect(providerChain).not.toHaveBeenCalled() + expect(credentials).not.toHaveBeenCalled() + expect(fetchMock).not.toHaveBeenCalled() + }) + + it.each(BEDROCK_RERANK_MODELS)( + 'sends a signed rerank request for %s', + async (model) => { + await rerank({ + ...rerankOptions(), + adapter: bedrockRerank(model, { region: 'us-west-2' }), + topN: 1, + modelOptions: { + additionalModelRequestFields: { max_tokens_per_doc: 512 }, + }, + }) + + expect(fetchMock).toHaveBeenCalledTimes(1) + const [url, init] = fetchMock.mock.calls[0] ?? [] + expect(url).toBe( + 'https://bedrock-agent-runtime.us-west-2.amazonaws.com/rerank', + ) + expect(init?.method).toBe('POST') + const headers = new Headers(init?.headers) + expect(headers.get('authorization')).toContain( + 'AWS4-HMAC-SHA256 Credential=AKIA_TEST/', + ) + expect(headers.get('authorization')).toContain( + '/us-west-2/bedrock/aws4_request', + ) + expect(headers.get('x-amz-security-token')).toBe('test-session') + expect(headers.get('content-type')).toBe('application/json') + expect(requestBody()).toEqual({ + queries: [{ type: 'TEXT', textQuery: { text: 'talk about rain' } }], + sources: [ + { + type: 'INLINE', + inlineDocumentSource: { + type: 'TEXT', + textDocument: { text: 'sunny day at the beach' }, + }, + }, + { + type: 'INLINE', + inlineDocumentSource: { + type: 'TEXT', + textDocument: { text: 'rainy afternoon in the city' }, + }, + }, + ], + rerankingConfiguration: { + type: 'BEDROCK_RERANKING_MODEL', + bedrockRerankingConfiguration: { + modelConfiguration: { + modelArn: `arn:aws:bedrock:us-west-2::foundation-model/${model}`, + additionalModelRequestFields: { max_tokens_per_doc: 512 }, + }, + numberOfResults: 1, + }, + }, + }) + }, + ) + + it('defaults to us-east-1 and explicitly requests every document', async () => { + await rerank(rerankOptions()) + expect(fetchMock.mock.calls[0]?.[0]).toBe( + 'https://bedrock-agent-runtime.us-east-1.amazonaws.com/rerank', + ) + expect(requestBody()).toMatchObject({ + rerankingConfiguration: { + bedrockRerankingConfiguration: { + modelConfiguration: { + modelArn: + 'arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0', + }, + numberOfResults: 2, + }, + }, + }) + expect(requestBody()).not.toHaveProperty( + 'rerankingConfiguration.bedrockRerankingConfiguration.modelConfiguration.additionalModelRequestFields', + ) + expect(requestBody()).not.toHaveProperty('nextToken') + }) + + it.each(['BEDROCK_API_KEY', 'AWS_BEARER_TOKEN_BEDROCK'])( + 'ignores %s and uses AWS credentials', + async (variable) => { + vi.stubEnv(variable, 'not-a-supported-rerank-key') + await rerank(rerankOptions()) + expect(credentials).toHaveBeenCalledOnce() + expect( + new Headers(fetchMock.mock.calls[0]?.[1]?.headers).get('authorization'), + ).toMatch(/^AWS4-HMAC-SHA256 /) + }, + ) + + it('signs custom headers and transport with a normalized endpoint override', async () => { + const transport = vi.fn().mockResolvedValue(bedrockResponse()) + await rerank({ + ...rerankOptions(), + adapter: bedrockRerank('cohere.rerank-v3-5:0', { + baseURL: 'https://proxy.example/bedrock///', + defaultHeaders: { + 'X-Gateway': 'test', + Authorization: 'must-be-replaced', + }, + fetch: transport, + }), + }) + expect(fetchMock).not.toHaveBeenCalled() + expect(transport.mock.calls[0]?.[0]).toBe( + 'https://proxy.example/bedrock/rerank', + ) + const headers = new Headers(transport.mock.calls[0]?.[1]?.headers) + expect(headers.get('x-gateway')).toBe('test') + expect(headers.get('authorization')).toMatch(/^AWS4-HMAC-SHA256 /) + expect(headers.get('authorization')).toContain('x-gateway') + }) + + it('maps scores and indices without claiming usage or mutating input', async () => { + const result = await rerank(rerankOptions()) + expect(result).toEqual({ + id: 'aws-request-1', + model: 'cohere.rerank-v3-5:0', + ranking: [ + { index: 1, score: 0.98, document: documents[1] }, + { index: 0, score: 0.12, document: documents[0] }, + ], + rerankedDocuments: [documents[1], documents[0]], + usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 }, + }) + expect(documents).toEqual([ + 'sunny day at the beach', + 'rainy afternoon in the city', + ]) + }) + + it('serializes object documents once and returns the originals', async () => { + const objects = [ + { id: 1, text: 'sun' }, + { id: 2, text: 'rain' }, + ] + const result = await rerank({ ...rerankOptions(), documents: objects }) + expect(requestBody()).toMatchObject({ + sources: [ + { + inlineDocumentSource: { + type: 'TEXT', + textDocument: { text: '{"id":1,"text":"sun"}' }, + }, + }, + { + inlineDocumentSource: { + type: 'TEXT', + textDocument: { text: '{"id":2,"text":"rain"}' }, + }, + }, + ], + }) + expect(result.ranking[0]?.document).toBe(objects[1]) + expect(result.rerankedDocuments[1]).toBe(objects[0]) + }) + + it('generates an id if AWS omits its request id', async () => { + fetchMock.mockResolvedValue(bedrockResponse({ results }, '')) + expect((await rerank(rerankOptions())).id).toMatch(/^bedrock-/) + }) + + it('collects pages in provider order using the same request and the first request id', async () => { + fetchMock + .mockResolvedValueOnce( + bedrockResponse({ results: [results[0]], nextToken: 'page-2' }), + ) + .mockResolvedValueOnce( + bedrockResponse({ results: [results[1]] }, 'aws-request-2'), + ) + const result = await rerank(rerankOptions()) + expect(result.id).toBe('aws-request-1') + expect(result.rerankedDocuments).toEqual([documents[1], documents[0]]) + const secondBody = requestBody(1) + if ( + typeof secondBody !== 'object' || + secondBody === null || + !('nextToken' in secondBody) + ) { + throw new Error('Missing continuation request body') + } + const { nextToken, ...initialBody } = secondBody + expect(initialBody).toEqual(requestBody()) + expect(nextToken).toBe('page-2') + expect(fetchMock).toHaveBeenCalledTimes(2) + }) + + it('stops at topN even when AWS returns a continuation token', async () => { + fetchMock.mockResolvedValue( + bedrockResponse({ results, nextToken: 'not-needed' }), + ) + const result = await rerank({ ...rerankOptions(), topN: 1 }) + expect(result.ranking).toHaveLength(1) + expect(fetchMock).toHaveBeenCalledTimes(1) + }) + + it('rejects repeated pagination tokens', async () => { + fetchMock + .mockResolvedValueOnce( + bedrockResponse({ results: [results[0]], nextToken: 'again' }), + ) + .mockResolvedValueOnce( + bedrockResponse({ results: [results[1]], nextToken: 'again' }), + ) + await expect(rerank({ ...rerankOptions(), topN: 3 })).rejects.toThrow( + 'pagination did not advance', + ) + expect(fetchMock).toHaveBeenCalledTimes(2) + }) + + it('rejects an empty continuation page', async () => { + fetchMock.mockResolvedValue( + bedrockResponse({ results: [], nextToken: 'again' }), + ) + await expect(rerank(rerankOptions())).rejects.toThrow( + 'pagination did not advance', + ) + expect(fetchMock).toHaveBeenCalledTimes(1) + }) + + it('accepts 1000 documents and rejects 1001 before resolving credentials', async () => { + await rerank({ + ...rerankOptions(), + documents: Array.from({ length: 1000 }, () => 'text'), + }) + expect(requestBody()).toMatchObject({ + rerankingConfiguration: { + bedrockRerankingConfiguration: { numberOfResults: 1000 }, + }, + }) + vi.clearAllMocks() + await expect( + rerank({ + ...rerankOptions(), + documents: Array.from({ length: 1001 }, () => 'text'), + }), + ).rejects.toThrow('between 1 and 1000 documents') + expect(credentials).not.toHaveBeenCalled() + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('lets core reject an empty document list', async () => { + await expect(rerank({ ...rerankOptions(), documents: [] })).rejects.toThrow( + 'at least one document', + ) + expect(fetchMock).not.toHaveBeenCalled() + }) + + it.each([0, -1, 1.5, 1001, NaN, Infinity])( + 'rejects invalid topN %s', + async (topN) => { + await expect(rerank({ ...rerankOptions(), topN })).rejects.toThrow( + 'topN must be an integer between 1 and 1000', + ) + expect(credentials).not.toHaveBeenCalled() + expect(fetchMock).not.toHaveBeenCalled() + }, + ) + + it.each([ + null, + {}, + { results: 'bad' }, + { results: [null] }, + { results: [{ index: 0.5, relevanceScore: 0.5 }] }, + { results: [{ index: '0', relevanceScore: 0.5 }] }, + { results: [{ index: 0, relevanceScore: '0.5' }] }, + { results: [{ index: 0, relevanceScore: null }] }, + { results: [{ index: 0 }] }, + { results, nextToken: null }, + { results, nextToken: '' }, + { results, nextToken: 'has spaces' }, + { results, nextToken: 'x'.repeat(2049) }, + ])('rejects malformed response %#', async (body) => { + fetchMock.mockResolvedValue(bedrockResponse(body)) + await expect(rerank(rerankOptions())).rejects.toThrow('unexpected shape') + }) + + it('rejects a non-finite JSON score', async () => { + fetchMock.mockResolvedValue( + new Response('{"results":[{"index":0,"relevanceScore":1e400}]}'), + ) + await expect(rerank(rerankOptions())).rejects.toThrow('unexpected shape') + }) + + it('propagates invalid JSON and lets core reject out-of-range indices', async () => { + fetchMock.mockResolvedValueOnce(new Response('not json')) + await expect(rerank(rerankOptions())).rejects.toBeInstanceOf(SyntaxError) + fetchMock.mockResolvedValueOnce( + bedrockResponse({ results: [{ index: 5, relevanceScore: 0.5 }] }), + ) + await expect(rerank(rerankOptions())).rejects.toThrow( + 'out-of-range index 5', + ) + }) + + it.each([403, 429, 500])( + 'reports HTTP %s and AWS error detail without retrying', + async (status) => { + fetchMock.mockResolvedValue( + new Response('{"message":"AWS failure"}', { status }), + ) + await expect(rerank(rerankOptions())).rejects.toThrow( + new RegExp(`Bedrock rerank request failed: ${status} .*AWS failure`), + ) + expect(fetchMock).toHaveBeenCalledOnce() + }, + ) + + it('propagates network and credential errors', async () => { + const networkError = new Error('connection failed') + fetchMock.mockRejectedValueOnce(networkError) + await expect(rerank(rerankOptions())).rejects.toBe(networkError) + fetchMock.mockClear() + const credentialError = new Error('credentials unavailable') + credentials.mockRejectedValueOnce(credentialError) + await expect(rerank(rerankOptions())).rejects.toBe(credentialError) + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('does not request anything for an already-aborted signal', async () => { + const controller = new AbortController() + controller.abort() + await expect( + rerank({ ...rerankOptions(), abortSignal: controller.signal }), + ).rejects.toMatchObject({ name: 'AbortError' }) + expect(credentials).not.toHaveBeenCalled() + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('forwards in-flight cancellation to fetch', async () => { + const controller = new AbortController() + const started = new Promise((resolve) => { + fetchMock.mockImplementation( + (_url, init) => + new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(init.signal?.reason), + { once: true }, + ) + resolve() + }), + ) + }) + const pending = rerank({ + ...rerankOptions(), + abortSignal: controller.signal, + }) + await started + expect(fetchMock.mock.calls[0]?.[1]?.signal).toBe(controller.signal) + controller.abort() + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + }) + + it('does not fetch the next page after cancellation', async () => { + const controller = new AbortController() + const response = bedrockResponse({ + results: [results[0]], + nextToken: 'page-2', + }) + const readJson = response.json.bind(response) + vi.spyOn(response, 'json').mockImplementation(async () => { + const page: unknown = await readJson() + controller.abort() + return page + }) + fetchMock.mockResolvedValue(response) + await expect( + rerank({ ...rerankOptions(), abortSignal: controller.signal }), + ).rejects.toMatchObject({ name: 'AbortError' }) + expect(fetchMock).toHaveBeenCalledOnce() + }) + + it('passes the core timeout signal to the pending request', async () => { + fetchMock.mockImplementation( + (_url, init) => + new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(init.signal?.reason), + { once: true }, + ) + }), + ) + await expect( + rerank({ ...rerankOptions(), timeout: 50 }), + ).rejects.toMatchObject({ name: 'TimeoutError' }) + expect(fetchMock).toHaveBeenCalledOnce() + expect(fetchMock.mock.calls[0]?.[1]?.signal?.aborted).toBe(true) + }) +}) diff --git a/packages/ai-bedrock/tests/rerank-types.test-d.ts b/packages/ai-bedrock/tests/rerank-types.test-d.ts new file mode 100644 index 0000000000..4671b2c315 --- /dev/null +++ b/packages/ai-bedrock/tests/rerank-types.test-d.ts @@ -0,0 +1,60 @@ +import { expectTypeOf } from 'vitest' +import { chat, rerank } from '@tanstack/ai' +import { bedrockRerank, bedrockText, resolveBedrockAuth } from '../src' +import type { + BedrockRerankAdapter, + BedrockRerankConfig, + BedrockRerankProviderOptions, +} from '../src' + +const config: BedrockRerankConfig = { region: 'us-west-2' } +const adapter = bedrockRerank('amazon.rerank-v1:0', config) +expectTypeOf(adapter).toEqualTypeOf< + BedrockRerankAdapter<'amazon.rerank-v1:0'> +>() +expectTypeOf( + adapter['~types'].providerOptions, +).toEqualTypeOf() + +const result = rerank({ + adapter, + query: 'rain', + documents: [{ id: 1, text: 'rain' }], + modelOptions: { additionalModelRequestFields: { max_tokens_per_doc: 512 } }, +}) +expectTypeOf['rerankedDocuments']>().toEqualTypeOf< + Array<{ id: number; text: string }> +>() + +// @ts-expect-error - chat model IDs cannot select a reranker +bedrockRerank('us.amazon.nova-pro-v1:0') +// @ts-expect-error - Bedrock rerank cannot use API keys +bedrockRerank('cohere.rerank-v3-5:0', { apiKey: 'key' }) +// @ts-expect-error - rerank always uses SigV4 +bedrockRerank('cohere.rerank-v3-5:0', { auth: 'apikey' }) + +rerank({ + adapter, + query: 'rain', + documents: ['rain'], + // @ts-expect-error - pagination is handled internally + modelOptions: { nextToken: 'token' }, +}) + +rerank({ + // @ts-expect-error - a chat adapter cannot rerank documents + adapter: bedrockText('us.amazon.nova-pro-v1:0'), + query: 'rain', + documents: ['rain'], +}) + +chat({ + threadId: 'rerank-type-test', + // @ts-expect-error - a rerank adapter cannot chat + adapter, + messages: [{ role: 'user', content: 'hello' }], +}) + +const auth = resolveBedrockAuth({ auth: 'sigv4' }, 'runtime') +expectTypeOf(auth.kind).toEqualTypeOf<'sigv4'>() +expectTypeOf(auth.credentials).toBeFunction() diff --git a/testing/e2e/global-setup.ts b/testing/e2e/global-setup.ts index 99f73170bf..e293991d79 100644 --- a/testing/e2e/global-setup.ts +++ b/testing/e2e/global-setup.ts @@ -3,6 +3,7 @@ import fs from 'fs' import http from 'http' import path from 'path' import { fileURLToPath } from 'url' +import { isDeepStrictEqual } from 'node:util' import { EventStreamCodec } from '@smithy/eventstream-codec' import type { Mountable } from '@copilotkit/aimock' @@ -178,6 +179,7 @@ export default async function globalSetup() { // which request sections ended with a `cachePoint` block and answers with // Bedrock's cache usage counters; the companion spec asserts both. mock.mount('/bedrock-converse-cache', bedrockConverseCacheMount()) + mock.mount('/bedrock-rerank', bedrockRerankMount()) // decide() usage on a failed answer mapping. aimock has no OpenAI Decisions // endpoint, so this mount hand-crafts a `/v1/decisions` response with valid @@ -2646,6 +2648,90 @@ function bedrockConverseCacheMount(): Mountable { } } +/** Validate the Bedrock Rerank wire format before serving stateless result pages. */ +function bedrockRerankMount(): Mountable { + return { + async handleRequest(req, res, pathname) { + if (req.method !== 'POST' || pathname !== '/rerank') return false + const body = await readJsonRequestBody(req) + const config = asRecord( + asRecord(body?.rerankingConfiguration)?.bedrockRerankingConfiguration, + ) + const modelArn = asRecord(config?.modelConfiguration)?.modelArn + const count = config?.numberOfResults + const token = body?.nextToken + const authorization = req.headers.authorization ?? '' + const modelAllowed = + modelArn === + 'arn:aws:bedrock:us-west-2::foundation-model/cohere.rerank-v3-5:0' || + modelArn === + 'arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0' + const expectedBody = { + queries: [{ type: 'TEXT', textQuery: { text: 'talk about rain' } }], + sources: [ + '{"id":"sun","text":"sunny day at the beach"}', + '{"id":"rain","text":"rainy afternoon in the city"}', + '{"id":"snow","text":"snow in the mountains"}', + ].map((text) => ({ + type: 'INLINE', + inlineDocumentSource: { type: 'TEXT', textDocument: { text } }, + })), + rerankingConfiguration: { + type: 'BEDROCK_RERANKING_MODEL', + bedrockRerankingConfiguration: { + modelConfiguration: { modelArn }, + numberOfResults: count, + }, + }, + ...(token === undefined ? {} : { nextToken: token }), + } + res.setHeader('Content-Type', 'application/json') + if ( + !modelAllowed || + (count !== 1 && count !== 3) || + (token !== undefined && (token !== 'rain-page-2' || count !== 3)) || + !isDeepStrictEqual(body, expectedBody) || + !authorization.startsWith( + 'AWS4-HMAC-SHA256 Credential=AKIA_BEDROCK_E2E/', + ) || + !authorization.includes('/us-west-2/bedrock/aws4_request') || + !authorization.includes('x-test-id') || + !req.headers['x-test-id'] || + !/^[a-f0-9]{64}$/.test(String(req.headers['x-amz-content-sha256'])) || + req.headers['content-type'] !== 'application/json' + ) { + res.statusCode = 400 + res.end( + JSON.stringify({ + message: 'Invalid Bedrock rerank request or signing scope', + }), + ) + return true + } + res.setHeader( + 'x-amzn-requestid', + token ? 'rerank-page-2' : 'rerank-page-1', + ) + res.end( + JSON.stringify( + token + ? { + results: [ + { index: 2, relevanceScore: 0.3 }, + { index: 0, relevanceScore: 0.1 }, + ], + } + : { + results: [{ index: 1, relevanceScore: 0.98 }], + ...(count === 3 ? { nextToken: 'rain-page-2' } : {}), + }, + ), + ) + return true + }, + } +} + function drainBody(req: http.IncomingMessage): Promise { return new Promise((resolve, reject) => { req.on('data', () => {}) diff --git a/testing/e2e/playwright.config.ts b/testing/e2e/playwright.config.ts index 7ffddd7cad..5e05a21edb 100644 --- a/testing/e2e/playwright.config.ts +++ b/testing/e2e/playwright.config.ts @@ -31,6 +31,15 @@ export default defineConfig({ globalTeardown: './global-teardown.ts', webServer: { command: 'pnpm run dev', + // Bedrock reranking signs requests to aimock, never to AWS. + env: { + AWS_ACCESS_KEY_ID: 'AKIA_BEDROCK_E2E', + AWS_SECRET_ACCESS_KEY: 'bedrock-e2e-secret', + AWS_SESSION_TOKEN: '', + AWS_PROFILE: '', + AWS_DEFAULT_PROFILE: '', + AWS_EC2_METADATA_DISABLED: 'true', + }, url: 'http://localhost:3010', reuseExistingServer: !process.env.CI, timeout: 120_000, diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index f77976d93e..31452f9bf7 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -73,6 +73,7 @@ import { Route as ApiAudioRouteImport } from './routes/api.audio' import { Route as ApiAzureOpenaiWireRouteImport } from './routes/api.azure-openai-wire' import { Route as ApiBedrockConverseCacheRouteImport } from './routes/api.bedrock-converse-cache' import { Route as ApiBedrockConverseTruncatedToolCallRouteImport } from './routes/api.bedrock-converse-truncated-tool-call' +import { Route as ApiBedrockRerankRouteImport } from './routes/api.bedrock-rerank' import { Route as ApiByokChatRouteImport } from './routes/api.byok-chat' import { Route as ApiByteplusEncryptedReasoningWireRouteImport } from './routes/api.byteplus-encrypted-reasoning-wire' import { Route as ApiByteplusSeedance1080pWireRouteImport } from './routes/api.byteplus-seedance-1080p-wire' @@ -557,6 +558,11 @@ const ApiBedrockConverseTruncatedToolCallRoute = path: '/api/bedrock-converse-truncated-tool-call', getParentRoute: () => rootRouteImport, } as any) +const ApiBedrockRerankRoute = ApiBedrockRerankRouteImport.update({ + id: '/api/bedrock-rerank', + path: '/api/bedrock-rerank', + getParentRoute: () => rootRouteImport, +} as any) const ApiByokChatRoute = ApiByokChatRouteImport.update({ id: '/api/byok-chat', path: '/api/byok-chat', @@ -1399,6 +1405,7 @@ export interface FileRoutesByFullPath { '/api/azure-openai-wire': typeof ApiAzureOpenaiWireRoute '/api/bedrock-converse-cache': typeof ApiBedrockConverseCacheRoute '/api/bedrock-converse-truncated-tool-call': typeof ApiBedrockConverseTruncatedToolCallRoute + '/api/bedrock-rerank': typeof ApiBedrockRerankRoute '/api/byok-chat': typeof ApiByokChatRoute '/api/byteplus-encrypted-reasoning-wire': typeof ApiByteplusEncryptedReasoningWireRoute '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute @@ -1613,6 +1620,7 @@ export interface FileRoutesByTo { '/api/azure-openai-wire': typeof ApiAzureOpenaiWireRoute '/api/bedrock-converse-cache': typeof ApiBedrockConverseCacheRoute '/api/bedrock-converse-truncated-tool-call': typeof ApiBedrockConverseTruncatedToolCallRoute + '/api/bedrock-rerank': typeof ApiBedrockRerankRoute '/api/byok-chat': typeof ApiByokChatRoute '/api/byteplus-encrypted-reasoning-wire': typeof ApiByteplusEncryptedReasoningWireRoute '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute @@ -1828,6 +1836,7 @@ export interface FileRoutesById { '/api/azure-openai-wire': typeof ApiAzureOpenaiWireRoute '/api/bedrock-converse-cache': typeof ApiBedrockConverseCacheRoute '/api/bedrock-converse-truncated-tool-call': typeof ApiBedrockConverseTruncatedToolCallRoute + '/api/bedrock-rerank': typeof ApiBedrockRerankRoute '/api/byok-chat': typeof ApiByokChatRoute '/api/byteplus-encrypted-reasoning-wire': typeof ApiByteplusEncryptedReasoningWireRoute '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute @@ -2044,6 +2053,7 @@ export interface FileRouteTypes { | '/api/azure-openai-wire' | '/api/bedrock-converse-cache' | '/api/bedrock-converse-truncated-tool-call' + | '/api/bedrock-rerank' | '/api/byok-chat' | '/api/byteplus-encrypted-reasoning-wire' | '/api/byteplus-seedance-1080p-wire' @@ -2258,6 +2268,7 @@ export interface FileRouteTypes { | '/api/azure-openai-wire' | '/api/bedrock-converse-cache' | '/api/bedrock-converse-truncated-tool-call' + | '/api/bedrock-rerank' | '/api/byok-chat' | '/api/byteplus-encrypted-reasoning-wire' | '/api/byteplus-seedance-1080p-wire' @@ -2472,6 +2483,7 @@ export interface FileRouteTypes { | '/api/azure-openai-wire' | '/api/bedrock-converse-cache' | '/api/bedrock-converse-truncated-tool-call' + | '/api/bedrock-rerank' | '/api/byok-chat' | '/api/byteplus-encrypted-reasoning-wire' | '/api/byteplus-seedance-1080p-wire' @@ -2687,6 +2699,7 @@ export interface RootRouteChildren { ApiAzureOpenaiWireRoute: typeof ApiAzureOpenaiWireRoute ApiBedrockConverseCacheRoute: typeof ApiBedrockConverseCacheRoute ApiBedrockConverseTruncatedToolCallRoute: typeof ApiBedrockConverseTruncatedToolCallRoute + ApiBedrockRerankRoute: typeof ApiBedrockRerankRoute ApiByokChatRoute: typeof ApiByokChatRoute ApiByteplusEncryptedReasoningWireRoute: typeof ApiByteplusEncryptedReasoningWireRoute ApiByteplusSeedance1080pWireRoute: typeof ApiByteplusSeedance1080pWireRoute @@ -3283,6 +3296,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiBedrockConverseTruncatedToolCallRouteImport parentRoute: typeof rootRouteImport } + '/api/bedrock-rerank': { + id: '/api/bedrock-rerank' + path: '/api/bedrock-rerank' + fullPath: '/api/bedrock-rerank' + preLoaderRoute: typeof ApiBedrockRerankRouteImport + parentRoute: typeof rootRouteImport + } '/api/byok-chat': { id: '/api/byok-chat' path: '/api/byok-chat' @@ -4447,6 +4467,7 @@ const rootRouteChildren: RootRouteChildren = { ApiBedrockConverseCacheRoute: ApiBedrockConverseCacheRoute, ApiBedrockConverseTruncatedToolCallRoute: ApiBedrockConverseTruncatedToolCallRoute, + ApiBedrockRerankRoute: ApiBedrockRerankRoute, ApiByokChatRoute: ApiByokChatRoute, ApiByteplusEncryptedReasoningWireRoute: ApiByteplusEncryptedReasoningWireRoute, diff --git a/testing/e2e/src/routes/api.bedrock-rerank.ts b/testing/e2e/src/routes/api.bedrock-rerank.ts new file mode 100644 index 0000000000..7d8575cedb --- /dev/null +++ b/testing/e2e/src/routes/api.bedrock-rerank.ts @@ -0,0 +1,43 @@ +import { createFileRoute } from '@tanstack/react-router' +import { rerank } from '@tanstack/ai' +import { bedrockRerank } from '@tanstack/ai-bedrock' + +const LLMOCK_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010' + +export const Route = createFileRoute('/api/bedrock-rerank')({ + server: { + handlers: { + POST: async ({ request }) => { + const params = new URL(request.url).searchParams + const model = params.get('model') + if ( + model !== 'cohere.rerank-v3-5:0' && + model !== 'amazon.rerank-v1:0' + ) { + return new Response('Unknown rerank model', { status: 400 }) + } + const documents = [ + { id: 'sun', text: 'sunny day at the beach' }, + { id: 'rain', text: 'rainy afternoon in the city' }, + { id: 'snow', text: 'snow in the mountains' }, + ] + const result = await rerank({ + adapter: bedrockRerank(model, { + region: 'us-west-2', + baseURL: `${LLMOCK_BASE}/bedrock-rerank`, + defaultHeaders: { 'X-Test-Id': params.get('testId') ?? '' }, + }), + query: 'talk about rain', + documents, + topN: params.get('top') === '1' ? 1 : undefined, + }) + return Response.json({ + ...result, + originalDocumentsPreserved: result.ranking.every( + ({ index, document }) => document === documents[index], + ), + }) + }, + }, + }, +}) diff --git a/testing/e2e/tests/bedrock-rerank.spec.ts b/testing/e2e/tests/bedrock-rerank.spec.ts new file mode 100644 index 0000000000..6a218fd68f --- /dev/null +++ b/testing/e2e/tests/bedrock-rerank.spec.ts @@ -0,0 +1,39 @@ +import { test, expect } from './fixtures' + +for (const model of ['cohere.rerank-v3-5:0', 'amazon.rerank-v1:0']) { + for (const top of ['1', 'all']) { + test(`bedrock rerank ${model}: ${top === '1' ? 'top result' : 'all results across pages'}`, async ({ + request, + testId, + }) => { + const params = new URLSearchParams({ model, top, testId }) + const response = await request.post(`/api/bedrock-rerank?${params}`) + expect(response.status(), await response.text()).toBe(200) + const expectedRanking = [ + { + index: 1, + score: 0.98, + document: { id: 'rain', text: 'rainy afternoon in the city' }, + }, + { + index: 2, + score: 0.3, + document: { id: 'snow', text: 'snow in the mountains' }, + }, + { + index: 0, + score: 0.1, + document: { id: 'sun', text: 'sunny day at the beach' }, + }, + ].slice(0, top === '1' ? 1 : 3) + expect(await response.json()).toEqual({ + id: 'rerank-page-1', + model, + ranking: expectedRanking, + rerankedDocuments: expectedRanking.map(({ document }) => document), + usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 }, + originalDocumentsPreserved: true, + }) + }) + } +} From 30ac15e8e60c0e4357439d1bdde16fdf65acfeb0 Mon Sep 17 00:00:00 2001 From: phrog Date: Mon, 9 Nov 2026 23:20:08 -0400 Subject: [PATCH 2/2] feat(ai-bedrock): allow empty rerank continuation pages --- packages/ai-bedrock/src/adapters/rerank.ts | 2 +- .../ai-bedrock/tests/rerank-adapter.test.ts | 25 +++++++++++++------ testing/e2e/global-setup.ts | 10 ++++++-- testing/e2e/tests/bedrock-rerank.spec.ts | 2 +- 4 files changed, 27 insertions(+), 12 deletions(-) diff --git a/packages/ai-bedrock/src/adapters/rerank.ts b/packages/ai-bedrock/src/adapters/rerank.ts index 87bc72bc20..cf2a9eacbb 100644 --- a/packages/ai-bedrock/src/adapters/rerank.ts +++ b/packages/ai-bedrock/src/adapters/rerank.ts @@ -152,7 +152,7 @@ export class BedrockRerankAdapter< ) nextToken = page.nextToken if (nextToken) { - if (seenTokens.has(nextToken) || page.results.length === 0) { + if (seenTokens.has(nextToken)) { throw new Error('Bedrock rerank pagination did not advance') } seenTokens.add(nextToken) diff --git a/packages/ai-bedrock/tests/rerank-adapter.test.ts b/packages/ai-bedrock/tests/rerank-adapter.test.ts index 866f88e03c..90840ec81a 100644 --- a/packages/ai-bedrock/tests/rerank-adapter.test.ts +++ b/packages/ai-bedrock/tests/rerank-adapter.test.ts @@ -290,14 +290,23 @@ describe('BedrockRerankAdapter', () => { expect(fetchMock).toHaveBeenCalledTimes(2) }) - it('rejects an empty continuation page', async () => { - fetchMock.mockResolvedValue( - bedrockResponse({ results: [], nextToken: 'again' }), - ) - await expect(rerank(rerankOptions())).rejects.toThrow( - 'pagination did not advance', - ) - expect(fetchMock).toHaveBeenCalledTimes(1) + it('follows an empty intermediate page to collect later results', async () => { + fetchMock + .mockResolvedValueOnce( + bedrockResponse({ results: [results[0]], nextToken: 'page-2' }), + ) + .mockResolvedValueOnce( + bedrockResponse({ results: [], nextToken: 'page-3' }, 'aws-request-2'), + ) + .mockResolvedValueOnce( + bedrockResponse({ results: [results[1]] }, 'aws-request-3'), + ) + const result = await rerank(rerankOptions()) + expect(result.id).toBe('aws-request-1') + expect(result.rerankedDocuments).toEqual([documents[1], documents[0]]) + expect(requestBody(1)).toMatchObject({ nextToken: 'page-2' }) + expect(requestBody(2)).toMatchObject({ nextToken: 'page-3' }) + expect(fetchMock).toHaveBeenCalledTimes(3) }) it('accepts 1000 documents and rejects 1001 before resolving credentials', async () => { diff --git a/testing/e2e/global-setup.ts b/testing/e2e/global-setup.ts index e293991d79..56cb8a273d 100644 --- a/testing/e2e/global-setup.ts +++ b/testing/e2e/global-setup.ts @@ -2689,7 +2689,9 @@ function bedrockRerankMount(): Mountable { if ( !modelAllowed || (count !== 1 && count !== 3) || - (token !== undefined && (token !== 'rain-page-2' || count !== 3)) || + (token !== undefined && + ((token !== 'rain-page-2' && token !== 'rain-page-3') || + count !== 3)) || !isDeepStrictEqual(body, expectedBody) || !authorization.startsWith( 'AWS4-HMAC-SHA256 Credential=AKIA_BEDROCK_E2E/', @@ -2710,8 +2712,12 @@ function bedrockRerankMount(): Mountable { } res.setHeader( 'x-amzn-requestid', - token ? 'rerank-page-2' : 'rerank-page-1', + token ? token.replace('rain-', 'rerank-') : 'rerank-page-1', ) + if (token === 'rain-page-2') { + res.end(JSON.stringify({ results: [], nextToken: 'rain-page-3' })) + return true + } res.end( JSON.stringify( token diff --git a/testing/e2e/tests/bedrock-rerank.spec.ts b/testing/e2e/tests/bedrock-rerank.spec.ts index 6a218fd68f..a1b09f8d21 100644 --- a/testing/e2e/tests/bedrock-rerank.spec.ts +++ b/testing/e2e/tests/bedrock-rerank.spec.ts @@ -2,7 +2,7 @@ import { test, expect } from './fixtures' for (const model of ['cohere.rerank-v3-5:0', 'amazon.rerank-v1:0']) { for (const top of ['1', 'all']) { - test(`bedrock rerank ${model}: ${top === '1' ? 'top result' : 'all results across pages'}`, async ({ + test(`bedrock rerank ${model}: ${top === '1' ? 'top result' : 'all results across an empty page'}`, async ({ request, testId, }) => {