From 434e262b8ccc15539c4e95755ca4190152adf2b5 Mon Sep 17 00:00:00 2001 From: Nate Mar Date: Tue, 18 Mar 2025 01:40:10 -0700 Subject: [PATCH] revert space_key change and add tests for arize integration --- .../logging_observability/LiteLLM_Arize.ipynb | 2 +- .../docs/observability/arize_integration.md | 4 +- docs/my-website/docs/proxy/config_settings.md | 2 +- docs/my-website/docs/proxy/logging.md | 2 +- litellm/integrations/arize/arize.py | 8 ++-- litellm/litellm_core_utils/litellm_logging.py | 6 +-- litellm/types/integrations/arize.py | 2 +- tests/local_testing/test_arize_ai.py | 4 +- .../test_arize_logging.py | 40 +++++++++++++++++++ .../test_unit_tests_init_callbacks.py | 2 +- 10 files changed, 56 insertions(+), 16 deletions(-) diff --git a/cookbook/logging_observability/LiteLLM_Arize.ipynb b/cookbook/logging_observability/LiteLLM_Arize.ipynb index d40b4c4bc4..82dfc1ceff 100644 --- a/cookbook/logging_observability/LiteLLM_Arize.ipynb +++ b/cookbook/logging_observability/LiteLLM_Arize.ipynb @@ -105,7 +105,7 @@ "import os\n", "from getpass import getpass\n", "\n", - "os.environ[\"ARIZE_SPACE_ID\"] = getpass(\"Enter your Arize space id: \")\n", + "os.environ[\"ARIZE_SPACE_KEY\"] = getpass(\"Enter your Arize space key: \")\n", "os.environ[\"ARIZE_API_KEY\"] = getpass(\"Enter your Arize API key: \")\n", "os.environ['OPENAI_API_KEY']= getpass(\"Enter your OpenAI API key: \")" ] diff --git a/docs/my-website/docs/observability/arize_integration.md b/docs/my-website/docs/observability/arize_integration.md index ebdf043466..62cf8404c0 100644 --- a/docs/my-website/docs/observability/arize_integration.md +++ b/docs/my-website/docs/observability/arize_integration.md @@ -29,7 +29,7 @@ litellm.callbacks = ["arize"] import litellm import os -os.environ["ARIZE_SPACE_ID"] = "" +os.environ["ARIZE_SPACE_KEY"] = "" os.environ["ARIZE_API_KEY"] = "" # LLM API Keys @@ -61,7 +61,7 @@ litellm_settings: callbacks: ["arize"] environment_variables: - ARIZE_SPACE_ID: "d0*****" + ARIZE_SPACE_KEY: "d0*****" ARIZE_API_KEY: "141a****" ARIZE_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize GRPC api endpoint ARIZE_HTTP_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize HTTP api endpoint. Set either this or ARIZE_ENDPOINT diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 21c95e5f40..9e24437449 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -302,7 +302,7 @@ router_settings: | AISPEND_API_KEY | API Key for AI Spend | ALLOWED_EMAIL_DOMAINS | List of email domains allowed for access | ARIZE_API_KEY | API key for Arize platform integration -| ARIZE_SPACE_ID | Space key for Arize platform +| ARIZE_SPACE_KEY | Space key for Arize platform | ARGILLA_BATCH_SIZE | Batch size for Argilla logging | ARGILLA_API_KEY | API key for Argilla platform | ARGILLA_SAMPLING_RATE | Sampling rate for Argilla logging diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index 3493d37391..6f2f7250b4 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -2000,7 +2000,7 @@ litellm_settings: callbacks: ["arize"] environment_variables: - ARIZE_SPACE_ID: "d0*****" + ARIZE_SPACE_KEY: "d0*****" ARIZE_API_KEY: "141a****" ARIZE_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize GRPC api endpoint ARIZE_HTTP_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize HTTP api endpoint. Set either this or ARIZE_ENDPOINT diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index b7e2bc4545..652957e1ee 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -40,11 +40,11 @@ class ArizeLogger: Raises: ValueError: If required environment variables are not set. """ - space_id = os.environ.get("ARIZE_SPACE_ID") + space_key = os.environ.get("ARIZE_SPACE_KEY") api_key = os.environ.get("ARIZE_API_KEY") - if not space_id: - raise ValueError("ARIZE_SPACE_ID not found in environment variables") + if not space_key: + raise ValueError("ARIZE_SPACE_KEY not found in environment variables") if not api_key: raise ValueError("ARIZE_API_KEY not found in environment variables") @@ -65,7 +65,7 @@ class ArizeLogger: endpoint = "https://otlp.arize.com/v1" return ArizeConfig( - space_id=space_id, + space_key=space_key, api_key=api_key, protocol=protocol, endpoint=endpoint, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1fc9c95458..0945c45491 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2654,7 +2654,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 ) os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - f"space_id={arize_config.space_id},api_key={arize_config.api_key}" + f"space_key={arize_config.space_key},api_key={arize_config.api_key}" ) for callback in _in_memory_loggers: if ( @@ -2899,8 +2899,8 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 elif logging_integration == "arize": from litellm.integrations.opentelemetry import OpenTelemetry - if "ARIZE_SPACE_ID" not in os.environ: - raise ValueError("ARIZE_SPACE_ID not found in environment variables") + if "ARIZE_SPACE_KEY" not in os.environ: + raise ValueError("ARIZE_SPACE_KEY not found in environment variables") if "ARIZE_API_KEY" not in os.environ: raise ValueError("ARIZE_API_KEY not found in environment variables") for callback in _in_memory_loggers: diff --git a/litellm/types/integrations/arize.py b/litellm/types/integrations/arize.py index d8c02f438c..b1559aafa8 100644 --- a/litellm/types/integrations/arize.py +++ b/litellm/types/integrations/arize.py @@ -8,7 +8,7 @@ else: Protocol = Any class ArizeConfig(BaseModel): - space_id: str + space_key: str api_key: str protocol: Protocol endpoint: str diff --git a/tests/local_testing/test_arize_ai.py b/tests/local_testing/test_arize_ai.py index 37ad4685b0..0caf1b5e46 100644 --- a/tests/local_testing/test_arize_ai.py +++ b/tests/local_testing/test_arize_ai.py @@ -34,7 +34,7 @@ async def test_async_otel_callback(): @pytest.fixture def mock_env_vars(monkeypatch): - monkeypatch.setenv("ARIZE_SPACE_ID", "test_space_id") + monkeypatch.setenv("ARIZE_SPACE_KEY", "test_space_key") monkeypatch.setenv("ARIZE_API_KEY", "test_api_key") @@ -44,7 +44,7 @@ def test_get_arize_config(mock_env_vars): """ config = ArizeLogger.get_arize_config() assert isinstance(config, ArizeConfig) - assert config.space_id == "test_space_id" + assert config.space_key == "test_space_key" assert config.api_key == "test_api_key" assert config.endpoint == "https://otlp.arize.com/v1" assert config.protocol == "otlp_grpc" diff --git a/tests/logging_callback_tests/test_arize_logging.py b/tests/logging_callback_tests/test_arize_logging.py index e69de29bb2..59257ad905 100644 --- a/tests/logging_callback_tests/test_arize_logging.py +++ b/tests/logging_callback_tests/test_arize_logging.py @@ -0,0 +1,40 @@ +import os +import sys +import time +from unittest.mock import Mock, patch + +from litellm.main import completion +import opentelemetry.exporter.otlp.proto.grpc.trace_exporter + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system-path + +import litellm + + +def test_arize_callback(): + litellm.callbacks = ["arize"] + os.environ["ARIZE_SPACE_KEY"] = "test_space_key" + os.environ["ARIZE_API_KEY"] = "test_api_key" + os.environ["ARIZE_ENDPOINT"] = "https://otlp.arize.com/v1" + + os.environ["OTEL_BSP_MAX_QUEUE_SIZE"] = "1" + os.environ["OTEL_BSP_MAX_EXPORT_BATCH_SIZE"] = "1" + os.environ["OTEL_BSP_SCHEDULE_DELAY_MILLIS"] = "1" + os.environ["OTEL_BSP_EXPORT_TIMEOUT_MILLIS"] = "5" + + with patch.object( + opentelemetry.exporter.otlp.proto.grpc.trace_exporter.OTLPSpanExporter, + 'export', + new=Mock() + ) as patched_export: + completion( + model="openai/test-model", + messages=[{"role": "user", "content": "arize test content"}], + stream=False, + mock_response="hello there!", + ) + + time.sleep(1) + assert patched_export.called diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index 654a4b0615..fcba3ebbc3 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py @@ -89,7 +89,7 @@ expected_env_vars = { "OPIK_API_KEY": "opik_api_key", "LANGTRACE_API_KEY": "langtrace_api_key", "LOGFIRE_TOKEN": "logfire_token", - "ARIZE_SPACE_ID": "arize_space_id", + "ARIZE_SPACE_KEY": "arize_space_key", "ARIZE_API_KEY": "arize_api_key", "PHOENIX_API_KEY": "phoenix_api_key", "ARGILLA_API_KEY": "argilla_api_key",