fixing arize tests

This commit is contained in:
yuneng-jiang
2026-01-23 23:13:21 -08:00
parent e1bb4ae280
commit 63166c3acc
@@ -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"}'