From ade2ae1977a2d1900b2d4966c13b86a82bf88939 Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sat, 14 Feb 2026 14:33:55 -0300 Subject: [PATCH 1/4] refactor: extract helper functions to reduce statement count Refactor add_guardrails_from_policy_engine by extracting two helper functions to reduce the statement count from 52 to under 50: - _match_and_track_policies: Handles policy matching and tracking - _apply_resolved_guardrails_to_metadata: Handles guardrail resolution Fixes Ruff PLR0915 linting error (too many statements). Co-Authored-By: Claude Sonnet 4.5 --- litellm/proxy/litellm_pre_call_utils.py | 147 ++++++++++++++---------- 1 file changed, 89 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index fa024cc33d..d211501ef9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1515,75 +1515,26 @@ def move_guardrails_to_metadata( ] = request_body_guardrail_config -def add_guardrails_from_policy_engine( +def _match_and_track_policies( data: dict, - metadata_variable_name: str, - user_api_key_dict: UserAPIKeyAuth, -) -> None: + context: "PolicyMatchContext", + request_body_policies: Any, +) -> tuple[list[str], dict[str, str]]: """ - Add guardrails from the policy engine based on request context. + Match policies via attachments and request body, track them in metadata. - This function: - 1. Extracts "policies" from request body (if present) for dynamic policy application - 2. Gets matching policies based on team_alias, key_alias, and model (via attachments) - 3. Combines dynamic policies with attachment-based policies - 4. Resolves guardrails from all policies (including inheritance) - 5. Adds guardrails to request metadata - 6. Tracks applied policies in metadata for response headers - 7. Removes "policies" from request body so it's not forwarded to LLM provider - - Args: - data: The request data to update - metadata_variable_name: The name of the metadata field in data - user_api_key_dict: The user's API key authentication info + Returns: + Tuple of (applied_policy_names, policy_reasons) """ from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.callback_utils import ( add_policy_sources_to_metadata, add_policy_to_applied_policies_header, ) - from litellm.proxy.common_utils.http_parsing_utils import ( - get_tags_from_request_body, - ) from litellm.proxy.policy_engine.attachment_registry import ( get_attachment_registry, ) from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - from litellm.proxy.policy_engine.policy_registry import get_policy_registry - from litellm.proxy.policy_engine.policy_resolver import PolicyResolver - from litellm.types.proxy.policy_engine import PolicyMatchContext - - # Extract dynamic policies from request body (if present) - # These will be combined with attachment-based policies - request_body_policies = data.pop("policies", None) - - registry = get_policy_registry() - verbose_proxy_logger.debug( - f"Policy engine: registry initialized={registry.is_initialized()}, " - f"policy_count={len(registry.get_all_policies())}" - ) - if not registry.is_initialized(): - verbose_proxy_logger.debug( - "Policy engine not initialized, skipping policy matching" - ) - return - - # Extract tags using the shared helper (handles metadata / litellm_metadata, - # top-level tags, deduplication, and type filtering). - - all_tags = get_tags_from_request_body(data) or None - - context = PolicyMatchContext( - team_alias=user_api_key_dict.team_alias, - key_alias=user_api_key_dict.key_alias, - model=data.get("model"), - tags=all_tags, - ) - - verbose_proxy_logger.debug( - f"Policy engine: matching policies for context team_alias={context.team_alias}, " - f"key_alias={context.key_alias}, model={context.model}, tags={context.tags}" - ) # Get matching policies via attachments (with match reasons for attribution) attachment_registry = get_attachment_registry() @@ -1591,7 +1542,6 @@ def add_guardrails_from_policy_engine( context ) matching_policy_names = [m["policy_name"] for m in matches_with_reasons] - # Build reasons map: {"hipaa-policy": "tag:healthcare", ...} policy_reasons = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons} verbose_proxy_logger.debug( @@ -1607,7 +1557,7 @@ def add_guardrails_from_policy_engine( ) if not all_policy_names: - return + return [], {} # Filter to only policies whose conditions match the context applied_policy_names = PolicyMatcher.get_policies_with_matching_conditions( @@ -1635,6 +1585,18 @@ def add_guardrails_from_policy_engine( request_data=data, policy_sources=applied_reasons ) + return applied_policy_names, policy_reasons + + +def _apply_resolved_guardrails_to_metadata( + data: dict, + metadata_variable_name: str, + context: "PolicyMatchContext", +) -> None: + """Apply resolved guardrails and pipelines to request metadata.""" + from litellm._logging import verbose_proxy_logger + from litellm.proxy.policy_engine.policy_resolver import PolicyResolver + # Resolve guardrails from matching policies resolved_guardrails = PolicyResolver.resolve_guardrails_for_context(context=context) @@ -1683,6 +1645,75 @@ def add_guardrails_from_policy_engine( ) +def add_guardrails_from_policy_engine( + data: dict, + metadata_variable_name: str, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """ + Add guardrails from the policy engine based on request context. + + This function: + 1. Extracts "policies" from request body (if present) for dynamic policy application + 2. Gets matching policies based on team_alias, key_alias, and model (via attachments) + 3. Combines dynamic policies with attachment-based policies + 4. Resolves guardrails from all policies (including inheritance) + 5. Adds guardrails to request metadata + 6. Tracks applied policies in metadata for response headers + 7. Removes "policies" from request body so it's not forwarded to LLM provider + + Args: + data: The request data to update + metadata_variable_name: The name of the metadata field in data + user_api_key_dict: The user's API key authentication info + """ + from litellm._logging import verbose_proxy_logger + from litellm.proxy.common_utils.http_parsing_utils import ( + get_tags_from_request_body, + ) + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import PolicyMatchContext + + # Extract dynamic policies from request body (if present) + request_body_policies = data.pop("policies", None) + + registry = get_policy_registry() + verbose_proxy_logger.debug( + f"Policy engine: registry initialized={registry.is_initialized()}, " + f"policy_count={len(registry.get_all_policies())}" + ) + if not registry.is_initialized(): + verbose_proxy_logger.debug( + "Policy engine not initialized, skipping policy matching" + ) + return + + # Extract tags and build context + all_tags = get_tags_from_request_body(data) or None + context = PolicyMatchContext( + team_alias=user_api_key_dict.team_alias, + key_alias=user_api_key_dict.key_alias, + model=data.get("model"), + tags=all_tags, + ) + + verbose_proxy_logger.debug( + f"Policy engine: matching policies for context team_alias={context.team_alias}, " + f"key_alias={context.key_alias}, model={context.model}, tags={context.tags}" + ) + + # Match and track policies + applied_policy_names, _ = _match_and_track_policies( + data, context, request_body_policies + ) + + if not applied_policy_names: + return + + # Apply resolved guardrails to metadata + _apply_resolved_guardrails_to_metadata(data, metadata_variable_name, context) + + def add_provider_specific_headers_to_request( data: dict, headers: dict, From 9e0af32d468a34f96f8912c83efa5a9c5c1872e6 Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sat, 14 Feb 2026 14:38:57 -0300 Subject: [PATCH 2/4] fix: add PolicyMatchContext to TYPE_CHECKING imports Add PolicyMatchContext to TYPE_CHECKING block to fix F821 undefined name errors in helper function type hints. Fixes: - proxy/litellm_pre_call_utils.py:1524:15: F821 - proxy/litellm_pre_call_utils.py:1596:15: F821 --- litellm/proxy/litellm_pre_call_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d211501ef9..0171b7efd7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -40,10 +40,12 @@ service_logger_obj = ServiceLogging() # used for tracking latency on OTEL if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig + from litellm.types.proxy.policy_engine import PolicyMatchContext ProxyConfig = _ProxyConfig else: ProxyConfig = Any + PolicyMatchContext = Any def parse_cache_control(cache_control): From 0f0b7a3f5046cbc8b6d9043fd1ec5fc5b7533c7e Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sat, 14 Feb 2026 14:41:23 -0300 Subject: [PATCH 3/4] fix: preserve original behavior - always call PolicyResolver Remove early return when applied_policy_names is empty to preserve original behavior where PolicyResolver.resolve_guardrails_for_context() is always called. The resolver does its own independent matching and inheritance resolution, which can apply guardrails from parent policies even when child policies fail their conditions. This ensures no guardrails are silently dropped due to the refactoring. Addresses Greptile review feedback on behavioral change. --- litellm/proxy/litellm_pre_call_utils.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0171b7efd7..bafd473d20 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1704,15 +1704,12 @@ def add_guardrails_from_policy_engine( f"key_alias={context.key_alias}, model={context.model}, tags={context.tags}" ) - # Match and track policies - applied_policy_names, _ = _match_and_track_policies( - data, context, request_body_policies - ) - - if not applied_policy_names: - return + # Match and track policies (may be empty, but resolution still happens) + _match_and_track_policies(data, context, request_body_policies) # Apply resolved guardrails to metadata + # Note: PolicyResolver does its own independent matching and inheritance resolution, + # so we always call it even if applied_policy_names is empty _apply_resolved_guardrails_to_metadata(data, metadata_variable_name, context) From 3b7cc1d1382a88e53529d61aaae9d48545f8cc04 Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Sat, 14 Feb 2026 14:42:54 -0300 Subject: [PATCH 4/4] docs: improve comment clarity per Greptile feedback Restructure comments to make the intent clearer: - First function: matches/tracks policies (may find nothing) - Second function: always runs independent resolution (may find inherited guardrails) This makes it clear why both functions must always be called. --- litellm/proxy/litellm_pre_call_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index bafd473d20..4d77af513a 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1704,12 +1704,12 @@ def add_guardrails_from_policy_engine( f"key_alias={context.key_alias}, model={context.model}, tags={context.tags}" ) - # Match and track policies (may be empty, but resolution still happens) + # Match and track policies based on attachments and request body _match_and_track_policies(data, context, request_body_policies) - # Apply resolved guardrails to metadata - # Note: PolicyResolver does its own independent matching and inheritance resolution, - # so we always call it even if applied_policy_names is empty + # Always resolve and apply guardrails, even if no policies matched above. + # PolicyResolver does its own independent matching and inheritance resolution, + # so guardrails can still be applied via inherited parent policies. _apply_resolved_guardrails_to_metadata(data, metadata_variable_name, context)