diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index b7b393ada6b..912936a8ecd 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -206,12 +206,14 @@ class _TextContainer(Protocol): _DeltaHolder = tuple[_GroupKey, _TextContainer, str] +_StreamGroupKey: TypeAlias = tuple[int | None, _GroupKey] + class _StreamFrame(TypedDict): """One raw event-stream frame plus the guardrailable texts it carries.""" raw: ReadOnly[bytes] - texts: ReadOnly[Sequence[tuple[_GroupKey, str]]] + texts: ReadOnly[Sequence[tuple[_StreamGroupKey, str]]] def _unpack_uint32(buffer: bytes) -> int: @@ -302,12 +304,14 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): frames.append({"raw": frame_raw, "texts": []}) continue - texts: list[tuple[_GroupKey, str]] = [] + texts: list[tuple[_StreamGroupKey, str]] = [] if event_type == "contentBlockDelta": try: payload_dict: dict[str, object] = _json.loads(payload_bytes) + block_index = payload_dict.get("contentBlockIndex") + # Group per content block: a guardrail edit that changes length must not move text between blocks. texts = [ - (group_key, container[key]) + ((block_index if isinstance(block_index, int) else None, group_key), container[key]) for group_key, container, key in _collect_stream_delta_text_holders(payload_dict.get("delta")) ] except Exception as e: @@ -321,9 +325,9 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): trailing_bytes: Final = body_bytes[offset:] - group_order: Final[list[_GroupKey]] = [] - group_members: Final[dict[_GroupKey, list[tuple[int, int]]]] = {} - group_texts: Final[dict[_GroupKey, list[str]]] = {} + group_order: Final[list[_StreamGroupKey]] = [] + group_members: Final[dict[_StreamGroupKey, list[tuple[int, int]]]] = {} + group_texts: Final[dict[_StreamGroupKey, list[str]]] = {} for frame_idx, frame in enumerate(frames): for local_idx, (group_key, text) in enumerate(frame["texts"]): if group_key not in group_members: diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index da7ed635dcb..57768b41943 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -10,9 +10,21 @@ Validates that: """ import copy -import pytest +from datetime import datetime +from typing import Literal from unittest.mock import AsyncMock, MagicMock +import pytest + +import litellm +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + from litellm.llms.bedrock.passthrough.guardrail_translation.handler import ( BedrockPassthroughGuardrailHandler, _extract_converse_texts, @@ -820,6 +832,52 @@ def _build_event_stream_frame(event_type: str, payload: dict) -> bytes: return prelude + prelude_crc_b + headers_bytes + payload_bytes + msg_crc_b +_TEST_CARD = "4111 1111 1111 1111" + + +class _CardMaskingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="card-mask", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.received_texts: list[list[str]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts", []) + self.received_texts.append(list(texts)) + return {**inputs, "texts": [text.replace(_TEST_CARD, "****") for text in texts]} + + +async def _run_post_call_guardrail_on_stream( + stream_bytes: bytes, guardrail: CustomGuardrail, monkeypatch: pytest.MonkeyPatch +) -> bytes: + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + logging_obj = LiteLLMLoggingObj( + model="claude", + messages=[], + stream=True, + call_type="allm_passthrough_route", + start_time=datetime(2026, 1, 1), + litellm_call_id="test-call", + function_id="test-function", + ) + return await BedrockPassthroughGuardrailHandler.de_anonymize_event_stream( + body_bytes=stream_bytes, + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + data={ + "model": "claude", + "custom_llm_provider": "bedrock", + "endpoint": "model/claude/converse-stream", + "litellm_logging_obj": logging_obj, + }, + ) + + class TestDeAnonymizeConverseStream: def _make_proxy_logging(self, mock_hook) -> MagicMock: proxy_logging_obj = MagicMock() @@ -870,6 +928,78 @@ class TestDeAnonymizeConverseStream: ] assert "".join(texts) == "John Doe works at Acme Corp" + @staticmethod + def _deltas_by_block(result: bytes, field: str) -> dict[int, str]: + import json + from botocore.eventstream import EventStreamBuffer + + buf = EventStreamBuffer() + buf.add_data(result) + blocks: dict[int, str] = {} + for msg in buf: + if msg.headers.get(":event-type") != "contentBlockDelta": + continue + payload = json.loads(msg.payload) + delta = payload["delta"] + value = delta["toolUse"]["input"] if field == "toolUse" else delta[field] + blocks[payload["contentBlockIndex"]] = blocks.get(payload["contentBlockIndex"], "") + value + return blocks + + @pytest.mark.asyncio + async def test_text_blocks_are_guardrailed_separately(self, monkeypatch): + stream_bytes = ( + _build_event_stream_frame("messageStart", {"role": "assistant"}) + + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": f"Card: {_TEST_CARD}."}} + ) + + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}) + + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Thanks for sailing with us!"}} + ) + + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}) + + _build_event_stream_frame("messageStop", {"stopReason": "end_turn"}) + ) + guardrail = _CardMaskingGuardrail() + + result = await _run_post_call_guardrail_on_stream(stream_bytes, guardrail, monkeypatch) + + assert guardrail.received_texts == [[f"Card: {_TEST_CARD}.", "Thanks for sailing with us!"]] + assert self._deltas_by_block(result, "text") == {0: "Card: ****.", 1: "Thanks for sailing with us!"} + + @pytest.mark.asyncio + async def test_parallel_tool_use_inputs_stay_in_their_own_block(self, monkeypatch): + import json + + first_input = f'{{"ship":"AU","cardNumber":"{_TEST_CARD}"}}' + second_input = '{"ship":"AU"}' + stream_bytes = ( + _build_event_stream_frame( + "contentBlockStart", + {"contentBlockIndex": 0, "start": {"toolUse": {"toolUseId": "tooluse_a", "name": "book"}}}, + ) + + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"toolUse": {"input": first_input}}} + ) + + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 0}) + + _build_event_stream_frame( + "contentBlockStart", + {"contentBlockIndex": 1, "start": {"toolUse": {"toolUseId": "tooluse_b", "name": "ship_info"}}}, + ) + + _build_event_stream_frame( + "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"toolUse": {"input": second_input}}} + ) + + _build_event_stream_frame("contentBlockStop", {"contentBlockIndex": 1}) + + _build_event_stream_frame("messageStop", {"stopReason": "tool_use"}) + ) + guardrail = _CardMaskingGuardrail() + + result = await _run_post_call_guardrail_on_stream(stream_bytes, guardrail, monkeypatch) + + assert guardrail.received_texts == [[first_input, second_input]] + inputs = self._deltas_by_block(result, "toolUse") + assert json.loads(inputs[0]) == {"ship": "AU", "cardNumber": "****"} + assert json.loads(inputs[1]) == {"ship": "AU"} + @pytest.mark.asyncio async def test_tokens_split_across_chunks_reassembled(self): import json