Record a chat span for every LLM call

The token-usage wrappers open one span per decorated invocation, so a
primary attempt, its retry and a fallback appear as siblings with the
provider that ran. Stream spans start on the first pull, carry tokens,
cost, time to first token and cache hits, and feed the GenAI metrics.

tests/tracing/test_llm_spans.py

@@ -0,0 +1,216 @@
+"""LLM calls become ``chat`` spans through the token-usage wrappers."""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+
+from docsgpt import tracing
+from docsgpt.cache import gen_cache, stream_cache
+from docsgpt.core.settings import settings
+from docsgpt.usage import gen_token_usage, stream_token_usage
+
+
+class _LLM:
+    provider_name = "openai"
+
+    def __init__(self, source=None):
+        self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
+        self.decoded_token = {"sub": "u1"}
+        self.user_api_key = None
+        self.agent_id = None
+        if source:
+            self._token_usage_source = source
+
+
+@pytest.fixture(autouse=True)
+def _env(monkeypatch):
+    monkeypatch.setattr(settings, "TRACES_ENABLED", True)
+    monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True)
+    with patch("docsgpt.usage._persist_call_usage", return_value=0.0012):
+        yield
+
+
+@pytest.fixture()
+def trace():
+    t = tracing.start_trace(source="stream", capture_otel_context=False)
+    with tracing.activate(t):
+        yield t
+
+
+@pytest.fixture()
+def metrics():
+    with patch("docsgpt.tracing.llm.record_llm_metrics") as rec:
+        yield rec
+
+
+class TestNonStreaming:
+    def test_span_with_usage_and_preview(self, trace, metrics):
+        @gen_token_usage
+        def _gen(self, model, messages, stream, tools, **kwargs):
+            return "the answer"
+
+        _gen(_LLM(), "gpt-4o", [{"role": "user", "content": "hi"}], False, None)
+        (span,) = trace.spans
+        assert span.kind == tracing.KIND_LLM
+        assert span.name == "chat gpt-4o"
+        assert span.status == "ok"
+        attrs = span.attributes
+        assert attrs["gen_ai.operation.name"] == "chat"
+        assert attrs["gen_ai.provider.name"] == "openai"
+        assert attrs["gen_ai.request.model"] == "gpt-4o"
+        assert attrs["gen_ai.usage.input_tokens"] > 0
+        assert attrs["gen_ai.usage.output_tokens"] > 0
+        assert attrs["docsgpt.token_source"] == "agent_stream"
+        assert attrs["docsgpt.cost_usd"] == 0.0012
+        assert span.previews["output"] == "the answer"
+        metrics.assert_called_once()
+        assert metrics.call_args.kwargs["error_type"] is None
+
+    def test_failure_marks_span_error(self, trace, metrics):
+        @gen_token_usage
+        def _gen(self, model, messages, stream, tools, **kwargs):
+            raise TimeoutError("slow")
+
+        with pytest.raises(TimeoutError):
+            _gen(_LLM(source="fallback"), "m", [], False, None)
+        (span,) = trace.spans
+        assert span.status == "error"
+        assert span.attributes["error.type"] == "TimeoutError"
+        assert span.attributes["docsgpt.token_source"] == "fallback"
+        assert metrics.call_args.kwargs["error_type"] == "TimeoutError"
+
+    def test_provider_reported_usage_is_flagged(self, trace, metrics):
+        llm = _LLM()
+
+        @gen_token_usage
+        def _gen(self, model, messages, stream, tools, **kwargs):
+            self._last_usage = {
+                "prompt_tokens": 100,
+                "completion_tokens": 7,
+                "prompt_tokens_details": {"cached_tokens": 40},
+            }
+            self._last_usage_claimed = False
+            return "x"
+
+        _gen(llm, "m", [], False, None)
+        attrs = trace.spans[0].attributes
+        assert attrs["gen_ai.usage.input_tokens"] == 100
+        assert attrs["gen_ai.usage.output_tokens"] == 7
+        assert attrs["gen_ai.usage.cache_read.input_tokens"] == 40
+        assert attrs["docsgpt.usage_estimated"] is False
+
+    def test_no_trace_no_span(self, metrics):
+        @gen_token_usage
+        def _gen(self, model, messages, stream, tools, **kwargs):
+            return "x"
+
+        assert _gen(_LLM(), "m", [], False, None) == "x"
+        metrics.assert_called_once()
+
+
+class TestStreaming:
+    def test_span_starts_on_first_next_not_on_call(self, trace, metrics):
+        @stream_token_usage
+        def _stream(self, model, messages, stream, tools, **kwargs):
+            yield "a"
+            yield {"type": "thought", "thought": "hmm"}
+            yield "b"
+
+        gen = _stream(_LLM(), "m", [], True, None)
+        assert trace.spans == []
+        assert list(gen) == ["a", {"type": "thought", "thought": "hmm"}, "b"]
+        (span,) = trace.spans
+        assert span.status == "ok"
+        assert span.attributes["docsgpt.ttft_ms"] is not None
+        assert span.attributes["docsgpt.stream"] is True
+        assert span.previews["output"] == "ab"
+
+    def test_abandoned_stream_is_cancelled(self, trace, metrics):
+        @stream_token_usage
+        def _stream(self, model, messages, stream, tools, **kwargs):
+            yield "a"
+            yield "b"
+
+        gen = _stream(_LLM(), "m", [], True, None)
+        next(gen)
+        gen.close()
+        assert trace.spans[0].status == "cancelled"
+
+    def test_failed_stream(self, trace, metrics):
+        @stream_token_usage
+        def _stream(self, model, messages, stream, tools, **kwargs):
+            yield "a"
+            raise ConnectionError("reset")
+
+        with pytest.raises(ConnectionError):
+            list(_stream(_LLM(), "m", [], True, None))
+        assert trace.spans[0].status == "error"
+
+    def test_primary_and_fallback_are_siblings(self, trace, metrics):
+        @stream_token_usage
+        def _primary(self, model, messages, stream, tools, **kwargs):
+            raise ConnectionError("down")
+            yield  # pragma: no cover
+
+        @stream_token_usage
+        def _fallback(self, model, messages, stream, tools, **kwargs):
+            yield "ok"
+
+        with tracing.span(tracing.KIND_AGENT, "agent"):
+            with pytest.raises(ConnectionError):
+                list(_primary(_LLM(), "m", [], True, None))
+            list(_fallback(_LLM(source="fallback"), "m2", [], True, None))
+        agent, primary, fallback = trace.spans
+        assert primary.parent_id == agent.id == fallback.parent_id
+        assert primary.status == "error"
+        assert fallback.attributes["docsgpt.token_source"] == "fallback"
+
+
+class _FakeRedis:
+    def __init__(self):
+        self.store = {}
+
+    def get(self, key):
+        return self.store.get(key)
+
+    def set(self, key, value, ex=None):
+        self.store[key] = value.encode("utf-8") if isinstance(value, str) else value
+
+    def delete(self, key):
+        self.store.pop(key, None)
+
+
+class TestCacheHits:
+    def test_gen_cache_hit_records_a_cached_span(self, trace, metrics):
+        redis = _FakeRedis()
+
+        @gen_cache
+        @gen_token_usage
+        def _gen(self, model, messages, stream, tools=None, **kwargs):
+            return "fresh"
+
+        with patch("docsgpt.cache.get_redis_instance", return_value=redis):
+            _gen(_LLM(), "m", [{"role": "user", "content": "q"}], False)
+            _gen(_LLM(), "m", [{"role": "user", "content": "q"}], False)
+        first, second = trace.spans
+        assert first.attributes.get("docsgpt.cache_hit") is None
+        assert second.attributes["docsgpt.cache_hit"] is True
+        assert second.status == "ok"
+
+    def test_stream_cache_hit_flags_the_open_span(self, trace, metrics, monkeypatch):
+        monkeypatch.setattr("docsgpt.cache.time.sleep", lambda _s: None)
+        redis = _FakeRedis()
+
+        @stream_token_usage
+        @stream_cache
+        def _stream(self, model, messages, stream, tools=None, **kwargs):
+            yield "fresh"
+
+        with patch("docsgpt.cache.get_redis_instance", return_value=redis):
+            list(_stream(_LLM(), "m", [{"role": "user", "content": "q"}], True, None))
+            list(_stream(_LLM(), "m", [{"role": "user", "content": "q"}], True, None))
+        first, second = trace.spans
+        assert first.attributes.get("docsgpt.cache_hit") is None
+        assert second.attributes["docsgpt.cache_hit"] is True