scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(guardrails): run unified guardrails on /v1/images/edits (#44195)

Browse source
This commit is contained in:
michelligabriele 2026-10-03 06:14:35 +02:00 • committed by GitHub
parent 5ce81631fe
commit 119179942e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 146 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/llms/openai/image_generation/guardrail_translation/__init__.py
View file

@ -10,6 +10,8 @@ from litellm.types.utils import CallTypes
guardrail_translation_mappings: Final = {
CallTypes.image_generation: OpenAIImageGenerationHandler,
CallTypes.aimage_generation: OpenAIImageGenerationHandler,
CallTypes.image_edit: OpenAIImageGenerationHandler,
CallTypes.aimage_edit: OpenAIImageGenerationHandler,
}
__all__ = ["OpenAIImageGenerationHandler", "guardrail_translation_mappings"]

1
tests/e2e/coverage_registry/guardrail.yaml
View file

@ -9,6 +9,7 @@
- {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"}
- {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"}
- {id: guardrail.litellm_content_filter.pre_call.blocks_video, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [videos], source: "test_key_guardrail_video_e2e.py", fail_before_fix: proven, rationale: "A content-filter guardrail attached to a key (metadata.guardrails) blocks a banned prompt on POST /v1/videos before the provider is called; before the fix the route's call type was unknown to the unified guardrail hook and the prompt went to the provider unscanned (LIT-6685)"}
- {id: guardrail.litellm_content_filter.pre_call.blocks_image_edit, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [images_edits], source: "test_key_guardrail_image_edit_e2e.py", rationale: "A content-filter guardrail attached to a key (metadata.guardrails) blocks a banned prompt on POST /v1/images/edits before the provider is called; before the fix aimage_edit had no guardrail translation mapping and the prompt went to the provider unscanned"}
- {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"}
- {id: guardrail.litellm_content_filter.pre_call.returns_guardrail_information, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "guardrails/test_guardrail_information_response_e2e.py", rationale: "Opt-in chat responses expose successful guardrail execution details"}
- {id: guardrail.litellm_content_filter.apply_endpoint.blocks, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail blocks banned content for customers that call the apply surface directly"}

17
tests/e2e/guardrails/guardrails_client.py
View file

@ -9,7 +9,7 @@ from collections.abc import Callable
from dataclasses import dataclass
from typing import Final, Literal
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, SLOW_PROVIDER_TIMEOUT_SECONDS, settle_propagation, unique_marker
from e2e_http import NoBody, Result, StreamingResponse, Success, unwrap
from lifecycle import ResourceManager
from models import (
@ -20,6 +20,8 @@ from models import (
ChatMetadata,
ChatResponse,
ChatTool,
ImageEditForm,
ImageGenerationResponse,
KeyGenerateBody,
KeyMetadata,
LiteLLMParamsBody,
@ -364,6 +366,19 @@ class GuardrailsClient:
response_type=VideoCreateResponse,
)
def edit_image(self, key: str, model: str, prompt: str, image: bytes) -> Result[ImageGenerationResponse]:
return self.proxy.transport.upload(
"/v1/images/edits",
headers=self.proxy.transport.bearer(key),
form=ImageEditForm(model=model, prompt=prompt),
filename="image.png",
content=image,
file_content_type="image/png",
file_field="image",
response_type=ImageGenerationResponse,
timeout=SLOW_PROVIDER_TIMEOUT_SECONDS,
)
def chat(
self,
key: str,

72
tests/e2e/guardrails/test_key_guardrail_image_edit_e2e.py Normal file
View file

@ -0,0 +1,72 @@
from __future__ import annotations
import base64
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import Success, UnknownApiError
from guardrails_client import GuardrailsClient, poll_until_blocked
from lifecycle import ResourceManager
from models import LiteLLMParamsBody
pytestmark = pytest.mark.e2e
CHAT_MODEL: Final = "gemini-2.5-flash"
IMAGE_BACKEND: Final = "openai/gpt-image-2.5-flare"
SOURCE_PNG: Final = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAEAAAABACAIAAAAlC+aJAAAAS0lEQVR42u3PMQ0AAAwDoPo3"
"3UrYvQQckD4XAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEB"
"AYHLAMpT0sIcNbcEAAAAAElFTkSuQmCC"
)
def _edit_prompt_with(banned_keyword: str) -> str:
return f"Turn this into a watercolor painting of a lighthouse. {banned_keyword}"
def _create_image_model(client: GuardrailsClient, resources: ResourceManager) -> str:
model_name = f"e2e-guard-image-edit-{unique_marker()}"
model_id = client.proxy.create_model(
model_name,
LiteLLMParamsBody(model=IMAGE_BACKEND, api_key="os.environ/OPENAI_API_KEY"),
provider_live=True,
)
resources.defer(lambda: client.proxy.delete_model(model_id))
return model_name
class TestKeyAttachedGuardrailOnImageEdits:
@pytest.mark.covers(
"guardrail.litellm_content_filter.pre_call.blocks_image_edit",
exercised_on=["images_edits"],
)
def test_key_attached_content_filter_blocks_banned_image_edit_prompt(
self, client: GuardrailsClient, resources: ResourceManager
) -> None:
banned = unique_marker()
guardrail_name = f"e2e-image-edit-filter-{banned}"
guardrail_id = client.create_content_filter_guardrail(guardrail_name, banned, default_on=False)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
key = client.create_key_with_guardrails(resources, [guardrail_name])
model = _create_image_model(client, resources)
synced = poll_until_blocked(lambda: client.chat(key, CHAT_MODEL, _edit_prompt_with(banned)))
assert isinstance(synced, UnknownApiError) and synced.status_code == 400, (
f"key guardrail {guardrail_name!r} never synced to the proxy on /chat/completions: {synced}"
)
result = client.edit_image(key, model, _edit_prompt_with(banned), SOURCE_PNG)
match result:
case UnknownApiError(status_code=status, body=body):
assert status == 400, f"expected a 400 guardrail block, got {status}: {body[:300]}"
assert "content blocked" in body.lower() or banned in body, (
f"block response missing content-filter reason: {body[:300]}"
)
case Success():
pytest.fail(
f"key-attached guardrail {guardrail_name!r} was skipped on /v1/images/edits: "
"the banned prompt reached the provider and an edited image came back"
)
case _:
pytest.fail(f"unexpected /v1/images/edits outcome for a banned prompt: {result}")

55
tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -1,5 +1,6 @@
"""Tests for unified guardrail."""
import io
import logging
from types import SimpleNamespace
from typing import TYPE_CHECKING, Final, Literal
@ -373,6 +374,37 @@ class TestUnifiedLLMGuardrails:
assert result["prompt"] == "a paper boat on a stream [GUARDRAILED]"
assert result["seconds"] == "4"
@pytest.mark.asyncio
@pytest.mark.parametrize("call_type", ["aimage_edit", "image_edit"])
async def test_image_edit_routes_scan_prompt_and_keep_rewrite(self, monkeypatch, call_type: str) -> None:
"""/v1/images/edits dispatches call_type="aimage_edit", which had no translation mapping,
so the hook returned the request unscanned. Runs against the discovered handler map."""
_patch_translation_mappings(monkeypatch, discover_guardrail_translation_mappings())
handler = UnifiedLLMGuardrails()
guardrail = RewritingGuardrail()
image = io.BytesIO(b"\x89PNG\r\n\x1a\n")
data = {
"guardrail_to_apply": guardrail,
"model": "gemini-3-pro-image",
"prompt": "a watercolor painting of a lighthouse",
"image": [image],
}
result = await handler.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
cache=DualCache(),
data=data,
call_type=call_type,
)
assert guardrail.event_history == [GuardrailEventHooks.pre_call]
assert [call["inputs"]["texts"] for call in guardrail.apply_calls] == [
["a watercolor painting of a lighthouse"]
]
assert guardrail.apply_calls[0]["inputs"]["model"] == "gemini-3-pro-image"
assert result["prompt"] == "a watercolor painting of a lighthouse [GUARDRAILED]"
assert result["image"] == [image]
class TestAsyncModerationHook:
@pytest.mark.asyncio
async def test_uses_mcp_event_type(self):
@ -419,6 +451,29 @@ class TestUnifiedLLMGuardrails:
assert guardrail.event_history == [GuardrailEventHooks.during_call]
@pytest.mark.asyncio
async def test_runs_for_image_edits(self, monkeypatch) -> None:
_patch_translation_mappings(monkeypatch, discover_guardrail_translation_mappings())
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
data = {
"guardrail_to_apply": guardrail,
"model": "gemini-3-pro-image",
"prompt": "a watercolor painting of a lighthouse",
"image": [io.BytesIO(b"\x89PNG\r\n\x1a\n")],
}
await handler.async_moderation_hook(
data=data,
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
call_type=CallTypes.aimage_edit.value,
)
assert guardrail.event_history == [GuardrailEventHooks.during_call]
assert [call["inputs"]["texts"] for call in guardrail.apply_calls] == [
["a watercolor painting of a lighthouse"]
]
class TestAsyncPostCallStreamingIteratorHook:
@pytest.mark.asyncio
async def test_streaming_content_not_lost_on_sampled_chunks(self, monkeypatch):