Merge 7a083353a8 into 1c61c2606e

litellm/integrations/custom_guardrail.py

@@ -30,6 +30,7 @@ from litellm.types.utils import (
     GuardrailStatus,
     GuardrailTracingDetail,
     LLMResponseTypes,
+    ModelResponse,
     StandardLoggingGuardrailInformation,
 )
 
@@ -883,7 +884,9 @@ class CustomGuardrail(CustomLogger):
         """logging_only: run apply_guardrail on copies of the logged request/response and record the verdict."""
         from litellm.llms import get_guardrail_translation_mapping
 
-        if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
+        if not self.uses_apply_guardrail_interface():
+            return kwargs, result
+        if not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
             return kwargs, result
         try:
             translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))()
@@ -900,8 +903,16 @@ class CustomGuardrail(CustomLogger):
             for key, value in (litellm_params.get("metadata") or {}).items()
             if key != "standard_logging_guardrail_information"
         }
+        response: Final = (
+            kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result
+        )
+        output_translation: Final = (
+            get_guardrail_translation_mapping(CallTypes.acompletion)()
+            if isinstance(response, ModelResponse)
+            else translation
+        )
         try:
-            await self._scan_logged_call(kwargs, result, translation, scratch_metadata)
+            await self._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata)
         except Exception as e:
             verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
         recorded: Final = scratch_metadata.get("standard_logging_guardrail_information")
@@ -918,8 +929,9 @@ class CustomGuardrail(CustomLogger):
     async def _scan_logged_call(
         self,
         kwargs: dict,  # mutable-ok: CustomLogger.async_logging_hook contract
-        result: object,
+        response: object | None,
         translation: "BaseTranslation",
+        output_translation: "BaseTranslation",
         scratch_metadata: dict,  # mutable-ok: apply_guardrail records its verdict into request metadata
     ) -> None:
         optional_params: Final = kwargs.get("optional_params") or {}
@@ -933,8 +945,10 @@ class CustomGuardrail(CustomLogger):
             "metadata": scratch_metadata,
         }
         await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
-        await translation.process_output_response(
-            response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request
+        if response is None:
+            return
+        await output_translation.process_output_response(
+            response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
         )
 
     def supports_scan_only_tool_results(self) -> bool:

litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py

@@ -1,11 +1,13 @@
+import time
 from collections.abc import AsyncGenerator, Mapping, Sequence
 from enum import Enum, auto
-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
@@ -52,6 +54,7 @@ from litellm.types.utils import (
     CallTypes,
     CallTypesLiteral,
     Choices,
+    GenericGuardrailAPIInputs,
     GuardrailStatus,
     ModelResponse,
     ModelResponseStream,
@@ -118,8 +121,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
     Supports:
     - Pre-call sanitization (sanitizeUserPrompt)
     - Post-call sanitization (sanitizeModelResponse)
+    - logging_only: scans the completed response after it reaches the client and
+      records the verdict in spend logs without blocking
     """
 
+    use_native_lifecycle_hooks: ClassVar[bool] = True
+
     @classmethod
     def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
         return [
@@ -128,6 +135,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
             GuardrailEventHooks.post_call,
             GuardrailEventHooks.pre_mcp_call,
             GuardrailEventHooks.during_mcp_call,
+            GuardrailEventHooks.logging_only,
         ]
 
     def __init__(
@@ -1096,6 +1104,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
             add_guardrail_to_applied_guardrails_header,
         )
 
+        if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True:
+            async for chunk in response:
+                yield chunk
+            return
+
         all_chunks: Final[Sequence[object]] = tuple([chunk async for chunk in response])
 
         if not all_chunks or self._is_terminal_error_stream(all_chunks):
@@ -1213,6 +1226,60 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         for chunk in all_chunks:
             yield chunk
 
+    @log_guardrail_information
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: "LiteLLMLoggingObj | None" = None,
+    ) -> GenericGuardrailAPIInputs:
+        content: Final = "\n".join(text for text in inputs.get("texts") or () if text)
+        if not content:
+            return inputs
+
+        source: Final[Literal["user_prompt", "model_response"]] = (
+            "user_prompt" if input_type == "request" else "model_response"
+        )
+        start_time: Final = time.time()
+        try:
+            armor_response: Final = await self.make_model_armor_request(
+                content=content, source=source, request_data=request_data
+            )
+        except (ModelArmorAPIError, httpx.HTTPError) as e:
+            error_end_time: Final = time.time()
+            self.add_standard_logging_guardrail_information_to_request_data(
+                guardrail_json_response=str(e),
+                request_data=request_data,
+                guardrail_status="guardrail_failed_to_respond",
+                guardrail_provider="model_armor",
+                start_time=start_time,
+                end_time=error_end_time,
+                duration=error_end_time - start_time,
+            )
+            return inputs
+
+        flagged: Final = self._should_block_content(armor_response, allow_sanitization=False)
+        end_time: Final = time.time()
+        self.add_standard_logging_guardrail_information_to_request_data(
+            guardrail_json_response=self._build_logging_response(armor_response),
+            request_data=request_data,
+            guardrail_status="guardrail_flagged" if flagged else "success",
+            guardrail_provider="model_armor",
+            start_time=start_time,
+            end_time=end_time,
+            duration=end_time - start_time,
+        )
+        if flagged and not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
+            raise HTTPException(
+                status_code=400,
+                detail=self._build_block_error_detail(
+                    "Response blocked by Model Armor" if input_type == "response" else "Content blocked by Model Armor",
+                    armor_response,
+                ),
+            )
+        return inputs
+
     @staticmethod
     def get_config_model() -> type["GuardrailConfigModel"] | None:
         """

tests/test_litellm/integrations/test_custom_guardrail.py

@@ -2625,7 +2625,7 @@ class TestLoggingOnlyApplyGuardrail:
         assert [e["guardrail_status"] for e in entries] == ["success"]
 
     @pytest.mark.asyncio
-    async def test_native_lifecycle_hook_guardrail_is_left_alone(self):
+    async def test_native_lifecycle_hook_guardrail_scans_in_logging_only(self):
         class _NativeHooks(_ApplyOnlyObserver):
             use_native_lifecycle_hooks = True
 
@@ -2634,9 +2634,9 @@ class TestLoggingOnlyApplyGuardrail:
 
         out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
 
-        assert guardrail.calls == []
-        assert out_kwargs is kwargs
+        assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])]
         assert out_response is response
+        assert out_kwargs["standard_logging_object"]["guardrail_information"]
 
     @pytest.mark.asyncio
     async def test_aresponses_scans_logged_messages_when_input_is_cleared(self):
@@ -2890,6 +2890,55 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook:
         assert len(_guardrail_entries(request_data)) == 1
 
 
+class _NativeLifecycleLoggingGuardrail(CustomGuardrail):
+    """Native lifecycle guardrail that also implements apply_guardrail, like the azure guards."""
+
+    use_native_lifecycle_hooks: ClassVar[bool] = True
+
+    def __init__(self):
+        from litellm.types.guardrails import GuardrailEventHooks
+
+        super().__init__(
+            guardrail_name="native-logging-guardrail",
+            event_hook=GuardrailEventHooks.logging_only,
+        )
+        self.calls: list = []
+
+    async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+        self.calls.append((input_type, list(inputs.get("texts") or [])))
+        return inputs
+
+
+@pytest.mark.asyncio
+async def test_native_lifecycle_guardrail_logging_only_scans_assembled_response():
+    """A use_native_lifecycle_hooks guardrail accepts mode logging_only and its
+    async_logging_hook scans kwargs["async_complete_streaming_response"], not the raw result."""
+    from litellm.types.utils import Choices, Message, ModelResponse
+
+    guardrail = _NativeLifecycleLoggingGuardrail()
+    assembled = ModelResponse(
+        choices=[Choices(message=Message(role="assistant", content="assembled stream text"))]
+    )
+    sentinel_result = object()
+    kwargs = {
+        "model": "gpt-5.4-mini",
+        "messages": [{"role": "user", "content": "hi"}],
+        "litellm_call_id": "call-1",
+        "litellm_params": {"metadata": {}},
+        "optional_params": {},
+        "standard_logging_object": {"guardrail_information": None},
+        "async_complete_streaming_response": assembled,
+    }
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=sentinel_result, call_type=CallTypes.acompletion.value
+    )
+
+    assert out_result is sentinel_result
+    assert ("response", ["assembled stream text"]) in guardrail.calls
+    assert out_kwargs["standard_logging_object"]["guardrail_information"]
+
+
 class TestPreCallHookResponseIsNotLoggedVerbatim:
     """Regression for LIT-6935: a pre_call hook returning the request payload leaked the prompt
     into ``guardrail_response`` and from there onto OTEL guardrail spans."""

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py

@@ -4929,3 +4929,356 @@ def test_every_responses_delta_event_is_in_the_scanned_set():
     }
     assert not missing
     assert "response.mcp_call_arguments.delta" in _RESPONSES_DELTA_EVENT_TYPES
+
+
+def _clean_armor_response() -> dict:
+    return {
+        "sanitizationResult": {
+            "filterMatchState": "NO_MATCH_FOUND",
+            "filterResults": {},
+        }
+    }
+
+
+def _flagged_armor_response() -> dict:
+    return {
+        "sanitizationResult": {
+            "filterMatchState": "MATCH_FOUND",
+            "filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}},
+        }
+    }
+
+
+def _logging_only_guardrail() -> ModelArmorGuardrail:
+    return ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-logging",
+        event_hook=GuardrailEventHooks.logging_only,
+    )
+
+
+def _logged_kwargs() -> dict:
+    return {
+        "model": "gpt-4o",
+        "messages": [{"role": "user", "content": "hi"}],
+        "litellm_call_id": "call-1",
+        "litellm_params": {"metadata": {}},
+        "optional_params": {},
+        "standard_logging_object": {"guardrail_information": None},
+    }
+
+
+def _chat_response(text: str) -> litellm.ModelResponse:
+    return litellm.ModelResponse(
+        choices=[
+            litellm.types.utils.Choices(
+                message=litellm.types.utils.Message(role="assistant", content=text)
+            )
+        ]
+    )
+
+
+def _stream_chunk(text: str) -> litellm.ModelResponseStream:
+    return litellm.ModelResponseStream(
+        choices=[
+            litellm.types.utils.StreamingChoices(
+                delta=litellm.types.utils.Delta(content=text)
+            )
+        ]
+    )
+
+
+def _metadata_entries(kwargs: dict) -> list:
+    return kwargs["standard_logging_object"].get("guardrail_information") or []
+
+
+def test_logging_only_mode_is_accepted_and_keeps_native_hooks():
+    guardrail = _logging_only_guardrail()
+    assert guardrail.event_hook == GuardrailEventHooks.logging_only
+    assert guardrail.use_native_lifecycle_hooks is True
+    assert GuardrailEventHooks.logging_only in ModelArmorGuardrail.get_supported_event_hooks()
+
+    post_call_guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-post",
+        event_hook=GuardrailEventHooks.post_call,
+    )
+    assert post_call_guardrail._deployment_hook_target() is post_call_guardrail
+
+
+@pytest.mark.asyncio
+async def test_logging_only_stream_yields_chunks_without_waiting_for_scan():
+    """A logging_only guardrail must pass stream chunks straight through; the scan happens
+    afterwards on the assembled response via async_logging_hook."""
+    guardrail = _logging_only_guardrail()
+    guardrail.make_model_armor_request = AsyncMock(
+        side_effect=AssertionError("logging_only must not scan the stream")
+    )
+
+    produced = 0
+
+    async def gen():
+        nonlocal produced
+        for i in range(3):
+            produced += 1
+            yield _stream_chunk(f"chunk-{i} ")
+
+    hook_iter = guardrail.async_post_call_streaming_iterator_hook(
+        user_api_key_dict=UserAPIKeyAuth(),
+        response=gen(),
+        request_data={"metadata": {}, "guardrails": ["model-armor-logging"]},
+    )
+    first = await hook_iter.__anext__()
+    assert produced == 1
+    chunks = [first]
+    async for chunk in hook_iter:
+        chunks.append(chunk)
+    assert len(chunks) == 3
+    guardrail.make_model_armor_request.assert_not_awaited()
+
+    guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response())
+    response = _chat_response("all clear")
+    kwargs = _logged_kwargs()
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+    assert out_result is response
+    entries = _metadata_entries(out_kwargs)
+    assert len(entries) >= 1
+    entry = entries[-1]
+    assert entry["guardrail_status"] == "success"
+    assert entry["guardrail_mode"] == "logging_only"
+    assert entry["guardrail_provider"] == "model_armor"
+
+
+@pytest.mark.asyncio
+async def test_logging_only_records_flagged_verdict_without_altering_response():
+    guardrail = _logging_only_guardrail()
+    guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response())
+    response = _chat_response("flagged output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    assert out_result is response
+    entries = _metadata_entries(out_kwargs)
+    assert entries[-1]["guardrail_status"] == "guardrail_flagged"
+    assert entries[-1]["guardrail_mode"] == "logging_only"
+
+
+@pytest.mark.asyncio
+async def test_logging_only_records_model_armor_api_error():
+    guardrail = _logging_only_guardrail()
+    guardrail.make_model_armor_request = AsyncMock(
+        side_effect=ModelArmorAPIError("Model Armor API error (upstream 500)")
+    )
+    response = _chat_response("some output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    assert out_result is response
+    entries = _metadata_entries(out_kwargs)
+    assert entries[-1]["guardrail_status"] == "guardrail_failed_to_respond"
+
+
+@pytest.mark.asyncio
+async def test_logging_only_scans_assembled_responses_api_stream():
+    """The terminal ResponseCompletedEvent is an envelope; the scan must run on the
+    assembled ResponsesAPIResponse kept in kwargs."""
+    from openai.types.responses import ResponseOutputMessage, ResponseOutputText
+
+    from litellm.types.llms.openai import (
+        ResponseCompletedEvent,
+        ResponsesAPIResponse,
+        ResponsesAPIStreamEvents,
+    )
+
+    assembled = ResponsesAPIResponse(
+        id="resp-1",
+        created_at=1700000000,
+        output=[
+            ResponseOutputMessage(
+                id="msg-1",
+                type="message",
+                role="assistant",
+                status="completed",
+                content=[
+                    ResponseOutputText(
+                        annotations=[], text="assembled output text", type="output_text"
+                    )
+                ],
+            )
+        ],
+    )
+    event = ResponseCompletedEvent(
+        type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=assembled
+    )
+
+    guardrail = _logging_only_guardrail()
+    guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response())
+    kwargs = _logged_kwargs()
+    del kwargs["messages"]
+    kwargs["input"] = "hello"
+    kwargs["async_complete_streaming_response"] = assembled
+
+    out_kwargs, _ = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=event, call_type="aresponses"
+    )
+
+    response_scans = [
+        call
+        for call in guardrail.make_model_armor_request.await_args_list
+        if call.kwargs.get("source") == "model_response"
+    ]
+    assert response_scans, "expected a model_response scan of the assembled response"
+    assert "assembled output text" in response_scans[0].kwargs["content"]
+    assert _metadata_entries(out_kwargs)
+
+
+@pytest.mark.asyncio
+async def test_logging_only_scans_anthropic_messages_model_response():
+    """/v1/messages logs a ModelResponse; the output scan must extract the assistant text."""
+    guardrail = _logging_only_guardrail()
+    guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response())
+    kwargs = _logged_kwargs()
+    kwargs["messages"] = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]
+    response = _chat_response("anthropic assembled text")
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="anthropic_messages"
+    )
+
+    assert out_result is response
+    response_scans = [
+        call
+        for call in guardrail.make_model_armor_request.await_args_list
+        if call.kwargs.get("source") == "model_response"
+    ]
+    assert response_scans
+    assert "anthropic assembled text" in response_scans[0].kwargs["content"]
+    assert _metadata_entries(out_kwargs)
+
+
+@pytest.mark.asyncio
+async def test_logging_only_skips_output_scan_when_no_assembled_response():
+    guardrail = _logging_only_guardrail()
+    guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response())
+    kwargs = _logged_kwargs()
+
+    await guardrail.async_logging_hook(kwargs=kwargs, result=None, call_type="acompletion")
+
+    sources = [call.kwargs.get("source") for call in guardrail.make_model_armor_request.await_args_list]
+    assert "model_response" not in sources
+
+
+@pytest.mark.asyncio
+async def test_native_post_call_mode_ignores_logging_hook():
+    guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-post",
+        event_hook=GuardrailEventHooks.post_call,
+    )
+    guardrail.make_model_armor_request = AsyncMock(return_value=_clean_armor_response())
+    response = _chat_response("some output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    assert out_kwargs is kwargs
+    assert out_result is response
+    guardrail.make_model_armor_request.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_records_flagged_without_raising():
+    guardrail = _logging_only_guardrail()
+    guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response())
+    request_data = {"metadata": {}}
+    inputs = {"texts": ["forbidden output"]}
+
+    result = await guardrail.apply_guardrail(
+        inputs=inputs,
+        request_data=request_data,
+        input_type="response",
+    )
+
+    assert result == inputs
+    entries = request_data["metadata"]["standard_logging_guardrail_information"]
+    assert entries[-1]["guardrail_status"] == "guardrail_flagged"
+
+
+@pytest.mark.asyncio
+async def test_logging_only_records_transport_error():
+    guardrail = _logging_only_guardrail()
+    guardrail.make_model_armor_request = AsyncMock(side_effect=httpx.ConnectError("boom"))
+    response = _chat_response("some output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, out_result = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    assert out_result is response
+    entries = _metadata_entries(out_kwargs)
+    failed = [e for e in entries if e["guardrail_status"] == "guardrail_failed_to_respond"]
+    assert failed
+    assert all(e["guardrail_provider"] == "model_armor" for e in failed)
+
+
+@pytest.mark.asyncio
+async def test_logging_only_flagged_prompt_still_scans_response():
+    """A flagged input scan must not abort the output scan; both verdicts are recorded."""
+    guardrail = _logging_only_guardrail()
+    guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response())
+    response = _chat_response("flagged output")
+    kwargs = _logged_kwargs()
+
+    out_kwargs, _ = await guardrail.async_logging_hook(
+        kwargs=kwargs, result=response, call_type="acompletion"
+    )
+
+    sources = [call.kwargs.get("source") for call in guardrail.make_model_armor_request.await_args_list]
+    assert sources == ["user_prompt", "model_response"]
+    entries = _metadata_entries(out_kwargs)
+    flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"]
+    assert len(flagged) == 2
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_raises_on_flagged_when_not_logging_only():
+    """The /guardrails/apply_guardrail endpoint calls apply_guardrail directly; a
+    non-logging_only instance must signal the block so flagged text is not returned as clean."""
+    guardrail = ModelArmorGuardrail(
+        template_id="test-template",
+        project_id="test-project",
+        location="us-central1",
+        guardrail_name="model-armor-pre",
+        event_hook=GuardrailEventHooks.pre_call,
+    )
+    guardrail.make_model_armor_request = AsyncMock(return_value=_flagged_armor_response())
+    request_data = {"metadata": {}}
+
+    with pytest.raises(HTTPException) as exc_info:
+        await guardrail.apply_guardrail(
+            inputs={"texts": ["forbidden prompt"]},
+            request_data=request_data,
+            input_type="request",
+        )
+
+    assert exc_info.value.status_code == 400
+    entries = request_data["metadata"]["standard_logging_guardrail_information"]
+    flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"]
+    assert len(flagged) == 1