diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 580e55273f6..bb59b4617af 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1979,6 +1979,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"] @@ -2021,6 +2022,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): tool_call_index_map: dict[int, int] | None = None, # mutable-ok: per-stream state, remapped in place replays_openai_history: bool = False, structured_output: bool = False, + reasoning_summary_streamed: bool = False, ) -> "ModelResponseStream": """ Translate a Responses API streaming chunk to OpenAI chat completion streaming format. @@ -2030,6 +2032,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): tool_call_index_map: Per-stream output_index -> sequential tool_call index map replays_openai_history: Record the completed output's item sequence for the next turn structured_output: Only the first message item became content + reasoning_summary_streamed: Whether earlier chunks already streamed reasoning summary text Returns: ModelResponseStream: OpenAI-formatted streaming chunk @@ -2210,13 +2213,15 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) else: raise ValueError(f"Chat provider: Invalid text delta {parsed_chunk}") + 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) if content_part: return ModelResponseStream( choices=[ StreamingChoices( - index=cast(int, parsed_chunk.get("summary_index")), + index=0, delta=Delta(reasoning_content=content_part), ) ] @@ -2317,10 +2322,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): tool_call_index_map=self._tool_call_index_map, replays_openai_history=self._replays_openai_history, structured_output=self._structured_output, + 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 8e1791c3c92..c69d9277d65 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 @@ -257,7 +257,7 @@ def apply_grounding_request_counts( class ChunkProcessor: - def __init__(self, chunks: list, messages: Sequence | None = None): + def __init__(self, chunks: list, messages: Sequence[object] | None = None): self.chunks = self._sort_chunks(chunks) self.messages = messages self.first_chunk = chunks[0] @@ -390,6 +390,30 @@ class ChunkProcessor: case _: return "assistant" + @staticmethod + def _finish_reason_of_chunks(chunks: Sequence["_BaseChunk"], unreported: str = "stop") -> object: + finish_reason = unreported + 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, unreported="tool_calls" if _carries_tool_calls(chunks) else "stop" + ), + message={"role": ChunkProcessor._get_role_from_chunks(chunks), "content": ""}, + ) + @staticmethod def _get_model_from_chunks(chunks: Sequence["_BaseChunk"], first_chunk_model: str) -> str: """ @@ -414,7 +438,11 @@ class ChunkProcessor: None, ) - def build_base_response(self, chunks: Sequence["_BaseChunk"]) -> ModelResponse: + def build_base_response( + self, + chunks: Sequence["_BaseChunk"], + choice_streams: Sequence[tuple[int, Sequence["_BaseChunk"]]], + ) -> ModelResponse: chunk = self.first_chunk id: Final = ChunkProcessor._get_chunk_id(chunks) object: Final = chunk["object"] @@ -426,16 +454,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( @@ -455,13 +474,17 @@ class ChunkProcessor: if service_tier is not None else MappingProxyType({}) ), - "choices": [ - { - "index": 0, - "message": {"role": role, "content": ""}, - "finish_reason": finish_reason, - } - ], + "choices": ( + [ChunkProcessor.build_base_choice(index, choice_chunks) for index, choice_chunks in choice_streams] + if len(choice_streams) > 1 + else [ + { + "index": 0, + "message": {"role": role, "content": ""}, + "finish_reason": finish_reason, + } + ] + ), "usage": { "prompt_tokens": 0, # Modify as needed "completion_tokens": 0, # Modify as needed @@ -1520,11 +1543,188 @@ 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 _carries_tool_calls(chunks: Sequence[_BaseChunk]) -> bool: + return any( + "tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] + for choice in chain.from_iterable(chunk.get("choices") or () for chunk in chunks) + ) + + +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 _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]]: + 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 + ] + + message_annotations, provider_annotations = ( + partition_annotations( + list(chain.from_iterable(chunk["choices"][0]["delta"]["annotations"] for chunk in annotation_chunks)) + ) + if annotation_chunks + else (None, ()) + ) + if message_annotations is not None: + yield "annotations", message_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 + ] + + 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 {} + provider_annotation_fields: Final = {"annotations": list(provider_annotations)} if provider_annotations else {} + combined_provider_fields: Final = { + key: value + for fields in (provider_annotation_fields, 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: @@ -1561,26 +1761,15 @@ def stream_chunk_builder( usage_model: Final = usage_tokenizer_model( model, get_base_model_from_metadata(logging_obj.model_call_details) if logging_obj is not None else None ) - # Initialize the response dictionary - response: Final = processor.build_base_response(chunks) + choice_streams: Final = _choice_streams(chunks) + response: Final = processor.build_base_response(chunks, 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) - - if is_simple_text_stream: - if simple_content_parts: - response["choices"][0]["message"]["content"] = "".join(simple_content_parts) + 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( chunks=chunks, @@ -1619,165 +1808,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[object]] = [] - for ac in annotation_chunks: - all_annotations.extend(ac["choices"][0]["delta"]["annotations"]) - message_annotations, provider_annotations = partition_annotations(all_annotations) - if message_annotations is not None: - response["choices"][0]["message"]["annotations"] = message_annotations - if provider_annotations: - response.choices[0].message.provider_specific_fields = { - **(response.choices[0].message.provider_specific_fields or {}), - "annotations": list(provider_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 = { - **(_choice.message.provider_specific_fields or {}), - **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/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 3070bedd30d..5c5204ef2f4 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -687,11 +687,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far: Sequence["ModelResponseStream"], litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> "ModelResponse": - """``stream_chunk_builder`` folds every choice of a stream into one index-0 - choice, so the stream is rebuilt one choice index at a time (every chunk - kept, its choices narrowed to that index, so usage-only chunks still - count) and the rebuilt choices are stitched into one response, each - carrying the index the stream gave it.""" choice_indices: Final = tuple( sorted(frozenset(choice.index for response in responses_so_far for choice in response.choices)) ) 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/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 85f476ecd62..d69fe4068e0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -246,20 +246,38 @@ class NomaGuardrail(CustomGuardrail): ) -> str | None: """Shared logic for processing LLM response checks""" - start_time: Final = datetime.now() - extra_data: Final = self.get_guardrail_dynamic_request_body_params(request_data) - if not isinstance(response, litellm.ModelResponse): return None - content = None - for choice in response.choices: - if isinstance(choice, litellm.Choices) and choice.message.content: - content = choice.message.content - break + checked_contents: Final = tuple( + [ + await self._process_choice_check(request_data, response, position, content, user_auth, event_type) + for position, content in self._choice_texts(response) + ] + ) + return next(iter(checked_contents), None) - if not content or not isinstance(content, str): - return None + @staticmethod + def _choice_texts(response: litellm.ModelResponse) -> tuple[tuple[int, str], ...]: + return tuple( + (position, choice.message.content) + for position, choice in enumerate(response.choices) + if isinstance(choice, litellm.Choices) + and isinstance(choice.message.content, str) + and choice.message.content + ) + + async def _process_choice_check( + self, + request_data: dict[str, object], + response: litellm.ModelResponse, + position: int, + content: str, + user_auth: UserAPIKeyAuth, + event_type: GuardrailEventHooks | None, + ) -> str: + start_time: Final = datetime.now() + extra_data: Final = self.get_guardrail_dynamic_request_body_params(request_data) payload: Final = { "input": [ @@ -306,7 +324,7 @@ class NomaGuardrail(CustomGuardrail): anonymized_content: Final = self._extract_anonymized_content(response_json, ASSISTANT_ROLE) if anonymized_content: # Replace the LLM response content with anonymized version - self._replace_llm_response_content(response, anonymized_content) + self._replace_llm_response_content(response, anonymized_content, position) verbose_proxy_logger.debug("Noma guardrail anonymized LLM response: %s", anonymized_content) return anonymized_content @@ -486,21 +504,14 @@ class NomaGuardrail(CustomGuardrail): messages[i]["content"] = anonymized_content break - def _replace_llm_response_content(self, response: LLMResponse, anonymized_content: str): - """ - Replace the LLM response content with anonymized version. - - Args: - response: The original LLM response - anonymized_content: The anonymized content to replace with - """ + def _replace_llm_response_content(self, response: LLMResponse, anonymized_content: str, position: int): + """Replace the content of the choice at ``position`` with its anonymized version.""" if not isinstance(response, litellm.ModelResponse): return - # Replace content in all choices - for choice in response.choices: - if isinstance(choice, litellm.Choices) and choice.message.content: - choice.message.content = anonymized_content + choice: Final = response.choices[position] + if isinstance(choice, litellm.Choices): + choice.message.content = anonymized_content async def _check_user_message_background( self, @@ -815,6 +826,7 @@ class NomaGuardrail(CustomGuardrail): ) if isinstance(assembled_model_response, ModelResponse): + texts_before_check: Final = self._choice_texts(assembled_model_response) try: processed_response: Final = await self._check_llm_response( request_data, @@ -828,7 +840,11 @@ class NomaGuardrail(CustomGuardrail): if self.block_failures: raise verbose_proxy_logger.error("Noma streaming post-call hook failed: %s", e) - for chunk in all_chunks: + if self._choice_texts(assembled_model_response) == texts_before_check: + for chunk in all_chunks: + yield chunk + return + async for chunk in MockResponseIterator(model_response=assembled_model_response): yield chunk return diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index fa92a0c90e0..e8c094c0c37 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -36,6 +36,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( ) from litellm.types.utils import ( CallTypesLiteral, + ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall, Choices, Function, @@ -744,6 +745,14 @@ 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 | ChatCompletionMessageCustomToolCall, + ) -> tuple[str | None, ...]: + if isinstance(tool_call, ChatCompletionMessageCustomToolCall): + return (tool_call.id, tool_call.type, tool_call.custom.name, tool_call.custom.input) + 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 +770,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 +783,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 +795,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/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index c01d07fb5c3..c88a8c1cc28 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -35,10 +35,7 @@ from litellm.types.llms.openai import ToolParam as ResponsesToolParam from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.utils import ( CallTypes, - ChatCompletionMessageCustomToolCall, - ChatCompletionMessageToolCall, Choices, - Message, ModelResponse, StandardLoggingMCPToolCall, ) @@ -587,29 +584,18 @@ class LiteLLM_Proxy_MCP_Handler: _extract_tool_calls_from_response = extract_tool_calls_from_response + @staticmethod + def _continued_choice(response: ModelResponse) -> Choices | None: + choices: Final = tuple(choice for choice in response.choices if isinstance(choice, Choices)) + return next((choice for choice in choices if choice.message.tool_calls), next(iter(choices), None)) + @staticmethod def extract_tool_calls_from_chat_response(response: ModelResponse) -> list[object]: - """Extract tool calls from a chat completion response.""" - tool_calls: Final[list[object]] = [] - - try: - for choice in response.choices: - message: Message | None = getattr(choice, "message", None) - if message is None: - continue - tool_call_entries: ( - Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None - ) = getattr(message, "tool_calls", None) - if tool_call_entries: - for tool_call in tool_call_entries: - if hasattr(tool_call, "model_dump"): - tool_calls.append(tool_call.model_dump()) - else: - tool_calls.append(tool_call) - except Exception: - verbose_logger.exception("Failed to extract tool calls from chat completion response") - - return tool_calls + """Extract the tool calls of the first choice that made any, since the follow-up turn continues it alone.""" + continued_choice: Final = LiteLLM_Proxy_MCP_Handler._continued_choice(response) + if continued_choice is None: + return [] + return [tool_call.model_dump() for tool_call in continued_choice.message.tool_calls or ()] _extract_tool_calls_from_chat_response = extract_tool_calls_from_chat_response @@ -1018,9 +1004,9 @@ class LiteLLM_Proxy_MCP_Handler: message_to_append: dict | None = None try: - first_choice: Final = response.choices[0] - if isinstance(first_choice, Choices) and getattr(first_choice, "message", None): - message_to_append = first_choice.message.model_dump(exclude_none=True) + continued_choice: Final = LiteLLM_Proxy_MCP_Handler._continued_choice(response) + if continued_choice is not None: + message_to_append = continued_choice.message.model_dump(exclude_none=True) # Ensure tool_calls have arguments field (required by OpenAI API) if message_to_append.get("tool_calls"): for tool_call in message_to_append["tool_calls"]: 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 d59f59703ad..4e6ed16e360 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 @@ -418,6 +418,45 @@ def test_openai_responses_chunk_parser_reasoning_summary(): assert delta.function_call is None +def test_bridged_reasoning_summary_parts_rebuild_into_the_one_choice_they_streamed_on(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + events = ( + {"type": "response.reasoning_summary_part.added", "output_index": 0, "summary_index": 0}, + { + "type": "response.reasoning_summary_text.delta", + "delta": "First part.", + "output_index": 0, + "summary_index": 0, + }, + {"type": "response.reasoning_summary_part.done", "output_index": 0, "summary_index": 0}, + {"type": "response.reasoning_summary_part.added", "output_index": 0, "summary_index": 1}, + { + "type": "response.reasoning_summary_text.delta", + "delta": "Second part.", + "output_index": 0, + "summary_index": 1, + }, + {"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. Next item.", "Answer") + ] + + def test_chunk_parser_string_output_text_delta_produces_text(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, 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 dd1f495b0d0..28587bf4749 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,118 @@ 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_gives_a_choice_without_a_finish_chunk_the_finish_reason_its_content_implies(): + def chunk(choice_index: int, delta: Delta, 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, finish_reason=finish_reason)], + ) + + def tool_call(call_id: str) -> list[ChatCompletionDeltaToolCall]: + return [ + ChatCompletionDeltaToolCall( + id=call_id, index=0, type="function", function=Function(name="lookup_fruit", arguments="{}") + ) + ] + + response = stream_chunk_builder( + chunks=[ + chunk(0, Delta(tool_calls=tool_call("call_a"))), + chunk(1, Delta(tool_calls=tool_call("call_b"))), + chunk(2, Delta(content="Kiwis are fuzzy.")), + chunk(0, Delta(), finish_reason="tool_calls"), + ] + ) + + assert [(choice.index, choice.finish_reason) for choice in response.choices] == [ + (0, "tool_calls"), + (1, "tool_calls"), + (2, "stop"), + ] + + +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_noma.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma.py index ee369f72f01..574550e1bdd 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma.py @@ -1,4 +1,5 @@ import copy +from itertools import chain import os from typing import cast from unittest.mock import AsyncMock, MagicMock, patch @@ -18,7 +19,7 @@ import litellm.proxy.guardrails.guardrail_hooks.noma.noma as noma_legacy_module from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import Choices, Message +from litellm.types.utils import Choices, Delta, Message, ModelResponseStream, StreamingChoices @pytest.fixture @@ -1848,12 +1849,207 @@ class TestNomaAnonymizationLogic: ) anonymize_guardrail._replace_llm_response_content( - response, "Your email is *******" + response, "Your email is *******", 0 ) assert response.choices[0].message.content == "Your email is *******" +def _noma_reply_per_text(verdict_for_text): + def reply(*args, **kwargs): + response = MagicMock() + response.json.return_value = verdict_for_text(kwargs["json"]["input"][0]["content"][0]["text"]) + response.raise_for_status = MagicMock() + return response + + return reply + + +class TestNomaMultiChoiceResponses: + @pytest.mark.asyncio + async def test_streamed_two_choice_reply_is_blocked_when_only_the_second_choice_is_flagged(self): + guardrail = 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, + ) + scanned_texts = [] + + def verdict_for_text(text): + scanned_texts.append(text) + return {"aggregatedScanResult": "launch code" in text, "scanResult": []} + + async def upstream(): + for index, part in ((0, "Kiwis are fuzzy."), (1, "The launch code is 0000.")): + yield ModelResponseStream( + id="resp_1", + choices=[StreamingChoices(delta=Delta(content=part), finish_reason="stop", index=index)], + created=1234567890, + model="gpt-4.1-mini", + object="chat.completion.chunk", + ) + + with patch.object(guardrail.async_handler, "post", side_effect=_noma_reply_per_text(verdict_for_text)): + with pytest.raises(NomaBlockedMessage): + async for _ 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": {}}, + ): + pass + + assert scanned_texts == ["Kiwis are fuzzy.", "The launch code is 0000."] + + @pytest.mark.asyncio + async def test_each_choice_gets_its_own_anonymized_text(self): + guardrail = NomaGuardrail( + api_key="test-api-key", + api_base="https://api.test.noma.security/", + anonymize_input=True, + monitor_mode=False, + block_failures=True, + guardrail_name="noma-post-call", + event_hook="post_call", + default_on=True, + ) + + def verdict_for_text(text): + return { + "aggregatedScanResult": False, + "scanResult": [ + { + "role": "assistant", + "type": "message", + "results": { + "anonymizedContent": {"anonymized": text.replace("test@example.com", "*******")}, + "sensitiveData": {"PII": {"result": False, "status": "SUCCESS"}}, + }, + } + ], + } + + response = ModelResponse( + id="resp_1", + choices=[ + Choices(index=0, finish_reason="stop", message=Message(role="assistant", content="Mail test@example.com")), + Choices( + index=1, finish_reason="stop", message=Message(role="assistant", content="Write test@example.com") + ), + ], + model="gpt-4.1-mini", + ) + + with patch.object(guardrail.async_handler, "post", side_effect=_noma_reply_per_text(verdict_for_text)): + await guardrail.async_post_call_success_hook( + data={"litellm_call_id": "call-1", "metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=response, + ) + + assert [choice.message.content for choice in response.choices] == ["Mail *******", "Write *******"] + + @pytest.mark.asyncio + async def test_fail_open_stream_keeps_redactions_made_before_a_later_choice_fails(self): + guardrail = NomaGuardrail( + api_key="test-api-key", + api_base="https://api.test.noma.security/", + anonymize_input=True, + monitor_mode=False, + block_failures=False, + guardrail_name="noma-post-call", + event_hook="post_call", + default_on=True, + ) + + def verdict_for_text(text): + if text.startswith("Write"): + raise httpx.ConnectError("Noma is unreachable") + return { + "aggregatedScanResult": False, + "scanResult": [ + { + "role": "assistant", + "type": "message", + "results": { + "anonymizedContent": {"anonymized": text.replace("test@example.com", "*******")}, + "sensitiveData": {"PII": {"result": False, "status": "SUCCESS"}}, + }, + } + ], + } + + async def upstream(): + for index, part in ((0, "Mail test@example.com"), (1, "Write back soon")): + yield ModelResponseStream( + id="resp_1", + choices=[StreamingChoices(delta=Delta(content=part), finish_reason="stop", index=index)], + created=1234567890, + model="gpt-4.1-mini", + object="chat.completion.chunk", + ) + + with patch.object(guardrail.async_handler, "post", side_effect=_noma_reply_per_text(verdict_for_text)): + streamed = [ + 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": {}}, + ) + ] + + streamed_choices = tuple(chain.from_iterable(part.choices for part in streamed)) + text_by_choice = { + index: "".join(choice.delta.content or "" for choice in streamed_choices if choice.index == index) + for index in (0, 1) + } + assert text_by_choice == {0: "Mail *******", 1: "Write back soon"} + + @pytest.mark.asyncio + async def test_fail_open_stream_replays_the_original_chunks_when_nothing_was_redacted(self): + guardrail = NomaGuardrail( + api_key="test-api-key", + api_base="https://api.test.noma.security/", + anonymize_input=True, + monitor_mode=False, + block_failures=False, + guardrail_name="noma-post-call", + event_hook="post_call", + default_on=True, + ) + originals = [ + ModelResponseStream( + id="resp_1", + choices=[StreamingChoices(delta=Delta(content=part), index=0)], + created=1234567890, + model="gpt-4.1-mini", + object="chat.completion.chunk", + ) + for part in ("Kiwis", " are", " fuzzy.") + ] + + async def upstream(): + for chunk in originals: + yield chunk + + with patch.object(guardrail.async_handler, "post", side_effect=httpx.ConnectError("Noma is unreachable")): + streamed = [ + 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": {}}, + ) + ] + + assert len(streamed) == len(originals) + assert all(sent is original for sent, original in zip(streamed, originals)) + + class TestNomaAnonymizationFlow: """Test full anonymization flow with real Noma response objects""" 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..21376068686 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,16 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ) from litellm.types.utils import ( + ChatCompletionCustomToolCallPayload, + ChatCompletionDeltaToolCall, + ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall, Choices, + Delta, + Function, ModelResponse, ModelResponseStream, + StreamingChoices, ) @@ -425,6 +432,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, @@ -866,6 +908,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( @@ -952,6 +1034,35 @@ class TestToolPermissionGuardrail: assert "Permission denied" in choice.message.content + def test_modify_response_with_permission_errors_keeps_custom_tool_call_in_another_choice(self): + denied_call = ChatCompletionMessageToolCall( + id="call_bash", type="function", function=Function(name="Bash", arguments="{}") + ) + custom_call = ChatCompletionMessageCustomToolCall( + id="call_patch", custom=ChatCompletionCustomToolCallPayload(name="apply_patch", input="*** Begin Patch") + ) + response = ModelResponse( + choices=[ + Choices(index=0, message={"role": "assistant", "content": None, "tool_calls": [denied_call]}), + Choices(index=1, message={"role": "assistant", "content": None, "tool_calls": [custom_call]}), + ] + ) + denied_tools = [ + ( + denied_call, + PermissionError(tool_name="Bash", rule_id="deny_bash", message="Tool 'Bash' denied by rule 'deny_bash'"), + ) + ] + + self.guardrail._modify_response_with_permission_errors(response, denied_tools) + + kept = [ + [call.id for call in choice.message.tool_calls or []] + for choice in response.choices + if isinstance(choice, Choices) + ] + assert kept == [[], ["call_patch"]] + class TestToolPermissionGuardrailIntegration: """Integration tests for Tool Permission Guardrail""" diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index 229d5cbe676..2f27622d77c 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -131,6 +131,86 @@ def test_extract_tool_calls_from_chat_response_handles_tool_calls(): assert tool_calls[0]["function"]["name"] == "foo" +def test_chat_follow_up_answers_only_tool_calls_in_its_assistant_turn(): + response = ModelResponse( + id="resp-n2", + choices=[ + { + "index": index, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": call_id, "type": "function", "function": {"name": "foo", "arguments": "{}"}} + ], + }, + "finish_reason": "tool_calls", + } + for index, call_id in enumerate(("call-a", "call-b")) + ], + model="gpt", + created=0, + object="chat.completion", + ) + tool_results = [ + {"tool_call_id": tool_call["id"], "name": tool_call["function"]["name"], "result": "done"} + for tool_call in LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response) + ] + + follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_messages_for_chat( + [{"role": "user", "content": "hi"}], response, tool_results + ) + + assistant_call_ids = [tool_call["id"] for tool_call in follow_up[1]["tool_calls"]] + answered_call_ids = [message["tool_call_id"] for message in follow_up if message["role"] == "tool"] + assert answered_call_ids == assistant_call_ids == ["call-a"] + + +def test_chat_follow_up_continues_the_choice_that_called_a_tool_when_choice_0_answered(): + response = ModelResponse( + id="resp-n2-second", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "no tool needed"}, + "finish_reason": "stop", + }, + { + "index": 1, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call-b", "type": "function", "function": {"name": "foo", "arguments": "{}"}} + ], + }, + "finish_reason": "tool_calls", + }, + ], + model="gpt", + created=0, + object="chat.completion", + ) + tool_results = [ + {"tool_call_id": tool_call["id"], "name": tool_call["function"]["name"], "result": "done"} + for tool_call in LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response) + ] + + follow_up = LiteLLM_Proxy_MCP_Handler.create_follow_up_messages_for_chat( + [{"role": "user", "content": "hi"}], response, tool_results + ) + + assistant_call_ids = [tool_call["id"] for tool_call in follow_up[1]["tool_calls"]] + answered_call_ids = [message["tool_call_id"] for message in follow_up if message["role"] == "tool"] + assert answered_call_ids == assistant_call_ids == ["call-b"] + + +def test_extract_tool_calls_from_chat_response_without_choices_finds_none(): + response = ModelResponse(id="resp-empty", choices=[], model="gpt", created=0, object="chat.completion") + + assert LiteLLM_Proxy_MCP_Handler.extract_tool_calls_from_chat_response(response) == [] + + def test_create_follow_up_messages_for_chat_appends_tool_results(): original_messages = [{"role": "user", "content": "hi"}] response = ModelResponse( diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 11104cf7bca..571f890a4e6 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -2831,9 +2831,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"}') ]