diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index 93b5be6b1b..5871bd5d61 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -1,11 +1,15 @@ import os from typing import TYPE_CHECKING, Any, Optional, Union +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.trace import SpanKind + from litellm._logging import verbose_logger from litellm.integrations.arize import _utils from litellm.integrations.arize._utils import ArizeOTELAttributes -from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig from litellm.integrations.opentelemetry import OpenTelemetry +from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute +from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -46,9 +50,6 @@ class ArizePhoenixLogger(OpenTelemetry): By creating our own provider we guarantee Arize Phoenix always gets its own exporter pipeline, regardless of initialisation order. """ - from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.trace import SpanKind - if tracer_provider is not None: # Explicitly supplied (e.g. in tests) — honour it. self.tracer = tracer_provider.get_tracer("litellm") @@ -81,8 +82,6 @@ class ArizePhoenixLogger(OpenTelemetry): @staticmethod def set_arize_phoenix_attributes(span: Span, kwargs, response_obj): - from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute - _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) # Dynamic project name: check metadata first, then fall back to env var config diff --git a/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py b/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py index 40b814e69e..a0dcf1c091 100644 --- a/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py +++ b/tests/test_litellm/integrations/arize/test_arize_otel_coexistence.py @@ -8,11 +8,9 @@ Covers the three root-cause fixes: 3. Arize loggers do NOT overwrite ``proxy_server.open_telemetry_logger``. """ -import os import unittest from unittest.mock import patch -import pytest from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter diff --git a/tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py b/tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py new file mode 100644 index 0000000000..4a5d901b74 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_image_generation_guardrails.py @@ -0,0 +1,293 @@ +""" +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"