From e7f1fa26ab6a1f3a7f8abc02fb0d0a6e249aa2d7 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 5 Jul 2025 10:19:29 -0700 Subject: [PATCH] UI - Azure Content Guardrails (#12341) * build(model_prices_and_context_window.json): remove 'supports_tool_choice' for specific mistral models Closes https://github.com/BerriAI/litellm/issues/11750 * feat: initial commit adding cleaner ui for azure text moderation guardrails * feat(guardrail_endpoints.py): add discoverable guardrail configs and improve converting base model to dict with types * fix(guardrail_provider_fields.tsx): render from api endpoint correctly * fix(guardrail_provider_fields.tsx): cleanup * refactor(guardrail_endpoints.py): refactor to handle dictionaries with literal - allows multiselect * feat(ui/): render dictionary with known keys correctly * feat(ui/): render optional params on separate page * style(ui/): style improvements to rendering optional params on the UI * feat(azure/prompt_shield.py): add azure prompt shield back on UI * fix(add_guardrail_form.tsx): fix form to handle updated api * fix(guardrail_optional_params.tsx): ensure values are nested correctly for writing to api * fix: fix linting error * feat(text_moderation.py): handle str to int conversion * fix(guardrail_info.tsx): only render pii settings if guardrail is presidio * fix(guardrail_info.tsx): add guardrail specific fields to update settings allows updating guardrail fields (e.g. severity threshold) post-create * fix(guardrail_endpoints.py): set guardrail_id in guardrail object ensures duplicate objects not created on guardrail update * fix(guardrail_endpoints.py): allow provider specific fields to be updated on patch update * refactor(guardrail_endpoints.py): remove duplicate info endpoint * fix(guardrail_endpoints.py): mask sensitive keys on returning via guardrail `/info` Prevent leaking keys * fix(guardrail_optional_params.tsx): return numerical input when numerical component used fixes issue where output was a str * fix(guardrail_optional_params.tsx): render dict keys correctly * fix(text_moderation.py): fix severity by category check * fix(proxy/utils.py): check if guardrail should run for post call streaming hook Prevents invalid guardrails from running if not requested * test: fix import * fix: fix linting error * test: update test * fix: fix tests * fix: fix code qa errors * fix(guardrail_endpoints.py): set max depth for function * test: update recursive_detector.py * test: update list * build: merge main * fix: fix ruff check errors --- litellm/integrations/custom_guardrail.py | 3 + litellm/litellm_core_utils/litellm_logging.py | 38 +-- litellm/litellm_core_utils/redact_messages.py | 10 +- litellm/llms/custom_httpx/http_handler.py | 2 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 1 + .../proxy/_experimental/mcp_server/server.py | 8 +- litellm/proxy/_new_secret_config.yaml | 3 - litellm/proxy/common_request_processing.py | 4 +- .../proxy/guardrails/guardrail_endpoints.py | 117 +++----- .../guardrail_hooks/azure/__init__.py | 22 +- .../guardrail_hooks/azure/text_moderation.py | 8 +- litellm/proxy/proxy_server.py | 8 +- litellm/proxy/utils.py | 59 ++-- litellm/types/guardrails.py | 59 ++-- .../code_coverage_tests/recursive_detector.py | 29 +- .../test_guardrails_config.py | 7 +- .../test_litellm_logging.py | 54 +++- .../auth/test_user_api_key_auth_mcp.py | 2 +- .../guardrails/test_guardrail_endpoints.py | 13 +- .../guardrails/add_guardrail_form.tsx | 1 - .../components/guardrails/guardrail_info.tsx | 253 +++++++++++++++--- .../guardrails/guardrail_optional_params.tsx | 45 +++- .../guardrails/guardrail_provider_fields.tsx | 1 - .../src/components/guardrails/index.ts | 5 - 24 files changed, 512 insertions(+), 240 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8d6182b7e0..22dee0cce5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -102,10 +102,13 @@ class CustomGuardrail(CustomLogger): for _guardrail in requested_guardrails: if isinstance(_guardrail, dict): if self.guardrail_name in _guardrail: + return True elif isinstance(_guardrail, str): if self.guardrail_name == _guardrail: + return True + return False def should_run_guardrail(self, data, event_type: GuardrailEventHooks) -> bool: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9ff7903540..42d91ab6d4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2911,31 +2911,37 @@ def _get_masked_values( ] return { k: ( - ( - v[: unmasked_length // 2] - + "*" * number_of_asterisks - + v[-unmasked_length // 2 :] - ) - if ( - isinstance(v, str) - and len(v) > unmasked_length - and number_of_asterisks is not None + # If ignore_sensitive_values is True, or if this key doesn't contain sensitive keywords, return original value + v + if ignore_sensitive_values + or not any( + sensitive_keyword in k.lower() + for sensitive_keyword in sensitive_keywords ) else ( + # Apply masking to sensitive keys ( v[: unmasked_length // 2] - + "*" * (len(v) - unmasked_length) + + "*" * number_of_asterisks + v[-unmasked_length // 2 :] ) - if (isinstance(v, str) and len(v) > unmasked_length) - else "*****" + if ( + isinstance(v, str) + and len(v) > unmasked_length + and number_of_asterisks is not None + ) + else ( + ( + v[: unmasked_length // 2] + + "*" * (len(v) - unmasked_length) + + v[-unmasked_length // 2 :] + ) + if (isinstance(v, str) and len(v) > unmasked_length) + else ("*****" if isinstance(v, str) else v) + ) ) ) for k, v in sensitive_object.items() - if not ignore_sensitive_values - or not any( - sensitive_keyword in k.lower() for sensitive_keyword in sensitive_keywords - ) } diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index a93eb30760..2ef7af5eaa 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -62,7 +62,9 @@ def perform_redaction(model_call_details: dict, result): elif hasattr(_streaming_response, "output"): # Handle ResponsesAPIResponse format for output_item in _streaming_response.output: - if hasattr(output_item, "content") and isinstance(output_item.content, list): + if hasattr(output_item, "content") and isinstance( + output_item.content, list + ): for content_part in output_item.content: if hasattr(content_part, "text"): content_part.text = "redacted-by-litellm" @@ -154,9 +156,9 @@ def _get_turn_off_message_logging_from_dynamic_params( handles boolean and string values of `turn_off_message_logging` """ - standard_callback_dynamic_params: Optional[ - StandardCallbackDynamicParams - ] = model_call_details.get("standard_callback_dynamic_params", None) + standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( + model_call_details.get("standard_callback_dynamic_params", None) + ) if standard_callback_dynamic_params: _turn_off_message_logging = standard_callback_dynamic_params.get( "turn_off_message_logging" diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 201874603e..cf2187153a 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -4,8 +4,8 @@ import ssl import time from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional, Union -import httpx import certifi +import httpx from aiohttp import ClientSession, TCPConnector from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport from httpx._types import RequestFiles diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 7fface60ed..fba7928e7a 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_logger from litellm.proxy._types import LiteLLM_TeamTable, SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + class MCPRequestHandler: """ Class to handle MCP request processing, including: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index de488e6ca7..930d099f83 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -17,16 +17,16 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) from litellm.proxy._experimental.mcp_server.utils import ( + LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, LITELLM_MCP_SERVER_VERSION, - LITELLM_MCP_SERVER_DESCRIPTION, normalize_server_name, + normalize_server_name, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.utils import StandardLoggingMCPToolCall from litellm.utils import client - # Check if MCP is available # "mcp" requires python 3.10 or higher, but several litellm users use python 3.8 # We're making this conditional import to avoid breaking users who use python 3.8. @@ -67,7 +67,9 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) - from litellm.proxy._experimental.mcp_server.utils import get_server_name_prefix_tool_mcp + from litellm.proxy._experimental.mcp_server.utils import ( + get_server_name_prefix_tool_mcp, + ) ###################################################### ############ MCP Tools List REST API Response Object # diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index afaac02c08..949116beaa 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -8,6 +8,3 @@ model_list: model: langfuse/langfuse-model prompt_id: test-chat-prompt prompt_version: 4 - -litellm_settings: - callbacks: ["langfuse_otel"] \ No newline at end of file diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3a33b3c5f9..8f18e5cbf2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -727,7 +727,9 @@ class ProxyBaseLLMRequestProcessing: ) ### CALL HOOKS ### - modify outgoing data chunk = await proxy_logging_obj.async_post_call_streaming_hook( - user_api_key_dict=user_api_key_dict, response=chunk + user_api_key_dict=user_api_key_dict, + response=chunk, + data=request_data, ) # Format chunk using helper function diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 7680bf394f..0680fa1260 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -8,6 +8,7 @@ from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry from litellm.types.guardrails import ( @@ -264,64 +265,6 @@ async def create_guardrail(request: CreateGuardrailRequest): raise HTTPException(status_code=500, detail=str(e)) -@router.get( - "/guardrails/{guardrail_id}", - tags=["Guardrails"], - dependencies=[Depends(user_api_key_auth)], -) -async def get_guardrail(guardrail_id: str): - """ - Get a guardrail by ID - - 👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start) - - Example Request: - ```bash - curl -X GET "http://localhost:4000/guardrails/123e4567-e89b-12d3-a456-426614174000" \\ - -H "Authorization: Bearer " - ``` - - Example Response: - ```json - { - "guardrail_id": "123e4567-e89b-12d3-a456-426614174000", - "guardrail_name": "my-bedrock-guard", - "litellm_params": { - "guardrail": "bedrock", - "mode": "pre_call", - "guardrailIdentifier": "ff6ujrregl1q", - "guardrailVersion": "DRAFT", - "default_on": true - }, - "guardrail_info": { - "description": "Bedrock content moderation guardrail" - }, - "created_at": "2023-11-09T12:34:56.789Z", - "updated_at": "2023-11-09T12:34:56.789Z" - } - ``` - """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - raise HTTPException(status_code=500, detail="Prisma client not initialized") - - try: - result = await GUARDRAIL_REGISTRY.get_guardrail_by_id_from_db( - guardrail_id=guardrail_id, prisma_client=prisma_client - ) - - if result is None: - raise HTTPException( - status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" - ) - return result - except HTTPException as e: - raise e - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) - - class UpdateGuardrailRequest(BaseModel): guardrail: Guardrail @@ -542,19 +485,13 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): litellm_params = LitellmParams( **dict(existing_guardrail.get("litellm_params", {})) ) - if ( - request.litellm_params is not None - and request.litellm_params.default_on is not None - ): - litellm_params.default_on = request.litellm_params.default_on - - if ( - request.litellm_params is not None - and request.litellm_params.pii_entities_config is not None - ): - litellm_params.pii_entities_config = ( - request.litellm_params.pii_entities_config + if request.litellm_params is not None: + requested_litellm_params = request.litellm_params.model_dump( + exclude_unset=True ) + litellm_params_dict = litellm_params.model_dump(exclude_unset=True) + litellm_params_dict.update(requested_litellm_params) + litellm_params = LitellmParams(**litellm_params_dict) # Update guardrail_info if provided guardrail_info = ( @@ -565,6 +502,7 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): # Create the guardrail object guardrail = Guardrail( + guardrail_id=guardrail_id, guardrail_name=guardrail_name or "", litellm_params=litellm_params, guardrail_info=guardrail_info, @@ -588,6 +526,11 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): raise HTTPException(status_code=500, detail=str(e)) +@router.get( + "/guardrails/{guardrail_id}", + tags=["Guardrails"], + dependencies=[Depends(user_api_key_auth)], +) @router.get( "/guardrails/{guardrail_id}/info", tags=["Guardrails"], @@ -625,6 +568,8 @@ async def get_guardrail_info(guardrail_id: str): } ``` """ + + from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client @@ -645,10 +590,24 @@ async def get_guardrail_info(guardrail_id: str): status_code=404, detail=f"Guardrail with ID {guardrail_id} not found" ) + litellm_params: Optional[Union[LitellmParams, dict]] = result.get( + "litellm_params" + ) + result_litellm_params_dict = ( + litellm_params.model_dump(exclude_none=True) + if isinstance(litellm_params, LitellmParams) + else litellm_params + ) or {} + masked_litellm_params_dict = _get_masked_values( + result_litellm_params_dict, + unmasked_length=4, + number_of_asterisks=4, + ) + return GuardrailInfoResponse( guardrail_id=result.get("guardrail_id"), guardrail_name=result.get("guardrail_name"), - litellm_params=dict(result.get("litellm_params") or {}), + litellm_params=masked_litellm_params_dict, guardrail_info=dict(result.get("guardrail_info") or {}), created_at=result.get("created_at"), updated_at=result.get("updated_at"), @@ -802,7 +761,17 @@ def _get_fields_from_model(model_class: Type[BaseModel]) -> Dict[str, Any]: """ import inspect - def _extract_fields_recursive(model: Type[BaseModel]) -> Dict[str, Any]: + def _extract_fields_recursive( + model: Type[BaseModel], + depth: int = 0, + ) -> Dict[str, Any]: + # Check if we've exceeded the maximum recursion depth + if depth > DEFAULT_MAX_RECURSE_DEPTH: + raise HTTPException( + status_code=400, + detail=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing model fields. Please check the model structure for excessive nesting.", + ) + fields = {} for field_name, field in model.model_fields.items(): @@ -840,7 +809,7 @@ def _get_fields_from_model(model_class: Type[BaseModel]) -> Dict[str, Any]: if is_basemodel_subclass: # Recursively get fields from the nested model nested_fields = _extract_fields_recursive( - cast(Type[BaseModel], field_annotation) + cast(Type[BaseModel], field_annotation), depth + 1 ) fields[field_name] = { "description": description, @@ -896,7 +865,7 @@ def _get_fields_from_model(model_class: Type[BaseModel]) -> Dict[str, Any]: return fields - return _extract_fields_recursive(model_class) + return _extract_fields_recursive(model_class, depth=0) @router.get( diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py index 9f69c71303..449dd42ba5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py @@ -29,18 +29,24 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" AzureContentSafetyTextModerationGuardrail, ] = AzureContentSafetyPromptShieldGuardrail( guardrail_name=guardrail_name, - api_key=litellm_params.api_key, - api_base=litellm_params.api_base, - default_on=litellm_params.default_on, - event_hook=litellm_params.mode, + **{ + **litellm_params.model_dump(exclude_none=True), + "api_key": litellm_params.api_key, + "api_base": litellm_params.api_base, + "default_on": litellm_params.default_on, + "event_hook": litellm_params.mode, + }, ) elif azure_guardrail == "text_moderations": azure_content_safety_guardrail = AzureContentSafetyTextModerationGuardrail( guardrail_name=guardrail_name, - api_key=litellm_params.api_key, - api_base=litellm_params.api_base, - default_on=litellm_params.default_on, - event_hook=litellm_params.mode, + **{ + **litellm_params.model_dump(exclude_none=True), + "api_key": litellm_params.api_key, + "api_base": litellm_params.api_base, + "default_on": litellm_params.default_on, + "event_hook": litellm_params.mode, + }, ) else: raise ValueError( diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index 1e7c294ed5..9c31e06202 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -91,7 +91,9 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr "outputType": kwargs.get("outputType") or "FourSeverityLevels", } - self.severity_threshold = severity_threshold + self.severity_threshold = ( + int(severity_threshold) if severity_threshold else None + ) self.severity_threshold_by_category = severity_threshold_by_category verbose_proxy_logger.info( @@ -124,6 +126,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr verbose_proxy_logger.debug( "Azure Text Moderation guard request: %s", request_body ) + response = await self.async_handler.post( url=f"{self.api_base}/contentsafety/text:analyze?api-version={self.api_version}", headers={ @@ -146,6 +149,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr - Check if general severity threshold set - If both none, use default_severity_threshold """ + if self.severity_threshold_by_category: for category in response["categoriesAnalysis"]: severity_category_threshold_item = ( @@ -153,7 +157,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr ) if ( severity_category_threshold_item is not None - and severity_category_threshold_item >= category["severity"] + and category["severity"] >= severity_category_threshold_item ): raise HTTPException( status_code=400, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6140d5cd93..aa2543361d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3097,7 +3097,9 @@ async def async_assistants_data_generator( async with response as chunk: ### CALL HOOKS ### - modify outgoing data chunk = await proxy_logging_obj.async_post_call_streaming_hook( - user_api_key_dict=user_api_key_dict, response=chunk + user_api_key_dict=user_api_key_dict, + response=chunk, + data=request_data, ) # chunk = chunk.model_dump_json(exclude_none=True) @@ -3156,7 +3158,9 @@ async def async_data_generator( ) ### CALL HOOKS ### - modify outgoing data chunk = await proxy_logging_obj.async_post_call_streaming_hook( - user_api_key_dict=user_api_key_dict, response=chunk + user_api_key_dict=user_api_key_dict, + response=chunk, + data=request_data, ) if isinstance(chunk, BaseModel): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5fddd47c07..cc26769e20 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1061,6 +1061,7 @@ class ProxyLogging: async def async_post_call_streaming_hook( self, + data: dict, response: Union[ ModelResponse, EmbeddingResponse, ImageResponse, ModelResponseStream ], @@ -1079,6 +1080,17 @@ class ProxyLogging: for callback in litellm.callbacks: try: _callback: Optional[CustomLogger] = None + if isinstance(callback, CustomGuardrail): + # Main - V2 Guardrails implementation + from litellm.types.guardrails import GuardrailEventHooks + + if ( + callback.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.post_call + ) + is not True + ): + continue if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( callback @@ -2478,15 +2490,23 @@ class PrismaClient: ) # Health Check Database Methods - def _validate_response_time(self, response_time_ms: Optional[float]) -> Optional[float]: + def _validate_response_time( + self, response_time_ms: Optional[float] + ) -> Optional[float]: """Validate and clean response time value""" if response_time_ms is None: return None try: value = float(response_time_ms) - return value if value == value and value not in (float('inf'), float('-inf')) else None + return ( + value + if value == value and value not in (float("inf"), float("-inf")) + else None + ) except (ValueError, TypeError): - verbose_proxy_logger.warning(f"Invalid response_time_ms value: {response_time_ms}") + verbose_proxy_logger.warning( + f"Invalid response_time_ms value: {response_time_ms}" + ) return None def _clean_details(self, details: Optional[dict]) -> Optional[dict]: @@ -2520,7 +2540,7 @@ class PrismaClient: "healthy_count": int(healthy_count), "unhealthy_count": int(unhealthy_count), } - + # Add optional fields using dict comprehension and helper methods optional_fields = { "error_message": str(error_message)[:500] if error_message else None, @@ -2529,15 +2549,19 @@ class PrismaClient: "checked_by": str(checked_by) if checked_by else None, "model_id": str(model_id) if model_id else None, } - + # Add only non-None optional fields - health_check_data.update({k: v for k, v in optional_fields.items() if v is not None}) - + health_check_data.update( + {k: v for k, v in optional_fields.items() if v is not None} + ) + verbose_proxy_logger.debug(f"Saving health check data: {health_check_data}") return await self.db.litellm_healthchecktable.create(data=health_check_data) - + except Exception as e: - verbose_proxy_logger.error(f"Error saving health check result for model {model_name}: {e}") + verbose_proxy_logger.error( + f"Error saving health check result for model {model_name}: {e}" + ) return None async def get_health_check_history( @@ -2577,13 +2601,13 @@ class PrismaClient: all_checks = await self.db.litellm_healthchecktable.find_many( order={"checked_at": "desc"} ) - + # Group by model_name and get the latest for each latest_checks = {} for check in all_checks: if check.model_name not in latest_checks: latest_checks[check.model_name] = check - + return list(latest_checks.values()) except Exception as e: verbose_proxy_logger.error(f"Error getting all latest health checks: {e}") @@ -2592,6 +2616,7 @@ class PrismaClient: ### HELPER FUNCTIONS ### + async def _cache_user_row(user_id: str, cache: DualCache, db: PrismaClient): """ Check if a user_id exists in cache, @@ -3012,7 +3037,7 @@ def _get_redoc_url() -> Optional[str]: - If NO_REDOC is True, return None. - Otherwise, default to "/redoc". """ - if (redoc_url := os.getenv("REDOC_URL")): + if redoc_url := os.getenv("REDOC_URL"): return redoc_url if str_to_bool(os.getenv("NO_REDOC")) is True: @@ -3029,7 +3054,7 @@ def _get_docs_url() -> Optional[str]: - If NO_DOCS is True, return None. - Otherwise, default to "/". """ - if (docs_url := os.getenv("DOCS_URL")): + if docs_url := os.getenv("DOCS_URL"): return docs_url if str_to_bool(os.getenv("NO_DOCS")) is True: @@ -3097,15 +3122,15 @@ def join_paths(base_path: str, route: str) -> str: # Remove trailing slashes from base_path and leading slashes from route base_path = base_path.rstrip("/") route = route.lstrip("/") - + # If base_path is empty, return route with leading slash if not base_path: return f"/{route}" if route else "/" - + # If route is empty, return just base_path if not route: return base_path - + # Join with single slash return f"{base_path}/{route}" @@ -3117,7 +3142,7 @@ def get_custom_url(request_base_url: str, route: Optional[str] = None) -> str: base_url = server_base_url else: base_url = request_base_url - + server_root_path = get_server_root_path() if route is not None: if server_root_path != "": diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 742dcf82b5..d7835db54b 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -339,16 +339,7 @@ class LassoGuardrailConfigModel(BaseModel): ) -class LitellmParams( - PresidioConfigModel, - BedrockGuardrailConfigModel, - LakeraV2GuardrailConfigModel, - LassoGuardrailConfigModel, -): - guardrail: str = Field(description="The type of guardrail integration to use") - mode: Union[str, List[str]] = Field( - description="When to apply the guardrail (pre_call, post_call, during_call, logging_only)" - ) +class BaseLitellmParams(BaseModel): # works for new and patch update guardrails api_key: Optional[str] = Field( default=None, description="API key for the guardrail service" ) @@ -395,6 +386,29 @@ class LitellmParams( default=None, description="Recipe for output (LLM response)" ) + model_config = ConfigDict(extra="allow", protected_namespaces=()) + + +class LitellmParams( + PresidioConfigModel, + BedrockGuardrailConfigModel, + LakeraV2GuardrailConfigModel, + LassoGuardrailConfigModel, + BaseLitellmParams, +): + guardrail: str = Field(description="The type of guardrail integration to use") + mode: Union[str, List[str]] = Field( + description="When to apply the guardrail (pre_call, post_call, during_call, logging_only)" + ) + + def __init__(self, **kwargs): + default_on = kwargs.pop("default_on", None) + if default_on is not None: + kwargs["default_on"] = default_on + else: + kwargs["default_on"] = False + super().__init__(**kwargs) + class Guardrail(TypedDict, total=False): guardrail_id: Optional[str] @@ -420,26 +434,10 @@ class DynamicGuardrailParams(TypedDict): extra_body: Dict[str, Any] -class GuardrailInfoLiteLLMParamsResponse(BaseModel): - """The returned LiteLLM Params object for /guardrails/list""" - - guardrail: str - mode: Union[str, List[str]] - default_on: Optional[bool] = False - pii_entities_config: Optional[Dict[PiiEntityType, PiiAction]] = None - - def __init__(self, **kwargs): - default_on = kwargs.get("default_on") - if default_on is None: - default_on = False - - super().__init__(**kwargs) - - class GuardrailInfoResponse(BaseModel): guardrail_id: Optional[str] = None guardrail_name: str - litellm_params: Optional[GuardrailInfoLiteLLMParamsResponse] = None + litellm_params: Optional[BaseLitellmParams] = None guardrail_info: Optional[Dict] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None @@ -480,12 +478,7 @@ class ApplyGuardrailResponse(BaseModel): response_text: str -class PatchGuardrailLitellmParams(BaseModel): - default_on: Optional[bool] = None - pii_entities_config: Optional[Dict[PiiEntityType, PiiAction]] = None - - class PatchGuardrailRequest(BaseModel): guardrail_name: Optional[str] = None - litellm_params: Optional[PatchGuardrailLitellmParams] = None + litellm_params: Optional[BaseLitellmParams] = None guardrail_info: Optional[Dict[str, Any]] = None diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index b2fe50bae1..ae8138f057 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -10,20 +10,21 @@ IGNORE_FUNCTIONS = [ "_check_for_os_environ_vars", "clean_message", "unpack_defs", - "convert_anyof_null_to_nullable", # has a set max depth + "convert_anyof_null_to_nullable", # has a set max depth "add_object_type", "strip_field", "_transform_prompt", "mask_dict", "_serialize", # we now set a max depth for this - "_sanitize_request_body_for_spend_logs_payload", # testing added for circular reference - "_sanitize_value", # testing added for circular reference - "set_schema_property_ordering", # testing added for infinite recursion - "process_items", # testing added for infinite recursion + max depth set. - "_can_object_call_model", # max depth set. - "encode_unserializable_types", # max depth set. - "filter_value_from_dict", # max depth set. - "normalize_json_schema_types", # max depth set. + "_sanitize_request_body_for_spend_logs_payload", # testing added for circular reference + "_sanitize_value", # testing added for circular reference + "set_schema_property_ordering", # testing added for infinite recursion + "process_items", # testing added for infinite recursion + max depth set. + "_can_object_call_model", # max depth set. + "encode_unserializable_types", # max depth set. + "filter_value_from_dict", # max depth set. + "normalize_json_schema_types", # max depth set. + "_extract_fields_recursive", # max depth set. ] @@ -85,6 +86,7 @@ def find_recursive_functions_in_directory(directory): ignored_recursive_functions[file_path] = ignored return recursive_functions, ignored_recursive_functions + if __name__ == "__main__": # Example usage # raise exception if any recursive functions are found, except for the ignored ones @@ -101,10 +103,9 @@ if __name__ == "__main__": # raise exception if any recursive functions are found for file, functions in recursive_functions.items(): print( - f"🚨 Unignored recursive functions found in {file}: {functions}. THIS IS REALLY BAD, it has caused CPU Usage spikes in the past. Only keep this if it's ABSOLUTELY necessary." - ) + f"🚨 Unignored recursive functions found in {file}: {functions}. THIS IS REALLY BAD, it has caused CPU Usage spikes in the past. Only keep this if it's ABSOLUTELY necessary." + ) file, functions = list(recursive_functions.items())[0] raise Exception( - f"🚨 Unignored recursive functions found include {file}: {functions}. THIS IS REALLY BAD, it has caused CPU Usage spikes in the past. Only keep this if it's ABSOLUTELY necessary." - ) - + f"🚨 Unignored recursive functions found include {file}: {functions}. THIS IS REALLY BAD, it has caused CPU Usage spikes in the past. Only keep this if it's ABSOLUTELY necessary." + ) diff --git a/tests/guardrails_tests/test_guardrails_config.py b/tests/guardrails_tests/test_guardrails_config.py index f76e3f5ec0..110c0a30a9 100644 --- a/tests/guardrails_tests/test_guardrails_config.py +++ b/tests/guardrails_tests/test_guardrails_config.py @@ -93,11 +93,14 @@ def test_guardrail_list_of_event_hooks(): def test_guardrail_info_response(): - from litellm.types.guardrails import GuardrailInfoResponse, LitellmParams, GuardrailInfoLiteLLMParamsResponse + from litellm.types.guardrails import ( + GuardrailInfoResponse, + LitellmParams, + ) guardrail_info = GuardrailInfoResponse( guardrail_name="aporia-pre-guard", - litellm_params=GuardrailInfoLiteLLMParamsResponse( + litellm_params=LitellmParams( guardrail="aporia", mode="pre_call", ), diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index d4512b1671..e9b962f189 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -11,9 +11,9 @@ sys.path.insert( import time +from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import set_callbacks -from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST @pytest.fixture @@ -343,4 +343,54 @@ def test_sentry_event_scrubber_initialization(monkeypatch): mock_init.assert_called_once() call_args = mock_init.call_args[1] assert call_args["event_scrubber"] == mock_event_scrubber_instance - assert call_args["send_default_pii"] is False \ No newline at end of file + assert call_args["send_default_pii"] is False + + +def test_get_masked_values(): + from litellm.litellm_core_utils.litellm_logging import _get_masked_values + + sensitive_object = { + "mode": "pre_call", + "api_key": "sensitive_api_key", + "payload": True, + "api_base": "sensitive_api_base", + "dev_info": True, + "metadata": None, + "breakdown": True, + "guardrail": "azure/text_moderations", + "default_on": False, + "guard_name": None, + "project_id": None, + "aws_role_name": None, + "lasso_user_id": None, + "aws_region_name": None, + "aws_profile_name": None, + "aws_session_name": None, + "aws_sts_endpoint": None, + "guardrailVersion": None, + "output_parse_pii": None, + "aws_access_key_id": None, + "aws_session_token": None, + "presidio_language": "en", + "mock_redacted_text": None, + "severity_threshold": "5", + "category_thresholds": None, + "guardrailIdentifier": None, + "pangea_input_recipe": None, + "pii_entities_config": {}, + "mask_request_content": None, + "pangea_output_recipe": None, + "aws_secret_access_key": None, + "detect_secrets_config": None, + "lasso_conversation_id": None, + "mask_response_content": None, + "aws_web_identity_token": None, + "presidio_analyzer_api_base": None, + "presidio_ad_hoc_recognizers": None, + "aws_bedrock_runtime_endpoint": None, + "presidio_anonymizer_api_base": None, + } + masked_values = _get_masked_values( + sensitive_object, unmasked_length=4, number_of_asterisks=4 + ) + assert masked_values["presidio_anonymizer_api_base"] is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 3d3f7f336d..67c7a9c122 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -15,7 +15,7 @@ sys.path.insert( from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) -from litellm.proxy._types import UserAPIKeyAuth, SpecialHeaders +from litellm.proxy._types import SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index c4d9f517db..cf4ab0152b 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -22,8 +22,9 @@ from litellm.proxy.guardrails.guardrail_registry import ( InMemoryGuardrailHandler, ) from litellm.types.guardrails import ( - GuardrailInfoLiteLLMParamsResponse, + BaseLitellmParams, GuardrailInfoResponse, + LitellmParams, ) # Mock data for testing @@ -98,7 +99,7 @@ async def test_list_guardrails_v2_with_db_and_config( ) assert db_guardrail.guardrail_name == "Test DB Guardrail" assert db_guardrail.guardrail_definition_location == "db" - assert isinstance(db_guardrail.litellm_params, GuardrailInfoLiteLLMParamsResponse) + assert isinstance(db_guardrail.litellm_params, BaseLitellmParams) # Check config guardrail config_guardrail = next( @@ -106,9 +107,7 @@ async def test_list_guardrails_v2_with_db_and_config( ) assert config_guardrail.guardrail_name == "Test Config Guardrail" assert config_guardrail.guardrail_definition_location == "config" - assert isinstance( - config_guardrail.litellm_params, GuardrailInfoLiteLLMParamsResponse - ) + assert isinstance(config_guardrail.litellm_params, BaseLitellmParams) @pytest.mark.asyncio @@ -120,7 +119,7 @@ async def test_get_guardrail_info_from_db(mocker, mock_prisma_client): assert response.guardrail_id == "test-db-guardrail" assert response.guardrail_name == "Test DB Guardrail" - assert isinstance(response.litellm_params, GuardrailInfoLiteLLMParamsResponse) + assert isinstance(response.litellm_params, BaseLitellmParams) assert response.guardrail_info == {"description": "Test guardrail from DB"} @@ -144,7 +143,7 @@ async def test_get_guardrail_info_from_config( assert response.guardrail_id == "test-config-guardrail" assert response.guardrail_name == "Test Config Guardrail" - assert isinstance(response.litellm_params, GuardrailInfoLiteLLMParamsResponse) + assert isinstance(response.litellm_params, BaseLitellmParams) assert response.guardrail_info == {"description": "Test guardrail from config"} diff --git a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx index 2298d5f2af..285b1c7673 100644 --- a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx @@ -6,7 +6,6 @@ import { GuardrailProviders, guardrail_provider_map, shouldRenderPIIConfigSettin import { createGuardrailCall, getGuardrailUISettings, getGuardrailProviderSpecificParams } from '../networking'; import PiiConfiguration from './pii_configuration'; import GuardrailProviderFields from './guardrail_provider_fields'; -import AzureTextModerationConfiguration from './azure_text_moderation_configuration'; import GuardrailOptionalParams from './guardrail_optional_params'; const { Title, Text, Link } = Typography; diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx index 9fdcb1badf..f4cbc1cd2f 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx @@ -15,9 +15,11 @@ import { } from "@tremor/react"; import { Button, Form, Input, Select, message, Tooltip, Divider } from "antd"; import { InfoCircleOutlined } from '@ant-design/icons'; -import { getGuardrailInfo, updateGuardrailCall, getGuardrailUISettings } from "@/components/networking"; -import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { getGuardrailInfo, updateGuardrailCall, getGuardrailUISettings, getGuardrailProviderSpecificParams } from "@/components/networking"; +import { getGuardrailLogoAndName, guardrail_provider_map } from "./guardrail_info_helpers"; import PiiConfiguration from "./pii_configuration"; +import GuardrailProviderFields from "./guardrail_provider_fields"; +import GuardrailOptionalParams from "./guardrail_optional_params"; export interface GuardrailInfoProps { guardrailId: string; @@ -26,6 +28,22 @@ export interface GuardrailInfoProps { isAdmin: boolean; } +interface ProviderParam { + param: string; + description: string; + required: boolean; + default_value?: string; + options?: string[]; + type?: string; + fields?: { [key: string]: ProviderParam }; + dict_key_options?: string[]; + dict_value_type?: string; +} + +interface ProviderParamsResponse { + [provider: string]: { [key: string]: ProviderParam }; +} + const GuardrailInfoView: React.FC = ({ guardrailId, onClose, @@ -33,6 +51,7 @@ const GuardrailInfoView: React.FC = ({ isAdmin }) => { const [guardrailData, setGuardrailData] = useState(null); + const [guardrailProviderSpecificParams, setGuardrailProviderSpecificParams] = useState(null); const [loading, setLoading] = useState(true); const [isEditing, setIsEditing] = useState(false); const [form] = Form.useForm(); @@ -89,21 +108,53 @@ const GuardrailInfoView: React.FC = ({ } }; + const fetchGuardrailProviderSpecificParams = async () => { + try { + if (!accessToken) return; + const response = await getGuardrailProviderSpecificParams(accessToken); + setGuardrailProviderSpecificParams(response); + } catch (error) { + console.error("Error fetching guardrail provider specific params:", error); + } + }; + const fetchGuardrailUISettings = async () => { try { if (!accessToken) return; const uiSettings = await getGuardrailUISettings(accessToken); setGuardrailSettings(uiSettings); + } catch (error) { console.error("Error fetching guardrail UI settings:", error); } }; + useEffect(() => { + fetchGuardrailProviderSpecificParams(); + }, [accessToken]); + useEffect(() => { fetchGuardrailInfo(); fetchGuardrailUISettings(); }, [guardrailId, accessToken]); + // Reset form when guardrail data or provider params change + useEffect(() => { + if (guardrailData && form) { + form.setFieldsValue({ + guardrail_name: guardrailData.guardrail_name, + ...guardrailData.litellm_params, + guardrail_info: guardrailData.guardrail_info + ? JSON.stringify(guardrailData.guardrail_info, null, 2) + : "", + // Include any optional_params if they exist + ...(guardrailData.litellm_params?.optional_params && { + optional_params: guardrailData.litellm_params.optional_params + }) + }); + } + }, [guardrailData, guardrailProviderSpecificParams, form]); + const handlePiiEntitySelect = (entity: string) => { setSelectedPiiEntities(prev => { if (prev.includes(entity)) { @@ -125,29 +176,119 @@ const GuardrailInfoView: React.FC = ({ try { if (!accessToken) return; - // Prepare update data object + // Prepare update data object - only include changed fields const updateData: any = { - guardrail_name: values.guardrail_name, - litellm_params: { - default_on: values.default_on, - }, - guardrail_info: values.guardrail_info ? JSON.parse(values.guardrail_info) : undefined + litellm_params: {} }; - // Only add PII entities config if we have selected entities - if (selectedPiiEntities.length > 0) { - // Create PII config object only with selected entities - const piiEntitiesConfig: {[key: string]: string} = {}; - selectedPiiEntities.forEach(entity => { - piiEntitiesConfig[entity] = selectedPiiActions[entity] || "MASK"; + // Only include guardrail_name if it has changed + if (values.guardrail_name !== guardrailData.guardrail_name) { + updateData.guardrail_name = values.guardrail_name; + } + + // Only include default_on if it has changed + if (values.default_on !== guardrailData.litellm_params?.default_on) { + updateData.litellm_params.default_on = values.default_on; + } + + // Only include guardrail_info if it has changed + const originalGuardrailInfo = guardrailData.guardrail_info; + const newGuardrailInfo = values.guardrail_info ? JSON.parse(values.guardrail_info) : undefined; + if (JSON.stringify(originalGuardrailInfo) !== JSON.stringify(newGuardrailInfo)) { + updateData.guardrail_info = newGuardrailInfo; + } + + // Only add PII entities config if there are changes + const originalPiiConfig = guardrailData.litellm_params?.pii_entities_config || {}; + const newPiiEntitiesConfig: {[key: string]: string} = {}; + + selectedPiiEntities.forEach(entity => { + newPiiEntitiesConfig[entity] = selectedPiiActions[entity] || "MASK"; + }); + + // Only update if PII config has changed + if (JSON.stringify(originalPiiConfig) !== JSON.stringify(newPiiEntitiesConfig)) { + updateData.litellm_params.pii_entities_config = newPiiEntitiesConfig; + } + + /****************************** + * Add provider-specific params (reusing logic from add_guardrail_form.tsx) + * ---------------------------------- + * The backend exposes exactly which extra parameters a provider + * accepts via `/guardrails/ui/provider_specific_params`. + * Instead of copying every unknown form field, we fetch the list for + * the selected provider and ONLY pass those recognised params. + ******************************/ + + // Get the current provider from the guardrail data + const currentProvider = Object.keys(guardrail_provider_map).find( + key => guardrail_provider_map[key] === guardrailData.litellm_params?.guardrail + ); + + console.log("values: ", JSON.stringify(values)); + console.log("currentProvider: ", currentProvider); + + // Use pre-fetched provider params to copy recognised params + if (guardrailProviderSpecificParams && currentProvider) { + const providerKey = guardrail_provider_map[currentProvider]?.toLowerCase(); + const providerSpecificParams = guardrailProviderSpecificParams[providerKey] || {}; + + const allowedParams = new Set(); + + console.log("providerSpecificParams: ", JSON.stringify(providerSpecificParams)); + + // Add root-level parameters (like api_key, api_base, api_version) + Object.keys(providerSpecificParams).forEach(paramName => { + if (paramName !== 'optional_params') { + allowedParams.add(paramName); + } }); - // Add to litellm_params only if we have entities - updateData.litellm_params.pii_entities_config = piiEntitiesConfig; - } else { - // If no entities selected, explicitly set to empty object - // This will clear any existing PII config - updateData.litellm_params.pii_entities_config = {}; + // Add nested parameters from optional_params.fields + if (providerSpecificParams.optional_params && + providerSpecificParams.optional_params.fields) { + Object.keys(providerSpecificParams.optional_params.fields).forEach(paramName => { + allowedParams.add(paramName); + }); + } + + console.log("allowedParams: ", allowedParams); + allowedParams.forEach((paramName) => { + // Check for both direct parameter name and nested optional_params object + let paramValue = values[paramName]; + if (paramValue === undefined || paramValue === null || paramValue === '') { + paramValue = values.optional_params?.[paramName]; + } + + // Get the original value for comparison + const originalValue = guardrailData.litellm_params?.[paramName]; + + // Check if the value has changed from the original + const hasChanged = JSON.stringify(paramValue) !== JSON.stringify(originalValue); + + // Include if value has changed and has a meaningful value, OR if user explicitly cleared a value + if (hasChanged) { + if (paramValue !== undefined && paramValue !== null && paramValue !== '') { + // User set a new value + updateData.litellm_params[paramName] = paramValue; + } else if (originalValue !== undefined && originalValue !== null && originalValue !== '') { + // User cleared an existing value - set to null to indicate removal + updateData.litellm_params[paramName] = null; + } + } + }); + } + + // Remove empty litellm_params object if no parameters were changed + if (Object.keys(updateData.litellm_params).length === 0) { + delete updateData.litellm_params; + } + + // Only proceed with update if there are actual changes + if (Object.keys(updateData).length === 0) { + message.info("No changes detected"); + setIsEditing(false); + return; } await updateGuardrailCall(accessToken, guardrailId, updateData); @@ -290,6 +431,10 @@ const GuardrailInfoView: React.FC = ({ guardrail_info: guardrailData.guardrail_info ? JSON.stringify(guardrailData.guardrail_info, null, 2) : "", + // Include any optional_params if they exist + ...(guardrailData.litellm_params?.optional_params && { + optional_params: guardrailData.litellm_params.optional_params + }) }} layout="vertical" > @@ -311,21 +456,59 @@ const GuardrailInfoView: React.FC = ({ - PII Protection -
- {guardrailSettings && ( - - )} -
- + {guardrailData.litellm_params?.guardrail === "presidio" && ( + <> + PII Protection +
+ {guardrailSettings && ( + + )} +
+ + )} + + Provider Settings + + {/* Provider-specific fields */} + guardrail_provider_map[key] === guardrailData.litellm_params?.guardrail + ) || null} + accessToken={accessToken} + providerParams={guardrailProviderSpecificParams} + /> + + {/* Optional parameters */} + {guardrailProviderSpecificParams && ( + (() => { + const currentProvider = Object.keys(guardrail_provider_map).find( + key => guardrail_provider_map[key] === guardrailData.litellm_params?.guardrail + ); + if (!currentProvider) return null; + + const providerKey = guardrail_provider_map[currentProvider]?.toLowerCase(); + const providerFields = guardrailProviderSpecificParams[providerKey]; + + if (!providerFields || !providerFields.optional_params) return null; + + return ( + + ); + })() + )} + Advanced Settings ; } interface DictFieldProps { field: ProviderParam; fieldKey: string; fullFieldKey: string | string[]; + value: any | null; } -const DictField: React.FC = ({ field, fieldKey, fullFieldKey }) => { +const DictField: React.FC = ({ field, fieldKey, fullFieldKey, value }) => { const [selectedEntries, setSelectedEntries] = React.useState>([]); const [availableKeys, setAvailableKeys] = React.useState(field.dict_key_options || []); + // Initialize selectedEntries and availableKeys based on existing value + React.useEffect(() => { + if (value && typeof value === 'object') { + const existingKeys = Object.keys(value); + const entries = existingKeys.map(key => ({ + key: key, + id: `${key}_${Date.now()}_${Math.random()}` + })); + setSelectedEntries(entries); + + const remainingKeys = (field.dict_key_options || []).filter(key => !existingKeys.includes(key)); + setAvailableKeys(remainingKeys); + } + }, [value, field.dict_key_options]); + const addEntry = (selectedKey: string) => { if (!selectedKey) return; @@ -59,6 +76,12 @@ const DictField: React.FC = ({ field, fieldKey, fullFieldKey }) { + if (value === null || value === undefined || value === '') return undefined; + const num = Number(value); + return isNaN(num) ? value : num; + } : undefined} > {field.dict_value_type === "number" ? ( = ({ field, fieldKey, fullFieldKey }) const GuardrailOptionalParams: React.FC = ({ optionalParams, - parentFieldKey + parentFieldKey, + values, }) => { const renderField = (fieldKey: string, field: ProviderParam) => { const fullFieldKey = `${parentFieldKey}.${fieldKey}`; - + const value = values?.[fieldKey]; + console.log("value", value); // Handle dict fields separately since they manage their own Form.Items if (field.type === "dict" && field.dict_key_options) { return ( @@ -128,6 +153,7 @@ const GuardrailOptionalParams: React.FC = ({ field={field} fieldKey={fieldKey} fullFieldKey={[parentFieldKey, fieldKey]} + value={value} /> ); @@ -145,11 +171,16 @@ const GuardrailOptionalParams: React.FC = ({ } rules={field.required ? [{ required: true, message: `${fieldKey} is required` }] : undefined} className="mb-0" + initialValue={value !== undefined ? value : field.default_value} + normalize={field.type === "number" ? (value) => { + if (value === null || value === undefined || value === '') return undefined; + const num = Number(value); + return isNaN(num) ? value : num; + } : undefined} > {field.type === "select" && field.options ? ( {field.options.map((option) => ( @@ -172,7 +202,6 @@ const GuardrailOptionalParams: React.FC = ({ ) : field.type === "bool" || field.type === "boolean" ? (