forked from phoenix/litellm-mirror
b75019c1a5
* arize use helper for get_arize_opentelemetry_config * use helper to get Arize OTEL config * arize add helpers for arize * docs allow using arize http endpoint * fix importing OTEL for Arize * use static methods for ArizeLogger * fix ArizeLogger tests
88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
import time
|
|
|
|
import pytest
|
|
from dotenv import load_dotenv
|
|
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
|
|
|
import litellm
|
|
from litellm._logging import verbose_logger, verbose_proxy_logger
|
|
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
|
|
from litellm.integrations.arize_ai import ArizeConfig, ArizeLogger
|
|
|
|
load_dotenv()
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_async_otel_callback():
|
|
litellm.set_verbose = True
|
|
|
|
verbose_proxy_logger.setLevel(logging.DEBUG)
|
|
verbose_logger.setLevel(logging.DEBUG)
|
|
litellm.success_callback = ["arize"]
|
|
|
|
await litellm.acompletion(
|
|
model="gpt-3.5-turbo",
|
|
messages=[{"role": "user", "content": "hi test from local arize"}],
|
|
mock_response="hello",
|
|
temperature=0.1,
|
|
user="OTEL_USER",
|
|
)
|
|
|
|
await asyncio.sleep(2)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_env_vars(monkeypatch):
|
|
monkeypatch.setenv("ARIZE_SPACE_KEY", "test_space_key")
|
|
monkeypatch.setenv("ARIZE_API_KEY", "test_api_key")
|
|
|
|
|
|
def test_get_arize_config(mock_env_vars):
|
|
"""
|
|
Use Arize default endpoint when no endpoints are provided
|
|
"""
|
|
config = ArizeLogger._get_arize_config()
|
|
assert isinstance(config, ArizeConfig)
|
|
assert config.space_key == "test_space_key"
|
|
assert config.api_key == "test_api_key"
|
|
assert config.grpc_endpoint == "https://otlp.arize.com/v1"
|
|
assert config.http_endpoint is None
|
|
|
|
|
|
def test_get_arize_config_with_endpoints(mock_env_vars, monkeypatch):
|
|
"""
|
|
Use provided endpoints when they are set
|
|
"""
|
|
monkeypatch.setenv("ARIZE_ENDPOINT", "grpc://test.endpoint")
|
|
monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "http://test.endpoint")
|
|
|
|
config = ArizeLogger._get_arize_config()
|
|
assert config.grpc_endpoint == "grpc://test.endpoint"
|
|
assert config.http_endpoint == "http://test.endpoint"
|
|
|
|
|
|
def test_get_arize_opentelemetry_config_grpc(mock_env_vars, monkeypatch):
|
|
"""
|
|
Use provided GRPC endpoint when it is set
|
|
"""
|
|
monkeypatch.setenv("ARIZE_ENDPOINT", "grpc://test.endpoint")
|
|
|
|
config = ArizeLogger.get_arize_opentelemetry_config()
|
|
assert isinstance(config, OpenTelemetryConfig)
|
|
assert config.exporter == "otlp_grpc"
|
|
assert config.endpoint == "grpc://test.endpoint"
|
|
|
|
|
|
def test_get_arize_opentelemetry_config_http(mock_env_vars, monkeypatch):
|
|
"""
|
|
Use provided HTTP endpoint when it is set
|
|
"""
|
|
monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "http://test.endpoint")
|
|
|
|
config = ArizeLogger.get_arize_opentelemetry_config()
|
|
assert isinstance(config, OpenTelemetryConfig)
|
|
assert config.exporter == "otlp_http"
|
|
assert config.endpoint == "http://test.endpoint"
|