From e8c5d51aeafb1b8c8aa92c4bb7c855305d26abc5 Mon Sep 17 00:00:00 2001 From: splendor023 Date: Thu, 13 Aug 2026 16:43:09 +0800 Subject: [PATCH] feat(guardrails): add aliyun security guardrail integration --- .../guardrail_hooks/aliyun/__init__.py | 94 + .../aliyun/aliyun_ai_guardrail.py | 959 ++++++++++ .../guardrails/guardrail_hooks/aliyun/base.py | 103 ++ litellm/types/guardrails.py | 52 + .../aliyun/aliyun_ai_guardrail.py | 165 ++ .../aliyun/test_aliyun_ai_guardrail.py | 1548 +++++++++++++++++ 6 files changed, 2921 insertions(+) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/aliyun/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/aliyun/test_aliyun_ai_guardrail.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/aliyun/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/aliyun/__init__.py new file mode 100644 index 00000000000..7ff7e7d4256 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/__init__.py @@ -0,0 +1,94 @@ +""" +Aliyun AI Security Guardrail Integration for LiteLLM +阿里云AI安全护栏集成 +This module provides integration with Aliyun's AI Security Guardrail service for: +- ContentModeration 内容合规检测 +- PromptAttack 提示词攻击检测 +- SensitiveData 敏感内容检测 +- ModelHallucination 模型幻觉 +- MaliciousUrl 恶意URL检测 +... +Documentation: https://help.aliyun.com/document_detail/2873209.html +""" + +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .aliyun_ai_guardrail import AliyunAIGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> AliyunAIGuardrail: + """ + Initialize an Aliyun AI Guardrail instance. + Credentials are configured in config.yaml (litellm_params) and support + os.environ/ references: + - access_key_id: Aliyun Access Key ID + - access_key_secret: Aliyun Access Key Secret + Args: + litellm_params: The LiteLLM parameters for the guardrail + guardrail: The guardrail configuration + Returns: + AliyunAIGuardrail instance + """ + import litellm + from litellm.secret_managers.main import get_secret_str + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError("Aliyun AI Guardrail: guardrail_name is required") + + level = getattr(litellm_params, "level", None) + max_text_length = getattr(litellm_params, "max_text_length", None) + stream_window_size = getattr(litellm_params, "stream_window_size", None) + stream_slide_step = getattr(litellm_params, "stream_slide_step", None) + stream_first_check_step = getattr(litellm_params, "stream_first_check_step", None) + region_id = getattr(litellm_params, "region_id", None) + service_input = getattr(litellm_params, "service_input", None) + service_output = getattr(litellm_params, "service_output", None) + service_mcp = getattr(litellm_params, "service_mcp", None) + + # Get credentials from config. These custom fields are not auto-resolved by + # guardrail_registry.py (only api_key/api_base are), so resolve os.environ/ + # references manually here. + access_key_id = getattr(litellm_params, "access_key_id", None) + access_key_secret = getattr(litellm_params, "access_key_secret", None) + if isinstance(access_key_id, str) and access_key_id.startswith("os.environ/"): + access_key_id = get_secret_str(access_key_id) + if isinstance(access_key_secret, str) and access_key_secret.startswith("os.environ/"): + access_key_secret = get_secret_str(access_key_secret) + + aliyun_guardrail = AliyunAIGuardrail( + guardrail_name=guardrail_name, + access_key_id=access_key_id, + access_key_secret=access_key_secret, + level=level, + max_text_length=max_text_length, + stream_window_size=stream_window_size, + stream_slide_step=stream_slide_step, + stream_first_check_step=stream_first_check_step, + region_id=region_id, + service_input=service_input, + service_output=service_output, + service_mcp=service_mcp, + default_on=litellm_params.default_on, + event_hook=litellm_params.mode, + ) + + litellm.logging_callback_manager.add_litellm_callback(aliyun_guardrail) + + return aliyun_guardrail + + +# Registry for guardrail initializers +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.ALIYUN_AI_GUARDRAIL.value: initialize_guardrail, +} + +# Registry for guardrail classes +guardrail_class_registry = { + SupportedGuardrailIntegrations.ALIYUN_AI_GUARDRAIL.value: AliyunAIGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py new file mode 100644 index 00000000000..dafbcc72e10 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py @@ -0,0 +1,959 @@ +#!/usr/bin/env python3 +""" +Aliyun AI Security Guardrail Integration for LiteLLM +阿里云AI安全护栏集成 +This guardrail scans prompts and responses using the Aliyun AI Security Guardrail API to detect: +- Content moderation violations +- Sensitive data (PII) +- Prompt injection attacks +- Malicious URLs +Documentation: https://help.aliyun.com/document_detail/2875413.html +Credentials: +Configured in config.yaml (litellm_params), support os.environ/ references: +- access_key_id: Aliyun Access Key ID +- access_key_secret: Aliyun Access Key Secret +""" + +from __future__ import annotations + +import asyncio +import base64 +import hashlib +import hmac +import json +import re +import uuid +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, AsyncGenerator, Literal +from urllib.parse import quote + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.proxy._types import UserAPIKeyAuth + +from .base import AliyunGuardrailBase + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.llms.openai import AllMessageValues + from litellm.types.mcp import MCPPostCallResponseObject + from litellm.types.proxy.guardrails.guardrail_hooks.aliyun.aliyun_ai_guardrail import ( + AliyunAIGuardrailResponse, + ) + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + from litellm.types.utils import EmbeddingResponse, ImageResponse, ModelResponse + +# Constants +ENCODING = "UTF-8" +ISO8601_DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ" +ALGORITHM = "HmacSHA1" + +# Region to endpoint mapping +REGION_ENDPOINTS = { + "cn-shanghai": "green-cip.cn-shanghai.aliyuncs.com", + "cn-beijing": "green-cip.cn-beijing.aliyuncs.com", + "cn-hangzhou": "green-cip.cn-hangzhou.aliyuncs.com", + "cn-shenzhen": "green-cip.cn-shenzhen.aliyuncs.com", + "cn-chengdu": "green-cip.cn-chengdu.aliyuncs.com", + "ap-southeast-1": "green-cip.ap-southeast-1.aliyuncs.com", + "eu-central-1": "green-cip.eu-central-1.aliyuncs.com", +} + +# Service codes for domestic (China) regions +SERVICE_INPUT_DOMESTIC = "query_security_check_pro" +SERVICE_OUTPUT_DOMESTIC = "response_security_check_pro" + +# Service codes for international regions +SERVICE_INPUT_INTERNATIONAL = "query_security_check_cb" +SERVICE_OUTPUT_INTERNATIONAL = "response_security_check_cb" + + +# Detection types +CONTENT_MODERATION_TYPE = "contentModeration" +PROMPT_ATTACK_TYPE = "promptAttack" +SENSITIVE_DATA_TYPE = "sensitiveData" +MALICIOUS_URL_TYPE = "maliciousUrl" +MODEL_HALLUCINATION_TYPE = "modelHallucination" +CUSTOM_LABEL_TYPE = "customLabel" + +# Suggestion returned by Aliyun when it has decided the content must be rejected +BLOCK_SUGGESTION = "block" + +# An explicit upstream block that carries no parseable severity is treated as the most +# severe level, so it is still weighed against the configured threshold rather than +# being silently downgraded to "none". +UNRESOLVED_BLOCK_LEVEL = "high" + + +def level_to_int(risk_level: str) -> int: + """ + Convert risk level string to integer for comparison. + Higher value = higher risk. + Supports both standard risk levels (none/low/medium/high) + and sensitive data levels (S0/S1/S2/S3/S4). + Mapping: + - none/S0 = 0 (no risk) + - low/S1 = 1 (low risk) + - medium/S2 = 2 (medium risk) + - high/S3/S4 = 3 (high risk) + """ + level_lower = risk_level.lower() if risk_level else "none" + level_map = { + # Standard risk levels + "none": 0, + "low": 1, + "medium": 2, + "high": 3, + # Sensitive data levels (mapped to standard levels) + "s0": 0, # No risk + "s1": 1, # Low risk + "s2": 2, # Medium risk + "s3": 3, # High risk + "s4": 3, # High risk (highest sensitive level) + } + return level_map.get(level_lower, 0) + + +# Protection level thresholds +# If detected_level >= threshold, then block +PROTECTION_LEVEL_THRESHOLD = { + "low": 1, # High protection: block low, medium, high (threshold=1, block if >=1) + "medium": 2, # Medium protection: block medium, high (threshold=2, block if >=2) + "high": 3, # Low protection: block high only (threshold=3, block if >=3) + "max": 99, # Observation mode: never block (threshold very high) +} + + +class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): + """ + LiteLLM Built-in Guardrail for Aliyun AI Security Guardrail. + This guardrail scans prompts and responses using the Aliyun AI Security Guardrail API to detect + malicious content, injection attempts, sensitive data, and policy violations. + Configuration: + guardrail_name: Name of the guardrail instance + access_key_id: Aliyun Access Key ID + access_key_secret: Aliyun Access Key Secret + region_id: Aliyun region ID (default: cn-shanghai) + default_on: Whether to enable by default + """ + + def __init__( + self, + guardrail_name: str, + access_key_id: str | None = None, + access_key_secret: str | None = None, + region_id: str | None = None, + level: str | None = None, + max_text_length: int | None = None, + stream_window_size: int | None = None, + stream_slide_step: int | None = None, + stream_first_check_step: int | None = None, + service_input: str | None = None, + service_output: str | None = None, + service_mcp: str | None = None, + **kwargs, + ): + """ + Initialize Aliyun AI Guardrail handler. + Credentials (access_key_id / access_key_secret) are passed in from config.yaml + via the guardrail loader. + Args: + region_id: Aliyun region ID (default: cn-shanghai) + level: Protection level for risk filtering + - "low": High protection, block all risks (low, medium, high, S1+) + - "medium": Medium protection, block medium and high risks (medium, high, S2+) + - "high": Low protection, block only high risks (high, S3+) + - "max": Observation mode, no blocking + service_input: Service code for input detection (default: query_security_check_pro) + service_output: Service code for output detection (default: response_security_check_pro) + service_mcp: Service code for MCP tool call detection, used by both + pre_mcp_call and post_mcp_call (default: query_security_check_pro) + """ + super().__init__( + guardrail_name=guardrail_name, + **kwargs, + ) + self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.access_key_id = access_key_id or "" + self.access_key_secret = access_key_secret or "" + self.region_id = region_id or "cn-shanghai" + if not self.access_key_id: + raise ValueError( + "Aliyun AI Guardrail: ak is required. Set access_key_id in config.yaml (supports os.environ/ reference)" + ) + if not self.access_key_secret: + raise ValueError( + "Aliyun AI Guardrail: sk is required. " + "Set access_key_secret in config.yaml (supports os.environ/ reference)" + ) + self.level = level or "medium" + if self.level not in PROTECTION_LEVEL_THRESHOLD: + raise ValueError( + f"Aliyun AI Guardrail: Invalid level '{self.level}'. " + f"Valid values are: {list(PROTECTION_LEVEL_THRESHOLD.keys())}" + ) + self.max_text_length = max_text_length or 2000 + self.endpoint = REGION_ENDPOINTS.get(self.region_id, REGION_ENDPOINTS["cn-shanghai"]) + self.service_url = f"https://{self.endpoint}" + self.service_input = service_input or SERVICE_INPUT_DOMESTIC + self.service_output = service_output or SERVICE_OUTPUT_DOMESTIC + self.service_mcp = service_mcp or SERVICE_INPUT_DOMESTIC + self.stream_window_size = stream_window_size or 500 + self.stream_slide_step = stream_slide_step or 300 + self.stream_first_check_step = stream_first_check_step or 50 + verbose_proxy_logger.info( + f"Initialized Aliyun AI Security Guardrail: {guardrail_name}, " + f"region: {self.region_id}, level: {self.level}, " + f"service_input: {self.service_input}, service_output: {self.service_output}, " + f"service_mcp: {self.service_mcp}" + ) + + @staticmethod + def get_config_model() -> type[GuardrailConfigModel] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.aliyun.aliyun_ai_guardrail import ( + AliyunAIGuardrailConfigModel, + ) + + return AliyunAIGuardrailConfigModel + + @staticmethod + def _format_iso8601_date() -> str: + """Format current timestamp in ISO8601 format.""" + return datetime.now(timezone.utc).strftime(ISO8601_DATE_FORMAT) + + @staticmethod + def _percent_encode(value: str | None) -> str: + """URL encode a value according to Aliyun signature requirements.""" + if value is None: + return "" + return quote(value.encode(ENCODING), safe="~").replace("+", "%20").replace("*", "%2A") + + def _create_signature(self, string_to_sign: str) -> str: + """Create HMAC-SHA1 signature for API request.""" + secret = self.access_key_secret + "&" + signature = hmac.new( + secret.encode(ENCODING), + string_to_sign.encode(ENCODING), + hashlib.sha1, + ).digest() + return base64.b64encode(signature).decode(ENCODING) + + def _create_string_to_sign(self, http_method: str, parameters: dict[str, str]) -> str: + """Create the string to sign for API request.""" + sorted_keys = sorted(parameters.keys()) + canonicalized_query_string = "" + for key in sorted_keys: + canonicalized_query_string += "&" + self._percent_encode(key) + "=" + self._percent_encode(parameters[key]) + string_to_sign = ( + http_method + "&" + self._percent_encode("/") + "&" + self._percent_encode(canonicalized_query_string[1:]) + ) + return string_to_sign + + def _split_text(self, text: str, max_length: int = 2000) -> list[str]: + """ + Split text into segments of maximum length, trying to preserve sentence boundaries. + Args: + text: Text to split + max_length: Maximum length of each segment + Returns: + List of text segments + """ + segments = [] + while len(text) > max_length: + chunk = text[:max_length] + match = None + for pattern in [r"[。！？；:\.?!]+"]: + matches = list(re.finditer(pattern, chunk)) + if matches: + match = matches[-1] + if match: + cut_point = match.end() + else: + cut_point = max_length + segments.append(text[:cut_point]) + text = text[cut_point:] + if text: + segments.append(text) + return segments + + async def async_make_request( + self, + text: str | None = None, + service_type: Literal["input", "output", "mcp"] = "input", + image_urls: list[str] | None = None, + ) -> AliyunAIGuardrailResponse: + """ + Make a request to the Aliyun AI Security Guardrail API. + Args: + text: Text to check (optional when only images are checked) + service_type: "input" for query_security_check, "output" for response_security_check, + "mcp" for MCP tool call check (uses service_mcp config) + image_urls: Public image URLs to check (optional) + Returns: + AliyunAIGuardrailResponse + """ + from litellm.types.proxy.guardrails.guardrail_hooks.aliyun.aliyun_ai_guardrail import ( + AliyunAIGuardrailResponse, + ) + + if service_type == "mcp": + service_code = self.service_mcp + elif service_type == "input": + service_code = self.service_input + else: + service_code = self.service_output + service_parameters: dict[str, Any] = {"requestFrom": "LiteLLM"} + if text: + service_parameters["content"] = text + if image_urls: + service_parameters["imageUrls"] = image_urls + parameters = { + "Action": "MultiModalGuard", + "Version": "2022-03-02", + "AccessKeyId": self.access_key_id, + "Timestamp": self._format_iso8601_date(), + "SignatureMethod": "HMAC-SHA1", + "SignatureVersion": "1.0", + "SignatureNonce": str(uuid.uuid4()), + "Format": "JSON", + "Service": service_code, + "ServiceParameters": json.dumps(service_parameters, ensure_ascii=False), + } + string_to_sign = self._create_string_to_sign("POST", parameters) + signature = self._create_signature(string_to_sign) + parameters["Signature"] = signature + verbose_proxy_logger.debug( + "Aliyun AI Guardrail request: service=%s, text_length=%d, image_count=%d", + service_code, + len(text) if text else 0, + len(image_urls) if image_urls else 0, + ) + response = await self.async_handler.post( + url=self.service_url, + data=parameters, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + timeout=30.0, + ) + body = response.json() + verbose_proxy_logger.debug("Aliyun AI Guardrail response: %s", body) + if response.status_code != 200: + raise HTTPException( + status_code=response.status_code, + detail={"error": f"Aliyun AI Guardrail request failed. Status: {response.status_code}, Body: {body}"}, + ) + if body.get("Code") != 200: + raise HTTPException( + status_code=400, + detail={ + "error": f"Aliyun AI Guardrail API error. Code: {body.get('Code')}, Message: {body.get('Message')}" + }, + ) + return AliyunAIGuardrailResponse( + RequestId=body.get("RequestId", ""), + Code=body.get("Code", 0), + Message=body.get("Message"), + Data=body.get("Data"), + ) + + def _should_block_by_level(self, detected_level: str) -> bool: + """ + Check if the detected risk level should trigger blocking based on protection level. + Logic: If detected_level_int >= threshold_int, then should block. + Args: + detected_level: Risk level from API response (none/low/medium/high or S0/S1/S2/S3/S4) + Returns: + True if should block, False otherwise + """ + threshold = PROTECTION_LEVEL_THRESHOLD.get(self.level, 99) + detected_int = level_to_int(detected_level) + return detected_int >= threshold + + def _resolve_detail_level(self, detail: dict[str, Any]) -> str: + """ + Resolve the risk level of a single Detail entry. + + MultiModalGuard reports severity in two shapes: the ``_pro`` service codes + return ``Detail[].Level``, while the documented response carries it as + ``Detail[].Result[].RiskLevel``. Honouring only the former downgrades the + latter to "none", which would let content Aliyun rejected pass through. + Args: + detail: A single entry of the response ``Detail`` list + Returns: + The risk level string to weigh against the configured threshold + """ + level = detail.get("Level") + if isinstance(level, str) and level.strip(): + return level + # Fall back to the highest RiskLevel reported across the individual results + resolved = "" + resolved_int = -1 + for result in detail.get("Result") or []: + if not isinstance(result, dict): + continue + risk_level = result.get("RiskLevel") + if not isinstance(risk_level, str) or not risk_level.strip(): + continue + risk_int = level_to_int(risk_level) + if risk_int > resolved_int: + resolved_int = risk_int + resolved = risk_level + if resolved: + return resolved + # No parseable severity: never downgrade an explicit block to "none" + if detail.get("Suggestion") == BLOCK_SUGGESTION: + return UNRESOLVED_BLOCK_LEVEL + return "none" + + def _parse_response_and_check( + self, + response: AliyunAIGuardrailResponse, + check_type: Literal["input", "output"], + ) -> dict[str, Any]: + """ + Parse the guardrail response and check if content should be blocked. + Blocking logic: + Check if detected level >= threshold based on protection level setting + Args: + response: The API response + check_type: "input" or "output" + Returns: + Dict with parsed results + Raises: + HTTPException if content should be blocked + """ + data = response.get("Data", {}) + if not data: + return {"flagged": False, "suggestion": "pass", "details": {}, "message": ""} + final_suggestion = data.get("Suggestion", "pass") + detail_list = data.get("Detail") or [] + details: dict[str, dict[str, Any]] = {} + desensitization = "" + should_block = False + blocked_type = "" + blocked_level = "" + block_message = "" + for detail in detail_list: + detection_type = detail.get("Type", "") + detected_level = self._resolve_detail_level(detail) + suggestion = detail.get("Suggestion", "pass") + results = detail.get("Result", []) + details[detection_type] = { + "level": detected_level, + "suggestion": suggestion, + "results": results, + } + if detection_type == SENSITIVE_DATA_TYPE and results: + for result in results: + ext = result.get("Ext", {}) + if ext and ext.get("Desensitization"): + desensitization = ext.get("Desensitization", "") + break + if not should_block and self._should_block_by_level(detected_level): + should_block = True + blocked_type = detection_type + blocked_level = detected_level + block_message = f"检测到{detection_type} (风险等级: {detected_level})" + # Aliyun rejected the content overall but no single detection could be attributed + # (e.g. an empty Detail list, or every entry reporting pass). Treat it as the most + # severe level so the decision is not lost, still subject to the threshold. + if not should_block and final_suggestion == BLOCK_SUGGESTION: + if self._should_block_by_level(UNRESOLVED_BLOCK_LEVEL): + should_block = True + blocked_level = UNRESOLVED_BLOCK_LEVEL + block_message = f"阿里云返回阻断建议 (Suggestion: {final_suggestion})" + verbose_proxy_logger.debug( + f"Aliyun AI Guardrail: level={self.level}, " + f"check_type={check_type}, should_block={should_block}, " + f"blocked_type={blocked_type}, blocked_level={blocked_level}" + ) + if should_block: + raise HTTPException( + status_code=400, + detail={ + "error": f"Aliyun AI Guardrail: {block_message}", + "type": check_type, + "details": details, + }, + ) + return { + "flagged": final_suggestion != "pass", + "suggestion": final_suggestion, + "desensitization": desensitization, + "details": details, + "message": block_message, + } + + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: Any, + data: dict[str, Any], + call_type: Literal[ + "completion", + "text_completion", + "embeddings", + "image_generation", + "moderation", + "audio_transcription", + "pass_through_endpoint", + "rerank", + "mcp_call", + "anthropic_messages", + "responses", + ], + ) -> dict[str, Any] | None: + """ + Pre-call hook to scan user prompts before sending to LLM. + Raises HTTPException if content should be blocked. + """ + verbose_proxy_logger.info( + "Aliyun AI Guardrail: Running pre-call prompt scan, on call_type: %s", + call_type, + ) + if call_type == "call_mcp_tool": + return await self._mcp_pre_call_check(data) + new_messages: list[AllMessageValues] | None = data.get("messages") + if new_messages is None: + verbose_proxy_logger.warning("Aliyun AI Guardrail: not running guardrail. No messages in data") + return data + user_prompt = self.get_user_prompt(new_messages) + image_urls = self.get_image_urls(new_messages) + if not user_prompt and not image_urls: + verbose_proxy_logger.warning("Aliyun AI Guardrail: No user prompt or image found") + return None + verbose_proxy_logger.info( + "Aliyun AI Guardrail: Pre-call scan started, prompt length: %d, image count: %d", + len(user_prompt) if user_prompt else 0, + len(image_urls), + ) + if user_prompt: + if len(user_prompt) > self.max_text_length: + segments = self._split_text(user_prompt, self.max_text_length) + else: + segments = [user_prompt] + else: + segments = [] + # Attach all image URLs to the first text segment so content + images + # are checked together; remaining segments carry text only. When there + # is no text, send a single image-only request. + payloads: list[tuple] = [] + if segments: + for idx, segment in enumerate(segments): + payloads.append((segment, image_urls if idx == 0 and image_urls else None)) + elif image_urls: + payloads.append((None, image_urls)) + semaphore = asyncio.Semaphore(5) # Max 5 concurrent requests, MultiModalGuard API limit is 20 + + async def check_with_semaphore(segment_text: str | None, segment_images: list[str] | None): + async with semaphore: + return await self.async_make_request(text=segment_text, service_type="input", image_urls=segment_images) + + responses = await asyncio.gather(*[check_with_semaphore(t, imgs) for t, imgs in payloads]) + for response in responses: + self._parse_response_and_check(response, check_type="input") + verbose_proxy_logger.info("Aliyun AI Guardrail: Pre-call scan passed") + return None + + # ================================================================ + # MCP-specific guardrail methods + # ================================================================ + + async def _mcp_pre_call_check(self, data: dict) -> dict | None: + """MCP pre-call: audit tool name + arguments before execution.""" + messages = data.get("messages", []) + content = messages[0].get("content", "") if messages else "" + verbose_proxy_logger.info( + "Aliyun AI Guardrail: ★ MCP pre-call check started, content: %s", + content, + ) + if not content: + return None + if len(content) > self.max_text_length: + segments = self._split_text(content, self.max_text_length) + else: + segments = [content] + semaphore = asyncio.Semaphore(5) + + async def check(text: str): + async with semaphore: + return await self.async_make_request(text=text, service_type="mcp") + + responses = await asyncio.gather(*[check(s) for s in segments]) + for resp in responses: + self._parse_response_and_check(resp, check_type="input") + verbose_proxy_logger.info("Aliyun AI Guardrail: ★ MCP pre-call check passed") + return None + + def _should_run_post_mcp_call(self) -> bool: + """Check if post_mcp_call is configured in event_hook. + + Since GuardrailEventHooks enum does not include post_mcp_call, + we cannot use should_run_guardrail(). This method manually checks + whether the user configured 'post_mcp_call' in the guardrail mode. + + Returns True if post_mcp_call should run: + - event_hook is None → run for all events + - event_hook is a list containing 'post_mcp_call' + - event_hook is a string equal to 'post_mcp_call' + - event_hook is a Mode with 'post_mcp_call' in tags or default + """ + from litellm.types.guardrails import Mode + + if self.event_hook is None: + return True + if isinstance(self.event_hook, list): + return "post_mcp_call" in self.event_hook + if isinstance(self.event_hook, Mode): + for tag_value in self.event_hook.tags.values(): + if isinstance(tag_value, list): + if "post_mcp_call" in tag_value: + return True + elif tag_value == "post_mcp_call": + return True + if self.event_hook.default: + default_list = ( + self.event_hook.default if isinstance(self.event_hook.default, list) else [self.event_hook.default] + ) + return "post_mcp_call" in default_list + return False + return self.event_hook == "post_mcp_call" + + async def async_post_mcp_tool_call_hook( + self, + kwargs: dict, + response_obj: MCPPostCallResponseObject, + start_time: datetime, + end_time: datetime, + ) -> MCPPostCallResponseObject | None: + """MCP post-call: audit tool output after execution. + + Raising from here does not block: the dispatcher treats every callback exception + as a non-blocking logging error and hands the untouched tool result back. Both + call sites also discard this hook's return value, so a violation has to be + written into the live tool result carried by ``kwargs["original_response"]``. + The replacement object is returned as well, to honour the hook's contract. + """ + # Since GuardrailEventHooks enum has no post_mcp_call, the framework + # always invokes this hook if implemented. We check config manually. + if not self._should_run_post_mcp_call(): + verbose_proxy_logger.info("Aliyun AI Guardrail: Skipping post_mcp_call — not configured in event_hook") + return None + verbose_proxy_logger.info("Aliyun AI Guardrail: ★ MCP post-call check started") + original_response = kwargs.get("original_response") + combined_text = "" + for candidate in (original_response, getattr(response_obj, "mcp_tool_call_response", None)): + if candidate is None: + continue + combined_text = self._extract_mcp_tool_text(candidate) + if combined_text: + break + if not combined_text: + return None + verbose_proxy_logger.info( + "Aliyun AI Guardrail: ★ MCP post-call check, response length: %d", + len(combined_text), + ) + if len(combined_text) > self.max_text_length: + segments = self._split_text(combined_text, self.max_text_length) + else: + segments = [combined_text] + semaphore = asyncio.Semaphore(5) + + async def check(text: str): + async with semaphore: + return await self.async_make_request(text=text, service_type="mcp") + + try: + responses = await asyncio.gather(*[check(s) for s in segments]) + for resp in responses: + self._parse_response_and_check(resp, check_type="output") + except HTTPException as e: + verbose_proxy_logger.warning( + "Aliyun AI Guardrail: ★ MCP post-call blocked — tool output replaced with the violation detail" + ) + return self._block_mcp_tool_output( + detail=e.detail, + response_obj=response_obj, + original_response=original_response, + ) + except Exception as e: + # Raising here is swallowed by the dispatcher as a non-blocking logging + # error, which would hand the unaudited tool output straight to the + # client. Fail closed, matching every other path of this integration. + verbose_proxy_logger.error( + f"Aliyun AI Guardrail: ★ MCP post-call check failed, blocking unaudited tool output: {str(e)}", + exc_info=True, + ) + return self._block_mcp_tool_output( + detail={"error": f"Aliyun AI Guardrail 调用失败，工具输出未经审核，已拦截: {str(e)}"}, + response_obj=response_obj, + original_response=original_response, + ) + verbose_proxy_logger.info("Aliyun AI Guardrail: ★ MCP post-call check passed") + return None + + @staticmethod + def _iter_mcp_content_items(payload: Any) -> list[Any]: + """ + Normalise an MCP tool result into its list of content items. + Args: + payload: A CallToolResult, its content list, or a raw response body + Returns: + The content items to audit, empty when none can be located + """ + if isinstance(payload, str): + return [payload] + content = getattr(payload, "content", None) + if isinstance(content, list): + return content + if isinstance(payload, dict): + inner = payload.get("content") + return inner if isinstance(inner, list) else [payload] + if isinstance(payload, list): + # MCPPostCallResponseObject declares mcp_tool_call_response as a list, so a + # CallToolResult handed to it is coerced by iterating the model into + # (field, value) pairs. Recover the real content instead of auditing reprs. + for item in payload: + if isinstance(item, tuple) and len(item) == 2 and item[0] == "content": + if isinstance(item[1], list): + return item[1] + return payload + return [] + + def _extract_mcp_tool_text(self, payload: Any) -> str: + """Collect the textual output of an MCP tool result.""" + texts: list[str] = [] + for item in self._iter_mcp_content_items(payload): + if isinstance(item, str): + text = item + elif isinstance(item, dict): + text = item.get("text", "") or "" + else: + text = getattr(item, "text", "") or "" + if text: + texts.append(text) + return "\n".join(texts) + + @staticmethod + def _replace_tool_output_in_place(target: Any, blocked_content: list[Any]) -> bool: + """Overwrite an MCP tool result's content with ``blocked_content``, in place.""" + if target is None: + return False + content = getattr(target, "content", None) + if isinstance(content, list): + content[:] = blocked_content + if hasattr(target, "isError"): + try: + target.isError = True + except (AttributeError, TypeError, ValueError): + pass + return True + if isinstance(target, list): + target[:] = blocked_content + return True + if isinstance(target, dict): + result = target.get("result") + if isinstance(result, dict) and isinstance(result.get("content"), list): + result["content"] = list(blocked_content) + return True + if isinstance(target.get("content"), list): + target["content"] = list(blocked_content) + return True + return False + + def _block_mcp_tool_output( + self, + detail: Any, + response_obj: MCPPostCallResponseObject, + original_response: Any, + ) -> MCPPostCallResponseObject: + """Replace blocked MCP tool output, both in place and as the returned object.""" + from litellm.types.llms.base import HiddenParams + from litellm.types.mcp import MCPPostCallResponseObject as _MCPPostCallResponseObject + from mcp.types import TextContent + + payload = detail if isinstance(detail, dict) else {"error": str(detail)} + blocked_content: list[Any] = [TextContent(type="text", text=json.dumps(payload, ensure_ascii=False))] + for target in (original_response, getattr(response_obj, "mcp_tool_call_response", None)): + self._replace_tool_output_in_place(target, blocked_content) + hidden_params = getattr(response_obj, "hidden_params", None) + return _MCPPostCallResponseObject( + mcp_tool_call_response=blocked_content, + hidden_params=hidden_params if isinstance(hidden_params, HiddenParams) else HiddenParams(), + ) + + @log_guardrail_information + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Any | ModelResponse | EmbeddingResponse | ImageResponse, + ) -> Any: + """ + Post-call hook to scan LLM responses. + Raises HTTPException if content should be blocked. + """ + from litellm.types.utils import Choices, ModelResponse + + if isinstance(response, ModelResponse) and response.choices: + # Scan every choice: n>1 responses would otherwise return unchecked content + content = "\n".join( + str(choice.message.content) + for choice in response.choices + if isinstance(choice, Choices) and choice.message.content + ) + if content: + verbose_proxy_logger.info( + "Aliyun AI Guardrail: Post-call scan started, response length: %d", + len(content), + ) + if len(content) > self.max_text_length: + segments = self._split_text(content, self.max_text_length) + else: + segments = [content] + semaphore = asyncio.Semaphore(5) # Max 5 concurrent requests + + async def check_with_semaphore(segment: str): + async with semaphore: + return await self.async_make_request(text=segment, service_type="output") + + responses = await asyncio.gather(*[check_with_semaphore(segment) for segment in segments]) + for guardrail_response in responses: + self._parse_response_and_check(guardrail_response, check_type="output") + verbose_proxy_logger.info("Aliyun AI Guardrail: Post-call scan passed") + return response + + @staticmethod + def _extract_stream_chunk_text(chunk: Any) -> str: + """ + Collect every text field a streaming chunk can carry. + Covers both chat completion chunks and Responses API streaming events. + Auditing only ``delta.content`` would release tool call arguments, + reasoning text and every /v1/responses chunk to the client unchecked. + Args: + chunk: A streaming chunk + Returns: + The concatenated text to audit, empty when the chunk carries none + """ + parts: list[str] = [] + + def collect_call(function: Any) -> None: + for field in ("name", "arguments"): + value = getattr(function, field, None) + if value: + parts.append(str(value)) + + # Responses API events carry text outside of choices: either a plain + # string delta, or the assembled output of a terminal event. + event_delta = getattr(chunk, "delta", None) + if isinstance(event_delta, str) and event_delta: + parts.append(event_delta) + for output_item in getattr(getattr(chunk, "response", None), "output", None) or []: + content_parts = getattr(output_item, "content", None) + if content_parts: + for content_part in content_parts: + text = getattr(content_part, "text", None) + if text: + parts.append(str(text)) + else: + text = getattr(output_item, "text", None) + if text: + parts.append(str(text)) + + for choice in getattr(chunk, "choices", None) or []: + delta = getattr(choice, "delta", None) + if delta is None: + continue + for field in ("content", "reasoning_content"): + value = getattr(delta, field, None) + if value: + parts.append(str(value)) + for call in getattr(delta, "tool_calls", None) or []: + collect_call(getattr(call, "function", None)) + collect_call(getattr(delta, "function_call", None)) + return "".join(parts) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + request_data: dict[str, Any], + ) -> AsyncGenerator[Any, None]: + """ + Process streaming response with sliding window guardrail checks. + This method implements sliding window guardrail checks based on + `stream_window_size` and `stream_slide_step`. + It triggers a guardrail API call when: + 1. Every `stream_slide_step` new chars accumulate since the last check. + 2. Stream ends and there's remaining unchecked content. + For example, if stream_window_size=2000, stream_slide_step=300: + - At 300 chars: check chars 0-300 (window: last 2000) + - At 600 chars: check chars 0-600 (window: last 2000) + - At 2100 chars: check chars 100-2100 (window slides forward) + - At 2400 chars: check chars 400-2400 (window slides forward) + - When stream ends with 2500 chars: check chars 500-2500 (final window) + """ + accumulated_text = "" + last_check_position = 0 # Position (total length) when last check was triggered + pending_chunks = [] # Buffer chunks until guardrail check passes + chunk_count = 0 + is_first_check = True # First check uses smaller threshold to reduce first-token latency + verbose_proxy_logger.info( + "Aliyun AI Guardrail: Streaming scan started, window=%d, step=%d, first_check_step=%d", + self.stream_window_size, + self.stream_slide_step, + self.stream_first_check_step, + ) + try: + async for chunk in response: + chunk_text = self._extract_stream_chunk_text(chunk) + accumulated_text += chunk_text + chunk_count += 1 + # Buffer the chunk, don't yield until guardrail check passes + pending_chunks.append(chunk) + current_length = len(accumulated_text) + new_chars_since_last_check = current_length - last_check_position + check_threshold = self.stream_first_check_step if is_first_check else self.stream_slide_step + if new_chars_since_last_check >= check_threshold: + start = max(0, current_length - self.stream_window_size) + text_to_check = accumulated_text[start:current_length] + guardrail_response = await self.async_make_request(text=text_to_check, service_type="output") + self._parse_response_and_check(guardrail_response, check_type="output") + verbose_proxy_logger.info( + "Aliyun AI Guardrail: Streaming check passed at position %d", current_length + ) + for pending_chunk in pending_chunks: + yield pending_chunk + pending_chunks.clear() + last_check_position = current_length + is_first_check = False + # Stream ended - check any remaining unchecked content with a final window + if len(accumulated_text) > last_check_position: + start = max(0, len(accumulated_text) - self.stream_window_size) + remaining_text = accumulated_text[start:] + guardrail_response = await self.async_make_request(text=remaining_text, service_type="output") + self._parse_response_and_check(guardrail_response, check_type="output") + verbose_proxy_logger.info( + "Aliyun AI Guardrail: Streaming scan completed, total length: %d", len(accumulated_text) + ) + for pending_chunk in pending_chunks: + yield pending_chunk + pending_chunks.clear() + except HTTPException as e: + error_detail = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} + verbose_proxy_logger.info("Aliyun AI Guardrail: Streaming blocked at position %d", len(accumulated_text)) + yield f"data: {json.dumps({'error': error_detail}, ensure_ascii=False)}\n\n" + return + except Exception as e: + verbose_proxy_logger.error(f"Aliyun AI Guardrail streaming error: {str(e)}", exc_info=True) + raise diff --git a/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py b/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py new file mode 100644 index 00000000000..8917a06df5b --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py @@ -0,0 +1,103 @@ +""" +Base class for Aliyun guardrails +阿里云护栏基类 +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from litellm.types.llms.openai import AllMessageValues + + +class AliyunGuardrailBase: + """ + Base class for Aliyun guardrails. + """ + + def get_user_prompt(self, messages: list[AllMessageValues]) -> str | None: + """ + Get the last consecutive block of messages from the user. + Example: + messages = [ + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm good, thank you!"}, + {"role": "user", "content": "What is the weather in Tokyo?"}, + ] + get_user_prompt(messages) -> "What is the weather in Tokyo?" + """ + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, + ) + + if not messages: + return None + # Iterate from the end to find the last consecutive block of user messages + user_messages = [] + for message in reversed(messages): + if message.get("role") == "user": + user_messages.append(message) + else: + # Stop when we hit a non-user message + break + if not user_messages: + return None + # Reverse to get the messages in chronological order + user_messages.reverse() + user_prompt = "" + for message in user_messages: + text_content = convert_content_list_to_str(message) + user_prompt += text_content + "\n" + result = user_prompt.strip() + return result if result else None + + def get_image_urls(self, messages: list[AllMessageValues]) -> list[str]: + """ + Extract image URLs from the last consecutive block of user messages. + Only publicly accessible http(s) URLs are collected (in order, + de-duplicated). Uses the same message range as ``get_user_prompt``. + Example: + messages = [ + {"role": "user", "content": [ + {"type": "text", "text": "what is in this image?"}, + {"type": "image_url", "image_url": {"url": "https://a.com/x.png"}}, + ]}, + ] + get_image_urls(messages) -> ["https://a.com/x.png"] + """ + if not messages: + return [] + # Iterate from the end to find the last consecutive block of user messages + user_messages = [] + for message in reversed(messages): + if message.get("role") == "user": + user_messages.append(message) + else: + break + if not user_messages: + return [] + user_messages.reverse() + image_urls: list[str] = [] + seen = set() + for message in user_messages: + content = message.get("content") + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict) or part.get("type") != "image_url": + continue + image_url = part.get("image_url") + url: str | None = None + if isinstance(image_url, dict): + url = image_url.get("url") + elif isinstance(image_url, str): + url = image_url + if not isinstance(url, str): + continue + if not (url.startswith("http://") or url.startswith("https://")): + continue + if url not in seen: + seen.add(url) + image_urls.append(url) + return image_urls diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c7cdfaad780..9cdc3d2e155 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -134,6 +134,7 @@ class SupportedGuardrailIntegrations(Enum): HEADROOM = "headroom" COMPRESR = "compresr" STRAIKER = "straiker" + ALIYUN_AI_GUARDRAIL = "aliyun_ai_guardrail" class Role(Enum): @@ -654,6 +655,55 @@ class JavelinGuardrailConfigModel(BaseModel): config: dict | None = Field(default=None, description="Additional configuration for the guardrail") +class AliyunAIGuardrailConfigModel(BaseModel): + """Configuration parameters for the Aliyun AI Security guardrail.""" + + level: str | None = Field( + default=None, + description="Protection level. 'low': block all risks (high protection), 'medium': block medium+high risks, 'high': block high only, 'max': observe mode. Default: medium", + ) + max_text_length: int | None = Field( + default=None, + description="Maximum text length for a single API call. Text longer than this will be split.", + ) + stream_window_size: int | None = Field( + default=None, + description="Sliding window size (in chars) for streaming output guardrail checks. Default: 500", + ) + stream_slide_step: int | None = Field( + default=None, + description="Sliding step (in chars) for streaming output guardrail checks. Default: 300", + ) + stream_first_check_step: int | None = Field( + default=None, + description="First check threshold (in chars) to reduce first-token latency. Default: 50", + ) + region_id: str | None = Field( + default=None, + description="Aliyun region ID. Default: cn-shanghai", + ) + service_input: str | None = Field( + default=None, + description="Service code for input (pre-call) detection. Default: query_security_check", + ) + service_output: str | None = Field( + default=None, + description="Service code for output (post-call) detection. Default: response_security_check", + ) + service_mcp: str | None = Field( + default=None, + description="Service code for MCP tool call detection (pre/post MCP call). Default: query_security_check", + ) + access_key_id: str | None = Field( + default=None, + description="Aliyun Access Key ID for the guardrail. Configure in config.yaml, supports os.environ/ reference", + ) + access_key_secret: str | None = Field( + default=None, + description="Aliyun Access Key Secret for the guardrail. Configure in config.yaml, supports os.environ/ reference", + ) + + class ContentFilterAction(str, Enum): """Action to take when content filter detects a match""" @@ -999,6 +1049,7 @@ class LitellmParams( QostodianNexusConfigModel, VigilGuardGuardrailConfigModel, SingulrGuardrailConfigModel, + AliyunAIGuardrailConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") mode: str | list[str] | Mode = Field( @@ -1126,3 +1177,4 @@ class PatchGuardrailRequest(BaseModel): guardrail_name: str | None = None litellm_params: BaseLitellmParams | None = None guardrail_info: dict[str, Any] | None = None + diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py b/litellm/types/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py new file mode 100644 index 00000000000..2a612ecd6ad --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py @@ -0,0 +1,165 @@ +""" +Type definitions for Aliyun AI Security Guardrail +阿里云AI安全护栏类型定义 +Aliyun AI Guardrail supports the following detection types: +- contentModeration: Content safety moderation +- sensitiveData: Sensitive data detection (PII, etc.) +- promptAttack: Prompt injection attack detection +- maliciousUrl: Malicious URL detection +""" + +from typing import Any, Dict, List, Literal, Optional + +from pydantic import BaseModel, Field +from typing_extensions import NotRequired, TypedDict + +from ..base import GuardrailConfigModel + + +# Response types +class AliyunAIGuardrailResponseDetailResultExt(TypedDict, total=False): + """Extended information in result""" + + Desensitization: Optional[str] # Desensitized text when action is mask + + +class AliyunAIGuardrailResponseDetailResult(TypedDict, total=False): + """Result item in detail""" + + Confidence: Optional[float] + Label: Optional[str] + Ext: Optional[AliyunAIGuardrailResponseDetailResultExt] + # Per-result risk level. This is the shape documented for MultiModalGuard; the + # ``_pro`` service codes report the severity on the parent Detail as ``Level`` + # instead, so both have to be honoured when deciding whether to block. + RiskLevel: Optional[str] + + +class AliyunAIGuardrailResponseDetail(TypedDict): + """Detail item in response data""" + + Type: str # contentModeration, sensitiveData, promptAttack, maliciousUrl + Suggestion: str # pass, block, mask + Result: List[AliyunAIGuardrailResponseDetailResult] + # Risk level as returned by the ``_pro`` service codes (none/low/medium/high, or + # S0-S4 for sensitiveData). Absent in the documented response shape, which carries + # the severity as Result[].RiskLevel. + Level: NotRequired[str] + + +class AliyunAIGuardrailResponseData(TypedDict, total=False): + """Response data from Aliyun AI Guardrail API""" + + Suggestion: str # Overall suggestion: pass, block, mask + Detail: Optional[List[AliyunAIGuardrailResponseDetail]] + + +class AliyunAIGuardrailResponse(TypedDict): + """Response from Aliyun AI Guardrail API""" + + RequestId: str + Code: int + Message: Optional[str] + Data: Optional[AliyunAIGuardrailResponseData] + + +# Suggestion type +AliyunAIGuardrailSuggestion = Literal["pass", "block", "watch"] + +# Detection type +AliyunAIGuardrailDetectionType = Literal["contentModeration", "sensitiveData", "promptAttack", "maliciousUrl"] + + +class AliyunAIGuardrailRequestParams(TypedDict, total=False): + """Request parameters for Aliyun AI Guardrail API""" + + Action: str + Version: str + AccessKeyId: str + Timestamp: str + SignatureMethod: str + SignatureVersion: str + SignatureNonce: str + Format: str + Service: str + ServiceParameters: str + Signature: str + + +# Risk level literals +AliyunRiskLevel = Literal["none", "low", "medium", "high"] + +# Protection level literals +AliyunProtectionLevel = Literal["low", "medium", "high", "max"] + + +# Configuration models +class AliyunAIGuardrailOptionalParams(BaseModel): + """ + Optional parameters for Aliyun AI Guardrail. + Credentials (access_key_id / access_key_secret) are configured + in config.yaml on the AliyunAIGuardrailConfigModel and support os.environ/ references. + """ + + level: Optional[AliyunProtectionLevel] = Field( + default="medium", + description="Protection level for risk filtering. 'low': block all risks (high protection), 'medium': block medium and high risks, 'high': block only high risks (low protection), 'max': observation mode (no blocking). Default: medium", + ) + max_text_length: Optional[int] = Field( + default=2000, + description="Maximum text length for a single API call. Text longer than this will be split.", + ) + stream_window_size: Optional[int] = Field( + default=500, + description="Sliding window size (in chars) for streaming output guardrail checks. Each check sends the most recent N chars to the API.", + ) + stream_slide_step: Optional[int] = Field( + default=300, + description="Sliding step (in chars) for streaming output guardrail checks. A check is triggered every time N new chars accumulate since the last check.", + ) + stream_first_check_step: Optional[int] = Field( + default=50, + description="First check threshold (in chars) for streaming output. The first guardrail check triggers earlier (at N chars) to reduce first-token latency, subsequent checks use stream_slide_step.", + ) + region_id: Optional[str] = Field( + default="cn-shanghai", + description="Aliyun region ID. Default: cn-shanghai", + ) + service_input: Optional[str] = Field( + default="query_security_check_pro", + description="Service code for input (pre-call) detection. Default: query_security_check_pro", + ) + service_output: Optional[str] = Field( + default="response_security_check_pro", + description="Service code for output (post-call) detection. Default: response_security_check_pro", + ) + service_mcp: Optional[str] = Field( + default="query_security_check_pro", + description="Service code for MCP tool call detection (pre_mcp_call and post_mcp_call). Default: query_security_check_pro", + ) + + +class AliyunAIGuardrailConfigModel(GuardrailConfigModel[AliyunAIGuardrailOptionalParams]): + """ + Configuration model for Aliyun AI Guardrail. + Credentials are configured in config.yaml and support os.environ/ references: + - access_key_id: Aliyun Access Key ID + - access_key_secret: Aliyun Access Key Secret + """ + + access_key_id: Optional[str] = Field( + default=None, + description="Aliyun Access Key ID. Configure in config.yaml, supports os.environ/ reference", + ) + access_key_secret: Optional[str] = Field( + default=None, + description="Aliyun Access Key Secret. Configure in config.yaml, supports os.environ/ reference", + ) + optional_params: AliyunAIGuardrailOptionalParams = Field( + default_factory=AliyunAIGuardrailOptionalParams, + description="Optional parameters for the Aliyun AI Guardrail", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Aliyun AI Security Guardrail" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/aliyun/test_aliyun_ai_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/aliyun/test_aliyun_ai_guardrail.py new file mode 100644 index 00000000000..d3833cd8644 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/aliyun/test_aliyun_ai_guardrail.py @@ -0,0 +1,1548 @@ +""" +Unit tests for Aliyun AI Security Guardrail. + +Tests cover: +- Registration in guardrail system (enum, initializer, registry) +- Constructor validation (credentials, level, configurable service codes) +- Helper functions (level_to_int, _split_text) +- Blocking logic (_should_block_by_level, _parse_response_and_check) +- Pre-call hook (text + multimodal image URL detection) +- Post-call hook (blocks violations in response, passes clean response) +- Image URL extraction (get_image_urls) +- ServiceParameters construction (text / image / mixed combos) +- Config model (get_config_model, ui_friendly_name) +""" + +import json +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.aliyun.aliyun_ai_guardrail import ( + CONTENT_MODERATION_TYPE, + PROMPT_ATTACK_TYPE, + SENSITIVE_DATA_TYPE, + AliyunAIGuardrail, + level_to_int, +) + +FAKE_AK = "test-access-key-id" +FAKE_SK = "test-access-key-secret" + +IMG_A = "https://example.com/a.png" +IMG_B = "http://example.com/b.jpg" +DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + + +def _make_guardrail(**kwargs) -> AliyunAIGuardrail: + defaults = dict( + guardrail_name="test-aliyun", + access_key_id=FAKE_AK, + access_key_secret=FAKE_SK, + level="medium", + ) + defaults.update(kwargs) + return AliyunAIGuardrail(**defaults) + + +def _make_aliyun_api_response( + suggestion: str = "pass", + detail: list = None, + code: int = 200, +) -> MagicMock: + """Build a mock httpx.Response mimicking Aliyun API output.""" + mock = MagicMock() + mock.status_code = 200 + mock.json.return_value = { + "Code": code, + "RequestId": "test-req-id", + "Message": None, + "Data": { + "Suggestion": suggestion, + "Detail": detail or [], + }, + } + return mock + + +def _make_detail( + detection_type: str = CONTENT_MODERATION_TYPE, + level: str = "high", + suggestion: str = "block", + results: list = None, +) -> dict: + """Build a Detail item. Pass level=None to omit the Level field entirely, + mimicking the response shape documented for MultiModalGuard, which carries + the severity as Result[].RiskLevel instead.""" + detail = { + "Type": detection_type, + "Suggestion": suggestion, + "Result": results or [], + } + if level is not None: + detail["Level"] = level + return detail + + +def _captured_service_parameters(mock_post: AsyncMock): + """Return (ServiceParameters dict, Service code) from a mocked post call.""" + _, kwargs = mock_post.call_args + params = kwargs["data"] + return json.loads(params["ServiceParameters"]), params["Service"] + + +# --------------------------------------------------------------------------- +# Registration tests +# --------------------------------------------------------------------------- + + +class TestAliyunGuardrailRegistration: + def test_supported_guardrail_enum_entry(self): + from litellm.types.guardrails import SupportedGuardrailIntegrations + + assert hasattr(SupportedGuardrailIntegrations, "ALIYUN_AI_GUARDRAIL") + assert SupportedGuardrailIntegrations.ALIYUN_AI_GUARDRAIL.value == "aliyun_ai_guardrail" + + def test_initialize_guardrail_function_exists(self): + from litellm.proxy.guardrails.guardrail_hooks.aliyun import ( + guardrail_initializer_registry, + initialize_guardrail, + ) + + assert initialize_guardrail is not None + assert "aliyun_ai_guardrail" in guardrail_initializer_registry + + def test_guardrail_class_registry_exists(self): + from litellm.proxy.guardrails.guardrail_hooks.aliyun import ( + guardrail_class_registry, + ) + + assert "aliyun_ai_guardrail" in guardrail_class_registry + assert guardrail_class_registry["aliyun_ai_guardrail"] is AliyunAIGuardrail + + def test_aliyun_in_global_registry(self): + from litellm.proxy.guardrails.guardrail_registry import ( + guardrail_initializer_registry, + ) + + assert "aliyun_ai_guardrail" in guardrail_initializer_registry + + def test_initialize_guardrail_creates_instance(self): + from litellm.proxy.guardrails.guardrail_hooks.aliyun import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="aliyun_ai_guardrail", + mode="pre_call", + level="medium", + access_key_id=FAKE_AK, + access_key_secret=FAKE_SK, + ) + guardrail_config = {"guardrail_name": "test-aliyun-guard"} + + with patch("litellm.logging_callback_manager.add_litellm_callback") as mock_add: + result = initialize_guardrail(litellm_params, guardrail_config) + + assert isinstance(result, AliyunAIGuardrail) + assert result.guardrail_name == "test-aliyun-guard" + assert result.level == "medium" + assert result.access_key_id == FAKE_AK + assert result.access_key_secret == FAKE_SK + mock_add.assert_called_once_with(result) + + def test_initialize_guardrail_resolves_os_environ_reference(self): + from litellm.proxy.guardrails.guardrail_hooks.aliyun import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="aliyun_ai_guardrail", + mode="pre_call", + access_key_id="os.environ/GUARD_ACCESS_KEY_ID", + access_key_secret="os.environ/GUARD_ACCESS_KEY_SECRET", + ) + guardrail_config = {"guardrail_name": "test-aliyun-guard"} + + with ( + patch.dict( + "os.environ", + { + "GUARD_ACCESS_KEY_ID": FAKE_AK, + "GUARD_ACCESS_KEY_SECRET": FAKE_SK, + }, + ), + patch("litellm.logging_callback_manager.add_litellm_callback"), + ): + result = initialize_guardrail(litellm_params, guardrail_config) + + assert result.access_key_id == FAKE_AK + assert result.access_key_secret == FAKE_SK + + def test_initialize_guardrail_forwards_service_mcp(self): + from litellm.proxy.guardrails.guardrail_hooks.aliyun import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="aliyun_ai_guardrail", + mode="pre_mcp_call", + access_key_id=FAKE_AK, + access_key_secret=FAKE_SK, + service_mcp="text_img_mix_guard", + ) + guardrail_config = {"guardrail_name": "test-aliyun-guard"} + + with patch("litellm.logging_callback_manager.add_litellm_callback"): + result = initialize_guardrail(litellm_params, guardrail_config) + + assert result.service_mcp == "text_img_mix_guard" + + +# --------------------------------------------------------------------------- +# Constructor tests +# --------------------------------------------------------------------------- + + +class TestAliyunGuardrailConstructor: + def test_init_with_explicit_credentials(self): + g = _make_guardrail() + assert g.access_key_id == FAKE_AK + assert g.access_key_secret == FAKE_SK + assert g.level == "medium" + assert g.region_id == "cn-shanghai" + + def test_init_credentials_from_config(self): + g = AliyunAIGuardrail( + guardrail_name="config-test", + access_key_id="cfg-ak", + access_key_secret="cfg-sk", + level="low", + ) + assert g.access_key_id == "cfg-ak" + assert g.access_key_secret == "cfg-sk" + # region defaults to cn-shanghai when not provided via config + assert g.region_id == "cn-shanghai" + + def test_init_raises_without_api_key(self): + with pytest.raises(ValueError, match="ak is required"): + AliyunAIGuardrail(guardrail_name="test") + + def test_init_raises_without_api_secret(self): + with pytest.raises(ValueError, match="sk is required"): + AliyunAIGuardrail(guardrail_name="test", access_key_id=FAKE_AK) + + def test_init_invalid_level_raises(self): + with pytest.raises(ValueError, match="Invalid level"): + _make_guardrail(level="invalid") + + def test_init_default_service_codes_domestic(self): + g = _make_guardrail() + assert g.service_input == "query_security_check_pro" + assert g.service_output == "response_security_check_pro" + + def test_init_service_codes_are_configurable(self): + g = _make_guardrail( + service_input="text_img_mix_guard", + service_output="response_security_check_cb", + ) + assert g.service_input == "text_img_mix_guard" + assert g.service_output == "response_security_check_cb" + + def test_init_region_is_configurable(self): + g = _make_guardrail(region_id="eu-central-1") + assert g.region_id == "eu-central-1" + + +# --------------------------------------------------------------------------- +# Helper function tests +# --------------------------------------------------------------------------- + + +class TestLevelToInt: + def test_standard_levels(self): + assert level_to_int("none") == 0 + assert level_to_int("low") == 1 + assert level_to_int("medium") == 2 + assert level_to_int("high") == 3 + + def test_sensitive_data_levels(self): + assert level_to_int("S0") == 0 + assert level_to_int("S1") == 1 + assert level_to_int("S2") == 2 + assert level_to_int("S3") == 3 + assert level_to_int("S4") == 3 + + def test_case_insensitive(self): + assert level_to_int("HIGH") == 3 + assert level_to_int("Low") == 1 + + def test_empty_string_defaults_to_zero(self): + assert level_to_int("") == 0 + assert level_to_int(None) == 0 + + def test_unknown_level_defaults_to_zero(self): + assert level_to_int("unknown") == 0 + + +class TestSplitText: + def test_short_text_returns_single_segment(self): + g = _make_guardrail() + result = g._split_text("short text", max_length=100) + assert result == ["short text"] + + def test_long_text_splits_at_sentence_boundary(self): + g = _make_guardrail() + text = "Hello world. This is a test. Final sentence." + result = g._split_text(text, max_length=20) + assert len(result) >= 2 + assert "".join(result) == text + + def test_long_text_without_boundary_splits_at_max_length(self): + g = _make_guardrail() + text = "a" * 100 + result = g._split_text(text, max_length=30) + assert len(result) >= 3 + assert "".join(result) == text + + def test_empty_text_returns_empty_list(self): + g = _make_guardrail() + result = g._split_text("", max_length=100) + assert result == [] + + +# --------------------------------------------------------------------------- +# Blocking logic tests +# --------------------------------------------------------------------------- + + +class TestShouldBlockByLevel: + def test_low_protection_blocks_all_risks(self): + g = _make_guardrail(level="low") + assert g._should_block_by_level("low") is True + assert g._should_block_by_level("medium") is True + assert g._should_block_by_level("high") is True + assert g._should_block_by_level("none") is False + + def test_medium_protection_blocks_medium_and_high(self): + g = _make_guardrail(level="medium") + assert g._should_block_by_level("low") is False + assert g._should_block_by_level("medium") is True + assert g._should_block_by_level("high") is True + + def test_high_protection_blocks_high_only(self): + g = _make_guardrail(level="high") + assert g._should_block_by_level("medium") is False + assert g._should_block_by_level("high") is True + + def test_max_observation_never_blocks(self): + g = _make_guardrail(level="max") + assert g._should_block_by_level("high") is False + assert g._should_block_by_level("S4") is False + + +class TestParseResponseAndCheck: + def test_blocks_when_level_meets_threshold(self): + g = _make_guardrail(level="medium") + response = { + "Data": { + "Suggestion": "block", + "Detail": [_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + }, + } + with pytest.raises(HTTPException) as exc_info: + g._parse_response_and_check(response, check_type="input") + assert exc_info.value.status_code == 400 + # Message uses the raw detection type returned by Aliyun + assert CONTENT_MODERATION_TYPE in str(exc_info.value.detail) + + def test_passes_when_level_below_threshold(self): + g = _make_guardrail(level="high") + response = { + "Data": { + "Suggestion": "pass", + "Detail": [_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="low")], + }, + } + result = g._parse_response_and_check(response, check_type="input") + assert result["flagged"] is False + + def test_empty_data_returns_pass(self): + g = _make_guardrail() + response = {"Data": {}} + result = g._parse_response_and_check(response, check_type="input") + assert result["flagged"] is False + assert result["suggestion"] == "pass" + + def test_extracts_desensitization_for_sensitive_data(self): + # level=max never blocks, so the parsed result (with desensitization) is returned + g = _make_guardrail(level="max") + response = { + "Data": { + "Suggestion": "mask", + "Detail": [ + { + "Type": SENSITIVE_DATA_TYPE, + "Level": "S2", + "Suggestion": "mask", + "Result": [{"Ext": {"Desensitization": "masked_text"}}], + }, + ], + }, + } + result = g._parse_response_and_check(response, check_type="input") + assert result["desensitization"] == "masked_text" + + def test_prompt_attack_block_message(self): + g = _make_guardrail(level="medium") + response = { + "Data": { + "Suggestion": "block", + "Detail": [_make_detail(detection_type=PROMPT_ATTACK_TYPE, level="medium")], + }, + } + with pytest.raises(HTTPException) as exc_info: + g._parse_response_and_check(response, check_type="input") + assert PROMPT_ATTACK_TYPE in str(exc_info.value.detail) + + +class TestRiskLevelResolution: + """MultiModalGuard reports severity in two shapes: Detail[].Level (returned by + the _pro service codes) and Detail[].Result[].RiskLevel (the documented shape). + Reading only the former silently downgrades the latter to "none" and lets + Aliyun's own block decision through. + """ + + def test_falls_back_to_result_risk_level_when_level_absent(self): + g = _make_guardrail(level="medium") + response = { + "Data": { + "Suggestion": "block", + "Detail": [ + _make_detail( + detection_type=CONTENT_MODERATION_TYPE, + level=None, + results=[{"Label": "violence", "Confidence": 99.5, "RiskLevel": "high"}], + ) + ], + }, + } + with pytest.raises(HTTPException) as exc_info: + g._parse_response_and_check(response, check_type="input") + assert CONTENT_MODERATION_TYPE in str(exc_info.value.detail) + + def test_result_risk_level_still_respects_threshold(self): + g = _make_guardrail(level="high") + response = { + "Data": { + "Suggestion": "pass", + "Detail": [ + _make_detail( + detection_type=CONTENT_MODERATION_TYPE, + level=None, + suggestion="pass", + results=[{"Label": "spam", "RiskLevel": "low"}], + ) + ], + }, + } + result = g._parse_response_and_check(response, check_type="input") + assert result["flagged"] is False + + def test_highest_result_risk_level_wins(self): + g = _make_guardrail(level="high") + response = { + "Data": { + "Suggestion": "block", + "Detail": [ + _make_detail( + detection_type=CONTENT_MODERATION_TYPE, + level=None, + results=[ + {"Label": "spam", "RiskLevel": "low"}, + {"Label": "violence", "RiskLevel": "high"}, + ], + ) + ], + }, + } + with pytest.raises(HTTPException): + g._parse_response_and_check(response, check_type="input") + + def test_detail_level_takes_precedence_over_result_risk_level(self): + g = _make_guardrail(level="medium") + response = { + "Data": { + "Suggestion": "pass", + "Detail": [ + _make_detail( + detection_type=CONTENT_MODERATION_TYPE, + level="none", + suggestion="pass", + results=[{"Label": "spam", "RiskLevel": "low"}], + ) + ], + }, + } + result = g._parse_response_and_check(response, check_type="input") + assert result["flagged"] is False + + def test_blocks_on_detail_block_suggestion_without_any_risk_level(self): + g = _make_guardrail(level="medium") + response = { + "Data": { + "Suggestion": "pass", + "Detail": [_make_detail(detection_type=PROMPT_ATTACK_TYPE, level=None, suggestion="block")], + }, + } + with pytest.raises(HTTPException) as exc_info: + g._parse_response_and_check(response, check_type="input") + assert PROMPT_ATTACK_TYPE in str(exc_info.value.detail) + + def test_blocks_on_overall_block_suggestion_without_any_risk_level(self): + g = _make_guardrail(level="medium") + response = { + "Data": { + "Suggestion": "block", + "Detail": [_make_detail(detection_type=CONTENT_MODERATION_TYPE, level=None, suggestion="pass")], + }, + } + with pytest.raises(HTTPException): + g._parse_response_and_check(response, check_type="input") + + def test_blocks_on_overall_block_suggestion_with_empty_detail(self): + g = _make_guardrail(level="medium") + response = {"Data": {"Suggestion": "block", "Detail": []}} + with pytest.raises(HTTPException): + g._parse_response_and_check(response, check_type="input") + + def test_observation_mode_never_blocks_on_block_suggestion(self): + """level=max is an explicit opt-in to logging only; fail-closed handling of an + unparseable severity must still honour it rather than bypass the threshold.""" + g = _make_guardrail(level="max") + response = { + "Data": { + "Suggestion": "block", + "Detail": [_make_detail(detection_type=CONTENT_MODERATION_TYPE, level=None, suggestion="block")], + }, + } + result = g._parse_response_and_check(response, check_type="input") + assert result["flagged"] is True + + def test_documented_response_shape_blocks(self): + """Regression guard for the exact payload in the integration README, which + carries no Detail[].Level at all.""" + g = _make_guardrail(level="medium") + response = { + "RequestId": "xxx", + "Code": 200, + "Data": { + "Suggestion": "block", + "Detail": [ + { + "Type": "contentModeration", + "Suggestion": "block", + "Result": [{"Label": "violence", "Confidence": 99.5, "RiskLevel": "high"}], + } + ], + }, + } + with pytest.raises(HTTPException) as exc_info: + g._parse_response_and_check(response, check_type="input") + assert exc_info.value.status_code == 400 + assert "contentModeration" in str(exc_info.value.detail) + + +# --------------------------------------------------------------------------- +# Config model tests +# --------------------------------------------------------------------------- + + +class TestConfigModel: + def test_get_config_model_returns_correct_type(self): + from litellm.types.proxy.guardrails.guardrail_hooks.aliyun.aliyun_ai_guardrail import ( + AliyunAIGuardrailConfigModel, + ) + + assert AliyunAIGuardrail.get_config_model() is AliyunAIGuardrailConfigModel + + def test_config_model_ui_friendly_name(self): + from litellm.types.proxy.guardrails.guardrail_hooks.aliyun.aliyun_ai_guardrail import ( + AliyunAIGuardrailConfigModel, + ) + + assert AliyunAIGuardrailConfigModel.ui_friendly_name() == "Aliyun AI Security Guardrail" + + def test_litellm_params_declares_all_service_codes(self): + from litellm.types.guardrails import LitellmParams + + for field in ("service_input", "service_output", "service_mcp"): + assert field in LitellmParams.model_fields + + +# --------------------------------------------------------------------------- +# Image URL extraction tests +# --------------------------------------------------------------------------- + + +class TestGetImageUrls: + def test_extracts_http_and_https_urls(self): + g = _make_guardrail() + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is in these?"}, + {"type": "image_url", "image_url": {"url": IMG_A}}, + {"type": "image_url", "image_url": {"url": IMG_B}}, + ], + } + ] + assert g.get_image_urls(messages) == [IMG_A, IMG_B] + + def test_skips_non_url_images(self): + g = _make_guardrail() + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + {"type": "image_url", "image_url": {"url": DATA_URI}}, + {"type": "image_url", "image_url": {"url": IMG_A}}, + ], + } + ] + assert g.get_image_urls(messages) == [IMG_A] + + def test_plain_text_returns_empty(self): + g = _make_guardrail() + messages = [{"role": "user", "content": "just text"}] + assert g.get_image_urls(messages) == [] + + def test_deduplicates_across_messages(self): + g = _make_guardrail() + messages = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": IMG_A}}]}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": IMG_A}}]}, + ] + assert g.get_image_urls(messages) == [IMG_A] + + def test_only_last_consecutive_user_block(self): + g = _make_guardrail() + messages = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": IMG_B}}]}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": IMG_A}}]}, + ] + assert g.get_image_urls(messages) == [IMG_A] + + def test_empty_messages_returns_empty(self): + g = _make_guardrail() + assert g.get_image_urls([]) == [] + + +# --------------------------------------------------------------------------- +# ServiceParameters construction tests +# --------------------------------------------------------------------------- + + +class TestServiceParametersConstruction: + @pytest.mark.asyncio + async def test_text_only(self): + g = _make_guardrail(service_input="query_security_check") + with patch.object( + g.async_handler, "post", new_callable=AsyncMock, return_value=_make_aliyun_api_response() + ) as mock_post: + await g.async_make_request(text="hello", service_type="input") + sp, service = _captured_service_parameters(mock_post) + assert sp == {"requestFrom": "LiteLLM", "content": "hello"} + assert service == "query_security_check" + + @pytest.mark.asyncio + async def test_text_and_images(self): + g = _make_guardrail(service_input="text_img_mix_guard") + with patch.object( + g.async_handler, "post", new_callable=AsyncMock, return_value=_make_aliyun_api_response() + ) as mock_post: + await g.async_make_request(text="hello", service_type="input", image_urls=[IMG_A, IMG_B]) + sp, service = _captured_service_parameters(mock_post) + assert sp == {"requestFrom": "LiteLLM", "content": "hello", "imageUrls": [IMG_A, IMG_B]} + assert service == "text_img_mix_guard" + + @pytest.mark.asyncio + async def test_images_only(self): + g = _make_guardrail(service_input="img_query_security_check") + with patch.object( + g.async_handler, "post", new_callable=AsyncMock, return_value=_make_aliyun_api_response() + ) as mock_post: + await g.async_make_request(service_type="input", image_urls=[IMG_A]) + sp, service = _captured_service_parameters(mock_post) + assert sp == {"requestFrom": "LiteLLM", "imageUrls": [IMG_A]} + assert service == "img_query_security_check" + + +# --------------------------------------------------------------------------- +# Pre-call hook tests (text + multimodal) +# --------------------------------------------------------------------------- + + +class TestPreCallHook: + @pytest.mark.asyncio + async def test_blocks_violation(self): + g = _make_guardrail(level="medium") + mock_api_response = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response): + with pytest.raises(HTTPException) as exc_info: + await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": [{"role": "user", "content": "违规内容"}]}, + call_type="completion", + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_passes_clean_content(self): + g = _make_guardrail(level="medium") + mock_api_response = _make_aliyun_api_response(suggestion="pass", detail=[]) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response): + result = await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": [{"role": "user", "content": "你好"}]}, + call_type="completion", + ) + assert result is None + + @pytest.mark.asyncio + async def test_no_messages_returns_data(self): + g = _make_guardrail() + result = await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={}, + call_type="completion", + ) + assert result == {} + + @pytest.mark.asyncio + async def test_no_user_prompt_returns_none(self): + g = _make_guardrail() + result = await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": [{"role": "system", "content": "system prompt only"}]}, + call_type="completion", + ) + assert result is None + + @pytest.mark.asyncio + async def test_long_text_splits_into_multiple_requests(self): + g = _make_guardrail(max_text_length=10, level="medium") + mock_api_response = _make_aliyun_api_response(suggestion="pass", detail=[]) + long_content = "This is a very long text that exceeds the max_text_length limit." + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response) as mock_post: + await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": [{"role": "user", "content": long_content}]}, + call_type="completion", + ) + assert mock_post.call_count >= 2 + + @pytest.mark.asyncio + async def test_blocks_image_violation(self): + g = _make_guardrail(level="medium", service_input="text_img_mix_guard") + mock_resp = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": {"url": IMG_A}}, + ], + } + ] + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_resp): + with pytest.raises(HTTPException) as exc_info: + await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": messages}, + call_type="completion", + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_image_only_request_sends_imageurls(self): + g = _make_guardrail(level="medium", service_input="img_query_security_check") + mock_resp = _make_aliyun_api_response(suggestion="pass", detail=[]) + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": IMG_A}}]}] + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_resp) as mock_post: + await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": messages}, + call_type="completion", + ) + assert mock_post.call_count == 1 + sp, _ = _captured_service_parameters(mock_post) + assert sp == {"requestFrom": "LiteLLM", "imageUrls": [IMG_A]} + + @pytest.mark.asyncio + async def test_first_segment_carries_images(self): + g = _make_guardrail(max_text_length=10, level="medium", service_input="text_img_mix_guard") + mock_resp = _make_aliyun_api_response(suggestion="pass", detail=[]) + long_text = "This is a very long text exceeding the limit for sure." + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": long_text}, + {"type": "image_url", "image_url": {"url": IMG_A}}, + ], + } + ] + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_resp) as mock_post: + await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": messages}, + call_type="completion", + ) + assert mock_post.call_count >= 2 + image_carrying = 0 + for call in mock_post.call_args_list: + sp = json.loads(call.kwargs["data"]["ServiceParameters"]) + if "imageUrls" in sp: + image_carrying += 1 + assert image_carrying == 1 + + +# --------------------------------------------------------------------------- +# Post-call hook tests +# --------------------------------------------------------------------------- + + +class TestPostCallHook: + @pytest.mark.asyncio + async def test_blocks_violation_in_response(self): + import litellm + + g = _make_guardrail(level="medium") + mock_api_response = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message(role="assistant", content="违规响应内容"), + ) + ], + ) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response): + with pytest.raises(HTTPException) as exc_info: + await g.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hi"}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=response, + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_passes_clean_response(self): + import litellm + + g = _make_guardrail(level="medium") + mock_api_response = _make_aliyun_api_response(suggestion="pass", detail=[]) + response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message(role="assistant", content="正常的回复内容"), + ) + ], + ) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response): + result = await g.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hi"}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=response, + ) + assert result is response + + @pytest.mark.asyncio + async def test_blocks_violation_in_later_choice(self): + import litellm + + g = _make_guardrail(level="medium") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + blocked = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + + async def block_only_violating_text(*args, **kwargs): + scanned = json.loads(kwargs["data"]["ServiceParameters"])["content"] + return blocked if "违规响应内容" in scanned else clean + + response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message(role="assistant", content="正常的回复内容"), + ), + litellm.Choices( + index=1, + message=litellm.Message(role="assistant", content="违规响应内容"), + ), + ], + ) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, side_effect=block_only_violating_text): + with pytest.raises(HTTPException) as exc_info: + await g.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hi"}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=response, + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_scans_content_from_every_choice(self): + import litellm + + g = _make_guardrail(level="medium", max_text_length=10000) + mock_api_response = _make_aliyun_api_response(suggestion="pass", detail=[]) + response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices(index=0, message=litellm.Message(role="assistant", content="第一个回复")), + litellm.Choices(index=1, message=litellm.Message(role="assistant", content="第二个回复")), + ], + ) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response) as mock_post: + await g.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hi"}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=response, + ) + + scanned = "".join( + json.loads(call.kwargs["data"]["ServiceParameters"])["content"] for call in mock_post.call_args_list + ) + assert "第一个回复" in scanned + assert "第二个回复" in scanned + + @pytest.mark.asyncio + async def test_non_model_response_passthrough(self): + g = _make_guardrail() + result = await g.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response="not a model response", + ) + assert result == "not a model response" + + @pytest.mark.asyncio + async def test_empty_content_response_passthrough(self): + import litellm + + g = _make_guardrail() + response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message(role="assistant", content=""), + ) + ], + ) + result = await g.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=response, + ) + assert result is response + + +# --------------------------------------------------------------------------- +# Post-MCP hook tests +# --------------------------------------------------------------------------- + + +def _make_call_tool_result(text: str = "tool output"): + from mcp.types import CallToolResult, TextContent + + return CallToolResult(content=[TextContent(type="text", text=text)], isError=False) + + +# --------------------------------------------------------------------------- +# Streaming hook tests +# --------------------------------------------------------------------------- + + +def _make_stream_chunk(content=None, tool_call_arguments=None): + """Build a ModelResponseStream chunk carrying content and/or tool call arguments.""" + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + tool_calls = None + if tool_call_arguments is not None: + tool_calls = [ + ChatCompletionDeltaToolCall( + id="call_1", + type="function", + index=0, + function=Function(name="send_message", arguments=tool_call_arguments), + ) + ] + return ModelResponseStream( + id="test-id", + choices=[StreamingChoices(index=0, delta=Delta(content=content, tool_calls=tool_calls))], + ) + + +async def _aiter(chunks): + for chunk in chunks: + yield chunk + + +# --------------------------------------------------------------------------- +# MCP pre-call hook tests +# --------------------------------------------------------------------------- + + +class TestMcpPreCallCheck: + @pytest.mark.asyncio + async def test_blocks_violating_tool_arguments(self): + g = _make_guardrail(level="medium") + blocked = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=blocked): + with pytest.raises(HTTPException) as exc_info: + await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": [{"role": "user", "content": "delete_all_files 违规参数"}]}, + call_type="call_mcp_tool", + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_uses_mcp_service_code(self): + g = _make_guardrail(level="medium", service_mcp="text_img_mix_guard") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) as mock_post: + result = await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": [{"role": "user", "content": "send_message hello"}]}, + call_type="call_mcp_tool", + ) + + assert result is None + service_parameters, service = _captured_service_parameters(mock_post) + assert service == "text_img_mix_guard" + assert service_parameters["content"] == "send_message hello" + + @pytest.mark.asyncio + async def test_empty_content_skips_check(self): + g = _make_guardrail() + with patch.object(g.async_handler, "post", new_callable=AsyncMock) as mock_post: + result = await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": [{"role": "user", "content": ""}]}, + call_type="call_mcp_tool", + ) + + assert result is None + mock_post.assert_not_called() + + @pytest.mark.asyncio + async def test_long_content_is_split_into_multiple_requests(self): + g = _make_guardrail(max_text_length=10) + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) as mock_post: + await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": [{"role": "user", "content": "一二三四五六七八九十。壹贰叁肆伍陆柒捌玖抾。"}]}, + call_type="call_mcp_tool", + ) + + assert mock_post.call_count > 1 + + +# --------------------------------------------------------------------------- +# post_mcp_call event hook gating tests +# --------------------------------------------------------------------------- + + +class TestShouldRunPostMcpCall: + def test_no_event_hook_runs(self): + g = _make_guardrail() + g.event_hook = None + assert g._should_run_post_mcp_call() is True + + def test_list_containing_post_mcp_call_runs(self): + g = _make_guardrail() + g.event_hook = ["pre_call", "post_mcp_call"] + assert g._should_run_post_mcp_call() is True + + def test_list_without_post_mcp_call_skips(self): + g = _make_guardrail() + g.event_hook = ["pre_call", "post_call"] + assert g._should_run_post_mcp_call() is False + + def test_plain_string_is_matched(self): + g = _make_guardrail() + g.event_hook = "post_mcp_call" + assert g._should_run_post_mcp_call() is True + g.event_hook = "pre_call" + assert g._should_run_post_mcp_call() is False + + def test_mode_tags_are_matched(self): + from litellm.types.guardrails import Mode + + g = _make_guardrail() + g.event_hook = Mode(tags={"team-a": ["pre_call", "post_mcp_call"]}) + assert g._should_run_post_mcp_call() is True + + g.event_hook = Mode(tags={"team-a": "post_mcp_call"}) + assert g._should_run_post_mcp_call() is True + + def test_mode_falls_back_to_default(self): + from litellm.types.guardrails import Mode + + g = _make_guardrail() + g.event_hook = Mode(tags={"team-a": "pre_call"}, default=["post_mcp_call"]) + assert g._should_run_post_mcp_call() is True + + g.event_hook = Mode(tags={"team-a": "pre_call"}, default="pre_call") + assert g._should_run_post_mcp_call() is False + + def test_mode_without_match_or_default_skips(self): + from litellm.types.guardrails import Mode + + g = _make_guardrail() + g.event_hook = Mode(tags={"team-a": "pre_call"}) + assert g._should_run_post_mcp_call() is False + + +# --------------------------------------------------------------------------- +# MCP payload shape handling tests +# --------------------------------------------------------------------------- + + +class TestIterMcpContentItems: + def test_plain_string_is_wrapped(self): + g = _make_guardrail() + assert g._iter_mcp_content_items("hello") == ["hello"] + + def test_object_with_content_list(self): + g = _make_guardrail() + payload = MagicMock() + payload.content = ["a", "b"] + assert g._iter_mcp_content_items(payload) == ["a", "b"] + + def test_dict_with_content_list(self): + g = _make_guardrail() + assert g._iter_mcp_content_items({"content": ["a"]}) == ["a"] + + def test_dict_without_content_returns_itself(self): + g = _make_guardrail() + payload = {"text": "no content key"} + assert g._iter_mcp_content_items(payload) == [payload] + + def test_coerced_tuple_pairs_recover_real_content(self): + g = _make_guardrail() + payload = [("meta", None), ("content", ["real"]), ("isError", False)] + assert g._iter_mcp_content_items(payload) == ["real"] + + def test_plain_list_passes_through(self): + g = _make_guardrail() + assert g._iter_mcp_content_items(["a", "b"]) == ["a", "b"] + + def test_unsupported_payload_returns_empty(self): + g = _make_guardrail() + assert g._iter_mcp_content_items(123) == [] + + +class TestReplaceToolOutputInPlace: + def test_none_target_is_rejected(self): + g = _make_guardrail() + assert g._replace_tool_output_in_place(None, ["blocked"]) is False + + def test_object_content_is_overwritten_and_flagged(self): + g = _make_guardrail() + target = MagicMock() + target.content = ["original"] + target.isError = False + assert g._replace_tool_output_in_place(target, ["blocked"]) is True + assert target.content == ["blocked"] + assert target.isError is True + + def test_list_target_is_overwritten(self): + g = _make_guardrail() + target = ["original"] + assert g._replace_tool_output_in_place(target, ["blocked"]) is True + assert target == ["blocked"] + + def test_nested_result_content_is_overwritten(self): + g = _make_guardrail() + target = {"result": {"content": ["original"]}} + assert g._replace_tool_output_in_place(target, ["blocked"]) is True + assert target["result"]["content"] == ["blocked"] + + def test_dict_content_is_overwritten(self): + g = _make_guardrail() + target = {"content": ["original"]} + assert g._replace_tool_output_in_place(target, ["blocked"]) is True + assert target["content"] == ["blocked"] + + def test_unsupported_shape_is_rejected(self): + g = _make_guardrail() + assert g._replace_tool_output_in_place({"unrelated": 1}, ["blocked"]) is False + + +class TestStreamingHook: + @pytest.mark.asyncio + async def test_scans_tool_call_arguments(self): + g = _make_guardrail(level="medium", stream_first_check_step=1, stream_slide_step=1) + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + chunk = _make_stream_chunk(tool_call_arguments='{"text": "违规工具参数"}') + + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) as mock_post: + async for _ in g.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=_aiter([chunk]), + request_data={}, + ): + pass + + scanned = "".join( + json.loads(call.kwargs["data"]["ServiceParameters"])["content"] for call in mock_post.call_args_list + ) + assert "违规工具参数" in scanned + + @pytest.mark.asyncio + async def test_blocks_violating_tool_call_arguments(self): + g = _make_guardrail(level="medium", stream_first_check_step=1, stream_slide_step=1) + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + blocked = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + + async def block_only_violating_text(*args, **kwargs): + scanned = json.loads(kwargs["data"]["ServiceParameters"])["content"] + return blocked if "违规工具参数" in scanned else clean + + chunk = _make_stream_chunk(tool_call_arguments='{"text": "违规工具参数"}') + + with patch.object(g.async_handler, "post", new_callable=AsyncMock, side_effect=block_only_violating_text): + emitted = [ + chunk + async for chunk in g.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=_aiter([chunk]), + request_data={}, + ) + ] + + # Streaming blocks by emitting an SSE error event instead of raising + assert not any(getattr(item, "choices", None) for item in emitted if hasattr(item, "choices")) + + @pytest.mark.asyncio + async def test_scans_plain_content(self): + g = _make_guardrail(level="medium", stream_first_check_step=1, stream_slide_step=1) + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + chunk = _make_stream_chunk(content="普通流式内容") + + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) as mock_post: + emitted = [ + item + async for item in g.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=_aiter([chunk]), + request_data={}, + ) + ] + + scanned = "".join( + json.loads(call.kwargs["data"]["ServiceParameters"])["content"] for call in mock_post.call_args_list + ) + assert "普通流式内容" in scanned + assert emitted == [chunk] + + @pytest.mark.asyncio + async def test_scans_responses_api_text_delta(self): + from litellm.types.llms.openai import OutputTextDeltaEvent + + g = _make_guardrail(level="medium", stream_first_check_step=1, stream_slide_step=1) + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + chunk = OutputTextDeltaEvent( + type="response.output_text.delta", + item_id="item_1", + output_index=0, + content_index=0, + delta="违规的 responses 输出", + ) + + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) as mock_post: + async for _ in g.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=_aiter([chunk]), + request_data={}, + ): + pass + + scanned = "".join( + json.loads(call.kwargs["data"]["ServiceParameters"])["content"] for call in mock_post.call_args_list + ) + assert "违规的 responses 输出" in scanned + + @pytest.mark.asyncio + async def test_blocks_violating_responses_api_text_delta(self): + from litellm.types.llms.openai import OutputTextDeltaEvent + + g = _make_guardrail(level="medium", stream_first_check_step=1, stream_slide_step=1) + blocked = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + chunk = OutputTextDeltaEvent( + type="response.output_text.delta", + item_id="item_1", + output_index=0, + content_index=0, + delta="违规的 responses 输出", + ) + + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=blocked): + emitted = [ + item + async for item in g.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=_aiter([chunk]), + request_data={}, + ) + ] + + assert chunk not in emitted + + @pytest.mark.asyncio + async def test_scans_responses_api_completed_event(self): + from litellm.types.llms.openai import ResponseCompletedEvent + + g = _make_guardrail(level="medium", stream_first_check_step=1, stream_slide_step=1) + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + content_part = MagicMock() + content_part.text = "完成事件里的违规内容" + output_item = MagicMock() + output_item.content = [content_part] + chunk = MagicMock(spec=ResponseCompletedEvent) + chunk.type = "response.completed" + chunk.response = MagicMock() + chunk.response.output = [output_item] + + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) as mock_post: + async for _ in g.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=_aiter([chunk]), + request_data={}, + ): + pass + + scanned = "".join( + json.loads(call.kwargs["data"]["ServiceParameters"])["content"] for call in mock_post.call_args_list + ) + assert "完成事件里的违规内容" in scanned + + +def _make_post_mcp_hook_args(tool_result): + """Mimic how litellm_logging dispatches the post-MCP hook. + + The live CallToolResult is stored in model_call_details["original_response"] and + the same object is wrapped into MCPPostCallResponseObject, whose + mcp_tool_call_response field is declared as a list - so pydantic coerces the + CallToolResult by iterating it into (key, value) tuples. + """ + from litellm.types.llms.base import HiddenParams + from litellm.types.mcp import MCPPostCallResponseObject + + kwargs = {"original_response": tool_result} + response_obj = MCPPostCallResponseObject(mcp_tool_call_response=tool_result, hidden_params=HiddenParams()) + return kwargs, response_obj + + +class TestPostMcpToolCallHook: + @pytest.mark.asyncio + async def test_audits_tool_text_not_pydantic_tuple_repr(self): + """The wrapped response degrades into (key, value) tuples, so auditing it + verbatim would send Python reprs - and MCP envelope fields - to Aliyun + instead of the tool's own output.""" + g = _make_guardrail(level="medium") + tool_result = _make_call_tool_result("工具返回的内容") + kwargs, response_obj = _make_post_mcp_hook_args(tool_result) + mock_api_response = _make_aliyun_api_response(suggestion="pass", detail=[]) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response) as mock_post: + await g.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + sp, service = _captured_service_parameters(mock_post) + assert sp["content"] == "工具返回的内容" + for envelope_field in ("isError", "structuredContent", "annotations", "TextContent"): + assert envelope_field not in sp["content"] + + @pytest.mark.asyncio + async def test_violation_replaces_tool_output_in_place(self): + """Both dispatch sites discard this hook's return value and hand the original + CallToolResult to the client, so the violation must be written into it.""" + g = _make_guardrail(level="medium") + tool_result = _make_call_tool_result("违规的工具输出") + kwargs, response_obj = _make_post_mcp_hook_args(tool_result) + mock_api_response = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response): + await g.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + remaining = " ".join(getattr(item, "text", "") for item in tool_result.content) + assert "违规的工具输出" not in remaining + assert CONTENT_MODERATION_TYPE in remaining + assert tool_result.isError is True + + @pytest.mark.asyncio + async def test_violation_returns_replacement_object_without_raising(self): + """Raising is a no-op here: the dispatcher catches every callback exception as a + non-blocking logging error and returns the untouched tool result.""" + from litellm.types.mcp import MCPPostCallResponseObject + + g = _make_guardrail(level="medium") + tool_result = _make_call_tool_result("违规的工具输出") + kwargs, response_obj = _make_post_mcp_hook_args(tool_result) + mock_api_response = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=PROMPT_ATTACK_TYPE, level="high")], + ) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response): + result = await g.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert isinstance(result, MCPPostCallResponseObject) + replaced = " ".join(getattr(item, "text", "") for item in result.mcp_tool_call_response) + assert PROMPT_ATTACK_TYPE in replaced + + @pytest.mark.asyncio + async def test_clean_output_leaves_tool_result_untouched(self): + g = _make_guardrail(level="medium") + tool_result = _make_call_tool_result("正常的工具输出") + kwargs, response_obj = _make_post_mcp_hook_args(tool_result) + mock_api_response = _make_aliyun_api_response(suggestion="pass", detail=[]) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=mock_api_response): + result = await g.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert result is None + assert tool_result.content[0].text == "正常的工具输出" + assert tool_result.isError is False + + @pytest.mark.asyncio + async def test_api_network_failure_fails_closed_in_place(self): + """A network error is swallowed by the dispatcher as a non-blocking logging + error, so unaudited tool output would reach the client untouched. The hook + must fail closed by replacing the live tool result instead of raising.""" + import httpx + + g = _make_guardrail(level="medium") + tool_result = _make_call_tool_result("未拦截将泄漏的敏感内容") + kwargs, response_obj = _make_post_mcp_hook_args(tool_result) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, side_effect=httpx.ConnectError("boom")): + await g.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + remaining = " ".join(getattr(item, "text", "") for item in tool_result.content) + assert "未拦截将泄漏的敏感内容" not in remaining + assert "未经审核" in remaining + assert tool_result.isError is True + + @pytest.mark.asyncio + async def test_api_network_failure_returns_replacement_object(self): + """The failure notice must also be returned as the documented replacement + object, mirroring the violation path.""" + import httpx + + from litellm.types.mcp import MCPPostCallResponseObject + + g = _make_guardrail(level="medium") + tool_result = _make_call_tool_result("未拦截将泄漏的敏感内容") + kwargs, response_obj = _make_post_mcp_hook_args(tool_result) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, side_effect=httpx.ConnectError("boom")): + result = await g.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert isinstance(result, MCPPostCallResponseObject) + replaced = " ".join(getattr(item, "text", "") for item in result.mcp_tool_call_response) + assert "未经审核" in replaced + + @pytest.mark.asyncio + async def test_skipped_when_post_mcp_call_not_configured(self): + g = _make_guardrail(level="medium", event_hook=["pre_call"]) + tool_result = _make_call_tool_result("违规的工具输出") + kwargs, response_obj = _make_post_mcp_hook_args(tool_result) + with patch.object(g.async_handler, "post", new_callable=AsyncMock) as mock_post: + result = await g.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert result is None + mock_post.assert_not_called() + assert tool_result.content[0].text == "违规的工具输出"