diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index 2386f80e819..63ead52baa6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Optional from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -8,9 +8,23 @@ from litellm.types.guardrails import Guardrail, LitellmParams +def _get_config_value(litellm_params: Any, optional_params: Any, attribute_name: str) -> Optional[Any]: + if optional_params is not None: + value = ( + optional_params.get(attribute_name) + if isinstance(optional_params, dict) + else getattr(optional_params, attribute_name, None) + ) + if value is not None: + return value + return getattr(litellm_params, attribute_name, None) + + def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): import litellm + optional_params = getattr(litellm_params, "optional_params", None) + _generic_guardrail_api_callback = GenericGuardrailAPI( api_base=litellm_params.api_base, api_key=litellm_params.api_key, @@ -22,6 +36,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, default_on=litellm_params.default_on, + streaming_end_of_stream_only=_get_config_value(litellm_params, optional_params, "streaming_end_of_stream_only"), + streaming_sampling_rate=_get_config_value(litellm_params, optional_params, "streaming_sampling_rate"), ) litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index df80ea09de0..dc519f56d1a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -33,6 +33,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel GUARDRAIL_NAME = "generic_guardrail_api" @@ -178,6 +179,8 @@ def __init__( unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", fail_on_error: Optional[bool] = True, extra_headers: Optional[list] = None, + streaming_end_of_stream_only: Optional[bool] = None, + streaming_sampling_rate: Optional[int] = None, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -209,6 +212,15 @@ def __init__( self.fail_on_error: bool = True if fail_on_error is None else fail_on_error + # Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook + # via getattr(guardrail_to_apply, "streaming_*", default). + self.streaming_end_of_stream_only: bool = ( + False if streaming_end_of_stream_only is None else streaming_end_of_stream_only + ) + if streaming_sampling_rate is not None and streaming_sampling_rate < 1: + raise ValueError(f"streaming_sampling_rate must be >= 1 (got {streaming_sampling_rate})") + self.streaming_sampling_rate: int = 5 if streaming_sampling_rate is None else streaming_sampling_rate + # Set supported event hooks if "supported_event_hooks" not in kwargs: kwargs["supported_event_hooks"] = [ @@ -470,3 +482,11 @@ async def apply_guardrail( return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj) except Exception as e: return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj, is_unreachable=False) + + @staticmethod + def get_config_model() -> Optional[type["GuardrailConfigModel"]]: + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIConfigModel, + ) + + return GenericGuardrailAPIConfigModel diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 28fb482b3af..d0ac8bb8998 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -1,7 +1,7 @@ from typing import Any, Dict, List, Literal, Optional, Union from pydantic import BaseModel, ConfigDict, Field -from typing_extensions import TYPE_CHECKING, TypedDict +from typing_extensions import TypedDict from litellm.types.llms.openai import ( AllMessageValues, @@ -60,6 +60,30 @@ class GenericGuardrailAPIOptionalParams(BaseModel): ), ) + streaming_end_of_stream_only: Optional[bool] = Field( + default=None, + description=( + "If False (default when unset), the guardrail runs on sampled chunks during " + "the stream at the cadence set by streaming_sampling_rate, and an in-flight " + "BLOCKED stops further chunks from streaming. If True, the guardrail runs " + "once at end of stream over the assembled response; lower cost and latency, " + "but flagged content has already streamed to the client before the terminal " + "block. Defaults are applied in GenericGuardrailAPI.__init__ when None so " + "unset optional_params does not shadow top-level litellm_params." + ), + ) + + streaming_sampling_rate: Optional[int] = Field( + default=None, + ge=1, + description=( + "When streaming_end_of_stream_only is False, the guardrail runs every Nth " + "streamed chunk. Ignored when streaming_end_of_stream_only is True. " + "Must be >= 1 when set. Defaults to 5 in GenericGuardrailAPI.__init__ " + "when None so unset optional_params does not shadow top-level litellm_params." + ), + ) + class GenericGuardrailAPIConfigModel( GuardrailConfigModel[GenericGuardrailAPIOptionalParams], diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 399442a5f71..791fdd4077c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -609,7 +609,6 @@ async def test_images_passed_in_request( request_data=mock_request_data_input, input_type="request", ) - result_texts = guardrailed_inputs.get("texts", []) result_images = guardrailed_inputs.get("images", None) # Verify API was called with images @@ -943,7 +942,7 @@ async def test_multimodal_message_serialization(self): guardrail.async_handler, "post", return_value=mock_response ) as mock_post: # This should not raise SerializationIterator error - result = await guardrail.apply_guardrail( + await guardrail.apply_guardrail( inputs={ "texts": ["What's in this image?"], "images": ["https://example.com/image.jpg"], @@ -1006,7 +1005,7 @@ def content_generator(): with patch.object( guardrail.async_handler, "post", return_value=mock_response ) as mock_post: - result = await guardrail.apply_guardrail( + await guardrail.apply_guardrail( inputs={ "texts": ["Hello", "World"], "structured_messages": messages_with_iterable, @@ -1023,6 +1022,717 @@ def content_generator(): assert isinstance(json_payload["structured_messages"], list) +def _make_stream_chunk(content: str, finish_reason=None): + """Build a real ModelResponseStream so the handler's isinstance checks pass.""" + from litellm.types.utils import Delta, ModelResponseStream + + return ModelResponseStream( + model="gpt-4", + choices=[ + litellm.StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + ) + ], + ) + + +def _make_assembled_model_response(content: str) -> ModelResponse: + return ModelResponse( + id="mock-response", + model="gpt-4", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message(role="assistant", content=content), + finish_reason="stop", + ) + ], + ) + + +def _mock_guardrail_post_response(action: str = "NONE", texts=None, blocked_reason=None): + mock_response = MagicMock() + payload = {"action": action} + if texts is not None: + payload["texts"] = texts + if blocked_reason is not None: + payload["blocked_reason"] = blocked_reason + mock_response.json.return_value = payload + mock_response.raise_for_status = MagicMock() + return mock_response + + +def _make_responses_stream_events(text: str): + """Minimal /v1/responses SSE event sequence ending in response.completed.""" + return ( + {"type": "response.created", "response": {"id": "resp_test"}}, + { + "type": "response.output_item.added", + "item": {"type": "message", "id": "msg_test"}, + }, + { + "type": "response.content_part.added", + "part": {"type": "output_text", "text": ""}, + }, + {"type": "response.output_text.delta", "delta": text}, + { + "type": "response.output_text.done", + "text": text, + }, + { + "type": "response.completed", + "response": { + "id": "resp_test", + "output": [ + { + "type": "message", + "id": "msg_test", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text}], + } + ], + "status": "completed", + }, + }, + ) + + +class TestGenericGuardrailAPIStreamingConfig: + """Streaming knobs on GenericGuardrailAPI and initialize_guardrail plumbing.""" + + def test_streaming_defaults(self): + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + ) + assert guardrail.streaming_end_of_stream_only is False + assert guardrail.streaming_sampling_rate == 5 + + def test_streaming_overrides(self): + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + streaming_end_of_stream_only=True, + streaming_sampling_rate=2, + ) + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 2 + + @pytest.mark.parametrize("invalid_rate", [0, -1, -5]) + def test_streaming_sampling_rate_rejects_non_positive(self, invalid_rate): + with pytest.raises(ValueError, match="streaming_sampling_rate must be >= 1"): + GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + streaming_sampling_rate=invalid_rate, + ) + + def test_optional_params_streaming_sampling_rate_ge_one(self): + from pydantic import ValidationError + + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIOptionalParams, + ) + + with pytest.raises(ValidationError): + GenericGuardrailAPIOptionalParams(streaming_sampling_rate=0) + + def test_get_config_model(self): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIConfigModel, + ) + + assert GenericGuardrailAPI.get_config_model() is GenericGuardrailAPIConfigModel + + def test_initialize_guardrail_forwards_streaming_flags(self): + from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="generic_guardrail_api", + mode="post_call", + api_base="https://api.test.guardrail.com", + default_on=False, + ) + # LitellmParams uses extra="allow" on the base; set streaming knobs dynamically + litellm_params.streaming_end_of_stream_only = False # type: ignore[attr-defined] + litellm_params.streaming_sampling_rate = 3 # type: ignore[attr-defined] + + guardrail_config = {"guardrail_name": "test-generic-streaming"} + + with patch( + "litellm.logging_callback_manager.add_litellm_callback" + ): + guardrail = initialize_guardrail(litellm_params, guardrail_config) + + assert guardrail.streaming_end_of_stream_only is False + assert guardrail.streaming_sampling_rate == 3 + + def test_initialize_guardrail_optional_params_defaults_do_not_shadow_top_level( + self, + ): + """Top-level streaming knobs win when optional_params only carries siblings.""" + from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIOptionalParams, + ) + + litellm_params = LitellmParams( + guardrail="generic_guardrail_api", + mode="post_call", + api_base="https://api.test.guardrail.com", + default_on=False, + ) + litellm_params.streaming_end_of_stream_only = True # type: ignore[attr-defined] + litellm_params.streaming_sampling_rate = 2 # type: ignore[attr-defined] + # Sibling optional_params only; streaming fields stay at Pydantic default None. + litellm_params.optional_params = GenericGuardrailAPIOptionalParams( # type: ignore[attr-defined] + additional_provider_specific_params={"tenant": "acme"}, + ) + + guardrail_config = {"guardrail_name": "test-generic-streaming-mixed"} + + with patch( + "litellm.logging_callback_manager.add_litellm_callback" + ): + guardrail = initialize_guardrail(litellm_params, guardrail_config) + + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 2 + + def test_initialize_guardrail_explicit_optional_params_streaming_wins(self): + from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIOptionalParams, + ) + + litellm_params = LitellmParams( + guardrail="generic_guardrail_api", + mode="post_call", + api_base="https://api.test.guardrail.com", + default_on=False, + ) + litellm_params.streaming_end_of_stream_only = False # type: ignore[attr-defined] + litellm_params.streaming_sampling_rate = 9 # type: ignore[attr-defined] + litellm_params.optional_params = GenericGuardrailAPIOptionalParams( # type: ignore[attr-defined] + streaming_end_of_stream_only=True, + streaming_sampling_rate=1, + ) + + guardrail_config = {"guardrail_name": "test-generic-streaming-nested-wins"} + + with patch( + "litellm.logging_callback_manager.add_litellm_callback" + ): + guardrail = initialize_guardrail(litellm_params, guardrail_config) + + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 1 + + def test_initialize_guardrail_dict_optional_params_streaming_wins(self): + """Guardrail API/UI delivers optional_params as a plain dict, not a model.""" + from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="generic_guardrail_api", + mode="post_call", + api_base="https://api.test.guardrail.com", + default_on=False, + ) + litellm_params.streaming_end_of_stream_only = False # type: ignore[attr-defined] + litellm_params.streaming_sampling_rate = 9 # type: ignore[attr-defined] + # Plain dict mirrors how configs arrive from the guardrail API/UI. + litellm_params.optional_params = { # type: ignore[attr-defined] + "streaming_end_of_stream_only": True, + "streaming_sampling_rate": 1, + } + + guardrail_config = {"guardrail_name": "test-generic-streaming-dict-optional"} + + with patch( + "litellm.logging_callback_manager.add_litellm_callback" + ): + guardrail = initialize_guardrail(litellm_params, guardrail_config) + + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 1 + + def test_initialize_guardrail_dict_optional_params_sibling_only_falls_through( + self, + ): + """Dict optional_params without streaming keys must not shadow top-level knobs.""" + from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="generic_guardrail_api", + mode="post_call", + api_base="https://api.test.guardrail.com", + default_on=False, + ) + litellm_params.streaming_end_of_stream_only = True # type: ignore[attr-defined] + litellm_params.streaming_sampling_rate = 2 # type: ignore[attr-defined] + litellm_params.optional_params = { # type: ignore[attr-defined] + "additional_provider_specific_params": {"tenant": "acme"}, + } + + guardrail_config = {"guardrail_name": "test-generic-streaming-dict-sibling"} + + with patch( + "litellm.logging_callback_manager.add_litellm_callback" + ): + guardrail = initialize_guardrail(litellm_params, guardrail_config) + + assert guardrail.streaming_end_of_stream_only is True + assert guardrail.streaming_sampling_rate == 2 + + +class TestGenericGuardrailAPIStreamingViaUnified: + """Streaming output checks routed through UnifiedLLMGuardrails.""" + + @pytest.mark.asyncio + async def test_streaming_safe_content_yields_all_chunks(self): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + ) + unified_guardrail = UnifiedLLMGuardrails() + + async def mock_stream(): + chunks_data = ["Hello", " ", "world", "!", " Goodbye"] + for i, content in enumerate(chunks_data): + yield _make_stream_chunk( + content, + finish_reason="stop" if i == len(chunks_data) - 1 else None, + ) + + mock_post = AsyncMock( + return_value=_mock_guardrail_post_response( + action="NONE", texts=["Hello world! Goodbye"] + ) + ) + + with ( + patch.object(guardrail.async_handler, "post", mock_post), + patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=_make_assembled_model_response("Hello world! Goodbye"), + ), + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-generic-guardrail"]}, + } + + chunks_received = 0 + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + chunks_received += 1 + + assert chunks_received == 5 + assert mock_post.await_count >= 1 + + @pytest.mark.asyncio + async def test_streaming_blocked_content_raises(self): + from litellm.exceptions import GuardrailRaisedException + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + streaming_sampling_rate=1, + ) + unified_guardrail = UnifiedLLMGuardrails() + + async def mock_stream(): + chunks_data = ["Hello", " ishaan", " here"] + for i, content in enumerate(chunks_data): + yield _make_stream_chunk( + content, + finish_reason="stop" if i == len(chunks_data) - 1 else None, + ) + + mock_post = AsyncMock( + return_value=_mock_guardrail_post_response( + action="BLOCKED", blocked_reason="Ishaan is not allowed" + ) + ) + + with ( + patch.object(guardrail.async_handler, "post", mock_post), + patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=_make_assembled_model_response("Hello ishaan here"), + ), + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-generic-guardrail"]}, + } + + with pytest.raises(GuardrailRaisedException) as exc_info: + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + assert "Ishaan is not allowed" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_streaming_default_uses_sampled_cadence(self): + """Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3.""" + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + ) + unified_guardrail = UnifiedLLMGuardrails() + + async def mock_stream(): + chunks_data = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"] + for i, content in enumerate(chunks_data): + yield _make_stream_chunk( + content, + finish_reason="stop" if i == len(chunks_data) - 1 else None, + ) + + mock_post = AsyncMock( + return_value=_mock_guardrail_post_response( + action="NONE", texts=["ABCDEFGHIJ"] + ) + ) + + with ( + patch.object(guardrail.async_handler, "post", mock_post), + patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=_make_assembled_model_response("ABCDEFGHIJ"), + ), + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-generic-guardrail"]}, + } + + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + assert mock_post.await_count == 3, ( + f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), " + f"got {mock_post.await_count}" + ) + for call in mock_post.await_args_list: + assert call.kwargs["json"]["input_type"] == "response" + + @pytest.mark.asyncio + async def test_streaming_end_of_stream_only_calls_guardrail_once(self): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + streaming_end_of_stream_only=True, + ) + unified_guardrail = UnifiedLLMGuardrails() + + async def mock_stream(): + chunks_data = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J"] + for i, content in enumerate(chunks_data): + yield _make_stream_chunk( + content, + finish_reason="stop" if i == len(chunks_data) - 1 else None, + ) + + mock_post = AsyncMock( + return_value=_mock_guardrail_post_response( + action="NONE", texts=["ABCDEFGHIJ"] + ) + ) + + with ( + patch.object(guardrail.async_handler, "post", mock_post), + patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=_make_assembled_model_response("ABCDEFGHIJ"), + ), + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-generic-guardrail"]}, + } + + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + assert mock_post.await_count == 1, ( + f"Expected exactly one guardrail call at end of stream, " + f"got {mock_post.await_count}" + ) + + @pytest.mark.asyncio + async def test_streaming_sampling_rate_override(self): + """sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls.""" + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + streaming_end_of_stream_only=False, + streaming_sampling_rate=2, + ) + unified_guardrail = UnifiedLLMGuardrails() + + async def mock_stream(): + chunks_data = ["A", "B", "C", "D", "E", "F"] + for i, content in enumerate(chunks_data): + yield _make_stream_chunk( + content, + finish_reason="stop" if i == len(chunks_data) - 1 else None, + ) + + mock_post = AsyncMock( + return_value=_mock_guardrail_post_response(action="NONE", texts=["ABCDEF"]) + ) + + with ( + patch.object(guardrail.async_handler, "post", mock_post), + patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=_make_assembled_model_response("ABCDEF"), + ), + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-generic-guardrail"]}, + } + + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + assert mock_post.await_count == 4, ( + f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), " + f"got {mock_post.await_count}" + ) + + @pytest.mark.asyncio + async def test_streaming_fail_open_on_unreachable_continues_stream(self): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + unreachable_fallback="fail_open", + streaming_end_of_stream_only=True, + ) + unified_guardrail = UnifiedLLMGuardrails() + + async def mock_stream(): + for i, content in enumerate(["A", "B", "C"]): + yield _make_stream_chunk( + content, finish_reason="stop" if i == 2 else None + ) + + mock_post = AsyncMock(side_effect=httpx.ConnectError("connection refused")) + + with ( + patch.object(guardrail.async_handler, "post", mock_post), + patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=_make_assembled_model_response("ABC"), + ), + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-generic-guardrail"]}, + } + + chunks_received = 0 + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + chunks_received += 1 + + assert chunks_received == 3 + + @pytest.mark.asyncio + async def test_responses_api_streaming_end_of_stream_only_calls_guardrail_once(self): + """/v1/responses path through unified hook; end-of-stream-only = one call.""" + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + streaming_end_of_stream_only=True, + ) + unified_guardrail = UnifiedLLMGuardrails() + + async def mock_responses_stream(): + for event in _make_responses_stream_events("Hello world"): + yield event + + mock_post = AsyncMock( + return_value=_mock_guardrail_post_response( + action="NONE", texts=["Hello world"] + ) + ) + + with patch.object(guardrail.async_handler, "post", mock_post): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/v1/responses" + ) + request_data = { + "input": "hi", + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-generic-guardrail"]}, + } + + events_received = 0 + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_responses_stream(), + request_data=request_data, + ): + events_received += 1 + + assert events_received == 6 + assert mock_post.await_count == 1, ( + f"Expected exactly one guardrail call at end of /v1/responses stream, " + f"got {mock_post.await_count}" + ) + assert mock_post.await_args.kwargs["json"]["input_type"] == "response" + + @pytest.mark.asyncio + async def test_responses_api_streaming_blocked_raises(self): + """Mid-stream BLOCKED on /v1/responses surfaces GuardrailRaisedException.""" + from litellm.exceptions import GuardrailRaisedException + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + streaming_sampling_rate=1, + ) + unified_guardrail = UnifiedLLMGuardrails() + + async def mock_responses_stream(): + for event in _make_responses_stream_events("blocked content"): + yield event + + mock_post = AsyncMock( + return_value=_mock_guardrail_post_response( + action="BLOCKED", blocked_reason="Responses content not allowed" + ) + ) + + with patch.object(guardrail.async_handler, "post", mock_post): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/v1/responses" + ) + request_data = { + "input": "hi", + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-generic-guardrail"]}, + } + + with pytest.raises(GuardrailRaisedException) as exc_info: + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_responses_stream(), + request_data=request_data, + ): + pass + + assert "Responses content not allowed" in str(exc_info.value) + class TestToolSupport: """Test tool handling in guardrail requests"""