mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(anthropic): join multiline SSE data before guardrail processing
This commit is contained in:
parent
0b58e88c61
commit
a28fe54e83
2 changed files with 32 additions and 3 deletions
|
|
@ -55,9 +55,9 @@ def _event_payload(block: bytes) -> tuple[str | None, Mapping[str, JsonValue] |
|
|||
text: Final = block.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return None, None
|
||||
lines: Final = tuple(reversed(text.splitlines()))
|
||||
event_type: Final = next((line[6:].strip() for line in lines if line.startswith("event:")), None)
|
||||
data_line: Final = next((line[5:].strip() for line in lines if line.startswith("data:")), None)
|
||||
lines: Final = tuple(text.splitlines())
|
||||
event_type: Final = next((line[6:].strip() for line in reversed(lines) if line.startswith("event:")), None)
|
||||
data_line: Final = "\n".join(line[5:].removeprefix(" ") for line in lines if line.startswith("data:"))
|
||||
if not data_line:
|
||||
return event_type, None
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -122,6 +122,35 @@ class TestDeAnonymizeEventStream:
|
|||
proxy_logging_obj.post_call_success_hook = mock_hook
|
||||
return proxy_logging_obj
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("line_end", [b"\n", b"\r\n", b"\r"])
|
||||
async def test_multiline_data_reaches_guardrail(self, line_end: bytes):
|
||||
frame = line_end.join(
|
||||
(
|
||||
b"event: content_block_delta",
|
||||
b'data: {"type":"content_block_delta","index":3,',
|
||||
b'data: "delta":{"type":"text_delta","text":"<PERSON_1>"}}',
|
||||
b"",
|
||||
b"",
|
||||
)
|
||||
)
|
||||
stop = _message_stop_frame(sep=line_end * 2)
|
||||
|
||||
async def hook(data, user_api_key_dict, response):
|
||||
assert response["content"] == [{"type": "text", "text": "<PERSON_1>"}]
|
||||
return {**response, "content": [{"type": "text", "text": "Alice"}]}
|
||||
|
||||
result = await AnthropicPassthroughGuardrailHandler.de_anonymize_event_stream(
|
||||
body_bytes=frame + stop,
|
||||
proxy_logging_obj=self._proxy(hook),
|
||||
user_api_key_dict=MagicMock(),
|
||||
data={},
|
||||
)
|
||||
|
||||
assert _text_delta(_parse_sse_blocks(result)[0]) == (3, "Alice")
|
||||
assert b"<PERSON_1>" not in result
|
||||
assert result.endswith(stop)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_crlf_framed_stream_still_invokes_guardrail(self):
|
||||
"""P1: CRLF frames must not merge so message_stop wins and deltas skip rewriting."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue