diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index d27875da58f..d1f7e530ce3 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -440,6 +440,10 @@ router_settings: | DAYS_IN_A_MONTH | Days in a month for calculation purposes. Default is 28 | DAYS_IN_A_WEEK | Days in a week for calculation purposes. Default is 7 | DAYS_IN_A_YEAR | Days in a year for calculation purposes. Default is 365 +| DYNAMOAI_API_KEY | API key for DynamoAI Guardrails service +| DYNAMOAI_API_BASE | Base URL for DynamoAI API. Default is https://api.dynamo.ai +| DYNAMOAI_MODEL_ID | Model ID for DynamoAI tracking/logging purposes +| DYNAMOAI_POLICY_IDS | Comma-separated list of DynamoAI policy IDs to apply | DD_BASE_URL | Base URL for Datadog integration | DATADOG_BASE_URL | (Alternative to DD_BASE_URL) Base URL for Datadog integration | _DATADOG_BASE_URL | (Alternative to DD_BASE_URL) Base URL for Datadog integration diff --git a/docs/my-website/docs/proxy/guardrails/dynamoai.md b/docs/my-website/docs/proxy/guardrails/dynamoai.md new file mode 100644 index 00000000000..532ae76ca08 --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/dynamoai.md @@ -0,0 +1,214 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# DynamoAI Guardrails + +LiteLLM supports DynamoAI guardrails for content moderation and policy enforcement on LLM inputs and outputs. + +## Quick Start + +### 1. Define Guardrails on your LiteLLM config.yaml + +Define your guardrails under the `guardrails` section: + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: gpt-4 + litellm_params: + model: openai/gpt-4 + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "dynamoai-guard" + litellm_params: + guardrail: dynamoai + mode: "pre_call" + api_key: os.environ/DYNAMOAI_API_KEY +``` + +#### Supported values for `mode` + +- `pre_call` - Run **before** LLM call, on **input** +- `post_call` - Run **after** LLM call, on **output** +- `during_call` - Run **during** LLM call, on **input**. Same as `pre_call` but runs in parallel as LLM call + +### 2. Set Environment Variables + +```bash +export DYNAMOAI_API_KEY="your-api-key" +# Optional: Set policy IDs via environment variable (comma-separated) +export DYNAMOAI_POLICY_IDS="policy-id-1,policy-id-2,policy-id-3" +``` + +### 3. Start LiteLLM Gateway + +```shell +litellm --config config.yaml --detailed_debug +``` + +### 4. Test Request + +**[Langchain, OpenAI SDK Usage Examples](../proxy/user_keys#request-format)** + + + + +```shell showLineNumbers title="Successful Request" +curl -i http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "What is the capital of France?"} + ], + "guardrails": ["dynamoai-guard"] + }' +``` + +**Response: HTTP 200 Success** + +Content passes all policy checks and is allowed through. + + + + + +```shell showLineNumbers title="Blocked Request" +curl -i http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "Content that violates policy"} + ], + "guardrails": ["dynamoai-guard"] + }' +``` + +**Expected Response on Block: HTTP 400 Error** + +```json showLineNumbers +{ + "error": { + "message": "Guardrail failed: 1 violation(s) detected\n\n- POLICY NAME:\n Action: BLOCK\n Method: TOXICITY\n Description: Policy description\n Policy ID: policy-id-123", + "type": "None", + "param": "None", + "code": "400" + } +} +``` + + + + +## Advanced Configuration + +### Specify Policy IDs + +Configure specific DynamoAI policies to apply: + +```yaml showLineNumbers title="config.yaml" +guardrails: + - guardrail_name: "dynamoai-policies" + litellm_params: + guardrail: dynamoai + mode: "pre_call" + api_key: os.environ/DYNAMOAI_API_KEY + policy_ids: + - "policy-id-1" + - "policy-id-2" + - "policy-id-3" +``` + +### Custom API Base + +Specify a custom DynamoAI API endpoint: + +```yaml showLineNumbers title="config.yaml" +guardrails: + - guardrail_name: "dynamoai-custom" + litellm_params: + guardrail: dynamoai + mode: "pre_call" + api_key: os.environ/DYNAMOAI_API_KEY + api_base: "https://custom.dynamo.ai" +``` + +### Model ID for Tracking + +Add a model ID for tracking and logging purposes: + +```yaml showLineNumbers title="config.yaml" +guardrails: + - guardrail_name: "dynamoai-tracked" + litellm_params: + guardrail: dynamoai + mode: "pre_call" + api_key: os.environ/DYNAMOAI_API_KEY + model_id: "gpt-4-production" +``` + +### Input and Output Guardrails + +Configure separate guardrails for input and output: + +```yaml showLineNumbers title="config.yaml" +guardrails: + # Input guardrail + - guardrail_name: "dynamoai-input" + litellm_params: + guardrail: dynamoai + mode: "pre_call" + api_key: os.environ/DYNAMOAI_API_KEY + + # Output guardrail + - guardrail_name: "dynamoai-output" + litellm_params: + guardrail: dynamoai + mode: "post_call" + api_key: os.environ/DYNAMOAI_API_KEY +``` + +## Configuration Options + +| Parameter | Type | Description | Default | +|-----------|------|-------------|---------| +| `api_key` | string | DynamoAI API key (required) | `DYNAMOAI_API_KEY` env var | +| `api_base` | string | DynamoAI API base URL | `https://api.dynamo.ai` | +| `policy_ids` | array | List of DynamoAI policy IDs to apply (optional) | `DYNAMOAI_POLICY_IDS` env var (comma-separated) | +| `model_id` | string | Model ID for tracking/logging | `DYNAMOAI_MODEL_ID` env var | +| `mode` | string | When to run: `pre_call`, `post_call`, or `during_call` | Required | + +## Observability + +DynamoAI guardrail logs include: + +- **guardrail_status**: `success`, `guardrail_intervened`, or `guardrail_failed_to_respond` +- **guardrail_provider**: `dynamoai` +- **guardrail_json_response**: Full API response with policy details +- **duration**: Time taken for guardrail check +- **start_time** and **end_time**: Timestamps + +These logs are available through your configured LiteLLM logging callbacks. + +## Error Handling + +The guardrail handles errors gracefully: + +- **API Failures**: Logs error and raises exception with status `guardrail_failed_to_respond` +- **Policy Violations**: Raises `ValueError` with detailed violation information +- **Invalid Configuration**: Raises `ValueError` on initialization if API key is missing + +## Current Limitations + +- Only the `BLOCK` action is currently supported +- `WARN`, `REDACT`, and `SANITIZE` actions are treated as success (pass through) + +## Support + +For more information about DynamoAI: +- Website: [https://dynamo.ai](https://dynamo.ai) +- Documentation: Contact DynamoAI for API documentation + diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 93c1af1ce2e..8f9be715b68 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -43,6 +43,7 @@ const sidebars = { "proxy/guardrails/lakera_ai", "proxy/guardrails/model_armor", "proxy/guardrails/noma_security", + "proxy/guardrails/dynamoai", "proxy/guardrails/openai_moderation", "proxy/guardrails/pangea", "proxy/guardrails/pillar_security", diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py new file mode 100644 index 00000000000..c1ffa337647 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/__init__.py @@ -0,0 +1,4 @@ +from .dynamoai import DynamoAIGuardrails + +__all__ = ["DynamoAIGuardrails"] + diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py new file mode 100644 index 00000000000..9ea5f4cc747 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -0,0 +1,520 @@ +# +-------------------------------------------------------------+ +# +# Use DynamoAI Guardrails for your LLM calls +# https://dynamo.ai +# +# +-------------------------------------------------------------+ + +import os +from datetime import datetime +from typing import ( + Any, + AsyncGenerator, + Dict, + List, + Literal, + Optional, + Type, + Union, +) + +import httpx + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.caching.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.dynamoai import ( + DynamoAIProcessedResult, + DynamoAIRequest, + DynamoAIResponse, +) +from litellm.types.utils import GuardrailStatus, ModelResponseStream + +GUARDRAIL_NAME = "dynamoai" + + +class DynamoAIGuardrails(CustomGuardrail): + """ + DynamoAI Guardrails integration for LiteLLM. + + Provides content moderation and policy enforcement using DynamoAI's guardrail API. + """ + def __init__( + self, + guardrail_name: str = "litellm_test", + api_key: Optional[str] = None, + api_base: Optional[str] = None, + model_id: str = "", + policy_ids: List[str] = [], + **kwargs, + ): + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) + + # Set API configuration + self.api_key = api_key or os.getenv("DYNAMOAI_API_KEY") + if not self.api_key: + raise ValueError( + "DynamoAI API key is required. Set DYNAMOAI_API_KEY environment variable or pass api_key parameter." + ) + + self.api_base = api_base or os.getenv( + "DYNAMOAI_API_BASE", "https://api.dynamo.ai" + ) + self.api_url = f"{self.api_base}/v1/moderation/analyze/" + + # Model ID for tracking/logging purposes + self.model_id = model_id or os.getenv("DYNAMOAI_MODEL_ID", "") + + # Policy IDs - get from parameter, env var, or use empty list + env_policy_ids = os.getenv("DYNAMOAI_POLICY_IDS", "") + self.policy_ids = policy_ids or (env_policy_ids.split(",") if env_policy_ids else []) + self.guardrail_name = guardrail_name + self.guardrail_provider = "dynamoai" + + # store kwargs as optional_params + self.optional_params = kwargs + + # Set supported event hooks + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.during_call, + ] + + super().__init__(guardrail_name=guardrail_name, **kwargs) + + verbose_proxy_logger.debug( + "DynamoAI Guardrail initialized with guardrail_name=%s, model_id=%s", + self.guardrail_name, + self.model_id, + ) + + async def _call_dynamoai_guardrails( + self, + messages: List[Dict[str, Any]], + text_type: str = "input", + request_data: Optional[dict] = None, + ) -> DynamoAIResponse: + """ + Call DynamoAI Guardrails API to analyze messages for policy violations. + + Args: + messages: List of messages to analyze + text_type: Type of text being analyzed ("input" or "output") + request_data: Optional request data for logging purposes + + Returns: + DynamoAIResponse: Response from the DynamoAI Guardrails API + """ + start_time = datetime.now() + + payload: DynamoAIRequest = { + "messages": messages, + } + + # Add optional fields if provided + if self.policy_ids: + payload["policyIds"] = self.policy_ids + if self.model_id: + payload["modelId"] = self.model_id + + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {self.api_key}" + } + + verbose_proxy_logger.debug( + "DynamoAI request to %s with payload=%s", + self.api_url, + payload, + ) + + try: + response = await self.async_handler.post( + url=self.api_url, + json=dict(payload), + headers=headers, + ) + response.raise_for_status() + response_json = response.json() + + end_time = datetime.now() + duration = (end_time - start_time).total_seconds() + + # Add guardrail information to request trace + if request_data: + guardrail_status = self._determine_guardrail_status(response_json) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response=response_json, + request_data=request_data, + guardrail_status=guardrail_status, + start_time=start_time.timestamp(), + end_time=end_time.timestamp(), + duration=duration, + ) + + return response_json + + except httpx.HTTPError as e: + end_time = datetime.now() + duration = (end_time - start_time).total_seconds() + + verbose_proxy_logger.error( + "DynamoAI API request failed: %s", str(e) + ) + + # Add guardrail information with failure status + if request_data: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response={"error": str(e)}, + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=end_time.timestamp(), + duration=duration, + ) + + raise + + def _process_dynamoai_guardrails_response( + self, response: DynamoAIResponse + ) -> DynamoAIProcessedResult: + """ + Process the response from the DynamoAI Guardrails API + + Args: + response: The response from the API with 'finalAction' and 'appliedPolicies' keys + + Returns: + DynamoAIProcessedResult: Processed response with detected violations + """ + final_action = response.get("finalAction", "NONE") + applied_policies = response.get("appliedPolicies", []) + + violations_detected: List[str] = [] + violation_details: Dict[str, Any] = {} + + # For now, only handle BLOCK action + if final_action == "BLOCK": + for applied_policy in applied_policies: + policy_info = applied_policy.get("policy", {}) + policy_outputs = applied_policy.get("outputs", {}) + + # Get policy name and action + policy_name = policy_info.get("name", "unknown") + + # Check for action in multiple places + policy_action = ( + applied_policy.get("action") or + (policy_outputs.get("action") if policy_outputs else None) or + "NONE" + ) + + # Only include policies with BLOCK action + if policy_action == "BLOCK": + violations_detected.append(policy_name) + violation_details[policy_name] = { + "policyId": policy_info.get("id"), + "action": policy_action, + "method": policy_info.get("method"), + "description": policy_info.get("description"), + "message": policy_outputs.get("message") if policy_outputs else None, + } + + return { + "violations_detected": violations_detected, + "violation_details": violation_details + } + + def _determine_guardrail_status( + self, response_json: DynamoAIResponse + ) -> GuardrailStatus: + """ + Determine the guardrail status based on DynamoAI API response. + + Returns: + "success": Content allowed through with no violations (finalAction is NONE) + "guardrail_intervened": Content blocked (finalAction is BLOCK) + "guardrail_failed_to_respond": Technical error or API failure + """ + try: + if not isinstance(response_json, dict): + return "guardrail_failed_to_respond" + + # Check for error in response + if response_json.get("error"): + return "guardrail_failed_to_respond" + + final_action = response_json.get("finalAction", "NONE") + + if final_action == "NONE": + return "success" + elif final_action == "BLOCK": + return "guardrail_intervened" + + # For now, treat other actions as success (WARN, REDACT, SANITIZE not implemented yet) + return "success" + + except Exception as e: + verbose_proxy_logger.error( + "Error determining DynamoAI guardrail status: %s", str(e) + ) + return "guardrail_failed_to_respond" + + def _create_error_message(self, processed_result: DynamoAIProcessedResult) -> str: + """ + Create a detailed error message from processed guardrail results. + + Args: + processed_result: Processed response with detected violations + + Returns: + Formatted error message string + """ + violations_detected = processed_result["violations_detected"] + violation_details = processed_result["violation_details"] + + error_message = f"Guardrail failed: {len(violations_detected)} violation(s) detected\n\n" + + for policy_name in violations_detected: + error_message += f"- {policy_name.upper()}:\n" + details = violation_details.get(policy_name, {}) + + # Format violation details + if details.get("action"): + error_message += f" Action: {details['action']}\n" + if details.get("method"): + error_message += f" Method: {details['method']}\n" + if details.get("description"): + error_message += f" Description: {details['description']}\n" + if details.get("message"): + error_message += f" Message: {details['message']}\n" + if details.get("policyId"): + error_message += f" Policy ID: {details['policyId']}\n" + error_message += "\n" + + return error_message.strip() + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: Literal[ + "completion", + "text_completion", + "embeddings", + "image_generation", + "moderation", + "audio_transcription", + "pass_through_endpoint", + "rerank", + "mcp_call", + "anthropic_messages", + ], + ) -> Union[Exception, str, dict, None]: + """ + Runs before the LLM API call + Runs on only Input + Use this if you want to MODIFY the input + """ + verbose_proxy_logger.debug("Running DynamoAI pre-call hook") + + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + + event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call + if self.should_run_guardrail(data=data, event_type=event_type) is not True: + return data + + _messages = data.get("messages") + if _messages: + result = await self._call_dynamoai_guardrails( + messages=_messages, + text_type="input", + request_data=data, + ) + + verbose_proxy_logger.debug("Guardrails async_pre_call_hook result=%s", result) + + # Process the guardrails response + processed_result = self._process_dynamoai_guardrails_response(result) + violations_detected = processed_result["violations_detected"] + + # If any violations are detected, raise an error + if violations_detected: + error_message = self._create_error_message(processed_result) + raise ValueError(error_message) + + # Add guardrail to applied guardrails header + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) + + return data + + async def async_moderation_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: Literal[ + "completion", + "embeddings", + "image_generation", + "moderation", + "audio_transcription", + "responses", + "mcp_call", + "anthropic_messages", + ], + ): + """ + Runs in parallel to LLM API call + Runs on only Input + + This can NOT modify the input, only used to reject or accept a call before going to LLM API + """ + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + + event_type: GuardrailEventHooks = GuardrailEventHooks.during_call + if self.should_run_guardrail(data=data, event_type=event_type) is not True: + return + + _messages = data.get("messages") + if _messages: + result = await self._call_dynamoai_guardrails( + messages=_messages, + text_type="input", + request_data=data, + ) + + verbose_proxy_logger.debug("Guardrails async_moderation_hook result=%s", result) + + # Process the guardrails response + processed_result = self._process_dynamoai_guardrails_response(result) + violations_detected = processed_result["violations_detected"] + + # If any violations are detected, raise an error + if violations_detected: + error_message = self._create_error_message(processed_result) + raise ValueError(error_message) + + # Add guardrail to applied guardrails header + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) + + return data + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response, + ): + """ + Runs on response from LLM API call + + It can be used to reject a response + + Uses DynamoAI guardrails to check the response for policy violations + """ + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + from litellm.types.guardrails import GuardrailEventHooks + + if ( + self.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.post_call + ) + is not True + ): + return + + verbose_proxy_logger.debug("async_post_call_success_hook response=%s", response) + + # Check if the ModelResponse has text content in its choices + # to avoid sending empty content to DynamoAI (e.g., during tool calls) + if isinstance(response, litellm.ModelResponse): + has_text_content = False + dynamoai_messages: List[Dict[str, Any]] = [] + + for choice in response.choices: + if isinstance(choice, litellm.Choices): + if choice.message.content and isinstance(choice.message.content, str): + has_text_content = True + dynamoai_messages.append({ + "role": choice.message.role or "assistant", + "content": choice.message.content + }) + + if not has_text_content: + verbose_proxy_logger.warning( + "DynamoAI: not running guardrail. No output text in response" + ) + return + + if dynamoai_messages: + result = await self._call_dynamoai_guardrails( + messages=dynamoai_messages, + text_type="output", + request_data=data, + ) + + verbose_proxy_logger.debug("Guardrails async_post_call_success_hook result=%s", result) + + # Process the guardrails response + processed_result = self._process_dynamoai_guardrails_response(result) + violations_detected = processed_result["violations_detected"] + + # If any violations are detected, raise an error + if violations_detected: + error_message = self._create_error_message(processed_result) + raise ValueError(error_message) + + # Add guardrail to applied guardrails header + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + request_data: dict, + ) -> AsyncGenerator[ModelResponseStream, None]: + """ + Passes the entire stream to the guardrail + + This is useful for guardrails that need to see the entire response, such as PII masking. + + Triggered by mode: 'post_call' + """ + async for item in response: + yield item + + @staticmethod + def get_config_model() -> Optional[Type[GuardrailConfigModel]]: + from litellm.types.proxy.guardrails.guardrail_hooks.dynamoai import ( + DynamoAIGuardrailConfigModel, + ) + + return DynamoAIGuardrailConfigModel + diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py b/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py new file mode 100644 index 00000000000..04123e3964d --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/dynamoai.py @@ -0,0 +1,118 @@ +# Type definitions for DynamoAI Guardrails API + +import enum +from typing import Any, Dict, List, Literal, Optional, TypedDict + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class DynamoAIMessage(TypedDict): + """Message structure for DynamoAI API""" + role: str + content: str + +class DynamoRequestMetadata(TypedDict): + endUserId: Optional[str] + +class DynamoTextType(str, enum.Enum): + MODEL_INPUT = "MODEL_INPUT" + MODEL_RESPONSE = "MODEL_RESPONSE" + + +class PolicyMethod(str, enum.Enum): + PII = "PII" + TOXICITY = "TOXICITY" + ALIGNMENT = "ALIGNMENT" + HALLUCINATION = "HALLUCINATION" + RAG_HALLUCINATION = "RAG_HALLUCINATION" + + +class PolicyApplicableTo(str, enum.Enum): + INPUT = "INPUT" + OUTPUT = "OUTPUT" + ALL = "ALL" + + +class DynamoAIRequest(TypedDict, total=False): + """Request structure for DynamoAI /moderation/analyze endpoint""" + messages: List[Dict[str, Any]] + textType: Optional[DynamoTextType] + policyIds: List[str] + modelId: Optional[str] + clientId: Optional[str] + metadata: Optional[DynamoRequestMetadata] + + +class PolicyInfo(TypedDict, total=False): + """Policy information from DynamoAI response""" + id: str + name: str + description: str + method: PolicyMethod + action: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"] + methodParams: Dict[str, Any] + decisionParams: Dict[str, Any] + applicableTo: PolicyApplicableTo + created_at: str + creatorId: str + + +class PolicyOutputs(TypedDict, total=False): + """Outputs from the policy""" + action: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"] + message: Optional[str] + + +class AppliedPolicyDto(TypedDict, total=False): + """Applied policy details from DynamoAI response""" + policy: PolicyInfo + outputs: Optional[Dict[str, Any]] + action: Optional[str] + + +class DynamoAIResponse(TypedDict, total=False): + """Response structure from DynamoAI /moderation/analyze endpoint""" + text: str + textType: DynamoTextType + finalAction: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"] + appliedPolicies: List[AppliedPolicyDto] + error: Optional[str] + + +class DynamoAIProcessedResult(TypedDict): + """Processed result from DynamoAI guardrail check""" + violations_detected: List[str] + violation_details: Dict[str, Any] + + + +class DynamoAIGuardrailConfigModel(GuardrailConfigModel): + """Configuration model for DynamoAI Guardrails""" + + api_key: Optional[str] = Field( + default=None, + description="API key for DynamoAI Guardrails. If not provided, the `DYNAMOAI_API_KEY` environment variable is checked.", + ) + api_base: Optional[str] = Field( + default=None, + description="Base URL for DynamoAI API. If not provided, the `DYNAMOAI_API_BASE` environment variable is checked, defaults to https://api.dynamo.ai", + ) + policy_ids: Optional[List[str]] = Field( + default=None, + description="List of DynamoAI policy IDs to apply. If not provided, the `DYNAMOAI_POLICY_IDS` environment variable is checked (comma-separated).", + ) + model_id: Optional[str] = Field( + default=None, + description="Model ID for tracking/logging purposes. If not provided, the `DYNAMOAI_MODEL_ID` environment variable is checked.", + ) + guardrail_name: Optional[str] = Field( + default=None, + description="Name of the guardrail for identification in logs and traces.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "DynamoAI Guardrails" + diff --git a/tests/guardrails_tests/test_dynamoai_guardrails.py b/tests/guardrails_tests/test_dynamoai_guardrails.py new file mode 100644 index 00000000000..65bb9e27dc7 --- /dev/null +++ b/tests/guardrails_tests/test_dynamoai_guardrails.py @@ -0,0 +1,126 @@ +""" +Test DynamoAI Guardrails integration +""" +import sys +import os +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.proxy.guardrails.guardrail_hooks.dynamoai import DynamoAIGuardrails +from litellm.proxy._types import UserAPIKeyAuth +from litellm.caching.caching import DualCache +from unittest.mock import AsyncMock, MagicMock + + +@pytest.mark.asyncio +async def test_dynamoai_blocks_content_with_block_action(): + """ + Test that DynamoAI guardrail blocks content when finalAction is BLOCK. + """ + # Create guardrail instance + guardrail = DynamoAIGuardrails( + guardrail_name="test-dynamoai", + api_key="test-api-key", + api_base="https://api.dynamo.ai", + ) + + # Mock the DynamoAI API response with BLOCK action + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "text": "This is harmful content", + "textType": "MODEL_INPUT", + "finalAction": "BLOCK", + "appliedPolicies": [ + { + "policy": { + "id": "policy-123", + "name": "Toxicity Policy", + "description": "Blocks toxic content", + "method": "TOXICITY", + }, + "outputs": { + "action": "BLOCK", + "message": "Content contains toxic language" + } + } + ] + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + request_data = { + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "This is harmful content"} + ], + } + + # Mock should_run_guardrail to return True + guardrail.should_run_guardrail = MagicMock(return_value=True) + + # Test that the guardrail raises ValueError for blocked content + with pytest.raises(ValueError) as exc_info: + await guardrail.async_pre_call_hook( + data=request_data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + cache=MagicMock(spec=DualCache), + ) + + # Verify the error message contains policy information + error_message = str(exc_info.value) + assert "Guardrail failed" in error_message + assert "TOXICITY POLICY" in error_message.upper() + assert "BLOCK" in error_message.upper() + + +@pytest.mark.asyncio +async def test_dynamoai_allows_content_with_none_action(): + """ + Test that DynamoAI guardrail allows content when finalAction is NONE. + """ + # Create guardrail instance + guardrail = DynamoAIGuardrails( + guardrail_name="test-dynamoai", + api_key="test-api-key", + api_base="https://api.dynamo.ai", + ) + + # Mock the DynamoAI API response with NONE action (no violations) + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "text": "Hello, how are you?", + "textType": "MODEL_INPUT", + "finalAction": "NONE", + "appliedPolicies": [] + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + request_data = { + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "Hello, how are you?"} + ], + } + + # Mock should_run_guardrail to return True + guardrail.should_run_guardrail = MagicMock(return_value=True) + + # Test that the guardrail allows the content (no exception raised) + result = await guardrail.async_pre_call_hook( + data=request_data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + cache=MagicMock(spec=DualCache), + ) + + # Should return the request data unchanged + assert result == request_data + + + + diff --git a/tests/guardrails_tests/test_lasso_guardrails.py b/tests/guardrails_tests/test_lasso_guardrails.py index 7c321054002..e25007652fb 100644 --- a/tests/guardrails_tests/test_lasso_guardrails.py +++ b/tests/guardrails_tests/test_lasso_guardrails.py @@ -103,37 +103,40 @@ async def test_callback(): } # Test violation detection + mock_response = Response( + json={ + "violations_detected": True, + "deputies": { + "jailbreak": True, + "custom-policies": False, + "sexual": False, + "hate": False, + "illegality": False, + "violence": False, + "pattern-detection": False, + }, + "deputies_predictions": { + "jailbreak": 0.923, + "custom-policies": 0.234, + "sexual": 0.145, + "hate": 0.156, + "illegality": 0.167, + "violence": 0.178, + "pattern-detection": 0.189, + }, + "findings": { + "jailbreak": [{"action": "BLOCK", "severity": "HIGH"}] + } + }, + status_code=200, + request=Request( + method="POST", url="https://server.lasso.security/gateway/v2/classify" + ), + ) + mock_response.raise_for_status = lambda: None + with pytest.raises(HTTPException) as excinfo: - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - return_value=Response( - json={ - "deputies": { - "jailbreak": True, - "custom-policies": False, - "sexual": False, - "hate": False, - "illegality": False, - "violence": False, - "pattern-detection": False, - }, - "deputies_predictions": { - "jailbreak": 0.923, - "custom-policies": 0.234, - "sexual": 0.145, - "hate": 0.156, - "illegality": 0.167, - "violence": 0.178, - "pattern-detection": 0.189, - }, - "violations_detected": True, - }, - status_code=200, - request=Request( - method="POST", url="https://server.lasso.security/gateway/v1/chat" - ), - ), - ): + with patch.object(lasso_guardrail.async_handler, "post", return_value=mock_response): await lasso_guardrail.async_pre_call_hook( data=data, cache=DualCache(), @@ -146,36 +149,37 @@ async def test_callback(): assert "jailbreak" in str(excinfo.value.detail) # Test no violation - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - return_value=Response( - json={ - "deputies": { - "jailbreak": False, - "custom-policies": False, - "sexual": False, - "hate": False, - "illegality": False, - "violence": False, - "pattern-detection": False, - }, - "deputies_predictions": { - "jailbreak": 0.123, - "custom-policies": 0.234, - "sexual": 0.145, - "hate": 0.156, - "illegality": 0.167, - "violence": 0.178, - "pattern-detection": 0.189, - }, - "violations_detected": False, + mock_response_no_violation = Response( + json={ + "violations_detected": False, + "deputies": { + "jailbreak": False, + "custom-policies": False, + "sexual": False, + "hate": False, + "illegality": False, + "violence": False, + "pattern-detection": False, }, - status_code=200, - request=Request( - method="POST", url="https://server.lasso.security/gateway/v1/chat" - ), + "deputies_predictions": { + "jailbreak": 0.123, + "custom-policies": 0.234, + "sexual": 0.145, + "hate": 0.156, + "illegality": 0.167, + "violence": 0.178, + "pattern-detection": 0.189, + }, + "findings": {} + }, + status_code=200, + request=Request( + method="POST", url="https://server.lasso.security/gateway/v2/classify" ), - ): + ) + mock_response_no_violation.raise_for_status = lambda: None + + with patch.object(lasso_guardrail.async_handler, "post", return_value=mock_response_no_violation): result = await lasso_guardrail.async_pre_call_hook( data=data, cache=DualCache(), @@ -231,10 +235,7 @@ async def test_api_error_handling(): } # Test handling of connection error - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - side_effect=Exception("Connection error"), - ): + with patch.object(lasso_guardrail.async_handler, "post", side_effect=Exception("Connection error")): # Expect the guardrail to raise a LassoGuardrailAPIError with pytest.raises(LassoGuardrailAPIError) as excinfo: await lasso_guardrail.async_pre_call_hook( @@ -249,10 +250,7 @@ async def test_api_error_handling(): assert "Connection error" in str(excinfo.value) # Test with a different error message - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - side_effect=Exception("API timeout"), - ): + with patch.object(lasso_guardrail.async_handler, "post", side_effect=Exception("API timeout")): # Expect the guardrail to raise a LassoGuardrailAPIError with pytest.raises(LassoGuardrailAPIError) as excinfo: await lasso_guardrail.async_pre_call_hook(