From e8c5d51aeafb1b8c8aa92c4bb7c855305d26abc5 Mon Sep 17 00:00:00 2001 From: splendor023 Date: Thu, 13 Aug 2026 16:43:09 +0800 Subject: [PATCH 1/5] 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 == "违规的工具输出" From b7291c3ea4818f3c5a2be8dc2e0294cd035acbc6 Mon Sep 17 00:00:00 2001 From: splendor023 Date: Fri, 14 Aug 2026 14:20:52 +0800 Subject: [PATCH 2/5] feat(guardrails): add aliyun security guardrail integration --- .../guardrail_hooks/aliyun/__init__.py | 59 +- .../aliyun/aliyun_ai_guardrail.py | 905 +++++++++++------- .../guardrails/guardrail_hooks/aliyun/base.py | 117 ++- litellm/types/guardrails.py | 1 - .../aliyun/aliyun_ai_guardrail.py | 91 +- .../aliyun/test_aliyun_ai_guardrail.py | 393 +++++++- 6 files changed, 1052 insertions(+), 514 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aliyun/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/aliyun/__init__.py index 7ff7e7d4256..1fa0e880584 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aliyun/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/__init__.py @@ -11,7 +11,7 @@ This module provides integration with Aliyun's AI Security Guardrail service for Documentation: https://help.aliyun.com/document_detail/2873209.html """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Final from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -21,6 +21,19 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams +def _resolve_os_environ_reference(value: str | None) -> str | None: + """Resolve an ``os.environ/`` reference. + + guardrail_registry.py only auto-resolves api_key/api_base, so the Aliyun + credential fields have to be resolved here. + """ + from litellm.secret_managers.main import get_secret_str + + if isinstance(value, str) and value.startswith("os.environ/"): + return get_secret_str(value) + return value + + def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> AliyunAIGuardrail: """ Initialize an Aliyun AI Guardrail instance. @@ -35,33 +48,27 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" AliyunAIGuardrail instance """ import litellm - from litellm.secret_managers.main import get_secret_str - guardrail_name = guardrail.get("guardrail_name") + guardrail_name: Final = 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) + level: Final = getattr(litellm_params, "level", None) + max_text_length: Final = getattr(litellm_params, "max_text_length", None) + stream_window_size: Final = getattr(litellm_params, "stream_window_size", None) + stream_slide_step: Final = getattr(litellm_params, "stream_slide_step", None) + stream_first_check_step: Final = getattr(litellm_params, "stream_first_check_step", None) + region_id: Final = getattr(litellm_params, "region_id", None) + service_input: Final = getattr(litellm_params, "service_input", None) + service_output: Final = getattr(litellm_params, "service_output", None) + service_mcp: Final = 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) + # These custom credential fields are not auto-resolved by guardrail_registry.py + # (only api_key/api_base are), so os.environ/ references are resolved here. + access_key_id: Final = _resolve_os_environ_reference(getattr(litellm_params, "access_key_id", None)) + access_key_secret: Final = _resolve_os_environ_reference(getattr(litellm_params, "access_key_secret", None)) - aliyun_guardrail = AliyunAIGuardrail( + aliyun_guardrail: Final = AliyunAIGuardrail( guardrail_name=guardrail_name, access_key_id=access_key_id, access_key_secret=access_key_secret, @@ -83,12 +90,14 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" return aliyun_guardrail -# Registry for guardrail initializers -guardrail_initializer_registry = { +# Registry for guardrail initializers. +# Plain dicts: guardrail_registry.py gates discovery on `isinstance(registry, dict)`, +# which a MappingProxyType would fail, silently skipping this guardrail's registration. +guardrail_initializer_registry: Final = { # mutable-ok: loader requires a real dict SupportedGuardrailIntegrations.ALIYUN_AI_GUARDRAIL.value: initialize_guardrail, } # Registry for guardrail classes -guardrail_class_registry = { +guardrail_class_registry: Final = { # mutable-ok: loader requires a real dict 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 index dafbcc72e10..27ce33e4ca7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 """ Aliyun AI Security Guardrail Integration for LiteLLM 阿里云AI安全护栏集成 @@ -23,8 +22,14 @@ import hmac import json import re import uuid +from collections.abc import AsyncGenerator, AsyncIterable, Iterable, Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, AsyncGenerator, Literal +from types import MappingProxyType + +# `Any` is needed for the heterogeneous LLM payloads this guardrail walks (request bodies, +# streaming chunks, MCP tool results); every field it actually reads is probed defensively. +# `cast` re-labels the upstream message iterator, whose element type is untyped upstream. +from typing import TYPE_CHECKING, Any, Final, Literal, cast # noqa: TID251 # see comment above from urllib.parse import quote from fastapi import HTTPException @@ -40,9 +45,17 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth +# Only the private iterator walks `messages` *and* the Responses API's `input` while keeping +# multimodal parts intact; the public helpers flatten to text, which would drop image URLs. +from litellm.proxy.guardrails._content_utils import ( + _iter_inspection_messages, # pyright: ignore[reportPrivateUsage] # see comment above +) +from litellm.types.utils import CallTypes + from .base import AliyunGuardrailBase if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import AllMessageValues from litellm.types.mcp import MCPPostCallResponseObject @@ -50,48 +63,68 @@ if TYPE_CHECKING: AliyunAIGuardrailResponse, ) from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel - from litellm.types.utils import EmbeddingResponse, ImageResponse, ModelResponse + from litellm.types.utils import CallTypesLiteral, LLMResponseTypes # Constants -ENCODING = "UTF-8" -ISO8601_DATE_FORMAT = "%Y-%m-%dT%H:%M:%SZ" -ALGORITHM = "HmacSHA1" +ENCODING: Final = "UTF-8" +ISO8601_DATE_FORMAT: Final = "%Y-%m-%dT%H:%M:%SZ" +ALGORITHM: Final = "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", -} +REGION_ENDPOINTS: Final = MappingProxyType( + { + "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_INPUT_DOMESTIC: Final = "query_security_check_pro" +SERVICE_OUTPUT_DOMESTIC: Final = "response_security_check_pro" # Service codes for international regions -SERVICE_INPUT_INTERNATIONAL = "query_security_check_cb" -SERVICE_OUTPUT_INTERNATIONAL = "response_security_check_cb" +SERVICE_INPUT_INTERNATIONAL: Final = "query_security_check_cb" +SERVICE_OUTPUT_INTERNATIONAL: Final = "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" +CONTENT_MODERATION_TYPE: Final = "contentModeration" +PROMPT_ATTACK_TYPE: Final = "promptAttack" +SENSITIVE_DATA_TYPE: Final = "sensitiveData" +MALICIOUS_URL_TYPE: Final = "maliciousUrl" +MODEL_HALLUCINATION_TYPE: Final = "modelHallucination" +CUSTOM_LABEL_TYPE: Final = "customLabel" # Suggestion returned by Aliyun when it has decided the content must be rejected -BLOCK_SUGGESTION = "block" +BLOCK_SUGGESTION: Final = "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" +UNRESOLVED_BLOCK_LEVEL: Final = "high" + +# Risk level to integer, for threshold comparison. Covers both the standard levels +# (none/low/medium/high) and the sensitiveData levels (S0-S4). +RISK_LEVEL_TO_INT: Final = MappingProxyType( + { + # 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) + } +) def level_to_int(risk_level: str) -> int: @@ -106,31 +139,28 @@ def level_to_int(risk_level: str) -> int: - 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) + level_lower: Final = risk_level.lower() if risk_level else "none" + return RISK_LEVEL_TO_INT.get(level_lower, 0) +# Sentence-ending punctuation, used to cut long text on a boundary instead of mid-sentence +SENTENCE_BOUNDARY_PATTERN: Final = r"[。！？；:\.?!]+" + +# Parsed result of a response Aliyun did not flag at all +PASS_RESULT: Final = MappingProxyType( + {"flagged": False, "suggestion": "pass", "details": MappingProxyType({}), "message": ""} +) + # 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) -} +PROTECTION_LEVEL_THRESHOLD: Final = MappingProxyType( + { + "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): @@ -160,8 +190,8 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): service_input: str | None = None, service_output: str | None = None, service_mcp: str | None = None, - **kwargs, - ): + **kwargs: object, # kwargs-ok: forwarded verbatim to CustomGuardrail's own **kwargs contract + ) -> None: """ Initialize Aliyun AI Guardrail handler. Credentials (access_key_id / access_key_secret) are passed in from config.yaml @@ -199,7 +229,7 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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())}" + f"Valid values are: {tuple(PROTECTION_LEVEL_THRESHOLD)}" ) self.max_text_length = max_text_length or 2000 self.endpoint = REGION_ENDPOINTS.get(self.region_id, REGION_ENDPOINTS["cn-shanghai"]) @@ -211,10 +241,16 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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}" + "Initialized Aliyun AI Security Guardrail: %s, " + "region: %s, level: %s, " + "service_input: %s, service_output: %s, " + "service_mcp: %s", + guardrail_name, + self.region_id, + self.level, + self.service_input, + self.service_output, + self.service_mcp, ) @staticmethod @@ -239,57 +275,75 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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: Final = self.access_key_secret + "&" + signature: Final = 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: + def _create_string_to_sign(self, http_method: str, parameters: Mapping[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 = ( + canonicalized_query_string: Final = "".join( + f"&{self._percent_encode(key)}={self._percent_encode(parameters[key])}" for key in sorted(parameters) + ) + return ( 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]: + @staticmethod + def _iter_text_segments(text: str, max_length: int) -> Iterator[str]: + """ + Yield ``text`` in segments of at most ``max_length`` characters. + Each cut lands on the last sentence boundary inside the window, so a + segment does not split a sentence in half when it can be avoided. + """ + start = 0 # rebind-ok: cursor advances by one emitted segment per iteration + while start < len(text): + if len(text) - start <= max_length: + yield text[start:] + return + window = text[start : start + max_length] + boundaries = tuple(re.finditer(SENTENCE_BOUNDARY_PATTERN, window)) + cut_point = boundaries[-1].end() if boundaries else max_length + yield window[:cut_point] + start += cut_point + + def _split_text(self, text: str, max_length: int = 2000) -> tuple[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 + The text segments, in order """ - 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 + return tuple(self._iter_text_segments(text, max_length)) + + @staticmethod + def _build_service_parameters(text: str | None, image_urls: Sequence[str] | None) -> Mapping[str, Any]: + """ + Build the ServiceParameters payload of one guardrail request. + Args: + text: The text to audit, when there is any + image_urls: The image URLs to audit, when there are any + Returns: + The payload, as a real dict because json.dumps rejects a MappingProxyType + """ + return dict( # mutable-ok: json.dumps rejects a MappingProxyType + ( + ("requestFrom", "LiteLLM"), + *((("content", text),) if text else ()), + *((("imageUrls", tuple(image_urls)),) if image_urls else ()), + ) + ) async def async_make_request( self, text: str | None = None, service_type: Literal["input", "output", "mcp"] = "input", - image_urls: list[str] | None = None, + image_urls: Sequence[str] | None = None, ) -> AliyunAIGuardrailResponse: """ Make a request to the Aliyun AI Security Guardrail API. @@ -305,18 +359,16 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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 = { + service_code: Final = ( + self.service_mcp + if service_type == "mcp" + else self.service_input + if service_type == "input" + else self.service_output + ) + service_parameters: Final = self._build_service_parameters(text, image_urls) + # httpx form-encodes `data=` from a dict, so this payload stays a real dict too + parameters: Final[dict[str, str]] = { # mutable-ok: httpx form-encodes `data=` from a dict "Action": "MultiModalGuard", "Version": "2022-03-02", "AccessKeyId": self.access_key_id, @@ -328,32 +380,33 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): "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 + string_to_sign: Final = self._create_string_to_sign("POST", parameters) + parameters["Signature"] = self._create_signature(string_to_sign) 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( + response: Final = await self.async_handler.post( url=self.service_url, data=parameters, - headers={"Content-Type": "application/x-www-form-urlencoded"}, + headers={"Content-Type": "application/x-www-form-urlencoded"}, # mutable-ok: httpx wants a dict timeout=30.0, ) - body = response.json() + body: Final = 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}"}, + detail={ # mutable-ok: HTTPException detail payload + "error": f"Aliyun AI Guardrail request failed. Status: {response.status_code}, Body: {body}" + }, ) if body.get("Code") != 200: raise HTTPException( status_code=400, - detail={ + detail={ # mutable-ok: HTTPException detail payload "error": f"Aliyun AI Guardrail API error. Code: {body.get('Code')}, Message: {body.get('Message')}" }, ) @@ -373,11 +426,11 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): Returns: True if should block, False otherwise """ - threshold = PROTECTION_LEVEL_THRESHOLD.get(self.level, 99) - detected_int = level_to_int(detected_level) + threshold: Final = PROTECTION_LEVEL_THRESHOLD.get(self.level, 99) + detected_int: Final = level_to_int(detected_level) return detected_int >= threshold - def _resolve_detail_level(self, detail: dict[str, Any]) -> str: + def _resolve_detail_level(self, detail: Mapping[str, Any]) -> str: """ Resolve the risk level of a single Detail entry. @@ -390,34 +443,80 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): Returns: The risk level string to weigh against the configured threshold """ - level = detail.get("Level") + level: Final = 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 + risk_levels: Final = tuple( + risk_level + for result in detail.get("Result") or () + if isinstance(result, dict) + for risk_level in (result.get("RiskLevel"),) + if isinstance(risk_level, str) and risk_level.strip() + ) + if risk_levels: + # max() keeps the first of equal-severity levels, as the previous scan did + return max(risk_levels, key=level_to_int) # No parseable severity: never downgrade an explicit block to "none" if detail.get("Suggestion") == BLOCK_SUGGESTION: return UNRESOLVED_BLOCK_LEVEL return "none" + @staticmethod + def _resolve_desensitization(detail_list: Sequence[Any]) -> str: + """ + Return the desensitized text reported for sensitive data. + The first ``Desensitization`` within one detail's results wins, and a later + sensitiveData detail overrides an earlier one. + Args: + detail_list: The response ``Detail`` entries + Returns: + The desensitized text, empty when none was reported + """ + resolved = "" # rebind-ok: a later sensitiveData detail overrides an earlier one + for detail in detail_list: + if detail.get("Type", "") != SENSITIVE_DATA_TYPE: + continue + for ext in (result.get("Ext") for result in detail.get("Result", ()) or ()): + if ext and ext.get("Desensitization"): + resolved = ext.get("Desensitization", "") + break + return resolved + + def _resolve_block_decision(self, detail_list: Sequence[Any], final_suggestion: str) -> tuple[str, str, str]: + """ + Decide whether the audited content must be blocked. + Args: + detail_list: The response ``Detail`` entries + final_suggestion: The overall ``Suggestion`` of the response + Returns: + (blocked_type, blocked_level, block_message); an empty message means pass + """ + for detail in detail_list: + detected_level = self._resolve_detail_level(detail) + if self._should_block_by_level(detected_level): + detection_type = detail.get("Type", "") + return detection_type, detected_level, 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 final_suggestion == BLOCK_SUGGESTION and self._should_block_by_level(UNRESOLVED_BLOCK_LEVEL): + return "", UNRESOLVED_BLOCK_LEVEL, f"阿里云返回阻断建议 (Suggestion: {final_suggestion})" + return "", "", "" + + def _summarise_detail(self, detail: Mapping[str, Any]) -> Mapping[str, Any]: + """Summarise one ``Detail`` entry for the parsed result payload.""" + return { # mutable-ok: returned payload + "level": self._resolve_detail_level(detail), + "suggestion": detail.get("Suggestion", "pass"), + "results": tuple(detail.get("Result") or ()), + } + def _parse_response_and_check( self, response: AliyunAIGuardrailResponse, check_type: Literal["input", "output"], - ) -> dict[str, Any]: + ) -> Mapping[str, Any]: """ Parse the guardrail response and check if content should be blocked. Blocking logic: @@ -426,65 +525,38 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): response: The API response check_type: "input" or "output" Returns: - Dict with parsed results + Mapping with parsed results Raises: HTTPException if content should be blocked """ - data = response.get("Data", {}) + data: Final = 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})" + return PASS_RESULT + final_suggestion: Final = data.get("Suggestion", "pass") + detail_list: Final = tuple(data.get("Detail") or ()) + details: Final[dict[str, Mapping[str, Any]]] = { # mutable-ok: returned payload + detail.get("Type", ""): self._summarise_detail(detail) for detail in detail_list + } + desensitization: Final = self._resolve_desensitization(detail_list) + blocked_type, blocked_level, block_message = self._resolve_block_decision(detail_list, 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}" + "Aliyun AI Guardrail: level=%s, check_type=%s, should_block=%s, blocked_type=%s, blocked_level=%s", + self.level, + check_type, + bool(block_message), + blocked_type, + blocked_level, ) - if should_block: + if block_message: raise HTTPException( status_code=400, - detail={ + detail={ # mutable-ok: HTTPException detail payload "error": f"Aliyun AI Guardrail: {block_message}", "type": check_type, "details": details, }, ) - return { + return { # mutable-ok: returned payload "flagged": final_suggestion != "pass", "suggestion": final_suggestion, "desensitization": desensitization, @@ -492,26 +564,37 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): "message": block_message, } + @staticmethod + def _build_input_payloads( + segments: Sequence[str], image_urls: Sequence[str] + ) -> tuple[tuple[str | None, Sequence[str] | None], ...]: + """ + Pair each text segment with the images to audit alongside it. + Every image URL rides with the first text segment so content and images are + checked together; later segments carry text only. With no text at all, a + single image-only request is produced. + Args: + segments: The text segments to audit + image_urls: The image URLs to audit + Returns: + (text, images) pairs, one per request to send + """ + if segments: + return tuple( + (segment, image_urls if idx == 0 and image_urls else None) for idx, segment in enumerate(segments) + ) + if image_urls: + return ((None, image_urls),) + return () + @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: + cache: DualCache, + data: dict[str, Any], # mutable-ok: CustomLogger's hook contract + call_type: CallTypesLiteral, + ) -> dict[str, Any] | None: # mutable-ok: CustomLogger's hook contract """ Pre-call hook to scan user prompts before sending to LLM. Raises HTTPException if content should be blocked. @@ -520,14 +603,18 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): "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: + if call_type == CallTypes.call_mcp_tool.value: + await self._mcp_pre_call_check(data) + return None + # Walks messages AND the Responses API's input, so /v1/responses is covered + new_messages: Final = cast( # cast-ok: the upstream iterator is typed as plain dicts + "Sequence[AllMessageValues]", tuple(_iter_inspection_messages(data)) + ) + if not new_messages: 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) + user_prompt: Final = self.get_user_prompt(new_messages) + image_urls: Final = 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 @@ -536,29 +623,18 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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 + # _split_text already returns a single segment when the text fits + segments: Final = self._split_text(user_prompt, self.max_text_length) if user_prompt else () + payloads: Final = self._build_input_payloads(segments, image_urls) + semaphore: Final = 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 def check_with_semaphore( + segment_text: str | None, segment_images: Sequence[str] | None + ) -> AliyunAIGuardrailResponse: 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]) + responses: Final = 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") @@ -568,31 +644,28 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): # MCP-specific guardrail methods # ================================================================ - async def _mcp_pre_call_check(self, data: dict) -> dict | None: + async def _mcp_pre_call_check(self, data: Mapping[str, Any]) -> None: """MCP pre-call: audit tool name + arguments before execution.""" - messages = data.get("messages", []) - content = messages[0].get("content", "") if messages else "" + messages: Final = data.get("messages", ()) + content: Final = 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) + return + # _split_text already returns a single segment when the text fits + segments: Final = self._split_text(content, self.max_text_length) + semaphore: Final = asyncio.Semaphore(5) - async def check(text: str): + async def check(text: str) -> AliyunAIGuardrailResponse: async with semaphore: return await self.async_make_request(text=text, service_type="mcp") - responses = await asyncio.gather(*[check(s) for s in segments]) + responses: Final = 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. @@ -621,16 +694,14 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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 + default_value: Final = self.event_hook.default + return "post_mcp_call" in (default_value if isinstance(default_value, list) else (default_value,)) return False return self.event_hook == "post_mcp_call" async def async_post_mcp_tool_call_hook( self, - kwargs: dict, + kwargs: Mapping[str, Any], response_obj: MCPPostCallResponseObject, start_time: datetime, end_time: datetime, @@ -649,32 +720,36 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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 + original_response: Final = kwargs.get("original_response") + # The live tool result is preferred; the wrapped copy is the fallback. + combined_text: Final = next( + ( + text + for text in ( + self._extract_mcp_tool_text(candidate) + for candidate in (original_response, getattr(response_obj, "mcp_tool_call_response", None)) + if candidate is not None + ) + if text + ), + "", + ) 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) + # _split_text already returns a single segment when the text fits + segments: Final = self._split_text(combined_text, self.max_text_length) + semaphore: Final = asyncio.Semaphore(5) - async def check(text: str): + async def check(text: str) -> AliyunAIGuardrailResponse: 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]) + responses: Final = 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: @@ -686,16 +761,19 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): response_obj=response_obj, original_response=original_response, ) - except Exception as e: + except Exception as e: # noqa: BLE001 - fail closed on any guardrail failure # 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)}", + "Aliyun AI Guardrail: ★ MCP post-call check failed, blocking unaudited tool output: %s", + e, exc_info=True, ) return self._block_mcp_tool_output( - detail={"error": f"Aliyun AI Guardrail 调用失败，工具输出未经审核，已拦截: {str(e)}"}, + detail={ # mutable-ok: HTTPException detail payload + "error": f"Aliyun AI Guardrail 调用失败，工具输出未经审核，已拦截: {e!s}" + }, # mutable-ok: detail payload response_obj=response_obj, original_response=original_response, ) @@ -703,138 +781,253 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): return None @staticmethod - def _iter_mcp_content_items(payload: Any) -> list[Any]: + def _find_coerced_field(payload: Sequence[object], field: str) -> object: """ - Normalise an MCP tool result into its list of content items. + Return one field of a CallToolResult that was coerced into (field, value) pairs. + Args: + payload: The coerced pairs + field: The field name to recover + Returns: + The field's value, or None when it is absent + """ + return next( + (entry[1] for entry in payload if isinstance(entry, tuple) and len(entry) == 2 and entry[0] == field), + None, + ) + + @classmethod + def _iter_mcp_content_items(cls, payload: object) -> tuple[object, ...]: + """ + Normalise an MCP tool result into its 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) + return (payload,) + content: Final = getattr(payload, "content", None) if isinstance(content, list): - return content + return tuple(content) if isinstance(payload, dict): - inner = payload.get("content") - return inner if isinstance(inner, list) else [payload] + inner: Final = payload.get("content") + return tuple(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 [] + coerced: Final = cls._find_coerced_field(payload, "content") + return tuple(coerced) if isinstance(coerced, list) else tuple(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) + @classmethod + def _extract_mcp_item_text(cls, item: object) -> str: + """ + Collect the text of a single MCP content item. + Args: + item: A content item, as a string, dict or MCP content model + Returns: + The item's text, empty when it carries none + """ + if isinstance(item, str): + return item + direct_text: Final = item.get("text") if isinstance(item, dict) else getattr(item, "text", None) + if direct_text: + return str(direct_text) + # An EmbeddedResource keeps its payload one level down, on the resource. + resource: Final = item.get("resource") if isinstance(item, dict) else getattr(item, "resource", None) + resource_text: Final = resource.get("text") if isinstance(resource, dict) else getattr(resource, "text", None) + return str(resource_text or "") + + @classmethod + def _extract_mcp_structured_content(cls, payload: object) -> str: + """ + Serialise a tool result's ``structuredContent``, when it carries one. + Args: + payload: A CallToolResult, a raw response body, or its coerced form + Returns: + The serialised structured payload, empty when there is none + """ + structured: Final = ( + payload.get("structuredContent") + if isinstance(payload, dict) + # MCPPostCallResponseObject coerces a CallToolResult into (field, value) pairs. + else cls._find_coerced_field(payload, "structuredContent") + if isinstance(payload, list) + else getattr(payload, "structuredContent", None) + ) + if structured is None: + return "" + if isinstance(structured, str): + return structured + return json.dumps(structured, ensure_ascii=False, default=str) + + def _extract_mcp_tool_text(self, payload: object) -> str: + """ + Collect the textual output of an MCP tool result. + A tool result carries text beyond ``content[].text``: ``structuredContent`` + holds an arbitrary JSON payload and an EmbeddedResource keeps its text on + ``resource.text``. Auditing only the plain text items would let a tool + return prohibited content in either field unchecked. + """ + item_texts: Final = tuple( + text for text in map(self._extract_mcp_item_text, self._iter_mcp_content_items(payload)) if text + ) + structured_text: Final = self._extract_mcp_structured_content(payload) + return "\n".join((*item_texts, structured_text) if structured_text else item_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.""" + def _replace_tool_output_in_place(target: object, blocked_content: Sequence[Any]) -> bool: + """ + Overwrite an MCP tool result's content with ``blocked_content``, in place. + The in-place stores are the point: both dispatch sites hand the caller's own + tool result to the client, so the violation has to be written into it. + """ if target is None: return False - content = getattr(target, "content", None) + content: Final = getattr(target, "content", None) if isinstance(content, list): content[:] = blocked_content if hasattr(target, "isError"): try: - target.isError = True + # Kept duck-typed on purpose: any MCP result shape carrying `isError` + # must be flagged, not just the SDK's CallToolResult. + target.isError = True # pyright: ignore[reportAttributeAccessIssue] # rebind-ok: flag the caller's result except (AttributeError, TypeError, ValueError): pass return True if isinstance(target, list): - target[:] = blocked_content + target[:] = blocked_content # rebind-ok: the caller's result must be overwritten return True if isinstance(target, dict): - result = target.get("result") + result: Final = target.get("result") if isinstance(result, dict) and isinstance(result.get("content"), list): - result["content"] = list(blocked_content) + result["content"] = list(blocked_content) # mutable-ok: the replaced field must stay a JSON list return True if isinstance(target.get("content"), list): - target["content"] = list(blocked_content) + blocked: Final = list(blocked_content) # mutable-ok: must stay a JSON list + target["content"] = blocked # rebind-ok: overwrite the caller's result return True return False def _block_mcp_tool_output( self, - detail: Any, + detail: object, response_obj: MCPPostCallResponseObject, - original_response: Any, + original_response: object, ) -> 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))] + from litellm.types.llms.base import HiddenParams + from litellm.types.mcp import MCPPostCallResponseObject as _MCPPostCallResponseObject + + payload: Final = detail if isinstance(detail, dict) else {"error": str(detail)} # mutable-ok: JSON payload + blocked_content: Final = (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) + hidden_params: Final = 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(), ) + @staticmethod + def _iter_function_call_text(function: object) -> Iterator[str]: + """Yield the name and arguments of a tool/function call.""" + for field in ("name", "arguments"): + value = getattr(function, field, None) + if value: + yield str(value) + + @staticmethod + def _iter_output_item_text(output_items: Iterable[object] | None) -> Iterator[str]: + """Yield the text carried by Responses API output items.""" + for output_item in output_items 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: + yield str(text) + else: + text = getattr(output_item, "text", None) + if text: + yield str(text) + + @classmethod + def _iter_completion_text(cls, message: object) -> Iterator[str]: + """ + Yield every text field a chat completion message or streaming delta carries. + Args: + message: A choice's ``message`` (non-streaming) or ``delta`` (streaming) + """ + for field in ("content", "reasoning_content"): + value = getattr(message, field, None) + if value: + yield str(value) + for call in getattr(message, "tool_calls", None) or (): + yield from cls._iter_function_call_text(getattr(call, "function", None)) + yield from cls._iter_function_call_text(getattr(message, "function_call", None)) + + @classmethod + def _extract_response_text(cls, response: object) -> str: + """ + Collect every text field a non-streaming response can carry. + Mirrors ``_extract_stream_chunk_text``: auditing only ``message.content`` + would release tool call arguments, reasoning text and every + /v1/responses body to the client unchecked whenever ``stream=False``. + Args: + response: A non-streaming response object + Returns: + The concatenated text to audit, empty when the response carries none + """ + return "\n".join( + ( + # Responses API bodies keep their text in output items, not in choices. + *cls._iter_output_item_text(getattr(response, "output", None)), + *( + text + for choice in getattr(response, "choices", None) or () + for text in cls._iter_completion_text(getattr(choice, "message", None)) + ), + ) + ) + @log_guardrail_information async def async_post_call_success_hook( self, - data: dict, + data: dict, # mutable-ok: CustomLogger's hook contract user_api_key_dict: UserAPIKeyAuth, - response: Any | ModelResponse | EmbeddingResponse | ImageResponse, - ) -> Any: + response: LLMResponseTypes, + ) -> LLMResponseTypes: """ 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 + # Every choice is scanned: n>1 responses would otherwise return unchecked content + content: Final = self._extract_response_text(response) + if content: + verbose_proxy_logger.info( + "Aliyun AI Guardrail: Post-call scan started, response length: %d", + len(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 + # _split_text already returns a single segment when the text fits + segments: Final = self._split_text(content, self.max_text_length) + semaphore: Final = 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") + async def check_with_semaphore(segment: str) -> AliyunAIGuardrailResponse: + 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") + responses: Final = 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: + @classmethod + def _extract_stream_chunk_text(cls, chunk: object) -> str: """ Collect every text field a streaming chunk can carry. Covers both chat completion chunks and Responses API streaming events. @@ -845,49 +1038,31 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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)) + event_delta: Final = getattr(chunk, "delta", None) + return "".join( + ( + *((event_delta,) if isinstance(event_delta, str) and event_delta else ()), + *cls._iter_output_item_text(getattr(getattr(chunk, "response", None), "output", None)), + *( + text + for choice in getattr(chunk, "choices", None) or () + for text in cls._iter_completion_text(getattr(choice, "delta", None)) + ), + ) + ) - 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) + @staticmethod + def _as_error_payload(detail: object) -> Mapping[str, Any]: + """Normalise an HTTPException detail into a JSON object.""" + return detail if isinstance(detail, dict) else {"message": str(detail)} # mutable-ok: JSON payload async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, - request_data: dict[str, Any], + response: AsyncIterable[Any], + request_data: Mapping[str, Any], ) -> AsyncGenerator[Any, None]: """ Process streaming response with sliding window guardrail checks. @@ -903,11 +1078,11 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): - 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 + # Sliding-window state: every name below advances as the stream is consumed. + accumulated_text = "" # rebind-ok: grows with each chunk + last_check_position = 0 # rebind-ok: position (total length) of the last check + pending_chunks: Final[list[Any]] = [] # mutable-ok: buffer held back until a check passes + is_first_check = True # rebind-ok: the first check uses a smaller threshold verbose_proxy_logger.info( "Aliyun AI Guardrail: Streaming scan started, window=%d, step=%d, first_check_step=%d", self.stream_window_size, @@ -918,7 +1093,6 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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) @@ -939,10 +1113,10 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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") + final_start: Final = max(0, len(accumulated_text) - self.stream_window_size) + remaining_text: Final = accumulated_text[final_start:] + final_response: Final = await self.async_make_request(text=remaining_text, service_type="output") + self._parse_response_and_check(final_response, check_type="output") verbose_proxy_logger.info( "Aliyun AI Guardrail: Streaming scan completed, total length: %d", len(accumulated_text) ) @@ -950,10 +1124,11 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): yield pending_chunk pending_chunks.clear() except HTTPException as e: - error_detail = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} + detail_payload: Final = self._as_error_payload(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" + payload: Final = json.dumps({"error": detail_payload}, ensure_ascii=False) # mutable-ok: SSE payload + yield f"data: {payload}\n\n" return except Exception as e: - verbose_proxy_logger.error(f"Aliyun AI Guardrail streaming error: {str(e)}", exc_info=True) + verbose_proxy_logger.error("Aliyun AI Guardrail streaming error: %s", 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 index 8917a06df5b..452e8638b8a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py @@ -5,7 +5,8 @@ Base class for Aliyun guardrails from __future__ import annotations -from typing import TYPE_CHECKING +from collections.abc import Iterator, Sequence +from typing import TYPE_CHECKING, Final if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues @@ -16,45 +17,69 @@ class AliyunGuardrailBase: Base class for Aliyun guardrails. """ - def get_user_prompt(self, messages: list[AllMessageValues]) -> str | None: + @staticmethod + def _iter_user_messages(messages: Sequence[AllMessageValues]) -> Iterator[AllMessageValues]: """ - Get the last consecutive block of messages from the user. + Yield every user message of the request, in order. + Restricting this to the trailing user block would let a caller hide a + prohibited turn behind an attacker-supplied assistant message. + """ + return (message for message in messages if message.get("role") == "user") + + @staticmethod + def _extract_image_url(part: object) -> str | None: + """ + Return the URL of an ``image_url`` content part. + Args: + part: A single content part of a message + Returns: + The URL string, or None when the part carries no image URL + """ + if not isinstance(part, dict) or part.get("type") != "image_url": + return None + image_url: Final = part.get("image_url") + if isinstance(image_url, dict): + url: Final = image_url.get("url") + return url if isinstance(url, str) else None + return image_url if isinstance(image_url, str) else None + + def get_user_prompt(self, messages: Sequence[AllMessageValues]) -> str | None: + """ + Collect the text of every user message in the request. + Scanning only the trailing user block would let a caller hide a + prohibited prompt behind an attacker-supplied assistant message, so all + user turns of the submitted request are audited. 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?" + get_user_prompt(messages) -> "Hello, how are you?\nWhat 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 + user_prompt: Final = "\n".join( + convert_content_list_to_str(message) for message in self._iter_user_messages(messages) + ).strip() + return user_prompt or None - def get_image_urls(self, messages: list[AllMessageValues]) -> list[str]: + def _iter_public_image_urls(self, messages: Sequence[AllMessageValues]) -> Iterator[str]: + """Yield the publicly reachable image URLs of every user message, in order.""" + for content in (message.get("content") for message in self._iter_user_messages(messages)): + if not isinstance(content, list): + continue + for url in (self._extract_image_url(part) for part in content): + # Only public http(s) URLs are reachable by the Aliyun API, so + # data: URIs and other inline payloads are skipped. + if url is not None and url.startswith(("http://", "https://")): + yield url + + def get_image_urls(self, messages: Sequence[AllMessageValues]) -> tuple[str, ...]: """ - Extract image URLs from the last consecutive block of user messages. + Extract image URLs from every user message in the request. Only publicly accessible http(s) URLs are collected (in order, de-duplicated). Uses the same message range as ``get_user_prompt``. Example: @@ -64,40 +89,8 @@ class AliyunGuardrailBase: {"type": "image_url", "image_url": {"url": "https://a.com/x.png"}}, ]}, ] - get_image_urls(messages) -> ["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 + # dict.fromkeys is the order-preserving dedup; it is transient and the + # result is frozen into a tuple before it leaves this method. + return tuple(dict.fromkeys(self._iter_public_image_urls(messages))) # mutable-ok: transient dedup, frozen here diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 9cdc3d2e155..9350da5af32 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1177,4 +1177,3 @@ 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 index 2a612ecd6ad..9d763602352 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py @@ -8,10 +8,11 @@ Aliyun AI Guardrail supports the following detection types: - maliciousUrl: Malicious URL detection """ -from typing import Any, Dict, List, Literal, Optional +from collections.abc import Sequence +from typing import Literal, TypeAlias from pydantic import BaseModel, Field -from typing_extensions import NotRequired, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from ..base import GuardrailConfigModel @@ -20,77 +21,79 @@ from ..base import GuardrailConfigModel class AliyunAIGuardrailResponseDetailResultExt(TypedDict, total=False): """Extended information in result""" - Desensitization: Optional[str] # Desensitized text when action is mask + Desensitization: ReadOnly[str | None] # Desensitized text when action is mask class AliyunAIGuardrailResponseDetailResult(TypedDict, total=False): """Result item in detail""" - Confidence: Optional[float] - Label: Optional[str] - Ext: Optional[AliyunAIGuardrailResponseDetailResultExt] + Confidence: ReadOnly[float | None] + Label: ReadOnly[str | None] + Ext: ReadOnly[AliyunAIGuardrailResponseDetailResultExt | None] # 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] + RiskLevel: ReadOnly[str | None] class AliyunAIGuardrailResponseDetail(TypedDict): """Detail item in response data""" - Type: str # contentModeration, sensitiveData, promptAttack, maliciousUrl - Suggestion: str # pass, block, mask - Result: List[AliyunAIGuardrailResponseDetailResult] + Type: ReadOnly[str] # contentModeration, sensitiveData, promptAttack, maliciousUrl + Suggestion: ReadOnly[str] # pass, block, mask + Result: ReadOnly[Sequence[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] + Level: ReadOnly[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]] + Suggestion: ReadOnly[str] # Overall suggestion: pass, block, mask + Detail: ReadOnly[Sequence[AliyunAIGuardrailResponseDetail] | None] class AliyunAIGuardrailResponse(TypedDict): """Response from Aliyun AI Guardrail API""" - RequestId: str - Code: int - Message: Optional[str] - Data: Optional[AliyunAIGuardrailResponseData] + RequestId: ReadOnly[str] + Code: ReadOnly[int] + Message: ReadOnly[str | None] + Data: ReadOnly[AliyunAIGuardrailResponseData | None] # Suggestion type -AliyunAIGuardrailSuggestion = Literal["pass", "block", "watch"] +AliyunAIGuardrailSuggestion: TypeAlias = Literal["pass", "block", "watch"] # Detection type -AliyunAIGuardrailDetectionType = Literal["contentModeration", "sensitiveData", "promptAttack", "maliciousUrl"] +AliyunAIGuardrailDetectionType: TypeAlias = 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 + Action: ReadOnly[str] + Version: ReadOnly[str] + AccessKeyId: ReadOnly[str] + Timestamp: ReadOnly[str] + SignatureMethod: ReadOnly[str] + SignatureVersion: ReadOnly[str] + SignatureNonce: ReadOnly[str] + Format: ReadOnly[str] + Service: ReadOnly[str] + ServiceParameters: ReadOnly[str] + Signature: ReadOnly[str] # Risk level literals -AliyunRiskLevel = Literal["none", "low", "medium", "high"] +AliyunRiskLevel: TypeAlias = Literal["none", "low", "medium", "high"] # Protection level literals -AliyunProtectionLevel = Literal["low", "medium", "high", "max"] +AliyunProtectionLevel: TypeAlias = Literal["low", "medium", "high", "max"] # Configuration models @@ -101,39 +104,39 @@ class AliyunAIGuardrailOptionalParams(BaseModel): in config.yaml on the AliyunAIGuardrailConfigModel and support os.environ/ references. """ - level: Optional[AliyunProtectionLevel] = Field( + level: AliyunProtectionLevel | None = 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( + max_text_length: int | None = 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( + stream_window_size: int | None = 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( + stream_slide_step: int | None = 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( + stream_first_check_step: int | None = 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( + region_id: str | None = Field( default="cn-shanghai", description="Aliyun region ID. Default: cn-shanghai", ) - service_input: Optional[str] = Field( + service_input: str | None = Field( default="query_security_check_pro", description="Service code for input (pre-call) detection. Default: query_security_check_pro", ) - service_output: Optional[str] = Field( + service_output: str | None = Field( default="response_security_check_pro", description="Service code for output (post-call) detection. Default: response_security_check_pro", ) - service_mcp: Optional[str] = Field( + service_mcp: str | None = 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", ) @@ -147,15 +150,15 @@ class AliyunAIGuardrailConfigModel(GuardrailConfigModel[AliyunAIGuardrailOptiona - access_key_secret: Aliyun Access Key Secret """ - access_key_id: Optional[str] = Field( + access_key_id: str | None = Field( default=None, description="Aliyun Access Key ID. Configure in config.yaml, supports os.environ/ reference", ) - access_key_secret: Optional[str] = Field( + access_key_secret: str | None = Field( default=None, description="Aliyun Access Key Secret. Configure in config.yaml, supports os.environ/ reference", ) - optional_params: AliyunAIGuardrailOptionalParams = Field( + optional_params: AliyunAIGuardrailOptionalParams | None = Field( default_factory=AliyunAIGuardrailOptionalParams, description="Optional parameters for the Aliyun AI 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 index d3833cd8644..3c28632cbd8 100644 --- 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 @@ -295,7 +295,7 @@ 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"] + assert result == ("short text",) def test_long_text_splits_at_sentence_boundary(self): g = _make_guardrail() @@ -311,10 +311,10 @@ class TestSplitText: assert len(result) >= 3 assert "".join(result) == text - def test_empty_text_returns_empty_list(self): + def test_empty_text_returns_no_segments(self): g = _make_guardrail() result = g._split_text("", max_length=100) - assert result == [] + assert result == () # --------------------------------------------------------------------------- @@ -591,6 +591,62 @@ class TestConfigModel: # --------------------------------------------------------------------------- +# --------------------------------------------------------------------------- +# Trailing non-user message must not hide the prompt from scanning +# --------------------------------------------------------------------------- + + +class TestTrailingAssistantMessage: + def test_user_text_survives_trailing_assistant_message(self): + g = _make_guardrail() + messages = [ + {"role": "user", "content": "违规的用户提问"}, + {"role": "assistant", "content": "攻击者伪造的回复"}, + ] + prompt = g.get_user_prompt(messages) + assert prompt is not None + assert "违规的用户提问" in prompt + + def test_images_survive_trailing_assistant_message(self): + g = _make_guardrail() + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": IMG_A}}], + }, + {"role": "assistant", "content": "攻击者伪造的回复"}, + ] + assert g.get_image_urls(messages) == (IMG_A,) + + @pytest.mark.asyncio + async def test_blocks_violation_despite_trailing_assistant_message(self): + 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"]).get("content", "") + return blocked if "违规的用户提问" in scanned else clean + + 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_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={ + "messages": [ + {"role": "user", "content": "违规的用户提问"}, + {"role": "assistant", "content": "攻击者伪造的回复"}, + ] + }, + call_type="completion", + ) + assert exc_info.value.status_code == 400 + + class TestGetImageUrls: def test_extracts_http_and_https_urls(self): g = _make_guardrail() @@ -604,7 +660,7 @@ class TestGetImageUrls: ], } ] - assert g.get_image_urls(messages) == [IMG_A, IMG_B] + assert g.get_image_urls(messages) == (IMG_A, IMG_B) def test_skips_non_url_images(self): g = _make_guardrail() @@ -618,12 +674,12 @@ class TestGetImageUrls: ], } ] - assert g.get_image_urls(messages) == [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) == [] + assert g.get_image_urls(messages) == () def test_deduplicates_across_messages(self): g = _make_guardrail() @@ -631,20 +687,20 @@ class TestGetImageUrls: {"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] + assert g.get_image_urls(messages) == (IMG_A,) - def test_only_last_consecutive_user_block(self): + def test_collects_images_from_every_user_message(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] + assert g.get_image_urls(messages) == (IMG_B, IMG_A) def test_empty_messages_returns_empty(self): g = _make_guardrail() - assert g.get_image_urls([]) == [] + assert g.get_image_urls([]) == () # --------------------------------------------------------------------------- @@ -693,6 +749,67 @@ class TestServiceParametersConstruction: class TestPreCallHook: + @pytest.mark.asyncio + async def test_scans_responses_api_string_input(self): + g = _make_guardrail(level="medium") + 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={"input": "违规的 responses 输入"}, + call_type="responses", + ) + + scanned = "".join( + json.loads(call.kwargs["data"]["ServiceParameters"]).get("content", "") for call in mock_post.call_args_list + ) + assert "违规的 responses 输入" in scanned + + @pytest.mark.asyncio + async def test_blocks_violating_responses_api_input(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={"input": "违规的 responses 输入"}, + call_type="responses", + ) + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_scans_responses_api_structured_input_with_image(self): + g = _make_guardrail(level="medium") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + data = { + "input": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "结构化输入文本"}, + {"type": "image_url", "image_url": {"url": IMG_A}}, + ], + } + ] + } + 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=data, + call_type="responses", + ) + + sent = [json.loads(call.kwargs["data"]["ServiceParameters"]) for call in mock_post.call_args_list] + assert any("结构化输入文本" in (sp.get("content") or "") for sp in sent) + assert any(IMG_A in (sp.get("imageUrls") or []) for sp in sent) + @pytest.mark.asyncio async def test_blocks_violation(self): g = _make_guardrail(level="medium") @@ -982,6 +1099,161 @@ class TestPostCallHook: assert result is response +def _make_tool_call_response(arguments: str, name: str = "send_email"): + """Build a non-streaming response whose only output is a tool call.""" + import litellm + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + return litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function(name=name, arguments=arguments), + ) + ], + ), + ) + ], + ) + + +def _make_responses_api_response(text: str): + """Build a non-streaming /v1/responses body carrying assistant text.""" + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.responses.main import GenericResponseOutputItem, OutputText + + return ResponsesAPIResponse( + id="resp_1", + created_at=0, + model="gpt-4o", + object="response", + output=[ + GenericResponseOutputItem( + type="message", + id="m1", + status="completed", + role="assistant", + content=[OutputText(type="output_text", text=text, annotations=[])], + ) + ], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + temperature=1.0, + top_p=1.0, + ) + + +class TestPostCallStructuredFields: + """The streaming path already audits tool calls, reasoning text and + /v1/responses output. Auditing only ``message.content`` here would let the + very same content reach the client unchecked whenever stream=False.""" + + @pytest.mark.asyncio + async def test_scans_tool_call_arguments(self): + g = _make_guardrail(level="medium") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + response = _make_tool_call_response('{"body": "违规的工具参数"}') + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) 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 + + @pytest.mark.asyncio + async def test_blocks_violation_in_tool_call_arguments(self): + g = _make_guardrail(level="medium") + blocked = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + response = _make_tool_call_response('{"body": "违规的工具参数"}') + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=blocked): + 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_reasoning_content(self): + import litellm + + g = _make_guardrail(level="medium") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + response = litellm.ModelResponse( + id="test-id", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", + content="正常的回复内容", + reasoning_content="推理过程里的违规内容", + ), + ) + ], + ) + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) 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 + + @pytest.mark.asyncio + async def test_scans_responses_api_output(self): + g = _make_guardrail(level="medium") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + response = _make_responses_api_response("响应体里的违规内容") + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) as mock_post: + await g.async_post_call_success_hook( + data={"input": "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 + + @pytest.mark.asyncio + async def test_blocks_violation_in_responses_api_output(self): + g = _make_guardrail(level="medium") + blocked = _make_aliyun_api_response( + suggestion="block", + detail=[_make_detail(detection_type=CONTENT_MODERATION_TYPE, level="high")], + ) + response = _make_responses_api_response("响应体里的违规内容") + with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=blocked): + with pytest.raises(HTTPException) as exc_info: + await g.async_post_call_success_hook( + data={"input": "hi"}, + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + response=response, + ) + assert exc_info.value.status_code == 400 + + # --------------------------------------------------------------------------- # Post-MCP hook tests # --------------------------------------------------------------------------- @@ -993,6 +1265,93 @@ def _make_call_tool_result(text: str = "tool output"): return CallToolResult(content=[TextContent(type="text", text=text)], isError=False) +class TestExtractMcpToolText: + """A tool result carries text outside of ``content[].text``. Auditing only that + field would release structured payloads and embedded resources unchecked.""" + + def test_collects_structured_content(self): + from mcp.types import CallToolResult + + g = _make_guardrail() + result = CallToolResult(content=[], structuredContent={"note": "结构化字段里的违规内容"}, isError=False) + assert "结构化字段里的违规内容" in g._extract_mcp_tool_text(result) + + def test_collects_structured_content_alongside_text(self): + from mcp.types import CallToolResult, TextContent + + g = _make_guardrail() + result = CallToolResult( + content=[TextContent(type="text", text="正常的工具输出")], + structuredContent={"note": "结构化字段里的违规内容"}, + isError=False, + ) + extracted = g._extract_mcp_tool_text(result) + assert "正常的工具输出" in extracted + assert "结构化字段里的违规内容" in extracted + + def test_collects_embedded_resource_text(self): + from mcp.types import CallToolResult, EmbeddedResource, TextResourceContents + + g = _make_guardrail() + result = CallToolResult( + content=[ + EmbeddedResource( + type="resource", + resource=TextResourceContents( + uri="file:///tmp/note.txt", + mimeType="text/plain", + text="内嵌资源里的违规内容", + ), + ) + ], + isError=False, + ) + assert "内嵌资源里的违规内容" in g._extract_mcp_tool_text(result) + + def test_collects_structured_content_from_dict_payload(self): + g = _make_guardrail() + payload = {"content": [], "structuredContent": {"note": "结构化字段里的违规内容"}} + assert "结构化字段里的违规内容" in g._extract_mcp_tool_text(payload) + + def test_collects_structured_content_from_coerced_tuple_list(self): + """MCPPostCallResponseObject coerces a CallToolResult into (field, value) pairs.""" + g = _make_guardrail() + payload = [("content", []), ("structuredContent", {"note": "结构化字段里的违规内容"}), ("isError", False)] + assert "结构化字段里的违规内容" in g._extract_mcp_tool_text(payload) + + @pytest.mark.asyncio + async def test_blocks_violation_in_structured_content(self): + from mcp.types import CallToolResult, TextContent + + g = _make_guardrail(level="medium") + tool_result = CallToolResult( + content=[TextContent(type="text", text="正常的工具输出")], + structuredContent={"note": "结构化字段里的违规内容"}, + isError=False, + ) + kwargs, response_obj = _make_post_mcp_hook_args(tool_result) + 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_structured_content(*args, **kwargs): + scanned = json.loads(kwargs["data"]["ServiceParameters"]).get("content", "") + return blocked if "结构化字段里的违规内容" in scanned else clean + + with patch.object(g.async_handler, "post", new_callable=AsyncMock, side_effect=block_only_structured_content): + 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 CONTENT_MODERATION_TYPE in remaining + assert tool_result.isError is True + + # --------------------------------------------------------------------------- # Streaming hook tests # --------------------------------------------------------------------------- @@ -1162,35 +1521,35 @@ class TestShouldRunPostMcpCall: class TestIterMcpContentItems: def test_plain_string_is_wrapped(self): g = _make_guardrail() - assert g._iter_mcp_content_items("hello") == ["hello"] + 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"] + 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"] + 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] + 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"] + 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"] + 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) == [] + assert g._iter_mcp_content_items(123) == () class TestReplaceToolOutputInPlace: From 11b742a4b5a25fc1572d7d4b434bacfc266ac344 Mon Sep 17 00:00:00 2001 From: splendor023 Date: Wed, 19 Aug 2026 20:36:17 +0800 Subject: [PATCH 3/5] fix(guardrails): fix structured content scanning and violation content blocking --- .../aliyun/aliyun_ai_guardrail.py | 7 +++++ .../guardrails/guardrail_hooks/aliyun/base.py | 8 +++-- .../aliyun/test_aliyun_ai_guardrail.py | 31 +++++++++++++++++-- 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py index 27ce33e4ca7..894ad3694ac 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py @@ -885,9 +885,14 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): """ if target is None: return False + from litellm.proxy._experimental.mcp_server.utils import ( + set_mcp_tool_result_structured_content, + ) + content: Final = getattr(target, "content", None) if isinstance(content, list): content[:] = blocked_content + set_mcp_tool_result_structured_content(target, None) if hasattr(target, "isError"): try: # Kept duck-typed on purpose: any MCP result shape carrying `isError` @@ -903,10 +908,12 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): result: Final = target.get("result") if isinstance(result, dict) and isinstance(result.get("content"), list): result["content"] = list(blocked_content) # mutable-ok: the replaced field must stay a JSON list + set_mcp_tool_result_structured_content(result, None) return True if isinstance(target.get("content"), list): blocked: Final = list(blocked_content) # mutable-ok: must stay a JSON list target["content"] = blocked # rebind-ok: overwrite the caller's result + set_mcp_tool_result_structured_content(target, None) return True return False diff --git a/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py b/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py index 452e8638b8a..1d1f6073511 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py @@ -26,6 +26,10 @@ class AliyunGuardrailBase: """ return (message for message in messages if message.get("role") == "user") + @staticmethod + def _iter_audited_text_messages(messages: Sequence[AllMessageValues]) -> Iterator[AllMessageValues]: + return (message for message in messages if message.get("role") in ("user", "tool")) + @staticmethod def _extract_image_url(part: object) -> str | None: """ @@ -35,7 +39,7 @@ class AliyunGuardrailBase: Returns: The URL string, or None when the part carries no image URL """ - if not isinstance(part, dict) or part.get("type") != "image_url": + if not isinstance(part, dict) or part.get("type") not in ("image_url", "input_image"): return None image_url: Final = part.get("image_url") if isinstance(image_url, dict): @@ -62,7 +66,7 @@ class AliyunGuardrailBase: ) user_prompt: Final = "\n".join( - convert_content_list_to_str(message) for message in self._iter_user_messages(messages) + convert_content_list_to_str(message) for message in self._iter_audited_text_messages(messages) ).strip() return user_prompt or None 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 index 3c28632cbd8..8b85b37d63a 100644 --- 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 @@ -792,8 +792,8 @@ class TestPreCallHook: { "role": "user", "content": [ - {"type": "text", "text": "结构化输入文本"}, - {"type": "image_url", "image_url": {"url": IMG_A}}, + {"type": "input_text", "text": "结构化输入文本"}, + {"type": "input_image", "image_url": IMG_A, "detail": "auto"}, ], } ] @@ -810,6 +810,32 @@ class TestPreCallHook: assert any("结构化输入文本" in (sp.get("content") or "") for sp in sent) assert any(IMG_A in (sp.get("imageUrls") or []) for sp in sent) + @pytest.mark.asyncio + async def test_scans_responses_api_function_call_output(self): + g = _make_guardrail(level="medium") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + data = { + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": "工具输出里的违规内容", + } + ] + } + 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=data, + call_type="responses", + ) + + scanned = "".join( + json.loads(call.kwargs["data"]["ServiceParameters"]).get("content", "") for call in mock_post.call_args_list + ) + assert "工具输出里的违规内容" in scanned + @pytest.mark.asyncio async def test_blocks_violation(self): g = _make_guardrail(level="medium") @@ -1350,6 +1376,7 @@ class TestExtractMcpToolText: remaining = " ".join(getattr(item, "text", "") for item in tool_result.content) assert CONTENT_MODERATION_TYPE in remaining assert tool_result.isError is True + assert tool_result.structuredContent is None # --------------------------------------------------------------------------- From 1c58fb525e70dcb14f570bd7a4fecff7890b0324 Mon Sep 17 00:00:00 2001 From: splendor023 Date: Thu, 20 Aug 2026 14:42:47 +0800 Subject: [PATCH 4/5] fix(guardrails): optimize streaming checks and function-call argument scanning --- .../aliyun/aliyun_ai_guardrail.py | 52 +++++++-- .../aliyun/test_aliyun_ai_guardrail.py | 110 ++++++++++++++++++ 2 files changed, 150 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py index 894ad3694ac..7a3742d50e7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/aliyun_ai_guardrail.py @@ -587,6 +587,18 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): return ((None, image_urls),) return () + @staticmethod + def _iter_responses_function_call_arguments(data: Mapping[str, object]) -> Iterator[str]: + input_value: Final = data.get("input") + if not isinstance(input_value, list): + return + for item in input_value: + if not isinstance(item, dict) or item.get("type") != "function_call": + continue + arguments = item.get("arguments") + if isinstance(arguments, str) and arguments: + yield arguments + @log_guardrail_information async def async_pre_call_hook( self, @@ -610,10 +622,14 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): new_messages: Final = cast( # cast-ok: the upstream iterator is typed as plain dicts "Sequence[AllMessageValues]", tuple(_iter_inspection_messages(data)) ) - if not new_messages: + function_call_arguments: Final = tuple(self._iter_responses_function_call_arguments(data)) + if not new_messages and not function_call_arguments: verbose_proxy_logger.warning("Aliyun AI Guardrail: not running guardrail. No messages in data") return data - user_prompt: Final = self.get_user_prompt(new_messages) + message_prompt: Final = self.get_user_prompt(new_messages) + user_prompt: Final = ( + "\n".join(text for text in (message_prompt, *function_call_arguments) if text).strip() or None + ) image_urls: Final = 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") @@ -649,8 +665,8 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): messages: Final = data.get("messages", ()) content: Final = messages[0].get("content", "") if messages else "" verbose_proxy_logger.info( - "Aliyun AI Guardrail: ★ MCP pre-call check started, content: %s", - content, + "Aliyun AI Guardrail: ★ MCP pre-call check started, content length: %d", + len(content), ) if not content: return @@ -1065,6 +1081,18 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): """Normalise an HTTPException detail into a JSON object.""" return detail if isinstance(detail, dict) else {"message": str(detail)} # mutable-ok: JSON payload + def _stream_check_windows(self, text: str, last_check_position: int) -> tuple[str, ...]: + current_length: Final = len(text) + window_size: Final = self.stream_window_size + latest_window_start: Final = max(0, current_length - window_size) + if current_length - last_check_position <= window_size: + return (text[latest_window_start:],) + step: Final = max(1, min(self.stream_slide_step, window_size)) + overlap: Final = window_size - step + first_window_start: Final = max(0, last_check_position - overlap) + window_starts: Final = (*range(first_window_start, latest_window_start, step), latest_window_start) + return tuple(text[start : min(start + window_size, current_length)] for start in window_starts) + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -1106,10 +1134,10 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): 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") + check_windows = self._stream_check_windows(accumulated_text, last_check_position) + for text_to_check in check_windows: + 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 ) @@ -1120,10 +1148,10 @@ class AliyunAIGuardrail(AliyunGuardrailBase, CustomGuardrail): is_first_check = False # Stream ended - check any remaining unchecked content with a final window if len(accumulated_text) > last_check_position: - final_start: Final = max(0, len(accumulated_text) - self.stream_window_size) - remaining_text: Final = accumulated_text[final_start:] - final_response: Final = await self.async_make_request(text=remaining_text, service_type="output") - self._parse_response_and_check(final_response, check_type="output") + final_windows: Final = self._stream_check_windows(accumulated_text, last_check_position) + for remaining_text in final_windows: + final_response = await self.async_make_request(text=remaining_text, service_type="output") + self._parse_response_and_check(final_response, check_type="output") verbose_proxy_logger.info( "Aliyun AI Guardrail: Streaming scan completed, total length: %d", len(accumulated_text) ) 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 index 8b85b37d63a..13367e75296 100644 --- 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 @@ -836,6 +836,33 @@ class TestPreCallHook: ) assert "工具输出里的违规内容" in scanned + @pytest.mark.asyncio + async def test_scans_responses_api_function_call_arguments(self): + g = _make_guardrail(level="medium") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + data = { + "input": [ + { + "type": "function_call", + "call_id": "call_1", + "name": "send_message", + "arguments": '{"text":"函数调用里的违规参数"}', + } + ] + } + 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=data, + call_type="responses", + ) + + scanned = "".join( + json.loads(call.kwargs["data"]["ServiceParameters"]).get("content", "") for call in mock_post.call_args_list + ) + assert "函数调用里的违规参数" in scanned + @pytest.mark.asyncio async def test_blocks_violation(self): g = _make_guardrail(level="medium") @@ -1455,6 +1482,29 @@ class TestMcpPreCallCheck: assert service == "text_img_mix_guard" assert service_parameters["content"] == "send_message hello" + @pytest.mark.asyncio + async def test_logs_mcp_content_length_without_content(self): + g = _make_guardrail(level="medium") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + sensitive_content = "send_message token=secret-value" + logger_path = "litellm.proxy.guardrails.guardrail_hooks.aliyun.aliyun_ai_guardrail.verbose_proxy_logger.info" + with ( + patch(logger_path) as mock_info, + patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean), + ): + await g.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test"), + cache=MagicMock(), + data={"messages": [{"role": "user", "content": sensitive_content}]}, + call_type="call_mcp_tool", + ) + + mock_info.assert_any_call( + "Aliyun AI Guardrail: ★ MCP pre-call check started, content length: %d", + len(sensitive_content), + ) + assert sensitive_content not in str(mock_info.call_args_list) + @pytest.mark.asyncio async def test_empty_content_skips_check(self): g = _make_guardrail() @@ -1686,6 +1736,66 @@ class TestStreamingHook: assert "普通流式内容" in scanned assert emitted == [chunk] + @pytest.mark.asyncio + async def test_blocks_violation_in_oversized_delta_prefix(self): + g = _make_guardrail( + level="medium", + stream_window_size=10, + stream_first_check_step=1, + stream_slide_step=6, + ) + 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")], + ) + chunk = _make_stream_chunk(content="违规前缀" + "正常内容" * 8) + + async def block_prefix(*args, **kwargs): + scanned = json.loads(kwargs["data"]["ServiceParameters"])["content"] + return blocked if "违规前缀" in scanned else clean + + with patch.object(g.async_handler, "post", new_callable=AsyncMock, side_effect=block_prefix): + 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_all_windows_before_emitting_oversized_delta_once(self): + g = _make_guardrail( + level="medium", + stream_window_size=10, + stream_first_check_step=1, + stream_slide_step=6, + ) + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + chunk = _make_stream_chunk(content="开头内容" + "中间内容" * 6 + "结尾内容") + + 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_windows = tuple( + json.loads(call.kwargs["data"]["ServiceParameters"])["content"] for call in mock_post.call_args_list + ) + assert all(len(window) <= g.stream_window_size for window in scanned_windows) + assert any("开头内容" in window for window in scanned_windows) + assert any("结尾内容" in window for window in scanned_windows) + assert emitted == [chunk] + @pytest.mark.asyncio async def test_scans_responses_api_text_delta(self): from litellm.types.llms.openai import OutputTextDeltaEvent From 9e333e58b27e6d69a58f1fc52b8bb999ef6f5aee Mon Sep 17 00:00:00 2001 From: splendor023 Date: Thu, 20 Aug 2026 15:37:15 +0800 Subject: [PATCH 5/5] fix(guardrails): support extracting public image URLs from tool messages. --- .../guardrails/guardrail_hooks/aliyun/base.py | 6 ++--- .../aliyun/test_aliyun_ai_guardrail.py | 26 +++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py b/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py index 1d1f6073511..b50884b3de2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py @@ -71,8 +71,8 @@ class AliyunGuardrailBase: return user_prompt or None def _iter_public_image_urls(self, messages: Sequence[AllMessageValues]) -> Iterator[str]: - """Yield the publicly reachable image URLs of every user message, in order.""" - for content in (message.get("content") for message in self._iter_user_messages(messages)): + """Yield publicly reachable image URLs from user and tool messages.""" + for content in (message.get("content") for message in self._iter_audited_text_messages(messages)): if not isinstance(content, list): continue for url in (self._extract_image_url(part) for part in content): @@ -83,7 +83,7 @@ class AliyunGuardrailBase: def get_image_urls(self, messages: Sequence[AllMessageValues]) -> tuple[str, ...]: """ - Extract image URLs from every user message in the request. + Extract image URLs from every user and tool message in the request. Only publicly accessible http(s) URLs are collected (in order, de-duplicated). Uses the same message range as ``get_user_prompt``. Example: 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 index 13367e75296..cc27f4cbf08 100644 --- 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 @@ -836,6 +836,32 @@ class TestPreCallHook: ) assert "工具输出里的违规内容" in scanned + @pytest.mark.asyncio + async def test_scans_responses_api_function_call_output_image(self): + g = _make_guardrail(level="medium") + clean = _make_aliyun_api_response(suggestion="pass", detail=[]) + data = { + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": [ + {"type": "input_image", "image_url": IMG_A, "detail": "auto"}, + ], + } + ] + } + 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=data, + call_type="responses", + ) + + sent = [json.loads(call.kwargs["data"]["ServiceParameters"]) for call in mock_post.call_args_list] + assert any(IMG_A in (service_parameters.get("imageUrls") or []) for service_parameters in sent) + @pytest.mark.asyncio async def test_scans_responses_api_function_call_arguments(self): g = _make_guardrail(level="medium")