diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 8d243e572fa..bfbf01d6329 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index 45e785fd53c..3d555707159 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -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) # ------------------------------------------------------------------ diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 59946d54dc7..7b1f7a62fd3 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -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)] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py index 4d1f254aef9..5ce78fd05d1 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py @@ -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( diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 2ee4c8c72c0..991a36ca41f 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -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( diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f38bbcf4404..ad286a62a67 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -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"}') ]