mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(streaming): scan the parsed chunks when a raw SSE frame sits among them
stream_chunk_builder raised on any str or bytes chunk, so a post_call guardrail that collected a stream with one stray frame turned a finished completion into a 500, or blocked it. Raw frames among parsed chunks are now unpacked. A frame with no data (a comment, an event name, [DONE]) adds nothing and a data line that carries a chat chunk is read as one. Any other frame still raises, whether it holds a line that is not a comment or a complete field, a data line that is not a chunk with a text delta per choice, or an error. Lines are split on every break str.splitlines knows, and a bare field name raises, so an unusual line break or a frame cut mid line raises or is scanned instead of being skipped. A stream of raw frames only is unchanged is_raw_sse_stream now means every chunk is a frame, so Bedrock, tool permission and Model Armor reach the builder. Model Armor still fails closed on a raw frame beside Responses events. PANW and Cisco no longer read the stream shape off the first chunk Refs #37873
This commit is contained in:
Vineeth Sai
parent
d1cde1002e
commit
f1e773fb5e
15 changed files with 814 additions and 12 deletions
|
|
@ -16,6 +16,8 @@ from litellm.constants import VLLM_CHOICE_FIELDS
|
|||
from litellm.litellm_core_utils.core_helpers import (
|
||||
PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS,
|
||||
UPSTREAM_USAGE_COST,
|
||||
as_object_list,
|
||||
as_str_mapping,
|
||||
chunk_usage,
|
||||
finite_cost,
|
||||
provider_reported_usage_cost,
|
||||
|
|
@ -1800,6 +1802,53 @@ def _assembled_message_fields(processor: ChunkProcessor, chunks: Sequence) -> It
|
|||
yield "provider_specific_fields", combined_provider_fields
|
||||
|
||||
|
||||
_RAW_FRAME_TYPES: Final = (str, bytes)
|
||||
_FRAME_LINE_STARTS: Final = (":", "data:", "event:", "id:", "retry:")
|
||||
_FRAME_PAYLOAD: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _is_chat_chunk(payload: Mapping[str, object]) -> bool:
|
||||
choices: Final = as_object_list(payload.get("choices"))
|
||||
return (
|
||||
"error" not in payload
|
||||
and choices is not None
|
||||
and all(
|
||||
(choice := as_str_mapping(raw_choice)) is not None
|
||||
and (delta := as_str_mapping(choice.get("delta"))) is not None
|
||||
and isinstance(delta.get("content"), (str, type(None)))
|
||||
for raw_choice in choices
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _frame_chunks(frame: str | bytes) -> tuple[Mapping[str, object], ...]:
|
||||
text: Final = frame if isinstance(frame, str) else frame.decode("utf-8")
|
||||
lines: Final = tuple(filter(None, (line.strip() for line in text.splitlines())))
|
||||
if not all(line.startswith(_FRAME_LINE_STARTS) for line in lines):
|
||||
raise ValueError("a raw frame beside parsed chunks carries a line that is not an SSE comment or field")
|
||||
payloads: Final = (
|
||||
payload
|
||||
for line in lines
|
||||
if line.startswith("data:") and (payload := line.removeprefix("data:").strip()) not in ("", "[DONE]")
|
||||
)
|
||||
chunks: Final = tuple(map(_FRAME_PAYLOAD.validate_json, payloads))
|
||||
if not all(map(_is_chat_chunk, chunks)):
|
||||
raise ValueError("a raw frame beside parsed chunks carries data that is not a chat completion chunk")
|
||||
return chunks
|
||||
|
||||
|
||||
def _with_raw_frames_unpacked(chunks: Sequence[object]) -> Sequence[object]:
|
||||
frames: Final = sum(isinstance(chunk, _RAW_FRAME_TYPES) for chunk in chunks)
|
||||
if frames in (0, len(chunks)):
|
||||
return chunks
|
||||
verbose_logger.warning("stream_chunk_builder: unpacking %d raw SSE frame(s) found among parsed chunks", frames)
|
||||
return tuple(
|
||||
chain.from_iterable(
|
||||
_frame_chunks(chunk) if isinstance(chunk, _RAW_FRAME_TYPES) else (chunk,) for chunk in chunks
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def stream_chunk_builder(
|
||||
chunks: list,
|
||||
messages: Sequence | None = None,
|
||||
|
|
@ -1832,7 +1881,7 @@ def _assembled_stream_response(
|
|||
if not chunks:
|
||||
return None
|
||||
|
||||
processor: Final = ChunkProcessor(list(chunks), messages)
|
||||
processor: Final = ChunkProcessor(list(_with_raw_frames_unpacked(chunks)), messages)
|
||||
chunks = processor.chunks
|
||||
|
||||
### BASE-CASE ###
|
||||
|
|
|
|||
|
|
@ -436,7 +436,7 @@ def _is_anthropic_event_type(event_type: object) -> bool:
|
|||
|
||||
|
||||
def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool:
|
||||
return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
|
||||
return len(all_chunks) > 0 and all(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
|
||||
|
||||
|
||||
def _joined_sse_bytes(all_chunks: Sequence[object]) -> bytes:
|
||||
|
|
|
|||
|
|
@ -477,7 +477,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail):
|
|||
if not all_chunks:
|
||||
return
|
||||
|
||||
if not isinstance(all_chunks[0], (ModelResponse, ModelResponseStream)):
|
||||
if not any(isinstance(chunk, (ModelResponse, ModelResponseStream)) for chunk in all_chunks):
|
||||
verbose_proxy_logger.warning(
|
||||
"Cisco AI Defense guardrail (%s): unsupported streaming chunk shape (%s) — failing closed.",
|
||||
self.guardrail_name,
|
||||
|
|
|
|||
|
|
@ -1003,7 +1003,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
isinstance(event_type := getattr(chunk, "type", None), str) and event_type.startswith("response.")
|
||||
for chunk in all_chunks
|
||||
):
|
||||
return _StreamSurface.RESPONSES
|
||||
has_raw_frame: Final = any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
|
||||
return _StreamSurface.OPAQUE_SSE if has_raw_frame else _StreamSurface.RESPONSES
|
||||
return _StreamSurface.CHAT_COMPLETIONS
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -1446,26 +1446,27 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
|
||||
verbose_proxy_logger.info("PANW Prisma AIRS: Running post-call streaming scan")
|
||||
|
||||
all_chunks: Final = []
|
||||
content_was_modified = False
|
||||
|
||||
try:
|
||||
start_time: Final = datetime.now()
|
||||
|
||||
# Collect all chunks
|
||||
async for chunk in response:
|
||||
all_chunks.append(chunk)
|
||||
all_chunks: Final = [chunk async for chunk in response]
|
||||
byte_frames: Final = tuple(chunk for chunk in all_chunks if isinstance(chunk, bytes))
|
||||
has_parsed_chunks: Final = any(
|
||||
isinstance(chunk, (ModelResponse, ModelResponseStream)) for chunk in all_chunks
|
||||
)
|
||||
|
||||
# Handle /v1/messages streaming: chunks are raw bytes (Anthropic SSE)
|
||||
if all_chunks and isinstance(all_chunks[0], bytes):
|
||||
text = self._extract_text_from_sse_bytes(all_chunks)
|
||||
if byte_frames and len(byte_frames) == len(all_chunks):
|
||||
text = self._extract_text_from_sse_bytes(byte_frames)
|
||||
await self._scan_raw_streaming_text(text, request_data, start_time)
|
||||
for chunk in all_chunks:
|
||||
yield chunk
|
||||
return
|
||||
|
||||
# Handle /v1/responses streaming: chunks are Pydantic events (not ModelResponse/ModelResponseStream)
|
||||
if all_chunks and not isinstance(all_chunks[0], (ModelResponse, ModelResponseStream)):
|
||||
if all_chunks and not has_parsed_chunks and not byte_frames:
|
||||
text = self._extract_text_from_streaming_events(all_chunks)
|
||||
await self._scan_raw_streaming_text(text, request_data, start_time)
|
||||
for chunk in all_chunks:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import base64
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
|
@ -2880,3 +2881,165 @@ def test_stream_chunk_builder_counts_a_custom_named_deployments_missing_usage_wi
|
|||
model="my-gpt-6", messages=messages
|
||||
)
|
||||
assert built.usage.prompt_tokens == litellm.token_counter(model=counted_as, messages=messages)
|
||||
|
||||
|
||||
def _streamed_text(content: str) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-raw-frames",
|
||||
created=1,
|
||||
model="gpt-5.4-mini",
|
||||
choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=content))],
|
||||
)
|
||||
|
||||
|
||||
_RAW_FRAMES_CARRYING_NO_CHUNK: Final = (
|
||||
pytest.param(b"data: [DONE]\n\n", id="done_bytes"),
|
||||
pytest.param("data: [DONE]\n\n", id="done_str"),
|
||||
pytest.param(b": keepalive\n\n", id="comment"),
|
||||
pytest.param(b"event: ping\n\n", id="event_only"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("position", [0, 1, 2])
|
||||
@pytest.mark.parametrize("frame", _RAW_FRAMES_CARRYING_NO_CHUNK)
|
||||
def test_stream_chunk_builder_assembles_the_parsed_chunks_beside_a_raw_frame(frame: str | bytes, position: int) -> None:
|
||||
parsed: Final = [_streamed_text("Hello"), _streamed_text(" world")]
|
||||
|
||||
built: Final = stream_chunk_builder(chunks=[*parsed[:position], frame, *parsed[position:]])
|
||||
|
||||
assert built is not None
|
||||
assert built.choices[0].message.content == "Hello world"
|
||||
|
||||
|
||||
def test_stream_chunk_builder_reads_the_chunk_a_raw_frame_carries() -> None:
|
||||
carried: Final = _openai_sse(
|
||||
{"model": "gpt-5.4-mini", "choices": [{"index": 0, "delta": {"content": " there"}, "finish_reason": None}]}
|
||||
)
|
||||
|
||||
built: Final = stream_chunk_builder(chunks=[_streamed_text("Hello"), carried.encode(), _streamed_text(" world")])
|
||||
|
||||
assert built is not None
|
||||
assert built.choices[0].message.content == "Hello there world"
|
||||
|
||||
|
||||
_SECRET_CHUNK: Final[dict[str, object]] = {
|
||||
"id": "chatcmpl-raw-frames",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "gpt-5.4-mini",
|
||||
"choices": [{"index": 0, "delta": {"content": "LEAKED-SECRET"}}],
|
||||
}
|
||||
|
||||
|
||||
def _data_line(payload: Mapping[str, object]) -> str:
|
||||
return f"data: {json.dumps(payload)}\n\n"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("line_break", ["\r", "\r\n", "\x0c", "\u2028"])
|
||||
def test_stream_chunk_builder_reads_a_data_line_after_any_line_break_an_sse_reader_honours(line_break: str) -> None:
|
||||
frame: Final = f"event: message{line_break}{_data_line(_SECRET_CHUNK)}".encode()
|
||||
|
||||
built: Final = stream_chunk_builder(chunks=[_streamed_text("Hello"), frame, _streamed_text(" world")])
|
||||
|
||||
assert built is not None
|
||||
assert built.choices[0].message.content == "HelloLEAKED-SECRET world"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"frame",
|
||||
[
|
||||
pytest.param(b'data: {"error": {"message": "blocked", "code": 400}}\n\n', id="error_payload"),
|
||||
pytest.param(b'data: {"type": "content_block_delta"}\n\n', id="payload_that_is_not_a_chat_chunk"),
|
||||
pytest.param(b"data: not json\n\n", id="payload_that_is_not_json"),
|
||||
pytest.param(_data_line({**_SECRET_CHUNK, "error": {"message": "blocked"}}), id="chunk_and_error"),
|
||||
pytest.param(
|
||||
_data_line({**_SECRET_CHUNK, "choices": [{"text": "LEAKED-SECRET"}]}), id="choice_without_a_delta"
|
||||
),
|
||||
pytest.param(
|
||||
_data_line({**_SECRET_CHUNK, "choices": [{"message": {"content": "LEAKED-SECRET"}}]}),
|
||||
id="choice_with_a_message",
|
||||
),
|
||||
*(
|
||||
pytest.param(
|
||||
_data_line({**_SECRET_CHUNK, "choices": [{"index": 0, "delta": {"content": content}}]}),
|
||||
id=f"delta_content_{name}",
|
||||
)
|
||||
for name, content in (
|
||||
("list", [{"type": "text", "text": "LEAKED-SECRET"}]),
|
||||
("int", 7),
|
||||
("dict", {"text": "LEAKED-SECRET"}),
|
||||
)
|
||||
),
|
||||
pytest.param(_data_line(_SECRET_CHUNK).replace("LEAKED", "LEAK\u2028ED"), id="line_break_inside_the_json"),
|
||||
pytest.param(json.dumps(_SECRET_CHUNK), id="chunk_without_a_data_field"),
|
||||
pytest.param("a plain text frame", id="plain_text_str"),
|
||||
pytest.param(b"a plain text frame", id="plain_text_bytes"),
|
||||
pytest.param("[DONE]", id="bare_done"),
|
||||
pytest.param(b"data\n\n", id="bare_data_field"),
|
||||
pytest.param(b"\xef\xbb\xbfdata: [DONE]\n\n", id="byte_order_mark"),
|
||||
pytest.param(bytearray(b"data: [DONE]\n\n"), id="bytearray"),
|
||||
],
|
||||
)
|
||||
def test_stream_chunk_builder_still_fails_on_a_raw_frame_it_cannot_account_for(frame: bytes | str) -> None:
|
||||
with pytest.raises(litellm.APIError):
|
||||
stream_chunk_builder(chunks=[_streamed_text("Hello"), frame])
|
||||
|
||||
|
||||
def _built_content(chunks: list[object]) -> str | None:
|
||||
try:
|
||||
built: Final = stream_chunk_builder(chunks=chunks)
|
||||
except litellm.APIError:
|
||||
return None
|
||||
return None if built is None else built.choices[0].message.content
|
||||
|
||||
|
||||
@pytest.mark.parametrize("line_break", ["\n", "\r\n", "\r"])
|
||||
def test_stream_chunk_builder_never_skips_a_data_line_when_its_frame_is_split_in_two(line_break: str) -> None:
|
||||
frame: Final = f"event: message{line_break}data: {json.dumps(_SECRET_CHUNK)}{line_break}{line_break}".encode()
|
||||
|
||||
skipped: Final = [
|
||||
at
|
||||
for at in range(1, len(frame))
|
||||
if (
|
||||
content := _built_content([_streamed_text("Hello"), frame[:at], frame[at:], _streamed_text("!")])
|
||||
)
|
||||
is not None
|
||||
and "LEAKED-SECRET" not in content
|
||||
]
|
||||
|
||||
assert skipped == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"frames",
|
||||
[
|
||||
pytest.param([b"data: [DONE]\n\n"], id="done_only"),
|
||||
pytest.param([_openai_sse({"choices": [{"index": 0, "delta": {"content": "Hi"}}]}).encode()], id="chunk_frame"),
|
||||
],
|
||||
)
|
||||
def test_stream_chunk_builder_still_fails_on_a_stream_of_raw_frames_only(frames: list[bytes]) -> None:
|
||||
with pytest.raises(litellm.APIError):
|
||||
stream_chunk_builder(chunks=frames)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("frames", [1, 3])
|
||||
def test_stream_chunk_builder_warns_with_the_number_of_raw_frames_it_unpacks(
|
||||
frames: int, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
stream: Final = [_streamed_text("Hello"), *[b"data: [DONE]\n\n"] * frames, _streamed_text(" world")]
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
built: Final = stream_chunk_builder(chunks=stream)
|
||||
|
||||
assert built is not None
|
||||
assert [record.getMessage() for record in caplog.records if "raw SSE frame" in record.getMessage()] == [
|
||||
f"stream_chunk_builder: unpacking {frames} raw SSE frame(s) found among parsed chunks"
|
||||
]
|
||||
|
||||
|
||||
def test_stream_chunk_builder_does_not_warn_when_no_raw_frame_is_present(caplog: pytest.LogCaptureFixture) -> None:
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
built: Final = stream_chunk_builder(chunks=[_streamed_text("Hello"), _streamed_text(" world")])
|
||||
|
||||
assert built is not None
|
||||
assert not [record for record in caplog.records if "raw SSE frame" in record.getMessage()]
|
||||
|
|
|
|||
|
|
@ -8168,3 +8168,76 @@ async def test_guardrail_secret_lookups_leave_the_event_loop_free(
|
|||
assert vault.read.is_set()
|
||||
assert not vault.read_on_loop.is_set()
|
||||
assert upstream.calls.last.request.url.host == "bedrock-runtime.us-west-2.amazonaws.com"
|
||||
|
||||
|
||||
def _streamed_text(content: str) -> litellm.ModelResponseStream:
|
||||
return litellm.ModelResponseStream(
|
||||
id="tid",
|
||||
choices=[
|
||||
litellm.types.utils.StreamingChoices(
|
||||
delta=litellm.types.utils.Delta(content=content, role="assistant"), index=0
|
||||
)
|
||||
],
|
||||
created=1,
|
||||
model="gpt-4o-mini",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
|
||||
def _post_call_guardrail() -> BedrockGuardrail:
|
||||
return BedrockGuardrail(
|
||||
guardrail_name="bedrock-raw-frame",
|
||||
guardrailIdentifier="test-id",
|
||||
guardrailVersion="DRAFT",
|
||||
event_hook=GuardrailEventHooks.post_call,
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
|
||||
async def _streamed(items: list[object]):
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("position", [0, 1, 2])
|
||||
async def test_streaming_post_call_scans_the_parsed_chunks_beside_a_raw_frame(position: int):
|
||||
guardrail: Final = _post_call_guardrail()
|
||||
parsed: Final = [_streamed_text("Hello"), _streamed_text(" world")]
|
||||
stream: Final = [*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]]
|
||||
|
||||
clean_scan: Final = {"action": "NONE", "assessments": [], "outputs": []}
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=clean_scan)) as mock_make:
|
||||
delivered = [
|
||||
chunk
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(), response=_streamed(stream), request_data={"messages": []}
|
||||
)
|
||||
]
|
||||
|
||||
scanned: Final = [
|
||||
call.kwargs["response"].choices[0].message.content
|
||||
for call in mock_make.call_args_list
|
||||
if call.kwargs.get("source") == "OUTPUT"
|
||||
]
|
||||
assert scanned == ["Hello world"]
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in delivered) == "Hello world"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_post_call_still_blocks_a_violating_response_that_carries_a_raw_frame():
|
||||
guardrail: Final = _post_call_guardrail()
|
||||
violation: Final = HTTPException(status_code=400, detail={"error": "Violated guardrail policy"})
|
||||
stream: Final = [_streamed_text("Hello"), b"data: [DONE]\n\n", _streamed_text(" world")]
|
||||
|
||||
with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=violation)):
|
||||
with pytest.raises(HTTPException) as blocked:
|
||||
_ = [
|
||||
chunk
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(), response=_streamed(stream), request_data={"messages": []}
|
||||
)
|
||||
]
|
||||
|
||||
assert blocked.value is violation
|
||||
|
|
|
|||
|
|
@ -3067,3 +3067,63 @@ def test_config_model_exposed():
|
|||
CiscoAIDefenseGuardrail.get_config_model() is CiscoAIDefenseGuardrailConfigModel
|
||||
)
|
||||
assert CiscoAIDefenseGuardrailConfigModel.ui_friendly_name() == "Cisco AI Defense"
|
||||
|
||||
|
||||
class TestCiscoAIDefenseRawFramesInAStream:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("position", [0, 1, 2])
|
||||
async def test_the_parsed_chunks_are_inspected_and_delivered_beside_a_raw_frame(self, position):
|
||||
g = _make_guardrail(event_hook=["pre_call", "post_call"])
|
||||
parsed = _make_streaming_chunks(["Hello", " world"])
|
||||
stream = [*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]]
|
||||
|
||||
received, post_mock = await _streaming_setup(g, stream, cisco_response=_safe_response())
|
||||
|
||||
inspected = [message["content"] for message in post_mock.call_args.kwargs["json"]["messages"]]
|
||||
assert inspected[-1] == "Hello world"
|
||||
assert received == stream
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_violation_withholds_the_chunks_when_a_raw_frame_sits_beside_them(self):
|
||||
g = _make_guardrail(event_hook=["pre_call", "post_call"])
|
||||
parsed = _make_streaming_chunks(["Here is your SSN: ", "123-45-6789."])
|
||||
|
||||
received, post_mock = await _streaming_setup(
|
||||
g, [parsed[0], b"data: [DONE]\n\n", parsed[1]], cisco_response=_violation_response()
|
||||
)
|
||||
|
||||
assert post_mock.called
|
||||
assert all(isinstance(chunk, str) and chunk.startswith("data: ") for chunk in received)
|
||||
assert any('"error"' in chunk and "Cisco AI Defense" in chunk for chunk in received)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_data_line_after_a_carriage_return_is_inspected_and_delivered(self):
|
||||
g = _make_guardrail(event_hook=["pre_call", "post_call"])
|
||||
parsed = _make_streaming_chunks(["Hello", " world"])
|
||||
carried = '{"created": 1, "model": "m", "choices": [{"index": 0, "delta": {"content": " LEAKED-SECRET"}}]}'
|
||||
stream = [parsed[0], f'event: message\rdata: {carried}\n\n'.encode(), parsed[1]]
|
||||
|
||||
received, post_mock = await _streaming_setup(g, stream, cisco_response=_safe_response())
|
||||
|
||||
inspected = [message["content"] for message in post_mock.call_args.kwargs["json"]["messages"]]
|
||||
delivered = "".join(chunk.choices[0].delta.content or "" for chunk in received)
|
||||
assert delivered == "Hello LEAKED-SECRET world"
|
||||
assert delivered in inspected
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("frame", [b"LEAKED-SECRET", b'data: {"choices": [{"text": "LEAKED-SECRET"}]}\n\n'])
|
||||
async def test_a_frame_that_cannot_be_read_is_not_inspected_or_delivered(self, frame):
|
||||
g = _make_guardrail(event_hook=["pre_call", "post_call"])
|
||||
post_mock = AsyncMock(return_value=_safe_response())
|
||||
delivered: list = []
|
||||
|
||||
with _patch_inspection_post(g, post_mock), pytest.raises(litellm.APIError):
|
||||
async for chunk in g.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=_aiter([*_make_streaming_chunks(["Hello"]), frame]),
|
||||
request_data={"messages": [{"role": "user", "content": "hi"}]},
|
||||
):
|
||||
delivered.append(chunk)
|
||||
|
||||
assert delivered == []
|
||||
post_mock.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -2803,3 +2803,54 @@ class TestRegistration:
|
|||
guardrail_class_registry["microsoft_purview"]
|
||||
is MicrosoftPurviewDLPGuardrail
|
||||
)
|
||||
|
||||
|
||||
class TestRawFramesInAStreamedResponse:
|
||||
@staticmethod
|
||||
def _streamed_text(content: str):
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-raw-frame",
|
||||
created=1,
|
||||
model="gpt-4o",
|
||||
choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=content))],
|
||||
)
|
||||
|
||||
async def _scan(self, items: list[object], check: AsyncMock) -> list[object]:
|
||||
guardrail: Final = _make_guardrail()
|
||||
|
||||
async def upstream():
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
with patch.object(guardrail, "_check_content", check):
|
||||
return [
|
||||
chunk
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
response=upstream(),
|
||||
request_data={"metadata": {"user_id": "user-123"}},
|
||||
)
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("position", [0, 1, 2])
|
||||
async def test_the_parsed_chunks_are_scanned_and_delivered_beside_a_raw_frame(self, position: int):
|
||||
parsed: Final = [self._streamed_text("Hello"), self._streamed_text(" world")]
|
||||
check: Final = AsyncMock(return_value={"policyActions": []})
|
||||
|
||||
delivered: Final = await self._scan([*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]], check)
|
||||
|
||||
assert [call.kwargs["text"] for call in check.call_args_list] == ["Hello world"]
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in delivered) == "Hello world"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_violation_is_raised_when_a_raw_frame_sits_beside_the_chunks(self):
|
||||
parsed: Final = [self._streamed_text("SSN: 123-45"), self._streamed_text("-6789")]
|
||||
blocked: Final = HTTPException(status_code=400, detail={"error": "Content blocked by policy"})
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await self._scan([parsed[0], b"data: [DONE]\n\n", parsed[1]], AsyncMock(side_effect=blocked))
|
||||
|
||||
assert raised.value is blocked
|
||||
|
|
|
|||
|
|
@ -5318,3 +5318,98 @@ async def test_apply_guardrail_raises_on_flagged_when_not_logging_only():
|
|||
entries = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"]
|
||||
assert len(flagged) == 1
|
||||
|
||||
|
||||
def _streamed_text(content: str) -> litellm.ModelResponseStream:
|
||||
return litellm.ModelResponseStream(
|
||||
choices=[litellm.types.utils.StreamingChoices(delta=litellm.types.utils.Delta(content=content))]
|
||||
)
|
||||
|
||||
|
||||
def _clean_scan_response() -> dict[str, object]:
|
||||
return {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("position", [0, 1, 2])
|
||||
async def test_streaming_hook_scans_the_parsed_chunks_beside_a_raw_frame(position: int):
|
||||
guardrail = ModelArmorGuardrail(
|
||||
template_id="test-template",
|
||||
project_id="test-project",
|
||||
location="us-central1",
|
||||
guardrail_name="model-armor-test",
|
||||
)
|
||||
guardrail.make_model_armor_request = AsyncMock(return_value=_clean_scan_response())
|
||||
guardrail.should_run_guardrail = Mock(return_value=True)
|
||||
parsed = [_streamed_text("Hello"), _streamed_text(" world")]
|
||||
stream = [*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]]
|
||||
|
||||
delivered = await _drain_surface_hook(guardrail, stream)
|
||||
|
||||
scanned = [call.kwargs["content"] for call in guardrail.make_model_armor_request.call_args_list]
|
||||
assert scanned == ["Hello world"]
|
||||
assert delivered == stream
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_blocks_a_flagged_response_that_carries_a_raw_frame():
|
||||
guardrail = ModelArmorGuardrail(
|
||||
template_id="test-template",
|
||||
project_id="test-project",
|
||||
location="us-central1",
|
||||
guardrail_name="model-armor-test",
|
||||
)
|
||||
guardrail.make_model_armor_request = AsyncMock(
|
||||
return_value={
|
||||
"sanitizationResult": {
|
||||
"filterMatchState": "MATCH_FOUND",
|
||||
"filterResults": {
|
||||
"sdp": {
|
||||
"sdpFilterResult": {
|
||||
"inspectResult": {
|
||||
"matchState": "MATCH_FOUND",
|
||||
"findings": [{"infoType": "PASSWORD", "likelihood": "VERY_LIKELY"}],
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
guardrail.should_run_guardrail = Mock(return_value=True)
|
||||
parsed = [_streamed_text("My password is "), _streamed_text("hunter2")]
|
||||
|
||||
delivered = await _drain_surface_hook(guardrail, [parsed[0], b"data: [DONE]\n\n", parsed[1]])
|
||||
|
||||
assert guardrail.make_model_armor_request.call_count == 1
|
||||
assert len(delivered) == 1
|
||||
assert isinstance(delivered[0], str)
|
||||
assert int(json.loads(delivered[0].removeprefix("data: "))["error"]["code"]) == 400
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"frame",
|
||||
[
|
||||
pytest.param(b"data: [DONE]\n\n", id="done_bytes"),
|
||||
pytest.param(b'data: {"type":"response.output_text.delta","delta":"LEAKED-SECRET"}\n\n', id="delta_bytes"),
|
||||
pytest.param('data: {"type":"response.output_text.delta","delta":"LEAKED-SECRET"}\n\n', id="delta_str"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("frame_first", [False, True], ids=["frame_last", "frame_first"])
|
||||
async def test_streaming_hook_fails_closed_on_a_raw_frame_beside_responses_events(
|
||||
frame: bytes | str, frame_first: bool
|
||||
):
|
||||
guardrail = _surface_guardrail()
|
||||
post = _armor_post_mock(_MODEL_ARMOR_CLEAN)
|
||||
events = _responses_api_events()
|
||||
stream = (frame, *events) if frame_first else (*events, frame)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", post):
|
||||
delivered = await _drain_surface_hook(guardrail, stream)
|
||||
|
||||
post.assert_not_called()
|
||||
assert len(delivered) == 1
|
||||
assert isinstance(delivered[0], str)
|
||||
assert "could not be assembled for scanning" in delivered[0]
|
||||
assert "LEAKED-SECRET" not in delivered[0]
|
||||
|
|
|
|||
|
|
@ -2494,3 +2494,70 @@ class TestNomaAnonymizationFlow:
|
|||
user_api_key_dict=mock_user_api_key_dict,
|
||||
response=llm_response,
|
||||
)
|
||||
|
||||
|
||||
class TestNomaRawFramesInAStream:
|
||||
@staticmethod
|
||||
def _guardrail():
|
||||
return NomaGuardrail(
|
||||
api_key="test-api-key",
|
||||
api_base="https://api.test.noma.security/",
|
||||
monitor_mode=False,
|
||||
block_failures=True,
|
||||
guardrail_name="noma-post-call",
|
||||
event_hook="post_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _streamed_text(content):
|
||||
return ModelResponseStream(
|
||||
id="resp_1",
|
||||
choices=[StreamingChoices(delta=Delta(content=content), index=0)],
|
||||
created=1234567890,
|
||||
model="gpt-4.1-mini",
|
||||
object="chat.completion.chunk",
|
||||
)
|
||||
|
||||
async def _scan(self, items, verdict_for_text):
|
||||
guardrail = self._guardrail()
|
||||
|
||||
async def upstream():
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", side_effect=_noma_reply_per_text(verdict_for_text)):
|
||||
return [
|
||||
chunk
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
|
||||
response=upstream(),
|
||||
request_data={"litellm_call_id": "call-1", "metadata": {}},
|
||||
)
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("position", [0, 1, 2])
|
||||
async def test_the_parsed_chunks_are_scanned_and_delivered_beside_a_raw_frame(self, position):
|
||||
scanned = []
|
||||
|
||||
def verdict_for_text(text):
|
||||
scanned.append(text)
|
||||
return {"aggregatedScanResult": False, "scanResult": []}
|
||||
|
||||
parsed = [self._streamed_text("Hello"), self._streamed_text(" world")]
|
||||
|
||||
delivered = await self._scan([*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]], verdict_for_text)
|
||||
|
||||
assert scanned == ["Hello world"]
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in delivered) == "Hello world"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_flagged_response_is_blocked_when_a_raw_frame_sits_beside_it(self):
|
||||
parsed = [self._streamed_text("The launch code"), self._streamed_text(" is 0000.")]
|
||||
|
||||
with pytest.raises(NomaBlockedMessage):
|
||||
await self._scan(
|
||||
[parsed[0], b"data: [DONE]\n\n", parsed[1]],
|
||||
lambda text: {"aggregatedScanResult": "launch code" in text, "scanResult": []},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ import logging
|
|||
from collections.abc import Mapping, Sequence
|
||||
from contextlib import AbstractContextManager
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -6598,5 +6599,101 @@ class TestPanwAirsServerToolCalls:
|
|||
assert message.tool_calls is None
|
||||
|
||||
|
||||
class TestRawFramesInAStreamedResponse:
|
||||
@staticmethod
|
||||
def _streamed_text(content: str) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
id="test_id",
|
||||
created=1234567890,
|
||||
model="gpt-4",
|
||||
choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=content))],
|
||||
)
|
||||
|
||||
async def _scan(self, stream: list[object], scan_result: dict[str, object]) -> tuple[list[object], AsyncMock]:
|
||||
handler = make_handler()
|
||||
|
||||
async def upstream():
|
||||
for item in stream:
|
||||
yield item
|
||||
|
||||
with patch.object(handler, "_call_panw_api", new_callable=AsyncMock, return_value=scan_result) as scan:
|
||||
delivered = [
|
||||
chunk
|
||||
async for chunk in handler.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
|
||||
response=upstream(),
|
||||
request_data={"model": "gpt-4", "messages": [], "litellm_call_id": "call-1"},
|
||||
)
|
||||
]
|
||||
return delivered, scan
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("position", [0, 1, 2])
|
||||
async def test_the_parsed_chunks_are_scanned_and_delivered_beside_a_raw_frame(self, position: int):
|
||||
parsed = [self._streamed_text("Hello"), self._streamed_text(" world")]
|
||||
|
||||
delivered, scan = await self._scan(
|
||||
[*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]], {"action": "allow", "category": "safe"}
|
||||
)
|
||||
|
||||
assert [call.kwargs["content"] for call in scan.call_args_list] == ["Hello world"]
|
||||
assert delivered == [*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_blocked_response_is_not_delivered_when_a_raw_frame_sits_beside_it(self):
|
||||
parsed = [self._streamed_text("Hello"), self._streamed_text(" world")]
|
||||
|
||||
delivered, scan = await self._scan(
|
||||
[parsed[0], b"data: [DONE]\n\n", parsed[1]], {"action": "block", "category": "malicious"}
|
||||
)
|
||||
|
||||
assert scan.call_count == 1
|
||||
assert all(not isinstance(chunk, (ModelResponseStream, bytes)) for chunk in delivered)
|
||||
assert any("error" in str(chunk) for chunk in delivered)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_stream_of_raw_frames_stays_blocked_when_a_str_frame_sits_beside_bytes(self):
|
||||
frames = [
|
||||
b'event: content_block_delta\ndata: {"type": "content_block_delta"}\n\n',
|
||||
"data: [DONE]\n\n",
|
||||
]
|
||||
|
||||
delivered, scan = await self._scan(frames, {"action": "allow", "category": "safe"})
|
||||
|
||||
assert scan.call_count == 0
|
||||
assert all(isinstance(chunk, str) and '"error"' in chunk for chunk in delivered)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_data_line_split_across_two_frames_is_blocked_not_delivered_unscanned(self):
|
||||
carried = json.dumps(
|
||||
{
|
||||
"id": "c1",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "gpt-4",
|
||||
"choices": [{"index": 0, "delta": {"content": "LEAKED-SECRET"}}],
|
||||
}
|
||||
)
|
||||
stream = [self._streamed_text("Hello"), b"data", f": {carried}\n\n".encode(), self._streamed_text(" world")]
|
||||
|
||||
delivered, scan = await self._scan(stream, {"action": "allow", "category": "safe"})
|
||||
|
||||
assert scan.call_count == 0
|
||||
assert all(isinstance(chunk, str) and '"error"' in chunk for chunk in delivered)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("frame_first", [False, True])
|
||||
async def test_a_bytes_frame_beside_responses_events_is_blocked_not_scanned_in_part(self, frame_first: bool):
|
||||
events = [SimpleNamespace(type="response.output_text.delta", delta="test content")]
|
||||
frame = b"data: [DONE]\n\n"
|
||||
|
||||
delivered, scan = await self._scan(
|
||||
[frame, *events] if frame_first else [*events, frame], {"action": "allow", "category": "safe"}
|
||||
)
|
||||
|
||||
assert scan.call_count == 0
|
||||
assert all(isinstance(chunk, str) and '"error"' in chunk for chunk in delivered)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
|
|
|
|||
|
|
@ -1087,6 +1087,15 @@ class TestRepelloAIStreaming:
|
|||
|
||||
return _gen()
|
||||
|
||||
@staticmethod
|
||||
def _stream_items(*contents):
|
||||
from litellm.types.utils import Delta, StreamingChoices
|
||||
|
||||
return tuple(
|
||||
ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=content))])
|
||||
for content in contents
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_passed_reemits_chunks(self, monkeypatch):
|
||||
guardrail = _guardrail(event_hook="post_call")
|
||||
|
|
@ -1106,6 +1115,58 @@ class TestRepelloAIStreaming:
|
|||
]
|
||||
assert len(out) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("position", [0, 1, 2])
|
||||
async def test_streaming_scans_the_parsed_chunks_beside_a_raw_frame(self, monkeypatch, position):
|
||||
guardrail = _guardrail(event_hook="post_call")
|
||||
captured = {}
|
||||
|
||||
async def capture(url, headers, json, **_kwargs):
|
||||
captured["json"] = json
|
||||
return _verdict_response("passed", url)
|
||||
|
||||
monkeypatch.setattr(guardrail.async_handler, "post", capture)
|
||||
stream = list(self._stream_items("unsafe ", "answer"))
|
||||
items = [*stream[:position], b"data: [DONE]\n\n", *stream[position:]]
|
||||
|
||||
async def upstream():
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
out = [
|
||||
chunk
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=upstream(),
|
||||
request_data={"messages": [{"role": "user", "content": "q"}]},
|
||||
)
|
||||
]
|
||||
|
||||
assert captured["json"]["scan_data"]["response"] == "unsafe answer"
|
||||
assert out == items
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_blocks_a_flagged_response_that_carries_a_raw_frame(self, monkeypatch):
|
||||
from litellm.proxy.proxy_server import StreamingCallbackError
|
||||
|
||||
guardrail = _guardrail(event_hook="post_call")
|
||||
monkeypatch.setattr(
|
||||
guardrail.async_handler, "post", _async_return(_verdict_response("blocked", ANALYZE_RESPONSE_URL))
|
||||
)
|
||||
stream = list(self._stream_items("unsafe ", "answer"))
|
||||
|
||||
async def upstream():
|
||||
for item in (stream[0], b"data: [DONE]\n\n", stream[1]):
|
||||
yield item
|
||||
|
||||
with pytest.raises(StreamingCallbackError):
|
||||
async for _ in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=upstream(),
|
||||
request_data={"messages": [{"role": "user", "content": "q"}]},
|
||||
):
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_blocked_raises(self, monkeypatch):
|
||||
from litellm.proxy.proxy_server import StreamingCallbackError
|
||||
|
|
|
|||
|
|
@ -4162,3 +4162,70 @@ class TestToolPermissionGuardrailAnthropicMessages:
|
|||
with patch.object(self.blocking, "should_run_guardrail", return_value=True):
|
||||
with pytest.raises(GuardrailRaisedException):
|
||||
await self._drain(self.blocking, [b"data: not-json\n\n", b"event: weird\n\n"])
|
||||
|
||||
|
||||
class TestToolPermissionGuardrailRawFramesInAChatStream:
|
||||
def setup_method(self):
|
||||
self.blocking = ToolPermissionGuardrail(
|
||||
guardrail_name="chat-block",
|
||||
rules=[
|
||||
{"id": "allow_bash", "tool_name": r"^Bash$", "decision": "allow"},
|
||||
{"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"},
|
||||
],
|
||||
default_action="deny",
|
||||
on_disallowed_action="block",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _tool_call_stream(tool_name: str) -> list[ModelResponseStream]:
|
||||
def chunk(delta: Delta, finish_reason: str | None = None) -> ModelResponseStream:
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-raw-frame",
|
||||
created=1,
|
||||
model="gpt-5.4-mini",
|
||||
choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)],
|
||||
)
|
||||
|
||||
call: Final = {
|
||||
"index": 0,
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": tool_name, "arguments": '{"file_path": "/etc/passwd"}'},
|
||||
}
|
||||
return [chunk(Delta(role="assistant", tool_calls=[call])), chunk(Delta(), finish_reason="tool_calls")]
|
||||
|
||||
async def _drain(self, items: list[object]) -> list[object]:
|
||||
async def _stream():
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
with patch.object(self.blocking, "should_run_guardrail", return_value=True):
|
||||
return [
|
||||
chunk
|
||||
async for chunk in self.blocking.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(), response=_stream(), request_data={}
|
||||
)
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("position", [0, 1, 2])
|
||||
async def test_denied_tool_call_is_blocked_whichever_way_a_raw_frame_sits_in_the_stream(self, position: int):
|
||||
parsed: Final = self._tool_call_stream("Read")
|
||||
|
||||
with pytest.raises(GuardrailRaisedException, match="denied by rule 'deny_read'"):
|
||||
await self._drain([*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_allowed_tool_call_is_delivered_when_a_raw_frame_sits_in_the_stream(self):
|
||||
parsed: Final = self._tool_call_stream("Bash")
|
||||
|
||||
delivered: Final = await self._drain([parsed[0], b"data: [DONE]\n\n", parsed[1]])
|
||||
|
||||
names: Final = [
|
||||
call.function.name
|
||||
for chunk in delivered
|
||||
if isinstance(chunk, ModelResponseStream)
|
||||
for call in (chunk.choices[0].delta.tool_calls or ())
|
||||
if call.function.name
|
||||
]
|
||||
assert names == ["Bash"]
|
||||
|
|
|
|||
|
|
@ -8,8 +8,9 @@ import pytest
|
|||
from litellm.proxy.guardrails.anthropic_sse import (
|
||||
anthropic_sse_chunks_from_response,
|
||||
assemble_anthropic_sse_stream,
|
||||
is_raw_sse_stream,
|
||||
)
|
||||
from litellm.types.utils import Choices, ModelResponse
|
||||
from litellm.types.utils import Choices, Delta, ModelResponse, ModelResponseStream, StreamingChoices
|
||||
|
||||
|
||||
def _frame(event: str, data: dict[str, object]) -> bytes:
|
||||
|
|
@ -429,3 +430,19 @@ def test_healthy_stream_with_a_tool_input_that_is_not_json_still_raises():
|
|||
|
||||
with pytest.raises(ValueError, match="Failed to parse tool call arguments"):
|
||||
anthropic_sse_chunks_from_response(assembled, upstream)
|
||||
|
||||
|
||||
_PARSED_CHUNK: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Hi"))])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("chunks", "is_raw"),
|
||||
[
|
||||
pytest.param([b"event: message_start\ndata: {}\n\n", "data: [DONE]\n\n"], True, id="frames_only"),
|
||||
pytest.param([_PARSED_CHUNK, _PARSED_CHUNK], False, id="parsed_chunks_only"),
|
||||
pytest.param([_PARSED_CHUNK, b"data: [DONE]\n\n", _PARSED_CHUNK], False, id="parsed_chunks_and_a_frame"),
|
||||
pytest.param([], False, id="empty"),
|
||||
],
|
||||
)
|
||||
def test_a_stream_is_raw_sse_only_when_every_chunk_is_a_frame(chunks: list[object], is_raw: bool):
|
||||
assert is_raw_sse_stream(chunks) is is_raw
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue