scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-09 08:18:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 67

fix(streaming): rebuild each choice of an n>1 chat stream on its own

Browse source 
stream_chunk_builder folded every choice of an n>1 chat stream into choice 0, so post-call guardrails such as tool_permission re-emitted one choice holding every choice's tool calls, and spend logs and logging callbacks kept only the first choice's text. Each choice is now assembled from its own chunks, while usage, cost, and hidden params still read the whole stream. Cisco AI Defense compares each choice's streamed text on its own so safe n>1 streams still pass through unchanged.
This commit is contained in:
mateo-berri 2026-10-28 17:41:19 -07:00
parent 66ddbaab86
commit 2ebb92c6e1
6 changed files with 432 additions and 205 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

387
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -384,6 +384,28 @@ class ChunkProcessor:
case _:
return "assistant"
@staticmethod
def _finish_reason_of_chunks(chunks: Sequence["_BaseChunk"]) -> object:
finish_reason = "stop"
for chunk in chunks:
if "choices" in chunk and len(chunk["choices"]) > 0:
chunk_finish_reason = None
if hasattr(chunk["choices"][0], "finish_reason"):
chunk_finish_reason = chunk["choices"][0].finish_reason
elif "finish_reason" in chunk["choices"][0]:
chunk_finish_reason = _CHOICE_FIELD_VALUE.validate_python(chunk["choices"][0]["finish_reason"])
if chunk_finish_reason is not None:
finish_reason = chunk_finish_reason
return finish_reason
@staticmethod
def build_base_choice(index: int, chunks: Sequence["_BaseChunk"]) -> Choices:
return Choices(
index=index,
finish_reason=ChunkProcessor._finish_reason_of_chunks(chunks),
message={"role": ChunkProcessor._get_role_from_chunks(chunks), "content": ""},
)
@staticmethod
def _get_model_from_chunks(chunks: Sequence["_BaseChunk"], first_chunk_model: str) -> str:
"""
@ -420,16 +442,7 @@ class ChunkProcessor:
service_tier: Final = ChunkProcessor._get_service_tier_from_chunks(chunks)
role: Final = ChunkProcessor._get_role_from_chunks(chunks)
finish_reason = "stop"
for chunk in chunks:
if "choices" in chunk and len(chunk["choices"]) > 0:
chunk_finish_reason = None
if hasattr(chunk["choices"][0], "finish_reason"):
chunk_finish_reason = chunk["choices"][0].finish_reason
elif "finish_reason" in chunk["choices"][0]:
chunk_finish_reason = _CHOICE_FIELD_VALUE.validate_python(chunk["choices"][0]["finish_reason"])
if chunk_finish_reason is not None:
finish_reason = chunk_finish_reason
finish_reason: Final = ChunkProcessor._finish_reason_of_chunks(chunks)
# Initialize the response dictionary
response = ModelResponse(
@ -1514,11 +1527,180 @@ def _bedrock_trace_fields(chunk: object) -> Mapping[str, object]:
return {field: fields[field] for field in BEDROCK_TRACE_FIELDS if field in fields}
def _stream_choice_index(choice: object) -> int:
index: Final = _field_of(choice, "index")
return index if isinstance(index, int) else 0
def _choice_streams(chunks: Sequence[_BaseChunk]) -> tuple[tuple[int, Sequence[_BaseChunk]], ...]:
indices: Final = sorted(
frozenset(map(_stream_choice_index, chain.from_iterable(chunk.get("choices") or () for chunk in chunks)))
)
if len(indices) <= 1:
return ((0, chunks),)
return tuple((index, _narrowed_to_choice(chunks, index)) for index in indices)
def _narrowed_to_choice(chunks: Sequence[_BaseChunk], index: int) -> tuple[_BaseChunk, ...]:
return tuple(
_BaseChunk(choices=[choice for choice in chunk.get("choices") or () if _stream_choice_index(choice) == index])
for chunk in chunks
)
def _joined_simple_text(chunks: Sequence[_BaseChunk]) -> str:
return "".join(_simple_text_part(chunk["choices"]) or "" for chunk in chunks if chunk.get("choices"))
def _assembled_message_fields(processor: ChunkProcessor, chunks: Sequence) -> Iterator[tuple[str, object]]:
tool_call_chunks: Final = [
chunk
for chunk in chunks
if any(
"tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] is not None
for choice in chunk.get("choices") or ()
)
]
if tool_call_chunks:
tool_calls_list: Final = processor.get_combined_tool_content(tool_call_chunks)
yield "content", None
yield "tool_calls", tool_calls_list
function_call_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "function_call" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["function_call"] is not None
]
if function_call_chunks:
yield "content", None
yield "function_call", processor.get_combined_function_call_content(function_call_chunks)
content_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "content" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["content"] is not None
]
if content_chunks:
yield "content", processor.get_combined_content(content_chunks)
thinking_blocks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "thinking_blocks" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["thinking_blocks"] is not None
]
if thinking_blocks:
yield "thinking_blocks", processor.get_combined_thinking_content(thinking_blocks)
reasoning_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "reasoning_content" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["reasoning_content"] is not None
]
if reasoning_chunks:
yield "reasoning_content", processor.get_combined_reasoning_content(reasoning_chunks)
typed_chunks: Final = cast(Sequence[object], chunks) # cast-ok: chunks are untyped upstream
reasoning_items: Final = tuple(chain.from_iterable(_chunk_reasoning_items(chunk) for chunk in typed_chunks))
if reasoning_items:
yield "reasoning_items", list(reasoning_items)
annotation_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "annotations" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["annotations"] is not None
]
if annotation_chunks:
# Merge annotations from ALL chunks — providers may spread
# them across multiple streaming chunks or send them only in
# the final chunk.
all_annotations: Final[list[object]] = []
for ac in annotation_chunks:
all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
yield "annotations", all_annotations
audio_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "audio" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["audio"] is not None
]
if audio_chunks:
yield "audio", processor.get_combined_audio_content(audio_chunks)
# Handle image chunks from models like gemini-2.5-flash-image
# See: https://github.com/BerriAI/litellm/issues/19478
image_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "images" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["images"] is not None
]
if image_chunks:
# Images come complete in a single chunk, collect all images from all chunks
all_images: Final = []
for chunk in image_chunks:
all_images.extend(chunk["choices"][0]["delta"]["images"])
yield "images", all_images
# Combine provider_specific_fields from streaming chunks (e.g., web_search_results, citations)
# See: https://github.com/BerriAI/litellm/issues/17737
provider_specific_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "provider_specific_fields" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["provider_specific_fields"] is not None
]
if provider_specific_chunks:
provider_field_dicts: Final = tuple(
fields
for chunk in provider_specific_chunks
for fields in (chunk["choices"][0]["delta"]["provider_specific_fields"],)
if isinstance(fields, dict)
)
streamed_citations: Final = tuple(
fields["citation"] for fields in provider_field_dicts if fields.get("citation") is not None
)
citation_fields: Final = (
{"citations": _joined_streamed_citations(streamed_citations)} if streamed_citations else {}
)
combined_provider_fields: Final = {
key: value
for fields in (citation_fields, *provider_field_dicts)
for key, value in fields.items()
if key != "citation"
}
if combined_provider_fields:
yield "provider_specific_fields", combined_provider_fields
def stream_chunk_builder(
chunks: list,
messages: Sequence | None = None,
start_time=None,
end_time=None,
start_time: object = None,
end_time: object = None,
logging_obj: Optional["Logging"] = None,
count_prompt_tokens: Callable[[], int] | None = None,
) -> ModelResponse | TextCompletionResponse | None:
@ -1552,26 +1734,21 @@ def stream_chunk_builder(
return stream_chunk_builder_text_completion(chunks=chunks, messages=messages)
model: Final = chunks[0]["model"]
# Initialize the response dictionary
response: Final = processor.build_base_response(chunks)
choice_streams: Final = _choice_streams(chunks)
if len(choice_streams) > 1:
response.choices = [
ChunkProcessor.build_base_choice(index, choice_chunks) for index, choice_chunks in choice_streams
]
rebuilt_choices: Final = tuple(zip(response.choices, (choice_chunks for _, choice_chunks in choice_streams)))
# Fast path for the common text-only streaming case:
# avoid repeated multi-pass list scans over chunks.
simple_content_parts: Final[list[str]] = []
is_simple_text_stream = True
for chunk in chunks:
if not chunk.get("choices"):
continue
if (part := _simple_text_part(chunk["choices"])) is None:
is_simple_text_stream = False
break
if part:
simple_content_parts.append(part)
is_simple_text_stream: Final = all(
_simple_text_part(chunk["choices"]) is not None for chunk in chunks if chunk.get("choices")
)
if is_simple_text_stream:
if simple_content_parts:
response["choices"][0]["message"]["content"] = "".join(simple_content_parts)
for choice, choice_chunks in rebuilt_choices:
if simple_text := _joined_simple_text(choice_chunks):
choice.message.content = simple_text
completion_output = get_content_from_model_response(response)
usage = processor.calculate_usage(
chunks=chunks,
@ -1610,155 +1787,9 @@ def stream_chunk_builder(
processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
return response
tool_call_chunks: Final = [
chunk
for chunk in chunks
if any(
"tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] is not None
for choice in chunk.get("choices") or ()
)
]
if len(tool_call_chunks) > 0:
tool_calls_list: Final = processor.get_combined_tool_content(tool_call_chunks)
_choice = response.choices[0]
_choice.message.content = None
_choice.message.tool_calls = tool_calls_list
function_call_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "function_call" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["function_call"] is not None
]
if len(function_call_chunks) > 0:
_choice = response.choices[0]
_choice.message.content = None
_choice.message.function_call = processor.get_combined_function_call_content(function_call_chunks)
content_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "content" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["content"] is not None
]
if len(content_chunks) > 0:
response["choices"][0]["message"]["content"] = processor.get_combined_content(content_chunks)
thinking_blocks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "thinking_blocks" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["thinking_blocks"] is not None
]
if len(thinking_blocks) > 0:
response["choices"][0]["message"]["thinking_blocks"] = processor.get_combined_thinking_content(
thinking_blocks
)
reasoning_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "reasoning_content" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["reasoning_content"] is not None
]
if len(reasoning_chunks) > 0:
response["choices"][0]["message"]["reasoning_content"] = processor.get_combined_reasoning_content(
reasoning_chunks
)
typed_chunks: Final = cast(Sequence[object], chunks) # cast-ok: chunks are untyped upstream
reasoning_items: Final = tuple(chain.from_iterable(_chunk_reasoning_items(chunk) for chunk in typed_chunks))
if reasoning_items:
response["choices"][0]["message"]["reasoning_items"] = list(reasoning_items)
annotation_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "annotations" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["annotations"] is not None
]
if len(annotation_chunks) > 0:
# Merge annotations from ALL chunks — providers may spread
# them across multiple streaming chunks or send them only in
# the final chunk.
all_annotations: Final[list] = []
for ac in annotation_chunks:
all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
response["choices"][0]["message"]["annotations"] = all_annotations
audio_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "audio" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["audio"] is not None
]
if len(audio_chunks) > 0:
_choice = response.choices[0]
_choice.message.audio = processor.get_combined_audio_content(audio_chunks)
# Handle image chunks from models like gemini-2.5-flash-image
# See: https://github.com/BerriAI/litellm/issues/19478
image_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "images" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["images"] is not None
]
if len(image_chunks) > 0:
# Images come complete in a single chunk, collect all images from all chunks
all_images: Final = []
for chunk in image_chunks:
all_images.extend(chunk["choices"][0]["delta"]["images"])
response["choices"][0]["message"]["images"] = all_images
# Combine provider_specific_fields from streaming chunks (e.g., web_search_results, citations)
# See: https://github.com/BerriAI/litellm/issues/17737
provider_specific_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "provider_specific_fields" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["provider_specific_fields"] is not None
]
if len(provider_specific_chunks) > 0:
provider_field_dicts: Final = tuple(
fields
for chunk in provider_specific_chunks
for fields in (chunk["choices"][0]["delta"]["provider_specific_fields"],)
if isinstance(fields, dict)
)
streamed_citations: Final = tuple(
fields["citation"] for fields in provider_field_dicts if fields.get("citation") is not None
)
citation_fields: Final = (
{"citations": _joined_streamed_citations(streamed_citations)} if streamed_citations else {}
)
combined_provider_fields: Final = {
key: value
for fields in (citation_fields, *provider_field_dicts)
for key, value in fields.items()
if key != "citation"
}
if combined_provider_fields:
_choice = response.choices[0]
_choice.message.provider_specific_fields = combined_provider_fields
for choice, choice_chunks in rebuilt_choices:
for field, value in _assembled_message_fields(processor, choice_chunks):
choice.message[field] = value
completion_output = get_content_from_model_response(response)

63
litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py
View file

@ -22,6 +22,7 @@ import os
from collections.abc import AsyncIterator, Mapping, Sequence
from dataclasses import dataclass, replace
from datetime import datetime
from itertools import chain
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
@ -120,6 +121,11 @@ class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
"""Base-class constructor options this guardrail forwards untouched to CustomGuardrail."""
def _streamed_choice_index(choice: object) -> int:
index: Final[object] = getattr(choice, "index", None)
return index if isinstance(index, int) else 0
class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail):
"""
Cisco AI Defense guardrail integration.
@ -585,29 +591,44 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail):
@classmethod
def _extract_streaming_chunk_scan_text(cls, chunks: Sequence[object]) -> str:
streamed_choices: Final = tuple(
choice
for choice in chain.from_iterable(getattr(chunk, "choices", None) or () for chunk in chunks)
if getattr(choice, "delta", None) is not None
)
choice_indices: Final = sorted(frozenset(map(_streamed_choice_index, streamed_choices)))
choice_texts: Final = (
cls._extract_streamed_deltas_scan_text(
tuple(
getattr(choice, "delta", None)
for choice in streamed_choices
if _streamed_choice_index(choice) == index
)
)
for index in choice_indices
)
return " ".join(text for text in choice_texts if text)
@classmethod
def _extract_streamed_deltas_scan_text(cls, deltas: Sequence[object]) -> str:
original_text = ""
argument_text = ""
for chunk in chunks:
choices = getattr(chunk, "choices", None) or []
for c in choices:
delta: object | None = getattr(c, "delta", None)
if delta is None:
continue
text = getattr(delta, "content", None)
if isinstance(text, str):
original_text += text
reasoning_text = " ".join(cls._extract_message_reasoning_parts(delta))
if reasoning_text:
original_text += reasoning_text
for tc in getattr(delta, "tool_calls", None) or []:
args = cls._extract_tool_call_arguments(tc)
if args:
argument_text += args
fc: object | None = getattr(delta, "function_call", None)
if fc is not None:
args = cls._extract_function_call_arguments(fc)
if args:
argument_text += args
for delta in deltas:
text = getattr(delta, "content", None)
if isinstance(text, str):
original_text += text
reasoning_text = " ".join(cls._extract_message_reasoning_parts(delta))
if reasoning_text:
original_text += reasoning_text
for tc in getattr(delta, "tool_calls", None) or []:
args = cls._extract_tool_call_arguments(tc)
if args:
argument_text += args
fc: object | None = getattr(delta, "function_call", None)
if fc is not None:
args = cls._extract_function_call_arguments(fc)
if args:
argument_text += args
return " ".join(part for part in (original_text, argument_text) if part)
# ------------------------------------------------------------------

79
tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py
View file

@ -1304,12 +1304,85 @@ def test_stream_chunk_builder_keeps_each_choices_tool_call_arguments_apart():
]
)
assert [(tool_call.id, tool_call.function.arguments) for tool_call in response.choices[0].message.tool_calls] == [
("call_a", '{"fruit": "persimmon"}'),
("call_b", '{"fruit": "durian"}'),
assert [
(choice.index, [(tool_call.id, tool_call.function.arguments) for tool_call in choice.message.tool_calls])
for choice in response.choices
] == [
(0, [("call_a", '{"fruit": "persimmon"}')]),
(1, [("call_b", '{"fruit": "durian"}')]),
]
def test_stream_chunk_builder_rebuilds_each_choice_of_a_multi_choice_text_stream_on_its_own():
def chunk(choice_index: int, content: str, finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-123",
object="chat.completion.chunk",
created=1234567890,
model="gpt-4.1-mini",
choices=[StreamingChoices(index=choice_index, delta=Delta(content=content), finish_reason=finish_reason)],
)
response = stream_chunk_builder(
chunks=[
chunk(0, "The persim"),
chunk(1, "The"),
chunk(1, " persimmon is"),
chunk(0, "mon is ripe.", finish_reason="stop"),
chunk(1, " ripe and", finish_reason="length"),
]
)
assert [(choice.index, choice.message.content, choice.finish_reason) for choice in response.choices] == [
(0, "The persimmon is ripe.", "stop"),
(1, "The persimmon is ripe and", "length"),
]
def test_stream_chunk_builder_keeps_each_choices_function_call_annotations_and_images_on_that_choice():
def chunk(choice_index: int, delta: Delta) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-123",
object="chat.completion.chunk",
created=1234567890,
model="gpt-4.1-mini",
choices=[StreamingChoices(index=choice_index, delta=delta, finish_reason=None)],
)
citation: Final = {
"type": "url_citation",
"url_citation": {"start_index": 0, "end_index": 5, "url": "https://example.com/kiwi", "title": "Kiwi"},
}
image: Final = {"type": "image_url", "index": 0, "image_url": {"url": "data:image/png;base64,a2l3aQ=="}}
response = stream_chunk_builder(
chunks=[
chunk(0, Delta(function_call={"name": "lookup_fruit", "arguments": '{"fruit": '})),
chunk(1, Delta(content="Kiwis are")),
chunk(0, Delta(function_call={"arguments": '"kiwi"}'})),
chunk(1, Delta(content=" fuzzy.", annotations=[citation])),
chunk(1, Delta(images=[image])),
]
)
function_choice, cited_choice = response.choices
assert (
function_choice.index,
function_choice.message.content,
function_choice.message.function_call.name,
function_choice.message.function_call.arguments,
getattr(function_choice.message, "annotations", None),
getattr(function_choice.message, "images", None),
) == (0, None, "lookup_fruit", '{"fruit": "kiwi"}', None, None)
assert (
cited_choice.index,
cited_choice.message.content,
cited_choice.message.function_call,
cited_choice.message.annotations,
cited_choice.message.images,
) == (1, "Kiwis are fuzzy.", None, [citation], [image])
def test_get_combined_tool_content_joins_many_dict_shaped_argument_fragments_in_order():
processor = ChunkProcessor.__new__(ChunkProcessor)
first_fragments = [f"a{i};" for i in range(300)]

56
tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py
View file

@ -2108,6 +2108,62 @@ class TestCiscoAIDefenseStreamingBypass:
f"Original: {chunks!r}, received: {received!r}"
)
@pytest.mark.asyncio
async def test_streaming_safe_n2_response_inspects_each_choice_and_yields_original_chunks(self):
g = _make_guardrail(event_hook=["pre_call", "post_call"])
chunks = [
ModelResponseStream(
id="resp_1",
choices=[StreamingChoices(delta=Delta(content=part), finish_reason=finish_reason, index=index)],
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
)
for index, part, finish_reason in (
(0, "The persim", None),
(1, "The", None),
(0, "mon is ripe.", "stop"),
(1, " persimmon is ripe.", "stop"),
)
]
received, post_mock = await _streaming_setup(g, chunks, cisco_response=_safe_response())
inspected = post_mock.call_args.kwargs["json"]["messages"]
assert [m["content"] for m in inspected if m["role"] == "assistant"] == [
"The persimmon is ripe.",
"The persimmon is ripe.",
]
assert received == chunks
@pytest.mark.asyncio
async def test_streaming_safe_n2_response_with_a_function_call_choice_yields_original_chunks(self):
g = _make_guardrail(event_hook=["pre_call", "post_call"])
chunks = [
ModelResponseStream(
id="resp_1",
choices=[StreamingChoices(delta=delta, finish_reason=None, index=index)],
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
)
for index, delta in (
(0, Delta(function_call={"name": "lookup_fruit", "arguments": '{"fruit": '})),
(1, Delta(content="Kiwis are")),
(0, Delta(function_call={"arguments": '"kiwi"}'})),
(1, Delta(content=" fuzzy.")),
)
]
received, post_mock = await _streaming_setup(g, chunks, cisco_response=_safe_response())
inspected = post_mock.call_args.kwargs["json"]["messages"]
assert [m["content"] for m in inspected if m["role"] == "assistant"] == [
'{"fruit": "kiwi"}',
"Kiwis are fuzzy.",
]
assert received == chunks
@pytest.mark.asyncio
async def test_streaming_redact_does_not_replay_tool_call_arguments(self):
g = _make_guardrail(

45
tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py
View file

@ -5,6 +5,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics)
import json
import logging
import re
from itertools import chain
from typing import Final
from unittest.mock import patch
@ -25,10 +26,14 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
PermissionError,
)
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
ChatCompletionMessageToolCall,
Choices,
Delta,
Function,
ModelResponse,
ModelResponseStream,
StreamingChoices,
)
@ -866,6 +871,46 @@ class TestToolPermissionGuardrail:
f"got: {chunks[0].choices[0].delta.content!r}"
)
@pytest.mark.asyncio
async def test_async_post_call_streaming_iterator_hook_re_emits_each_choice_of_an_n2_tool_call_stream(self):
def chunk(choice_index: int, tool_call: ChatCompletionDeltaToolCall) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-n2-tools",
created=1700000000,
model="gpt-4.1-mini",
object="chat.completion.chunk",
choices=[StreamingChoices(index=choice_index, delta=Delta(tool_calls=[tool_call]))],
)
def fragment(arguments: str, name: str | None = None) -> ChatCompletionDeltaToolCall:
return ChatCompletionDeltaToolCall(
id="call_shared" if name else None,
index=0,
type="function",
function=Function(name=name, arguments=arguments),
)
async def _stream():
for streamed in (
chunk(0, fragment("", name="Bash")),
chunk(1, fragment("", name="Bash")),
chunk(0, fragment('{"command": "ls"}')),
chunk(1, fragment('{"command": "pwd"}')),
):
yield streamed
re_emitted: Final = [
chunk
async for chunk in self.guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(), response=_stream(), request_data={}
)
]
assert [
(choice.index, [(call.function.name, call.function.arguments) for call in choice.delta.tool_calls])
for choice in chain.from_iterable(streamed.choices for streamed in re_emitted)
] == [(0, [("Bash", '{"command": "ls"}')]), (1, [("Bash", '{"command": "pwd"}')])]
def test_modify_response_with_permission_errors(self):
# Setup a response with one tool_call
tool_call = ChatCompletionMessageToolCall(

7
tests/unit/test_main.py
View file

@ -2843,9 +2843,10 @@ def test_stream_chunk_builder_keeps_tool_calls_carried_only_by_a_later_choice_of
response = stream_chunk_builder(chunks=chunks)
tool_calls = response.choices[0].message.tool_calls
assert tool_calls is not None
assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [
text_choice, tool_choice = response.choices
assert (text_choice.index, text_choice.message.content, text_choice.message.tool_calls) == (0, "hello world", None)
assert tool_choice.index == 1
assert [(call.id, call.function.name, call.function.arguments) for call in tool_choice.message.tool_calls] == [
("call_1", "lookup_fruit", '{"fruit":"kiwi"}')
]