revert space_key change and add tests for arize integration

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: \")"
       ]

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

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

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

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,

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:

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

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"

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

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",