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

address review: normalize timeout consistently, fix sync converse path, add tests

Browse source 
Made-with: Cursor
This commit is contained in:
pradyyadav 2026-03-13 01:02:07 +05:30
parent bd0e92926a
commit 24b768e271
5 changed files with 403 additions and 567 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -42,7 +42,9 @@ def make_sync_call(
_params: dict = {}
if timeout is not None:
_params["timeout"] = timeout
client = _get_httpx_client(params=_params if _params else None)
client = _get_httpx_client(
params=_params if _params else None
)
response = client.post(
api_base,
@ -78,9 +80,7 @@ def make_sync_call(
)
else:
decoder = AWSEventStreamDecoder(model=model, json_mode=json_mode)
completion_stream = decoder.iter_bytes(
response.iter_bytes(chunk_size=stream_chunk_size)
)
completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size))
# LOGGING
logging_obj.post_call(
@ -134,7 +134,7 @@ class BedrockConverseLLM(BaseAWSLLM):
endpoint_url=api_base,
data=data,
headers=headers,
api_key=api_key,
api_key=api_key
)
## LOGGING
@ -195,7 +195,7 @@ class BedrockConverseLLM(BaseAWSLLM):
headers=headers,
)
data = json.dumps(request_data)
prepped = self.get_request_headers(
credentials=credentials,
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
@ -203,7 +203,7 @@ class BedrockConverseLLM(BaseAWSLLM):
endpoint_url=api_base,
data=data,
headers=headers,
api_key=api_key,
api_key=api_key
)
## LOGGING
@ -289,7 +289,7 @@ class BedrockConverseLLM(BaseAWSLLM):
_stripped = _model_for_id
for rp in ["bedrock/converse/", "bedrock/", "converse/"]:
if _stripped.startswith(rp):
_stripped = _stripped[len(rp) :]
_stripped = _stripped[len(rp):]
break
# Strip embedded region prefix (e.g. "bedrock/us-east-1/model" -> "model")
# and capture it so it can be used as aws_region_name below.
@ -305,10 +305,7 @@ class BedrockConverseLLM(BaseAWSLLM):
break
modelId = self.encode_model_id(model_id=_model_for_id)
# Inject region extracted from model path so _get_aws_region_name picks it up
if (
_region_from_model is not None
and "aws_region_name" not in optional_params
):
if _region_from_model is not None and "aws_region_name" not in optional_params:
optional_params["aws_region_name"] = _region_from_model
fake_stream = litellm.AmazonConverseConfig().should_fake_stream(
@ -318,6 +315,7 @@ class BedrockConverseLLM(BaseAWSLLM):
custom_llm_provider="bedrock",
)
### SET REGION NAME ###
aws_region_name = self._get_aws_region_name(
optional_params=optional_params,
@ -375,7 +373,7 @@ class BedrockConverseLLM(BaseAWSLLM):
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
# Filter beta headers in HTTP headers before making the request
headers = update_headers_with_filtered_beta(
headers=headers, provider="bedrock_converse"
@ -421,7 +419,7 @@ class BedrockConverseLLM(BaseAWSLLM):
timeout=timeout,
client=client,
credentials=credentials,
api_key=api_key,
api_key=api_key
) # type: ignore
## TRANSFORMATION ##
@ -434,7 +432,7 @@ class BedrockConverseLLM(BaseAWSLLM):
headers=extra_headers,
)
data = json.dumps(_data)
prepped = self.get_request_headers(
credentials=credentials,
aws_region_name=aws_region_name,
@ -442,7 +440,7 @@ class BedrockConverseLLM(BaseAWSLLM):
endpoint_url=proxy_endpoint_url,
data=data,
headers=headers,
api_key=api_key,
api_key=api_key
)
## LOGGING

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

@ -202,11 +202,7 @@ async def make_call(
if client is None:
_params: dict = {}
if (
logging_obj
and logging_obj.litellm_params
and logging_obj.litellm_params.get("ssl_verify")
):
if logging_obj and logging_obj.litellm_params and logging_obj.litellm_params.get("ssl_verify"):
_params["ssl_verify"] = logging_obj.litellm_params.get("ssl_verify")
if timeout is not None:
_params["timeout"] = timeout
@ -417,9 +413,9 @@ class BedrockLLM(BaseAWSLLM):
# Claude 3+ indicators (all use Messages API)
messages_api_indicators = [
"claude-3", # Claude 3.x models
"claude-opus-4", # Claude Opus 4
"claude-sonnet-4", # Claude Sonnet 4
"claude-3", # Claude 3.x models
"claude-opus-4", # Claude Opus 4
"claude-sonnet-4", # Claude Sonnet 4
"claude-haiku-4", # Claude Haiku 4
]

151
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"]
@ -510,19 +510,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"parameters" in _openai_function_object
and _openai_function_object["parameters"] is not None
and isinstance(_openai_function_object["parameters"], dict)
):
if supports_response_json_schema(model):
# Gemini 2.0+: minimal transform (resolve $ref only)
_openai_function_object[
"parameters"
] = _build_vertex_schema_for_gemini_2(
_openai_function_object["parameters"]
)
else:
# Gemini 1.5: full OpenAPI-style transform
_openai_function_object["parameters"] = _build_vertex_schema(
_openai_function_object["parameters"]
)
): # OPENAI accepts JSON Schema, Google accepts OpenAPI schema.
_openai_function_object["parameters"] = _build_vertex_schema(
_openai_function_object["parameters"]
)
openai_function_object = _openai_function_object
@ -641,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()
@ -811,9 +802,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
# Check if this is gemini-3-flash which supports MINIMAL thinking level
# Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview, etc.
is_gemini3flash = model and (
"gemini-3-flash" in model.lower() or "gemini-3.1-flash" in model.lower()
"gemini-3-flash" in model.lower()
or "gemini-3.1-flash" in model.lower()
)
is_gemini31pro = model and (
"gemini-3.1-pro-preview" in model.lower()
)
is_gemini31pro = model and ("gemini-3.1-pro-preview" in model.lower())
if reasoning_effort == "minimal":
if is_gemini3flash:
return {"thinkingLevel": "minimal", "includeThoughts": True}
@ -1096,16 +1090,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
@ -1114,11 +1108,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)
@ -1236,25 +1230,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"IMAGE_PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for prohibited image content.",
}
_GEMINI_FINISH_REASON_KEYS = frozenset(
{
"STOP",
"MAX_TOKENS",
"SAFETY",
"RECITATION",
"FINISH_REASON_UNSPECIFIED",
"MALFORMED_FUNCTION_CALL",
"LANGUAGE",
"OTHER",
"BLOCKLIST",
"PROHIBITED_CONTENT",
"SPII",
"IMAGE_SAFETY",
"IMAGE_PROHIBITED_CONTENT",
"TOO_MANY_TOOL_CALLS",
"MALFORMED_RESPONSE",
}
)
_GEMINI_FINISH_REASON_KEYS = frozenset({
"STOP", "MAX_TOKENS", "SAFETY", "RECITATION", "FINISH_REASON_UNSPECIFIED",
"MALFORMED_FUNCTION_CALL", "LANGUAGE", "OTHER", "BLOCKLIST",
"PROHIBITED_CONTENT", "SPII", "IMAGE_SAFETY", "IMAGE_PROHIBITED_CONTENT",
"TOO_MANY_TOOL_CALLS", "MALFORMED_RESPONSE",
})
@staticmethod
def get_finish_reason_mapping() -> Dict[str, OpenAIChatCompletionFinishReason]:
@ -1477,10 +1458,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
@ -2290,37 +2271,35 @@ 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(
"trafficType"
)
if traffic_type:
model_response._hidden_params.setdefault(
"provider_specific_fields", {}
)["traffic_type"] = traffic_type
model_response._hidden_params.setdefault("provider_specific_fields", {})["traffic_type"] = traffic_type
except Exception as e:
raise VertexAIError(
@ -2407,11 +2386,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()
@ -2990,11 +2965,7 @@ class ModelResponseIterator:
# to correctly set finish_reason="tool_calls" per the OpenAI spec.
if not self.has_seen_tool_calls:
for choice in model_response.choices:
if (
hasattr(choice, "delta")
and choice.delta
and choice.delta.tool_calls
):
if hasattr(choice, "delta") and choice.delta and choice.delta.tool_calls:
self.has_seen_tool_calls = True
break
@ -3010,10 +2981,8 @@ class ModelResponseIterator:
if self.has_seen_tool_calls:
mapped_finish_reason = "tool_calls"
else:
mapped_finish_reason = (
VertexGeminiConfig._check_finish_reason(
None, finish_reason_str
)
mapped_finish_reason = VertexGeminiConfig._check_finish_reason(
None, finish_reason_str
)
choice = StreamingChoices(
finish_reason=mapped_finish_reason,
@ -3046,9 +3015,7 @@ class ModelResponseIterator:
"trafficType"
)
if traffic_type:
model_response._hidden_params.setdefault(
"provider_specific_fields", {}
)["traffic_type"] = traffic_type
model_response._hidden_params.setdefault("provider_specific_fields", {})["traffic_type"] = traffic_type
setattr(model_response, "usage", usage) # type: ignore

22
tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py
View file

@ -249,9 +249,9 @@ def test_bedrock_invoke_async_streaming_passes_timeout_to_make_call():
make_call_partial = captured_partial.get("make_call")
assert make_call_partial is not None
assert (
make_call_partial.keywords.get("timeout") == timeout
), "timeout must be forwarded via partial() to make_call()"
assert make_call_partial.keywords.get("timeout") == timeout, (
"timeout must be forwarded via partial() to make_call()"
)
def test_bedrock_converse_async_streaming_passes_timeout_to_make_call():
@ -311,9 +311,9 @@ def test_bedrock_converse_async_streaming_passes_timeout_to_make_call():
asyncio.run(run())
_, kwargs = mock_make_call.call_args
assert (
kwargs.get("timeout") == timeout
), "timeout must be forwarded to make_call() in BedrockConverseLLM.async_streaming()"
assert kwargs.get("timeout") == timeout, (
"timeout must be forwarded to make_call() in BedrockConverseLLM.async_streaming()"
)
def test_bedrock_converse_sync_make_sync_call_passes_timeout_to_client_post():
@ -338,7 +338,9 @@ def test_bedrock_converse_sync_make_sync_call_passes_timeout_to_client_post():
timeout = httpx.Timeout(4.0)
with patch("litellm.llms.bedrock.chat.converse_handler.AWSEventStreamDecoder"):
with patch(
"litellm.llms.bedrock.chat.converse_handler.AWSEventStreamDecoder"
):
make_sync_call(
client=mock_client,
api_base="https://example.com",
@ -351,6 +353,6 @@ def test_bedrock_converse_sync_make_sync_call_passes_timeout_to_client_post():
)
_, kwargs = mock_client.post.call_args
assert (
kwargs.get("timeout") == timeout
), "timeout must be forwarded to client.post() in converse make_sync_call()"
assert kwargs.get("timeout") == timeout, (
"timeout must be forwarded to client.post() in converse make_sync_call()"
)

755
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

File diff suppressed because it is too large Load diff