mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-17 05:07:23 +08:00
feat(llm): add AWS Bedrock reranker connector (#16960)
## 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)
This commit is contained in:
VincentLambert
@@ -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"
|
||||
|
||||
|
||||
178
test/unit_test/rag/llm/test_bedrock_rerank.py
Normal file
178
test/unit_test/rag/llm/test_bedrock_rerank.py
Normal file
@@ -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])
|
||||
Reference in New Issue
Block a user