scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

fix(bedrock): pass timeout to bedrock rerank http client

Browse source
This commit is contained in:
Ishaan Jaffer 2026-02-24 08:45:49 -08:00
parent 4652c73259
commit 4711cbb85e
4 changed files with 160 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/llms/bedrock/rerank/handler.py
View file

@ -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

95
litellm/model_prices_and_context_window_backup.json
View file

@ -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,

1
litellm/rerank_api/main.py
View file

@ -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,

77
tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py
View file

@ -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.