[Fix] Enhance MidStreamFallbackError to preserve original status code and attributes

- Updated MidStreamFallbackError to retrieve and maintain the original status code from the wrapped exception.
- Ensured that message, request, and response fields remain consistent after calling the parent constructor.
- Added unit tests to verify the correct propagation of status codes and attributes in various scenarios.
This commit is contained in:
Dharamendra Kumar 2026-02-26 13:01:14 -08:00 • committed by Sameer Kankute
parent faf79c5c72
commit 9d1f3ae2eb
2 changed files with 60 additions and 2 deletions

View file

@ -955,7 +955,8 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore
generated_content: str = "",
is_pre_first_chunk: bool = False,
):
self.status_code = 503 # Service Unavailable
original_status = getattr(original_exception, "status_code", None)
self.status_code = int(original_status) if original_status is not None else 503
self.message = f"litellm.MidStreamFallbackError: {message}"
self.model = model
self.llm_provider = llm_provider
@ -978,7 +979,14 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore
else:
self.response = response
# Call the parent constructor
# Save the original attributes before they are overridden by ServiceUnavailableError
_saved_response = self.response
_saved_request = getattr(self.response, "request", None) or httpx.Request(
method="POST", url=f"https://{llm_provider}.com/v1/"
)
_saved_message = self.message
# Call the parent constructor (which hardcodes status_code=503 and modifies the response object)
super().__init__(
message=self.message,
llm_provider=llm_provider,
@ -988,6 +996,14 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore
max_retries=self.max_retries,
num_retries=self.num_retries,
)
# Restore the propagated status and original response/request objects
self.status_code = int(original_status) if original_status is not None else 503
self.response = _saved_response
self.request = _saved_request
self.message = _saved_message
self.args = (_saved_message,)
self.args = (_saved_message,)
def __str__(self):
_message = self.message

View file

@ -1422,4 +1422,46 @@ async def test_exception_bubbling_up(sync_mode, stream_mode, model):
assert exc_info.value.type == "invalid_request_error"
def test_midstream_fallback_error_status_code_propagation():
"""
MidStreamFallbackError should preserve the original status code and keep
message/request/response fields consistent after super().__init__().
"""
import litellm
from httpx import Request, Response
# 1) Wrapping a 429 should preserve the 429 status
original_req = Request("POST", "https://api.openai.com/v1/chat/completions")
original_resp = Response(status_code=429, request=original_req)
rate_limit_error = litellm.RateLimitError(
message="Rate limit exceeded",
llm_provider="openai",
model="gpt-4o-mini",
response=original_resp,
)
midstream_error = litellm.exceptions.MidStreamFallbackError(
message="stream broke",
model="gpt-4o-mini",
llm_provider="openai",
original_exception=rate_limit_error,
)
assert midstream_error.status_code == 429
assert midstream_error.response.status_code == 429
assert str(midstream_error.response.request.url) == "https://openai.com/v1/"
assert midstream_error.message == "litellm.MidStreamFallbackError: stream broke"
assert midstream_error.args == ("litellm.MidStreamFallbackError: stream broke",)
# 2) With no original exception, should default to 503
midstream_fallback = litellm.exceptions.MidStreamFallbackError(
message="stream broke without original",
model="gpt-4o-mini",
llm_provider="openai",
original_exception=None,
)
assert midstream_fallback.status_code == 503
assert midstream_fallback.response.status_code == 503
assert str(midstream_fallback.response.request.url) == "https://openai.com/v1/"