diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index e74ab81c97a..098f60c10e3 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -30,10 +30,10 @@ class BedrockRerankHandler(BaseAWSLLM): async def arerank( self, prepared_request: BedrockPreparedRequest, - documents: Sequence[str | dict[str, Any]], - return_documents: bool, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, + documents: Sequence[str | dict[str, Any]] | None = None, + return_documents: bool = True, ): if client is None: client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py index e1a4ab96cc7..fc013990f8a 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_return_documents.py @@ -165,3 +165,45 @@ async def test_bedrock_rerank_return_documents_false_async(): assert response.results is not None for result in response.results: assert "document" not in result, f"return_documents=False must omit document, got {result}" + + +def test_bedrock_rerank_dict_document_without_text_key_skips_back_fill(): + from litellm.llms.bedrock.rerank.transformation import BedrockRerankConfig + + json_documents = [ + {"title": "zero", "body": "json zero body"}, + {"title": "one", "body": "json one body"}, + {"title": "two", "body": "json two body"}, + ] + response = BedrockRerankConfig()._transform_response( + BEDROCK_RESPONSE, documents=json_documents, return_documents=True + ) + assert response.results is not None + for result in response.results: + assert "document" not in result, f"dict documents without 'text' key have no back-fill target, got {result}" + + +def test_bedrock_rerank_non_integer_index_skips_back_fill(): + from litellm.llms.bedrock.rerank.transformation import BedrockRerankConfig + + response = BedrockRerankConfig()._transform_response( + {"results": [{"index": "2", "relevanceScore": 0.9}]}, + documents=DOCUMENTS, + return_documents=True, + ) + assert response.results is not None + assert len(response.results) == 1 + assert "document" not in response.results[0] + + +def test_bedrock_rerank_out_of_range_index_skips_back_fill(): + from litellm.llms.bedrock.rerank.transformation import BedrockRerankConfig + + response = BedrockRerankConfig()._transform_response( + {"results": [{"index": 99, "relevanceScore": 0.9}]}, + documents=DOCUMENTS, + return_documents=True, + ) + assert response.results is not None + assert len(response.results) == 1 + assert "document" not in response.results[0]