From 63166c3accba05e7fb03145c0a329b111cda3cee Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 23 Jan 2026 23:13:21 -0800
Subject: [PATCH] fixing arize tests

---
 .../integrations/arize/test_arize_utils.py    | 21 ++++++-------------
 1 file changed, 6 insertions(+), 15 deletions(-)

diff --git a/tests/test_litellm/integrations/arize/test_arize_utils.py b/tests/test_litellm/integrations/arize/test_arize_utils.py
index 6eb6e67a08..9a9f3d5afc 100644
--- a/tests/test_litellm/integrations/arize/test_arize_utils.py
+++ b/tests/test_litellm/integrations/arize/test_arize_utils.py
@@ -84,7 +84,7 @@ def test_arize_set_attributes():
     ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
 
     # Validate that the expected number of attributes were set
-    assert span.set_attribute.call_count == 28
+    assert span.set_attribute.call_count == 26
 
     # Metadata attached to the span
     span.set_attribute.assert_any_call(
@@ -108,7 +108,8 @@ def test_arize_set_attributes():
     # Response metadata
     span.set_attribute.assert_any_call("llm.response.id", "chatcmpl-ID")
     span.set_attribute.assert_any_call("llm.response.model", "gpt-4o")
-    span.set_attribute.assert_any_call(SpanAttributes.OPENINFERENCE_SPAN_KIND, "LLM")
+    # Span kind is set to TOOL when tools are present
+    span.set_attribute.assert_any_call(SpanAttributes.OPENINFERENCE_SPAN_KIND, "TOOL")
 
     # Request message content and metadata
     span.set_attribute.assert_any_call(
@@ -125,14 +126,14 @@ def test_arize_set_attributes():
 
     # Tool call definitions and function names
     span.set_attribute.assert_any_call(
-        f"{SpanAttributes.LLM_TOOLS}.0.{SpanAttributes.TOOL_NAME}", "get_weather"
+        f"{SpanAttributes.LLM_TOOLS}.0.name", "get_weather"
     )
     span.set_attribute.assert_any_call(
-        f"{SpanAttributes.LLM_TOOLS}.0.{SpanAttributes.TOOL_DESCRIPTION}",
+        f"{SpanAttributes.LLM_TOOLS}.0.description",
         "Fetches weather details.",
     )
     span.set_attribute.assert_any_call(
-        f"{SpanAttributes.LLM_TOOLS}.0.{SpanAttributes.TOOL_PARAMETERS}",
+        f"{SpanAttributes.LLM_TOOLS}.0.parameters",
         json.dumps(
             {
                 "type": "object",
@@ -144,16 +145,6 @@ def test_arize_set_attributes():
         ),
     )
 
-    # Tool calls captured from optional_params
-    span.set_attribute.assert_any_call(
-        f"{MessageAttributes.MESSAGE_TOOL_CALLS}.0.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}",
-        "get_weather",
-    )
-    span.set_attribute.assert_any_call(
-        f"{MessageAttributes.MESSAGE_TOOL_CALLS}.1.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}",
-        "get_stock_price",
-    )
-
     # Invocation parameters
     span.set_attribute.assert_any_call(
         SpanAttributes.LLM_INVOCATION_PARAMETERS, '{"user": "test_user"}'