Noma non blocking on monitor mode

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:

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