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