From e5d68e522220f02e3271a69d829a8d9e5faef61e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 25 Jul 2025 15:37:27 -0700 Subject: [PATCH] feat(proxy/utils.py): support model level guardrails on stream event enables guardrails to work with streaming --- litellm/integrations/custom_guardrail.py | 1 + .../out/model_hub_table/index.html | 1 + litellm/proxy/_new_secret_config.yaml | 7 -- litellm/proxy/utils.py | 79 ++++++++++++++++++- 4 files changed, 80 insertions(+), 8 deletions(-) create mode 100644 litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 5bf84bd8ce..b679235433 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -129,6 +129,7 @@ class CustomGuardrail(CustomLogger): self, requested_guardrails: Union[List[str], List[Dict[str, DynamicGuardrailParams]]], ) -> bool: + for _guardrail in requested_guardrails: if isinstance(_guardrail, dict): if self.guardrail_name in _guardrail: diff --git a/litellm/proxy/_experimental/out/model_hub_table/index.html b/litellm/proxy/_experimental/out/model_hub_table/index.html new file mode 100644 index 0000000000..155cf9b802 --- /dev/null +++ b/litellm/proxy/_experimental/out/model_hub_table/index.html @@ -0,0 +1 @@ +LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index c249da0534..c9d8c4f7f5 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -10,13 +10,6 @@ model_list: model: openai/gpt-4o guardrails: - - guardrail_name: "presidio-pii" - litellm_params: - guardrail: presidio # supported values: "aporia", "bedrock", "lakera", "presidio" - mode: "pre_call" - presidio_language: "en" # optional: set default language for PII analysis - pii_entities_config: - PERSON: "BLOCK" # Will mask credit card numbers - guardrail_name: azure-text-moderation litellm_params: guardrail: azure/text_moderations diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 655b50cfc2..02b7e14229 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1077,6 +1077,8 @@ class ProxyLogging: Covers: 1. /chat/completions """ + from litellm.proxy.proxy_server import llm_router + response_str: Optional[str] = None if isinstance(response, (ModelResponse, ModelResponseStream)): response_str = litellm.get_response_string(response_obj=response) @@ -1088,9 +1090,15 @@ class ProxyLogging: # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks + ## CHECK FOR MODEL-LEVEL GUARDRAILS + modified_data = _check_and_merge_model_level_guardrails( + data=data, llm_router=llm_router + ) + if ( callback.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call + data=modified_data, + event_type=GuardrailEventHooks.post_call, ) is not True ): @@ -3012,6 +3020,74 @@ def _to_ns(dt): return int(dt.timestamp() * 1e9) +def _check_and_merge_model_level_guardrails( + data: dict, llm_router: Optional[Router] +) -> dict: + """ + Check if the model has guardrails defined and merge them with existing guardrails in the request data. + + Args: + data: The request data dict + llm_router: The LLM router instance to get deployment info from + + Returns: + Modified data dict with merged guardrails (if any model-level guardrails exist) + """ + if llm_router is None: + return data + + # Get the model ID from the data + metadata = data.get("metadata") or {} + model_info = metadata.get("model_info") or {} + model_id = model_info.get("id", None) + + if model_id is None: + return data + + # Check if the model has guardrails + deployment = llm_router.get_deployment(model_id=model_id) + if deployment is None: + return data + + model_level_guardrails = deployment.litellm_params.get("guardrails") + + if model_level_guardrails is None: + return data + + # Merge model-level guardrails with existing ones + return _merge_guardrails_with_existing(data, model_level_guardrails) + + +def _merge_guardrails_with_existing(data: dict, model_level_guardrails: Any) -> dict: + """ + Merge model-level guardrails with any existing guardrails in the request data. + + Args: + data: The request data dict + model_level_guardrails: Guardrails defined at the model level + + Returns: + Modified data dict with merged guardrails in metadata + """ + modified_data = data.copy() + metadata = modified_data.setdefault("metadata", {}) + existing_guardrails = metadata.get("guardrails", []) + + # Ensure existing_guardrails is a list + if not isinstance(existing_guardrails, list): + existing_guardrails = [existing_guardrails] if existing_guardrails else [] + + # Ensure model_level_guardrails is a list + if not isinstance(model_level_guardrails, list): + model_level_guardrails = ( + [model_level_guardrails] if model_level_guardrails else [] + ) + + # Combine existing and model-level guardrails + metadata["guardrails"] = list(set(existing_guardrails + model_level_guardrails)) + return modified_data + + def get_error_message_str(e: Exception) -> str: error_message = "" if isinstance(e, HTTPException): @@ -3194,6 +3270,7 @@ def is_valid_api_key(key: str) -> bool: - Length between 20 and 100 characters """ import re + if not isinstance(key, str): return False if 3 <= len(key) <= 100: