mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-17 04:17:43 +00:00
294 lines
9.6 KiB
Python
294 lines
9.6 KiB
Python
"""
|
|
Tests that guardrails (post_call_success_hook) fire for image generation requests.
|
|
|
|
The /images/generations endpoint in proxy/image_endpoints/endpoints.py calls
|
|
proxy_logging_obj.post_call_success_hook after a successful image generation.
|
|
These tests verify:
|
|
1. CustomGuardrail.async_post_call_success_hook is invoked for image generation.
|
|
2. A guardrail can inspect and transform the image response.
|
|
3. A guardrail that raises blocks the response (exception propagates).
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from typing import Any, Optional
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(0, os.path.abspath("../../../.."))
|
|
|
|
import litellm
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
from litellm.types.utils import ImageObject, ImageResponse
|
|
|
|
|
|
def _make_image_response(**kwargs) -> ImageResponse:
|
|
"""Helper to build a minimal ImageResponse for tests."""
|
|
return ImageResponse(
|
|
data=[ImageObject(url="https://example.com/img.png")],
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 1. Hook is invoked for image generation responses
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TrackingGuardrail(CustomGuardrail):
|
|
"""Guardrail that records whether it was called and with what args."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="tracking_guardrail",
|
|
default_on=True,
|
|
event_hook=GuardrailEventHooks.post_call,
|
|
)
|
|
self.called = False
|
|
self.received_data: Optional[dict] = None
|
|
self.received_response: Optional[Any] = None
|
|
|
|
async def async_post_call_success_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: Any,
|
|
) -> Any:
|
|
self.called = True
|
|
self.received_data = data
|
|
self.received_response = response
|
|
return response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_call_success_hook_invoked_for_image_generation():
|
|
"""
|
|
Verify that a default-on guardrail's async_post_call_success_hook is
|
|
called when ProxyLogging.post_call_success_hook is invoked with an
|
|
ImageResponse (the same path used by the /images/generations endpoint).
|
|
"""
|
|
guardrail = TrackingGuardrail()
|
|
image_response = _make_image_response()
|
|
|
|
with patch("litellm.callbacks", [guardrail]):
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
data = {"model": "dall-e-3", "prompt": "A sunset over mountains"}
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
result = await proxy_logging.post_call_success_hook(
|
|
data=data,
|
|
response=image_response,
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
assert guardrail.called is True, "Guardrail hook was not invoked for image generation"
|
|
assert guardrail.received_data is not None
|
|
assert guardrail.received_data["model"] == "dall-e-3"
|
|
assert isinstance(guardrail.received_response, ImageResponse)
|
|
# The response should be passed through unchanged
|
|
assert result is image_response
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 2. Guardrail can transform image generation response
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TransformingGuardrail(CustomGuardrail):
|
|
"""Guardrail that replaces the image URL in the response."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="transforming_guardrail",
|
|
default_on=True,
|
|
event_hook=GuardrailEventHooks.post_call,
|
|
)
|
|
|
|
async def async_post_call_success_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: Any,
|
|
) -> Any:
|
|
# Return a modified image response (e.g., watermarked URL)
|
|
return ImageResponse(
|
|
data=[ImageObject(url="https://example.com/watermarked.png")],
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_guardrail_can_transform_image_response():
|
|
"""
|
|
Verify that a guardrail can replace the ImageResponse returned to the client.
|
|
"""
|
|
guardrail = TransformingGuardrail()
|
|
original_response = _make_image_response()
|
|
|
|
with patch("litellm.callbacks", [guardrail]):
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
data = {"model": "dall-e-3", "prompt": "A sunset"}
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
result = await proxy_logging.post_call_success_hook(
|
|
data=data,
|
|
response=original_response,
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
assert result is not original_response
|
|
assert isinstance(result, ImageResponse)
|
|
assert result.data[0].url == "https://example.com/watermarked.png"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 3. Guardrail that raises blocks the image response
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class BlockingGuardrail(CustomGuardrail):
|
|
"""Guardrail that raises on unsafe image prompts."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="blocking_guardrail",
|
|
default_on=True,
|
|
event_hook=GuardrailEventHooks.post_call,
|
|
)
|
|
|
|
async def async_post_call_success_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: Any,
|
|
) -> Any:
|
|
raise ValueError("Image content blocked by guardrail")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_guardrail_exception_propagates_for_image_generation():
|
|
"""
|
|
Verify that an exception raised in a guardrail's post_call_success_hook
|
|
propagates up (the proxy endpoint wraps this in an error response).
|
|
"""
|
|
guardrail = BlockingGuardrail()
|
|
|
|
with patch("litellm.callbacks", [guardrail]):
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
data = {"model": "dall-e-3", "prompt": "Something unsafe"}
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
with pytest.raises(ValueError, match="Image content blocked by guardrail"):
|
|
await proxy_logging.post_call_success_hook(
|
|
data=data,
|
|
response=_make_image_response(),
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 4. Non-guardrail CustomLogger also fires for image generation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TrackingLogger(CustomLogger):
|
|
"""Plain CustomLogger (not a guardrail) that tracks invocations."""
|
|
|
|
def __init__(self):
|
|
self.called = False
|
|
self.received_response = None
|
|
|
|
async def async_post_call_success_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: Any,
|
|
) -> Any:
|
|
self.called = True
|
|
self.received_response = response
|
|
return response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_logger_post_call_success_hook_fires_for_image_generation():
|
|
"""
|
|
Verify that a plain CustomLogger (non-guardrail) callback also has its
|
|
async_post_call_success_hook invoked for image generation responses.
|
|
"""
|
|
logger = TrackingLogger()
|
|
image_response = _make_image_response()
|
|
|
|
with patch("litellm.callbacks", [logger]):
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
data = {"model": "dall-e-3", "prompt": "A cat"}
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
result = await proxy_logging.post_call_success_hook(
|
|
data=data,
|
|
response=image_response,
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
assert logger.called is True
|
|
assert isinstance(logger.received_response, ImageResponse)
|
|
assert result is image_response
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 5. Guardrail with should_run_guardrail=False is skipped
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class OptInGuardrail(CustomGuardrail):
|
|
"""Guardrail that is NOT default_on, so it only runs if explicitly requested."""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="opt_in_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.post_call,
|
|
)
|
|
self.called = False
|
|
|
|
async def async_post_call_success_hook(
|
|
self,
|
|
data: dict,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
response: Any,
|
|
) -> Any:
|
|
self.called = True
|
|
return response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_default_guardrail_skipped_for_image_generation():
|
|
"""
|
|
Verify that a guardrail with default_on=False is NOT invoked for image
|
|
generation unless the request explicitly enables it.
|
|
"""
|
|
guardrail = OptInGuardrail()
|
|
|
|
with patch("litellm.callbacks", [guardrail]):
|
|
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
|
|
|
# No guardrails key in data -> should_run_guardrail returns False
|
|
data = {"model": "dall-e-3", "prompt": "A sunset"}
|
|
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
|
|
|
await proxy_logging.post_call_success_hook(
|
|
data=data,
|
|
response=_make_image_response(),
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
assert guardrail.called is False, "Opt-in guardrail should not fire without explicit request"
|