From 0bb4f876c0a652ed27ce42d6f494f989da4d857a Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 14 Feb 2026 20:02:50 -0800 Subject: [PATCH] removed unused import --- litellm/proxy/litellm_pre_call_utils.py | 171 +++++++++--------------- 1 file changed, 60 insertions(+), 111 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 7f49c389334..f3fc65aa4a1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1,6 +1,5 @@ import asyncio import copy -import logging import time from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union @@ -42,12 +41,10 @@ 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): @@ -1539,26 +1536,75 @@ def move_guardrails_to_metadata( ] = request_body_guardrail_config -def _match_and_track_policies( +def add_guardrails_from_policy_engine( data: dict, - context: "PolicyMatchContext", - request_body_policies: Any, -) -> tuple[list[str], dict[str, str]]: + metadata_variable_name: str, + user_api_key_dict: UserAPIKeyAuth, +) -> None: """ - Match policies via attachments and request body, track them in metadata. + Add guardrails from the policy engine based on request context. - Returns: - Tuple of (applied_policy_names, policy_reasons) + 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.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() @@ -1566,6 +1612,7 @@ def _match_and_track_policies( 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( @@ -1581,7 +1628,7 @@ def _match_and_track_policies( ) 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( @@ -1609,18 +1656,6 @@ def _match_and_track_policies( 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) @@ -1628,40 +1663,20 @@ def _apply_resolved_guardrails_to_metadata( f"Policy engine: resolved guardrails: {resolved_guardrails}" ) - # Resolve pipelines from matching policies - pipelines = PolicyResolver.resolve_pipelines_for_context(context=context) + if not resolved_guardrails: + return # Add resolved guardrails to request metadata if metadata_variable_name not in data: data[metadata_variable_name] = {} - # Track pipeline-managed guardrails to exclude from independent execution - pipeline_managed_guardrails: set = set() - if pipelines: - pipeline_managed_guardrails = PolicyResolver.get_pipeline_managed_guardrails( - pipelines - ) - data[metadata_variable_name]["_guardrail_pipelines"] = pipelines - data[metadata_variable_name]["_pipeline_managed_guardrails"] = ( - pipeline_managed_guardrails - ) - verbose_proxy_logger.debug( - f"Policy engine: resolved {len(pipelines)} pipeline(s), " - f"managed guardrails: {pipeline_managed_guardrails}" - ) - - if not resolved_guardrails and not pipelines: - return - existing_guardrails = data[metadata_variable_name].get("guardrails", []) if not isinstance(existing_guardrails, list): existing_guardrails = [] # Combine existing guardrails with policy-resolved guardrails (no duplicates) - # Exclude pipeline-managed guardrails from the flat list combined = set(existing_guardrails) combined.update(resolved_guardrails) - combined -= pipeline_managed_guardrails data[metadata_variable_name]["guardrails"] = list(combined) verbose_proxy_logger.debug( @@ -1669,72 +1684,6 @@ def _apply_resolved_guardrails_to_metadata( ) -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 based on attachments and request body - _match_and_track_policies(data, context, request_body_policies) - - # 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) - - def add_provider_specific_headers_to_request( data: dict, headers: dict,