From e2046bd59c4deaa77d354490e69fbc45352d0901 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Jun 2024 16:13:44 -0700 Subject: [PATCH 1/2] feat - laker return orig response from lakera api --- enterprise/enterprise_hooks/lakera_ai.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/enterprise/enterprise_hooks/lakera_ai.py b/enterprise/enterprise_hooks/lakera_ai.py index dd37ae2c10..2a4ad418b3 100644 --- a/enterprise/enterprise_hooks/lakera_ai.py +++ b/enterprise/enterprise_hooks/lakera_ai.py @@ -114,7 +114,11 @@ class _ENTERPRISE_lakeraAI_Moderation(CustomLogger): if flagged == True: raise HTTPException( - status_code=400, detail={"error": "Violated content safety policy"} + status_code=400, + detail={ + "error": "Violated content safety policy", + "lakera_ai_response": _json_response, + }, ) pass From 0223e52b214fd91c6952d5e281d373768d70cde2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Jun 2024 16:14:26 -0700 Subject: [PATCH 2/2] test - lakera ai detection --- .../tests/test_lakera_ai_prompt_injection.py | 32 ++++++++++++------- 1 file changed, 21 insertions(+), 11 deletions(-) diff --git a/litellm/tests/test_lakera_ai_prompt_injection.py b/litellm/tests/test_lakera_ai_prompt_injection.py index 6227eabaa3..3e328c8244 100644 --- a/litellm/tests/test_lakera_ai_prompt_injection.py +++ b/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