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

fix(streaming): forward timeout to make_sync_call() for Bedrock Converse and Vertex AI

Browse source 
The original fix addressed async streaming only. The synchronous
make_sync_call() in both converse_handler.py and
vertex_and_google_ai_studio_gemini.py still dropped timeout silently,
defaulting to 600s regardless of configuration.

- Add timeout param to make_sync_call() in both files
- Forward timeout to client.post() in both make_sync_call() functions
- Pass timeout from sync call sites (partial and direct) in both
  converse_handler.py and vertex_and_google_ai_studio_gemini.py
- Fix sys.path.insert in tests to use __file__-relative paths
- Add sync streaming timeout tests for both providers (7 new tests)

Fixes BerriAI#23375

Made-with: Cursor
This commit is contained in:
netbrah 2026-03-19 17:26:52 -04:00
parent 1439728dc4
commit 20917f69bb
5 changed files with 200 additions and 65 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -33,6 +33,7 @@ def make_sync_call(
json_mode: Optional[bool] = False,
fake_stream: bool = False,
stream_chunk_size: int = 1024,
timeout: Optional[Union[float, httpx.Timeout]] = None,
):
if client is None:
client = _get_httpx_client() # Create a new client if none provided
@ -43,6 +44,7 @@ def make_sync_call(
data=data,
stream=not fake_stream,
logging_obj=logging_obj,
timeout=timeout,
)
if response.status_code != 200:
@ -333,9 +335,9 @@ class BedrockConverseLLM(BaseAWSLLM):
aws_external_id = optional_params.pop("aws_external_id", None)
optional_params.pop("aws_region_name", None)
litellm_params[
"aws_region_name"
] = aws_region_name # [DO NOT DELETE] important for async calls
litellm_params["aws_region_name"] = (
aws_region_name # [DO NOT DELETE] important for async calls
)
credentials: Credentials = self.get_credentials(
aws_access_key_id=aws_access_key_id,
@ -473,6 +475,7 @@ class BedrockConverseLLM(BaseAWSLLM):
json_mode=json_mode,
fake_stream=fake_stream,
stream_chunk_size=stream_chunk_size,
timeout=timeout,
)
streaming_response = CustomStreamWrapper(
completion_stream=completion_stream,

36
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -200,11 +200,13 @@ async def make_call(
if client is None:
client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.BEDROCK,
params={"ssl_verify": logging_obj.litellm_params.get("ssl_verify")}
if logging_obj
and logging_obj.litellm_params
and logging_obj.litellm_params.get("ssl_verify")
else None,
params=(
{"ssl_verify": logging_obj.litellm_params.get("ssl_verify")}
if logging_obj
and logging_obj.litellm_params
and logging_obj.litellm_params.get("ssl_verify")
else None
),
) # Create a new client if none provided
response = await client.post(
@ -295,11 +297,13 @@ def make_sync_call(
try:
if client is None:
client = _get_httpx_client(
params={"ssl_verify": logging_obj.litellm_params.get("ssl_verify")}
if logging_obj
and logging_obj.litellm_params
and logging_obj.litellm_params.get("ssl_verify")
else None
params=(
{"ssl_verify": logging_obj.litellm_params.get("ssl_verify")}
if logging_obj
and logging_obj.litellm_params
and logging_obj.litellm_params.get("ssl_verify")
else None
)
)
response = client.post(
@ -549,9 +553,9 @@ class BedrockLLM(BaseAWSLLM):
content=None,
)
model_response.choices[0].message = _message # type: ignore
model_response._hidden_params[
"original_response"
] = outputText # allow user to access raw anthropic tool calling response
model_response._hidden_params["original_response"] = (
outputText # allow user to access raw anthropic tool calling response
)
if (
_is_function_call is True
and stream is not None
@ -884,9 +888,9 @@ class BedrockLLM(BaseAWSLLM):
): # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in
inference_params[k] = v
if stream is True:
inference_params[
"stream"
] = True # cohere requires stream = True in inference params
inference_params["stream"] = (
True # cohere requires stream = True in inference params
)
data = json.dumps({"prompt": prompt, **inference_params})
elif provider == "anthropic":
if self.is_claude_messages_api_model(model):

92
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -498,9 +498,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
value = _remove_strict_from_schema(value)
for tool in value:
openai_function_object: Optional[
ChatCompletionToolParamFunctionChunk
] = None
openai_function_object: Optional[ChatCompletionToolParamFunctionChunk] = (
None
)
if "function" in tool: # tools list
_openai_function_object = ChatCompletionToolParamFunctionChunk( # type: ignore
**tool["function"]
@ -632,15 +632,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
_tools_list.append(search_tool)
if googleSearchRetrieval is not None:
retrieval_tool = Tools()
retrieval_tool[
VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value
] = googleSearchRetrieval
retrieval_tool[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = (
googleSearchRetrieval
)
_tools_list.append(retrieval_tool)
if enterpriseWebSearch is not None:
enterprise_tool = Tools()
enterprise_tool[
VertexToolName.ENTERPRISE_WEB_SEARCH.value
] = enterpriseWebSearch
enterprise_tool[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = (
enterpriseWebSearch
)
_tools_list.append(enterprise_tool)
if code_execution is not None:
code_tool = Tools()
@ -1087,16 +1087,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
param_description="thinking_budget",
)
if VertexGeminiConfig._is_gemini_3_or_newer(model):
optional_params[
"thinkingConfig"
] = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
effort_value, model
optional_params["thinkingConfig"] = (
VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
effort_value, model
)
)
else:
optional_params[
"thinkingConfig"
] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(
effort_value, model
optional_params["thinkingConfig"] = (
VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(
effort_value, model
)
)
elif param == "thinking":
# Validate no conflict with thinking_level
@ -1105,11 +1105,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
param_name="thinking",
param_description="thinking_budget",
)
optional_params[
"thinkingConfig"
] = VertexGeminiConfig._map_thinking_param(
cast(AnthropicThinkingParam, value),
model=model,
optional_params["thinkingConfig"] = (
VertexGeminiConfig._map_thinking_param(
cast(AnthropicThinkingParam, value),
model=model,
)
)
elif param == "modalities" and isinstance(value, list):
response_modalities = self.map_response_modalities(value)
@ -1468,10 +1468,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
_tool_response_chunk["provider_specific_fields"] = { # type: ignore
"thought_signature": thought_signature
}
_tool_response_chunk[
"id"
] = _encode_tool_call_id_with_signature(
_tool_response_chunk["id"] or "", thought_signature
_tool_response_chunk["id"] = (
_encode_tool_call_id_with_signature(
_tool_response_chunk["id"] or "", thought_signature
)
)
_tools.append(_tool_response_chunk)
cumulative_tool_call_idx += 1
@ -2281,28 +2281,28 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
## ADD METADATA TO RESPONSE ##
setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata)
model_response._hidden_params[
"vertex_ai_grounding_metadata"
] = grounding_metadata
model_response._hidden_params["vertex_ai_grounding_metadata"] = (
grounding_metadata
)
setattr(
model_response, "vertex_ai_url_context_metadata", url_context_metadata
)
model_response._hidden_params[
"vertex_ai_url_context_metadata"
] = url_context_metadata
model_response._hidden_params["vertex_ai_url_context_metadata"] = (
url_context_metadata
)
setattr(model_response, "vertex_ai_safety_results", safety_ratings)
model_response._hidden_params[
"vertex_ai_safety_results"
] = safety_ratings # older approach - maintaining to prevent regressions
model_response._hidden_params["vertex_ai_safety_results"] = (
safety_ratings # older approach - maintaining to prevent regressions
)
## ADD CITATION METADATA ##
setattr(model_response, "vertex_ai_citation_metadata", citation_metadata)
model_response._hidden_params[
"vertex_ai_citation_metadata"
] = citation_metadata # older approach - maintaining to prevent regressions
model_response._hidden_params["vertex_ai_citation_metadata"] = (
citation_metadata # older approach - maintaining to prevent regressions
)
## ADD TRAFFIC TYPE ##
traffic_type = completion_response.get("usageMetadata", {}).get(
@ -2391,7 +2391,12 @@ async def make_call(
try:
response = await client.post(
api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj, timeout=timeout
api_base,
headers=headers,
data=data,
stream=True,
logging_obj=logging_obj,
timeout=timeout,
)
response.raise_for_status()
except httpx.HTTPStatusError as e:
@ -2433,6 +2438,7 @@ def make_sync_call(
model: str,
messages: list,
logging_obj,
timeout: Optional[Union[float, httpx.Timeout]] = None,
):
if gemini_client is not None:
client = gemini_client
@ -2440,7 +2446,12 @@ def make_sync_call(
client = HTTPHandler() # Create a new client if none provided
response = 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,
)
if response.status_code != 200 and response.status_code != 201:
@ -2859,6 +2870,7 @@ class VertexLLM(VertexBase):
messages=messages,
logging_obj=logging_obj,
headers=headers,
timeout=timeout,
),
model=model,
custom_llm_provider="vertex_ai_beta",

53
tests/test_litellm/llms/bedrock/chat/test_bedrock_streaming_timeout.py
View file

@ -1,6 +1,6 @@
"""
Verify that timeout is forwarded from async_streaming() through make_call()
to client.post() for Bedrock streaming requests.
Verify that timeout is forwarded through make_call() (async) and
make_sync_call() (sync) to client.post() for Bedrock streaming requests.
Regression test for https://github.com/BerriAI/litellm/issues/23375
"""
@ -14,7 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
sys.path.insert(
0, os.path.abspath("../../../../..")
0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..")
)
@ -81,3 +81,50 @@ def test_bedrock_make_call_partial_includes_timeout():
timeout=0.5,
)
assert bound.keywords["timeout"] == 0.5
# --- Sync path: make_sync_call (converse_handler) ---
def _run_bedrock_make_sync_call(**extra_kwargs):
"""Helper to call bedrock make_sync_call with mocked dependencies."""
from litellm.llms.bedrock.chat.converse_handler import make_sync_call
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.iter_bytes = MagicMock(return_value=iter([b"chunk"]))
mock_client = MagicMock()
mock_client.post = MagicMock(return_value=mock_response)
mock_logging = MagicMock()
with patch("litellm.llms.bedrock.chat.converse_handler.AWSEventStreamDecoder"):
make_sync_call(
client=mock_client,
api_base="https://bedrock.us-east-1.amazonaws.com/model/converse",
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_sync_call_forwards_timeout_to_client_post():
mock_client = _run_bedrock_make_sync_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_sync_call_timeout_defaults_to_none():
mock_client = _run_bedrock_make_sync_call()
assert mock_client.post.call_args.kwargs.get("timeout") is None
def test_bedrock_make_sync_call_forwards_httpx_timeout_object():
timeout_obj = httpx.Timeout(5.0, connect=2.0)
mock_client = _run_bedrock_make_sync_call(timeout=timeout_obj)
assert mock_client.post.call_args.kwargs.get("timeout") is timeout_obj

75
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_streaming_timeout.py
View file

@ -1,6 +1,6 @@
"""
Verify that timeout is forwarded from async_streaming() through make_call()
to client.post() for Vertex AI Gemini streaming requests.
Verify that timeout is forwarded through make_call() (async) and
make_sync_call() (sync) to client.post() for Vertex AI Gemini streaming requests.
Regression test for https://github.com/BerriAI/litellm/issues/23375
"""
@ -14,7 +14,7 @@ from unittest.mock import AsyncMock, MagicMock
import httpx
sys.path.insert(
0, os.path.abspath("../../../../..")
0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..")
)
@ -85,3 +85,72 @@ def test_vertex_make_call_partial_includes_timeout():
timeout=0.5,
)
assert bound.keywords["timeout"] == 0.5
# --- Sync path: make_sync_call ---
def _run_vertex_make_sync_call(**extra_kwargs):
"""Helper to call vertex make_sync_call with mocked dependencies."""
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
make_sync_call,
)
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.iter_lines = MagicMock(return_value=iter([]))
mock_client = MagicMock()
mock_client.post = MagicMock(return_value=mock_response)
mock_logging = MagicMock()
make_sync_call(
client=None,
gemini_client=mock_client,
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_sync_call_forwards_timeout_to_client_post():
mock_client = _run_vertex_make_sync_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_sync_call_timeout_defaults_to_none():
mock_client = _run_vertex_make_sync_call()
assert mock_client.post.call_args.kwargs.get("timeout") is None
def test_vertex_make_sync_call_forwards_httpx_timeout_object():
timeout_obj = httpx.Timeout(5.0, connect=2.0)
mock_client = _run_vertex_make_sync_call(timeout=timeout_obj)
assert mock_client.post.call_args.kwargs.get("timeout") is timeout_obj
def test_vertex_make_sync_call_partial_includes_timeout():
"""Verify that partial(make_sync_call, ..., timeout=X) binds the timeout arg."""
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
make_sync_call,
)
bound = partial(
make_sync_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