知微
fa45d8aa5f
- health_checklist.json: 192.168.1.122→node122
- ocr_client.py: docstring IP→node122
- docs/market-data-requirements.md: IP→node122
- 所有API调用通过ProxyHandler({})绕过系统代理
Privoxy对node122:18003返回500，直连正常
799 lines
30 KiB
Python
799 lines
30 KiB
Python
# litellm/proxy/guardrails/guardrail_registry.py
|
|
|
|
import importlib
|
|
import os
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Dict, List, Literal, Optional, Set, Type, cast
|
|
|
|
from pydantic import ValidationError
|
|
|
|
import litellm
|
|
from litellm import Router
|
|
from litellm._logging import verbose_proxy_logger
|
|
from litellm._uuid import uuid
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
|
from litellm.proxy.guardrails.guardrail_hooks.grayswan import (
|
|
GraySwanGuardrail,
|
|
)
|
|
from litellm.proxy.guardrails.guardrail_hooks.grayswan import (
|
|
initialize_guardrail as initialize_grayswan,
|
|
)
|
|
from litellm.proxy.types_utils.utils import get_instance_fn
|
|
from litellm.proxy.utils import PrismaClient
|
|
from litellm.repositories.table_repositories import GuardrailsRepository
|
|
from litellm.secret_managers.main import get_secret
|
|
from litellm.types.guardrails import (
|
|
Guardrail,
|
|
GuardrailEventHooks,
|
|
LakeraCategoryThresholds,
|
|
LitellmParams,
|
|
SupportedGuardrailIntegrations,
|
|
)
|
|
|
|
from .guardrail_hooks.llm_as_a_judge import (
|
|
initialize_guardrail as initialize_llm_as_a_judge,
|
|
)
|
|
from .guardrail_initializers import (
|
|
initialize_bedrock,
|
|
initialize_hide_secrets,
|
|
initialize_lakera,
|
|
initialize_lakera_v2,
|
|
initialize_presidio,
|
|
initialize_tool_permission,
|
|
)
|
|
|
|
guardrail_initializer_registry = {
|
|
SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock,
|
|
SupportedGuardrailIntegrations.LAKERA.value: initialize_lakera,
|
|
SupportedGuardrailIntegrations.LAKERA_V2.value: initialize_lakera_v2,
|
|
SupportedGuardrailIntegrations.PRESIDIO.value: initialize_presidio,
|
|
SupportedGuardrailIntegrations.HIDE_SECRETS.value: initialize_hide_secrets,
|
|
SupportedGuardrailIntegrations.TOOL_PERMISSION.value: initialize_tool_permission,
|
|
SupportedGuardrailIntegrations.GRAYSWAN.value: initialize_grayswan,
|
|
SupportedGuardrailIntegrations.LLM_AS_A_JUDGE.value: initialize_llm_as_a_judge,
|
|
}
|
|
|
|
guardrail_class_registry: Dict[str, Type[CustomGuardrail]] = {
|
|
SupportedGuardrailIntegrations.GRAYSWAN.value: GraySwanGuardrail
|
|
}
|
|
|
|
|
|
def get_guardrail_initializer_from_hooks():
|
|
"""
|
|
Get guardrail initializers by discovering them from the guardrail_hooks directory structure.
|
|
|
|
Scans the guardrail_hooks directory for subdirectories containing __init__.py files
|
|
with either guardrail_initializer_registry or initialize_guardrail functions.
|
|
|
|
Returns:
|
|
Dict[str, Callable]: A dictionary mapping guardrail types to their initializer functions
|
|
"""
|
|
discovered_initializers = {}
|
|
|
|
try:
|
|
# Get the path to the guardrail_hooks directory
|
|
current_dir = os.path.dirname(__file__)
|
|
hooks_dir = os.path.join(current_dir, "guardrail_hooks")
|
|
|
|
if not os.path.exists(hooks_dir):
|
|
verbose_proxy_logger.debug("guardrail_hooks directory not found")
|
|
return discovered_initializers
|
|
|
|
# Scan each subdirectory in guardrail_hooks
|
|
for item in os.listdir(hooks_dir):
|
|
item_path = os.path.join(hooks_dir, item)
|
|
|
|
# Skip files and __pycache__ directories
|
|
if not os.path.isdir(item_path) or item.startswith("__"):
|
|
continue
|
|
|
|
# Check if the directory has an __init__.py file
|
|
init_file = os.path.join(item_path, "__init__.py")
|
|
if not os.path.exists(init_file):
|
|
continue
|
|
|
|
module_path = f"litellm.proxy.guardrails.guardrail_hooks.{item}"
|
|
try:
|
|
# Import the module
|
|
verbose_proxy_logger.debug(f"Discovering guardrails in: {module_path}")
|
|
|
|
module = importlib.import_module(module_path)
|
|
|
|
# Check for guardrail_initializer_registry dictionary
|
|
if hasattr(module, "guardrail_initializer_registry"):
|
|
registry = getattr(module, "guardrail_initializer_registry")
|
|
if isinstance(registry, dict):
|
|
discovered_initializers.update(registry)
|
|
verbose_proxy_logger.debug(
|
|
f"Found guardrail_initializer_registry in {module_path}: {list(registry.keys())}"
|
|
)
|
|
|
|
# Check for standalone initialize_guardrail function (fallback for directory-based guardrails)
|
|
elif hasattr(module, "initialize_guardrail"):
|
|
# For directories with just initialize_guardrail, use the directory name as the key
|
|
initialize_fn = getattr(module, "initialize_guardrail")
|
|
discovered_initializers[item] = initialize_fn
|
|
verbose_proxy_logger.debug(
|
|
f"Found initialize_guardrail function in {module_path}"
|
|
)
|
|
|
|
except ImportError as e:
|
|
verbose_proxy_logger.error(f"Could not import {module_path}: {e}")
|
|
continue
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(f"Error processing {module_path}: {e}")
|
|
continue
|
|
|
|
verbose_proxy_logger.debug(
|
|
f"Discovered {len(discovered_initializers)} guardrail initializers: {list(discovered_initializers.keys())}"
|
|
)
|
|
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(f"Error discovering guardrail initializers: {e}")
|
|
|
|
return discovered_initializers
|
|
|
|
|
|
def get_guardrail_class_from_hooks():
|
|
"""
|
|
Get guardrail classes by discovering them from the guardrail_hooks directory structure.
|
|
"""
|
|
"""
|
|
Get guardrail initializers by discovering them from the guardrail_hooks directory structure.
|
|
|
|
Scans the guardrail_hooks directory for subdirectories containing __init__.py files
|
|
with either guardrail_initializer_registry or initialize_guardrail functions.
|
|
|
|
Returns:
|
|
Dict[str, Callable]: A dictionary mapping guardrail types to their initializer functions
|
|
"""
|
|
discovered_classes = {}
|
|
|
|
try:
|
|
# Get the path to the guardrail_hooks directory
|
|
current_dir = os.path.dirname(__file__)
|
|
hooks_dir = os.path.join(current_dir, "guardrail_hooks")
|
|
|
|
if not os.path.exists(hooks_dir):
|
|
verbose_proxy_logger.debug("guardrail_hooks directory not found")
|
|
return discovered_classes
|
|
|
|
# Scan each subdirectory in guardrail_hooks
|
|
for item in os.listdir(hooks_dir):
|
|
item_path = os.path.join(hooks_dir, item)
|
|
|
|
# Skip files and __pycache__ directories
|
|
if not os.path.isdir(item_path) or item.startswith("__"):
|
|
continue
|
|
|
|
# Check if the directory has an __init__.py file
|
|
init_file = os.path.join(item_path, "__init__.py")
|
|
|
|
if not os.path.exists(init_file):
|
|
continue
|
|
|
|
module_path = f"litellm.proxy.guardrails.guardrail_hooks.{item}"
|
|
|
|
try:
|
|
# Import the module
|
|
verbose_proxy_logger.debug(f"Discovering guardrails in: {module_path}")
|
|
|
|
module = importlib.import_module(module_path)
|
|
|
|
# Check for guardrail_initializer_registry dictionary
|
|
if hasattr(module, "guardrail_class_registry"):
|
|
registry = getattr(module, "guardrail_class_registry")
|
|
if isinstance(registry, dict):
|
|
discovered_classes.update(registry)
|
|
|
|
except ImportError as e:
|
|
verbose_proxy_logger.debug(f"Could not import {module_path}: {e}")
|
|
continue
|
|
except Exception as e:
|
|
verbose_proxy_logger.exception(f"Error processing {module_path}: {e}")
|
|
continue
|
|
|
|
except Exception as e:
|
|
verbose_proxy_logger.error(f"Error discovering guardrail initializers: {e}")
|
|
|
|
return discovered_classes
|
|
|
|
|
|
guardrail_class_registry.update(get_guardrail_class_from_hooks())
|
|
|
|
|
|
# Merge with dynamically discovered guardrail initializers
|
|
_discovered_initializers = get_guardrail_initializer_from_hooks()
|
|
|
|
guardrail_initializer_registry.update(_discovered_initializers)
|
|
|
|
|
|
class GuardrailRegistry:
|
|
"""
|
|
Registry for guardrails
|
|
|
|
Handles adding, removing, and getting guardrails in DB + in memory
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
###########################################################
|
|
########### In memory management helpers for guardrails ###########
|
|
############################################################
|
|
def get_initialized_guardrail_callback(
|
|
self, guardrail_name: str
|
|
) -> Optional[CustomGuardrail]:
|
|
"""
|
|
Returns the initialized guardrail callback for a given guardrail name
|
|
"""
|
|
active_guardrails = (
|
|
litellm.logging_callback_manager.get_custom_loggers_for_type(
|
|
callback_type=CustomGuardrail
|
|
)
|
|
)
|
|
for active_guardrail in active_guardrails:
|
|
if isinstance(active_guardrail, CustomGuardrail):
|
|
if active_guardrail.guardrail_name == guardrail_name:
|
|
return active_guardrail
|
|
return None
|
|
|
|
###########################################################
|
|
########### DB management helpers for guardrails ###########
|
|
############################################################
|
|
async def add_guardrail_to_db(
|
|
self, guardrail: Guardrail, prisma_client: PrismaClient
|
|
):
|
|
"""
|
|
Add a guardrail to the database
|
|
"""
|
|
try:
|
|
guardrail_name = guardrail.get("guardrail_name")
|
|
# Properly serialize LitellmParams Pydantic model to dict
|
|
litellm_params_obj: Any = guardrail.get("litellm_params", {})
|
|
if hasattr(litellm_params_obj, "model_dump"):
|
|
litellm_params_dict = litellm_params_obj.model_dump()
|
|
else:
|
|
litellm_params_dict = (
|
|
dict(litellm_params_obj) if litellm_params_obj else {}
|
|
)
|
|
litellm_params: str = safe_dumps(litellm_params_dict)
|
|
guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {}))
|
|
|
|
# Create guardrail in DB
|
|
created_guardrail = await GuardrailsRepository(prisma_client).table.create(
|
|
data={
|
|
"guardrail_name": guardrail_name,
|
|
"litellm_params": litellm_params,
|
|
"guardrail_info": guardrail_info,
|
|
"created_at": datetime.now(timezone.utc),
|
|
"updated_at": datetime.now(timezone.utc),
|
|
}
|
|
)
|
|
|
|
# Add guardrail_id to the returned guardrail object
|
|
guardrail_dict = dict(guardrail)
|
|
guardrail_dict["guardrail_id"] = created_guardrail.guardrail_id
|
|
|
|
return guardrail_dict
|
|
except Exception as e:
|
|
raise Exception(f"Error adding guardrail to DB: {str(e)}")
|
|
|
|
async def delete_guardrail_from_db(
|
|
self, guardrail_id: str, prisma_client: PrismaClient
|
|
):
|
|
"""
|
|
Delete a guardrail from the database
|
|
"""
|
|
try:
|
|
# Delete from DB
|
|
await GuardrailsRepository(prisma_client).table.delete(
|
|
where={"guardrail_id": guardrail_id}
|
|
)
|
|
|
|
return {"message": f"Guardrail {guardrail_id} deleted successfully"}
|
|
except Exception as e:
|
|
raise Exception(f"Error deleting guardrail from DB: {str(e)}")
|
|
|
|
async def update_guardrail_in_db(
|
|
self, guardrail_id: str, guardrail: Guardrail, prisma_client: PrismaClient
|
|
):
|
|
"""
|
|
Update a guardrail in the database
|
|
"""
|
|
try:
|
|
guardrail_name = guardrail.get("guardrail_name")
|
|
# Properly serialize LitellmParams Pydantic model to dict
|
|
litellm_params_obj: Any = guardrail.get("litellm_params", {})
|
|
if hasattr(litellm_params_obj, "model_dump"):
|
|
litellm_params_dict = litellm_params_obj.model_dump()
|
|
else:
|
|
litellm_params_dict = (
|
|
dict(litellm_params_obj) if litellm_params_obj else {}
|
|
)
|
|
litellm_params: str = safe_dumps(litellm_params_dict)
|
|
guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {}))
|
|
|
|
# Update in DB
|
|
updated_guardrail = await GuardrailsRepository(prisma_client).table.update(
|
|
where={"guardrail_id": guardrail_id},
|
|
data={
|
|
"guardrail_name": guardrail_name,
|
|
"litellm_params": litellm_params,
|
|
"guardrail_info": guardrail_info,
|
|
"updated_at": datetime.now(timezone.utc),
|
|
},
|
|
)
|
|
|
|
# Convert to dict and return
|
|
return dict(updated_guardrail)
|
|
except Exception as e:
|
|
raise Exception(f"Error updating guardrail in DB: {str(e)}")
|
|
|
|
@staticmethod
|
|
async def get_all_guardrails_from_db(
|
|
prisma_client: PrismaClient,
|
|
) -> List[Guardrail]:
|
|
"""
|
|
Get all active guardrails from the database.
|
|
Only rows with status == "active" are returned (pending_review and rejected are excluded).
|
|
"""
|
|
try:
|
|
guardrails_from_db = await GuardrailsRepository(
|
|
prisma_client
|
|
).table.find_many(
|
|
where={"status": "active"},
|
|
order={"created_at": "desc"},
|
|
)
|
|
|
|
guardrails: List[Guardrail] = []
|
|
for guardrail in guardrails_from_db:
|
|
guardrails.append(Guardrail(**(dict(guardrail)))) # type: ignore
|
|
|
|
return guardrails
|
|
except Exception as e:
|
|
raise Exception(f"Error getting guardrails from DB: {str(e)}")
|
|
|
|
async def get_guardrail_by_id_from_db(
|
|
self, guardrail_id: str, prisma_client: PrismaClient
|
|
) -> Optional[Guardrail]:
|
|
"""
|
|
Get a guardrail by its ID from the database
|
|
"""
|
|
try:
|
|
guardrail = await GuardrailsRepository(prisma_client).table.find_unique(
|
|
where={"guardrail_id": guardrail_id}
|
|
)
|
|
|
|
if not guardrail:
|
|
return None
|
|
|
|
return Guardrail(**(dict(guardrail))) # type: ignore
|
|
except Exception as e:
|
|
raise Exception(f"Error getting guardrail from DB: {str(e)}")
|
|
|
|
async def get_guardrail_by_name_from_db(
|
|
self, guardrail_name: str, prisma_client: PrismaClient
|
|
) -> Optional[Guardrail]:
|
|
"""
|
|
Get a guardrail by its name from the database
|
|
"""
|
|
try:
|
|
guardrail = await GuardrailsRepository(prisma_client).table.find_unique(
|
|
where={"guardrail_name": guardrail_name}
|
|
)
|
|
|
|
if not guardrail:
|
|
return None
|
|
|
|
return Guardrail(**(dict(guardrail))) # type: ignore
|
|
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
|
|
"""
|
|
|
|
self._sources: Dict[str, Literal["db", "config"]] = {}
|
|
"""
|
|
Guardrail id to provenance marker. "db" entries are reconciled against
|
|
the DB on each polling tick; "config" entries are owned by proxy_config.yaml
|
|
and never deleted by reconciliation.
|
|
"""
|
|
|
|
def initialize_guardrail(
|
|
self,
|
|
guardrail: Guardrail,
|
|
config_file_path: Optional[str] = None,
|
|
llm_router: Optional["Router"] = None,
|
|
source: Literal["db", "config"] = "config",
|
|
) -> 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())
|
|
guardrail["guardrail_id"] = guardrail_id
|
|
if guardrail_id in self.IN_MEMORY_GUARDRAILS:
|
|
verbose_proxy_logger.debug(
|
|
"guardrail_id already exists in IN_MEMORY_GUARDRAILS"
|
|
)
|
|
# Honor the caller's source even on the early-return path so a
|
|
# racing polling tick or a hot-reload of config can correct an
|
|
# entry's provenance.
|
|
self._sources[guardrail_id] = source
|
|
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)
|
|
|
|
if isinstance(litellm_params_data, dict):
|
|
litellm_params = LitellmParams(**litellm_params_data)
|
|
else:
|
|
litellm_params = 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:
|
|
# Try to call with llm_router first, fall back to without if it fails
|
|
import inspect
|
|
|
|
sig = inspect.signature(initializer)
|
|
if "llm_router" in sig.parameters:
|
|
custom_guardrail_callback = initializer(
|
|
litellm_params, guardrail, llm_router # type: ignore
|
|
)
|
|
else:
|
|
custom_guardrail_callback = initializer(litellm_params, guardrail)
|
|
elif isinstance(guardrail_type, str) and "." in guardrail_type:
|
|
custom_guardrail_callback = self.initialize_custom_guardrail(
|
|
guardrail=cast(dict, guardrail),
|
|
guardrail_type=guardrail_type,
|
|
litellm_params=litellm_params,
|
|
config_file_path=config_file_path,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported guardrail: {guardrail_type}")
|
|
|
|
if custom_guardrail_callback is not None:
|
|
setattr(
|
|
custom_guardrail_callback,
|
|
"skip_system_message_in_guardrail",
|
|
getattr(litellm_params, "skip_system_message_in_guardrail", None),
|
|
)
|
|
setattr(
|
|
custom_guardrail_callback,
|
|
"skip_tool_message_in_guardrail",
|
|
getattr(litellm_params, "skip_tool_message_in_guardrail", None),
|
|
)
|
|
|
|
parsed_guardrail = Guardrail(
|
|
guardrail_id=guardrail.get("guardrail_id"),
|
|
guardrail_name=guardrail["guardrail_name"],
|
|
litellm_params=litellm_params,
|
|
)
|
|
|
|
# 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
|
|
self._sources[guardrail_id] = source
|
|
|
|
return parsed_guardrail
|
|
|
|
def initialize_custom_guardrail(
|
|
self,
|
|
guardrail: Dict,
|
|
guardrail_type: str,
|
|
litellm_params: LitellmParams,
|
|
config_file_path: Optional[str] = None,
|
|
) -> Optional[CustomGuardrail]:
|
|
"""
|
|
Initialize a Custom Guardrail from a python file or module path
|
|
|
|
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"
|
|
)
|
|
|
|
verbose_proxy_logger.debug(
|
|
"Initializing custom guardrail: %s",
|
|
guardrail_type,
|
|
)
|
|
|
|
_guardrail_class = get_instance_fn(
|
|
guardrail_type, config_file_path=config_file_path
|
|
)
|
|
|
|
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
|
|
|
|
# Extract additional params from litellm_params to pass to custom guardrail
|
|
# This matches the behavior of other guardrail initializers (e.g., initialize_lakera)
|
|
# and aligns with the documented behavior for custom guardrails
|
|
if hasattr(litellm_params, "model_dump"):
|
|
extra_params = litellm_params.model_dump(exclude_none=True)
|
|
else:
|
|
extra_params = dict(litellm_params) if litellm_params else {}
|
|
|
|
# Remove params that are handled explicitly or are internal
|
|
for key in ["guardrail", "mode", "default_on"]:
|
|
extra_params.pop(key, None)
|
|
|
|
_guardrail_callback = _guardrail_class(
|
|
guardrail_name=guardrail["guardrail_name"],
|
|
event_hook=mode,
|
|
default_on=default_on,
|
|
**extra_params,
|
|
)
|
|
litellm.logging_callback_manager.add_litellm_callback(_guardrail_callback) # type: ignore
|
|
|
|
return _guardrail_callback
|
|
|
|
def update_in_memory_guardrail(
|
|
self,
|
|
guardrail_id: str,
|
|
guardrail: Guardrail,
|
|
source: Literal["db", "config"] = "db",
|
|
) -> 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
|
|
self._sources[guardrail_id] = source
|
|
|
|
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 and remove from litellm callbacks.
|
|
|
|
The callback is purged from every callback list, not just
|
|
litellm.callbacks: request handling promotes guardrail callbacks into the
|
|
success/failure/async lists, so removing it from only litellm.callbacks
|
|
leaves the old instance stranded in those lists on every re-initialization.
|
|
"""
|
|
# Remove from in-memory storage
|
|
self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None)
|
|
self._sources.pop(guardrail_id, None)
|
|
|
|
custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.pop(
|
|
guardrail_id, None
|
|
)
|
|
if custom_guardrail_callback is None:
|
|
return
|
|
|
|
litellm.logging_callback_manager.remove_callback_from_all_lists(
|
|
custom_guardrail_callback
|
|
)
|
|
|
|
def list_in_memory_guardrails(self) -> List[Guardrail]:
|
|
"""
|
|
List all guardrails in memory
|
|
"""
|
|
return list(self.IN_MEMORY_GUARDRAILS.values())
|
|
|
|
def get_guardrail_by_id(self, guardrail_id: str) -> Optional[Guardrail]:
|
|
"""
|
|
Get a guardrail by its ID from memory
|
|
"""
|
|
return self.IN_MEMORY_GUARDRAILS.get(guardrail_id)
|
|
|
|
def get_source(self, guardrail_id: str) -> Optional[Literal["db", "config"]]:
|
|
"""
|
|
Return the provenance of an in-memory guardrail.
|
|
"""
|
|
return self._sources.get(guardrail_id)
|
|
|
|
def reconcile_db_guardrails(self, db_guardrail_ids: Set[str]) -> List[str]:
|
|
"""
|
|
Drop in-memory entries that originated from the DB but are no longer
|
|
present in db_guardrail_ids. Config-loaded guardrails are never touched.
|
|
|
|
Called by the periodic DB polling tick so that a guardrail deleted
|
|
on another pod is eventually purged from this pod's memory + callbacks.
|
|
"""
|
|
stale_ids = [
|
|
guardrail_id
|
|
for guardrail_id, source in self._sources.items()
|
|
if source == "db" and guardrail_id not in db_guardrail_ids
|
|
]
|
|
for guardrail_id in stale_ids:
|
|
verbose_proxy_logger.info(
|
|
"Reconcile: removing stale DB-backed guardrail '%s' from memory "
|
|
"(deleted in DB by another pod)",
|
|
guardrail_id,
|
|
)
|
|
self.delete_in_memory_guardrail(guardrail_id)
|
|
return stale_ids
|
|
|
|
@staticmethod
|
|
def _normalize_litellm_params_for_comparison(
|
|
params: Optional[Any],
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Render litellm_params to a canonical dict so an in-memory LitellmParams and
|
|
the raw dict loaded from the DB compare equal when they describe the same
|
|
config. The in-memory side is a LitellmParams whose model_dump() carries
|
|
every field default and coerces enums, while the DB side is the raw stored
|
|
dict holding only the keys originally provided. Comparing those two shapes
|
|
directly never matches, so each DB poll would re-initialize the guardrail
|
|
forever; normalizing both through LitellmParams keeps the diff meaningful.
|
|
"""
|
|
if params is None:
|
|
return None
|
|
if isinstance(params, LitellmParams):
|
|
return params.model_dump()
|
|
if isinstance(params, dict):
|
|
try:
|
|
return LitellmParams(**params).model_dump()
|
|
except ValidationError as e:
|
|
verbose_proxy_logger.warning(
|
|
f"Could not normalize guardrail litellm_params for comparison; "
|
|
f"treating the guardrail as changed. Error: {e}"
|
|
)
|
|
return params
|
|
return params
|
|
|
|
def _has_guardrail_params_changed(
|
|
self, guardrail_id: str, new_guardrail: Guardrail
|
|
) -> bool:
|
|
"""
|
|
Check if guardrail params or name have changed compared to in-memory version.
|
|
Returns True if params/name changed or guardrail doesn't exist in memory.
|
|
"""
|
|
existing = self.IN_MEMORY_GUARDRAILS.get(guardrail_id)
|
|
if existing is None:
|
|
return True
|
|
|
|
# Compare guardrail_name
|
|
if existing.get("guardrail_name") != new_guardrail.get("guardrail_name"):
|
|
return True
|
|
|
|
# Compare litellm_params
|
|
existing_dict = self._normalize_litellm_params_for_comparison(
|
|
existing.get("litellm_params")
|
|
)
|
|
new_dict = self._normalize_litellm_params_for_comparison(
|
|
new_guardrail.get("litellm_params")
|
|
)
|
|
|
|
# Compare and identify specific differences
|
|
changed_fields = {}
|
|
if existing_dict is not None and new_dict is not None:
|
|
all_keys = set(existing_dict.keys()) | set(new_dict.keys())
|
|
for key in all_keys:
|
|
old_val = existing_dict.get(key)
|
|
new_val = new_dict.get(key)
|
|
if old_val != new_val:
|
|
changed_fields[key] = {"old": old_val, "new": new_val}
|
|
elif existing_dict != new_dict:
|
|
changed_fields = {"litellm_params": {"old": existing_dict, "new": new_dict}}
|
|
|
|
# Log differences if any found
|
|
if changed_fields:
|
|
verbose_proxy_logger.debug(
|
|
f"Guardrail params changed. Differences: {changed_fields}"
|
|
)
|
|
|
|
# Return True if any fields changed
|
|
return len(changed_fields) > 0
|
|
|
|
def reinitialize_guardrail(
|
|
self,
|
|
guardrail: Guardrail,
|
|
config_file_path: Optional[str] = None,
|
|
source: Literal["db", "config"] = "config",
|
|
) -> Optional[Guardrail]:
|
|
"""
|
|
Force re-initialization of a guardrail even if it exists in memory.
|
|
Removes old callback from litellm.callbacks and creates fresh instance.
|
|
"""
|
|
guardrail_id = guardrail.get("guardrail_id")
|
|
if not guardrail_id:
|
|
verbose_proxy_logger.error(
|
|
"Cannot reinitialize guardrail without guardrail_id"
|
|
)
|
|
return None
|
|
|
|
# Remove from memory if exists (also removes from callbacks)
|
|
if guardrail_id in self.IN_MEMORY_GUARDRAILS:
|
|
self.delete_in_memory_guardrail(guardrail_id)
|
|
|
|
# Initialize fresh (will add new callback to litellm.callbacks)
|
|
return self.initialize_guardrail(
|
|
guardrail=guardrail, config_file_path=config_file_path, source=source
|
|
)
|
|
|
|
def sync_guardrail_from_db(
|
|
self, guardrail: Guardrail, config_file_path: Optional[str] = None
|
|
) -> Optional[Guardrail]:
|
|
"""
|
|
Sync a guardrail from DB - initializes if new, re-initializes if changed.
|
|
This is the method to call during DB polling.
|
|
"""
|
|
guardrail_id = guardrail.get("guardrail_id")
|
|
if not guardrail_id:
|
|
verbose_proxy_logger.error("Cannot sync guardrail without guardrail_id")
|
|
return None
|
|
|
|
if self._has_guardrail_params_changed(guardrail_id, guardrail):
|
|
guardrail_name = guardrail.get("guardrail_name", "Unknown")
|
|
verbose_proxy_logger.info(
|
|
f"Guardrail '{guardrail_name}' (ID: {guardrail_id}) params changed, re-initializing..."
|
|
)
|
|
return self.reinitialize_guardrail(
|
|
guardrail=guardrail,
|
|
config_file_path=config_file_path,
|
|
source="db",
|
|
)
|
|
|
|
# Params unchanged but the entry is still DB-backed; make sure the
|
|
# source marker reflects that even if it was previously set differently
|
|
# (e.g. a config entry whose UUID later collided with a DB row).
|
|
self._sources[guardrail_id] = "db"
|
|
return self.IN_MEMORY_GUARDRAILS.get(guardrail_id)
|
|
|
|
|
|
########################################################
|
|
# In Memory Guardrail Handler for LiteLLM Proxy
|
|
########################################################
|
|
IN_MEMORY_GUARDRAIL_HANDLER = InMemoryGuardrailHandler()
|
|
########################################################
|