diff --git a/enterprise/enterprise_hooks/openai_moderation.py b/enterprise/enterprise_hooks/openai_moderation.py new file mode 100644 index 00000000000..0fa375fb250 --- /dev/null +++ b/enterprise/enterprise_hooks/openai_moderation.py @@ -0,0 +1,68 @@ +# +-------------------------------------------------------------+ +# +# Use OpenAI /moderations for your LLM calls +# +# +-------------------------------------------------------------+ +# Thank you users! We ❤️ you! - Krrish & Ishaan + +import sys, os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +from typing import Optional, Literal, Union +import litellm, traceback, sys, uuid +from litellm.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.integrations.custom_logger import CustomLogger +from fastapi import HTTPException +from litellm._logging import verbose_proxy_logger +from litellm.utils import ( + ModelResponse, + EmbeddingResponse, + ImageResponse, + StreamingChoices, +) +from datetime import datetime +import aiohttp, asyncio +from litellm._logging import verbose_proxy_logger + +litellm.set_verbose = True + + +class _ENTERPRISE_OpenAI_Moderation(CustomLogger): + def __init__(self): + self.model_name = ( + litellm.openai_moderations_model_name or "text-moderation-latest" + ) # pass the model_name you initialized on litellm.Router() + pass + + #### CALL HOOKS - proxy only #### + + async def async_moderation_hook( ### 👈 KEY CHANGE ### + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: Literal["completion", "embeddings", "image_generation"], + ): + if "messages" in data and isinstance(data["messages"], list): + text = "" + for m in data["messages"]: # assume messages is a list + if "content" in m and isinstance(m["content"], str): + text += m["content"] + + from litellm.proxy.proxy_server import llm_router + + if llm_router is None: + return + + moderation_response = await llm_router.amoderation( + model=self.model_name, input=text + ) + + verbose_proxy_logger.debug("Moderation response: %s", moderation_response) + if moderation_response.results[0].flagged == True: + raise HTTPException( + status_code=403, detail={"error": "Violated content safety policy"} + ) + pass diff --git a/litellm/__init__.py b/litellm/__init__.py index 92610afd9de..2af65f790c2 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -97,6 +97,7 @@ ssl_verify: bool = True disable_streaming_logging: bool = False ### GUARDRAILS ### llamaguard_model_name: Optional[str] = None +openai_moderations_model_name: Optional[str] = None presidio_ad_hoc_recognizers: Optional[str] = None google_moderation_confidence_threshold: Optional[float] = None llamaguard_unsafe_content_categories: Optional[str] = None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b52c9b249e5..4045c7d914c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2313,6 +2313,18 @@ class ProxyConfig: llama_guard_object = _ENTERPRISE_LlamaGuard() imported_list.append(llama_guard_object) + elif ( + isinstance(callback, str) + and callback == "openai_moderations" + ): + from enterprise.enterprise_hooks.openai_moderation import ( + _ENTERPRISE_OpenAI_Moderation, + ) + + openai_moderations_object = ( + _ENTERPRISE_OpenAI_Moderation() + ) + imported_list.append(openai_moderations_object) elif ( isinstance(callback, str) and callback == "google_text_moderation" diff --git a/litellm/tests/test_openai_moderations_hook.py b/litellm/tests/test_openai_moderations_hook.py new file mode 100644 index 00000000000..745f188fc2d --- /dev/null +++ b/litellm/tests/test_openai_moderations_hook.py @@ -0,0 +1,76 @@ +# What is this? +## This tests the llm guard integration + +# What is this? +## Unit test for presidio pii masking +import sys, os, asyncio, time, random +from datetime import datetime +import traceback +from dotenv import load_dotenv + +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.enterprise.enterprise_hooks.openai_moderation import ( + _ENTERPRISE_OpenAI_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 + +### UNIT TESTS FOR OpenAI Moderation ### + + +@pytest.mark.asyncio +async def test_openai_moderation_error_raising(): + """ + Tests to see OpenAI Moderation raises an error for a flagged response + """ + + openai_mod = _ENTERPRISE_OpenAI_Moderation() + litellm.openai_moderations_model_name = "text-moderation-latest" + _api_key = "sk-12345" + _api_key = hash_token("sk-12345") + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) + local_cache = DualCache() + + from litellm.proxy.proxy_server import llm_router + + llm_router = litellm.Router( + model_list=[ + { + "model_name": "text-moderation-latest", + "litellm_params": { + "model": "text-moderation-latest", + "api_key": os.environ["OPENAI_API_KEY"], + }, + } + ] + ) + + setattr(litellm.proxy.proxy_server, "llm_router", llm_router) + + try: + await openai_mod.async_moderation_hook( + data={ + "messages": [ + { + "role": "user", + "content": "fuck off you're the worst", + } + ] + }, + user_api_key_dict=user_api_key_dict, + 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