From f6bc4d0bf949da3240382807e74f904ebc164e67 Mon Sep 17 00:00:00 2001 From: Tom Alon Date: Wed, 10 Sep 2025 11:58:15 +0300 Subject: [PATCH] Noma non blocking on monitor mode --- .../guardrails/guardrail_hooks/noma/noma.py | 228 +++++++++--- .../guardrails/guardrail_hooks/test_noma.py | 325 +++++++++++++++++- 2 files changed, 502 insertions(+), 51 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index ed5929f056..3bbd183e49 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -5,6 +5,7 @@ # # +-------------------------------------------------------------+ +import asyncio import copy import os from typing import Any, Dict, Literal, Optional, Union @@ -24,6 +25,10 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import EmbeddingResponse, ImageResponse +# Type aliases +MessageRole = Literal["user", "assistant"] +LLMResponse = Union[Any, ModelResponse, EmbeddingResponse, ImageResponse] + class NomaBlockedMessage(HTTPException): """Exception raised when Noma guardrail blocks a message""" @@ -164,6 +169,138 @@ class NomaGuardrail(CustomGuardrail): super().__init__(**kwargs) + def _create_background_noma_check( + self, + coro, + ) -> None: + """Create a background task for Noma API calls without blocking the main flow""" + try: + asyncio.create_task(coro) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to create background Noma task: {str(e)}" + ) + + async def _process_user_message_check( + self, + request_data: dict, + user_auth: UserAPIKeyAuth, + ) -> Optional[str]: + """Shared logic for processing user message checks""" + extra_data = self.get_guardrail_dynamic_request_body_params(request_data) + + user_message = await self._extract_user_message(request_data) + if not user_message: + return None + + payload = {"request": {"text": user_message}} + response_json = await self._call_noma_api( + payload=payload, + llm_request_id=None, + request_data=request_data, + user_auth=user_auth, + extra_data=extra_data, + ) + + if self.monitor_mode: + await self._handle_verdict_background("user", user_message, response_json) + else: + await self._check_verdict("user", user_message, response_json) + + return user_message + + async def _process_llm_response_check( + self, + request_data: dict, + response: LLMResponse, + user_auth: UserAPIKeyAuth, + ) -> Optional[str]: + """Shared logic for processing LLM response checks""" + extra_data = self.get_guardrail_dynamic_request_body_params(request_data) + + if not isinstance(response, litellm.ModelResponse): + return None + + content = None + for choice in response.choices: + if isinstance(choice, litellm.Choices) and choice.message.content: + content = choice.message.content + break + + if not content or not isinstance(content, str): + return None + + payload = {"response": {"text": content}} + + response_json = await self._call_noma_api( + payload=payload, + llm_request_id=response.id, + request_data=request_data, + user_auth=user_auth, + extra_data=extra_data, + ) + + if self.monitor_mode: + await self._handle_verdict_background("assistant", content, response_json) + else: + await self._check_verdict("assistant", content, response_json) + + return content + + async def _check_user_message_background( + self, + request_data: dict, + user_auth: UserAPIKeyAuth, + ) -> None: + """Check user message in background for monitor mode - non-blocking""" + try: + await self._process_user_message_check(request_data, user_auth) + except Exception as e: + verbose_proxy_logger.error( + f"Noma background user message check failed: {str(e)}" + ) + + async def _check_llm_response_background( + self, + request_data: dict, + response: LLMResponse, + user_auth: UserAPIKeyAuth, + ) -> None: + """Check LLM response in background for monitor mode - non-blocking""" + try: + await self._process_llm_response_check(request_data, response, user_auth) + except Exception as e: + verbose_proxy_logger.error( + f"Noma background response check failed: {str(e)}" + ) + + async def _handle_verdict_background( + self, + type: MessageRole, + message: str, + response_json: dict, + ) -> None: + """Handle verdict from Noma API in background - logging only, never blocks""" + try: + if not response_json.get("verdict", True): + msg = str.format( + "Noma guardrail blocked {type} message: {message}", + type=type, + message=message, + ) + verbose_proxy_logger.warning(msg) + else: + msg = str.format( + "Noma guardrail allowed {type} message: {message}", + type=type, + message=message, + ) + verbose_proxy_logger.info(msg) + except Exception as e: + verbose_proxy_logger.error( + f"Noma background verdict handling failed: {str(e)}" + ) + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -191,6 +328,18 @@ class NomaGuardrail(CustomGuardrail): ): return data + # In monitor mode, run Noma check in background and return immediately + if self.monitor_mode: + try: + self._create_background_noma_check( + self._check_user_message_background(data, user_api_key_dict) + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to start background Noma pre-call check: {str(e)}" + ) + return data + try: return await self._check_user_message(data, user_api_key_dict) except NomaBlockedMessage: @@ -198,7 +347,7 @@ class NomaGuardrail(CustomGuardrail): except Exception as e: verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}") - if self.block_failures and not self.monitor_mode: + if self.block_failures: raise return data @@ -220,6 +369,18 @@ class NomaGuardrail(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data + # In monitor mode, run Noma check in background and return immediately + if self.monitor_mode: + try: + self._create_background_noma_check( + self._check_user_message_background(data, user_api_key_dict) + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to start background Noma moderation check: {str(e)}" + ) + return data + try: return await self._check_user_message(data, user_api_key_dict) except NomaBlockedMessage: @@ -227,7 +388,7 @@ class NomaGuardrail(CustomGuardrail): except Exception as e: verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}") - if self.block_failures and not self.monitor_mode: + if self.block_failures: raise return data @@ -235,19 +396,33 @@ class NomaGuardrail(CustomGuardrail): self, data: dict, user_api_key_dict: UserAPIKeyAuth, - response: Union[Any, ModelResponse, EmbeddingResponse, ImageResponse], + response: LLMResponse, ): event_type: GuardrailEventHooks = GuardrailEventHooks.post_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: return response + # In monitor mode, run Noma check in background and return immediately + if self.monitor_mode: + try: + self._create_background_noma_check( + self._check_llm_response_background( + data, response, user_api_key_dict + ) + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to start background Noma post-call check: {str(e)}" + ) + return response + try: return await self._check_llm_response(data, response, user_api_key_dict) except NomaBlockedMessage: raise except Exception as e: verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}") - if self.block_failures and not self.monitor_mode: + if self.block_failures: raise return response @@ -257,55 +432,24 @@ class NomaGuardrail(CustomGuardrail): user_auth: UserAPIKeyAuth, ) -> Union[Exception, str, dict, None]: """Check user message for policy violations""" - extra_data = self.get_guardrail_dynamic_request_body_params(request_data) - - user_message = await self._extract_user_message(request_data) + user_message = await self._process_user_message_check(request_data, user_auth) if not user_message: return request_data - payload = {"request": {"text": user_message}} - response_json = await self._call_noma_api( - payload=payload, - llm_request_id=None, - request_data=request_data, - user_auth=user_auth, - extra_data=extra_data, - ) - await self._check_verdict("user", user_message, response_json) - return request_data async def _check_llm_response( self, request_data: dict, - response: Union[Any, ModelResponse, EmbeddingResponse, ImageResponse], + response: LLMResponse, user_auth: UserAPIKeyAuth, ) -> Union[Exception, ModelResponse, Any]: """Check LLM response for policy violations""" - extra_data = self.get_guardrail_dynamic_request_body_params(request_data) - - if not isinstance(response, litellm.ModelResponse): - return response - - content = None - for choice in response.choices: - if isinstance(choice, litellm.Choices) and choice.message.content: - content = choice.message.content - break - - if not content or not isinstance(content, str): - return response - - payload = {"response": {"text": content}} - - response_json = await self._call_noma_api( - payload=payload, - llm_request_id=response.id, - request_data=request_data, - user_auth=user_auth, - extra_data=extra_data, + content = await self._process_llm_response_check( + request_data, response, user_auth ) - await self._check_verdict("assistant", content, response_json) + if not content: + return response return response @@ -371,7 +515,7 @@ class NomaGuardrail(CustomGuardrail): async def _check_verdict( self, - type: Literal["user", "assistant"], + type: MessageRole, message: str, response_json: dict, ) -> None: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py index aeea5f81b1..f1e91db7d5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py @@ -50,7 +50,6 @@ def mock_user_api_key_dict(): soft_budget=None, tpm_limit=None, rpm_limit=None, - parallel_request_limit=None, metadata={}, max_parallel_requests=None, allowed_cache_controls=[], @@ -291,15 +290,38 @@ class TestNomaGuardrailHooks: default_on=True, ) - mock_response = MagicMock() - mock_response.json.return_value = { - "verdict": False, - "originalResponse": {"prompt": {"harmfulContent": {"result": True}}}, - } - mock_response.raise_for_status = MagicMock() + with patch.object( + guardrail, "_create_background_noma_check" + ) as mock_create_background: + # Should return immediately without waiting for API call + result = await guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key_dict, + cache=MagicMock(), + data=mock_request_data, + call_type="completion", + ) - with patch.object(guardrail.async_handler, "post", return_value=mock_response): - # Should not raise exception in monitor mode + assert result == mock_request_data + # Verify background task was created + mock_create_background.assert_called_once() + + @pytest.mark.asyncio + async def test_pre_call_hook_monitor_mode_background_task_failure( + self, mock_user_api_key_dict, mock_request_data + ): + """Test pre-call hook in monitor mode when background task creation fails""" + guardrail = NomaGuardrail( + api_key="test-key", + monitor_mode=True, + guardrail_name="test-guardrail", + event_hook="pre_call", + default_on=True, + ) + + with patch.object( + guardrail, "_create_background_noma_check", side_effect=Exception("Task creation failed") + ): + # Should still return successfully even if background task creation fails result = await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, cache=MagicMock(), @@ -457,6 +479,291 @@ class TestNomaGuardrailHooks: assert message is None +class TestBackgroundProcessing: + """Test the new background processing functionality""" + + @pytest.fixture + def monitor_mode_guardrail(self): + """Create a guardrail with monitor mode enabled""" + return NomaGuardrail( + api_key="test-api-key", + api_base="https://api.test.noma.security/", + application_id="test-app", + monitor_mode=True, # Enable monitor mode + block_failures=True, + guardrail_name="test-noma-guardrail", + event_hook="pre_call", + default_on=True, + ) + + @pytest.mark.asyncio + async def test_process_user_message_check_monitor_mode( + self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test shared helper method in monitor mode""" + mock_response = MagicMock() + mock_response.json.return_value = { + "verdict": False, + "originalResponse": {"prompt": {"harmfulContent": {"result": True}}}, + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + monitor_mode_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + with patch.object( + monitor_mode_guardrail, "_handle_verdict_background" + ) as mock_handle_verdict: + result = await monitor_mode_guardrail._process_user_message_check( + mock_request_data, mock_user_api_key_dict + ) + + assert result == "Hello, how are you?" + mock_post.assert_called_once() + mock_handle_verdict.assert_called_once_with( + "user", "Hello, how are you?", mock_response.json.return_value + ) + + @pytest.mark.asyncio + async def test_process_user_message_check_non_monitor_mode( + self, noma_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test shared helper method in non-monitor mode""" + mock_response = MagicMock() + mock_response.json.return_value = {"verdict": True} + mock_response.raise_for_status = MagicMock() + + with patch.object( + noma_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + with patch.object( + noma_guardrail, "_check_verdict" + ) as mock_check_verdict: + result = await noma_guardrail._process_user_message_check( + mock_request_data, mock_user_api_key_dict + ) + + assert result == "Hello, how are you?" + mock_post.assert_called_once() + mock_check_verdict.assert_called_once_with( + "user", "Hello, how are you?", mock_response.json.return_value + ) + + @pytest.mark.asyncio + async def test_process_llm_response_check_monitor_mode( + self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test LLM response processing in monitor mode""" + from litellm.types.utils import Choices, Message + + response = ModelResponse( + id="test-response-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="I'm doing well, thank you!", role="assistant" + ), + ) + ], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion", + system_fingerprint=None, + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + ) + + mock_api_response = MagicMock() + mock_api_response.json.return_value = {"verdict": True} + mock_api_response.raise_for_status = MagicMock() + + with patch.object( + monitor_mode_guardrail.async_handler, "post", return_value=mock_api_response + ) as mock_post: + with patch.object( + monitor_mode_guardrail, "_handle_verdict_background" + ) as mock_handle_verdict: + result = await monitor_mode_guardrail._process_llm_response_check( + mock_request_data, response, mock_user_api_key_dict + ) + + assert result == "I'm doing well, thank you!" + mock_post.assert_called_once() + mock_handle_verdict.assert_called_once_with( + "assistant", "I'm doing well, thank you!", mock_api_response.json.return_value + ) + + @pytest.mark.asyncio + async def test_check_user_message_background( + self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test background user message check method""" + with patch.object( + monitor_mode_guardrail, "_process_user_message_check" + ) as mock_process: + await monitor_mode_guardrail._check_user_message_background( + mock_request_data, mock_user_api_key_dict + ) + + mock_process.assert_called_once_with(mock_request_data, mock_user_api_key_dict) + + @pytest.mark.asyncio + async def test_check_user_message_background_exception_handling( + self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test background user message check handles exceptions gracefully""" + with patch.object( + monitor_mode_guardrail, "_process_user_message_check", + side_effect=Exception("API failed") + ): + # Should not raise exception, just log error + await monitor_mode_guardrail._check_user_message_background( + mock_request_data, mock_user_api_key_dict + ) + + @pytest.mark.asyncio + async def test_check_llm_response_background( + self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test background LLM response check method""" + from litellm.types.utils import Choices, Message + + response = ModelResponse( + id="test-response-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="Test response", role="assistant"), + ) + ], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion", + system_fingerprint=None, + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + ) + + with patch.object( + monitor_mode_guardrail, "_process_llm_response_check" + ) as mock_process: + await monitor_mode_guardrail._check_llm_response_background( + mock_request_data, response, mock_user_api_key_dict + ) + + mock_process.assert_called_once_with( + mock_request_data, response, mock_user_api_key_dict + ) + + @pytest.mark.asyncio + async def test_handle_verdict_background_blocked(self, monitor_mode_guardrail): + """Test background verdict handling for blocked content""" + response_json = { + "verdict": False, + "originalResponse": {"prompt": {"harmfulContent": {"result": True}}}, + } + + with patch("litellm._logging.verbose_proxy_logger.warning") as mock_warning: + await monitor_mode_guardrail._handle_verdict_background( + "user", "test message", response_json + ) + + mock_warning.assert_called_once() + assert "blocked user message" in mock_warning.call_args[0][0] + + @pytest.mark.asyncio + async def test_handle_verdict_background_allowed(self, monitor_mode_guardrail): + """Test background verdict handling for allowed content""" + response_json = {"verdict": True} + + with patch("litellm._logging.verbose_proxy_logger.info") as mock_info: + await monitor_mode_guardrail._handle_verdict_background( + "assistant", "test response", response_json + ) + + mock_info.assert_called_once() + assert "allowed assistant message" in mock_info.call_args[0][0] + + @pytest.mark.asyncio + async def test_create_background_noma_check(self, monitor_mode_guardrail): + """Test background task creation""" + async def dummy_coroutine(): + return "completed" + + with patch("asyncio.create_task") as mock_create_task: + monitor_mode_guardrail._create_background_noma_check(dummy_coroutine()) + mock_create_task.assert_called_once() + + @pytest.mark.asyncio + async def test_create_background_noma_check_exception(self, monitor_mode_guardrail): + """Test background task creation with exception handling""" + async def dummy_coroutine(): + return "completed" + + with patch("asyncio.create_task", side_effect=Exception("Task creation failed")): + # Should not raise exception, just log error + monitor_mode_guardrail._create_background_noma_check(dummy_coroutine()) + + @pytest.mark.asyncio + async def test_moderation_hook_monitor_mode( + self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test moderation hook in monitor mode""" + # Update event hook to during_call + monitor_mode_guardrail.event_hook = "during_call" + + with patch.object( + monitor_mode_guardrail, "_create_background_noma_check" + ) as mock_create_background: + result = await monitor_mode_guardrail.async_moderation_hook( + data=mock_request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", + ) + + assert result == mock_request_data + mock_create_background.assert_called_once() + + @pytest.mark.asyncio + async def test_post_call_success_hook_monitor_mode( + self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test post-call success hook in monitor mode""" + from litellm.types.utils import Choices, Message + + # Update event hook to post_call + monitor_mode_guardrail.event_hook = "post_call" + + response = ModelResponse( + id="test-response-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="Test response", role="assistant"), + ) + ], + created=1234567890, + model="gpt-3.5-turbo", + object="chat.completion", + system_fingerprint=None, + usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + ) + + with patch.object( + monitor_mode_guardrail, "_create_background_noma_check" + ) as mock_create_background: + result = await monitor_mode_guardrail.async_post_call_success_hook( + data=mock_request_data, + user_api_key_dict=mock_user_api_key_dict, + response=response, + ) + + assert result == response + mock_create_background.assert_called_once() + + class TestIntegration: @pytest.mark.asyncio async def test_full_guardrail_flow(self):