mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(bedrock-passthrough)!: return AWS-shaped errors so AWS SDKs read the code and message (#46522)
* fix(bedrock-passthrough): return AWS-shaped errors so AWS SDKs read the code and message A guardrail block on a /bedrock route now answers with x-amzn-errortype ValidationException and a top-level message, and an error Bedrock returned (such as a 429 ThrottlingException) is passed through with its status, x-amzn-errortype and body * fix(bedrock-passthrough): keep 5xx guardrail failures and failure-hook replacements unchanged * test(bedrock-passthrough): guardrail vendor failures and unknown models keep the litellm error * fix(bedrock-passthrough): only answer guardrail exceptions that blocked content as ValidationException * fix(bedrock-passthrough): pass Bedrock's content-type through with its error body * refactor(bedrock-passthrough): wrap the guardrail ProxyException instead of copying its fields Unit tests drive the error response with the real request processor and ProxyLogging * test(bedrock-passthrough): cover the direct model route, invoke, and auth, budget and rate-limit errors * fix(bedrock-passthrough): pass through only errors a Bedrock model operation returned A guardrail service's HTTP error that carries x-amzn-errortype keeps the litellm error * refactor(bedrock-passthrough): mark Bedrock's errors where the passthrough raises them instead of matching the URL A guardrail service's HTTP error never goes through the passthrough call, so it keeps the litellm error whatever its URL --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
yucheng-berri
GitHub
parent
0620f1a515
commit
b5870df967
8 changed files with 872 additions and 26 deletions
|
|
@ -1705,6 +1705,7 @@ AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:tran
|
|||
AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com"
|
||||
AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com"
|
||||
AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key"
|
||||
AWS_ERROR_TYPE_HEADER: Final = "x-amzn-errortype"
|
||||
AZURE_SPEECH_SHORT_AUDIO_MODEL: Final = "short-audio"
|
||||
AZURE_SPEECH_BATCH_MODEL: Final = "batch-transcription"
|
||||
AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL: Final = "fast-transcription"
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import weakref
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Generator, Iterator
|
||||
from functools import partial
|
||||
from types import TracebackType
|
||||
|
|
@ -26,6 +27,23 @@ from litellm.utils import client
|
|||
base_llm_http_handler = BaseLLMHTTPHandler()
|
||||
from .utils import BasePassthroughUtils
|
||||
|
||||
_provider_errors: Final[weakref.WeakSet[httpx.HTTPStatusError]] = weakref.WeakSet()
|
||||
|
||||
|
||||
def raise_for_provider_status(response: httpx.Response) -> None:
|
||||
"""``response.raise_for_status()`` for the provider's response to a passthrough request, recording the raised
|
||||
error as the provider's own."""
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as error:
|
||||
_provider_errors.add(error)
|
||||
raise
|
||||
|
||||
|
||||
def is_provider_error(error: httpx.HTTPStatusError) -> bool:
|
||||
"""Whether ``error`` is the provider's own error response to a passthrough request."""
|
||||
return error in _provider_errors
|
||||
|
||||
|
||||
async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[bytes, bytes]:
|
||||
async for chunk in iterable:
|
||||
|
|
@ -112,7 +130,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]):
|
|||
self.status_code = self._response.status_code
|
||||
self._initialized = True
|
||||
try:
|
||||
self._response.raise_for_status()
|
||||
raise_for_provider_status(self._response)
|
||||
self._iterator = _as_async_generator(self._response.aiter_bytes())
|
||||
except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic
|
||||
try:
|
||||
|
|
@ -626,7 +644,7 @@ async def _async_passthrough_request(
|
|||
else:
|
||||
response: Final = await response_result
|
||||
await response.aread()
|
||||
response.raise_for_status()
|
||||
raise_for_provider_status(response)
|
||||
return response
|
||||
else:
|
||||
# Fallback for sync-like behavior (shouldn't happen in async path)
|
||||
|
|
|
|||
|
|
@ -3814,6 +3814,16 @@ class ModelAccessDeniedProxyException(ProxyException):
|
|||
return self.internal_message.replace("\r", "").replace("\n", "")
|
||||
|
||||
|
||||
class AwsErrorTypeProxyException(Exception):
|
||||
"""A ProxyException answered to an AWS SDK, which reads the error code from the ``x-amzn-errortype`` header and
|
||||
the message from a top-level ``message`` field."""
|
||||
|
||||
def __init__(self, error: ProxyException, aws_error_type: str) -> None:
|
||||
super().__init__(error.message)
|
||||
self.error: Final = error
|
||||
self.aws_error_type: Final = aws_error_type
|
||||
|
||||
|
||||
from litellm.types.proxy.common_proxy_errors import ( # noqa: E402 # public re-export
|
||||
CommonProxyErrors as CommonProxyErrors, # noqa: PLC0414 # public re-export
|
||||
)
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm import get_llm_provider
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS,
|
||||
AWS_ERROR_TYPE_HEADER,
|
||||
AZURE_SPEECH_BATCH_PATH_PREFIX,
|
||||
AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN,
|
||||
AZURE_SPEECH_CUSTOM_LLM_PROVIDER,
|
||||
|
|
@ -45,7 +46,8 @@ from litellm.constants import (
|
|||
BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES,
|
||||
TYPESAFE_DANGER_ZONE_HEADER_PREFIX,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
||||
from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers
|
||||
from litellm.llms.azure.passthrough.transformation import (
|
||||
|
|
@ -67,7 +69,7 @@ from litellm.llms.openai.common_utils import OpenAIError as LiteLLMOpenAIError
|
|||
from litellm.llms.openai.workload_identity import get_workload_identity_bearer_token_for_api_base
|
||||
from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, validate_oss_request
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
||||
from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
|
|
@ -80,7 +82,11 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul
|
|||
user_api_key_auth_websocket,
|
||||
user_api_key_auth_websocket_for_model,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import model_deployments_support_stream_options, open_sse_before_first_byte
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
model_deployments_support_stream_options,
|
||||
open_sse_before_first_byte,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
|
||||
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
_safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
|
|
@ -1064,6 +1070,68 @@ def extract_model_from_bedrock_endpoint(endpoint: str) -> str:
|
|||
_extract_model_from_bedrock_endpoint: Final = extract_model_from_bedrock_endpoint
|
||||
|
||||
|
||||
def _guardrail_refusal_status(error: Exception) -> int | None:
|
||||
"""The status a guardrail blocked the request with. None for a request check that answers with the same 400, 403
|
||||
or 422, for a guardrail that failed closed because it could not get a verdict, and for a guardrail block raised
|
||||
with a plain-string detail.
|
||||
|
||||
Guardrail hooks stamp ``guardrail_name`` onto the dict detail of an HTTPException they raise."""
|
||||
match error:
|
||||
case HTTPException():
|
||||
detail: Final[object] = error.detail
|
||||
stamped: Final = isinstance(detail, dict) and "guardrail_name" in detail
|
||||
return error.status_code if stamped and is_guardrail_intervention(error) else None
|
||||
case GuardrailRaisedException(blocked_content=True) | BlockedPiiEntityError():
|
||||
return error.status_code
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _bedrock_error_response(error: Exception, handled: HTTPException) -> Response | None:
|
||||
"""Bedrock's own error response, for an error Bedrock answered the passthrough request with an AWS error
|
||||
type."""
|
||||
if not isinstance(error, httpx.HTTPStatusError) or not is_provider_error(error):
|
||||
return None
|
||||
upstream: Final = error.response
|
||||
if AWS_ERROR_TYPE_HEADER not in upstream.headers:
|
||||
return None
|
||||
passed_through: Final = {
|
||||
name: upstream.headers[name] for name in (AWS_ERROR_TYPE_HEADER, "content-type") if name in upstream.headers
|
||||
}
|
||||
return Response(
|
||||
content=upstream.content,
|
||||
status_code=upstream.status_code,
|
||||
headers={**(handled.headers or {}), **passed_through},
|
||||
)
|
||||
|
||||
|
||||
async def bedrock_passthrough_error_response(
|
||||
processor: ProxyBaseLLMRequestProcessing,
|
||||
error: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLoggingType,
|
||||
) -> Response:
|
||||
"""Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block or an error
|
||||
Bedrock returned the way Bedrock does, so an AWS SDK reads its code and message. Anything else is raised as the
|
||||
shared handler raises it."""
|
||||
try:
|
||||
await processor.handle_llm_api_exception(
|
||||
e=error,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except HTTPException as handled:
|
||||
response: Final = _bedrock_error_response(error, handled)
|
||||
if response is None:
|
||||
raise
|
||||
return response
|
||||
except ProxyException as handled:
|
||||
refusal_status: Final = _guardrail_refusal_status(error)
|
||||
if refusal_status is not None and str(refusal_status) == handled.code:
|
||||
raise AwsErrorTypeProxyException(handled, "ValidationException") from handled
|
||||
raise
|
||||
|
||||
|
||||
async def handle_bedrock_passthrough_router_model(
|
||||
model: str,
|
||||
endpoint: str,
|
||||
|
|
@ -1148,11 +1216,8 @@ async def handle_bedrock_passthrough_router_model(
|
|||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
# Use common exception handling
|
||||
raise await base_llm_response_processor.handle_llm_api_exception(
|
||||
e=e,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
return await bedrock_passthrough_error_response(
|
||||
base_llm_response_processor, e, user_api_key_dict, proxy_logging_obj
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1351,10 +1416,8 @@ async def bedrock_llm_proxy_route(
|
|||
|
||||
return result
|
||||
except Exception as e:
|
||||
raise await base_llm_response_processor.handle_llm_api_exception(
|
||||
e=e,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
return await bedrock_passthrough_error_response(
|
||||
base_llm_response_processor, e, user_api_key_dict, proxy_logging_obj
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -67,6 +67,7 @@ from litellm.constants import (
|
|||
AIOHTTP_NEEDS_CLEANUP_CLOSED,
|
||||
AIOHTTP_TTL_DNS_CACHE,
|
||||
AUDIO_SPEECH_CHUNK_SIZE,
|
||||
AWS_ERROR_TYPE_HEADER,
|
||||
BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME,
|
||||
BASE_MCP_ROUTE,
|
||||
DAILY_TAG_SPEND_BATCH_MULTIPLIER,
|
||||
|
|
@ -97,6 +98,7 @@ from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
|||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.proxy._types import (
|
||||
UI_TEAM_ID,
|
||||
AwsErrorTypeProxyException,
|
||||
CallbackDelete,
|
||||
CallInfo,
|
||||
CommonProxyErrors,
|
||||
|
|
@ -2061,6 +2063,17 @@ async def openai_exception_handler(request: Request, exc: ProxyException):
|
|||
)
|
||||
|
||||
|
||||
@app.exception_handler(AwsErrorTypeProxyException)
|
||||
async def aws_error_type_exception_handler(request: Request, exc: AwsErrorTypeProxyException) -> Response:
|
||||
"""The ProxyException response plus the ``x-amzn-errortype`` header and top-level ``message`` an AWS SDK reads."""
|
||||
response: Final = await openai_exception_handler(request, exc.error)
|
||||
return JSONResponse(
|
||||
status_code=response.status_code,
|
||||
content={"message": exc.error.message, **JSON_OBJECT.validate_json(bytes(response.body))},
|
||||
headers={**exc.error.headers, AWS_ERROR_TYPE_HEADER: exc.aws_error_type},
|
||||
)
|
||||
|
||||
|
||||
@app.exception_handler(StarletteHTTPException)
|
||||
async def otlp_http_exception_handler(connection: Request | WebSocket, exc: StarletteHTTPException) -> Response | None:
|
||||
if isinstance(connection, WebSocket):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,443 @@
|
|||
"""Errors on the Bedrock passthrough reach the AWS SDK with their message (LIT-9617).
|
||||
|
||||
boto3 reads an error's code from the ``x-amzn-errortype`` header and its message from the body's ``message``. A
|
||||
guardrail block on ``/bedrock/model/<model>/...`` that happens before any frame was sent reaches the caller as a
|
||||
``ValidationException`` carrying the guardrail's message, and an error Bedrock itself returned, such as a 429
|
||||
``ThrottlingException``, reaches the caller with Bedrock's status, ``x-amzn-errortype`` and body."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import boto3
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from botocore.config import Config
|
||||
from botocore.exceptions import ClientError
|
||||
from integration._support.budget import case_timeout_seconds
|
||||
from integration._support.client import Gateway, JsonValue
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.upstream import aws_event_stream_frame
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
|
||||
ROUTER_MODEL: Final = "aws-errors"
|
||||
EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
|
||||
POST_CALL: Final = "aws-errors-post-call"
|
||||
PRE_CALL: Final = "aws-errors-pre-call"
|
||||
BEDROCK: Final = "aws-errors-bedrock"
|
||||
VENDOR_DOWN: Final = "aws-errors-vendor-down"
|
||||
CATO: Final = "aws-errors-cato"
|
||||
UNGUARDED: Final = "unguarded"
|
||||
OVER_BUDGET: Final = "over-budget"
|
||||
ONE_RPM: Final = "one-rpm"
|
||||
BLOCKED_WORD: Final = "FORBIDDEN"
|
||||
BLOCK_REASON: Final = "synthetic policy says no"
|
||||
BEDROCK_BLOCK_MESSAGE: Final = "Violated guardrail policy"
|
||||
THROTTLE_MESSAGE: Final = "Too many requests, please wait before trying again."
|
||||
THROTTLE_TYPE: Final = "ThrottlingException:http://internal.amazon.com/coral/com.amazon.bedrock/"
|
||||
INVALID_MESSAGE: Final = "The provided model identifier is invalid."
|
||||
INVALID_TYPE: Final = "ValidationException:http://internal.amazon.com/coral/com.amazon.bedrock/"
|
||||
ANTHROPIC_THROTTLE: Final = b'{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}'
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UpstreamError:
|
||||
status: int
|
||||
errortype: str
|
||||
message: str
|
||||
|
||||
@property
|
||||
def body(self) -> bytes:
|
||||
return json.dumps({"message": self.message}).encode()
|
||||
|
||||
@property
|
||||
def code(self) -> str:
|
||||
return self.errortype.split(":", 1)[0]
|
||||
|
||||
|
||||
THROTTLED: Final = UpstreamError(429, THROTTLE_TYPE, THROTTLE_MESSAGE)
|
||||
INVALID: Final = UpstreamError(400, INVALID_TYPE, INVALID_MESSAGE)
|
||||
UPSTREAM_ERRORS: Final = MappingProxyType({"throttled": THROTTLED, "invalid": INVALID})
|
||||
ROUTES: Final = (pytest.param(ROUTER_MODEL, id="router"), pytest.param(MODEL_ID, id="direct"))
|
||||
|
||||
|
||||
def upstream_marker(kind: str) -> str:
|
||||
"""A prompt the scripted Bedrock answers with the ``UPSTREAM_ERRORS`` entry named ``kind``."""
|
||||
return f"upstream-{kind}-{uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def upstream_error_of(marker: str) -> UpstreamError | None:
|
||||
prefix, _, rest = marker.partition("-")
|
||||
return UPSTREAM_ERRORS.get(rest.partition("-")[0]) if prefix == "upstream" else None
|
||||
|
||||
|
||||
def event_frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes:
|
||||
return aws_event_stream_frame(
|
||||
{":event-type": event_type, ":content-type": "application/json", ":message-type": "event"},
|
||||
json.dumps(payload, separators=(",", ":")).encode(),
|
||||
)
|
||||
|
||||
|
||||
def converse_frames(text: str) -> tuple[bytes, ...]:
|
||||
return (
|
||||
event_frame("messageStart", {"role": "assistant"}),
|
||||
event_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": text}}),
|
||||
event_frame("contentBlockStop", {"contentBlockIndex": 0}),
|
||||
event_frame("messageStop", {"stopReason": "end_turn"}),
|
||||
event_frame("metadata", {"usage": {"inputTokens": 9, "outputTokens": 4, "totalTokens": 13}, "metrics": {}}),
|
||||
)
|
||||
|
||||
|
||||
def converse_reply(text: str) -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": text}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 9, "outputTokens": 4, "totalTokens": 13},
|
||||
"metrics": {"latencyMs": 1},
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
def invoke_body(marker: str) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": marker}]}],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def marker_of(request: Request) -> str:
|
||||
return str(json.loads(request.body)["messages"][0]["content"][0]["text"])
|
||||
|
||||
|
||||
def bedrock_texts(request: Request) -> tuple[str, ...]:
|
||||
return tuple(str(item["text"]["text"]) for item in json.loads(request.body)["content"])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
gateway: Gateway
|
||||
provider: Wire
|
||||
vendor: Wire
|
||||
keys: Mapping[str, str]
|
||||
|
||||
def provider_calls(self, marker: str) -> int:
|
||||
return sum(1 for request in self.provider.drain() if marker.encode() in request.body)
|
||||
|
||||
def client(self, guardrail: str, attempts: int = 1): # noqa: ANN201 # boto3 clients are untyped
|
||||
client: Final = boto3.client(
|
||||
"bedrock-runtime",
|
||||
region_name="us-east-1",
|
||||
endpoint_url=str(self.gateway.client.base_url).rstrip("/") + "/bedrock",
|
||||
aws_access_key_id="AKIASYNTHETICCALLER",
|
||||
aws_secret_access_key="synthetic-caller-secret",
|
||||
config=Config(
|
||||
read_timeout=case_timeout_seconds(),
|
||||
retries={"total_max_attempts": attempts, "mode": "standard"},
|
||||
),
|
||||
)
|
||||
key: Final = self.keys[guardrail]
|
||||
|
||||
def authenticate(request: object, **_: object) -> None:
|
||||
getattr(request, "headers")["x-litellm-api-key"] = key
|
||||
|
||||
client.meta.events.register("before-send.bedrock-runtime.*", authenticate)
|
||||
return client
|
||||
|
||||
def call(
|
||||
self, guardrail: str, action: str, marker: str, attempts: int = 1, model: str = ROUTER_MODEL
|
||||
) -> ClientError | None:
|
||||
client: Final = self.client(guardrail, attempts)
|
||||
messages: Final = [{"role": "user", "content": [{"text": marker}]}]
|
||||
try:
|
||||
if action == "converse-stream":
|
||||
for _ in client.converse_stream(modelId=model, messages=messages)["stream"]:
|
||||
pass
|
||||
elif action == "invoke":
|
||||
client.invoke_model(modelId=model, body=invoke_body(marker))
|
||||
else:
|
||||
client.converse(modelId=model, messages=messages)
|
||||
except ClientError as error:
|
||||
return error
|
||||
return None
|
||||
|
||||
def raw(self, guardrail: str, action: str, marker: str) -> httpx.Response:
|
||||
return self.gateway.client.post(
|
||||
f"/bedrock/model/{ROUTER_MODEL}/{action}",
|
||||
json={"messages": [{"role": "user", "content": [{"text": marker}]}]},
|
||||
headers={"x-litellm-api-key": self.keys[guardrail]},
|
||||
timeout=case_timeout_seconds(),
|
||||
)
|
||||
|
||||
|
||||
def _generic_guardrail(name: str, mode: str, vendor: Wire, **settings: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {
|
||||
"guardrail": "generic_guardrail_api",
|
||||
"mode": mode,
|
||||
"default_on": False,
|
||||
"api_base": f"{vendor.url}/{name}",
|
||||
"api_key": "synthetic-guardrail-key",
|
||||
**settings,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
def provider(request: Request) -> Reply:
|
||||
if request.target == "/v1/messages":
|
||||
return Reply(status=429, body=ANTHROPIC_THROTTLE)
|
||||
assert request.target in tuple(
|
||||
f"/model/{MODEL_ID}/{action}" for action in ("converse-stream", "converse", "invoke")
|
||||
), request.target
|
||||
marker: Final = marker_of(request)
|
||||
error: Final = upstream_error_of(marker)
|
||||
if error is not None:
|
||||
return Reply(status=error.status, body=error.body, headers={"x-amzn-errortype": error.errortype})
|
||||
text: Final = f"The answer to {marker} is {BLOCKED_WORD}"
|
||||
if request.target.endswith("/converse-stream"):
|
||||
return Reply(content_type=EVENT_STREAM, chunks=converse_frames(text))
|
||||
return Reply(body=converse_reply(text))
|
||||
|
||||
def vendor(request: Request) -> Reply:
|
||||
if request.target.startswith(f"/{BEDROCK}/"):
|
||||
if not any(BLOCKED_WORD in text for text in bedrock_texts(request)):
|
||||
return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode())
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"outputs": [{"text": "Sorry, the model cannot answer this."}],
|
||||
"assessments": [
|
||||
{"wordPolicy": {"customWords": [{"match": BLOCKED_WORD, "action": "BLOCKED"}]}}
|
||||
],
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
if request.target.startswith(f"/{VENDOR_DOWN}"):
|
||||
return Reply(status=500, body=b'{"detail":"vendor is down"}')
|
||||
if request.target.startswith(f"/{CATO}/"):
|
||||
return Reply(
|
||||
status=THROTTLED.status, body=THROTTLED.body, headers={"x-amzn-errortype": THROTTLED.errortype}
|
||||
)
|
||||
texts: Final = [str(text) for text in json.loads(request.body)["texts"]]
|
||||
if any(BLOCKED_WORD in text for text in texts):
|
||||
return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": BLOCK_REASON}).encode())
|
||||
return Reply(body=json.dumps({"action": "NONE"}).encode())
|
||||
|
||||
with httpx.Client(base_url=os.environ["INTEGRATION_PROXY_URL"], timeout=30, trust_env=False) as root_client:
|
||||
root: Final = Gateway(
|
||||
root_client,
|
||||
os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"),
|
||||
os.environ["INTEGRATION_UPSTREAM_URL"],
|
||||
)
|
||||
directory: Final = tmp_path_factory.mktemp("bedrock-aws-errors")
|
||||
with wire_server(provider) as upstream, wire_server(vendor) as sink:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["guardrails"] = [
|
||||
_generic_guardrail(POST_CALL, "post_call", sink, streaming_buffer_until_moderated=True),
|
||||
_generic_guardrail(PRE_CALL, "pre_call", sink),
|
||||
_generic_guardrail(VENDOR_DOWN, "pre_call", sink),
|
||||
{
|
||||
"guardrail_name": CATO,
|
||||
"litellm_params": {
|
||||
"guardrail": "cato_networks",
|
||||
"mode": "pre_call",
|
||||
"default_on": False,
|
||||
"api_base": f"{sink.url}/{CATO}",
|
||||
"api_key": "synthetic-cato-key",
|
||||
},
|
||||
},
|
||||
{
|
||||
"guardrail_name": BEDROCK,
|
||||
"litellm_params": {
|
||||
"guardrail": "bedrock",
|
||||
"mode": "post_call",
|
||||
"default_on": False,
|
||||
"guardrailIdentifier": "aws-errors",
|
||||
"guardrailVersion": "1",
|
||||
"aws_access_key_id": "AKIASYNTHETICGUARD",
|
||||
"aws_secret_access_key": "synthetic-guardrail-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_bedrock_runtime_endpoint": f"{sink.url}/{BEDROCK}",
|
||||
},
|
||||
},
|
||||
]
|
||||
config["model_list"] = [
|
||||
{
|
||||
"model_name": name,
|
||||
"litellm_params": {
|
||||
"model": f"bedrock/{model}",
|
||||
"api_base": upstream.url,
|
||||
"aws_access_key_id": "AKIASYNTHETICPROVIDER",
|
||||
"aws_secret_access_key": "synthetic-provider-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
}
|
||||
for name, model in ((ROUTER_MODEL, MODEL_ID), ("anthropic.*", "anthropic.*"))
|
||||
]
|
||||
config["router_settings"] = {**config.get("router_settings", {}), "num_retries": 0}
|
||||
config_path: Final = directory / "bedrock-aws-errors.yaml"
|
||||
config_path.write_text(yaml.safe_dump(config))
|
||||
overrides: Final = {"ANTHROPIC_API_BASE": upstream.url, "ANTHROPIC_API_KEY": "sk-ant-synthetic"}
|
||||
with owned_proxy_process(root, directory, overrides, config=config_path, workers=2) as owned:
|
||||
with owned.gateway.scenario() as scenario:
|
||||
yield Rig(
|
||||
owned.gateway,
|
||||
upstream,
|
||||
sink,
|
||||
{
|
||||
UNGUARDED: scenario.key(),
|
||||
OVER_BUDGET: scenario.key(max_budget=0),
|
||||
ONE_RPM: scenario.key(rpm_limit=1),
|
||||
**{
|
||||
name: scenario.key(guardrails=[name])
|
||||
for name in (POST_CALL, PRE_CALL, BEDROCK, VENDOR_DOWN, CATO)
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
BLOCKS: Final = (
|
||||
pytest.param(POST_CALL, "converse-stream", BLOCK_REASON, 1, id="post-call-held-converse-stream"),
|
||||
pytest.param(POST_CALL, "converse", BLOCK_REASON, 1, id="post-call-converse"),
|
||||
pytest.param(PRE_CALL, "converse-stream", BLOCK_REASON, 0, id="pre-call-converse-stream"),
|
||||
pytest.param(PRE_CALL, "converse", BLOCK_REASON, 0, id="pre-call-converse"),
|
||||
pytest.param(BEDROCK, "converse-stream", BEDROCK_BLOCK_MESSAGE, 1, id="bedrock-guardrail-converse-stream"),
|
||||
pytest.param(BEDROCK, "converse", BEDROCK_BLOCK_MESSAGE, 1, id="bedrock-guardrail-converse"),
|
||||
pytest.param(PRE_CALL, "invoke", BLOCK_REASON, 0, id="pre-call-invoke"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ROUTES)
|
||||
@pytest.mark.parametrize(("guardrail", "action", "message", "provider_calls"), BLOCKS)
|
||||
def test_guardrail_block_reaches_boto3_as_a_validation_exception_with_the_guardrail_message(
|
||||
rig: Rig, guardrail: str, action: str, message: str, provider_calls: int, model: str
|
||||
) -> None:
|
||||
marker: Final = f"{BLOCKED_WORD}-{uuid.uuid4().hex}"
|
||||
error: Final = rig.call(guardrail, action, marker, model=model)
|
||||
assert error is not None, f"{guardrail} {action} was not blocked"
|
||||
assert error.response["ResponseMetadata"]["HTTPStatusCode"] == 400, error.response
|
||||
assert (error.response["Error"]["Code"], error.response["Error"]["Message"]) == (
|
||||
"ValidationException",
|
||||
message,
|
||||
), error.response
|
||||
assert rig.provider_calls(marker) == provider_calls
|
||||
|
||||
|
||||
@pytest.mark.parametrize("action", ("converse-stream", "converse"))
|
||||
def test_guardrail_block_body_adds_the_aws_message_and_keeps_the_litellm_error(rig: Rig, action: str) -> None:
|
||||
marker: Final = f"{BLOCKED_WORD}-{uuid.uuid4().hex}"
|
||||
response: Final = rig.raw(POST_CALL, action, marker)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.headers.get("x-amzn-errortype") == "ValidationException", (dict(response.headers), response.text)
|
||||
assert response.headers.get("content-type", "").startswith("application/json"), dict(response.headers)
|
||||
body: Final = response.json()
|
||||
assert body["message"] == BLOCK_REASON, body
|
||||
assert body["error"]["message"] == BLOCK_REASON, body
|
||||
assert response.headers.get("x-litellm-call-id"), dict(response.headers)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", tuple(UPSTREAM_ERRORS))
|
||||
@pytest.mark.parametrize("action", ("converse-stream", "converse"))
|
||||
@pytest.mark.parametrize("guardrail", (UNGUARDED, POST_CALL))
|
||||
def test_bedrock_error_reaches_the_caller_as_bedrock_sent_it(rig: Rig, kind: str, action: str, guardrail: str) -> None:
|
||||
upstream: Final = UPSTREAM_ERRORS[kind]
|
||||
marker: Final = upstream_marker(kind)
|
||||
response: Final = rig.raw(guardrail, action, marker)
|
||||
assert response.status_code == upstream.status, response.text
|
||||
assert response.headers.get("x-amzn-errortype") == upstream.errortype, (dict(response.headers), response.text)
|
||||
assert response.content == upstream.body, response.text
|
||||
assert response.headers.get("content-type") == "application/json", dict(response.headers)
|
||||
assert response.headers.get("x-litellm-call-id"), dict(response.headers)
|
||||
assert rig.provider_calls(marker) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ROUTES)
|
||||
@pytest.mark.parametrize("action", ("converse-stream", "converse", "invoke"))
|
||||
@pytest.mark.parametrize("guardrail", (UNGUARDED, POST_CALL))
|
||||
def test_boto3_sees_bedrock_throttling_and_retries_it_as_against_bedrock(
|
||||
rig: Rig, action: str, guardrail: str, model: str
|
||||
) -> None:
|
||||
marker: Final = upstream_marker("throttled")
|
||||
error: Final = rig.call(guardrail, action, marker, attempts=3, model=model)
|
||||
assert error is not None, f"{action} was not throttled"
|
||||
assert (error.response["Error"]["Code"], error.response["Error"]["Message"]) == (
|
||||
THROTTLED.code,
|
||||
THROTTLE_MESSAGE,
|
||||
), error.response
|
||||
assert error.response["ResponseMetadata"]["HTTPStatusCode"] == 429, error.response
|
||||
assert error.response["ResponseMetadata"]["RetryAttempts"] == 2, error.response
|
||||
assert rig.provider_calls(marker) == 3
|
||||
|
||||
|
||||
def converse_with_key(rig: Rig, model: str, key: str) -> httpx.Response:
|
||||
return rig.gateway.client.post(
|
||||
f"/bedrock/model/{model}/converse",
|
||||
json={"messages": [{"role": "user", "content": [{"text": f"plain-{uuid.uuid4().hex}"}]}]},
|
||||
headers={"x-litellm-api-key": key},
|
||||
timeout=case_timeout_seconds(),
|
||||
)
|
||||
|
||||
|
||||
def assert_litellm_error(response: httpx.Response, status: int) -> None:
|
||||
assert response.status_code == status, response.text
|
||||
assert "x-amzn-errortype" not in response.headers, dict(response.headers)
|
||||
assert set(response.json()) == {"error"}, response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("guardrail", "model", "status"),
|
||||
(
|
||||
pytest.param(VENDOR_DOWN, ROUTER_MODEL, 500, id="guardrail-vendor-down"),
|
||||
pytest.param(UNGUARDED, "not-a-configured-model", 400, id="unknown-model"),
|
||||
pytest.param(OVER_BUDGET, ROUTER_MODEL, 422, id="over-budget"),
|
||||
),
|
||||
)
|
||||
def test_errors_that_are_not_guardrail_blocks_keep_the_litellm_error(
|
||||
rig: Rig, guardrail: str, model: str, status: int
|
||||
) -> None:
|
||||
assert_litellm_error(converse_with_key(rig, model, rig.keys[guardrail]), status)
|
||||
|
||||
|
||||
def test_guardrail_service_error_shaped_like_aws_keeps_the_litellm_error(rig: Rig) -> None:
|
||||
response: Final = converse_with_key(rig, ROUTER_MODEL, rig.keys[CATO])
|
||||
assert response.status_code == THROTTLED.status, response.text
|
||||
assert "x-amzn-errortype" not in response.headers, dict(response.headers)
|
||||
assert response.content != THROTTLED.body, response.text
|
||||
|
||||
|
||||
def test_unknown_key_keeps_the_litellm_error(rig: Rig) -> None:
|
||||
assert_litellm_error(converse_with_key(rig, ROUTER_MODEL, f"sk-unknown-{uuid.uuid4().hex}"), 401)
|
||||
|
||||
|
||||
def test_litellm_rate_limit_keeps_the_litellm_error(rig: Rig) -> None:
|
||||
admitted: Final = converse_with_key(rig, ROUTER_MODEL, rig.keys[ONE_RPM])
|
||||
assert admitted.status_code == 200, admitted.text
|
||||
assert_litellm_error(converse_with_key(rig, ROUTER_MODEL, rig.keys[ONE_RPM]), 429)
|
||||
|
||||
|
||||
def test_other_provider_passthrough_errors_carry_no_aws_error_type(rig: Rig) -> None:
|
||||
response: Final = rig.gateway.client.post(
|
||||
"/anthropic/v1/messages",
|
||||
json={"model": "claude-sonnet-5", "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}]},
|
||||
headers={"x-litellm-api-key": rig.keys[UNGUARDED]},
|
||||
timeout=case_timeout_seconds(),
|
||||
)
|
||||
assert response.status_code == 429, response.text
|
||||
assert "x-amzn-errortype" not in response.headers, dict(response.headers)
|
||||
assert response.content == ANTHROPIC_THROTTLE, response.text
|
||||
|
|
@ -24,10 +24,15 @@ from starlette.datastructures import FormData
|
|||
|
||||
import litellm
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from tests._master_key import MASTER_KEY as SHARED_MASTER_KEY
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.types.utils import CallTypesLiteral
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException
|
||||
from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.passthrough.main import raise_for_provider_status
|
||||
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
|
|
@ -40,6 +45,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
|||
azure_proxy_route,
|
||||
azure_speech_proxy_route,
|
||||
bedrock_llm_proxy_route,
|
||||
bedrock_passthrough_error_response,
|
||||
bedrock_proxy_route,
|
||||
create_pass_through_route,
|
||||
cursor_proxy_route,
|
||||
|
|
@ -1658,17 +1664,7 @@ class TestBedrockLLMProxyRoute:
|
|||
|
||||
bedrock_error_message = '{"message":"ContentBlock object at messages.0.content.0 must set one of the following keys: text, image, toolUse, toolResult, document, video."}'
|
||||
|
||||
# Create a mock httpx.Response for the error
|
||||
mock_error_response = Mock(spec=httpx.Response)
|
||||
mock_error_response.status_code = 400
|
||||
mock_error_response.aread = AsyncMock(return_value=bedrock_error_message.encode("utf-8"))
|
||||
|
||||
# Create the HTTPStatusError
|
||||
mock_http_error = httpx.HTTPStatusError(
|
||||
message="Bad Request",
|
||||
request=Mock(spec=httpx.Request),
|
||||
response=mock_error_response,
|
||||
)
|
||||
mock_http_error = _bedrock_status_error(400, bedrock_error_message.encode("utf-8"))
|
||||
|
||||
# Create mocks for all required parameters
|
||||
mock_request = MagicMock(spec=Request)
|
||||
|
|
@ -1926,6 +1922,284 @@ class TestBedrockLLMProxyRoute:
|
|||
assert "Blocked by guardrail" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
_THROTTLE_TYPE: Final = "ThrottlingException:http://internal.amazon.com/coral/com.amazon.bedrock/"
|
||||
_THROTTLE_BODY: Final = b'{"message":"Too many requests, please wait before trying again."}'
|
||||
|
||||
|
||||
def _status_error(
|
||||
status_code: int,
|
||||
content: bytes,
|
||||
headers: Mapping[str, str],
|
||||
url: str,
|
||||
raise_for_status: Callable[[httpx.Response], None],
|
||||
) -> httpx.HTTPStatusError:
|
||||
request: Final = httpx.Request("POST", url)
|
||||
response: Final = httpx.Response(status_code, headers=dict(headers), content=content, request=request)
|
||||
with pytest.raises(httpx.HTTPStatusError) as raised:
|
||||
raise_for_status(response)
|
||||
return raised.value
|
||||
|
||||
|
||||
def _bedrock_status_error(
|
||||
status_code: int,
|
||||
content: bytes,
|
||||
headers: Mapping[str, str] = MappingProxyType({}),
|
||||
url: str = "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse",
|
||||
) -> httpx.HTTPStatusError:
|
||||
"""An error Bedrock answered the passthrough request with."""
|
||||
return _status_error(status_code, content, headers, url, raise_for_provider_status)
|
||||
|
||||
|
||||
def _guardrail_service_status_error(
|
||||
status_code: int, content: bytes, headers: Mapping[str, str]
|
||||
) -> httpx.HTTPStatusError:
|
||||
"""An error a guardrail service's own HTTP call raised, from a URL shaped like a Bedrock model call."""
|
||||
return _status_error(
|
||||
status_code,
|
||||
content,
|
||||
headers,
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse",
|
||||
httpx.Response.raise_for_status,
|
||||
)
|
||||
|
||||
|
||||
def _guardrail_detail(message: str) -> dict[str, str]:
|
||||
return {"error": message, "guardrail_name": "block-demo"}
|
||||
|
||||
|
||||
class _ReplacesFailures(CustomLogger):
|
||||
def __init__(self, replacement: HTTPException) -> None:
|
||||
super().__init__()
|
||||
self.replacement: Final = replacement
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict,
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> HTTPException | None:
|
||||
return self.replacement
|
||||
|
||||
|
||||
class TestBedrockPassthroughAwsErrors:
|
||||
"""A guardrail block or an error Bedrock returned reaches an AWS SDK caller with its code and message."""
|
||||
|
||||
@staticmethod
|
||||
async def _answer(error: Exception) -> Response:
|
||||
return await bedrock_passthrough_error_response(
|
||||
ProxyBaseLLMRequestProcessing(data={}),
|
||||
error,
|
||||
UserAPIKeyAuth(api_key="test-key"),
|
||||
ProxyLogging(user_api_key_cache=DualCache()),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_error_reaches_the_caller_with_bedrocks_status_header_and_body(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
error: Final = _bedrock_status_error(
|
||||
429,
|
||||
_THROTTLE_BODY,
|
||||
{
|
||||
"x-amzn-errortype": _THROTTLE_TYPE,
|
||||
"x-amzn-requestid": "req-1",
|
||||
"content-type": "application/x-amz-json-1.1",
|
||||
},
|
||||
)
|
||||
|
||||
response: Final = await self._answer(error)
|
||||
|
||||
assert response.status_code == 429
|
||||
assert response.body == _THROTTLE_BODY
|
||||
assert response.headers["x-amzn-errortype"] == _THROTTLE_TYPE
|
||||
assert response.headers["content-type"] == "application/x-amz-json-1.1"
|
||||
assert response.headers["llm_provider-x-amzn-requestid"] == "req-1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_arn_model_error_reaches_the_caller_as_bedrock_sent_it(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
error: Final = _bedrock_status_error(
|
||||
429,
|
||||
_THROTTLE_BODY,
|
||||
{"x-amzn-errortype": _THROTTLE_TYPE},
|
||||
url="https://bedrock-runtime.us-east-1.amazonaws.com/model/"
|
||||
"arn%3Aaws%3Abedrock%3Aus-east-1%3A123%3Aapplication-inference-profile%2Fabc123/converse-stream?trace=a/b",
|
||||
)
|
||||
|
||||
response: Final = await self._answer(error)
|
||||
|
||||
assert (response.status_code, response.body) == (429, _THROTTLE_BODY)
|
||||
assert response.headers["x-amzn-errortype"] == _THROTTLE_TYPE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("error", "status_code", "body"),
|
||||
(
|
||||
pytest.param(
|
||||
_bedrock_status_error(502, b"<html>bad gateway</html>"),
|
||||
502,
|
||||
"<html>bad gateway</html>",
|
||||
id="bedrock-without-error-type",
|
||||
),
|
||||
pytest.param(
|
||||
_guardrail_service_status_error(429, _THROTTLE_BODY, {"x-amzn-errortype": _THROTTLE_TYPE}),
|
||||
429,
|
||||
_THROTTLE_BODY.decode(),
|
||||
id="guardrail-service-with-error-type",
|
||||
),
|
||||
),
|
||||
)
|
||||
async def test_http_error_bedrock_did_not_answer_with_an_aws_error_type_keeps_the_litellm_error(
|
||||
self, monkeypatch: pytest.MonkeyPatch, error: httpx.HTTPStatusError, status_code: int, body: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await self._answer(error)
|
||||
|
||||
assert exc_info.value.status_code == status_code
|
||||
assert exc_info.value.detail == {"error": body}
|
||||
assert "x-amzn-errortype" not in (exc_info.value.headers or {})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("error", "replacement", "code"),
|
||||
(
|
||||
pytest.param(
|
||||
_bedrock_status_error(429, _THROTTLE_BODY, {"x-amzn-errortype": _THROTTLE_TYPE}),
|
||||
HTTPException(status_code=503, detail={"error": "rewritten by a callback"}),
|
||||
"503",
|
||||
id="bedrock-error",
|
||||
),
|
||||
pytest.param(
|
||||
HTTPException(status_code=400, detail=_guardrail_detail("Violated guardrail policy")),
|
||||
HTTPException(status_code=429, detail={"error": "rewritten by a callback"}),
|
||||
"429",
|
||||
id="guardrail-block",
|
||||
),
|
||||
),
|
||||
)
|
||||
async def test_error_replaced_by_a_failure_hook_keeps_the_replacement(
|
||||
self, monkeypatch: pytest.MonkeyPatch, error: Exception, replacement: HTTPException, code: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [_ReplacesFailures(replacement)])
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await self._answer(error)
|
||||
|
||||
assert (exc_info.value.message, exc_info.value.code) == ("rewritten by a callback", code)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("error", "message", "status_code"),
|
||||
(
|
||||
pytest.param(
|
||||
GuardrailRaisedException(
|
||||
guardrail_name="block-demo",
|
||||
message="synthetic policy says no",
|
||||
should_wrap_with_default_message=False,
|
||||
blocked_content=True,
|
||||
),
|
||||
"synthetic policy says no",
|
||||
"400",
|
||||
id="guardrail-raised",
|
||||
),
|
||||
pytest.param(
|
||||
HTTPException(status_code=400, detail=_guardrail_detail("Violated guardrail policy")),
|
||||
"Violated guardrail policy",
|
||||
"400",
|
||||
id="guardrail-http-400",
|
||||
),
|
||||
pytest.param(
|
||||
HTTPException(status_code=403, detail=_guardrail_detail("blocked by policy")),
|
||||
"blocked by policy",
|
||||
"403",
|
||||
id="guardrail-http-403",
|
||||
),
|
||||
pytest.param(
|
||||
BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="pii"),
|
||||
"Blocked entity detected: EMAIL_ADDRESS by Guardrail: pii. This entity is not allowed to be used in this request.",
|
||||
"400",
|
||||
id="blocked-pii",
|
||||
),
|
||||
),
|
||||
)
|
||||
async def test_guardrail_block_is_raised_with_its_aws_error_type(
|
||||
self, monkeypatch: pytest.MonkeyPatch, error: Exception, message: str, status_code: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
||||
with pytest.raises(AwsErrorTypeProxyException) as exc_info:
|
||||
await self._answer(error)
|
||||
|
||||
assert exc_info.value.aws_error_type == "ValidationException"
|
||||
assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("error", "message", "code"),
|
||||
(
|
||||
pytest.param(
|
||||
HTTPException(status_code=400, detail={"error": "model not found"}),
|
||||
"model not found",
|
||||
"400",
|
||||
id="request-check-400",
|
||||
),
|
||||
pytest.param(
|
||||
HTTPException(status_code=403, detail="cache control not allowed"),
|
||||
"cache control not allowed",
|
||||
"403",
|
||||
id="request-check-403",
|
||||
),
|
||||
pytest.param(
|
||||
HTTPException(status_code=429, detail=_guardrail_detail("vendor rate limited")),
|
||||
"vendor rate limited",
|
||||
"429",
|
||||
id="429",
|
||||
),
|
||||
pytest.param(
|
||||
GuardrailRaisedException(
|
||||
guardrail_name="fail-closed",
|
||||
message="guardrail vendor unreachable",
|
||||
should_wrap_with_default_message=False,
|
||||
),
|
||||
"guardrail vendor unreachable",
|
||||
"400",
|
||||
id="guardrail-vendor-unreachable-400",
|
||||
),
|
||||
pytest.param(
|
||||
GuardrailRaisedException(
|
||||
guardrail_name="fail-closed",
|
||||
message="guardrail vendor unreachable",
|
||||
should_wrap_with_default_message=False,
|
||||
status_code=503,
|
||||
),
|
||||
"guardrail vendor unreachable",
|
||||
"503",
|
||||
id="guardrail-vendor-unreachable-503",
|
||||
),
|
||||
pytest.param(
|
||||
ModifyResponseException(message="replaced", model="bedrock-model", request_data={}),
|
||||
"replaced",
|
||||
"500",
|
||||
id="modify-response",
|
||||
),
|
||||
),
|
||||
)
|
||||
async def test_errors_that_are_not_guardrail_blocks_keep_the_litellm_error(
|
||||
self, monkeypatch: pytest.MonkeyPatch, error: Exception, message: str, code: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await self._answer(error)
|
||||
|
||||
assert type(exc_info.value) is ProxyException
|
||||
assert (exc_info.value.message, exc_info.value.code) == (message, code)
|
||||
|
||||
|
||||
class TestBedrockAgentRuntimePassthroughToggle:
|
||||
AGENT_RUNTIME_ENDPOINT: Final = "knowledgebases/KB1234567/retrieve"
|
||||
MODEL_ENDPOINT: Final = "model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/converse"
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.caching.redis_cluster_cache import RedisClusterCache
|
|||
from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import (
|
||||
AwsErrorTypeProxyException,
|
||||
LitellmUserRoles,
|
||||
ModelAccessDeniedProxyException,
|
||||
ProxyErrorTypes,
|
||||
|
|
@ -43,7 +44,7 @@ from litellm.proxy._types import (
|
|||
from litellm.proxy.auth.login_throttle import LoginThrottle
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash
|
||||
from litellm.proxy.proxy_server import app, initialize, openai_exception_handler
|
||||
from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler
|
||||
from litellm.utils import _invalidate_model_cost_lowercase_map
|
||||
from tests._master_key import MASTER_KEY
|
||||
|
||||
|
|
@ -11858,6 +11859,29 @@ async def test_openai_exception_handler_no_denial_log_for_plain_proxy_exception(
|
|||
assert [r for r in caplog.records if r.levelname == "WARNING"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aws_error_type_exception_handler_adds_the_aws_message_and_header():
|
||||
error: Final = ProxyException(
|
||||
message="synthetic policy says no",
|
||||
type="invalid_request_error",
|
||||
param=None,
|
||||
code=400,
|
||||
headers={"x-litellm-call-id": "call-1"},
|
||||
)
|
||||
|
||||
response: Final = await aws_error_type_exception_handler(
|
||||
_http_request_scope(), AwsErrorTypeProxyException(error, "ValidationException")
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert response.headers["x-amzn-errortype"] == "ValidationException"
|
||||
assert response.headers["x-litellm-call-id"] == "call-1"
|
||||
assert json.loads(response.body) == {
|
||||
"message": "synthetic policy says no",
|
||||
"error": {"message": "synthetic policy says no", "type": "invalid_request_error", "param": None, "code": "400"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_model_access_denial_logs_sanitized_internal_message(caplog):
|
||||
reservation = {"reserved_cost": 0.0, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue