mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(anthropic): answer an upstream Anthropic 529 as 529 overloaded_error on /v1/messages (#46063)
* fix(anthropic): answer a 503 as 529 overloaded_error on /v1/messages
Anthropic's error contract has no 503, so a client built against it (Claude Code) treats a 503 from /v1/messages as a plain 5xx and retries it ten times instead of taking its overload path. A Bedrock 503 now leaves /v1/messages and /v1/messages/count_tokens as 529 overloaded_error, for both a JSON answer and a stream that fails before its first byte.
An Anthropic 529 overloaded_error now maps to litellm.ServiceUnavailableError (503) instead of InternalServerError (500), so OpenAI-shaped routes answer it as 503, the status OpenAI documents for an overloaded engine, and router retry and allowed-fails policies key it as ServiceUnavailableError.
* test(proxy): expect overloaded_error on a 503 stream error after a keepalive ping
* fix(proxy): record the 529 the Anthropic routes answer and type the chat stream error frame
The messages route stamps the answered status on the exception before post_call_failure_hook runs, so the spend log row, the OTel server span, the failed-request metric, and alerting record the 529 the client saw instead of the deployment's 503, while the deployment failure metric and cooldown keep the deployment's own status. A chat completions stream error frame names the type its status stands for instead of the null .type ServiceUnavailableError carries. An Anthropic-shaped 503 envelope passes through as 529 overloaded_error through a validated TypedDict, and count_tokens answers every ProxyException as the Anthropic envelope. Five integration tests that pinned 503 on the messages route now expect 529
* test(proxy): expect the top-level Anthropic envelope from count_tokens errors
count_tokens now answers a ProxyException as a JSONResponse carrying the
Anthropic error envelope instead of raising HTTPException, which FastAPI
wrapped in {"detail": ...}. The three legacy error-format tests now read
the response body and check that no detail wrap is left.
* test(proxy): cover the 529 overloaded_error answer across routes, cooldown, and retry budgets
* fix(anthropic): keep Anthropic overloads as InternalServerError and answer 529 from the provider status
An upstream Anthropic 529 stays litellm.InternalServerError (status 500) so
InternalServerErrorRetries, InternalServerErrorAllowedFails, SDK handlers,
and the OpenAI-shaped routes behave as on main. The error now carries
provider_status_code, and the /v1/messages routes answer 529
overloaded_error from it.
* refactor(proxy): unwrap a mid-stream fallback error once instead of recursing for the provider status
* fix(anthropic): keep count_tokens' detail body, end keepalive streams with an Anthropic error event, and log any 529 as an overload
* test(integration): cover the audit round 3 cells and expect 529 on the count_tokens 503 cases
* fix(proxy): re-raise an unchanged stream HTTPException as itself and walk nested mid-stream fallbacks for the provider status
* fix(proxy): walk nested mid-stream fallbacks without recursion
* fix(anthropic): preserve 503s and answer upstream 529s on messages routes
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(anthropic): answer an upstream Anthropic 529 as 529 on count_tokens
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(anthropic): use integration coverage for keepalive errors
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mateo Wang
GitHub
parent
2a9fad611f
commit
d5ca88900c
23 changed files with 2068 additions and 49 deletions
|
|
@ -867,6 +867,7 @@ class BadGatewayError(openai.APIStatusError):
|
|||
class InternalServerError(openai.InternalServerError):
|
||||
response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response
|
||||
request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request
|
||||
provider_status_code: int | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -878,8 +879,10 @@ class InternalServerError(openai.InternalServerError):
|
|||
max_retries: int | None = None,
|
||||
num_retries: int | None = None,
|
||||
body: object | None = None,
|
||||
provider_status_code: int | None = None,
|
||||
):
|
||||
self.status_code = 500
|
||||
self.provider_status_code = provider_status_code
|
||||
self.message = f"litellm.InternalServerError: {message}"
|
||||
self.llm_provider = llm_provider
|
||||
self.model = model
|
||||
|
|
|
|||
|
|
@ -133,6 +133,8 @@ _STATUS_CODE_MAP: Final[Mapping[str, str]] = MappingProxyType(
|
|||
}
|
||||
)
|
||||
|
||||
_OVERLOADED_STATUS_CODE: Final = "529"
|
||||
|
||||
_INTERNAL_STATE_EXCEPTIONS: Final[tuple[type[BaseException], ...]] = (
|
||||
TypeError,
|
||||
KeyError,
|
||||
|
|
@ -201,6 +203,8 @@ def normalize_error(exc: Exception | None, status_code: str, message: str) -> st
|
|||
)
|
||||
if by_message is not None:
|
||||
return by_message
|
||||
if status_code == _OVERLOADED_STATUS_CODE:
|
||||
return PROVIDER_OVERLOADED
|
||||
by_class: Final = _classify_by_class(exc)
|
||||
if by_class is not None:
|
||||
return by_class
|
||||
|
|
|
|||
|
|
@ -605,6 +605,11 @@ def _map_openai_exception(
|
|||
)
|
||||
|
||||
|
||||
def _anthropic_overload_status(original_exception: _ProviderHTTPException) -> int | None:
|
||||
upstream_status: Final = getattr(original_exception, "status_code", None)
|
||||
return 529 if isinstance(upstream_status, int) and upstream_status >= 500 else None
|
||||
|
||||
|
||||
def _map_anthropic_exception(
|
||||
*,
|
||||
model: str,
|
||||
|
|
@ -630,6 +635,7 @@ def _map_anthropic_exception(
|
|||
message=f"AnthropicError - {error_str}",
|
||||
model=model,
|
||||
llm_provider="anthropic",
|
||||
provider_status_code=_anthropic_overload_status(original_exception),
|
||||
)
|
||||
if "Invalid API Key" in error_str:
|
||||
raise AuthenticationError(
|
||||
|
|
@ -694,6 +700,7 @@ def _map_anthropic_exception(
|
|||
llm_provider="anthropic",
|
||||
model=model,
|
||||
response=getattr(original_exception, "response", None),
|
||||
provider_status_code=original_exception.status_code,
|
||||
)
|
||||
elif original_exception.status_code == 502:
|
||||
raise BadGatewayError(
|
||||
|
|
|
|||
30
litellm/proxy/anthropic_endpoints/answered_status.py
Normal file
30
litellm/proxy/anthropic_endpoints/answered_status.py
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
from typing import Final, TypeVar
|
||||
|
||||
from litellm.exceptions import InternalServerError, MidStreamFallbackError
|
||||
|
||||
_Raised = TypeVar("_Raised", bound=Exception)
|
||||
|
||||
|
||||
def _provider_status_code(exc: Exception) -> int | None:
|
||||
unwrapped = exc # rebind-ok: the recursion gate requires iterative wrapper traversal
|
||||
while isinstance(unwrapped, MidStreamFallbackError) and unwrapped.original_exception is not None:
|
||||
unwrapped = unwrapped.original_exception
|
||||
match unwrapped:
|
||||
case InternalServerError(provider_status_code=int() as provider_status):
|
||||
return provider_status
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _copy_carrying(exc: _Raised, attribute: str, value: int | str) -> _Raised:
|
||||
answered: Final = type(exc).__new__(type(exc), *exc.args)
|
||||
answered.__dict__.update({**exc.__dict__, attribute: value})
|
||||
answered.__cause__ = exc.__cause__
|
||||
answered.__context__ = exc.__context__
|
||||
answered.__suppress_context__ = exc.__suppress_context__
|
||||
return answered.with_traceback(exc.__traceback__)
|
||||
|
||||
|
||||
def with_answered_status(exc: Exception) -> Exception:
|
||||
provider_status_code: Final = _provider_status_code(exc)
|
||||
return exc if provider_status_code is None else _copy_carrying(exc, "status_code", provider_status_code)
|
||||
|
|
@ -22,6 +22,7 @@ from litellm.llms.anthropic.messages.context_management import (
|
|||
AnthropicContextManagementError,
|
||||
)
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.anthropic_endpoints.answered_status import with_answered_status
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
|
|
@ -225,15 +226,18 @@ async def anthropic_response(
|
|||
request_id=request.headers.get("x-request-id"),
|
||||
)
|
||||
return JSONResponse(status_code=e.status_code, content=body)
|
||||
except Exception as e:
|
||||
except Exception as raised:
|
||||
answered: Final = with_answered_status(raised)
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=base_llm_response_processor.data
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=answered,
|
||||
request_data=base_llm_response_processor.data,
|
||||
)
|
||||
log_llm_api_exception(e, base_llm_response_processor.litellm_call_id)
|
||||
log_llm_api_exception(answered, base_llm_response_processor.litellm_call_id)
|
||||
|
||||
if isinstance(e, ProxyException):
|
||||
if isinstance(answered, ProxyException):
|
||||
return _anthropic_error_json_response(
|
||||
with_litellm_call_id(e, base_llm_response_processor.litellm_call_id), request
|
||||
with_litellm_call_id(answered, base_llm_response_processor.litellm_call_id), request
|
||||
)
|
||||
|
||||
# Extract model_id from request metadata (same as success path)
|
||||
|
|
@ -250,20 +254,20 @@ async def anthropic_response(
|
|||
response_cost=0,
|
||||
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
|
||||
request_data=data,
|
||||
timeout=getattr(e, "timeout", None),
|
||||
timeout=getattr(answered, "timeout", None),
|
||||
litellm_logging_obj=None,
|
||||
)
|
||||
|
||||
if isinstance(e, HTTPException):
|
||||
return _anthropic_error_json_response(proxy_exception_from_http_exception(e, headers), request)
|
||||
if isinstance(answered, HTTPException):
|
||||
return _anthropic_error_json_response(proxy_exception_from_http_exception(answered, headers), request)
|
||||
|
||||
error_msg: Final = f"{e}"
|
||||
error_msg: Final = f"{answered}"
|
||||
return _anthropic_error_json_response(
|
||||
ProxyException(
|
||||
message=getattr(e, "message", error_msg),
|
||||
type=openai_error_type(e, error_status_code(e, 500)),
|
||||
param=openai_error_param(e),
|
||||
code=error_status_code(e, 500),
|
||||
message=getattr(answered, "message", error_msg),
|
||||
type=openai_error_type(answered, error_status_code(answered, 500)),
|
||||
param=openai_error_param(answered),
|
||||
code=error_status_code(answered, 500),
|
||||
headers=headers,
|
||||
),
|
||||
request,
|
||||
|
|
@ -383,8 +387,9 @@ async def count_tokens(
|
|||
except ProxyException as e:
|
||||
return _anthropic_error_json_response(e, request)
|
||||
except Exception as e:
|
||||
log_llm_api_exception(e, litellm_call_id)
|
||||
return _anthropic_error_json_response(proxy_exception_from_route_error(e), request)
|
||||
answered: Final = with_answered_status(e)
|
||||
log_llm_api_exception(answered, litellm_call_id)
|
||||
return _anthropic_error_json_response(proxy_exception_from_route_error(answered), request)
|
||||
else:
|
||||
return {
|
||||
"input_tokens": token_response.total_tokens,
|
||||
|
|
|
|||
|
|
@ -272,6 +272,7 @@ if TYPE_CHECKING:
|
|||
ProxyConfig = _ProxyConfig
|
||||
else:
|
||||
ProxyConfig = Any
|
||||
from litellm.proxy.anthropic_endpoints.answered_status import with_answered_status
|
||||
from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
|
||||
AnthropicStreamModelRestamper,
|
||||
)
|
||||
|
|
@ -1272,6 +1273,21 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]:
|
|||
return f"data: {json.dumps({'error': error_obj}, default=str)}\n\n", "data: [DONE]\n\n"
|
||||
|
||||
|
||||
def _openai_late_failure_frames(exc: Exception) -> tuple[str, ...]:
|
||||
_, error_obj = sse_error_payload(exc)
|
||||
return _sse_error_frames(error_obj)
|
||||
|
||||
|
||||
def _anthropic_late_failure_frames(exc: Exception) -> tuple[str, ...]:
|
||||
answered: Final = proxy_exception_from_route_error(with_answered_status(exc))
|
||||
return (
|
||||
anthropic_error_sse_frame(
|
||||
status_code=error_status_code(answered, status.HTTP_500_INTERNAL_SERVER_ERROR),
|
||||
raw_message=answered.message,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
|
||||
"""`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384)."""
|
||||
return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS})
|
||||
|
|
@ -1520,6 +1536,7 @@ async def open_sse_before_first_byte(
|
|||
ping_interval_seconds: float | str | None,
|
||||
media_type: str = "text/event-stream",
|
||||
on_late_failure: Callable[[Exception], Awaitable[HTTPException | None]] | None = None,
|
||||
late_failure_frames: Callable[[Exception], tuple[str, ...]] = _openai_late_failure_frames,
|
||||
) -> _LateResponseT | StreamingResponse:
|
||||
"""Write SSE keepalive comments while the upstream LLM call is still in flight.
|
||||
|
||||
|
|
@ -1565,8 +1582,7 @@ async def open_sse_before_first_byte(
|
|||
# would never fire and the failure would go unaudited. The hook
|
||||
# also gets to sanitize what reaches the client, by returning or
|
||||
# raising a replacement, so its answer decides the frame.
|
||||
_, error_obj = sse_error_payload(await _sanitized_late_failure(exc, on_late_failure))
|
||||
for frame in _sse_error_frames(error_obj):
|
||||
for frame in late_failure_frames(await _sanitized_late_failure(exc, on_late_failure)):
|
||||
yield frame.encode()
|
||||
return
|
||||
async for chunk in _relay_late_response(produced):
|
||||
|
|
@ -2982,13 +2998,14 @@ class ProxyBaseLLMRequestProcessing:
|
|||
idle. See ``open_sse_before_first_byte``; unwrapped unless an operator sets
|
||||
``litellm_settings.sse_keepalive_ping_interval_seconds``.
|
||||
"""
|
||||
answers_as_anthropic: Final = route_type == "anthropic_messages"
|
||||
|
||||
async def _audit_late_failure(exc: Exception) -> HTTPException | None:
|
||||
# Once a keepalive is on the wire this can no longer raise, so the
|
||||
# caller's `except` never runs its own post_call_failure_hook.
|
||||
return await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=exc,
|
||||
original_exception=with_answered_status(exc) if answers_as_anthropic else exc,
|
||||
request_data=self.data,
|
||||
dispatched_call_type=self.route_type,
|
||||
)
|
||||
|
|
@ -3017,6 +3034,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
),
|
||||
ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router),
|
||||
on_late_failure=_audit_late_failure,
|
||||
late_failure_frames=_anthropic_late_failure_frames if answers_as_anthropic else _openai_late_failure_frames,
|
||||
)
|
||||
|
||||
async def _process_llm_request(
|
||||
|
|
@ -4641,6 +4659,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request: Request | None = None,
|
||||
flush_tail: Callable[[], bytes] | None = None,
|
||||
seal_open_frame: Callable[[bytes], str] | None = None,
|
||||
answer_status: Callable[[Exception], Exception] | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Shared streaming data generator: runs proxy iterator hook, per-chunk hook,
|
||||
|
|
@ -4656,6 +4675,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
relays raw upstream bytes, so an upstream that hangs up mid-frame leaves the
|
||||
client inside an open frame, where an error frame would be swallowed or
|
||||
misparsed instead of raised.
|
||||
|
||||
``answer_status`` gives a failure the status the route answers it with
|
||||
before the failure hook records it.
|
||||
"""
|
||||
verbose_proxy_logger.debug("inside generator")
|
||||
# Resolve per-stream (not per-chunk) whether the heavy per-chunk path
|
||||
|
|
@ -4752,6 +4774,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.proxy_server.async_data_generator(): Exception occured - %s", e
|
||||
)
|
||||
answered: Final = e if answer_status is None else answer_status(e)
|
||||
already_reported: Final = (
|
||||
isinstance(stream_logging_obj, LiteLLMLoggingObj)
|
||||
and not stream_logging_obj.claim_stream_failure_report()
|
||||
|
|
@ -4761,24 +4784,25 @@ class ProxyBaseLLMRequestProcessing:
|
|||
if already_reported
|
||||
else await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
original_exception=answered,
|
||||
request_data=request_data,
|
||||
)
|
||||
)
|
||||
if transformed_exception is not None:
|
||||
e = transformed_exception
|
||||
failure: Final = answered if transformed_exception is None else transformed_exception
|
||||
verbose_proxy_logger.debug(
|
||||
"An error occurred: %s",
|
||||
e,
|
||||
failure,
|
||||
)
|
||||
|
||||
if isinstance(e, HTTPException):
|
||||
raise e
|
||||
stream_error_status: Final = error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
if isinstance(failure, HTTPException):
|
||||
if failure is e:
|
||||
raise
|
||||
raise failure from e
|
||||
stream_error_status: Final = error_status_code(failure, status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
proxy_exception: Final = ProxyException(
|
||||
message=redact_internal_details_from_client_message(getattr(e, "message", str(e))),
|
||||
type=openai_error_type(e, stream_error_status),
|
||||
param=openai_error_param(e),
|
||||
message=redact_internal_details_from_client_message(getattr(failure, "message", str(failure))),
|
||||
type=openai_error_type(failure, stream_error_status),
|
||||
param=openai_error_param(failure),
|
||||
code=stream_error_status,
|
||||
)
|
||||
stream_completed = True
|
||||
|
|
@ -4834,6 +4858,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request=request,
|
||||
flush_tail=None if restamper is None else restamper.flush,
|
||||
seal_open_frame=seal_open_sse_frame,
|
||||
answer_status=with_answered_status,
|
||||
)
|
||||
|
||||
@overload
|
||||
|
|
|
|||
|
|
@ -198,7 +198,7 @@ def _detect_loop(history: list[str], new_calls: Sequence[Mapping[str, object]])
|
|||
return False
|
||||
|
||||
|
||||
_EXHAUSTION_STATUSES: Final = {408, 413, 429, 503, 504}
|
||||
_EXHAUSTION_STATUSES: Final = {408, 413, 429, 503, 504, 529}
|
||||
|
||||
_EXHAUSTION_KEYWORDS: Final = (
|
||||
"context length",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,724 @@
|
|||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import unquote
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.anthropic_sse import (
|
||||
LIFECYCLE,
|
||||
error_frame,
|
||||
event_type,
|
||||
event_types,
|
||||
message_id,
|
||||
message_json,
|
||||
message_stream,
|
||||
parse_sse,
|
||||
status_reply,
|
||||
stream_reply,
|
||||
user_prompt,
|
||||
)
|
||||
from integration._support.client import (
|
||||
JSON_OBJECT,
|
||||
Gateway,
|
||||
JsonValue,
|
||||
Scenario,
|
||||
eventually,
|
||||
object_value,
|
||||
string_value,
|
||||
)
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.openai_wire import openai_error
|
||||
from integration._support.upstream import aws_event_stream_frame
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from redis import Redis
|
||||
|
||||
_MODEL: Final = "claude-sonnet-4-5-20250929"
|
||||
_API_KEY: Final = "synthetic-anthropic-key"
|
||||
_BEDROCK_MODEL: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
_BEDROCK_TOKEN: Final = "synthetic-bedrock-bearer"
|
||||
_BEDROCK_503: Final = b'{"message": "Bedrock is unable to process your request."}'
|
||||
_CONVERSE_MODEL: Final = f"bedrock/converse/{_BEDROCK_MODEL}"
|
||||
_CONVERSE_STREAM_TARGET: Final = f"/model/{_BEDROCK_MODEL}/converse-stream"
|
||||
_OPENAI_BACKEND: Final = "gpt-4o-mini"
|
||||
_CHAT_ONLY_BACKEND: Final = "audit-chat-only"
|
||||
_CODEX_BACKEND: Final = "gpt-5.3-codex"
|
||||
_TEXT: Final = "Hello"
|
||||
_KEEPALIVE_SECONDS: Final = 0.2
|
||||
_LATE_ANSWER_SECONDS: Final = 1.0
|
||||
_OVERLOAD_CLASS: Final = "503_PROVIDER_OVERLOADED"
|
||||
|
||||
Shape = Literal["stream", "non-stream"]
|
||||
|
||||
|
||||
def _prompt() -> str:
|
||||
return "overloaded-" + uuid.uuid4().hex
|
||||
|
||||
|
||||
def _always(reply: Reply) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST", request
|
||||
return reply
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _late(reply: Reply, seconds: float) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST", request
|
||||
time.sleep(seconds)
|
||||
return reply
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _anthropic_sdk(gateway: Gateway) -> anthropic.Anthropic:
|
||||
return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0, timeout=60)
|
||||
|
||||
|
||||
def _async_anthropic_sdk(gateway: Gateway) -> anthropic.AsyncAnthropic:
|
||||
return anthropic.AsyncAnthropic(
|
||||
base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0, timeout=60
|
||||
)
|
||||
|
||||
|
||||
def _openai_sdk(gateway: Gateway) -> openai.OpenAI:
|
||||
return openai.OpenAI(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0, timeout=60)
|
||||
|
||||
|
||||
def _async_openai_sdk(gateway: Gateway) -> openai.AsyncOpenAI:
|
||||
return openai.AsyncOpenAI(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0, timeout=60)
|
||||
|
||||
|
||||
def _messages_body(model: str, prompt: str, *, stream: bool) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": model,
|
||||
"max_tokens": 16,
|
||||
"stream": stream,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"num_retries": 0,
|
||||
}
|
||||
|
||||
|
||||
def _anthropic_error(response: httpx.Response) -> dict[str, JsonValue]:
|
||||
assert response.headers["content-type"].startswith("application/json"), dict(response.headers)
|
||||
payload: Final = object_value(json.loads(response.content))
|
||||
assert payload["type"] == "error", response.text
|
||||
return object_value(payload["error"])
|
||||
|
||||
|
||||
def _openai_error(response: httpx.Response) -> dict[str, JsonValue]:
|
||||
return object_value(object_value(json.loads(response.content))["error"])
|
||||
|
||||
|
||||
def _assert_overloaded(response: httpx.Response, *, message: str) -> None:
|
||||
assert response.status_code == 529, response.text
|
||||
error: Final = _anthropic_error(response)
|
||||
assert error["type"] == "overloaded_error", response.text
|
||||
assert message in string_value(error["message"]), response.text
|
||||
assert "content_block_delta" not in response.text, response.text
|
||||
|
||||
|
||||
def _assert_api_error(response: httpx.Response, *, status: int, message: str) -> None:
|
||||
assert response.status_code == status, response.text
|
||||
error: Final = _anthropic_error(response)
|
||||
assert error["type"] == "api_error", response.text
|
||||
assert message in string_value(error["message"]), response.text
|
||||
|
||||
|
||||
def _assert_internal_server_error(response: httpx.Response) -> None:
|
||||
assert response.status_code == 500, response.text
|
||||
error: Final = _openai_error(response)
|
||||
assert error["code"] == "500", response.text
|
||||
assert "InternalServerError" in string_value(error["message"]), response.text
|
||||
|
||||
|
||||
def _one_post(wire: Wire, target: str) -> None:
|
||||
assert [(request.method, unquote(request.target)) for request in wire.drain()] == [("POST", target)]
|
||||
|
||||
|
||||
def _spend_rows(model: str, expected: int) -> list[dict[str, JsonValue]]:
|
||||
return eventually(
|
||||
lambda: read_rows(
|
||||
"SELECT status, metadata->'error_information'->>'error_code' AS code, "
|
||||
"metadata->'error_information'->>'error_class' AS error_class, "
|
||||
"metadata->'error_information'->>'normalized_error' AS normalized_error "
|
||||
'FROM "LiteLLM_SpendLogs" WHERE model_group=%s',
|
||||
(model,),
|
||||
),
|
||||
lambda rows: len(rows) >= expected,
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def _assert_logged_failure(model: str, *, rows: int, error_class: str, code: str = "529") -> None:
|
||||
assert (
|
||||
_spend_rows(model, rows)
|
||||
== [{"status": "failure", "code": code, "error_class": error_class, "normalized_error": _OVERLOAD_CLASS}] * rows
|
||||
)
|
||||
|
||||
|
||||
def test_messages_529_answers_overloaded_error_to_the_sync_sdk_and_logs_each_call_once(gateway: Gateway) -> None:
|
||||
with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
sdk: Final = _anthropic_sdk(gateway)
|
||||
prompt: Final = _prompt()
|
||||
for _ in range(2):
|
||||
with pytest.raises(anthropic.APIStatusError) as caught:
|
||||
sdk.messages.create(
|
||||
model=model,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
extra_body={"num_retries": 0},
|
||||
)
|
||||
assert caught.value.status_code == 529, caught.value.response.text
|
||||
_assert_overloaded(caught.value.response, message="scripted 529")
|
||||
assert [unquote(request.target) for request in wire.drain()] == ["/v1/messages"] * 2
|
||||
_assert_logged_failure(model, rows=2, error_class="InternalServerError")
|
||||
|
||||
|
||||
async def test_messages_529_answers_overloaded_error_to_the_async_sdk(gateway: Gateway) -> None:
|
||||
with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
with pytest.raises(anthropic.APIStatusError) as caught:
|
||||
await _async_anthropic_sdk(gateway).messages.create(
|
||||
model=model,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": _prompt()}],
|
||||
extra_body={"num_retries": 0},
|
||||
)
|
||||
assert caught.value.status_code == 529, caught.value.response.text
|
||||
_assert_overloaded(caught.value.response, message="scripted 529")
|
||||
_one_post(wire, "/v1/messages")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_status", "error_class", "error_type"),
|
||||
((529, "InternalServerError", "overloaded_error"), (503, "ServiceUnavailableError", "api_error")),
|
||||
)
|
||||
def test_messages_stream_rejected_before_its_first_byte_answers_the_provider_status_as_json(
|
||||
gateway: Gateway, upstream_status: int, error_class: str, error_type: str
|
||||
) -> None:
|
||||
with wire_server(_always(status_reply(upstream_status))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = gateway.request("POST", "/v1/messages", _messages_body(model, _prompt(), stream=True))
|
||||
assert response.status_code == upstream_status, response.text
|
||||
error: Final = _anthropic_error(response)
|
||||
assert error["type"] == error_type, response.text
|
||||
assert f"scripted {upstream_status}" in string_value(error["message"]), response.text
|
||||
_one_post(wire, "/v1/messages")
|
||||
_assert_logged_failure(model, rows=1, error_class=error_class, code=str(upstream_status))
|
||||
|
||||
|
||||
def test_messages_stream_529_raises_529_in_the_sync_sdk(gateway: Gateway) -> None:
|
||||
with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
with pytest.raises(anthropic.APIStatusError) as caught:
|
||||
with _anthropic_sdk(gateway).messages.stream(
|
||||
model=model,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": _prompt()}],
|
||||
extra_body={"num_retries": 0},
|
||||
) as stream:
|
||||
stream.get_final_message()
|
||||
assert caught.value.status_code == 529, caught.value.response.text
|
||||
_assert_overloaded(caught.value.response, message="scripted 529")
|
||||
_one_post(wire, "/v1/messages")
|
||||
|
||||
|
||||
def _error_frames(text: str) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
object_value(json.loads(line.removeprefix("data:").strip()))
|
||||
for line in text.splitlines()
|
||||
if line.startswith("data:") and '"error"' in line
|
||||
)
|
||||
|
||||
|
||||
def test_messages_529_after_the_keepalive_opened_the_response_ends_with_an_overloaded_error_event(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
with (
|
||||
wire_server(_late(status_reply(529), _LATE_ANSWER_SECONDS)) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, keepalive_seconds=_KEEPALIVE_SECONDS
|
||||
)
|
||||
response: Final = gateway.request("POST", "/v1/messages", _messages_body(model, _prompt(), stream=True))
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), dict(response.headers)
|
||||
assert response.text.startswith(":"), response.text
|
||||
(failure,) = (event for event in parse_sse(response.text) if event.data and event.event != "ping")
|
||||
assert failure.event == "error", response.text
|
||||
assert object_value(failure.data["error"])["type"] == "overloaded_error", response.text
|
||||
assert "[DONE]" not in response.text, response.text
|
||||
_one_post(wire, "/v1/messages")
|
||||
|
||||
|
||||
def _late_bedrock_deployment(scenario: Scenario, wire: Wire) -> str:
|
||||
return scenario.model(
|
||||
model=f"bedrock/{_BEDROCK_MODEL}",
|
||||
api_key=_BEDROCK_TOKEN,
|
||||
aws_region_name="us-east-1",
|
||||
api_base=wire.url,
|
||||
keepalive_seconds=_KEEPALIVE_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _one_bedrock_post(wire: Wire) -> None:
|
||||
(request,) = wire.drain()
|
||||
assert unquote(request.target).startswith(f"/model/{_BEDROCK_MODEL}/"), request
|
||||
|
||||
|
||||
def test_messages_bedrock_503_after_the_keepalive_opened_the_response_ends_with_an_api_error_event(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
with (
|
||||
wire_server(_late(Reply(status=503, body=_BEDROCK_503), _LATE_ANSWER_SECONDS)) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = _late_bedrock_deployment(scenario, wire)
|
||||
response: Final = gateway.request("POST", "/v1/messages", _messages_body(model, _prompt(), stream=True))
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.text.startswith(":"), response.text
|
||||
assert "[DONE]" not in response.text, response.text
|
||||
(failure,) = (event for event in parse_sse(response.text) if event.data and event.event != "ping")
|
||||
assert failure.event == "error", response.text
|
||||
error: Final = object_value(failure.data["error"])
|
||||
assert error["type"] == "api_error", response.text
|
||||
assert "Bedrock is unable to process your request." in string_value(error["message"]), response.text
|
||||
_one_bedrock_post(wire)
|
||||
_assert_logged_failure(model, rows=1, error_class="ServiceUnavailableError", code="503")
|
||||
|
||||
|
||||
def test_chat_completions_bedrock_503_after_the_keepalive_opened_the_response_keeps_the_openai_error_frame(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
with (
|
||||
wire_server(_late(Reply(status=503, body=_BEDROCK_503), _LATE_ANSWER_SECONDS)) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = _late_bedrock_deployment(scenario, wire)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": _prompt()}], "stream": True, "num_retries": 0},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.text.startswith(":"), response.text
|
||||
assert "event: error" not in response.text, response.text
|
||||
(frame,) = _error_frames(response.text)
|
||||
assert object_value(frame["error"])["code"] == "503", response.text
|
||||
assert response.text.rstrip().endswith("data: [DONE]"), response.text
|
||||
_one_bedrock_post(wire)
|
||||
_assert_logged_failure(model, rows=1, error_class="ServiceUnavailableError", code="503")
|
||||
|
||||
|
||||
def test_chat_completions_overloaded_frame_after_content_ends_with_a_500_error_frame(gateway: Gateway) -> None:
|
||||
prompt: Final = _prompt()
|
||||
upstream: Final = (*message_stream(prompt, _MODEL, _TEXT)[:3], error_frame(529, "Overloaded"))
|
||||
with wire_server(_always(stream_reply(upstream))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": prompt}], "stream": True, "num_retries": 0},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _TEXT in response.text, response.text
|
||||
(frame,) = _error_frames(response.text)
|
||||
error: Final = object_value(frame["error"])
|
||||
assert error["code"] == "500", response.text
|
||||
_one_post(wire, "/v1/messages")
|
||||
|
||||
|
||||
def _overloaded_before_content(prompt: str) -> Reply:
|
||||
return stream_reply((*message_stream(prompt, _MODEL, _TEXT)[:1], error_frame(529, "Overloaded")))
|
||||
|
||||
|
||||
def test_messages_overloaded_frame_before_content_without_a_retry_budget_is_relayed_as_an_overloaded_event(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
prompt: Final = _prompt()
|
||||
with wire_server(_always(_overloaded_before_content(prompt))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = gateway.request("POST", "/v1/messages", _messages_body(model, prompt, stream=True))
|
||||
assert response.status_code == 200, response.text
|
||||
events: Final = parse_sse(response.text)
|
||||
assert event_types(events) == ("message_start", "error"), response.text
|
||||
assert message_id(events) == prompt, response.text
|
||||
assert events[1].data == {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}
|
||||
_one_post(wire, "/v1/messages")
|
||||
assert _spend_rows(model, 1) == [
|
||||
{"status": "success", "code": None, "error_class": None, "normalized_error": None}
|
||||
]
|
||||
|
||||
|
||||
def test_chat_completions_overloaded_frame_before_content_answers_500(gateway: Gateway) -> None:
|
||||
prompt: Final = _prompt()
|
||||
with wire_server(_always(_overloaded_before_content(prompt))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": prompt}], "stream": True, "num_retries": 0},
|
||||
)
|
||||
_assert_internal_server_error(response)
|
||||
_one_post(wire, "/v1/messages")
|
||||
|
||||
|
||||
def test_chat_completions_529_answers_500_to_the_sync_sdk(gateway: Gateway) -> None:
|
||||
with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
with pytest.raises(openai.APIStatusError) as caught:
|
||||
_openai_sdk(gateway).chat.completions.create(
|
||||
model=model, messages=[{"role": "user", "content": _prompt()}], extra_body={"num_retries": 0}
|
||||
)
|
||||
assert caught.value.status_code == 500, caught.value.response.text
|
||||
_assert_internal_server_error(caught.value.response)
|
||||
_one_post(wire, "/v1/messages")
|
||||
_assert_logged_failure(model, rows=1, error_class="InternalServerError", code="500")
|
||||
|
||||
|
||||
async def test_chat_completions_stream_529_answers_500_to_the_async_sdk(gateway: Gateway) -> None:
|
||||
with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
with pytest.raises(openai.APIStatusError) as caught:
|
||||
await _async_openai_sdk(gateway).chat.completions.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": _prompt()}],
|
||||
stream=True,
|
||||
extra_body={"num_retries": 0},
|
||||
)
|
||||
assert caught.value.status_code == 500, caught.value.response.text
|
||||
_assert_internal_server_error(caught.value.response)
|
||||
_one_post(wire, "/v1/messages")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ("non-stream", "stream"))
|
||||
def test_responses_529_answers_500(gateway: Gateway, shape: Shape) -> None:
|
||||
with wire_server(_always(status_reply(529))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/responses", {"model": model, "input": _prompt(), "stream": shape == "stream", "num_retries": 0}
|
||||
)
|
||||
_assert_internal_server_error(response)
|
||||
_one_post(wire, "/v1/messages")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ("non-stream", "stream"))
|
||||
def test_messages_bedrock_503_answers_503_api_error(gateway: Gateway, shape: Shape) -> None:
|
||||
with wire_server(_always(Reply(status=503, body=_BEDROCK_503))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"bedrock/{_BEDROCK_MODEL}", api_key=_BEDROCK_TOKEN, aws_region_name="us-east-1", api_base=wire.url
|
||||
)
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/messages", _messages_body(model, _prompt(), stream=shape == "stream")
|
||||
)
|
||||
_assert_api_error(response, status=503, message="Bedrock is unable to process your request.")
|
||||
(request,) = wire.drain()
|
||||
assert unquote(request.target).startswith(f"/model/{_BEDROCK_MODEL}/"), request
|
||||
_assert_logged_failure(model, rows=1, error_class="ServiceUnavailableError", code="503")
|
||||
|
||||
|
||||
def _converse_event(event: str, payload: Mapping[str, JsonValue]) -> bytes:
|
||||
return aws_event_stream_frame(
|
||||
{":event-type": event, ":content-type": "application/json", ":message-type": "event"},
|
||||
json.dumps(payload).encode(),
|
||||
)
|
||||
|
||||
|
||||
def _converse_503_after_text() -> Reply:
|
||||
unavailable: Final = aws_event_stream_frame(
|
||||
{
|
||||
":exception-type": "serviceUnavailableException",
|
||||
":content-type": "application/json",
|
||||
":message-type": "exception",
|
||||
},
|
||||
_BEDROCK_503,
|
||||
)
|
||||
frames: Final = (
|
||||
_converse_event("messageStart", {"role": "assistant"}),
|
||||
_converse_event("contentBlockDelta", {"delta": {"text": _TEXT}, "contentBlockIndex": 0}),
|
||||
unavailable,
|
||||
)
|
||||
return Reply(body=b"".join(frames), content_type="application/vnd.amazon.eventstream")
|
||||
|
||||
|
||||
def _converse_deployment(scenario: Scenario, wire: Wire) -> str:
|
||||
return scenario.model(model=_CONVERSE_MODEL, api_key=_BEDROCK_TOKEN, aws_region_name="us-east-1", api_base=wire.url)
|
||||
|
||||
|
||||
def test_messages_bedrock_stream_503_after_text_ends_with_an_api_error_event(gateway: Gateway) -> None:
|
||||
with wire_server(_always(_converse_503_after_text())) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _converse_deployment(scenario, wire)
|
||||
response: Final = gateway.request("POST", "/v1/messages", _messages_body(model, _prompt(), stream=True))
|
||||
assert response.status_code == 200, response.text
|
||||
events: Final = parse_sse(response.text)
|
||||
types: Final = event_types(events)
|
||||
assert "message_stop" not in types, response.text
|
||||
assert types.index("content_block_delta") < types.index("error"), response.text
|
||||
assert _TEXT in response.text, response.text
|
||||
(failure,) = (event for event in events if event_type(event) == "error")
|
||||
error: Final = object_value(failure.data["error"])
|
||||
assert error["type"] == "api_error", response.text
|
||||
assert "Bedrock is unable to process your request." in string_value(error["message"]), response.text
|
||||
_one_post(wire, _CONVERSE_STREAM_TARGET)
|
||||
_assert_logged_failure(model, rows=1, error_class="ServiceUnavailableError", code="503")
|
||||
|
||||
|
||||
def test_chat_completions_bedrock_stream_503_after_text_ends_with_a_503_error_frame(gateway: Gateway) -> None:
|
||||
with wire_server(_always(_converse_503_after_text())) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _converse_deployment(scenario, wire)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": _prompt()}], "stream": True, "num_retries": 0},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _TEXT in response.text, response.text
|
||||
(frame,) = _error_frames(response.text)
|
||||
error: Final = object_value(frame["error"])
|
||||
assert error["code"] == "503", response.text
|
||||
_one_post(wire, _CONVERSE_STREAM_TARGET)
|
||||
_assert_logged_failure(model, rows=1, error_class="ServiceUnavailableError", code="503")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ("non-stream", "stream"))
|
||||
def test_chat_bridge_503_answers_503_api_error_on_messages(gateway: Gateway, shape: Shape) -> None:
|
||||
with wire_server(_always(openai_error(503))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"hosted_vllm/{_CHAT_ONLY_BACKEND}", api_base=wire.url + "/v1", api_key="synthetic-vllm-key"
|
||||
)
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/messages", _messages_body(model, _prompt(), stream=shape == "stream")
|
||||
)
|
||||
_assert_api_error(response, status=503, message="scripted 503")
|
||||
_one_post(wire, "/v1/chat/completions")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ("non-stream", "stream"))
|
||||
def test_responses_bridge_503_answers_503_api_error_on_messages(gateway: Gateway, shape: Shape) -> None:
|
||||
with wire_server(_always(openai_error(503))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"openai/{_CODEX_BACKEND}", api_base=wire.url + "/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/messages", _messages_body(model, _prompt(), stream=shape == "stream")
|
||||
)
|
||||
_assert_api_error(response, status=503, message="scripted 503")
|
||||
_one_post(wire, "/v1/responses")
|
||||
|
||||
|
||||
_SETTLED_ANSWERS: Final = 4
|
||||
|
||||
|
||||
def _messages_to(gateway: Gateway, model: str) -> tuple[tuple[str, httpx.Response], ...]:
|
||||
prompts: Final = tuple(_prompt() for _ in range(_SETTLED_ANSWERS))
|
||||
return tuple(
|
||||
(prompt, gateway.request("POST", "/v1/messages", _messages_body(model, prompt, stream=False)))
|
||||
for prompt in prompts
|
||||
)
|
||||
|
||||
|
||||
def _failure_rows_of(models: tuple[str, str], expected: int) -> list[dict[str, JsonValue]]:
|
||||
return eventually(
|
||||
lambda: read_rows(
|
||||
"SELECT metadata->'error_information'->>'error_code' AS code, "
|
||||
"metadata->'error_information'->>'error_class' AS error_class, "
|
||||
"metadata->'error_information'->>'normalized_error' AS normalized_error "
|
||||
'FROM "LiteLLM_SpendLogs" WHERE status = %s AND model_group IN (%s, %s)',
|
||||
("failure", *models),
|
||||
),
|
||||
lambda rows: len(rows) >= expected,
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _paused_deployment(gateway: Gateway) -> Generator[tuple[Wire, str, str]]:
|
||||
unpaused: Final = Reply(body=message_json("msg_unpaused", _MODEL, _TEXT))
|
||||
with wire_server(_always(unpaused)) as wire, gateway.scenario() as scenario:
|
||||
identity: Final = f"audit-paused-{uuid.uuid4().hex}"
|
||||
name: Final = scenario.model(
|
||||
model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, model_info={"id": identity}
|
||||
)
|
||||
paused: Final = gateway.request("PATCH", f"/model/{identity}/update", {"blocked": True})
|
||||
assert paused.status_code == 200, paused.text
|
||||
yield wire, identity, name
|
||||
|
||||
|
||||
def _assert_refused_without_an_upstream_call(
|
||||
gateway: Gateway, wire: Wire, target: str, *, status: int, error_type: str, message: str
|
||||
) -> None:
|
||||
answers: Final = eventually(
|
||||
lambda: _messages_to(gateway, target),
|
||||
lambda settled: all(response.status_code != 200 for _, response in settled),
|
||||
seconds=90,
|
||||
)
|
||||
texts: Final = [response.text for _, response in answers]
|
||||
assert [response.status_code for _, response in answers] == [status] * _SETTLED_ANSWERS, texts
|
||||
errors: Final = [_anthropic_error(response) for _, response in answers]
|
||||
assert [error["type"] for error in errors] == [error_type] * _SETTLED_ANSWERS, texts
|
||||
assert all(message in string_value(error["message"]) for error in errors), texts
|
||||
reached: Final = {user_prompt(JSON_OBJECT.validate_json(request.body)) for request in wire.drain()}
|
||||
assert reached.isdisjoint(prompt for prompt, _ in answers), reached
|
||||
|
||||
|
||||
def test_messages_to_a_paused_deployment_named_by_its_id_answer_503_api_error(gateway: Gateway) -> None:
|
||||
with _paused_deployment(gateway) as (wire, identity, name):
|
||||
_assert_refused_without_an_upstream_call(
|
||||
gateway, wire, identity, status=503, error_type="api_error", message="is currently paused"
|
||||
)
|
||||
rows: Final = _failure_rows_of((identity, name), _SETTLED_ANSWERS)
|
||||
assert {(row["code"], row["error_class"], row["normalized_error"]) for row in rows} == {
|
||||
("503", "ServiceUnavailableError", _OVERLOAD_CLASS)
|
||||
}, rows
|
||||
|
||||
|
||||
def test_messages_to_a_paused_model_group_named_by_its_name_answer_403_permission_error(gateway: Gateway) -> None:
|
||||
with _paused_deployment(gateway) as (wire, _, name):
|
||||
_assert_refused_without_an_upstream_call(
|
||||
gateway, wire, name, status=403, error_type="permission_error", message="Model is blocked"
|
||||
)
|
||||
|
||||
|
||||
_MALFORMED_529_BODIES: Final = {
|
||||
"empty": b"",
|
||||
"text-5k": b"x" * 5000,
|
||||
"html": b"<html><body><h1>529</h1></body></html>",
|
||||
"api_error-type": b'{"type": "error", "error": {"type": "api_error", "message": "capacity"}}',
|
||||
}
|
||||
|
||||
|
||||
def _overloaded_for_sad_prompts(body: bytes) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
prompt: Final = user_prompt(object_value(json.loads(request.body)))
|
||||
if prompt.startswith("sad:"):
|
||||
return Reply(status=529, body=body, content_type="text/plain" if body[:1] != b"{" else "application/json")
|
||||
return Reply(body=message_json(f"msg_{prompt}", _MODEL, _TEXT))
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body", _MALFORMED_529_BODIES.keys())
|
||||
def test_messages_529_with_a_malformed_body_still_answers_overloaded_error_and_unrelated_calls_keep_working(
|
||||
gateway: Gateway, body: str
|
||||
) -> None:
|
||||
with wire_server(_overloaded_for_sad_prompts(_MALFORMED_529_BODIES[body])) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
sad: Final = gateway.request("POST", "/v1/messages", _messages_body(model, "sad:" + _prompt(), stream=False))
|
||||
assert sad.status_code == 529, sad.text
|
||||
assert _anthropic_error(sad)["type"] == "overloaded_error", sad.text
|
||||
prompt: Final = _prompt()
|
||||
healthy: Final = gateway.request("POST", "/v1/messages", _messages_body(model, prompt, stream=False))
|
||||
assert healthy.status_code == 200, healthy.text
|
||||
assert object_value(json.loads(healthy.content))["id"] == f"msg_{prompt}", healthy.text
|
||||
assert [unquote(request.target) for request in wire.drain()] == ["/v1/messages"] * 2
|
||||
assert sorted(_spend_rows(model, 2), key=lambda row: str(row["status"])) == [
|
||||
{
|
||||
"status": "failure",
|
||||
"code": "529",
|
||||
"error_class": "InternalServerError",
|
||||
"normalized_error": _OVERLOAD_CLASS,
|
||||
},
|
||||
{"status": "success", "code": None, "error_class": None, "normalized_error": None},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("upstream_status", "error_type"), ((500, "api_error"), (429, "rate_limit_error")))
|
||||
def test_messages_other_upstream_statuses_keep_their_status_and_type(
|
||||
gateway: Gateway, upstream_status: int, error_type: str
|
||||
) -> None:
|
||||
with wire_server(_always(status_reply(upstream_status))) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = gateway.request("POST", "/v1/messages", _messages_body(model, _prompt(), stream=True))
|
||||
assert response.status_code == upstream_status, response.text
|
||||
assert _anthropic_error(response)["type"] == error_type, response.text
|
||||
_one_post(wire, "/v1/messages")
|
||||
|
||||
|
||||
def _anthropic_rejection(status: int, error_type: str, message: str) -> Reply:
|
||||
return Reply(
|
||||
status=status, body=json.dumps({"type": "error", "error": {"type": error_type, "message": message}}).encode()
|
||||
)
|
||||
|
||||
|
||||
_REJECTIONS_SAYING_OVERLOADED: Final = {
|
||||
"400-tool-name": _anthropic_rejection(
|
||||
400, "invalid_request_error", "tools.0.name: 'Overloaded' is not a valid tool name"
|
||||
),
|
||||
"404-model": _anthropic_rejection(404, "not_found_error", "model: claude-Overloaded"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ("non-stream", "stream"))
|
||||
@pytest.mark.parametrize("rejection", _REJECTIONS_SAYING_OVERLOADED.keys())
|
||||
def test_messages_4xx_whose_text_says_overloaded_is_not_answered_as_an_overload(
|
||||
gateway: Gateway, rejection: str, shape: Shape
|
||||
) -> None:
|
||||
with wire_server(_always(_REJECTIONS_SAYING_OVERLOADED[rejection])) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/messages", _messages_body(model, _prompt(), stream=shape == "stream")
|
||||
)
|
||||
assert response.status_code == 500, response.text
|
||||
error: Final = _anthropic_error(response)
|
||||
assert error["type"] == "api_error", response.text
|
||||
assert "Overloaded" in string_value(error["message"]), response.text
|
||||
_one_post(wire, "/v1/messages")
|
||||
assert _spend_rows(model, 1) == [
|
||||
{
|
||||
"status": "failure",
|
||||
"code": "500",
|
||||
"error_class": "InternalServerError",
|
||||
"normalized_error": _OVERLOAD_CLASS,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _string_values(cache: Redis) -> tuple[bytes, ...]:
|
||||
keys: Final = tuple(key for key in cache.scan_iter(count=1000) if cache.type(key) == b"string")
|
||||
return tuple(value for value in cache.mget(keys) if value is not None) if keys else ()
|
||||
|
||||
|
||||
def _cached_somewhere(served_id: str) -> bool:
|
||||
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
|
||||
return any(served_id.encode() in value for value in _string_values(cache))
|
||||
|
||||
|
||||
def _healthy_until(overloaded: threading.Event) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if overloaded.is_set():
|
||||
return status_reply(529)
|
||||
prompt: Final = user_prompt(object_value(json.loads(request.body)))
|
||||
return stream_reply(message_stream(f"msg_{prompt}", _MODEL, _TEXT))
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def test_messages_cache_hit_is_served_while_the_upstream_is_overloaded(gateway: Gateway) -> None:
|
||||
overloaded: Final = threading.Event()
|
||||
with wire_server(_healthy_until(overloaded)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY)
|
||||
prompt: Final = _prompt()
|
||||
body: Final = _messages_body(model, prompt, stream=True)
|
||||
first: Final = gateway.request("POST", "/v1/messages", body)
|
||||
assert first.status_code == 200, first.text
|
||||
assert event_types(parse_sse(first.text)) == LIFECYCLE, first.text
|
||||
assert message_id(parse_sse(first.text)) == f"msg_{prompt}", first.text
|
||||
eventually(lambda: _cached_somewhere(f"msg_{prompt}"), bool, seconds=30)
|
||||
overloaded.set()
|
||||
replay: Final = gateway.request("POST", "/v1/messages", body)
|
||||
assert replay.status_code == 200, replay.text
|
||||
assert message_id(parse_sse(replay.text)) == f"msg_{prompt}", replay.text
|
||||
fresh: Final = gateway.request("POST", "/v1/messages", _messages_body(model, _prompt(), stream=True))
|
||||
_assert_overloaded(fresh, message="scripted 529")
|
||||
assert [unquote(request.target) for request in wire.drain()] == ["/v1/messages"] * 2
|
||||
|
|
@ -135,6 +135,13 @@ def _drop_or_500(request: Request, marker: str) -> Reply:
|
|||
return openai_error(500)
|
||||
|
||||
|
||||
def _overloaded_while(outage: threading.Event) -> Callable[[Request, str], Reply]:
|
||||
def first_attempt(request: Request, marker: str) -> Reply:
|
||||
return status_reply(529) if outage.is_set() else _served(request, marker, 1)
|
||||
|
||||
return first_attempt
|
||||
|
||||
|
||||
def _outage(statuses: Mapping[str, int]) -> Callable[[Request, str], Reply]:
|
||||
def first_attempt(request: Request, marker: str) -> Reply:
|
||||
if request.target == "/v1/messages":
|
||||
|
|
@ -165,7 +172,7 @@ class _Upstream:
|
|||
return self.first_attempt(request, marker)
|
||||
|
||||
|
||||
def _config(wire: Wire, directory: Path, models: _Models) -> Path:
|
||||
def _config(wire: Wire, directory: Path, models: _Models, *, retries: int = 1) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["model_list"] = [
|
||||
{
|
||||
|
|
@ -174,7 +181,7 @@ def _config(wire: Wire, directory: Path, models: _Models) -> Path:
|
|||
"model": f"anthropic/{_ANTHROPIC_BACKEND}",
|
||||
"api_base": wire.url,
|
||||
"api_key": _ANTHROPIC_KEY,
|
||||
"num_retries": 1,
|
||||
"num_retries": retries,
|
||||
},
|
||||
},
|
||||
{
|
||||
|
|
@ -183,7 +190,7 @@ def _config(wire: Wire, directory: Path, models: _Models) -> Path:
|
|||
"model": f"openai/{_OPENAI_BACKEND}",
|
||||
"api_base": wire.url + "/v1",
|
||||
"api_key": _OPENAI_KEY,
|
||||
"num_retries": 1,
|
||||
"num_retries": retries,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
|
@ -408,3 +415,68 @@ async def test_outage_on_every_first_attempt_is_absorbed_by_the_deployment_budge
|
|||
_assert_completed_on_the_second_attempt(item)
|
||||
_assert_each_served_id_landed_exactly_once(models, served)
|
||||
assert all(upstream.attempts.count(call.marker) == 2 for call in calls)
|
||||
|
||||
|
||||
def _messages_rows(model: str, expected: int) -> list[dict[str, JsonValue]]:
|
||||
return eventually(
|
||||
lambda: read_rows(
|
||||
"SELECT request_id, status, metadata->'error_information'->>'error_code' AS code "
|
||||
'FROM "LiteLLM_SpendLogs" WHERE model_group=%s',
|
||||
(model,),
|
||||
),
|
||||
lambda found: len(found) >= expected,
|
||||
seconds=90,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(300)
|
||||
async def test_overloaded_outage_during_a_burst_with_a_worker_sigkill_answers_529_and_lands_every_call_once(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
models: Final = _models()
|
||||
calls: Final = _calls((("messages", True, 12), ("messages", False, 6)))
|
||||
release: Final = threading.Event()
|
||||
held: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
outage: Final = threading.Event()
|
||||
outage.set()
|
||||
upstream: Final = _Upstream(Attempts(), _overloaded_while(outage), held, release)
|
||||
with wire_server(upstream) as wire:
|
||||
config: Final = _config(wire, tmp_path, models, retries=0)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
base_url: Final = str(candidate.client.base_url)
|
||||
workers: Final = eventually(
|
||||
lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
|
||||
lambda pids: len(pids) == 2,
|
||||
seconds=30,
|
||||
)
|
||||
burst: Final = asyncio.create_task(
|
||||
_burst(base_url, candidate.key, models, calls, tolerate_transport_errors=True)
|
||||
)
|
||||
await asyncio.to_thread(eventually, held.qsize, lambda size: size == len(calls), 90)
|
||||
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers})
|
||||
assert sum(held_by.values()) == len(calls), held_by
|
||||
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
|
||||
with owned.replacing_worker(victim_pid) as kill:
|
||||
kill.freeze()
|
||||
await asyncio.to_thread(kill)
|
||||
release.set()
|
||||
served: Final = await burst
|
||||
assert held_by[survivor_pid] >= 5, held_by
|
||||
assert len(served) == held_by[survivor_pid], (held_by, len(served))
|
||||
for item in served:
|
||||
assert item.status == 529, (item.call, item.text)
|
||||
assert "overloaded_error" in item.text, (item.call, item.text)
|
||||
assert "content_block_delta" not in item.text, (item.call, item.text)
|
||||
outage.clear()
|
||||
follow_up: Final = _Call(endpoint="messages", stream=True, marker=uuid.uuid4().hex)
|
||||
(answered,) = await _burst(base_url, candidate.key, models, (follow_up,))
|
||||
assert answered.status == 200, answered.text
|
||||
assert message_id(parse_sse(answered.text)) == f"msg_{follow_up.marker}_a1", answered.text
|
||||
rows: Final = _messages_rows(models.messages, len(served) + 1)
|
||||
assert sorted(string_value(row["status"]) for row in rows) == ["failure"] * len(served) + ["success"], rows
|
||||
assert {row["code"] for row in rows if row["status"] == "failure"} == {"529"}, rows
|
||||
assert [row["request_id"] for row in rows if row["status"] == "success"] == [
|
||||
f"msg_{follow_up.marker}_a1"
|
||||
], rows
|
||||
assert all(upstream.attempts.count(item.call.marker) == 1 for item in (*served, answered))
|
||||
|
|
|
|||
|
|
@ -397,14 +397,14 @@ def test_retry_rejected_before_it_opens_counts_against_the_same_budget(gateway:
|
|||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1)
|
||||
response: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "deployment"))
|
||||
_assert_rejected_before_content(response, wire, status=500, attempts=2, error="error")
|
||||
_assert_rejected_before_content(response, wire, status=529, attempts=2, error="overloaded_error")
|
||||
rows: Final = eventually(
|
||||
lambda: _request_rows(gateway, response.headers["x-litellm-call-id"]),
|
||||
lambda values: len(values) >= 1,
|
||||
seconds=70,
|
||||
)
|
||||
rejected: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "request", budget=0))
|
||||
_assert_rejected_before_content(rejected, wire, status=500, attempts=1, error="error")
|
||||
_assert_rejected_before_content(rejected, wire, status=529, attempts=1, error="overloaded_error")
|
||||
rejection_rows: Final = eventually(
|
||||
lambda: _request_rows(gateway, rejected.headers["x-litellm-call-id"]),
|
||||
lambda values: len(values) >= 1,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,404 @@
|
|||
import gzip
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import unquote
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.anthropic_sse import error_body
|
||||
from integration._support.client import (
|
||||
JSON_OBJECT,
|
||||
Gateway,
|
||||
JsonValue,
|
||||
eventually,
|
||||
gateway_from_environment,
|
||||
object_value,
|
||||
string_value,
|
||||
)
|
||||
from integration._support.otlp_sink import Span, SpanSinks, recorded_spans, spans_for_trace
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.prometheus_series import Sample, scrape
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_MODEL: Final = "claude-sonnet-4-5-20250929"
|
||||
_BEDROCK_MODEL: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
_BEDROCK_503: Final = b'{"message": "Bedrock is unable to process your request."}'
|
||||
_MESSAGES_ANTHROPIC: Final = "audit-signals-messages-anthropic"
|
||||
_MESSAGES_BEDROCK: Final = "audit-signals-messages-bedrock"
|
||||
_CHAT_ANTHROPIC: Final = "audit-signals-chat-anthropic"
|
||||
_TRACED_MESSAGES: Final = "audit-signals-traced-messages"
|
||||
_TRACED_CHAT: Final = "audit-signals-traced-chat"
|
||||
_DATADOG_MESSAGES_ANTHROPIC: Final = "audit-signals-datadog-messages-anthropic"
|
||||
_DATADOG_MESSAGES_BEDROCK: Final = "audit-signals-datadog-messages-bedrock"
|
||||
_DATADOG_CHAT_ANTHROPIC: Final = "audit-signals-datadog-chat-anthropic"
|
||||
_GUARDED_MESSAGES: Final = "audit-signals-guarded-messages"
|
||||
_GUARDED_CHAT: Final = "audit-signals-guarded-chat"
|
||||
_AKTO_DOWN: Final = "audit-signals-akto-down"
|
||||
_AKTO_UNREACHABLE: Final = "Akto guardrail service unreachable"
|
||||
_LATE_PASS_THROUGH_BACKEND: Final = "claude-late-pass-through"
|
||||
_KEEPALIVE_SECONDS: Final = 0.2
|
||||
_LATE_ANSWER_SECONDS: Final = 1.0
|
||||
_OVERLOADED: Final = Reply(status=529, body=error_body(529, "Overloaded"))
|
||||
_SERVER: Final = 2
|
||||
|
||||
AuditConfigWriter = Callable[[Path, Mapping[str, JsonValue]], Path]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
proxy: Gateway
|
||||
wire: Wire
|
||||
spans: str
|
||||
datadog: Wire
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
assert request.method == "POST", request
|
||||
if unquote(request.target).startswith("/model/"):
|
||||
return Reply(status=503, body=_BEDROCK_503)
|
||||
if json.loads(request.body).get("model") == _LATE_PASS_THROUGH_BACKEND:
|
||||
time.sleep(_LATE_ANSWER_SECONDS)
|
||||
return _OVERLOADED
|
||||
|
||||
|
||||
def _datadog_intake(request: Request) -> Reply:
|
||||
assert request.target == "/api/v2/logs", request.target
|
||||
return Reply(status=202, body=b"{}")
|
||||
|
||||
|
||||
def _config(directory: Path, tracing: Path, wire: Wire) -> Path:
|
||||
configuration: Final = JSON_OBJECT.validate_python(yaml.safe_load(tracing.read_text()))
|
||||
anthropic: Final = {"model": f"anthropic/{_MODEL}", "api_base": wire.url, "api_key": "synthetic-anthropic-key"}
|
||||
bedrock: Final = {
|
||||
"model": f"bedrock/{_BEDROCK_MODEL}",
|
||||
"api_base": wire.url,
|
||||
"api_key": "synthetic-bedrock-bearer",
|
||||
"aws_region_name": "us-east-1",
|
||||
}
|
||||
deployments: Final = (
|
||||
(_MESSAGES_ANTHROPIC, anthropic),
|
||||
(_MESSAGES_BEDROCK, bedrock),
|
||||
(_CHAT_ANTHROPIC, anthropic),
|
||||
(_TRACED_MESSAGES, anthropic),
|
||||
(_TRACED_CHAT, anthropic),
|
||||
(_DATADOG_MESSAGES_ANTHROPIC, anthropic),
|
||||
(_DATADOG_MESSAGES_BEDROCK, bedrock),
|
||||
(_DATADOG_CHAT_ANTHROPIC, anthropic),
|
||||
(_GUARDED_MESSAGES, anthropic),
|
||||
(_GUARDED_CHAT, anthropic),
|
||||
)
|
||||
akto_down: Final = {
|
||||
"guardrail_name": _AKTO_DOWN,
|
||||
"litellm_params": {
|
||||
"guardrail": "akto",
|
||||
"mode": "pre_call",
|
||||
"akto_base_url": "http://127.0.0.1:9",
|
||||
"akto_api_key": "synthetic-akto-key",
|
||||
"unreachable_fallback": "fail_closed",
|
||||
},
|
||||
}
|
||||
path: Final = directory / "anthropic-overloaded-failure-signals.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**configuration,
|
||||
"model_list": [{"model_name": name, "litellm_params": params} for name, params in deployments],
|
||||
"router_settings": {**object_value(configuration.get("router_settings", {})), "num_retries": 0},
|
||||
"litellm_settings": {
|
||||
**object_value(configuration["litellm_settings"]),
|
||||
"sse_keepalive_ping_interval_seconds": _KEEPALIVE_SECONDS,
|
||||
},
|
||||
"guardrails": [akto_down],
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(
|
||||
tmp_path_factory: pytest.TempPathFactory, audit_sinks: SpanSinks, otel_audit_config: AuditConfigWriter
|
||||
) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("anthropic-overloaded-failure-signals")
|
||||
multiprocess: Final = directory / "prometheus"
|
||||
multiprocess.mkdir()
|
||||
tracing: Final = otel_audit_config(directory, {"callbacks": ["otel", "prometheus", "datadog"]})
|
||||
with (
|
||||
wire_server(_respond) as wire,
|
||||
wire_server(_datadog_intake) as datadog,
|
||||
gateway_from_environment() as shared,
|
||||
owned_proxy(
|
||||
shared,
|
||||
directory,
|
||||
{
|
||||
"OTEL_EXPORTER": "http/json",
|
||||
"OTEL_ENDPOINT": audit_sinks.operator,
|
||||
"PROMETHEUS_MULTIPROC_DIR": str(multiprocess),
|
||||
"DD_API_KEY": "synthetic-datadog-key",
|
||||
"DD_SITE": "datadog.invalid",
|
||||
"DD_BASE_URL": datadog.url,
|
||||
"DD_BATCH_SIZE": "1",
|
||||
"ANTHROPIC_API_BASE": wire.url,
|
||||
"ANTHROPIC_API_KEY": "synthetic-anthropic-key",
|
||||
},
|
||||
config=_config(directory, tracing, wire),
|
||||
workers=2,
|
||||
) as owned,
|
||||
):
|
||||
yield _Rig(owned, wire, audit_sinks.operator, datadog)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def v2_rig(
|
||||
tmp_path_factory: pytest.TempPathFactory, audit_sinks: SpanSinks, otel_audit_config: AuditConfigWriter
|
||||
) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("anthropic-overloaded-failure-signals-v2")
|
||||
tracing: Final = otel_audit_config(directory, {"callbacks": ["otel"]})
|
||||
with (
|
||||
wire_server(_respond) as wire,
|
||||
wire_server(_datadog_intake) as datadog,
|
||||
gateway_from_environment() as shared,
|
||||
owned_proxy(
|
||||
shared,
|
||||
directory,
|
||||
{"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, "LITELLM_OTEL_V2": "1"},
|
||||
config=_config(directory, tracing, wire),
|
||||
workers=2,
|
||||
) as owned,
|
||||
):
|
||||
yield _Rig(owned, wire, audit_sinks.operator, datadog)
|
||||
|
||||
|
||||
def _send(
|
||||
rig: _Rig,
|
||||
route: str,
|
||||
model: str,
|
||||
trace_id: str | None = None,
|
||||
*,
|
||||
key: str | None = None,
|
||||
extra: Mapping[str, JsonValue] | None = None,
|
||||
) -> httpx.Response:
|
||||
traceparent: Final = {} if trace_id is None else {"traceparent": f"00-{trace_id}-{uuid.uuid4().hex[:16]}-01"}
|
||||
return rig.proxy.request(
|
||||
"POST",
|
||||
route,
|
||||
{
|
||||
"model": model,
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": f"signals-{uuid.uuid4().hex}"}],
|
||||
"num_retries": 0,
|
||||
**(extra or {}),
|
||||
},
|
||||
key=key,
|
||||
headers={"anthropic-version": "2023-06-01", **traceparent},
|
||||
)
|
||||
|
||||
|
||||
def _series(samples: Sequence[Sample], name: str, model: str) -> tuple[Sample, ...]:
|
||||
return tuple(sample for sample in samples if sample.name == name and sample.labels.get("requested_model") == model)
|
||||
|
||||
|
||||
def _until(rig: _Rig, name: str, model: str) -> tuple[Sample, ...]:
|
||||
return eventually(lambda: _series(scrape(rig.proxy), name, model), bool, seconds=30)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "model", "answered", "deployment_status", "deployment_class"),
|
||||
(
|
||||
pytest.param("/v1/messages", _MESSAGES_ANTHROPIC, 529, "500", "InternalServerError", id="messages-anthropic"),
|
||||
pytest.param("/v1/messages", _MESSAGES_BEDROCK, 503, "503", "ServiceUnavailableError", id="messages-bedrock"),
|
||||
pytest.param("/v1/chat/completions", _CHAT_ANTHROPIC, 500, "500", "InternalServerError", id="chat-anthropic"),
|
||||
),
|
||||
)
|
||||
def test_the_failed_request_metric_counts_the_answered_status_and_the_deployment_metric_the_deployment_status(
|
||||
rig: _Rig, route: str, model: str, answered: int, deployment_status: str, deployment_class: str
|
||||
) -> None:
|
||||
rig.wire.drain()
|
||||
response: Final = _send(rig, route, model)
|
||||
assert response.status_code == answered, response.text
|
||||
assert len(rig.wire.drain()) == 1
|
||||
failed: Final = _until(rig, "litellm_proxy_failed_requests_metric_total", model)
|
||||
assert [(sample.labels["exception_status"], sample.labels["route"], sample.value) for sample in failed] == [
|
||||
(str(answered), route, 1.0)
|
||||
], failed
|
||||
deployment: Final = _until(rig, "litellm_deployment_failure_responses_total", model)
|
||||
assert [(sample.labels["exception_status"], sample.value) for sample in deployment] == [(deployment_status, 1.0)], (
|
||||
deployment
|
||||
)
|
||||
assert all(deployment_class in sample.labels["exception_class"] for sample in deployment), deployment
|
||||
|
||||
|
||||
def _server_spans(rig: _Rig, since: int, trace_id: str) -> tuple[Span, ...]:
|
||||
traced: Final = spans_for_trace(recorded_spans(rig.spans, since)[1], trace_id)
|
||||
return tuple(span for span in traced if span["kind"] == _SERVER)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "model", "answered"),
|
||||
(
|
||||
pytest.param("/v1/messages", _TRACED_MESSAGES, 529, id="messages"),
|
||||
pytest.param("/v1/chat/completions", _TRACED_CHAT, 500, id="chat"),
|
||||
),
|
||||
)
|
||||
def test_the_server_span_carries_the_status_the_route_answered(
|
||||
rig: _Rig, route: str, model: str, answered: int
|
||||
) -> None:
|
||||
rig.wire.drain()
|
||||
since, _ = recorded_spans(rig.spans)
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
response: Final = _send(rig, route, model, trace_id)
|
||||
assert response.status_code == answered, response.text
|
||||
assert len(rig.wire.drain()) == 1
|
||||
server: Final = eventually(
|
||||
lambda: _server_spans(rig, since, trace_id),
|
||||
lambda spans: any("error.code" in span["attributes"] for span in spans),
|
||||
seconds=40,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
assert [
|
||||
(str(span["attributes"].get("http.response.status_code")), str(span["attributes"].get("error.code")))
|
||||
for span in server
|
||||
] == [(str(answered), str(answered))], server
|
||||
|
||||
|
||||
def _inflated(request: Request) -> bytes:
|
||||
return gzip.decompress(request.body) if request.headers.get("content-encoding") == "gzip" else request.body
|
||||
|
||||
|
||||
def _datadog_events(batches: Sequence[Request]) -> Iterator[JsonValue]:
|
||||
for batch in batches:
|
||||
yield from json.loads(_inflated(batch))
|
||||
|
||||
|
||||
def _datadog_messages(batches: Sequence[Request]) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
JSON_OBJECT.validate_json(string_value(object_value(event)["message"])) for event in _datadog_events(batches)
|
||||
)
|
||||
|
||||
|
||||
def _keyed_by(message: Mapping[str, JsonValue], alias: str) -> bool:
|
||||
proxy_side: Final = message.get("user_api_key_dict")
|
||||
llm_side: Final = message.get("metadata")
|
||||
return any(isinstance(side, dict) and side.get("user_api_key_alias") == alias for side in (proxy_side, llm_side))
|
||||
|
||||
|
||||
def _datadog_status_codes(rig: _Rig, alias: str, expected: int) -> tuple[tuple[str, str], ...]:
|
||||
batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches
|
||||
|
||||
def delivered() -> tuple[tuple[str, str], ...]:
|
||||
batches.extend(rig.datadog.drain())
|
||||
return tuple(
|
||||
sorted(
|
||||
("proxy", str(message["status_code"]))
|
||||
if "user_api_key_dict" in message
|
||||
else ("llm", str(object_value(message["error_information"])["error_code"]))
|
||||
for message in _datadog_messages(batches)
|
||||
if _keyed_by(message, alias)
|
||||
)
|
||||
)
|
||||
|
||||
return eventually(delivered, lambda codes: len(codes) >= expected, seconds=40, return_last_on_timeout=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "model", "answered", "deployment_status"),
|
||||
(
|
||||
pytest.param("/v1/messages", _DATADOG_MESSAGES_ANTHROPIC, 529, "500", id="messages-anthropic"),
|
||||
pytest.param("/v1/messages", _DATADOG_MESSAGES_BEDROCK, 503, "503", id="messages-bedrock"),
|
||||
pytest.param("/v1/chat/completions", _DATADOG_CHAT_ANTHROPIC, 500, "500", id="chat-anthropic"),
|
||||
),
|
||||
)
|
||||
def test_the_datadog_proxy_failure_log_carries_the_answered_status_and_the_llm_failure_log_the_deployment_status(
|
||||
rig: _Rig, route: str, model: str, answered: int, deployment_status: str
|
||||
) -> None:
|
||||
rig.wire.drain()
|
||||
rig.datadog.drain()
|
||||
alias: Final = f"audit-signals-datadog-{uuid.uuid4().hex}"
|
||||
with rig.proxy.scenario() as scenario:
|
||||
response: Final = _send(rig, route, model, key=scenario.key(key_alias=alias, models=[model]))
|
||||
assert response.status_code == answered, response.text
|
||||
assert len(rig.wire.drain()) == 1
|
||||
assert _datadog_status_codes(rig, alias, 2) == (("llm", deployment_status), ("proxy", str(answered)))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "model", "answered"),
|
||||
(
|
||||
pytest.param("/v1/messages", _GUARDED_MESSAGES, 503, id="messages"),
|
||||
pytest.param("/v1/chat/completions", _GUARDED_CHAT, 503, id="chat"),
|
||||
),
|
||||
)
|
||||
def test_a_fail_closed_guardrail_outage_is_answered_and_recorded_with_the_routes_status(
|
||||
rig: _Rig, route: str, model: str, answered: int
|
||||
) -> None:
|
||||
rig.wire.drain()
|
||||
rig.datadog.drain()
|
||||
alias: Final = f"audit-signals-guarded-{uuid.uuid4().hex}"
|
||||
with rig.proxy.scenario() as scenario:
|
||||
response: Final = _send(
|
||||
rig, route, model, key=scenario.key(key_alias=alias, models=[model]), extra={"guardrails": [_AKTO_DOWN]}
|
||||
)
|
||||
assert response.status_code == answered, response.text
|
||||
assert _AKTO_UNREACHABLE in response.text, response.text
|
||||
assert rig.wire.drain() == ()
|
||||
failed: Final = _until(rig, "litellm_proxy_failed_requests_metric_total", model)
|
||||
assert [(sample.labels["exception_status"], sample.labels["route"], sample.value) for sample in failed] == [
|
||||
(str(answered), route, 1.0)
|
||||
], failed
|
||||
assert _datadog_status_codes(rig, alias, 2) == (("llm", str(answered)), ("proxy", str(answered)))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "model", "answered"),
|
||||
(
|
||||
pytest.param("/v1/messages", _TRACED_MESSAGES, 529, id="messages-anthropic"),
|
||||
pytest.param("/v1/messages", _MESSAGES_BEDROCK, 503, id="messages-bedrock"),
|
||||
pytest.param("/v1/chat/completions", _TRACED_CHAT, 500, id="chat-anthropic"),
|
||||
),
|
||||
)
|
||||
def test_the_otel_v2_server_span_carries_the_status_the_route_answered(
|
||||
v2_rig: _Rig, route: str, model: str, answered: int
|
||||
) -> None:
|
||||
v2_rig.wire.drain()
|
||||
since, _ = recorded_spans(v2_rig.spans)
|
||||
trace_id: Final = uuid.uuid4().hex
|
||||
response: Final = _send(v2_rig, route, model, trace_id)
|
||||
assert response.status_code == answered, response.text
|
||||
assert len(v2_rig.wire.drain()) == 1
|
||||
server: Final = eventually(
|
||||
lambda: _server_spans(v2_rig, since, trace_id),
|
||||
lambda spans: any("litellm.provider.error.code" in span["attributes"] for span in spans),
|
||||
seconds=40,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
assert [str(span["attributes"].get("litellm.provider.error.code")) for span in server] == [str(answered)], server
|
||||
|
||||
|
||||
def test_an_anthropic_pass_through_stream_that_fails_after_the_keepalive_opened_relays_the_upstream_body(
|
||||
rig: _Rig,
|
||||
) -> None:
|
||||
rig.wire.drain()
|
||||
response: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/anthropic/v1/messages",
|
||||
{
|
||||
"model": _LATE_PASS_THROUGH_BACKEND,
|
||||
"max_tokens": 16,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": f"signals-{uuid.uuid4().hex}"}],
|
||||
},
|
||||
headers={"anthropic-version": "2023-06-01"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
pings, _, relayed = response.text.rpartition("\n\n")
|
||||
assert pings and set(pings.split("\n")) <= {"", ": ping"}, response.text
|
||||
assert relayed == _OVERLOADED.body.decode(), response.text
|
||||
assert [unquote(request.target) for request in rig.wire.drain()] == ["/v1/messages"]
|
||||
|
|
@ -0,0 +1,121 @@
|
|||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, gateway_from_environment
|
||||
from integration._support.forward_proxy import ForwardProxy, tunnelling_forward_proxy
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.tls import server_context, write_self_signed_cert
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
_LOCATION: Final = "us-central1"
|
||||
_HOST: Final = f"{_LOCATION}-aiplatform.googleapis.com"
|
||||
_PROJECT: Final = "audit-vertex-pass-through"
|
||||
_STREAMING_MODEL: Final = "gemini-3.8-flash"
|
||||
_UNAVAILABLE_MODEL: Final = "gemini-3.8-flash-unavailable"
|
||||
_VENDOR_TOKEN: Final = "Bearer synthetic-vertex-token"
|
||||
_CHUNKS: Final = (
|
||||
b'data: {"candidates": [{"content": {"role": "model", "parts": [{"text": "Hel"}]}}]}\r\n\r\n',
|
||||
b'data: {"candidates": [{"content": {"role": "model", "parts": [{"text": "lo"}]}, "finishReason": "STOP"}],'
|
||||
b' "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}}\r\n\r\n',
|
||||
)
|
||||
_UNAVAILABLE_MESSAGE: Final = "The service is currently unavailable."
|
||||
_UNAVAILABLE_BODY: Final = json.dumps(
|
||||
{"error": {"code": 503, "message": _UNAVAILABLE_MESSAGE, "status": "UNAVAILABLE"}}
|
||||
).encode()
|
||||
_SCRUBBED: Final = frozenset(
|
||||
{
|
||||
"HTTPS_PROXY",
|
||||
"HTTP_PROXY",
|
||||
"ALL_PROXY",
|
||||
"NO_PROXY",
|
||||
"SSL_CERT_FILE",
|
||||
"DEFAULT_VERTEXAI_PROJECT",
|
||||
"DEFAULT_VERTEXAI_LOCATION",
|
||||
"DEFAULT_GOOGLE_APPLICATION_CREDENTIALS",
|
||||
"GOOGLE_APPLICATION_CREDENTIALS",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Vertex:
|
||||
wire: Wire
|
||||
tunnel: ForwardProxy
|
||||
proxy: Gateway
|
||||
|
||||
def stream(self, model: str, marker: str) -> httpx.Response:
|
||||
return self.proxy.request(
|
||||
"POST",
|
||||
f"/vertex_ai/v1/{_model_path(model)}:streamGenerateContent",
|
||||
{"contents": [{"role": "user", "parts": [{"text": marker}]}]},
|
||||
headers={"Authorization": _VENDOR_TOKEN, "x-litellm-api-key": f"Bearer {self.proxy.key}"},
|
||||
)
|
||||
|
||||
|
||||
def _model_path(model: str) -> str:
|
||||
return f"projects/{_PROJECT}/locations/{_LOCATION}/publishers/google/models/{model}"
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
if request.target.startswith(f"/v1/{_model_path(_UNAVAILABLE_MODEL)}:"):
|
||||
return Reply(status=503, body=_UNAVAILABLE_BODY)
|
||||
if request.target.startswith(f"/v1/{_model_path(_STREAMING_MODEL)}:"):
|
||||
return Reply(content_type="text/event-stream", chunks=_CHUNKS)
|
||||
return Reply(status=404, body=b'{"error": "no scripted Vertex model"}')
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vertex(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Vertex]:
|
||||
directory: Final = tmp_path_factory.mktemp("vertex-pass-through-stream")
|
||||
cert, key = write_self_signed_cert(directory, (_HOST,))
|
||||
with (
|
||||
wire_server(_respond, tls=server_context(cert, key)) as wire,
|
||||
tunnelling_forward_proxy(int(wire.url.rsplit(":", 1)[1])) as tunnel,
|
||||
gateway_from_environment() as rig,
|
||||
owned_proxy(
|
||||
rig,
|
||||
directory,
|
||||
{"HTTPS_PROXY": tunnel.url, "NO_PROXY": "127.0.0.1,localhost", "SSL_CERT_FILE": str(cert)},
|
||||
remove_environment=tuple(name for name in os.environ if name.upper() in _SCRUBBED),
|
||||
workers=2,
|
||||
) as proxy,
|
||||
):
|
||||
yield _Vertex(wire, tunnel, proxy)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def fresh_edges(vertex: _Vertex) -> None:
|
||||
vertex.wire.drain()
|
||||
vertex.tunnel.drain()
|
||||
|
||||
|
||||
def test_vertex_pass_through_stream_relays_every_sse_chunk_byte_for_byte(vertex: _Vertex) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = vertex.stream(_STREAMING_MODEL, marker)
|
||||
sent: Final = vertex.wire.drain()
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.headers["content-type"].startswith("text/event-stream"), response.headers
|
||||
assert response.content == b"".join(_CHUNKS), response.text
|
||||
assert [request.target for request in sent] == [
|
||||
f"/v1/{_model_path(_STREAMING_MODEL)}:streamGenerateContent?alt=sse"
|
||||
], sent
|
||||
assert sent[0].headers["authorization"] == _VENDOR_TOKEN
|
||||
assert json.loads(sent[0].body) == {"contents": [{"role": "user", "parts": [{"text": marker}]}]}
|
||||
assert set(vertex.tunnel.targets()) == {f"{_HOST}:443"}
|
||||
|
||||
|
||||
def test_vertex_pass_through_stream_relays_an_upstream_503_as_503(vertex: _Vertex) -> None:
|
||||
response: Final = vertex.stream(_UNAVAILABLE_MODEL, uuid.uuid4().hex)
|
||||
|
||||
assert response.status_code == 503, response.text
|
||||
assert response.content == _UNAVAILABLE_BODY, response.text
|
||||
assert [request.target for request in vertex.wire.drain()] == [
|
||||
f"/v1/{_model_path(_UNAVAILABLE_MODEL)}:streamGenerateContent?alt=sse"
|
||||
]
|
||||
|
|
@ -0,0 +1,247 @@
|
|||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import unquote
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.anthropic_sse import (
|
||||
LIFECYCLE,
|
||||
error_body,
|
||||
event_types,
|
||||
message_id,
|
||||
message_stream,
|
||||
parse_sse,
|
||||
status_reply,
|
||||
stream_reply,
|
||||
user_prompt,
|
||||
)
|
||||
from integration._support.client import (
|
||||
Gateway,
|
||||
JsonValue,
|
||||
gateway_from_environment,
|
||||
object_value,
|
||||
string_value,
|
||||
)
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
_OVERLOADED: Final = "claude-overloaded"
|
||||
_HEALTHY: Final = "claude-healthy"
|
||||
_COUNTED: Final = "claude-counted"
|
||||
_COUNTED_OVERLOADED: Final = "claude-counted-overloaded"
|
||||
_COUNTED_RATE_LIMITED: Final = "claude-counted-rate-limited"
|
||||
_CHAT_ONLY_BACKEND: Final = "audit-chat-only"
|
||||
_API_KEY: Final = "synthetic-anthropic-key"
|
||||
_TEXT: Final = "Hello"
|
||||
_COOLDOWN_GROUP: Final = "audit-overloaded-cooldown"
|
||||
_COUNT_DISABLED_GROUP: Final = "audit-count-tokens-disabled"
|
||||
_COUNT_UPSTREAM_GROUP: Final = "audit-count-tokens-upstream-503"
|
||||
_COUNT_OVERLOADED_GROUP: Final = "audit-count-tokens-upstream-529"
|
||||
_COUNT_RATE_LIMITED_GROUP: Final = "audit-count-tokens-upstream-429"
|
||||
_GATEWAY: Final = "/claude_code_gateway"
|
||||
_GATEWAY_GROUPS: Final = {
|
||||
"non-stream": "audit-gateway-overloaded-non-stream",
|
||||
"stream": "audit-gateway-overloaded-stream",
|
||||
}
|
||||
_COUNT_STATUSES: Final = {_COUNTED: 503, _COUNTED_OVERLOADED: 529, _COUNTED_RATE_LIMITED: 429}
|
||||
_COUNT_TOKENS_TARGET: Final = "/v1/messages/count_tokens"
|
||||
_UPSTREAM_URL_PLACEHOLDER: Final = "upstream-url"
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
|
||||
def _served_id(marker: str) -> str:
|
||||
return f"msg_{_HEALTHY}_{marker}"
|
||||
|
||||
|
||||
def _respond(request: Request) -> Reply:
|
||||
body: Final = object_value(json.loads(request.body))
|
||||
backend: Final = str(body["model"])
|
||||
if unquote(request.target) == _COUNT_TOKENS_TARGET:
|
||||
status: Final = _COUNT_STATUSES[backend]
|
||||
return Reply(status=status, body=error_body(status, f"scripted {status}"))
|
||||
if backend == _OVERLOADED:
|
||||
return status_reply(529)
|
||||
if backend == _HEALTHY:
|
||||
return stream_reply(message_stream(_served_id(user_prompt(body)), backend, _TEXT))
|
||||
return Reply(status=404, body=b'{"error": "no scripted backend"}')
|
||||
|
||||
|
||||
def _deployment(name: str, model: str, api_base: str, api_key: str, **params: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"model_name": name, "litellm_params": {"model": model, "api_base": api_base, "api_key": api_key, **params}}
|
||||
|
||||
|
||||
def _config(wire: Wire, directory: Path) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["model_list"] = [
|
||||
_deployment(_COOLDOWN_GROUP, f"anthropic/{_OVERLOADED}", _UPSTREAM_URL_PLACEHOLDER, _API_KEY, order=1),
|
||||
_deployment(_COOLDOWN_GROUP, f"anthropic/{_HEALTHY}", _UPSTREAM_URL_PLACEHOLDER, _API_KEY, order=2),
|
||||
_deployment(
|
||||
_COUNT_DISABLED_GROUP,
|
||||
f"hosted_vllm/{_CHAT_ONLY_BACKEND}",
|
||||
_UPSTREAM_URL_PLACEHOLDER + "/v1",
|
||||
"synthetic-vllm-key",
|
||||
),
|
||||
_deployment(_COUNT_UPSTREAM_GROUP, f"anthropic/{_COUNTED}", _UPSTREAM_URL_PLACEHOLDER, _API_KEY),
|
||||
_deployment(_COUNT_OVERLOADED_GROUP, f"anthropic/{_COUNTED_OVERLOADED}", _UPSTREAM_URL_PLACEHOLDER, _API_KEY),
|
||||
_deployment(
|
||||
_COUNT_RATE_LIMITED_GROUP, f"anthropic/{_COUNTED_RATE_LIMITED}", _UPSTREAM_URL_PLACEHOLDER, _API_KEY
|
||||
),
|
||||
*(
|
||||
_deployment(group, f"anthropic/{_OVERLOADED}", _UPSTREAM_URL_PLACEHOLDER, _API_KEY)
|
||||
for group in _GATEWAY_GROUPS.values()
|
||||
),
|
||||
]
|
||||
config["router_settings"] = {
|
||||
"num_retries": 0,
|
||||
"allowed_fails_policy": {"InternalServerErrorAllowedFails": 0},
|
||||
"cooldown_time": 120,
|
||||
"redis_host": os.environ["REDIS_HOST"],
|
||||
"redis_port": int(os.environ["REDIS_PORT"]),
|
||||
}
|
||||
config["litellm_settings"]["disable_token_counter"] = True
|
||||
config["general_settings"]["enable_claude_code_gateway"] = True
|
||||
path: Final = directory / "anthropic-overloaded-cooldown.yaml"
|
||||
path.write_text(yaml.safe_dump(config).replace(_UPSTREAM_URL_PLACEHOLDER, wire.url))
|
||||
return path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Rig:
|
||||
proxy: Gateway
|
||||
wire: Wire
|
||||
|
||||
def messages(self, path: str, model: str, marker: str, **extra: JsonValue) -> httpx.Response:
|
||||
return self.proxy.request(
|
||||
"POST",
|
||||
path,
|
||||
{
|
||||
"model": model,
|
||||
"max_tokens": 16,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
**extra,
|
||||
},
|
||||
headers={"anthropic-version": "2023-06-01"},
|
||||
)
|
||||
|
||||
def count_tokens(self, model: str, prefix: str = "") -> httpx.Response:
|
||||
return self.proxy.request(
|
||||
"POST",
|
||||
prefix + _COUNT_TOKENS_TARGET,
|
||||
{"model": model, "messages": [{"role": "user", "content": "count me"}]},
|
||||
headers={"anthropic-version": "2023-06-01"},
|
||||
)
|
||||
|
||||
def targets_called(self) -> tuple[str, ...]:
|
||||
return tuple(unquote(request.target) for request in self.wire.drain())
|
||||
|
||||
def backends_called(self) -> tuple[str, ...]:
|
||||
return tuple(str(object_value(json.loads(request.body))["model"]) for request in self.wire.drain())
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("anthropic-overloaded-cooldown")
|
||||
with gateway_from_environment() as gateway, wire_server(_respond) as wire:
|
||||
overrides: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": _API_KEY}
|
||||
with owned_proxy(gateway, directory, overrides, config=_config(wire, directory), workers=2) as proxy:
|
||||
yield _Rig(proxy, wire)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fresh_wire(rig: _Rig) -> None:
|
||||
rig.wire.drain()
|
||||
|
||||
|
||||
def _assert_overloaded_json(response: httpx.Response) -> None:
|
||||
assert response.status_code == 529, response.text
|
||||
payload: Final = object_value(json.loads(response.content))
|
||||
assert object_value(payload["error"])["type"] == "overloaded_error", response.text
|
||||
|
||||
|
||||
def _count_tokens_error(response: httpx.Response, status: int) -> dict[str, JsonValue]:
|
||||
assert response.status_code == status, response.text
|
||||
envelope: Final = object_value(json.loads(response.content))
|
||||
assert envelope["type"] == "error", response.text
|
||||
return object_value(envelope["error"])
|
||||
|
||||
|
||||
def test_an_overloaded_deployment_is_cooled_down_by_the_internal_server_error_allowed_fails_policy(rig: _Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
first: Final = rig.messages("/v1/messages", _COOLDOWN_GROUP, marker, num_retries=0, disable_fallbacks=True)
|
||||
_assert_overloaded_json(first)
|
||||
assert rig.backends_called() == (_OVERLOADED,)
|
||||
second: Final = rig.messages("/v1/messages", _COOLDOWN_GROUP, marker, num_retries=0, disable_fallbacks=True)
|
||||
assert second.status_code == 200, second.text
|
||||
events: Final = parse_sse(second.text)
|
||||
assert event_types(events) == LIFECYCLE, events
|
||||
assert message_id(events) == _served_id(marker), events
|
||||
assert rig.backends_called() == (_HEALTHY,)
|
||||
|
||||
|
||||
def test_count_tokens_disabled_without_a_provider_counter_answers_503_api_error(rig: _Rig) -> None:
|
||||
response: Final = rig.count_tokens(_COUNT_DISABLED_GROUP)
|
||||
error: Final = _count_tokens_error(response, 503)
|
||||
assert error["type"] == "api_error", response.text
|
||||
assert "Token counting is disabled" in string_value(error["message"]), response.text
|
||||
assert rig.targets_called() == ()
|
||||
|
||||
|
||||
def test_count_tokens_whose_upstream_answers_503_stays_503_api_error(rig: _Rig) -> None:
|
||||
response: Final = rig.count_tokens(_COUNT_UPSTREAM_GROUP)
|
||||
error: Final = _count_tokens_error(response, 503)
|
||||
assert error["type"] == "api_error", response.text
|
||||
assert "scripted 503" in string_value(error["message"]), response.text
|
||||
assert rig.targets_called() == (_COUNT_TOKENS_TARGET,)
|
||||
@pytest.mark.parametrize(
|
||||
("group", "status", "error_type"),
|
||||
(
|
||||
pytest.param(_COUNT_OVERLOADED_GROUP, 529, "overloaded_error", id="529"),
|
||||
pytest.param(_COUNT_RATE_LIMITED_GROUP, 429, "rate_limit_error", id="429"),
|
||||
),
|
||||
)
|
||||
def test_count_tokens_keeps_other_upstream_statuses_and_their_error_body(
|
||||
rig: _Rig, group: str, status: int, error_type: str
|
||||
) -> None:
|
||||
response: Final = rig.count_tokens(group)
|
||||
error: Final = _count_tokens_error(response, status)
|
||||
assert error["type"] == error_type, response.text
|
||||
assert f"scripted {status}" in string_value(error["message"]), response.text
|
||||
assert rig.targets_called() == (_COUNT_TOKENS_TARGET,)
|
||||
|
||||
|
||||
def test_pass_through_messages_stream_forwards_an_upstream_529_verbatim(rig: _Rig) -> None:
|
||||
response: Final = rig.messages("/anthropic/v1/messages", _OVERLOADED, uuid.uuid4().hex)
|
||||
assert response.status_code == 529, response.text
|
||||
assert response.content == error_body(529, "scripted 529"), response.text
|
||||
assert rig.backends_called() == (_OVERLOADED,)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", _GATEWAY_GROUPS.keys())
|
||||
def test_the_claude_code_gateway_messages_route_answers_an_upstream_529_as_overloaded_error(
|
||||
rig: _Rig, shape: str
|
||||
) -> None:
|
||||
response: Final = rig.messages(
|
||||
f"{_GATEWAY}/v1/messages",
|
||||
_GATEWAY_GROUPS[shape],
|
||||
uuid.uuid4().hex,
|
||||
stream=shape == "stream",
|
||||
num_retries=0,
|
||||
)
|
||||
_assert_overloaded_json(response)
|
||||
assert rig.backends_called() == (_OVERLOADED,)
|
||||
|
||||
|
||||
def test_the_claude_code_gateway_count_tokens_route_answers_an_upstream_503_as_503(rig: _Rig) -> None:
|
||||
response: Final = rig.count_tokens(_COUNT_UPSTREAM_GROUP, prefix=_GATEWAY)
|
||||
error: Final = _count_tokens_error(response, 503)
|
||||
assert error["type"] == "api_error", response.text
|
||||
assert "scripted 503" in string_value(error["message"]), response.text
|
||||
assert rig.targets_called() == (_COUNT_TOKENS_TARGET,)
|
||||
|
|
@ -45,10 +45,19 @@ _FALLBACK_GROUP: Final = "audit-fallback"
|
|||
_CW_FALLBACK_GROUP: Final = "audit-cw-fallback"
|
||||
_LONELY: Final = "audit-lonely"
|
||||
_POLICY: Final = "audit-policy"
|
||||
_SERVICE_UNAVAILABLE_RETRY_POLICY: Final = "audit-service-unavailable-retry-policy"
|
||||
_INTERNAL_RETRY_POLICY: Final = "audit-internal-retry-policy"
|
||||
_UPSTREAM_URL_PLACEHOLDER: Final = "upstream-url"
|
||||
|
||||
Behavior = Literal[
|
||||
"drop-once", "drop-always", "hold-ping", "drop-then-too-long", "drop-then-401", "overloaded-frames-fallback-503"
|
||||
"drop-once",
|
||||
"drop-always",
|
||||
"hold-ping",
|
||||
"drop-then-too-long",
|
||||
"drop-then-401",
|
||||
"overloaded-frames-fallback-529",
|
||||
"overloaded-frames",
|
||||
"overloaded-once",
|
||||
]
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
|
@ -81,13 +90,15 @@ def _primary_reply(
|
|||
return Reply(status=400, body=error_body(400, "prompt is too long: 250000 tokens > 200000 maximum"))
|
||||
case "drop-then-401":
|
||||
return dropped if attempt == 1 else status_reply(401)
|
||||
case "overloaded-frames-fallback-503":
|
||||
case "overloaded-frames-fallback-529" | "overloaded-frames":
|
||||
return stream_reply((error_frame(529, "scripted overloaded"),))
|
||||
case "overloaded-once":
|
||||
return status_reply(529) if attempt == 1 else stream_reply(full)
|
||||
raise AssertionError(behavior)
|
||||
|
||||
|
||||
def _fallback_reply(behavior: str, full: tuple[bytes, bytes, bytes, bytes]) -> Reply:
|
||||
if behavior == "overloaded-frames-fallback-503":
|
||||
if behavior == "overloaded-frames-fallback-529":
|
||||
return status_reply(503)
|
||||
return stream_reply(full)
|
||||
|
||||
|
|
@ -141,13 +152,19 @@ def _config(wire: Wire, directory: Path) -> Path:
|
|||
_deployment(_CW_FALLBACK_GROUP, _CONTEXT_WINDOW),
|
||||
_deployment(_LONELY, _PRIMARY, num_retries=1),
|
||||
_deployment(_POLICY, _PRIMARY),
|
||||
_deployment(_SERVICE_UNAVAILABLE_RETRY_POLICY, _PRIMARY),
|
||||
_deployment(_INTERNAL_RETRY_POLICY, _PRIMARY),
|
||||
]
|
||||
config["router_settings"] = {
|
||||
"num_retries": 1,
|
||||
"disable_cooldowns": True,
|
||||
"fallbacks": [{_WITH_FALLBACKS: [_FALLBACK_GROUP]}],
|
||||
"context_window_fallbacks": [{_WITH_FALLBACKS: [_CW_FALLBACK_GROUP]}],
|
||||
"model_group_retry_policy": {_POLICY: {"DefaultRetries": 2}},
|
||||
"model_group_retry_policy": {
|
||||
_POLICY: {"DefaultRetries": 2},
|
||||
_SERVICE_UNAVAILABLE_RETRY_POLICY: {"ServiceUnavailableErrorRetries": 1, "InternalServerErrorRetries": 0},
|
||||
_INTERNAL_RETRY_POLICY: {"InternalServerErrorRetries": 1, "ServiceUnavailableErrorRetries": 0},
|
||||
},
|
||||
}
|
||||
path: Final = directory / "messages-retry-budget-sources.yaml"
|
||||
path.write_text(yaml.safe_dump(config).replace(_UPSTREAM_URL_PLACEHOLDER, wire.url))
|
||||
|
|
@ -290,10 +307,34 @@ def test_a_retry_rejected_with_401_ends_the_retries_and_reaches_the_client_uncha
|
|||
assert rig.attempts_on(_PRIMARY, marker) == 2
|
||||
|
||||
|
||||
def test_overloaded_frames_whose_fallback_fails_answer_the_mapped_internal_server_error(rig: _Rig) -> None:
|
||||
def test_overloaded_frames_whose_fallback_fails_answer_529_overloaded_error(rig: _Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = rig.stream(_WITH_FALLBACKS, "overloaded-frames-fallback-503", marker)
|
||||
assert response.status_code == 500, response.text
|
||||
response: Final = rig.stream(_WITH_FALLBACKS, "overloaded-frames-fallback-529", marker)
|
||||
assert response.status_code == 529, response.text
|
||||
assert "overloaded_error" in response.text, response.text
|
||||
assert "content_block_delta" not in response.text, response.text
|
||||
assert rig.attempts_on(_PRIMARY, marker) == 2
|
||||
assert rig.attempts_on(_FALLBACK, marker) == 2
|
||||
|
||||
|
||||
def test_overloaded_frames_that_spend_the_budget_without_a_fallback_relay_the_last_attempts_frame(rig: _Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = rig.stream(_LONELY, "overloaded-frames", marker)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.text == error_frame(529, "scripted overloaded").decode(), response.text
|
||||
assert rig.attempts_on(_PRIMARY, marker) == 2
|
||||
|
||||
|
||||
def test_a_529_is_retried_under_the_internal_server_error_retry_policy(rig: _Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
_assert_completed(rig.stream(_INTERNAL_RETRY_POLICY, "overloaded-once", marker), _served_id(_PRIMARY, marker, 2))
|
||||
assert rig.attempts_on(_PRIMARY, marker) == 2
|
||||
|
||||
|
||||
def test_a_529_is_not_retried_under_the_service_unavailable_error_retry_policy(rig: _Rig) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = rig.stream(_SERVICE_UNAVAILABLE_RETRY_POLICY, "overloaded-once", marker)
|
||||
assert response.status_code == 529, response.text
|
||||
assert "overloaded_error" in response.text, response.text
|
||||
assert "content_block_delta" not in response.text, response.text
|
||||
assert rig.attempts_on(_PRIMARY, marker) == 1
|
||||
|
|
|
|||
|
|
@ -217,6 +217,22 @@ def test_unknown_exception_falls_back_to_status_then_unclassified() -> None:
|
|||
assert normalize_error(Exception("x"), "", "x") == "UNCLASSIFIED"
|
||||
|
||||
|
||||
def test_a_529_the_anthropic_routes_answered_is_an_overload() -> None:
|
||||
paused = _proxy_exc("Model is paused", "service_unavailable_error", 529)
|
||||
assert StandardLoggingPayloadSetup.get_error_information(paused)["normalized_error"] == "503_PROVIDER_OVERLOADED"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("body", ["", "<html>bad gateway</html>", '{"type":"error","error":{"type":"api_error"}}'])
|
||||
def test_an_internal_server_error_answered_as_529_is_an_overload_whatever_its_body_says(body: str) -> None:
|
||||
overloaded = litellm.InternalServerError(body, llm_provider="anthropic", model="claude", provider_status_code=529)
|
||||
overloaded.status_code = 529
|
||||
internal = litellm.InternalServerError(body, llm_provider="anthropic", model="claude")
|
||||
assert (
|
||||
StandardLoggingPayloadSetup.get_error_information(overloaded)["normalized_error"],
|
||||
StandardLoggingPayloadSetup.get_error_information(internal)["normalized_error"],
|
||||
) == ("503_PROVIDER_OVERLOADED", "500_PROVIDER_INTERNAL_ERROR")
|
||||
|
||||
|
||||
def test_budget_exceeded_error_with_custom_wording_is_still_a_budget_error() -> None:
|
||||
exc = litellm.BudgetExceededError(current_cost=2.0, max_budget=1.0, message="Spending cap reached for key")
|
||||
assert StandardLoggingPayloadSetup.get_error_information(exc)["normalized_error"] == "429_BUDGET_EXCEEDED"
|
||||
|
|
|
|||
|
|
@ -1230,12 +1230,12 @@ def test_openai_incorrect_api_key_keeps_the_proxy_authentication_error_type() ->
|
|||
|
||||
|
||||
class _UpstreamHTTPError(Exception):
|
||||
def __init__(self, status_code: int) -> None:
|
||||
super().__init__("upstream failure")
|
||||
self.message = "upstream failure"
|
||||
def __init__(self, status_code: int, message: str = "upstream failure") -> None:
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
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=message)
|
||||
|
||||
|
||||
UPSTREAM_STATUS_CODES = (400, 401, 402, 403, 404, 408, 422, 429, 500, 503)
|
||||
|
|
@ -1433,6 +1433,48 @@ def test_an_already_mapped_litellm_exception_passes_through_untouched(provider,
|
|||
assert returned is already_mapped
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"message",
|
||||
['{"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}', "upstream failure"],
|
||||
)
|
||||
def test_an_anthropic_529_stays_an_internal_server_error_that_carries_the_529(message, quiet_exception_mapping):
|
||||
with pytest.raises(litellm.InternalServerError) as raised:
|
||||
exception_type(
|
||||
model="claude-opus-5-5",
|
||||
original_exception=_UpstreamHTTPError(status_code=529, message=message),
|
||||
custom_llm_provider="anthropic",
|
||||
)
|
||||
|
||||
assert (type(raised.value), raised.value.status_code, raised.value.provider_status_code) == (
|
||||
litellm.InternalServerError,
|
||||
500,
|
||||
529,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status_code", [400, 404])
|
||||
def test_an_anthropic_client_error_naming_overloaded_carries_no_overload_status(status_code, quiet_exception_mapping):
|
||||
with pytest.raises(litellm.InternalServerError) as raised:
|
||||
exception_type(
|
||||
model="claude-opus-5-5",
|
||||
original_exception=_UpstreamHTTPError(status_code=status_code, message="model: claude-Overloaded"),
|
||||
custom_llm_provider="anthropic",
|
||||
)
|
||||
|
||||
assert raised.value.provider_status_code is None
|
||||
|
||||
|
||||
def test_an_anthropic_500_carries_its_own_status_as_the_provider_status(quiet_exception_mapping):
|
||||
with pytest.raises(litellm.InternalServerError) as raised:
|
||||
exception_type(
|
||||
model="claude-opus-5-5",
|
||||
original_exception=_UpstreamHTTPError(status_code=500),
|
||||
custom_llm_provider="anthropic",
|
||||
)
|
||||
|
||||
assert (raised.value.status_code, raised.value.provider_status_code) == (500, 500)
|
||||
|
||||
|
||||
@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):
|
||||
|
|
|
|||
|
|
@ -4605,6 +4605,7 @@ def test_anthropic_error_frame_exception_classes_an_overloaded_frame_as_internal
|
|||
error: Final = anthropic_error_frame_exception("overloaded_error", "Overloaded", 503, "claude-sonnet-4-5")
|
||||
|
||||
assert type(error) is litellm.InternalServerError
|
||||
assert error.provider_status_code == 529
|
||||
|
||||
|
||||
_NON_NUMERIC_BUDGETS: Final = (True, False, None, "", "lots", [4096])
|
||||
|
|
|
|||
116
tests/unit/proxy/anthropic_endpoints/test_answered_status.py
Normal file
116
tests/unit/proxy/anthropic_endpoints/test_answered_status.py
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
from typing import Final
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.anthropic_endpoints.answered_status import with_answered_status
|
||||
|
||||
|
||||
def _anthropic_internal_server_error(provider_status_code: int | None) -> litellm.InternalServerError:
|
||||
return litellm.InternalServerError(
|
||||
message="AnthropicError - Overloaded",
|
||||
llm_provider="anthropic",
|
||||
model="claude-opus-5-5",
|
||||
provider_status_code=provider_status_code,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("code", [400, 500, 503])
|
||||
def test_a_proxy_exception_passes_through_unchanged(code: int) -> None:
|
||||
exc: Final = ProxyException(message="budget service unavailable", type="None", param=None, code=code)
|
||||
|
||||
answered_exc: Final = with_answered_status(exc)
|
||||
|
||||
assert answered_exc is exc
|
||||
assert exc.code == str(code)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("exc", "answered", "copy"),
|
||||
[
|
||||
pytest.param(
|
||||
litellm.ServiceUnavailableError(message="BedrockException", llm_provider="bedrock", model="m"),
|
||||
503,
|
||||
False,
|
||||
id="service-unavailable-503",
|
||||
),
|
||||
pytest.param(_anthropic_internal_server_error(529), 529, True, id="anthropic-overload"),
|
||||
pytest.param(_anthropic_internal_server_error(500), 500, True, id="anthropic-500"),
|
||||
pytest.param(_anthropic_internal_server_error(None), 500, False, id="internal-server-error"),
|
||||
pytest.param(
|
||||
MidStreamFallbackError(
|
||||
message="fallbacks failed",
|
||||
model="m",
|
||||
llm_provider="anthropic",
|
||||
original_exception=_anthropic_internal_server_error(529),
|
||||
),
|
||||
529,
|
||||
True,
|
||||
id="fallback-over-an-anthropic-overload",
|
||||
),
|
||||
pytest.param(
|
||||
MidStreamFallbackError(
|
||||
message="fallbacks failed",
|
||||
model="m",
|
||||
llm_provider="anthropic",
|
||||
original_exception=MidStreamFallbackError(
|
||||
message="retry failed before content",
|
||||
model="m",
|
||||
llm_provider="anthropic",
|
||||
original_exception=_anthropic_internal_server_error(529),
|
||||
),
|
||||
),
|
||||
529,
|
||||
True,
|
||||
id="nested-fallback-over-an-anthropic-overload",
|
||||
),
|
||||
pytest.param(
|
||||
litellm.RateLimitError(message="slow down", llm_provider="anthropic", model="m"),
|
||||
429,
|
||||
False,
|
||||
id="rate-limit",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_a_litellm_exception_answers_the_status_anthropic_gives_its_failure(
|
||||
exc: openai.APIStatusError, answered: int, copy: bool
|
||||
) -> None:
|
||||
raised_status: Final = exc.status_code
|
||||
|
||||
answered_exc: Final = with_answered_status(exc)
|
||||
|
||||
assert (answered_exc is not exc) is copy
|
||||
assert (type(answered_exc), answered_exc.status_code, str(answered_exc)) == (type(exc), answered, str(exc))
|
||||
assert exc.status_code == raised_status
|
||||
|
||||
|
||||
def _raise_a_bedrock_overload_caused_by_a_reset() -> None:
|
||||
try:
|
||||
raise ConnectionError("upstream reset")
|
||||
except ConnectionError as upstream:
|
||||
raise litellm.InternalServerError(
|
||||
message="AnthropicError - Overloaded",
|
||||
llm_provider="anthropic",
|
||||
model="m",
|
||||
provider_status_code=529,
|
||||
) from upstream
|
||||
|
||||
|
||||
def test_the_answered_copy_keeps_the_raised_exceptions_cause_and_traceback() -> None:
|
||||
with pytest.raises(litellm.InternalServerError) as raised:
|
||||
_raise_a_bedrock_overload_caused_by_a_reset()
|
||||
|
||||
answered_exc: Final = with_answered_status(raised.value)
|
||||
|
||||
assert isinstance(raised.value.__cause__, ConnectionError)
|
||||
assert (answered_exc.__cause__, answered_exc.__traceback__) == (raised.value.__cause__, raised.value.__traceback__)
|
||||
|
||||
|
||||
def test_an_exception_without_a_status_passes_through_unchanged() -> None:
|
||||
exc: Final = ValueError("no status")
|
||||
|
||||
assert with_answered_status(exc) is exc
|
||||
assert not hasattr(exc, "status_code")
|
||||
|
|
@ -17,6 +17,7 @@ from fastapi.responses import StreamingResponse
|
|||
from fastapi.testclient import TestClient
|
||||
from starlette.types import Message
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
|
|
@ -302,7 +303,7 @@ class TestProxyExceptionAnthropicEnvelope:
|
|||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
request = MagicMock()
|
||||
request: Final = MagicMock()
|
||||
request.headers = {}
|
||||
|
||||
response = ep._anthropic_error_json_response(
|
||||
|
|
@ -319,10 +320,135 @@ class TestProxyExceptionAnthropicEnvelope:
|
|||
assert fallback.status_code == 500
|
||||
assert json.loads(fallback.body)["error"]["type"] == "api_error"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"unavailable, expected_status, expected_error_type, hook_receives_copy",
|
||||
[
|
||||
pytest.param(
|
||||
litellm.ServiceUnavailableError(
|
||||
message='BedrockException - {"message": "Bedrock is unable to process your request."}',
|
||||
llm_provider="bedrock",
|
||||
model="global.anthropic.claude-opus-5-5",
|
||||
),
|
||||
503,
|
||||
"api_error",
|
||||
False,
|
||||
id="bedrock-503",
|
||||
),
|
||||
pytest.param(
|
||||
litellm.InternalServerError(
|
||||
message='AnthropicError - {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}',
|
||||
llm_provider="anthropic",
|
||||
model="claude-opus-5-5",
|
||||
provider_status_code=529,
|
||||
),
|
||||
529,
|
||||
"overloaded_error",
|
||||
True,
|
||||
id="anthropic-529",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_anthropic_response_records_the_answered_status(
|
||||
self,
|
||||
unavailable: litellm.ServiceUnavailableError | litellm.InternalServerError,
|
||||
expected_status: int,
|
||||
expected_error_type: str,
|
||||
hook_receives_copy: bool,
|
||||
) -> None:
|
||||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
raised_status: Final = unavailable.status_code
|
||||
request: Final = MagicMock()
|
||||
request.headers = {}
|
||||
|
||||
with (
|
||||
patch.object(ep, "read_request_body", new=AsyncMock(return_value={})),
|
||||
patch.object(
|
||||
ep.ProxyBaseLLMRequestProcessing,
|
||||
"base_process_llm_request",
|
||||
new=AsyncMock(side_effect=unavailable),
|
||||
),
|
||||
patch.object(proxy_server, "proxy_logging_obj") as mock_logging,
|
||||
):
|
||||
mock_logging.post_call_failure_hook = AsyncMock()
|
||||
response = await ep.anthropic_response(
|
||||
fastapi_response=MagicMock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
recorded: Final = mock_logging.post_call_failure_hook.await_args.kwargs["original_exception"]
|
||||
assert recorded is not None
|
||||
assert (recorded is not unavailable) is hook_receives_copy
|
||||
assert recorded.status_code == expected_status
|
||||
assert unavailable.status_code == raised_status
|
||||
assert response.status_code == expected_status
|
||||
assert json.loads(response.body)["error"]["type"] == expected_error_type
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"unavailable, expected_status, expected_error_type",
|
||||
[
|
||||
pytest.param(
|
||||
litellm.ServiceUnavailableError(
|
||||
message="Bedrock is unable to process your request.",
|
||||
llm_provider="bedrock",
|
||||
model="global.anthropic.claude-opus-5-5",
|
||||
),
|
||||
503,
|
||||
"api_error",
|
||||
id="bedrock-503",
|
||||
),
|
||||
pytest.param(
|
||||
litellm.InternalServerError(
|
||||
message='AnthropicError - {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}',
|
||||
llm_provider="anthropic",
|
||||
model="claude-opus-5-5",
|
||||
provider_status_code=529,
|
||||
),
|
||||
529,
|
||||
"overloaded_error",
|
||||
id="anthropic-529",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_count_tokens_answers_provider_errors(
|
||||
self,
|
||||
unavailable: litellm.ServiceUnavailableError | litellm.InternalServerError,
|
||||
expected_status: int,
|
||||
expected_error_type: str,
|
||||
) -> None:
|
||||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
request = MagicMock()
|
||||
request.headers = {}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
ep,
|
||||
"read_request_body",
|
||||
new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}),
|
||||
),
|
||||
patch.object(
|
||||
proxy_server, "count_request_tokens", new=AsyncMock(side_effect=unavailable)
|
||||
),
|
||||
):
|
||||
response: Final = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
|
||||
body: Final = json.loads(response.body)
|
||||
assert response.status_code == expected_status
|
||||
assert body["type"] == "error"
|
||||
assert body["error"]["type"] == expected_error_type
|
||||
|
||||
@pytest.mark.parametrize("status_code", [402, 409, 422, 504])
|
||||
def test_a_failure_inside_the_handler_keeps_the_api_error_type_for_a_status_outside_the_table(
|
||||
self, status_code: int
|
||||
):
|
||||
) -> None:
|
||||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
|
|
@ -334,7 +460,6 @@ class TestProxyExceptionAnthropicEnvelope:
|
|||
|
||||
assert response.status_code == status_code
|
||||
assert json.loads(response.body)["error"] == {"type": "api_error", "message": "scripted"}
|
||||
|
||||
@staticmethod
|
||||
def _call_id_error_response(general_settings, provider_specific_fields=None):
|
||||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
|
|
@ -428,7 +553,7 @@ class TestHttpExceptionDictDetail:
|
|||
) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
|
||||
):
|
||||
mock_logging.post_call_failure_hook = AsyncMock()
|
||||
response = await ep.anthropic_response(
|
||||
response: Final = await ep.anthropic_response(
|
||||
fastapi_response=MagicMock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
|
|
|
|||
|
|
@ -4618,6 +4618,34 @@ class TestStreamCloseOnDisconnect:
|
|||
|
||||
assert upstream.aclosed
|
||||
|
||||
async def test_async_streaming_data_generator_reraises_an_unchanged_http_exception_without_chaining_it_to_itself(
|
||||
self,
|
||||
):
|
||||
blocked = HTTPException(status_code=400, detail="Content blocked by guardrail")
|
||||
|
||||
class FailingUpstream:
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise blocked
|
||||
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
response=FailingUpstream(),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(),
|
||||
request_data={"model": "mock-model"},
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()),
|
||||
serialize_chunk=lambda c: "data: x\n\n",
|
||||
serialize_error=lambda e: "data: error\n\n",
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await gen.__anext__()
|
||||
|
||||
assert raised.value is blocked
|
||||
assert raised.value.__cause__ is not blocked
|
||||
|
||||
async def test_async_streaming_data_generator_closes_the_guardrail_chain_on_client_disconnect(
|
||||
self,
|
||||
):
|
||||
|
|
|
|||
|
|
@ -81,6 +81,12 @@ def test_exhaustion_fires_on_infra_signal(fixture):
|
|||
assert state.exhaustion_count >= 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [503, 529])
|
||||
def test_exhaustion_fires_on_an_overloaded_answer(status):
|
||||
state, _ = _replay([{"user_content": "fix the build", "response_status": status}])
|
||||
assert state.exhaustion_count == 1
|
||||
|
||||
|
||||
def test_no_signals_on_clean_session():
|
||||
state, _ = _replay(_load("clean_no_signals"))
|
||||
assert state.misalignment_count == 0
|
||||
|
|
|
|||
|
|
@ -225,6 +225,7 @@ describe("statusCodeLabel", () => {
|
|||
[401, "Unauthorized"],
|
||||
[429, "Rate Limited"],
|
||||
[503, "Service Unavailable"],
|
||||
[529, "Overloaded"],
|
||||
[418, "Client Error"],
|
||||
[599, "Server Error"],
|
||||
] as const)("labels status %s", (code, label) => {
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ const STATUS_LABELS: Record<number, string> = {
|
|||
502: "Bad Gateway",
|
||||
503: "Service Unavailable",
|
||||
504: "Gateway Timeout",
|
||||
529: "Overloaded",
|
||||
};
|
||||
|
||||
export const statusCodeLabel = (code: number): string => {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue