mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat: ui improvements
This commit is contained in:
Krrish Dholakia
parent
c68ee52a4d
commit
9bd4ae3df4
8 changed files with 515 additions and 301 deletions
|
|
@ -10,15 +10,10 @@ import litellm
|
|||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.proxy._types import (
|
||||
AddTeamCallback,
|
||||
CommonProxyErrors,
|
||||
LitellmDataForBackendLLMCall,
|
||||
LitellmUserRoles,
|
||||
SpecialHeaders,
|
||||
TeamCallbackMetadata,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors,
|
||||
LitellmDataForBackendLLMCall,
|
||||
LitellmUserRoles, SpecialHeaders,
|
||||
TeamCallbackMetadata, UserAPIKeyAuth)
|
||||
|
||||
# Cache special headers as a frozenset for O(1) lookup performance
|
||||
_SPECIAL_HEADERS_CACHE = frozenset(
|
||||
|
|
@ -27,12 +22,9 @@ _SPECIAL_HEADERS_CACHE = frozenset(
|
|||
from litellm.router import Router
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
|
||||
from litellm.types.services import ServiceTypes
|
||||
from litellm.types.utils import (
|
||||
LlmProviders,
|
||||
ProviderSpecificHeader,
|
||||
StandardLoggingUserAPIKeyMetadata,
|
||||
SupportedCacheControls,
|
||||
)
|
||||
from litellm.types.utils import (LlmProviders, ProviderSpecificHeader,
|
||||
StandardLoggingUserAPIKeyMetadata,
|
||||
SupportedCacheControls)
|
||||
|
||||
service_logger_obj = ServiceLogging() # used for tracking latency on OTEL
|
||||
|
||||
|
|
@ -661,8 +653,7 @@ class LiteLLMProxyRequestSetup:
|
|||
return data
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ManagementEndpoint_MetadataFields,
|
||||
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
|
||||
)
|
||||
LiteLLM_ManagementEndpoint_MetadataFields_Premium)
|
||||
|
||||
# ignore any special fields
|
||||
added_metadata = {}
|
||||
|
|
@ -1125,7 +1116,7 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
|||
data["litellm_disabled_callbacks"] = disabled_callbacks
|
||||
|
||||
# Guardrails from key/team metadata and policy engine
|
||||
move_guardrails_to_metadata(
|
||||
await move_guardrails_to_metadata(
|
||||
data=data,
|
||||
_metadata_variable_name=_metadata_variable_name,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -1458,7 +1449,7 @@ def _add_guardrails_from_policies_in_metadata(
|
|||
)
|
||||
|
||||
|
||||
def move_guardrails_to_metadata(
|
||||
async def move_guardrails_to_metadata(
|
||||
data: dict,
|
||||
_metadata_variable_name: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -1487,7 +1478,8 @@ def move_guardrails_to_metadata(
|
|||
|
||||
# Only check policy engine if no local config (avoid import + registry lookup)
|
||||
if not (has_key_config or has_team_config or has_request_config):
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import \
|
||||
get_policy_registry
|
||||
|
||||
if not get_policy_registry().is_initialized():
|
||||
# Nothing configured anywhere - clean up request body fields and return
|
||||
|
|
@ -1515,7 +1507,7 @@ def move_guardrails_to_metadata(
|
|||
#########################################################################################
|
||||
# Add guardrails from policy engine based on team/key/model context
|
||||
#########################################################################################
|
||||
add_guardrails_from_policy_engine(
|
||||
await add_guardrails_from_policy_engine(
|
||||
data=data,
|
||||
metadata_variable_name=_metadata_variable_name,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -1549,10 +1541,29 @@ def move_guardrails_to_metadata(
|
|||
] = request_body_guardrail_config
|
||||
|
||||
|
||||
def _is_policy_version_id(s: str) -> bool:
|
||||
"""Return True if string is a policy version ID (starts with policy_<uuid> prefix)."""
|
||||
from litellm.proxy.policy_engine.policy_registry import \
|
||||
POLICY_VERSION_ID_PREFIX
|
||||
|
||||
return isinstance(s, str) and s.startswith(POLICY_VERSION_ID_PREFIX)
|
||||
|
||||
|
||||
def _extract_policy_id(s: str) -> Optional[str]:
|
||||
"""Extract raw UUID from policy_<uuid> string, or None if not a valid version ID."""
|
||||
from litellm.proxy.policy_engine.policy_registry import \
|
||||
POLICY_VERSION_ID_PREFIX
|
||||
|
||||
if not _is_policy_version_id(s):
|
||||
return None
|
||||
return s[len(POLICY_VERSION_ID_PREFIX) :].strip() or None
|
||||
|
||||
|
||||
def _match_and_track_policies(
|
||||
data: dict,
|
||||
context: "PolicyMatchContext",
|
||||
request_body_policies: Any,
|
||||
policies_override: Optional[Dict[str, Any]] = None,
|
||||
) -> tuple[list[str], dict[str, str]]:
|
||||
"""
|
||||
Match policies via attachments and request body, track them in metadata.
|
||||
|
|
@ -1562,10 +1573,9 @@ def _match_and_track_policies(
|
|||
"""
|
||||
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.policy_engine.attachment_registry import get_attachment_registry
|
||||
add_policy_sources_to_metadata, add_policy_to_applied_policies_header)
|
||||
from litellm.proxy.policy_engine.attachment_registry import \
|
||||
get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
|
||||
|
||||
# Get matching policies via attachments (with match reasons for attribution)
|
||||
|
|
@ -1595,6 +1605,7 @@ def _match_and_track_policies(
|
|||
applied_policy_names = PolicyMatcher.get_policies_with_matching_conditions(
|
||||
policy_names=list(all_policy_names),
|
||||
context=context,
|
||||
policies=policies_override,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -1622,20 +1633,30 @@ def _apply_resolved_guardrails_to_metadata(
|
|||
data: dict,
|
||||
metadata_variable_name: str,
|
||||
context: "PolicyMatchContext",
|
||||
policy_names: Optional[List[str]] = None,
|
||||
policies: Optional[Dict[str, Any]] = None,
|
||||
) -> 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)
|
||||
resolved_guardrails = PolicyResolver.resolve_guardrails_for_context(
|
||||
context=context,
|
||||
policies=policies,
|
||||
policy_names=policy_names,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"Policy engine: resolved guardrails: {resolved_guardrails}"
|
||||
)
|
||||
|
||||
# Resolve pipelines from matching policies
|
||||
pipelines = PolicyResolver.resolve_pipelines_for_context(context=context)
|
||||
pipelines = PolicyResolver.resolve_pipelines_for_context(
|
||||
context=context,
|
||||
policies=policies,
|
||||
policy_names=policy_names,
|
||||
)
|
||||
|
||||
# Add resolved guardrails to request metadata
|
||||
if metadata_variable_name not in data:
|
||||
|
|
@ -1675,7 +1696,7 @@ def _apply_resolved_guardrails_to_metadata(
|
|||
)
|
||||
|
||||
|
||||
def add_guardrails_from_policy_engine(
|
||||
async def add_guardrails_from_policy_engine(
|
||||
data: dict,
|
||||
metadata_variable_name: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -1685,12 +1706,13 @@ def add_guardrails_from_policy_engine(
|
|||
|
||||
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
|
||||
2. Supports policy_<uuid> in policies to execute a specific version (e.g. published)
|
||||
3. Gets matching policies based on team_alias, key_alias, and model (via attachments)
|
||||
4. Combines dynamic policies with attachment-based policies
|
||||
5. Resolves guardrails from all policies (including inheritance)
|
||||
6. Adds guardrails to request metadata
|
||||
7. Tracks applied policies in metadata for response headers
|
||||
8. Removes "policies" from request body so it's not forwarded to LLM provider
|
||||
|
||||
Args:
|
||||
data: The request data to update
|
||||
|
|
@ -1698,12 +1720,13 @@ def add_guardrails_from_policy_engine(
|
|||
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.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)
|
||||
request_body_policies_raw = data.pop("policies", None)
|
||||
|
||||
registry = get_policy_registry()
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -1730,13 +1753,69 @@ def add_guardrails_from_policy_engine(
|
|||
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)
|
||||
# Separate policy names from policy version IDs (policy_<uuid>)
|
||||
request_body_names: List[str] = []
|
||||
request_body_version_ids: List[str] = []
|
||||
if request_body_policies_raw and isinstance(request_body_policies_raw, list):
|
||||
for item in request_body_policies_raw:
|
||||
if not isinstance(item, str):
|
||||
continue
|
||||
if _is_policy_version_id(item):
|
||||
policy_id = _extract_policy_id(item)
|
||||
if policy_id:
|
||||
request_body_version_ids.append(policy_id)
|
||||
else:
|
||||
request_body_names.append(item)
|
||||
|
||||
# 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)
|
||||
# Fetch policy versions by ID from DB
|
||||
merged_policies: Dict[str, Any] = dict(registry.get_all_policies())
|
||||
fetched_policy_names: List[str] = []
|
||||
if request_body_version_ids:
|
||||
try:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is not None:
|
||||
for policy_id in request_body_version_ids:
|
||||
result = await registry.get_policy_by_id_for_request(
|
||||
policy_id=policy_id,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
if result is not None:
|
||||
pname, policy = result
|
||||
merged_policies[pname] = policy
|
||||
fetched_policy_names.append(pname)
|
||||
verbose_proxy_logger.debug(
|
||||
f"Policy engine: loaded version by ID policy_{policy_id} -> {pname}"
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Policy engine: policy version {policy_id} not found, skipping"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Policy engine: failed to fetch policy versions by ID: {e}"
|
||||
)
|
||||
|
||||
# Build request body list: names + policy names from fetched versions
|
||||
request_body_policies = request_body_names + fetched_policy_names
|
||||
|
||||
# Match and track policies (with merged_policies when we have version overrides)
|
||||
applied_policy_names, _ = _match_and_track_policies(
|
||||
data,
|
||||
context,
|
||||
request_body_policies,
|
||||
policies_override=merged_policies if request_body_version_ids else None,
|
||||
)
|
||||
|
||||
# Resolve and apply guardrails. Use applied_policy_names so request-body policies
|
||||
# (names + version IDs) are included. Use merged_policies when we have version overrides.
|
||||
_apply_resolved_guardrails_to_metadata(
|
||||
data,
|
||||
metadata_variable_name,
|
||||
context,
|
||||
policy_names=applied_policy_names if applied_policy_names else None,
|
||||
policies=merged_policies if request_body_version_ids else None,
|
||||
)
|
||||
|
||||
|
||||
def add_provider_specific_headers_to_request(
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ by policy_attachments (see AttachmentRegistry).
|
|||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.types.proxy.policy_engine import (GuardrailPipeline, PipelineStep,
|
||||
|
|
@ -24,6 +24,9 @@ from litellm.types.proxy.policy_engine import (GuardrailPipeline, PipelineStep,
|
|||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
# Prefix for policy version IDs in request body. Use policy_<uuid> to execute a specific version.
|
||||
POLICY_VERSION_ID_PREFIX = "policy_"
|
||||
|
||||
|
||||
def _row_to_policy_db_response(row: Any) -> PolicyDBResponse:
|
||||
"""Build PolicyDBResponse from a Prisma LiteLLM_PolicyTable row."""
|
||||
|
|
@ -463,6 +466,44 @@ class PolicyRegistry:
|
|||
verbose_proxy_logger.exception(f"Error getting policy from DB: {e}")
|
||||
raise Exception(f"Error getting policy from DB: {str(e)}")
|
||||
|
||||
async def get_policy_by_id_for_request(
|
||||
self,
|
||||
policy_id: str,
|
||||
prisma_client: "PrismaClient",
|
||||
) -> Optional[Tuple[str, Policy]]:
|
||||
"""
|
||||
Fetch a policy version by ID from the DB and convert to Policy for resolution.
|
||||
|
||||
Used when the request body specifies policy_<uuid> to execute a specific version
|
||||
(e.g. published or draft) instead of production.
|
||||
|
||||
Args:
|
||||
policy_id: The policy version ID (raw UUID, no prefix)
|
||||
prisma_client: The Prisma client instance
|
||||
|
||||
Returns:
|
||||
(policy_name, Policy) if found, None otherwise
|
||||
"""
|
||||
response = await self.get_policy_by_id_from_db(
|
||||
policy_id=policy_id, prisma_client=prisma_client
|
||||
)
|
||||
if response is None:
|
||||
return None
|
||||
policy = self._parse_policy(
|
||||
response.policy_name,
|
||||
{
|
||||
"inherit": response.inherit,
|
||||
"description": response.description,
|
||||
"guardrails": {
|
||||
"add": response.guardrails_add,
|
||||
"remove": response.guardrails_remove,
|
||||
},
|
||||
"condition": response.condition,
|
||||
"pipeline": response.pipeline,
|
||||
},
|
||||
)
|
||||
return (response.policy_name, policy)
|
||||
|
||||
async def get_all_policies_from_db(
|
||||
self,
|
||||
prisma_client: "PrismaClient",
|
||||
|
|
|
|||
|
|
@ -11,12 +11,9 @@ Handles:
|
|||
from typing import Dict, List, Optional, Set, Tuple
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.types.proxy.policy_engine import (
|
||||
GuardrailPipeline,
|
||||
Policy,
|
||||
PolicyMatchContext,
|
||||
ResolvedPolicy,
|
||||
)
|
||||
from litellm.types.proxy.policy_engine import (GuardrailPipeline, Policy,
|
||||
PolicyMatchContext,
|
||||
ResolvedPolicy)
|
||||
|
||||
|
||||
class PolicyResolver:
|
||||
|
|
@ -90,7 +87,8 @@ class PolicyResolver:
|
|||
Returns:
|
||||
ResolvedPolicy with final guardrails list
|
||||
"""
|
||||
from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator
|
||||
from litellm.proxy.policy_engine.condition_evaluator import \
|
||||
ConditionEvaluator
|
||||
|
||||
inheritance_chain = PolicyResolver.resolve_inheritance_chain(
|
||||
policy_name=policy_name, policies=policies
|
||||
|
|
@ -134,12 +132,13 @@ class PolicyResolver:
|
|||
def resolve_guardrails_for_context(
|
||||
context: PolicyMatchContext,
|
||||
policies: Optional[Dict[str, Policy]] = None,
|
||||
policy_names: Optional[List[str]] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Resolve the final list of guardrails for a request context.
|
||||
|
||||
This:
|
||||
1. Finds all policies that match the context via policy_attachments
|
||||
1. Finds all policies that match the context via policy_attachments (or policy_names if provided)
|
||||
2. Resolves each policy's guardrails (including inheritance)
|
||||
3. Evaluates model conditions
|
||||
4. Combines all guardrails (union)
|
||||
|
|
@ -147,12 +146,14 @@ class PolicyResolver:
|
|||
Args:
|
||||
context: The request context
|
||||
policies: Dictionary of all policies (if None, uses global registry)
|
||||
policy_names: If provided, use this list instead of attachment matching
|
||||
|
||||
Returns:
|
||||
List of guardrail names to apply
|
||||
"""
|
||||
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_registry import \
|
||||
get_policy_registry
|
||||
|
||||
if policies is None:
|
||||
registry = get_policy_registry()
|
||||
|
|
@ -160,8 +161,12 @@ class PolicyResolver:
|
|||
return []
|
||||
policies = registry.get_all_policies()
|
||||
|
||||
# Get matching policies via attachments
|
||||
matching_policy_names = PolicyMatcher.get_matching_policies(context=context)
|
||||
# Use provided policy names or get matching policies via attachments
|
||||
matching_policy_names = (
|
||||
policy_names
|
||||
if policy_names is not None
|
||||
else PolicyMatcher.get_matching_policies(context=context)
|
||||
)
|
||||
|
||||
if not matching_policy_names:
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -195,6 +200,7 @@ class PolicyResolver:
|
|||
def resolve_pipelines_for_context(
|
||||
context: PolicyMatchContext,
|
||||
policies: Optional[Dict[str, Policy]] = None,
|
||||
policy_names: Optional[List[str]] = None,
|
||||
) -> List[Tuple[str, GuardrailPipeline]]:
|
||||
"""
|
||||
Resolve pipelines from matching policies for a request context.
|
||||
|
|
@ -206,12 +212,14 @@ class PolicyResolver:
|
|||
Args:
|
||||
context: The request context
|
||||
policies: Dictionary of all policies (if None, uses global registry)
|
||||
policy_names: If provided, use this list instead of attachment matching
|
||||
|
||||
Returns:
|
||||
List of (policy_name, GuardrailPipeline) tuples
|
||||
"""
|
||||
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_registry import \
|
||||
get_policy_registry
|
||||
|
||||
if policies is None:
|
||||
registry = get_policy_registry()
|
||||
|
|
@ -219,7 +227,11 @@ class PolicyResolver:
|
|||
return []
|
||||
policies = registry.get_all_policies()
|
||||
|
||||
matching_policy_names = PolicyMatcher.get_matching_policies(context=context)
|
||||
matching_policy_names = (
|
||||
policy_names
|
||||
if policy_names is not None
|
||||
else PolicyMatcher.get_matching_policies(context=context)
|
||||
)
|
||||
if not matching_policy_names:
|
||||
return []
|
||||
|
||||
|
|
@ -269,7 +281,8 @@ class PolicyResolver:
|
|||
Returns:
|
||||
Dictionary mapping policy names to ResolvedPolicy objects
|
||||
"""
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import \
|
||||
get_policy_registry
|
||||
|
||||
if policies is None:
|
||||
registry = get_policy_registry()
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import copy
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import Request
|
||||
|
|
@ -11,16 +11,11 @@ from fastapi import Request
|
|||
import litellm
|
||||
from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
KeyAndTeamLoggingSettings,
|
||||
LiteLLMProxyRequestSetup,
|
||||
_get_dynamic_logging_metadata,
|
||||
_get_enforced_params,
|
||||
_get_metadata_variable_name,
|
||||
_update_model_if_key_alias_exists,
|
||||
add_guardrails_from_policy_engine,
|
||||
add_litellm_data_to_request,
|
||||
check_if_token_is_service_account,
|
||||
)
|
||||
KeyAndTeamLoggingSettings, LiteLLMProxyRequestSetup,
|
||||
_get_dynamic_logging_metadata, _get_enforced_params,
|
||||
_get_metadata_variable_name, _update_model_if_key_alias_exists,
|
||||
add_guardrails_from_policy_engine, add_litellm_data_to_request,
|
||||
check_if_token_is_service_account)
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
|
|
@ -159,7 +154,8 @@ def test_get_enforced_params(
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_parses_string_metadata():
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import \
|
||||
add_litellm_data_to_request
|
||||
|
||||
# Setup
|
||||
request_mock = MagicMock(spec=Request)
|
||||
|
|
@ -205,7 +201,8 @@ async def test_add_litellm_data_to_request_parses_string_metadata():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_user_spend_and_budget():
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import \
|
||||
add_litellm_data_to_request
|
||||
|
||||
request_mock = MagicMock(spec=Request)
|
||||
request_mock.url.path = "/v1/completions"
|
||||
|
|
@ -243,7 +240,8 @@ async def test_add_litellm_data_to_request_user_spend_and_budget():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_audio_transcription_multipart():
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import \
|
||||
add_litellm_data_to_request
|
||||
|
||||
# Setup request mock for /v1/audio/transcriptions
|
||||
request_mock = MagicMock(spec=Request)
|
||||
|
|
@ -308,7 +306,8 @@ async def test_add_litellm_data_to_request_disabled_callbacks():
|
|||
"""
|
||||
Test that litellm_disabled_callbacks from key metadata is properly added to the request data.
|
||||
"""
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import \
|
||||
add_litellm_data_to_request
|
||||
|
||||
# Setup mock request
|
||||
request_mock = MagicMock(spec=Request)
|
||||
|
|
@ -361,7 +360,8 @@ async def test_add_litellm_data_to_request_disabled_callbacks_empty():
|
|||
"""
|
||||
Test that litellm_disabled_callbacks is not added when it's empty.
|
||||
"""
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import \
|
||||
add_litellm_data_to_request
|
||||
|
||||
# Setup mock request
|
||||
request_mock = MagicMock(spec=Request)
|
||||
|
|
@ -413,7 +413,8 @@ async def test_add_litellm_data_to_request_disabled_callbacks_not_present():
|
|||
"""
|
||||
Test that litellm_disabled_callbacks is not added when it's not present in metadata.
|
||||
"""
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import \
|
||||
add_litellm_data_to_request
|
||||
|
||||
# Setup mock request
|
||||
request_mock = MagicMock(spec=Request)
|
||||
|
|
@ -465,7 +466,8 @@ async def test_add_litellm_data_to_request_disabled_callbacks_invalid_type():
|
|||
"""
|
||||
Test that litellm_disabled_callbacks is not added when it's not a list.
|
||||
"""
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import \
|
||||
add_litellm_data_to_request
|
||||
|
||||
# Setup mock request
|
||||
request_mock = MagicMock(spec=Request)
|
||||
|
|
@ -517,7 +519,8 @@ async def test_add_litellm_data_to_request_disabled_callbacks_with_logging_setti
|
|||
"""
|
||||
Test that litellm_disabled_callbacks works correctly alongside logging settings.
|
||||
"""
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import \
|
||||
add_litellm_data_to_request
|
||||
|
||||
# Setup mock request
|
||||
request_mock = MagicMock(spec=Request)
|
||||
|
|
@ -1027,7 +1030,8 @@ from unittest.mock import AsyncMock
|
|||
from fastapi.responses import Response
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_request_processing import \
|
||||
ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
|
@ -1403,7 +1407,8 @@ async def test_embedding_header_forwarding_with_model_group():
|
|||
importlib.reload(pre_call_utils_module)
|
||||
|
||||
# Re-import the function after reload to get the fresh version
|
||||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||||
from litellm.proxy.litellm_pre_call_utils import \
|
||||
add_litellm_data_to_request
|
||||
|
||||
# Setup mock request for embeddings
|
||||
request_mock = MagicMock(spec=Request)
|
||||
|
|
@ -1531,18 +1536,17 @@ async def test_embedding_header_forwarding_without_model_group_config():
|
|||
litellm.model_group_settings = original_model_group_settings
|
||||
|
||||
|
||||
def test_add_guardrails_from_policy_engine():
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_guardrails_from_policy_engine():
|
||||
"""
|
||||
Test that add_guardrails_from_policy_engine adds guardrails from matching policies
|
||||
and tracks applied policies in metadata.
|
||||
"""
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.attachment_registry import \
|
||||
get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.types.proxy.policy_engine import (
|
||||
Policy,
|
||||
PolicyAttachment,
|
||||
PolicyGuardrails,
|
||||
)
|
||||
from litellm.types.proxy.policy_engine import (Policy, PolicyAttachment,
|
||||
PolicyGuardrails)
|
||||
|
||||
# Setup test data
|
||||
data = {
|
||||
|
|
@ -1578,7 +1582,7 @@ def test_add_guardrails_from_policy_engine():
|
|||
attachment_registry._initialized = True
|
||||
|
||||
# Call the function
|
||||
add_guardrails_from_policy_engine(
|
||||
await add_guardrails_from_policy_engine(
|
||||
data=data,
|
||||
metadata_variable_name="metadata",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -1601,11 +1605,12 @@ def test_add_guardrails_from_policy_engine():
|
|||
attachment_registry._initialized = False
|
||||
|
||||
|
||||
def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_from_data():
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_from_data():
|
||||
"""
|
||||
Test that add_guardrails_from_policy_engine accepts dynamic 'policies' from the request body
|
||||
and removes them to prevent forwarding to the LLM provider.
|
||||
|
||||
|
||||
This is critical because 'policies' is a LiteLLM proxy-specific parameter that should
|
||||
not be sent to the actual LLM API (e.g., OpenAI, Anthropic, etc.).
|
||||
"""
|
||||
|
|
@ -1631,7 +1636,7 @@ def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_fro
|
|||
policy_registry._initialized = False
|
||||
|
||||
# Call the function - should accept dynamic policies and not raise an error
|
||||
add_guardrails_from_policy_engine(
|
||||
await add_guardrails_from_policy_engine(
|
||||
data=data,
|
||||
metadata_variable_name="metadata",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -1646,3 +1651,69 @@ def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_fro
|
|||
assert "messages" in data
|
||||
assert data["messages"] == [{"role": "user", "content": "Hello"}]
|
||||
assert "metadata" in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_guardrails_from_policy_engine_policy_version_by_id():
|
||||
"""
|
||||
Test that add_guardrails_from_policy_engine executes a specific policy version
|
||||
when policy_<uuid> is passed in the request body.
|
||||
"""
|
||||
from litellm.proxy.policy_engine.attachment_registry import \
|
||||
get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails
|
||||
|
||||
policy_version_uuid = "12345678-1234-5678-1234-567812345678"
|
||||
policy_version_ref = f"policy_{policy_version_uuid}"
|
||||
|
||||
# Policy from the specific version (e.g. published) - different guardrail than production
|
||||
published_version_policy = Policy(
|
||||
guardrails=PolicyGuardrails(add=["published_version_guardrail"]),
|
||||
)
|
||||
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"policies": [policy_version_ref],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
team_alias="test-team",
|
||||
key_alias="test-key",
|
||||
)
|
||||
|
||||
policy_registry = get_policy_registry()
|
||||
policy_registry._policies = {}
|
||||
policy_registry._initialized = True
|
||||
|
||||
attachment_registry = get_attachment_registry()
|
||||
attachment_registry._attachments = []
|
||||
attachment_registry._initialized = True
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
|
||||
with patch.object(
|
||||
policy_registry,
|
||||
"get_policy_by_id_for_request",
|
||||
new_callable=AsyncMock,
|
||||
return_value=("test-policy-from-version", published_version_policy),
|
||||
):
|
||||
await add_guardrails_from_policy_engine(
|
||||
data=data,
|
||||
metadata_variable_name="metadata",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Verify guardrails from the specific version were applied
|
||||
assert "metadata" in data
|
||||
assert "guardrails" in data["metadata"]
|
||||
assert "published_version_guardrail" in data["metadata"]["guardrails"]
|
||||
assert "policies" not in data
|
||||
|
||||
# Clean up
|
||||
policy_registry._policies = {}
|
||||
policy_registry._initialized = False
|
||||
|
|
|
|||
|
|
@ -8,9 +8,10 @@ import {
|
|||
} from "@/data/compliancePrompts";
|
||||
import {
|
||||
getGuardrailsList,
|
||||
getPoliciesList,
|
||||
testPoliciesAndGuardrails,
|
||||
} from "@/components/networking";
|
||||
import PolicySelector, { getPolicyOptionEntries } from "@/components/policies/PolicySelector";
|
||||
import { Policy } from "@/components/policies/types";
|
||||
import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion";
|
||||
import {
|
||||
AlertTriangle,
|
||||
|
|
@ -98,11 +99,6 @@ interface QuickTestMessage {
|
|||
type ResultFilter = "all" | "matches" | "mismatches" | "pending";
|
||||
type RightPanelTab = "quick-test" | "batch-results";
|
||||
|
||||
interface PolicyOption {
|
||||
id: string;
|
||||
name: string;
|
||||
}
|
||||
|
||||
interface GuardrailOption {
|
||||
id: string;
|
||||
name: string;
|
||||
|
|
@ -132,11 +128,10 @@ export default function ComplianceUI({
|
|||
}: ComplianceUIProps) {
|
||||
const frameworks = getFrameworks();
|
||||
|
||||
const [policyOptions, setPolicyOptions] = useState<PolicyOption[]>([]);
|
||||
const [policyValueToLabel, setPolicyValueToLabel] = useState<Map<string, string>>(new Map());
|
||||
const [guardrailOptions, setGuardrailOptions] = useState<GuardrailOption[]>([]);
|
||||
const [selectedPolicies, setSelectedPolicies] = useState<string[]>([]);
|
||||
const [selectedGuardrails, setSelectedGuardrails] = useState<string[]>([]);
|
||||
const [showPolicyDropdown, setShowPolicyDropdown] = useState(false);
|
||||
const [showGuardrailDropdown, setShowGuardrailDropdown] = useState(false);
|
||||
|
||||
const [selectedPromptIds, setSelectedPromptIds] = useState<Set<string>>(new Set());
|
||||
|
|
@ -164,20 +159,16 @@ export default function ComplianceUI({
|
|||
const [expandedResults, setExpandedResults] = useState<Set<string>>(new Set());
|
||||
const batchAbortControllerRef = useRef<AbortController | null>(null);
|
||||
|
||||
const handlePoliciesLoaded = useCallback((policies: Policy[]) => {
|
||||
const entries = getPolicyOptionEntries(policies);
|
||||
setPolicyValueToLabel(new Map(entries.map((e) => [e.value, e.label])));
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
if (!accessToken) return;
|
||||
const fetchOptions = async () => {
|
||||
const fetchGuardrails = async () => {
|
||||
try {
|
||||
const [policiesRes, guardrailsRes] = await Promise.all([
|
||||
getPoliciesList(accessToken).catch(() => ({ policies: [] })),
|
||||
getGuardrailsList(accessToken).catch(() => ({ guardrails: [] })),
|
||||
]);
|
||||
setPolicyOptions(
|
||||
(policiesRes.policies || []).map((p: { policy_name: string; policy_id?: string }) => ({
|
||||
id: p.policy_id ?? p.policy_name,
|
||||
name: p.policy_name,
|
||||
}))
|
||||
);
|
||||
const guardrailsRes = await getGuardrailsList(accessToken).catch(() => ({ guardrails: [] }));
|
||||
setGuardrailOptions(
|
||||
(guardrailsRes.guardrails || []).map((g: { guardrail_name: string }) => ({
|
||||
id: g.guardrail_name,
|
||||
|
|
@ -186,11 +177,10 @@ export default function ComplianceUI({
|
|||
}))
|
||||
);
|
||||
} catch {
|
||||
setPolicyOptions([]);
|
||||
setGuardrailOptions([]);
|
||||
}
|
||||
};
|
||||
fetchOptions();
|
||||
fetchGuardrails();
|
||||
}, [accessToken]);
|
||||
|
||||
useEffect(() => {
|
||||
|
|
@ -282,12 +272,6 @@ export default function ComplianceUI({
|
|||
|
||||
const deselectAll = () => setSelectedPromptIds(new Set());
|
||||
|
||||
const togglePolicy = (id: string) => {
|
||||
setSelectedPolicies((prev) =>
|
||||
prev.includes(id) ? prev.filter((p) => p !== id) : [...prev, id]
|
||||
);
|
||||
};
|
||||
|
||||
const toggleGuardrail = (id: string) => {
|
||||
setSelectedGuardrails((prev) =>
|
||||
prev.includes(id) ? prev.filter((g) => g !== id) : [...prev, id]
|
||||
|
|
@ -766,76 +750,13 @@ export default function ComplianceUI({
|
|||
<label className="text-[11px] font-medium text-gray-500 uppercase tracking-wide mb-1.5 block">
|
||||
Policies
|
||||
</label>
|
||||
<div className="relative">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setShowPolicyDropdown(!showPolicyDropdown);
|
||||
setShowGuardrailDropdown(false);
|
||||
}}
|
||||
className="w-full flex items-center justify-between border border-gray-200 rounded-lg px-3 py-2 text-sm text-left hover:border-gray-300 transition-colors"
|
||||
>
|
||||
<span
|
||||
className={
|
||||
selectedPolicies.length > 0 ? "text-gray-700" : "text-gray-400"
|
||||
}
|
||||
>
|
||||
{selectedPolicies.length > 0
|
||||
? `${selectedPolicies.length} selected`
|
||||
: "None selected"}
|
||||
</span>
|
||||
<ChevronDown className="w-4 h-4 text-gray-400" />
|
||||
</button>
|
||||
{showPolicyDropdown && (
|
||||
<div className="absolute z-30 top-full left-0 right-0 mt-1 bg-white border border-gray-200 rounded-lg shadow-lg py-1 max-h-52 overflow-y-auto">
|
||||
{policyOptions.length === 0 ? (
|
||||
<div className="px-3 py-2 text-xs text-gray-500">
|
||||
No policies available. Create policies in the Policies page.
|
||||
</div>
|
||||
) : (
|
||||
policyOptions.map((policy) => (
|
||||
<button
|
||||
key={policy.id}
|
||||
type="button"
|
||||
onClick={() => togglePolicy(policy.id)}
|
||||
className="w-full flex items-center gap-2.5 px-3 py-2 text-sm text-left hover:bg-gray-50"
|
||||
>
|
||||
<div
|
||||
className={`w-4 h-4 rounded border flex items-center justify-center flex-shrink-0 ${selectedPolicies.includes(policy.id) ? "bg-blue-500 border-blue-500" : "border-gray-300"}`}
|
||||
>
|
||||
{selectedPolicies.includes(policy.id) && (
|
||||
<Check className="w-3 h-3 text-white" />
|
||||
)}
|
||||
</div>
|
||||
<span className="text-gray-700">{policy.name}</span>
|
||||
</button>
|
||||
))
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
{selectedPolicies.length > 0 && (
|
||||
<div className="flex flex-wrap gap-1 mt-1.5">
|
||||
{selectedPolicies.map((id) => {
|
||||
const p = policyOptions.find((x) => x.id === id);
|
||||
return (
|
||||
<span
|
||||
key={id}
|
||||
className="inline-flex items-center gap-1 text-[11px] bg-blue-50 text-blue-700 px-1.5 py-0.5 rounded font-medium"
|
||||
>
|
||||
{p?.name}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => togglePolicy(id)}
|
||||
className="hover:text-blue-900"
|
||||
aria-label="Remove"
|
||||
>
|
||||
<X className="w-2.5 h-2.5" />
|
||||
</button>
|
||||
</span>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
{accessToken && (
|
||||
<PolicySelector
|
||||
value={selectedPolicies}
|
||||
onChange={setSelectedPolicies}
|
||||
accessToken={accessToken}
|
||||
onPoliciesLoaded={handlePoliciesLoaded}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
|
||||
|
|
@ -852,10 +773,7 @@ export default function ComplianceUI({
|
|||
<div className="relative">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setShowGuardrailDropdown(!showGuardrailDropdown);
|
||||
setShowPolicyDropdown(false);
|
||||
}}
|
||||
onClick={() => setShowGuardrailDropdown(!showGuardrailDropdown)}
|
||||
className="w-full flex items-center justify-between border border-gray-200 rounded-lg px-3 py-2 text-sm text-left hover:border-gray-300 transition-colors"
|
||||
>
|
||||
<span
|
||||
|
|
@ -1342,17 +1260,14 @@ export default function ComplianceUI({
|
|||
<span className="text-[11px] font-medium text-gray-500">
|
||||
Testing against:
|
||||
</span>
|
||||
{selectedPolicies.map((id) => {
|
||||
const p = policyOptions.find((x) => x.id === id);
|
||||
return (
|
||||
<span
|
||||
key={id}
|
||||
className="text-[11px] bg-blue-50 text-blue-700 px-2 py-0.5 rounded font-medium"
|
||||
>
|
||||
{p?.name}
|
||||
</span>
|
||||
);
|
||||
})}
|
||||
{selectedPolicies.map((id) => (
|
||||
<span
|
||||
key={id}
|
||||
className="text-[11px] bg-blue-50 text-blue-700 px-2 py-0.5 rounded font-medium"
|
||||
>
|
||||
{policyValueToLabel.get(id) ?? id}
|
||||
</span>
|
||||
))}
|
||||
{selectedGuardrails.map((id) => {
|
||||
const g = guardrailOptions.find((x) => x.id === id);
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -3,20 +3,53 @@ import { Select } from "antd";
|
|||
import { Policy } from "./types";
|
||||
import { getPoliciesList } from "../networking";
|
||||
|
||||
/** Prefix for policy version IDs in request body; must match backend POLICY_VERSION_ID_PREFIX. */
|
||||
export const POLICY_VERSION_ID_PREFIX = "policy_";
|
||||
|
||||
/** Build the value sent in the request body: policy_<uuid> so backend executes this exact version. */
|
||||
export function policyVersionRef(policyId: string): string {
|
||||
return `${POLICY_VERSION_ID_PREFIX}${policyId}`;
|
||||
}
|
||||
|
||||
/** Build select options from policies (filter non-draft, label with name/version/status). */
|
||||
export function getPolicyOptionEntries(policies: Policy[]): { value: string; label: string }[] {
|
||||
return policies
|
||||
.filter((policy) => (policy.version_status ?? "draft") !== "draft")
|
||||
.map((policy) => {
|
||||
const versionNum = policy.version_number ?? 1;
|
||||
const status = policy.version_status ?? "draft";
|
||||
const label = `${policy.policy_name} — v${versionNum} (${status})${
|
||||
policy.description ? ` — ${policy.description}` : ""
|
||||
}`;
|
||||
const isProduction = status === "production";
|
||||
return {
|
||||
label,
|
||||
value: isProduction
|
||||
? policy.policy_name
|
||||
: policy.policy_id
|
||||
? policyVersionRef(policy.policy_id)
|
||||
: policy.policy_name,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
interface PolicySelectorProps {
|
||||
onChange: (selectedPolicies: string[]) => void;
|
||||
value?: string[];
|
||||
className?: string;
|
||||
accessToken: string;
|
||||
disabled?: boolean;
|
||||
/** Called after policies are loaded; use to build value→label map for display elsewhere. */
|
||||
onPoliciesLoaded?: (policies: Policy[]) => void;
|
||||
}
|
||||
|
||||
const PolicySelector: React.FC<PolicySelectorProps> = ({
|
||||
onChange,
|
||||
value,
|
||||
className,
|
||||
accessToken,
|
||||
disabled
|
||||
const PolicySelector: React.FC<PolicySelectorProps> = ({
|
||||
onChange,
|
||||
value,
|
||||
className,
|
||||
accessToken,
|
||||
disabled,
|
||||
onPoliciesLoaded,
|
||||
}) => {
|
||||
const [policies, setPolicies] = useState<Policy[]>([]);
|
||||
const [loading, setLoading] = useState(false);
|
||||
|
|
@ -28,10 +61,9 @@ const PolicySelector: React.FC<PolicySelectorProps> = ({
|
|||
setLoading(true);
|
||||
try {
|
||||
const response = await getPoliciesList(accessToken);
|
||||
console.log("Policies response:", response);
|
||||
if (response.policies) {
|
||||
console.log("Policies data:", response.policies);
|
||||
setPolicies(response.policies);
|
||||
onPoliciesLoaded?.(response.policies);
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Error fetching policies:", error);
|
||||
|
|
@ -41,10 +73,9 @@ const PolicySelector: React.FC<PolicySelectorProps> = ({
|
|||
};
|
||||
|
||||
fetchPolicies();
|
||||
}, [accessToken]);
|
||||
}, [accessToken, onPoliciesLoaded]);
|
||||
|
||||
const handlePolicyChange = (selectedValues: string[]) => {
|
||||
console.log("Selected policies:", selectedValues);
|
||||
onChange(selectedValues);
|
||||
};
|
||||
|
||||
|
|
@ -53,19 +84,17 @@ const PolicySelector: React.FC<PolicySelectorProps> = ({
|
|||
<Select
|
||||
mode="multiple"
|
||||
disabled={disabled}
|
||||
placeholder={disabled ? "Setting policies is a premium feature." : "Select policies"}
|
||||
placeholder={
|
||||
disabled
|
||||
? "Setting policies is a premium feature."
|
||||
: "Select policies (production or published versions)"
|
||||
}
|
||||
onChange={handlePolicyChange}
|
||||
value={value}
|
||||
loading={loading}
|
||||
className={className}
|
||||
allowClear
|
||||
options={policies.map((policy) => {
|
||||
console.log("Mapping policy:", policy);
|
||||
return {
|
||||
label: `${policy.policy_name}${policy.description ? ` - ${policy.description}` : ""}`,
|
||||
value: policy.policy_name,
|
||||
};
|
||||
})}
|
||||
options={getPolicyOptionEntries(policies)}
|
||||
optionFilterProp="label"
|
||||
showSearch
|
||||
style={{ width: "100%" }}
|
||||
|
|
|
|||
|
|
@ -452,11 +452,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
|
|||
onEdit={(policy) => {
|
||||
setEditingPolicy(policy);
|
||||
setSelectedPolicyId(null);
|
||||
if (policy.pipeline) {
|
||||
setShowFlowBuilder(true);
|
||||
} else {
|
||||
setIsAddPolicyModalVisible(true);
|
||||
}
|
||||
setShowFlowBuilder(true);
|
||||
}}
|
||||
accessToken={accessToken}
|
||||
isAdmin={isAdmin}
|
||||
|
|
|
|||
|
|
@ -7,9 +7,22 @@ import { GuardrailPipeline, PipelineStep, PipelineTestResult, PolicyCreateReques
|
|||
import { Guardrail } from "../guardrails/types";
|
||||
import { testPipelineCall, listPolicyVersions, createPolicyVersion, updatePolicyVersionStatus } from "../networking";
|
||||
import NotificationsManager from "../molecules/notifications_manager";
|
||||
import { getComplianceDatasetPrompts } from "../../data/compliancePrompts";
|
||||
import {
|
||||
getComplianceDatasetPrompts,
|
||||
getFrameworks,
|
||||
} from "../../data/compliancePrompts";
|
||||
import type { CompliancePrompt } from "../../data/compliancePrompts";
|
||||
|
||||
const TEST_SOURCE_QUICK = "quick_chat";
|
||||
const TEST_SOURCE_ALL = "__all__";
|
||||
|
||||
function getPromptsForTestSource(source: string): CompliancePrompt[] {
|
||||
if (source === TEST_SOURCE_QUICK) return [];
|
||||
if (source === TEST_SOURCE_ALL) return getComplianceDatasetPrompts();
|
||||
const fw = getFrameworks().find((f) => f.name === source);
|
||||
return fw ? fw.categories.flatMap((c) => c.prompts) : [];
|
||||
}
|
||||
|
||||
const { Text } = Typography;
|
||||
|
||||
const ACTION_OPTIONS = [
|
||||
|
|
@ -603,18 +616,28 @@ function complianceMatchExpected(expected: "pass" | "fail", terminalAction: stri
|
|||
return terminalAction === "block";
|
||||
}
|
||||
|
||||
const testSourceOptions = [
|
||||
{ value: TEST_SOURCE_QUICK, label: "Quick chat (custom message)" },
|
||||
...getFrameworks().map((f) => ({ value: f.name, label: f.name })),
|
||||
{ value: TEST_SOURCE_ALL, label: "All compliance datasets" },
|
||||
];
|
||||
|
||||
const PipelineTestPanel: React.FC<PipelineTestPanelProps> = ({
|
||||
pipeline,
|
||||
accessToken,
|
||||
onClose,
|
||||
}) => {
|
||||
const [testSource, setTestSource] = useState<string>(TEST_SOURCE_QUICK);
|
||||
const [testMessage, setTestMessage] = useState("Hello, can you help me?");
|
||||
const [isRunning, setIsRunning] = useState(false);
|
||||
const [result, setResult] = useState<PipelineTestResult | null>(null);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [complianceRunning, setComplianceRunning] = useState(false);
|
||||
const [complianceResults, setComplianceResults] = useState<ComplianceRunEntry[]>([]);
|
||||
|
||||
const isQuickChat = testSource === TEST_SOURCE_QUICK;
|
||||
const promptsForSource = getPromptsForTestSource(testSource);
|
||||
const isDataset = promptsForSource.length > 0;
|
||||
|
||||
const handleRunTest = async () => {
|
||||
if (!accessToken) return;
|
||||
|
||||
|
|
@ -624,40 +647,29 @@ const PipelineTestPanel: React.FC<PipelineTestPanelProps> = ({
|
|||
return;
|
||||
}
|
||||
|
||||
setError(null);
|
||||
setIsRunning(true);
|
||||
setResult(null);
|
||||
setError(null);
|
||||
setComplianceResults([]);
|
||||
|
||||
try {
|
||||
const data = await testPipelineCall(
|
||||
accessToken,
|
||||
pipeline,
|
||||
[{ role: "user", content: testMessage }]
|
||||
);
|
||||
setResult(data);
|
||||
} catch (e) {
|
||||
setError(e instanceof Error ? e.message : String(e));
|
||||
} finally {
|
||||
setIsRunning(false);
|
||||
}
|
||||
};
|
||||
|
||||
const handleRunComplianceDataset = async () => {
|
||||
if (!accessToken) return;
|
||||
|
||||
const emptySteps = pipeline.steps.filter((s) => !s.guardrail);
|
||||
if (emptySteps.length > 0) {
|
||||
setError("All steps must have a guardrail selected");
|
||||
if (isQuickChat) {
|
||||
try {
|
||||
const data = await testPipelineCall(
|
||||
accessToken,
|
||||
pipeline,
|
||||
[{ role: "user", content: testMessage }]
|
||||
);
|
||||
setResult(data);
|
||||
} catch (e) {
|
||||
setError(e instanceof Error ? e.message : String(e));
|
||||
} finally {
|
||||
setIsRunning(false);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
setError(null);
|
||||
setResult(null);
|
||||
setComplianceRunning(true);
|
||||
const prompts = getComplianceDatasetPrompts();
|
||||
const entries: ComplianceRunEntry[] = [];
|
||||
|
||||
for (const prompt of prompts) {
|
||||
for (const prompt of promptsForSource) {
|
||||
try {
|
||||
const data = await testPipelineCall(accessToken, pipeline, [
|
||||
{ role: "user", content: prompt.prompt },
|
||||
|
|
@ -674,9 +686,8 @@ const PipelineTestPanel: React.FC<PipelineTestPanelProps> = ({
|
|||
});
|
||||
}
|
||||
}
|
||||
|
||||
setComplianceResults(entries);
|
||||
setComplianceRunning(false);
|
||||
setIsRunning(false);
|
||||
};
|
||||
|
||||
return (
|
||||
|
|
@ -720,40 +731,60 @@ const PipelineTestPanel: React.FC<PipelineTestPanelProps> = ({
|
|||
{/* Input section */}
|
||||
<div style={{ padding: 16, borderBottom: "1px solid #e5e7eb" }}>
|
||||
<label style={{ fontSize: 12, fontWeight: 500, color: "#6b7280", display: "block", marginBottom: 6 }}>
|
||||
Test Message
|
||||
Test with
|
||||
</label>
|
||||
<textarea
|
||||
value={testMessage}
|
||||
onChange={(e) => setTestMessage(e.target.value)}
|
||||
placeholder="Enter a test message..."
|
||||
rows={3}
|
||||
style={{
|
||||
width: "100%",
|
||||
border: "1px solid #d1d5db",
|
||||
borderRadius: 6,
|
||||
padding: "8px 10px",
|
||||
fontSize: 13,
|
||||
resize: "vertical",
|
||||
fontFamily: "inherit",
|
||||
}}
|
||||
<Select
|
||||
value={testSource}
|
||||
onChange={setTestSource}
|
||||
options={testSourceOptions}
|
||||
style={{ width: "100%", marginBottom: 12 }}
|
||||
size="middle"
|
||||
/>
|
||||
{isQuickChat && (
|
||||
<>
|
||||
<label style={{ fontSize: 12, fontWeight: 500, color: "#6b7280", display: "block", marginBottom: 6 }}>
|
||||
Message
|
||||
</label>
|
||||
<textarea
|
||||
value={testMessage}
|
||||
onChange={(e) => setTestMessage(e.target.value)}
|
||||
placeholder="Enter a test message..."
|
||||
rows={3}
|
||||
style={{
|
||||
width: "100%",
|
||||
border: "1px solid #d1d5db",
|
||||
borderRadius: 6,
|
||||
padding: "8px 10px",
|
||||
fontSize: 13,
|
||||
resize: "vertical",
|
||||
fontFamily: "inherit",
|
||||
}}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
{isDataset && (
|
||||
<div
|
||||
style={{
|
||||
fontSize: 12,
|
||||
color: "#6b7280",
|
||||
padding: "8px 10px",
|
||||
backgroundColor: "#f9fafb",
|
||||
borderRadius: 6,
|
||||
marginBottom: 8,
|
||||
}}
|
||||
>
|
||||
{testSource === TEST_SOURCE_ALL
|
||||
? "Run pipeline against all compliance prompts (EU AI Act, GDPR, Topic Blocking, Airline, etc.)."
|
||||
: `Run pipeline against ${promptsForSource.length} prompts from "${testSource}".`}
|
||||
</div>
|
||||
)}
|
||||
<Button
|
||||
onClick={handleRunTest}
|
||||
loading={isRunning}
|
||||
disabled={complianceRunning}
|
||||
style={{ marginTop: 8, width: "100%" }}
|
||||
>
|
||||
Run Test
|
||||
</Button>
|
||||
<Button
|
||||
onClick={handleRunComplianceDataset}
|
||||
loading={complianceRunning}
|
||||
disabled={isRunning}
|
||||
style={{ marginTop: 8, width: "100%" }}
|
||||
variant="secondary"
|
||||
>
|
||||
Test pipeline (compliance dataset)
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
{/* Results section */}
|
||||
|
|
@ -967,8 +998,7 @@ const PipelineTestPanel: React.FC<PipelineTestPanelProps> = ({
|
|||
|
||||
{!result && !error && complianceResults.length === 0 && (
|
||||
<div style={{ textAlign: "center", color: "#9ca3af", fontSize: 13, marginTop: 24 }}>
|
||||
Enter a test message and click "Run Test" or "Test pipeline (compliance dataset)" to
|
||||
execute the pipeline
|
||||
Choose a test source above (quick chat or a compliance dataset) and click "Run Test"
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
|
@ -1044,11 +1074,22 @@ const PolicyVersionsSidebar: React.FC<PolicyVersionsSidebarProps> = ({
|
|||
color: "#6b7280",
|
||||
letterSpacing: "0.06em",
|
||||
display: "block",
|
||||
marginBottom: 12,
|
||||
marginBottom: 4,
|
||||
}}
|
||||
>
|
||||
Versions
|
||||
</span>
|
||||
<span
|
||||
style={{
|
||||
fontSize: 11,
|
||||
color: "#6b7280",
|
||||
lineHeight: 1.4,
|
||||
display: "block",
|
||||
marginBottom: 12,
|
||||
}}
|
||||
>
|
||||
Production = the version used when anyone calls this policy by name.
|
||||
</span>
|
||||
<Button
|
||||
onClick={onNewVersion}
|
||||
disabled={!accessToken || isCreatingVersion}
|
||||
|
|
@ -1115,25 +1156,50 @@ const PolicyVersionsSidebar: React.FC<PolicyVersionsSidebarProps> = ({
|
|||
{(canPublish || canPromote) && (
|
||||
<div style={{ marginTop: 12, paddingTop: 12, borderTop: "1px solid #e5e7eb" }}>
|
||||
{canPublish && (
|
||||
<Button
|
||||
variant="secondary"
|
||||
onClick={onPublish}
|
||||
disabled={!accessToken || isUpdatingStatus}
|
||||
loading={isUpdatingStatus}
|
||||
style={{ width: "100%", marginBottom: canPromote ? 8 : 0 }}
|
||||
>
|
||||
Publish
|
||||
</Button>
|
||||
<>
|
||||
<Button
|
||||
variant="secondary"
|
||||
onClick={onPublish}
|
||||
disabled={!accessToken || isUpdatingStatus}
|
||||
loading={isUpdatingStatus}
|
||||
style={{ width: "100%", marginBottom: 8 }}
|
||||
>
|
||||
Publish
|
||||
</Button>
|
||||
<span
|
||||
style={{
|
||||
fontSize: 11,
|
||||
color: "#6b7280",
|
||||
lineHeight: 1.4,
|
||||
display: "block",
|
||||
marginBottom: canPromote ? 8 : 0,
|
||||
}}
|
||||
>
|
||||
Published versions can be tested in the Playground before promoting to production.
|
||||
</span>
|
||||
</>
|
||||
)}
|
||||
{canPromote && (
|
||||
<Button
|
||||
onClick={onPromoteToProduction}
|
||||
disabled={!accessToken || isUpdatingStatus}
|
||||
loading={isUpdatingStatus}
|
||||
style={{ width: "100%" }}
|
||||
>
|
||||
Promote to production
|
||||
</Button>
|
||||
<>
|
||||
<Button
|
||||
onClick={onPromoteToProduction}
|
||||
disabled={!accessToken || isUpdatingStatus}
|
||||
loading={isUpdatingStatus}
|
||||
style={{ width: "100%", marginBottom: 8 }}
|
||||
>
|
||||
Promote to production
|
||||
</Button>
|
||||
<span
|
||||
style={{
|
||||
fontSize: 11,
|
||||
color: "#6b7280",
|
||||
lineHeight: 1.4,
|
||||
display: "block",
|
||||
}}
|
||||
>
|
||||
This version will be used when anyone calls this policy by name.
|
||||
</span>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
|
@ -1264,6 +1330,8 @@ export const FlowBuilderPage: React.FC<FlowBuilderPageProps> = ({
|
|||
const newPolicy = await createPolicyVersion(accessToken, editingPolicy.policy_name);
|
||||
NotificationsManager.success("New draft version created");
|
||||
onVersionCreated?.(newPolicy);
|
||||
const list = await listPolicyVersions(accessToken, editingPolicy.policy_name);
|
||||
setVersions(list.versions ?? []);
|
||||
} catch (error) {
|
||||
NotificationsManager.fromBackend(
|
||||
"Failed to create version: " + (error instanceof Error ? error.message : String(error))
|
||||
|
|
@ -1282,7 +1350,9 @@ export const FlowBuilderPage: React.FC<FlowBuilderPageProps> = ({
|
|||
setIsUpdatingStatus(true);
|
||||
try {
|
||||
const updated = await updatePolicyVersionStatus(accessToken, editingPolicy.policy_id, "published");
|
||||
NotificationsManager.success("Version published");
|
||||
NotificationsManager.success(
|
||||
"Version published. You can test it in the Playground by selecting this version in the Policies dropdown."
|
||||
);
|
||||
const list = await listPolicyVersions(accessToken, editingPolicy.policy_name ?? "");
|
||||
setVersions(list.versions ?? []);
|
||||
onVersionStatusUpdated?.(updated);
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue