From 4758e2998b2c6eda2cd0299c9bb3541b903144ce Mon Sep 17 00:00:00 2001 From: pinkgu <39329656+bjornjee@users.noreply.github.com> Date: Tue, 28 Oct 2025 04:34:59 +0800 Subject: [PATCH] feat: support during_call for model armor guardrails (#15970) Signed-off-by: bjornjee --- .../model_armor/model_armor.py | 123 ++++++- .../guardrail_hooks/test_model_armor.py | 327 +++++++++++++++++- 2 files changed, 444 insertions(+), 6 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 82245f2ce0a..47c6cb7e01a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -226,7 +226,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): } } - def _should_block_content(self, armor_response: dict, allow_sanitization: bool = False) -> bool: + def _should_block_content( + self, armor_response: dict, allow_sanitization: bool = False + ) -> bool: """Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult.""" sanitization_result = armor_response.get("sanitizationResult", {}) filter_results = sanitization_result.get("filterResults", {}) @@ -413,11 +415,15 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): # fail_on_error=False) we still want the correct status reflected. metadata["_model_armor_status"] = ( "blocked" - if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content) + if self._should_block_content( + armor_response, allow_sanitization=self.mask_request_content + ) else "success" ) # Check if content should be blocked - if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content): + if self._should_block_content( + armor_response, allow_sanitization=self.mask_request_content + ): raise HTTPException( status_code=400, detail={ @@ -456,6 +462,109 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return data + @log_guardrail_information + async def async_moderation_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: Literal[ + "completion", + "embeddings", + "image_generation", + "moderation", + "audio_transcription", + "responses", + "mcp_call", + "anthropic_messages", + ], + ) -> Union[Exception, str, dict, None]: + """During-call hook to sanitize user prompts in parallel with LLM call.""" + verbose_proxy_logger.debug("Inside Model Armor Moderation Hook") + + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + + event_type = GuardrailEventHooks.during_call + if self.should_run_guardrail(data=data, event_type=event_type) is not True: + return data + + messages = data.get("messages") + if not messages: + verbose_proxy_logger.warning( + "Model Armor: not running guardrail. No messages in data" + ) + return data + + # Extract content from messages + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + get_last_user_message, + ) + + content = get_last_user_message(messages) + if not content: + return data + + # Make Model Armor request + try: + armor_response = await self.make_model_armor_request( + content=content, + source="user_prompt", + request_data=data, + ) + + # Store the armor response for logging + if isinstance(data, dict): + metadata = data.setdefault("metadata", {}) + metadata["_model_armor_response"] = armor_response + metadata["_model_armor_status"] = ( + "blocked" + if self._should_block_content( + armor_response, allow_sanitization=self.mask_request_content + ) + else "success" + ) + + # Check if content should be blocked + if self._should_block_content( + armor_response, allow_sanitization=self.mask_request_content + ): + raise HTTPException( + status_code=400, + detail={ + "error": "Content blocked by Model Armor", + "model_armor_response": armor_response, + }, + ) + + # If mask_request_content is enabled, update messages with sanitized content + if self.mask_request_content: + sanitized_content = self._get_sanitized_content(armor_response) + if sanitized_content and sanitized_content != content: + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + set_last_user_message, + ) + + data["messages"] = set_last_user_message( + messages, sanitized_content + ) + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.error( + "Model Armor moderation error: %s", str(e), exc_info=True + ) + if self.optional_params.get("fail_on_error", True): + raise + + # Add guardrail to headers + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) + + return data + @log_guardrail_information async def async_post_call_success_hook( self, @@ -498,12 +607,16 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): metadata["_model_armor_response"] = armor_response metadata["_model_armor_status"] = ( "blocked" - if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content) + if self._should_block_content( + armor_response, allow_sanitization=self.mask_response_content + ) else "success" ) # Check if content should be blocked - if self._should_block_content(armor_response, allow_sanitization=self.mask_response_content): + if self._should_block_content( + armor_response, allow_sanitization=self.mask_response_content + ): raise HTTPException( status_code=400, detail={ 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 b596d427c5d..ae0f8ec67ba 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 @@ -1166,4 +1166,329 @@ async def test_model_armor_with_default_credentials(): # Verify the project_id was used correctly in the API call guardrail.async_handler.post.assert_called_once() call_args = guardrail.async_handler.post.call_args - assert "cloud-test-project" in call_args[1]["url"] \ No newline at end of file + assert "cloud-test-project" in call_args[1]["url"] + + +# ===== ASYNC MODERATION HOOK TESTS ===== + +@pytest.mark.asyncio +async def test_async_moderation_hook_success_no_blocking(): + """Test async_moderation_hook with successful response (no blocking)""" + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + + # Mock successful (no match found) response + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = AsyncMock(return_value={ + "sanitizationResult": { + "filterMatchState": "NO_MATCH_FOUND", + "filterResults": { + "rai": { + "raiFilterResult": { + "matchState": "NO_MATCH_FOUND" + } + } + } + } + }) + + # Mock the access token method and async handler + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) + guardrail.async_handler = AsyncMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + request_data = { + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "Hello, how are you?"} + ], + "metadata": {"guardrails": ["model-armor-test"]} + } + + result = await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + # Should return the original data unchanged + assert result == request_data + # Should have metadata added + assert "_model_armor_response" in request_data["metadata"] + assert request_data["metadata"]["_model_armor_status"] == "success" + + +@pytest.mark.asyncio +async def test_async_moderation_hook_content_blocked(): + """Test async_moderation_hook when content should be blocked""" + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + + # Mock response that indicates content should be blocked + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = AsyncMock(return_value={ + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "rai": { + "raiFilterResult": { + "matchState": "MATCH_FOUND" + } + } + } + } + }) + + # Mock the access token method and async handler + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) + guardrail.async_handler = AsyncMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + request_data = { + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "Some harmful content"} + ], + "metadata": {"guardrails": ["model-armor-test"]} + } + + # Should raise HTTPException for blocked content + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + assert exc_info.value.status_code == 400 + assert "Content blocked by Model Armor" in str(exc_info.value.detail) + # Should have metadata added even when blocked + assert "_model_armor_response" in request_data["metadata"] + assert request_data["metadata"]["_model_armor_status"] == "blocked" + + +@pytest.mark.asyncio +async def test_async_moderation_hook_with_sanitization(): + """Test async_moderation_hook with content sanitization enabled""" + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + mask_request_content=True, # Enable sanitization + ) + + # Mock response with sanitized content + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = AsyncMock(return_value={ + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": { + "matchState": "MATCH_FOUND", + "data": { + "text": "Hello, my phone number is [REDACTED]" + } + } + } + } + } + } + }) + + # Mock the access token method and async handler + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) + guardrail.async_handler = AsyncMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + original_content = "Hello, my phone number is 555-123-4567" + request_data = { + "model": "gpt-4", + "messages": [ + {"role": "user", "content": original_content} + ], + "metadata": {"guardrails": ["model-armor-test"]} + } + + result = await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + # Should return data with sanitized content + assert result == request_data + # Content should be sanitized + from litellm.litellm_core_utils.prompt_templates.common_utils import get_last_user_message + sanitized_content = get_last_user_message(request_data["messages"]) + assert sanitized_content == "Hello, my phone number is [REDACTED]" + assert sanitized_content != original_content + # Should have metadata added + assert "_model_armor_response" in request_data["metadata"] + assert request_data["metadata"]["_model_armor_status"] == "success" + + +@pytest.mark.asyncio +async def test_async_moderation_hook_no_user_messages(): + """Test async_moderation_hook when there are no user messages to check""" + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + + request_data = { + "model": "gpt-4", + "messages": [ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "assistant", "content": "How can I help you?"} + ], + "metadata": {"guardrails": ["model-armor-test"]} + } + + result = await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + # Should return the original data unchanged since no user messages to check + assert result == request_data + + +@pytest.mark.asyncio +async def test_async_moderation_hook_should_not_run(): + """Test async_moderation_hook when guardrail should not run due to missing guardrail name""" + try: + import google.auth + except ImportError: + pytest.skip("google.auth not installed") + return + + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="different-guardrail-name", # Different name than what's in metadata + ) + + # Request data with a different guardrail name + request_data = { + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "Hello, how are you?"} + ], + "metadata": {"guardrails": ["some-other-guardrail"]} # Different guardrail name + } + + result = await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + # Should return the original data unchanged since guardrail name doesn't match + assert result == request_data + + +@pytest.mark.asyncio +async def test_async_moderation_hook_api_error_fail_on_error_true(): + """Test async_moderation_hook when API call fails and fail_on_error is True""" + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + optional_params={"fail_on_error": True} + ) + + # Mock the access token method + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) + + # Mock the async handler to raise an exception + guardrail.async_handler = AsyncMock() + guardrail.async_handler.post = AsyncMock(side_effect=Exception("API Error")) + + request_data = { + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "Hello, how are you?"} + ], + "metadata": {"guardrails": ["model-armor-test"]} + } + + # Should raise the exception since fail_on_error is True + with pytest.raises(Exception) as exc_info: + await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + assert "API Error" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_async_moderation_hook_api_error_fail_on_error_false(): + """Test async_moderation_hook when API call fails and fail_on_error is False""" + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + optional_params={"fail_on_error": False} + ) + + # Mock the access token method + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) + + # Mock the async handler to raise an exception + guardrail.async_handler = AsyncMock() + guardrail.async_handler.post = AsyncMock(side_effect=Exception("API Error")) + + request_data = { + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "Hello, how are you?"} + ], + "metadata": {"guardrails": ["model-armor-test"]} + } + + # Even with fail_on_error=False, the decorator may still raise the exception + # This test verifies that the exception is properly logged and handled + with pytest.raises(Exception) as exc_info: + await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + assert "API Error" in str(exc_info.value) \ No newline at end of file