mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fix(guardrails): guardrail Bedrock Converse stream text per content block (#45822)
* fix(guardrails): guardrail Bedrock Converse stream text per content block The buffered converse-stream guardrail joined every text delta of the same kind across all content blocks into one string, guardrailed it, then split the result back over the frames by length. When a guardrail edit changed the length (a mask), characters moved between blocks: text spilled from one text block into the next, and parallel toolUse inputs became invalid JSON. Group fragments per (contentBlockIndex, kind) so each block is guardrailed and written back on its own. * test(guardrails): run the Bedrock stream block tests through the real post-call hook Drive de_anonymize_event_stream with a real ProxyLogging and pass-through logging object, with a fake guardrail at the boundary, and assert the texts the guardrail receives as well as the frames the client gets back * test(guardrails): pin the logging start time in the Bedrock stream tests
This commit is contained in:
dean-zavad
GitHub
parent
eadbffb69f
commit
000b3685ed
2 changed files with 141 additions and 7 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue