mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
fix(proxy): stream a guardrail block on /v1/messages as Anthropic events (#49010)
* fix(proxy): stream a guardrail block on /v1/messages as Anthropic events * fix(proxy): keep post-call guardrails off the streamed messages block The streamed block is sent as-is, like the Responses endpoint's block, so post-call guardrails and streaming hooks never scan the proxy's own block message. Tests cover a post-call rail on the stream and the call id on the JSON block * test(proxy): expect Anthropic events from a streamed messages guardrail block The legacy guardrail suite pinned the old single bare data: line, which is the malformed block this branch replaces. The test now checks the six Anthropic stream events and the block text, so it fails if the bare line comes back * test(proxy): block messages through a real pre-call guardrail in endpoint tests The block tests replaced ProxyBaseLLMRequestProcessing.base_process_llm_request on the class to raise the block. They now send a real request through the endpoint with a pre-call guardrail that raises the passthrough block, and check that the x-litellm-call-id header carries the request's call id --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tessellark
GitHub
parent
444423a5ca
commit
41602e318d
5 changed files with 287 additions and 52 deletions
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue