test - lakera ai detection

litellm/tests/test_lakera_ai_prompt_injection.py

@@ -1,10 +1,16 @@
 # What is this?
 ## This tests the Lakera AI integration
 
-import sys, os, asyncio, time, random
-from datetime import datetime
+import asyncio
+import os
+import random
+import sys
+import time
 import traceback
+from datetime import datetime
+
 from dotenv import load_dotenv
+from fastapi import HTTPException
 
 load_dotenv()
 import os
@@ -12,17 +18,19 @@ import os
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
+import logging
+
 import pytest
+
 import litellm
+from litellm import Router, mock_completion
+from litellm._logging import verbose_proxy_logger
+from litellm.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.enterprise.enterprise_hooks.lakera_ai import (
     _ENTERPRISE_lakeraAI_Moderation,
 )
-from litellm import Router, mock_completion
 from litellm.proxy.utils import ProxyLogging, hash_token
-from litellm.proxy._types import UserAPIKeyAuth
-from litellm.caching import DualCache
-from litellm._logging import verbose_proxy_logger
-import logging
 
 verbose_proxy_logger.setLevel(logging.DEBUG)
 
@@ -55,10 +63,12 @@ async def test_lakera_prompt_injection_detection():
             call_type="completion",
         )
         pytest.fail(f"Should have failed")
-    except Exception as e:
-        print("Got exception: ", e)
-        assert "Violated content safety policy" in str(e)
-        pass
+    except HTTPException as http_exception:
+        print("http exception details=", http_exception.detail)
+
+        # Assert that the laker ai response is in the exception raise
+        assert "lakera_ai_response" in http_exception.detail
+        assert "Violated content safety policy" in str(http_exception)
 
 
 @pytest.mark.asyncio