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.

docsgpt/agents/base.py

@@ -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(

docsgpt/agents/research_agent.py

@@ -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,

docsgpt/agents/tool_executor.py

@@ -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 = {

docsgpt/logging.py

@@ -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
         },

docsgpt/tracing/llm.py

@@ -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,

docsgpt/tracing/retrieval.py

@@ -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",

docsgpt/usage.py

@@ -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):