scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(guardrails): support extracting public image URLs from tool messages.

Browse source
This commit is contained in:
splendor023 2026-08-20 15:37:15 +08:00
parent 1c58fb525e
commit 9e333e58b2
2 changed files with 29 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/guardrails/guardrail_hooks/aliyun/base.py
View file

@ -71,8 +71,8 @@ class AliyunGuardrailBase:
return user_prompt or None
def _iter_public_image_urls(self, messages: Sequence[AllMessageValues]) -> Iterator[str]:
"""Yield the publicly reachable image URLs of every user message, in order."""
for content in (message.get("content") for message in self._iter_user_messages(messages)):
"""Yield publicly reachable image URLs from user and tool messages."""
for content in (message.get("content") for message in self._iter_audited_text_messages(messages)):
if not isinstance(content, list):
continue
for url in (self._extract_image_url(part) for part in content):
@ -83,7 +83,7 @@ class AliyunGuardrailBase:
def get_image_urls(self, messages: Sequence[AllMessageValues]) -> tuple[str, ...]:
"""
Extract image URLs from every user message in the request.
Extract image URLs from every user and tool message in the request.
Only publicly accessible http(s) URLs are collected (in order,
de-duplicated). Uses the same message range as ``get_user_prompt``.
Example:

26
tests/test_litellm/proxy/guardrails/guardrail_hooks/aliyun/test_aliyun_ai_guardrail.py
View file

@ -836,6 +836,32 @@ class TestPreCallHook:
)
assert "工具输出里的违规内容" in scanned
@pytest.mark.asyncio
async def test_scans_responses_api_function_call_output_image(self):
g = _make_guardrail(level="medium")
clean = _make_aliyun_api_response(suggestion="pass", detail=[])
data = {
"input": [
{
"type": "function_call_output",
"call_id": "call_1",
"output": [
{"type": "input_image", "image_url": IMG_A, "detail": "auto"},
],
}
]
}
with patch.object(g.async_handler, "post", new_callable=AsyncMock, return_value=clean) as mock_post:
await g.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test"),
cache=MagicMock(),
data=data,
call_type="responses",
)
sent = [json.loads(call.kwargs["data"]["ServiceParameters"]) for call in mock_post.call_args_list]
assert any(IMG_A in (service_parameters.get("imageUrls") or []) for service_parameters in sent)
@pytest.mark.asyncio
async def test_scans_responses_api_function_call_arguments(self):
g = _make_guardrail(level="medium")