mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(guardrails): scan and re-emit raw Anthropic SSE streams in the bedrock post-call hook (#36598)
* fix(guardrails): scan and re-emit raw Anthropic SSE streams in the bedrock post-call hook * fix(guardrails): keep upstream id and model on a blocked Anthropic stream * fix(guardrails): deliver a blocked Anthropic stream as an error frame * fix(guardrails): deliver an unscannable Anthropic stream as an error frame * fix(guardrails): emit the guardrail block detail as JSON in the stream error frame * fix(guardrails): deliver an Anthropic block through the shared block-SSE builder * fix(guardrails): keep the shared SSE assembler behavior-identical for existing callers * fix(guardrails): keep the stream error message a string and drop an unreachable branch * chore(guardrails): drop a comment that repeated its own docstring * fix(guardrails): let bedrock service failures keep their status instead of framing them as blocks * fix(guardrails): key the streamed block decision on status, not detail shape InvokeGuardrailChecks details a Mapping on its 500 for an unparseable response, so a detail-shape test read that outage as a policy block and framed it as a 200 guardrail_error. Both block sites raise 400, so gate on the status too. * refactor(guardrails): narrow the SSE error-frame helper to the input it actually takes Both callers pass a string, so the Mapping overload and its json.dumps branch were unreachable. Folds the block branch's narrative comment into the rebind suppressions that already carry a reason.
This commit is contained in:
yucheng-berri
GitHub
parent
3d76dfc72e
commit
6f84c468d4
6 changed files with 533 additions and 67 deletions
|
|
@ -21,6 +21,17 @@ def _coerce_interval(ping_interval_seconds: float | str | None) -> float | None:
|
|||
return interval
|
||||
|
||||
|
||||
def keepalive_ping_has_fired(elapsed_seconds: float, ping_interval_seconds: float | str | None) -> bool:
|
||||
"""Whether a keepalive ping has already gone out, which flushes the response headers.
|
||||
|
||||
A caller that discovers a failure after that point cannot raise its way to the client, since
|
||||
the status line is already on the wire. With pings disabled nothing flushes early, so a raise
|
||||
still carries its real status.
|
||||
"""
|
||||
interval: Final = _coerce_interval(ping_interval_seconds)
|
||||
return interval is not None and elapsed_seconds >= interval
|
||||
|
||||
|
||||
def wrap_sse_stream_with_keepalive_pings(
|
||||
stream: AsyncGenerator[str, None],
|
||||
ping_interval_seconds: float | str | None,
|
||||
|
|
|
|||
125
litellm/proxy/guardrails/anthropic_sse.py
Normal file
125
litellm/proxy/guardrails/anthropic_sse.py
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
"""Anthropic SSE <-> ModelResponse conversion for guardrail streaming hooks.
|
||||
|
||||
`/v1/messages` streams reach a guardrail's `async_post_call_streaming_iterator_hook` as raw SSE
|
||||
frames rather than chunk objects, which `stream_chunk_builder` cannot assemble. These helpers let a
|
||||
hook scan such a stream, and re-emit it when the guardrail rewrote the response.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
from litellm.types.utils import Choices, ModelResponse
|
||||
|
||||
|
||||
def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool:
|
||||
return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
|
||||
|
||||
|
||||
def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None:
|
||||
raw: Final = b"".join(
|
||||
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
|
||||
for chunk in all_chunks
|
||||
if isinstance(chunk, (str, bytes))
|
||||
)
|
||||
try:
|
||||
return raw.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return None
|
||||
|
||||
|
||||
def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None:
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
||||
AnthropicPassthroughLoggingHandler,
|
||||
)
|
||||
|
||||
return next(
|
||||
(
|
||||
message
|
||||
for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses
|
||||
if (event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None # pyright: ignore[reportPrivateUsage] # same parser the assembler uses; a private import beats forking SSE parsing
|
||||
and event_data.get("type") == "message_start"
|
||||
and isinstance(message := event_data.get("message"), dict)
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def assemble_anthropic_sse_stream(
|
||||
all_chunks: Sequence[object], *, restore_identity: bool = False
|
||||
) -> ModelResponse | None:
|
||||
"""Assemble raw Anthropic SSE frames into a ModelResponse.
|
||||
|
||||
``restore_identity`` stamps the upstream message id and model onto the result, which the
|
||||
assembler does not carry through. It is off by default so callers that re-emit the assembled
|
||||
response keep the wire shape they had before this helper was shared. The writes land on a
|
||||
freshly built object that is unreachable from caller state until returned.
|
||||
"""
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
||||
AnthropicPassthroughLoggingHandler,
|
||||
)
|
||||
|
||||
sse_stream: Final = _joined_sse_stream(all_chunks)
|
||||
if sse_stream is None:
|
||||
return None
|
||||
message_start: Final = _anthropic_message_start(sse_stream)
|
||||
if message_start is None:
|
||||
return None
|
||||
model: Final = message_start.get("model") if restore_identity else None
|
||||
try:
|
||||
assembled: Final = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only SSE-to-ModelResponse assembler; reimplementing it here would fork the parser
|
||||
all_chunks=(sse_stream,),
|
||||
litellm_logging_obj=None, # pyright: ignore[reportArgumentType] # only forwarded to stream_chunk_builder, which accepts None
|
||||
model=model if isinstance(model, str) else "",
|
||||
)
|
||||
except Exception: # noqa: BLE001 # stream_chunk_builder re-raises every assembly failure as litellm.APIError
|
||||
return None
|
||||
if not isinstance(assembled, ModelResponse):
|
||||
return None
|
||||
if not restore_identity:
|
||||
return assembled
|
||||
message_id: Final = message_start.get("id")
|
||||
if isinstance(message_id, str):
|
||||
assembled.id = message_id
|
||||
if isinstance(model, str) and model:
|
||||
assembled.model = model
|
||||
return assembled
|
||||
|
||||
|
||||
def model_response_text(response: ModelResponse) -> str:
|
||||
"""Assistant text of a response, used to detect whether a guardrail rewrote it."""
|
||||
return "".join(
|
||||
choice.message.content
|
||||
for choice in response.choices
|
||||
if isinstance(choice, Choices) # pyright: ignore[reportUnnecessaryIsInstance] # runtime choices can be StreamingChoices
|
||||
and isinstance(choice.message.content, str)
|
||||
)
|
||||
|
||||
|
||||
def anthropic_sse_error_frames(message: str) -> tuple[bytes, ...]:
|
||||
"""Anthropic error event, for a failure discovered after the response headers were flushed.
|
||||
|
||||
Once a keepalive ping has been sent a raise cannot reach the client, so the failure has to
|
||||
travel as a frame.
|
||||
"""
|
||||
body: Final = json.dumps(message)
|
||||
return (
|
||||
f'event: error\ndata: {{"type": "error", "error": {{"type": "guardrail_error", '
|
||||
f'"message": {body}}}}}\n\n'.encode(),
|
||||
)
|
||||
|
||||
|
||||
def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]:
|
||||
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
|
||||
LiteLLMAnthropicMessagesAdapter,
|
||||
)
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
|
||||
anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(
|
||||
response=assembled
|
||||
)
|
||||
return tuple(FakeAnthropicMessagesStreamIterator(response=anthropic_response).chunks)
|
||||
|
|
@ -14,6 +14,7 @@ import copy
|
|||
import json
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from itertools import accumulate, groupby
|
||||
|
|
@ -30,6 +31,7 @@ from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
|
|||
from litellm.exceptions import ModifyResponseException
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
effective_scan_only_tool_results_for_guardrail,
|
||||
)
|
||||
|
|
@ -39,6 +41,15 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_request_processing import _serialize_http_exception_detail
|
||||
from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired
|
||||
from litellm.proxy.guardrails.anthropic_sse import (
|
||||
anthropic_sse_chunks_from_response,
|
||||
anthropic_sse_error_frames,
|
||||
assemble_anthropic_sse_stream,
|
||||
is_raw_sse_stream,
|
||||
model_response_text,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
|
||||
|
|
@ -2578,14 +2589,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
from litellm.types.utils import TextCompletionResponse
|
||||
|
||||
# Collect all chunks to process them together
|
||||
started_at: Final = time.monotonic()
|
||||
all_chunks: Final[list[ModelResponseStream]] = []
|
||||
async for chunk in response:
|
||||
all_chunks.append(chunk)
|
||||
|
||||
assembled_model_response: ModelResponse | TextCompletionResponse | None = stream_chunk_builder(
|
||||
chunks=all_chunks,
|
||||
# /v1/messages arrives as SSE frames, which stream_chunk_builder cannot assemble
|
||||
raw_sse: Final = is_raw_sse_stream(all_chunks)
|
||||
assembled_model_response: ModelResponse | TextCompletionResponse | None = (
|
||||
assemble_anthropic_sse_stream(all_chunks, restore_identity=True)
|
||||
if raw_sse
|
||||
else stream_chunk_builder(chunks=all_chunks)
|
||||
)
|
||||
if isinstance(assembled_model_response, ModelResponse):
|
||||
pre_guardrail_text: Final = model_response_text(assembled_model_response)
|
||||
_pre_block_response: Final = assembled_model_response
|
||||
####################################################################
|
||||
########## 1. Make Bedrock Apply Guardrail API request ##########
|
||||
#
|
||||
|
|
@ -2609,7 +2627,32 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
request_data=request_data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
except HTTPException as block_exc:
|
||||
block_detail: Final = block_exc.detail
|
||||
# A policy block is the only 400 carrying a structured detail; a service failure
|
||||
# either details a plain string or reports a non-400 status. Re-raising a service
|
||||
# failure keeps its real status, but only while the headers are unflushed: past the
|
||||
# first keepalive ping the raise reaches nobody, so it has to travel as a frame too
|
||||
is_block: Final = raw_sse and block_exc.status_code == 400 and isinstance(block_detail, Mapping)
|
||||
headers_flushed: Final = keepalive_ping_has_fired(
|
||||
time.monotonic() - started_at, litellm.anthropic_sse_ping_interval_seconds
|
||||
)
|
||||
if not raw_sse or (not is_block and not headers_flushed):
|
||||
raise
|
||||
block_message, _ = _serialize_http_exception_detail(block_detail)
|
||||
for error_frame in anthropic_sse_error_frames(
|
||||
block_message if is_block else f"{block_exc.status_code}: {block_message}"
|
||||
):
|
||||
yield error_frame
|
||||
return
|
||||
except ModifyResponseException as e:
|
||||
if raw_sse:
|
||||
e.model = _pre_block_response.model or e.model # rebind-ok: exc.model defaults to the guardrail
|
||||
if e.original_response is None:
|
||||
e.original_response = _pre_block_response # rebind-ok: the block builder reads usage off this
|
||||
for block_chunk in AnthropicMessagesHandler().build_block_sse_chunks(e, stream_started=False):
|
||||
yield block_chunk
|
||||
return
|
||||
# Preserve upstream usage from the LLM call we already
|
||||
# consumed. Non-streaming blocks carry it via
|
||||
# ModifyResponseException.original_response +
|
||||
|
|
@ -2642,11 +2685,29 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################################
|
||||
########## 3. Return the (potentially masked) chunks ##########
|
||||
#########################################################################
|
||||
if raw_sse:
|
||||
for sse_chunk in (
|
||||
anthropic_sse_chunks_from_response(assembled_model_response)
|
||||
if model_response_text(assembled_model_response) != pre_guardrail_text
|
||||
else all_chunks
|
||||
):
|
||||
yield sse_chunk
|
||||
return
|
||||
|
||||
mock_response: Final = MockResponseIterator(model_response=assembled_model_response)
|
||||
|
||||
# Return the reconstructed stream
|
||||
async for chunk in mock_response:
|
||||
yield chunk
|
||||
elif raw_sse:
|
||||
# Forwarding an unscannable stream would silently disable the guardrail, so fail closed.
|
||||
# A raise cannot reach the client once a keepalive ping has flushed the headers, so the
|
||||
# refusal travels as a frame, matching how a block is delivered above
|
||||
for error_frame in anthropic_sse_error_frames(
|
||||
f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it"
|
||||
):
|
||||
yield error_frame
|
||||
return
|
||||
else:
|
||||
for chunk in all_chunks:
|
||||
yield chunk
|
||||
|
|
|
|||
|
|
@ -17,6 +17,11 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
from litellm.proxy.guardrails.anthropic_sse import (
|
||||
anthropic_sse_chunks_from_response,
|
||||
assemble_anthropic_sse_stream,
|
||||
is_raw_sse_stream,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
|
||||
PermissionError,
|
||||
|
|
@ -870,7 +875,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
|
|||
all_chunks.append(chunk)
|
||||
|
||||
assembled_model_response: Final[ModelResponse | TextCompletionResponse | None] = (
|
||||
stream_chunk_builder(chunks=all_chunks) if not self._is_raw_sse_stream(all_chunks) else None
|
||||
stream_chunk_builder(chunks=all_chunks) if not is_raw_sse_stream(all_chunks) else None
|
||||
)
|
||||
if isinstance(assembled_model_response, ModelResponse):
|
||||
denied_tools = self._check_assembled_stream(assembled_model_response)
|
||||
|
|
@ -883,9 +888,9 @@ class ToolPermissionGuardrail(CustomGuardrail):
|
|||
yield chunk
|
||||
return
|
||||
|
||||
anthropic_response: Final = self._assemble_anthropic_stream(all_chunks)
|
||||
anthropic_response: Final = assemble_anthropic_sse_stream(all_chunks)
|
||||
if anthropic_response is None:
|
||||
if self._is_raw_sse_stream(all_chunks):
|
||||
if is_raw_sse_stream(all_chunks):
|
||||
raise GuardrailRaisedException(
|
||||
guardrail_name=self.guardrail_name,
|
||||
message=(
|
||||
|
|
@ -904,13 +909,9 @@ class ToolPermissionGuardrail(CustomGuardrail):
|
|||
return
|
||||
|
||||
self._modify_response_with_permission_errors(anthropic_response, anthropic_denials)
|
||||
for sse_chunk in self._rewritten_anthropic_sse_chunks(anthropic_response):
|
||||
for sse_chunk in anthropic_sse_chunks_from_response(anthropic_response):
|
||||
yield sse_chunk
|
||||
|
||||
@staticmethod
|
||||
def _is_raw_sse_stream(all_chunks: Sequence[Any]) -> bool:
|
||||
return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
|
||||
|
||||
def _check_assembled_stream(
|
||||
self, assembled: ModelResponse
|
||||
) -> tuple[tuple[ChatCompletionMessageToolCall, PermissionError], ...]:
|
||||
|
|
@ -924,60 +925,3 @@ class ToolPermissionGuardrail(CustomGuardrail):
|
|||
if not denied_tools:
|
||||
verbose_proxy_logger.debug("Tool Permission Guardrail Post-Call Hook: All tools allowed")
|
||||
return denied_tools
|
||||
|
||||
@staticmethod
|
||||
def _joined_sse_stream(all_chunks: Sequence[Any]) -> str | None:
|
||||
raw: Final = b"".join(
|
||||
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
|
||||
for chunk in all_chunks
|
||||
if isinstance(chunk, (str, bytes))
|
||||
)
|
||||
try:
|
||||
return raw.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _has_anthropic_message_start(sse_stream: str) -> bool:
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
||||
AnthropicPassthroughLoggingHandler,
|
||||
)
|
||||
|
||||
return any(
|
||||
(event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None # pyright: ignore[reportPrivateUsage] # same parser the assembler uses; a private import beats forking SSE parsing
|
||||
and event_data.get("type") == "message_start"
|
||||
for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _assemble_anthropic_stream(all_chunks: Sequence[Any]) -> ModelResponse | None:
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
||||
AnthropicPassthroughLoggingHandler,
|
||||
)
|
||||
|
||||
sse_stream: Final = ToolPermissionGuardrail._joined_sse_stream(all_chunks)
|
||||
if sse_stream is None or not ToolPermissionGuardrail._has_anthropic_message_start(sse_stream):
|
||||
return None
|
||||
try:
|
||||
assembled = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only SSE-to-ModelResponse assembler; reimplementing it here would fork the parser
|
||||
all_chunks=(sse_stream,),
|
||||
litellm_logging_obj=None, # pyright: ignore[reportArgumentType] # only forwarded to stream_chunk_builder, which accepts None
|
||||
model="",
|
||||
)
|
||||
except (AttributeError, TypeError, ValueError, json.JSONDecodeError):
|
||||
return None
|
||||
return assembled if isinstance(assembled, ModelResponse) else None
|
||||
|
||||
@staticmethod
|
||||
def _rewritten_anthropic_sse_chunks(assembled: ModelResponse) -> tuple[bytes, ...]:
|
||||
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
|
||||
LiteLLMAnthropicMessagesAdapter,
|
||||
)
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
|
||||
anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(
|
||||
response=assembled
|
||||
)
|
||||
return tuple(FakeAnthropicMessagesStreamIterator(response=anthropic_response).chunks)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ sys.path.insert(0, os.path.abspath("../../../../../.."))
|
|||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
|
||||
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
||||
|
|
@ -2846,6 +2847,292 @@ async def test_apply_guardrail_propagates_modify_response_on_block():
|
|||
assert exc_info.value.message == "Sorry, the model cannot answer this question."
|
||||
|
||||
|
||||
_ANTHROPIC_SSE_CHUNKS = (
|
||||
b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_1","type":"message",'
|
||||
b'"role":"assistant","model":"claude","content":[],"usage":{"input_tokens":5,"output_tokens":0}}}\n\n',
|
||||
b'event: content_block_start\ndata: {"type":"content_block_start","index":0,'
|
||||
b'"content_block":{"type":"text","text":""}}\n\n',
|
||||
b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,'
|
||||
b'"delta":{"type":"text_delta","text":"my ssn is 123-45-6789"}}\n\n',
|
||||
b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n',
|
||||
b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn"},'
|
||||
b'"usage":{"output_tokens":9}}\n\n',
|
||||
b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
|
||||
)
|
||||
|
||||
|
||||
async def _anthropic_sse_stream():
|
||||
for chunk in _ANTHROPIC_SSE_CHUNKS:
|
||||
yield chunk
|
||||
|
||||
|
||||
async def _drain_streaming_hook(
|
||||
guardrail: BedrockGuardrail, request_data: dict[str, object] | None = None
|
||||
) -> list[object]:
|
||||
return [
|
||||
chunk
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=_anthropic_sse_stream(),
|
||||
request_data=request_data
|
||||
if request_data is not None
|
||||
else {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "what is my ssn"}]},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _sse_guardrail(**kwargs: object) -> BedrockGuardrail:
|
||||
return BedrockGuardrail(
|
||||
guardrail_name="bedrock-sse",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
event_hook=GuardrailEventHooks.post_call,
|
||||
default_on=True,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_scans_raw_anthropic_sse_instead_of_crashing():
|
||||
"""A /v1/messages stream arrives as raw SSE frames and must be assembled, then scanned.
|
||||
|
||||
Regression for `500 Error building chunks for logging/streaming usage calculation`:
|
||||
stream_chunk_builder subscripts each chunk, which raises TypeError on bytes.
|
||||
"""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {"action": "NONE"}
|
||||
delivered = await _drain_streaming_hook(guardrail)
|
||||
|
||||
mock_api.assert_called_once()
|
||||
kwargs = mock_api.call_args.kwargs
|
||||
assert kwargs["source"] == "OUTPUT"
|
||||
assert "my ssn is 123-45-6789" in str(kwargs["response"].choices[0].message.content)
|
||||
assert kwargs["messages"] == [{"role": "user", "content": "what is my ssn"}]
|
||||
assert tuple(delivered) == _ANTHROPIC_SSE_CHUNKS
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_emits_masked_text_for_raw_anthropic_sse():
|
||||
"""Masking must reach the client on /v1/messages, with mask_response_content unset.
|
||||
|
||||
The assembled path masks regardless of the flag, so forwarding the original frames here
|
||||
would ship exactly the text the guardrail redacted.
|
||||
"""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"outputs": [{"text": "my ssn is {SSN}"}],
|
||||
}
|
||||
delivered = await _drain_streaming_hook(guardrail)
|
||||
|
||||
body = b"".join(delivered)
|
||||
assert b"{SSN}" in body
|
||||
assert b"123-45-6789" not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_block_stream_keeps_upstream_identity():
|
||||
"""A blocked stream must carry the same id and model as the mask path, not the proxy alias."""
|
||||
guardrail = _sse_guardrail(disable_exception_on_block=True)
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.side_effect = ModifyResponseException(
|
||||
message="Sorry, the model cannot answer this question.",
|
||||
model="my-proxy-alias",
|
||||
request_data={},
|
||||
)
|
||||
delivered = await _drain_streaming_hook(guardrail)
|
||||
|
||||
body = b"".join(delivered)
|
||||
# the shared block builder mints a new message id: the block is not the upstream message
|
||||
assert b'"id": "msg_' in body
|
||||
assert b'"model": "claude"' in body
|
||||
assert b"my-proxy-alias" not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_reraises_guardrail_service_failures():
|
||||
"""A Bedrock outage must keep its status, not be reported to the caller as a guardrail decision.
|
||||
|
||||
A policy block is the only 400 detailing a Mapping.
|
||||
"""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.side_effect = HTTPException(
|
||||
status_code=500, detail="Bedrock guardrail throttle retries exhausted"
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _drain_streaming_hook(guardrail)
|
||||
|
||||
assert exc.value.status_code == 500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_frames_a_service_failure_once_a_keepalive_ping_flushed_the_headers():
|
||||
"""Past the ping the status line is already on the wire, so a raise reaches the client as nothing.
|
||||
|
||||
The failure has to travel as a frame instead, carrying its real status in the message.
|
||||
"""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
with (
|
||||
patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api,
|
||||
patch.object(litellm, "anthropic_sse_ping_interval_seconds", 0.0001),
|
||||
):
|
||||
mock_api.side_effect = HTTPException(status_code=503, detail="Bedrock is unavailable")
|
||||
delivered = await _drain_streaming_hook(guardrail)
|
||||
|
||||
body = b"".join(delivered).decode()
|
||||
frame = next(line for line in body.splitlines() if line.startswith("data: "))
|
||||
message = json.loads(frame[6:])["error"]["message"]
|
||||
assert message == "503: Bedrock is unavailable"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_reraises_a_service_failure_that_details_a_mapping():
|
||||
"""InvokeGuardrailChecks details a Mapping on its 500, so detail shape alone cannot mean "block"."""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.side_effect = HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": "Bedrock InvokeGuardrailChecks returned an unexpected response shape"},
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _drain_streaming_hook(guardrail)
|
||||
|
||||
assert exc.value.status_code == 500
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_block_error_frame_message_is_a_string():
|
||||
"""AnthropicErrorDetail.message is typed str, built by the proxy's own detail serializer."""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.side_effect = HTTPException(
|
||||
status_code=400, detail={"error": "Violated guardrail policy", "guardrailIdentifier": "gid"}
|
||||
)
|
||||
delivered = await _drain_streaming_hook(guardrail)
|
||||
|
||||
frame = next(line for line in b"".join(delivered).decode().splitlines() if line.startswith("data: "))
|
||||
message = json.loads(frame[6:])["error"]["message"]
|
||||
# AnthropicErrorDetail.message is typed str, and the proxy's own serializer produces the
|
||||
# readable message rather than a repr of the detail dict
|
||||
assert isinstance(message, str)
|
||||
assert message == "Violated guardrail policy"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_fails_closed_when_raw_sse_cannot_be_assembled():
|
||||
"""An unscannable stream must not be delivered: forwarding it silently disables the guardrail."""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
async def _unparseable_stream():
|
||||
yield b'data: {"type":"content_block_delta"}\n\n'
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
delivered = [
|
||||
chunk
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=_unparseable_stream(),
|
||||
request_data={"model": "claude-sonnet-4-5"},
|
||||
)
|
||||
]
|
||||
|
||||
mock_api.assert_not_called()
|
||||
body = b"".join(delivered)
|
||||
# a raise cannot reach the client once a keepalive ping has flushed the headers
|
||||
assert b"event: error" in body
|
||||
assert b"could not be assembled" in body
|
||||
assert b"content_block_delta" not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_fails_closed_when_assembler_raises_api_error():
|
||||
"""stream_chunk_builder re-raises assembly failures as litellm.APIError; it must not escape.
|
||||
|
||||
That exception message is the exact 500 this fix exists to remove.
|
||||
"""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_provider_handlers."
|
||||
"anthropic_passthrough_logging_handler.AnthropicPassthroughLoggingHandler."
|
||||
"_build_complete_streaming_response",
|
||||
side_effect=litellm.APIError(
|
||||
status_code=500,
|
||||
message="Error building chunks for logging/streaming usage calculation",
|
||||
llm_provider="",
|
||||
model="",
|
||||
),
|
||||
):
|
||||
delivered = await _drain_streaming_hook(guardrail)
|
||||
|
||||
assert b"event: error" in b"".join(delivered)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_preserves_message_id_and_model_when_re_emitting():
|
||||
"""A rewritten stream must still look like the upstream Anthropic response."""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"outputs": [{"text": "my ssn is {SSN}"}],
|
||||
}
|
||||
delivered = await _drain_streaming_hook(guardrail)
|
||||
|
||||
body = b"".join(delivered)
|
||||
assert b'"id": "msg_1"' in body
|
||||
assert b"unknown-model" not in body
|
||||
assert b'"model": "claude"' in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_blocks_raw_anthropic_sse_on_violation():
|
||||
"""A block on the extracted text must stop the stream rather than deliver it."""
|
||||
guardrail = _sse_guardrail()
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.side_effect = HTTPException(status_code=400, detail={"error": "Violated guardrail policy"})
|
||||
delivered = await _drain_streaming_hook(guardrail)
|
||||
|
||||
body = b"".join(delivered)
|
||||
# a keepalive ping may already have flushed the headers, so the block has to travel as a frame
|
||||
assert b"event: error" in body
|
||||
assert b"Violated guardrail policy" in body
|
||||
assert b"123-45-6789" not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_yields_synthetic_block_stream_for_raw_anthropic_sse():
|
||||
"""disable_exception_on_block must keep behaving as a stream, not an SSE 500 frame."""
|
||||
guardrail = _sse_guardrail(disable_exception_on_block=True)
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.side_effect = ModifyResponseException(
|
||||
message="Sorry, the model cannot answer this question.",
|
||||
model="claude",
|
||||
request_data={},
|
||||
)
|
||||
delivered = await _drain_streaming_hook(guardrail)
|
||||
|
||||
body = b"".join(delivered)
|
||||
assert b"Sorry, the model cannot answer this question." in body
|
||||
assert b"123-45-6789" not in body
|
||||
# the upstream call was already paid for, so the block frame must still report its usage
|
||||
assert b'"input_tokens": 5' in body
|
||||
assert b'"output_tokens": 9' in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_post_call_block_yields_synthetic_stream_not_raise():
|
||||
"""LIT-4186 regression: with disable_exception_on_block=True, streaming
|
||||
|
|
|
|||
|
|
@ -1215,6 +1215,44 @@ class TestToolPermissionGuardrailAnthropicMessages:
|
|||
)
|
||||
assert '"stop_reason": "tool_use"' not in body
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewrite_mode_keeps_the_stream_identity_it_had_before_the_shared_helper(self):
|
||||
"""Well-formed SSE must round-trip exactly as it did before the helpers were shared.
|
||||
|
||||
The shared module can stamp the upstream message id and model onto the assembled response
|
||||
for callers that ask for it; this path never did, and a client reads those bytes.
|
||||
"""
|
||||
with patch.object(self.rewriting, "should_run_guardrail", return_value=True):
|
||||
out = await self._drain(self.rewriting, self._sse_chunks("Read"))
|
||||
|
||||
body = b"".join(c if isinstance(c, bytes) else str(c).encode() for c in out).decode()
|
||||
message_start = next(
|
||||
json.loads(line[6:])
|
||||
for line in body.splitlines()
|
||||
if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_start"
|
||||
)["message"]
|
||||
assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id"
|
||||
assert message_start["model"] == "unknown-model", "the rewritten stream must not adopt the upstream model"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_start_without_a_dict_message_fails_closed(self):
|
||||
"""Malformed SSE must not be forwarded unscanned.
|
||||
|
||||
The shared assembler requires message_start.message to be a dict; the private helper it
|
||||
replaced accepted anything, and assembled a response from it.
|
||||
"""
|
||||
events = [
|
||||
{"type": "message_start", "message": "not-a-dict"},
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}},
|
||||
{"type": "message_stop"},
|
||||
]
|
||||
chunks = [f"event: {e['type']}\ndata: {json.dumps(e)}\n\n".encode() for e in events]
|
||||
|
||||
with patch.object(self.rewriting, "should_run_guardrail", return_value=True):
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await self._drain(self.rewriting, chunks)
|
||||
|
||||
def _resplit(self, chunks, size=7):
|
||||
joined = b"".join(chunks)
|
||||
return [joined[i : i + size] for i in range(0, len(joined), size)]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue