tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-20 04:23:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

[QA/Guardrails] Allow editing guardrail entity configs (#10918)

Browse Source 
* allow editing guardrail entity

* fix: allow patching guardrail entity config

* fix: test init guardrails

* fix: update in mem guardrail params

* fix: update in mem guardrail params

* allow editing guardrail entities on ui

* fix pii entity config

* fixes for unselect/select entities

* fix: test fixes
This commit is contained in:
Ishaan Jaff
2025-05-17 16:27:02 -07:00
committed by GitHub
parent d79cd70a1a
commit 50a6f289cd
12 changed files with 392 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/custom_guardrail.py
+7
View File
@@ -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):
"""
litellm/proxy/guardrails/guardrail_endpoints.py
+28 -6
View File
@@ -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
litellm/proxy/guardrails/guardrail_hooks/presidio.py
+8
View File
@@ -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
litellm/proxy/guardrails/guardrail_initializers.py
+11
View File
@@ -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
litellm/proxy/guardrails/guardrail_registry.py
+189 -2
View File
@@ -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()
########################################################
litellm/proxy/guardrails/init_guardrails.py
+4 -129
View File
@@ -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
litellm/proxy/proxy_config.yaml
-5
View File
@@ -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
litellm/proxy/proxy_server.py
+9 -5
View File
@@ -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()
litellm/types/guardrails.py
+4 -1
View File
@@ -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):
tests/litellm/proxy/guardrails/test_init_guardrails.py
+3 -2
View File
@@ -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,
)
ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx
+126 -4
View File
@@ -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<GuardrailInfoProps> = ({
const [loading, setLoading] = useState(true);
const [isEditing, setIsEditing] = useState(false);
const [form] = Form.useForm();
const [selectedPiiEntities, setSelectedPiiEntities] = useState<string[]>([]);
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<GuardrailInfoProps> = ({
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<GuardrailInfoProps> = ({
}
};
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<GuardrailInfoProps> = ({
</Card>
</Grid>
{guardrailData.litellm_params?.pii_entities_config && Object.keys(guardrailData.litellm_params.pii_entities_config).length > 0 && (
<Card className="mt-6">
<div className="flex justify-between items-center">
<Text className="font-medium">PII Protection</Text>
<Badge color="blue">
{Object.keys(guardrailData.litellm_params.pii_entities_config).length} PII entities configured
</Badge>
</div>
</Card>
)}
{guardrailData.guardrail_info && Object.keys(guardrailData.guardrail_info).length > 0 && (
<Card className="mt-6">
<Text>Guardrail Info</Text>
@@ -216,7 +310,23 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({
<Select.Option value={false}>No</Select.Option>
</Select>
</Form.Item>
<Divider orientation="left">PII Protection</Divider>
<div className="mb-6">
{guardrailSettings && (
<PiiConfiguration
entities={guardrailSettings.supported_entities}
actions={guardrailSettings.supported_actions}
selectedEntities={selectedPiiEntities}
selectedActions={selectedPiiActions}
onEntitySelect={handlePiiEntitySelect}
onActionSelect={handlePiiActionSelect}
entityCategories={guardrailSettings.pii_entity_categories}
/>
)}
</div>
<Divider orientation="left">Advanced Settings</Divider>
<Form.Item
label="Guardrail Information"
name="guardrail_info"
@@ -257,6 +367,18 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({
{guardrailData.litellm_params?.default_on ? "Yes" : "No"}
</Badge>
</div>
{guardrailData.litellm_params?.pii_entities_config && Object.keys(guardrailData.litellm_params.pii_entities_config).length > 0 && (
<div>
<Text className="font-medium">PII Protection</Text>
<div className="mt-2">
<Badge color="blue">
{Object.keys(guardrailData.litellm_params.pii_entities_config).length} PII entities configured
</Badge>
</div>
</div>
)}
<div>
<Text className="font-medium">Created At</Text>
<div>{formatDate(guardrailData.created_at)}</div>
ui/litellm-dashboard/src/components/guardrails/pii_configuration.tsx
+3
View File
@@ -52,6 +52,9 @@ const PiiConfiguration: React.FC<PiiConfigurationProps> = ({
// 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);
});