diff --git a/docsgpt/agents/base.py b/docsgpt/agents/base.py index 002e7fa6..3dea7254 100644 --- a/docsgpt/agents/base.py +++ b/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( diff --git a/docsgpt/agents/research_agent.py b/docsgpt/agents/research_agent.py index 2abb6cdb..c4858ec0 100644 --- a/docsgpt/agents/research_agent.py +++ b/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, diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 2c7b446d..8e2b65e1 100644 --- a/docsgpt/agents/tool_executor.py +++ b/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 = { diff --git a/docsgpt/logging.py b/docsgpt/logging.py index 9689b1b2..38ddf6da 100644 --- a/docsgpt/logging.py +++ b/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 }, diff --git a/docsgpt/tracing/llm.py b/docsgpt/tracing/llm.py index 596818dd..556fc373 100644 --- a/docsgpt/tracing/llm.py +++ b/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, diff --git a/docsgpt/tracing/retrieval.py b/docsgpt/tracing/retrieval.py index f9ce7103..374fd207 100644 --- a/docsgpt/tracing/retrieval.py +++ b/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", diff --git a/docsgpt/usage.py b/docsgpt/usage.py index 307cd4cd..0f48dfbc 100644 --- a/docsgpt/usage.py +++ b/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): diff --git a/frontend/src/locale/de.json b/frontend/src/locale/de.json index 2c1bceab..05e2770e 100644 --- a/frontend/src/locale/de.json +++ b/frontend/src/locale/de.json @@ -501,6 +501,7 @@ "fields": { "status": "Status", "duration": "Dauer", + "modelTime": "Modellzeit", "provider": "Anbieter", "model": "Modell", "inputTokens": "Eingabe-Tokens", diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index 6ee9dfaf..9bb4018f 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -506,6 +506,7 @@ "fields": { "status": "Status", "duration": "Duration", + "modelTime": "Model time", "provider": "Provider", "model": "Model", "inputTokens": "Input tokens", diff --git a/frontend/src/locale/es.json b/frontend/src/locale/es.json index d0aa427c..e5d9617c 100644 --- a/frontend/src/locale/es.json +++ b/frontend/src/locale/es.json @@ -501,6 +501,7 @@ "fields": { "status": "Estado", "duration": "Duración", + "modelTime": "Tiempo del modelo", "provider": "Proveedor", "model": "Modelo", "inputTokens": "Tokens de entrada", diff --git a/frontend/src/locale/jp.json b/frontend/src/locale/jp.json index 33a05be7..e0f4b035 100644 --- a/frontend/src/locale/jp.json +++ b/frontend/src/locale/jp.json @@ -500,6 +500,7 @@ "fields": { "status": "ステータス", "duration": "所要時間", + "modelTime": "モデル処理時間", "provider": "プロバイダー", "model": "モデル", "inputTokens": "入力トークン", diff --git a/frontend/src/locale/ru.json b/frontend/src/locale/ru.json index 5e724a0b..ac6233ce 100644 --- a/frontend/src/locale/ru.json +++ b/frontend/src/locale/ru.json @@ -527,6 +527,7 @@ "fields": { "status": "Статус", "duration": "Длительность", + "modelTime": "Время модели", "provider": "Провайдер", "model": "Модель", "inputTokens": "Входные токены", diff --git a/frontend/src/locale/zh-TW.json b/frontend/src/locale/zh-TW.json index fbc8b567..beb8acf7 100644 --- a/frontend/src/locale/zh-TW.json +++ b/frontend/src/locale/zh-TW.json @@ -500,6 +500,7 @@ "fields": { "status": "狀態", "duration": "耗時", + "modelTime": "模型耗時", "provider": "供應商", "model": "模型", "inputTokens": "輸入詞元", diff --git a/frontend/src/locale/zh.json b/frontend/src/locale/zh.json index c1db1fee..540c40a4 100644 --- a/frontend/src/locale/zh.json +++ b/frontend/src/locale/zh.json @@ -500,6 +500,7 @@ "fields": { "status": "状态", "duration": "耗时", + "modelTime": "模型耗时", "provider": "提供商", "model": "模型", "inputTokens": "输入词元", diff --git a/frontend/src/settings/traces/TraceSpanDetails.tsx b/frontend/src/settings/traces/TraceSpanDetails.tsx index 09554057..146acef1 100644 --- a/frontend/src/settings/traces/TraceSpanDetails.tsx +++ b/frontend/src/settings/traces/TraceSpanDetails.tsx @@ -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 ( diff --git a/tests/test_logging.py b/tests/test_logging.py index 784b9234..69642c82 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -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 diff --git a/tests/tracing/test_llm_spans.py b/tests/tracing/test_llm_spans.py index 312e65c9..897f0a89 100644 --- a/tests/tracing/test_llm_spans.py +++ b/tests/tracing/test_llm_spans.py @@ -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 diff --git a/tests/tracing/test_tool_spans.py b/tests/tracing/test_tool_spans.py index 899b6a8b..54ef8b40 100644 --- a/tests/tracing/test_tool_spans.py +++ b/tests/tracing/test_tool_spans.py @@ -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"