diff --git a/docsgpt/cache.py b/docsgpt/cache.py index 4ee069de..820b20c2 100644 --- a/docsgpt/cache.py +++ b/docsgpt/cache.py @@ -8,6 +8,7 @@ from threading import Lock import redis from docsgpt.core.settings import settings +from docsgpt.tracing.llm import CACHE_HIT_ATTR, record_cached_gen from docsgpt.utils import get_hash logger = logging.getLogger(__name__) @@ -240,6 +241,7 @@ def gen_cache(func): if cached_response: decoded = cached_response.decode("utf-8") if not _is_stream_payload(decoded): + record_cached_gen(self, model, decoded) return decoded except Exception as e: logger.error(f"Error getting cached response: {e}", exc_info=True) @@ -295,6 +297,12 @@ def stream_cache(func): if cached_chunks is not None: logger.info(f"Cache hit for stream key: {cache_key}") + # ``stream_token_usage`` wraps this cache and owns + # the call's span; flag it as served from cache. + try: + setattr(self, CACHE_HIT_ATTR, True) + except AttributeError: + pass for chunk in cached_chunks: yield chunk time.sleep(0.03) # Simulate streaming delay diff --git a/docsgpt/tracing/llm.py b/docsgpt/tracing/llm.py new file mode 100644 index 00000000..596818dd --- /dev/null +++ b/docsgpt/tracing/llm.py @@ -0,0 +1,119 @@ +"""``chat`` spans and GenAI metrics for LLM calls. + +Called from the token-usage wrappers in ``docsgpt/usage.py``: one span per +decorated invocation, so a primary attempt, its same-provider retry and a +fallback each get their own span with the provider that actually ran. +""" + +from __future__ import annotations + +from typing import Any, Dict, Iterable, Optional + +from docsgpt.tracing import core +from docsgpt.tracing.otel import provider_name, record_llm_metrics + +#: LLM attribute a cache wrapper sets when it served the call from Redis. +CACHE_HIT_ATTR = "_trace_cache_hit" + + +def start_llm_span(llm: Any, model: Optional[str], *, stream: bool, tools: Any = None): + """Open a ``chat {model}`` span for a call on ``llm`` (no-op without a trace).""" + if core.current_trace() is None: + return core.NOOP_SPAN + attributes = { + "gen_ai.operation.name": "chat", + "gen_ai.provider.name": provider_name(getattr(llm, "provider_name", None)), + "gen_ai.request.model": str(model) if model else None, + "docsgpt.token_source": getattr(llm, "_token_usage_source", None) or "agent_stream", + "docsgpt.stream": bool(stream), + "docsgpt.tool_count": len(tools) if tools else None, + } + return core.start_span( + core.KIND_LLM, + f"chat {model}" if model else "chat", + attributes={k: v for k, v in attributes.items() if v is not None}, + ) + + +def output_text(chunks: Iterable[Any]) -> str: + """Concatenate the text deltas of a streamed or returned response.""" + if isinstance(chunks, str): + return chunks + return "".join(chunk for chunk in chunks if isinstance(chunk, str)) + + +def finish_llm_call( + span: Any, + llm: Any, + model: Optional[str], + call_usage: Dict[str, Any], + *, + duration_ms: int, + error: Optional[BaseException], + completed: bool = True, + ttft_ms: Optional[int] = None, + cost_usd: Optional[float] = None, + estimated: bool = True, + output: Optional[str] = None, +) -> None: + """Close the call's span and record the GenAI client metrics. + + Args: + span: The span from :func:`start_llm_span`. + llm: The LLM instance that ran the call. + model: Model the call was made with. + call_usage: Final token counts (provider-reported when available). + duration_ms: Provider time for the call. + error: The exception the call raised, if any. + completed: False when a stream was abandoned before it finished. + ttft_ms: Time to first streamed chunk. + cost_usd: Cost priced for the call, when known. + estimated: True when token counts are local estimates. + output: Response text for the preview. + """ + cache_hit = bool(getattr(llm, CACHE_HIT_ATTR, False)) + try: + setattr(llm, CACHE_HIT_ATTR, False) + except AttributeError: + pass + record_llm_metrics( + provider=getattr(llm, "provider_name", None), + model=str(model) if model else None, + input_tokens=call_usage.get("prompt_tokens", 0), + output_tokens=call_usage.get("generated_tokens", 0), + duration_s=max(duration_ms, 0) / 1000.0, + error_type=type(error).__name__ if error is not None else None, + ) + if not span: + return + if output: + span.preview("output", output) + span.end( + None if completed or error is not None else core.STATUS_CANCELLED, + error=error, + attributes={ + "gen_ai.usage.input_tokens": int(call_usage.get("prompt_tokens") or 0), + "gen_ai.usage.output_tokens": int(call_usage.get("generated_tokens") or 0), + "gen_ai.usage.cache_read.input_tokens": call_usage.get("cached_tokens"), + "gen_ai.usage.cache_creation.input_tokens": call_usage.get("cache_write_tokens"), + "docsgpt.usage_estimated": estimated, + "docsgpt.provider_ms": duration_ms, + "docsgpt.ttft_ms": ttft_ms, + "docsgpt.cost_usd": cost_usd if isinstance(cost_usd, (int, float)) else None, + "docsgpt.cache_hit": True if cache_hit else None, + }, + ) + + +def record_cached_gen(llm: Any, model: Optional[str], output: Optional[str]) -> None: + """Record a non-streaming call answered from the response cache. + + The gen cache wraps the usage wrapper, so a hit never reaches it; this + records the zero-cost call so the trace still shows it happened. + """ + span = start_llm_span(llm, model, stream=False) + if not span: + return + if output: + span.preview("output", output) + span.end(attributes={"docsgpt.cache_hit": True}) diff --git a/docsgpt/usage.py b/docsgpt/usage.py index 8ac6c436..307cd4cd 100644 --- a/docsgpt/usage.py +++ b/docsgpt/usage.py @@ -3,6 +3,7 @@ import time from typing import Any, Dict from docsgpt.pricing import compute_cost_usd +from docsgpt.tracing.llm import finish_llm_call, output_text, start_llm_span from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository from docsgpt.storage.db.session import db_session from docsgpt.utils import num_tokens_from_object_or_list, num_tokens_from_string @@ -108,9 +109,12 @@ def _persist_call_usage(llm, call_usage, *, duration_ms=None, ttft_ms=None): duration_ms: Wall-clock for the call, measured by the wrapper. ttft_ms: Time to the first streamed chunk; None for a non-streaming call and for a stream that failed before yielding anything. + + Returns: + The call's priced cost in USD, or None when no row was written. """ if call_usage["prompt_tokens"] == 0 and call_usage["generated_tokens"] == 0: - return + return None decoded_token = getattr(llm, "decoded_token", None) user_id = ( decoded_token.get("sub") if isinstance(decoded_token, dict) else None @@ -126,7 +130,7 @@ def _persist_call_usage(llm, call_usage, *, duration_ms=None, ttft_ms=None): "source": getattr(llm, "_token_usage_source", "agent_stream"), }, ) - return + return None model_id = getattr(llm, "_canonical_model_id", None) # Bring-your-own models run on the user's own provider key: recorded, never priced. if getattr(llm, "_is_byom", False): @@ -161,6 +165,7 @@ def _persist_call_usage(llm, call_usage, *, duration_ms=None, ttft_ms=None): ) except Exception: logger.exception("token_usage persist failed") + return cost def _call_cost_usd(model_id, call_usage) -> float: @@ -257,8 +262,10 @@ def gen_token_usage(func): usage_attachments=usage_attachments, **kwargs, ) + span = start_llm_span(self, model, stream=False, tools=tools) started_at = time.monotonic() error: BaseException | None = None + result = None try: result = func(self, model, messages, stream, tools, **kwargs) call_usage["generated_tokens"] += _count_tokens(result) @@ -268,11 +275,23 @@ def gen_token_usage(func): raise finally: duration_ms = int((time.monotonic() - started_at) * 1000) + estimated_usage = call_usage call_usage = _prefer_provider_usage(self, call_usage) self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"] self.token_usage["generated_tokens"] += call_usage["generated_tokens"] # A non-streaming call has no first-token moment; ttft stays NULL. - _persist_call_usage(self, call_usage, duration_ms=duration_ms) + cost = _persist_call_usage(self, call_usage, duration_ms=duration_ms) + finish_llm_call( + span, + self, + model, + call_usage, + duration_ms=duration_ms, + error=error, + cost_usd=cost, + estimated=call_usage is estimated_usage, + output=result if isinstance(result, str) else None, + ) emit = getattr(self, "_emit_gen_finished_log", None) if callable(emit): try: @@ -314,6 +333,10 @@ def stream_token_usage(func): # non-streaming durations under one p50. provider_seconds = 0.0 error: BaseException | None = None + completed = False + # This body runs on the first ``next()``, not at ``gen_stream()`` + # time, so the span starts when the provider call really does. + span = start_llm_span(self, model, stream=True, tools=tools) try: result = func(self, model, messages, stream, tools, **kwargs) stream_iter = iter(result) @@ -323,6 +346,7 @@ def stream_token_usage(func): r = next(stream_iter) except StopIteration: provider_seconds += time.monotonic() - pull_started + completed = True break provider_seconds += time.monotonic() - pull_started if first_chunk_at is None: @@ -346,12 +370,26 @@ def stream_token_usage(func): ) for line in batch: call_usage["generated_tokens"] += _count_tokens(line) + estimated_usage = call_usage call_usage = _prefer_provider_usage(self, call_usage) self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"] self.token_usage["generated_tokens"] += call_usage["generated_tokens"] - _persist_call_usage( + cost = _persist_call_usage( self, call_usage, duration_ms=duration_ms, ttft_ms=ttft_ms ) + finish_llm_call( + span, + self, + model, + call_usage, + duration_ms=duration_ms, + error=error, + completed=completed, + ttft_ms=ttft_ms, + cost_usd=cost, + estimated=call_usage is estimated_usage, + output=output_text(batch), + ) emit = getattr(self, "_emit_stream_finished_log", None) if callable(emit): try: diff --git a/tests/tracing/test_llm_spans.py b/tests/tracing/test_llm_spans.py new file mode 100644 index 00000000..312e65c9 --- /dev/null +++ b/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