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: