diff --git a/docs/my-website/docs/adding_provider/adding_guardrail_support.md b/docs/my-website/docs/adding_provider/adding_guardrail_support.md new file mode 100644 index 0000000000..2646b626ab --- /dev/null +++ b/docs/my-website/docs/adding_provider/adding_guardrail_support.md @@ -0,0 +1,412 @@ +# Adding Guardrail Support to Endpoints + +This guide explains how to add guardrail translation support to new LiteLLM endpoints (e.g., Chat Completions, Responses API, etc.). + +## When to Add Guardrail Support + +Add guardrail support when: +- You're creating a new LiteLLM endpoint (e.g., a new API format) +- You want to enable guardrails for an existing endpoint that doesn't support them +- You need custom text extraction logic for a specific message format + +## Directory Structure + +Guardrail handlers follow this structure: + +``` +litellm/llms/{provider}/{endpoint}/guardrail_translation/ +├── __init__.py # Exports handler and registers call types +├── handler.py # Main handler implementation +└── README.md # Documentation (optional but recommended) +``` + +### Example Structures + +**OpenAI Chat Completions:** +``` +litellm/llms/openai/chat/guardrail_translation/ +├── __init__.py +├── handler.py +└── README.md +``` + +**OpenAI Responses API:** +``` +litellm/llms/openai/responses/guardrail_translation/ +├── __init__.py +├── handler.py +└── README.md +``` + +**Anthropic Messages:** +``` +litellm/llms/anthropic/chat/guardrail_translation/ +├── __init__.py +└── handler.py +``` + +## Step-by-Step Implementation + +### Step 1: Create the Handler Class + +Create `handler.py` that inherits from `BaseTranslation`: + +```python +""" +{Provider} {Endpoint} Handler for Unified Guardrails + +This module provides guardrail translation support for {Provider}'s {Endpoint} format. +""" + +import asyncio +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.utils import ModelResponse # Or appropriate response type + + +class MyEndpointHandler(BaseTranslation): + """ + Handler for processing {Endpoint} with guardrails. + + This class provides methods to: + 1. Process input (pre-call hook) + 2. Process output response (post-call hook) + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process input by applying guardrails to text content. + + Args: + data: Request data dictionary + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified data with guardrails applied + """ + # Your implementation here + pass + + async def process_output_response( + self, + response: Any, # Use appropriate response type + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process output response by applying guardrails to text content. + + Args: + response: API response object + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified response with guardrails applied + """ + # Your implementation here + pass +``` + +### Step 2: Implement Core Methods + +#### A. Process Input Messages + +Extract text from input, apply guardrails, and map back: + +```python +async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", +) -> Any: + """Process input messages by applying guardrails to text content.""" + # 1. Get input data from request + messages = data.get("messages") # or appropriate field + if messages is None: + return data + + # 2. Extract text and create tasks + tasks = [] + task_mappings: List[Tuple[int, Optional[int]]] = [] + + for msg_idx, message in enumerate(messages): + await self._extract_input_text_and_create_tasks( + message=message, + msg_idx=msg_idx, + tasks=tasks, + task_mappings=task_mappings, + guardrail_to_apply=guardrail_to_apply, + ) + + # 3. Run all guardrail tasks in parallel + if tasks: + responses = await asyncio.gather(*tasks) + + # 4. Map responses back to original structure + await self._apply_guardrail_responses_to_input( + messages=messages, + responses=responses, + task_mappings=task_mappings, + ) + + return data +``` + +#### B. Process Output Response + +Extract text from response, apply guardrails, and update: + +```python +async def process_output_response( + self, + response: "ModelResponse", + guardrail_to_apply: "CustomGuardrail", +) -> Any: + """Process output response by applying guardrails to text content.""" + # 1. Check if response has text to process + if not self._has_text_content(response): + return response + + # 2. Extract text and create tasks + tasks = [] + task_mappings: List[Tuple[int, Optional[int]]] = [] + + for idx, item in enumerate(response.choices): # or appropriate field + await self._extract_output_text_and_create_tasks( + item=item, + idx=idx, + tasks=tasks, + task_mappings=task_mappings, + guardrail_to_apply=guardrail_to_apply, + ) + + # 3. Run all guardrail tasks in parallel + if tasks: + responses = await asyncio.gather(*tasks) + + # 4. Update response with guardrailed text + await self._apply_guardrail_responses_to_output( + response=response, + responses=responses, + task_mappings=task_mappings, + ) + + return response +``` + +### Step 3: Create Helper Methods + +Implement helper methods for text extraction and mapping: + +```python +async def _extract_input_text_and_create_tasks( + self, + message: Dict[str, Any], + msg_idx: int, + tasks: List, + task_mappings: List[Tuple[int, Optional[int]]], + guardrail_to_apply: "CustomGuardrail", +) -> None: + """Extract text content from a message and create guardrail tasks.""" + content = message.get("content") + if content is None: + return + + if isinstance(content, str): + # Simple string content + tasks.append(guardrail_to_apply.apply_guardrail(text=content)) + task_mappings.append((msg_idx, None)) + elif isinstance(content, list): + # List content (e.g., multimodal) + for content_idx, content_item in enumerate(content): + if isinstance(content_item, dict): + text_str = content_item.get("text") + if text_str: + tasks.append(guardrail_to_apply.apply_guardrail(text=text_str)) + task_mappings.append((msg_idx, int(content_idx))) + +async def _apply_guardrail_responses_to_input( + self, + messages: List[Dict[str, Any]], + responses: List[str], + task_mappings: List[Tuple[int, Optional[int]]], +) -> None: + """Apply guardrail responses back to input messages.""" + for task_idx, guardrail_response in enumerate(responses): + msg_idx, content_idx = task_mappings[task_idx] + + if content_idx is None: + # String content + messages[msg_idx]["content"] = guardrail_response + else: + # List content + messages[msg_idx]["content"][content_idx]["text"] = guardrail_response + +def _has_text_content(self, response: Any) -> bool: + """Check if response has any text content to process.""" + # Implement based on your response structure + return True # or appropriate logic +``` + +### Step 4: Register the Handler + +Create `__init__.py` to register the handler with call types: + +```python +"""My Endpoint handler for Unified Guardrails.""" + +from litellm.llms.{provider}/{endpoint}/guardrail_translation.handler import ( + MyEndpointHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.my_endpoint: MyEndpointHandler, + CallTypes.amy_endpoint: MyEndpointHandler, # async version if applicable +} + +__all__ = ["guardrail_translation_mappings"] +``` + +**Important:** Make sure your `CallTypes` are defined in `litellm/types/utils.py`. + +### Step 5: Add Documentation + +Create `README.md` with usage examples and format details: + +```markdown +# {Provider} {Endpoint} Guardrail Translation Handler + +Handler for processing {Provider}'s {Endpoint} with guardrails. + +## Overview + +This handler processes {Endpoint} input/output by: +1. Extracting text from messages/responses +2. Applying guardrails to text content +3. Mapping guardrailed text back to original structure + +## Data Format + +### Input Format +```json +{ + "field": "value", + "messages": [...] +} +``` + +### Output Format +```json +{ + "field": "value", + "output": [...] +} +``` + +## Usage + +The handler is automatically discovered and applied when guardrails are used with this endpoint. + +```bash +curl -X POST 'http://localhost:4000/{my_endpoint}' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "guardrails": ["test"] +}' + +``` +## Extension + +Override these methods to customize behavior: +- `_extract_input_text_and_create_tasks()`: Custom text extraction +- `_apply_guardrail_responses_to_input()`: Custom response mapping +- `_has_text_content()`: Custom content detection +``` + +### Step 6: Add Unit Tests + +Create comprehensive tests in `tests/test_litellm/llms/{provider}/{endpoint}/`: + +```python +""" +Unit tests for {Provider} {Endpoint} Guardrail Translation Handler +""" + +import os +import sys +import pytest + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms import get_guardrail_translation_mapping +from litellm.llms.{provider}.{endpoint}.guardrail_translation.handler import ( + MyEndpointHandler, +) +from litellm.types.utils import CallTypes + + +class MockGuardrail(CustomGuardrail): + """Mock guardrail for testing""" + + async def apply_guardrail(self, text: str) -> str: + return f"{text} [GUARDRAILED]" + + +class TestHandlerDiscovery: + """Test that the handler is properly discovered""" + + def test_handler_discovered(self): + handler_class = get_guardrail_translation_mapping(CallTypes.my_endpoint) + assert handler_class == MyEndpointHandler + + +class TestInputProcessing: + """Test input processing functionality""" + + @pytest.mark.asyncio + async def test_process_simple_input(self): + handler = MyEndpointHandler() + guardrail = MockGuardrail(guardrail_name="test") + + data = {"messages": [{"role": "user", "content": "Hello"}]} + result = await handler.process_input_messages(data, guardrail) + + assert result["messages"][0]["content"] == "Hello [GUARDRAILED]" + + +class TestOutputProcessing: + """Test output processing functionality""" + + @pytest.mark.asyncio + async def test_process_simple_output(self): + handler = MyEndpointHandler() + guardrail = MockGuardrail(guardrail_name="test") + + # Create mock response + response = create_mock_response() + result = await handler.process_output_response(response, guardrail) + + # Assert guardrail was applied + assert "GUARDRAILED" in get_response_text(result) +``` + +## Support + +For questions or issues: +- Check existing handler implementations for examples +- Review the base translation class documentation +- Create an issue on GitHub with the `guardrails` label + diff --git a/docs/my-website/docs/anthropic_unified.md b/docs/my-website/docs/anthropic_unified.md index 62dea7051b..9981547ce1 100644 --- a/docs/my-website/docs/anthropic_unified.md +++ b/docs/my-website/docs/anthropic_unified.md @@ -10,14 +10,14 @@ Use LiteLLM to call all your LLM APIs in the Anthropic `v1/messages` format. | Feature | Supported | Notes | |-------|-------|-------| -| Cost Tracking | ✅ | | -| Logging | ✅ | works across all integrations | +| Cost Tracking | ✅ | Works with all supported models | +| Logging | ✅ | Works across all integrations | | End-user Tracking | ✅ | | | Streaming | ✅ | | -| Fallbacks | ✅ | between supported models | -| Loadbalancing | ✅ | between supported models | -| Guardrails | ✅ | | -| Support llm providers | **All LiteLLM supported providers** | `openai`, `anthropic`, `bedrock`, `vertex_ai`, `gemini`, `azure`, `azure_ai`, etc. | +| Fallbacks | ✅ | Works between supported models | +| Loadbalancing | ✅ | Works between supported models | +| Guardrails | ✅ | Applies to input and output text (non-streaming only) | +| Supported Providers | **All LiteLLM supported providers** | `openai`, `anthropic`, `bedrock`, `vertex_ai`, `gemini`, `azure`, `azure_ai`, etc. | ## Usage --- diff --git a/docs/my-website/docs/audio_transcription.md b/docs/my-website/docs/audio_transcription.md index 8cbc567180..fd55cc66e9 100644 --- a/docs/my-website/docs/audio_transcription.md +++ b/docs/my-website/docs/audio_transcription.md @@ -7,12 +7,13 @@ import TabItem from '@theme/TabItem'; | Feature | Supported | Notes | |-------|-------|-------| -| Cost Tracking | ✅ | | -| Logging | ✅ | works across all integrations | +| Cost Tracking | ✅ | Works with all supported models | +| Logging | ✅ | Works across all integrations | | End-user Tracking | ✅ | | -| Fallbacks | ✅ | between supported models | -| Loadbalancing | ✅ | between supported models | -| Support llm providers | `openai`, `azure`, `vertex_ai`, `gemini`, `deepgram`, `groq`, `fireworks_ai` | | +| Fallbacks | ✅ | Works between supported models | +| Loadbalancing | ✅ | Works between supported models | +| Guardrails | ✅ | Applies to output transcribed text (non-streaming only) | +| Supported Providers | `openai`, `azure`, `vertex_ai`, `gemini`, `deepgram`, `groq`, `fireworks_ai` | | ## Quick Start diff --git a/docs/my-website/docs/image_generation.md b/docs/my-website/docs/image_generation.md index 8cd5803aa6..b4eaef3652 100644 --- a/docs/my-website/docs/image_generation.md +++ b/docs/my-website/docs/image_generation.md @@ -5,6 +5,18 @@ import TabItem from '@theme/TabItem'; # Image Generations +## Overview + +| Feature | Supported | Notes | +|---------|-----------|-------| +| Cost Tracking | ✅ | Works with all supported models | +| Logging | ✅ | Works across all integrations | +| End-user Tracking | ✅ | | +| Fallbacks | ✅ | Works between supported models | +| Loadbalancing | ✅ | Works between supported models | +| Guardrails | ✅ | Applies to input prompts (non-streaming only) | +| Supported Providers | OpenAI, Azure, Google AI Studio, Vertex AI, AWS Bedrock, Recraft, Xinference, Nscale | | + ## Quick Start ### LiteLLM Python SDK diff --git a/docs/my-website/docs/providers/openai/text_to_speech.md b/docs/my-website/docs/providers/openai/text_to_speech.md index 34cd0f069e..a4aeb9e525 100644 --- a/docs/my-website/docs/providers/openai/text_to_speech.md +++ b/docs/my-website/docs/providers/openai/text_to_speech.md @@ -4,6 +4,18 @@ import TabItem from '@theme/TabItem'; # OpenAI - Text-to-speech +## Overview + +| Feature | Supported | Notes | +|---------|-----------|-------| +| Cost Tracking | ✅ | Works with all supported models | +| Logging | ✅ | Works across all integrations | +| End-user Tracking | ✅ | | +| Fallbacks | ✅ | Works between supported models | +| Loadbalancing | ✅ | Works between supported models | +| Guardrails | ✅ | Applies to input text | +| Supported Models | tts-1, tts-1-hd, gpt-4o-mini-tts | | + ## **LiteLLM Python SDK Usage** ### Quick Start diff --git a/docs/my-website/docs/rerank.md b/docs/my-website/docs/rerank.md index 72f35c6cec..ec0592f31f 100644 --- a/docs/my-website/docs/rerank.md +++ b/docs/my-website/docs/rerank.md @@ -6,6 +6,18 @@ LiteLLM Follows the [cohere api request / response for the rerank api](https://c ::: +## Overview + +| Feature | Supported | Notes | +|---------|-----------|-------| +| Cost Tracking | ✅ | Works with all supported models | +| Logging | ✅ | Works across all integrations | +| End-user Tracking | ✅ | | +| Fallbacks | ✅ | Works between supported models | +| Loadbalancing | ✅ | Works between supported models | +| Guardrails | ✅ | Applies to input query only (not documents) | +| Supported Providers | Cohere, Together AI, Azure AI, DeepInfra, Nvidia NIM, Infinity | | + ## **LiteLLM Python SDK Usage** ### Quick Start diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index 7b7db32fc6..fe21266958 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -17,7 +17,7 @@ Requests to /chat/completions may be bridged here automatically when the provide | Image Generation Streaming | ✅ | Progressive image generation with partial images (1-3) | | Fallbacks | ✅ | Works between supported models | | Loadbalancing | ✅ | Works between supported models | -| Guardrails | ✅ | | +| Guardrails | ✅ | Applies to input and output text (non-streaming only) | | Supported operations | Create a response, Get a response, Delete a response | | | Supported LiteLLM Versions | 1.63.8+ | | | Supported LLM providers | **All LiteLLM supported providers** | `openai`, `anthropic`, `bedrock`, `vertex_ai`, `gemini`, `azure`, `azure_ai` etc. | diff --git a/docs/my-website/docs/text_completion.md b/docs/my-website/docs/text_completion.md index cbf2db00a0..234494c2dd 100644 --- a/docs/my-website/docs/text_completion.md +++ b/docs/my-website/docs/text_completion.md @@ -3,6 +3,19 @@ import TabItem from '@theme/TabItem'; # /completions +## Overview + +| Feature | Supported | Notes | +|---------|-----------|-------| +| Cost Tracking | ✅ | Works with all supported models | +| Logging | ✅ | Works across all integrations | +| End-user Tracking | ✅ | | +| Streaming | ✅ | | +| Fallbacks | ✅ | Works between supported models | +| Loadbalancing | ✅ | Works between supported models | +| Guardrails | ✅ | Applies to input prompts and output text (non-streaming only) | +| Supported Providers | All Chat Completion Providers | | + ### Usage diff --git a/docs/my-website/docs/text_to_speech.md b/docs/my-website/docs/text_to_speech.md index 2c6a3aa058..c530e70e4b 100644 --- a/docs/my-website/docs/text_to_speech.md +++ b/docs/my-website/docs/text_to_speech.md @@ -4,18 +4,17 @@ import TabItem from '@theme/TabItem'; # /audio/speech -## Overview +## Overview -| Feature | Supported | Notes | -|-------|-------|-------| -| Cost Tracking | ✅ | | -| Logging | ✅ | works across all integrations | +| Feature | Supported | Notes | +|---------|-----------|-------| +| Cost Tracking | ✅ | Works with all supported models | +| Logging | ✅ | Works across all integrations | | End-user Tracking | ✅ | | -| Fallbacks | ✅ | between supported models | -| Loadbalancing | ✅ | between supported models | -| Guardrails | ❌ Please make an [issue if you need this feature](https://github.com/BerriAI/litellm/issues/new) | | -| Support llm providers | | `openai`, `azure`, `azure_ai`, `vertex_ai`, `gemini`, etc. | - +| Fallbacks | ✅ | Works between supported models | +| Loadbalancing | ✅ | Works between supported models | +| Guardrails | ✅ | Applies to input text (non-streaming only) | +| Supported Providers | OpenAI, Azure OpenAI, Vertex AI | | ## **LiteLLM Python SDK Usage** ### Quick Start diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index e2753932f5..f0e311c303 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -709,7 +709,8 @@ const sidebars = { label: "Adding Providers", items: [ "adding_provider/directory_structure", - "adding_provider/new_rerank_provider"], + "adding_provider/new_rerank_provider", + "adding_provider/adding_guardrail_support"], }, "extras/contributing", "contributing", diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 18973add86..1676e7bb1e 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -1,8 +1,16 @@ -from typing import TYPE_CHECKING, Optional +import importlib +import os +from typing import TYPE_CHECKING, Dict, Optional, Type + +from litellm._logging import verbose_logger +from litellm.types.utils import CallTypes from . import * if TYPE_CHECKING: + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) from litellm.types.utils import ModelInfo, Usage @@ -33,3 +41,120 @@ def get_cost_for_web_search_request( return cost_per_web_search_request_vertex_ai(usage=usage, model_info=model_info) else: return None + + +def discover_guardrail_translation_mappings() -> ( + Dict[CallTypes, Type["BaseTranslation"]] +): + """ + Discover guardrail translation mappings by scanning the llms directory structure. + + Scans for modules with guardrail_translation_mappings dictionaries and aggregates them. + + Returns: + Dict[CallTypes, Type[BaseTranslation]]: A dictionary mapping call types to their translation handler classes + """ + discovered_mappings: Dict[CallTypes, Type["BaseTranslation"]] = {} + + try: + # Get the path to the llms directory + current_dir = os.path.dirname(__file__) + llms_dir = current_dir + + if not os.path.exists(llms_dir): + verbose_logger.debug("llms directory not found") + return discovered_mappings + + # Recursively scan for guardrail_translation directories + for root, dirs, files in os.walk(llms_dir): + # Skip __pycache__ and base_llm directories + dirs[:] = [d for d in dirs if not d.startswith("__") and d != "base_llm"] + + # Check if this is a guardrail_translation directory with __init__.py + if ( + os.path.basename(root) == "guardrail_translation" + and "__init__.py" in files + ): + # Build the module path relative to litellm + rel_path = os.path.relpath(root, os.path.dirname(llms_dir)) + module_path = "litellm." + rel_path.replace(os.sep, ".") + + try: + # Import the module + verbose_logger.debug( + f"Discovering guardrail translations in: {module_path}" + ) + + module = importlib.import_module(module_path) + + # Check for guardrail_translation_mappings dictionary + if hasattr(module, "guardrail_translation_mappings"): + mappings = getattr(module, "guardrail_translation_mappings") + if isinstance(mappings, dict): + discovered_mappings.update(mappings) + verbose_logger.debug( + f"Found guardrail_translation_mappings in {module_path}: {list(mappings.keys())}" + ) + + except ImportError as e: + verbose_logger.error(f"Could not import {module_path}: {e}") + continue + except Exception as e: + verbose_logger.error(f"Error processing {module_path}: {e}") + continue + + verbose_logger.debug( + f"Discovered {len(discovered_mappings)} guardrail translation mappings: {list(discovered_mappings.keys())}" + ) + + except Exception as e: + verbose_logger.error(f"Error discovering guardrail translation mappings: {e}") + + return discovered_mappings + + +# Cache the discovered mappings +endpoint_guardrail_translation_mappings: Optional[ + Dict[CallTypes, Type["BaseTranslation"]] +] = None + + +def load_guardrail_translation_mappings(): + global endpoint_guardrail_translation_mappings + if endpoint_guardrail_translation_mappings is None: + endpoint_guardrail_translation_mappings = ( + discover_guardrail_translation_mappings() + ) + return endpoint_guardrail_translation_mappings + + +def get_guardrail_translation_mapping(call_type: CallTypes) -> Type["BaseTranslation"]: + """ + Get the guardrail translation handler for a given call type. + + Args: + call_type: The type of call (e.g., completion, acompletion, anthropic_messages) + + Returns: + The translation handler class for the given call type + + Raises: + ValueError: If no translation mapping exists for the given call type + """ + global endpoint_guardrail_translation_mappings + + # Lazy load the mappings on first access + if endpoint_guardrail_translation_mappings is None: + endpoint_guardrail_translation_mappings = ( + discover_guardrail_translation_mappings() + ) + + # Get the translation handler class for the call type + if call_type not in endpoint_guardrail_translation_mappings: + raise ValueError( + f"No guardrail translation mapping found for call_type: {call_type}. " + f"Available mappings: {list(endpoint_guardrail_translation_mappings.keys())}" + ) + + # Return the handler class directly + return endpoint_guardrail_translation_mappings[call_type] diff --git a/litellm/llms/anthropic/chat/guardrail_translation/__init__.py b/litellm/llms/anthropic/chat/guardrail_translation/__init__.py new file mode 100644 index 0000000000..ab327ee9f2 --- /dev/null +++ b/litellm/llms/anthropic/chat/guardrail_translation/__init__.py @@ -0,0 +1,10 @@ +from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.anthropic_messages: AnthropicMessagesHandler, +} + +__all__ = ["guardrail_translation_mappings"] diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py new file mode 100644 index 0000000000..cc0b260998 --- /dev/null +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -0,0 +1,270 @@ +""" +Anthropic Message Handler for Unified Guardrails + +This module provides a class-based handler for Anthropic-format messages. +The class methods can be overridden for custom behavior. + +Pattern Overview: +----------------- +1. Extract text content from messages/responses (both string and list formats) +2. Create async tasks to apply guardrails to each text segment +3. Track mappings to know where each response belongs +4. Apply guardrail responses back to the original structure +""" + +import asyncio +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + AnthropicResponseTextBlock, + ) + + +class AnthropicMessagesHandler(BaseTranslation): + """ + Handler for processing Anthropic messages with guardrails. + + This class provides methods to: + 1. Process input messages (pre-call hook) + 2. Process output responses (post-call hook) + + Methods can be overridden to customize behavior for different message formats. + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process input messages by applying guardrails to text content. + """ + messages = data.get("messages") + if messages is None: + return data + + tasks = [] + task_mappings: List[Tuple[int, Optional[int]]] = [] + # Track (message_index, content_index) for each task + # content_index is None for string content, int for list content + + # Step 1: Extract all text content and create guardrail tasks + for msg_idx, message in enumerate(messages): + await self._extract_input_text_and_create_tasks( + message=message, + msg_idx=msg_idx, + tasks=tasks, + task_mappings=task_mappings, + guardrail_to_apply=guardrail_to_apply, + ) + + # Step 2: Run all guardrail tasks in parallel + responses = await asyncio.gather(*tasks) + + # Step 3: Map guardrail responses back to original message structure + await self._apply_guardrail_responses_to_input( + messages=messages, + responses=responses, + task_mappings=task_mappings, + ) + + verbose_proxy_logger.debug( + "Anthropic Messages: Processed input messages: %s", messages + ) + + return data + + async def _extract_input_text_and_create_tasks( + self, + message: Dict[str, Any], + msg_idx: int, + tasks: List, + task_mappings: List[Tuple[int, Optional[int]]], + guardrail_to_apply: "CustomGuardrail", + ) -> None: + """ + Extract text content from a message and create guardrail tasks. + + Override this method to customize text extraction logic. + """ + content = message.get("content", None) + if content is None: + return + + if isinstance(content, str): + # Simple string content + tasks.append(guardrail_to_apply.apply_guardrail(text=content)) + task_mappings.append((msg_idx, None)) + + elif isinstance(content, list): + # List content (e.g., multimodal with text and images) + for content_idx, content_item in enumerate(content): + text_str = content_item.get("text", None) + if text_str is None: + continue + tasks.append(guardrail_to_apply.apply_guardrail(text=text_str)) + task_mappings.append((msg_idx, int(content_idx))) + + async def _apply_guardrail_responses_to_input( + self, + messages: List[Dict[str, Any]], + responses: List[str], + task_mappings: List[Tuple[int, Optional[int]]], + ) -> None: + """ + Apply guardrail responses back to input messages. + + Override this method to customize how responses are applied. + """ + for task_idx, guardrail_response in enumerate(responses): + mapping = task_mappings[task_idx] + msg_idx = cast(int, mapping[0]) + content_idx_optional = cast(Optional[int], mapping[1]) + + content = messages[msg_idx].get("content", None) + if content is None: + continue + + if isinstance(content, str) and content_idx_optional is None: + # Replace string content with guardrail response + messages[msg_idx]["content"] = guardrail_response + + elif isinstance(content, list) and content_idx_optional is not None: + # Replace specific text item in list content + messages[msg_idx]["content"][content_idx_optional][ + "text" + ] = guardrail_response + + async def process_output_response( + self, + response: "AnthropicMessagesResponse", + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process output response by applying guardrails to text content. + + Args: + response: Anthropic MessagesResponse object + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified response with guardrail applied to content + + Response Format Support: + - List content: response.content = [{"type": "text", "text": "text here"}, ...] + """ + # Step 0: Check if response has any text content to process + if not self._has_text_content(response): + verbose_proxy_logger.warning( + "Anthropic Messages: No text content in response, skipping guardrail" + ) + return response + + tasks = [] + task_mappings: List[Tuple[int, Optional[int]]] = [] + # Track (choice_index, content_index) for each task + + response_content = response.get("content", []) + if not response_content: + return response + # Step 1: Extract all text content from response choices + for content_idx, content_block in enumerate(response_content): + # Check if this is a text block by checking the 'type' field + if isinstance(content_block, dict) and content_block.get("type") == "text": + # Cast to dict to handle the union type properly + await self._extract_output_text_and_create_tasks( + content_block=cast(Dict[str, Any], content_block), + content_idx=content_idx, + tasks=tasks, + task_mappings=task_mappings, + guardrail_to_apply=guardrail_to_apply, + ) + + # Step 2: Run all guardrail tasks in parallel + responses = await asyncio.gather(*tasks) + + # Step 3: Map guardrail responses back to original response structure + await self._apply_guardrail_responses_to_output( + response=response, + responses=responses, + task_mappings=task_mappings, + ) + + verbose_proxy_logger.debug( + "Anthropic Messages: Processed output response: %s", response + ) + + return response + + def _has_text_content(self, response: "AnthropicMessagesResponse") -> bool: + """ + Check if response has any text content to process. + + Override this method to customize text content detection. + """ + response_content = response.get("content", []) + if not response_content: + return False + for content_block in response_content: + # Check if this is a text block by checking the 'type' field + if isinstance(content_block, dict) and content_block.get("type") == "text": + content_text = content_block.get("text") + if content_text and isinstance(content_text, str): + return True + return False + + async def _extract_output_text_and_create_tasks( + self, + content_block: Dict[str, Any], + content_idx: int, + tasks: List, + task_mappings: List[Tuple[int, Optional[int]]], + guardrail_to_apply: "CustomGuardrail", + ) -> None: + """ + Extract text content from a response choice and create guardrail tasks. + + Override this method to customize text extraction logic. + """ + content_text = content_block.get("text") + if content_text and isinstance(content_text, str): + # Simple string content + tasks.append(guardrail_to_apply.apply_guardrail(text=content_text)) + task_mappings.append((content_idx, None)) + + async def _apply_guardrail_responses_to_output( + self, + response: "AnthropicMessagesResponse", + responses: List[str], + task_mappings: List[Tuple[int, Optional[int]]], + ) -> None: + """ + Apply guardrail responses back to output response. + + Override this method to customize how responses are applied. + """ + for task_idx, guardrail_response in enumerate(responses): + mapping = task_mappings[task_idx] + content_idx = cast(int, mapping[0]) + + response_content = response.get("content", []) + if not response_content: + continue + + # Get the content block at the index + if content_idx >= len(response_content): + continue + + content_block = response_content[content_idx] + + # Verify it's a text block and update the text field + if isinstance(content_block, dict) and content_block.get("type") == "text": + # Cast to dict to handle the union type properly for assignment + content_block = cast("AnthropicResponseTextBlock", content_block) + content_block["text"] = guardrail_response diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py new file mode 100644 index 0000000000..4599af1b74 --- /dev/null +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -0,0 +1,23 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + + +class BaseTranslation(ABC): + @abstractmethod + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + pass + + @abstractmethod + async def process_output_response( + self, + response: Any, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + pass diff --git a/litellm/llms/cohere/rerank/guardrail_translation/README.md b/litellm/llms/cohere/rerank/guardrail_translation/README.md new file mode 100644 index 0000000000..e77e5a74dd --- /dev/null +++ b/litellm/llms/cohere/rerank/guardrail_translation/README.md @@ -0,0 +1,229 @@ +# Cohere Rerank Guardrail Translation Handler + +Handler for processing the rerank endpoint (`/v1/rerank`) with guardrails. + +## Overview + +This handler processes rerank requests by: +1. Extracting the query text from the request +2. Applying guardrails to the query +3. Updating the request with the guardrailed query +4. Returning the output unchanged (rankings are not text) + +Note: Documents are not processed by guardrails as they represent the corpus +being searched, not user input. Only the query is guardrailed. + +## Data Format + +### Input Format + +**With String Documents:** +```json +{ + "model": "rerank-english-v3.0", + "query": "What is the capital of France?", + "documents": [ + "Paris is the capital of France.", + "Berlin is the capital of Germany.", + "Madrid is the capital of Spain." + ], + "top_n": 2 +} +``` + +**With Dict Documents:** +```json +{ + "model": "rerank-english-v3.0", + "query": "What is the capital of France?", + "documents": [ + {"text": "Paris is the capital of France.", "id": "doc1"}, + {"text": "Berlin is the capital of Germany.", "id": "doc2"}, + {"text": "Madrid is the capital of Spain.", "id": "doc3"} + ], + "top_n": 2 +} +``` + +### Output Format + +```json +{ + "id": "rerank-abc123", + "results": [ + {"index": 0, "relevance_score": 0.98}, + {"index": 2, "relevance_score": 0.12} + ], + "meta": { + "billed_units": {"search_units": 1} + } +} +``` + +## Usage + +The handler is automatically discovered and applied when guardrails are used with the rerank endpoint. + +### Example: Using Guardrails with Rerank + +```bash +curl -X POST 'http://localhost:4000/v1/rerank' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "rerank-english-v3.0", + "query": "What is machine learning?", + "documents": [ + "Machine learning is a subset of AI.", + "Deep learning uses neural networks.", + "Python is a programming language." + ], + "guardrails": ["content_filter"], + "top_n": 2 +}' +``` + +The guardrail will be applied to the query only (not the documents). + +### Example: PII Masking in Query + +```bash +curl -X POST 'http://localhost:4000/v1/rerank' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "rerank-english-v3.0", + "query": "Find documents about John Doe from john@example.com", + "documents": [ + "Document 1 content here.", + "Document 2 content here.", + "Document 3 content here." + ], + "guardrails": ["mask_pii"], + "top_n": 3 +}' +``` + +The query will be masked to: "Find documents about [NAME_REDACTED] from [EMAIL_REDACTED]" + +### Example: Mixed Document Types + +```bash +curl -X POST 'http://localhost:4000/v1/rerank' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "rerank-english-v3.0", + "query": "Technical documentation", + "documents": [ + {"text": "This is document 1", "metadata": {"source": "wiki"}}, + {"text": "This is document 2", "metadata": {"source": "docs"}}, + "This is document 3 as a plain string" + ], + "guardrails": ["content_moderation"] +}' +``` + +## Implementation Details + +### Input Processing + +- **Query Field**: `query` (string) + - Processing: Apply guardrail to query text + - Result: Updated query + +- **Documents Field**: `documents` (list) + - Processing: Not processed (corpus being searched, not user input) + - Result: Unchanged + +### Output Processing + +- **Processing**: Not applicable (output contains relevance scores, not text) +- **Result**: Response returned unchanged + +## Use Cases + +1. **PII Protection**: Remove PII from queries before reranking +2. **Content Filtering**: Filter inappropriate content from search queries +3. **Compliance**: Ensure queries meet requirements +4. **Data Sanitization**: Clean up query text before semantic search operations + +## Extension + +Override these methods to customize behavior: + +- `process_input_messages()`: Customize how query is processed +- `process_output_response()`: Currently a no-op, but can be overridden if needed + +## Supported Call Types + +- `CallTypes.rerank` - Synchronous rerank +- `CallTypes.arerank` - Asynchronous rerank + +## Notes + +- Only the query is processed by guardrails +- Documents are not processed (they represent the corpus, not user input) +- Output processing is a no-op since rankings don't contain text +- Both sync and async call types use the same handler +- Works with all rerank providers (Cohere, Together AI, etc.) + +## Common Patterns + +### PII Masking in Search + +```python +import litellm + +response = litellm.rerank( + model="rerank-english-v3.0", + query="Find info about john@example.com", + documents=[ + "Document 1 content.", + "Document 2 content.", + "Document 3 content." + ], + guardrails=["mask_pii"], + top_n=2 +) + +# Query will have PII masked +# query becomes: "Find info about [EMAIL_REDACTED]" +print(response.results) +``` + +### Content Filtering + +```python +import litellm + +response = litellm.rerank( + model="rerank-english-v3.0", + query="Search query here", + documents=[ + {"text": "Document 1 content", "id": "doc1"}, + {"text": "Document 2 content", "id": "doc2"}, + ], + guardrails=["content_filter"], +) +``` + +### Async Rerank with Guardrails + +```python +import litellm +import asyncio + +async def rerank_with_guardrails(): + response = await litellm.arerank( + model="rerank-english-v3.0", + query="Technical query", + documents=["Doc 1", "Doc 2", "Doc 3"], + guardrails=["sanitize"], + top_n=2 + ) + return response + +result = asyncio.run(rerank_with_guardrails()) +``` + diff --git a/litellm/llms/cohere/rerank/guardrail_translation/__init__.py b/litellm/llms/cohere/rerank/guardrail_translation/__init__.py new file mode 100644 index 0000000000..70b580facf --- /dev/null +++ b/litellm/llms/cohere/rerank/guardrail_translation/__init__.py @@ -0,0 +1,11 @@ +"""Cohere Rerank handler for Unified Guardrails.""" + +from litellm.llms.cohere.rerank.guardrail_translation.handler import CohereRerankHandler +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.rerank: CohereRerankHandler, + CallTypes.arerank: CohereRerankHandler, +} + +__all__ = ["guardrail_translation_mappings", "CohereRerankHandler"] diff --git a/litellm/llms/cohere/rerank/guardrail_translation/handler.py b/litellm/llms/cohere/rerank/guardrail_translation/handler.py new file mode 100644 index 0000000000..0c5e50dc41 --- /dev/null +++ b/litellm/llms/cohere/rerank/guardrail_translation/handler.py @@ -0,0 +1,90 @@ +""" +Cohere Rerank Handler for Unified Guardrails + +This module provides guardrail translation support for the rerank endpoint. +The handler processes only the 'query' parameter for guardrails. +""" + +from typing import TYPE_CHECKING, Any + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.rerank import RerankResponse + + +class CohereRerankHandler(BaseTranslation): + """ + Handler for processing rerank requests with guardrails. + + This class provides methods to: + 1. Process input query (pre-call hook) + 2. Process output response (post-call hook) - not applicable for rerank + + The handler specifically processes: + - The 'query' parameter (string) + + Note: Documents are not processed by guardrails as they are the corpus + being searched, not user input. + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process input query by applying guardrails. + + Args: + data: Request data dictionary containing 'query' + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified data with guardrails applied to query only + """ + # Process query only + query = data.get("query") + if query is not None and isinstance(query, str): + guardrailed_query = await guardrail_to_apply.apply_guardrail(text=query) + data["query"] = guardrailed_query + + verbose_proxy_logger.debug( + "Rerank: Applied guardrail to query. " + "Original length: %d, New length: %d", + len(query), + len(guardrailed_query), + ) + else: + verbose_proxy_logger.debug( + "Rerank: No query to process or query is not a string" + ) + + return data + + async def process_output_response( + self, + response: "RerankResponse", + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process output response - not applicable for rerank. + + Rerank responses contain relevance scores and indices, not text, + so there's nothing to apply guardrails to. This method returns + the response unchanged. + + Args: + response: Rerank response object with rankings + guardrail_to_apply: The guardrail instance (unused) + + Returns: + Unmodified response (rankings don't need text guardrails) + """ + verbose_proxy_logger.debug( + "Rerank: Output processing not applicable " + "(output contains relevance scores, not text)" + ) + return response diff --git a/litellm/llms/openai/chat/guardrail_translation/README.md b/litellm/llms/openai/chat/guardrail_translation/README.md new file mode 100644 index 0000000000..05e3b55e54 --- /dev/null +++ b/litellm/llms/openai/chat/guardrail_translation/README.md @@ -0,0 +1,3 @@ +Translation of OpenAI `/chat/completions` input and output to a custom guardrail. + +This enables guardrails to be applied to OpenAI `/chat/completions` requests and responses. \ No newline at end of file diff --git a/litellm/llms/openai/chat/guardrail_translation/__init__.py b/litellm/llms/openai/chat/guardrail_translation/__init__.py new file mode 100644 index 0000000000..b0682aa475 --- /dev/null +++ b/litellm/llms/openai/chat/guardrail_translation/__init__.py @@ -0,0 +1,12 @@ +"""OpenAI Chat Completions message handler for Unified Guardrails.""" + +from litellm.llms.openai.chat.guardrail_translation.handler import ( + OpenAIChatCompletionsHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.completion: OpenAIChatCompletionsHandler, + CallTypes.acompletion: OpenAIChatCompletionsHandler, +} +__all__ = ["guardrail_translation_mappings"] diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py new file mode 100644 index 0000000000..12d8cfad23 --- /dev/null +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -0,0 +1,280 @@ +""" +OpenAI Chat Completions Message Handler for Unified Guardrails + +This module provides a class-based handler for OpenAI-format chat completions. +The class methods can be overridden for custom behavior. + +Pattern Overview: +----------------- +1. Extract text content from messages/responses (both string and list formats) +2. Create async tasks to apply guardrails to each text segment +3. Track mappings to know where each response belongs +4. Apply guardrail responses back to the original structure + +This pattern can be replicated for other message formats (e.g., Anthropic). +""" + +import asyncio +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import Choices + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.utils import ModelResponse + + +class OpenAIChatCompletionsHandler(BaseTranslation): + """ + Handler for processing OpenAI chat completions messages with guardrails. + + This class provides methods to: + 1. Process input messages (pre-call hook) + 2. Process output responses (post-call hook) + + Methods can be overridden to customize behavior for different message formats. + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process input messages by applying guardrails to text content. + """ + messages = data.get("messages") + if messages is None: + return data + + tasks = [] + task_mappings: List[Tuple[int, Optional[int]]] = [] + # Track (message_index, content_index) for each task + # content_index is None for string content, int for list content + + # Step 1: Extract all text content and create guardrail tasks + for msg_idx, message in enumerate(messages): + await self._extract_input_text_and_create_tasks( + message=message, + msg_idx=msg_idx, + tasks=tasks, + task_mappings=task_mappings, + guardrail_to_apply=guardrail_to_apply, + ) + + # Step 2: Run all guardrail tasks in parallel + responses = await asyncio.gather(*tasks) + + # Step 3: Map guardrail responses back to original message structure + await self._apply_guardrail_responses_to_input( + messages=messages, + responses=responses, + task_mappings=task_mappings, + ) + + verbose_proxy_logger.debug( + "OpenAI Chat Completions: Processed input messages: %s", messages + ) + + return data + + async def _extract_input_text_and_create_tasks( + self, + message: Dict[str, Any], + msg_idx: int, + tasks: List, + task_mappings: List[Tuple[int, Optional[int]]], + guardrail_to_apply: "CustomGuardrail", + ) -> None: + """ + Extract text content from a message and create guardrail tasks. + + Override this method to customize text extraction logic. + """ + content = message.get("content", None) + if content is None: + return + + if isinstance(content, str): + # Simple string content + tasks.append(guardrail_to_apply.apply_guardrail(text=content)) + task_mappings.append((msg_idx, None)) + + elif isinstance(content, list): + # List content (e.g., multimodal with text and images) + for content_idx, content_item in enumerate(content): + text_str = content_item.get("text", None) + if text_str is None: + continue + tasks.append(guardrail_to_apply.apply_guardrail(text=text_str)) + task_mappings.append((msg_idx, int(content_idx))) + + async def _apply_guardrail_responses_to_input( + self, + messages: List[Dict[str, Any]], + responses: List[str], + task_mappings: List[Tuple[int, Optional[int]]], + ) -> None: + """ + Apply guardrail responses back to input messages. + + Override this method to customize how responses are applied. + """ + for task_idx, guardrail_response in enumerate(responses): + mapping = task_mappings[task_idx] + msg_idx = cast(int, mapping[0]) + content_idx_optional = cast(Optional[int], mapping[1]) + + content = messages[msg_idx].get("content", None) + if content is None: + continue + + if isinstance(content, str) and content_idx_optional is None: + # Replace string content with guardrail response + messages[msg_idx]["content"] = guardrail_response + + elif isinstance(content, list) and content_idx_optional is not None: + # Replace specific text item in list content + messages[msg_idx]["content"][content_idx_optional][ + "text" + ] = guardrail_response + + async def process_output_response( + self, + response: "ModelResponse", + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process output response by applying guardrails to text content. + + Args: + response: LiteLLM ModelResponse object + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified response with guardrail applied to content + + Response Format Support: + - String content: choice.message.content = "text here" + - List content: choice.message.content = [{"type": "text", "text": "text here"}, ...] + """ + # Step 0: Check if response has any text content to process + if not self._has_text_content(response): + verbose_proxy_logger.warning( + "OpenAI Chat Completions: No text content in response, skipping guardrail" + ) + return response + + tasks = [] + task_mappings: List[Tuple[int, Optional[int]]] = [] + # Track (choice_index, content_index) for each task + + # Step 1: Extract all text content from response choices + for choice_idx, choice in enumerate(response.choices): + await self._extract_output_text_and_create_tasks( + choice=choice, + choice_idx=choice_idx, + tasks=tasks, + task_mappings=task_mappings, + guardrail_to_apply=guardrail_to_apply, + ) + + # Step 2: Run all guardrail tasks in parallel + responses = await asyncio.gather(*tasks) + + # Step 3: Map guardrail responses back to original response structure + await self._apply_guardrail_responses_to_output( + response=response, + responses=responses, + task_mappings=task_mappings, + ) + + verbose_proxy_logger.debug( + "OpenAI Chat Completions: Processed output response: %s", response + ) + + return response + + def _has_text_content(self, response: "ModelResponse") -> bool: + """ + Check if response has any text content to process. + + Override this method to customize text content detection. + """ + for choice in response.choices: + if isinstance(choice, litellm.Choices): + if choice.message.content and isinstance(choice.message.content, str): + return True + return False + + async def _extract_output_text_and_create_tasks( + self, + choice: Any, + choice_idx: int, + tasks: List, + task_mappings: List[Tuple[int, Optional[int]]], + guardrail_to_apply: "CustomGuardrail", + ) -> None: + """ + Extract text content from a response choice and create guardrail tasks. + + Override this method to customize text extraction logic. + """ + if not isinstance(choice, litellm.Choices): + return + + verbose_proxy_logger.debug( + "OpenAI Chat Completions: Processing choice: %s", choice + ) + + if choice.message.content and isinstance(choice.message.content, str): + # Simple string content + tasks.append( + guardrail_to_apply.apply_guardrail(text=choice.message.content) + ) + task_mappings.append((choice_idx, None)) + + elif choice.message.content and isinstance(choice.message.content, list): + # List content (e.g., multimodal response) + for content_idx, content_item in enumerate(choice.message.content): + content_text = content_item.get("text") + if content_text: + tasks.append(guardrail_to_apply.apply_guardrail(text=content_text)) + task_mappings.append((choice_idx, int(content_idx))) + + async def _apply_guardrail_responses_to_output( + self, + response: "ModelResponse", + responses: List[str], + task_mappings: List[Tuple[int, Optional[int]]], + ) -> None: + """ + Apply guardrail responses back to output response. + + Override this method to customize how responses are applied. + """ + for task_idx, guardrail_response in enumerate(responses): + mapping = task_mappings[task_idx] + choice_idx = cast(int, mapping[0]) + content_idx_optional = cast(Optional[int], mapping[1]) + + content = cast(Choices, response.choices[choice_idx]).message.content + if content is None: + continue + + if isinstance(content, str) and content_idx_optional is None: + # Replace string content with guardrail response + cast(Choices, response.choices[choice_idx]).message.content = ( + guardrail_response + ) + + elif isinstance(content, list) and content_idx_optional is not None: + # Replace specific text item in list content + cast(Choices, response.choices[choice_idx]).message.content[ # type: ignore + content_idx_optional + ][ + "text" + ] = guardrail_response diff --git a/litellm/llms/openai/completion/guardrail_translation/README.md b/litellm/llms/openai/completion/guardrail_translation/README.md new file mode 100644 index 0000000000..93762206c4 --- /dev/null +++ b/litellm/llms/openai/completion/guardrail_translation/README.md @@ -0,0 +1,158 @@ +# OpenAI Text Completion Guardrail Translation Handler + +Handler for processing OpenAI's text completion endpoint (`/v1/completions`) with guardrails. + +## Overview + +This handler processes text completion requests by: +1. Extracting the text prompt(s) from the request +2. Applying guardrails to the prompt text(s) +3. Updating the request with the guardrailed prompt(s) +4. Applying guardrails to the completion output text + +## Data Format + +### Input Format + +**Single Prompt:** +```json +{ + "model": "gpt-3.5-turbo-instruct", + "prompt": "Say this is a test", + "max_tokens": 7, + "temperature": 0 +} +``` + +**Multiple Prompts (Batch):** +```json +{ + "model": "gpt-3.5-turbo-instruct", + "prompt": [ + "Tell me a joke", + "Write a poem" + ], + "max_tokens": 50 +} +``` + +### Output Format + +```json +{ + "id": "cmpl-uqkvlQyYK7bGYrRHQ0eXlWi7", + "object": "text_completion", + "created": 1589478378, + "model": "gpt-3.5-turbo-instruct", + "choices": [ + { + "text": "\n\nThis is indeed a test", + "index": 0, + "logprobs": null, + "finish_reason": "length" + } + ], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12 + } +} +``` + +## Usage + +The handler is automatically discovered and applied when guardrails are used with the text completion endpoint. + +### Example: Using Guardrails with Text Completion + +```bash +curl -X POST 'http://localhost:4000/v1/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "gpt-3.5-turbo-instruct", + "prompt": "Say this is a test", + "guardrails": ["content_moderation"], + "max_tokens": 7 +}' +``` + +The guardrail will be applied to both: +- **Input**: The prompt text before sending to the LLM +- **Output**: The completion text in the response + +### Example: PII Masking in Prompts and Completions + +```bash +curl -X POST 'http://localhost:4000/v1/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "gpt-3.5-turbo-instruct", + "prompt": "My name is John Doe and my email is john@example.com", + "guardrails": ["mask_pii"], + "metadata": { + "guardrails": ["mask_pii"] + } +}' +``` + +### Example: Batch Prompts with Guardrails + +```bash +curl -X POST 'http://localhost:4000/v1/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "gpt-3.5-turbo-instruct", + "prompt": [ + "Tell me about AI", + "What is machine learning?" + ], + "guardrails": ["content_filter"], + "max_tokens": 100 +}' +``` + +## Implementation Details + +### Input Processing + +- **Field**: `prompt` (string or list of strings) +- **Processing**: + - String prompts: Apply guardrail directly + - List prompts: Apply guardrail to each string in the list +- **Result**: Updated prompt(s) in request + +### Output Processing + +- **Field**: `choices[*].text` (string) +- **Processing**: Applies guardrail to each completion text +- **Result**: Updated completion texts in response + +### Supported Prompt Types + +1. **String**: Single prompt as a string +2. **List of Strings**: Multiple prompts for batch completion +3. **List of Lists**: Token-based prompts (passed through unchanged) + +## Extension + +Override these methods to customize behavior: + +- `process_input_messages()`: Customize how prompts are processed +- `process_output_response()`: Customize how completion texts are processed + +## Supported Call Types + +- `CallTypes.text_completion` - Synchronous text completion +- `CallTypes.atext_completion` - Asynchronous text completion + +## Notes + +- The handler processes both input prompts and output completion texts +- List prompts are processed individually (each string in the list) +- Non-string prompt items (e.g., token lists) are passed through unchanged +- Both sync and async call types use the same handler + diff --git a/litellm/llms/openai/completion/guardrail_translation/__init__.py b/litellm/llms/openai/completion/guardrail_translation/__init__.py new file mode 100644 index 0000000000..51e43c4593 --- /dev/null +++ b/litellm/llms/openai/completion/guardrail_translation/__init__.py @@ -0,0 +1,13 @@ +"""OpenAI Text Completion handler for Unified Guardrails.""" + +from litellm.llms.openai.completion.guardrail_translation.handler import ( + OpenAITextCompletionHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.text_completion: OpenAITextCompletionHandler, + CallTypes.atext_completion: OpenAITextCompletionHandler, +} + +__all__ = ["guardrail_translation_mappings", "OpenAITextCompletionHandler"] diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py new file mode 100644 index 0000000000..b5db730620 --- /dev/null +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -0,0 +1,137 @@ +""" +OpenAI Text Completion Handler for Unified Guardrails + +This module provides guardrail translation support for OpenAI's text completion endpoint. +The handler processes the 'prompt' parameter for guardrails. +""" + +from typing import TYPE_CHECKING, Any + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.utils import TextCompletionResponse + + +class OpenAITextCompletionHandler(BaseTranslation): + """ + Handler for processing OpenAI text completion requests with guardrails. + + This class provides methods to: + 1. Process input prompt (pre-call hook) + 2. Process output response (post-call hook) + + The handler specifically processes the 'prompt' parameter which can be: + - A single string + - A list of strings (for batch completions) + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process input prompt by applying guardrails to text content. + + Args: + data: Request data dictionary containing 'prompt' parameter + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified data with guardrails applied to prompt + """ + prompt = data.get("prompt") + if prompt is None: + verbose_proxy_logger.debug( + "OpenAI Text Completion: No prompt found in request data" + ) + return data + + if isinstance(prompt, str): + # Single string prompt + guardrailed_prompt = await guardrail_to_apply.apply_guardrail(text=prompt) + data["prompt"] = guardrailed_prompt + + verbose_proxy_logger.debug( + "OpenAI Text Completion: Applied guardrail to string prompt. " + "Original length: %d, New length: %d", + len(prompt), + len(guardrailed_prompt), + ) + + elif isinstance(prompt, list): + # List of string prompts (batch completion) + guardrailed_prompts = [] + for idx, p in enumerate(prompt): + if isinstance(p, str): + guardrailed_p = await guardrail_to_apply.apply_guardrail(text=p) + guardrailed_prompts.append(guardrailed_p) + verbose_proxy_logger.debug( + "OpenAI Text Completion: Applied guardrail to prompt[%d]. " + "Original length: %d, New length: %d", + idx, + len(p), + len(guardrailed_p), + ) + else: + # For non-string items (e.g., token lists), keep unchanged + guardrailed_prompts.append(p) + verbose_proxy_logger.debug( + "OpenAI Text Completion: Skipping guardrail for prompt[%d] " + "(not a string, type: %s)", + idx, + type(p), + ) + + data["prompt"] = guardrailed_prompts + + else: + verbose_proxy_logger.warning( + "OpenAI Text Completion: Unexpected prompt type: %s. Expected string or list.", + type(prompt), + ) + + return data + + async def process_output_response( + self, + response: "TextCompletionResponse", + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process output response by applying guardrails to completion text. + + Args: + response: Text completion response object + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified response with guardrails applied to completion text + """ + if not hasattr(response, "choices") or not response.choices: + verbose_proxy_logger.debug( + "OpenAI Text Completion: No choices in response to process" + ) + return response + + # Apply guardrails to each choice's text + for idx, choice in enumerate(response.choices): + if hasattr(choice, "text") and isinstance(choice.text, str): + original_text = choice.text + guardrailed_text = await guardrail_to_apply.apply_guardrail( + text=original_text + ) + choice.text = guardrailed_text + + verbose_proxy_logger.debug( + "OpenAI Text Completion: Applied guardrail to choice[%d] text. " + "Original length: %d, New length: %d", + idx, + len(original_text), + len(guardrailed_text), + ) + + return response diff --git a/litellm/llms/openai/image_generation/__init__.py b/litellm/llms/openai/image_generation/__init__.py index eb2a0576b6..e20c80f20b 100644 --- a/litellm/llms/openai/image_generation/__init__.py +++ b/litellm/llms/openai/image_generation/__init__.py @@ -5,11 +5,17 @@ from litellm.llms.base_llm.image_generation.transformation import ( from .dall_e_2_transformation import DallE2ImageGenerationConfig from .dall_e_3_transformation import DallE3ImageGenerationConfig from .gpt_transformation import GPTImageGenerationConfig +from .guardrail_translation import ( + OpenAIImageGenerationHandler, + guardrail_translation_mappings, +) __all__ = [ "DallE2ImageGenerationConfig", "DallE3ImageGenerationConfig", "GPTImageGenerationConfig", + "OpenAIImageGenerationHandler", + "guardrail_translation_mappings", ] diff --git a/litellm/llms/openai/image_generation/guardrail_translation/README.md b/litellm/llms/openai/image_generation/guardrail_translation/README.md new file mode 100644 index 0000000000..fcbd2d154d --- /dev/null +++ b/litellm/llms/openai/image_generation/guardrail_translation/README.md @@ -0,0 +1,106 @@ +# OpenAI Image Generation Guardrail Translation Handler + +Handler for processing OpenAI's image generation endpoint with guardrails. + +## Overview + +This handler processes image generation requests by: +1. Extracting the text prompt from the request +2. Applying guardrails to the prompt text +3. Updating the request with the guardrailed prompt + +## Data Format + +### Input Format + +```json +{ + "model": "dall-e-3", + "prompt": "A cute baby sea otter", + "n": 1, + "size": "1024x1024", + "quality": "standard" +} +``` + +### Output Format + +```json +{ + "created": 1589478378, + "data": [ + { + "url": "https://...", + "revised_prompt": "A cute baby sea otter..." + } + ] +} +``` + +## Usage + +The handler is automatically discovered and applied when guardrails are used with the image generation endpoint. + +### Example: Using Guardrails with Image Generation + +```bash +curl -X POST 'http://localhost:4000/v1/images/generations' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "dall-e-3", + "prompt": "A cute baby sea otter wearing a hat", + "guardrails": ["content_moderation"], + "size": "1024x1024" +}' +``` + +The guardrail will be applied to the prompt text before the image generation request is sent to the provider. + +### Example: PII Masking in Prompts + +```bash +curl -X POST 'http://localhost:4000/v1/images/generations' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "dall-e-3", + "prompt": "Generate an image of John Doe at john@example.com", + "guardrails": ["mask_pii"], + "metadata": { + "guardrails": ["mask_pii"] + } +}' +``` + +## Implementation Details + +### Input Processing + +- **Field**: `prompt` (string) +- **Processing**: Applies guardrail to prompt text +- **Result**: Updated prompt in request + +### Output Processing + +- **Processing**: Not applicable (images don't contain text to guardrail) +- **Result**: Response returned unchanged + +## Extension + +Override these methods to customize behavior: + +- `process_input_messages()`: Customize how the prompt is processed +- `process_output_response()`: Add custom processing for image metadata if needed + +## Supported Call Types + +- `CallTypes.image_generation` - Synchronous image generation +- `CallTypes.aimage_generation` - Asynchronous image generation + +## Notes + +- The handler only processes the `prompt` parameter +- Output processing is a no-op since images don't contain text +- Both sync and async call types use the same handler + diff --git a/litellm/llms/openai/image_generation/guardrail_translation/__init__.py b/litellm/llms/openai/image_generation/guardrail_translation/__init__.py new file mode 100644 index 0000000000..1fba2a3692 --- /dev/null +++ b/litellm/llms/openai/image_generation/guardrail_translation/__init__.py @@ -0,0 +1,13 @@ +"""OpenAI Image Generation handler for Unified Guardrails.""" + +from litellm.llms.openai.image_generation.guardrail_translation.handler import ( + OpenAIImageGenerationHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.image_generation: OpenAIImageGenerationHandler, + CallTypes.aimage_generation: OpenAIImageGenerationHandler, +} + +__all__ = ["guardrail_translation_mappings", "OpenAIImageGenerationHandler"] diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py new file mode 100644 index 0000000000..5fcb5278f0 --- /dev/null +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -0,0 +1,93 @@ +""" +OpenAI Image Generation Handler for Unified Guardrails + +This module provides guardrail translation support for OpenAI's image generation endpoint. +The handler processes the 'prompt' parameter for guardrails. +""" + +from typing import TYPE_CHECKING, Any + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.utils import ImageResponse + + +class OpenAIImageGenerationHandler(BaseTranslation): + """ + Handler for processing OpenAI image generation requests with guardrails. + + This class provides methods to: + 1. Process input prompt (pre-call hook) + 2. Process output response (post-call hook) - typically not needed for images + + The handler specifically processes the 'prompt' parameter which contains + the text description for image generation. + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process input prompt by applying guardrails to text content. + + Args: + data: Request data dictionary containing 'prompt' parameter + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified data with guardrails applied to prompt + """ + prompt = data.get("prompt") + if prompt is None: + verbose_proxy_logger.debug( + "OpenAI Image Generation: No prompt found in request data" + ) + return data + + # Apply guardrail to the prompt + if isinstance(prompt, str): + guardrailed_prompt = await guardrail_to_apply.apply_guardrail(text=prompt) + data["prompt"] = guardrailed_prompt + + verbose_proxy_logger.debug( + "OpenAI Image Generation: Applied guardrail to prompt. " + "Original length: %d, New length: %d", + len(prompt), + len(guardrailed_prompt), + ) + else: + verbose_proxy_logger.debug( + "OpenAI Image Generation: Unexpected prompt type: %s. Expected string.", + type(prompt), + ) + + return data + + async def process_output_response( + self, + response: "ImageResponse", + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process output response - typically not needed for image generation. + + Image responses don't contain text to apply guardrails to, so this + method returns the response unchanged. This is provided for completeness + and can be overridden if needed for custom image metadata processing. + + Args: + response: Image generation response object + guardrail_to_apply: The guardrail instance to apply + + Returns: + Unmodified response (images don't need text guardrails) + """ + verbose_proxy_logger.debug( + "OpenAI Image Generation: Output processing not needed for image responses" + ) + return response diff --git a/litellm/llms/openai/responses/guardrail_translation/README.md b/litellm/llms/openai/responses/guardrail_translation/README.md new file mode 100644 index 0000000000..bc1bd6f4f2 --- /dev/null +++ b/litellm/llms/openai/responses/guardrail_translation/README.md @@ -0,0 +1,119 @@ +# OpenAI Responses API Guardrail Translation Handler + +This module provides guardrail translation support for the OpenAI Responses API format. + +## Overview + +The `OpenAIResponsesHandler` class handles the translation of guardrail operations for both input and output of the Responses API. It follows the same pattern as the Chat Completions handler but is adapted for the Responses API's specific data structures. + +## Responses API Format + +### Input Format +The Responses API accepts input in two formats: + +1. **String input**: Simple text string + ```python + {"input": "Hello world", "model": "gpt-4"} + ``` + +2. **List input**: Array of message objects (ResponseInputParam) + ```python + { + "input": [ + { + "role": "user", + "content": "Hello", # Can be string or list of content items + "type": "message" + } + ], + "model": "gpt-4" + } + ``` + +### Output Format +The Responses API returns a `ResponsesAPIResponse` object with: + +```python +{ + "id": "resp_123", + "output": [ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Assistant response", + "annotations": [] + } + ] + } + ] +} +``` + +## Usage + +The handler is automatically discovered and registered for `CallTypes.responses` and `CallTypes.aresponses`. + +### Example + +```python +from litellm.llms import get_guardrail_translation_mapping +from litellm.types.utils import CallTypes + +# Get the handler +handler_class = get_guardrail_translation_mapping(CallTypes.responses) +handler = handler_class() + +# Process input +data = {"input": "User message", "model": "gpt-4"} +processed_data = await handler.process_input_messages(data, guardrail_instance) + +# Process output +response = await litellm.aresponses(**processed_data) +processed_response = await handler.process_output_response(response, guardrail_instance) +``` + +## Key Methods + +### `process_input_messages(data, guardrail_to_apply)` +Processes input data by: +1. Handling both string and list input formats +2. Extracting text content from messages +3. Applying guardrails to text content in parallel +4. Mapping guardrail responses back to the original structure + +### `process_output_response(response, guardrail_to_apply)` +Processes output response by: +1. Extracting text from output items' content +2. Applying guardrails to all text content in parallel +3. Replacing original text with guardrailed versions + +## Extending the Handler + +The handler can be customized by overriding these methods: + +- `_extract_input_text_and_create_tasks()`: Customize input text extraction logic +- `_apply_guardrail_responses_to_input()`: Customize how guardrail responses are applied to input +- `_extract_output_text_and_create_tasks()`: Customize output text extraction logic +- `_apply_guardrail_responses_to_output()`: Customize how guardrail responses are applied to output +- `_has_text_content()`: Customize text content detection + +## Testing + +Comprehensive tests are available in `tests/llm_translation/test_openai_responses_guardrail_handler.py`: + +```bash +pytest tests/llm_translation/test_openai_responses_guardrail_handler.py -v +``` + +## Implementation Details + +- **Parallel Processing**: All text content is processed in parallel using `asyncio.gather()` +- **Mapping Tracking**: Uses tuples to track the location of each text segment for accurate replacement +- **Type Safety**: Handles both Pydantic objects and dict representations +- **Multimodal Support**: Properly handles mixed content with text and other media types + diff --git a/litellm/llms/openai/responses/guardrail_translation/__init__.py b/litellm/llms/openai/responses/guardrail_translation/__init__.py new file mode 100644 index 0000000000..d2d9e5375c --- /dev/null +++ b/litellm/llms/openai/responses/guardrail_translation/__init__.py @@ -0,0 +1,12 @@ +"""OpenAI Responses API handler for Unified Guardrails.""" + +from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.responses: OpenAIResponsesHandler, + CallTypes.aresponses: OpenAIResponsesHandler, +} +__all__ = ["guardrail_translation_mappings"] diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py new file mode 100644 index 0000000000..40b1d73ddc --- /dev/null +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -0,0 +1,332 @@ +""" +OpenAI Responses API Handler for Unified Guardrails + +This module provides a class-based handler for OpenAI Responses API format. +The class methods can be overridden for custom behavior. + +Pattern Overview: +----------------- +1. Extract text content from input/output (both string and list formats) +2. Create async tasks to apply guardrails to each text segment +3. Track mappings to know where each response belongs +4. Apply guardrail responses back to the original structure + +Responses API Format: +--------------------- +Input: Union[str, List[Dict]] where each dict has: + - role: str + - content: Union[str, List[Dict]] (can have text items) + - type: str (e.g., "message") + +Output: response.output is List[GenericResponseOutputItem] where each has: + - type: str (e.g., "message") + - id: str + - status: str + - role: str + - content: List[OutputText] where OutputText has: + - type: str (e.g., "output_text") + - text: str +""" + +import asyncio +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.responses.main import GenericResponseOutputItem, OutputText + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.llms.openai import ResponseInputParam + from litellm.types.utils import ResponsesAPIResponse + + +class OpenAIResponsesHandler(BaseTranslation): + """ + Handler for processing OpenAI Responses API with guardrails. + + This class provides methods to: + 1. Process input (pre-call hook) + 2. Process output response (post-call hook) + + Methods can be overridden to customize behavior for different message formats. + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process input by applying guardrails to text content. + + Handles both string input and list of message objects. + """ + input_data: Optional[Union[str, "ResponseInputParam"]] = data.get("input") + if input_data is None: + return data + + # Handle simple string input + if isinstance(input_data, str): + guardrail_response = await guardrail_to_apply.apply_guardrail( + text=input_data + ) + data["input"] = guardrail_response + verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") + return data + + # Handle list input (ResponseInputParam) + if not isinstance(input_data, list): + return data + + tasks = [] + task_mappings: List[Tuple[int, Optional[int]]] = [] + # Track (message_index, content_index) for each task + # content_index is None for string content, int for list content + + # Step 1: Extract all text content and create guardrail tasks + for msg_idx, message in enumerate(input_data): + await self._extract_input_text_and_create_tasks( + message=message, + msg_idx=msg_idx, + tasks=tasks, + task_mappings=task_mappings, + guardrail_to_apply=guardrail_to_apply, + ) + + # Step 2: Run all guardrail tasks in parallel + if tasks: + responses = await asyncio.gather(*tasks) + + # Step 3: Map guardrail responses back to original input structure + await self._apply_guardrail_responses_to_input( + messages=input_data, + responses=responses, + task_mappings=task_mappings, + ) + + verbose_proxy_logger.debug( + "OpenAI Responses API: Processed input messages: %s", input_data + ) + + return data + + async def _extract_input_text_and_create_tasks( + self, + message: Dict[str, Any], + msg_idx: int, + tasks: List, + task_mappings: List[Tuple[int, Optional[int]]], + guardrail_to_apply: "CustomGuardrail", + ) -> None: + """ + Extract text content from an input message and create guardrail tasks. + + Override this method to customize text extraction logic. + """ + content = message.get("content", None) + if content is None: + return + + if isinstance(content, str): + # Simple string content + tasks.append(guardrail_to_apply.apply_guardrail(text=content)) + task_mappings.append((msg_idx, None)) + + elif isinstance(content, list): + # List content (e.g., multimodal with text and images) + for content_idx, content_item in enumerate(content): + if isinstance(content_item, dict): + text_str = content_item.get("text", None) + if text_str is not None: + tasks.append(guardrail_to_apply.apply_guardrail(text=text_str)) + task_mappings.append((msg_idx, int(content_idx))) + + async def _apply_guardrail_responses_to_input( + self, + messages: List[Dict[str, Any]], + responses: List[str], + task_mappings: List[Tuple[int, Optional[int]]], + ) -> None: + """ + Apply guardrail responses back to input messages. + + Override this method to customize how responses are applied. + """ + for task_idx, guardrail_response in enumerate(responses): + mapping = task_mappings[task_idx] + msg_idx = cast(int, mapping[0]) + content_idx_optional = cast(Optional[int], mapping[1]) + + content = messages[msg_idx].get("content", None) + if content is None: + continue + + if isinstance(content, str) and content_idx_optional is None: + # Replace string content with guardrail response + messages[msg_idx]["content"] = guardrail_response + + elif isinstance(content, list) and content_idx_optional is not None: + # Replace specific text item in list content + if isinstance(messages[msg_idx]["content"][content_idx_optional], dict): + messages[msg_idx]["content"][content_idx_optional][ + "text" + ] = guardrail_response + + async def process_output_response( + self, + response: "ResponsesAPIResponse", + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process output response by applying guardrails to text content. + + Args: + response: LiteLLM ResponsesAPIResponse object + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified response with guardrail applied to content + + Response Format Support: + - response.output is a list of output items + - Each output item has a content list with OutputText objects + - Each OutputText object has a text field + """ + # Step 0: Check if response has any text content to process + if not self._has_text_content(response): + verbose_proxy_logger.warning( + "OpenAI Responses API: No text content in response, skipping guardrail" + ) + return response + + tasks = [] + task_mappings: List[Tuple[int, int]] = [] + # Track (output_item_index, content_index) for each task + + # Step 1: Extract all text content from response output + for output_idx, output_item in enumerate(response.output): + await self._extract_output_text_and_create_tasks( + output_item=output_item, + output_idx=output_idx, + tasks=tasks, + task_mappings=task_mappings, + guardrail_to_apply=guardrail_to_apply, + ) + + # Step 2: Run all guardrail tasks in parallel + if tasks: + responses = await asyncio.gather(*tasks) + + # Step 3: Map guardrail responses back to original response structure + await self._apply_guardrail_responses_to_output( + response=response, + responses=responses, + task_mappings=task_mappings, + ) + + verbose_proxy_logger.debug( + "OpenAI Responses API: Processed output response: %s", response + ) + + return response + + def _has_text_content(self, response: "ResponsesAPIResponse") -> bool: + """ + Check if response has any text content to process. + + Override this method to customize text content detection. + """ + if not hasattr(response, "output") or response.output is None: + return False + + for output_item in response.output: + if isinstance(output_item, (GenericResponseOutputItem, dict)): + content = ( + output_item.content + if isinstance(output_item, GenericResponseOutputItem) + else output_item.get("content", []) + ) + if content: + for content_item in content: + # Check if it's an OutputText with text + if isinstance(content_item, OutputText): + if content_item.text: + return True + elif isinstance(content_item, dict): + if content_item.get("text"): + return True + return False + + async def _extract_output_text_and_create_tasks( + self, + output_item: Any, + output_idx: int, + tasks: List, + task_mappings: List[Tuple[int, int]], + guardrail_to_apply: "CustomGuardrail", + ) -> None: + """ + Extract text content from a response output item and create guardrail tasks. + + Override this method to customize text extraction logic. + """ + # Handle both GenericResponseOutputItem and dict + if isinstance(output_item, GenericResponseOutputItem): + content = output_item.content + elif isinstance(output_item, dict): + content = output_item.get("content", []) + else: + return + + if not content: + return + + verbose_proxy_logger.debug( + "OpenAI Responses API: Processing output item: %s", output_item + ) + + # Iterate through content items (list of OutputText objects) + for content_idx, content_item in enumerate(content): + # Handle both OutputText objects and dicts + if isinstance(content_item, OutputText): + text_content = content_item.text + elif isinstance(content_item, dict): + text_content = content_item.get("text") + else: + continue + + if text_content: + tasks.append(guardrail_to_apply.apply_guardrail(text=text_content)) + task_mappings.append((output_idx, int(content_idx))) + + async def _apply_guardrail_responses_to_output( + self, + response: "ResponsesAPIResponse", + responses: List[str], + task_mappings: List[Tuple[int, int]], + ) -> None: + """ + Apply guardrail responses back to output response. + + Override this method to customize how responses are applied. + """ + for task_idx, guardrail_response in enumerate(responses): + mapping = task_mappings[task_idx] + output_idx = cast(int, mapping[0]) + content_idx = cast(int, mapping[1]) + + output_item = response.output[output_idx] + + # Handle both GenericResponseOutputItem and dict + if isinstance(output_item, GenericResponseOutputItem): + content_item = output_item.content[content_idx] + if isinstance(content_item, OutputText): + content_item.text = guardrail_response + elif isinstance(content_item, dict): + content_item["text"] = guardrail_response + elif isinstance(output_item, dict): + content = output_item.get("content", []) + if content and content_idx < len(content): + if isinstance(content[content_idx], dict): + content[content_idx]["text"] = guardrail_response diff --git a/litellm/llms/openai/speech/guardrail_translation/README.md b/litellm/llms/openai/speech/guardrail_translation/README.md new file mode 100644 index 0000000000..52e89ffa92 --- /dev/null +++ b/litellm/llms/openai/speech/guardrail_translation/README.md @@ -0,0 +1,178 @@ +# OpenAI Text-to-Speech Guardrail Translation Handler + +Handler for processing OpenAI's text-to-speech endpoint (`/v1/audio/speech`) with guardrails. + +## Overview + +This handler processes text-to-speech requests by: +1. Extracting the input text from the request +2. Applying guardrails to the input text +3. Updating the request with the guardrailed text +4. Returning the output unchanged (audio is binary, not text) + +## Data Format + +### Input Format + +```json +{ + "model": "tts-1", + "input": "The quick brown fox jumped over the lazy dog.", + "voice": "alloy", + "response_format": "mp3", + "speed": 1.0 +} +``` + +### Output Format + +The output is binary audio data (MP3, WAV, etc.), not text, so it cannot be guardrailed. + +## Usage + +The handler is automatically discovered and applied when guardrails are used with the text-to-speech endpoint. + +### Example: Using Guardrails with Text-to-Speech + +```bash +curl -X POST 'http://localhost:4000/v1/audio/speech' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "tts-1", + "input": "The quick brown fox jumped over the lazy dog.", + "voice": "alloy", + "guardrails": ["content_moderation"] +}' \ +--output speech.mp3 +``` + +The guardrail will be applied to the input text before the text-to-speech conversion. + +### Example: PII Masking in TTS Input + +```bash +curl -X POST 'http://localhost:4000/v1/audio/speech' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "tts-1", + "input": "Please call John Doe at john@example.com", + "voice": "nova", + "guardrails": ["mask_pii"] +}' \ +--output speech.mp3 +``` + +The audio will say: "Please call [NAME_REDACTED] at [EMAIL_REDACTED]" + +### Example: Content Filtering Before TTS + +```bash +curl -X POST 'http://localhost:4000/v1/audio/speech' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer your-api-key' \ +-d '{ + "model": "tts-1-hd", + "input": "This is the text that will be spoken", + "voice": "shimmer", + "guardrails": ["content_filter"] +}' \ +--output speech.mp3 +``` + +## Implementation Details + +### Input Processing + +- **Field**: `input` (string) +- **Processing**: Applies guardrail to input text +- **Result**: Updated input text in request + +### Output Processing + +- **Processing**: Not applicable (audio is binary data) +- **Result**: Response returned unchanged + +## Use Cases + +1. **PII Protection**: Remove personally identifiable information before converting to speech +2. **Content Filtering**: Remove inappropriate content before TTS conversion +3. **Compliance**: Ensure text meets requirements before voice synthesis +4. **Text Sanitization**: Clean up text before audio generation + +## Extension + +Override these methods to customize behavior: + +- `process_input_messages()`: Customize how input text is processed +- `process_output_response()`: Currently a no-op, but can be overridden if needed + +## Supported Call Types + +- `CallTypes.speech` - Synchronous text-to-speech +- `CallTypes.aspeech` - Asynchronous text-to-speech + +## Notes + +- Only the input text is processed by guardrails +- Output processing is a no-op since audio cannot be text-guardrailed +- Both sync and async call types use the same handler +- Works with all TTS models (tts-1, tts-1-hd, etc.) +- Works with all voice options + +## Common Patterns + +### Remove PII Before TTS + +```python +import litellm +from pathlib import Path + +speech_file_path = Path(__file__).parent / "speech.mp3" +response = litellm.speech( + model="tts-1", + voice="alloy", + input="Hi, this is John Doe calling from john@company.com", + guardrails=["mask_pii"], +) +response.stream_to_file(speech_file_path) +# Audio will have PII masked +``` + +### Content Moderation Before TTS + +```python +import litellm +from pathlib import Path + +speech_file_path = Path(__file__).parent / "speech.mp3" +response = litellm.speech( + model="tts-1-hd", + voice="nova", + input="Your text here", + guardrails=["content_moderation"], +) +response.stream_to_file(speech_file_path) +``` + +### Async TTS with Guardrails + +```python +import litellm +import asyncio +from pathlib import Path + +async def generate_speech(): + speech_file_path = Path(__file__).parent / "speech.mp3" + response = await litellm.aspeech( + model="tts-1", + voice="echo", + input="Text to convert to speech", + guardrails=["pii_mask"], + ) + response.stream_to_file(speech_file_path) + +asyncio.run(generate_speech()) +``` + diff --git a/litellm/llms/openai/speech/guardrail_translation/__init__.py b/litellm/llms/openai/speech/guardrail_translation/__init__.py new file mode 100644 index 0000000000..ef7d50f861 --- /dev/null +++ b/litellm/llms/openai/speech/guardrail_translation/__init__.py @@ -0,0 +1,13 @@ +"""OpenAI Text-to-Speech handler for Unified Guardrails.""" + +from litellm.llms.openai.speech.guardrail_translation.handler import ( + OpenAITextToSpeechHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings = { + CallTypes.speech: OpenAITextToSpeechHandler, + CallTypes.aspeech: OpenAITextToSpeechHandler, +} + +__all__ = ["guardrail_translation_mappings", "OpenAITextToSpeechHandler"] diff --git a/litellm/llms/openai/speech/guardrail_translation/handler.py b/litellm/llms/openai/speech/guardrail_translation/handler.py new file mode 100644 index 0000000000..aa049801d1 --- /dev/null +++ b/litellm/llms/openai/speech/guardrail_translation/handler.py @@ -0,0 +1,93 @@ +""" +OpenAI Text-to-Speech Handler for Unified Guardrails + +This module provides guardrail translation support for OpenAI's text-to-speech endpoint. +The handler processes the 'input' text parameter (output is audio, so no text to guardrail). +""" + +from typing import TYPE_CHECKING, Any + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.llms.openai import HttpxBinaryResponseContent + + +class OpenAITextToSpeechHandler(BaseTranslation): + """ + Handler for processing OpenAI text-to-speech requests with guardrails. + + This class provides methods to: + 1. Process input text (pre-call hook) + + Note: Output processing is not applicable since the output is audio (binary), + not text. Only the input text is processed. + """ + + async def process_input_messages( + self, + data: dict, + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process input text by applying guardrails. + + Args: + data: Request data dictionary containing 'input' parameter + guardrail_to_apply: The guardrail instance to apply + + Returns: + Modified data with guardrails applied to input text + """ + input_text = data.get("input") + if input_text is None: + verbose_proxy_logger.debug( + "OpenAI Text-to-Speech: No input text found in request data" + ) + return data + + if isinstance(input_text, str): + guardrailed_input = await guardrail_to_apply.apply_guardrail( + text=input_text + ) + data["input"] = guardrailed_input + + verbose_proxy_logger.debug( + "OpenAI Text-to-Speech: Applied guardrail to input text. " + "Original length: %d, New length: %d", + len(input_text), + len(guardrailed_input), + ) + else: + verbose_proxy_logger.debug( + "OpenAI Text-to-Speech: Unexpected input type: %s. Expected string.", + type(input_text), + ) + + return data + + async def process_output_response( + self, + response: "HttpxBinaryResponseContent", + guardrail_to_apply: "CustomGuardrail", + ) -> Any: + """ + Process output - not applicable for text-to-speech. + + The output is audio (binary data), not text, so there's nothing to apply + guardrails to. This method returns the response unchanged. + + Args: + response: Binary audio response + guardrail_to_apply: The guardrail instance (unused) + + Returns: + Unmodified response (audio data doesn't need text guardrails) + """ + verbose_proxy_logger.debug( + "OpenAI Text-to-Speech: Output processing not applicable " + "(output is audio data, not text)" + ) + return response diff --git a/litellm/llms/openai/transcriptions/guardrail_translation/README.md b/litellm/llms/openai/transcriptions/guardrail_translation/README.md new file mode 100644 index 0000000000..08e5b6f85c --- /dev/null +++ b/litellm/llms/openai/transcriptions/guardrail_translation/README.md @@ -0,0 +1,159 @@ +# OpenAI Audio Transcription Guardrail Translation Handler + +Handler for processing OpenAI's audio transcription endpoint (`/v1/audio/transcriptions`) with guardrails. + +## Overview + +This handler processes audio transcription responses by: +1. Applying guardrails to the transcribed text output +2. Returning the input unchanged (since input is an audio file, not text) + +## Data Format + +### Input Format + +The input is an audio file, which cannot be guardrailed (it's binary data, not text). + +```json +{ + "model": "whisper-1", + "file": "