From bda9eecd45da1b4d3af20ba092abd0ffac405b8b Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 5 Jul 2025 18:44:00 -0700 Subject: [PATCH] Litellm dev 07 05 2025 p3 (#12349) * refactor(aim.py): refactor to support adding aim guardrails on UI * fix(base.py): add ui_friendly_name to config model * feat(ui/): support loading new guardrails from backend api call removes need to onboard each guardrail to ui * fix: don't show optional params if not set and don't show ui_friendly_name (internal param0 * fix(ui/add_guardrail_form.tsx): ensure dynamic provider value is used * fix(ui/): just one-time update the provider map dictionary * fix(ui/): show masked api base / api key on guardrail update * refactor(aporia_ai/): refactor to show on UI * feat(aporia_ai/): add aporia ai guardrail to UI * refactor(guardrails_ai/): refactor to add via UI * refactor(lasso.py): refactor to enable adding lasso guardrails via UI * feat(pangea.py): add pangea guardrail on UI * feat(panw): add panw prisma airs through UI * test: update tests * fix: fix ruff linting error * test: update tests * fix: add missing docs * fix: fix guardrail init * fix: suppress linting errors * fix(proxy_server.py): fix linting error --- docs/my-website/docs/proxy/config_settings.md | 9 +- litellm/proxy/common_utils/callback_utils.py | 2 +- .../proxy/guardrails/guardrail_endpoints.py | 235 ++++++++++-------- .../guardrail_hooks/aim/__init__.py | 34 +++ .../guardrail_hooks/{ => aim}/aim.py | 16 +- .../guardrail_hooks/aporia_ai/__init__.py | 33 +++ .../{ => aporia_ai}/aporia_ai.py | 18 +- .../guardrail_hooks/guardrails_ai/__init__.py | 39 +++ .../{ => guardrails_ai}/guardrails_ai.py | 25 +- .../guardrail_hooks/lasso/__init__.py | 35 +++ .../guardrail_hooks/{ => lasso}/lasso.py | 22 +- .../guardrail_hooks/pangea/__init__.py | 38 +++ .../guardrail_hooks/{ => pangea}/pangea.py | 18 +- .../panw_prisma_airs/__init__.py | 42 ++++ .../panw_prisma_airs.py | 40 ++- .../guardrails/guardrail_initializers.py | 110 -------- .../proxy/guardrails/guardrail_registry.py | 31 +-- litellm/proxy/guardrails/init_guardrails.py | 4 +- litellm/proxy/proxy_server.py | 3 +- litellm/types/guardrails.py | 18 +- .../proxy/guardrails/guardrail_hooks/aim.py | 20 ++ .../guardrails/guardrail_hooks/aporia_ai.py | 20 ++ .../azure/azure_prompt_shield.py | 4 +- .../azure/azure_text_moderation.py | 4 + .../proxy/guardrails/guardrail_hooks/base.py | 9 +- .../guardrail_hooks/guardrails_ai.py | 21 ++ .../proxy/guardrails/guardrail_hooks/lasso.py | 33 +++ .../guardrails/guardrail_hooks/pangea.py | 33 +++ .../guardrail_hooks/panw_prisma_airs.py | 25 ++ .../guardrails_tests/test_lasso_guardrails.py | 179 +++++++------ tests/local_testing/test_aim_guardrails.py | 20 +- .../guardrails/guardrail_hooks/test_pangea.py | 2 +- .../guardrail_hooks/test_panw_prisma_airs.py | 80 +++--- .../guardrails/test_guardrail_endpoints.py | 75 +++++- .../guardrails/test_guardrail_registry.py | 17 ++ .../guardrails/add_guardrail_form.tsx | 11 +- .../guardrails/edit_guardrail_form.tsx | 4 +- .../components/guardrails/guardrail_info.tsx | 1 + .../guardrails/guardrail_info_helpers.tsx | 85 +++++-- .../guardrails/guardrail_provider_fields.tsx | 30 ++- .../guardrail_provider_specific_fields.tsx | 6 +- 41 files changed, 1031 insertions(+), 420 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/aim/__init__.py rename litellm/proxy/guardrails/guardrail_hooks/{ => aim}/aim.py (95%) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/aporia_ai/__init__.py rename litellm/proxy/guardrails/guardrail_hooks/{ => aporia_ai}/aporia_ai.py (92%) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/__init__.py rename litellm/proxy/guardrails/guardrail_hooks/{ => guardrails_ai}/guardrails_ai.py (83%) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/lasso/__init__.py rename litellm/proxy/guardrails/guardrail_hooks/{ => lasso}/lasso.py (89%) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/pangea/__init__.py rename litellm/proxy/guardrails/guardrail_hooks/{ => pangea}/pangea.py (96%) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/__init__.py rename litellm/proxy/guardrails/guardrail_hooks/{ => panw_prisma_airs}/panw_prisma_airs.py (89%) create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/aim.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/aporia_ai.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/guardrails_ai.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/lasso.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/pangea.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py create mode 100644 tests/test_litellm/proxy/guardrails/test_guardrail_registry.py diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 7cd0082164..4333b0afc1 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -484,6 +484,7 @@ router_settings: | GOOGLE_CLIENT_ID | Client ID for Google OAuth | GOOGLE_CLIENT_SECRET | Client secret for Google OAuth | GOOGLE_KMS_RESOURCE_NAME | Name of the resource in Google KMS +| GUARDRAILS_AI_API_BASE | Base URL for Guardrails AI API | HEALTH_CHECK_TIMEOUT_SECONDS | Timeout in seconds for health checks. Default is 60 | HF_API_BASE | Base URL for Hugging Face API | HCP_VAULT_ADDR | Address for [Hashicorp Vault Secret Manager](../secret.md#hashicorp-vault) @@ -522,6 +523,10 @@ router_settings: | LANGSMITH_PROJECT | Project name for Langsmith integration | LANGSMITH_SAMPLING_RATE | Sampling rate for Langsmith logging | LANGTRACE_API_KEY | API key for Langtrace service +| LASSO_API_BASE | Base URL for Lasso API +| LASSO_API_KEY | API key for Lasso service +| LASSO_USER_ID | User ID for Lasso service +| LASSO_CONVERSATION_ID | Conversation ID for Lasso service | LENGTH_OF_LITELLM_GENERATED_KEY | Length of keys generated by LiteLLM. Default is 16 | LITERAL_API_KEY | API key for Literal integration | LITERAL_API_URL | API URL for Literal service @@ -538,6 +543,7 @@ router_settings: | LITELLM_LICENSE | License key for LiteLLM usage | LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM | LITELLM_LOG | Enable detailed logging for LiteLLM +| LITELLM_MASTER_KEY | Master key for proxy authentication | LITELLM_MODE | Operating mode for LiteLLM (e.g., production, development) | LITELLM_RATE_LIMIT_WINDOW_SIZE | Rate limit window size for LiteLLM. Default is 60 | LITELLM_SALT_KEY | Salt key for encryption in LiteLLM @@ -595,6 +601,8 @@ router_settings: | OTEL_SERVICE_NAME | Service name identifier for OpenTelemetry | OTEL_TRACER_NAME | Tracer name for OpenTelemetry tracing | PAGERDUTY_API_KEY | API key for PagerDuty Alerting +| PANW_PRISMA_AIRS_API_KEY | API key for PANW Prisma AIRS service +| PANW_PRISMA_AIRS_API_BASE | Base URL for PANW Prisma AIRS service | PHOENIX_API_KEY | API key for Arize Phoenix | PHOENIX_COLLECTOR_ENDPOINT | API endpoint for Arize Phoenix | PHOENIX_COLLECTOR_HTTP_ENDPOINT | API http endpoint for Arize Phoenix @@ -612,7 +620,6 @@ router_settings: | PROXY_BUDGET_RESCHEDULER_MAX_TIME | Maximum time in seconds to wait before checking database for budget resets. Default is 605 | PROXY_BUDGET_RESCHEDULER_MIN_TIME | Minimum time in seconds to wait before checking database for budget resets. Default is 597 | PROXY_LOGOUT_URL | URL for logging out of the proxy service -| LITELLM_MASTER_KEY | Master key for proxy authentication | QDRANT_API_BASE | Base URL for Qdrant API | QDRANT_API_KEY | API key for Qdrant service | QDRANT_SCALAR_QUANTILE | Scalar quantile for Qdrant operations. Default is 0.99 diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 39498202b6..a2afb30f05 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -123,7 +123,7 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 lakera_moderations_object = lakeraAI_Moderation(**init_params) imported_list.append(lakera_moderations_object) elif isinstance(callback, str) and callback == "aporia_prompt_injection": - from litellm.proxy.guardrails.guardrail_hooks.aporia_ai import ( + from litellm.proxy.guardrails.guardrail_hooks.aporia_ai.aporia_ai import ( AporiaGuardrail, ) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 0680fa1260..465b2f5074 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2,7 +2,8 @@ CRUD ENDPOINTS FOR GUARDRAILS """ -from typing import Any, Dict, List, Optional, Type, Union, cast +import inspect +from typing import Any, Dict, List, Optional, Type, TypeVar, Union, cast from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel @@ -755,115 +756,135 @@ def _get_list_element_options(field_annotation: Any) -> Optional[List[str]]: return None +def _extract_fields_recursive( + model: Type[BaseModel], + depth: int = 0, +) -> Dict[str, Any]: + + # Check if we've exceeded the maximum recursion depth + if depth > DEFAULT_MAX_RECURSE_DEPTH: + raise HTTPException( + status_code=400, + detail=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing model fields. Please check the model structure for excessive nesting.", + ) + + fields = {} + + for field_name, field in model.model_fields.items(): + # Skip optional_params if it's not meaningfully overridden + if field_name == "optional_params": + field_annotation = field.annotation + if field_annotation is None: + continue + # Check if the annotation is still a generic TypeVar (not specialized) + if isinstance(field_annotation, TypeVar) or ( + hasattr(field_annotation, "__origin__") + and field_annotation.__origin__ is TypeVar + ): + # Skip this field as it's not meaningfully overridden + continue + # Also skip if it's a generic type that wasn't specialized + if hasattr(field_annotation, "__name__") and field_annotation.__name__ in ( + "T", + "TypeVar", + ): + continue + + # Get field metadata + description = field.description or field_name + + # Check if this field is required + required = field.is_required() + + # Check if the field annotation is a BaseModel subclass + field_annotation = field.annotation + + # Handle Optional types and get the actual type + if field_annotation is None: + continue + + if ( + hasattr(field_annotation, "__origin__") + and field_annotation.__origin__ is Union + and hasattr(field_annotation, "__args__") + ): + # For Optional[BaseModel], get the non-None type + args = field_annotation.__args__ + non_none_args = [arg for arg in args if arg is not type(None)] + if non_none_args: + field_annotation = non_none_args[0] + + # Check if this is a BaseModel subclass + is_basemodel_subclass = ( + inspect.isclass(field_annotation) + and issubclass(field_annotation, BaseModel) + and field_annotation is not BaseModel + ) + + if is_basemodel_subclass: + # Recursively get fields from the nested model + nested_fields = _extract_fields_recursive( + cast(Type[BaseModel], field_annotation), depth + 1 + ) + fields[field_name] = { + "description": description, + "required": required, + "type": "nested", + "fields": nested_fields, + } + else: + # Determine the field type from annotation + field_type = _get_field_type_from_annotation(field_annotation) + + # Check for custom UI type override + field_json_schema_extra = getattr(field, "json_schema_extra", {}) + if field_json_schema_extra and "ui_type" in field_json_schema_extra: + field_type = field_json_schema_extra["ui_type"].value + elif field_json_schema_extra and "type" in field_json_schema_extra: + field_type = field_json_schema_extra["type"] + + # Add the field to the dictionary + field_dict = { + "description": description, + "required": required, + "type": field_type, + } + + # Extract options from type annotations + if field_type == "dict": + # For Dict[Literal[...], T] types, extract key options + dict_key_options = _get_dict_key_options(field_annotation) + if dict_key_options: + field_dict["dict_key_options"] = dict_key_options + + # Extract value type for the dict values + dict_value_type = _get_dict_value_type(field_annotation) + field_dict["dict_value_type"] = dict_value_type + + elif field_type == "array": + # For List[Literal[...]] types, extract element options + list_element_options = _get_list_element_options(field_annotation) + if list_element_options: + field_dict["options"] = list_element_options + field_dict["type"] = "multiselect" + + # Add options if they exist in json_schema_extra (this takes precedence) + if field_json_schema_extra and "options" in field_json_schema_extra: + field_dict["options"] = field_json_schema_extra["options"] + + # Add default value if it exists + if field.default is not None and field.default is not ...: + field_dict["default_value"] = field.default + + fields[field_name] = field_dict + + return fields + + def _get_fields_from_model(model_class: Type[BaseModel]) -> Dict[str, Any]: """ Get the fields from a Pydantic model as a nested dictionary structure """ - import inspect - - def _extract_fields_recursive( - model: Type[BaseModel], - depth: int = 0, - ) -> Dict[str, Any]: - # Check if we've exceeded the maximum recursion depth - if depth > DEFAULT_MAX_RECURSE_DEPTH: - raise HTTPException( - status_code=400, - detail=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing model fields. Please check the model structure for excessive nesting.", - ) - - fields = {} - - for field_name, field in model.model_fields.items(): - # Get field metadata - description = field.description or field_name - - # Check if this field is required - required = field.is_required() - - # Check if the field annotation is a BaseModel subclass - field_annotation = field.annotation - - # Handle Optional types and get the actual type - if field_annotation is None: - continue - - if ( - hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is Union - and hasattr(field_annotation, "__args__") - ): - # For Optional[BaseModel], get the non-None type - args = field_annotation.__args__ - non_none_args = [arg for arg in args if arg is not type(None)] - if non_none_args: - field_annotation = non_none_args[0] - - # Check if this is a BaseModel subclass - is_basemodel_subclass = ( - inspect.isclass(field_annotation) - and issubclass(field_annotation, BaseModel) - and field_annotation is not BaseModel - ) - - if is_basemodel_subclass: - # Recursively get fields from the nested model - nested_fields = _extract_fields_recursive( - cast(Type[BaseModel], field_annotation), depth + 1 - ) - fields[field_name] = { - "description": description, - "required": required, - "type": "nested", - "fields": nested_fields, - } - else: - # Determine the field type from annotation - field_type = _get_field_type_from_annotation(field_annotation) - - # Check for custom UI type override - field_json_schema_extra = getattr(field, "json_schema_extra", {}) - if field_json_schema_extra and "ui_type" in field_json_schema_extra: - field_type = field_json_schema_extra["ui_type"].value - elif field_json_schema_extra and "type" in field_json_schema_extra: - field_type = field_json_schema_extra["type"] - - # Add the field to the dictionary - field_dict = { - "description": description, - "required": required, - "type": field_type, - } - - # Extract options from type annotations - if field_type == "dict": - # For Dict[Literal[...], T] types, extract key options - dict_key_options = _get_dict_key_options(field_annotation) - if dict_key_options: - field_dict["dict_key_options"] = dict_key_options - - # Extract value type for the dict values - dict_value_type = _get_dict_value_type(field_annotation) - field_dict["dict_value_type"] = dict_value_type - - elif field_type == "array": - # For List[Literal[...]] types, extract element options - list_element_options = _get_list_element_options(field_annotation) - if list_element_options: - field_dict["options"] = list_element_options - field_dict["type"] = "multiselect" - - # Add options if they exist in json_schema_extra (this takes precedence) - if field_json_schema_extra and "options" in field_json_schema_extra: - field_dict["options"] = field_json_schema_extra["options"] - - # Add default value if it exists - if field.default is not None and field.default is not ...: - field_dict["default_value"] = field.default - - fields[field_name] = field_dict - - return fields return _extract_fields_recursive(model_class, depth=0) @@ -944,6 +965,8 @@ async def get_provider_specific_params(): if guardrail_config_model: fields = _get_fields_from_model(guardrail_config_model) + ui_friendly_name = guardrail_config_model.ui_friendly_name() + fields["ui_friendly_name"] = ui_friendly_name provider_params[guardrail_name] = fields return provider_params diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/aim/__init__.py new file mode 100644 index 0000000000..3e00d87d61 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/__init__.py @@ -0,0 +1,34 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .aim import AimGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + from litellm.proxy.guardrails.guardrail_hooks.aim import AimGuardrail + + _aim_callback = AimGuardrail( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_aim_callback) + + return _aim_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.AIM.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.AIM.value: AimGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py similarity index 95% rename from litellm/proxy/guardrails/guardrail_hooks/aim.py rename to litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 66a0aff6cb..317db76e37 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -7,7 +7,7 @@ import asyncio import json import os -from typing import Any, AsyncGenerator, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, AsyncGenerator, Literal, Optional, Type, Union from fastapi import HTTPException from pydantic import BaseModel @@ -22,7 +22,6 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.proxy_server import StreamingCallbackError from litellm.types.utils import ( Choices, EmbeddingResponse, @@ -31,6 +30,9 @@ from litellm.types.utils import ( ModelResponseStream, ) +if TYPE_CHECKING: + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + class AimGuardrailMissingSecrets(Exception): pass @@ -313,6 +315,8 @@ class AimGuardrail(CustomGuardrail): if result.get("done"): return if blocking_message := result.get("blocking_message"): + from litellm.proxy.proxy_server import StreamingCallbackError + raise StreamingCallbackError(blocking_message) verbose_proxy_logger.error( f"Unknown message received from AIM: {result}" @@ -334,3 +338,11 @@ class AimGuardrail(CustomGuardrail): def _set_dlp_entities(self, entities: list[dict]) -> None: self.dlp_entities = entities[: self._max_dlp_entities] + + @staticmethod + def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: + from litellm.types.proxy.guardrails.guardrail_hooks.aim import ( + AimGuardrailConfigModel, + ) + + return AimGuardrailConfigModel diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/__init__.py new file mode 100644 index 0000000000..e47267a1c7 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/__init__.py @@ -0,0 +1,33 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .aporia_ai import AporiaGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + _aporia_callback = AporiaGuardrail( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_aporia_callback) + + return _aporia_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.APORIA.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.APORIA.value: AporiaGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py similarity index 92% rename from litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py rename to litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py index 3c39b90b0a..8c61c0c104 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py @@ -13,7 +13,7 @@ sys.path.insert( ) # Adds the parent directory to the system path import json import sys -from typing import Any, List, Literal, Optional +from typing import TYPE_CHECKING, Any, List, Literal, Optional, Type from fastapi import HTTPException @@ -31,13 +31,15 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata from litellm.types.guardrails import GuardrailEventHooks litellm.set_verbose = True GUARDRAIL_NAME = "aporia" +if TYPE_CHECKING: + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + class AporiaGuardrail(CustomGuardrail): def __init__( @@ -194,12 +196,16 @@ class AporiaGuardrail(CustomGuardrail): from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, ) + from litellm.proxy.guardrails.guardrail_helpers import ( + should_proceed_based_on_metadata, + ) event_type: GuardrailEventHooks = GuardrailEventHooks.during_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: return # old implementation - backwards compatibility + if ( await should_proceed_based_on_metadata( data=data, @@ -226,3 +232,11 @@ class AporiaGuardrail(CustomGuardrail): "Aporia AI: not running guardrail. No messages in data" ) pass + + @staticmethod + def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: + from litellm.types.proxy.guardrails.guardrail_hooks.aporia_ai import ( + AporiaGuardrailConfigModel, + ) + + return AporiaGuardrailConfigModel diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/__init__.py new file mode 100644 index 0000000000..885a88334d --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/__init__.py @@ -0,0 +1,39 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .guardrails_ai import GuardrailsAI + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + if litellm_params.guard_name is None: + raise Exception( + "GuardrailsAIException - Please pass the Guardrails AI guard name via 'litellm_params::guard_name'" + ) + + _guardrails_ai_callback = GuardrailsAI( + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + guard_name=litellm_params.guard_name, + ) + litellm.logging_callback_manager.add_litellm_callback(_guardrails_ai_callback) + + return _guardrails_ai_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.GUARDRAILS_AI.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.GUARDRAILS_AI.value: GuardrailsAI, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py similarity index 83% rename from litellm/proxy/guardrails/guardrail_hooks/guardrails_ai.py rename to litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index 1a2c5a217b..48027ec1c6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -6,7 +6,8 @@ # Thank you for using Litellm! - Krrish & Ishaan import json -from typing import Optional, TypedDict +import os +from typing import TYPE_CHECKING, Optional, Type, TypedDict from fastapi import HTTPException @@ -20,11 +21,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.callback_utils import ( - add_guardrail_to_applied_guardrails_header, -) from litellm.types.guardrails import GuardrailEventHooks +if TYPE_CHECKING: + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + class GuardrailsAIResponse(TypedDict): callId: str @@ -45,7 +46,9 @@ class GuardrailsAI(CustomGuardrail): "GuardrailsAIException - Please pass the Guardrails AI guard name via 'litellm_params::guard_name'" ) # store kwargs as optional_params - self.guardrails_ai_api_base = api_base or "http://0.0.0.0:8000" + self.guardrails_ai_api_base = ( + api_base or os.getenv("GUARDRAILS_AI_API_BASE") or "http://0.0.0.0:8000" + ) self.guardrails_ai_guard_name = guard_name self.optional_params = kwargs supported_event_hooks = [GuardrailEventHooks.post_call] @@ -94,6 +97,10 @@ class GuardrailsAI(CustomGuardrail): It can be used to reject a response """ + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + event_type: GuardrailEventHooks = GuardrailEventHooks.post_call if self.should_run_guardrail(data=data, event_type=event_type) is not True: return @@ -112,3 +119,11 @@ class GuardrailsAI(CustomGuardrail): ) return + + @staticmethod + def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: + from litellm.types.proxy.guardrails.guardrail_hooks.guardrails_ai import ( + GuardrailsAIGuardrailConfigModel, + ) + + return GuardrailsAIGuardrailConfigModel diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/__init__.py new file mode 100644 index 0000000000..f68d61e7e5 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/__init__.py @@ -0,0 +1,35 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .lasso import LassoGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + _lasso_callback = LassoGuardrail( + guardrail_name=guardrail.get("guardrail_name", ""), + api_key=litellm_params.api_key, + api_base=litellm_params.api_base, + user_id=litellm_params.lasso_user_id, + conversation_id=litellm_params.lasso_conversation_id, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_lasso_callback) + + return _lasso_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.LASSO.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.LASSO.value: LassoGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py similarity index 89% rename from litellm/proxy/guardrails/guardrail_hooks/lasso.py rename to litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 4debe61ce4..00350d5a36 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -6,7 +6,7 @@ # +-------------------------------------------------------------+ import os -from typing import Any, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union from fastapi import HTTPException @@ -22,6 +22,9 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth +if TYPE_CHECKING: + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + class LassoGuardrailMissingSecrets(Exception): pass @@ -37,6 +40,7 @@ class LassoGuardrail(CustomGuardrail): def __init__( self, lasso_api_key: Optional[str] = None, + api_key: Optional[str] = None, api_base: Optional[str] = None, user_id: Optional[str] = None, conversation_id: Optional[str] = None, @@ -45,7 +49,7 @@ class LassoGuardrail(CustomGuardrail): self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback ) - self.lasso_api_key = lasso_api_key or os.environ.get("LASSO_API_KEY") + self.lasso_api_key = lasso_api_key or api_key or os.environ.get("LASSO_API_KEY") self.user_id = user_id or os.environ.get("LASSO_USER_ID") self.conversation_id = conversation_id or os.environ.get( "LASSO_CONVERSATION_ID" @@ -58,7 +62,9 @@ class LassoGuardrail(CustomGuardrail): ) raise LassoGuardrailMissingSecrets(msg) - self.api_base = api_base or "https://server.lasso.security/gateway/v2/classify" + self.api_base = ( + api_base or os.getenv("LASSO_API_BASE") or "https://server.lasso.security" + ) super().__init__(**kwargs) @log_guardrail_information @@ -169,7 +175,7 @@ class LassoGuardrail(CustomGuardrail): """Call the Lasso API and return the response.""" verbose_proxy_logger.debug(f"Sending request to Lasso API: {payload}") response = await self.async_handler.post( - url=self.api_base, + url=f"{self.api_base}/gateway/v2/classify", headers=headers, json=payload, timeout=10.0, @@ -203,3 +209,11 @@ class LassoGuardrail(CustomGuardrail): if is_violated: violated_deputies.append(deputy) return violated_deputies + + @staticmethod + def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: + from litellm.types.proxy.guardrails.guardrail_hooks.lasso import ( + LassoGuardrailConfigModel, + ) + + return LassoGuardrailConfigModel diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/__init__.py new file mode 100644 index 0000000000..181633065e --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/__init__.py @@ -0,0 +1,38 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .pangea import PangeaHandler + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("Pangea guardrail name is required") + + _pangea_callback = PangeaHandler( + guardrail_name=guardrail_name, + pangea_input_recipe=litellm_params.pangea_input_recipe, + pangea_output_recipe=litellm_params.pangea_output_recipe, + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_pangea_callback) + + return _pangea_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.PANGEA.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.PANGEA.value: PangeaHandler, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py similarity index 96% rename from litellm/proxy/guardrails/guardrail_hooks/pangea.py rename to litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 0159955d81..07e9fa760f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -1,6 +1,6 @@ # litellm/proxy/guardrails/guardrail_hooks/pangea.py import os -from typing import Any, Optional, Protocol +from typing import TYPE_CHECKING, Any, Optional, Protocol, Type from fastapi import HTTPException @@ -10,7 +10,6 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) - from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -22,6 +21,9 @@ from litellm.proxy.common_utils.callback_utils import ( from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import LLMResponseTypes, ModelResponse, TextCompletionResponse +if TYPE_CHECKING: + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + class PangeaGuardrailMissingSecrets(Exception): """Custom exception for missing Pangea secrets.""" @@ -30,10 +32,10 @@ class PangeaGuardrailMissingSecrets(Exception): class _Transformer(Protocol): - def get_messages(self) -> list[dict]: + def get_messages(self) -> list[dict]: # noqa: E704 ... - def update_original_body(self, prompt_messages: list[dict]) -> Any: + def update_original_body(self, prompt_messages: list[dict]) -> Any: # noqa: E704 ... @@ -397,3 +399,11 @@ class PangeaHandler(CustomGuardrail): "exceptions": str(e), }, ) from e + + @staticmethod + def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: + from litellm.types.proxy.guardrails.guardrail_hooks.pangea import ( + PangeaGuardrailConfigModel, + ) + + return PangeaGuardrailConfigModel diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/__init__.py new file mode 100644 index 0000000000..747ef3cc91 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/__init__.py @@ -0,0 +1,42 @@ +from typing import TYPE_CHECKING, Optional, cast + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .panw_prisma_airs import PanwPrismaAirsHandler + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + guardrail_name = guardrail.get("guardrail_name") + profile_name = cast(Optional[str], getattr(litellm_params, "profile_name", None)) + if not litellm_params.api_key: + raise ValueError("PANW Prisma AIRS: api_key is required") + if not profile_name: + raise ValueError("PANW Prisma AIRS: profile_name is required") + if not guardrail_name: + raise ValueError("PANW Prisma AIRS: guardrail_name is required") + + _panw_callback = PanwPrismaAirsHandler( + **{ + **litellm_params.model_dump(), + "guardrail_name": guardrail_name, + "default_on": litellm_params.default_on or False, + } + ) + litellm.logging_callback_manager.add_litellm_callback(_panw_callback) + + return _panw_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.PANW_PRISMA_AIRS.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.PANW_PRISMA_AIRS.value: PanwPrismaAirsHandler, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py similarity index 89% rename from litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py rename to litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 9e65efae5f..f75f6db82a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -4,8 +4,9 @@ PANW Prisma AIRS Built-in Guardrail for LiteLLM """ +import os import uuid -from typing import Any, Dict, List, Literal, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, cast from fastapi import HTTPException @@ -21,6 +22,9 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ModelResponse +if TYPE_CHECKING: + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + class PanwPrismaAirsHandler(CustomGuardrail): """ @@ -52,8 +56,12 @@ class PanwPrismaAirsHandler(CustomGuardrail): super().__init__(guardrail_name=guardrail_name, default_on=default_on, **kwargs) # Store configuration - self.api_key = api_key - self.api_base = api_base + self.api_key = api_key or os.getenv("PANW_PRISMA_AIRS_API_KEY") + self.api_base = ( + api_base + or os.getenv("PANW_PRISMA_AIRS_API_BASE") + or "https://service.api.aisecurity.paloaltonetworks.com" + ) self.profile_name = profile_name verbose_proxy_logger.info( @@ -98,14 +106,15 @@ class PanwPrismaAirsHandler(CustomGuardrail): def _extract_response_text(self, response: ModelResponse) -> str: """Extract text from LLM response.""" try: + from litellm.types.utils import Choices + if ( hasattr(response, "choices") and response.choices and len(response.choices) > 0 and hasattr(response.choices[0], "message") - and hasattr(response.choices[0].message, "content") ): - return response.choices[0].message.content or "" + return cast(Choices, response.choices[0]).message.content or "" except (AttributeError, IndexError): verbose_proxy_logger.error( "PANW Prisma AIRS: Error extracting response text" @@ -132,9 +141,9 @@ class PanwPrismaAirsHandler(CustomGuardrail): "tr_id": transaction_id, "ai_profile": {"profile_name": self.profile_name}, "metadata": { - "app_user": metadata.get("user", "litellm_user") - if metadata - else "litellm_user", + "app_user": ( + metadata.get("user", "litellm_user") if metadata else "litellm_user" + ), "ai_model": metadata.get("model", "unknown") if metadata else "unknown", "source": "litellm_builtin_guardrail", }, @@ -157,7 +166,10 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) response = await async_client.post( - self.api_base, headers=headers, json=payload, timeout=10.0 + f"{self.api_base}/v1/scan/sync/request", + headers=headers, + json=payload, + timeout=10.0, ) response.raise_for_status() @@ -231,6 +243,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): "image_generation", "moderation", "audio_transcription", + "pass_through_endpoint", + "rerank", ], ) -> Optional[Dict[str, Any]]: """ @@ -328,3 +342,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): raise HTTPException(status_code=400, detail=error_detail) return response + + @staticmethod + def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: + from litellm.types.proxy.guardrails.guardrail_hooks.panw_prisma_airs import ( + PanwPrismaAirsGuardrailConfigModel, + ) + + return PanwPrismaAirsGuardrailConfigModel diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index bb4fb2b883..6282478a76 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -4,23 +4,6 @@ from litellm.proxy._types import CommonProxyErrors from litellm.types.guardrails import * -def initialize_aporia( - litellm_params: LitellmParams, - guardrail: Guardrail, -): - from litellm.proxy.guardrails.guardrail_hooks.aporia_ai import AporiaGuardrail - - _aporia_callback = AporiaGuardrail( - api_base=litellm_params.api_base, - api_key=litellm_params.api_key, - guardrail_name=guardrail.get("guardrail_name", ""), - event_hook=litellm_params.mode, - 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): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, @@ -83,21 +66,6 @@ def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail): return _lakera_v2_callback -def initialize_aim(litellm_params: LitellmParams, guardrail: Guardrail): - from litellm.proxy.guardrails.guardrail_hooks.aim import AimGuardrail - - _aim_callback = AimGuardrail( - api_base=litellm_params.api_base, - api_key=litellm_params.api_key, - guardrail_name=guardrail.get("guardrail_name", ""), - event_hook=litellm_params.mode, - default_on=litellm_params.default_on, - ) - 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 ( _OPTIONAL_PresidioPIIMasking, @@ -152,81 +120,3 @@ def initialize_hide_secrets(litellm_params: LitellmParams, guardrail: Guardrail) ) litellm.logging_callback_manager.add_litellm_callback(_secret_detection_object) return _secret_detection_object - - -def initialize_guardrails_ai(litellm_params, guardrail): - from litellm.proxy.guardrails.guardrail_hooks.guardrails_ai import GuardrailsAI - - _guard_name = litellm_params.guard_name - if not _guard_name: - raise Exception( - "GuardrailsAIException - Please pass the Guardrails AI guard name via 'litellm_params::guard_name'" - ) - - _guardrails_ai_callback = GuardrailsAI( - api_base=litellm_params.api_base, - guard_name=_guard_name, - guardrail_name=SupportedGuardrailIntegrations.GURDRAILS_AI.value, - default_on=litellm_params.default_on, - ) - litellm.logging_callback_manager.add_litellm_callback(_guardrails_ai_callback) - - return _guardrails_ai_callback - - -def initialize_pangea(litellm_params, guardrail): - from litellm.proxy.guardrails.guardrail_hooks.pangea import PangeaHandler - - _pangea_callback = PangeaHandler( - guardrail_name=guardrail["guardrail_name"], - pangea_input_recipe=litellm_params.pangea_input_recipe, - pangea_output_recipe=litellm_params.pangea_output_recipe, - api_base=litellm_params.api_base, - api_key=litellm_params.api_key, - default_on=litellm_params.default_on, - ) - litellm.logging_callback_manager.add_litellm_callback(_pangea_callback) - - return _pangea_callback - - -def initialize_lasso( - litellm_params: LitellmParams, - guardrail: Guardrail, -): - from litellm.proxy.guardrails.guardrail_hooks.lasso import LassoGuardrail - - _lasso_callback = LassoGuardrail( - guardrail_name=guardrail.get("guardrail_name", ""), - lasso_api_key=litellm_params.api_key, - api_base=litellm_params.api_base, - user_id=litellm_params.lasso_user_id, - conversation_id=litellm_params.lasso_conversation_id, - event_hook=litellm_params.mode, - default_on=litellm_params.default_on, - ) - litellm.logging_callback_manager.add_litellm_callback(_lasso_callback) - - return _lasso_callback - - -def initialize_panw_prisma_airs(litellm_params, guardrail): - from litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs import ( - PanwPrismaAirsHandler, - ) - - if not litellm_params.api_key: - raise ValueError("PANW Prisma AIRS: api_key is required") - if not litellm_params.profile_name: - raise ValueError("PANW Prisma AIRS: profile_name is required") - - _panw_callback = PanwPrismaAirsHandler( - guardrail_name=guardrail.get("guardrail_name", "panw_prisma_airs"), # Use .get() with default - api_key=litellm_params.api_key, - api_base=litellm_params.api_base or "https://service.api.aisecurity.paloaltonetworks.com/v1/scan/sync/request", - profile_name=litellm_params.profile_name, - default_on=litellm_params.default_on, - ) - litellm.logging_callback_manager.add_litellm_callback(_panw_callback) - - return _panw_callback diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 4fec8e4830..a0b1c16051 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -21,31 +21,19 @@ from litellm.types.guardrails import ( ) from .guardrail_initializers import ( - initialize_aim, - initialize_aporia, initialize_bedrock, - initialize_guardrails_ai, initialize_hide_secrets, initialize_lakera, initialize_lakera_v2, - initialize_lasso, - initialize_pangea, - initialize_panw_prisma_airs, initialize_presidio, ) guardrail_initializer_registry = { - SupportedGuardrailIntegrations.APORIA.value: initialize_aporia, SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock, SupportedGuardrailIntegrations.LAKERA.value: initialize_lakera, SupportedGuardrailIntegrations.LAKERA_V2.value: initialize_lakera_v2, - SupportedGuardrailIntegrations.AIM.value: initialize_aim, SupportedGuardrailIntegrations.PRESIDIO.value: initialize_presidio, SupportedGuardrailIntegrations.HIDE_SECRETS.value: initialize_hide_secrets, - SupportedGuardrailIntegrations.GURDRAILS_AI.value: initialize_guardrails_ai, - SupportedGuardrailIntegrations.PANGEA.value: initialize_pangea, - SupportedGuardrailIntegrations.LASSO.value: initialize_lasso, - SupportedGuardrailIntegrations.PANW_PRISMA_AIRS.value: initialize_panw_prisma_airs, } guardrail_class_registry: Dict[str, Type[CustomGuardrail]] = {} @@ -111,7 +99,7 @@ def get_guardrail_initializer_from_hooks(): ) except ImportError as e: - verbose_proxy_logger.debug(f"Could not import {module_path}: {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}") @@ -183,7 +171,7 @@ def get_guardrail_class_from_hooks(): verbose_proxy_logger.debug(f"Could not import {module_path}: {e}") continue except Exception as e: - verbose_proxy_logger.error(f"Error processing {module_path}: {e}") + verbose_proxy_logger.exception(f"Error processing {module_path}: {e}") continue except Exception as e: @@ -323,7 +311,7 @@ class GuardrailRegistry: guardrails: List[Guardrail] = [] for guardrail in guardrails_from_db: - guardrails.append(Guardrail(**(dict(guardrail)))) + guardrails.append(Guardrail(**(dict(guardrail)))) # type: ignore return guardrails except Exception as e: @@ -343,7 +331,7 @@ class GuardrailRegistry: if not guardrail: return None - return Guardrail(**(dict(guardrail))) + return Guardrail(**(dict(guardrail))) # type: ignore except Exception as e: raise Exception(f"Error getting guardrail from DB: {str(e)}") @@ -361,7 +349,7 @@ class GuardrailRegistry: if not guardrail: return None - return Guardrail(**(dict(guardrail))) + return Guardrail(**(dict(guardrail))) # type: ignore except Exception as e: raise Exception(f"Error getting guardrail from DB: {str(e)}") @@ -384,7 +372,7 @@ class InMemoryGuardrailHandler: def initialize_guardrail( self, - guardrail: Dict, + guardrail: Guardrail, config_file_path: Optional[str] = None, ) -> Optional[Guardrail]: """ @@ -404,7 +392,10 @@ class InMemoryGuardrailHandler: litellm_params_data = guardrail["litellm_params"] verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data) - litellm_params = LitellmParams(**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 @@ -433,7 +424,7 @@ class InMemoryGuardrailHandler: custom_guardrail_callback = initializer(litellm_params, guardrail) elif isinstance(guardrail_type, str) and "." in guardrail_type: custom_guardrail_callback = self.initialize_custom_guardrail( - guardrail=guardrail, + guardrail=cast(dict, guardrail), guardrail_type=guardrail_type, litellm_params=litellm_params, config_file_path=config_file_path, diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index 6009fce31e..0ed48d671d 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional +from typing import Dict, List, Optional, cast import litellm from litellm._logging import verbose_proxy_logger @@ -83,7 +83,7 @@ def init_guardrails_v2( for guardrail in all_guardrails: initialized_guardrail = IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( - guardrail=guardrail, + guardrail=cast(Guardrail, guardrail), config_file_path=config_file_path, ) if initialized_guardrail: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index aa2543361d..3e15a2ecf5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2828,7 +2828,7 @@ class ProxyConfig: ) for guardrail in guardrails_in_db: IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( - guardrail=dict(guardrail), + guardrail=cast(Guardrail, guardrail), ) except Exception as e: verbose_proxy_logger.exception( @@ -2978,6 +2978,7 @@ async def initialize( # noqa: PLR0915 ): global user_model, user_api_base, user_debug, user_detailed_debug, user_user_max_tokens, user_request_timeout, user_temperature, user_telemetry, user_headers, experimental, llm_model_list, llm_router, general_settings, master_key, user_custom_auth, prisma_client from litellm.proxy.common_utils.banner import show_banner + show_banner() if os.getenv("LITELLM_DONT_SHOW_FEEDBACK_BOX", "").lower() != "true": generate_feedback_box() diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index d7835db54b..5f8816351d 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -22,7 +22,7 @@ guardrails: class SupportedGuardrailIntegrations(Enum): APORIA = "aporia" BEDROCK = "bedrock" - GURDRAILS_AI = "guardrails_ai" + GUARDRAILS_AI = "guardrails_ai" LAKERA = "lakera" LAKERA_V2 = "lakera_v2" PRESIDIO = "presidio" @@ -409,11 +409,23 @@ class LitellmParams( kwargs["default_on"] = False super().__init__(**kwargs) + def __contains__(self, key): + # Define custom behavior for the 'in' operator + return hasattr(self, key) + + def get(self, key, default=None): + # Custom .get() method to access attributes with a default value if the attribute doesn't exist + return getattr(self, key, default) + + def __getitem__(self, key): + # Allow dictionary-style access to attributes + return getattr(self, key) + class Guardrail(TypedDict, total=False): guardrail_id: Optional[str] - guardrail_name: str - litellm_params: LitellmParams + guardrail_name: Required[str] + litellm_params: Required[LitellmParams] guardrail_info: Optional[Dict] created_at: Optional[datetime] updated_at: Optional[datetime] diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/aim.py b/litellm/types/proxy/guardrails/guardrail_hooks/aim.py new file mode 100644 index 0000000000..4329225e36 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/aim.py @@ -0,0 +1,20 @@ +from typing import Optional + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class AimGuardrailConfigModel(GuardrailConfigModel): + api_key: Optional[str] = Field( + default=None, + description="The API key for the Aim guardrail. If not provided, the `AIM_API_KEY` environment variable is checked.", + ) + api_base: Optional[str] = Field( + default=None, + description="The API base for the Aim guardrail. Default is https://api.aim.security. Also checks if the `AIM_API_BASE` environment variable is set.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "AIM Guardrail" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/aporia_ai.py b/litellm/types/proxy/guardrails/guardrail_hooks/aporia_ai.py new file mode 100644 index 0000000000..26247d7509 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/aporia_ai.py @@ -0,0 +1,20 @@ +from typing import Optional + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class AporiaGuardrailConfigModel(GuardrailConfigModel): + api_key: Optional[str] = Field( + default=None, + description="The API key for the Aporia guardrail. If not provided, the `APORIA_API_KEY` environment variable is checked.", + ) + api_base: Optional[str] = Field( + default=None, + description="The API base for the Aporia guardrail. If not provided, the `APORIA_API_BASE` environment variable is checked.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Aporia AI" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py index 6e04bb8299..9612c6c48d 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py @@ -29,4 +29,6 @@ class AzurePromptShieldGuardrailConfigModel( AzureContentSafetyConfigModel, GuardrailConfigModel, ): - pass + @staticmethod + def ui_friendly_name() -> str: + return "Azure Content Safety Prompt Shield" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py index 16eebb0892..308b2d68ac 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py @@ -80,3 +80,7 @@ class AzureContentSafetyTextModerationConfigModel( optional_params: AzureTextModerationOptionalParams = Field( description="Optional parameters for the Azure Content Safety Text Moderation guardrail", ) + + @staticmethod + def ui_friendly_name() -> str: + return "Azure Content Safety Text Moderation" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/base.py b/litellm/types/proxy/guardrails/guardrail_hooks/base.py index c346784edb..d2607e8912 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/base.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/base.py @@ -1,3 +1,4 @@ +from abc import ABC, abstractmethod from typing import Generic, TypeVar from pydantic import BaseModel, Field @@ -5,9 +6,15 @@ from pydantic import BaseModel, Field T = TypeVar("T", bound=BaseModel) -class GuardrailConfigModel(BaseModel, Generic[T]): +class GuardrailConfigModel(BaseModel, Generic[T], ABC): """Base model for guardrail configuration""" optional_params: T = Field( description="Optional parameters for the guardrail", ) + + @staticmethod + @abstractmethod + def ui_friendly_name() -> str: + """UI-friendly name for the guardrail""" + pass diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/guardrails_ai.py b/litellm/types/proxy/guardrails/guardrail_hooks/guardrails_ai.py new file mode 100644 index 0000000000..cb5e85661e --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/guardrails_ai.py @@ -0,0 +1,21 @@ +from typing import Optional + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class GuardrailsAIGuardrailConfigModel(GuardrailConfigModel): + api_base: Optional[str] = Field( + default=None, + description="The API base for the Guardrails AI guardrail. Defaults to http://0.0.0.0:8000, the `GUARDRAILS_AI_API_BASE` environment variable is checked.", + ) + + guard_name: Optional[str] = Field( + default=None, + description="The name of the Guardrails AI guardrail. Required for the Guardrails AI guardrail.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Guardrails AI" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/lasso.py b/litellm/types/proxy/guardrails/guardrail_hooks/lasso.py new file mode 100644 index 0000000000..4fe183c181 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/lasso.py @@ -0,0 +1,33 @@ +from typing import Optional + +from pydantic import BaseModel, Field + +from .base import GuardrailConfigModel + + +class LassoGuardrailConfigModelOptionalParams(BaseModel): + user_id: Optional[str] = Field( + default=None, + description="The user ID for the Lasso guardrail. If not provided, the `LASSO_USER_ID` environment variable is checked.", + ) + conversation_id: Optional[str] = Field( + default=None, + description="The conversation ID for the Lasso guardrail. If not provided, the `LASSO_CONVERSATION_ID` environment variable is checked.", + ) + + +class LassoGuardrailConfigModel( + GuardrailConfigModel[LassoGuardrailConfigModelOptionalParams] +): + api_key: Optional[str] = Field( + default=None, + description="The API key for the Lasso guardrail. If not provided, the `LASSO_API_KEY` environment variable is checked.", + ) + api_base: Optional[str] = Field( + default=None, + description="The API base for the Lasso guardrail. Default is https://server.lasso.security. Also checks if the `LASSO_API_BASE` environment variable is set.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Lasso Guardrail" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/pangea.py b/litellm/types/proxy/guardrails/guardrail_hooks/pangea.py new file mode 100644 index 0000000000..820da727bb --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/pangea.py @@ -0,0 +1,33 @@ +from typing import Optional + +from pydantic import BaseModel, Field + +from .base import GuardrailConfigModel + + +class PangeaGuardrailConfigModelOptionalParams(BaseModel): + pangea_input_recipe: Optional[str] = Field( + default=None, + description="The Pangea input recipe for the Pangea guardrail. Used for pre-call hook.", + ) + pangea_output_recipe: Optional[str] = Field( + default=None, + description="The Pangea output recipe for the Pangea guardrail. Used for post-call hook.", + ) + + +class PangeaGuardrailConfigModel( + GuardrailConfigModel[PangeaGuardrailConfigModelOptionalParams] +): + api_key: Optional[str] = Field( + default=None, + description="The Pangea API key. Reads from PANGEA_API_KEY env var if None.", + ) + api_base: Optional[str] = Field( + default=None, + description="The Pangea API base URL. Defaults to https://ai-guard.aws.us.pangea.cloud. Also checks if the PANGEA_API_BASE env var is set.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Pangea Guardrail" diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py b/litellm/types/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py new file mode 100644 index 0000000000..2d728f7076 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/panw_prisma_airs.py @@ -0,0 +1,25 @@ +from typing import Optional + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class PanwPrismaAirsGuardrailConfigModel(GuardrailConfigModel): + api_key: Optional[str] = Field( + default=None, + description="The API key for the PANW Prisma AIRS guardrail. If not provided, the `PANW_PRISMA_AIRS_API_KEY` environment variable is checked.", + ) + api_base: Optional[str] = Field( + default=None, + description="The API base for the PANW Prisma AIRS guardrail. Defaults to https://service.api.aisecurity.paloaltonetworks.com. If not provided, the `PANW_PRISMA_AIRS_API_BASE` environment variable is checked.", + ) + + profile_name: str = Field( + default="default", + description="PANW Prisma AIRS security profile name. Required.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "PANW Prisma AIRS" diff --git a/tests/guardrails_tests/test_lasso_guardrails.py b/tests/guardrails_tests/test_lasso_guardrails.py index c12968ed50..7c32105400 100644 --- a/tests/guardrails_tests/test_lasso_guardrails.py +++ b/tests/guardrails_tests/test_lasso_guardrails.py @@ -8,9 +8,15 @@ import pytest from litellm import DualCache from litellm.proxy.proxy_server import UserAPIKeyAuth -from litellm.proxy.guardrails.guardrail_hooks.lasso import LassoGuardrailMissingSecrets, LassoGuardrail, LassoGuardrailAPIError +from litellm.proxy.guardrails.guardrail_hooks.lasso.lasso import ( + LassoGuardrailMissingSecrets, + LassoGuardrail, + LassoGuardrailAPIError, +) -sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path import litellm from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 @@ -35,7 +41,7 @@ def test_lasso_guard_config(): ], config_file_path="", ) - + # Clean up del os.environ["LASSO_API_KEY"] @@ -43,12 +49,14 @@ def test_lasso_guard_config(): def test_lasso_guard_config_no_api_key(): litellm.set_verbose = True litellm.guardrail_name_config_map = {} - + # Ensure LASSO_API_KEY is not in environment if "LASSO_API_KEY" in os.environ: del os.environ["LASSO_API_KEY"] - - with pytest.raises(LassoGuardrailMissingSecrets, match="Couldn't get Lasso api key"): + + with pytest.raises( + LassoGuardrailMissingSecrets, match="Couldn't get Lasso api key" + ): init_guardrails_v2( all_guardrails=[ { @@ -69,7 +77,7 @@ async def test_callback(): # Set environment variable for testing os.environ["LASSO_API_KEY"] = "test-key" os.environ["LASSO_USER_ID"] = "test-user" - os.environ["LASSO_CONVERSATION_ID"] = "test-conversation" + os.environ["LASSO_CONVERSATION_ID"] = "test-conversation" init_guardrails_v2( all_guardrails=[ { @@ -82,7 +90,9 @@ async def test_callback(): } ], ) - lasso_guardrails = litellm.logging_callback_manager.get_custom_loggers_for_type(LassoGuardrail) + lasso_guardrails = litellm.logging_callback_manager.get_custom_loggers_for_type( + LassoGuardrail + ) print("found lasso guardrails", lasso_guardrails) lasso_guardrail = lasso_guardrails[0] @@ -98,73 +108,83 @@ async def test_callback(): "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", return_value=Response( json={ - "deputies": { - "jailbreak": True, - "custom-policies": False, - "sexual": False, - "hate": False, - "illegality": False, - "violence": False, - "pattern-detection": False - }, - "deputies_predictions": { - "jailbreak": 0.923, - "custom-policies": 0.234, - "sexual": 0.145, - "hate": 0.156, - "illegality": 0.167, - "violence": 0.178, - "pattern-detection": 0.189 - }, - "violations_detected": True - }, - status_code=200, - request=Request(method="POST", url="https://server.lasso.security/gateway/v1/chat"), - ), - ): - await lasso_guardrail.async_pre_call_hook( - data=data, cache=DualCache(), user_api_key_dict=UserAPIKeyAuth(), call_type="completion" - ) - - # Check for the correct error message - assert "Violated Lasso guardrail policy" in str(excinfo.value.detail) - assert "jailbreak" in str(excinfo.value.detail) - - # Test no violation - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - return_value=Response( - json={ "deputies": { - "jailbreak": False, + "jailbreak": True, "custom-policies": False, "sexual": False, "hate": False, "illegality": False, "violence": False, - "pattern-detection": False + "pattern-detection": False, }, "deputies_predictions": { - "jailbreak": 0.123, + "jailbreak": 0.923, "custom-policies": 0.234, "sexual": 0.145, "hate": 0.156, "illegality": 0.167, "violence": 0.178, - "pattern-detection": 0.189 + "pattern-detection": 0.189, }, - "violations_detected": False + "violations_detected": True, }, + status_code=200, + request=Request( + method="POST", url="https://server.lasso.security/gateway/v1/chat" + ), + ), + ): + await lasso_guardrail.async_pre_call_hook( + data=data, + cache=DualCache(), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + # Check for the correct error message + assert "Violated Lasso guardrail policy" in str(excinfo.value.detail) + assert "jailbreak" in str(excinfo.value.detail) + + # Test no violation + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=Response( + json={ + "deputies": { + "jailbreak": False, + "custom-policies": False, + "sexual": False, + "hate": False, + "illegality": False, + "violence": False, + "pattern-detection": False, + }, + "deputies_predictions": { + "jailbreak": 0.123, + "custom-policies": 0.234, + "sexual": 0.145, + "hate": 0.156, + "illegality": 0.167, + "violence": 0.178, + "pattern-detection": 0.189, + }, + "violations_detected": False, + }, status_code=200, - request=Request(method="POST", url="https://server.lasso.security/gateway/v1/chat"), + request=Request( + method="POST", url="https://server.lasso.security/gateway/v1/chat" + ), ), ): result = await lasso_guardrail.async_pre_call_hook( - data=data, cache=DualCache(), user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + data=data, + cache=DualCache(), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", ) - + assert result == data # Should return the original data unchanged - + # Clean up del os.environ["LASSO_API_KEY"] del os.environ["LASSO_USER_ID"] @@ -175,21 +195,22 @@ async def test_callback(): async def test_empty_messages(): """Test handling of empty messages""" os.environ["LASSO_API_KEY"] = "test-key" - + lasso_guardrail = LassoGuardrail( - guardrail_name="test-guard", - event_hook="pre_call", - default_on=True + guardrail_name="test-guard", event_hook="pre_call", default_on=True ) - + data = {"messages": []} - + result = await lasso_guardrail.async_pre_call_hook( - data=data, cache=DualCache(), user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + data=data, + cache=DualCache(), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", ) - + assert result == data - + # Clean up del os.environ["LASSO_API_KEY"] @@ -198,48 +219,52 @@ async def test_empty_messages(): async def test_api_error_handling(): """Test handling of API errors""" os.environ["LASSO_API_KEY"] = "test-key" - + lasso_guardrail = LassoGuardrail( - guardrail_name="test-guard", - event_hook="pre_call", - default_on=True + guardrail_name="test-guard", event_hook="pre_call", default_on=True ) - + data = { "messages": [ {"role": "user", "content": "Hello, how are you?"}, ] } - + # Test handling of connection error with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - side_effect=Exception("Connection error") + side_effect=Exception("Connection error"), ): # Expect the guardrail to raise a LassoGuardrailAPIError with pytest.raises(LassoGuardrailAPIError) as excinfo: await lasso_guardrail.async_pre_call_hook( - data=data, cache=DualCache(), user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + data=data, + cache=DualCache(), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", ) - + # Verify the error message assert "Failed to verify request safety with Lasso API" in str(excinfo.value) assert "Connection error" in str(excinfo.value) - + # Test with a different error message with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - side_effect=Exception("API timeout") + side_effect=Exception("API timeout"), ): # Expect the guardrail to raise a LassoGuardrailAPIError with pytest.raises(LassoGuardrailAPIError) as excinfo: await lasso_guardrail.async_pre_call_hook( - data=data, cache=DualCache(), user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + data=data, + cache=DualCache(), + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", ) - + # Verify the error message for the second test assert "Failed to verify request safety with Lasso API" in str(excinfo.value) assert "API timeout" in str(excinfo.value) - + # Clean up del os.environ["LASSO_API_KEY"] diff --git a/tests/local_testing/test_aim_guardrails.py b/tests/local_testing/test_aim_guardrails.py index 95614094b1..b11a424ac3 100644 --- a/tests/local_testing/test_aim_guardrails.py +++ b/tests/local_testing/test_aim_guardrails.py @@ -10,7 +10,7 @@ from fastapi.exceptions import HTTPException from httpx import Request, Response from litellm import DualCache -from litellm.proxy.guardrails.guardrail_hooks.aim import ( +from litellm.proxy.guardrails.guardrail_hooks.aim.aim import ( AimGuardrail, AimGuardrailMissingSecrets, ) @@ -298,7 +298,7 @@ async def test_post_call_stream__all_chunks_are_valid(monkeypatch, length: int): yield websocket_mock monkeypatch.setattr( - "litellm.proxy.guardrails.guardrail_hooks.aim.connect", connect_mock + "litellm.proxy.guardrails.guardrail_hooks.aim.aim.connect", connect_mock ) results = [] @@ -316,6 +316,8 @@ async def test_post_call_stream__all_chunks_are_valid(monkeypatch, length: int): @pytest.mark.asyncio async def test_post_call_stream__blocked_chunks(monkeypatch): + from litellm.proxy.proxy_server import StreamingCallbackError + init_guardrails_v2( all_guardrails=[ { @@ -357,17 +359,27 @@ async def test_post_call_stream__blocked_chunks(monkeypatch): yield websocket_mock monkeypatch.setattr( - "litellm.proxy.guardrails.guardrail_hooks.aim.connect", connect_mock + "litellm.proxy.guardrails.guardrail_hooks.aim.aim.connect", connect_mock ) results = [] - with pytest.raises(StreamingCallbackError, match="Jailbreak detected"): + # For async generators, we need to manually iterate and catch the exception + exception_caught = False + try: async for result in aim_guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), response=llm_response(), request_data=data, ): results.append(result) + except StreamingCallbackError: + exception_caught = True + except Exception as e: + print("INSIDE EXCEPTION") + raise e + + # Assert that the exception was caught + assert exception_caught, "StreamingCallbackError should have been raised" # Chunks that were received before the blocking message should be returned as usual. assert len(results) == 1 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_pangea.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_pangea.py index 6f3d480fdd..78a686f672 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_pangea.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_pangea.py @@ -4,7 +4,7 @@ import httpx import pytest from fastapi import HTTPException -from litellm.proxy.guardrails.guardrail_hooks.pangea import ( +from litellm.proxy.guardrails.guardrail_hooks.pangea.pangea import ( PangeaGuardrailMissingSecrets, PangeaHandler, ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 8d68cc8685..e4c9d056dd 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -3,25 +3,24 @@ Test suite for PANW AIRS Guardrail Integration This test file follows LiteLLM's testing patterns and covers: - Guardrail initialization -- Prompt scanning (blocking and allowing) +- Prompt scanning (blocking and allowing) - Response scanning - Error handling - Configuration validation """ -import pytest -from unittest.mock import AsyncMock, patch, MagicMock -from fastapi import HTTPException from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs import ( PanwPrismaAirsHandler, + initialize_guardrail, ) -from litellm.proxy.guardrails.guardrail_initializers import ( - initialize_panw_prisma_airs, -) -from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.utils import ModelResponse, Choices, Message +from litellm.types.utils import Choices, Message, ModelResponse class TestPanwAirsInitialization: @@ -42,18 +41,22 @@ class TestPanwAirsInitialization: assert handler.api_base == "https://test.panw.com/api" assert handler.profile_name == "test_profile" - def test_initialize_panw_prisma_airs_function(self): - """Test the initialize_panw_prisma_airs function.""" - litellm_params = SimpleNamespace( + def test_initialize_guardrail_function(self): + """Test the initialize_guardrail function.""" + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="panw_prisma_airs", + mode="pre_call", api_key="test_key", - profile_name="test_profile", + profile_name="test_profile", api_base="https://test.panw.com/api", default_on=True, ) guardrail_config = {"guardrail_name": "test_guardrail"} with patch("litellm.logging_callback_manager.add_litellm_callback"): - handler = initialize_panw_prisma_airs(litellm_params, guardrail_config) + handler = initialize_guardrail(litellm_params, guardrail_config) assert isinstance(handler, PanwPrismaAirsHandler) assert handler.guardrail_name == "test_guardrail" @@ -64,12 +67,12 @@ class TestPanwAirsInitialization: profile_name="test_profile", api_base=None, default_on=True, - api_key=None # Missing API key + api_key=None, # Missing API key ) guardrail_config = {"guardrail_name": "test_guardrail"} with pytest.raises(ValueError, match="api_key is required"): - initialize_panw_prisma_airs(litellm_params, guardrail_config) + initialize_guardrail(litellm_params, guardrail_config) def test_missing_profile_name_raises_error(self): """Test that missing profile name raises ValueError.""" @@ -77,12 +80,12 @@ class TestPanwAirsInitialization: api_key="test_key", api_base=None, default_on=True, - profile_name=None # Missing profile name + profile_name=None, # Missing profile name ) guardrail_config = {"guardrail_name": "test_guardrail"} with pytest.raises(ValueError, match="profile_name is required"): - initialize_panw_prisma_airs(litellm_params, guardrail_config) + initialize_guardrail(litellm_params, guardrail_config) class TestPanwAirsPromptScanning: @@ -331,7 +334,7 @@ class TestPanwAirsAPIIntegration: mock_response.raise_for_status.return_value = None with patch( - "litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs.get_async_httpx_client" + "litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs.panw_prisma_airs.get_async_httpx_client" ) as mock_client: mock_async_client = AsyncMock() mock_async_client.post = AsyncMock(return_value=mock_response) @@ -351,7 +354,7 @@ class TestPanwAirsAPIIntegration: """Test API error handling (fail closed).""" # Mock the HTTP client to raise an exception with patch( - "litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs.get_async_httpx_client" + "litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs.panw_prisma_airs.get_async_httpx_client" ) as mock_client: mock_async_client = AsyncMock() mock_async_client.post = AsyncMock(side_effect=Exception("API Error")) @@ -374,7 +377,7 @@ class TestPanwAirsAPIIntegration: mock_response.raise_for_status.return_value = None with patch( - "litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs.get_async_httpx_client" + "litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs.panw_prisma_airs.get_async_httpx_client" ) as mock_client: mock_async_client = AsyncMock() mock_async_client.post = AsyncMock(return_value=mock_response) @@ -403,7 +406,11 @@ class TestPanwAirsConfiguration: def test_default_api_base(self): """Test that default API base is set correctly.""" - litellm_params = SimpleNamespace( + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="panw_prisma_airs", + mode="pre_call", api_key="test_key", profile_name="test_profile", api_base=None, # No api_base provided @@ -412,17 +419,18 @@ class TestPanwAirsConfiguration: guardrail_config = {"guardrail_name": "test"} with patch("litellm.logging_callback_manager.add_litellm_callback"): - handler = initialize_panw_prisma_airs(litellm_params, guardrail_config) + handler = initialize_guardrail(litellm_params, guardrail_config) - assert ( - handler.api_base - == "https://service.api.aisecurity.paloaltonetworks.com/v1/scan/sync/request" - ) + assert handler.api_base == "https://service.api.aisecurity.paloaltonetworks.com" def test_custom_api_base(self): """Test custom API base configuration.""" + from litellm.types.guardrails import LitellmParams + custom_base = "https://custom.panw.com/api/v2/scan" - litellm_params = SimpleNamespace( + litellm_params = LitellmParams( + guardrail="panw_prisma_airs", + mode="pre_call", api_key="test_key", profile_name="test_profile", api_base=custom_base, @@ -431,24 +439,30 @@ class TestPanwAirsConfiguration: guardrail_config = {"guardrail_name": "test"} with patch("litellm.logging_callback_manager.add_litellm_callback"): - handler = initialize_panw_prisma_airs(litellm_params, guardrail_config) + handler = initialize_guardrail(litellm_params, guardrail_config) assert handler.api_base == custom_base def test_default_guardrail_name(self): """Test default guardrail name.""" - litellm_params = SimpleNamespace( + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="panw_prisma_airs", + mode="pre_call", api_key="test_key", profile_name="test_profile", api_base=None, default_on=True, ) - guardrail_config = {} # No guardrail_name + guardrail_config = { + "guardrail_name": "test_guardrail", + } # No guardrail_name with patch("litellm.logging_callback_manager.add_litellm_callback"): - handler = initialize_panw_prisma_airs(litellm_params, guardrail_config) + handler = initialize_guardrail(litellm_params, guardrail_config) - assert handler.guardrail_name == "panw_prisma_airs" + assert handler.guardrail_name == "test_guardrail" if __name__ == "__main__": diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index cf4ab0152b..8567e84950 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -178,9 +178,11 @@ def test_get_provider_specific_params(): AzureContentSafetyTextModerationGuardrail, ) - fields = _get_fields_from_model( - AzureContentSafetyTextModerationGuardrail.get_config_model() - ) + config_model = AzureContentSafetyTextModerationGuardrail.get_config_model() + if config_model is None: + pytest.skip("Azure config model not available") + + fields = _get_fields_from_model(config_model) print("FIELDS", fields) # Test that we get the expected nested structure @@ -231,3 +233,70 @@ def test_get_provider_specific_params(): assert ( nested_fields["outputType"]["type"] == "select" ) # Literal type should be select + + +def test_optional_params_not_returned_when_not_overridden(): + """Test that optional_params is not returned when the config model doesn't override it""" + from typing import Optional + + from pydantic import BaseModel, Field + + from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + + class TestGuardrailConfig(GuardrailConfigModel): + api_key: Optional[str] = Field( + default=None, + description="Test API key", + ) + api_base: Optional[str] = Field( + default=None, + description="Test API base", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Test Guardrail" + + # Get fields from the model + fields = _get_fields_from_model(TestGuardrailConfig) + print("FIELDS", fields) + assert "optional_params" not in fields + + +def test_optional_params_returned_when_properly_overridden(): + """Test that optional_params IS returned when the config model properly overrides it""" + from typing import Optional + + from pydantic import BaseModel, Field + + from litellm.proxy.guardrails.guardrail_endpoints import _get_fields_from_model + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + + # Create specific optional params model + class SpecificOptionalParams(BaseModel): + threshold: Optional[float] = Field( + default=0.5, description="Detection threshold" + ) + categories: Optional[List[str]] = Field( + default=None, description="Categories to check" + ) + + # Create a config model that DOES override optional_params with a specific type + class TestGuardrailConfigWithOptionalParams( + GuardrailConfigModel[SpecificOptionalParams] + ): + api_key: Optional[str] = Field( + default=None, + description="Test API key", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Test Guardrail With Optional Params" + + # Get fields from the model + fields = _get_fields_from_model(TestGuardrailConfigWithOptionalParams) + + print("FIELDS", fields) + assert "optional_params" in fields diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py new file mode 100644 index 0000000000..f05ec653fb --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -0,0 +1,17 @@ +from litellm.proxy.guardrails.guardrail_registry import ( + get_guardrail_initializer_from_hooks, +) + + +def test_get_guardrail_initializer_from_hooks(): + initializers = get_guardrail_initializer_from_hooks() + print(f"initializers: {initializers}") + assert "aim" in initializers + + +def test_guardrail_class_registry(): + from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry + + print(f"guardrail_class_registry: {guardrail_class_registry}") + assert "aim" in guardrail_class_registry + assert "aporia" in guardrail_class_registry diff --git a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx index 285b1c7673..f68decf942 100644 --- a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx @@ -2,7 +2,7 @@ import React, { useState, useEffect } from 'react'; import { Card, Form, Typography, Select, Input, Switch, Tooltip, Modal, message, Divider, Space, Tag, Image, Steps } from 'antd'; import { Button, TextInput } from '@tremor/react'; import type { FormInstance } from 'antd'; -import { GuardrailProviders, guardrail_provider_map, shouldRenderPIIConfigSettings, guardrailLogoMap } from './guardrail_info_helpers'; +import { GuardrailProviders, guardrail_provider_map, shouldRenderPIIConfigSettings, guardrailLogoMap, populateGuardrailProviders, populateGuardrailProviderMap, getGuardrailProviders } from './guardrail_info_helpers'; import { createGuardrailCall, getGuardrailUISettings, getGuardrailProviderSpecificParams } from '../networking'; import PiiConfiguration from './pii_configuration'; import GuardrailProviderFields from './guardrail_provider_fields'; @@ -95,6 +95,10 @@ const AddGuardrailForm: React.FC = ({ setGuardrailSettings(uiSettings); setProviderParams(providerParamsResp); + + // Populate dynamic providers from API response + populateGuardrailProviders(providerParamsResp); + populateGuardrailProviderMap(providerParamsResp); } catch (error) { console.error('Error fetching guardrail data:', error); message.error('Failed to load guardrail configuration'); @@ -288,6 +292,7 @@ const AddGuardrailForm: React.FC = ({ // Use pre-fetched provider params to copy recognised params if (providerParams && selectedProvider) { const providerKey = guardrail_provider_map[selectedProvider]?.toLowerCase(); + console.log("providerKey: ", providerKey); const providerSpecificParams = providerParams[providerKey] || {}; const allowedParams = new Set(); @@ -368,7 +373,7 @@ const AddGuardrailForm: React.FC = ({ dropdownRender={menu => menu} showSearch={true} > - {Object.entries(GuardrailProviders).map(([key, value]) => ( + {Object.entries(getGuardrailProviders()).map(([key, value]) => (