litellm/tests/test_litellm/integrations/arize/test_arize_phoenix.py

import unittest
from unittest.mock import MagicMock, patch

import pytest

from litellm.integrations.arize.arize_phoenix import (
    ArizePhoenixConfig,
    ArizePhoenixLogger,
)
from litellm.integrations.arize._utils import ArizeOTELAttributes


class TestArizePhoenixConfig(unittest.TestCase):
    @patch.dict(
        "os.environ",
        {
            "PHOENIX_API_KEY": "test_api_key",
            "PHOENIX_COLLECTOR_HTTP_ENDPOINT": "http://test.endpoint",
        },
    )
    def test_get_arize_phoenix_config_http(self):
        # Call the function to get the configuration
        config = ArizePhoenixLogger.get_arize_phoenix_config()

        # Verify the configuration - now uses standard Authorization Bearer format
        self.assertEqual(
            config.otlp_auth_headers, "Authorization=Bearer test_api_key"
        )
        self.assertEqual(config.endpoint, "http://test.endpoint/v1/traces")
        self.assertEqual(config.protocol, "otlp_http")

    @patch.dict(
        "os.environ",
        {
            "PHOENIX_API_KEY": "test_api_key",
            "PHOENIX_COLLECTOR_ENDPOINT": "grpc://test.endpoint",
        },
    )
    def test_get_arize_phoenix_config_grpc(self):
        # Call the function to get the configuration
        config = ArizePhoenixLogger.get_arize_phoenix_config()

        # Verify the configuration - now uses standard Authorization Bearer format
        self.assertEqual(
            config.otlp_auth_headers, "Authorization=Bearer test_api_key"
        )
        self.assertEqual(config.endpoint, "grpc://test.endpoint")
        self.assertEqual(config.protocol, "otlp_grpc")

    @patch.dict(
        "os.environ",
        {
            "PHOENIX_API_KEY": "test_api_key",
            "PHOENIX_COLLECTOR_ENDPOINT": "http://localhost:6006",
        },
    )
    def test_get_arize_phoenix_config_http_local(self):
        # Test with local Phoenix instance
        config = ArizePhoenixLogger.get_arize_phoenix_config()

        # Should automatically append /v1/traces to local endpoint
        self.assertEqual(
            config.otlp_auth_headers, "Authorization=Bearer test_api_key"
        )
        self.assertEqual(config.endpoint, "http://localhost:6006/v1/traces")
        self.assertEqual(config.protocol, "otlp_http")

    @patch.dict(
        "os.environ",
        {
            "PHOENIX_COLLECTOR_ENDPOINT": "http://localhost:4317",
        },
        clear=True
    )
    def test_get_arize_phoenix_config_grpc_no_api_key(self):
        # Test gRPC endpoint detection and no API key (for local development)
        config = ArizePhoenixLogger.get_arize_phoenix_config()

        # No API key should be fine for local development
        self.assertIsNone(config.otlp_auth_headers)
        self.assertEqual(config.endpoint, "http://localhost:4317")
        self.assertEqual(config.protocol, "otlp_grpc")

    @patch.dict("os.environ", {}, clear=True)
    def test_get_arize_phoenix_config_defaults_to_local(self):
        # Test that it defaults to local Phoenix when no config is provided
        config = ArizePhoenixLogger.get_arize_phoenix_config()

        # Should default to localhost
        self.assertEqual(config.endpoint, "http://localhost:6006/v1/traces")
        self.assertEqual(config.protocol, "otlp_http")
        # No auth headers when no API key is provided for local instance
        self.assertIsNone(config.otlp_auth_headers)



@pytest.mark.parametrize(
    "env_vars, expected_headers, expected_endpoint, expected_protocol",
    [
        pytest.param(
            {"PHOENIX_API_KEY": "test_api_key"},
            "Authorization=Bearer test_api_key",
            "http://localhost:6006/v1/traces",
            "otlp_http",
            id="default to http protocol and self-hosted Phoenix endpoint",
        ),
        pytest.param(
            {"PHOENIX_COLLECTOR_HTTP_ENDPOINT": "", "PHOENIX_API_KEY": "test_api_key"},
            "Authorization=Bearer test_api_key",
            "http://localhost:6006/v1/traces",
            "otlp_http",
            id="empty string/unset endpoint will default to http protocol and self-hosted Phoenix endpoint",
        ),
        pytest.param(
            {"PHOENIX_COLLECTOR_HTTP_ENDPOINT": "http://localhost:4318", "PHOENIX_COLLECTOR_ENDPOINT": "http://localhost:4317", "PHOENIX_API_KEY": "test_api_key"},
            "Authorization=Bearer test_api_key",
            "http://localhost:4318/v1/traces",
            "otlp_http",
            id="prioritize http if both endpoints are set",
        ),
        pytest.param(
            {"PHOENIX_COLLECTOR_ENDPOINT": "https://localhost:6006", "PHOENIX_API_KEY": "test_api_key"},
            "Authorization=Bearer test_api_key",
            "https://localhost:6006/v1/traces",
            "otlp_http",
            id="custom https endpoint treated as http",
        ),
        pytest.param(
            {"PHOENIX_COLLECTOR_ENDPOINT": "https://localhost:6006"},
            None,
            "https://localhost:6006/v1/traces",
            "otlp_http",
            id="custom https endpoint with no auth treated as http",
        ),
        pytest.param(
            {"PHOENIX_COLLECTOR_ENDPOINT": "grpc://localhost:6006", "PHOENIX_API_KEY": "test_api_key"},
            "Authorization=Bearer test_api_key",
            "grpc://localhost:6006",
            "otlp_grpc",
            id="explicit grpc endpoint with grpc:// prefix",
        ),
        pytest.param(
            {"PHOENIX_COLLECTOR_ENDPOINT": "http://localhost:4317"},
            None,
            "http://localhost:4317",
            "otlp_grpc",
            id="grpc endpoint with standard grpc port 4317",
        ),
        pytest.param(
            {"PHOENIX_COLLECTOR_HTTP_ENDPOINT": "https://localhost:6006", "PHOENIX_API_KEY": "test_api_key"},
            "Authorization=Bearer test_api_key",
            "https://localhost:6006/v1/traces",
            "otlp_http",
            id="custom http endpoint",
        ),
    ],
)
def test_get_arize_phoenix_config(monkeypatch, env_vars, expected_headers, expected_endpoint, expected_protocol):
    # Clear all Phoenix-related env vars first to ensure clean state
    for key in ["PHOENIX_API_KEY", "PHOENIX_COLLECTOR_ENDPOINT", "PHOENIX_COLLECTOR_HTTP_ENDPOINT"]:
        monkeypatch.delenv(key, raising=False)
    
    for key, value in env_vars.items():
        monkeypatch.setenv(key, value)

    config = ArizePhoenixLogger.get_arize_phoenix_config()

    assert isinstance(config, ArizePhoenixConfig)
    assert config.otlp_auth_headers == expected_headers
    assert config.endpoint == expected_endpoint
    assert config.protocol == expected_protocol

@pytest.mark.parametrize(
    "env_vars",
    [
        pytest.param(
            {"PHOENIX_COLLECTOR_ENDPOINT": "https://app.phoenix.arize.com/v1/traces"},
            id="missing api_key with explicit Arize Phoenix Cloud endpoint"
        ),
        pytest.param(
            {"PHOENIX_COLLECTOR_HTTP_ENDPOINT": "https://app.phoenix.arize.com/v1/traces"},
            id="missing api_key with HTTP Arize Phoenix Cloud endpoint"
        ),
    ],
)
def test_get_arize_phoenix_config_expection_on_missing_api_key(monkeypatch, env_vars):
    # Clear all Phoenix-related env vars first to ensure clean state
    for key in ["PHOENIX_API_KEY", "PHOENIX_COLLECTOR_ENDPOINT", "PHOENIX_COLLECTOR_HTTP_ENDPOINT"]:
        monkeypatch.delenv(key, raising=False)
    
    for key, value in env_vars.items():
        monkeypatch.setenv(key, value)

    with pytest.raises(ValueError, match="PHOENIX_API_KEY must be set when using Phoenix Cloud"):
        ArizePhoenixLogger.get_arize_phoenix_config()



# ---------------------------------------------------------------------------
# Dynamic project naming from metadata
# ---------------------------------------------------------------------------


class TestGetDynamicProjectName:
    """Tests for _get_dynamic_project_name extraction logic."""

    def test_extracts_from_standard_logging_object_metadata(self):
        kwargs = {
            "standard_logging_object": {
                "metadata": {"phoenix_project_name": "my-project"},
            }
        }
        assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) == "my-project"

    def test_extracts_from_litellm_params_metadata(self):
        kwargs = {
            "litellm_params": {
                "metadata": {"phoenix_project_name": "sdk-project"},
            }
        }
        assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) == "sdk-project"

    def test_returns_none_when_no_metadata(self):
        assert ArizePhoenixLogger._get_dynamic_project_name({}) is None

    def test_non_dict_standard_logging_object_does_not_raise(self):
        """isinstance(dict) guard prevents AttributeError on non-dict payloads."""
        kwargs = {"standard_logging_object": "not-a-dict"}
        assert ArizePhoenixLogger._get_dynamic_project_name(kwargs) is None


class TestDynamicProjectNameOnSpan:
    """set_arize_phoenix_attributes sets openinference.project.name on the span."""

    @patch.dict("os.environ", {"PHOENIX_PROJECT_NAME": "env-fallback"}, clear=False)
    @patch("litellm.integrations.arize._utils.set_attributes")
    def test_dynamic_name_sets_span_attribute(self, _mock_set_attrs):
        span = MagicMock()
        kwargs = {
            "standard_logging_object": {
                "metadata": {"phoenix_project_name": "dynamic-proj"},
            }
        }
        ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj=None)

        span.set_attribute.assert_called_once_with("openinference.project.name", "dynamic-proj")

    @patch.dict("os.environ", {"PHOENIX_PROJECT_NAME": "env-project"}, clear=False)
    @patch("litellm.integrations.arize._utils.set_attributes")
    def test_falls_back_to_env_var_when_no_dynamic_name(self, _mock_set_attrs):
        span = MagicMock()
        ArizePhoenixLogger.set_arize_phoenix_attributes(span, {}, response_obj=None)

        span.set_attribute.assert_called_once_with("openinference.project.name", "env-project")


if __name__ == "__main__":
    unittest.main()