scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(bedrock): pass timeout param to bedrock rerank http client (#22021)

Browse source 
* fix(bedrock): pass timeout to bedrock rerank http client

* refactor: extract large functions to fix PLR0915 ruff lint errors
This commit is contained in:
Ishaan Jaff 2026-02-24 09:32:11 -08:00 • committed by Sameer Kankute
parent b950d997ae
commit d86d49f4ed
7 changed files with 427 additions and 251 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

248
litellm/litellm_core_utils/realtime_streaming.py
View file

@ -203,6 +203,129 @@ class RealTimeStreaming:
return True
return False
async def _handle_provider_config_message(self, raw_response) -> None:
"""Process a backend message when a provider_config is set (transformed path)."""
returned_object = self.provider_config.transform_realtime_response( # type: ignore[union-attr]
raw_response,
self.model,
self.logging_obj,
realtime_response_transform_input={
"session_configuration_request": self.session_configuration_request,
"current_output_item_id": self.current_output_item_id,
"current_response_id": self.current_response_id,
"current_delta_chunks": self.current_delta_chunks,
"current_conversation_id": self.current_conversation_id,
"current_item_chunks": self.current_item_chunks,
"current_delta_type": self.current_delta_type,
},
)
transformed_response = returned_object["response"]
self.current_output_item_id = returned_object["current_output_item_id"]
self.current_response_id = returned_object["current_response_id"]
self.current_delta_chunks = returned_object["current_delta_chunks"]
self.current_conversation_id = returned_object["current_conversation_id"]
self.current_item_chunks = returned_object["current_item_chunks"]
self.current_delta_type = returned_object["current_delta_type"]
self.session_configuration_request = returned_object["session_configuration_request"]
events = (
transformed_response
if isinstance(transformed_response, list)
else [transformed_response]
)
for event in events:
## GUARDRAIL: inject create_response=false on session.created
if isinstance(event, dict) and event.get("type") == "session.created":
if self._has_realtime_guardrails():
await self.backend_ws.send(
json.dumps(
{
"type": "session.update",
"session": {
"turn_detection": {
"type": "server_vad",
"create_response": False,
}
},
}
)
)
for event in events:
event_str = json.dumps(event)
## GUARDRAIL: run on transcription events in provider_config path too
if (
isinstance(event, dict)
and event.get("type")
== "conversation.item.input_audio_transcription.completed"
):
transcript = event.get("transcript", "")
self.store_message(event_str)
await self.websocket.send_text(event_str)
blocked = await self.run_realtime_guardrails(
transcript, item_id=event.get("item_id")
)
if not blocked:
await self.backend_ws.send(
json.dumps({"type": "response.create"})
)
continue
## LOGGING
self.store_message(event_str)
await self.websocket.send_text(event_str)
async def _handle_raw_backend_message(self, raw_response) -> bool:
"""Process a backend message without provider_config (raw path).
Returns True if the caller should skip the default store+forward (i.e. continue the loop).
"""
try:
event_obj = json.loads(raw_response)
if event_obj.get("type") == "session.created":
# If any realtime guardrails are registered, proactively
# set create_response=false so the LLM never auto-responds
# before our guardrail has a chance to run.
if self._has_realtime_guardrails():
await self.backend_ws.send(
json.dumps(
{
"type": "session.update",
"session": {
"turn_detection": {
"type": "server_vad",
"create_response": False,
}
},
}
)
)
verbose_logger.debug(
"[realtime guardrail] injected create_response=false into session"
)
if (
event_obj.get("type")
== "conversation.item.input_audio_transcription.completed"
):
transcript = event_obj.get("transcript", "")
## LOGGING — must happen before continue below
self.store_message(raw_response)
# Forward transcript to client so user sees what they said
await self.websocket.send_text(raw_response)
blocked = await self.run_realtime_guardrails(
transcript,
item_id=event_obj.get("item_id"),
)
if not blocked:
# Clean — trigger LLM response
await self.backend_ws.send(
json.dumps({"type": "response.create"})
)
return True
except (json.JSONDecodeError, AttributeError):
pass
return False
async def backend_to_client_send_messages(self):
import websockets
@ -216,128 +339,11 @@ class RealTimeStreaming:
raw_response = await self.backend_ws.recv() # type: ignore[assignment]
if self.provider_config:
returned_object = self.provider_config.transform_realtime_response(
raw_response,
self.model,
self.logging_obj,
realtime_response_transform_input={
"session_configuration_request": self.session_configuration_request,
"current_output_item_id": self.current_output_item_id,
"current_response_id": self.current_response_id,
"current_delta_chunks": self.current_delta_chunks,
"current_conversation_id": self.current_conversation_id,
"current_item_chunks": self.current_item_chunks,
"current_delta_type": self.current_delta_type,
},
)
transformed_response = returned_object["response"]
self.current_output_item_id = returned_object[
"current_output_item_id"
]
self.current_response_id = returned_object["current_response_id"]
self.current_delta_chunks = returned_object["current_delta_chunks"]
self.current_conversation_id = returned_object[
"current_conversation_id"
]
self.current_item_chunks = returned_object["current_item_chunks"]
self.current_delta_type = returned_object["current_delta_type"]
self.session_configuration_request = returned_object[
"session_configuration_request"
]
events = (
transformed_response
if isinstance(transformed_response, list)
else [transformed_response]
)
for event in events:
## GUARDRAIL: inject create_response=false on session.created
if isinstance(event, dict) and event.get("type") == "session.created":
if self._has_realtime_guardrails():
await self.backend_ws.send(
json.dumps(
{
"type": "session.update",
"session": {
"turn_detection": {
"type": "server_vad",
"create_response": False,
}
},
}
)
)
for event in events:
event_str = json.dumps(event)
## GUARDRAIL: run on transcription events in provider_config path too
if (
isinstance(event, dict)
and event.get("type")
== "conversation.item.input_audio_transcription.completed"
):
transcript = event.get("transcript", "")
self.store_message(event_str)
await self.websocket.send_text(event_str)
blocked = await self.run_realtime_guardrails(
transcript, item_id=event.get("item_id")
)
if not blocked:
await self.backend_ws.send(
json.dumps({"type": "response.create"})
)
continue
## LOGGING
self.store_message(event_str)
await self.websocket.send_text(event_str)
await self._handle_provider_config_message(raw_response)
else:
## GUARDRAIL: intercept transcription events before triggering LLM
try:
event_obj = json.loads(raw_response)
if event_obj.get("type") == "session.created":
# If any realtime guardrails are registered, proactively
# set create_response=false so the LLM never auto-responds
# before our guardrail has a chance to run.
if self._has_realtime_guardrails():
await self.backend_ws.send(
json.dumps(
{
"type": "session.update",
"session": {
"turn_detection": {
"type": "server_vad",
"create_response": False,
}
},
}
)
)
verbose_logger.debug(
"[realtime guardrail] injected create_response=false into session"
)
if (
event_obj.get("type")
== "conversation.item.input_audio_transcription.completed"
):
transcript = event_obj.get("transcript", "")
## LOGGING — must happen before continue below
self.store_message(raw_response)
# Forward transcript to client so user sees what they said
await self.websocket.send_text(raw_response)
blocked = await self.run_realtime_guardrails(
transcript,
item_id=event_obj.get("item_id"),
)
if not blocked:
# Clean — trigger LLM response
await self.backend_ws.send(
json.dumps({"type": "response.create"})
)
continue
except (json.JSONDecodeError, AttributeError):
pass
handled = await self._handle_raw_backend_message(raw_response)
if handled:
continue
## LOGGING
self.store_message(raw_response)
await self.websocket.send_text(raw_response)

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,

122
litellm/proxy/health_check.py
View file

@ -11,7 +11,7 @@ from typing import List, Optional
import litellm
logger = logging.getLogger(__name__)
from litellm.constants import HEALTH_CHECK_TIMEOUT_SECONDS, DEFAULT_HEALTH_CHECK_PROMPT
from litellm.constants import DEFAULT_HEALTH_CHECK_PROMPT, HEALTH_CHECK_TIMEOUT_SECONDS
ILLEGAL_DISPLAY_PARAMS = [
"messages",
@ -98,6 +98,66 @@ async def run_with_timeout(task, timeout):
return {"error": "Timeout exceeded"}
async def _run_model_health_check(model: dict):
litellm_params = model["litellm_params"]
model_info = model.get("model_info", {})
mode = model_info.get("mode", None)
litellm_params = _update_litellm_params_for_health_check(model_info, litellm_params)
timeout = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS
return await run_with_timeout(
litellm.ahealth_check(
litellm_params,
mode=mode,
prompt=DEFAULT_HEALTH_CHECK_PROMPT,
input=["test from litellm"],
),
timeout,
)
async def _run_health_checks_with_bounded_concurrency(
models: list, concurrency_limit: int
) -> tuple[list, int]:
"""
Run health checks with at most `concurrency_limit` active tasks.
Preserves result ordering to match `models`.
"""
results: list = [None] * len(models)
tasks_to_index: dict[asyncio.Task, int] = {}
model_iter = iter(enumerate(models))
peak_in_flight = 0
def _schedule_next() -> bool:
nonlocal peak_in_flight
try:
idx, next_model = next(model_iter)
except StopIteration:
return False
task = asyncio.create_task(_run_model_health_check(next_model))
tasks_to_index[task] = idx
peak_in_flight = max(peak_in_flight, len(tasks_to_index))
return True
for _ in range(min(concurrency_limit, len(models))):
_schedule_next()
while tasks_to_index:
done, _ = await asyncio.wait(
set(tasks_to_index.keys()),
return_when=asyncio.FIRST_COMPLETED,
)
for task in done:
idx = tasks_to_index.pop(task)
try:
results[idx] = task.result()
except Exception as e:
results[idx] = e
_schedule_next()
return results, peak_in_flight
async def _perform_health_check(
model_list: list,
details: Optional[bool] = True,
@ -115,66 +175,6 @@ async def _perform_health_check(
cycle_id = instrumentation_context.get("cycle_id", "unknown")
source = instrumentation_context.get("source", "unknown")
async def _run_model_health_check(model: dict):
litellm_params = model["litellm_params"]
model_info = model.get("model_info", {})
mode = model_info.get("mode", None)
litellm_params = _update_litellm_params_for_health_check(
model_info, litellm_params
)
timeout = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS
return await run_with_timeout(
litellm.ahealth_check(
litellm_params,
mode=mode,
prompt=DEFAULT_HEALTH_CHECK_PROMPT,
input=["test from litellm"],
),
timeout,
)
async def _run_health_checks_with_bounded_concurrency(
models: list, concurrency_limit: int
) -> tuple[list, int]:
"""
Run health checks with at most `concurrency_limit` active tasks.
Preserves result ordering to match `models`.
"""
results: list = [None] * len(models)
tasks_to_index: dict[asyncio.Task, int] = {}
model_iter = iter(enumerate(models))
peak_in_flight = 0
def _schedule_next() -> bool:
nonlocal peak_in_flight
try:
idx, next_model = next(model_iter)
except StopIteration:
return False
task = asyncio.create_task(_run_model_health_check(next_model))
tasks_to_index[task] = idx
peak_in_flight = max(peak_in_flight, len(tasks_to_index))
return True
for _ in range(min(concurrency_limit, len(models))):
_schedule_next()
while tasks_to_index:
done, _ = await asyncio.wait(
set(tasks_to_index.keys()),
return_when=asyncio.FIRST_COMPLETED,
)
for task in done:
idx = tasks_to_index.pop(task)
try:
results[idx] = task.result()
except Exception as e:
results[idx] = e
_schedule_next()
return results, peak_in_flight
dispatch_mode = "unbounded"
peak_in_flight = 0
if isinstance(max_concurrency, int) and max_concurrency > 0:

127
litellm/proxy/proxy_server.py
View file

@ -713,7 +713,7 @@ async def _initialize_shared_aiohttp_session():
try:
from aiohttp import ClientSession, TCPConnector
connector_kwargs = {
connector_kwargs: Dict[str, Any] = {
"keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT,
"ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE,
}
@ -1959,6 +1959,65 @@ def _rss_mb_for_log() -> str:
return f"{rss_mb:.2f}"
async def _run_direct_health_check_with_instrumentation(
model_list: list,
details: Optional[bool],
max_concurrency: Optional[int],
instrumentation_context: dict,
):
try:
return await perform_health_check(
model_list=model_list,
details=details,
max_concurrency=max_concurrency,
instrumentation_context=instrumentation_context,
)
except TypeError as e:
if "instrumentation_context" not in str(e):
raise
# Backward compatibility for monkeypatched or wrapped callables
# that do not accept instrumentation_context.
return await perform_health_check(
model_list=model_list,
details=details,
max_concurrency=max_concurrency,
)
def _schedule_background_health_check_db_save(
prisma_client,
shared_health_manager,
model_list: list,
healthy_endpoints: list,
unhealthy_endpoints: list,
):
"""Fire-and-forget: persist health check results to DB if prisma is available."""
if prisma_client is None:
return
import time as time_module
from litellm.proxy.health_endpoints._health_endpoints import (
_save_background_health_checks_to_db,
)
checked_by = (
shared_health_manager.pod_id
if shared_health_manager is not None
else "background_health_check"
)
start_time = time_module.time()
asyncio.create_task(
_save_background_health_checks_to_db(
prisma_client,
model_list,
healthy_endpoints,
unhealthy_endpoints,
start_time,
checked_by=checked_by,
)
)
async def _run_background_health_check():
"""
Periodically run health checks in the background on the endpoints.
@ -2053,25 +2112,6 @@ async def _run_background_health_check():
"cycle_id": cycle_id,
}
async def _run_direct_health_check_with_instrumentation():
try:
return await perform_health_check(
model_list=_llm_model_list,
details=health_check_details,
max_concurrency=health_check_concurrency,
instrumentation_context=instrumentation_context,
)
except TypeError as e:
if "instrumentation_context" not in str(e):
raise
# Backward compatibility for monkeypatched or wrapped callables
# that do not accept instrumentation_context.
return await perform_health_check(
model_list=_llm_model_list,
details=health_check_details,
max_concurrency=health_check_concurrency,
)
# Use shared health check if available, otherwise fall back to direct health check
# Convert health_check_details to bool for perform_shared_health_check (defaults to True if None)
details_bool = (
@ -2094,11 +2134,21 @@ async def _run_background_health_check():
str(e),
)
healthy_endpoints, unhealthy_endpoints = (
await _run_direct_health_check_with_instrumentation()
await _run_direct_health_check_with_instrumentation(
_llm_model_list,
health_check_details,
health_check_concurrency,
instrumentation_context,
)
)
else:
healthy_endpoints, unhealthy_endpoints = (
await _run_direct_health_check_with_instrumentation()
await _run_direct_health_check_with_instrumentation(
_llm_model_list,
health_check_details,
health_check_concurrency,
instrumentation_context,
)
)
# Update the global variable with the health check results
@ -2127,32 +2177,13 @@ async def _run_background_health_check():
)
# Save background health checks to database (non-blocking)
if prisma_client is not None:
import time as time_module
from litellm.proxy.health_endpoints._health_endpoints import (
_save_background_health_checks_to_db,
)
# Use pod_id or a system identifier for checked_by if shared health check is enabled
checked_by = None
if shared_health_manager is not None:
checked_by = shared_health_manager.pod_id
else:
# Use a system identifier for background health checks
checked_by = "background_health_check"
start_time = time_module.time()
asyncio.create_task(
_save_background_health_checks_to_db(
prisma_client,
_llm_model_list,
healthy_endpoints,
unhealthy_endpoints,
start_time,
checked_by=checked_by,
)
)
_schedule_background_health_check_db_save(
prisma_client,
shared_health_manager,
_llm_model_list,
healthy_endpoints,
unhealthy_endpoints,
)
await asyncio.sleep(health_check_interval)

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.