diff --git a/tests/test_litellm/tracing/test_store.py b/tests/test_litellm/tracing/test_store.py new file mode 100644 index 00000000000..40f44e11800 --- /dev/null +++ b/tests/test_litellm/tracing/test_store.py @@ -0,0 +1,349 @@ +""" +Tests for the pure read-side helpers in litellm/tracing/store.py (no ClickHouse needed). +""" + +import os +import sys +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +sys.path.insert(0, os.path.abspath("../../..")) + +import pytest + +from litellm.tracing.store import ( + ClickHouseTraceStore, + agent_nodes, + decode_cursor, + encode_cursor, + span_from_row, + trace_from_rows, + trace_summary_from_row, +) + +T0 = 1_790_742_989_000_000_000 # ns +MS = 1_000_000 + + +def _row( + span_id: str, + parent: str, + name: str, + type_: str, + agent: str, + start_ms: float = 0, + duration_ms: float = 10, + status: str = "STATUS_CODE_OK", + **extra: Any, +) -> dict[str, Any]: + """A TRACE_SPANS_SQL result row; `s_*` columns empty unless the span joined a spend log.""" + return { + "span_id": span_id, + "parent_span_id": parent, + "name": name, + "type": type_, + "agent": agent, + "status": status, + "start_ns": T0 + int(start_ms * MS), + "duration_ns": int(duration_ms * MS), + "service": "agent-demo", + "input_preview": f"input of {name}", + "model": "", + "input_tokens": 0, + "output_tokens": 0, + "litellm_request_id": "", + "s_request_id": "", + **extra, + } + + +def _llm_row(span_id: str, parent: str, agent: str, request_id: str, spend: float, start_ms: float = 1) -> dict: + return _row( + span_id, + parent, + "ChatOpenAI", + "llm", + agent, + start_ms=start_ms, + duration_ms=100, + model="claude-sonnet-4-5", + input_tokens=100, + output_tokens=20, + litellm_request_id=request_id, + s_request_id=request_id, + s_model="claude-sonnet-4-5", + s_model_group="claude-sonnet-4-5", + s_provider="anthropic", + s_api_base="https://api.anthropic.com", + s_key_alias="research-bot", + s_team_alias="research-agents", + s_spend=spend, + s_prompt_tokens=100, + s_completion_tokens=20, + s_cache_read_tokens=5, + s_cache_write_tokens=0, + s_start_ms=T0 // MS, + s_end_ms=T0 // MS + 100, + s_ttft_start_ms=T0 // MS + 30, + s_status="success", + ) + + +def _deep_agent_rows(researcher_invocations: int = 1) -> list[dict[str, Any]]: + """root agent -> llm, task tool -> researcher subagent (N times) -> llm + search_docs tool.""" + rows = [ + _row("root", "", "deep_research_agent", "agent", "deep_research_agent", duration_ms=1000), + _llm_row("llm-root", "root", "deep_research_agent", "chatcmpl-root", 0.01), + _row("task", "root", "task", "tool", "deep_research_agent", start_ms=200, duration_ms=700), + ] + for i in range(researcher_invocations): + rows += [ + _row(f"res-{i}", "task", "researcher", "agent", "researcher", start_ms=201, duration_ms=5), + _llm_row(f"res-llm-{i}", f"res-{i}", "researcher", f"chatcmpl-res-{i}", 0.001, start_ms=202), + _row(f"res-tool-{i}", f"res-{i}", "search_docs", "tool", "researcher", start_ms=203, duration_ms=1), + _row(f"res-mw-{i}", f"res-{i}", "FilesystemMiddleware.wrap_model_call", "framework", "researcher"), + ] + return rows + + +# ---------------------------------------------------------------- trace_from_rows + + +def test_empty_rows_is_none(): + assert trace_from_rows("abc", []) is None + + +def test_litellm_request_joined_onto_llm_spans_only(): + trace = trace_from_rows("t1", _deep_agent_rows()) + assert trace is not None + spans = {s["span_id"]: s for s in trace["spans"]} + llm = spans["llm-root"]["litellm"] + assert llm == { + "request_id": "chatcmpl-root", + "model": "claude-sonnet-4-5", + "model_group": "claude-sonnet-4-5", + "provider": "anthropic", + "key_alias": "research-bot", + "team_alias": "research-agents", + "spend": 0.01, + "prompt_tokens": 100, + "completion_tokens": 20, + "cache_read_tokens": 5, + "cache_write_tokens": 0, + "latency_ms": 100, + "ttft_ms": 30, + "status": "success", + } + assert spans["task"]["litellm"] is None + assert spans["root"]["litellm"] is None + + +def test_llm_span_without_spend_log_has_no_litellm(): + rows = [_row("root", "", "a", "agent", "a"), _row("llm", "root", "ChatOpenAI", "llm", "a", litellm_request_id="x")] + trace = trace_from_rows("t1", rows) + assert trace is not None + assert trace["spans"][1]["litellm"] is None + assert trace["summary"]["spend"] == 0 + + +def test_summary_totals(): + trace = trace_from_rows("t1", _deep_agent_rows()) + assert trace is not None + summary = trace["summary"] + assert summary["trace_id"] == "t1" + assert summary["name"] == "deep_research_agent" + assert summary["service"] == "agent-demo" + assert summary["input_preview"] == "input of deep_research_agent" + assert summary["status"] == "ok" + assert summary["span_count"] == 7 + assert summary["agent_count"] == 2 + assert summary["llm_calls"] == 2 + assert summary["tool_calls"] == 2 + assert summary["error_count"] == 0 + assert (summary["input_tokens"], summary["output_tokens"]) == (200, 40) + assert summary["spend"] == pytest.approx(0.011) + assert summary["models"] == ["claude-sonnet-4-5"] + assert summary["duration_ms"] == 1000 + assert summary["start_time"].startswith("2026-09-30T") + + +def test_error_count_counts_error_spans(): + rows = _deep_agent_rows() + rows[2]["status"] = "STATUS_CODE_ERROR" + trace = trace_from_rows("t1", rows) + assert trace is not None + assert trace["summary"]["error_count"] == 1 + assert trace["summary"]["status"] == "ok" # root span status; the UI uses error_count for "failed" + assert trace["spans"][2]["status"] == "error" + + +def test_offsets_are_relative_to_trace_start_in_ms(): + trace = trace_from_rows("t1", _deep_agent_rows()) + assert trace is not None + spans = {s["span_id"]: s for s in trace["spans"]} + assert spans["root"]["start_offset_ms"] == 0 + assert spans["task"]["start_offset_ms"] == 200 + assert spans["task"]["duration_ms"] == 700 + assert spans["root"]["parent_span_id"] is None + assert spans["task"]["parent_span_id"] == "root" + + +def test_span_from_row_optional_fields(): + span = span_from_row(_row("s", "", "x", "chain", "a", status="STATUS_CODE_UNSET"), T0) + assert (span["model"], span["parent_span_id"], span["status"], span["litellm"]) == (None, None, "unset", None) + + +def test_ttft_is_none_without_completion_start_time(): + row = _llm_row("l", "", "a", "req", 0.1) + row["s_ttft_start_ms"] = 0 + assert span_from_row(row, T0)["litellm"]["ttft_ms"] is None # type: ignore[index] + + +# ---------------------------------------------------------------- agent_nodes + + +def test_agent_nodes_parent_and_per_agent_counts(): + trace = trace_from_rows("t1", _deep_agent_rows()) + assert trace is not None + assert trace["agents"] == [ + { + "name": "deep_research_agent", + "parent_agent": None, + "invocations": 1, + "llm_calls": 1, + "tool_calls": 1, + "spend": 0.01, + "duration_ms": 1000, + }, + { + "name": "researcher", + "parent_agent": "deep_research_agent", + "invocations": 1, + "llm_calls": 1, + "tool_calls": 1, + "spend": 0.001, + "duration_ms": 5, + }, + ] + + +def test_200_subagent_invocations_aggregate_into_one_node(): + trace = trace_from_rows("t1", _deep_agent_rows(researcher_invocations=200)) + assert trace is not None + assert [a["name"] for a in trace["agents"]] == ["deep_research_agent", "researcher"] + researcher = trace["agents"][1] + assert researcher["parent_agent"] == "deep_research_agent" + assert researcher["invocations"] == 200 + assert researcher["llm_calls"] == 200 + assert researcher["tool_calls"] == 200 + assert researcher["spend"] == pytest.approx(0.2) + assert researcher["duration_ms"] == pytest.approx(1000) + assert trace["summary"]["agent_count"] == 2 + assert trace["summary"]["span_count"] == 3 + 4 * 200 + + +def test_parent_agent_skips_same_name_ancestors(): + """A recursive agent (researcher -> researcher) still reports the nearest *different* agent.""" + rows = [ + _row("root", "", "lead", "agent", "lead"), + _row("r1", "root", "researcher", "agent", "researcher"), + _row("r2", "r1", "researcher", "agent", "researcher"), + ] + spans = [span_from_row(r, T0) for r in rows] + nodes = {n["name"]: n for n in agent_nodes(spans)} + assert nodes["researcher"]["parent_agent"] == "lead" + assert nodes["researcher"]["invocations"] == 2 + + +def test_agent_nodes_ignores_spans_of_unknown_agents(): + spans = [span_from_row(_row("t", "", "tool", "tool", "ghost"), T0)] + assert agent_nodes(spans) == [] + + +# ---------------------------------------------------------------- list helpers + + +def test_cursor_round_trip(): + cursor = encode_cursor(1790742989377, "4bad42b84e9de3ba46fc870185f8f023") + assert decode_cursor(cursor) == (1790742989377, "4bad42b84e9de3ba46fc870185f8f023") + assert decode_cursor(None) == (0, "") + assert decode_cursor("") == (0, "") + + +def test_trace_summary_from_row(): + summary = trace_summary_from_row( + { + "trace_id": "t1", + "name": "deep_research_agent", + "service": "agent-demo", + "input_preview": "hi", + "start_ms": 1790742989377, + "duration_ms": 51385, + "status": "STATUS_CODE_OK", + "span_count": "126", + "agent_count": "2", + "llm_calls": "7", + "tool_calls": "26", + "error_count": "1", + "input_tokens": "30175", + "output_tokens": "2620", + "models": ["claude-sonnet-4-5"], + "spend": None, + } + ) + assert summary["status"] == "ok" + assert (summary["span_count"], summary["error_count"], summary["spend"]) == (126, 1, 0.0) + assert summary["start_time"] == "2026-09-30T04:36:29.377000+00:00" + + +@pytest.mark.asyncio +async def test_list_traces_sets_next_cursor_on_full_page(): + client = MagicMock() + row = { + "trace_id": "t2", + "name": "a", + "service": "s", + "input_preview": "", + "start_ms": 1000, + "duration_ms": 1, + "status": "STATUS_CODE_OK", + "span_count": 1, + "agent_count": 1, + "llm_calls": 0, + "tool_calls": 0, + "error_count": 0, + "input_tokens": 0, + "output_tokens": 0, + "models": [], + "spend": 0, + } + client.query = AsyncMock(return_value=[row, {**row, "trace_id": "t1", "start_ms": 900}]) + store = ClickHouseTraceStore(client) + scope = {"team_ids": ["team-a"], "api_key_hash": ""} + + page = await store.list_traces(scope, 0, 2000, limit=2) # type: ignore[arg-type] + assert [t["trace_id"] for t in page["data"]] == ["t2", "t1"] + assert page["next_cursor"] is not None + assert decode_cursor(page["next_cursor"]) == (900, "t1") + params = client.query.call_args.args[1] + assert params["team_ids"] == ["team-a"] and params["limit"] == 2 and params["cursor_ms"] == 0 + + page = await store.list_traces(scope, 0, 2000, cursor=page["next_cursor"], limit=3) # type: ignore[arg-type] + assert page["next_cursor"] is None + assert client.query.call_args.args[1]["cursor_trace_id"] == "t1" + + +@pytest.mark.asyncio +async def test_get_span_not_found_and_found(): + client = MagicMock() + client.query = AsyncMock(return_value=[]) + store = ClickHouseTraceStore(client) + scope = {"team_ids": [], "api_key_hash": ""} + assert await store.get_span("t", "s", scope) is None # type: ignore[arg-type] + client.query = AsyncMock(return_value=[{"span_id": "s", "input": "i", "output": "o", "attributes": {"k": "v"}}]) + assert await store.get_span("t", "s", scope) == { # type: ignore[arg-type] + "span_id": "s", + "input": "i", + "output": "o", + "attributes": {"k": "v"}, + }