mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
Preserve main's per-block streaming fixes and regression coverage
This commit is contained in:
moyai-devin-berriai[bot]
GitHub
parent
1431cbb038
commit
ab903cba6a
2 changed files with 140 additions and 7 deletions
|
|
@ -37,7 +37,7 @@ if TYPE_CHECKING:
|
|||
from litellm.llms.pass_through.guardrail_translation.handler import (
|
||||
PassThroughEndpointHandler,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
|
||||
|
|
@ -584,12 +584,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:
|
||||
|
|
@ -680,12 +682,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:
|
||||
|
|
@ -699,9 +703,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:
|
||||
|
|
|
|||
|
|
@ -9,6 +9,17 @@ Validates that:
|
|||
- Non-converse endpoints skip guardrail execution
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
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 typing import Literal
|
||||
import litellm
|
||||
|
||||
import copy
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -1436,6 +1447,78 @@ class TestDeAnonymizeConverseStream:
|
|||
hook_spy.assert_not_called()
|
||||
assert result is stream_bytes
|
||||
|
||||
@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_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!"}
|
||||
|
||||
@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
|
||||
|
||||
|
||||
class TestConverseRewriteRegressions:
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1697,3 +1780,49 @@ class TestConverseRewriteRegressions:
|
|||
result = await BedrockPassthroughGuardrailHandler().process_input_messages(data, guard)
|
||||
assert result == data
|
||||
guard.apply_guardrail.assert_not_called()
|
||||
|
||||
|
||||
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,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
_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]}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue