From 4711cbb85e14a166acb01ecb946ec16ccf77cc58 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 24 Feb 2026 08:45:49 -0800 Subject: [PATCH] fix(bedrock): pass timeout to bedrock rerank http client --- litellm/llms/bedrock/rerank/handler.py | 8 +- ...odel_prices_and_context_window_backup.json | 95 +++++++++++++++---- litellm/rerank_api/main.py | 1 + .../test_bedrock_rerank_header_forwarding.py | 77 +++++++++++++++ 4 files changed, 160 insertions(+), 21 deletions(-) diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 06f1e9e86c9..37167e7c330 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -29,12 +29,13 @@ class BedrockRerankHandler(BaseAWSLLM): async def arerank( self, prepared_request: BedrockPreparedRequest, + timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[AsyncHTTPHandler] = None, ): if client is None: client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK) try: - response = await client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"]) + response = await client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"], timeout=timeout) response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code @@ -56,6 +57,7 @@ class BedrockRerankHandler(BaseAWSLLM): return_documents: Optional[bool] = True, max_chunks_per_doc: Optional[int] = None, _is_async: Optional[bool] = False, + timeout: Optional[Union[float, httpx.Timeout]] = None, api_base: Optional[str] = None, extra_headers: Optional[dict] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, @@ -89,12 +91,12 @@ class BedrockRerankHandler(BaseAWSLLM): ) if _is_async: - return self.arerank(prepared_request, client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None) # type: ignore + return self.arerank(prepared_request, timeout=timeout, client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None) # type: ignore if client is None or not isinstance(client, HTTPHandler): client = _get_httpx_client() try: - response = client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"]) + response = client.post(url=prepared_request["endpoint_url"], headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"], timeout=timeout) response.raise_for_status() except httpx.HTTPStatusError as err: error_code = err.response.status_code diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3909ce4c8b0..624714feb87 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26555,65 +26555,124 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "perplexity/preset/fast-search": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_preset": true, + "supports_function_calling": true + }, "perplexity/preset/pro-search": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_preset": true + "supports_preset": true, + "supports_function_calling": true }, - "perplexity/openai/gpt-4o": { + "perplexity/preset/deep-research": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_preset": true, + "supports_function_calling": true }, - "perplexity/openai/gpt-4o-mini": { + "perplexity/preset/advanced-deep-research": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_preset": true, + "supports_function_calling": true }, "perplexity/openai/gpt-5.2": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": true + "supports_reasoning": true, + "supports_function_calling": true }, - "perplexity/anthropic/claude-3-5-sonnet-20241022": { + "perplexity/openai/gpt-5.1": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/anthropic/claude-3-5-haiku-20241022": { + "perplexity/openai/gpt-5-mini": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/google/gemini-2.0-flash-exp": { + "perplexity/anthropic/claude-opus-4-6": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/google/gemini-2.0-flash-thinking-exp": { + "perplexity/anthropic/claude-opus-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": true + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/xai/grok-2-1212": { + "perplexity/anthropic/claude-sonnet-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true }, - "perplexity/xai/grok-2-vision-1212": { + "perplexity/anthropic/claude-haiku-4-5": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, - "supports_reasoning": false + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-3-pro-preview": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-3-flash-preview": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-2.5-pro": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/google/gemini-2.5-flash": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/xai/grok-4-1-fast-non-reasoning": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true + }, + "perplexity/perplexity/sonar": { + "litellm_provider": "perplexity", + "mode": "responses", + "supports_web_search": true, + "supports_reasoning": false, + "supports_function_calling": true }, "publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": { "input_cost_per_token": 0.0, diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index f47fd6323f0..871b18062ff 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -365,6 +365,7 @@ def rerank( # noqa: PLR0915 max_chunks_per_doc=max_chunks_per_doc, _is_async=_is_async, optional_params=optional_params.model_dump(exclude_unset=True), + timeout=optional_params.timeout, api_base=api_base, extra_headers=merged_headers, logging_obj=litellm_logging_obj, diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index a8ac680908e..d6dc4bfa48d 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -237,6 +237,83 @@ async def test_bedrock_rerank_header_forwarding_async(model): pytest.fail(f"Failed to forward headers to {model}: {str(e)}") +def test_bedrock_rerank_timeout_sync(): + """ + Test that the timeout parameter is passed through to the HTTP client for Bedrock rerank (sync). + """ + client = HTTPHandler() + model = "bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0" + mock_credentials_info = create_mock_credentials() + + with patch.object(client, "post") as mock_post, \ + patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \ + patch("botocore.auth.SigV4Auth") as mock_sigv4: + + mock_sigv4.return_value = MagicMock() + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(bedrock_rerank_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_response.raise_for_status = lambda: None + mock_post.return_value = mock_response + + litellm.rerank( + model=model, + query=test_query, + documents=test_documents, + top_n=3, + client=client, + timeout=0.001, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + + assert mock_post.called + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs.get("timeout") == 0.001, ( + f"Expected timeout=0.001, got timeout={call_kwargs.get('timeout')}" + ) + + +@pytest.mark.asyncio +async def test_bedrock_rerank_timeout_async(): + """ + Test that the timeout parameter is passed through to the HTTP client for Bedrock rerank (async). + """ + client = AsyncHTTPHandler() + model = "bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0" + mock_credentials_info = create_mock_credentials() + + with patch.object(client, "post", new_callable=AsyncMock) as mock_post, \ + patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \ + patch("botocore.auth.SigV4Auth") as mock_sigv4: + + mock_sigv4.return_value = MagicMock() + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.text = json.dumps(bedrock_rerank_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_response.raise_for_status = lambda: None + mock_post.return_value = mock_response + + await litellm.arerank( + model=model, + query=test_query, + documents=test_documents, + top_n=3, + client=client, + timeout=0.001, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + + assert mock_post.called + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs.get("timeout") == 0.001, ( + f"Expected timeout=0.001, got timeout={call_kwargs.get('timeout')}" + ) + + def test_bedrock_rerank_extra_headers_and_headers_merge(): """ Test that both extra_headers and headers parameters are correctly merged for Bedrock rerank.