From 8f87e1912053fb170c05d312f5da18a7264f7824 Mon Sep 17 00:00:00 2001 From: VincentLambert Date: Thu, 16 Jul 2026 05:06:20 +0200 Subject: [PATCH] feat(llm): add AWS Bedrock reranker connector (#16960) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What Adds a reranker connector for the **Bedrock** factory, which previously offered chat/embedding/CV models but no reranker — selecting a Bedrock rerank model raised `Factory not in rerank model`. ## How `BedrockRerank` calls the `bedrock-agent-runtime` Rerank API. It reuses the same JSON key protocol as `BedrockEmbed` (`auth_mode` / `bedrock_region` / `bedrock_ak` / `bedrock_sk`, with `access_key_secret` / `iam_role` / `assume_role` modes). Documents are truncated to the model window (Cohere Rerank v3.5 ~2k of its shared 4k window, Amazon Rerank v1 8k) on top of Bedrock's own internal truncation. Scores are returned in `[0, 1]`, so the shared `Base.similarity` normalization applies unchanged. Verified against `amazon.rerank-v1:0` and `cohere.rerank-v3-5:0` in `eu-central-1`. > Note: this PR adds the connector only. Bedrock rerank models can be selected by > adding the relevant entries to `conf/llm_factories.json` under the Bedrock > provider; that catalog change is intentionally left out of this PR. ## Tests `test/unit_test/rag/llm/test_bedrock_rerank.py` — boto3 is mocked (no AWS call): score-by-index mapping, per-model document truncation, model ARN construction, auth-mode validation and the empty-input short-circuit. `pytest` green alongside the existing reranker normalization suite. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- rag/llm/rerank_model.py | 107 +++++++++++ test/unit_test/rag/llm/test_bedrock_rerank.py | 178 ++++++++++++++++++ 2 files changed, 285 insertions(+) create mode 100644 test/unit_test/rag/llm/test_bedrock_rerank.py diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py index 5d6171ab21..7778b08b04 100644 --- a/rag/llm/rerank_model.py +++ b/rag/llm/rerank_model.py @@ -15,6 +15,7 @@ # import json import logging +import time from abc import ABC from urllib.parse import urljoin from typing import Tuple, List @@ -298,6 +299,112 @@ class CoHereRerank(Base): return rank, token_count +# Reranker connector for AWS Bedrock, calling the bedrock-agent-runtime Rerank +# API (e.g. amazon.rerank-v1:0, cohere.rerank-v3-5:0). The JSON key protocol +# (auth_mode / bedrock_region / bedrock_ak / bedrock_sk) mirrors BedrockEmbed in +# embedding_model.py. +class BedrockRerank(Base): + _FACTORY_NAME = "Bedrock" + + # Hard limits of the bedrock-agent-runtime Rerank API: each document text + # (RerankTextDocument.text) is capped at 32,000 characters, and a single + # request accepts at most 1,000 sources / numberOfResults. + _MAX_DOC_CHARS = 32000 + _MAX_SOURCES = 1000 + + def __init__(self, key, model_name, **kwargs): + import boto3 + + key = json.loads(key) + mode = key.get("auth_mode") + if not mode: + logging.error("Bedrock auth_mode is not provided in the key") + raise ValueError("Bedrock auth_mode must be provided in the key") + + self.bedrock_region = key.get("bedrock_region") + self.model_name = model_name + # On-demand foundation-model ARN; works for amazon.rerank-v1:0 / cohere.rerank-*. + self.model_arn = f"arn:aws:bedrock:{self.bedrock_region}::foundation-model/{self.model_name}" + # Per-document truncation guard sized to the model window. Cohere Rerank + # v3.5 shares a ~4k window between query and document (~2048 for docs); + # Amazon Rerank v1 handles 32k, but chunks are small so a generous cap + # just bounds pathological payloads. Bedrock also truncates internally. + self.doc_max_tokens = 2048 if self.model_name.split(".")[0] == "cohere" else 8192 + + # Rerank lives on the bedrock-agent-runtime service, not bedrock-runtime. + if mode == "access_key_secret": + self.client = boto3.client( + service_name="bedrock-agent-runtime", + region_name=self.bedrock_region, + aws_access_key_id=key.get("bedrock_ak"), + aws_secret_access_key=key.get("bedrock_sk"), + ) + elif mode == "iam_role": + sts_client = boto3.client("sts", region_name=self.bedrock_region) + resp = sts_client.assume_role(RoleArn=key.get("aws_role_arn"), RoleSessionName="BedrockSession") + creds = resp["Credentials"] + self.client = boto3.client( + service_name="bedrock-agent-runtime", + region_name=self.bedrock_region, + aws_access_key_id=creds["AccessKeyId"], + aws_secret_access_key=creds["SecretAccessKey"], + aws_session_token=creds["SessionToken"], + ) + else: # assume_role: default AWS credential chain + self.client = boto3.client("bedrock-agent-runtime", region_name=self.bedrock_region) + + def _compute_rank(self, query: str, texts: List) -> Tuple[np.ndarray, int]: + # Truncate to the model token window, then enforce the API's hard 32k-char + # per-text limit (a longer RerankTextQuery / RerankTextDocument is rejected). + query = query[: self._MAX_DOC_CHARS] + texts = [truncate(t, self.doc_max_tokens)[: self._MAX_DOC_CHARS] for t in texts] + # Bedrock does not report token usage; count locally like CoHereRerank. + token_count = num_tokens_from_string(query) + sum(num_tokens_from_string(t) for t in texts) + + rank = np.zeros(len(texts), dtype=float) + result_count = 0 + started = time.perf_counter() + # Both `sources` and `numberOfResults` are capped at 1,000 per request; + # rerank in batches and map each score back to its global position. + for offset in range(0, len(texts), self._MAX_SOURCES): + batch = texts[offset : offset + self._MAX_SOURCES] + sources = [{"type": "INLINE", "inlineDocumentSource": {"type": "TEXT", "textDocument": {"text": t}}} for t in batch] + reranking_config = { + "type": "BEDROCK_RERANKING_MODEL", + "bedrockRerankingConfiguration": { + "numberOfResults": len(batch), + "modelConfiguration": {"modelArn": self.model_arn}, + }, + } + # Drain paginated results: the API may split a batch across nextToken pages. + next_token = None + while True: + request = {"queries": [{"type": "TEXT", "textQuery": {"text": query}}], "sources": sources, "rerankingConfiguration": reranking_config} + if next_token: + request["nextToken"] = next_token + res = self.client.rerank(**request) + try: + for d in res.get("results", []): + rank[offset + d["index"]] = d["relevanceScore"] + result_count += 1 + except (KeyError, IndexError, TypeError) as _e: + log_exception(_e, res) + next_token = res.get("nextToken") + if not next_token: + break + # Safe diagnostics only: no query, document text or credentials. + logging.debug( + "BedrockRerank model=%s region=%s sources=%d tokens=%d results=%d elapsed=%.3fs", + self.model_name, + self.bedrock_region, + len(texts), + token_count, + result_count, + time.perf_counter() - started, + ) + return rank, token_count + + class TogetherAIRerank(Base): _FACTORY_NAME = "TogetherAI" diff --git a/test/unit_test/rag/llm/test_bedrock_rerank.py b/test/unit_test/rag/llm/test_bedrock_rerank.py new file mode 100644 index 0000000000..777308286b --- /dev/null +++ b/test/unit_test/rag/llm/test_bedrock_rerank.py @@ -0,0 +1,178 @@ +# +# Copyright 2025 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Unit tests for the AWS Bedrock reranker connector. + +``BedrockRerank`` talks to the ``bedrock-agent-runtime`` Rerank API via boto3. +These tests patch ``boto3.client`` so no AWS call is made, and verify the +score-by-index mapping, per-model document truncation and key/ARN handling. +""" + +import json +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from common.token_utils import num_tokens_from_string +from rag.llm.rerank_model import BedrockRerank + +pytestmark = pytest.mark.p1 + +KEY = json.dumps( + { + "auth_mode": "access_key_secret", + "bedrock_region": "eu-central-1", + "bedrock_ak": "AKIA_TEST", + "bedrock_sk": "secret_test", + } +) + + +def _rerank_response(scores): + """Bedrock returns results out of order; ``index`` maps back to the source doc.""" + return {"results": [{"index": i, "relevanceScore": s} for i, s in scores]} + + +def _make(model_name="amazon.rerank-v1:0", key=KEY): + """Instantiate the connector with ``boto3.client`` patched; return (model, client).""" + with patch("boto3.client") as client_factory: + client = MagicMock() + client_factory.return_value = client + mdl = BedrockRerank(key, model_name) + return mdl, client + + +def test_scores_are_mapped_back_by_index(): + mdl, client = _make() + # Response deliberately out of source order. + client.rerank.return_value = _rerank_response([(2, 0.9), (0, 0.1), (1, 0.5)]) + rank, _ = mdl.similarity("q", ["a", "b", "c"]) + assert np.allclose(rank, [0.1, 0.5, 0.9]) + + +def test_model_arn_is_built_from_region_and_name(): + mdl, _ = _make(model_name="cohere.rerank-v3-5:0") + assert mdl.model_arn == "arn:aws:bedrock:eu-central-1::foundation-model/cohere.rerank-v3-5:0" + + +def test_doc_window_depends_on_model(): + cohere, _ = _make(model_name="cohere.rerank-v3-5:0") + amazon, _ = _make(model_name="amazon.rerank-v1:0") + assert cohere.doc_max_tokens == 2048 + assert amazon.doc_max_tokens == 8192 + + +def test_documents_are_truncated_before_send(): + mdl, client = _make(model_name="cohere.rerank-v3-5:0") # cap 2048 + client.rerank.return_value = _rerank_response([(0, 0.5)]) + mdl.similarity("q", ["mot " * 5000]) # > 2048 tokens + sent = client.rerank.call_args.kwargs["sources"][0]["inlineDocumentSource"]["textDocument"]["text"] + assert num_tokens_from_string(sent) <= 2048 + + +def test_number_of_results_covers_all_documents(): + mdl, client = _make() + client.rerank.return_value = _rerank_response([(0, 0.3), (1, 0.6)]) + mdl.similarity("q", ["a", "b"]) + cfg = client.rerank.call_args.kwargs["rerankingConfiguration"]["bedrockRerankingConfiguration"] + assert cfg["numberOfResults"] == 2 + assert cfg["modelConfiguration"]["modelArn"].endswith("amazon.rerank-v1:0") + + +def test_missing_auth_mode_raises(): + with patch("boto3.client"): + with pytest.raises(ValueError): + BedrockRerank(json.dumps({"bedrock_region": "eu-central-1"}), "amazon.rerank-v1:0") + + +def test_access_key_secret_mode_wires_the_client(): + with patch("boto3.client") as client_factory: + client_factory.return_value = MagicMock() + BedrockRerank(KEY, "amazon.rerank-v1:0") + kwargs = client_factory.call_args.kwargs + assert kwargs["service_name"] == "bedrock-agent-runtime" + assert kwargs["region_name"] == "eu-central-1" + assert kwargs["aws_access_key_id"] == "AKIA_TEST" + assert kwargs["aws_secret_access_key"] == "secret_test" + + +@pytest.mark.parametrize("query, texts", [("", ["a"]), ("q", [])]) +def test_empty_input_short_circuits_without_calling_bedrock(query, texts): + mdl, client = _make() + rank, tokens = mdl.similarity(query, texts) + assert tokens == 0 + assert rank.size == len(texts) + client.rerank.assert_not_called() + + +def test_document_is_capped_at_32000_chars(): + # RerankTextDocument.text is hard-capped at 32,000 chars; a longer doc would + # otherwise be rejected by the API. Amazon's 8k-token window can exceed it. + mdl, client = _make(model_name="amazon.rerank-v1:0") + client.rerank.return_value = _rerank_response([(0, 0.5)]) + mdl.similarity("q", ["word " * 30000]) # 150k chars -> ~41k after token truncation + sent = client.rerank.call_args.kwargs["sources"][0]["inlineDocumentSource"]["textDocument"]["text"] + assert len(sent) == 32000 + + +def test_query_is_capped_at_32000_chars(): + mdl, client = _make(model_name="amazon.rerank-v1:0") + client.rerank.return_value = _rerank_response([(0, 0.5)]) + mdl.similarity("q " * 30000, ["doc"]) # oversize query + sent = client.rerank.call_args.kwargs["queries"][0]["textQuery"]["text"] + assert len(sent) == 32000 + + +def test_short_document_is_not_char_capped(): + mdl, client = _make(model_name="amazon.rerank-v1:0") + client.rerank.return_value = _rerank_response([(0, 0.5)]) + mdl.similarity("q", ["short document"]) + sent = client.rerank.call_args.kwargs["sources"][0]["inlineDocumentSource"]["textDocument"]["text"] + assert sent == "short document" + + +def test_batches_requests_over_the_1000_source_limit(): + mdl, client = _make() + + def _fake_rerank(**kwargs): + n = len(kwargs["sources"]) + return {"results": [{"index": i, "relevanceScore": 0.5} for i in range(n)]} + + client.rerank.side_effect = _fake_rerank + n = 2001 + rank, _ = mdl.similarity("q", ["doc"] * n) + assert rank.shape == (n,) + assert client.rerank.call_count == 3 # 1000 + 1000 + 1 + for call in client.rerank.call_args_list: + cfg = call.kwargs["rerankingConfiguration"]["bedrockRerankingConfiguration"] + sources = call.kwargs["sources"] + assert len(sources) <= 1000 + assert cfg["numberOfResults"] == len(sources) + + +def test_paginated_results_are_all_consumed(): + mdl, client = _make() + # First page returns one score + a nextToken; second page returns the rest. + pages = [ + {"results": [{"index": 0, "relevanceScore": 0.9}], "nextToken": "tok"}, + {"results": [{"index": 1, "relevanceScore": 0.4}, {"index": 2, "relevanceScore": 0.1}]}, + ] + client.rerank.side_effect = pages + rank, _ = mdl.similarity("q", ["a", "b", "c"]) + assert client.rerank.call_count == 2 + assert client.rerank.call_args_list[1].kwargs["nextToken"] == "tok" + assert np.allclose(rank, [0.9, 0.4, 0.1])