diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 3ce819439cb..a825ce22b25 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1236,6 +1236,275 @@ async def get_provider_specific_params(): return provider_params +class TestCustomCodeGuardrailRequest(BaseModel): + """Request model for testing custom code guardrails.""" + + custom_code: str + """The Python-like code containing the apply_guardrail function.""" + + test_input: Dict[str, Any] + """The test input to pass to the guardrail. Should contain 'texts', optionally 'images', 'tools', etc.""" + + input_type: str = "request" + """Whether this is a 'request' or 'response' input type.""" + + request_data: Optional[Dict[str, Any]] = None + """Optional mock request_data (model, user_id, team_id, metadata, etc.).""" + + +class TestCustomCodeGuardrailResponse(BaseModel): + """Response model for testing custom code guardrails.""" + + success: bool + """Whether the test executed successfully (no errors).""" + + result: Optional[Dict[str, Any]] = None + """The guardrail result: action (allow/block/modify), reason, modified_texts, etc.""" + + error: Optional[str] = None + """Error message if execution failed.""" + + error_type: Optional[str] = None + """Type of error: 'compilation' or 'execution'.""" + + +@router.post( + "/guardrails/test_custom_code", + tags=["Guardrails"], + dependencies=[Depends(user_api_key_auth)], + response_model=TestCustomCodeGuardrailResponse, +) +async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest): + """ + Test custom code guardrail logic without creating a guardrail. + + This endpoint allows admins to experiment with custom code guardrails by: + 1. Compiling the provided code in a sandbox + 2. Executing the apply_guardrail function with test input + 3. Returning the result (allow/block/modify) + + 👉 [Custom Code Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/custom_code_guardrail) + + Example Request: + ```bash + curl -X POST "http://localhost:4000/guardrails/test_custom_code" \\ + -H "Authorization: Bearer " \\ + -H "Content-Type: application/json" \\ + -d '{ + "custom_code": "def apply_guardrail(inputs, request_data, input_type):\\n for text in inputs[\\"texts\\"]:\\n if regex_match(text, r\\"\\\\d{3}-\\\\d{2}-\\\\d{4}\\"):\\n return block(\\"SSN detected\\")\\n return allow()", + "test_input": { + "texts": ["My SSN is 123-45-6789"] + }, + "input_type": "request" + }' + ``` + + Example Success Response (blocked): + ```json + { + "success": true, + "result": { + "action": "block", + "reason": "SSN detected" + }, + "error": null, + "error_type": null + } + ``` + + Example Success Response (allowed): + ```json + { + "success": true, + "result": { + "action": "allow" + }, + "error": null, + "error_type": null + } + ``` + + Example Success Response (modified): + ```json + { + "success": true, + "result": { + "action": "modify", + "texts": ["My SSN is [REDACTED]"] + }, + "error": null, + "error_type": null + } + ``` + + Example Error Response (compilation error): + ```json + { + "success": false, + "result": null, + "error": "Syntax error in custom code: invalid syntax (, line 1)", + "error_type": "compilation" + } + ``` + """ + import concurrent.futures + import re + + from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import ( + get_custom_code_primitives, + ) + + # Security validation patterns + FORBIDDEN_PATTERNS = [ + # Import statements + (r"\bimport\s+", "import statements are not allowed"), + (r"\bfrom\s+\w+\s+import\b", "from...import statements are not allowed"), + (r"__import__\s*\(", "__import__() is not allowed"), + # Dangerous builtins + (r"\bexec\s*\(", "exec() is not allowed"), + (r"\beval\s*\(", "eval() is not allowed"), + (r"\bcompile\s*\(", "compile() is not allowed"), + (r"\bopen\s*\(", "open() is not allowed"), + (r"\bgetattr\s*\(", "getattr() is not allowed"), + (r"\bsetattr\s*\(", "setattr() is not allowed"), + (r"\bdelattr\s*\(", "delattr() is not allowed"), + (r"\bglobals\s*\(", "globals() is not allowed"), + (r"\blocals\s*\(", "locals() is not allowed"), + (r"\bvars\s*\(", "vars() is not allowed"), + (r"\bdir\s*\(", "dir() is not allowed"), + (r"\bbreakpoint\s*\(", "breakpoint() is not allowed"), + (r"\binput\s*\(", "input() is not allowed"), + # Dangerous dunder access + (r"__builtins__", "__builtins__ access is not allowed"), + (r"__globals__", "__globals__ access is not allowed"), + (r"__code__", "__code__ access is not allowed"), + (r"__subclasses__", "__subclasses__ access is not allowed"), + (r"__bases__", "__bases__ access is not allowed"), + (r"__mro__", "__mro__ access is not allowed"), + (r"__class__", "__class__ access is not allowed"), + (r"__dict__", "__dict__ access is not allowed"), + (r"__getattribute__", "__getattribute__ access is not allowed"), + (r"__reduce__", "__reduce__ access is not allowed"), + (r"__reduce_ex__", "__reduce_ex__ access is not allowed"), + # OS/system access + (r"\bos\.", "os module access is not allowed"), + (r"\bsys\.", "sys module access is not allowed"), + (r"\bsubprocess\.", "subprocess module access is not allowed"), + ] + + EXECUTION_TIMEOUT_SECONDS = 5 + + try: + # Step 0: Security validation - check for forbidden patterns + code = request.custom_code + for pattern, error_msg in FORBIDDEN_PATTERNS: + if re.search(pattern, code): + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Security violation: {error_msg}", + error_type="compilation", + ) + + # Step 1: Compile the custom code with restricted environment + exec_globals = get_custom_code_primitives().copy() + + # Remove access to builtins to prevent escape + exec_globals["__builtins__"] = {} + + try: + exec(compile(request.custom_code, "", "exec"), exec_globals) + except SyntaxError as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Syntax error in custom code: {e}", + error_type="compilation", + ) + except Exception as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Failed to compile custom code: {e}", + error_type="compilation", + ) + + # Step 2: Verify apply_guardrail function exists + if "apply_guardrail" not in exec_globals: + return TestCustomCodeGuardrailResponse( + success=False, + error="Custom code must define an 'apply_guardrail' function. " + "Expected signature: apply_guardrail(inputs, request_data, input_type)", + error_type="compilation", + ) + + apply_fn = exec_globals["apply_guardrail"] + if not callable(apply_fn): + return TestCustomCodeGuardrailResponse( + success=False, + error="'apply_guardrail' must be a callable function", + error_type="compilation", + ) + + # Step 3: Prepare test inputs + test_inputs = request.test_input + if "texts" not in test_inputs: + test_inputs["texts"] = [] + + # Prepare mock request_data + mock_request_data = request.request_data or {} + safe_request_data = { + "model": mock_request_data.get("model", "test-model"), + "user_id": mock_request_data.get("user_id"), + "team_id": mock_request_data.get("team_id"), + "end_user_id": mock_request_data.get("end_user_id"), + "metadata": mock_request_data.get("metadata", {}), + } + + # Step 4: Execute the function with timeout protection + + def execute_guardrail(): + return apply_fn(test_inputs, safe_request_data, request.input_type) + + try: + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(execute_guardrail) + try: + result = future.result(timeout=EXECUTION_TIMEOUT_SECONDS) + except concurrent.futures.TimeoutError: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Execution timeout: code took longer than {EXECUTION_TIMEOUT_SECONDS} seconds", + error_type="execution", + ) + except Exception as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Execution error: {e}", + error_type="execution", + ) + + # Step 5: Validate and return result + if not isinstance(result, dict): + return TestCustomCodeGuardrailResponse( + success=True, + result={ + "action": "allow", + "warning": f"Expected dict result, got {type(result).__name__}. Treating as allow.", + }, + ) + + return TestCustomCodeGuardrailResponse( + success=True, + result=result, + ) + + except Exception as e: + verbose_proxy_logger.exception(f"Error testing custom code guardrail: {e}") + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Unexpected error: {e}", + error_type="execution", + ) + + @router.post("/guardrails/apply_guardrail", response_model=ApplyGuardrailResponse) @router.post("/apply_guardrail", response_model=ApplyGuardrailResponse) async def apply_guardrail( diff --git a/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx b/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx index 147696921b3..73d96c590ec 100644 --- a/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx @@ -9,7 +9,7 @@ import { CaretRightOutlined, SaveOutlined, } from "@ant-design/icons"; -import { createGuardrailCall } from "../../networking"; +import { createGuardrailCall, testCustomCodeGuardrail } from "../../networking"; import NotificationsManager from "../../molecules/notifications_manager"; const { Panel } = Collapse; @@ -226,28 +226,60 @@ const CustomCodeModal: React.FC = ({ } }; - // Test guardrail (placeholder - would need backend endpoint) + // Test guardrail using backend endpoint const handleTest = async () => { + if (!accessToken) { + setTestResult({ error: "No access token available" }); + return; + } + setIsTesting(true); setTestResult(null); - - // Simulate test - in real implementation, call backend - setTimeout(() => { + + try { + // Parse test input JSON + let parsedInput; try { - JSON.parse(testInput); - // Mock result based on code content - if (code.includes("block(")) { - setTestResult({ action: "block", reason: "Test blocked" }); - } else if (code.includes("modify(")) { - setTestResult({ action: "modify", modified_texts: ["[REDACTED]"] }); - } else { - setTestResult({ action: "allow" }); - } + parsedInput = JSON.parse(testInput); } catch (e) { setTestResult({ error: "Invalid test input JSON" }); + setIsTesting(false); + return; } + + // Ensure texts array exists + if (!parsedInput.texts) { + parsedInput.texts = []; + } + + const response = await testCustomCodeGuardrail(accessToken, { + custom_code: code, + test_input: parsedInput, + input_type: mode as "request" | "response", + request_data: { + model: "test-model", + metadata: {}, + }, + }); + + if (response.success && response.result) { + setTestResult(response.result); + } else if (response.error) { + setTestResult({ + error: response.error, + error_type: response.error_type, + }); + } else { + setTestResult({ error: "Unknown error occurred" }); + } + } catch (error) { + console.error("Failed to test custom code:", error); + setTestResult({ + error: error instanceof Error ? error.message : "Failed to test custom code", + }); + } finally { setIsTesting(false); - }, 500); + } }; const lineCount = code.split("\n").length; @@ -384,13 +416,28 @@ const CustomCodeModal: React.FC = ({ "text-blue-600" }`}> {testResult.error ? ( - <> {testResult.error} + <> + + + {testResult.error_type && [{testResult.error_type}] } + {testResult.error} + + ) : testResult.action === "allow" ? ( <> Allowed ) : testResult.action === "block" ? ( <> Blocked: {testResult.reason} + ) : testResult.action === "modify" ? ( + <> + Modified + {testResult.texts && testResult.texts.length > 0 && ( + + → {testResult.texts[0].substring(0, 50)}{testResult.texts[0].length > 50 ? "..." : ""} + + )} + ) : ( - <> Modified + <> {testResult.action || "Unknown"} )} )} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 5cc5216e132..ca4c16a781f 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -7607,6 +7607,89 @@ export const applyGuardrail = async ( } }; +export interface TestCustomCodeGuardrailRequest { + custom_code: string; + test_input: { + texts: string[]; + images?: string[]; + tools?: Record[]; + tool_calls?: Record[]; + structured_messages?: Record[]; + model?: string; + }; + input_type?: "request" | "response"; + request_data?: { + model?: string; + user_id?: string; + team_id?: string; + end_user_id?: string; + metadata?: Record; + }; +} + +export interface TestCustomCodeGuardrailResponse { + success: boolean; + result?: { + action: "allow" | "block" | "modify"; + reason?: string; + texts?: string[]; + images?: string[]; + tool_calls?: Record[]; + detection_info?: Record; + warning?: string; + }; + error?: string; + error_type?: "compilation" | "execution"; +} + +export const testCustomCodeGuardrail = async ( + accessToken: string, + request: TestCustomCodeGuardrailRequest +): Promise => { + try { + const url = proxyBaseUrl + ? `${proxyBaseUrl}/guardrails/test_custom_code` + : `/guardrails/test_custom_code`; + + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(request), + }); + + if (!response.ok) { + const errorData = await response.text(); + let errorMessage = "Failed to test custom code guardrail"; + + try { + const errorJson = JSON.parse(errorData); + if (errorJson.error?.message) { + errorMessage = errorJson.error.message; + } else if (errorJson.detail) { + errorMessage = errorJson.detail; + } else if (errorJson.message) { + errorMessage = errorJson.message; + } + } catch (e) { + errorMessage = errorData || errorMessage; + } + + handleError(errorData); + throw new Error(errorMessage); + } + + const data = await response.json(); + console.log("Test custom code guardrail response:", data); + return data; + } catch (error) { + console.error("Failed to test custom code guardrail:", error); + throw error; + } +}; + export const validateBlockedWordsFile = async (accessToken: string, fileContent: string) => { try { const url = proxyBaseUrl