mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
[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:
parent
faf79c5c72
commit
9d1f3ae2eb
2 changed files with 60 additions and 2 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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/"
|
||||
Loading…
Add table
Reference in a new issue