From ba74e6d9d246b7b71d8ab71e2e3ef69e0ddfa748 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 6 Feb 2026 17:34:32 -0800 Subject: [PATCH] Add http support to custom code guardrails + Unified guardrails for MCP + Agent guardrail support (#20619) * fix: fix styling * fix(custom_code_guardrail.py): add http support for custom code guardrails allows users to call external guardrails on litellm with minimal code changes (no custom handlers) Test guardrail integrations more easily * feat(a2a/): add guardrails for agent interactions allows the same guardrails for llm's to be applied to agents as well * fix(a2a/): support passing guardrails to a2a from the UI * style(code-editor): allow editing custom code guardrails on ui + add examples of pre/post calls for custom code guardrails * feat(mcp/): support custom code guardrails for mcp calls allows custom code guardrails to work on mcp input * feat(chatui.tsx): support guardrails on mcp tool calls on playground --- .../proxy/guardrails/custom_code_guardrail.md | 58 ++- litellm/integrations/custom_guardrail.py | 1 + .../a2a/chat/guardrail_translation/README.md | 155 ++++++++ .../chat/guardrail_translation/__init__.py | 11 + .../a2a/chat/guardrail_translation/handler.py | 315 ++++++++++++++++ .../guardrail_translation/handler.py | 105 +++--- .../mcp_server/mcp_server_manager.py | 24 +- .../mcp_server/rest_endpoints.py | 87 +++-- .../proxy/agent_endpoints/a2a_endpoints.py | 33 +- litellm/proxy/common_request_processing.py | 16 +- .../custom_code/custom_code_guardrail.py | 33 +- .../guardrail_hooks/custom_code/primitives.py | 230 ++++++++++++ .../unified_guardrail/unified_guardrail.py | 2 + .../custom_code/CustomCodeModal.tsx | 277 +++++++++++--- .../components/guardrails/guardrail_info.tsx | 58 ++- .../mcp_tools/MCPToolArgumentsForm.tsx | 345 ++++++++++++++++++ .../src/components/networking.tsx | 20 +- .../components/playground/chat_ui/ChatUI.tsx | 324 ++++++++++++---- .../playground/chat_ui/chatConstants.ts | 1 + .../chat_ui/mode_endpoint_mapping.tsx | 1 + .../playground/llm_calls/a2a_send_message.tsx | 6 + 21 files changed, 1862 insertions(+), 240 deletions(-) create mode 100644 litellm/llms/a2a/chat/guardrail_translation/README.md create mode 100644 litellm/llms/a2a/chat/guardrail_translation/__init__.py create mode 100644 litellm/llms/a2a/chat/guardrail_translation/handler.py create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/MCPToolArgumentsForm.tsx diff --git a/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md index cb24614449..8cbc247ae5 100644 --- a/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md +++ b/docs/my-website/docs/proxy/guardrails/custom_code_guardrail.md @@ -61,15 +61,23 @@ curl -X POST http://localhost:4000/chat/completions \ ### Function Signature -Your code must define an `apply_guardrail` function: +Your code must define an `apply_guardrail` function. It can be either sync or async: ```python +# Sync version def apply_guardrail(inputs, request_data, input_type): # inputs: see table below # request_data: {"model": "...", "user_id": "...", "team_id": "...", "metadata": {...}} # input_type: "request" or "response" return allow() # or block() or modify() + +# Async version (recommended when using HTTP primitives) +async def apply_guardrail(inputs, request_data, input_type): + response = await http_post("https://api.example.com/check", body={"text": inputs["texts"][0]}) + if response["success"] and response["body"].get("flagged"): + return block("Content flagged") + return allow() ``` ### `inputs` Parameter @@ -145,6 +153,29 @@ def apply_guardrail(inputs, request_data, input_type): | `char_count(text)` | Count characters | | `lower(text)` / `upper(text)` / `trim(text)` | String transforms | +### HTTP Requests (Async) + +Make async HTTP requests to external APIs for additional validation or content moderation. + +| Function | Description | +|----------|-------------| +| `await http_request(url, method, headers, body, timeout)` | General async HTTP request | +| `await http_get(url, headers, timeout)` | Async GET request | +| `await http_post(url, body, headers, timeout)` | Async POST request | + +**Response format:** +```python +{ + "status_code": 200, # HTTP status code + "body": {...}, # Response body (parsed JSON or string) + "headers": {...}, # Response headers + "success": True, # True if status code is 2xx + "error": None # Error message if request failed +} +``` + +**Note:** When using HTTP primitives, define your function as `async def apply_guardrail(...)` for non-blocking execution. + ## Examples ### Block PII (SSN) @@ -213,6 +244,29 @@ def apply_guardrail(inputs, request_data, input_type): return allow() ``` +### Call External Moderation API (Async) + +```python +async def apply_guardrail(inputs, request_data, input_type): + # Call an external moderation API + for text in inputs["texts"]: + response = await http_post( + "https://api.example.com/moderate", + body={"text": text, "user_id": request_data["user_id"]}, + headers={"Authorization": "Bearer YOUR_API_KEY"}, + timeout=10 + ) + + if not response["success"]: + # API call failed - decide whether to allow or block + return allow() + + if response["body"].get("flagged"): + return block(response["body"].get("reason", "Content flagged")) + + return allow() +``` + ### Combine Multiple Checks ```python @@ -241,8 +295,8 @@ Custom code runs in a restricted environment: - ❌ No `import` statements - ❌ No file I/O -- ❌ No network access - ❌ No `exec()` or `eval()` +- ✅ HTTP requests via built-in `http_request`, `http_get`, `http_post` primitives - ✅ Only LiteLLM-provided primitives available ## Per-Request Usage diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1652ec2aa0..bbd55a59bc 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -268,6 +268,7 @@ class CustomGuardrail(CustomLogger): """ Returns the guardrail(s) to be run from the metadata or root """ + if "guardrails" in data: return data["guardrails"] metadata = data.get("litellm_metadata") or data.get("metadata", {}) diff --git a/litellm/llms/a2a/chat/guardrail_translation/README.md b/litellm/llms/a2a/chat/guardrail_translation/README.md new file mode 100644 index 0000000000..1e18f5cda3 --- /dev/null +++ b/litellm/llms/a2a/chat/guardrail_translation/README.md @@ -0,0 +1,155 @@ +# A2A Protocol Guardrail Translation Handler + +Handler for processing A2A (Agent-to-Agent) Protocol messages with guardrails. + +## Overview + +This handler processes A2A JSON-RPC 2.0 input/output by: +1. Extracting text from message parts (`kind: "text"`) +2. Applying guardrails to text content +3. Mapping guardrailed text back to original structure + +## A2A Protocol Format + +### Input Format (JSON-RPC 2.0) + +```json +{ + "jsonrpc": "2.0", + "id": "request-id", + "method": "message/send", + "params": { + "message": { + "kind": "message", + "messageId": "...", + "role": "user", + "parts": [ + {"kind": "text", "text": "Hello, my SSN is 123-45-6789"} + ] + }, + "metadata": { + "guardrails": ["block-ssn"] + } + } +} +``` + +### Output Formats + +The handler supports multiple A2A response formats: + +**Direct message:** +```json +{ + "result": { + "kind": "message", + "parts": [{"kind": "text", "text": "Response text"}] + } +} +``` + +**Nested message:** +```json +{ + "result": { + "message": { + "parts": [{"kind": "text", "text": "Response text"}] + } + } +} +``` + +**Task with artifacts:** +```json +{ + "result": { + "kind": "task", + "artifacts": [ + {"parts": [{"kind": "text", "text": "Artifact text"}]} + ] + } +} +``` + +**Task with status message:** +```json +{ + "result": { + "kind": "task", + "status": { + "message": { + "parts": [{"kind": "text", "text": "Status message"}] + } + } + } +} +``` + +**Streaming artifact-update:** +```json +{ + "result": { + "kind": "artifact-update", + "artifact": { + "parts": [{"kind": "text", "text": "Streaming text"}] + } + } +} +``` + +## Usage + +The handler is automatically discovered and applied when guardrails are used with A2A endpoints. + +### Via LiteLLM Proxy + +```bash +curl -X POST 'http://localhost:4000/a2a/my-agent' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "jsonrpc": "2.0", + "id": "1", + "method": "message/send", + "params": { + "message": { + "kind": "message", + "messageId": "msg-1", + "role": "user", + "parts": [{"kind": "text", "text": "Hello, my SSN is 123-45-6789"}] + }, + "metadata": { + "guardrails": ["block-ssn"] + } + } +}' +``` + +### Specifying Guardrails + +Guardrails can be specified in the A2A request via the `metadata.guardrails` field: + +```json +{ + "params": { + "message": {...}, + "metadata": { + "guardrails": ["block-ssn", "pii-filter"] + } + } +} +``` + +## Extension + +Override these methods to customize behavior: + +- `_extract_texts_from_result()`: Custom text extraction from A2A responses +- `_extract_texts_from_parts()`: Custom text extraction from message parts +- `_apply_text_to_path()`: Custom application of guardrailed text + +## Call Types + +This handler is registered for: +- `CallTypes.send_message`: Synchronous A2A message sending +- `CallTypes.asend_message`: Asynchronous A2A message sending diff --git a/litellm/llms/a2a/chat/guardrail_translation/__init__.py b/litellm/llms/a2a/chat/guardrail_translation/__init__.py new file mode 100644 index 0000000000..13c2067748 --- /dev/null +++ b/litellm/llms/a2a/chat/guardrail_translation/__init__.py @@ -0,0 +1,11 @@ +"""A2A Protocol handler for Unified Guardrails.""" + +from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.send_message: A2AGuardrailHandler, + CallTypes.asend_message: A2AGuardrailHandler, +} + +__all__ = ["guardrail_translation_mappings"] diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py new file mode 100644 index 0000000000..770453f2de --- /dev/null +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -0,0 +1,315 @@ +""" +A2A Protocol Handler for Unified Guardrails + +This module provides guardrail translation support for A2A (Agent-to-Agent) Protocol. +It handles both JSON-RPC 2.0 input requests and output responses, extracting text +from message parts and applying guardrails. + +A2A Protocol Format: +- Input: JSON-RPC 2.0 with params.message.parts containing text parts +- Output: JSON-RPC 2.0 with result containing message/artifact parts +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth + + +class A2AGuardrailHandler(BaseTranslation): + """ + Handler for processing A2A Protocol messages with guardrails. + + This class provides methods to: + 1. Process input messages (pre-call hook) - extracts text from A2A message parts + 2. Process output responses (post-call hook) - extracts text from A2A response parts + + A2A Message Format: + - Input: params.message.parts[].text (where kind == "text") + - Output: result.message.parts[].text or result.artifacts[].parts[].text + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> Any: + """ + Process A2A input messages by applying guardrails to text content. + + Extracts text from A2A message parts and applies guardrails. + + Args: + data: The A2A JSON-RPC 2.0 request data + guardrail_to_apply: The guardrail instance to apply + litellm_logging_obj: Optional logging object + + Returns: + Modified data with guardrails applied to text content + """ + # A2A request format: { "params": { "message": { "parts": [...] } } } + params = data.get("params", {}) + message = params.get("message", {}) + parts = message.get("parts", []) + + if not parts: + verbose_proxy_logger.debug("A2A: No parts in message, skipping guardrail") + return data + + texts_to_check: List[str] = [] + text_part_indices: List[int] = [] # Track which parts contain text + + # Step 1: Extract text from all text parts + for part_idx, part in enumerate(parts): + if part.get("kind") == "text": + text = part.get("text", "") + if text: + texts_to_check.append(text) + text_part_indices.append(part_idx) + + # Step 2: Apply guardrail to all texts in batch + if texts_to_check: + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + + # Pass the structured A2A message to guardrails + inputs["structured_messages"] = [message] + + # Include agent model info if available + model = data.get("model") + if model: + inputs["model"] = model + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + + guardrailed_texts = guardrailed_inputs.get("texts", []) + + # Step 3: Apply guardrailed text back to original parts + if guardrailed_texts and len(guardrailed_texts) == len(text_part_indices): + for task_idx, part_idx in enumerate(text_part_indices): + parts[part_idx]["text"] = guardrailed_texts[task_idx] + + verbose_proxy_logger.debug("A2A: Processed input message: %s", message) + + return data + + async def process_output_response( + self, + response: Any, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + ) -> Any: + """ + Process A2A output response by applying guardrails to text content. + + Handles multiple A2A response formats: + - Direct message: {"result": {"kind": "message", "parts": [...]}} + - Nested message: {"result": {"message": {"parts": [...]}}} + - Task with artifacts: {"result": {"kind": "task", "artifacts": [{"parts": [...]}]}} + - Task with status message: {"result": {"kind": "task", "status": {"message": {"parts": [...]}}}} + + Args: + response: A2A JSON-RPC 2.0 response dict or object + guardrail_to_apply: The guardrail instance to apply + litellm_logging_obj: Optional logging object + user_api_key_dict: User API key metadata + + Returns: + Modified response with guardrails applied to text content + """ + # Handle both dict and Pydantic model responses + if hasattr(response, "model_dump"): + response_dict = response.model_dump() + is_pydantic = True + elif isinstance(response, dict): + response_dict = response + is_pydantic = False + else: + verbose_proxy_logger.warning( + "A2A: Unknown response type %s, skipping guardrail", type(response) + ) + return response + + result = response_dict.get("result", {}) + if not result or not isinstance(result, dict): + verbose_proxy_logger.debug("A2A: No result in response, skipping guardrail") + return response + + # Find all text-containing parts in the response + texts_to_check: List[str] = [] + # Each mapping is (path_to_parts_list, part_index) + # path_to_parts_list is a tuple of keys to navigate to the parts list + task_mappings: List[Tuple[Tuple[str, ...], int]] = [] + + # Extract texts from all possible locations + self._extract_texts_from_result( + result=result, + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + if not texts_to_check: + verbose_proxy_logger.debug("A2A: No text content in response") + return response + + # Step 2: Apply guardrail to all texts in batch + # Create a request_data dict with response info and user API key metadata + request_data: dict = {"response": response_dict} + + # Add user API key metadata with prefixed keys + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) + if user_metadata: + request_data["litellm_metadata"] = user_metadata + + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + + guardrailed_texts = guardrailed_inputs.get("texts", []) + + # Step 3: Apply guardrailed text back to original response + if guardrailed_texts and len(guardrailed_texts) == len(task_mappings): + for task_idx, (path, part_idx) in enumerate(task_mappings): + self._apply_text_to_path( + result=result, + path=path, + part_idx=part_idx, + text=guardrailed_texts[task_idx], + ) + + verbose_proxy_logger.debug("A2A: Processed output response") + + # Update the original response + if is_pydantic: + # For Pydantic models, we need to update the underlying dict + # and the model will reflect the changes + response_dict["result"] = result + return response + else: + response["result"] = result + return response + + def _extract_texts_from_result( + self, + result: Dict[str, Any], + texts_to_check: List[str], + task_mappings: List[Tuple[Tuple[str, ...], int]], + ) -> None: + """ + Extract text from all possible locations in an A2A result. + + Handles multiple response formats: + 1. Direct message with parts: {"parts": [...]} + 2. Nested message: {"message": {"parts": [...]}} + 3. Task with artifacts: {"artifacts": [{"parts": [...]}]} + 4. Task with status message: {"status": {"message": {"parts": [...]}}} + 5. Streaming artifact-update: {"artifact": {"parts": [...]}} + """ + # Case 1: Direct parts in result (direct message) + if "parts" in result: + self._extract_texts_from_parts( + parts=result["parts"], + path=("parts",), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + # Case 2: Nested message + message = result.get("message") + if message and isinstance(message, dict) and "parts" in message: + self._extract_texts_from_parts( + parts=message["parts"], + path=("message", "parts"), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + # Case 3: Streaming artifact-update (singular artifact) + artifact = result.get("artifact") + if artifact and isinstance(artifact, dict) and "parts" in artifact: + self._extract_texts_from_parts( + parts=artifact["parts"], + path=("artifact", "parts"), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + # Case 4: Task with status message + status = result.get("status", {}) + if isinstance(status, dict): + status_message = status.get("message") + if ( + status_message + and isinstance(status_message, dict) + and "parts" in status_message + ): + self._extract_texts_from_parts( + parts=status_message["parts"], + path=("status", "message", "parts"), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + # Case 5: Task with artifacts (plural, array) + artifacts = result.get("artifacts", []) + if artifacts and isinstance(artifacts, list): + for artifact_idx, art in enumerate(artifacts): + if isinstance(art, dict) and "parts" in art: + self._extract_texts_from_parts( + parts=art["parts"], + path=("artifacts", str(artifact_idx), "parts"), + texts_to_check=texts_to_check, + task_mappings=task_mappings, + ) + + def _extract_texts_from_parts( + self, + parts: List[Dict[str, Any]], + path: Tuple[str, ...], + texts_to_check: List[str], + task_mappings: List[Tuple[Tuple[str, ...], int]], + ) -> None: + """Extract text from message parts.""" + for part_idx, part in enumerate(parts): + if part.get("kind") == "text": + text = part.get("text", "") + if text: + texts_to_check.append(text) + task_mappings.append((path, part_idx)) + + def _apply_text_to_path( + self, + result: Dict[Union[str, int], Any], + path: Tuple[str, ...], + part_idx: int, + text: str, + ) -> None: + """Apply guardrailed text back to the specified path in the result.""" + # Navigate to the parts list + current = result + for key in path: + if key.isdigit(): + # Array index + current = current[int(key)] + else: + current = current[key] + + # Update the text in the part + current[part_idx]["text"] = text diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 4d53ae7059..14bbb82808 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -1,26 +1,37 @@ """ MCP Guardrail Handler for Unified Guardrails. -This handler works with the synthetic "messages" payload generated by -`ProxyLogging._convert_mcp_to_llm_format`, which always produces a single user -message whose `content` string encodes the MCP tool name and arguments. The -handler simply feeds that text through the configured guardrail and writes the -result back onto the message. +Converts an MCP call_tool (name + arguments) into a single OpenAI-compatible +tool_call and passes it to apply_guardrail. Works with the synthetic payload +from ProxyLogging._convert_mcp_to_llm_format. + +Note: For MCP tool definitions (schema) -> OpenAI tools=[], see +litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_tool +when you have a full MCP Tool from list_tools. Here we only have the call +payload (name + arguments) so we just build the tool_call. """ from typing import TYPE_CHECKING, Any, Dict, Optional +from mcp.types import Tool as MCPTool + from litellm._logging import verbose_proxy_logger +from litellm.experimental_mcp_client.tools import transform_mcp_tool_to_openai_tool from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.llms.openai import ( + ChatCompletionToolParam, + ChatCompletionToolParamFunctionChunk, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail from mcp.types import CallToolResult + from litellm.integrations.custom_guardrail import CustomGuardrail + class MCPGuardrailTranslationHandler(BaseTranslation): - """Guardrail translation handler for MCP tool calls.""" + """Guardrail translation handler for MCP tool calls (passes a single tool_call to guardrail).""" async def process_input_messages( self, @@ -28,56 +39,51 @@ class MCPGuardrailTranslationHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, ) -> Dict[str, Any]: - messages = data.get("messages") - if not isinstance(messages, list) or not messages: - verbose_proxy_logger.debug("MCP Guardrail: No messages to process") + mcp_tool_name = data.get("mcp_tool_name") or data.get("name") + mcp_arguments = data.get("mcp_arguments") or data.get("arguments") + mcp_tool_description = data.get("mcp_tool_description") or data.get( + "description" + ) + if mcp_arguments is None or not isinstance(mcp_arguments, dict): + mcp_arguments = {} + + if not mcp_tool_name: + verbose_proxy_logger.debug("MCP Guardrail: mcp_tool_name missing") return data - first_message = messages[0] - content: Optional[str] = None - if isinstance(first_message, dict): - content = first_message.get("content") - else: - content = getattr(first_message, "content", None) + # Convert MCP input via transform_mcp_tool_to_openai_tool, then map to litellm + # ChatCompletionToolParam (openai SDK type has incompatible strict/cache_control). + mcp_tool = MCPTool( + name=mcp_tool_name, + description=mcp_tool_description or "", + inputSchema={}, # Call payload has no schema; guardrail gets args from request_data + ) + openai_tool = transform_mcp_tool_to_openai_tool(mcp_tool) + fn = openai_tool["function"] + tool_def: ChatCompletionToolParam = { + "type": "function", + "function": ChatCompletionToolParamFunctionChunk( + name=fn["name"], + description=fn.get("description") or "", + parameters=fn.get("parameters") + or { + "type": "object", + "properties": {}, + "additionalProperties": False, + }, + strict=fn.get("strict", False) or False, # Default to False if None + ), + } + inputs: GenericGuardrailAPIInputs = GenericGuardrailAPIInputs( + tools=[tool_def], + ) - if not isinstance(content, str): - verbose_proxy_logger.debug( - "MCP Guardrail: Message content missing or not a string", - ) - return data - - inputs = GenericGuardrailAPIInputs(texts=[content]) - # Include model information if available - model = data.get("model") - if model: - inputs["model"] = model - guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, ) - guardrailed_texts = ( - guardrailed_inputs.get("texts", []) if guardrailed_inputs else [] - ) - - if guardrailed_texts: - new_content = guardrailed_texts[0] - if isinstance(first_message, dict): - first_message["content"] = new_content - else: - setattr(first_message, "content", new_content) - - verbose_proxy_logger.debug( - "MCP Guardrail: Updated content for tool %s", - data.get("mcp_tool_name"), - ) - else: - verbose_proxy_logger.debug( - "MCP Guardrail: Guardrail returned no text updates for tool %s", - data.get("mcp_tool_name"), - ) - return data async def process_output_response( @@ -87,7 +93,6 @@ class MCPGuardrailTranslationHandler(BaseTranslation): litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, ) -> Any: - # Not implemented: MCP guardrail translation never calls this path today. verbose_proxy_logger.debug( "MCP Guardrail: Output processing not implemented for MCP tools", ) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 4c17a2ff3e..f5b0f152fd 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -65,12 +65,13 @@ from litellm.types.mcp_server.mcp_server_manager import ( from litellm.types.utils import CallTypes try: - from mcp.shared.tool_name_validation import ( # type: ignore - SEP_986_URL, - validate_tool_name, + from mcp.shared.tool_name_validation import ( + validate_tool_name, # type: ignore[reportAssignmentType] ) + from mcp.shared.tool_name_validation import SEP_986_URL except ImportError: from pydantic import BaseModel + SEP_986_URL = "https://github.com/modelcontextprotocol/protocol/blob/main/proposals/0001-tool-name-validation.md" class ToolNameValidationResult(BaseModel): @@ -469,12 +470,12 @@ class MCPServerManager: ) # Update tool name to server name mapping (for both prefixed and base names) - self.tool_name_to_mcp_server_name_mapping[ - base_tool_name - ] = server_prefix - self.tool_name_to_mcp_server_name_mapping[ - prefixed_tool_name - ] = server_prefix + self.tool_name_to_mcp_server_name_mapping[base_tool_name] = ( + server_prefix + ) + self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = ( + server_prefix + ) registered_count += 1 verbose_logger.debug( @@ -1929,7 +1930,9 @@ class MCPServerManager: ) async def _call_tool_via_client(client, params): - return await client.call_tool(params, host_progress_callback=host_progress_callback) + return await client.call_tool( + params, host_progress_callback=host_progress_callback + ) tasks.append( asyncio.create_task(_call_tool_via_client(client, call_tool_params)) @@ -1967,7 +1970,6 @@ class MCPServerManager: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, host_progress_callback: Optional[Callable] = None, - ) -> CallToolResult: """ Call a tool with the given name and arguments diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d93f852f22..eb47fb4f60 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -12,6 +12,7 @@ from litellm.proxy._experimental.mcp_server.utils import merge_mcp_headers from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.mcp import MCPAuth +from litellm.types.utils import CallTypes MCP_AVAILABLE: bool = True try: @@ -28,6 +29,7 @@ router = APIRouter( if MCP_AVAILABLE: from mcp.types import Tool as MCPTool + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -96,6 +98,35 @@ if MCP_AVAILABLE: return _create_tool_response_objects(tools, server.mcp_info) + async def _resolve_allowed_mcp_servers_for_tool_call( + user_api_key_dict: UserAPIKeyAuth, + server_id: str, + ) -> List[MCPServer]: + """Resolve allowed MCP servers for the given user and validate server_id access.""" + auth_contexts = await build_effective_auth_contexts(user_api_key_dict) + allowed_server_ids_set = set() + for auth_context in auth_contexts: + servers = await global_mcp_server_manager.get_allowed_mcp_servers( + user_api_key_auth=auth_context + ) + allowed_server_ids_set.update(servers) + if server_id not in allowed_server_ids_set: + raise HTTPException( + status_code=403, + detail={ + "error": "access_denied", + "message": f"The key is not allowed to access server {server_id}", + }, + ) + allowed_mcp_servers: List[MCPServer] = [] + for allowed_server_id in allowed_server_ids_set: + server = global_mcp_server_manager.get_mcp_server_by_id( + allowed_server_id + ) + if server is not None: + allowed_mcp_servers.append(server) + return allowed_mcp_servers + ######################################################## @router.get("/tools/list", dependencies=[Depends(user_api_key_auth)]) async def list_tool_rest_api( @@ -261,7 +292,14 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) - from litellm.proxy.proxy_server import add_litellm_data_to_request, proxy_config + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + from litellm.proxy.proxy_server import ( + general_settings, + proxy_config, + proxy_logging_obj, + ) try: data = await request.json() @@ -289,11 +327,16 @@ if MCP_AVAILABLE: tool_arguments = data.get("arguments") - data = await add_litellm_data_to_request( - data=data, - request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, + proxy_base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) + data, logging_obj = ( + await proxy_base_llm_response_processor.common_processing_pre_call_logic( + request=request, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + route_type=CallTypes.call_mcp_tool.value, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + ) ) # FIX: Extract MCP auth headers from request @@ -322,35 +365,9 @@ if MCP_AVAILABLE: if "metadata" in data and "user_api_key_auth" in data["metadata"]: data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"] - # Get all auth contexts - auth_contexts = await build_effective_auth_contexts(user_api_key_dict) - - # Collect allowed server IDs from all contexts - allowed_server_ids_set = set() - for auth_context in auth_contexts: - servers = await global_mcp_server_manager.get_allowed_mcp_servers( - user_api_key_auth=auth_context - ) - allowed_server_ids_set.update(servers) - - # Check if the specified server_id is allowed - if server_id not in allowed_server_ids_set: - raise HTTPException( - status_code=403, - detail={ - "error": "access_denied", - "message": f"The key is not allowed to access server {server_id}", - }, - ) - - # Build allowed_mcp_servers list (only include allowed servers) - allowed_mcp_servers: List[MCPServer] = [] - for allowed_server_id in allowed_server_ids_set: - server = global_mcp_server_manager.get_mcp_server_by_id( - allowed_server_id - ) - if server is not None: - allowed_mcp_servers.append(server) + allowed_mcp_servers = await _resolve_allowed_mcp_servers_for_tool_call( + user_api_key_dict, server_id + ) # Call execute_mcp_tool directly (permission checks already done) result = await execute_mcp_tool( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 24727aacd7..12b2d5c4df 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -14,6 +14,7 @@ from fastapi.responses import JSONResponse, StreamingResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.utils import all_litellm_params router = APIRouter() @@ -75,10 +76,7 @@ async def _handle_stream_message( return StreamingResponse(_error_stream(), media_type="application/x-ndjson") - from a2a.types import ( - MessageSendParams, - SendStreamingMessageRequest, - ) + from a2a.types import MessageSendParams, SendStreamingMessageRequest async def stream_response(): try: @@ -208,16 +206,17 @@ async def invoke_agent_a2a( from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, ) - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ( general_settings, proxy_config, + proxy_logging_obj, version, ) body = {} try: body = await request.json() + verbose_proxy_logger.debug(f"A2A request for agent '{agent_id}': {body}") # Validate JSON-RPC format @@ -230,6 +229,16 @@ async def invoke_agent_a2a( method = body.get("method") params = body.get("params", {}) + if params: + # extract any litellm params from the params - eg. 'guardrails' + params_to_remove = [] + for key, value in params.items(): + if key in all_litellm_params: + params_to_remove.append(key) + body[key] = value + for key in params_to_remove: + params.pop(key) + if not A2A_SDK_AVAILABLE: return _jsonrpc_error( request_id, @@ -283,12 +292,18 @@ async def invoke_agent_a2a( ) # Add litellm data (user_api_key, user_id, team_id, etc.) - data = await add_litellm_data_to_request( - data=body, + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + + processor = ProxyBaseLLMRequestProcessing(data=body) + data, logging_obj = await processor.common_processing_pre_call_logic( request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + route_type="asend_message", version=version, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3e07904577..f33b241226 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -267,7 +267,9 @@ def _override_openai_response_model( hidden_params = getattr(response_obj, "_hidden_params", {}) or {} if isinstance(hidden_params, dict): fallback_headers = hidden_params.get("additional_headers", {}) or {} - attempted_fallbacks = fallback_headers.get("x-litellm-attempted-fallbacks", None) + attempted_fallbacks = fallback_headers.get( + "x-litellm-attempted-fallbacks", None + ) if attempted_fallbacks is not None and attempted_fallbacks > 0: # A fallback occurred - preserve the actual model that was used verbose_proxy_logger.debug( @@ -517,6 +519,8 @@ class ProxyBaseLLMRequestProcessing: "aget_interaction", "adelete_interaction", "acancel_interaction", + "asend_message", + "call_mcp_tool", ], version: Optional[str] = None, user_model: Optional[str] = None, @@ -837,7 +841,9 @@ class ProxyBaseLLMRequestProcessing: # aliasing/routing, but the OpenAI-compatible response `model` field should reflect # what the client sent. if requested_model_from_client: - self.data["_litellm_client_requested_model"] = requested_model_from_client + self.data["_litellm_client_requested_model"] = ( + requested_model_from_client + ) if route_type == "allm_passthrough_route": # Check if response is an async generator if self._is_streaming_response(response): @@ -1409,9 +1415,9 @@ class ProxyBaseLLMRequestProcessing: # Add cache-related fields to **params (handled by Usage.__init__) if cache_creation_input_tokens is not None: - usage_kwargs[ - "cache_creation_input_tokens" - ] = cache_creation_input_tokens + usage_kwargs["cache_creation_input_tokens"] = ( + cache_creation_input_tokens + ) if cache_read_input_tokens is not None: usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index a0ca324411..68f9dfd7ab 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -5,7 +5,7 @@ This module provides a guardrail that executes user-defined Python-like code to implement custom guardrail logic. The code runs in a sandboxed environment with access to LiteLLM-provided primitives for common guardrail operations. -Example custom code: +Example custom code (sync): def apply_guardrail(inputs, request_data, input_type): '''Block messages containing SSNs''' @@ -13,8 +13,22 @@ Example custom code: if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): return block("Social Security Number detected") return allow() + +Example custom code (async with HTTP): + + async def apply_guardrail(inputs, request_data, input_type): + '''Call external moderation API''' + for text in inputs["texts"]: + response = await http_post( + "https://api.example.com/moderate", + body={"text": text} + ) + if response["success"] and response["body"].get("flagged"): + return block("Content flagged by moderation API") + return allow() """ +import asyncio import threading from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast @@ -101,6 +115,9 @@ class CustomCodeGuardrail(CustomGuardrail): GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call, + GuardrailEventHooks.pre_mcp_call, + GuardrailEventHooks.during_mcp_call, + GuardrailEventHooks.logging_only, ] super().__init__( @@ -175,6 +192,13 @@ class CustomCodeGuardrail(CustomGuardrail): This method calls the user-defined apply_guardrail function and processes its result to determine the appropriate action. + The user-defined function can be either sync or async: + - Sync: def apply_guardrail(inputs, request_data, input_type): ... + - Async: async def apply_guardrail(inputs, request_data, input_type): ... + + Async functions are recommended when using http_request, http_get, or + http_post primitives to avoid blocking the event loop. + Args: inputs: Dictionary containing texts, images, tool_calls request_data: The original request data with metadata @@ -188,6 +212,7 @@ class CustomCodeGuardrail(CustomGuardrail): HTTPException: If content is blocked CustomCodeExecutionError: If execution fails """ + if self._compiled_function is None: if self._compile_error: raise CustomCodeExecutionError( @@ -201,9 +226,13 @@ class CustomCodeGuardrail(CustomGuardrail): # Prepare request_data with safe subset of information safe_request_data = self._prepare_safe_request_data(request_data) - # Execute the custom function + # Execute the custom function - handle both sync and async functions result = self._compiled_function(inputs, safe_request_data, input_type) + # If the function is async (returns a coroutine), await it + if asyncio.iscoroutine(result): + result = await result + # Process the result return self._process_result( result=result, diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index 695e59977c..de7690635d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -10,7 +10,11 @@ import re from typing import Any, Dict, List, Optional, Tuple, Type, Union from urllib.parse import urlparse +import httpx + from litellm._logging import verbose_proxy_logger +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.llms.custom_http import httpxSpecialProvider # ============================================================================= # Result Types - Used by Starlark code to return guardrail decisions @@ -349,6 +353,228 @@ def get_url_domain(url: str) -> Optional[str]: return None +# ============================================================================= +# HTTP Request Primitives (Async) +# ============================================================================= + +# Default timeout for HTTP requests (in seconds) +_HTTP_DEFAULT_TIMEOUT = 30.0 + +# Maximum allowed timeout (in seconds) +_HTTP_MAX_TIMEOUT = 60.0 + + +def _http_error_response(error: str) -> Dict[str, Any]: + """Create a standardized error response for HTTP requests.""" + return { + "status_code": 0, + "body": None, + "headers": {}, + "success": False, + "error": error, + } + + +def _http_success_response(response: httpx.Response) -> Dict[str, Any]: + """Create a standardized success response from an httpx Response.""" + parsed_body: Any + try: + parsed_body = response.json() + except (json.JSONDecodeError, ValueError): + parsed_body = response.text + + return { + "status_code": response.status_code, + "body": parsed_body, + "headers": dict(response.headers), + "success": 200 <= response.status_code < 300, + "error": None, + } + + +def _prepare_http_body( + body: Optional[Any], +) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + """Prepare body arguments for HTTP request - returns (json_body, data_body).""" + if body is None: + return None, None + if isinstance(body, dict): + return body, None + if isinstance(body, list): + return None, json.dumps(body) + if isinstance(body, str): + return None, body + return None, str(body) + + +async def http_request( + url: str, + method: str = "GET", + headers: Optional[Dict[str, str]] = None, + body: Optional[Any] = None, + timeout: Optional[float] = None, +) -> Dict[str, Any]: + """ + Make an async HTTP request to an external service. + + This function allows custom guardrails to call external APIs for + additional validation, content moderation, or data enrichment. + + Uses LiteLLM's global cached AsyncHTTPHandler for connection pooling + and better performance. + + Args: + url: The URL to request + method: HTTP method (GET, POST, PUT, DELETE, PATCH). Defaults to GET. + headers: Optional dict of HTTP headers + body: Optional request body (will be JSON-encoded if dict/list) + timeout: Optional timeout in seconds (default: 30, max: 60) + + Returns: + Dict containing: + - status_code: HTTP status code + - body: Response body (parsed as JSON if possible, otherwise string) + - headers: Response headers as dict + - success: True if status code is 2xx + - error: Error message if request failed, None otherwise + + Example: + # Simple GET request + response = await http_request("https://api.example.com/check") + if response["success"]: + data = response["body"] + + # POST request with JSON body + response = await http_request( + "https://api.example.com/moderate", + method="POST", + headers={"Authorization": "Bearer token"}, + body={"text": "content to check"} + ) + """ + # Validate URL + if not is_valid_url(url): + return _http_error_response(f"Invalid URL: {url}") + + # Validate and normalize method + method = method.upper() + allowed_methods = {"GET", "POST", "PUT", "DELETE", "PATCH"} + if method not in allowed_methods: + return _http_error_response( + f"Invalid HTTP method: {method}. Allowed: {', '.join(allowed_methods)}" + ) + + # Apply timeout limits + if timeout is None: + timeout = _HTTP_DEFAULT_TIMEOUT + else: + timeout = min(max(0.1, timeout), _HTTP_MAX_TIMEOUT) + + # Get the global cached async HTTP client + client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + params={"timeout": httpx.Timeout(timeout=timeout, connect=5.0)}, + ) + + try: + response = await _execute_http_request( + client, method, url, headers, body, timeout + ) + return _http_success_response(response) + + except httpx.TimeoutException as e: + verbose_proxy_logger.warning(f"Custom code http_request timeout: {e}") + return _http_error_response(f"Request timeout after {timeout}s") + except httpx.HTTPStatusError as e: + # Return the response even for non-2xx status codes + return _http_success_response(e.response) + except httpx.RequestError as e: + verbose_proxy_logger.warning(f"Custom code http_request error: {e}") + return _http_error_response(f"Request failed: {str(e)}") + except Exception as e: + verbose_proxy_logger.warning(f"Custom code http_request unexpected error: {e}") + return _http_error_response(f"Unexpected error: {str(e)}") + + +async def _execute_http_request( + client: Any, + method: str, + url: str, + headers: Optional[Dict[str, str]], + body: Optional[Any], + timeout: float, +) -> httpx.Response: + """Execute the HTTP request using the appropriate client method.""" + json_body, data_body = _prepare_http_body(body) + + if method == "GET": + return await client.get(url=url, headers=headers) + elif method == "POST": + return await client.post( + url=url, headers=headers, json=json_body, data=data_body, timeout=timeout + ) + elif method == "PUT": + return await client.put( + url=url, headers=headers, json=json_body, data=data_body, timeout=timeout + ) + elif method == "DELETE": + return await client.delete( + url=url, headers=headers, json=json_body, data=data_body, timeout=timeout + ) + elif method == "PATCH": + return await client.patch( + url=url, headers=headers, json=json_body, data=data_body, timeout=timeout + ) + else: + raise ValueError(f"Unsupported HTTP method: {method}") + + +async def http_get( + url: str, + headers: Optional[Dict[str, str]] = None, + timeout: Optional[float] = None, +) -> Dict[str, Any]: + """ + Make an async HTTP GET request. + + Convenience wrapper around http_request for GET requests. + + Args: + url: The URL to request + headers: Optional dict of HTTP headers + timeout: Optional timeout in seconds + + Returns: + Same as http_request + """ + return await http_request(url=url, method="GET", headers=headers, timeout=timeout) + + +async def http_post( + url: str, + body: Optional[Any] = None, + headers: Optional[Dict[str, str]] = None, + timeout: Optional[float] = None, +) -> Dict[str, Any]: + """ + Make an async HTTP POST request. + + Convenience wrapper around http_request for POST requests. + + Args: + url: The URL to request + body: Optional request body (will be JSON-encoded if dict/list) + headers: Optional dict of HTTP headers + timeout: Optional timeout in seconds + + Returns: + Same as http_request + """ + return await http_request( + url=url, method="POST", headers=headers, body=body, timeout=timeout + ) + + # ============================================================================= # Code Detection Primitives # ============================================================================= @@ -575,6 +801,10 @@ def get_custom_code_primitives() -> Dict[str, Any]: "is_valid_url": is_valid_url, "all_urls_valid": all_urls_valid, "get_url_domain": get_url_domain, + # HTTP (async) + "http_request": http_request, + "http_get": http_get, + "http_post": http_post, # Code detection "detect_code": detect_code, "detect_code_languages": detect_code_languages, diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index f07f65d10f..cc05358baf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -51,6 +51,7 @@ class UnifiedLLMGuardrails(CustomLogger): Runs on only Input Use this if you want to MODIFY the input """ + global endpoint_guardrail_translation_mappings from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -66,6 +67,7 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.pre_mcp_call + if ( guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True 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 73d96c590e..1b8a8c034b 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, testCustomCodeGuardrail } from "../../networking"; +import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "../../networking"; import NotificationsManager from "../../molecules/notifications_manager"; const { Panel } = Collapse; @@ -19,7 +19,7 @@ const { TextArea } = Input; const CODE_TEMPLATES = { empty: { name: "Empty Template", - code: `def apply_guardrail(inputs, request_data, input_type): + code: `async def apply_guardrail(inputs, request_data, input_type): # inputs: {texts, images, tools, tool_calls, structured_messages, model} # request_data: {model, user_id, team_id, end_user_id, metadata} # input_type: "request" or "response" @@ -68,6 +68,27 @@ const CODE_TEMPLATES = { return block("Response missing required fields") return allow()`, }, + externalAPI: { + name: "External API Check (async)", + code: `async def apply_guardrail(inputs, request_data, input_type): + # Call an external moderation API (async for non-blocking) + for text in inputs["texts"]: + response = await http_post( + "https://api.example.com/moderate", + body={"text": text, "user_id": request_data["user_id"]}, + headers={"Authorization": "Bearer YOUR_API_KEY"}, + timeout=10 + ) + + if not response["success"]: + # API call failed, allow by default or block + return allow() + + if response["body"].get("flagged"): + return block(response["body"].get("reason", "Content flagged")) + + return allow()`, + }, }; // Available primitives organized by category @@ -77,6 +98,11 @@ const PRIMITIVES = { { name: "block(reason)", desc: "Reject with message" }, { name: "modify(texts=[], images=[], tool_calls=[])", desc: "Transform content" }, ], + "HTTP Requests (async)": [ + { name: "await http_request(url, method, headers, body)", desc: "Make async HTTP request" }, + { name: "await http_get(url, headers)", desc: "Async GET request" }, + { name: "await http_post(url, body, headers)", desc: "Async POST request" }, + ], "Regex Functions": [ { name: "regex_match(text, pattern)", desc: "Returns True if pattern found" }, { name: "regex_replace(text, pattern, replacement)", desc: "Replace all matches" }, @@ -111,13 +137,30 @@ const MODE_OPTIONS = [ { value: "post_call", label: "post_call (Response)" }, { value: "during_call", label: "during_call (Parallel)" }, { value: "logging_only", label: "logging_only" }, + { value: "pre_mcp_call", label: "pre_mcp_call (Before MCP Tool Call)" }, + { value: "post_mcp_call", label: "post_mcp_call (After MCP Tool Call)" }, + { value: "during_mcp_call", label: "during_mcp_call (During MCP Tool Call)" }, ]; +// Data for editing an existing guardrail +export interface EditGuardrailData { + guardrail_id: string; + guardrail_name: string; + litellm_params: { + mode?: string | string[]; + default_on?: boolean; + custom_code?: string; + [key: string]: any; + }; +} + interface CustomCodeModalProps { visible: boolean; onClose: () => void; onSuccess: () => void; accessToken: string | null; + /** If provided, the modal will be in edit mode */ + editData?: EditGuardrailData | null; } const CustomCodeModal: React.FC = ({ @@ -125,16 +168,72 @@ const CustomCodeModal: React.FC = ({ onClose, onSuccess, accessToken, + editData, }) => { + const isEditMode = !!editData; const [guardrailName, setGuardrailName] = useState(""); - const [mode, setMode] = useState("pre_call"); + const [mode, setMode] = useState(["pre_call"]); const [defaultOn, setDefaultOn] = useState(false); const [selectedTemplate, setSelectedTemplate] = useState("empty"); const [code, setCode] = useState(CODE_TEMPLATES.empty.code); const [isSaving, setIsSaving] = useState(false); const [isTesting, setIsTesting] = useState(false); const [testExpanded, setTestExpanded] = useState(false); - const [testInput, setTestInput] = useState('{"texts": ["Hello, my SSN is 123-45-6789"], "images": [], "tools": [], "tool_calls": [], "structured_messages": [], "model": "gpt-4"}'); + + // Test input examples for pre_call and post_call + const TEST_INPUT_EXAMPLES = { + pre_call: { + name: "Pre-call (Request)", + data: { + texts: ["Hello, my SSN is 123-45-6789"], + images: [], + tools: [ + { + type: "function", + function: { + name: "get_weather", + description: "Get the current weather in a location", + parameters: { + type: "object", + properties: { + location: { type: "string", description: "City name" } + }, + required: ["location"] + } + } + } + ], + tool_calls: [], + structured_messages: [ + { role: "system", content: "You are a helpful assistant." }, + { role: "user", content: "Hello, my SSN is 123-45-6789" } + ], + model: "gpt-4" + } + }, + post_call: { + name: "Post-call (Response)", + data: { + texts: ["The weather in San Francisco is 72°F and sunny."], + images: [], + tools: [], + tool_calls: [ + { + id: "call_abc123", + type: "function", + function: { + name: "get_weather", + arguments: "{\"location\": \"San Francisco\"}" + } + } + ], + structured_messages: [], + model: "gpt-4" + } + } + }; + + const [testInput, setTestInput] = useState(JSON.stringify(TEST_INPUT_EXAMPLES.pre_call.data, null, 2)); const [testResult, setTestResult] = useState(null); const [copiedPrimitive, setCopiedPrimitive] = useState(null); const textareaRef = useRef(null); @@ -145,18 +244,35 @@ const CustomCodeModal: React.FC = ({ setCode(CODE_TEMPLATES[templateKey as keyof typeof CODE_TEMPLATES].code); }; - // Reset form when modal opens + // Normalize mode from API (string or string[]) to string[] + const normalizeMode = (m: string | string[] | undefined): string[] => { + if (m === undefined || m === null) return ["pre_call"]; + if (Array.isArray(m)) return m.length ? m : ["pre_call"]; + return [m]; + }; + + // Reset form when modal opens or editData changes useEffect(() => { if (visible) { - setGuardrailName(""); - setMode("pre_call"); - setDefaultOn(false); - setSelectedTemplate("empty"); - setCode(CODE_TEMPLATES.empty.code); + if (editData) { + // Edit mode: populate with existing data + setGuardrailName(editData.guardrail_name || ""); + setMode(normalizeMode(editData.litellm_params?.mode)); + setDefaultOn(editData.litellm_params?.default_on || false); + setCode(editData.litellm_params?.custom_code || CODE_TEMPLATES.empty.code); + setSelectedTemplate(""); // No template selected in edit mode + } else { + // Create mode: reset to defaults + setGuardrailName(""); + setMode(["pre_call"]); + setDefaultOn(false); + setSelectedTemplate("empty"); + setCode(CODE_TEMPLATES.empty.code); + } setTestResult(null); setTestExpanded(false); } - }, [visible]); + }, [visible, editData]); // Copy primitive to clipboard const copyPrimitive = async (primitive: string) => { @@ -184,7 +300,7 @@ const CustomCodeModal: React.FC = ({ } }; - // Save guardrail + // Save guardrail (create or update) const handleSave = async () => { if (!guardrailName.trim()) { NotificationsManager.fromBackend("Please enter a guardrail name"); @@ -201,25 +317,53 @@ const CustomCodeModal: React.FC = ({ setIsSaving(true); try { - const guardrailData = { - guardrail_name: guardrailName, - litellm_params: { - guardrail: "custom_code", - mode: mode, - default_on: defaultOn, - custom_code: code, - }, - guardrail_info: {}, - }; + if (isEditMode && editData) { + // Update existing guardrail + const updateData: any = { + litellm_params: { + custom_code: code, + }, + }; - await createGuardrailCall(accessToken, guardrailData); - NotificationsManager.success("Custom code guardrail created successfully"); + // Only include changed fields + if (guardrailName !== editData.guardrail_name) { + updateData.guardrail_name = guardrailName; + } + const existingMode = normalizeMode(editData.litellm_params?.mode); + const modeChanged = + mode.length !== existingMode.length || + mode.some((m, i) => m !== existingMode[i]); + if (modeChanged) { + updateData.litellm_params.mode = mode; + } + if (defaultOn !== editData.litellm_params?.default_on) { + updateData.litellm_params.default_on = defaultOn; + } + + await updateGuardrailCall(accessToken, editData.guardrail_id, updateData); + NotificationsManager.success("Custom code guardrail updated successfully"); + } else { + // Create new guardrail + const guardrailData = { + guardrail_name: guardrailName, + litellm_params: { + guardrail: "custom_code", + mode: mode, + default_on: defaultOn, + custom_code: code, + }, + guardrail_info: {}, + }; + + await createGuardrailCall(accessToken, guardrailData); + NotificationsManager.success("Custom code guardrail created successfully"); + } onSuccess(); onClose(); } catch (error) { - console.error("Failed to create guardrail:", error); + console.error("Failed to save guardrail:", error); NotificationsManager.fromBackend( - "Failed to create guardrail: " + (error instanceof Error ? error.message : String(error)) + `Failed to ${isEditMode ? "update" : "create"} guardrail: ` + (error instanceof Error ? error.message : String(error)) ); } finally { setIsSaving(false); @@ -252,10 +396,20 @@ const CustomCodeModal: React.FC = ({ parsedInput.texts = []; } + // Use first request-like or response-like mode for test input_type + const requestModes = ["pre_call", "pre_mcp_call"]; + const responseModes = ["post_call", "post_mcp_call"]; + const testInputType: "request" | "response" = + mode.some((m) => requestModes.includes(m)) + ? "request" + : mode.some((m) => responseModes.includes(m)) + ? "response" + : "request"; + const response = await testCustomCodeGuardrail(accessToken, { custom_code: code, test_input: parsedInput, - input_type: mode as "request" | "response", + input_type: testInputType, request_data: { model: "test-model", metadata: {}, @@ -289,7 +443,7 @@ const CustomCodeModal: React.FC = ({ open={visible} onCancel={onClose} footer={null} - width={1200} + width={1400} className="custom-code-modal" closable={true} destroyOnClose @@ -297,7 +451,9 @@ const CustomCodeModal: React.FC = ({
{/* Header */}
-

Create Custom Guardrail

+

+ {isEditMode ? "Edit Custom Guardrail" : "Create Custom Guardrail"} +

Define custom logic using Python-like syntax

@@ -311,14 +467,16 @@ const CustomCodeModal: React.FC = ({ placeholder="e.g., block-pii-custom" />
-
- +
+