Skip trace work when disabled; share the agent span builder

With TRACES_ENABLED off, LLM calls record no GenAI metrics, and a streamed
answer is joined into a preview only when a span keeps it. Agent runs and
continuations build their invoke_agent span in one place, which now reads
the model from model_id. The trace panel shows a stream's model time next
to how long it was open. Adds missing type hints.
This commit is contained in:
arc53-machine committed 2026-09-23 22:09:58 +01:00
1 parent 699906a015
commit 0bbbac0eb2
18 files changed
+88 -35

No files matched your search

+2 -17
View File
@@ -32,7 +32,7 @@ from docsgpt.guardrails.stream import StreamingOutputGuard
from docsgpt.guardrails.types import Action, Stage, resolve_tool_result
from docsgpt.llm.handlers.handler_creator import LLMHandlerCreator
from docsgpt.llm.llm_creator import LLMCreator
from docsgpt.logging import build_stack_data, log_activity, LogContext
from docsgpt.logging import build_stack_data, log_activity, LogContext, start_agent_span
logger = logging.getLogger(__name__)
@@ -504,22 +504,7 @@ class BaseAgent(ABC):
pending_tool_calls: The pending tool call descriptors from the pause.
tool_actions: Client-provided actions resolving the pending calls.
"""
label = type(self).__name__
span = tracing.start_span(
tracing.KIND_AGENT,
f"invoke_agent {label}",
attributes={
k: v
for k, v in {
"gen_ai.operation.name": "invoke_agent",
"gen_ai.agent.id": str(self.agent_id) if getattr(self, "agent_id", None) else None,
"gen_ai.request.model": getattr(self, "model_id", None),
"docsgpt.agent_type": label,
"docsgpt.continuation": True,
}.items()
if v is not None
},
)
span = start_agent_span(self, continuation=True)
completed = False
try:
yield from self._gen_continuation_inner(
+2 -2
View File
@@ -2,7 +2,7 @@ import json
import logging
import os
import time
from typing import Dict, Generator, List, Optional
from typing import Any, Dict, Generator, List, Optional
from docsgpt import tracing
from docsgpt.agents.base import BaseAgent
@@ -53,7 +53,7 @@ SYNTHESIS_PROMPT = _load_prompt("synthesis.txt")
def _phase_span(phase: str, **attributes):
def _phase_span(phase: str, **attributes: Any) -> Any:
"""Open a trace span for one research phase (clarify, plan, a step, synthesis)."""
return tracing.start_span(
tracing.KIND_STEP,
+1 -1
View File
@@ -30,7 +30,7 @@ from docsgpt.storage.db.session import db_readonly, db_session
logger = logging.getLogger(__name__)
def record_tool_span_start(call: Any, **attributes: Any):
def record_tool_span_start(call: Any, **attributes: Any) -> Any:
"""Open an ``execute_tool`` span for ``call`` (no-op without an active trace)."""
llm_name = getattr(call, "name", None) or "unknown"
base = {
+35 -4
View File
@@ -140,7 +140,7 @@ def log_activity() -> Callable:
# The outermost agent run names the trace's activity, which is how
# a webhook/system Logs row (keyed by activity_id) finds its trace.
tracing.bind_if_unset(activity_id=activity_id)
span = _start_agent_span(args[0], agent_id=agent_id, model=model, endpoint=endpoint)
span = start_agent_span(args[0], agent_id=agent_id, model=model, endpoint=endpoint)
error: BaseException | None = None
completed = False
@@ -170,9 +170,39 @@ def log_activity() -> Callable:
return decorator
def _start_agent_span(agent: Any, *, agent_id: Any, model: Any, endpoint: str):
"""Open the ``invoke_agent`` span for one agent run (no-op without a trace)."""
def start_agent_span(
agent: Any,
*,
agent_id: Any = None,
model: Any = None,
endpoint: Optional[str] = None,
continuation: bool = False,
) -> Any:
"""Open the ``invoke_agent`` span for one agent run (no-op without a trace).
The one builder for agent spans: ``@log_activity`` uses it for every run
and ``gen_continuation`` for a resumed one. Values not passed are read off
the agent.
Args:
agent: The agent instance.
agent_id: Overrides ``agent.agent_id``.
model: Overrides the agent's model (``gpt_model``, ``model`` or ``model_id``).
endpoint: Overrides ``agent.endpoint``.
continuation: True for a run resumed after tool approval.
Returns:
The span, or a no-op span without an active trace.
"""
label = type(agent).__name__
agent_id = agent_id or getattr(agent, "agent_id", None)
model = (
model
or getattr(agent, "gpt_model", None)
or getattr(agent, "model", None)
or getattr(agent, "model_id", None)
)
endpoint = endpoint or getattr(agent, "endpoint", None)
return tracing.start_span(
tracing.KIND_AGENT,
f"invoke_agent {label}",
@@ -183,7 +213,8 @@ def _start_agent_span(agent: Any, *, agent_id: Any, model: Any, endpoint: str):
"gen_ai.agent.id": str(agent_id) if agent_id else None,
"gen_ai.request.model": str(model) if model else None,
"docsgpt.agent_type": label,
"docsgpt.endpoint": endpoint or None,
"docsgpt.endpoint": str(endpoint) if endpoint else None,
"docsgpt.continuation": True if continuation else None,
}.items()
if v is not None
},
+12 -6
View File
@@ -7,8 +7,9 @@ fallback each get their own span with the provider that actually ran.
from __future__ import annotations
from typing import Any, Dict, Iterable, Optional
from typing import Any, Dict, Iterable, Optional, Union
from docsgpt.core.settings import settings
from docsgpt.tracing import core
from docsgpt.tracing.otel import provider_name, record_llm_metrics
@@ -16,7 +17,9 @@ from docsgpt.tracing.otel import provider_name, record_llm_metrics
CACHE_HIT_ATTR = "_trace_cache_hit"
def start_llm_span(llm: Any, model: Optional[str], *, stream: bool, tools: Any = None):
def start_llm_span(
llm: Any, model: Optional[str], *, stream: bool, tools: Any = None
) -> Union[core.Span, Any]:
"""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
@@ -54,7 +57,7 @@ def finish_llm_call(
ttft_ms: Optional[int] = None,
cost_usd: Optional[float] = None,
estimated: bool = True,
output: Optional[str] = None,
output: Union[str, Iterable[Any], None] = None,
) -> None:
"""Close the call's span and record the GenAI client metrics.
@@ -69,13 +72,16 @@ def finish_llm_call(
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.
output: The response text, or the streamed chunks to join into it;
joined only when the span keeps a preview.
"""
cache_hit = bool(getattr(llm, CACHE_HIT_ATTR, False))
try:
setattr(llm, CACHE_HIT_ATTR, False)
except AttributeError:
pass
if not settings.TRACES_ENABLED:
return
record_llm_metrics(
provider=getattr(llm, "provider_name", None),
model=str(model) if model else None,
@@ -86,8 +92,8 @@ def finish_llm_call(
)
if not span:
return
if output:
span.preview("output", output)
if output and settings.TRACES_CAPTURE_CONTENT:
span.preview("output", output_text(output))
span.end(
None if completed or error is not None else core.STATUS_CANCELLED,
error=error,
+7 -3
View File
@@ -12,7 +12,9 @@ _SNIPPET_CHARS = 240
_MAX_PREVIEW_CHUNKS = 20
def start_retrieval_span(name: str, *, sources: Optional[Iterable[Any]] = None, **attributes: Any):
def start_retrieval_span(
name: str, *, sources: Optional[Iterable[Any]] = None, **attributes: Any
) -> Any:
"""Open a ``retrieval`` span (a container: embeddings and searches nest under it)."""
source_ids = [str(s) for s in (sources or []) if s]
base = {
@@ -64,7 +66,7 @@ def describe_documents(span: Any, docs: Optional[List[Dict[str, Any]]], *, query
)
def start_embedding_span(model: Optional[str], *, inputs: int = 1, **attributes: Any):
def start_embedding_span(model: Optional[str], *, inputs: int = 1, **attributes: Any) -> Any:
"""Open an ``embeddings`` span for a query embedding (leaf)."""
base = {
"gen_ai.operation.name": "embeddings",
@@ -79,7 +81,9 @@ def start_embedding_span(model: Optional[str], *, inputs: int = 1, **attributes:
)
def start_source_search_span(source_id: Any, *, top_k: Optional[int] = None, **attributes: Any):
def start_source_search_span(
source_id: Any, *, top_k: Optional[int] = None, **attributes: Any
) -> Any:
"""Open a per-source vector ``search`` span (leaf)."""
base = {
"gen_ai.operation.name": "search",
+2 -2
View File
@@ -3,7 +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.tracing.llm import finish_llm_call, 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
@@ -388,7 +388,7 @@ def stream_token_usage(func):
ttft_ms=ttft_ms,
cost_usd=cost,
estimated=call_usage is estimated_usage,
output=output_text(batch),
output=batch,
)
emit = getattr(self, "_emit_stream_finished_log", None)
if callable(emit):
+1
View File
@@ -501,6 +501,7 @@
"fields": {
"status": "Status",
"duration": "Dauer",
"modelTime": "Modellzeit",
"provider": "Anbieter",
"model": "Modell",
"inputTokens": "Eingabe-Tokens",
+1
View File
@@ -506,6 +506,7 @@
"fields": {
"status": "Status",
"duration": "Duration",
"modelTime": "Model time",
"provider": "Provider",
"model": "Model",
"inputTokens": "Input tokens",
+1
View File
@@ -501,6 +501,7 @@
"fields": {
"status": "Estado",
"duration": "Duración",
"modelTime": "Tiempo del modelo",
"provider": "Proveedor",
"model": "Modelo",
"inputTokens": "Tokens de entrada",
+1
View File
@@ -500,6 +500,7 @@
"fields": {
"status": "ステータス",
"duration": "所要時間",
"modelTime": "モデル処理時間",
"provider": "プロバイダー",
"model": "モデル",
"inputTokens": "入力トークン",
+1
View File
@@ -527,6 +527,7 @@
"fields": {
"status": "Статус",
"duration": "Длительность",
"modelTime": "Время модели",
"provider": "Провайдер",
"model": "Модель",
"inputTokens": "Входные токены",
+1
View File
@@ -500,6 +500,7 @@
"fields": {
"status": "狀態",
"duration": "耗時",
"modelTime": "模型耗時",
"provider": "供應商",
"model": "模型",
"inputTokens": "輸入詞元",
+1
View File
@@ -500,6 +500,7 @@
"fields": {
"status": "状态",
"duration": "耗时",
"modelTime": "模型耗时",
"provider": "提供商",
"model": "模型",
"inputTokens": "输入词元",
@@ -53,6 +53,10 @@ export default function TraceSpanDetails({ span }: { span: TraceSpan }) {
if (typeof output === 'number')
push(f('outputTokens'), formatTokens(output));
push(f('cachedTokens'), a['gen_ai.usage.cache_read.input_tokens']);
// A stream's span covers the whole time it was open, including the
// app's own work between chunks; model time is the provider's share.
if (a['docsgpt.stream'] && typeof a['docsgpt.provider_ms'] === 'number')
push(f('modelTime'), formatDurationMs(a['docsgpt.provider_ms']));
if (typeof a['docsgpt.ttft_ms'] === 'number')
push(f('timeToFirstToken'), formatDurationMs(a['docsgpt.ttft_ms']));
if (
+2
View File
@@ -437,6 +437,8 @@ class TestLogActivityTraceSpan:
assert agent_span.attributes["gen_ai.operation.name"] == "invoke_agent"
assert agent_span.attributes["gen_ai.agent.id"] == "agent-1"
assert agent_span.attributes["docsgpt.answer_chars"] == 2
# BaseAgent keeps its model in ``model_id``; the span reads it too.
assert agent_span.attributes["gen_ai.request.model"] == "gpt-4o"
assert agent_span.attributes["docsgpt.source_count"] == 1
assert llm_span.parent_id == agent_span.id
assert trace.activity_id is not None
+12
View File
@@ -110,6 +110,18 @@ class TestNonStreaming:
metrics.assert_called_once()
class TestDisabled:
def test_no_metrics_when_tracing_is_off(self, metrics, monkeypatch):
monkeypatch.setattr(settings, "TRACES_ENABLED", False)
@gen_token_usage
def _gen(self, model, messages, stream, tools, **kwargs):
return "x"
_gen(_LLM(), "m", [], False, None)
metrics.assert_not_called()
class TestStreaming:
def test_span_starts_on_first_next_not_on_call(self, trace, metrics):
@stream_token_usage
+2
View File
@@ -245,6 +245,8 @@ class TestContinuation:
tool_span = trace.spans[1]
assert agent_span.kind == tracing.KIND_AGENT
assert agent_span.attributes["docsgpt.continuation"] is True
assert agent_span.attributes["gen_ai.request.model"] == "gpt-4"
assert agent_span.name == "invoke_agent ClassicAgent"
assert agent_span.status == "ok"
assert tool_span.parent_id == agent_span.id
assert tool_span.status == "denied"