mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
removed unused import
This commit is contained in:
shivam
parent
dd7ed80bc7
commit
0bb4f876c0
1 changed files with 60 additions and 111 deletions
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue