Repository: BerriAI/litellm
Pull request #9338: Arize integration Fix
Author: login: nate-mar
Created: 2025-03-18T11:00:10Z | Merged: 2025-03-18T21:28:25Z

## Title

Fixes broken Arize <> LiteLLM Router integration

## Pre-Submission checklist

**Please complete all items before asking a LiteLLM maintainer to review your PR**

- [x] I have Added testing in the `tests/litellm/` directory, **Adding at least 1 test is a hard requirement** - [see details](https://docs.litellm.ai/docs/extras/contributing_code)
- [x] I have added a screenshot of my new test passing locally 
- [x] My PR passes all unit tests on (`make test-unit`)[https://docs.litellm.ai/docs/extras/contributing_code]
- [ ] My PR's scope is as isolated as possible, it only solves 1 specific problem


## Type
🐛 Bug Fix
✅ Test

## Changes

Fixes two important breaking issues:
1. A regression during the adding of support for Phoenix caused an import error that prevented the Arize callback from working. The code was attempting to import an invalid package and throwing an error preventing the integration from working.
2. Even after getting past the import error, the Arize spans do not successfully get ingested into the Arize platform and shown due to the fact that the metadata is not JSON serializable and Arize currently expects a valid json string. This PR fixes this by using the standard logging payload object, grabbing the metadata from that and dumping it as a json string, rather than trying to convert a python dictionary to it's python string representation

Tests:
<img width="1552" alt="Screenshot 2025-03-18 at 4 12 30 AM" src="https://github.com/user-attachments/assets/c1af7c23-cd6c-491e-bd5d-8f80cf291a37" />
<img width="1558" alt="Screenshot 2025-03-18 at 4 11 03 AM" src="https://github.com/user-attachments/assets/d6d047ef-ca3c-4329-a8a4-6803e2840d0a" />


Arize:
<img width="1530" alt="Screenshot 2025-03-18 at 3 47 18 AM" src="https://github.com/user-attachments/assets/ca7651e5-b96b-472d-8e91-511e7e643a89" />



Diff:
diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py
index 9921d47affe..4697245f574 100644
--- a/litellm/integrations/arize/_utils.py
+++ b/litellm/integrations/arize/_utils.py
@@ -12,20 +12,23 @@
 
 
 def set_attributes(span: Span, kwargs, response_obj):
-    from openinference.semconv.trace import (
+    from litellm.integrations._types.open_inference import (
         MessageAttributes,
         OpenInferenceSpanKindValues,
         SpanAttributes,
     )
 
     try:
-        litellm_params = kwargs.get("litellm_params", {}) or {}
-
+        standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get(
+            "standard_logging_object"
+        )
+        
         #############################################
         ############ LLM CALL METADATA ##############
         #############################################
-        metadata = litellm_params.get("metadata", {}) or {}
-        span.set_attribute(SpanAttributes.METADATA, str(metadata))
+        
+        if standard_logging_payload and (metadata := standard_logging_payload["metadata"]):
+            span.set_attribute(SpanAttributes.METADATA, json.dumps(metadata))
 
         #############################################
         ########## LLM Request Attributes ###########
@@ -62,9 +65,6 @@ def set_attributes(span: Span, kwargs, response_obj):
                     msg.get("content", ""),
                 )
 
-        standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get(
-            "standard_logging_object"
-        )
         if standard_logging_payload and (model_params := standard_logging_payload["model_parameters"]):
             # The Generative AI Provider: Azure, OpenAI, etc.
             span.set_attribute(
diff --git a/tests/local_testing/test_arize_ai.py b/tests/local_testing/test_arize_ai.py
index 9d2f27ca30c..b68981ffd58 100644
--- a/tests/local_testing/test_arize_ai.py
+++ b/tests/local_testing/test_arize_ai.py
@@ -1,11 +1,14 @@
 import asyncio
+import json
 import logging
 
+from litellm import Choices
 import pytest
 from dotenv import load_dotenv
 
 import litellm
 from litellm._logging import verbose_logger, verbose_proxy_logger
+from litellm.integrations._types.open_inference import SpanAttributes
 from litellm.integrations.arize.arize import ArizeConfig, ArizeLogger
 
 load_dotenv()
@@ -58,3 +61,40 @@ def test_get_arize_config_with_endpoints(mock_env_vars, monkeypatch):
     config = ArizeLogger.get_arize_config()
     assert config.endpoint == "grpc://test.endpoint"
     assert config.protocol == "otlp_grpc"
+
+
+def test_arize_set_attributes():
+    """
+    Test setting attributes for Arize
+    """
+    from unittest.mock import MagicMock
+    from litellm.types.utils import ModelResponse
+
+    span = MagicMock()
+    kwargs = {
+        "role": "user",
+        "content": "simple arize test",
+        "model": "gpt-4o",
+        "messages": [{"role": "user", "content": "basic arize test"}],
+        "standard_logging_object": {"model_parameters": {"user": "test_user"}, "metadata": {"key": "value", "key2": None}},
+    }
+    response_obj = ModelResponse(usage={"total_tokens": 100, "completion_tokens": 60, "prompt_tokens": 40}, 
+                                 choices=[Choices(message={"role": "assistant", "content": "response content"})])
+
+    ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
+
+    assert span.set_attribute.call_count == 14
+    span.set_attribute.assert_any_call(SpanAttributes.METADATA, json.dumps({"key": "value", "key2": None}))
+    span.set_attribute.assert_any_call(SpanAttributes.LLM_MODEL_NAME, "gpt-4o")
+    span.set_attribute.assert_any_call(SpanAttributes.OPENINFERENCE_SPAN_KIND, "LLM")
+    span.set_attribute.assert_any_call(SpanAttributes.INPUT_VALUE, "basic arize test")
+    span.set_attribute.assert_any_call("llm.input_messages.0.message.role", "user")
+    span.set_attribute.assert_any_call("llm.input_messages.0.message.content", "basic arize test")
+    span.set_attribute.assert_any_call(SpanAttributes.LLM_INVOCATION_PARAMETERS, '{"user": "test_user"}')
+    span.set_attribute.assert_any_call(SpanAttributes.USER_ID, "test_user")
+    span.set_attribute.assert_any_call(SpanAttributes.OUTPUT_VALUE, "response content")
+    span.set_attribute.assert_any_call("llm.output_messages.0.message.role", "assistant")
+    span.set_attribute.assert_any_call("llm.output_messages.0.message.content", "response content")
+    span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 100)
+    span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 60)
+    span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 40)
diff --git a/tests/logging_callback_tests/test_arize_logging.py b/tests/logging_callback_tests/test_arize_logging.py
new file mode 100644
index 00000000000..27f067d7281
--- /dev/null
+++ b/tests/logging_callback_tests/test_arize_logging.py
@@ -0,0 +1,42 @@
+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"
+
+    # Set the batch span processor to quickly flush after a span has been added
+    # This is to ensure that the span is exported before the test ends
+    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) # Wait for the batch span processor to flush
+        assert patched_export.called


Reviews:
user: login: ishaan-jaff
body: reviewed 
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T17:08:31Z
user: login: nate-mar
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T17:32:35Z

Review comments:
diff_hunk: @@ -12,20 +12,23 @@
 
 
 def set_attributes(span: Span, kwargs, response_obj):
-    from openinference.semconv.trace import (
+    from litellm.integrations._types.open_inference import (
         MessageAttributes,
         OpenInferenceSpanKindValues,
         SpanAttributes,
     )
 
     try:
-        litellm_params = kwargs.get("litellm_params", {}) or {}
-
+        standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get(
+            "standard_logging_object"
+        )
+        
         #############################################
         ############ LLM CALL METADATA ##############
         #############################################
-        metadata = litellm_params.get("metadata", {}) or {}
-        span.set_attribute(SpanAttributes.METADATA, str(metadata))
+        
+        if standard_logging_payload and (metadata := standard_logging_payload["metadata"]):
+            span.set_attribute(SpanAttributes.METADATA, json.dumps(metadata))
path: litellm/integrations/arize/_utils.py
user: login: ishaan-jaff
body: can you use `safe_dumps` in litellm_core_utils here 
created_at: 2025-03-18T17:07:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -58,3 +61,40 @@ def test_get_arize_config_with_endpoints(mock_env_vars, monkeypatch):
     config = ArizeLogger.get_arize_config()
     assert config.endpoint == "grpc://test.endpoint"
     assert config.protocol == "otlp_grpc"
+
+
+def test_arize_set_attributes():
path: tests/local_testing/test_arize_ai.py
user: login: ishaan-jaff
body: add this here too, tests/logging_callback_tests/test_arize_logging.py 

let's move all future arize tests to this file: tests/logging_callback_tests/test_arize_logging.py 
created_at: 2025-03-18T17:08:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -58,3 +61,40 @@ def test_get_arize_config_with_endpoints(mock_env_vars, monkeypatch):
     config = ArizeLogger.get_arize_config()
     assert config.endpoint == "grpc://test.endpoint"
     assert config.protocol == "otlp_grpc"
+
+
+def test_arize_set_attributes():
path: tests/local_testing/test_arize_ai.py
user: login: nate-mar
body: got it will do
created_at: 2025-03-18T17:32:35Z
author_association: CONTRIBUTOR

PR comments:
user: login: vercel[bot]
created_at: 2025-03-18T11:00:14Z
body: [vc]: #hPn8MMLay2ujRkHfjm54UCTRW58lrVZY2iEAkfWLO54=:eyJpc01vbm9yZXBvIjp0cnVlLCJ0eXBlIjoiZ2l0aHViIiwicHJvamVjdHMiOlt7Im5hbWUiOiJsaXRlbGxtIiwicm9vdERpcmVjdG9yeSI6ImRvY3MvbXktd2Vic2l0ZSIsImluc3BlY3RvclVybCI6Imh0dHBzOi8vdmVyY2VsLmNvbS9jbGVya2llYWkvbGl0ZWxsbS9CeG0yZGk1MkdFalhxY3hBdEFSR3pTSlJydGJCIiwicHJldmlld1VybCI6ImxpdGVsbG0tZ2l0LWZvcmstbmF0ZS1tYXItYXJpemUtaW50ZWdyYXRpb24tZml4ZXMtY2xlcmtpZWFpLnZlcmNlbC5hcHAiLCJuZXh0Q29tbWl0U3RhdHVzIjoiREVQTE9ZRUQiLCJsaXZlRmVlZGJhY2siOnsicmVzb2x2ZWQiOjAsInVucmVzb2x2ZWQiOjAsInRvdGFsIjowLCJsaW5rIjoibGl0ZWxsbS1naXQtZm9yay1uYXRlLW1hci1hcml6ZS1pbnRlZ3JhdGlvbi1maXhlcy1jbGVya2llYWkudmVyY2VsLmFwcCJ9fV19
**The latest updates on your projects**. Learn more about [Vercel for Git ↗︎](https://vercel.link/github-learn-more)

| Name | Status | Preview | Comments | Updated (UTC) |
| :--- | :----- | :------ | :------- | :------ |
| **litellm** | ✅ Ready ([Inspect](https://vercel.com/clerkieai/litellm/Bxm2di52GEjXqcxAtARGzSJRrtbB)) | [Visit Preview](https://litellm-git-fork-nate-mar-arize-integration-fixes-clerkieai.vercel.app) | 💬 [**Add feedback**](https://vercel.live/open-feedback/litellm-git-fork-nate-mar-arize-integration-fixes-clerkieai.vercel.app?via=pr-comment-feedback-link) | Mar 18, 2025 11:05am |


author_association: NONE
user: login: ishaan-jaff
created_at: 2025-03-18T21:28:21Z
body: will take over from here 
author_association: CONTRIBUTOR