mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Litellm dev 02 23 2026 p2 (#22118)
* feat(add-new-block_code_execution-guardrail): prevent agent from executing code * feat: working block code execution guardrail * feat(ui/): add new guardrails to guardrail_garden * fix: fix greptile feedback * feat: minor cleanup * docs: cleanup agent instructions * feat: cleanup ui * Update litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix: remove unused code * fixing build * fix: greptile fixes * fix: fix linting errors * fix: address greptile comments * fixing build and tests --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: yuneng-jiang <yuneng.jiang@gmail.com>
This commit is contained in:
parent
edd00c025c
commit
21ba9bfe89
17 changed files with 1237 additions and 142 deletions
21
AGENTS.md
21
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
|
||||
- Consider backward compatibility impact
|
||||
- For proxy/MCP code, confirm no new duplicate or N+1 DB queries in request paths
|
||||
10
CLAUDE.md
10
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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -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"
|
||||
|
|
@ -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<ToolPoliciesProps> = ({ accessToken }) => {
|
|||
}
|
||||
}, [accessToken]);
|
||||
|
||||
useEffect(() => { load(); }, [load]);
|
||||
useEffect(() => {
|
||||
load();
|
||||
}, [load]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isLiveTail) return;
|
||||
|
|
@ -146,9 +128,7 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ 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<ToolPoliciesProps> = ({ 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<ToolPoliciesProps> = ({ accessToken }) => {
|
|||
<div className="p-6 w-full">
|
||||
<h1 className="text-2xl font-semibold text-gray-900 mb-6">Tool Policies</h1>
|
||||
<div className="bg-white rounded-lg shadow w-full max-w-full box-border">
|
||||
|
||||
{/* Toolbar */}
|
||||
<div className="border-b px-6 py-4 w-full max-w-full box-border">
|
||||
<div className="flex flex-col md:flex-row items-start md:items-center justify-between space-y-4 md:space-y-0 w-full max-w-full box-border">
|
||||
|
|
@ -256,16 +237,29 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ 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);
|
||||
}}
|
||||
/>
|
||||
<svg className="absolute left-2.5 top-2.5 h-4 w-4 text-gray-500" fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
||||
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0z" />
|
||||
<svg
|
||||
className="absolute left-2.5 top-2.5 h-4 w-4 text-gray-500"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
viewBox="0 0 24 24"
|
||||
>
|
||||
<path
|
||||
strokeLinecap="round"
|
||||
strokeLinejoin="round"
|
||||
strokeWidth={2}
|
||||
d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0z"
|
||||
/>
|
||||
</svg>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center gap-2">
|
||||
<span className="text-sm font-medium text-gray-900">Live Tail</span>
|
||||
<Switch color="green" checked={isLiveTail} onChange={setIsLiveTail} />
|
||||
<Switch checked={isLiveTail} onChange={setIsLiveTail} />
|
||||
</div>
|
||||
|
||||
<button
|
||||
|
|
@ -273,8 +267,18 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
|
|||
disabled={isButtonLoading}
|
||||
className="flex items-center gap-1.5 px-3 py-2 text-sm border rounded-md hover:bg-gray-50 disabled:opacity-60"
|
||||
>
|
||||
<svg className={`w-4 h-4 ${isButtonLoading ? "animate-spin" : ""}`} fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
||||
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15" />
|
||||
<svg
|
||||
className={`w-4 h-4 ${isButtonLoading ? "animate-spin" : ""}`}
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
viewBox="0 0 24 24"
|
||||
>
|
||||
<path
|
||||
strokeLinecap="round"
|
||||
strokeLinejoin="round"
|
||||
strokeWidth={2}
|
||||
d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15"
|
||||
/>
|
||||
</svg>
|
||||
{isButtonLoading ? "Fetching" : "Fetch"}
|
||||
</button>
|
||||
|
|
@ -282,14 +286,27 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
|
|||
|
||||
<div className="flex items-center gap-4 text-sm text-gray-600 whitespace-nowrap">
|
||||
<span>
|
||||
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
|
||||
</span>
|
||||
<span>
|
||||
Page {currentPage} of {totalPages}
|
||||
</span>
|
||||
<span>Page {currentPage} of {totalPages}</span>
|
||||
<div className="flex gap-1">
|
||||
<button onClick={() => setCurrentPage((p) => Math.max(1, p - 1))} disabled={currentPage === 1}
|
||||
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40">Previous</button>
|
||||
<button onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))} disabled={currentPage === totalPages}
|
||||
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40">Next</button>
|
||||
<button
|
||||
onClick={() => setCurrentPage((p) => Math.max(1, p - 1))}
|
||||
disabled={currentPage === 1}
|
||||
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40"
|
||||
>
|
||||
Previous
|
||||
</button>
|
||||
<button
|
||||
onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))}
|
||||
disabled={currentPage === totalPages}
|
||||
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40"
|
||||
>
|
||||
Next
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
|
@ -309,7 +326,9 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
|
|||
{isLiveTail && (
|
||||
<div className="bg-green-50 border-b border-green-100 px-6 py-2 flex items-center justify-between">
|
||||
<span className="text-sm text-green-700">Auto-refreshing every 15 seconds</span>
|
||||
<button onClick={() => setIsLiveTail(false)} className="text-xs text-green-600 underline">Stop</button>
|
||||
<button onClick={() => setIsLiveTail(false)} className="text-xs text-green-600 underline">
|
||||
Stop
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
|
|
@ -321,20 +340,34 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
|
|||
<Table className="[&_td]:py-0.5 [&_th]:py-1 w-full">
|
||||
<TableHead>
|
||||
<TableRow>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Discovered" field="created_at" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Tool Name" field="tool_name" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Policy" field="call_policy" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="# Calls" field="call_count" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Team Name" field="team_id" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">
|
||||
<SortHeader label="Discovered" field="created_at" />
|
||||
</TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">
|
||||
<SortHeader label="Tool Name" field="tool_name" />
|
||||
</TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">
|
||||
<SortHeader label="Policy" field="call_policy" />
|
||||
</TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">
|
||||
<SortHeader label="# Calls" field="call_count" />
|
||||
</TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">
|
||||
<SortHeader label="Team Name" field="team_id" />
|
||||
</TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">Key Hash</TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Key Name" field="key_alias" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">
|
||||
<SortHeader label="Key Name" field="key_alias" />
|
||||
</TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">Origin</TableHeaderCell>
|
||||
</TableRow>
|
||||
</TableHead>
|
||||
<TableBody>
|
||||
{loading ? (
|
||||
<TableRow>
|
||||
<TableCell colSpan={8} className="h-8 text-center text-gray-500">Loading tools…</TableCell>
|
||||
<TableCell colSpan={8} className="h-8 text-center text-gray-500">
|
||||
Loading tools…
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
) : paginated.length === 0 ? (
|
||||
<TableRow>
|
||||
|
|
@ -397,12 +430,25 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
|
|||
{/* Bottom pagination (only when > 1 page) */}
|
||||
{totalPages > 1 && (
|
||||
<div className="border-t px-6 py-3 flex items-center justify-between text-sm text-gray-600">
|
||||
<span>Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length}</span>
|
||||
<span>
|
||||
Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of{" "}
|
||||
{sorted.length}
|
||||
</span>
|
||||
<div className="flex gap-1">
|
||||
<button onClick={() => setCurrentPage((p) => Math.max(1, p - 1))} disabled={currentPage === 1}
|
||||
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40">Previous</button>
|
||||
<button onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))} disabled={currentPage === totalPages}
|
||||
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40">Next</button>
|
||||
<button
|
||||
onClick={() => setCurrentPage((p) => Math.max(1, p - 1))}
|
||||
disabled={currentPage === 1}
|
||||
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40"
|
||||
>
|
||||
Previous
|
||||
</button>
|
||||
<button
|
||||
onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))}
|
||||
disabled={currentPage === totalPages}
|
||||
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40"
|
||||
>
|
||||
Next
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -166,12 +166,16 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
|
|||
|
||||
// Set provider
|
||||
setSelectedProvider(preset.provider);
|
||||
form.setFieldsValue({
|
||||
const baseValues: Record<string, any> = {
|
||||
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<AddGuardrailFormProps> = ({ visible, onClose, a
|
|||
const handleProviderChange = (value: string) => {
|
||||
setSelectedProvider(value);
|
||||
// Reset form fields that are provider-specific
|
||||
form.setFieldsValue({
|
||||
const resetValues: Record<string, any> = {
|
||||
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([]);
|
||||
|
|
|
|||
|
|
@ -148,6 +148,18 @@ export const GUARDRAIL_PRESETS: Record<string, GuardrailPreset> = {
|
|||
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: {
|
||||
|
|
|
|||
|
|
@ -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[] = [
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ export const guardrail_provider_map: Record<string, string> = {
|
|||
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
|
||||
|
|
|
|||
|
|
@ -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<GuardrailProviderFieldsProps> = ({
|
|||
);
|
||||
}
|
||||
|
||||
const percentageInitialValue =
|
||||
field.type === "percentage" && (fieldValue === undefined || fieldValue === null)
|
||||
? (field.default_value ?? 0.5)
|
||||
: undefined;
|
||||
|
||||
return (
|
||||
<Form.Item
|
||||
key={fullFieldKey}
|
||||
|
|
@ -161,6 +169,7 @@ const GuardrailProviderFields: React.FC<GuardrailProviderFieldsProps> = ({
|
|||
label={fieldKey}
|
||||
tooltip={field.description}
|
||||
rules={field.required ? [{ required: true, message: `${fieldKey} is required` }] : undefined}
|
||||
initialValue={percentageInitialValue}
|
||||
>
|
||||
{field.type === "select" && field.options ? (
|
||||
<Select placeholder={field.description} defaultValue={fieldValue || field.default_value}>
|
||||
|
|
@ -186,6 +195,17 @@ const GuardrailProviderFields: React.FC<GuardrailProviderFieldsProps> = ({
|
|||
<Select.Option value="true">True</Select.Option>
|
||||
<Select.Option value="false">False</Select.Option>
|
||||
</Select>
|
||||
) : field.type === "percentage" && field.min != null && field.max != null ? (
|
||||
<Slider
|
||||
min={field.min}
|
||||
max={field.max}
|
||||
step={field.step ?? 0.1}
|
||||
marks={{
|
||||
[field.min]: "0%",
|
||||
[(field.min + field.max) / 2]: "50%",
|
||||
[field.max]: "100%",
|
||||
}}
|
||||
/>
|
||||
) : field.type === "number" ? (
|
||||
<NumericalInput
|
||||
step={1}
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ vi.mock("../networking", () => ({
|
|||
soft_budget: null,
|
||||
}),
|
||||
fetchMCPAccessGroups: vi.fn().mockResolvedValue([]),
|
||||
getAgentsList: vi.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
vi.mock("../molecules/notifications_manager", () => ({
|
||||
|
|
|
|||
|
|
@ -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<CreateKeyProps> = ({ 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<CreateKeyProps> = ({ team, teams, data, addKey }) => {
|
|||
<Radio value="you">You</Radio>
|
||||
<Radio value="service_account">Service Account</Radio>
|
||||
{userRole === "Admin" && <Radio value="another_user">Another User</Radio>}
|
||||
<Radio value="agent">Agent <Tag color="purple">New</Tag></Radio>
|
||||
<Radio value="agent">
|
||||
Agent <Tag color="purple">New</Tag>
|
||||
</Radio>
|
||||
</Radio.Group>
|
||||
</Form.Item>
|
||||
|
||||
|
|
@ -1005,9 +1006,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ 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<CreateKeyProps> = ({ team, teams, data, addKey }) => {
|
|||
className="mt-4"
|
||||
help="Select access groups to assign to this key"
|
||||
>
|
||||
<AccessGroupSelector
|
||||
placeholder="Select access groups (optional)"
|
||||
/>
|
||||
<AccessGroupSelector placeholder="Select access groups (optional)" />
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
|
|
@ -1297,7 +1294,11 @@ const CreateKey: React.FC<CreateKeyProps> = ({ 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
|
||||
}
|
||||
/>
|
||||
</div>
|
||||
</AccordionBody>
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue