mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
abd239f903
* fix(otel): label retrieval and agent metrics correctly and emit gen_ai.provider.name The GenAI metric attribute builder mapped only chat, text completion, embedding, responses and MCP tool calls to an operation name, so vector-store searches and A2A agent sends fell through to the "chat" default. Their duration and cost then landed in the same series a Grafana GenAI dashboard reads chat latency off, with no way to tell them apart. Both now map to the operation names the convention defines for them, retrieval and invoke_agent, and an unmapped call type says so at debug instead of silently becoming chat. The provider label used gen_ai.system, which the convention deprecated in favor of gen_ai.provider.name; the dashboards built on that vocabulary find nothing under the old key. Metrics now carry gen_ai.provider.name with the semconv provider value (bedrock -> aws.bedrock) via the resolve_provider helper the span path already uses, and keep dual-emitting gen_ai.system with its raw value so a dashboard already querying it keeps matching. A request litellm cannot attribute to a provider gets no provider label at all rather than a placeholder "Unknown" that minted a permanent series nobody can act on. Resolves LIT-4954 Resolves LIT-4959 * fix(otel): map the rest of the vector-store call types off the chat default Mapping only the search left the store lifecycle (create, retrieve, list, update, delete) and the file operations (create, list, retrieve, content, update, delete) falling through to chat, so vector-store admin traffic kept polluting the same series a dashboard reads chat latency off. A live run confirmed it: all 20 metric datapoints from a create, retrieve, list, file-list and delete came out labelled chat. The convention names no operation for vector-store management, so these take vendor values under the litellm. prefix, litellm.vector_store_management and litellm.vector_store_file_management, one per REST resource. Its note on gen_ai.operation.name directs instrumentation to use a system-specific name when no predefined value applies, which is the same allowance resolve_provider already relies on for unmapped providers. Excluding them from the GenAI metrics altogether was the alternative; it deletes series an operator may be watching today and is far harder to reverse than a rename, so it stays available as a follow-up rather than being decided here. Mapping them onto the semconv memory store family was rejected: litellm vector stores hold documents, not agent memory records, and borrowing those names would put document admin calls into whatever charts agent-memory operations, which is the bug this fixes. /rag/query reaches the same recorder and is the same operation as a vector-store search, so query and aquery map to retrieval too; leaving them would have left the defect alive on a second retrieval surface. /rag/ingest is a write with no semconv equivalent and no retrieval or agent confusion, so it is left for the RAG owners to name. Resolves LIT-4954 * fix(otel): give the streaming A2A path a call type so it labels as invoke_agent The streaming logging object is built by hand and never runs through update_environment_variables, the only place call_type reaches model_call_details, so every streamed agent turn arrived at the recorder with no call type and fell back to chat. Stamp it, and map the streaming spelling alongside the non-streaming ones.
135 lines
4.4 KiB
Python
135 lines
4.4 KiB
Python
"""Tests for litellm/a2a_protocol/main.py non-streaming send behavior."""
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("a2a.compat.v0_3.conversions")
|
|
|
|
from a2a.compat.v0_3 import conversions as _conv
|
|
from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest
|
|
|
|
from litellm.a2a_protocol.main import _send_message
|
|
|
|
|
|
def _request() -> SendMessageRequest:
|
|
params = MessageSendParams(
|
|
message={
|
|
"messageId": "m1",
|
|
"role": "user",
|
|
"parts": [{"kind": "text", "text": "hi"}],
|
|
}
|
|
)
|
|
return SendMessageRequest(id="r1", params=params)
|
|
|
|
|
|
def _message_stream_response():
|
|
sr = _conv.pb2_v10.StreamResponse()
|
|
sr.message.message_id = "reply-1"
|
|
sr.message.role = _conv.pb2_v10.Role.ROLE_AGENT
|
|
sr.message.parts.add().text = "hello back"
|
|
return sr
|
|
|
|
|
|
def _status_update_stream_response():
|
|
sr = _conv.pb2_v10.StreamResponse()
|
|
sr.status_update.task_id = "t1"
|
|
sr.status_update.context_id = "c1"
|
|
return sr
|
|
|
|
|
|
class _FakeClient:
|
|
def __init__(self, *events):
|
|
self._events = events
|
|
|
|
async def send_message(self, _pb_request):
|
|
for event in self._events:
|
|
yield event
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_message_returns_message_result():
|
|
response = await _send_message(_FakeClient(_message_stream_response()), _request())
|
|
result = response.root.result
|
|
assert type(result).__name__ == "Message"
|
|
assert response.root.id == "r1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_message_rejects_update_event_final_with_runtime_error():
|
|
with pytest.raises(RuntimeError, match="Message or Task"):
|
|
await _send_message(_FakeClient(_status_update_stream_response()), _request())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_trace_id_prefers_logging_trace_id():
|
|
"""The streaming X-LiteLLM-Trace-Id must use the logging object's trace id (same
|
|
as the non-streaming path), not the JSON-RPC request id, so traces correlate."""
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from a2a.compat.v0_3.types import (
|
|
MessageSendParams,
|
|
SendStreamingMessageRequest,
|
|
)
|
|
|
|
from litellm.a2a_protocol import main as a2a_main
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
|
|
request = SendStreamingMessageRequest(
|
|
id="rpc-1",
|
|
params=MessageSendParams(
|
|
message={
|
|
"messageId": "m1",
|
|
"role": "user",
|
|
"parts": [{"kind": "text", "text": "hi"}],
|
|
}
|
|
),
|
|
)
|
|
logging_obj = MagicMock(spec=Logging)
|
|
logging_obj.litellm_trace_id = "trace-from-logging"
|
|
|
|
captured: dict = {}
|
|
|
|
async def _capture(*, base_url, extra_headers=None, streaming=False, **_):
|
|
captured["extra_headers"] = extra_headers
|
|
raise RuntimeError("stop")
|
|
|
|
with patch.object(
|
|
a2a_main, "create_a2a_client", new=AsyncMock(side_effect=_capture)
|
|
):
|
|
with pytest.raises(RuntimeError, match="stop"):
|
|
async for _ in a2a_main.asend_message_streaming(
|
|
request=request,
|
|
api_base="http://upstream.local",
|
|
litellm_logging_obj=logging_obj,
|
|
):
|
|
pass
|
|
|
|
assert captured["extra_headers"]["X-LiteLLM-Trace-Id"] == "trace-from-logging"
|
|
|
|
|
|
def test_streaming_logging_obj_carries_call_type_into_model_call_details():
|
|
"""The streaming logging object is built by hand rather than through
|
|
``update_environment_variables``, which is the only place ``call_type`` normally
|
|
reaches ``model_call_details``. Callbacks read the call type from there, so
|
|
without this the streamed turn arrives at every logger with no call type at all
|
|
and OTel's GenAI metrics label it ``chat`` instead of ``invoke_agent``."""
|
|
from a2a.compat.v0_3.types import MessageSendParams, SendStreamingMessageRequest
|
|
|
|
from litellm.a2a_protocol.main import _build_streaming_logging_obj
|
|
|
|
request = SendStreamingMessageRequest(
|
|
id="rpc-call-type",
|
|
params=MessageSendParams(
|
|
message={"messageId": "m1", "role": "user", "parts": [{"kind": "text", "text": "hi"}]}
|
|
),
|
|
)
|
|
|
|
logging_obj = _build_streaming_logging_obj(
|
|
request=request,
|
|
agent_name="some-agent",
|
|
agent_id=None,
|
|
litellm_params=None,
|
|
metadata=None,
|
|
proxy_server_request=None,
|
|
)
|
|
|
|
assert logging_obj.model_call_details["call_type"] == "asend_message_streaming"
|