diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 0bc9481563..a82eed8eb8 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -6,6 +6,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.types.guardrails import ( DynamicGuardrailParams, GuardrailEventHooks, + LitellmParams, PiiEntityType, ) from litellm.types.utils import StandardLoggingGuardrailInformation @@ -314,6 +315,12 @@ class CustomGuardrail(CustomLogger): # Mask the content return content_string[:start_index] + mask_string + content_string[end_index:] + def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: + """ + Update the guardrails litellm params in memory + """ + pass + def log_guardrail_information(func): """ diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 1fdcf36fcc..b9a7d85f20 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -410,6 +410,7 @@ async def delete_guardrail(guardrail_id: str): } ``` """ + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -429,6 +430,11 @@ async def delete_guardrail(guardrail_id: str): result = await GUARDRAIL_REGISTRY.delete_guardrail_from_db( guardrail_id=guardrail_id, prisma_client=prisma_client ) + + # delete in memory guardrail + IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail( + guardrail_id=guardrail_id, + ) return result except HTTPException as e: raise e @@ -487,6 +493,7 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): } ``` """ + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -510,7 +517,7 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): else existing_guardrail.get("guardrail_name") ) - # Update litellm_params if default_on is provided + # Update litellm_params if default_on is provided or pii_entities_config is provided litellm_params = LitellmParams( **dict(existing_guardrail.get("litellm_params", {})) ) @@ -520,6 +527,14 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): ): litellm_params.default_on = request.litellm_params.default_on + if ( + request.litellm_params is not None + and request.litellm_params.pii_entities_config is not None + ): + litellm_params.pii_entities_config = ( + request.litellm_params.pii_entities_config + ) + # Update guardrail_info if provided guardrail_info = ( request.guardrail_info @@ -528,15 +543,22 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): ) # Create the guardrail object + guardrail = Guardrail( + guardrail_name=guardrail_name or "", + litellm_params=litellm_params, + guardrail_info=guardrail_info, + ) result = await GUARDRAIL_REGISTRY.update_guardrail_in_db( guardrail_id=guardrail_id, - guardrail=Guardrail( - guardrail_name=guardrail_name or "", - litellm_params=litellm_params, - guardrail_info=guardrail_info, - ), + guardrail=guardrail, prisma_client=prisma_client, ) + + # update in memory guardrail + IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail( + guardrail_id=guardrail_id, + guardrail=guardrail, + ) return result except HTTPException as e: raise e diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 0b7e38ad6c..8052dfd3a0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -25,6 +25,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import ( GuardrailEventHooks, + LitellmParams, PiiAction, PiiEntityType, PresidioPerRequestConfig, @@ -575,3 +576,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): request_data={}, ) return text + + def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: + """ + Update the guardrails litellm params in memory + """ + if litellm_params.pii_entities_config: + self.pii_entities_config = litellm_params.pii_entities_config diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 55d746f1b6..468df07ea4 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -17,6 +17,7 @@ def initialize_aporia( default_on=litellm_params.default_on, ) litellm.logging_callback_manager.add_litellm_callback(_aporia_callback) + return _aporia_callback def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): @@ -44,6 +45,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint, ) litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback) + return _bedrock_callback def initialize_lakera(litellm_params: LitellmParams, guardrail: Guardrail): @@ -58,6 +60,7 @@ def initialize_lakera(litellm_params: LitellmParams, guardrail: Guardrail): default_on=litellm_params.default_on, ) litellm.logging_callback_manager.add_litellm_callback(_lakera_callback) + return _lakera_callback def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail): @@ -76,6 +79,7 @@ def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail): dev_info=litellm_params.dev_info, ) litellm.logging_callback_manager.add_litellm_callback(_lakera_v2_callback) + return _lakera_v2_callback def initialize_aim(litellm_params: LitellmParams, guardrail: Guardrail): @@ -90,6 +94,8 @@ def initialize_aim(litellm_params: LitellmParams, guardrail: Guardrail): ) litellm.logging_callback_manager.add_litellm_callback(_aim_callback) + return _aim_callback + def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): from litellm.proxy.guardrails.guardrail_hooks.presidio import ( @@ -121,6 +127,8 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail): ) litellm.logging_callback_manager.add_litellm_callback(_success_callback) + return _presidio_callback + def initialize_hide_secrets(litellm_params: LitellmParams, guardrail: Guardrail): from litellm_enterprise.enterprise_callbacks.secret_detection import ( @@ -134,6 +142,7 @@ def initialize_hide_secrets(litellm_params: LitellmParams, guardrail: Guardrail) default_on=litellm_params.default_on, ) litellm.logging_callback_manager.add_litellm_callback(_secret_detection_object) + return _secret_detection_object def initialize_guardrails_ai(litellm_params, guardrail): @@ -152,3 +161,5 @@ def initialize_guardrails_ai(litellm_params, guardrail): default_on=litellm_params.default_on, ) litellm.logging_callback_manager.add_litellm_callback(_guardrails_ai_callback) + + return _guardrails_ai_callback diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 1f853f3c02..07a577e254 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -1,13 +1,24 @@ # litellm/proxy/guardrails/guardrail_registry.py +import importlib +import os +import uuid from datetime import datetime, timezone -from typing import List, Optional +from typing import Dict, List, Optional, cast import litellm +from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.utils import PrismaClient -from litellm.types.guardrails import Guardrail, SupportedGuardrailIntegrations +from litellm.secret_managers.main import get_secret +from litellm.types.guardrails import ( + Guardrail, + GuardrailEventHooks, + LakeraCategoryThresholds, + LitellmParams, + SupportedGuardrailIntegrations, +) from .guardrail_initializers import ( initialize_aim, @@ -195,3 +206,179 @@ class GuardrailRegistry: return Guardrail(**(dict(guardrail))) except Exception as e: raise Exception(f"Error getting guardrail from DB: {str(e)}") + + +class InMemoryGuardrailHandler: + """ + Class that handles initializing guardrails and adding them to the CallbackManager + """ + + def __init__(self): + self.IN_MEMORY_GUARDRAILS: Dict[str, Guardrail] = {} + """ + Guardrail id to Guardrail object mapping + """ + + self.guardrail_id_to_custom_guardrail: Dict[str, Optional[CustomGuardrail]] = {} + """ + Guardrail id to CustomGuardrail object mapping + """ + + def initialize_guardrail( + self, + guardrail: Dict, + config_file_path: Optional[str] = None, + ) -> Optional[Guardrail]: + """ + Initialize a guardrail from a dictionary and add it to the litellm callback manager + + Returns a Guardrail object if the guardrail is initialized successfully + """ + guardrail_id = guardrail.get("guardrail_id") or str(uuid.uuid4()) + if guardrail_id in self.IN_MEMORY_GUARDRAILS: + verbose_proxy_logger.debug( + "guardrail_id already exists in IN_MEMORY_GUARDRAILS" + ) + return self.IN_MEMORY_GUARDRAILS[guardrail_id] + + custom_guardrail_callback: Optional[CustomGuardrail] = None + litellm_params_data = guardrail["litellm_params"] + verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data) + + litellm_params = LitellmParams(**litellm_params_data) + + if ( + "category_thresholds" in litellm_params_data + and litellm_params_data["category_thresholds"] + ): + lakera_category_thresholds = LakeraCategoryThresholds( + **litellm_params_data["category_thresholds"] + ) + litellm_params.category_thresholds = lakera_category_thresholds + + if litellm_params.api_key and litellm_params.api_key.startswith("os.environ/"): + litellm_params.api_key = str(get_secret(litellm_params.api_key)) + + if litellm_params.api_base and litellm_params.api_base.startswith( + "os.environ/" + ): + litellm_params.api_base = str(get_secret(litellm_params.api_base)) + + guardrail_type = litellm_params.guardrail + if guardrail_type is None: + raise ValueError("guardrail_type is required") + + initializer = guardrail_initializer_registry.get(guardrail_type) + + if initializer: + custom_guardrail_callback = initializer(litellm_params, guardrail) + elif isinstance(guardrail_type, str) and "." in guardrail_type: + self.initialize_custom_guardrail( + guardrail=guardrail, + guardrail_type=guardrail_type, + litellm_params=litellm_params, + config_file_path=config_file_path, + ) + else: + raise ValueError(f"Unsupported guardrail: {guardrail_type}") + + parsed_guardrail = Guardrail( + guardrail_id=guardrail.get("guardrail_id"), + guardrail_name=guardrail["guardrail_name"], + litellm_params=litellm_params, + ) + + guardrail_id = parsed_guardrail.get("guardrail_id") or str(uuid.uuid4()) + + # store references to the guardrail in memory + self.IN_MEMORY_GUARDRAILS[guardrail_id] = parsed_guardrail + self.guardrail_id_to_custom_guardrail[guardrail_id] = custom_guardrail_callback + + return parsed_guardrail + + def initialize_custom_guardrail( + self, + guardrail: Dict, + guardrail_type: str, + litellm_params: LitellmParams, + config_file_path: Optional[str] = None, + ) -> None: + """ + Initialize a Custom Guardrail from a python file + + This initializes it by adding it to the litellm callback manager + """ + if not config_file_path: + raise Exception( + "GuardrailsAIException - Please pass the config_file_path to initialize_guardrails_v2" + ) + + _file_name, _class_name = guardrail_type.split(".") + verbose_proxy_logger.debug( + "Initializing custom guardrail: %s, file_name: %s, class_name: %s", + guardrail_type, + _file_name, + _class_name, + ) + + directory = os.path.dirname(config_file_path) + module_file_path = os.path.join(directory, _file_name) + ".py" + + spec = importlib.util.spec_from_file_location(_class_name, module_file_path) # type: ignore + if not spec: + raise ImportError( + f"Could not find a module specification for {module_file_path}" + ) + + module = importlib.util.module_from_spec(spec) # type: ignore + spec.loader.exec_module(module) # type: ignore + _guardrail_class = getattr(module, _class_name) + + mode = litellm_params.mode + if mode is None: + raise ValueError( + f"mode is required for guardrail {guardrail_type} please set mode to one of the following: {', '.join(GuardrailEventHooks)}" + ) + + default_on = litellm_params.default_on + _guardrail_callback = _guardrail_class( + guardrail_name=guardrail["guardrail_name"], + event_hook=mode, + default_on=default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_guardrail_callback) # type: ignore + + def update_in_memory_guardrail( + self, guardrail_id: str, guardrail: Guardrail + ) -> None: + """ + Update a guardrail in memory + + - updates the guardrail in memory + - updates the guardrail params in litellm.callback_manager + """ + self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail + + custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.get( + guardrail_id + ) + if custom_guardrail_callback: + updated_litellm_params = cast( + LitellmParams, guardrail.get("litellm_params", {}) + ) + custom_guardrail_callback.update_in_memory_litellm_params( + litellm_params=updated_litellm_params + ) + + def delete_in_memory_guardrail(self, guardrail_id: str) -> None: + """ + Delete a guardrail in memory + """ + self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None) + + +######################################################## +# In Memory Guardrail Handler for LiteLLM Proxy +######################################################## +IN_MEMORY_GUARDRAIL_HANDLER = InMemoryGuardrailHandler() +######################################################## diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index 2864f0d268..70c6fe65ad 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -1,23 +1,11 @@ -import importlib -import os from typing import Dict, List, Optional import litellm -from litellm import get_secret from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy # v2 implementation -from litellm.types.guardrails import ( - Guardrail, - GuardrailEventHooks, - GuardrailItem, - GuardrailItemSpec, - LakeraCategoryThresholds, - LitellmParams, -) - -from .guardrail_registry import guardrail_initializer_registry +from litellm.types.guardrails import Guardrail, GuardrailItem, GuardrailItemSpec all_guardrails: List[GuardrailItem] = [] @@ -89,10 +77,12 @@ def init_guardrails_v2( all_guardrails: List[Dict], config_file_path: Optional[str] = None, ): + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + guardrail_list: List[Guardrail] = [] for guardrail in all_guardrails: - initialized_guardrail = InitializeGuardrails.initialize_guardrail( + initialized_guardrail = IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( guardrail=guardrail, config_file_path=config_file_path, ) @@ -100,118 +90,3 @@ def init_guardrails_v2( guardrail_list.append(initialized_guardrail) verbose_proxy_logger.info(f"\nGuardrail List:{guardrail_list}\n") - - -class InitializeGuardrails: - """ - Class that handles initializing guardrails and adding them to the CallbackManager - """ - - @staticmethod - def initialize_guardrail( - guardrail: Dict, - config_file_path: Optional[str] = None, - ) -> Optional[Guardrail]: - """ - Initialize a guardrail from a dictionary and add it to the litellm callback manager - - Returns a Guardrail object if the guardrail is initialized successfully - """ - litellm_params_data = guardrail["litellm_params"] - verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data) - - litellm_params = LitellmParams(**litellm_params_data) - - if ( - "category_thresholds" in litellm_params_data - and litellm_params_data["category_thresholds"] - ): - lakera_category_thresholds = LakeraCategoryThresholds( - **litellm_params_data["category_thresholds"] - ) - litellm_params.category_thresholds = lakera_category_thresholds - - if litellm_params.api_key and litellm_params.api_key.startswith("os.environ/"): - litellm_params.api_key = str(get_secret(litellm_params.api_key)) - - if litellm_params.api_base and litellm_params.api_base.startswith( - "os.environ/" - ): - litellm_params.api_base = str(get_secret(litellm_params.api_base)) - - guardrail_type = litellm_params.guardrail - if guardrail_type is None: - raise ValueError("guardrail_type is required") - - initializer = guardrail_initializer_registry.get(guardrail_type) - - if initializer: - initializer(litellm_params, guardrail) - elif isinstance(guardrail_type, str) and "." in guardrail_type: - InitializeGuardrails.initialize_custom_guardrail( - guardrail=guardrail, - guardrail_type=guardrail_type, - litellm_params=litellm_params, - config_file_path=config_file_path, - ) - else: - raise ValueError(f"Unsupported guardrail: {guardrail_type}") - - parsed_guardrail = Guardrail( - guardrail_name=guardrail["guardrail_name"], - litellm_params=litellm_params, - ) - - return parsed_guardrail - - @staticmethod - def initialize_custom_guardrail( - guardrail: Dict, - guardrail_type: str, - litellm_params: LitellmParams, - config_file_path: Optional[str] = None, - ) -> None: - """ - Initialize a Custom Guardrail from a python file - - This initializes it by adding it to the litellm callback manager - """ - if not config_file_path: - raise Exception( - "GuardrailsAIException - Please pass the config_file_path to initialize_guardrails_v2" - ) - - _file_name, _class_name = guardrail_type.split(".") - verbose_proxy_logger.debug( - "Initializing custom guardrail: %s, file_name: %s, class_name: %s", - guardrail_type, - _file_name, - _class_name, - ) - - directory = os.path.dirname(config_file_path) - module_file_path = os.path.join(directory, _file_name) + ".py" - - spec = importlib.util.spec_from_file_location(_class_name, module_file_path) # type: ignore - if not spec: - raise ImportError( - f"Could not find a module specification for {module_file_path}" - ) - - module = importlib.util.module_from_spec(spec) # type: ignore - spec.loader.exec_module(module) # type: ignore - _guardrail_class = getattr(module, _class_name) - - mode = litellm_params.mode - if mode is None: - raise ValueError( - f"mode is required for guardrail {guardrail_type} please set mode to one of the following: {', '.join(GuardrailEventHooks)}" - ) - - default_on = litellm_params.default_on - _guardrail_callback = _guardrail_class( - guardrail_name=guardrail["guardrail_name"], - event_hook=mode, - default_on=default_on, - ) - litellm.logging_callback_manager.add_litellm_callback(_guardrail_callback) # type: ignore diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index bd179efbe3..d33d932759 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -5,11 +5,6 @@ model_list: api_key: any_key api_base: https://exampleopenaiendpoint-production.up.railway.app/ -litellm_settings: - callbacks: - - "langfuse" - - "arize_phoenix" - general_settings: store_prompts_in_spend_logs: true \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4b81dc1fab..1e4c4adee0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -143,7 +143,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.proxy._experimental.mcp_server.server import router as mcp_router -from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) @@ -198,6 +197,7 @@ from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES from litellm.proxy.credential_endpoints.endpoints import router as credential_router +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config @@ -2707,10 +2707,10 @@ class ProxyConfig: async def _init_guardrails_in_db(self, prisma_client: PrismaClient): from litellm.proxy.guardrails.guardrail_registry import ( + IN_MEMORY_GUARDRAIL_HANDLER, Guardrail, GuardrailRegistry, ) - from litellm.proxy.guardrails.init_guardrails import InitializeGuardrails try: guardrails_in_db: List[ @@ -2722,7 +2722,7 @@ class ProxyConfig: "guardrails from the DB %s", str(guardrails_in_db) ) for guardrail in guardrails_in_db: - InitializeGuardrails.initialize_guardrail( + IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( guardrail=dict(guardrail), ) except Exception as e: @@ -3302,7 +3302,9 @@ class ProxyStartupEvent: if general_settings.get("maximum_spend_logs_retention_period") is not None: spend_log_cleanup = SpendLogCleanup() # Get the interval from config or default to 1 day - retention_interval = general_settings.get("maximum_spend_logs_retention_interval", "1d") + retention_interval = general_settings.get( + "maximum_spend_logs_retention_interval", "1d" + ) try: interval_seconds = duration_in_seconds(retention_interval) scheduler.add_job( @@ -3312,7 +3314,9 @@ class ProxyStartupEvent: args=[prisma_client], ) except ValueError: - verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value") + verbose_proxy_logger.error( + "Invalid maximum_spend_logs_retention_interval value" + ) scheduler.start() diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 11b8aceee7..8f7114837d 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -312,7 +312,7 @@ class LitellmParams( LakeraV2GuardrailConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") - mode: str = Field( + mode: Union[str, List[str]] = Field( description="When to apply the guardrail (pre_call, post_call, during_call, logging_only)" ) api_key: Optional[str] = Field( @@ -354,6 +354,7 @@ class LitellmParams( class Guardrail(TypedDict, total=False): + guardrail_id: Optional[str] guardrail_name: str litellm_params: LitellmParams guardrail_info: Optional[Dict] @@ -382,6 +383,7 @@ class GuardrailLiteLLMParamsResponse(BaseModel): guardrail: str mode: Union[str, List[str]] default_on: bool = Field(default=False) + pii_entities_config: Optional[Dict[PiiEntityType, PiiAction]] = None def __init__(self, **kwargs): default_on = kwargs.get("default_on") @@ -436,6 +438,7 @@ class ApplyGuardrailResponse(BaseModel): class PatchGuardrailLitellmParams(BaseModel): default_on: Optional[bool] = None + pii_entities_config: Optional[Dict[PiiEntityType, PiiAction]] = None class PatchGuardrailRequest(BaseModel): diff --git a/tests/litellm/proxy/guardrails/test_init_guardrails.py b/tests/litellm/proxy/guardrails/test_init_guardrails.py index 782f9d69ae..a511229942 100644 --- a/tests/litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/litellm/proxy/guardrails/test_init_guardrails.py @@ -9,7 +9,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.proxy.guardrails.init_guardrails import InitializeGuardrails +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -30,7 +30,8 @@ def test_initialize_presidio_guardrail(): } # Call the initialize_guardrail method - result = InitializeGuardrails.initialize_guardrail( + guardrail_handler = InMemoryGuardrailHandler() + result = guardrail_handler.initialize_guardrail( guardrail=test_guardrail, ) diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx index eae0097b98..9fdcb1badf 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx @@ -13,10 +13,11 @@ import { TabPanels, TextInput, } from "@tremor/react"; -import { Button, Form, Input, Select, message, Tooltip } from "antd"; +import { Button, Form, Input, Select, message, Tooltip, Divider } from "antd"; import { InfoCircleOutlined } from '@ant-design/icons'; -import { getGuardrailInfo, updateGuardrailCall } from "@/components/networking"; +import { getGuardrailInfo, updateGuardrailCall, getGuardrailUISettings } from "@/components/networking"; import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import PiiConfiguration from "./pii_configuration"; export interface GuardrailInfoProps { guardrailId: string; @@ -35,6 +36,17 @@ const GuardrailInfoView: React.FC = ({ const [loading, setLoading] = useState(true); const [isEditing, setIsEditing] = useState(false); const [form] = Form.useForm(); + const [selectedPiiEntities, setSelectedPiiEntities] = useState([]); + const [selectedPiiActions, setSelectedPiiActions] = useState<{[key: string]: string}>({}); + const [guardrailSettings, setGuardrailSettings] = useState<{ + supported_entities: string[]; + supported_actions: string[]; + pii_entity_categories: Array<{ + category: string; + entities: string[]; + }>; + supported_modes: string[]; + } | null>(null); const fetchGuardrailInfo = async () => { try { @@ -42,6 +54,33 @@ const GuardrailInfoView: React.FC = ({ if (!accessToken) return; const response = await getGuardrailInfo(accessToken, guardrailId); setGuardrailData(response); + + // Initialize PII configuration from guardrail data + if (response.litellm_params?.pii_entities_config) { + const piiConfig = response.litellm_params.pii_entities_config; + + // Clear previous selections + setSelectedPiiEntities([]); + setSelectedPiiActions({}); + + // Only if there are entities configured + if (Object.keys(piiConfig).length > 0) { + const entities: string[] = []; + const actions: {[key: string]: string} = {}; + + Object.entries(piiConfig).forEach(([entity, action]: [string, any]) => { + entities.push(entity); + actions[entity] = typeof action === 'string' ? action : "MASK"; + }); + + setSelectedPiiEntities(entities); + setSelectedPiiActions(actions); + } + } else { + // Clear selections if no PII config exists + setSelectedPiiEntities([]); + setSelectedPiiActions({}); + } } catch (error) { message.error("Failed to load guardrail information"); console.error("Error fetching guardrail info:", error); @@ -50,23 +89,67 @@ const GuardrailInfoView: React.FC = ({ } }; + const fetchGuardrailUISettings = async () => { + try { + if (!accessToken) return; + const uiSettings = await getGuardrailUISettings(accessToken); + setGuardrailSettings(uiSettings); + } catch (error) { + console.error("Error fetching guardrail UI settings:", error); + } + }; + useEffect(() => { fetchGuardrailInfo(); + fetchGuardrailUISettings(); }, [guardrailId, accessToken]); + const handlePiiEntitySelect = (entity: string) => { + setSelectedPiiEntities(prev => { + if (prev.includes(entity)) { + return prev.filter(e => e !== entity); + } else { + return [...prev, entity]; + } + }); + }; + + const handlePiiActionSelect = (entity: string, action: string) => { + setSelectedPiiActions(prev => ({ + ...prev, + [entity]: action + })); + }; const handleGuardrailUpdate = async (values: any) => { try { if (!accessToken) return; - const updateData = { + // Prepare update data object + const updateData: any = { guardrail_name: values.guardrail_name, litellm_params: { - default_on: values.default_on + default_on: values.default_on, }, guardrail_info: values.guardrail_info ? JSON.parse(values.guardrail_info) : undefined }; + // Only add PII entities config if we have selected entities + if (selectedPiiEntities.length > 0) { + // Create PII config object only with selected entities + const piiEntitiesConfig: {[key: string]: string} = {}; + selectedPiiEntities.forEach(entity => { + piiEntitiesConfig[entity] = selectedPiiActions[entity] || "MASK"; + }); + + // Add to litellm_params only if we have entities + updateData.litellm_params.pii_entities_config = piiEntitiesConfig; + } else { + // If no entities selected, explicitly set to empty object + // This will clear any existing PII config + updateData.litellm_params.pii_entities_config = {}; + } + await updateGuardrailCall(accessToken, guardrailId, updateData); message.success("Guardrail updated successfully"); fetchGuardrailInfo(); @@ -152,6 +235,17 @@ const GuardrailInfoView: React.FC = ({ + {guardrailData.litellm_params?.pii_entities_config && Object.keys(guardrailData.litellm_params.pii_entities_config).length > 0 && ( + +
+ PII Protection + + {Object.keys(guardrailData.litellm_params.pii_entities_config).length} PII entities configured + +
+ + )} + {guardrailData.guardrail_info && Object.keys(guardrailData.guardrail_info).length > 0 && ( Guardrail Info @@ -216,7 +310,23 @@ const GuardrailInfoView: React.FC = ({ No + + PII Protection +
+ {guardrailSettings && ( + + )} +
+ Advanced Settings = ({ {guardrailData.litellm_params?.default_on ? "Yes" : "No"} + + {guardrailData.litellm_params?.pii_entities_config && Object.keys(guardrailData.litellm_params.pii_entities_config).length > 0 && ( +
+ PII Protection +
+ + {Object.keys(guardrailData.litellm_params.pii_entities_config).length} PII entities configured + +
+
+ )} +
Created At
{formatDate(guardrailData.created_at)}
diff --git a/ui/litellm-dashboard/src/components/guardrails/pii_configuration.tsx b/ui/litellm-dashboard/src/components/guardrails/pii_configuration.tsx index 23873803fe..ffe2543ea6 100644 --- a/ui/litellm-dashboard/src/components/guardrails/pii_configuration.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/pii_configuration.tsx @@ -52,6 +52,9 @@ const PiiConfiguration: React.FC = ({ // Unselect all entities const handleUnselectAll = () => { + // Instead of iterating through each entity and toggling, + // we'll directly set the selected entities to an empty array + // This is more reliable and ensures a clean slate selectedEntities.forEach(entity => { onEntitySelect(entity); });