mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-08 11:11:14 +00:00
fix(streaming): normalize status code extraction to prevent 4xx errors from triggering mid-stream fallback (#18698)
在流式处理错误时,添加状态码标准化逻辑,确保 4xx 客户端错误直接抛出而不是被包装成 MidStreamFallbackError。 - 新增 _normalize_status_code 函数用于从异常对象提取状态码 - 优先从异常的 status_code 属性获取,其次从 response.status_code 获取 - 当映射异常或原始异常的状态码在 400-499 范围内时,直接抛出映射异常 - 添加单元测试验证 Vertex AI 400 错误正确抛出为 BadRequestError - 确保流式处理中的客户端错误能够正确传播,而不会触发回退机制
This commit is contained in:
@@ -2000,24 +2000,56 @@ class CustomStreamWrapper:
|
||||
)
|
||||
## Map to OpenAI Exception
|
||||
try:
|
||||
raise exception_type(
|
||||
mapped_exception = exception_type(
|
||||
model=self.model,
|
||||
custom_llm_provider=self.custom_llm_provider,
|
||||
original_exception=e,
|
||||
completion_kwargs={},
|
||||
extra_kwargs={},
|
||||
)
|
||||
except Exception as e:
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
except Exception as mapping_error:
|
||||
mapped_exception = mapping_error
|
||||
|
||||
raise MidStreamFallbackError(
|
||||
message=str(e),
|
||||
model=self.model,
|
||||
llm_provider=self.custom_llm_provider or "anthropic",
|
||||
original_exception=e,
|
||||
generated_content=self.response_uptil_now,
|
||||
is_pre_first_chunk=not self.sent_first_chunk,
|
||||
)
|
||||
def _normalize_status_code(exc: Exception) -> Optional[int]:
|
||||
"""
|
||||
Best-effort status_code extraction.
|
||||
Uses status_code on the exception, then falls back to the response.
|
||||
"""
|
||||
try:
|
||||
code = getattr(exc, "status_code", None)
|
||||
if code is not None:
|
||||
return int(code)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
response = getattr(exc, "response", None)
|
||||
if response is not None:
|
||||
try:
|
||||
status_code = getattr(response, "status_code", None)
|
||||
if status_code is not None:
|
||||
return int(status_code)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
mapped_status_code = _normalize_status_code(mapped_exception)
|
||||
original_status_code = _normalize_status_code(e)
|
||||
|
||||
if mapped_status_code is not None and 400 <= mapped_status_code < 500:
|
||||
raise mapped_exception
|
||||
if original_status_code is not None and 400 <= original_status_code < 500:
|
||||
raise mapped_exception
|
||||
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
raise MidStreamFallbackError(
|
||||
message=str(mapped_exception),
|
||||
model=self.model,
|
||||
llm_provider=self.custom_llm_provider or "anthropic",
|
||||
original_exception=mapped_exception,
|
||||
generated_content=self.response_uptil_now,
|
||||
is_pre_first_chunk=not self.sent_first_chunk,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _strip_sse_data_from_chunk(chunk: Optional[str]) -> Optional[str]:
|
||||
|
||||
@@ -691,6 +691,29 @@ async def test_streaming_completion_start_time(logging_obj: Logging):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_streaming_bad_request_not_midstream(logging_obj: Logging):
|
||||
"""Ensure Vertex bad request errors surface as 400, not mid-stream fallbacks."""
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
|
||||
async def _raise_bad_request(**kwargs):
|
||||
raise VertexAIError(status_code=400, message="invalid maxOutputTokens", headers=None)
|
||||
|
||||
response = CustomStreamWrapper(
|
||||
completion_stream=None,
|
||||
model="gemini-3-pro-preview",
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="vertex_ai_beta",
|
||||
make_call=_raise_bad_request,
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BadRequestError) as excinfo:
|
||||
await response.__anext__()
|
||||
|
||||
assert getattr(excinfo.value, "status_code", None) == 400
|
||||
assert "invalid maxOutputTokens" in str(excinfo.value)
|
||||
|
||||
|
||||
def test_streaming_handler_with_created_time_propagation(
|
||||
initialized_custom_stream_wrapper: CustomStreamWrapper, logging_obj: Logging
|
||||
):
|
||||
|
||||
@@ -0,0 +1,355 @@
|
||||
"""
|
||||
Test cases for functionCall args serialization in Vertex AI Gemini.
|
||||
|
||||
This test file specifically tests the edge cases where Vertex AI might return
|
||||
functionCall args in unexpected formats that could lead to invalid JSON strings
|
||||
like: {"x":"x"}{"a":"a"}
|
||||
"""
|
||||
import json
|
||||
from typing import List, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
from litellm.types.llms.vertex_ai import HttpxPartType
|
||||
|
||||
|
||||
class TestFunctionCallArgsSerialization:
|
||||
"""Test cases for functionCall args serialization edge cases."""
|
||||
|
||||
def test_normal_dict_args(self):
|
||||
"""Test normal case: args is a dict."""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": {"location": "Boston", "unit": "celsius"},
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
assert tools[0]["function"]["name"] == "get_weather"
|
||||
|
||||
# Verify arguments is a valid JSON string
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
assert isinstance(arguments, str)
|
||||
# Should be valid JSON
|
||||
parsed = json.loads(arguments)
|
||||
assert parsed == {"location": "Boston", "unit": "celsius"}
|
||||
|
||||
def test_none_args(self):
|
||||
"""Test case: args is None."""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": None,
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
# Should serialize None to "null" or empty dict
|
||||
assert isinstance(arguments, str)
|
||||
parsed = json.loads(arguments)
|
||||
# json.dumps(None) returns "null"
|
||||
assert parsed is None or parsed == {}
|
||||
|
||||
def test_args_as_string_valid_json(self):
|
||||
"""Test case: args is already a valid JSON string."""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": '{"location": "Boston"}', # String, not dict
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
# If args is a string, json.dumps will double-encode it
|
||||
# This would result in: "{\"location\": \"Boston\"}"
|
||||
assert isinstance(arguments, str)
|
||||
# This is the problematic case - string gets double-encoded
|
||||
# The result would be a JSON string containing a JSON string
|
||||
parsed = json.loads(arguments)
|
||||
# If it's double-encoded, parsed would be a string, not a dict
|
||||
if isinstance(parsed, str):
|
||||
# Double-encoded case
|
||||
inner_parsed = json.loads(parsed)
|
||||
assert inner_parsed == {"location": "Boston"}
|
||||
else:
|
||||
# Normal case (shouldn't happen if args is string)
|
||||
assert parsed == {"location": "Boston"}
|
||||
|
||||
def test_args_as_string_invalid_json_concatenated(self):
|
||||
"""Test case: args is a string with concatenated JSON objects (the bug case).
|
||||
|
||||
When args is a string like '{"x":"x"}{"a":"a"}', json.dumps() will serialize it
|
||||
as a JSON string, resulting in: "{\"x\":\"x\"}{\"a\":\"a\"}"
|
||||
This is a valid JSON string (the outer quotes), but the content inside is invalid JSON.
|
||||
When you try to parse the inner content, it fails.
|
||||
"""
|
||||
# This simulates the case where Vertex might return something like:
|
||||
# args = '{"x":"x"}{"a":"a"}' # Two JSON objects concatenated
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": '{"x":"x"}{"a":"a"}', # Invalid concatenated JSON
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
assert isinstance(arguments, str)
|
||||
|
||||
# json.dumps() on a string will escape it, so we get:
|
||||
# arguments = '"{\\"x\\":\\"x\\"}{\\"a\\":\\"a\\"}"'
|
||||
# This is a valid JSON string (the outer quotes), but the inner content is invalid
|
||||
parsed_outer = json.loads(arguments)
|
||||
assert isinstance(parsed_outer, str)
|
||||
|
||||
# The inner string is invalid JSON (two objects concatenated)
|
||||
# This is the bug: the inner content cannot be parsed as valid JSON
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
json.loads(parsed_outer)
|
||||
|
||||
# The arguments string would be: "{\"x\":\"x\"}{\"a\":\"a\"}"
|
||||
# Which when parsed gives: '{"x":"x"}{"a":"a"}' (invalid JSON)
|
||||
|
||||
def test_args_as_array(self):
|
||||
"""Test case: args is an array (unexpected but possible)."""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": [{"x": "x"}, {"a": "a"}], # Array of objects
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
assert isinstance(arguments, str)
|
||||
# Should serialize array correctly
|
||||
parsed = json.loads(arguments)
|
||||
assert parsed == [{"x": "x"}, {"a": "a"}]
|
||||
|
||||
def test_args_missing_key(self):
|
||||
"""Test case: args key is missing from functionCall.
|
||||
|
||||
This will raise a KeyError because the code directly accesses part["functionCall"]["args"]
|
||||
without checking if the key exists. This is a bug that should be fixed.
|
||||
"""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
# args key missing
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
# This should raise KeyError because args key is missing
|
||||
with pytest.raises(KeyError):
|
||||
VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
def test_multiple_function_calls(self):
|
||||
"""Test case: multiple function calls in parts."""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": {"location": "Boston"},
|
||||
}
|
||||
},
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_time",
|
||||
"args": {"timezone": "EST"},
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 2
|
||||
assert tools[0]["function"]["name"] == "get_weather"
|
||||
assert tools[1]["function"]["name"] == "get_time"
|
||||
|
||||
# Both should have valid JSON arguments
|
||||
args1 = json.loads(tools[0]["function"]["arguments"])
|
||||
args2 = json.loads(tools[1]["function"]["arguments"])
|
||||
assert args1 == {"location": "Boston"}
|
||||
assert args2 == {"timezone": "EST"}
|
||||
|
||||
def test_args_with_vertex_protobuf_format(self):
|
||||
"""Test case: args in Vertex protobuf format with string_value, etc."""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": {
|
||||
"location": {"string_value": "Boston, MA"},
|
||||
"unit": {"string_value": "celsius"},
|
||||
},
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
assert isinstance(arguments, str)
|
||||
# Should serialize the nested structure correctly
|
||||
parsed = json.loads(arguments)
|
||||
assert "location" in parsed
|
||||
assert "unit" in parsed
|
||||
|
||||
def test_args_as_empty_dict(self):
|
||||
"""Test case: args is an empty dict."""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": {},
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
assert isinstance(arguments, str)
|
||||
parsed = json.loads(arguments)
|
||||
assert parsed == {}
|
||||
|
||||
def test_args_with_special_characters(self):
|
||||
"""Test case: args contains special characters that need escaping."""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": {
|
||||
"location": 'Boston, MA "downtown"',
|
||||
"note": "Line 1\nLine 2",
|
||||
},
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
assert isinstance(arguments, str)
|
||||
# Should handle special characters correctly
|
||||
parsed = json.loads(arguments)
|
||||
assert parsed["location"] == 'Boston, MA "downtown"'
|
||||
assert parsed["note"] == "Line 1\nLine 2"
|
||||
|
||||
def test_args_as_list_of_strings_that_look_like_json(self):
|
||||
"""Test case: args is a list containing strings that look like JSON objects."""
|
||||
# This could potentially cause issues if not handled correctly
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "get_weather",
|
||||
"args": ['{"x":"x"}', '{"a":"a"}'], # List of JSON strings
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
assert isinstance(arguments, str)
|
||||
# Should serialize list correctly
|
||||
parsed = json.loads(arguments)
|
||||
assert isinstance(parsed, list)
|
||||
assert parsed == ['{"x":"x"}', '{"a":"a"}']
|
||||
|
||||
def test_args_as_dict_with_nested_structures(self):
|
||||
"""Test case: args contains nested dicts and lists."""
|
||||
parts: List[HttpxPartType] = [
|
||||
{
|
||||
"functionCall": {
|
||||
"name": "complex_function",
|
||||
"args": {
|
||||
"nested": {"key": "value"},
|
||||
"list": [1, 2, 3],
|
||||
"mixed": [{"a": 1}, {"b": 2}],
|
||||
},
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
function, tools, idx = VertexGeminiConfig._transform_parts(
|
||||
parts=parts, cumulative_tool_call_idx=0, is_function_call=False
|
||||
)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
arguments = tools[0]["function"]["arguments"]
|
||||
assert isinstance(arguments, str)
|
||||
parsed = json.loads(arguments)
|
||||
assert parsed["nested"] == {"key": "value"}
|
||||
assert parsed["list"] == [1, 2, 3]
|
||||
assert parsed["mixed"] == [{"a": 1}, {"b": 2}]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
@@ -10,11 +10,13 @@ from pydantic import BaseModel
|
||||
|
||||
import litellm
|
||||
from litellm import ModelResponse, completion
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
from litellm.types.llms.vertex_ai import UsageMetadata
|
||||
from litellm.types.utils import ChoiceLogprobs, Usage
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
|
||||
def test_top_logprobs():
|
||||
@@ -1605,6 +1607,39 @@ def test_vertex_ai_annotation_streaming_events():
|
||||
assert "Weather information" in annotation["url_citation"]["title"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_ai_streaming_bad_request_is_not_wrapped():
|
||||
class DummyLogging:
|
||||
def __init__(self):
|
||||
self.model_call_details = {"litellm_params": {}}
|
||||
self.optional_params = {}
|
||||
self.messages = []
|
||||
self.completion_start_time = None
|
||||
self.stream_options = None
|
||||
|
||||
def failure_handler(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
async def async_failure_handler(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
async def failing_make_call(client=None, **kwargs):
|
||||
raise VertexAIError(status_code=400, message="bad input", headers={})
|
||||
|
||||
stream = CustomStreamWrapper(
|
||||
completion_stream=None,
|
||||
make_call=failing_make_call,
|
||||
model="gemini-3-pro-preview",
|
||||
logging_obj=DummyLogging(),
|
||||
custom_llm_provider="vertex_ai_beta",
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
await stream.__anext__()
|
||||
|
||||
assert getattr(exc_info.value, "status_code", None) == 400
|
||||
|
||||
|
||||
def test_vertex_ai_annotation_conversion():
|
||||
"""
|
||||
Test the conversion of Vertex AI grounding metadata to OpenAI annotations.
|
||||
|
||||
Reference in New Issue
Block a user