diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index ef46ae5c189..5fc24e5ce18 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -151,6 +151,7 @@ class BedrockConverseLLM(BaseAWSLLM): fake_stream=fake_stream, json_mode=json_mode, stream_chunk_size=stream_chunk_size, + timeout=timeout, ) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 1077731779d..2ef3eb9bc02 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -194,6 +194,7 @@ async def make_call( json_mode: Optional[bool] = False, bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None, stream_chunk_size: int = 1024, + timeout: Optional[Union[float, httpx.Timeout]] = None, ): try: if client is None: @@ -212,6 +213,7 @@ async def make_call( data=data, stream=not fake_stream, logging_obj=logging_obj, + timeout=timeout, ) if response.status_code != 200: @@ -1240,6 +1242,7 @@ class BedrockLLM(BaseAWSLLM): logging_obj=logging_obj, fake_stream=True if "ai21" in api_base else False, stream_chunk_size=stream_chunk_size, + timeout=timeout, ), model=model, custom_llm_provider="bedrock", diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 3f1bccaccfc..86dd78c2a16 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -2380,6 +2380,7 @@ async def make_call( model: str, messages: list, logging_obj, + timeout: Optional[Union[float, httpx.Timeout]] = None, ): if gemini_client is not None: client = gemini_client @@ -2390,7 +2391,7 @@ async def make_call( try: response = await client.post( - api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj + api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj, timeout=timeout ) response.raise_for_status() except httpx.HTTPStatusError as e: @@ -2565,6 +2566,7 @@ class VertexLLM(VertexBase): model=model, messages=messages, logging_obj=logging_obj, + timeout=timeout, ), model=model, custom_llm_provider="vertex_ai_beta", diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_streaming_timeout.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_streaming_timeout.py new file mode 100644 index 00000000000..f127b207ca6 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_streaming_timeout.py @@ -0,0 +1,83 @@ +""" +Verify that timeout is forwarded from async_streaming() through make_call() +to client.post() for Bedrock streaming requests. + +Regression test for https://github.com/BerriAI/litellm/issues/23375 +""" + +import asyncio +import os +import sys +from functools import partial +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) + + +def _run_bedrock_make_call(**extra_kwargs): + """Helper to call bedrock make_call with mocked dependencies.""" + from litellm.llms.bedrock.chat.invoke_handler import make_call + + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.aiter_bytes = MagicMock(return_value=AsyncMock()) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + + mock_logging = MagicMock() + mock_logging.litellm_params = {} + + with patch("litellm.llms.bedrock.chat.invoke_handler.AWSEventStreamDecoder"): + asyncio.run( + make_call( + client=mock_client, + api_base="https://bedrock.us-east-1.amazonaws.com/model/invoke", + headers={"Content-Type": "application/json"}, + data='{"prompt": "test"}', + model="anthropic.claude-3-sonnet", + messages=[{"role": "user", "content": "test"}], + logging_obj=mock_logging, + **extra_kwargs, + ) + ) + return mock_client + + +def test_bedrock_make_call_forwards_timeout_to_client_post(): + mock_client = _run_bedrock_make_call(timeout=0.1) + mock_client.post.assert_called_once() + assert mock_client.post.call_args.kwargs.get("timeout") == 0.1 + + +def test_bedrock_make_call_timeout_defaults_to_none(): + mock_client = _run_bedrock_make_call() + assert mock_client.post.call_args.kwargs.get("timeout") is None + + +def test_bedrock_make_call_forwards_httpx_timeout_object(): + timeout_obj = httpx.Timeout(5.0, connect=2.0) + mock_client = _run_bedrock_make_call(timeout=timeout_obj) + assert mock_client.post.call_args.kwargs.get("timeout") is timeout_obj + + +def test_bedrock_make_call_partial_includes_timeout(): + """Verify that partial(make_call, ..., timeout=X) binds the timeout arg.""" + from litellm.llms.bedrock.chat.invoke_handler import make_call + + bound = partial( + make_call, + client=None, + api_base="https://example.com", + headers={}, + data="{}", + model="test", + messages=[], + logging_obj=MagicMock(), + timeout=0.5, + ) + assert bound.keywords["timeout"] == 0.5 diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_streaming_timeout.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_streaming_timeout.py new file mode 100644 index 00000000000..4c6fae22660 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_streaming_timeout.py @@ -0,0 +1,87 @@ +""" +Verify that timeout is forwarded from async_streaming() through make_call() +to client.post() for Vertex AI Gemini streaming requests. + +Regression test for https://github.com/BerriAI/litellm/issues/23375 +""" + +import asyncio +import os +import sys +from functools import partial +from unittest.mock import AsyncMock, MagicMock + +import httpx + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) + + +def _run_vertex_make_call(**extra_kwargs): + """Helper to call vertex make_call with mocked dependencies.""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + make_call, + ) + + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + mock_response.aiter_lines = MagicMock(return_value=AsyncMock()) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + + mock_logging = MagicMock() + + asyncio.run( + make_call( + client=mock_client, + gemini_client=None, + api_base="https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/l/publishers/google/models/gemini:streamGenerateContent", + headers={"Authorization": "Bearer token"}, + data='{"contents": []}', + model="gemini-2.5-flash", + messages=[{"role": "user", "content": "test"}], + logging_obj=mock_logging, + **extra_kwargs, + ) + ) + return mock_client + + +def test_vertex_make_call_forwards_timeout_to_client_post(): + mock_client = _run_vertex_make_call(timeout=0.1) + mock_client.post.assert_called_once() + assert mock_client.post.call_args.kwargs.get("timeout") == 0.1 + + +def test_vertex_make_call_timeout_defaults_to_none(): + mock_client = _run_vertex_make_call() + assert mock_client.post.call_args.kwargs.get("timeout") is None + + +def test_vertex_make_call_forwards_httpx_timeout_object(): + timeout_obj = httpx.Timeout(5.0, connect=2.0) + mock_client = _run_vertex_make_call(timeout=timeout_obj) + assert mock_client.post.call_args.kwargs.get("timeout") is timeout_obj + + +def test_vertex_make_call_partial_includes_timeout(): + """Verify that partial(make_call, ..., timeout=X) binds the timeout arg.""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + make_call, + ) + + bound = partial( + make_call, + gemini_client=None, + api_base="https://example.com", + headers={}, + data="{}", + model="test", + messages=[], + logging_obj=MagicMock(), + timeout=0.5, + ) + assert bound.keywords["timeout"] == 0.5