mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
fix(streaming): keep provider headers on MidStreamFallbackError
A 429 raised while opening a stream is wrapped in MidStreamFallbackError, which dropped the litellm_response_headers exception_type() attaches to the wrapped error. retry-after and x-request-id reached callers on the non-streaming path but not on the streaming one, so retry backoff, provider request id logging and proxy error headers lost them.
This commit is contained in:
parent
90e4962c81
commit
abb8aa1672
2 changed files with 54 additions and 0 deletions
|
|
@ -1184,6 +1184,13 @@ class MidStreamFallbackError(ServiceUnavailableError):
|
|||
self.message = _saved_message
|
||||
self.args = (_saved_message,)
|
||||
|
||||
# Keep the provider headers (retry-after, x-request-id, ...) that exception_type()
|
||||
# attached to the wrapped exception, so retry backoff, logging and the proxy error
|
||||
# response still see them once the failure is wrapped for fallback
|
||||
original_headers: Final = getattr(original_exception, "litellm_response_headers", None)
|
||||
if original_headers is not None:
|
||||
self.litellm_response_headers = original_headers
|
||||
|
||||
def __str__(self):
|
||||
_message = self.message
|
||||
if self.num_retries:
|
||||
|
|
|
|||
|
|
@ -849,6 +849,53 @@ def test_sync_streaming_rate_limit_triggers_midstream_fallback(logging_obj: Logg
|
|||
assert excinfo.value.generated_content == ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_rate_limit_midstream_fallback_keeps_provider_headers(
|
||||
logging_obj: Logging,
|
||||
):
|
||||
"""MidStreamFallbackError keeps the provider headers the non-streaming RateLimitError carries.
|
||||
|
||||
retry-after and x-request-id are what retry backoff, logging and the proxy error
|
||||
response read, so wrapping a 429 for fallback must not drop them.
|
||||
"""
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
|
||||
provider_headers = {"retry-after": "3", "x-request-id": "req-abc123"}
|
||||
|
||||
async def _raise_rate_limit(**kwargs):
|
||||
raise VertexAIError(status_code=429, message="Resource exhausted.", headers=provider_headers)
|
||||
|
||||
response = CustomStreamWrapper(
|
||||
completion_stream=None,
|
||||
model="gemini-3-flash-preview",
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="vertex_ai_beta",
|
||||
make_call=_raise_rate_limit,
|
||||
)
|
||||
|
||||
with pytest.raises(MidStreamFallbackError) as excinfo:
|
||||
await response.__anext__()
|
||||
|
||||
assert excinfo.value.status_code == 429
|
||||
assert excinfo.value.litellm_response_headers.get("retry-after") == "3"
|
||||
assert excinfo.value.litellm_response_headers.get("x-request-id") == "req-abc123"
|
||||
|
||||
|
||||
def test_midstream_fallback_error_without_provider_headers_adds_no_attribute():
|
||||
"""No provider headers means no attribute: Router checks hasattr() before response.headers."""
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
midstream_error = MidStreamFallbackError(
|
||||
message="stream broke",
|
||||
model="gpt-4o-mini",
|
||||
llm_provider="openai",
|
||||
original_exception=Exception("boom"),
|
||||
)
|
||||
|
||||
assert not hasattr(midstream_error, "litellm_response_headers")
|
||||
|
||||
|
||||
def test_sync_streaming_bad_request_not_midstream(logging_obj: Logging):
|
||||
"""Ensure __next__ raises BadRequestError (400) directly, not MidStreamFallbackError.
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue