added tests

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

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

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"