diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index b356cd4a369..f8a7c38ef79 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -185,11 +185,7 @@ jobs: env: RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }} - - name: Test native routes from the installed wheel - run: | - wheel=(dist/*.whl) - uv run --isolated --no-project --with "${wheel[0]}" \ - python tests/unit/rust_bridge/native_route_wheel_test.py "${wheel[0]}" + - run: python tests/unit/rust_bridge/native_route_wheel_test.py dist/*.whl - name: Run pytest tests/test_litellm_rust with the compiled extension run: make test-rust-extension diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 6280e636700..88c4dbfdce6 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional @@ -31,12 +31,8 @@ class StreamTransformSink: single call, not shared state. """ - mutated_text_per_choice: Mapping[int, str] = field(default_factory=dict) - holdback_per_choice: Mapping[int, int] = field(default_factory=dict) - - def record(self, texts: Mapping[int, str], holdback: Mapping[int, int]) -> None: - self.mutated_text_per_choice = texts - self.holdback_per_choice = holdback + mutated_text_per_choice: dict[int, str] = field(default_factory=dict) + holdback_per_choice: dict[int, int] = field(default_factory=dict) @dataclass(frozen=True, slots=True) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index d18fe36e3df..a12f5f164e4 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -18,11 +18,9 @@ import json import time import uuid from collections.abc import Iterator, Mapping, Sequence -from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast -from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -48,7 +46,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_items, unappliable_request_rewrite, ) -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( coerce_stream_holdback_value, ) @@ -170,8 +168,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) - guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", ()) + guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", []) guardrailed_tools: Final = guardrailed_inputs.get("tools") if guardrailed_tools is not None: data["tools"] = ( @@ -260,12 +258,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name).""" names: Final[list[str]] = [] - for tool in data.get("tools") or (): + for tool in data.get("tools") or []: if isinstance(tool, dict) and tool.get("type") == "function": fn = tool.get("function") if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - for fn in data.get("functions") or (): + for fn in data.get("functions") or []: if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) return names @@ -389,7 +387,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict[str, object] | None = None, + request_data: dict | None = None, ) -> ModelResponse: """ Process output response by applying guardrails to text content. @@ -440,7 +438,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # create a standalone dict (SDK / direct-call path). if request_data is None: request_data = {"response": response} - request_data.setdefault("response", response) + else: + if "response" not in request_data: + request_data["response"] = response self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) @@ -460,7 +460,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) + guardrailed_texts: Final = guardrailed_inputs.get("texts", []) returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") guardrailed_tool_calls: Final[list[dict[str, object]]] = ( cast(list[dict[str, object]], returned_tool_calls) @@ -494,7 +494,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict[str, object] | None = None, + request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: @@ -550,14 +550,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict[str, object] | None, + request_data: dict | None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """Block-only streaming path: run the guardrail so an in-flight BLOCK can terminate the stream. Text rewrites are not propagated to the client here (see ``_process_streaming_transform`` for the incremental_diff path) unless ``deliver_ended_stream_rewrites`` opts the ended-stream branch in.""" - has_stream_ended: Final = deliver_ended_stream_rewrites or self._first_choice_has_finished(responses_so_far) + has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) if has_stream_ended: await self._process_ended_stream( @@ -623,7 +623,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) + guardrailed_texts: Final = guardrailed_inputs.get("texts", []) # Step 4: Apply guardrailed text back to all streaming chunks # For each choice, replace the combined text across all chunks @@ -656,36 +656,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): model_response: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) pre_guardrail_texts: Final = self._string_choice_contents(model_response) pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response) - inspection_responses: Final = ( - tuple( - model_response.model_copy(update=MappingProxyType({"choices": [choice]})) - for choice in model_response.choices - ) - if pre_guardrail_tool_calls and len(model_response.choices) > 1 - else (model_response,) + await self.process_output_response( + response=model_response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, ) - inspection_request_data: Final = request_data if request_data is not None else {} - try: - for inspection_response in inspection_responses: - inspection_request_data["response"] = inspection_response - inspection_request_data["responses"] = ( - [ - self._narrowed_to_choice(chunk, inspection_response.choices[0].index) - for chunk in responses_so_far - ] - if len(inspection_response.choices) == 1 - else responses_so_far - ) - await self.process_output_response( - response=inspection_response, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=inspection_request_data, - ) - finally: - inspection_request_data["response"] = model_response - inspection_request_data["responses"] = responses_so_far if not deliver_ended_stream_rewrites: return await self._write_ended_stream_text_rewrites( @@ -786,7 +763,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict[str, object] | None, + request_data: dict | None, sink: StreamTransformSink, ) -> None: """Run the guardrail over the raw accumulated text and report the @@ -798,7 +775,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ raw_by_index: Final = self._accumulate_string_content_by_choice_index(responses_so_far) if not raw_by_index: - sink.record(MappingProxyType({}), MappingProxyType({})) + sink.mutated_text_per_choice = {} + sink.holdback_per_choice = {} return # Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits @@ -810,72 +788,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): indices: Final = sorted(raw_by_index.keys()) texts_to_check: Final = [raw_by_index[i] for i in indices] - resolved_request_data: Final[dict[str, object]] = ( - request_data if request_data is not None else {"responses": responses_so_far} - ) - if "responses" not in resolved_request_data: - resolved_request_data["responses"] = responses_so_far - self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) + if request_data is None: + request_data = {"responses": responses_so_far} + elif "responses" not in request_data: + request_data["responses"] = responses_so_far + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) - if self._streamed_tool_call_fingerprints(responses_so_far): - assembled: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) - resolved_request_data["response"] = assembled - if len(assembled.choices) > 1: - choice_rounds: Final = tuple( - ( - choice, - StreamTransformSink(), - [self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far], - ) - for choice in assembled.choices - ) - try: - for choice, choice_sink, choice_chunks in choice_rounds: - resolved_request_data["response"] = assembled.model_copy( - update=MappingProxyType({"choices": [choice]}) - ) - resolved_request_data["responses"] = choice_chunks - await self._process_streaming_transform( - responses_so_far=choice_chunks, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=resolved_request_data, - sink=choice_sink, - ) - finally: - resolved_request_data["response"] = assembled - resolved_request_data["responses"] = responses_so_far - merged_texts: Final = MappingProxyType( - dict( - chain.from_iterable( - choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds - ) - ) - ) - merged_holdback: Final = MappingProxyType( - dict( - chain.from_iterable( - choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds - ) - ) - ) - sink.record(merged_texts, merged_holdback) - return - tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) - inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( - tuple( - MappingProxyType({"index": index, **converted}) - for index, tool_call in enumerate(tool_calls) - if (converted := self._convert_tool_call_to_dict(tool_call)) is not None - ) - ) if responses_so_far and getattr(responses_so_far[0], "model", None): inputs["model"] = responses_so_far[0].model guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, - request_data=resolved_request_data, + request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, ) @@ -895,14 +819,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): len(texts_to_check), ) - holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or () - transformed_texts: Final = MappingProxyType( - {idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)} - ) - transformed_holdback: Final = MappingProxyType( - {indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)} - ) - sink.record(transformed_texts, transformed_holdback) + holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or [] + sink.mutated_text_per_choice = { + idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts) + } + sink.holdback_per_choice = { + indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) + } def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream)) @@ -1222,18 +1145,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def _function_tool_call_fragments( responses_so_far: Sequence["ModelResponseStream"], ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]: + """Group the stream's function tool-call fragments by their tool-call index, in + the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping + only the indices the builder keeps (an id and a name somewhere in the stream).""" fragments: Final = tuple( - (choice.index, tool_call) + tool_call for response in responses_so_far for choice in response.choices for tool_call in choice.delta.tool_calls or () if isinstance(tool_call, ChatCompletionDeltaToolCall) ) - identified: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.id) - named: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.function.name) + identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id) + named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name) return tuple( - tuple(fragment for choice, fragment in fragments if (choice, fragment.index) == key) - for key in sorted(identified & named) + tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named) ) def _write_ended_stream_tool_call_rewrites( @@ -1243,10 +1168,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation): pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...], guardrail_name: str, ) -> None: + """Write ended-stream guardrail tool-call rewrites back across the buffered + chunks: the rewritten name and full arguments land in the tool call's first + fragment and the arguments of its later fragments are blanked, mirroring the + text write-back. A rewrite on a stream carrying more than one distinct choice + index, or whose fragments do not line up with the rebuilt tool calls, is + reported as undeliverable, so the pipeline executor discards it and releases + the original chunks.""" post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response) if post_guardrail_tool_calls == pre_guardrail_tool_calls: return + stream_choice_indices: Final = frozenset( + choice.index for response in responses_so_far for choice in response.choices + ) fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) + if len(stream_choice_indices) != 1: + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite( + guardrail_name, + f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written " + "back on single-choice streams", + ) if len(fragments_by_tool_call) != len(post_guardrail_tool_calls): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite @@ -1269,7 +1212,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], - guardrailed_texts: Sequence[str], + guardrailed_texts: list[str], task_mappings: list[tuple[int, int | None]], ) -> None: """ @@ -1287,9 +1230,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Override this method to customize how responses are applied to streaming responses. """ # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx) - guardrail_map: Final = MappingProxyType( - {task_mappings[index]: text for index, text in enumerate(guardrailed_texts)} - ) + guardrail_map: Final[dict[tuple[int, int | None], str]] = {} + for task_idx, guardrail_response in enumerate(guardrailed_texts): + mapping = task_mappings[task_idx] + choice_idx = cast(int, mapping[0]) + content_idx_optional = cast(int | None, mapping[1]) + guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response # Track which choices we've already set the guardrailed text for # Key: (choice_idx, content_idx), Value: boolean (True if already set) @@ -1299,20 +1245,34 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # choice by its index field: on n>1 streams a chunk usually carries one # choice at list position 0 whose index names the logical choice. for response in responses: - for choice, message in ( - (choice, choice.delta if isinstance(choice, StreamingChoices) else choice.message) - for choice in response.choices - if isinstance(choice, (StreamingChoices, Choices)) - ): - content = message.content + for choice in response.choices: + if isinstance(choice, litellm.StreamingChoices): + content = choice.delta.content + elif isinstance(choice, litellm.Choices): + content = choice.message.content + else: + continue + if content is None: continue if isinstance(content, str): + # String content str_key: tuple[int, int | None] = (choice.index, None) if str_key in guardrail_map: - message.content = guardrail_map[str_key] if str_key not in already_set else "" - already_set[str_key] = True + if str_key not in already_set: + # First chunk - set the complete guardrailed text + if isinstance(choice, litellm.StreamingChoices): + choice.delta.content = guardrail_map[str_key] + elif isinstance(choice, litellm.Choices): + choice.message.content = guardrail_map[str_key] + already_set[str_key] = True + else: + # Subsequent chunks - clear the content + if isinstance(choice, litellm.StreamingChoices): + choice.delta.content = "" + elif isinstance(choice, litellm.Choices): + choice.message.content = "" elif isinstance(content, list): # List content - handle each content item diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index f2c553b52e7..94a9d7c2bdf 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import BaseModel, TypeAdapter, ValidationError +from pydantic import BaseModel, ValidationError from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH @@ -2318,26 +2318,6 @@ async def test_custom_code_guardrail( ) -_GUARDRAIL_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object]) - - -def _metadata_fields(value: object) -> Mapping[str, object]: - try: - return _GUARDRAIL_METADATA_ADAPTER.validate_python(value) - except ValidationError: - return {} - - -def _guardrail_request_metadata(caller: object, proxy: object) -> Mapping[str, object]: - caller_fields: Final = tuple( - (key, value) for key, value in _metadata_fields(caller).items() if not key.startswith("user_api_key_") - ) - key_identity: Final = tuple( - (key, value) for key, value in _metadata_fields(proxy).items() if key.startswith("user_api_key_") - ) - return dict((*caller_fields, *key_identity)) - - def _execution_timeout_response(timeout: float) -> TestCustomCodeGuardrailResponse: return TestCustomCodeGuardrailResponse( success=False, @@ -2517,14 +2497,9 @@ async def apply_guardrail( if litellm_logging_obj is not None: _patch_logging_obj_for_guardrail(litellm_logging_obj, request) - metadata: Final = _guardrail_request_metadata(request.metadata, data.get("metadata")) request_data: Final[dict] = { - key: value - for key, value in ( - ("messages", request.messages), - ("metadata", metadata if request.metadata is not None or metadata else None), - ) - if value is not None + **({"messages": request.messages} if request.messages is not None else {}), + **({"metadata": request.metadata} if request.metadata is not None else {}), } _input_type: Final = _resolve_guardrail_input_type(active_guardrail, request.input_type) guardrailed_inputs: Final = await active_guardrail.apply_guardrail( diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 9134a1cd0f3..a7f813d953e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime import anyio from fastapi import HTTPException -from pydantic import InstanceOf, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -50,7 +49,6 @@ if TYPE_CHECKING: A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME: Final = "unified_llm_guardrails" -GUARDRAIL_ADAPTER: Final[TypeAdapter[CustomGuardrail | None]] = TypeAdapter(InstanceOf[CustomGuardrail] | None) _RequestData: TypeAlias = dict[str, object] @@ -515,9 +513,9 @@ class UnifiedLLMGuardrails(CustomLogger): self, *, reference_chunk: object, - mutated_text_per_choice: Mapping[int, str], + mutated_text_per_choice: dict[int, str], emitted_text_per_choice: dict[int, str], - holdback_per_choice: Mapping[int, int], + holdback_per_choice: dict[int, int], finish_reason_per_choice: dict[int, str | None], held_chars_per_choice: dict[int, int], is_final: bool, @@ -753,13 +751,46 @@ class UnifiedLLMGuardrails(CustomLogger): saw_tool_calls = False saw_text_content = False + tool_calls_released = False # rebind-ok: set once a raw tool call reaches the client unscanned + end_of_stream_inspection_started = False # rebind-ok: set once the end-of-stream inspection owns the verdict try: async for item in response: + # v1 transforms only text. A chunk carrying tool_calls is passed + # through raw so function-calling turns are not dropped, but ONLY + # its tool-call fields are forwarded: content is stripped so any + # response text (in the same delta, or in another choice of an n>1 + # chunk) can never bypass the transform. The original chunk is kept + # in responses_so_far so its text is still accumulated + redacted + + # emitted as synthetic deltas, and so the guardrail inspects the + # assembled tool calls at end of stream (see the block inspection + # below), matching block_only. finish_reason rides on the raw + # tool-only chunk, so it is not recorded for the text flush. if self._chunk_has_tool_calls(item): saw_tool_calls = True responses_so_far.append(item) last_chunk = item + # Fix #3 — flush accumulated text BEFORE the tool-call + # passthrough. Without this, a stream of text chunks that + # hasn't yet hit a sampled round can be trailed by a + # tool-call chunk carrying finish_reason="tool_calls"; an + # SSE-compliant client stops reading at that finish_reason + # and drops the end-of-stream text flush that would follow. + if saw_text_content: + async for out in _round(item, is_final=False): + yield out + # Fix #1 — pass finish_reason_per_choice into the + # passthrough so a mixed content+tool_call chunk defers its + # finish_reason to the final text terminator (see the + # _tool_call_passthrough_chunk docstring). + tool_only = self._tool_call_passthrough_chunk( + item, + finish_reason_per_choice=finish_reason_per_choice, + held_choices=_held_choices(held_chars_per_choice), + ) + responses_yielded.append(tool_only) + tool_calls_released = True + yield tool_only continue if self._is_trailing_metadata_chunk(item): @@ -777,62 +808,39 @@ class UnifiedLLMGuardrails(CustomLogger): # sampled round here would guardrail the same content twice. if ( not end_of_stream_only - and not saw_tool_calls and not self._chunk_has_finish_reason(item) and chunk_counter % sampling_rate == 0 ): async for out in _round(item, is_final=False): yield out + # v1 does not transform streamed tool calls, but they must still go + # through the guardrail's block decision. Run the block_only inspection + # over the full assembled response so tool calls cannot bypass it. + # + # Pass a deep copy of responses_so_far — the block path routes through + # ``_process_streaming_block_only`` which mutates ``delta.content`` + # in-place on the chunk objects it receives. For an n>1 chunk carrying + # text on one choice and tool_calls (with finish_reason) on another, + # ``has_stream_ended`` reads ``choices[0]`` alone and can miss the + # terminal signal, letting the block path rewrite the raw accumulator. + # The subsequent final ``_round`` would then re-read the already-mutated + # text, producing double-application for a non-idempotent guardrail or a + # ``stream_transform_underflow`` 400 from mismatched prefixes. A shallow + # list copy wouldn't help — the mutation is on the chunk objects + # themselves — so we deepcopy. + end_of_stream_inspection_started = True if saw_tool_calls: - inspected_responses: Final = copy.deepcopy(responses_so_far) - async for out in self._inspect_full_response( + async for out in self._inspect_full_response_for_block( endpoint_translation=endpoint_translation, guardrail_to_apply=guardrail_to_apply, request_data=request_data, user_api_key_dict=user_api_key_dict, - responses_so_far=inspected_responses, + responses_so_far=copy.deepcopy(responses_so_far), responses_yielded=responses_yielded, ): yield out - if saw_text_content: - async for out in _round(last_chunk, is_final=False): - yield out - tool_chunks: Final = tuple( - self._tool_call_passthrough_chunk( - buffered_item, - finish_reason_per_choice=finish_reason_per_choice, - held_choices=_held_choices(held_chars_per_choice), - ) - for buffered_item in inspected_responses - if self._chunk_has_tool_calls(buffered_item) - ) - - async def checked_tail() -> AsyncGenerator[object, None]: - try: - async for tail_chunk in self._emit_stream_tail( - last_chunk=last_chunk, - final_round=_round, - responses_so_far=responses_so_far, - responses_yielded=responses_yielded, - ): - yield tail_chunk - except _StreamTerminated as exc: - yield exc - - tail_chunks: Final = tuple([chunk async for chunk in checked_tail()]) - if tail_chunks and isinstance(tail_chunks[-1], _StreamTerminated): - for error_chunk in tail_chunks[:-1]: - yield error_chunk - return - for tool_only in tool_chunks: - responses_yielded.append(tool_only) - yield tool_only - for tail_chunk in tail_chunks: - yield tail_chunk - return - async for out in self._emit_stream_tail( last_chunk=last_chunk, final_round=_round, @@ -842,6 +850,37 @@ class UnifiedLLMGuardrails(CustomLogger): yield out except _StreamTerminated: return + except (GeneratorExit, asyncio.CancelledError): + await self._scan_uninspected_tool_calls_after_disconnect( + uninspected=tool_calls_released and not end_of_stream_inspection_started and not terminated.is_set(), + endpoint_translation=endpoint_translation, + responses_released=responses_yielded, + guardrail_to_apply=guardrail_to_apply, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + raise + + @staticmethod + async def _scan_uninspected_tool_calls_after_disconnect( + *, + uninspected: bool, + endpoint_translation: _EndpointTranslation, + responses_released: Sequence[object], + guardrail_to_apply: CustomGuardrail, + user_api_key_dict: UserAPIKeyAuth, + request_data: _RequestData, + ) -> None: + if not uninspected: + return + await UnifiedLLMGuardrails._scan_released_stream_after_disconnect( + endpoint_translation=endpoint_translation, + responses_released=responses_released, + last_scan_key=None, + guardrail_to_apply=guardrail_to_apply, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) async def _emit_stream_tail( self, @@ -860,21 +899,29 @@ class UnifiedLLMGuardrails(CustomLogger): responses_yielded.append(trailing) yield trailing - async def _inspect_full_response( + async def _inspect_full_response_for_block( self, *, endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, - request_data: _RequestData, + request_data: dict, user_api_key_dict: UserAPIKeyAuth, responses_so_far: Sequence[object], responses_yielded: Sequence[object], ) -> AsyncGenerator[object, None]: + """Run the block-only guardrail inspection over the full assembled + response (text + tool calls) so nothing bypasses the block decision. + + The guardrail's returned transforms are discarded here (v1 does not + transform tool calls); only its block decision matters. A block is + surfaced the same way as elsewhere: ModifyResponseException terminates the + stream via the shared block handler; a GenericGuardrailAPI block raises and + propagates, matching block_only. + """ from litellm.integrations.custom_guardrail import ModifyResponseException - from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite try: - with anyio.CancelScope(shield=True): + with anyio.CancelScope(shield=bool(responses_yielded)): await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, guardrail_to_apply=guardrail_to_apply, @@ -882,10 +929,7 @@ class UnifiedLLMGuardrails(CustomLogger): user_api_key_dict=user_api_key_dict, request_data=request_data, stream_transform_sink=None, - deliver_ended_stream_rewrites=True, ) - except UndeliverableStreamRewrite as exc: - raise HTTPException(status_code=400, detail="Guardrail stream rewrite could not be applied") from exc except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far @@ -1069,15 +1113,12 @@ class UnifiedLLMGuardrails(CustomLogger): # litellm.integrations.custom_guardrail. from litellm.integrations.custom_guardrail import ModifyResponseException - resolved_guardrail: Final[CustomGuardrail | None] = ( - guardrail_to_apply - if guardrail_to_apply is not None - else GUARDRAIL_ADAPTER.validate_python(request_data.pop("guardrail_to_apply", None)) - ) + if guardrail_to_apply is None: + guardrail_to_apply = request_data.pop("guardrail_to_apply", None) typed_request_data: Final[_RequestData] = request_data def _streaming_flag(name: str, default: object) -> Any: - return self.resolve_streaming_flag(resolved_guardrail, name, default) + return self.resolve_streaming_flag(guardrail_to_apply, name, default) sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). @@ -1100,30 +1141,30 @@ class UnifiedLLMGuardrails(CustomLogger): if ( buffer_until_moderated - and resolved_guardrail is not None - and getattr(resolved_guardrail, "mask_response_content", False) + and guardrail_to_apply is not None + and getattr(guardrail_to_apply, "mask_response_content", False) ): verbose_proxy_logger.warning( "UnifiedLLMGuardrails: streaming_buffer_until_moderated is disabled for %s " "because mask_response_content=True -- buffered replay would release " "unredacted original chunks instead of the moderated output.", - resolved_guardrail.guardrail_name, + guardrail_to_apply.guardrail_name, ) buffer_until_moderated = False if buffer_until_moderated and not release_on_scan: end_of_stream_only = True - if resolved_guardrail is None: + if guardrail_to_apply is None: async for item in response: yield item return event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.post_call - if resolved_guardrail.should_run_guardrail(data=request_data, event_type=event_type) is not True: + if guardrail_to_apply.should_run_guardrail(data=request_data, event_type=event_type) is not True: verbose_proxy_logger.debug( "UnifiedLLMGuardrails: Post-call streaming scanning disabled for %s", - resolved_guardrail.guardrail_name, + guardrail_to_apply.guardrail_name, ) async for item in response: yield item @@ -1143,7 +1184,7 @@ class UnifiedLLMGuardrails(CustomLogger): if transform_call_type is not None: async with contextlib.aclosing( self._run_incremental_transform_stream( - guardrail_to_apply=resolved_guardrail, + guardrail_to_apply=guardrail_to_apply, response=response, request_data=typed_request_data, user_api_key_dict=user_api_key_dict, @@ -1160,7 +1201,7 @@ class UnifiedLLMGuardrails(CustomLogger): "UnifiedLLMGuardrails: streaming_transform_mode=incremental_diff is only supported " "for the OpenAI chat completions streaming path with a resolvable request route; " "falling back to block_only for %s", - getattr(resolved_guardrail, "guardrail_name", None), + getattr(guardrail_to_apply, "guardrail_name", None), ) # Infer call type from first chunk @@ -1232,7 +1273,7 @@ class UnifiedLLMGuardrails(CustomLogger): verbose_proxy_logger.debug( "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round", chunk_counter, - resolved_guardrail.guardrail_name, + guardrail_to_apply.guardrail_name, ) if buffer_until_moderated: if hold_window: @@ -1252,7 +1293,7 @@ class UnifiedLLMGuardrails(CustomLogger): "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, sampling_rate, - resolved_guardrail.guardrail_name, + guardrail_to_apply.guardrail_name, ) original_items = ( @@ -1262,7 +1303,7 @@ class UnifiedLLMGuardrails(CustomLogger): try: await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, - guardrail_to_apply=resolved_guardrail, + guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=request_data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, request_data=request_data, @@ -1308,7 +1349,7 @@ class UnifiedLLMGuardrails(CustomLogger): verbose_proxy_logger.debug( "Holding %s buffered chunks for guardrail %s: this round could not scan the whole window", len(withheld_items), - resolved_guardrail.guardrail_name, + guardrail_to_apply.guardrail_name, ) withheld_items[:] = original_items continue @@ -1328,7 +1369,7 @@ class UnifiedLLMGuardrails(CustomLogger): verbose_proxy_logger.debug( "Processing final streaming response with all %s chunks for guardrail %s", len(responses_so_far), - resolved_guardrail.guardrail_name, + guardrail_to_apply.guardrail_name, ) endpoint_translation = mappings[CallTypes(call_type)]() @@ -1345,7 +1386,7 @@ class UnifiedLLMGuardrails(CustomLogger): if _is_redundant_scan(end_scan_key, last_scan_key): verbose_proxy_logger.debug( "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all", - resolved_guardrail.guardrail_name, + guardrail_to_apply.guardrail_name, ) for buffered_item in buffered_items or (): yield buffered_item @@ -1358,7 +1399,7 @@ class UnifiedLLMGuardrails(CustomLogger): with anyio.CancelScope(shield=chunks_yielded): await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, - guardrail_to_apply=resolved_guardrail, + guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=request_data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, request_data=request_data, @@ -1402,13 +1443,13 @@ class UnifiedLLMGuardrails(CustomLogger): chunks_yielded and not verdict_settled and translation_class is not None - and isinstance(resolved_guardrail, CustomGuardrail) + and isinstance(guardrail_to_apply, CustomGuardrail) ): await self._scan_released_stream_after_disconnect( endpoint_translation=translation_class(), responses_released=responses_yielded, last_scan_key=last_scan_key, - guardrail_to_apply=resolved_guardrail, + guardrail_to_apply=guardrail_to_apply, user_api_key_dict=user_api_key_dict, request_data=typed_request_data, ) diff --git a/tests/unit/interactions/test_openapi_compliance.py b/tests/unit/interactions/test_openapi_compliance.py index a0387f18db2..247d02298aa 100644 --- a/tests/unit/interactions/test_openapi_compliance.py +++ b/tests/unit/interactions/test_openapi_compliance.py @@ -10,18 +10,11 @@ Run with: pytest tests/unit/interactions/test_openapi_compliance.py -v import json import os import re -from collections.abc import Mapping -from typing import Any, Dict, Final +from typing import Any, Dict from unittest.mock import MagicMock, patch import httpx import pytest -from jsonschema import Draft202012Validator -from pydantic import TypeAdapter - -from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig -from litellm.types.interactions import InteractionInput -from litellm.types.router import GenericLiteLLMParams from openapi_core import OpenAPI OPENAPI_SPEC_URL = "https://ai.google.dev/static/api/interactions.openapi.json" @@ -45,6 +38,25 @@ def _load_openapi_spec_dict() -> Dict[str, Any]: ) +def _model_create_request_schema(spec_dict: Dict[str, Any]) -> Dict[str, Any]: + schemas = spec_dict["components"]["schemas"] + create_path = next(path for path in spec_dict["paths"] if path.endswith("/interactions")) + body_schema = spec_dict["paths"][create_path]["post"]["requestBody"]["content"]["application/json"]["schema"] + variants = [schemas[option["$ref"].split("/")[-1]] for option in body_schema.get("oneOf", []) if "$ref" in option] + return next(variant for variant in variants if "model" in variant.get("properties", {})) + + +def _interaction_resource_path(spec_dict: Dict[str, Any], method: str) -> str | None: + return next( + ( + path + for path, methods in spec_dict["paths"].items() + if re.search(r"/interactions/\{[^}]+\}$", path) and method in methods + ), + None, + ) + + def _declared_type_value(variant_schema: Dict[str, Any]) -> Any: """The single `type` value a union variant pins, whether spelled as a const or a 1-item enum.""" type_property = variant_schema.get("properties", {}).get("type", {}) @@ -64,58 +76,59 @@ def openapi_spec(spec_dict: Dict[str, Any]) -> OpenAPI: return OpenAPI.from_dict(spec_dict) -@pytest.fixture(scope="module") -def model_request_schema(spec_dict: Mapping[str, object]) -> Mapping[str, object]: - objects: Final = TypeAdapter(Mapping[str, Mapping[str, object]]) - components: Final = TypeAdapter(Mapping[str, object]).validate_python(spec_dict["components"]) - schemas: Final = objects.validate_python(components["schemas"]) - paths: Final = objects.validate_python(spec_dict["paths"]) - operation: Final = next(methods["post"] for path, methods in paths.items() if path.endswith("/interactions")) - post: Final = TypeAdapter(Mapping[str, object]).validate_python(operation) - body: Final = TypeAdapter(Mapping[str, object]).validate_python(post["requestBody"]) - content: Final = objects.validate_python(body["content"]) - schema: Final = TypeAdapter(Mapping[str, object]).validate_python(content["application/json"]["schema"]) - variants: Final = TypeAdapter(tuple[Mapping[str, str], ...]).validate_python(schema["oneOf"]) - return next( - schemas[variant["$ref"].rsplit("/", 1)[-1]] - for variant in variants - if "model" in TypeAdapter(Mapping[str, object]).validate_python( - schemas[variant["$ref"].rsplit("/", 1)[-1]]["properties"] - ) - ) - - class TestRequestCompliance: - def test_create_model_interaction_request_schema(self, model_request_schema: Mapping[str, object]) -> None: - config: Final = GoogleAIStudioInteractionsConfig() - payload: Final = config.transform_request( - model="test-model", agent=None, input="test input", - optional_params={"stream": True, "response_mime_type": "application/json"}, - litellm_params=GenericLiteLLMParams(api_key="test-key"), headers={}, - ) - properties: Final = TypeAdapter(Mapping[str, object]).validate_python(model_request_schema["properties"]) - assert payload == { - "model": "test-model", "input": "test input", "stream": True, - "response_format": {"type": "text", "mime_type": "application/json"}, - } - assert payload.keys() <= properties.keys() - assert set(config.get_supported_params("test-model")) - {"agent", "response_mime_type"} <= properties.keys() + """Tests that our request bodies match the OpenAPI spec.""" - @pytest.mark.parametrize("input_value", ["test input", [{"type": "text", "text": "test input"}]]) - def test_input_types_match_spec( - self, spec_dict: Mapping[str, object], model_request_schema: Mapping[str, object], - input_value: InteractionInput, - ) -> None: - payload: Final = GoogleAIStudioInteractionsConfig().transform_request( - model="test-model", agent=None, input=input_value, optional_params={}, - litellm_params=GenericLiteLLMParams(api_key="test-key"), headers={}, - ) - properties: Final = TypeAdapter(Mapping[str, Mapping[str, object]]).validate_python( - model_request_schema["properties"] - ) - validator: Final = Draft202012Validator(spec_dict).evolve(schema=properties["input"]) - assert not tuple(validator.iter_errors(payload["input"])) - assert payload["input"] == input_value + def test_create_model_interaction_request_schema(self, spec_dict): + schema = _model_create_request_schema(spec_dict) + + assert "model" in schema["required"] + assert "input" in schema["properties"] + + # Check our supported optional fields exist in spec + our_optional_fields = [ + "tools", + "system_instruction", + "generation_config", + "stream", + "store", + "background", + "response_modalities", + "response_format", + "response_mime_type", + "previous_interaction_id", + ] + + spec_properties = schema["properties"] + for field in our_optional_fields: + assert field in spec_properties, f"Field '{field}' not in OpenAPI spec" + print(f"✓ Field '{field}' exists in spec") + + def test_input_types_match_spec(self, spec_dict): + """Verify input field supports string, Content, Content[], Turn[].""" + schema = _model_create_request_schema(spec_dict) + input_schema = schema["properties"]["input"] + + # The input property may be inline oneOf or a $ref to InteractionsInput + if "$ref" in input_schema: + ref_name = input_schema["$ref"].split("/")[-1] + input_schema = spec_dict["components"]["schemas"][ref_name] + + # Should be oneOf with multiple types + assert "oneOf" in input_schema + + input_types = [] + for option in input_schema["oneOf"]: + if option.get("type") == "string": + input_types.append("string") + elif option.get("type") == "array": + input_types.append("array") + elif "$ref" in option: + input_types.append(option["$ref"]) + + print(f"Input supports types: {input_types}") + assert "string" in input_types, "Input should support string" + assert "array" in input_types, "Input should support array" def test_content_variants_are_identified_by_their_type_field(self, spec_dict): """Verify a Content part can be told apart by its `type`, however the spec spells that. @@ -312,24 +325,19 @@ class TestEndpointCompliance: assert create_path is not None, "POST /interactions endpoint not found" print(f"✓ Create endpoint: POST {create_path}") - @pytest.mark.parametrize("method", ["get", "delete"]) - def test_interaction_item_url_matches_spec(self, spec_dict: Mapping[str, object], method: str) -> None: - config: Final = GoogleAIStudioInteractionsConfig() - transform: Final = ( - config.transform_get_interaction_request if method == "get" else config.transform_delete_interaction_request - ) - url, body = transform( - interaction_id="test-interaction", api_base="https://example.com", - litellm_params=GenericLiteLLMParams(api_key="test-key"), headers={}, - ) - path: Final = httpx.URL(url).path - paths: Final = TypeAdapter(Mapping[str, Mapping[str, object]]).validate_python(spec_dict["paths"]) - assert any( - re.fullmatch(re.sub(r"\{[^}]+\}", "[^/]+", template), path) and method in operations - for template, operations in paths.items() - ), path - assert path.endswith("/test-interaction") - assert body == {} + def test_get_endpoint_exists(self, spec_dict): + """Verify GET /interactions/{id} endpoint exists.""" + get_path = _interaction_resource_path(spec_dict, "get") + + assert get_path is not None, "GET /interactions/{id} endpoint not found" + print(f"✓ Get endpoint: GET {get_path}") + + def test_delete_endpoint_exists(self, spec_dict): + """Verify DELETE /interactions/{id} endpoint exists.""" + delete_path = _interaction_resource_path(spec_dict, "delete") + + assert delete_path is not None, "DELETE /interactions/{id} endpoint not found" + print(f"✓ Delete endpoint: DELETE {delete_path}") if __name__ == "__main__": diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index bffe6e4c123..5d71cbe9afc 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -7,7 +7,7 @@ with guardrail transformations, including tool calls. import json from collections.abc import Mapping -from typing import Any, Final, Literal, Optional +from typing import Any, Literal, Optional import pytest @@ -1395,65 +1395,29 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert [ - [(tool_call["id"], tool_call["function"]["arguments"]) for tool_call in inputs["tool_calls"]] - for inputs in guardrail.seen_inputs - ] == [[("call_1", '{"fruit": "persimmon"}')], [("call_2", '{"fruit": "durian"}')]] + (tool_call["id"], tool_call["function"]["arguments"]) + for tool_call in guardrail.seen_inputs[-1]["tool_calls"] + ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')] @pytest.mark.asyncio - @pytest.mark.parametrize("transform", [False, True]) - async def test_each_choice_refreshes_shared_guardrail_response_context(self, transform: bool) -> None: - from fastapi import HTTPException - from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink - from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - class ContextGuardrail(CustomGuardrail): - async def apply_guardrail( - self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], - input_type: Literal["request", "response"], logging_obj: object = None, - ) -> GenericGuardrailAPIInputs: - response: Final = request_data["response"] - assert isinstance(response, ModelResponse) - request_data["visited_choices"] = (*request_data.get("visited_choices", ()), response.choices[0].index) - if response.choices[0].message.content == "forbidden": - raise HTTPException(status_code=400, detail="later choice blocked") - return inputs + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_tool_call_stream_chunks() - chunks: Final = [ModelResponseStream(choices=[StreamingChoices( - index=index, delta=Delta(content=text, tool_calls=[{ - "index": 0, "id": f"call_{index}", "type": "function", - "function": {"name": "contact", "arguments": "{}"}, - }]), finish_reason="tool_calls", - ) for index, text in enumerate(("allowed", "forbidden"))])] - request_data: Final = {"metadata": {"trace": "retained"}} - with pytest.raises(HTTPException, match="later choice blocked"): - await OpenAIChatCompletionsHandler().process_output_streaming_response( - responses_so_far=chunks, guardrail_to_apply=ContextGuardrail(guardrail_name="context"), - request_data=request_data, deliver_ended_stream_rewrites=True, - stream_transform_sink=StreamTransformSink() if transform else None, + with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, ) - assert request_data["visited_choices"] == (0, 1) - assert request_data["metadata"]["trace"] == "retained" - assert request_data["responses"] is chunks - restored: Final = request_data["response"] - assert isinstance(restored, ModelResponse) - assert tuple(choice.index for choice in restored.choices) == (0, 1) - @pytest.mark.asyncio - async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: - handler: Final = OpenAIChatCompletionsHandler() - chunks: Final = self._two_choice_tool_call_stream_chunks() - await handler.process_output_streaming_response( - responses_so_far=chunks, - guardrail_to_apply=MockGuardrail(guardrail_name="test"), - litellm_logging_obj=None, - deliver_ended_stream_rewrites=True, + assert raised.value.guardrail_name == "test" + assert raised.value.reason == ( + "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams" ) - arguments: Final = tuple( - "".join(call.function.arguments for chunk in chunks for choice in chunk.choices - if choice.index == index for call in choice.delta.tool_calls or ()) - for index in range(2) - ) - assert arguments == ('{"fruit": "PERSIMMON"}', '{"fruit": "DURIAN"}') @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 4dd2f2cb6e3..3efbe6475b2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -8,8 +8,6 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Literal -from pydantic import TypeAdapter - import anyio import pytest @@ -52,7 +50,6 @@ from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ( CallTypes, CallTypesLiteral, - ChatCompletionMessageToolCall, Delta, GenericGuardrailAPIInputs, ModelResponse, @@ -1011,42 +1008,6 @@ def _delta_text(item): return item.choices[0].delta.content or "" -class _ToolRedactingGuardrail(CustomGuardrail): - def __init__(self, require_tool_context: bool = False) -> None: - super().__init__(guardrail_name="tool-redactor", event_hook=GuardrailEventHooks.post_call, default_on=True) - self.streaming_transform_mode = "incremental_diff" - self.streaming_sampling_rate = 1 - self.require_tool_context = require_tool_context - - async def apply_guardrail( - self, - inputs: GenericGuardrailAPIInputs, - request_data: dict[str, object], - input_type: Literal["request", "response"], - logging_obj: object | None = None, - ) -> GenericGuardrailAPIInputs: - calls: Final = TypeAdapter(tuple[ChatCompletionMessageToolCall, ...]).validate_python( - inputs.get("tool_calls", ()) - ) - input_texts: Final = inputs.get("texts", ()) - texts: Final = tuple( - "checked:" + text.replace("SECRET", "MASKED") - if not self.require_tool_context or (calls and index == len(input_texts) - 1) else text - for index, text in enumerate(input_texts) - ) - return { - **inputs, - "texts": list(texts), - "stream_holdback_chars": [len(text) for text in texts], - "tool_calls": [ - call.model_copy(update={"function": call.function.model_copy(update={ - "arguments": call.function.arguments.replace("SECRET", "MASKED"), - })}) - for call in calls - ], - } - - class TestStreamingTransform: """Streaming text-transformation (incremental_diff) path on the OpenAI chat completions streaming surface.""" @@ -1055,103 +1016,6 @@ class TestStreamingTransform: def _use_openai_handler_mapping(self, monkeypatch): _patch_translation_mappings(monkeypatch, {CallTypes.acompletion: OpenAIChatCompletionsHandler}) - @pytest.mark.asyncio - @pytest.mark.parametrize("require_tool_context", [False, True]) - @pytest.mark.parametrize("include_text", [False, True]) - @pytest.mark.parametrize("tool_count", [1, 2]) - async def test_buffered_tool_arguments_are_rewritten_before_delivery( - self, include_text: bool, tool_count: int, require_tool_context: bool - ) -> None: - chunks: Final = ( - *([_stream_chunk("hello SECRET")] if include_text else []), - ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[ - {"index": index, "id": f"call_{index}", "type": "function", - "function": {"name": "contact", "arguments": '{"contact":"SEC'}} - for index in range(tool_count) - ]))]), - ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[ - {"index": index, "function": {"arguments": 'RET"}'}} for index in range(tool_count) - ]), finish_reason="tool_calls")]), - ModelResponseStream(choices=[], usage={"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}), - ) - out: Final = await _drive_stream( - UnifiedLLMGuardrails(), _ToolRedactingGuardrail(require_tool_context=require_tool_context), chunks - ) - calls: Final = tuple( - call for chunk in out for choice in chunk.choices for call in choice.delta.tool_calls or () - ) - for index in range(tool_count): - arguments: Final = "".join(call.function.arguments or "" for call in calls if call.index == index) - assert arguments == '{"contact":"MASKED"}' - assert next(call.id for call in calls if call.index == index and call.id) == f"call_{index}" - assert "".join(_delta_text(chunk) for chunk in out) == ("checked:hello MASKED" if include_text else "") - assert any(choice.finish_reason == "tool_calls" for chunk in out for choice in chunk.choices) - assert out[-1].usage.total_tokens == 18 - assert all("SECRET" not in chunk.model_dump_json() for chunk in out) - - @pytest.mark.asyncio - @pytest.mark.parametrize("tool_choice_index", [0, 1]) - async def test_tool_dependent_text_rewrites_keep_choice_context(self, tool_choice_index: int) -> None: - chunks: Final = ( - ModelResponseStream(choices=[StreamingChoices( - index=index, delta=Delta(content="SECRET" if index == tool_choice_index else "plain"), - ) for index in (1, 0)]), - ModelResponseStream(choices=[StreamingChoices( - index=tool_choice_index, - delta=Delta(tool_calls=[{ - "index": 0, "id": "call_context", "type": "function", - "function": {"name": "contact", "arguments": '{"contact":"SECRET"}'}, - }]), finish_reason="tool_calls", - )]), - ) - out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(True), chunks) - for index in (0, 1): - text: Final = "".join( - choice.delta.content or "" for chunk in out for choice in chunk.choices if choice.index == index - ) - assert text == ("checked:MASKED" if index == tool_choice_index else "plain") - assert all("SECRET" not in chunk.model_dump_json() for chunk in out) - - @pytest.mark.asyncio - async def test_tool_rewrites_keep_completion_choices_separate(self) -> None: - async def response() -> AsyncIterator[ModelResponseStream]: - yield ModelResponseStream(choices=[StreamingChoices( - index=index, - delta=Delta(tool_calls=[{"index": 0, "id": f"call_{index}", "type": "function", - "function": {"name": "contact", "arguments": '{"contact":"SECRET"}'}}]), - finish_reason="tool_calls", - ) for index in range(2)]) - - iterator: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( - user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), - response=response(), request_data={"guardrail_to_apply": _ToolRedactingGuardrail()}, - ) - chunks: Final = tuple([chunk async for chunk in iterator]) - calls: Final = tuple( - (choice.index, call.id, call.function.arguments) - for chunk in chunks for choice in chunk.choices for call in choice.delta.tool_calls or () - ) - assert calls == ((0, "call_0", '{"contact":"MASKED"}'), (1, "call_1", '{"contact":"MASKED"}')) - - @pytest.mark.asyncio - async def test_undeliverable_rewrite_is_a_closed_failure(self) -> None: - from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - - class UndeliverableTranslation(OpenAIChatCompletionsHandler): - async def process_output_streaming_response( - self, *args: object, **kwargs: object - ) -> list[ModelResponseStream]: - raise UndeliverableStreamRewrite("tool-redactor", "unmappable tool fragments") - - iterator: Final = UnifiedLLMGuardrails()._inspect_full_response( - endpoint_translation=UndeliverableTranslation(), guardrail_to_apply=_ToolRedactingGuardrail(), - request_data={}, user_api_key_dict=UserAPIKeyAuth(), responses_so_far=(), responses_yielded=(), - ) - with pytest.raises(unified_module.HTTPException) as error: - await anext(iterator) - assert error.value.status_code == 400 - assert error.value.detail == "Guardrail stream rewrite could not be applied" - @pytest.mark.asyncio async def test_block_only_drops_text_rewrites(self): """Default block_only: the guardrail's uppercasing never reaches the @@ -1572,15 +1436,14 @@ class TestStreamingTransform: assert out[-2].choices[0].finish_reason == "stop" @pytest.mark.asyncio - @pytest.mark.parametrize(("content", "allowed_scans"), [(None, 0), ("proposal", 0), ("proposal", 1)]) - async def test_tool_call_blocking_guardrail_is_enforced(self, content: str | None, allowed_scans: int): + async def test_tool_call_blocking_guardrail_is_enforced(self): """A guardrail that blocks on tool calls must terminate the incremental_diff stream: tool calls go through the block decision, not bypass it.""" from litellm.exceptions import GuardrailRaisedException class _ToolCallBlocker(_StreamingTextGuardrail): async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): - if input_type == "response" and self.response_calls >= allowed_scans: + if input_type == "response" and inputs.get("tool_calls"): raise GuardrailRaisedException( guardrail_name="tc-block", message="blocked tool call", @@ -1593,7 +1456,7 @@ class TestStreamingTransform: StreamingChoices( index=0, delta=Delta( - content=content, + content=None, tool_calls=[ { "index": 0, @@ -1608,21 +1471,8 @@ class TestStreamingTransform: ], ) - async def upstream(): - yield tool_chunk - - stream: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"), - response=upstream(), - request_data={"guardrail_to_apply": _ToolCallBlocker(), "model": "gpt-4"}, - ) - async def consume_checked_stream() -> None: - async for chunk in stream: - assert isinstance(chunk, ModelResponseStream) - assert all(not choice.delta.tool_calls and choice.finish_reason is None for choice in chunk.choices) - with pytest.raises(GuardrailRaisedException): - await consume_checked_stream() + await _drive_stream(UnifiedLLMGuardrails(), _ToolCallBlocker(), [tool_chunk]) @pytest.mark.asyncio async def test_mixed_content_and_tool_call_chunk_does_not_leak_text(self): @@ -2840,7 +2690,7 @@ class TestStreamingClientDisconnectScan: yield _stream_chunk(" tail", finish_reason="stop") @pytest.mark.asyncio - async def test_incremental_diff_block_withholds_tool_calls_and_does_not_inspect_again(self): + async def test_closing_while_an_incremental_diff_block_is_delivered_does_not_inspect_the_tool_call_again(self): guardrail = _MarkerBlockingStreamingTextGuardrail() async def upstream() -> AsyncIterator[ModelResponseStream]: @@ -2854,16 +2704,12 @@ class TestStreamingClientDisconnectScan: released = [await stream.__anext__(), await stream.__anext__()] await stream.aclose() - assert all(isinstance(item, bytes) for item in released) - assert b"get_weather" not in b"".join(released) - assert b"Paris" not in b"".join(released) - assert guardrail.received_texts == [["BLOCKME tail"]], guardrail.received_texts - assert [[call["function"]["name"] for call in calls] for calls in guardrail.received_tool_calls] == [ - ["get_weather"] - ] + assert [call.function.name for call in released[0].choices[0].delta.tool_calls] == ["get_weather"] + assert guardrail.received_texts == [["BLOCKME"]], guardrail.received_texts + assert guardrail.received_tool_calls == [], guardrail.received_tool_calls @pytest.mark.asyncio - async def test_closing_after_a_buffered_tool_call_inspects_held_text_before_delivery(self): + async def test_closing_after_a_released_tool_call_under_incremental_diff_scans_no_held_back_text(self): guardrail = _StreamingTextGuardrail(holdback_schedule=[len("held secret")] * 2) async def upstream() -> AsyncIterator[ModelResponseStream]: @@ -2872,14 +2718,12 @@ class TestStreamingClientDisconnectScan: yield tool_chunk stream = self._guarded_stream(guardrail, upstream()) - received: Final = tuple([await stream.__anext__() for _ in range(2)]) - scans_before_close: Final = len(guardrail.received_texts) + received = await stream.__anext__() await stream.aclose() - assert [call.function.name for call in received[-1].choices[0].delta.tool_calls] == ["get_weather"] + assert [call.function.name for call in received.choices[0].delta.tool_calls] == ["get_weather"] assert guardrail.received_tool_calls, guardrail.received_texts - assert guardrail.received_texts[0] == ["held secret"] - assert len(guardrail.received_texts) == scans_before_close + assert all("held secret" not in text for text in guardrail.received_texts[-1]), guardrail.received_texts def _responses_delta(sequence_number, text): diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py index 3d5087d6d1e..13049093d6c 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py @@ -2127,7 +2127,7 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker): } -def _patch_apply_guardrail_env(mocker, guardrail_result, processed_data=None): +def _patch_apply_guardrail_env(mocker, guardrail_result): mock_guardrail = mocker.Mock() mock_guardrail.apply_guardrail = AsyncMock(return_value=guardrail_result) @@ -2140,7 +2140,7 @@ def _patch_apply_guardrail_env(mocker, guardrail_result, processed_data=None): mock_logging_obj.model_call_details = {} mock_processor = mocker.Mock() mock_processor.common_processing_pre_call_logic = AsyncMock( - return_value=(processed_data or {"guardrail_name": "test-guardrail"}, mock_logging_obj) + return_value=({"guardrail_name": "test-guardrail"}, mock_logging_obj) ) mocker.patch( "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", @@ -2211,68 +2211,6 @@ async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker): ) -@pytest.mark.asyncio -async def test_apply_guardrail_replaces_caller_identity_with_the_authenticated_key(mocker): - """Identity fields come from the proxy's own sanitized metadata, never from - the caller, so a request cannot impersonate another key or probe its policy.""" - mock_guardrail = _patch_apply_guardrail_env( - mocker, - {"texts": ["ok"]}, - processed_data={ - "guardrail_name": "test-guardrail", - "metadata": { - "route": "/apply_guardrail", - "user_api_key_alias": "billing-app", - "user_api_key_user_id": "u-1", - }, - }, - ) - - request = ApplyGuardrailRequest( - guardrail_name="test-guardrail", - text="hello", - metadata={ - "forbidden_topics": ["tax"], - "user_api_key_alias": "someone-else", - "user_api_key_user_email": "victim@example.com", - }, - ) - await apply_guardrail( - fastapi_request=mocker.Mock(), - request=request, - user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app", user_id="u-1"), - ) - - forwarded = mock_guardrail.apply_guardrail.await_args.kwargs["request_data"]["metadata"] - assert forwarded == { - "forbidden_topics": ["tax"], - "user_api_key_alias": "billing-app", - "user_api_key_user_id": "u-1", - } - - -@pytest.mark.asyncio -async def test_apply_guardrail_attaches_key_identity_when_caller_sends_no_metadata(mocker): - """A caller that sends no metadata still gets the authenticated identity - forwarded, the same shape every LLM route gives a guardrail.""" - mock_guardrail = _patch_apply_guardrail_env( - mocker, - {"texts": ["ok"]}, - processed_data={"guardrail_name": "test-guardrail", "metadata": {"user_api_key_alias": "billing-app"}}, - ) - - request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello") - await apply_guardrail( - fastapi_request=mocker.Mock(), - request=request, - user_api_key_dict=UserAPIKeyAuth(key_alias="billing-app"), - ) - - assert mock_guardrail.apply_guardrail.await_args.kwargs["request_data"] == { - "metadata": {"user_api_key_alias": "billing-app"} - } - - @pytest.mark.asyncio async def test_apply_guardrail_omits_metadata_when_not_sent(mocker): """Without metadata, request_data stays empty (backward-compatible)."""