diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index d371aab4581..31b0f06251d 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -432,27 +432,30 @@ class AnthropicMessagesHandler(BaseTranslation): message, _ = serialize_http_exception_detail(exc.detail) return tuple(anthropic_sse_error_frames(message)) - def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: + def build_block_response(self, exc: "ModifyResponseException") -> "AnthropicMessagesResponse": import uuid - from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( - FakeAnthropicMessagesStreamIterator, - ) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_response_usage, ) from litellm.types.utils import AnthropicMessagesResponse - block_response: Final = AnthropicMessagesResponse( + return AnthropicMessagesResponse( id=f"msg_{uuid.uuid4()}", type="message", role="assistant", content=[{"type": "text", "text": exc.message}], model=exc.model, stop_reason="end_turn", - usage=blocked_response_usage(getattr(exc, "original_response", None)), + usage=blocked_response_usage(exc.original_response), ) - return list(FakeAnthropicMessagesStreamIterator(response=block_response)) + + def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: + from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + + return list(FakeAnthropicMessagesStreamIterator(response=self.build_block_response(exc))) def _block_continuation_chunks( self, exc: "ModifyResponseException", responses_so_far: Sequence[object] diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1d88e133cf9..1a2e610d5c4 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -2,7 +2,7 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ -from collections.abc import Mapping +from collections.abc import AsyncGenerator, Mapping from types import MappingProxyType from typing import Final @@ -19,12 +19,10 @@ from litellm.anthropic_interface.exceptions import ( ) from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.litellm_core_utils.tool_fields import first_validation_error +from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.anthropic.messages.context_management import ( AnthropicContextManagementError, ) -from litellm.llms.base_llm.guardrail_translation.utils import ( - blocked_response_usage as _blocked_response_usage, -) from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ( @@ -44,6 +42,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg from litellm.proxy.common_utils.openai_error_payload import ( LITELLM_CALL_ID_HEADER, error_status_code, + litellm_call_id_headers, openai_error_param, openai_error_type, with_litellm_call_id, @@ -58,6 +57,8 @@ from litellm.router import Router router: Final = APIRouter() +_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) + def _with_provider_specific_fields(exc: ProxyException, detail: AnthropicErrorDetail) -> AnthropicErrorDetail: if not exc.provider_specific_fields: @@ -229,44 +230,22 @@ async def anthropic_response( request_data=_data, ) - # Create Anthropic-formatted response with violation message - import uuid + block_headers: Final = litellm_call_id_headers(base_llm_response_processor.litellm_call_id) or _NO_HEADERS + block_handler: Final = AnthropicMessagesHandler() - from litellm.types.utils import AnthropicMessagesResponse + if data.get("stream") is True: + block_chunks: Final = block_handler.build_block_sse_chunks(e) - # Report the blocked LLM response's real token usage (carried on the - # exception) instead of discarding it; zero for pre-call blocks. - _usage: Final = _blocked_response_usage(e.original_response) - - _anthropic_response: Final = AnthropicMessagesResponse( - id=f"msg_{uuid.uuid4()}", - type="message", - role="assistant", - content=[{"type": "text", "text": e.message}], - model=e.model, - stop_reason="end_turn", - usage=_usage, - ) - - if data.get("stream", None) is not None and data["stream"] is True: - # For streaming, use the standard SSE data generator - async def _passthrough_stream_generator(): - yield _anthropic_response - - selected_data_generator: Final = ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=_passthrough_stream_generator(), - user_api_key_dict=user_api_key_dict, - request_data=_data, - proxy_logging_obj=proxy_logging_obj, - ) + async def _blocked_stream() -> AsyncGenerator[str, None]: + for chunk in block_chunks: + yield chunk.decode() return await create_response( - generator=selected_data_generator, - media_type="text/event-stream", - headers={}, + generator=_blocked_stream(), media_type="text/event-stream", headers=block_headers ) - return _anthropic_response + fastapi_response.headers.update(block_headers) + return block_handler.build_block_response(e) except AnthropicContextManagementError as e: if e.status_code >= 500: # Server-side polyfill failures hit the failure hook for spend/alert diff --git a/tests/integration/messages_endpoint/test_guardrail_block_stream.py b/tests/integration/messages_endpoint/test_guardrail_block_stream.py new file mode 100644 index 00000000000..d487f81c757 --- /dev/null +++ b/tests/integration/messages_endpoint/test_guardrail_block_stream.py @@ -0,0 +1,129 @@ +import uuid +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import anthropic +import httpx +import pytest +import yaml +from anthropic.types import Message +from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, wire_server + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_DENIAL: Final = "This request is not allowed by policy." +_CALL_ID_HEADER: Final = "x-litellm-call-id" + + +@dataclass(frozen=True, slots=True) +class _Rail: + gateway: Gateway + names: tuple[str, str] + + +@dataclass(frozen=True, slots=True) +class _Ask: + model: str + rails: tuple[str, str] + call_id: str + + +@pytest.fixture(scope="module") +def denying_rail(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rail]: + directory: Final = tmp_path_factory.mktemp("messages-guardrail-block") + name: Final = f"messages-deny-{uuid.uuid4().hex}" + post_call_name: Final = f"messages-deny-replies-{uuid.uuid4().hex}" + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = directory / "messages-guardrail-block.yaml" + config.write_text( + yaml.safe_dump( + { + **stock, + "guardrails": [ + { + "guardrail_name": name, + "litellm_params": { + "guardrail": "custom_code", + "mode": "pre_call", + "default_on": False, + "custom_code": ( + f"def apply_guardrail(inputs, request_data, input_type):\n return block({_DENIAL!r})\n" + ), + }, + }, + { + "guardrail_name": post_call_name, + "litellm_params": { + "guardrail": "custom_code", + "mode": "post_call", + "default_on": False, + "custom_code": ( + "def apply_guardrail(inputs, request_data, input_type):\n" + " return block('A post-call rail scanned the block message.')\n" + ), + }, + }, + ], + } + ) + ) + with gateway_from_environment() as shared, owned_proxy(shared, directory, {}, config=config) as candidate: + yield _Rail(candidate, (name, post_call_name)) + + +def _provider_must_not_answer(request: Request) -> Reply: + return Reply(status=500, body=f"unexpected {request.method} {request.target}".encode()) + + +def _streamed(client: anthropic.Anthropic, ask: _Ask) -> tuple[Message, httpx.Headers]: + with client.messages.stream( + model=ask.model, + max_tokens=32, + messages=[{"role": "user", "content": "count from 1 to 100"}], + extra_body={"guardrails": list(ask.rails)}, + extra_headers={_CALL_ID_HEADER: ask.call_id}, + ) as stream: + return stream.get_final_message(), stream.response.headers + + +def _created(client: anthropic.Anthropic, ask: _Ask) -> tuple[Message, httpx.Headers]: + raw: Final = client.messages.with_raw_response.create( + model=ask.model, + max_tokens=32, + messages=[{"role": "user", "content": "count from 1 to 100"}], + extra_body={"guardrails": list(ask.rails)}, + extra_headers={_CALL_ID_HEADER: ask.call_id}, + ) + return raw.parse(), raw.headers + + +@pytest.mark.parametrize("send", (_streamed, _created), ids=("stream", "json")) +def test_anthropic_sdk_reads_a_pre_call_block_as_a_finished_message_past_post_call_rails( + denying_rail: _Rail, send: Callable[[anthropic.Anthropic, _Ask], tuple[Message, httpx.Headers]] +) -> None: + call_id: Final = f"messages-block-{uuid.uuid4().hex}" + with ( + wire_server(_provider_must_not_answer) as provider, + denying_rail.gateway.scenario() as scenario, + httpx.Client(timeout=30, trust_env=False) as transport, + anthropic.Anthropic( + api_key=denying_rail.gateway.key, + base_url=str(denying_rail.gateway.client.base_url).rstrip("/"), + max_retries=0, + http_client=transport, + ) as client, + ): + model: Final = scenario.model( + model="anthropic/claude-haiku-5-5", api_base=provider.url, api_key="synthetic-anthropic-key" + ) + message, headers = send(client, _Ask(model, denying_rail.names, call_id)) + sent: Final = provider.drain() + + assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", _DENIAL)], message + assert message.stop_reason == "end_turn", message + assert headers.get(_CALL_ID_HEADER) == call_id, dict(headers) + assert sent == (), sent diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 8f83287696c..537eb2bbf69 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -1902,11 +1902,20 @@ def test_messages_pre_call_denial_streams_message(gateway: Gateway, tmp_path: Pa ) assert response.status_code == 200, response.text assert response.headers["content-type"].startswith("text/event-stream"), response.text - lines: Final = tuple(line for line in response.text.split("\n") if line.startswith("data: ")) - assert len(lines) == 1, response.text - body: Final = json.loads(lines[0].removeprefix("data: ")) - assert body["type"] == "message" and body["role"] == "assistant", body - assert body["content"] == [{"type": "text", "text": _RESPONSES_DENIAL}], body + events: Final = tuple( + json.loads(line.removeprefix("data: ")) for line in response.text.split("\n") if line.startswith("data: ") + ) + assert [event["type"] for event in events] == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ], response.text + assert events[0]["message"]["role"] == "assistant", events[0] + assert events[2]["delta"] == {"type": "text_delta", "text": _RESPONSES_DENIAL}, events[2] + assert events[4]["delta"]["stop_reason"] == "end_turn", events[4] @pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_writes_zero_spend_row") diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 037b0a6c29a..e3e625be975 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -5,18 +5,25 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objec import json import logging import unittest -from collections.abc import Callable, Mapping +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from fastapi import Request, Response +from fastapi.responses import StreamingResponse from fastapi.testclient import TestClient +from starlette.types import Message from litellm import Router from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from litellm.types.utils import TokenCountResponse +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ModelResponseStream, TokenCountResponse class TestAnthropicEndpoints(unittest.TestCase): @@ -70,11 +77,76 @@ class TestAnthropicEndpoints(unittest.TestCase): assert mock_safe_dumps.call_count == 2 # Called twice, once for each dict object +_BLOCK_CALL_ID: Final = "messages-block" +_BLOCK_TEXT: Final = "blocked by guardrail" + + +class _BlockEveryRequest(CustomGuardrail): + async def async_pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str + ) -> None: + self.raise_passthrough_exception(violation_message=_BLOCK_TEXT, request_data=data) + + +class _RewriteEveryStreamedReply(CustomGuardrail): + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[ModelResponseStream, None]: + async for _ in response: + yield ModelResponseStream(choices=[]) + + +def _messages_request(body: Mapping[str, object]) -> Request: + async def receive() -> Message: + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/v1/messages", + "query_string": b"", + "headers": [(b"content-type", b"application/json"), (b"x-litellm-call-id", _BLOCK_CALL_ID.encode())], + }, + receive, + ) + + +async def _blocked_messages_reply( + body: Mapping[str, object], fastapi_response: Response, rails: tuple[CustomGuardrail, ...] +) -> tuple[object, str]: + import litellm + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.utils import ProxyLogging + + blocker: Final = _BlockEveryRequest( + guardrail_name="deny-list", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + with ( + patch.object( + proxy_server, "proxy_logging_obj", new=ProxyLogging(user_api_key_cache=DualCache()) + ), # test-quality-ok: module global imported at call time; no injection seam + patch.object(litellm, "callbacks", new=[blocker, *rails]), # test-quality-ok: guardrails register through this global + ): + response = await ep.anthropic_response( + fastapi_response=fastapi_response, + request=_messages_request(body), + user_api_key_dict=UserAPIKeyAuth(), + ) + if not isinstance(response, StreamingResponse): + return response, "" + streamed: Final = b"".join( + [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator] + ).decode() + return response, streamed + + class TestBlockedResponseUsage: """Blocked responses report the blocked LLM response's real usage.""" def test_uses_original_response_usage(self): - from litellm.proxy.anthropic_endpoints.endpoints import _blocked_response_usage + from litellm.llms.base_llm.guardrail_translation.utils import blocked_response_usage as _blocked_response_usage # original_response is the AnthropicMessagesResponse the LLM produced # before the guardrail blocked it; its usage is real. @@ -85,7 +157,7 @@ class TestBlockedResponseUsage: } def test_zero_usage_when_no_original_response(self): - from litellm.proxy.anthropic_endpoints.endpoints import _blocked_response_usage + from litellm.llms.base_llm.guardrail_translation.utils import blocked_response_usage as _blocked_response_usage # Pre-call blocks never invoked the LLM -> nothing consumed. assert _blocked_response_usage(None) == { @@ -131,6 +203,49 @@ class TestBlockedResponseUsage: assert response["usage"] == {"input_tokens": 12, "output_tokens": 5} mock_logging.post_call_failure_hook.assert_awaited_once() + @pytest.mark.asyncio + @pytest.mark.parametrize( + "rails", + ((), (_RewriteEveryStreamedReply(guardrail_name="post-call-rewrite", default_on=True),)), + ids=("no_rails", "post_call_rail"), + ) + async def test_streaming_block_answers_anthropic_events_with_the_call_id( + self, rails: tuple[CustomGuardrail, ...] + ) -> None: + response, streamed = await _blocked_messages_reply( + {"model": "claude-haiku-5-5", "stream": True}, Response(), rails + ) + + frames: Final = tuple(frame for frame in streamed.split("\n\n") if frame.strip()) + events: Final = tuple( + (frame.split("\n")[0].removeprefix("event: "), json.loads(frame.split("\n")[1].removeprefix("data: "))) + for frame in frames + ) + assert isinstance(response, StreamingResponse) + assert response.media_type == "text/event-stream" + assert response.headers["x-litellm-call-id"] == _BLOCK_CALL_ID + assert [name for name, _ in events] == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + assert all(payload["type"] == name for name, payload in events) + assert events[2][1]["delta"] == {"type": "text_delta", "text": _BLOCK_TEXT} + assert events[4][1]["delta"]["stop_reason"] == "end_turn" + + @pytest.mark.asyncio + async def test_json_block_answers_with_the_call_id(self) -> None: + fastapi_response: Final = Response() + + response, _ = await _blocked_messages_reply({"model": "claude-haiku-5-5"}, fastapi_response, ()) + + assert fastapi_response.headers["x-litellm-call-id"] == _BLOCK_CALL_ID + assert isinstance(response, Mapping) + assert response["content"] == [{"type": "text", "text": _BLOCK_TEXT}] + class TestProxyExceptionAnthropicEnvelope: @pytest.mark.asyncio