From 135ec00b27ac2452611032336c743c57f64ccae7 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 08:06:09 +0000 Subject: [PATCH 1/5] feat(model_armor): logging_only mode scans completed streams after delivery Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 25 +- .../model_armor/model_armor.py | 68 ++++- .../integrations/test_custom_guardrail.py | 57 +++- .../guardrail_hooks/test_model_armor.py | 289 ++++++++++++++++++ 4 files changed, 428 insertions(+), 11 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 77bf4820a1a..ae6d646b5ae 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -30,6 +30,7 @@ from litellm.types.utils import ( GuardrailStatus, GuardrailTracingDetail, LLMResponseTypes, + ModelResponse, StandardLoggingGuardrailInformation, ) @@ -602,9 +603,7 @@ class CustomGuardrail(CustomLogger): supported_event_hooks: list[GuardrailEventHooks], ) -> None: allowed_hooks: Final = frozenset(supported_event_hooks) | ( - frozenset((GuardrailEventHooks.logging_only,)) - if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks - else frozenset() + frozenset((GuardrailEventHooks.logging_only,)) if self.uses_apply_guardrail_interface() else frozenset() ) def _validate_event_hook_list_is_in_supported_event_hooks( @@ -883,7 +882,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))() @@ -922,6 +923,8 @@ class CustomGuardrail(CustomLogger): translation: "BaseTranslation", scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata ) -> None: + from litellm.llms import get_guardrail_translation_mapping + optional_params: Final = kwargs.get("optional_params") or {} scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input")) scratch_request: Final = { @@ -933,8 +936,18 @@ 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 + response: Final = ( + kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result + ) + if response is None: + return + output_translation: Final = ( + get_guardrail_translation_mapping(CallTypes.acompletion)() + if isinstance(response, ModelResponse) + else translation + ) + 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: 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 fde40111d49..01b05227a91 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/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, Optional 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,59 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): for chunk in all_chunks: yield chunk + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = 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 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, + ) + raise + + 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: + raise HTTPException( + status_code=400, + detail=self._build_block_error_detail( + "Response blocked by Model Armor" if input_type == "response" else "Violated content safety policy", + armor_response, + ), + ) + return inputs + @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: """ diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index ddc8439a83a..c0463210445 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2404,7 +2404,7 @@ def test_logging_only_requires_framework_support_or_explicit_declaration( event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, ) -> None: supported: Final = [GuardrailEventHooks.pre_call] - if guardrail_type is _InheritedApplyGuardrail: + if guardrail_type is not CustomGuardrail: guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) assert guardrail.event_hook == event_hook assert supported == [GuardrailEventHooks.pre_call] @@ -2566,7 +2566,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 @@ -2575,9 +2575,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): @@ -2829,3 +2829,52 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: assert response.choices[0].message.content == "filtered response" assert "guardrail_to_apply" not in request_data 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"] 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 47089b7b1b1..a6ed4e14616 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 @@ -4929,3 +4929,292 @@ 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_raises_on_flagged_content(): + guardrail = _logging_only_guardrail() + 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 output"]}, + request_data=request_data, + input_type="response", + ) + + assert exc_info.value.status_code == 400 + entries = request_data["metadata"]["standard_logging_guardrail_information"] + assert entries[-1]["guardrail_status"] == "guardrail_flagged" From a5cc65f1a9b5f7a674d98f2e45460790952724e5 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 08:10:35 +0000 Subject: [PATCH 2/5] refactor(custom_guardrail): resolve logging_only output translation in async_logging_hook Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index ae6d646b5ae..305a5c20764 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -901,8 +901,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") @@ -919,12 +927,11 @@ 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: - from litellm.llms import get_guardrail_translation_mapping - optional_params: Final = kwargs.get("optional_params") or {} scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input")) scratch_request: Final = { @@ -936,16 +943,8 @@ class CustomGuardrail(CustomLogger): "metadata": scratch_metadata, } await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self) - response: Final = ( - kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result - ) if response is None: return - output_translation: Final = ( - get_guardrail_translation_mapping(CallTypes.acompletion)() - if isinstance(response, ModelResponse) - else translation - ) await output_translation.process_output_response( response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request ) From f417d7f739fca2b8abed6ad3d4a399e39b2cfee7 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 08:49:11 +0000 Subject: [PATCH 3/5] fix(model_armor): record logging_only verdicts without raising so both scans run Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_armor/model_armor.py | 12 +---- .../guardrail_hooks/test_model_armor.py | 53 ++++++++++++++++--- 2 files changed, 47 insertions(+), 18 deletions(-) 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 01b05227a91..127a7ea6786 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1245,7 +1245,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): armor_response: Final = await self.make_model_armor_request( content=content, source=source, request_data=request_data ) - except ModelArmorAPIError as e: + 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), @@ -1256,7 +1256,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=error_end_time, duration=error_end_time - start_time, ) - raise + return inputs flagged: Final = self._should_block_content(armor_response, allow_sanitization=False) end_time: Final = time.time() @@ -1269,14 +1269,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=end_time, duration=end_time - start_time, ) - if flagged: - raise HTTPException( - status_code=400, - detail=self._build_block_error_detail( - "Response blocked by Model Armor" if input_type == "response" else "Violated content safety policy", - armor_response, - ), - ) return inputs @staticmethod 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 a6ed4e14616..bd358e84148 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 @@ -5203,18 +5203,55 @@ async def test_native_post_call_mode_ignores_logging_hook(): @pytest.mark.asyncio -async def test_apply_guardrail_raises_on_flagged_content(): +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"]} - with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail( - inputs={"texts": ["forbidden output"]}, - request_data=request_data, - input_type="response", - ) + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) - assert exc_info.value.status_code == 400 + 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 From 39b916f13f3330b5099ff45ede27e59f7278f53a Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 09:07:06 +0000 Subject: [PATCH 4/5] fix(model_armor): decorate apply_guardrail with log_guardrail_information Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/guardrails/guardrail_hooks/model_armor/model_armor.py | 1 + 1 file changed, 1 insertion(+) 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 127a7ea6786..a833de6096d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1226,6 +1226,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): for chunk in all_chunks: yield chunk + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, From 2ca29a9a9157d52149a4c2cd38e98f89c7f57e30 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 11 Sep 2026 09:33:57 +0000 Subject: [PATCH 5/5] fix(model_armor): gate apply_guardrail raise to non-logging_only and require native guardrails to declare logging_only Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 4 ++- .../model_armor/model_armor.py | 12 +++++++-- .../integrations/test_custom_guardrail.py | 2 +- .../guardrail_hooks/test_model_armor.py | 27 +++++++++++++++++++ 4 files changed, 41 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 305a5c20764..407445b2828 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -603,7 +603,9 @@ class CustomGuardrail(CustomLogger): supported_event_hooks: list[GuardrailEventHooks], ) -> None: allowed_hooks: Final = frozenset(supported_event_hooks) | ( - frozenset((GuardrailEventHooks.logging_only,)) if self.uses_apply_guardrail_interface() else frozenset() + frozenset((GuardrailEventHooks.logging_only,)) + if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks + else frozenset() ) def _validate_event_hook_list_is_in_supported_event_hooks( 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 a833de6096d..a0563a7a1c9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,7 +1,7 @@ import time from collections.abc import AsyncGenerator, Mapping, Sequence from enum import Enum, auto -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal import httpx from fastapi import HTTPException @@ -1232,7 +1232,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], - logging_obj: Optional["LiteLLMLoggingObj"] = None, + logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: content: Final = "\n".join(text for text in inputs.get("texts") or () if text) if not content: @@ -1270,6 +1270,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): 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 diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index c0463210445..bdb7fad21b3 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2404,7 +2404,7 @@ def test_logging_only_requires_framework_support_or_explicit_declaration( event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, ) -> None: supported: Final = [GuardrailEventHooks.pre_call] - if guardrail_type is not CustomGuardrail: + if guardrail_type is _InheritedApplyGuardrail: guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) assert guardrail.event_hook == event_hook assert supported == [GuardrailEventHooks.pre_call] 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 bd358e84148..126e162fec8 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 @@ -5255,3 +5255,30 @@ async def test_logging_only_flagged_prompt_still_scans_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