mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
1515 lines
53 KiB
Python
1515 lines
53 KiB
Python
import json
|
|
from typing import Optional
|
|
|
|
# Adds the grandparent directory to sys.path to allow importing project modules
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.integrations._types.open_inference import (
|
|
MessageAttributes,
|
|
SpanAttributes,
|
|
ToolCallAttributes,
|
|
)
|
|
from litellm.integrations.arize.arize import ArizeLogger
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.types.utils import Choices, StandardCallbackDynamicParams
|
|
|
|
|
|
def test_arize_set_attributes():
|
|
"""
|
|
Test setting attributes for Arize, including all custom LLM attributes.
|
|
Ensures that the correct span attributes are being added during a request.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import ModelResponse
|
|
|
|
span = MagicMock() # Mocked tracing span to test attribute setting
|
|
|
|
# Construct kwargs to simulate a real LLM request scenario
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "Basic Request Content"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {"user": "test_user"},
|
|
"metadata": {"key_1": "value_1", "key_2": None},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {
|
|
"max_tokens": "100",
|
|
"temperature": "1",
|
|
"top_p": "5",
|
|
"stream": False,
|
|
"user": "test_user",
|
|
"tools": [
|
|
{
|
|
"function": {
|
|
"name": "get_weather",
|
|
"description": "Fetches weather details.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {
|
|
"type": "string",
|
|
"description": "City name",
|
|
}
|
|
},
|
|
"required": ["location"],
|
|
},
|
|
}
|
|
}
|
|
],
|
|
"functions": [{"name": "get_weather"}, {"name": "get_stock_price"}],
|
|
},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
|
|
# Simulated LLM response object
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 100, "completion_tokens": 60, "prompt_tokens": 40},
|
|
choices=[
|
|
Choices(message={"role": "assistant", "content": "Basic Response Content"})
|
|
],
|
|
model="gpt-4o",
|
|
id="chatcmpl-ID",
|
|
)
|
|
|
|
# Apply attribute setting via ArizeLogger
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
|
|
# Validate that the expected number of attributes were set.
|
|
# OPENINFERENCE_SPAN_KIND is written exactly once (defensively, before
|
|
# the main attribute pipeline) so a partial failure cannot blank it.
|
|
# Per the OpenInference spec, a chat completion that passes `tools=[...]`
|
|
# is still an LLM span — not TOOL (TOOL is reserved for actual tool
|
|
# execution by application code).
|
|
assert span.set_attribute.call_count == 26
|
|
|
|
# Metadata attached to the span
|
|
span.set_attribute.assert_any_call(
|
|
SpanAttributes.METADATA, json.dumps({"key_1": "value_1", "key_2": None})
|
|
)
|
|
|
|
# Basic LLM information
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_MODEL_NAME, "gpt-4o")
|
|
span.set_attribute.assert_any_call("llm.request.type", "completion")
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_PROVIDER, "openai")
|
|
|
|
# LLM generation parameters
|
|
span.set_attribute.assert_any_call("llm.request.max_tokens", "100")
|
|
span.set_attribute.assert_any_call("llm.request.temperature", "1")
|
|
span.set_attribute.assert_any_call("llm.request.top_p", "5")
|
|
|
|
# Streaming and user info
|
|
span.set_attribute.assert_any_call("llm.is_streaming", "False")
|
|
span.set_attribute.assert_any_call("llm.user", "test_user")
|
|
|
|
# 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 kind stays LLM even when tools are passed (OpenInference spec).
|
|
span.set_attribute.assert_any_call(SpanAttributes.OPENINFERENCE_SPAN_KIND, "LLM")
|
|
# And TOOL must never be written for an LLM chat completion call.
|
|
span_kind_writes = [
|
|
c.args[1]
|
|
for c in span.set_attribute.call_args_list
|
|
if c.args[0] == SpanAttributes.OPENINFERENCE_SPAN_KIND
|
|
]
|
|
assert "TOOL" not in span_kind_writes
|
|
|
|
# Request message content and metadata
|
|
span.set_attribute.assert_any_call(
|
|
SpanAttributes.INPUT_VALUE, "Basic Request Content"
|
|
)
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}",
|
|
"user",
|
|
)
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENT}",
|
|
"Basic Request Content",
|
|
)
|
|
|
|
# Tool call definitions and function names
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_TOOLS}.0.name", "get_weather"
|
|
)
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_TOOLS}.0.description",
|
|
"Fetches weather details.",
|
|
)
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_TOOLS}.0.parameters",
|
|
json.dumps(
|
|
{
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {"type": "string", "description": "City name"}
|
|
},
|
|
"required": ["location"],
|
|
}
|
|
),
|
|
)
|
|
|
|
# Invocation parameters
|
|
span.set_attribute.assert_any_call(
|
|
SpanAttributes.LLM_INVOCATION_PARAMETERS, '{"user": "test_user"}'
|
|
)
|
|
|
|
# User ID
|
|
span.set_attribute.assert_any_call(SpanAttributes.USER_ID, "test_user")
|
|
|
|
# Output message content
|
|
span.set_attribute.assert_any_call(
|
|
SpanAttributes.OUTPUT_VALUE, "Basic Response Content"
|
|
)
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}",
|
|
"assistant",
|
|
)
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENT}",
|
|
"Basic Response Content",
|
|
)
|
|
|
|
# Token counts
|
|
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)
|
|
|
|
|
|
def test_arize_set_attributes_responses_api():
|
|
"""
|
|
Test setting attributes for Responses API with mixed output (reasoning + message).
|
|
Verifies that multiple output types are correctly handled.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
from litellm.types.llms.openai import (
|
|
ResponsesAPIResponse,
|
|
ResponseAPIUsage,
|
|
OutputTokensDetails,
|
|
)
|
|
from openai.types.responses import (
|
|
ResponseReasoningItem,
|
|
ResponseOutputMessage,
|
|
ResponseOutputText,
|
|
)
|
|
from openai.types.responses.response_reasoning_item import Summary
|
|
|
|
span = MagicMock() # Mocked tracing span to test attribute setting
|
|
|
|
# Construct kwargs to simulate a real LLM request scenario
|
|
kwargs = {
|
|
"model": "o3-mini",
|
|
"messages": [{"role": "user", "content": "What is the answer?"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {"user": "test_user", "stream": True},
|
|
"metadata": {"key_1": "value_1", "key_2": None},
|
|
"call_type": "responses",
|
|
},
|
|
"optional_params": {
|
|
"max_tokens": "100",
|
|
"temperature": "1",
|
|
"top_p": "5",
|
|
"stream": True,
|
|
"user": "test_user",
|
|
},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
|
|
# Simulate Responses API response with mixed output
|
|
response_obj = ResponsesAPIResponse(
|
|
id="response-123",
|
|
created_at=1625247600,
|
|
output=[
|
|
ResponseReasoningItem(
|
|
id="reasoning-001",
|
|
type="reasoning",
|
|
summary=[
|
|
Summary(text="First, I need to analyze...", type="summary_text")
|
|
],
|
|
),
|
|
ResponseOutputMessage(
|
|
id="msg-001",
|
|
type="message",
|
|
role="assistant",
|
|
status="completed",
|
|
content=[
|
|
ResponseOutputText(
|
|
annotations=[],
|
|
text="The answer is 42",
|
|
type="output_text",
|
|
)
|
|
],
|
|
),
|
|
],
|
|
usage=ResponseAPIUsage(
|
|
input_tokens=120,
|
|
output_tokens=250,
|
|
total_tokens=370,
|
|
output_tokens_details=OutputTokensDetails(reasoning_tokens=180),
|
|
),
|
|
)
|
|
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
|
|
# Verify reasoning summary was set (index 0)
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_REASONING_SUMMARY}",
|
|
"First, I need to analyze...",
|
|
)
|
|
|
|
# Verify message content was set (index 1)
|
|
span.set_attribute.assert_any_call(SpanAttributes.OUTPUT_VALUE, "The answer is 42")
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.1.{MessageAttributes.MESSAGE_CONTENT}",
|
|
"The answer is 42",
|
|
)
|
|
span.set_attribute.assert_any_call(
|
|
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.1.{MessageAttributes.MESSAGE_ROLE}",
|
|
"assistant",
|
|
)
|
|
|
|
# Verify token counts including reasoning tokens
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 370)
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 250)
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 120)
|
|
span.set_attribute.assert_any_call(
|
|
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 180
|
|
)
|
|
|
|
|
|
def test_set_usage_outputs_pydantic_completion_usage():
|
|
"""
|
|
Regression test for https://github.com/BerriAI/litellm/issues/13672
|
|
|
|
`_set_usage_outputs` previously called `usage.get(...)` which crashes when
|
|
`usage` is a plain Pydantic model (e.g. openai.types.completion_usage.CompletionUsage)
|
|
that does not implement dict-style `.get()`. Same crash for nested
|
|
`output_tokens_details` / `completion_tokens_details`.
|
|
|
|
The function must:
|
|
1. Read total/prompt/completion tokens from a Pydantic usage without `.get`.
|
|
2. Read reasoning_tokens from `completion_tokens_details` (chat completions API)
|
|
OR `output_tokens_details` (responses API), even when those nested objects
|
|
are Pydantic models without `.get`.
|
|
3. Not raise AttributeError; not call span.record_exception.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
|
|
from openai.types.completion_usage import (
|
|
CompletionTokensDetails,
|
|
CompletionUsage,
|
|
)
|
|
|
|
from litellm.integrations.arize._utils import _set_usage_outputs
|
|
|
|
span = MagicMock()
|
|
|
|
# Plain OpenAI Pydantic model — has no `.get()`
|
|
usage = CompletionUsage(
|
|
completion_tokens=60,
|
|
prompt_tokens=40,
|
|
total_tokens=100,
|
|
completion_tokens_details=CompletionTokensDetails(reasoning_tokens=25),
|
|
)
|
|
assert not hasattr(usage, "get"), "precondition: CompletionUsage must lack .get"
|
|
|
|
response_obj = {"usage": usage}
|
|
|
|
# Must not raise
|
|
_set_usage_outputs(span, response_obj, SpanAttributes)
|
|
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 100)
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 40)
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 60)
|
|
# reasoning_tokens for chat completions live in completion_tokens_details
|
|
span.set_attribute.assert_any_call(
|
|
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 25
|
|
)
|
|
|
|
|
|
def test_set_usage_outputs_pydantic_response_api_usage():
|
|
"""
|
|
Same crash also affects Responses API with `output_tokens_details` as a
|
|
Pydantic model that lacks `.get()`. Verifies the responses-API path.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.integrations.arize._utils import _set_usage_outputs
|
|
from litellm.types.llms.openai import OutputTokensDetails
|
|
|
|
# Build an object that mimics openai ResponsesAPI usage but lacks `.get`
|
|
# (uses a plain class — not BaseLiteLLMOpenAIResponseObject)
|
|
class PlainResponsesUsage:
|
|
def __init__(self):
|
|
self.total_tokens = 370
|
|
self.input_tokens = 120
|
|
self.output_tokens = 250
|
|
self.output_tokens_details = OutputTokensDetails(reasoning_tokens=180)
|
|
|
|
usage = PlainResponsesUsage()
|
|
assert not hasattr(usage, "get")
|
|
|
|
span = MagicMock()
|
|
response_obj = {"usage": usage}
|
|
|
|
_set_usage_outputs(span, response_obj, SpanAttributes)
|
|
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, 370)
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_PROMPT, 120)
|
|
span.set_attribute.assert_any_call(SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, 250)
|
|
span.set_attribute.assert_any_call(
|
|
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, 180
|
|
)
|
|
|
|
|
|
class TestArizeLogger(CustomLogger):
|
|
"""
|
|
Custom logger implementation to capture standard_callback_dynamic_params.
|
|
Used to verify that dynamic config keys are being passed to callbacks.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.standard_callback_dynamic_params: Optional[
|
|
StandardCallbackDynamicParams
|
|
] = None
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
# Capture dynamic params and print them for verification
|
|
print("logged kwargs", json.dumps(kwargs, indent=4, default=str))
|
|
self.standard_callback_dynamic_params = kwargs.get(
|
|
"standard_callback_dynamic_params"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arize_dynamic_params():
|
|
"""
|
|
Test to ensure that dynamic Arize keys (API key and space key)
|
|
are received inside the callback logger at runtime.
|
|
"""
|
|
test_arize_logger = TestArizeLogger()
|
|
litellm.callbacks = [test_arize_logger]
|
|
|
|
# Perform a mocked async completion call to trigger logging
|
|
await litellm.acompletion(
|
|
model="gpt-4o",
|
|
messages=[{"role": "user", "content": "Basic Request Content"}],
|
|
mock_response="test",
|
|
arize_api_key="test_api_key_dynamic",
|
|
arize_space_key="test_space_key_dynamic",
|
|
)
|
|
|
|
# Allow for async propagation
|
|
await asyncio.sleep(2)
|
|
|
|
# Assert dynamic parameters were received in the callback
|
|
assert test_arize_logger.standard_callback_dynamic_params is not None
|
|
assert (
|
|
test_arize_logger.standard_callback_dynamic_params.get("arize_api_key")
|
|
== "test_api_key_dynamic"
|
|
)
|
|
assert (
|
|
test_arize_logger.standard_callback_dynamic_params.get("arize_space_key")
|
|
== "test_space_key_dynamic"
|
|
)
|
|
|
|
|
|
def test_construct_dynamic_arize_headers():
|
|
"""
|
|
Test the construct_dynamic_arize_headers method with various input scenarios.
|
|
Ensures that dynamic Arize headers are properly constructed from callback parameters.
|
|
"""
|
|
from litellm.types.utils import StandardCallbackDynamicParams
|
|
|
|
# Test with all parameters present
|
|
dynamic_params_full = StandardCallbackDynamicParams(
|
|
arize_api_key="test_api_key", arize_space_id="test_space_id"
|
|
)
|
|
arize_logger = ArizeLogger()
|
|
|
|
headers = arize_logger.construct_dynamic_otel_headers(dynamic_params_full)
|
|
expected_headers = {"api_key": "test_api_key", "arize-space-id": "test_space_id"}
|
|
assert headers == expected_headers
|
|
|
|
# Test with only space_id
|
|
dynamic_params_space_id_only = StandardCallbackDynamicParams(
|
|
arize_space_id="test_space_id"
|
|
)
|
|
|
|
headers = arize_logger.construct_dynamic_otel_headers(dynamic_params_space_id_only)
|
|
expected_headers = {"arize-space-id": "test_space_id"}
|
|
assert headers == expected_headers
|
|
|
|
# Test with empty parameters dict
|
|
dynamic_params_empty = StandardCallbackDynamicParams()
|
|
|
|
headers = arize_logger.construct_dynamic_otel_headers(dynamic_params_empty)
|
|
assert headers == {}
|
|
|
|
# test with space key and api key
|
|
dynamic_params_space_key_and_api_key = StandardCallbackDynamicParams(
|
|
arize_space_key="test_space_key", arize_api_key="test_api_key"
|
|
)
|
|
headers = arize_logger.construct_dynamic_otel_headers(
|
|
dynamic_params_space_key_and_api_key
|
|
)
|
|
expected_headers = {"arize-space-id": "test_space_key", "api_key": "test_api_key"}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additive rendering-enhancement tests. None of these assert that previously
|
|
# emitted attributes were removed or changed — they only assert that the new
|
|
# attributes appear in their respective scenarios.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _collect_calls(span):
|
|
"""Helper: return dict[attr_name] = value of all set_attribute calls."""
|
|
out = {}
|
|
for call in span.set_attribute.call_args_list:
|
|
args = call.args
|
|
if len(args) >= 2:
|
|
out[args[0]] = args[1]
|
|
return out
|
|
|
|
|
|
def test_arize_emits_cache_tokens_openai_style():
|
|
"""OpenAI prompt_tokens_details.cached_tokens → cache_read attr."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.integrations.arize._utils import _set_usage_outputs
|
|
|
|
span = MagicMock()
|
|
response_obj = {
|
|
"usage": {
|
|
"total_tokens": 100,
|
|
"completion_tokens": 60,
|
|
"prompt_tokens": 40,
|
|
"prompt_tokens_details": {"cached_tokens": 32, "audio_tokens": 8},
|
|
}
|
|
}
|
|
_set_usage_outputs(span, response_obj, SpanAttributes)
|
|
attrs = _collect_calls(span)
|
|
assert attrs[SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ] == 32
|
|
assert attrs[SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO] == 8
|
|
|
|
|
|
def test_arize_emits_cache_tokens_anthropic_style():
|
|
"""Anthropic/Bedrock cache_read_input_tokens / cache_creation_input_tokens."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.integrations.arize._utils import _set_usage_outputs
|
|
|
|
span = MagicMock()
|
|
response_obj = {
|
|
"usage": {
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"cache_read_input_tokens": 80,
|
|
"cache_creation_input_tokens": 20,
|
|
}
|
|
}
|
|
_set_usage_outputs(span, response_obj, SpanAttributes)
|
|
attrs = _collect_calls(span)
|
|
assert attrs[SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ] == 80
|
|
assert attrs[SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE] == 20
|
|
|
|
|
|
def test_arize_emits_no_cache_tokens_when_absent():
|
|
"""Regression guard: when no cache fields exist, no cache attrs emitted."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.integrations.arize._utils import _set_usage_outputs
|
|
|
|
span = MagicMock()
|
|
response_obj = {
|
|
"usage": {"total_tokens": 10, "completion_tokens": 4, "prompt_tokens": 6}
|
|
}
|
|
_set_usage_outputs(span, response_obj, SpanAttributes)
|
|
attrs = _collect_calls(span)
|
|
assert SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ not in attrs
|
|
assert SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE not in attrs
|
|
|
|
|
|
def test_passthrough_call_type_resolves_to_llm_span_kind():
|
|
"""`allm_passthrough_route` should map to LLM (was UNKNOWN before fix)."""
|
|
from litellm.integrations._types.open_inference import OpenInferenceSpanKindValues
|
|
from litellm.integrations.arize._utils import _infer_open_inference_span_kind
|
|
|
|
assert (
|
|
_infer_open_inference_span_kind("allm_passthrough_route")
|
|
== OpenInferenceSpanKindValues.LLM.value
|
|
)
|
|
assert (
|
|
_infer_open_inference_span_kind("llm_passthrough_route")
|
|
== OpenInferenceSpanKindValues.LLM.value
|
|
)
|
|
|
|
|
|
def test_arize_chat_completion_with_tools_stays_llm_span_kind():
|
|
"""Regression guard against the old `TOOL` override: a chat completion
|
|
that passes `tools=[...]` AND returns `tool_calls` must remain LLM."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "weather?"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {
|
|
"tools": [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"description": "weather",
|
|
"parameters": {"type": "object", "properties": {}},
|
|
},
|
|
}
|
|
]
|
|
},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 10, "completion_tokens": 4, "prompt_tokens": 6},
|
|
choices=[
|
|
Choices(
|
|
message={
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_x",
|
|
"type": "function",
|
|
"function": {"name": "get_weather", "arguments": "{}"},
|
|
}
|
|
],
|
|
}
|
|
)
|
|
],
|
|
model="gpt-4o",
|
|
id="r-toolkind",
|
|
)
|
|
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
span_kind_writes = [
|
|
c.args[1]
|
|
for c in span.set_attribute.call_args_list
|
|
if c.args[0] == SpanAttributes.OPENINFERENCE_SPAN_KIND
|
|
]
|
|
assert span_kind_writes, "span.kind must be written"
|
|
assert all(v == "LLM" for v in span_kind_writes)
|
|
assert "TOOL" not in span_kind_writes
|
|
|
|
|
|
def test_arize_emits_assistant_tool_calls_on_output_message():
|
|
"""Assistant tool_calls should surface as MESSAGE_TOOL_CALLS.* attrs."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "weather?"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 10, "completion_tokens": 4, "prompt_tokens": 6},
|
|
choices=[
|
|
Choices(
|
|
message={
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_abc",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "SF"}',
|
|
},
|
|
}
|
|
],
|
|
}
|
|
)
|
|
],
|
|
model="gpt-4o",
|
|
id="chatcmpl-1",
|
|
)
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
attrs = _collect_calls(span)
|
|
base = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_TOOL_CALLS}.0"
|
|
assert attrs[f"{base}.{ToolCallAttributes.TOOL_CALL_ID}"] == "call_abc"
|
|
assert (
|
|
attrs[f"{base}.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}"] == "get_weather"
|
|
)
|
|
assert (
|
|
attrs[f"{base}.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}"]
|
|
== '{"location": "SF"}'
|
|
)
|
|
|
|
|
|
def test_arize_output_value_falls_back_to_tool_calls_summary():
|
|
"""When the assistant returns no text content but did request tool
|
|
calls, OUTPUT_VALUE should contain a JSON summary so Arize's Output
|
|
pane shows something."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "weather?"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 10, "completion_tokens": 4, "prompt_tokens": 6},
|
|
choices=[
|
|
Choices(
|
|
message={
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_abc",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "SF"}',
|
|
},
|
|
}
|
|
],
|
|
}
|
|
)
|
|
],
|
|
model="gpt-4o",
|
|
id="r-tc-out",
|
|
)
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
attrs = _collect_calls(span)
|
|
|
|
# OUTPUT_VALUE should contain the tool_call name + arguments JSON
|
|
out = attrs[SpanAttributes.OUTPUT_VALUE]
|
|
assert "tool_calls" in out
|
|
assert "get_weather" in out
|
|
assert "SF" in out
|
|
|
|
|
|
def test_arize_output_value_unchanged_when_content_present():
|
|
"""Regression guard: when content is non-empty, OUTPUT_VALUE must be
|
|
exactly that content (no summary written)."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 4, "completion_tokens": 2, "prompt_tokens": 2},
|
|
choices=[
|
|
Choices(
|
|
message={
|
|
"role": "assistant",
|
|
"content": "hello world",
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_x",
|
|
"type": "function",
|
|
"function": {"name": "n", "arguments": "{}"},
|
|
}
|
|
],
|
|
}
|
|
)
|
|
],
|
|
model="gpt-4o",
|
|
id="r-content",
|
|
)
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
attrs = _collect_calls(span)
|
|
assert attrs[SpanAttributes.OUTPUT_VALUE] == "hello world"
|
|
|
|
|
|
def test_arize_emits_tool_call_id_and_name_on_input_tool_message():
|
|
"""A tool-result input message should expose tool_call_id + name."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [
|
|
{"role": "user", "content": "weather?"},
|
|
{
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_abc",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "SF"}',
|
|
},
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": "call_abc",
|
|
"name": "get_weather",
|
|
"content": "sunny, 72F",
|
|
},
|
|
],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 10, "completion_tokens": 4, "prompt_tokens": 6},
|
|
choices=[Choices(message={"role": "assistant", "content": "It's sunny."})],
|
|
model="gpt-4o",
|
|
id="chatcmpl-2",
|
|
)
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
attrs = _collect_calls(span)
|
|
# Assistant tool_call surfaces on input msg index 1
|
|
assistant_base = f"{SpanAttributes.LLM_INPUT_MESSAGES}.1.{MessageAttributes.MESSAGE_TOOL_CALLS}.0"
|
|
assert attrs[f"{assistant_base}.{ToolCallAttributes.TOOL_CALL_ID}"] == "call_abc"
|
|
# Tool message at index 2
|
|
tool_prefix = f"{SpanAttributes.LLM_INPUT_MESSAGES}.2"
|
|
assert (
|
|
attrs[f"{tool_prefix}.{MessageAttributes.MESSAGE_TOOL_CALL_ID}"] == "call_abc"
|
|
)
|
|
assert attrs[f"{tool_prefix}.{MessageAttributes.MESSAGE_NAME}"] == "get_weather"
|
|
|
|
|
|
def test_arize_emits_multimodal_input_contents():
|
|
"""List-shaped content should populate MESSAGE_CONTENTS.* alongside the
|
|
legacy MESSAGE_CONTENT (which stays for back-compat)."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "What is in this image?"},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": "https://example.com/cat.png"},
|
|
},
|
|
],
|
|
}
|
|
],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 10, "completion_tokens": 4, "prompt_tokens": 6},
|
|
choices=[Choices(message={"role": "assistant", "content": "A cat."})],
|
|
model="gpt-4o",
|
|
id="chatcmpl-img",
|
|
)
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
attrs = _collect_calls(span)
|
|
base = f"{SpanAttributes.LLM_INPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENTS}"
|
|
assert attrs[f"{base}.0.message_content.type"] == "text"
|
|
assert attrs[f"{base}.0.message_content.text"] == "What is in this image?"
|
|
assert attrs[f"{base}.1.message_content.type"] == "image"
|
|
assert (
|
|
attrs[f"{base}.1.message_content.image.image.url"]
|
|
== "https://example.com/cat.png"
|
|
)
|
|
|
|
|
|
def test_arize_emits_session_and_user_attrs_from_metadata():
|
|
"""end_user_id → SESSION_ID; user_api_key_user_id → USER_ID (only when
|
|
optional_params.user/model_params.user absent)."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {
|
|
"user_api_key_user_id": "user_42",
|
|
"user_api_key_end_user_id": "session_99",
|
|
"user_api_key_team_id": "team_7",
|
|
"user_api_key_team_alias": "alpha",
|
|
"user_api_key_alias": "key_alpha",
|
|
},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 4, "completion_tokens": 2, "prompt_tokens": 2},
|
|
choices=[Choices(message={"role": "assistant", "content": "hello"})],
|
|
model="gpt-4o",
|
|
id="r1",
|
|
)
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
attrs = _collect_calls(span)
|
|
assert attrs[SpanAttributes.SESSION_ID] == "session_99"
|
|
assert attrs[SpanAttributes.USER_ID] == "user_42"
|
|
assert attrs["litellm.team_id"] == "team_7"
|
|
assert attrs["litellm.team_alias"] == "alpha"
|
|
assert attrs["litellm.key_alias"] == "key_alpha"
|
|
|
|
|
|
def test_arize_does_not_use_trace_id_as_session_id_fallback():
|
|
"""SESSION_ID must NOT fall back to trace_id (one session-per-request
|
|
would distort Arize Session analytics). trace_id is emitted under its
|
|
own `litellm.trace_id` key instead.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "completion",
|
|
"trace_id": "trace-xyz-123",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 4, "completion_tokens": 2, "prompt_tokens": 2},
|
|
choices=[Choices(message={"role": "assistant", "content": "hi"})],
|
|
model="gpt-4o",
|
|
id="r-trace",
|
|
)
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
attrs = _collect_calls(span)
|
|
|
|
# SESSION_ID must NOT be derived from trace_id.
|
|
assert SpanAttributes.SESSION_ID not in attrs
|
|
# trace_id surfaces under its own key.
|
|
assert attrs["litellm.trace_id"] == "trace-xyz-123"
|
|
|
|
|
|
def test_arize_does_not_overwrite_user_id_from_optional_params():
|
|
"""If optional_params.user is set, metadata USER_ID must NOT overwrite."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {"user": "from_model_params"},
|
|
"metadata": {"user_api_key_user_id": "from_metadata"},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {"user": "from_optional_params"},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 4, "completion_tokens": 2, "prompt_tokens": 2},
|
|
choices=[Choices(message={"role": "assistant", "content": "hello"})],
|
|
model="gpt-4o",
|
|
id="r2",
|
|
)
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
user_id_writes = [
|
|
c.args[1]
|
|
for c in span.set_attribute.call_args_list
|
|
if c.args[0] == SpanAttributes.USER_ID
|
|
]
|
|
assert "from_metadata" not in user_id_writes
|
|
|
|
|
|
def test_arize_emits_response_cost():
|
|
"""StandardLoggingPayload.response_cost → llm.cost.total (+ legacy llm.response.cost)."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "completion",
|
|
"response_cost": 0.0012345,
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
usage={"total_tokens": 4, "completion_tokens": 2, "prompt_tokens": 2},
|
|
choices=[Choices(message={"role": "assistant", "content": "hello"})],
|
|
model="gpt-4o",
|
|
id="r3",
|
|
)
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
attrs = _collect_calls(span)
|
|
assert attrs["llm.cost.total"] == 0.0012345
|
|
assert attrs["llm.response.cost"] == 0.0012345 # legacy key still emitted
|
|
|
|
|
|
def test_arize_passthrough_bedrock_anthropic_normalization():
|
|
"""Bedrock-Anthropic passthrough: input/output text must be set so the
|
|
span renders something other than raw provider attrs."""
|
|
from unittest.mock import MagicMock
|
|
|
|
span = MagicMock()
|
|
bedrock_response_body = {
|
|
"id": "msg_01",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "The capital of France is Paris."}],
|
|
"model": "anthropic.claude-sonnet-4-v1:0",
|
|
"stop_reason": "end_turn",
|
|
"usage": {"input_tokens": 18, "output_tokens": 12},
|
|
}
|
|
|
|
class FakeHttpxResponse:
|
|
"""Minimal httpx.Response stand-in: has `.text` and no `.get`."""
|
|
|
|
def __init__(self, body):
|
|
self.text = json.dumps(body)
|
|
|
|
response_obj = FakeHttpxResponse(bedrock_response_body)
|
|
kwargs = {
|
|
"model": "anthropic.claude-sonnet-4-v1:0",
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": json.dumps({"messages": [{"role": "user", "content": "?"}]}),
|
|
}
|
|
],
|
|
"additional_args": {
|
|
"complete_input_dict": {
|
|
"anthropic_version": "bedrock-2023-05-31",
|
|
"max_tokens": 64,
|
|
"messages": [
|
|
{"role": "user", "content": "What is the capital of France?"}
|
|
],
|
|
}
|
|
},
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "allm_passthrough_route",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "bedrock"},
|
|
}
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
attrs = _collect_calls(span)
|
|
|
|
# Input rendering
|
|
assert attrs[SpanAttributes.INPUT_VALUE] == "What is the capital of France?"
|
|
msg0 = f"{SpanAttributes.LLM_INPUT_MESSAGES}.0"
|
|
assert attrs[f"{msg0}.{MessageAttributes.MESSAGE_ROLE}"] == "user"
|
|
assert (
|
|
attrs[f"{msg0}.{MessageAttributes.MESSAGE_CONTENT}"]
|
|
== "What is the capital of France?"
|
|
)
|
|
|
|
# Output rendering (Anthropic content[].text)
|
|
assert attrs[SpanAttributes.OUTPUT_VALUE] == "The capital of France is Paris."
|
|
out0 = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0"
|
|
assert attrs[f"{out0}.{MessageAttributes.MESSAGE_ROLE}"] == "assistant"
|
|
assert (
|
|
attrs[f"{out0}.{MessageAttributes.MESSAGE_CONTENT}"]
|
|
== "The capital of France is Paris."
|
|
)
|
|
|
|
# Token counts (Bedrock input_tokens/output_tokens) — extracted via
|
|
# coercion of the non-dict response.
|
|
assert attrs[SpanAttributes.LLM_TOKEN_COUNT_PROMPT] == 18
|
|
assert attrs[SpanAttributes.LLM_TOKEN_COUNT_COMPLETION] == 12
|
|
|
|
# Span kind defended even though the call_type is a passthrough variant.
|
|
span_kind_writes = [
|
|
c.args[1]
|
|
for c in span.set_attribute.call_args_list
|
|
if c.args[0] == SpanAttributes.OPENINFERENCE_SPAN_KIND
|
|
]
|
|
assert span_kind_writes # at least one
|
|
assert all(v == "LLM" for v in span_kind_writes)
|
|
|
|
|
|
def test_arize_passthrough_call_type_does_not_run_on_chat_completion():
|
|
"""Guard: passthrough normalizer must not fire for normal chat calls.
|
|
|
|
If it did, it could double-write input/output for ordinary completions.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.integrations.arize._utils import _maybe_normalize_passthrough
|
|
|
|
span = MagicMock()
|
|
_maybe_normalize_passthrough(
|
|
span,
|
|
{
|
|
"additional_args": {
|
|
"complete_input_dict": {"messages": [{"role": "user", "content": "x"}]}
|
|
}
|
|
},
|
|
{"choices": [{"message": {"role": "assistant", "content": "y"}}]},
|
|
{"choices": [{"message": {"role": "assistant", "content": "y"}}]},
|
|
{"call_type": "completion"},
|
|
)
|
|
assert span.set_attribute.call_count == 0
|
|
|
|
|
|
def test_arize_passthrough_skipped_when_message_redaction_enabled():
|
|
"""Security guard: when message-logging redaction is enabled, the
|
|
passthrough normalizer must NOT export the raw prompt (read from
|
|
`complete_input_dict`, which bypasses central redaction) to the span.
|
|
"""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.integrations.arize._utils import _maybe_normalize_passthrough
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"additional_args": {
|
|
"complete_input_dict": {
|
|
"messages": [
|
|
{"role": "user", "content": "Patient John Doe, SSN 123-45-6789"}
|
|
]
|
|
}
|
|
},
|
|
# Enables redaction via the dynamic-param path inside
|
|
# should_redact_message_logging(), without touching globals.
|
|
"standard_callback_dynamic_params": {"turn_off_message_logging": True},
|
|
}
|
|
_maybe_normalize_passthrough(
|
|
span,
|
|
kwargs,
|
|
{"content": [{"type": "text", "text": "secret response"}]},
|
|
{"content": [{"type": "text", "text": "secret response"}]},
|
|
{"call_type": "allm_passthrough_route"},
|
|
)
|
|
# Nothing — neither input nor output — should be written to the span.
|
|
assert span.set_attribute.call_count == 0
|
|
|
|
|
|
def test_arize_coerce_response_obj_passes_dicts_through_untouched():
|
|
"""Regression guard for the BaseModel/dict path."""
|
|
from litellm.integrations.arize._utils import _coerce_response_obj_for_attrs
|
|
|
|
d = {"id": "x", "model": "m"}
|
|
assert _coerce_response_obj_for_attrs(d) is d
|
|
|
|
class HasGet:
|
|
def get(self, *a, **k): # noqa: D401
|
|
return None
|
|
|
|
obj = HasGet()
|
|
assert _coerce_response_obj_for_attrs(obj) is obj
|
|
|
|
assert _coerce_response_obj_for_attrs(None) is None
|
|
|
|
|
|
def test_arize_coerce_response_obj_parses_httpx_like():
|
|
"""httpx.Response-like objects without `.get` should JSON-decode."""
|
|
from litellm.integrations.arize._utils import _coerce_response_obj_for_attrs
|
|
|
|
class FakeHttpxResponse:
|
|
text = '{"id": "msg_1", "model": "claude"}'
|
|
|
|
parsed = _coerce_response_obj_for_attrs(FakeHttpxResponse())
|
|
assert parsed == {"id": "msg_1", "model": "claude"}
|
|
|
|
|
|
def test_arize_coerce_response_obj_returns_original_on_bad_json():
|
|
from litellm.integrations.arize._utils import _coerce_response_obj_for_attrs
|
|
|
|
class BadJson:
|
|
text = "not-json"
|
|
|
|
obj = BadJson()
|
|
assert _coerce_response_obj_for_attrs(obj) is obj
|
|
|
|
|
|
def test_arize_mcp_call_tool_result_does_not_break_attribute_setting():
|
|
"""`call_mcp_tool` logs the MCP SDK's `CallToolResult`, a Pydantic model
|
|
with no `.get`. It used to raise inside `_set_request_attributes`, aborting
|
|
the whole attribute block (input messages, invocation params, outputs)."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from mcp.types import CallToolResult, TextContent
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "MCP: get_weather",
|
|
"standard_logging_object": {
|
|
"model_parameters": {"user": "u-1"},
|
|
"metadata": {},
|
|
"call_type": "call_mcp_tool",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "mcp"},
|
|
}
|
|
response_obj = CallToolResult(
|
|
content=[TextContent(type="text", text="sunny, 21C")], isError=False
|
|
)
|
|
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
|
|
span.record_exception.assert_not_called()
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.OPENINFERENCE_SPAN_KIND] == "TOOL"
|
|
assert written["llm.request.type"] == "call_mcp_tool"
|
|
# Emitted after the old crash point, so absent before the fix.
|
|
assert written[SpanAttributes.LLM_INVOCATION_PARAMETERS] == '{"user": "u-1"}'
|
|
assert written[SpanAttributes.USER_ID] == "u-1"
|
|
|
|
|
|
def test_arize_coerce_response_obj_dumps_pydantic_without_get():
|
|
from mcp.types import CallToolResult, TextContent
|
|
|
|
from litellm.integrations.arize._utils import _coerce_response_obj_for_attrs
|
|
|
|
result = CallToolResult(content=[TextContent(type="text", text="hi")], isError=False)
|
|
coerced = _coerce_response_obj_for_attrs(result)
|
|
|
|
assert isinstance(coerced, dict)
|
|
assert coerced["isError"] is False
|
|
assert coerced["content"][0]["text"] == "hi"
|
|
|
|
|
|
def test_arize_request_attributes_survive_uncoercible_response_obj():
|
|
"""A response object that is neither dict-like nor coercible (binary
|
|
passthrough body, SDK object) must not abort attribute setting."""
|
|
from unittest.mock import MagicMock
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
|
|
class Opaque:
|
|
pass
|
|
|
|
ArizeLogger.set_arize_attributes(span, kwargs, Opaque())
|
|
|
|
span.record_exception.assert_not_called()
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written["llm.provider"] == "openai"
|
|
|
|
|
|
def _mcp_kwargs(mcp_tool_call_metadata=None, **overrides):
|
|
return {
|
|
"model": "MCP: get_weather",
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {
|
|
"mcp_tool_call_metadata": mcp_tool_call_metadata
|
|
or {
|
|
"name": "get_weather",
|
|
"arguments": {"city": "Seoul"},
|
|
"namespaced_tool_name": "weather-mcp/get_weather",
|
|
}
|
|
},
|
|
"call_type": "call_mcp_tool",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "mcp"},
|
|
**overrides,
|
|
}
|
|
|
|
|
|
def test_arize_mcp_tool_span_renders_name_input_and_output():
|
|
"""`call_mcp_tool` spans have no messages/choices, so Input and Output came
|
|
out blank. Render them from mcp_tool_call_metadata + CallToolResult."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from mcp.types import CallToolResult, TextContent
|
|
|
|
span = MagicMock()
|
|
response_obj = CallToolResult(
|
|
content=[TextContent(type="text", text="sunny, 21C")], isError=False
|
|
)
|
|
|
|
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.TOOL_NAME] == "get_weather"
|
|
assert written[SpanAttributes.INPUT_VALUE] == '{"city": "Seoul"}'
|
|
assert written[SpanAttributes.INPUT_MIME_TYPE] == "application/json"
|
|
assert written[SpanAttributes.OUTPUT_VALUE] == "sunny, 21C"
|
|
assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "text/plain"
|
|
|
|
|
|
def test_arize_mcp_tool_span_serializes_non_text_content():
|
|
"""Image/resource results have no text part, so fall back to JSON."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from mcp.types import CallToolResult, ImageContent
|
|
|
|
span = MagicMock()
|
|
response_obj = CallToolResult(
|
|
content=[ImageContent(type="image", data="Zm9v", mimeType="image/png")],
|
|
isError=False,
|
|
)
|
|
|
|
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json"
|
|
assert "image/png" in written[SpanAttributes.OUTPUT_VALUE]
|
|
|
|
|
|
def test_arize_mcp_tool_span_respects_message_redaction():
|
|
"""Tool arguments and results are user content. With redaction on, only the
|
|
tool name may reach the span."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from mcp.types import CallToolResult, TextContent
|
|
|
|
span = MagicMock()
|
|
response_obj = CallToolResult(
|
|
content=[TextContent(type="text", text="SSN 123-45-6789")], isError=False
|
|
)
|
|
|
|
ArizeLogger.set_arize_attributes(
|
|
span,
|
|
_mcp_kwargs(standard_callback_dynamic_params={"turn_off_message_logging": True}),
|
|
response_obj,
|
|
)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.TOOL_NAME] == "get_weather"
|
|
assert SpanAttributes.INPUT_VALUE not in written
|
|
assert SpanAttributes.OUTPUT_VALUE not in written
|
|
|
|
|
|
def test_arize_non_mcp_span_gets_no_tool_name():
|
|
"""The MCP emitter must not fire on ordinary completions."""
|
|
from unittest.mock import MagicMock
|
|
|
|
from litellm.types.utils import Choices, ModelResponse
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "hi"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {"mcp_tool_call_metadata": {"name": "get_weather"}},
|
|
"call_type": "completion",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "openai"},
|
|
}
|
|
response_obj = ModelResponse(
|
|
choices=[Choices(message={"role": "assistant", "content": "hello"})],
|
|
model="gpt-4o",
|
|
id="r-1",
|
|
)
|
|
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert SpanAttributes.TOOL_NAME not in written
|
|
assert written[SpanAttributes.OUTPUT_VALUE] == "hello"
|
|
|
|
|
|
def test_arize_mcp_tool_span_renders_empty_arguments():
|
|
from unittest.mock import MagicMock
|
|
|
|
from mcp.types import CallToolResult, TextContent
|
|
|
|
span = MagicMock()
|
|
kwargs = _mcp_kwargs(mcp_tool_call_metadata={"name": "ping", "arguments": {}})
|
|
response_obj = CallToolResult(content=[TextContent(type="text", text="pong")], isError=False)
|
|
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.INPUT_VALUE] == "{}"
|
|
assert written[SpanAttributes.INPUT_MIME_TYPE] == "application/json"
|
|
|
|
|
|
def test_arize_mcp_tool_span_renders_empty_content():
|
|
from unittest.mock import MagicMock
|
|
|
|
from mcp.types import CallToolResult
|
|
|
|
span = MagicMock()
|
|
response_obj = CallToolResult(content=[], isError=False)
|
|
|
|
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.OUTPUT_VALUE] == "[]"
|
|
assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json"
|
|
|
|
|
|
def test_arize_mcp_tool_span_falls_back_to_structured_content():
|
|
from unittest.mock import MagicMock
|
|
|
|
from mcp.types import CallToolResult
|
|
|
|
span = MagicMock()
|
|
response_obj = CallToolResult(content=[], structuredContent={"temp_c": 21}, isError=False)
|
|
|
|
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.OUTPUT_VALUE] == '{"temp_c": 21}'
|
|
assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json"
|
|
|
|
|
|
def test_arize_list_mcp_tools_response_does_not_break_attribute_setting():
|
|
from unittest.mock import MagicMock
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "MCP: list_tools",
|
|
"messages": [{"role": "user", "content": "list"}],
|
|
"standard_logging_object": {
|
|
"model_parameters": {},
|
|
"metadata": {},
|
|
"call_type": "list_mcp_tools",
|
|
},
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "mcp"},
|
|
}
|
|
|
|
ArizeLogger.set_arize_attributes(span, kwargs, [{"name": "get_weather"}])
|
|
|
|
span.record_exception.assert_not_called()
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written["llm.input_messages.0.message.content"] == "list"
|
|
|
|
|
|
def test_arize_mcp_tool_span_serializes_mixed_text_and_media():
|
|
from unittest.mock import MagicMock
|
|
|
|
from mcp.types import CallToolResult, ImageContent, TextContent
|
|
|
|
span = MagicMock()
|
|
response_obj = CallToolResult(
|
|
content=[
|
|
TextContent(type="text", text="see image"),
|
|
ImageContent(type="image", data="Zm9v", mimeType="image/png"),
|
|
],
|
|
isError=False,
|
|
)
|
|
|
|
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json"
|
|
assert "see image" in written[SpanAttributes.OUTPUT_VALUE]
|
|
assert "image/png" in written[SpanAttributes.OUTPUT_VALUE]
|
|
|
|
|
|
def test_arize_mcp_tool_span_without_response_object_keeps_name_and_input():
|
|
from unittest.mock import MagicMock
|
|
|
|
span = MagicMock()
|
|
|
|
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), None)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.TOOL_NAME] == "get_weather"
|
|
assert written[SpanAttributes.INPUT_VALUE] == '{"city": "Seoul"}'
|
|
assert SpanAttributes.OUTPUT_VALUE not in written
|
|
|
|
|
|
def test_arize_mcp_tool_span_without_content_emits_no_output():
|
|
from unittest.mock import MagicMock
|
|
|
|
span = MagicMock()
|
|
|
|
ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), {"isError": False})
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert written[SpanAttributes.TOOL_NAME] == "get_weather"
|
|
assert SpanAttributes.OUTPUT_VALUE not in written
|
|
|
|
|
|
def test_arize_mcp_emitter_is_inert_without_a_standard_logging_object():
|
|
from unittest.mock import MagicMock
|
|
|
|
span = MagicMock()
|
|
kwargs = {
|
|
"model": "MCP: get_weather",
|
|
"optional_params": {},
|
|
"litellm_params": {"custom_llm_provider": "mcp"},
|
|
}
|
|
|
|
ArizeLogger.set_arize_attributes(span, kwargs, None)
|
|
|
|
written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list}
|
|
assert SpanAttributes.TOOL_NAME not in written
|