scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Address CR feedback: fix auth dependency duplication, correct logger wording, clean imports

Browse source
This commit is contained in:
Harshit28j 2026-02-24 13:46:18 +05:30
parent ff38cd8c38
commit 2d2d7ced1a
4 changed files with 44 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -2,6 +2,7 @@
CRUD ENDPOINTS FOR GUARDRAILS
"""
import concurrent.futures
import inspect
from typing import Any, Dict, List, Optional, Type, TypeVar, Union, cast
@ -11,9 +12,16 @@ from pydantic import BaseModel
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
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,
@ -243,7 +251,6 @@ class CreateGuardrailRequest(BaseModel):
@router.post(
"/guardrails",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def create_guardrail(
request: CreateGuardrailRequest,
@ -298,7 +305,6 @@ async def create_guardrail(
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._types import LitellmUserRoles
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
@ -342,7 +348,6 @@ class UpdateGuardrailRequest(BaseModel):
@router.put(
"/guardrails/{guardrail_id}",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_guardrail(
guardrail_id: str,
@ -398,7 +403,6 @@ async def update_guardrail(
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._types import LitellmUserRoles
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
@ -450,7 +454,6 @@ async def update_guardrail(
@router.delete(
"/guardrails/{guardrail_id}",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_guardrail(
guardrail_id: str,
@ -476,7 +479,6 @@ async def delete_guardrail(
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._types import LitellmUserRoles
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
@ -526,7 +528,6 @@ async def delete_guardrail(
@router.patch(
"/guardrails/{guardrail_id}",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def patch_guardrail(
guardrail_id: str,
@ -580,7 +581,6 @@ async def patch_guardrail(
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._types import LitellmUserRoles
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
@ -1406,7 +1406,6 @@ class TestCustomCodeGuardrailResponse(BaseModel):
@router.post(
"/guardrails/test_custom_code",
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
response_model=TestCustomCodeGuardrailResponse,
)
async def test_custom_code_guardrail(
@ -1485,12 +1484,6 @@ async def test_custom_code_guardrail(
}
```
"""
import concurrent.futures
from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import (
get_custom_code_primitives,
)
from litellm.proxy._types import LitellmUserRoles
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
@ -1502,10 +1495,6 @@ async def test_custom_code_guardrail(
try:
# Step 0: Security validation - check for forbidden patterns
from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import (
validate_custom_code,
CustomCodeValidationError,
)
try:
validate_custom_code(request.custom_code)

4
litellm/proxy/guardrails/guardrail_hooks/custom_code/code_validator.py
View file

@ -37,6 +37,10 @@ FORBIDDEN_PATTERNS: List[Tuple[str, str]] = [
(r"\bos\.", "os module access is not allowed"),
(r"\bsys\.", "sys module access is not allowed"),
(r"\bsubprocess\.", "subprocess module access is not allowed"),
(r"\bshutil\.", "shutil module access is not allowed"),
(r"\bctypes\.", "ctypes module access is not allowed"),
(r"\bsocket\.", "socket module access is not allowed"),
(r"\bpickle\.", "pickle module access is not allowed"),
]

73
litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
View file

@ -144,6 +144,33 @@ class CustomCodeGuardrail(CustomGuardrail):
"""Returns the config model for the UI."""
return CustomCodeGuardrailConfigModel
def _do_compile(self) -> None:
"""Internal compilation method without lock. Expected to run inside _compile_lock."""
# Create a restricted execution environment
# Only include our safe primitives
exec_globals = get_custom_code_primitives().copy()
# CRITICAL: Restrict __builtins__ to prevent sandbox escape
exec_globals["__builtins__"] = {}
# Execute the user code in the restricted environment
exec(compile(self.custom_code, "<guardrail>", "exec"), exec_globals)
# Extract the apply_guardrail function
if "apply_guardrail" not in exec_globals:
raise CustomCodeCompilationError(
"Custom code must define an 'apply_guardrail' function. "
"Expected signature: apply_guardrail(inputs, request_data, input_type)"
)
apply_fn = exec_globals["apply_guardrail"]
if not callable(apply_fn):
raise CustomCodeCompilationError(
"'apply_guardrail' must be a callable function"
)
self._compiled_function = apply_fn
def _compile_custom_code(self) -> None:
"""
Compile the custom code and extract the apply_guardrail function.
@ -161,30 +188,8 @@ class CustomCodeGuardrail(CustomGuardrail):
except CustomCodeValidationError as e:
raise CustomCodeCompilationError(str(e)) from e
# Create a restricted execution environment
# Only include our safe primitives
exec_globals = get_custom_code_primitives().copy()
# CRITICAL: Restrict __builtins__ to prevent sandbox escape
exec_globals["__builtins__"] = {}
# Execute the user code in the restricted environment
exec(compile(self.custom_code, "<guardrail>", "exec"), exec_globals)
# Extract the apply_guardrail function
if "apply_guardrail" not in exec_globals:
raise CustomCodeCompilationError(
"Custom code must define an 'apply_guardrail' function. "
"Expected signature: apply_guardrail(inputs, request_data, input_type)"
)
apply_fn = exec_globals["apply_guardrail"]
if not callable(apply_fn):
raise CustomCodeCompilationError(
"'apply_guardrail' must be a callable function"
)
self._compiled_function = apply_fn
# Step 2: Compile logic
self._do_compile()
verbose_proxy_logger.debug(
f"Custom code guardrail '{self.guardrail_name}' compiled successfully"
)
@ -415,25 +420,7 @@ class CustomCodeGuardrail(CustomGuardrail):
try:
self.custom_code = new_code
# Inline compilation instead of calling _compile_custom_code()
# to avoid deadlock (re-acquiring self._compile_lock)
exec_globals = get_custom_code_primitives().copy()
exec_globals["__builtins__"] = {}
exec(compile(self.custom_code, "<guardrail>", "exec"), exec_globals)
if "apply_guardrail" not in exec_globals:
raise CustomCodeCompilationError(
"Custom code must define an 'apply_guardrail' function. "
"Expected signature: apply_guardrail(inputs, request_data, input_type)"
)
apply_fn = exec_globals["apply_guardrail"]
if not callable(apply_fn):
raise CustomCodeCompilationError(
"'apply_guardrail' must be a callable function"
)
self._compiled_function = apply_fn
self._do_compile()
verbose_proxy_logger.info(
f"Custom code guardrail '{self.guardrail_name}': Code updated successfully"
)

2
litellm/proxy/proxy_server.py
View file

@ -2994,7 +2994,7 @@ class ProxyConfig:
litellm_master_key_hash = hash_token(master_key)
else:
verbose_proxy_logger.critical(
"LITELLM_MASTER_KEY is not set! All unauthenticated requests will be treated as INTERNAL_USER. This is insecure for production."
"LITELLM_MASTER_KEY is not set! All requests will be treated as INTERNAL_USER with no admin access. Set LITELLM_MASTER_KEY for production use."
)
### USER API KEY CACHE IN-MEMORY TTL ###
user_api_key_cache_ttl = general_settings.get(