Guardrail API V2 - user api key metadata, session id, specify input type (request/response), image support

cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py

@@ -424,17 +424,20 @@ async def apply_guardrail(
 
 
 class LitellmBasicGuardrailRequest(BaseModel):
-    text: str
-    request_body: Dict[str, Any] = Field(default_factory=dict)
+    texts: List[str]
+    images: Optional[List[str]] = None
+    request_data: Dict[str, Any] = Field(default_factory=dict)
     additional_provider_specific_params: Dict[str, Any] = Field(default_factory=dict)
+    input_type: Literal["request", "response"]
 
 
 class LitellmBasicGuardrailResponse(BaseModel):
     action: Literal[
         "BLOCKED", "NONE", "GUARDRAIL_INTERVENED"
     ]  # BLOCKED = litellm will raise an error, NONE = litellm will continue, GUARDRAIL_INTERVENED = litellm will continue, but the text was modified by the guardrail
     blocked_reason: Optional[str] = None  # only if action is BLOCKED, otherwise None
-    text: Optional[str] = None
+    texts: Optional[List[str]] = None
+    images: Optional[List[str]] = None
 
 
 @app.post(
@@ -457,14 +460,17 @@ async def beta_litellm_basic_guardrail_api(
         LitellmBasicGuardrailResponse with analysis results
     """
     print(f"request: {request}")
-    if "ishaan" in request.text.lower():
+    if any("ishaan" in text.lower() for text in request.texts):
         return LitellmBasicGuardrailResponse(
             action="BLOCKED", blocked_reason="Ishaan is not allowed"
         )
-    elif "pii_value" in request.text:
+    elif any("pii_value" in text for text in request.texts):
         return LitellmBasicGuardrailResponse(
             action="GUARDRAIL_INTERVENED",
-            text=request.text.replace("pii_value", "pii_value_redacted"),
+            texts=[
+                text.replace("pii_value", "pii_value_redacted")
+                for text in request.texts
+            ],
         )
     return LitellmBasicGuardrailResponse(action="NONE")
 

docs/my-website/docs/adding_provider/generic_guardrail_api.md

@@ -17,15 +17,16 @@ The **Generic Guardrail API** lets you integrate with LiteLLM **instantly** by i
 1. **No PR Needed** - Deploy and integrate immediately
 2. **Universal Support** - Works across ALL LiteLLM endpoints (chat, embeddings, image generation, etc.)
 3. **Simple Contract** - One endpoint, three response types
-4. **Custom Parameters** - Pass provider-specific params via config
-5. **Full Control** - You own and maintain your guardrail API
+4. **Multi-Modal Support** - Handle both text and images in requests/responses
+5. **Custom Parameters** - Pass provider-specific params via config
+6. **Full Control** - You own and maintain your guardrail API
 
 ## How It Works
 
-1. LiteLLM extracts text from any request (chat messages, embeddings, image prompts, etc.)
-2. Sends extracted text + original request to your API endpoint
+1. LiteLLM extracts text and images from any request (chat messages, embeddings, image prompts, etc.)
+2. Sends extracted content + metadata to your API endpoint
 3. Your API responds with: `BLOCKED`, `NONE`, or `GUARDRAIL_INTERVENED`
-4. LiteLLM enforces the decision
+4. LiteLLM enforces the decision and applies any modifications
 
 ## API Contract
 
@@ -37,8 +38,21 @@ Implement `POST /beta/litellm_basic_guardrail_api`
 
 ```json
 {
-  "text": "extracted text from the request",
-  "request_body": {},  // full original request for context
+  "texts": ["extracted text from the request"],  // array of text strings
+  "images": ["base64_encoded_image_data"],  // optional array of images
+  "request_data": {
+    "user_api_key_hash": "hash of the litellm virtual key used",
+    "user_api_key_alias": "alias of the litellm virtual key used",
+    "user_api_key_user_id": "user id associated with the litellm virtual key used",
+    "user_api_key_user_email": "user email associated with the litellm virtual key used",
+    "user_api_key_team_id": "team id associated with the litellm virtual key used",
+    "user_api_key_team_alias": "team alias associated with the litellm virtual key used",
+    "user_api_key_end_user_id": "end user id associated with the litellm virtual key used",
+    "user_api_key_org_id": "org id associated with the litellm virtual key used"
+  },
+  "input_type": "request",  // "request" or "response"
+  "litellm_call_id": "unique_call_id",  // the call id of the individual LLM call
+  "litellm_trace_id": "trace_id",  // the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation
   "additional_provider_specific_params": {
     // your custom params from config
   }
@@ -51,14 +65,15 @@ Implement `POST /beta/litellm_basic_guardrail_api`
 {
   "action": "BLOCKED" | "NONE" | "GUARDRAIL_INTERVENED",
   "blocked_reason": "why content was blocked",  // required if action=BLOCKED
-  "text": "modified text"  // required if action=GUARDRAIL_INTERVENED
+  "texts": ["modified text"],  // optional array of modified text strings
+  "images": ["modified_base64_image"]  // optional array of modified images
 }
 ```
 
 **Actions:**
 - `BLOCKED` - LiteLLM raises error and blocks request
 - `NONE` - Request proceeds unchanged  
-- `GUARDRAIL_INTERVENED` - Request proceeds with modified text
+- `GUARDRAIL_INTERVENED` - Request proceeds with modified texts/images (provide `texts` and/or `images` fields)
 
 ## LiteLLM Configuration
 
@@ -116,27 +131,34 @@ See [mock_bedrock_guardrail_server.py](https://github.com/BerriAI/litellm/blob/m
 ```python
 from fastapi import FastAPI
 from pydantic import BaseModel
+from typing import List, Optional, Dict, Any
 
 app = FastAPI()
 
 class GuardrailRequest(BaseModel):
-    text: str
-    request_body: dict
-    additional_provider_specific_params: dict
+    texts: List[str]
+    images: Optional[List[str]] = None
+    request_data: Dict[str, Any]
+    input_type: str  # "request" or "response"
+    litellm_call_id: Optional[str] = None
+    litellm_trace_id: Optional[str] = None
+    additional_provider_specific_params: Dict[str, Any]
 
 class GuardrailResponse(BaseModel):
     action: str  # BLOCKED, NONE, or GUARDRAIL_INTERVENED
-    blocked_reason: str | None = None
-    text: str | None = None
+    blocked_reason: Optional[str] = None
+    texts: Optional[List[str]] = None
+    images: Optional[List[str]] = None
 
 @app.post("/beta/litellm_basic_guardrail_api")
 async def apply_guardrail(request: GuardrailRequest):
     # Your guardrail logic here
-    if "badword" in request.text.lower():
-        return GuardrailResponse(
-            action="BLOCKED",
-            blocked_reason="Content contains prohibited terms"
-        )
+    for text in request.texts:
+        if "badword" in text.lower():
+            return GuardrailResponse(
+                action="BLOCKED",
+                blocked_reason="Content contains prohibited terms"
+            )
 
     return GuardrailResponse(action="NONE")
 ```

litellm/integrations/custom_guardrail.py

@@ -1,5 +1,16 @@
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Type, Union, get_args
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+    get_args,
+)
 
 from litellm._logging import verbose_logger
 from litellm.caching import DualCache
@@ -20,6 +31,8 @@
     StandardLoggingGuardrailInformation,
 )
 
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 dc = DualCache()
 
 
@@ -437,30 +450,31 @@ def _append_guardrail_info(container: dict) -> None:
 
     async def apply_guardrail(
         self,
-        text: str,
-        language: Optional[str] = None,
-        entities: Optional[List[PiiEntityType]] = None,
-        request_data: Optional[dict] = None,
-    ) -> str:
+        texts: List[str],
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional["LiteLLMLoggingObj"] = None,
+        images: Optional[List[str]] = None,
+    ) -> Tuple[List[str], Optional[List[str]]]:
         """
         Apply your guardrail logic to the given text
 
         Args:
-            text: The text to apply the guardrail to
-            language: The language of the text
-            entities: The entities to mask, optional
-            request_data: The request data dictionary to store guardrail metadata
+            texts: The texts to apply the guardrail to
+            images: The images to apply the guardrail to
+            request_data: The request data dictionary - containing user api key metadata (e.g. user_id, team_id, etc.)
+            input_type: The type of input to apply the guardrail to - "request" or "response"
 
         Any of the custom guardrails can override this method to provide custom guardrail logic
 
-        Returns the text with the guardrail applied
+        Returns the texts with the guardrail applied and the images with the guardrail applied (if any)
 
         Raises:
             Exception:
                 - If the guardrail raises an exception
 
         """
-        return text
+        return texts, images
 
     def _process_response(
         self,