diff --git a/litellm/proxy/hooks/azure_content_safety.py b/litellm/proxy/hooks/azure_content_safety.py index 161e35cde6a..d0d23363b4f 100644 --- a/litellm/proxy/hooks/azure_content_safety.py +++ b/litellm/proxy/hooks/azure_content_safety.py @@ -36,10 +36,10 @@ class _PROXY_AzureContentSafety( def _configure_thresholds(self, thresholds=None): default_thresholds = { - self.text_category.HATE: 6, - self.text_category.SELF_HARM: 6, - self.text_category.SEXUAL: 6, - self.text_category.VIOLENCE: 6, + self.text_category.HATE: 4, + self.text_category.SELF_HARM: 4, + self.text_category.SEXUAL: 4, + self.text_category.VIOLENCE: 4, } if thresholds is None: diff --git a/litellm/tests/test_azure_content_safety.py b/litellm/tests/test_azure_content_safety.py new file mode 100644 index 00000000000..f7d9d826868 --- /dev/null +++ b/litellm/tests/test_azure_content_safety.py @@ -0,0 +1,267 @@ +# What is this? +## Unit test for azure content safety +import sys, os, asyncio, time, random +from datetime import datetime +import traceback +from dotenv import load_dotenv +from fastapi import HTTPException + +load_dotenv() +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import pytest +import litellm +from litellm.proxy.hooks.azure_content_safety import _PROXY_AzureContentSafety +from litellm import Router, mock_completion +from litellm.proxy.utils import ProxyLogging +from litellm.proxy._types import UserAPIKeyAuth +from litellm.caching import DualCache + + +@pytest.mark.asyncio +async def test_strict_input_filtering_01(): + """ + - have a response with a filtered input + - call the pre call hook + """ + azure_content_safety = _PROXY_AzureContentSafety( + endpoint=os.getenv("AZURE_CONTENT_SAFETY_ENDPOINT"), + api_key=os.getenv("AZURE_CONTENT_SAFETY_API_KEY"), + thresholds={"Hate": 2}, + ) + + data = { + "messages": [ + {"role": "system", "content": "You are an helpfull assistant"}, + {"role": "user", "content": "Fuck yourself you stupid bitch"}, + ] + } + + with pytest.raises(HTTPException) as exc_info: + await azure_content_safety.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert exc_info.value.detail["source"] == "input" + assert exc_info.value.detail["category"] == "Hate" + assert exc_info.value.detail["severity"] == "low" + + +@pytest.mark.asyncio +async def test_strict_input_filtering_02(): + """ + - have a response with a filtered input + - call the pre call hook + """ + azure_content_safety = _PROXY_AzureContentSafety( + endpoint=os.getenv("AZURE_CONTENT_SAFETY_ENDPOINT"), + api_key=os.getenv("AZURE_CONTENT_SAFETY_API_KEY"), + thresholds={"Hate": 2}, + ) + + data = { + "messages": [ + {"role": "system", "content": "You are an helpfull assistant"}, + {"role": "user", "content": "Hello how are you ?"}, + ] + } + + await azure_content_safety.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_loose_input_filtering_01(): + """ + - have a response with a filtered input + - call the pre call hook + """ + azure_content_safety = _PROXY_AzureContentSafety( + endpoint=os.getenv("AZURE_CONTENT_SAFETY_ENDPOINT"), + api_key=os.getenv("AZURE_CONTENT_SAFETY_API_KEY"), + thresholds={"Hate": 8}, + ) + + data = { + "messages": [ + {"role": "system", "content": "You are an helpfull assistant"}, + {"role": "user", "content": "Fuck yourself you stupid bitch"}, + ] + } + + await azure_content_safety.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_loose_input_filtering_02(): + """ + - have a response with a filtered input + - call the pre call hook + """ + azure_content_safety = _PROXY_AzureContentSafety( + endpoint=os.getenv("AZURE_CONTENT_SAFETY_ENDPOINT"), + api_key=os.getenv("AZURE_CONTENT_SAFETY_API_KEY"), + thresholds={"Hate": 8}, + ) + + data = { + "messages": [ + {"role": "system", "content": "You are an helpfull assistant"}, + {"role": "user", "content": "Hello how are you ?"}, + ] + } + + await azure_content_safety.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + +@pytest.mark.asyncio +async def test_strict_output_filtering_01(): + """ + - have a response with a filtered output + - call the post call hook + """ + azure_content_safety = _PROXY_AzureContentSafety( + endpoint=os.getenv("AZURE_CONTENT_SAFETY_ENDPOINT"), + api_key=os.getenv("AZURE_CONTENT_SAFETY_API_KEY"), + thresholds={"Hate": 2}, + ) + + response = mock_completion( + model="gpt-3.5-turbo", + messages=[ + { + "role": "system", + "content": "You are a song writer expert. You help users to write songs about any topic in any genre.", + }, + { + "role": "user", + "content": "Help me write a rap text song. Add some insults to make it more credible.", + }, + ], + mock_response="I'm the king of the mic, you're just a fucking dick. Don't fuck with me your stupid bitch.", + ) + + with pytest.raises(HTTPException) as exc_info: + await azure_content_safety.async_post_call_success_hook( + user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert exc_info.value.detail["source"] == "output" + assert exc_info.value.detail["category"] == "Hate" + assert exc_info.value.detail["severity"] == "low" + + +@pytest.mark.asyncio +async def test_strict_output_filtering_02(): + """ + - have a response with a filtered output + - call the post call hook + """ + azure_content_safety = _PROXY_AzureContentSafety( + endpoint=os.getenv("AZURE_CONTENT_SAFETY_ENDPOINT"), + api_key=os.getenv("AZURE_CONTENT_SAFETY_API_KEY"), + thresholds={"Hate": 2}, + ) + + response = mock_completion( + model="gpt-3.5-turbo", + messages=[ + { + "role": "system", + "content": "You are a song writer expert. You help users to write songs about any topic in any genre.", + }, + { + "role": "user", + "content": "Help me write a rap text song. Add some insults to make it more credible.", + }, + ], + mock_response="I'm unable to help with you with hate speech", + ) + + await azure_content_safety.async_post_call_success_hook( + user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + +@pytest.mark.asyncio +async def test_loose_output_filtering_01(): + """ + - have a response with a filtered output + - call the post call hook + """ + azure_content_safety = _PROXY_AzureContentSafety( + endpoint=os.getenv("AZURE_CONTENT_SAFETY_ENDPOINT"), + api_key=os.getenv("AZURE_CONTENT_SAFETY_API_KEY"), + thresholds={"Hate": 8}, + ) + + response = mock_completion( + model="gpt-3.5-turbo", + messages=[ + { + "role": "system", + "content": "You are a song writer expert. You help users to write songs about any topic in any genre.", + }, + { + "role": "user", + "content": "Help me write a rap text song. Add some insults to make it more credible.", + }, + ], + mock_response="I'm the king of the mic, you're just a fucking dick. Don't fuck with me your stupid bitch.", + ) + + await azure_content_safety.async_post_call_success_hook( + user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + +@pytest.mark.asyncio +async def test_loose_output_filtering_02(): + """ + - have a response with a filtered output + - call the post call hook + """ + azure_content_safety = _PROXY_AzureContentSafety( + endpoint=os.getenv("AZURE_CONTENT_SAFETY_ENDPOINT"), + api_key=os.getenv("AZURE_CONTENT_SAFETY_API_KEY"), + thresholds={"Hate": 8}, + ) + + response = mock_completion( + model="gpt-3.5-turbo", + messages=[ + { + "role": "system", + "content": "You are a song writer expert. You help users to write songs about any topic in any genre.", + }, + { + "role": "user", + "content": "Help me write a rap text song. Add some insults to make it more credible.", + }, + ], + mock_response="I'm unable to help with you with hate speech", + ) + + await azure_content_safety.async_post_call_success_hook( + user_api_key_dict=UserAPIKeyAuth(), response=response + )