diff --git a/litellm/__init__.py b/litellm/__init__.py index 5e108a455c..83bd98c463 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -55,9 +55,12 @@ baseten_key: Optional[str] = None aleph_alpha_key: Optional[str] = None nlp_cloud_key: Optional[str] = None use_client: bool = False +### GUARDRAILS ### llamaguard_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 +################## logging: bool = True caching: bool = ( False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 diff --git a/litellm/proxy/hooks/presidio_pii_masking.py b/litellm/proxy/hooks/presidio_pii_masking.py index a11ea7f55d..99e501589f 100644 --- a/litellm/proxy/hooks/presidio_pii_masking.py +++ b/litellm/proxy/hooks/presidio_pii_masking.py @@ -9,7 +9,7 @@ from typing import Optional, Literal, Union -import litellm, traceback, sys, uuid +import litellm, traceback, sys, uuid, json from litellm.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.integrations.custom_logger import CustomLogger @@ -27,6 +27,7 @@ import aiohttp, asyncio class _OPTIONAL_PresidioPIIMasking(CustomLogger): user_api_key_cache = None + ad_hoc_recognizers = None # Class variables or attributes def __init__( @@ -40,6 +41,22 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): if mock_testing == True: # for testing purposes only return + ad_hoc_recognizers = litellm.presidio_ad_hoc_recognizers + if ad_hoc_recognizers is not None: + try: + with open(ad_hoc_recognizers, "r") as file: + self.ad_hoc_recognizers = json.load(file) + except FileNotFoundError: + raise Exception(f"File not found. file_path={ad_hoc_recognizers}") + except json.JSONDecodeError as e: + raise Exception( + f"Error decoding JSON file: {str(e)}, file_path={ad_hoc_recognizers}" + ) + except Exception as e: + raise Exception( + f"An error occurred: {str(e)}, file_path={ad_hoc_recognizers}" + ) + self.presidio_analyzer_api_base = litellm.get_secret( "PRESIDIO_ANALYZER_API_BASE", None ) @@ -78,6 +95,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): analyze_url = f"{self.presidio_analyzer_api_base}analyze" verbose_proxy_logger.debug(f"Making request to: {analyze_url}") analyze_payload = {"text": text, "language": "en"} + if self.ad_hoc_recognizers is not None: + analyze_payload["ad_hoc_recognizers"] = self.ad_hoc_recognizers redacted_text = None async with session.post( analyze_url, json=analyze_payload