feat(model_armor): stream post_call chunks as they are scanned

Add streaming_buffer_until_moderated so post_call Model Armor can forward chunks while it scans, instead of holding every token until generation finishes.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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()
 

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)
 

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 (

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"""

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()