diff --git a/AGENTS.md b/AGENTS.md index bfd44304d55..a4be57ed301 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -49,6 +49,17 @@ LiteLLM is a unified interface for 100+ LLMs that: - Test provider-specific functionality thoroughly - Consider adding load tests for performance-critical changes +4. **Code Style**: Follow [Google's Python Style Guide](https://google.github.io/styleguide/pyguide.html) for Python; the project also uses Black, Ruff, and MyPy. + +### Performance and Database + +- **Hot paths:** Do not add duplicate DB queries in proxy/MCP auth or other high-frequency request paths. Reuse already-loaded auth/context (key, team, end-user objects) instead of re-querying. +- **List endpoints:** Avoid N+1 queries when implementing list endpoints (e.g. agents, MCP servers). Use batch loads, `include`/joins, or a single query with needed relations so DB round-trips stay constant. +- **Pre-merge:** For proxy and MCP code, verify that new DB usage in request-handling paths does not introduce duplicate or N+1 queries. + +### Database Migrations +- **Do not auto-create or write `migration.sql` files.** If schema changes are needed or a migration is missing, surface the error to the user and instruct them to run `ci_cd/run_migration.py` to generate new schemas/migrations. + ### MAKING CODE CHANGES FOR THE UI (IGNORE FOR BACKEND) 1. **Tremor is DEPRECATED, do not use Tremor components in new features/changes** @@ -174,8 +185,11 @@ When opening issues or pull requests, follow these templates: 3. **Rate Limits**: Respect provider rate limits in tests 4. **Memory Usage**: Be mindful of memory usage in streaming scenarios 5. **Dependencies**: Keep dependencies minimal and well-justified -6. **UI/Backend Contract Mismatch**: When adding a new entity type to the UI, always check whether the backend endpoint accepts a single value or an array. Match the UI control accordingly (single-select vs. multi-select) to avoid silently dropping user selections -7. **Missing Tests for New Entity Types**: When adding a new entity type (e.g., in `EntityUsage`, `UsageViewSelect`), always add corresponding tests in the existing test files and update any icon/component mocks +6. **Duplicate DB queries in hot paths**: Adding redundant DB calls in MCP auth, proxy auth, or other high-frequency paths causes performance degradation at scale. +7. **N+1 queries in list endpoints**: Looping over a list and performing a separate DB query per item (e.g. loading agents or related entities) — use batch/joined queries instead. +8. **Auto-creating migration.sql**: Do not generate or edit migration SQL files. If migrations are needed, tell the user to run `ci_cd/run_migration.py` to generate new schemas. +9. **UI/Backend Contract Mismatch**: When adding a new entity type to the UI, always check whether the backend endpoint accepts a single value or an array. Match the UI control accordingly (single-select vs. multi-select) to avoid silently dropping user selections +10. **Missing Tests for New Entity Types**: When adding a new entity type (e.g., in `EntityUsage`, `UsageViewSelect`), always add corresponding tests in the existing test files and update any icon/component mocks ## HELPFUL RESOURCES @@ -189,4 +203,5 @@ When opening issues or pull requests, follow these templates: - Check similar provider implementations - Ensure comprehensive test coverage - Update documentation appropriately -- Consider backward compatibility impact \ No newline at end of file +- Consider backward compatibility impact +- For proxy/MCP code, confirm no new duplicate or N+1 DB queries in request paths \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md index 3b597fb8a90..789a09469f0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -91,6 +91,13 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components: - Async/await patterns throughout - Type hints required for all public APIs - **Avoid imports within methods** — place all imports at the top of the file (module-level). Inline imports inside functions/methods make dependencies harder to trace and hurt readability. The only exception is avoiding circular imports where absolutely necessary. +- Follow [Google's Python Style Guide](https://google.github.io/styleguide/pyguide.html) for Python style and structure; the repo uses Black/Ruff/MyPy on top of that. + +### Performance and Database + +- **Hot paths:** Avoid adding duplicate or redundant DB queries in hot request paths (e.g. proxy auth, MCP request path). Each extra query per request multiplies load at scale. Prefer reusing already-fetched auth/context (e.g. key/team/end-user objects) instead of re-querying. +- **List endpoints and N+1:** When implementing list endpoints (e.g. agents, MCP servers), avoid N+1 patterns: do not loop over a list and perform a separate DB query per item. Use batch loads, `include`/joins, or a single query with the needed relations so the number of DB round-trips is constant. +- **Review before merging:** For proxy and MCP code, check that new DB access in request-handling paths does not introduce duplicate or N+1 queries. ### Testing Strategy - Unit tests in `tests/test_litellm/` @@ -103,8 +110,7 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components: - When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select ### Database Migrations -- Prisma handles schema migrations -- Migration files auto-generated with `prisma migrate dev` +- Prisma handles schema migrations. **Do not auto-create or write `migration.sql` files.** If schema changes are needed or a migration is missing, surface the error to the user and instruct them to run `ci_cd/run_migration.py` to generate new schemas/migrations. - Always test migrations against both PostgreSQL and SQLite ### Enterprise Features diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 20f6e6f1d39..7394901a446 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -14,35 +14,28 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import ( - CustomCodeValidationError, - validate_custom_code, -) -from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import ( - get_custom_code_primitives, -) -from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router -from litellm.types.guardrails import ( - PII_ENTITY_CATEGORIES_MAP, - ApplyGuardrailRequest, - ApplyGuardrailResponse, - BaseLitellmParams, - BedrockGuardrailConfigModel, - Guardrail, - GuardrailEventHooks, - GuardrailInfoResponse, - GuardrailUIAddGuardrailSettings, - LakeraV2GuardrailConfigModel, - ListGuardrailsResponse, - LitellmParams, - PatchGuardrailRequest, - PiiAction, - PiiEntityType, - PresidioPresidioConfigModelUserInterface, - SupportedGuardrailIntegrations, - ToolPermissionGuardrailConfigModel, -) + CustomCodeValidationError, validate_custom_code) +from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import \ + get_custom_code_primitives +from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry +from litellm.proxy.guardrails.usage_endpoints import \ + router as guardrails_usage_router +from litellm.types.guardrails import (PII_ENTITY_CATEGORIES_MAP, + ApplyGuardrailRequest, + ApplyGuardrailResponse, + BaseLitellmParams, + BedrockGuardrailConfigModel, Guardrail, + GuardrailEventHooks, + GuardrailInfoResponse, + GuardrailUIAddGuardrailSettings, + LakeraV2GuardrailConfigModel, + ListGuardrailsResponse, LitellmParams, + PatchGuardrailRequest, PiiAction, + PiiEntityType, + PresidioPresidioConfigModelUserInterface, + SupportedGuardrailIntegrations, + ToolPermissionGuardrailConfigModel) #### GUARDRAILS ENDPOINTS #### @@ -161,7 +154,8 @@ async def list_guardrails_v2(): ``` """ 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.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -303,7 +297,8 @@ async def create_guardrail( } ``` """ - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: @@ -401,7 +396,8 @@ async def update_guardrail( } ``` """ - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: @@ -477,7 +473,8 @@ async def delete_guardrail( } ``` """ - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: @@ -579,7 +576,8 @@ async def patch_guardrail( } ``` """ - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: @@ -707,7 +705,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.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client from litellm.types.guardrails import GUARDRAIL_DEFINITION_LOCATION @@ -782,10 +781,8 @@ async def get_guardrail_ui_settings(): - Content filter settings (patterns and categories) """ from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.patterns import ( - PATTERN_CATEGORIES, - get_available_content_categories, - get_pattern_metadata, - ) + PATTERN_CATEGORIES, get_available_content_categories, + get_pattern_metadata) # Convert the PII_ENTITY_CATEGORIES_MAP to the format expected by the UI category_maps = [] @@ -1205,11 +1202,22 @@ def _build_field_dict( # Add options if they exist in json_schema_extra (this takes precedence) if field_json_schema_extra and "options" in field_json_schema_extra: field_dict["options"] = field_json_schema_extra["options"] + elif field_type == "select": + # For Literal types, populate options so the UI can render a dropdown + literal_options = _extract_literal_values(field_annotation) + if literal_options: + field_dict["options"] = literal_options # Add default value if it exists if field.default is not None and field.default is not ...: field_dict["default_value"] = field.default + # Copy min, max, step from json_schema_extra for number/percentage inputs + if field_json_schema_extra: + for key in ("min", "max", "step", "default_value"): + if key in field_json_schema_extra: + field_dict[key] = field_json_schema_extra[key] + return field_dict @@ -1357,7 +1365,8 @@ async def get_provider_specific_params(): } ### get the config model for the guardrail - go through the registry and get the config model for the guardrail - from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry + from litellm.proxy.guardrails.guardrail_registry import \ + guardrail_class_registry for guardrail_name, guardrail_class in guardrail_class_registry.items(): guardrail_config_model = guardrail_class.get_config_model() @@ -1485,6 +1494,7 @@ async def test_custom_code_guardrail( ``` """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: raise HTTPException( status_code=403, diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py new file mode 100644 index 00000000000..c3691d80db5 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py @@ -0,0 +1,92 @@ +"""Block Code Execution guardrail: blocks or masks fenced code blocks by language.""" + +from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union, cast + +from litellm.types.guardrails import (GuardrailEventHooks, + SupportedGuardrailIntegrations) + +from .block_code_execution import BlockCodeExecutionGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + +# Default: run on both request and response (and during_call is supported too) +DEFAULT_EVENT_HOOKS = [ + GuardrailEventHooks.pre_call.value, + GuardrailEventHooks.post_call.value, +] + + +def _get_param( + litellm_params: "LitellmParams", + guardrail: "Guardrail", + key: str, + default: Any = None, +) -> Any: + """Get a param from litellm_params, with fallback to raw guardrail litellm_params (for extra fields not on LitellmParams).""" + value = getattr(litellm_params, key, default) + if value is not None: + return value + raw = guardrail.get("litellm_params") + if isinstance(raw, dict) and key in raw: + return raw[key] + return default + + +def initialize_guardrail( + litellm_params: "LitellmParams", + guardrail: "Guardrail", +) -> BlockCodeExecutionGuardrail: + """Initialize the Block Code Execution guardrail from config.""" + import litellm + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError( + "Block Code Execution guardrail requires a guardrail_name" + ) + + blocked_languages: Optional[List[str]] = cast( + Optional[List[str]], + _get_param(litellm_params, guardrail, "blocked_languages"), + ) + action = cast( + Literal["block", "mask"], + _get_param(litellm_params, guardrail, "action", "block"), + ) + confidence_threshold = float( + cast( + Union[int, float, str], + _get_param(litellm_params, guardrail, "confidence_threshold", 0.5), + ) + ) + mode = _get_param(litellm_params, guardrail, "mode") + event_hook = cast( + Optional[Union[Literal["pre_call", "post_call", "during_call"], List[str]]], + mode if mode is not None else DEFAULT_EVENT_HOOKS, + ) + + instance = BlockCodeExecutionGuardrail( + guardrail_name=guardrail_name, + blocked_languages=blocked_languages, + action=action, + confidence_threshold=confidence_threshold, + event_hook=event_hook, + default_on=bool(_get_param(litellm_params, guardrail, "default_on", False)), + ) + litellm.logging_callback_manager.add_litellm_callback(instance) + return instance + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.BLOCK_CODE_EXECUTION.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.BLOCK_CODE_EXECUTION.value: BlockCodeExecutionGuardrail, +} + +__all__ = [ + "BlockCodeExecutionGuardrail", + "initialize_guardrail", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py new file mode 100644 index 00000000000..53997029418 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py @@ -0,0 +1,363 @@ +""" +Block Code Execution guardrail. + +Detects markdown fenced code blocks in request/response content and blocks or masks them +when the language is in the blocked list (or all blocks when list is empty). Supports +confidence scoring and a tunable threshold (only block when confidence >= threshold). +""" + +import re +from datetime import datetime +from typing import (TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Literal, + Optional, Tuple, Union, cast) + +from fastapi import HTTPException + +from litellm.integrations.custom_guardrail import (CustomGuardrail, + ModifyResponseException) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.base import \ + GuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import ( + CodeBlockActionTaken, CodeBlockDetection) +from litellm.types.utils import (GenericGuardrailAPIInputs, GuardrailStatus, + GuardrailTracingDetail, ModelResponseStream) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import \ + Logging as LiteLLMLoggingObj + +# Language tag aliases (normalize to canonical for comparison) +LANGUAGE_ALIASES: Dict[str, str] = { + "js": "javascript", + "py": "python", +} + +# Tags that indicate non-executable / plain text (lower confidence when block-all) +NON_EXECUTABLE_TAGS: frozenset = frozenset( + {"text", "plaintext", "plain", "markdown", "md", "output", "result"} +) + +# Regex: fenced code block with optional language tag. Handles ```lang\n...\n``` +# Content between fences; does not handle nested ``` inside body (documented edge case). +FENCED_BLOCK_RE = re.compile(r"```(\w*)\n(.*?)```", re.DOTALL) + + +def _normalize_escaped_newlines(text: str) -> str: + """ + Replace literal escaped newlines (backslash + n or backslash + r) with real newlines. + API/JSON payloads sometimes deliver newlines as the two-character sequence \\n. + Applied whenever \\n or \\r appear, including in mixed content with real newlines. + """ + if not text: + return text + if "\\n" not in text and "\\r" not in text: + return text + # Order matters: replace \r\n first so we don't produce extra \n from \r then \n + text = text.replace("\\r\\n", "\n") + text = text.replace("\\n", "\n") + text = text.replace("\\r", "\n") + return text + + +def _normalize_language(tag: str) -> str: + """Normalize language tag (lowercase, resolve aliases).""" + tag = (tag or "").strip().lower() + return LANGUAGE_ALIASES.get(tag, tag) + + +def _is_blocked_language( + tag: str, + blocked_languages: Optional[List[str]], + block_all: bool, +) -> bool: + """True if this language tag should be considered blocked.""" + normalized = _normalize_language(tag) + if block_all: + # Block all: only allow through if it's explicitly non-executable (we still block but with lower confidence) + return True + # When block_all is False, caller guarantees blocked_languages is non-empty. + if not blocked_languages: + return True + normalized_list = [_normalize_language(t) for t in blocked_languages] + return normalized in normalized_list + + +def _confidence_for_block( + tag: str, + block_all: bool, + tag_in_blocked_list: bool, +) -> float: + """Return confidence in [0, 1] for this code block detection.""" + normalized = _normalize_language(tag) + if tag_in_blocked_list: + return 1.0 + if block_all: + # Explicit non-executable tags (e.g. text, plaintext) get lower confidence + if normalized in NON_EXECUTABLE_TAGS: + return 0.5 + # Untagged or other tags in block-all mode: treat as executable, high confidence + return 1.0 + return 0.0 + + +class BlockCodeExecutionGuardrail(CustomGuardrail): + """ + Guardrail that detects fenced code blocks (markdown ```) and blocks or masks them + when the language is in the blocked list (or all when list is empty/None). + Supports confidence threshold: only block when confidence >= confidence_threshold. + """ + + MASK_PLACEHOLDER = "[CODE_BLOCK_REDACTED]" + + def __init__( + self, + guardrail_name: Optional[str] = None, + blocked_languages: Optional[List[str]] = None, + action: Literal["block", "mask"] = "block", + confidence_threshold: float = 0.5, + event_hook: Optional[ + Union[Literal["pre_call", "post_call", "during_call"], List[str]] + ] = None, + default_on: bool = False, + **kwargs: Any, + ) -> None: + # Normalize to type expected by CustomGuardrail + _event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks]]] = ( + None + ) + if event_hook is not None: + if isinstance(event_hook, list): + _event_hook = [ + GuardrailEventHooks(h) if isinstance(h, str) else h + for h in event_hook + ] + else: + _event_hook = GuardrailEventHooks(event_hook) + super().__init__( + guardrail_name=guardrail_name or "block_code_execution", + supported_event_hooks=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.during_call, + ], + event_hook=_event_hook + or [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + default_on=default_on, + **kwargs, + ) + self.blocked_languages = blocked_languages + self.block_all = blocked_languages is None or len(blocked_languages) == 0 + self.action = action + self.confidence_threshold = max(0.0, min(1.0, confidence_threshold)) + + @staticmethod + def get_config_model() -> Optional[type[GuardrailConfigModel]]: + from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import \ + BlockCodeExecutionGuardrailConfigModel + + return BlockCodeExecutionGuardrailConfigModel + + def _find_blocks( + self, text: str + ) -> List[Tuple[int, int, str, str, float, CodeBlockActionTaken]]: + """ + Find all fenced code blocks in text. Returns list of + (start, end, language_tag, block_content, confidence, action_taken). + """ + results: List[ + Tuple[int, int, str, str, float, CodeBlockActionTaken] + ] = [] + for m in FENCED_BLOCK_RE.finditer(text): + tag = (m.group(1) or "").strip() + body = m.group(2) + tag_in_list = not self.block_all and _normalize_language(tag) in [ + _normalize_language(t) for t in (self.blocked_languages or []) + ] + is_blocked = _is_blocked_language( + tag, self.blocked_languages, self.block_all + ) + confidence = _confidence_for_block(tag, self.block_all, tag_in_list) + if not is_blocked: + action_taken: CodeBlockActionTaken = "allow" + elif confidence >= self.confidence_threshold: + action_taken = "block" + else: + action_taken = "log_only" + results.append( + (m.start(), m.end(), tag or "(none)", body, confidence, action_taken) + ) + return results + + def _scan_text( + self, + text: str, + detections: Optional[List[CodeBlockDetection]] = None, + ) -> Tuple[str, bool]: + """ + Scan one text: find blocks, apply block/mask/allow by confidence. + Returns (modified_text, should_raise). + """ + if not text: + return text, False + text = _normalize_escaped_newlines(text) + blocks = self._find_blocks(text) + if not blocks: + return text, False + + should_raise = False + last_end = 0 + parts: List[str] = [] + for start, end, tag, _body, confidence, action_taken in blocks: + if detections is not None: + detections.append( + cast( + CodeBlockDetection, + { + "type": "code_block", + "language": tag, + "confidence": round(confidence, 2), + "action_taken": action_taken, + }, + ) + ) + + if action_taken == "block" and self.action == "block": + should_raise = True + parts.append(text[last_end:start]) + if action_taken == "block": + parts.append(self.MASK_PLACEHOLDER) + else: + parts.append(text[start:end]) + last_end = end + + parts.append(text[last_end:]) + new_text = "".join(parts) + return new_text, should_raise + + def _raise_block_error( + self, language: str, is_output: bool, request_data: dict + ) -> None: + msg = f"Content blocked: executable code block detected (language: {language})" + if is_output: + raise HTTPException( + status_code=400, + detail={ + "error": msg, + "guardrail": self.guardrail_name, + "language": language, + }, + ) + self.raise_passthrough_exception( + violation_message=msg, + request_data=request_data, + detection_info={"language": language}, + ) + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + start_time = datetime.now() + detections: List[CodeBlockDetection] = [] + status: GuardrailStatus = "success" + exception_str = "" + + try: + texts = inputs.get("texts", []) + if not texts: + return inputs + + is_output = input_type == "response" + processed: List[str] = [] + for text in texts: + new_text, should_raise = self._scan_text(text, detections) + processed.append(new_text) + if should_raise: + # Determine language from first blocking detection + lang = "unknown" + for d in detections: + if d.get("action_taken") == "block": + lang = d.get("language", "unknown") + break + self._raise_block_error(lang, is_output, request_data) + + inputs["texts"] = processed + return inputs + except HTTPException: + status = "guardrail_intervened" + raise + except ModifyResponseException: + status = "guardrail_intervened" + raise + except Exception as e: + status = "guardrail_failed_to_respond" + exception_str = str(e) + raise + finally: + guardrail_response: Union[List[dict], str] = [dict(d) for d in detections] + if status != "success" and not detections: + guardrail_response = exception_str + max_confidence: Optional[float] = None + for d in detections: + c = d.get("confidence") + if c is not None and (max_confidence is None or c > max_confidence): + max_confidence = c + tracing_kw: Dict[str, Any] = { + "guardrail_id": self.guardrail_name, + "detection_method": "fenced_code_block", + "match_details": guardrail_response, + } + if max_confidence is not None: + tracing_kw["confidence_score"] = max_confidence + event_type = ( + GuardrailEventHooks.pre_call + if input_type == "request" + else GuardrailEventHooks.post_call + ) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider="block_code_execution", + guardrail_json_response=guardrail_response, + request_data=request_data, + guardrail_status=status, + start_time=start_time.timestamp(), + end_time=datetime.now().timestamp(), + duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, + tracing_detail=GuardrailTracingDetail(**tracing_kw), # type: ignore[typeddict-item] + ) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: Any, + response: Any, + request_data: dict, + ) -> AsyncGenerator[ModelResponseStream, None]: + """Accumulate streamed content and block as soon as a complete fenced code block is detected (before yielding that chunk).""" + accumulated = "" + async for item in response: + if isinstance(item, ModelResponseStream) and item.choices: + delta_content = "" + for choice in item.choices: + if hasattr(choice, "delta") and choice.delta: + content = getattr(choice.delta, "content", None) + if content and isinstance(content, str): + delta_content += content + accumulated += delta_content + # Check after every chunk so we block before yielding the chunk that completes a blocked block + normalized = _normalize_escaped_newlines(accumulated) + blocks = self._find_blocks(normalized) + for _start, _end, _tag, _body, confidence, action_taken in blocks: + if ( + action_taken == "block" + and confidence >= self.confidence_threshold + ): + lang = _tag or "unknown" + self._raise_block_error(lang, True, request_data) + yield item diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index b9c99eaabfb..d38300f8b0d 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -5,24 +5,20 @@ from typing import Any, Dict, List, Literal, Optional, Union from pydantic import BaseModel, ConfigDict, Field, field_validator from typing_extensions import Required, TypedDict -from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import ( - EnkryptAIGuardrailConfigs, -) -from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import ( - GraySwanGuardrailConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.ibm import ( - IBMGuardrailsBaseConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( - ContentFilterCategoryConfig, -) -from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import ( - QualifireGuardrailConfigModel, -) -from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( - ToolPermissionGuardrailConfigModel, -) +from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import \ + BlockCodeExecutionGuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import \ + EnkryptAIGuardrailConfigs +from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import \ + GraySwanGuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.ibm import \ + IBMGuardrailsBaseConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import \ + ContentFilterCategoryConfig +from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import \ + QualifireGuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import \ + ToolPermissionGuardrailConfigModel """ Pydantic object defining how to set guardrails on litellm proxy @@ -73,6 +69,7 @@ class SupportedGuardrailIntegrations(Enum): CUSTOM_CODE = "custom_code" SEMANTIC_GUARD = "semantic_guard" MCP_END_USER_PERMISSION = "mcp_end_user_permission" + BLOCK_CODE_EXECUTION = "block_code_execution" class Role(Enum): @@ -707,6 +704,7 @@ class LitellmParams( EnkryptAIGuardrailConfigs, IBMGuardrailsBaseConfigModel, QualifireGuardrailConfigModel, + BlockCodeExecutionGuardrailConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") mode: Union[str, List[str], Mode] = Field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py b/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py new file mode 100644 index 00000000000..0db706fd20c --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py @@ -0,0 +1,75 @@ +"""Types for the Block Code Execution guardrail.""" + +from typing import Any, List, Literal, Optional, TypedDict, cast + +from pydantic import Field + +from .base import GuardrailConfigModel + +CodeBlockActionTaken = Literal["block", "allow", "log_only"] + +# Supported language tags for the blocked_languages multiselect dropdown +BLOCKED_LANGUAGES_OPTIONS = [ + "python", + "javascript", + "js", + "bash", + "sh", + "ruby", + "go", + "java", + "csharp", + "php", + "c", + "cpp", + "rust", + "sql", +] + + +class CodeBlockDetection(TypedDict, total=False): + """Detection output for a single fenced code block (for tracing/logging).""" + + type: Literal["code_block"] + language: str + confidence: float + action_taken: CodeBlockActionTaken + evidence: Optional[str] + snippet: Optional[str] + + +class BlockCodeExecutionGuardrailConfigModel(GuardrailConfigModel): + """Configuration for the Block Code Execution guardrail.""" + + blocked_languages: Optional[List[str]] = Field( + default=None, + description="Language tags to block (e.g. python, javascript, bash). Empty or None = block all fenced code blocks.", + json_schema_extra=cast( + Any, + {"type": "multiselect", "options": BLOCKED_LANGUAGES_OPTIONS}, + ), + ) + action: Literal["block", "mask"] = Field( + default="block", + description="'block' raises an error; 'mask' replaces the code block with a placeholder.", + ) + confidence_threshold: float = Field( + default=0.5, + ge=0.0, + le=1.0, + description="Only block or mask when detection confidence >= this value; below threshold, allow or log_only.", + json_schema_extra=cast( + Any, + { + "type": "percentage", + "min": 0.0, + "max": 1.0, + "step": 0.1, + "default_value": 0.5, + }, + ), + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Block Code Execution" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py new file mode 100644 index 00000000000..0fcafc5b9ed --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py @@ -0,0 +1,429 @@ +"""Tests for the Block Code Execution guardrail.""" + +import pytest +from fastapi import HTTPException + +from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.proxy.guardrails.guardrail_hooks.block_code_execution import ( + DEFAULT_EVENT_HOOKS, BlockCodeExecutionGuardrail, initialize_guardrail) +from litellm.proxy.guardrails.guardrail_hooks.block_code_execution.block_code_execution import \ + _normalize_escaped_newlines +from litellm.types.guardrails import GuardrailEventHooks + + +class TestBlockCodeExecutionGuardrail: + """Test BlockCodeExecutionGuardrail detection and actions.""" + + def test_detects_python_block_when_in_blocked_list(self): + """Text with ```python block is detected when python is in blocked_languages.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + confidence_threshold=0.7, + ) + blocks = guardrail._find_blocks("Here is code:\n```python\nprint(1)\n```\nDone.") + assert len(blocks) == 1 + _start, _end, tag, _body, confidence, action_taken = blocks[0] + assert tag == "python" + assert confidence == 1.0 + assert action_taken == "block" + + def test_block_all_when_blocked_languages_empty(self): + """When blocked_languages is empty, any fenced block is blocked (block all).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=[], + confidence_threshold=0.7, + ) + blocks = guardrail._find_blocks("```\nfoo\n```") + assert len(blocks) == 1 + _start, _end, _tag, _body, confidence, action_taken = blocks[0] + assert action_taken == "block" + assert confidence in (0.5, 1.0) + + def test_no_block_when_language_not_in_list(self): + """When language is not in blocked_languages, block is not triggered.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + confidence_threshold=0.7, + ) + blocks = guardrail._find_blocks("```text\nplain output\n```") + assert len(blocks) == 1 + _start, _end, _tag, _body, confidence, action_taken = blocks[0] + assert action_taken == "allow" + assert confidence == 0.0 + + def test_confidence_below_threshold_allows(self): + """When confidence < confidence_threshold, action_taken is log_only and we do not block.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=[], # block all + confidence_threshold=0.9, + ) + # Block with no tag or plaintext tag gets confidence 0.5 + blocks = guardrail._find_blocks("```text\nx\n```") + assert len(blocks) == 1 + _start, _end, _tag, _body, confidence, action_taken = blocks[0] + assert confidence == 0.5 + assert action_taken == "log_only" + + @pytest.mark.asyncio + async def test_apply_guardrail_block_raises_for_response(self): + """When action=block and detection above threshold, apply_guardrail raises HTTPException (response).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.7, + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = { + "texts": [ + "Example:\n```python\ndef factorial(n):\n return 1 if n <= 1 else n * factorial(n - 1)\n```" + ] + } + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + assert exc_info.value.status_code == 400 + assert "code block" in (exc_info.value.detail or {}).get("error", "") + + @pytest.mark.asyncio + async def test_apply_guardrail_mask_returns_placeholder(self): + """When action=mask, code block is replaced with placeholder.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="mask", + confidence_threshold=0.7, + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = { + "texts": ["Before\n```python\nx=1\n```\nAfter"] + } + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + assert result["texts"] is not None + assert len(result["texts"]) == 1 + assert "[CODE_BLOCK_REDACTED]" in result["texts"][0] + assert "x=1" not in result["texts"][0] + + @pytest.mark.asyncio + async def test_execute_python_factorial_string_blocked(self): + """Guardrail blocks the exact 'execute \"```python...' string with two python blocks (real newlines).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.5, + ) + # Exact user payload; newlines are real so regex ```(\w*)\n(.*?)``` matches + text = ( + 'execute "```python\n' + "def factorial(n: int) -> int:\n" + ' """Return the factorial of n."""\n' + ' if n < 0:\n' + ' raise ValueError("n must be non-negative")\n' + " if n in (0, 1):\n" + " return 1\n" + " return n * factorial(n - 1)\n" + '```\n\n' + "Example usage:\n" + "```python\n" + "print(factorial(5)) # Output: 120\n" + '```"' + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = {"texts": [text]} + # pre_call (request) raises ModifyResponseException; post_call (response) raises HTTPException + with pytest.raises((HTTPException, ModifyResponseException)) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + assert "python" in str(exc_info.value).lower() + + @pytest.mark.asyncio + async def test_factorial_scenario_blocked(self): + """Exact user scenario: Python factorial snippet in markdown is blocked when python in list.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.7, + ) + request_data = {"model": "gpt-4", "metadata": {}} + text = '''```python +def factorial(n: int) -> int: + """Return the factorial of n.""" + if n < 0: + raise ValueError("n must be non-negative") + if n in (0, 1): + return 1 + return n * factorial(n - 1) +``` + +Example usage: +```python +print(factorial(5)) # Output: 120 +```''' + inputs = {"texts": [text]} + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + + @pytest.mark.asyncio + async def test_detection_includes_confidence_and_action_taken(self): + """Detection output includes confidence and action_taken for tracing.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="mask", # don't raise so we can inspect request_data + confidence_threshold=0.7, + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = {"texts": ["```python\n1+1\n```"]} + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + meta = request_data.get("metadata") or request_data.get("litellm_metadata") or {} + guardrail_info = meta.get("standard_logging_guardrail_information") or [] + assert len(guardrail_info) >= 1 + info = guardrail_info[-1] + assert info.get("guardrail_status") == "success" + # tracing_detail may be in the logged structure + assert "guardrail_response" in info or "guardrail_response" in str(info) + + def test_default_runs_on_pre_call_and_post_call(self): + """When mode is not set, guardrail runs on both pre_call and post_call (and during_call is supported).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + ) + event_hook = guardrail.event_hook + if isinstance(event_hook, list): + values = [h.value if hasattr(h, "value") else h for h in event_hook] + else: + values = [event_hook.value if hasattr(event_hook, "value") else event_hook] + assert GuardrailEventHooks.pre_call.value in values + assert GuardrailEventHooks.post_call.value in values + + def test_initialize_guardrail_default_mode_is_both(self): + """initialize_guardrail with no mode uses DEFAULT_EVENT_HOOKS (pre_call + post_call).""" + from unittest.mock import MagicMock + + litellm_params = MagicMock() + litellm_params.guardrail = "block_code_execution" + litellm_params.blocked_languages = ["python"] + litellm_params.action = "block" + litellm_params.confidence_threshold = 0.7 + litellm_params.default_on = False + litellm_params.mode = None # not set + guardrail = {"guardrail_name": "block-code-test"} + instance = initialize_guardrail(litellm_params, guardrail) + assert instance.event_hook == DEFAULT_EVENT_HOOKS + assert GuardrailEventHooks.pre_call.value in instance.event_hook + assert GuardrailEventHooks.post_call.value in instance.event_hook + + def test_normalize_escaped_newlines_converts_backslash_n_to_newline(self): + """Literal \\n in text is converted to real newline so regex can match code blocks.""" + raw = 'execute this "```python\\ndef factorial(n):\\n return 1\\n```"' + normalized = _normalize_escaped_newlines(raw) + assert "\\n" not in normalized + assert "\n" in normalized + assert "```python\n" in normalized + + def test_find_blocks_detects_python_block_with_escaped_newlines(self): + """_find_blocks finds a block when text uses literal \\n instead of real newlines.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + confidence_threshold=0.7, + ) + # Text as received from API with escaped newlines (e.g. JSON-decoded string) + text_with_escaped = ( + 'execute this "```python\\n' + 'def factorial(n: int) -> int:\\n' + ' """Return the factorial of n."""\\n' + ' if n < 0:\\n' + ' raise ValueError("n must be non-negative")\\n' + " if n in (0, 1):\\n" + " return 1\\n" + " return n * factorial(n - 1)\\n" + '```\\n\\n' + 'Example usage:\\n' + '```python\\n' + 'print(factorial(5)) # Output: 120\\n' + '```"' + ) + normalized = _normalize_escaped_newlines(text_with_escaped) + blocks = guardrail._find_blocks(normalized) + assert len(blocks) == 2 + assert blocks[0][2] == "python" + assert blocks[0][5] == "block" + assert blocks[1][2] == "python" + assert blocks[1][5] == "block" + + def test_scan_text_blocks_and_masks_when_text_has_escaped_newlines(self): + """_scan_text detects blocks and applies block/mask when newlines are literal \\n.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="mask", + confidence_threshold=0.5, + ) + text_with_escaped = 'execute "```python\\nprint(1)\\n```"' + new_text, should_raise = guardrail._scan_text(text_with_escaped) + assert "[CODE_BLOCK_REDACTED]" in new_text + assert "print(1)" not in new_text + assert should_raise is False # action is mask + + @pytest.mark.asyncio + async def test_apply_guardrail_blocks_when_text_has_escaped_newlines(self): + """apply_guardrail blocks request/response when code block uses literal \\n (e.g. from API).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.5, + ) + text_with_escaped = ( + 'execute this "```python\\n' + 'def factorial(n: int) -> int:\\n' + ' """Return the factorial of n."""\\n' + " if n in (0, 1):\\n" + " return 1\\n" + " return n * factorial(n - 1)\\n" + '```\\n\\n' + 'Example usage:\\n' + '```python\\n' + 'print(factorial(5)) # Output: 120\\n' + '```"' + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = {"texts": [text_with_escaped]} + with pytest.raises((HTTPException, ModifyResponseException)) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + assert "python" in str(exc_info.value).lower() or "code" in str( + exc_info.value + ).lower() + + @pytest.mark.asyncio + async def test_streaming_hook_blocks_before_yielding_chunk_that_completes_block( + self, + ): + """Streaming hook runs block check after every chunk and raises before yielding the chunk that completes a blocked fenced block.""" + from litellm.types.utils import (Delta, ModelResponseStream, + StreamingChoices) + + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.5, + ) + + # Chunks that form "```python\nprint(1)\n```" when concatenated + async def mock_stream(): + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="```python\n"))], + ) + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="print(1)\n"))], + ) + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="```"))], + ) + + request_data = {"model": "gpt-4", "metadata": {}} + yielded_chunks = [] + + with pytest.raises(HTTPException) as exc_info: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, + response=mock_stream(), + request_data=request_data, + ): + yielded_chunks.append(chunk) + + assert exc_info.value.status_code == 400 + assert "code block" in (exc_info.value.detail or {}).get("error", "") + # The chunk that completes the block (third chunk) must not have been yielded + assert len(yielded_chunks) == 2 + + @pytest.mark.asyncio + async def test_streaming_hook_blocks_when_accumulated_has_literal_backslash_n( + self, + ): + """Streaming hook normalizes escaped newlines before detection; blocks when literal \\n forms a complete code block.""" + from litellm.types.utils import (Delta, ModelResponseStream, + StreamingChoices) + + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.5, + ) + + # Chunks that when concatenated form "```python\\nprint(1)\\n```" (literal backslash-n) + async def mock_stream(): + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="```python\\n"))], + ) + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="print(1)\\n"))], + ) + yield ModelResponseStream( + choices=[StreamingChoices(delta=Delta(content="```"))], + ) + + request_data = {"model": "gpt-4", "metadata": {}} + yielded_chunks = [] + + with pytest.raises(HTTPException) as exc_info: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, + response=mock_stream(), + request_data=request_data, + ): + yielded_chunks.append(chunk) + + assert exc_info.value.status_code == 400 + assert "code block" in (exc_info.value.detail or {}).get("error", "") + # Block is detected after normalization; chunk that completes the block not yielded + assert len(yielded_chunks) == 2 + + def test_normalize_escaped_newlines_mixed_content_detects_block(self): + """Mixed content (real newlines and literal \\n) is normalized so code block is detected.""" + # Text with real newline then a fence using literal \n after language tag + mixed = "line1\n```py\\nprint(1)\\n```" + normalized = _normalize_escaped_newlines(mixed) + assert "```py\n" in normalized + assert "print(1)\n" in normalized + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python", "py"], + confidence_threshold=0.5, + ) + blocks = guardrail._find_blocks(normalized) + assert len(blocks) == 1 + assert blocks[0][2] == "py" + assert blocks[0][5] == "block" diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index f69d9cf0c47..860093ceadb 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -2,14 +2,7 @@ import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; import { Select, Switch, Tooltip } from "antd"; -import { - Table, - TableHead, - TableHeaderCell, - TableBody, - TableRow, - TableCell, -} from "@tremor/react"; +import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react"; import { TimeCell } from "./view_logs/time_cell"; import { TableHeaderSortDropdown } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; import type { SortState } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; @@ -17,14 +10,13 @@ import FilterComponent, { FilterOption } from "./molecules/filter"; import { fetchToolsList, updateToolPolicy, ToolRow } from "./networking"; const POLICY_OPTIONS = [ - { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, - { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, + { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, + { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, ] as const; type PolicyValue = "trusted" | "blocked"; -const policyStyle = (p: string) => - POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; +const policyStyle = (p: string) => POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at" | "call_count"; @@ -56,18 +48,6 @@ const PolicySelect: React.FC<{ minWidth: 110, fontWeight: 500, }} - styles={{ - selector: { - backgroundColor: style.bg, - borderColor: style.border, - color: style.color, - borderRadius: 999, - fontSize: 11, - fontWeight: 600, - paddingLeft: 8, - paddingRight: 4, - }, - }} popupMatchSelectWidth={false} options={POLICY_OPTIONS.map((o) => ({ value: o.value, @@ -133,7 +113,9 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { } }, [accessToken]); - useEffect(() => { load(); }, [load]); + useEffect(() => { + load(); + }, [load]); useEffect(() => { if (!isLiveTail) return; @@ -146,9 +128,7 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { setSaving(toolName); try { await updateToolPolicy(accessToken, toolName, newPolicy); - setTools((prev) => - prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t)) - ); + setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t))); } catch (e: any) { alert(`Failed to update policy: ${e.message}`); } finally { @@ -178,12 +158,14 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { }; // Build unique team/key options from loaded data - const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map( - (v) => ({ label: v as string, value: v as string }) - ); - const keyAliasOptions = Array.from(new Set(tools.map((t) => t.key_alias).filter(Boolean))).map( - (v) => ({ label: v as string, value: v as string }) - ); + const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map((v) => ({ + label: v as string, + value: v as string, + })); + const keyAliasOptions = Array.from(new Set(tools.map((t) => t.key_alias).filter(Boolean))).map((v) => ({ + label: v as string, + value: v as string, + })); const filterOptions: FilterOption[] = [ { @@ -245,7 +227,6 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {

Tool Policies

- {/* Toolbar */}
@@ -256,16 +237,29 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { placeholder="Search by Tool Name" className="w-full px-3 py-2 pl-8 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500" value={searchTerm} - onChange={(e) => { setSearchTerm(e.target.value); setCurrentPage(1); }} + onChange={(e) => { + setSearchTerm(e.target.value); + setCurrentPage(1); + }} /> - - + +
Live Tail - +
@@ -282,14 +286,27 @@ export const ToolPolicies: React.FC = ({ accessToken }) => {
- Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results + Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} -{" "} + {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results + + + Page {currentPage} of {totalPages} - Page {currentPage} of {totalPages}
- - + +
@@ -309,7 +326,9 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { {isLiveTail && (
Auto-refreshing every 15 seconds - +
)} @@ -321,20 +340,34 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { - - - - - + + + + + + + + + + + + + + + Key Hash - + + + Origin {loading ? ( - Loading tools… + + Loading tools… + ) : paginated.length === 0 ? ( @@ -397,12 +430,25 @@ export const ToolPolicies: React.FC = ({ accessToken }) => { {/* Bottom pagination (only when > 1 page) */} {totalPages > 1 && (
- Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length} + + Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of{" "} + {sorted.length} +
- - + +
)} diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx index 9a603252286..410d7510171 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx @@ -713,7 +713,7 @@ describe("UsagePage", () => { // Admin should see the user selector select element with the placeholder attribute const userSelects = screen.getAllByRole("combobox"); const userSelect = userSelects.find( - (el) => el.getAttribute("placeholder") === "All Users (Global View)", + (el) => el.getAttribute("placeholder") === "Select user to filter...", ); expect(userSelect).toBeDefined(); }); @@ -828,7 +828,7 @@ describe("UsagePage", () => { // Non-admin should not see the user selector const userSelects = screen.getAllByRole("combobox"); const userSelect = userSelects.find( - (el) => el.getAttribute("placeholder") === "All Users (Global View)", + (el) => el.getAttribute("placeholder") === "Select user to filter...", ); expect(userSelect).toBeUndefined(); }); 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 183763667f9..f8c8cc17aef 100644 --- a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx @@ -166,12 +166,16 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a // Set provider setSelectedProvider(preset.provider); - form.setFieldsValue({ + const baseValues: Record = { provider: preset.provider, guardrail_name: preset.guardrailNameSuggestion, mode: preset.mode, default_on: preset.defaultOn, - }); + }; + if (preset.provider === "BlockCodeExecution") { + baseValues.confidence_threshold = 0.5; + } + form.setFieldsValue(baseValues); // Pre-select content category if specified if (preset.categoryName && guardrailSettings.content_filter_settings?.content_categories) { @@ -195,11 +199,15 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a const handleProviderChange = (value: string) => { setSelectedProvider(value); // Reset form fields that are provider-specific - form.setFieldsValue({ + const resetValues: Record = { config: undefined, presidio_analyzer_api_base: undefined, presidio_anonymizer_api_base: undefined, - }); + }; + if (value === "BlockCodeExecution") { + resetValues.confidence_threshold = 0.5; + } + form.setFieldsValue(resetValues); // Reset PII selections when changing provider setSelectedEntities([]); diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_configs.ts b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_configs.ts index a6a142bdc37..7a1b5314d33 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_configs.ts +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_configs.ts @@ -148,6 +148,18 @@ export const GUARDRAIL_PRESETS: Record = { mode: "pre_call", defaultOn: false, }, + block_code_execution: { + provider: "BlockCodeExecution", + guardrailNameSuggestion: "Block Code Execution", + mode: "pre_call", + defaultOn: false, + }, + cf_competitor_intent: { + provider: "LitellmContentFilter", + guardrailNameSuggestion: "Competitor Name Blocking", + mode: "pre_call", + defaultOn: false, + }, // ── Partner Guardrails ── presidio: { diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_data.ts b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_data.ts index 2535bd19503..53ccb32c184 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_data.ts +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_garden_data.ts @@ -213,6 +213,24 @@ export const LITELLM_CONTENT_FILTER_CARDS: GuardrailCardInfo[] = [ logo: `${ASSET_PREFIX}litellm_logo.jpg`, tags: ["Keywords", "Blocklist"], }, + { + id: "block_code_execution", + name: "Block Code Execution", + description: "Detects markdown fenced code blocks in requests and responses. Block or mask executable code (e.g. Python, JavaScript, Bash) by language with configurable confidence.", + category: "litellm", + subcategory: "Code Safety", + logo: `${ASSET_PREFIX}litellm_logo.jpg`, + tags: ["Code", "Safety", "Prompt Injection"], + }, + { + id: "cf_competitor_intent", + name: "Competitor Name Blocking", + description: "Block or reframe competitor comparison and ranking intent. Detect when users ask to compare or recommend competitors (airline or generic competitor lists).", + category: "litellm", + subcategory: "Content Category", + logo: `${ASSET_PREFIX}litellm_logo.jpg`, + tags: ["Content Category", "Competitor", "Topic Blocker"], + }, ]; export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [ diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx index 040ef518e5e..d957be4306b 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx @@ -47,6 +47,7 @@ export const guardrail_provider_map: Record = { Lakera: "lakera_v2", LitellmContentFilter: "litellm_content_filter", ToolPermission: "tool_permission", + BlockCodeExecution: "block_code_execution", }; // Function to populate provider map from API response - updates the original map diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_provider_fields.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_provider_fields.tsx index e2077518ab3..2bc381c8e8f 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_provider_fields.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_provider_fields.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Form, Select, Spin, Input } from "antd"; +import { Form, Select, Spin, Input, Slider } from "antd"; import { guardrail_provider_map, populateGuardrailProviders, @@ -20,12 +20,15 @@ interface ProviderParam { param: string; description: string; required: boolean; - default_value?: string; + default_value?: string | number; options?: string[]; type?: string; fields?: { [key: string]: ProviderParam }; dict_key_options?: string[]; dict_value_type?: string; + min?: number; + max?: number; + step?: number; } interface ProviderParamsResponse { @@ -154,6 +157,11 @@ const GuardrailProviderFields: React.FC = ({ ); } + const percentageInitialValue = + field.type === "percentage" && (fieldValue === undefined || fieldValue === null) + ? (field.default_value ?? 0.5) + : undefined; + return ( = ({ label={fieldKey} tooltip={field.description} rules={field.required ? [{ required: true, message: `${fieldKey} is required` }] : undefined} + initialValue={percentageInitialValue} > {field.type === "select" && field.options ? ( + ) : field.type === "percentage" && field.min != null && field.max != null ? ( + ) : field.type === "number" ? ( ({ soft_budget: null, }), fetchMCPAccessGroups: vi.fn().mockResolvedValue([]), + getAgentsList: vi.fn().mockResolvedValue([]), })); vi.mock("../molecules/notifications_manager", () => ({ diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 714507da1aa..fc80c394ae6 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -8,10 +8,10 @@ import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, Tex import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd"; import debounce from "lodash/debounce"; import React, { useCallback, useEffect, useState } from "react"; -import { CopyToClipboard } from "react-copy-to-clipboard"; import { rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; import { mapDisplayToInternalNames } from "../callback_info_helpers"; +import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; import SchemaFormFields from "../common_components/check_openapi_schema"; import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; @@ -20,7 +20,6 @@ import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSel import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings"; import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion"; -import AccessGroupSelector from "../common_components/AccessGroupSelector"; import TeamDropdown from "../common_components/team_dropdown"; import { CreateUserButton } from "../CreateUserButton"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; @@ -40,10 +39,10 @@ import { proxyBaseUrl, userFilterUICall, } from "../networking"; +import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; import NumericalInput from "../shared/numerical_input"; import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; import { simplifyKeyGenerateError } from "./utils"; -import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; const { Option } = Select; @@ -299,7 +298,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { formValues.user_id = userID; } else if (keyOwner === "agent") { if (!selectedAgentId) { - message.error("Please select an agent"); + NotificationsManager.error("Please select an agent"); return; } formValues.agent_id = selectedAgentId; @@ -559,7 +558,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { You Service Account {userRole === "Admin" && Another User} - Agent New + + Agent New + @@ -1005,9 +1006,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { style={{ width: "100%" }} disabled={!premiumUser} placeholder={ - !premiumUser - ? "Premium feature - Upgrade to set policies by key" - : "Select or enter policies" + !premiumUser ? "Premium feature - Upgrade to set policies by key" : "Select or enter policies" } options={policiesList.map((name) => ({ value: name, label: name }))} /> @@ -1059,9 +1058,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { className="mt-4" help="Select access groups to assign to this key" > - + = ({ team, teams, data, addKey }) => { accessToken={accessToken || ""} value={routerSettings || undefined} onChange={setRouterSettings} - modelData={userModels.length > 0 ? { data: userModels.map((model) => ({ model_name: model })) } : undefined} + modelData={ + userModels.length > 0 + ? { data: userModels.map((model) => ({ model_name: model })) } + : undefined + } />