diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index c01f6db0aa9..62f390089dc 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1897,6 +1897,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state self._structured_output: Final = structured_output self._first_message_item_done: bool = False + self._reasoning_summary_streamed: bool = False def _handle_string_chunk( self, str_line: Union[str, "BaseModel"] @@ -1937,6 +1938,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): def translate_responses_chunk_to_openai_stream( parsed_chunk: dict | BaseModel, tool_call_index_map: dict[int, int] | None = None, # mutable-ok: per-stream state, remapped in place + reasoning_summary_streamed: bool = False, ) -> "ModelResponseStream": """ Translate a Responses API streaming chunk to OpenAI chat completion streaming format. @@ -1944,6 +1946,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): Args: parsed_chunk: Dict containing the Responses API event chunk tool_call_index_map: Per-stream output_index -> sequential tool_call index map + reasoning_summary_streamed: Whether earlier chunks already streamed reasoning summary text Returns: ModelResponseStream: OpenAI-formatted streaming chunk @@ -2124,7 +2127,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) else: raise ValueError(f"Chat provider: Invalid text delta {parsed_chunk}") - elif event_type == "response.reasoning_summary_part.added" and parsed_chunk.get("summary_index"): + elif event_type == "response.reasoning_summary_part.added" and reasoning_summary_streamed: return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(reasoning_content=" "))]) elif event_type == "response.reasoning_summary_text.delta": content_part = parsed_chunk.get("delta", None) @@ -2218,11 +2221,15 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): self._empty_content_chunk() if self._is_text_delta_after_structured_answer(chunk) else OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( - chunk, tool_call_index_map=self._tool_call_index_map + chunk, + tool_call_index_map=self._tool_call_index_map, + reasoning_summary_streamed=self._reasoning_summary_streamed, ) ) if self._structured_output and _is_message_item_done(chunk): self._first_message_item_done = True + if _responses_event_type(chunk) == "response.reasoning_summary_text.delta" and chunk.get("delta"): + self._reasoning_summary_streamed = True return self._with_served_service_tier(self._with_stream_scoped_id(translated)) def _is_text_delta_after_structured_answer(self, chunk: Mapping[str, object]) -> bool: diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index e438c81437e..27e59cc913d 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -3,7 +3,7 @@ import builtins import time from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass, replace -from itertools import chain, groupby +from itertools import chain, groupby, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional, Protocol, TypeAlias, TypedDict, TypeVar, Union, cast @@ -1566,8 +1566,10 @@ def _narrowed_to_choice(chunks: Sequence[_BaseChunk], index: int) -> tuple[_Base ) -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 _simple_text(chunks: Sequence[_BaseChunk]) -> str | None: + choices_per_chunk: Final = tuple(chunk["choices"] for chunk in chunks if chunk.get("choices")) + parts: Final = tuple(takewhile(lambda part: part is not None, map(_simple_text_part, choices_per_chunk))) + return "".join(filter(None, parts)) if len(parts) == len(choices_per_chunk) else None def _assembled_message_fields(processor: ChunkProcessor, chunks: Sequence) -> Iterator[tuple[str, object]]: @@ -1759,12 +1761,10 @@ def stream_chunk_builder( response: Final = processor.build_base_response(chunks, choice_streams) rebuilt_choices: Final = tuple(zip(response.choices, (choice_chunks for _, choice_chunks in choice_streams))) - 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: - for choice, choice_chunks in rebuilt_choices: - if simple_text := _joined_simple_text(choice_chunks): + simple_texts: Final = tuple(_simple_text(choice_chunks) for _, choice_chunks in choice_streams) + if None not in simple_texts: + for (choice, _), simple_text in zip(rebuilt_choices, simple_texts): + if simple_text: choice.message.content = simple_text completion_output = get_content_from_model_response(response) usage = processor.calculate_usage( diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index fa92a0c90e0..078a626b7ff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -744,6 +744,10 @@ class ToolPermissionGuardrail(CustomGuardrail): return ToolResult(tool_use_id=tool_call.id, content=error_message, is_error=True) + @staticmethod + def _judged_tool_call(tool_call: ChatCompletionMessageToolCall) -> tuple[str | None, ...]: + return (tool_call.id, tool_call.type, tool_call.function.name, tool_call.function.arguments) + def _modify_response_with_permission_errors( self, response: ModelResponse, @@ -761,11 +765,10 @@ class ToolPermissionGuardrail(CustomGuardrail): verbose_proxy_logger.info("Blocking %s unauthorized tool uses", len(denied_tools)) - # Create a mapping of tool_use_id to error result - error_results: Final = {} - for tool_use, error in denied_tools: - error_result = self._create_permission_error_result(tool_use, error) - error_results[tool_use.id] = error_result + error_results: Final = { + self._judged_tool_call(tool_use): self._create_permission_error_result(tool_use, error) + for tool_use, error in denied_tools + } # Modify the response content for choice_index, choice in enumerate(response.choices): @@ -775,9 +778,8 @@ class ToolPermissionGuardrail(CustomGuardrail): # Rewrite tool_calls for tool_call in choice.message.tool_calls or []: - tool_call_id = tool_call.id - if tool_call_id in error_results: - error_result = error_results[tool_call_id] + error_result = error_results.get(self._judged_tool_call(tool_call)) + if error_result is not None: error_messages.append(error_result.content) else: filtered_tool_calls.append(tool_call) @@ -788,7 +790,7 @@ class ToolPermissionGuardrail(CustomGuardrail): getattr(choice.message, "function_call", None), choice_index ) if legacy_tool_call is not None: - legacy_error_result = error_results.get(legacy_tool_call.id) + legacy_error_result = error_results.get(self._judged_tool_call(legacy_tool_call)) if legacy_error_result is not None: choice.message.function_call = None error_messages.append(legacy_error_result.content) diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 494401edcd9..5468cc0ce47 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -440,13 +440,20 @@ def test_bridged_reasoning_summary_parts_rebuild_into_the_one_choice_they_stream "output_index": 0, "summary_index": 1, }, - {"type": "response.output_text.delta", "delta": "Answer", "output_index": 1, "content_index": 0}, + {"type": "response.reasoning_summary_part.added", "output_index": 1, "summary_index": 0}, + { + "type": "response.reasoning_summary_text.delta", + "delta": "Next item.", + "output_index": 1, + "summary_index": 0, + }, + {"type": "response.output_text.delta", "delta": "Answer", "output_index": 2, "content_index": 0}, ) rebuilt = litellm.stream_chunk_builder(chunks=[iterator.chunk_parser(event) for event in events]) assert [(choice.index, choice.message.reasoning_content, choice.message.content) for choice in rebuilt.choices] == [ - (0, "First part. Second part.", "Answer") + (0, "First part. Second part. Next item.", "Answer") ] 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 991a36ca41f..b82cdc11d80 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -430,6 +430,41 @@ class TestToolPermissionGuardrail: assert isinstance(choice.message.content, str) assert "berri" in choice.message.content + @pytest.mark.asyncio + async def test_rewrite_denies_only_the_choice_whose_call_is_denied_when_choices_share_a_call_id(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="choice-guardrail", + rules=[{"id": "deny_bash", "tool_name": "Bash", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + response = ModelResponse( + choices=[ + Choices( + index=index, + finish_reason="tool_calls", + message={ + "tool_calls": [ + {"id": "call_shared", "type": "function", "function": {"name": name, "arguments": "{}"}} + ] + }, + ) + for index, name in enumerate(("Bash", "Read")) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["choice-guardrail"]}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + kept: Final = [ + [tool_call.function.name for tool_call in choice.message.tool_calls or ()] + for choice in response.choices + if isinstance(choice, Choices) + ] + assert kept == [[], ["Read"]] + @pytest.mark.asyncio async def test_async_post_call_success_hook_missing_arguments_blocks_param_rule( self,