diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index db6b474872f..7185e0e6fd6 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -124,10 +124,10 @@ def _is_text_delta_event(event: str) -> bool: def _text_delta_frame(text: str, index: object) -> bytes: - payload: Final = { + payload: Final = { # mutable-ok: json.dumps serializes a dict "type": "content_block_delta", "index": index if isinstance(index, int) else 0, - "delta": {"type": "text_delta", "text": text}, + "delta": {"type": "text_delta", "text": text}, # mutable-ok: json.dumps serializes a dict } return f"event: content_block_delta\ndata: {json.dumps(payload)}\n\n".encode() diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py index eda505e2453..7b9cf308ba6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -7,6 +9,8 @@ from .model_armor import ModelArmorGuardrail if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams +_NO_EXTRAS: Final[Mapping[str, object]] = MappingProxyType({}) # mutable-ok: MappingProxyType needs a dict + def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): import litellm @@ -14,6 +18,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" ModelArmorGuardrail, ) + streaming_params: Final[Mapping[str, object]] = litellm_params.model_extra or _NO_EXTRAS _model_armor_callback: Final = ModelArmorGuardrail( guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, @@ -28,6 +33,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" fail_on_error=litellm_params.fail_on_error, skip_unscannable_attachments=litellm_params.skip_unscannable_attachments, sanitize_error_detail=litellm_params.sanitize_error_detail, + streaming_buffer_until_moderated=streaming_params.get("streaming_buffer_until_moderated"), + streaming_sampling_rate=streaming_params.get("streaming_sampling_rate"), ) litellm.logging_callback_manager.add_litellm_callback(_model_armor_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 6de02c71e3e..de581857b5f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,10 +1,11 @@ from collections.abc import AsyncGenerator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal import httpx from fastapi import HTTPException if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel import json @@ -15,6 +16,7 @@ from litellm.caching import DualCache from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.custom_guardrail import ( CustomGuardrail, + ModifyResponseException, log_guardrail_information, ) from litellm.litellm_core_utils.core_helpers import ( @@ -37,6 +39,7 @@ from litellm.types.utils import ( CallTypes, CallTypesLiteral, Choices, + GenericGuardrailAPIInputs, GuardrailStatus, ModelResponse, ModelResponseStream, @@ -82,6 +85,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): - Post-call sanitization (sanitizeModelResponse) """ + # apply_guardrail only exists to scan streamed chunks: file scanning, masking and the MCP + # hooks need the native lifecycle hooks, so they must not be routed to the shared runner + use_native_lifecycle_hooks: ClassVar[bool] = True + @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: return [ @@ -123,6 +130,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): self.credentials = credentials self.api_endpoint = api_endpoint self.sanitize_error_detail = sanitize_error_detail is not False + self.streaming_buffer_until_moderated = kwargs.get("streaming_buffer_until_moderated") is not False + self.streaming_sampling_rate = kwargs.get("streaming_sampling_rate") or 5 # Store optional params self.optional_params = kwargs @@ -188,6 +197,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: super().update_in_memory_litellm_params(litellm_params) self.sanitize_error_detail = self.sanitize_error_detail is not False + self.streaming_buffer_until_moderated = self.streaming_buffer_until_moderated is not False + self.streaming_sampling_rate = self.streaming_sampling_rate or 5 def _log_request_debug( self, @@ -831,6 +842,60 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return response + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: CustomGuardrail.apply_guardrail signature + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + """Scan bare texts, which lets the shared runner drive Model Armor over a live stream.""" + texts: Final = inputs.get("texts") or () + scanned: Final = [ # mutable-ok: GenericGuardrailAPIInputs.texts is a list + await self._scan_text(text, input_type=input_type, request_data=request_data) if text else text + for text in texts + ] + return {**inputs, "texts": scanned} # mutable-ok: GenericGuardrailAPIInputs is a plain TypedDict + + async def _scan_text( + self, + text: str, + input_type: Literal["request", "response"], + request_data: dict, # mutable-ok: forwarded to make_model_armor_request + ) -> str: + masking_enabled: Final = self.mask_request_content if input_type == "request" else self.mask_response_content + try: + armor_response: Final = await self.make_model_armor_request( + content=text, + source="user_prompt" if input_type == "request" else "model_response", + request_data=request_data, + ) + except ModelArmorAPIError as e: + self._raise_if_fail_closed(e) + return text + + if self._should_block_content(armor_response, allow_sanitization=masking_enabled): + message: Final = f"{input_type.capitalize()} blocked by Model Armor" + if input_type == "request": + raise HTTPException( + status_code=400, + detail=self._build_block_error_detail(message, armor_response), + ) + raise ModifyResponseException( + message=message, + model=request_data.get("model") or "unknown", + request_data=request_data, + guardrail_name=self.guardrail_name or GUARDRAIL_NAME, + original_response=request_data.get("response"), + ) + + sanitized: Final = self._get_sanitized_content(armor_response) if masking_enabled else None + return sanitized or text + + def _streams_incrementally(self) -> bool: + """Sanitization needs the whole response, so only allow/block configs can scan mid-stream.""" + return not self.streaming_buffer_until_moderated and not self.mask_response_content + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -839,6 +904,21 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): ) -> AsyncGenerator[ModelResponseStream, None]: """Process streaming response chunks.""" + if self._streams_incrementally(): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + async for streamed_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=response, + request_data=request_data, + guardrail_to_apply=self, + buffer_until_moderated_default=False, + ): + yield streamed_chunk + return + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.main import stream_chunk_builder from litellm.proxy.guardrails.anthropic_sse import ( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py b/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py index ec7e1595215..5f0b59a6924 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/model_armor.py @@ -26,6 +26,26 @@ class ModelArmorGuardrailConfigModel(GuardrailConfigModel): ), ) + streaming_buffer_until_moderated: bool | None = Field( + default=True, + description=( + "True (default) withholds every streamed chunk until Model Armor has scanned the " + "assembled response, so nothing unscanned reaches the client but time to first token " + "grows to the full generation time. False streams chunks as they arrive and scans them " + "as the response grows, at the cadence set by streaming_sampling_rate, terminating the " + "stream when a scan blocks. Ignored when mask_response_content is True, since " + "sanitization needs the whole response." + ), + ) + streaming_sampling_rate: int | None = Field( + default=5, + ge=1, + description=( + "When streaming_buffer_until_moderated is False, scan every Nth streamed chunk. Lower " + "values catch violations sooner, at the cost of more Model Armor calls." + ), + ) + @staticmethod def ui_friendly_name() -> str: """Return the UI-friendly name for Model Armor guardrail""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 8f9966c6ecd..becc79797a6 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -790,6 +790,211 @@ async def test_streaming_hook_passes_through_responses_api_events(): assert tuple(delivered) == events +async def _iter_chunks(chunks: tuple[object, ...]): + for chunk in chunks: + yield chunk + + +async def _chunks_produced_before_first_delivery(guardrail: ModelArmorGuardrail) -> int: + """How much of the upstream stream is consumed before the client sees anything. + + 1 means the guardrail streams as it scans; the full chunk count means it buffers the + whole response first, which is what makes time to first token equal generation time. + """ + produced: list[object] = [] + + async def _stream(): + for chunk in _ANTHROPIC_SSE_CHUNKS: + produced.append(chunk) + yield chunk + + delivered = guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/messages"), + response=_stream(), + request_data={ + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "what is my ssn"}], + "metadata": {"guardrails": ["model-armor-test"]}, + }, + ) + try: + await delivered.__anext__() + finally: + await delivered.aclose() + return len(produced) + + +def test_config_disables_streaming_buffer(): + from litellm.proxy.guardrails.guardrail_hooks.model_armor import initialize_guardrail + from litellm.types.guardrails import LitellmParams + + guardrail = initialize_guardrail( + LitellmParams( + guardrail="model_armor", + mode="post_call", + template_id="test-template", + project_id="test-project", + streaming_buffer_until_moderated=False, + streaming_sampling_rate=2, + ), + {"guardrail_name": "model-armor-test"}, + ) + + assert guardrail._streams_incrementally() is True + assert guardrail.streaming_sampling_rate == 2 + + +def test_default_config_buffers_streams(): + from litellm.proxy.guardrails.guardrail_hooks.model_armor import initialize_guardrail + from litellm.types.guardrails import LitellmParams + + guardrail = initialize_guardrail( + LitellmParams( + guardrail="model_armor", + mode="post_call", + template_id="test-template", + project_id="test-project", + ), + {"guardrail_name": "model-armor-default"}, + ) + + assert guardrail._streams_incrementally() is False + assert guardrail.streaming_sampling_rate == 5 + + +@pytest.mark.asyncio +async def test_streaming_hook_buffers_whole_response_by_default(): + guardrail = _sse_armor_guardrail() + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_api_response({"filterMatchState": "NO_MATCH_FOUND"})), + ): + produced = await _chunks_produced_before_first_delivery(guardrail) + + assert produced == len(_ANTHROPIC_SSE_CHUNKS) + + +@pytest.mark.asyncio +async def test_streaming_hook_streams_while_scanning_when_buffering_disabled(): + guardrail = _sse_armor_guardrail( + streaming_buffer_until_moderated=False, + streaming_sampling_rate=1, + ) + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_api_response({"filterMatchState": "NO_MATCH_FOUND"})), + ): + produced = await _chunks_produced_before_first_delivery(guardrail) + + assert produced == 1 + + +@pytest.mark.asyncio +async def test_streaming_hook_keeps_buffering_when_masking_responses(): + """Sanitization needs the assembled response, so masking configs cannot stream as they scan.""" + guardrail = _sse_armor_guardrail( + mask_response_content=True, + streaming_buffer_until_moderated=False, + ) + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_api_response({"filterMatchState": "NO_MATCH_FOUND"})), + ): + produced = await _chunks_produced_before_first_delivery(guardrail) + + assert produced == len(_ANTHROPIC_SSE_CHUNKS) + + +@pytest.mark.asyncio +async def test_chat_completions_streams_while_scanning_when_buffering_disabled(): + guardrail = _sse_armor_guardrail( + streaming_buffer_until_moderated=False, + streaming_sampling_rate=1, + ) + produced: list[object] = [] + + async def _stream(): + for text in ("hello ", "world"): + chunk = litellm.ModelResponseStream( + choices=[ + litellm.types.utils.StreamingChoices(delta=litellm.types.utils.Delta(content=text)) + ] + ) + produced.append(chunk) + yield chunk + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_api_response({"filterMatchState": "NO_MATCH_FOUND"})), + ): + delivered = guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + response=_stream(), + request_data={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"guardrails": ["model-armor-test"]}, + }, + ) + first = await delivered.__anext__() + assert len(produced) == 1 + assert first.choices[0].delta.content == "hello " + await delivered.aclose() + + +@pytest.mark.asyncio +async def test_incremental_streaming_blocks_before_flagged_text_is_delivered(): + guardrail = _sse_armor_guardrail( + streaming_buffer_until_moderated=False, + streaming_sampling_rate=1, + ) + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock( + return_value=_armor_api_response( + { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "inspectResult": { + "matchState": "MATCH_FOUND", + "findings": [{"infoType": "US_SOCIAL_SECURITY_NUMBER"}], + } + } + } + }, + } + ) + ), + ): + delivered = [ + chunk + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/messages"), + response=_iter_chunks(_ANTHROPIC_SSE_CHUNKS), + request_data={ + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "what is my ssn"}], + "metadata": {"guardrails": ["model-armor-test"]}, + }, + ) + ] + + body = b"".join(chunk if isinstance(chunk, bytes) else str(chunk).encode() for chunk in delivered) + assert b"123-45-6789" not in body + assert b"Model Armor" in body + + @pytest.mark.asyncio async def test_streaming_hook_fails_closed_on_unparseable_raw_sse(): guardrail = _sse_armor_guardrail()