diff --git a/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md b/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md index 9c654cd1560..884a7397bde 100644 --- a/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md +++ b/docs/my-website/docs/adding_provider/simple_guardrail_tutorial.md @@ -101,12 +101,11 @@ model_list: - model_name: gpt-4 litellm_params: model: gpt-4 - api_key: os.environ/OPENAI_API_KEY + api_key: os.environ/OPENAI_API_KEY -litellm_settings: - guardrails: +guardrails: - guardrail_name: my_guardrail - litellm_params: + litellm_params: guardrail: my_guardrail mode: during_call api_key: os.environ/MY_GUARDRAIL_API_KEY diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6aeb51d5817..485bee4f191 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21488,6 +21488,20 @@ "supports_tool_choice": true, "supports_web_search": true }, + "moonshot/kimi-k2.5": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "moonshot", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://platform.moonshot.ai/docs/pricing/chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 2e-06, diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 13eeae14485..6f527e268b2 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -14,3 +14,14 @@ model_list: litellm_params: model: openai/gpt-4.1-mini +guardrails: + - guardrail_name: redact-ssn + litellm_params: + guardrail: custom_code + mode: pre_call + custom_code: | + def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"\d{3}-\d{2}-\d{4}"): + return block("SSN detected in message") + return allow() \ No newline at end of file diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py new file mode 100644 index 00000000000..747b188feea --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py @@ -0,0 +1,65 @@ +"""Custom code guardrail integration for LiteLLM. + +This module allows users to write custom guardrail logic using Python-like code +that runs in a sandboxed environment with access to LiteLLM-provided primitives. +""" + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .custom_code_guardrail import CustomCodeGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> CustomCodeGuardrail: + """ + Initialize a custom code guardrail. + + Args: + litellm_params: Configuration parameters including the custom code + guardrail: The guardrail configuration dict + + Returns: + CustomCodeGuardrail instance + """ + import litellm + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("Custom code guardrail requires a guardrail_name") + + # Get the custom code from litellm_params + custom_code = getattr(litellm_params, "custom_code", None) + if not custom_code: + raise ValueError( + "Custom code guardrail requires 'custom_code' in litellm_params" + ) + + custom_code_guardrail = CustomCodeGuardrail( + guardrail_name=guardrail_name, + custom_code=custom_code, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(custom_code_guardrail) + return custom_code_guardrail + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.CUSTOM_CODE.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.CUSTOM_CODE.value: CustomCodeGuardrail, +} + +__all__ = [ + "CustomCodeGuardrail", + "initialize_guardrail", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py new file mode 100644 index 00000000000..a0ca324411c --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -0,0 +1,372 @@ +""" +Custom code guardrail for LiteLLM. + +This module provides a guardrail that executes user-defined Python-like code +to implement custom guardrail logic. The code runs in a sandboxed environment +with access to LiteLLM-provided primitives for common guardrail operations. + +Example custom code: + + def apply_guardrail(inputs, request_data, input_type): + '''Block messages containing SSNs''' + for text in inputs["texts"]: + if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): + return block("Social Security Number detected") + return allow() +""" + +import threading +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +from litellm.types.utils import GenericGuardrailAPIInputs + +from .primitives import get_custom_code_primitives + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +class CustomCodeGuardrailError(Exception): + """Raised when custom code guardrail execution fails.""" + + def __init__(self, message: str, details: Optional[Dict[str, Any]] = None) -> None: + super().__init__(message) + self.details = details or {} + + +class CustomCodeCompilationError(CustomCodeGuardrailError): + """Raised when custom code fails to compile.""" + + +class CustomCodeExecutionError(CustomCodeGuardrailError): + """Raised when custom code fails during execution.""" + + +class CustomCodeGuardrailConfigModel(GuardrailConfigModel): + """Configuration parameters for the custom code guardrail.""" + + custom_code: str + """The Python-like code containing the apply_guardrail function.""" + + +class CustomCodeGuardrail(CustomGuardrail): + """ + Guardrail that executes user-defined Python-like code. + + The code runs in a sandboxed environment that provides: + - Access to LiteLLM primitives (regex_match, json_parse, etc.) + - No file I/O or network access + - No imports allowed + + Users write an `apply_guardrail(inputs, request_data, input_type)` function + that returns one of: + - allow() - let the request/response through + - block(reason) - reject with a message + - modify(texts=...) - transform the content + + Example: + def apply_guardrail(inputs, request_data, input_type): + for text in inputs["texts"]: + if regex_match(text, r"password"): + return block("Sensitive content detected") + return allow() + """ + + def __init__( + self, + custom_code: str, + guardrail_name: Optional[str] = "custom_code", + **kwargs: Any, + ) -> None: + """ + Initialize the custom code guardrail. + + Args: + custom_code: The source code containing apply_guardrail function + guardrail_name: Name of this guardrail instance + **kwargs: Additional arguments passed to CustomGuardrail + """ + self.custom_code = custom_code + self._compiled_function: Optional[Any] = None + self._compile_lock = threading.Lock() + self._compile_error: Optional[str] = None + + supported_event_hooks = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + ] + + super().__init__( + guardrail_name=guardrail_name, + supported_event_hooks=supported_event_hooks, + **kwargs, + ) + + # Compile the code on initialization + self._compile_custom_code() + + @staticmethod + def get_config_model() -> Optional[Type[GuardrailConfigModel]]: + """Returns the config model for the UI.""" + return CustomCodeGuardrailConfigModel + + def _compile_custom_code(self) -> None: + """ + Compile the custom code and extract the apply_guardrail function. + + The code runs in a sandboxed environment with only the allowed primitives. + """ + with self._compile_lock: + if self._compiled_function is not None: + return + + try: + # Create a restricted execution environment + # Only include our safe primitives + exec_globals = get_custom_code_primitives().copy() + + # Execute the user code in the restricted environment + exec(compile(self.custom_code, "", "exec"), exec_globals) + + # Extract the apply_guardrail function + if "apply_guardrail" not in exec_globals: + raise CustomCodeCompilationError( + "Custom code must define an 'apply_guardrail' function. " + "Expected signature: apply_guardrail(inputs, request_data, input_type)" + ) + + apply_fn = exec_globals["apply_guardrail"] + if not callable(apply_fn): + raise CustomCodeCompilationError( + "'apply_guardrail' must be a callable function" + ) + + self._compiled_function = apply_fn + verbose_proxy_logger.debug( + f"Custom code guardrail '{self.guardrail_name}' compiled successfully" + ) + + except SyntaxError as e: + self._compile_error = f"Syntax error in custom code: {e}" + raise CustomCodeCompilationError(self._compile_error) from e + except CustomCodeCompilationError: + raise + except Exception as e: + self._compile_error = f"Failed to compile custom code: {e}" + raise CustomCodeCompilationError(self._compile_error) from e + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """ + Apply the custom code guardrail to the inputs. + + This method calls the user-defined apply_guardrail function and + processes its result to determine the appropriate action. + + Args: + inputs: Dictionary containing texts, images, tool_calls + request_data: The original request data with metadata + input_type: "request" for pre-call, "response" for post-call + logging_obj: Optional logging object + + Returns: + GenericGuardrailAPIInputs - possibly modified + + Raises: + HTTPException: If content is blocked + CustomCodeExecutionError: If execution fails + """ + if self._compiled_function is None: + if self._compile_error: + raise CustomCodeExecutionError( + f"Custom code guardrail not compiled: {self._compile_error}" + ) + raise CustomCodeExecutionError("Custom code guardrail not compiled") + + try: + # Prepare inputs dict for the function + + # Prepare request_data with safe subset of information + safe_request_data = self._prepare_safe_request_data(request_data) + + # Execute the custom function + result = self._compiled_function(inputs, safe_request_data, input_type) + + # Process the result + return self._process_result( + result=result, + inputs=inputs, + request_data=request_data, + input_type=input_type, + ) + + except HTTPException: + # Re-raise HTTP exceptions (from block action) + raise + except Exception as e: + verbose_proxy_logger.error( + f"Custom code guardrail '{self.guardrail_name}' execution error: {e}" + ) + raise CustomCodeExecutionError( + f"Custom code guardrail execution failed: {e}", + details={ + "guardrail_name": self.guardrail_name, + "input_type": input_type, + }, + ) from e + + def _prepare_safe_request_data(self, request_data: dict) -> Dict[str, Any]: + """ + Prepare a safe subset of request_data for code execution. + + This filters out sensitive information and provides only what's + needed for guardrail logic. + + Args: + request_data: The full request data + + Returns: + Safe subset of request data + """ + return { + "model": request_data.get("model"), + "user_id": request_data.get("user_api_key_user_id"), + "team_id": request_data.get("user_api_key_team_id"), + "end_user_id": request_data.get("user_api_key_end_user_id"), + "metadata": request_data.get("metadata", {}), + } + + def _process_result( + self, + result: Any, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + ) -> GenericGuardrailAPIInputs: + """ + Process the result from the custom code function. + + Args: + result: The return value from apply_guardrail + inputs: The original inputs + request_data: The request data + input_type: "request" or "response" + + Returns: + GenericGuardrailAPIInputs - possibly modified + + Raises: + HTTPException: If action is "block" + """ + if not isinstance(result, dict): + verbose_proxy_logger.warning( + f"Custom code guardrail '{self.guardrail_name}': " + f"Expected dict result, got {type(result).__name__}. Treating as allow." + ) + return inputs + + action = result.get("action", "allow") + + if action == "allow": + verbose_proxy_logger.debug( + f"Custom code guardrail '{self.guardrail_name}': Allowing {input_type}" + ) + return inputs + + elif action == "block": + reason = result.get("reason", "Blocked by custom code guardrail") + detection_info = result.get("detection_info", {}) + + verbose_proxy_logger.info( + f"Custom code guardrail '{self.guardrail_name}': Blocking {input_type} - {reason}" + ) + + is_output = input_type == "response" + + # For pre-call, raise passthrough exception to return synthetic response + if not is_output: + self.raise_passthrough_exception( + violation_message=reason, + request_data=request_data, + detection_info=detection_info, + ) + + # For post-call, raise HTTP exception + raise HTTPException( + status_code=400, + detail={ + "error": reason, + "guardrail": self.guardrail_name, + "detection_info": detection_info, + }, + ) + + elif action == "modify": + verbose_proxy_logger.debug( + f"Custom code guardrail '{self.guardrail_name}': Modifying {input_type}" + ) + + # Apply modifications + modified_inputs = dict(inputs) + + if "texts" in result and result["texts"] is not None: + modified_inputs["texts"] = result["texts"] + + if "images" in result and result["images"] is not None: + modified_inputs["images"] = result["images"] + + if "tool_calls" in result and result["tool_calls"] is not None: + modified_inputs["tool_calls"] = result["tool_calls"] + + return cast(GenericGuardrailAPIInputs, modified_inputs) + + else: + verbose_proxy_logger.warning( + f"Custom code guardrail '{self.guardrail_name}': " + f"Unknown action '{action}'. Treating as allow." + ) + return inputs + + def update_custom_code(self, new_code: str) -> None: + """ + Update the custom code and recompile. + + This method allows hot-reloading of guardrail logic without + restarting the server. + + Args: + new_code: The new source code + + Raises: + CustomCodeCompilationError: If the new code fails to compile + """ + with self._compile_lock: + # Reset state + old_function = self._compiled_function + old_code = self.custom_code + self._compiled_function = None + self._compile_error = None + + try: + self.custom_code = new_code + self._compile_custom_code() + verbose_proxy_logger.info( + f"Custom code guardrail '{self.guardrail_name}': Code updated successfully" + ) + except CustomCodeCompilationError: + # Rollback on failure + self.custom_code = old_code + self._compiled_function = old_function + raise diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py new file mode 100644 index 00000000000..2cdb65ca5d3 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -0,0 +1,587 @@ +""" +Built-in primitives provided to custom code guardrails. + +These functions are injected into the custom code execution environment +and provide safe, sandboxed functionality for common guardrail operations. +""" + +import json +import re +from typing import Any, Dict, List, Optional +from urllib.parse import urlparse + +from litellm._logging import verbose_proxy_logger + +# ============================================================================= +# Result Types - Used by Starlark code to return guardrail decisions +# ============================================================================= + + +def allow() -> Dict[str, Any]: + """ + Allow the request/response to proceed unchanged. + + Returns: + Dict indicating the request should be allowed + """ + return {"action": "allow"} + + +def block( + reason: str, detection_info: Optional[Dict[str, Any]] = None +) -> Dict[str, Any]: + """ + Block the request/response with a reason. + + Args: + reason: Human-readable reason for blocking + detection_info: Optional additional detection metadata + + Returns: + Dict indicating the request should be blocked + """ + result: Dict[str, Any] = {"action": "block", "reason": reason} + if detection_info: + result["detection_info"] = detection_info + return result + + +def modify( + texts: Optional[List[str]] = None, + images: Optional[List[Any]] = None, + tool_calls: Optional[List[Any]] = None, +) -> Dict[str, Any]: + """ + Modify the request/response content. + + Args: + texts: Modified text content (if None, keeps original) + images: Modified image content (if None, keeps original) + tool_calls: Modified tool calls (if None, keeps original) + + Returns: + Dict indicating the content should be modified + """ + result: Dict[str, Any] = {"action": "modify"} + if texts is not None: + result["texts"] = texts + if images is not None: + result["images"] = images + if tool_calls is not None: + result["tool_calls"] = tool_calls + return result + + +# ============================================================================= +# Regex Primitives +# ============================================================================= + + +def regex_match(text: str, pattern: str, flags: int = 0) -> bool: + """ + Check if a regex pattern matches anywhere in the text. + + Args: + text: The text to search in + pattern: The regex pattern to match + flags: Optional regex flags (default: 0) + + Returns: + True if pattern matches, False otherwise + """ + try: + return bool(re.search(pattern, text, flags)) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_match error: {e}") + return False + + +def regex_match_all(text: str, pattern: str, flags: int = 0) -> bool: + """ + Check if a regex pattern matches the entire text. + + Args: + text: The text to match + pattern: The regex pattern + flags: Optional regex flags + + Returns: + True if pattern matches entire text, False otherwise + """ + try: + return bool(re.fullmatch(pattern, text, flags)) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_match_all error: {e}") + return False + + +def regex_replace(text: str, pattern: str, replacement: str, flags: int = 0) -> str: + """ + Replace all occurrences of a pattern in text. + + Args: + text: The text to modify + pattern: The regex pattern to find + replacement: The replacement string + flags: Optional regex flags + + Returns: + The text with replacements applied + """ + try: + return re.sub(pattern, replacement, text, flags=flags) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_replace error: {e}") + return text + + +def regex_find_all(text: str, pattern: str, flags: int = 0) -> List[str]: + """ + Find all occurrences of a pattern in text. + + Args: + text: The text to search + pattern: The regex pattern to find + flags: Optional regex flags + + Returns: + List of all matches + """ + try: + return re.findall(pattern, text, flags) + except re.error as e: + verbose_proxy_logger.warning(f"Starlark regex_find_all error: {e}") + return [] + + +# ============================================================================= +# JSON Primitives +# ============================================================================= + + +def json_parse(text: str) -> Optional[Any]: + """ + Parse a JSON string into a Python object. + + Args: + text: The JSON string to parse + + Returns: + Parsed Python object, or None if parsing fails + """ + try: + return json.loads(text) + except (json.JSONDecodeError, TypeError) as e: + verbose_proxy_logger.debug(f"Starlark json_parse error: {e}") + return None + + +def json_stringify(obj: Any) -> str: + """ + Convert a Python object to a JSON string. + + Args: + obj: The object to serialize + + Returns: + JSON string representation + """ + try: + return json.dumps(obj) + except (TypeError, ValueError) as e: + verbose_proxy_logger.warning(f"Starlark json_stringify error: {e}") + return "" + + +def json_schema_valid(obj: Any, schema: Dict[str, Any]) -> bool: + """ + Validate an object against a JSON schema. + + Args: + obj: The object to validate + schema: The JSON schema to validate against + + Returns: + True if valid, False otherwise + """ + try: + # Try to import jsonschema, fall back to basic validation if not available + try: + import jsonschema + + jsonschema.validate(instance=obj, schema=schema) + return True + except ImportError: + # Basic validation without jsonschema library + return _basic_json_schema_validate(obj, schema) + except Exception as validation_error: + # Catch jsonschema.ValidationError and other validation errors + if "ValidationError" in type(validation_error).__name__: + return False + raise + except Exception as e: + verbose_proxy_logger.warning(f"Custom code json_schema_valid error: {e}") + return False + + +def _basic_json_schema_validate(obj: Any, schema: Dict[str, Any]) -> bool: + """ + Basic JSON schema validation without external library. + Handles: type, required, properties + """ + # Check type + schema_type = schema.get("type") + if schema_type: + type_map = { + "object": dict, + "array": list, + "string": str, + "number": (int, float), + "integer": int, + "boolean": bool, + "null": type(None), + } + expected_type = type_map.get(schema_type) + if expected_type and not isinstance(obj, expected_type): + return False + + # Check required fields + if isinstance(obj, dict): + required = schema.get("required", []) + for field in required: + if field not in obj: + return False + + # Check properties + properties = schema.get("properties", {}) + for prop_name, prop_schema in properties.items(): + if prop_name in obj: + if not _basic_json_schema_validate(obj[prop_name], prop_schema): + return False + + return True + + +# ============================================================================= +# URL Primitives +# ============================================================================= + + +# Common URL pattern for extraction +_URL_PATTERN = re.compile( + r"https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+[^\s]*", re.IGNORECASE +) + + +def extract_urls(text: str) -> List[str]: + """ + Extract all URLs from text. + + Args: + text: The text to search for URLs + + Returns: + List of URLs found in the text + """ + return _URL_PATTERN.findall(text) + + +def is_valid_url(url: str) -> bool: + """ + Check if a URL is syntactically valid. + + Args: + url: The URL to validate + + Returns: + True if the URL is valid, False otherwise + """ + try: + result = urlparse(url) + return all([result.scheme, result.netloc]) + except Exception: + return False + + +def all_urls_valid(text: str) -> bool: + """ + Check if all URLs in text are valid. + + Args: + text: The text containing URLs + + Returns: + True if all URLs are valid (or no URLs), False otherwise + """ + urls = extract_urls(text) + return all(is_valid_url(url) for url in urls) + + +def get_url_domain(url: str) -> Optional[str]: + """ + Extract the domain from a URL. + + Args: + url: The URL to parse + + Returns: + The domain, or None if invalid + """ + try: + result = urlparse(url) + return result.netloc if result.netloc else None + except Exception: + return None + + +# ============================================================================= +# Code Detection Primitives +# ============================================================================= + + +# Common code patterns for detection +_CODE_PATTERNS = { + "sql": [ + r"\b(SELECT|INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE)\b.*\b(FROM|INTO|TABLE|SET|WHERE)\b", + r"\b(SELECT)\s+[\w\*,\s]+\s+FROM\s+\w+", + r"\b(INSERT\s+INTO|UPDATE\s+\w+\s+SET|DELETE\s+FROM)\b", + ], + "python": [ + r"^\s*(def|class|import|from|if|for|while|try|except|with)\s+", + r"^\s*@\w+", # decorators + r"\b(print|len|range|str|int|float|list|dict|set)\s*\(", + ], + "javascript": [ + r"\b(function|const|let|var|class|import|export)\s+", + r"=>", # arrow functions + r"\b(console\.(log|error|warn))\s*\(", + ], + "typescript": [ + r":\s*(string|number|boolean|any|void|never)\b", + r"\b(interface|type|enum)\s+\w+", + r"<[A-Z]\w*>", # generics + ], + "java": [ + r"\b(public|private|protected)\s+(static\s+)?(class|void|int|String)\b", + r"\bSystem\.(out|err)\.print", + ], + "go": [ + r"\bfunc\s+\w+\s*\(", + r"\b(package|import)\s+", + r":=", # short variable declaration + ], + "rust": [ + r"\b(fn|let|mut|impl|struct|enum|pub|mod)\s+", + r"->", # return type + r"\b(println!|format!)\s*\(", + ], + "shell": [ + r"^#!.*\b(bash|sh|zsh)\b", + r"\b(echo|grep|sed|awk|cat|ls|cd|mkdir|rm)\s+", + r"\$\{?\w+\}?", # variable expansion + ], + "html": [ + r"<\s*(html|head|body|div|span|p|a|img|script|style)\b[^>]*>", + r"", + ], + "css": [ + r"\{[^}]*:\s*[^}]+;[^}]*\}", + r"@(media|keyframes|import|font-face)\b", + ], +} + + +def detect_code(text: str) -> bool: + """ + Check if text contains code of any language. + + Args: + text: The text to check + + Returns: + True if code is detected, False otherwise + """ + return len(detect_code_languages(text)) > 0 + + +def detect_code_languages(text: str) -> List[str]: + """ + Detect which programming languages are present in text. + + Args: + text: The text to analyze + + Returns: + List of detected language names + """ + detected = [] + for lang, patterns in _CODE_PATTERNS.items(): + for pattern in patterns: + try: + if re.search(pattern, text, re.IGNORECASE | re.MULTILINE): + detected.append(lang) + break # Only add each language once + except re.error: + continue + return detected + + +def contains_code_language(text: str, languages: List[str]) -> bool: + """ + Check if text contains code from specific languages. + + Args: + text: The text to check + languages: List of language names to check for + + Returns: + True if any of the specified languages are detected + """ + detected = detect_code_languages(text) + return any(lang.lower() in [d.lower() for d in detected] for lang in languages) + + +# ============================================================================= +# Text Utility Primitives +# ============================================================================= + + +def contains(text: str, substring: str) -> bool: + """ + Check if text contains a substring. + + Args: + text: The text to search in + substring: The substring to find + + Returns: + True if substring is found, False otherwise + """ + return substring in text + + +def contains_any(text: str, substrings: List[str]) -> bool: + """ + Check if text contains any of the given substrings. + + Args: + text: The text to search in + substrings: List of substrings to find + + Returns: + True if any substring is found, False otherwise + """ + return any(s in text for s in substrings) + + +def contains_all(text: str, substrings: List[str]) -> bool: + """ + Check if text contains all of the given substrings. + + Args: + text: The text to search in + substrings: List of substrings to find + + Returns: + True if all substrings are found, False otherwise + """ + return all(s in text for s in substrings) + + +def word_count(text: str) -> int: + """ + Count the number of words in text. + + Args: + text: The text to count words in + + Returns: + Number of words + """ + return len(text.split()) + + +def char_count(text: str) -> int: + """ + Count the number of characters in text. + + Args: + text: The text to count characters in + + Returns: + Number of characters + """ + return len(text) + + +def lower(text: str) -> str: + """Convert text to lowercase.""" + return text.lower() + + +def upper(text: str) -> str: + """Convert text to uppercase.""" + return text.upper() + + +def trim(text: str) -> str: + """Remove leading and trailing whitespace.""" + return text.strip() + + +# ============================================================================= +# Primitives Registry +# ============================================================================= + + +def get_custom_code_primitives() -> Dict[str, Any]: + """ + Get all primitives to inject into the custom code environment. + + Returns: + Dict of function name to function + """ + return { + # Result types + "allow": allow, + "block": block, + "modify": modify, + # Regex + "regex_match": regex_match, + "regex_match_all": regex_match_all, + "regex_replace": regex_replace, + "regex_find_all": regex_find_all, + # JSON + "json_parse": json_parse, + "json_stringify": json_stringify, + "json_schema_valid": json_schema_valid, + # URL + "extract_urls": extract_urls, + "is_valid_url": is_valid_url, + "all_urls_valid": all_urls_valid, + "get_url_domain": get_url_domain, + # Code detection + "detect_code": detect_code, + "detect_code_languages": detect_code_languages, + "contains_code_language": contains_code_language, + # Text utilities + "contains": contains, + "contains_any": contains_any, + "contains_all": contains_all, + "word_count": word_count, + "char_count": char_count, + "lower": lower, + "upper": upper, + "trim": trim, + # Python builtins (safe subset) + "len": len, + "str": str, + "int": int, + "float": float, + "bool": bool, + "list": list, + "dict": dict, + "True": True, + "False": False, + "None": None, + } diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index ca22049720e..74ccb34ca6e 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -14,14 +14,14 @@ from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import ( from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( IBMGuardrailsBaseConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( - ToolPermissionGuardrailConfigModel, +from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( + ContentFilterCategoryConfig, ) from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( QualifireGuardrailConfigModel, ) -from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( - ContentFilterCategoryConfig, +from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( + ToolPermissionGuardrailConfigModel, ) """ @@ -68,6 +68,7 @@ class SupportedGuardrailIntegrations(Enum): PROMPT_SECURITY = "prompt_security" GENERIC_GUARDRAIL_API = "generic_guardrail_api" QUALIFIRE = "qualifire" + CUSTOM_CODE = "custom_code" class Role(Enum): @@ -296,13 +297,7 @@ class PresidioConfigModel(PresidioPresidioConfigModelUserInterface): pii_entities_config: Optional[Dict[Union[PiiEntityType, str], PiiAction]] = Field( default=None, description="Configuration for PII entity types and actions" ) - presidio_filter_scope: Literal["input", "output", "both"] = Field( - default="both", - description=( - "Where to apply Presidio checks: 'input' runs on user → model traffic, " - "'output' runs on model → user traffic, and 'both' applies to both." - ), - ) + presidio_score_thresholds: Optional[Dict[Union[PiiEntityType, str], float]] = Field( default=None, description=( @@ -656,6 +651,12 @@ class BaseLitellmParams( description="Additional provider-specific parameters for generic guardrail APIs", ) + # Custom code guardrail params + custom_code: Optional[str] = Field( + default=None, + description="Python-like code containing the apply_guardrail function for custom guardrail logic", + ) + model_config = ConfigDict(extra="allow", protected_namespaces=())