mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 18:15:33 +00:00
arc53-machine
9f7f0b2f1e
A failed trace-summary lookup now leaves the Logs page intact without chips. Tool-call counts include only calls that ran, not their paused, denied or skipped records. A local guardrail that fires unchanged on every streamed segment is recorded once, so it cannot use up the span cap.
398 lines
15 KiB
Python
398 lines
15 KiB
Python
"""Tests for the execution-trace recorder in ``docsgpt.tracing``."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
import pytest
|
|
|
|
from docsgpt import tracing
|
|
from docsgpt.core.settings import settings
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _tracing_on(monkeypatch):
|
|
monkeypatch.setattr(settings, "TRACES_ENABLED", True)
|
|
monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True)
|
|
monkeypatch.setattr(settings, "TRACES_MAX_SPANS", 500)
|
|
monkeypatch.setattr(settings, "TRACES_PREVIEW_CHARS", 2000)
|
|
|
|
|
|
def _by_name(trace):
|
|
return {s.name: s for s in trace.spans}
|
|
|
|
|
|
class TestNoActiveTrace:
|
|
def test_span_calls_are_noops_without_a_trace(self):
|
|
assert tracing.current_trace() is None
|
|
with tracing.span(tracing.KIND_LLM, "chat gpt") as s:
|
|
s.set(foo=1)
|
|
s.preview("output", "hi")
|
|
handle = tracing.start_span(tracing.KIND_TOOL, "tool")
|
|
handle.end(status="error")
|
|
tracing.bind(message_id="m")
|
|
|
|
def test_start_trace_returns_none_when_disabled(self, monkeypatch):
|
|
monkeypatch.setattr(settings, "TRACES_ENABLED", False)
|
|
assert tracing.start_trace(source="stream") is None
|
|
with tracing.activate(None):
|
|
assert tracing.current_trace() is None
|
|
|
|
|
|
class TestNesting:
|
|
def test_containers_push_and_leaves_record_parent(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
with tracing.span(tracing.KIND_AGENT, "agent"):
|
|
with tracing.span(tracing.KIND_LLM, "llm-1"):
|
|
pass
|
|
with tracing.span(tracing.KIND_TOOL, "tool"):
|
|
with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"):
|
|
with tracing.span(tracing.KIND_EMBEDDING, "embed"):
|
|
pass
|
|
with tracing.span(tracing.KIND_LLM, "llm-2"):
|
|
pass
|
|
spans = _by_name(trace)
|
|
assert spans["agent"].parent_id is None
|
|
assert spans["llm-1"].parent_id == spans["agent"].id
|
|
assert spans["tool"].parent_id == spans["agent"].id
|
|
assert spans["retrieval"].parent_id == spans["tool"].id
|
|
assert spans["embed"].parent_id == spans["retrieval"].id
|
|
assert spans["llm-2"].parent_id == spans["agent"].id
|
|
assert all(s.status == "ok" for s in trace.spans)
|
|
|
|
def test_leaf_span_does_not_become_a_parent(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
llm = tracing.start_span(tracing.KIND_LLM, "llm")
|
|
with tracing.span(tracing.KIND_TOOL, "tool"):
|
|
pass
|
|
llm.end()
|
|
spans = _by_name(trace)
|
|
assert spans["tool"].parent_id is None
|
|
|
|
def test_explicit_parent_wins(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
outer = tracing.start_span(tracing.KIND_RETRIEVAL, "outer")
|
|
with tracing.span(tracing.KIND_AGENT, "agent"):
|
|
child = tracing.start_span(tracing.KIND_SEARCH, "child", parent=outer)
|
|
child.end()
|
|
outer.end()
|
|
assert _by_name(trace)["child"].parent_id == _by_name(trace)["outer"].id
|
|
|
|
def test_interleaved_generators_nest_by_start_order(self):
|
|
"""An agent generator suspended at yield keeps its children nested."""
|
|
trace = tracing.start_trace(source="stream")
|
|
|
|
def llm_stream():
|
|
handle = tracing.start_span(tracing.KIND_LLM, "llm")
|
|
try:
|
|
yield "a"
|
|
yield "b"
|
|
finally:
|
|
handle.end()
|
|
|
|
def agent_gen():
|
|
handle = tracing.start_span(tracing.KIND_AGENT, "agent")
|
|
try:
|
|
yield from llm_stream()
|
|
with tracing.span(tracing.KIND_TOOL, "tool"):
|
|
pass
|
|
yield "c"
|
|
finally:
|
|
handle.end()
|
|
|
|
with tracing.activate(trace):
|
|
assert list(agent_gen()) == ["a", "b", "c"]
|
|
spans = _by_name(trace)
|
|
assert spans["llm"].parent_id == spans["agent"].id
|
|
assert spans["tool"].parent_id == spans["agent"].id
|
|
|
|
def test_out_of_order_end_cancels_abandoned_children(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
agent = tracing.start_span(tracing.KIND_AGENT, "agent")
|
|
tool = tracing.start_span(tracing.KIND_TOOL, "tool")
|
|
agent.end()
|
|
after = tracing.start_span(tracing.KIND_LLM, "after")
|
|
after.end()
|
|
tool.end() # late end is ignored
|
|
spans = _by_name(trace)
|
|
assert spans["tool"].status == "cancelled"
|
|
assert spans["agent"].status == "ok"
|
|
assert spans["after"].parent_id is None
|
|
|
|
def test_exception_marks_span_error_and_propagates(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
with pytest.raises(ValueError):
|
|
with tracing.span(tracing.KIND_TOOL, "tool"):
|
|
raise ValueError("boom")
|
|
span = trace.spans[0]
|
|
assert span.status == "error"
|
|
assert span.attributes["error.type"] == "ValueError"
|
|
assert span.error == "boom"
|
|
|
|
def test_generator_exit_marks_span_cancelled(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
|
|
def gen():
|
|
with tracing.span(tracing.KIND_AGENT, "agent"):
|
|
yield 1
|
|
yield 2
|
|
|
|
with tracing.activate(trace):
|
|
g = gen()
|
|
next(g)
|
|
g.close()
|
|
assert trace.spans[0].status == "cancelled"
|
|
|
|
|
|
class TestFinish:
|
|
def test_finish_cancels_open_spans_and_ignores_late_ends(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
open_span = tracing.start_span(tracing.KIND_AGENT, "agent")
|
|
trace.finish()
|
|
assert open_span.status == "cancelled"
|
|
duration = open_span.duration_ms
|
|
open_span.end(status="ok")
|
|
assert open_span.status == "cancelled"
|
|
assert open_span.duration_ms == duration
|
|
|
|
def test_finish_status_defaults_to_error_when_a_top_level_span_failed(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
tracing.start_span(tracing.KIND_AGENT, "agent").end(status="error")
|
|
trace.finish()
|
|
assert trace.status == "error"
|
|
|
|
def test_explicit_finish_status(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
trace.finish(status="paused")
|
|
assert trace.status == "paused"
|
|
|
|
def test_spans_after_finish_are_dropped(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
trace.finish()
|
|
with tracing.activate(trace):
|
|
tracing.start_span(tracing.KIND_LLM, "late").end()
|
|
assert trace.spans == []
|
|
|
|
def test_summary_rolls_up_llm_tool_retrieval(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
with tracing.span(tracing.KIND_AGENT, "agent"):
|
|
with tracing.span(
|
|
tracing.KIND_LLM,
|
|
"chat",
|
|
attributes={
|
|
"gen_ai.usage.input_tokens": 10,
|
|
"gen_ai.usage.output_tokens": 4,
|
|
},
|
|
):
|
|
pass
|
|
with tracing.span(tracing.KIND_RETRIEVAL, "r"):
|
|
with tracing.span(tracing.KIND_RETRIEVAL, "inner"):
|
|
pass
|
|
with tracing.span(tracing.KIND_TOOL, "t") as t:
|
|
t.end(status="error")
|
|
trace.finish()
|
|
summary = trace.summary()
|
|
assert summary["llm_calls"] == 1
|
|
assert summary["tool_calls"] == 1
|
|
assert summary["input_tokens"] == 10
|
|
assert summary["output_tokens"] == 4
|
|
assert summary["retrieval_calls"] == 1
|
|
assert summary["errors"] == 1
|
|
assert summary["retrieval_ms"] >= 0
|
|
|
|
|
|
class TestCap:
|
|
def test_span_cap_counts_dropped(self, monkeypatch):
|
|
monkeypatch.setattr(settings, "TRACES_MAX_SPANS", 10)
|
|
trace = tracing.start_trace(source="graph_extraction")
|
|
with tracing.activate(trace):
|
|
for i in range(15):
|
|
with tracing.span(tracing.KIND_LLM, f"llm-{i}"):
|
|
pass
|
|
assert len(trace.spans) == 10
|
|
assert trace.dropped_spans == 5
|
|
|
|
|
|
class TestThreads:
|
|
def test_wrap_carries_trace_and_parent_into_pool_threads(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
|
|
def work(i):
|
|
with tracing.span(tracing.KIND_SEARCH, f"search-{i}"):
|
|
with tracing.span(tracing.KIND_EMBEDDING, f"embed-{i}"):
|
|
pass
|
|
return threading.get_ident()
|
|
|
|
with tracing.activate(trace):
|
|
with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"):
|
|
with ThreadPoolExecutor(max_workers=3) as pool:
|
|
list(pool.map(tracing.wrap(work), range(3)))
|
|
spans = _by_name(trace)
|
|
for i in range(3):
|
|
assert spans[f"search-{i}"].parent_id == spans["retrieval"].id
|
|
assert spans[f"embed-{i}"].parent_id == spans["retrieval"].id
|
|
|
|
def test_pool_threads_without_wrap_see_no_trace(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
seen = []
|
|
with tracing.activate(trace):
|
|
t = threading.Thread(target=lambda: seen.append(tracing.current_trace()))
|
|
t.start()
|
|
t.join()
|
|
assert seen == [None]
|
|
|
|
def test_wrap_without_trace_is_passthrough(self):
|
|
fn = tracing.wrap(lambda x: x + 1)
|
|
assert fn(1) == 2
|
|
|
|
|
|
class TestBind:
|
|
def test_bind_sets_ids_on_active_trace(self):
|
|
trace = tracing.start_trace(source="stream", request_id="r1")
|
|
with tracing.activate(trace):
|
|
tracing.bind(message_id="m1", conversation_id="c1", unknown="x")
|
|
assert trace.request_id == "r1"
|
|
assert trace.message_id == "m1"
|
|
assert trace.conversation_id == "c1"
|
|
|
|
def test_bind_if_unset_keeps_first_value(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
tracing.bind_if_unset(activity_id="a1")
|
|
tracing.bind_if_unset(activity_id="a2")
|
|
assert trace.activity_id == "a1"
|
|
|
|
|
|
class TestPreviews:
|
|
def test_preview_redacts_secrets_and_bounds_strings(self, monkeypatch):
|
|
monkeypatch.setattr(settings, "TRACES_PREVIEW_CHARS", 100)
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
with tracing.span(tracing.KIND_TOOL, "tool") as s:
|
|
s.preview("arguments", {"api_key": "sk-123", "q": "x" * 500})
|
|
s.preview("result", "y" * 500 + "\x00")
|
|
preview = trace.spans[0].previews
|
|
assert preview["arguments"]["api_key"] == "[REDACTED]"
|
|
assert len(preview["arguments"]["q"]) < 200
|
|
assert preview["result"].endswith("…")
|
|
assert "\x00" not in preview["result"]
|
|
|
|
def test_preview_skipped_when_capture_disabled(self, monkeypatch):
|
|
monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", False)
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
with tracing.span(tracing.KIND_TOOL, "tool") as s:
|
|
s.preview("result", "secret stuff")
|
|
assert trace.spans[0].previews == {}
|
|
|
|
def test_guardrail_trigger_strips_all_previews(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
with tracing.span(tracing.KIND_TOOL, "tool") as s:
|
|
s.preview("result", "leaked")
|
|
tracing.mark_content_blocked()
|
|
trace.finish()
|
|
record = trace.to_record()
|
|
assert all(not s.get("preview") for s in record["spans"])
|
|
assert trace.content_blocked is True
|
|
|
|
def test_huge_structures_collapse_to_a_string(self, monkeypatch):
|
|
monkeypatch.setattr(settings, "TRACES_PREVIEW_CHARS", 100)
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
with tracing.span(tracing.KIND_TOOL, "tool") as s:
|
|
s.preview("result", [{"k": "v" * 50} for _ in range(100)])
|
|
value = trace.spans[0].previews["result"]
|
|
assert isinstance(value, str)
|
|
assert len(value) <= 101
|
|
|
|
|
|
class TestRecord:
|
|
def test_to_record_shape(self):
|
|
trace = tracing.start_trace(
|
|
source="stream", request_id="r", user_id="u", agent_id="a"
|
|
)
|
|
with tracing.activate(trace):
|
|
with tracing.span(
|
|
tracing.KIND_LLM, "chat m", attributes={"gen_ai.request.model": "m"}
|
|
):
|
|
pass
|
|
trace.finish()
|
|
record = trace.to_record()
|
|
assert record["source"] == "stream"
|
|
assert record["request_id"] == "r"
|
|
assert record["status"] == "ok"
|
|
assert record["span_count"] == 1
|
|
span = record["spans"][0]
|
|
assert set(span) >= {
|
|
"id", "parent_id", "kind", "name", "status", "offset_ms",
|
|
"duration_ms", "attributes",
|
|
}
|
|
assert span["attributes"]["gen_ai.request.model"] == "m"
|
|
assert span["offset_ms"] >= 0
|
|
|
|
|
|
class TestOutcome:
|
|
def test_outcome_is_used_when_no_status_given(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
trace.outcome = "paused"
|
|
trace.finish()
|
|
assert trace.status == "paused"
|
|
|
|
def test_explicit_status_beats_outcome(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
trace.outcome = "paused"
|
|
trace.finish(status="error")
|
|
assert trace.status == "error"
|
|
|
|
|
|
class TestRecordQuery:
|
|
def test_first_span_query_is_copied_into_summary(self):
|
|
trace = tracing.start_trace(source="search")
|
|
with tracing.activate(trace):
|
|
with tracing.span(tracing.KIND_RETRIEVAL, "retrieval") as s:
|
|
s.preview("query", "how do I deploy")
|
|
trace.finish()
|
|
assert trace.to_record()["summary"]["query"] == "how do I deploy"
|
|
|
|
def test_query_is_dropped_when_content_is_blocked(self):
|
|
trace = tracing.start_trace(source="search")
|
|
with tracing.activate(trace):
|
|
with tracing.span(tracing.KIND_RETRIEVAL, "retrieval") as s:
|
|
s.preview("query", "leaky")
|
|
tracing.mark_content_blocked()
|
|
trace.finish()
|
|
assert "query" not in trace.to_record()["summary"]
|
|
|
|
|
|
class TestToolCount:
|
|
def test_only_tool_calls_that_ran_are_counted(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
tracing.start_span(tracing.KIND_TOOL, "paused").end(tracing.STATUS_PENDING)
|
|
tracing.start_span(tracing.KIND_TOOL, "denied").end(tracing.STATUS_DENIED)
|
|
tracing.start_span(tracing.KIND_TOOL, "skipped").end(tracing.STATUS_SKIPPED)
|
|
tracing.start_span(tracing.KIND_TOOL, "ran").end()
|
|
tracing.start_span(tracing.KIND_TOOL, "failed").end(tracing.STATUS_ERROR)
|
|
trace.finish()
|
|
assert trace.summary()["tool_calls"] == 2
|
|
|
|
|
|
class TestFirstOccurrence:
|
|
def test_key_is_new_once_per_trace(self):
|
|
trace = tracing.start_trace(source="stream")
|
|
with tracing.activate(trace):
|
|
assert tracing.first_occurrence(("a", 1)) is True
|
|
assert tracing.first_occurrence(("a", 1)) is False
|
|
assert tracing.first_occurrence(("a", 2)) is True
|
|
assert tracing.first_occurrence(("a", 3)) is False # no active trace
|