mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(sdk): carry a litellm_proxy error's headers on e.response, not e.headers
A mapped litellm_proxy exception now attaches an httpx.Response that carries the proxy's response headers whenever the handler attached a header-less synthetic one, on every status branch and on the relay path. BadRequestError keeps its base-class contract: .headers stays the proxy-supplied channel, so the proxy edge keeps forwarding an upstream proxy's headers under the llm_provider- prefix and the date and server edge change is no longer needed.
This commit is contained in:
mateo-berri
parent
f4f1e2eace
commit
6a635cbb64
5 changed files with 223 additions and 431 deletions
|
|
@ -1974,11 +1974,7 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
|
|||
}
|
||||
)
|
||||
|
||||
ORIGIN_SERVER_HEADERS: Final[frozenset[str]] = frozenset({"date", "server"})
|
||||
|
||||
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = (
|
||||
HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS | ORIGIN_SERVER_HEADERS
|
||||
)
|
||||
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS
|
||||
|
||||
STRINGIFIED_NONE: Final[str] = "None"
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@
|
|||
## LiteLLM versions of the OpenAI Exception Types
|
||||
|
||||
import enum
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -226,7 +226,6 @@ class BadRequestError(openai.BadRequestError):
|
|||
max_retries: int | None = None,
|
||||
num_retries: int | None = None,
|
||||
body: dict | None = None,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
):
|
||||
self.status_code = 400
|
||||
self.message = f"litellm.BadRequestError: {message}"
|
||||
|
|
@ -235,9 +234,6 @@ class BadRequestError(openai.BadRequestError):
|
|||
self.litellm_debug_info = litellm_debug_info
|
||||
self.max_retries = max_retries
|
||||
self.num_retries = num_retries
|
||||
self.headers = (
|
||||
{k: str(v) for k, v in headers.items()} if headers else None # mutable-ok: the proxy updates it in place
|
||||
)
|
||||
# Use response if it's a valid httpx.Response with a request, otherwise use minimal error response
|
||||
# Note: We check _request (not .request property) to avoid RuntimeError when _request is None
|
||||
if (
|
||||
|
|
@ -628,7 +624,6 @@ class ContentPolicyViolationError(BadRequestError):
|
|||
litellm_debug_info: str | None = None,
|
||||
provider_specific_fields: dict | None = None,
|
||||
body: dict | None = None,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
):
|
||||
self.status_code = 400
|
||||
self.message = f"litellm.ContentPolicyViolationError: {message}"
|
||||
|
|
@ -643,7 +638,6 @@ class ContentPolicyViolationError(BadRequestError):
|
|||
response=response,
|
||||
litellm_debug_info=self.litellm_debug_info,
|
||||
body=body,
|
||||
headers=headers,
|
||||
) # Call the base class constructor with the parameters it needs
|
||||
|
||||
def __str__(self):
|
||||
|
|
|
|||
|
|
@ -216,7 +216,6 @@ def extract_and_raise_litellm_exception(
|
|||
model: str,
|
||||
custom_llm_provider: str,
|
||||
body: object | None = None,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
):
|
||||
"""
|
||||
Covers scenario where litellm sdk calling proxy.
|
||||
|
|
@ -237,9 +236,7 @@ def extract_and_raise_litellm_exception(
|
|||
message=error_str,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
**_accepted_init_kwargs(
|
||||
raised_exception_obj, MappingProxyType({"response": response, "body": body, "headers": headers})
|
||||
),
|
||||
**_accepted_init_kwargs(raised_exception_obj, MappingProxyType({"response": response, "body": body})),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -253,13 +250,20 @@ class _ProviderHTTPException(Protocol):
|
|||
llm_provider: str
|
||||
|
||||
|
||||
def _litellm_proxy_response_headers(
|
||||
def _litellm_proxy_response(
|
||||
original_exception: _ProviderHTTPException, custom_llm_provider: str
|
||||
) -> Mapping[str, str] | None:
|
||||
if custom_llm_provider != "litellm_proxy":
|
||||
return None
|
||||
) -> httpx.Response | None:
|
||||
response: Final = getattr(original_exception, "response", None)
|
||||
if custom_llm_provider != "litellm_proxy" or not isinstance(response, httpx.Response) or response.headers:
|
||||
return response
|
||||
headers: Final = getattr(original_exception, "headers", None)
|
||||
return headers if isinstance(headers, Mapping) else None
|
||||
if not isinstance(headers, Mapping) or not headers:
|
||||
return response
|
||||
return httpx.Response(
|
||||
status_code=response.status_code,
|
||||
headers={str(k): str(v) for k, v in headers.items()},
|
||||
request=getattr(original_exception, "request", None),
|
||||
)
|
||||
|
||||
|
||||
def _map_openai_exception(
|
||||
|
|
@ -272,7 +276,7 @@ def _map_openai_exception(
|
|||
exception_provider: str,
|
||||
extra_information: str,
|
||||
) -> None:
|
||||
upstream_headers: Final = _litellm_proxy_response_headers(original_exception, custom_llm_provider)
|
||||
response: Final = _litellm_proxy_response(original_exception, custom_llm_provider)
|
||||
# custom_llm_provider is openai, make it OpenAI
|
||||
message = get_error_message(error_obj=original_exception)
|
||||
if message is None:
|
||||
|
|
@ -301,14 +305,14 @@ def _map_openai_exception(
|
|||
message=f"RateLimitError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str):
|
||||
raise ContextWindowExceededError(
|
||||
message=f"ContextWindowExceededError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "invalid_request_error" in error_str and "model_not_found" in error_str:
|
||||
|
|
@ -316,7 +320,7 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "A timeout occurred" in error_str:
|
||||
|
|
@ -335,10 +339,9 @@ def _map_openai_exception(
|
|||
message=f"ContentPolicyViolationError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
headers=upstream_headers,
|
||||
)
|
||||
elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str:
|
||||
helpful_message: Final = (
|
||||
|
|
@ -356,20 +359,18 @@ def _map_openai_exception(
|
|||
message=helpful_message,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
headers=upstream_headers,
|
||||
)
|
||||
elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str:
|
||||
raise BadRequestError(
|
||||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
headers=upstream_headers,
|
||||
)
|
||||
elif (
|
||||
"Web server is returning an unknown error" in error_str
|
||||
|
|
@ -385,7 +386,7 @@ def _map_openai_exception(
|
|||
message=f"RateLimitError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif (
|
||||
|
|
@ -396,7 +397,7 @@ def _map_openai_exception(
|
|||
message=f"AuthenticationError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "Mistral API raised a streaming error" in error_str:
|
||||
|
|
@ -415,17 +416,16 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
headers=upstream_headers,
|
||||
)
|
||||
elif original_exception.status_code == 401:
|
||||
raise AuthenticationError(
|
||||
message=f"AuthenticationError: {exception_provider} - {message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 404:
|
||||
|
|
@ -433,7 +433,7 @@ def _map_openai_exception(
|
|||
message=f"NotFoundError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 408:
|
||||
|
|
@ -448,17 +448,16 @@ def _map_openai_exception(
|
|||
message=f"{exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
body=getattr(original_exception, "body", None),
|
||||
headers=upstream_headers,
|
||||
)
|
||||
elif original_exception.status_code == 429:
|
||||
raise RateLimitError(
|
||||
message=f"RateLimitError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 500:
|
||||
|
|
@ -466,7 +465,7 @@ def _map_openai_exception(
|
|||
message=f"InternalServerError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 502:
|
||||
|
|
@ -474,7 +473,7 @@ def _map_openai_exception(
|
|||
message=f"BadGatewayError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 503:
|
||||
|
|
@ -482,7 +481,7 @@ def _map_openai_exception(
|
|||
message=f"ServiceUnavailableError: {exception_provider} - {message}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif original_exception.status_code == 504: # gateway timeout error
|
||||
|
|
@ -2439,12 +2438,11 @@ def exception_type(
|
|||
custom_llm_provider == "litellm_proxy"
|
||||
): # handle special case where calling litellm proxy + exception str contains error message
|
||||
extract_and_raise_litellm_exception(
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=_litellm_proxy_response(mappable_exception, custom_llm_provider),
|
||||
error_str=error_str,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
body=getattr(original_exception, "body", None),
|
||||
headers=_litellm_proxy_response_headers(mappable_exception, custom_llm_provider),
|
||||
)
|
||||
if (
|
||||
custom_llm_provider == "openai"
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
|
|
@ -178,9 +177,7 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in error_strings:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is True, f"Should detect policy violation in: {error_str}"
|
||||
|
||||
def test_is_azure_content_policy_violation_error_case_insensitive(self):
|
||||
|
|
@ -194,12 +191,8 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in error_strings:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
assert (
|
||||
result is True
|
||||
), f"Should detect policy violation in uppercase: {error_str}"
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is True, f"Should detect policy violation in uppercase: {error_str}"
|
||||
|
||||
def test_is_azure_content_policy_violation_error_with_non_policy_errors(self):
|
||||
"""Test that non-policy violation errors are not detected as policy violations"""
|
||||
|
|
@ -216,12 +209,8 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in error_strings:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
assert (
|
||||
result is False
|
||||
), f"Should NOT detect policy violation in: {error_str}"
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is False, f"Should NOT detect policy violation in: {error_str}"
|
||||
|
||||
def test_is_azure_content_policy_violation_error_with_partial_matches(self):
|
||||
"""Test that partial keyword matches work correctly"""
|
||||
|
|
@ -234,9 +223,7 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in positive_cases:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is True, f"Should detect policy violation in: {error_str}"
|
||||
|
||||
# These should not match even though they contain similar words
|
||||
|
|
@ -248,12 +235,8 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in negative_cases:
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
assert (
|
||||
result is False
|
||||
), f"Should NOT detect policy violation in: {error_str}"
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
|
||||
assert result is False, f"Should NOT detect policy violation in: {error_str}"
|
||||
|
||||
|
||||
gemini_context_window_test_cases = [
|
||||
|
|
@ -271,12 +254,8 @@ gemini_context_window_test_cases = [
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error_message, should_raise_context_window", gemini_context_window_test_cases
|
||||
)
|
||||
def test_gemini_context_window_error_mapping(
|
||||
error_message, should_raise_context_window
|
||||
):
|
||||
@pytest.mark.parametrize("error_message, should_raise_context_window", gemini_context_window_test_cases)
|
||||
def test_gemini_context_window_error_mapping(error_message, should_raise_context_window):
|
||||
"""
|
||||
Tests that the exception_type function correctly maps Gemini's
|
||||
context window exceeded errors to litellm.ContextWindowExceededError.
|
||||
|
|
@ -421,9 +400,7 @@ vertex_rate_limit_test_cases = [
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error_message, should_raise_rate_limit", vertex_rate_limit_test_cases
|
||||
)
|
||||
@pytest.mark.parametrize("error_message, should_raise_rate_limit", vertex_rate_limit_test_cases)
|
||||
def test_vertex_ai_rate_limit_error_mapping(error_message, should_raise_rate_limit):
|
||||
"""
|
||||
Tests that the exception_type function correctly maps Vertex AI's
|
||||
|
|
@ -458,10 +435,7 @@ class TestGetBodyErrorCode:
|
|||
"""Unit tests for _get_body_error_code helper."""
|
||||
|
||||
def test_parses_int_code(self):
|
||||
body = (
|
||||
'{"error":{"message":"high demand","type":"upstream_error",'
|
||||
'"param":"","code":429}}'
|
||||
)
|
||||
body = '{"error":{"message":"high demand","type":"upstream_error","param":"","code":429}}'
|
||||
assert _get_body_error_code(body) == 429
|
||||
|
||||
def test_parses_string_code(self):
|
||||
|
|
@ -498,8 +472,7 @@ gemini_body_code_429_test_cases = [
|
|||
),
|
||||
(
|
||||
503,
|
||||
'{"error":{"message":"upstream unavailable","type":"upstream_error",'
|
||||
'"param":"","code":429}}',
|
||||
'{"error":{"message":"upstream unavailable","type":"upstream_error","param":"","code":429}}',
|
||||
litellm.RateLimitError,
|
||||
"HTTP 503 envelope with body code:429 -> RateLimitError",
|
||||
),
|
||||
|
|
@ -769,9 +742,7 @@ class _UpstreamHTTPError(Exception):
|
|||
self.message = "upstream failure"
|
||||
self.status_code = status_code
|
||||
self.request = httpx.Request("POST", "https://api.example.com/v1/chat/completions")
|
||||
self.response = httpx.Response(
|
||||
status_code=status_code, request=self.request, text="upstream failure"
|
||||
)
|
||||
self.response = httpx.Response(status_code=status_code, request=self.request, text="upstream failure")
|
||||
|
||||
|
||||
UPSTREAM_STATUS_CODES = (400, 401, 403, 404, 408, 422, 429, 500, 503)
|
||||
|
|
@ -892,15 +863,13 @@ PROVIDERS_WITHOUT_A_HANDLER = tuple(
|
|||
|
||||
MINIMAX_401_BODY = (
|
||||
'{"type":"error","error":{"type":"authorized_error","message":"login fail: Please carry the API secret key '
|
||||
"in the 'Authorization' field of the request header (1004)\",\"http_code\":\"401\"},"
|
||||
'in the \'Authorization\' field of the request header (1004)","http_code":"401"},'
|
||||
'"request_id":"06ddc9ba97ee6340e38f10e09787f547"}'
|
||||
)
|
||||
|
||||
|
||||
def _expected_for(provider: str, status_code: int) -> tuple[type[Exception], int]:
|
||||
return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get(
|
||||
status_code, OPENAI_SHAPED[status_code]
|
||||
)
|
||||
return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get(status_code, OPENAI_SHAPED[status_code])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -910,9 +879,7 @@ def quiet_exception_mapping(monkeypatch):
|
|||
|
||||
@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES)
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_an_upstream_status_maps_to_one_exception_per_provider(
|
||||
provider, status_code, quiet_exception_mapping
|
||||
):
|
||||
def test_an_upstream_status_maps_to_one_exception_per_provider(provider, status_code, quiet_exception_mapping):
|
||||
expected_class, expected_status = _expected_for(provider, status_code)
|
||||
|
||||
with pytest.raises(openai.APIError) as raised:
|
||||
|
|
@ -928,9 +895,7 @@ def test_an_upstream_status_maps_to_one_exception_per_provider(
|
|||
|
||||
@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES)
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(
|
||||
provider, status_code, quiet_exception_mapping
|
||||
):
|
||||
def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(provider, status_code, quiet_exception_mapping):
|
||||
with pytest.raises(openai.APIError) as raised:
|
||||
exception_type(
|
||||
model="test-model",
|
||||
|
|
@ -943,12 +908,8 @@ def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_an_already_mapped_litellm_exception_passes_through_untouched(
|
||||
provider, quiet_exception_mapping
|
||||
):
|
||||
already_mapped = litellm.RateLimitError(
|
||||
message="already mapped", llm_provider=provider, model="test-model"
|
||||
)
|
||||
def test_an_already_mapped_litellm_exception_passes_through_untouched(provider, quiet_exception_mapping):
|
||||
already_mapped = litellm.RateLimitError(message="already mapped", llm_provider=provider, model="test-model")
|
||||
|
||||
returned = exception_type(
|
||||
model="test-model",
|
||||
|
|
@ -961,9 +922,7 @@ def test_an_already_mapped_litellm_exception_passes_through_untouched(
|
|||
|
||||
@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES)
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITHOUT_A_HANDLER)
|
||||
def test_a_provider_without_a_handler_maps_by_the_upstream_status(
|
||||
provider, status_code, quiet_exception_mapping
|
||||
):
|
||||
def test_a_provider_without_a_handler_maps_by_the_upstream_status(provider, status_code, quiet_exception_mapping):
|
||||
expected_class, expected_status = STATUS_KEYED[status_code]
|
||||
|
||||
with pytest.raises(openai.APIError) as raised:
|
||||
|
|
@ -1015,9 +974,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_is_a_connection_error(q
|
|||
assert "boom" in raised.value.message
|
||||
|
||||
|
||||
def _raise_and_map(
|
||||
model: str | None, original_exception: Exception, custom_llm_provider: str | None
|
||||
) -> None:
|
||||
def _raise_and_map(model: str | None, original_exception: Exception, custom_llm_provider: str | None) -> None:
|
||||
"""Calls exception_type() from inside the except block, as litellm/main.py does,
|
||||
so traceback.format_exc() has a real stack."""
|
||||
try:
|
||||
|
|
@ -1058,9 +1015,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_message_keeps_traceback
|
|||
|
||||
|
||||
CONTEXT_WINDOW_MESSAGE = "This model's maximum context length is 4096 tokens."
|
||||
CONTENT_POLICY_MESSAGE = (
|
||||
'{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}'
|
||||
)
|
||||
CONTENT_POLICY_MESSAGE = '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}'
|
||||
TIMEOUT_MESSAGE = "Request timed out."
|
||||
|
||||
PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW = (
|
||||
|
|
@ -1103,15 +1058,11 @@ class _UpstreamErrorWithMessage(_UpstreamHTTPError):
|
|||
super().__init__(status_code=status_code)
|
||||
self.args = (message,)
|
||||
self.message = message
|
||||
self.response = httpx.Response(
|
||||
status_code=status_code, request=self.request, text=message
|
||||
)
|
||||
self.response = httpx.Response(status_code=status_code, request=self.request, text=message)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(
|
||||
provider, quiet_exception_mapping
|
||||
):
|
||||
def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping):
|
||||
if provider in PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW:
|
||||
expected_class, expected_status = litellm.ContextWindowExceededError, 400
|
||||
else:
|
||||
|
|
@ -1129,9 +1080,7 @@ def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(
|
||||
provider, quiet_exception_mapping
|
||||
):
|
||||
def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping):
|
||||
if provider in PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK:
|
||||
expected_class, expected_status = litellm.ContentPolicyViolationError, 400
|
||||
else:
|
||||
|
|
@ -1149,9 +1098,7 @@ def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
|
||||
def test_a_timed_out_request_is_a_timeout_for_every_provider(
|
||||
provider, quiet_exception_mapping
|
||||
):
|
||||
def test_a_timed_out_request_is_a_timeout_for_every_provider(provider, quiet_exception_mapping):
|
||||
with pytest.raises(litellm.Timeout) as raised:
|
||||
exception_type(
|
||||
model="test-model",
|
||||
|
|
@ -1442,9 +1389,7 @@ def _openai_handler_error(
|
|||
_PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)]
|
||||
)
|
||||
@pytest.mark.parametrize(("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)])
|
||||
def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int):
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
exception_type(
|
||||
|
|
@ -1457,12 +1402,10 @@ def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, s
|
|||
|
||||
assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project"
|
||||
assert exc_info.value.body["type"] == error_type
|
||||
assert exc_info.value.headers == _PROXY_HEADERS
|
||||
assert dict(exc_info.value.response.headers) == _PROXY_HEADERS
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError]
|
||||
)
|
||||
@pytest.mark.parametrize("relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError])
|
||||
def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]):
|
||||
message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}"
|
||||
|
||||
|
|
@ -1477,7 +1420,7 @@ def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_clas
|
|||
|
||||
assert type(exc_info.value) is relayed_class
|
||||
assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project"
|
||||
assert exc_info.value.headers == _PROXY_HEADERS
|
||||
assert dict(exc_info.value.response.headers) == _PROXY_HEADERS
|
||||
|
||||
|
||||
def test_openai_compatible_vendor_400_keeps_body_but_not_headers():
|
||||
|
|
@ -1491,4 +1434,4 @@ def test_openai_compatible_vendor_400_keeps_body_but_not_headers():
|
|||
)
|
||||
|
||||
assert exc_info.value.body["type"] == "vendor_specific_error"
|
||||
assert exc_info.value.headers is None
|
||||
assert not exc_info.value.response.headers
|
||||
|
|
|
|||
|
|
@ -127,16 +127,12 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
assert json.loads(result.body) == guardrailed_body
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch):
|
||||
"""The guardrail JSON path must forward upstream response headers (e.g.
|
||||
x-amzn-requestid) alongside the x-litellm-* headers, matching the
|
||||
non-guardrail passthrough path, while dropping length headers that no
|
||||
longer match the rewritten body."""
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(
|
||||
data={"custom_llm_provider": "bedrock"}
|
||||
)
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
|
|
@ -176,14 +172,10 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
assert result.headers["content-length"] == str(len(result.body))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch):
|
||||
"""The guardrail event-stream branch must also forward upstream response
|
||||
headers alongside the x-litellm-* headers."""
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(
|
||||
data={"custom_llm_provider": "bedrock"}
|
||||
)
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
|
|
@ -225,15 +217,11 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
assert result.headers["x-litellm-call-id"] == "test-call-id"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(self, monkeypatch):
|
||||
"""Guardrailed non-streaming passthrough responses must include headers
|
||||
injected by post_call_response_headers_hook, matching the headers a
|
||||
non-guardrailed passthrough response would carry."""
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(
|
||||
data={"custom_llm_provider": "bedrock"}
|
||||
)
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"})
|
||||
monkeypatch.setattr(
|
||||
processing_obj,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
|
|
@ -252,9 +240,7 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
return kwargs["response"]
|
||||
|
||||
proxy_logging_obj.post_call_success_hook = fake_post_call_success_hook
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value={"x-litellm-custom": "from-hook"}
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-litellm-custom": "from-hook"})
|
||||
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=upstream,
|
||||
|
|
@ -378,9 +364,7 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
json.dumps(persisted_body)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(self, monkeypatch):
|
||||
"""arm_pre_call must run before pre_call_hook: an auto router's own compression
|
||||
policy has to be in `data["metadata"]` (naming the model-side guardrail so it
|
||||
runs even if it isn't default_on) by the time guardrails see the request."""
|
||||
|
|
@ -2191,16 +2175,10 @@ class TestCommonRequestProcessingHelpers:
|
|||
|
||||
def _stringified_none_paths(node: object, path: str = "error") -> tuple[str, ...]:
|
||||
if isinstance(node, dict):
|
||||
return tuple(
|
||||
found
|
||||
for key, value in node.items()
|
||||
for found in _stringified_none_paths(value, f"{path}.{key}")
|
||||
)
|
||||
return tuple(found for key, value in node.items() for found in _stringified_none_paths(value, f"{path}.{key}"))
|
||||
if isinstance(node, (list, tuple)):
|
||||
return tuple(
|
||||
found
|
||||
for index, value in enumerate(node)
|
||||
for found in _stringified_none_paths(value, f"{path}[{index}]")
|
||||
found for index, value in enumerate(node) for found in _stringified_none_paths(value, f"{path}[{index}]")
|
||||
)
|
||||
return (path,) if node == "None" else ()
|
||||
|
||||
|
|
@ -2947,9 +2925,7 @@ class TestStreamingOverheadHeader:
|
|||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_id="test-call-id",
|
||||
hidden_params={},
|
||||
litellm_logging_obj=self._timing_logging_obj(
|
||||
{"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}
|
||||
),
|
||||
litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}),
|
||||
)
|
||||
|
||||
assert headers["x-litellm-response-duration-ms"] == "500.0"
|
||||
|
|
@ -2968,9 +2944,7 @@ class TestStreamingOverheadHeader:
|
|||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_id="test-call-id",
|
||||
hidden_params={},
|
||||
litellm_logging_obj=self._timing_logging_obj(
|
||||
{"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}
|
||||
),
|
||||
litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}),
|
||||
read_timing_from_logging_obj=False,
|
||||
)
|
||||
|
||||
|
|
@ -2991,9 +2965,7 @@ class TestStreamingOverheadHeader:
|
|||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_id="test-call-id",
|
||||
hidden_params={"_response_ms": 300.0},
|
||||
litellm_logging_obj=self._timing_logging_obj(
|
||||
{"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}
|
||||
),
|
||||
litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}),
|
||||
)
|
||||
|
||||
assert headers["x-litellm-response-duration-ms"] == "300.0"
|
||||
|
|
@ -3034,9 +3006,7 @@ class TestStreamingOverheadHeader:
|
|||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_id="test-call-id",
|
||||
hidden_params={"_response_ms": 300.0, "litellm_overhead_time_ms": 7.5},
|
||||
litellm_logging_obj=self._timing_logging_obj(
|
||||
{"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}
|
||||
),
|
||||
litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}),
|
||||
)
|
||||
|
||||
assert headers["x-litellm-response-duration-ms"] == "300.0"
|
||||
|
|
@ -3488,9 +3458,7 @@ class TestStreamCloseOnDisconnect:
|
|||
finally:
|
||||
closed.set()
|
||||
|
||||
response = _UpstreamClosingStreamingResponse(
|
||||
body(), media_type="text/event-stream"
|
||||
)
|
||||
response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream")
|
||||
|
||||
async def receive():
|
||||
await asyncio.Event().wait()
|
||||
|
|
@ -3521,9 +3489,7 @@ class TestStreamCloseOnDisconnect:
|
|||
finally:
|
||||
closed.set()
|
||||
|
||||
response = _UpstreamClosingStreamingResponse(
|
||||
body(), media_type="text/event-stream"
|
||||
)
|
||||
response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream")
|
||||
|
||||
async def receive():
|
||||
await disconnected.wait()
|
||||
|
|
@ -3594,9 +3560,7 @@ class TestStreamCloseOnDisconnect:
|
|||
finally:
|
||||
inner_closed.set()
|
||||
|
||||
response = await create_response(
|
||||
generator=wrapped(), media_type="text/event-stream", headers={}
|
||||
)
|
||||
response = await create_response(generator=wrapped(), media_type="text/event-stream", headers={})
|
||||
|
||||
async def receive():
|
||||
await asyncio.Event().wait()
|
||||
|
|
@ -3828,9 +3792,7 @@ class TestStreamCloseOnDisconnect:
|
|||
|
||||
with pytest.raises(_ClientDisconnectedBeforeFirstChunk):
|
||||
await asyncio.wait_for(
|
||||
_buffer_first_chunk_honoring_disconnect(
|
||||
AcloseRaises(), request=self._request_that_disconnects()
|
||||
),
|
||||
_buffer_first_chunk_honoring_disconnect(AcloseRaises(), request=self._request_that_disconnects()),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
|
|
@ -3846,9 +3808,7 @@ class TestStreamCloseOnDisconnect:
|
|||
|
||||
with pytest.raises(_ClientDisconnectedBeforeFirstChunk):
|
||||
await asyncio.wait_for(
|
||||
_buffer_first_chunk_honoring_disconnect(
|
||||
blocking_gen(), request=self._request_that_disconnects()
|
||||
),
|
||||
_buffer_first_chunk_honoring_disconnect(blocking_gen(), request=self._request_that_disconnects()),
|
||||
timeout=5,
|
||||
)
|
||||
assert closed.is_set()
|
||||
|
|
@ -3864,9 +3824,7 @@ class TestHandleLLMApiExceptionRetryAfter:
|
|||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test")
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value=callback_headers or {}
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {})
|
||||
|
||||
try:
|
||||
await processor._handle_llm_api_exception(
|
||||
|
|
@ -3918,9 +3876,7 @@ class TestHandleLLMApiExceptionRetryAfter:
|
|||
enable_pre_call_checks=False,
|
||||
cooldown_list=[],
|
||||
)
|
||||
proxy_exc = await self._invoke(
|
||||
exc, callback_headers={"retry-after": "", "x-custom": "1"}
|
||||
)
|
||||
proxy_exc = await self._invoke(exc, callback_headers={"retry-after": "", "x-custom": "1"})
|
||||
assert proxy_exc.headers["retry-after"] == "43"
|
||||
assert proxy_exc.headers["x-custom"] == "1"
|
||||
|
||||
|
|
@ -4063,18 +4019,6 @@ class TestHandleLLMApiExceptionFramingHeaders:
|
|||
assert proxy_exc.headers["x-custom-safe"] == "1"
|
||||
assert proxy_exc.headers["x-request-id"] == "abc-123"
|
||||
|
||||
async def test_strips_the_date_and_server_headers_of_an_upstream_litellm_proxy(self):
|
||||
exc = litellm.BadRequestError(
|
||||
message="Content blocked",
|
||||
llm_provider="litellm_proxy",
|
||||
model="claude-haiku-4-5",
|
||||
headers={"date": "Sun, 13 Sep 2026 08:43:51 GMT", "server": "uvicorn", "x-request-id": "abc-123"},
|
||||
)
|
||||
proxy_exc = await self._invoke(exc)
|
||||
assert "date" not in proxy_exc.headers
|
||||
assert "server" not in proxy_exc.headers
|
||||
assert proxy_exc.headers["x-request-id"] == "abc-123"
|
||||
|
||||
|
||||
class TestAsyncStreamingDataGeneratorFastPath:
|
||||
"""Fast/slow path branching in async_streaming_data_generator."""
|
||||
|
|
@ -4161,9 +4105,7 @@ class TestDisconnectGatherCleanup:
|
|||
return Request(scope={"type": "http", "headers": []}, receive=receive)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_process_llm_request_raises_499_on_client_disconnect(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_base_process_llm_request_raises_499_on_client_disconnect(self, monkeypatch):
|
||||
"""With cancel_on_disconnect enabled, base_process_llm_request returns 499."""
|
||||
import asyncio
|
||||
|
||||
|
|
@ -4192,9 +4134,7 @@ class TestDisconnectGatherCleanup:
|
|||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await processing_obj.base_process_llm_request(
|
||||
|
|
@ -4212,9 +4152,7 @@ class TestDisconnectGatherCleanup:
|
|||
assert "disconnected" in exc_info.value.detail.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(self, monkeypatch):
|
||||
import asyncio
|
||||
|
||||
import litellm.proxy.common_request_processing as cpr
|
||||
|
|
@ -4239,9 +4177,7 @@ class TestDisconnectGatherCleanup:
|
|||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
|
||||
monkeypatch.setattr(
|
||||
cpr,
|
||||
"route_request",
|
||||
|
|
@ -4302,9 +4238,7 @@ class TestDisconnectGatherCleanup:
|
|||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await processing_obj.base_process_llm_request(
|
||||
|
|
@ -4355,9 +4289,7 @@ class TestDisconnectGatherCleanup:
|
|||
assert task.done()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_process_llm_request_preserves_llm_error_after_gather(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_base_process_llm_request_preserves_llm_error_after_gather(self, monkeypatch):
|
||||
import litellm.proxy.common_request_processing as cpr
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
|
|
@ -4386,9 +4318,7 @@ class TestDisconnectGatherCleanup:
|
|||
"common_processing_pre_call_logic",
|
||||
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
|
||||
)
|
||||
monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.is_disconnected = AsyncMock(return_value=False)
|
||||
|
|
@ -4425,19 +4355,13 @@ class TestStreamingClientDisconnectLogging:
|
|||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
|
||||
|
||||
assert recorded is True
|
||||
assert request_data["metadata"]["client_disconnected"] is True
|
||||
assert request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
assert (
|
||||
request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
)
|
||||
assert (
|
||||
mock_logging_obj.model_call_details["litellm_params"]["metadata"][
|
||||
"error_information"
|
||||
]["error_code"]
|
||||
mock_logging_obj.model_call_details["litellm_params"]["metadata"]["error_information"]["error_code"]
|
||||
== "499"
|
||||
)
|
||||
|
||||
|
|
@ -4451,9 +4375,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
mock_request.is_disconnected = AsyncMock(return_value=False)
|
||||
request_data = {"metadata": {}}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
|
||||
|
||||
assert recorded is False
|
||||
assert "client_disconnected" not in request_data["metadata"]
|
||||
|
|
@ -4478,22 +4400,12 @@ class TestStreamingClientDisconnectLogging:
|
|||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
|
||||
|
||||
assert recorded is True
|
||||
assert request_data["metadata"]["client_disconnected"] is True
|
||||
assert (
|
||||
mock_logging_obj.model_call_details["litellm_params"]["metadata"][
|
||||
"client_disconnected"
|
||||
]
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
mock_logging_obj.model_call_details["metadata"]["client_disconnected"]
|
||||
is True
|
||||
)
|
||||
assert mock_logging_obj.model_call_details["litellm_params"]["metadata"]["client_disconnected"] is True
|
||||
assert mock_logging_obj.model_call_details["metadata"]["client_disconnected"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_streaming_client_disconnect_handles_none_request_data_metadata(self):
|
||||
|
|
@ -4509,15 +4421,11 @@ class TestStreamingClientDisconnectLogging:
|
|||
"litellm_params": {"metadata": None},
|
||||
}
|
||||
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(
|
||||
mock_request, request_data
|
||||
)
|
||||
recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
|
||||
|
||||
assert recorded is True
|
||||
assert request_data["metadata"]["client_disconnected"] is True
|
||||
assert (
|
||||
request_data["litellm_params"]["metadata"]["client_disconnected"] is True
|
||||
)
|
||||
assert request_data["litellm_params"]["metadata"]["client_disconnected"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_client_disconnect_metadata_none_returns_early(self):
|
||||
|
|
@ -4528,9 +4436,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
_apply_client_disconnect_metadata(None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(self, monkeypatch):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
|
@ -4562,9 +4468,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
assert request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(self, monkeypatch):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
|
@ -4594,9 +4498,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
assert "client_disconnected" not in request_data["metadata"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_streaming_data_generator_records_499_on_early_aclose(
|
||||
self, monkeypatch
|
||||
):
|
||||
async def test_async_streaming_data_generator_records_499_on_early_aclose(self, monkeypatch):
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
)
|
||||
|
|
@ -4611,9 +4513,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
yield {"choices": [{"delta": {"content": " there"}}]}
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.async_post_call_streaming_iterator_hook = (
|
||||
mock_streaming_iterator
|
||||
)
|
||||
mock_proxy_logging.async_post_call_streaming_iterator_hook = mock_streaming_iterator
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
|
|
@ -4624,9 +4524,7 @@ class TestStreamingClientDisconnectLogging:
|
|||
"model": "gemini-2.0-flash",
|
||||
"metadata": {},
|
||||
"litellm_params": {"metadata": {}},
|
||||
"litellm_logging_obj": MagicMock(
|
||||
model_call_details={"metadata": {}, "litellm_params": {}}
|
||||
),
|
||||
"litellm_logging_obj": MagicMock(model_call_details={"metadata": {}, "litellm_params": {}}),
|
||||
}
|
||||
|
||||
gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
|
|
@ -4645,6 +4543,8 @@ class TestStreamingClientDisconnectLogging:
|
|||
assert request_data["metadata"]["error_information"]["error_code"] == "499"
|
||||
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
|
||||
|
||||
class TestCancelOnDisconnect:
|
||||
"""
|
||||
Coverage for the opt-in `general_settings.cancel_on_disconnect` flag:
|
||||
|
|
@ -4671,23 +4571,17 @@ class TestCancelOnDisconnect:
|
|||
llm_call = asyncio.get_running_loop().create_future()
|
||||
disconnect_event = asyncio.Event()
|
||||
|
||||
await _cancel_llm_call_on_client_disconnect(
|
||||
request, llm_call, disconnect_event
|
||||
)
|
||||
await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
|
||||
|
||||
assert llm_call.cancelled()
|
||||
assert disconnect_event.is_set()
|
||||
|
||||
async def test_monitor_is_noop_while_client_stays_connected(self):
|
||||
request = self._request(
|
||||
[{"type": "http.request", "body": b"", "more_body": False}]
|
||||
)
|
||||
request = self._request([{"type": "http.request", "body": b"", "more_body": False}])
|
||||
llm_call = asyncio.get_running_loop().create_future()
|
||||
disconnect_event = asyncio.Event()
|
||||
|
||||
monitor = asyncio.create_task(
|
||||
_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
|
||||
)
|
||||
monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event))
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert not monitor.done()
|
||||
|
|
@ -4706,9 +4600,7 @@ class TestCancelOnDisconnect:
|
|||
llm_call = asyncio.get_running_loop().create_future()
|
||||
disconnect_event = asyncio.Event()
|
||||
|
||||
await _cancel_llm_call_on_client_disconnect(
|
||||
request, llm_call, disconnect_event
|
||||
)
|
||||
await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
|
||||
|
||||
assert not llm_call.cancelled()
|
||||
assert not disconnect_event.is_set()
|
||||
|
|
@ -4723,9 +4615,7 @@ class TestCancelOnDisconnect:
|
|||
with pytest.raises(asyncio.CancelledError):
|
||||
await _await_llm_call_cancelling_on_disconnect(request, llm_call)
|
||||
|
||||
async def _drive_base_process_llm_request(
|
||||
self, monkeypatch, general_settings: dict, llm_call, request: Request
|
||||
):
|
||||
async def _drive_base_process_llm_request(self, monkeypatch, general_settings: dict, llm_call, request: Request):
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
|
@ -4734,9 +4624,7 @@ class TestCancelOnDisconnect:
|
|||
logging_obj._on_deferred_stream_complete = None
|
||||
logging_obj.cost_breakdown = None
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(
|
||||
data={"model": "fake-model", "litellm_logging_obj": logging_obj}
|
||||
)
|
||||
processor = ProxyBaseLLMRequestProcessing(data={"model": "fake-model", "litellm_logging_obj": logging_obj})
|
||||
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
|
||||
|
|
@ -4744,9 +4632,7 @@ class TestCancelOnDisconnect:
|
|||
proxy_logging_obj.post_call_success_hook = AsyncMock(
|
||||
side_effect=lambda data, user_api_key_dict, response: response
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
|
||||
async def fake_route_request(**kwargs):
|
||||
return llm_call()
|
||||
|
|
@ -4825,9 +4711,7 @@ class TestCancelOnDisconnect:
|
|||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await processor._handle_llm_api_exception(
|
||||
e=HTTPException(
|
||||
status_code=499, detail="Client disconnected the request"
|
||||
),
|
||||
e=HTTPException(status_code=499, detail="Client disconnected the request"),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
|
@ -4893,7 +4777,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
|
|||
proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock())
|
||||
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", capture_hook)
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=httpx_response,
|
||||
|
|
@ -4943,7 +4829,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
|
|||
proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock())
|
||||
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", non_dict_hook)
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=httpx_response,
|
||||
|
|
@ -4981,7 +4869,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
|
|||
hook_spy = AsyncMock()
|
||||
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy)
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=httpx_response,
|
||||
|
|
@ -5022,7 +4912,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
|
|||
hook_spy = AsyncMock()
|
||||
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy)
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=httpx_response,
|
||||
|
|
@ -5134,7 +5026,9 @@ class TestEventStreamAllmPassthroughRoute:
|
|||
"content-length": "99",
|
||||
}
|
||||
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
|
||||
):
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={})
|
||||
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
|
||||
response=mock_response,
|
||||
|
|
@ -5165,9 +5059,7 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
de-anonymized.
|
||||
"""
|
||||
|
||||
def _build_processing_obj(
|
||||
self, custom_llm_provider: str, endpoint: str = ""
|
||||
) -> ProxyBaseLLMRequestProcessing:
|
||||
def _build_processing_obj(self, custom_llm_provider: str, endpoint: str = "") -> ProxyBaseLLMRequestProcessing:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "call-123"
|
||||
logging_obj.cost_breakdown = None
|
||||
|
|
@ -5254,14 +5146,17 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
processing_obj = self._build_processing_obj("anthropic")
|
||||
chunks = [b"chunk-1", b"chunk-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
),
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
|
|
@ -5270,27 +5165,27 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
assert streamed == chunks
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_converse_stream_is_buffered_through_handler(
|
||||
self, monkeypatch
|
||||
):
|
||||
processing_obj = self._build_processing_obj(
|
||||
"bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream"
|
||||
)
|
||||
async def test_bedrock_converse_stream_is_buffered_through_handler(self, monkeypatch):
|
||||
processing_obj = self._build_processing_obj("bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream")
|
||||
chunks = [b"raw-1", b"raw-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
), patch(
|
||||
"litellm.llms.bedrock.passthrough.guardrail_translation.handler."
|
||||
"BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
|
||||
new=AsyncMock(return_value=b"modified-body"),
|
||||
) as mock_handler:
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
),
|
||||
patch(
|
||||
"litellm.llms.bedrock.passthrough.guardrail_translation.handler."
|
||||
"BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
|
||||
new=AsyncMock(return_value=b"modified-body"),
|
||||
) as mock_handler,
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
assert isinstance(result, Response)
|
||||
|
|
@ -5306,19 +5201,23 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
)
|
||||
chunks = [b"raw-1", b"raw-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
), patch(
|
||||
"litellm.llms.bedrock.passthrough.guardrail_translation.handler."
|
||||
"BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
|
||||
new=AsyncMock(return_value=b"modified-body"),
|
||||
) as mock_handler:
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=True,
|
||||
),
|
||||
patch(
|
||||
"litellm.llms.bedrock.passthrough.guardrail_translation.handler."
|
||||
"BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
|
||||
new=AsyncMock(return_value=b"modified-body"),
|
||||
) as mock_handler,
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
assert isinstance(result, StreamingResponse)
|
||||
|
|
@ -5340,14 +5239,17 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
)
|
||||
chunks = [b"raw-1", b"raw-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=False,
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=False,
|
||||
),
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
|
|
@ -5366,14 +5268,17 @@ class TestAllmPassthroughStreamingProviderGate:
|
|||
processing_obj = self._build_processing_obj("anthropic")
|
||||
chunks = [b"chunk-1", b"chunk-2"]
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=False,
|
||||
with (
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
"_has_post_call_guardrails_for_passthrough",
|
||||
return_value=False,
|
||||
),
|
||||
):
|
||||
result = await self._run(processing_obj, monkeypatch, chunks)
|
||||
|
||||
|
|
@ -5821,9 +5726,7 @@ class TestCostHeadersForCallsPricedAtZero:
|
|||
fastapi_response = Response()
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj})
|
||||
|
||||
with patch.object(
|
||||
ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False
|
||||
):
|
||||
with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False):
|
||||
await processing_obj.base_process_llm_request(
|
||||
request=MagicMock(spec=Request, headers={}),
|
||||
fastapi_response=fastapi_response,
|
||||
|
|
@ -5894,9 +5797,7 @@ class TestCostHeadersForCallsPricedAtZero:
|
|||
assert breakdown.tool_usage_cost == 0.0
|
||||
|
||||
def test_cost_breakdown_stays_empty_for_an_inference_call(self):
|
||||
breakdown = _get_cost_breakdown_from_logging_obj(
|
||||
litellm_logging_obj=self._logging_obj(call_type="acompletion")
|
||||
)
|
||||
breakdown = _get_cost_breakdown_from_logging_obj(litellm_logging_obj=self._logging_obj(call_type="acompletion"))
|
||||
|
||||
assert breakdown == CostBreakdownHeaderValues()
|
||||
|
||||
|
|
@ -5923,7 +5824,6 @@ class TestCostHeadersForCallsPricedAtZero:
|
|||
|
||||
|
||||
class TestPreCallWithFallbacksOnLocalRateLimit:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fallback_triggered_on_local_rate_limit(self):
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
|
|
@ -6075,9 +5975,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}]
|
||||
|
||||
user_api_key_dict = MagicMock()
|
||||
user_api_key_dict.router_settings = {
|
||||
"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]
|
||||
}
|
||||
user_api_key_dict.router_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]}
|
||||
|
||||
with patch.object(
|
||||
processor,
|
||||
|
|
@ -6108,9 +6006,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(
|
||||
data={"model": "gpt-4", "disable_fallbacks": True}
|
||||
)
|
||||
processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "disable_fallbacks": True})
|
||||
|
||||
async def mock_pre_call_logic(**kwargs):
|
||||
raise ProxyRateLimitError(
|
||||
|
|
@ -6236,9 +6132,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
|
||||
# Real per-key per-model TPM limiter + a key carrying the customer's
|
||||
# `model_tpm_limit` metadata (only the primary is capped).
|
||||
limiter = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache())
|
||||
)
|
||||
limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="sk-lit3890",
|
||||
metadata={"model_tpm_limit": {primary_model: 100}},
|
||||
|
|
@ -6246,10 +6140,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
|
||||
# Pre-seed the primary's per-model token counter at the cap so the very
|
||||
# next request trips it. The counter key uses the *hashed* api_key.
|
||||
counter_key = (
|
||||
f"{user_api_key_dict.api_key}::{primary_model}"
|
||||
f"::{precise_minute}::request_count"
|
||||
)
|
||||
counter_key = f"{user_api_key_dict.api_key}::{primary_model}::{precise_minute}::request_count"
|
||||
await limiter.internal_usage_cache.async_set_cache(
|
||||
key=counter_key,
|
||||
value={"current_requests": 0, "current_tpm": 100, "current_rpm": 0},
|
||||
|
|
@ -6280,9 +6171,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
mock_router = MagicMock()
|
||||
mock_router.fallbacks = [{primary_model: [fallback_model]}]
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock
|
||||
):
|
||||
with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock):
|
||||
with patch.object(
|
||||
processor,
|
||||
"common_processing_pre_call_logic",
|
||||
|
|
@ -6312,9 +6201,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
|
|||
|
||||
# Sanity-check the premise: the limiter genuinely raises a
|
||||
# ProxyRateLimitError for the capped primary under the frozen clock.
|
||||
with patch(
|
||||
"litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock
|
||||
):
|
||||
with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock):
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await limiter.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -6675,16 +6562,12 @@ class TestStreamingClientDisconnectBilling:
|
|||
prompt_tokens=1000,
|
||||
completion_tokens=10,
|
||||
total_tokens=1010,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
cached_tokens=500
|
||||
),
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
event = await self._bill_and_collect_success_event(
|
||||
append_openai_style_cached_usage_chunk
|
||||
)
|
||||
event = await self._bill_and_collect_success_event(append_openai_style_cached_usage_chunk)
|
||||
|
||||
usage = event["response_obj"].usage
|
||||
assert getattr(usage, "cache_read_input_tokens", None) == 500
|
||||
|
|
@ -7454,9 +7337,7 @@ class TestInjectCostIntoUsageDict:
|
|||
logging_obj.model_call_details["custom_llm_provider"] = "anthropic"
|
||||
assert logging_obj.cost_breakdown is None
|
||||
|
||||
model_response = ModelResponse(
|
||||
usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)
|
||||
)
|
||||
model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224))
|
||||
cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj)
|
||||
|
||||
assert cost is not None and cost > 0
|
||||
|
|
@ -7485,9 +7366,7 @@ class TestInjectCostIntoUsageDict:
|
|||
)
|
||||
existing = logging_obj.cost_breakdown
|
||||
|
||||
model_response = ModelResponse(
|
||||
usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)
|
||||
)
|
||||
model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224))
|
||||
ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj)
|
||||
|
||||
assert logging_obj.cost_breakdown is existing
|
||||
|
|
@ -7782,9 +7661,7 @@ def test_ttft_keepalive_interval_only_arms_for_a_streaming_request(request_data,
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream_requested, expect_ping", [(True, True), (False, False)])
|
||||
async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(
|
||||
stream_requested, expect_ping
|
||||
):
|
||||
async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(stream_requested, expect_ping):
|
||||
"""The wiring, not the helper: every route funnels through this method, and the
|
||||
whole time-to-first-token is spent inside the call it wraps."""
|
||||
|
||||
|
|
@ -7930,9 +7807,7 @@ async def test_a_late_failure_is_reported_to_the_failure_hook():
|
|||
async def record(exc):
|
||||
audited.append(exc)
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=record
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=record)
|
||||
collected = await _drain(response)
|
||||
|
||||
assert [type(exc).__name__ for exc in audited] == ["HTTPException"]
|
||||
|
|
@ -7949,9 +7824,7 @@ async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame():
|
|||
async def broken_hook(exc):
|
||||
raise RuntimeError("the audit backend is down")
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook)
|
||||
collected = await _drain(response)
|
||||
|
||||
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
|
||||
|
|
@ -8005,9 +7878,7 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke
|
|||
[(0, False), (None, True)],
|
||||
ids=["operator-hard-disabled-this-deployment", "deployment-says-nothing"],
|
||||
)
|
||||
async def test_base_process_llm_request_honours_a_deployment_hard_disable(
|
||||
deployment_keepalive, expect_ping
|
||||
):
|
||||
async def test_base_process_llm_request_honours_a_deployment_hard_disable(deployment_keepalive, expect_ping):
|
||||
"""`keepalive_seconds: 0` is documented as a disable a request cannot lift. The
|
||||
funnel has to hand its router to the gate for that to hold before the upstream
|
||||
has answered, since no deployment has served the request yet."""
|
||||
|
|
@ -8053,9 +7924,7 @@ async def test_a_hook_returning_a_replacement_decides_what_the_client_sees():
|
|||
async def sanitize(exc):
|
||||
return HTTPException(status_code=502, detail="upstream unavailable")
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize)
|
||||
collected = await _drain(response)
|
||||
|
||||
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
|
||||
|
|
@ -8094,9 +7963,7 @@ async def test_a_hook_that_returns_nothing_leaves_the_real_error_intact():
|
|||
async def audit_only(exc):
|
||||
return None
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only)
|
||||
collected = await _drain(response)
|
||||
|
||||
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
|
||||
|
|
@ -8113,9 +7980,7 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug():
|
|||
async def broken_hook(exc):
|
||||
raise RuntimeError("the audit backend is down")
|
||||
|
||||
response = await open_sse_before_first_byte(
|
||||
slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook
|
||||
)
|
||||
response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook)
|
||||
collected = await _drain(response)
|
||||
|
||||
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
|
||||
|
|
@ -8328,9 +8193,7 @@ class TestStreamingResponseHeadersFollowFallback:
|
|||
proxy_logging_obj.post_call_success_hook = AsyncMock(
|
||||
side_effect=lambda data, user_api_key_dict, response: response
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value={"x-callback-header": "kept"}
|
||||
)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-callback-header": "kept"})
|
||||
|
||||
async def fake_route_request(**kwargs):
|
||||
async def call():
|
||||
|
|
@ -8338,9 +8201,7 @@ class TestStreamingResponseHeadersFollowFallback:
|
|||
|
||||
return call()
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm.proxy.common_request_processing, "route_request", fake_route_request
|
||||
)
|
||||
monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request)
|
||||
|
||||
result = await processor.base_process_llm_request(
|
||||
request=Request(scope={"type": "http", "headers": []}),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue