Merge pull request #3802 from BerriAI/litellm_add_moderations_safety_check_proxy

[Feat]- Proxy Add OpenAI Content Moderation Pre call hook

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

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

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"

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