mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 14:14:50 +00:00
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:
1 parent
699906a015
commit
0bbbac0eb2
18 files changed
+88
-35
No files matched your search
+2
-17
@@ -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,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,
|
||||
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
@@ -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):
|
||||
|
||||
@@ -501,6 +501,7 @@
|
||||
"fields": {
|
||||
"status": "Status",
|
||||
"duration": "Dauer",
|
||||
"modelTime": "Modellzeit",
|
||||
"provider": "Anbieter",
|
||||
"model": "Modell",
|
||||
"inputTokens": "Eingabe-Tokens",
|
||||
|
||||
@@ -506,6 +506,7 @@
|
||||
"fields": {
|
||||
"status": "Status",
|
||||
"duration": "Duration",
|
||||
"modelTime": "Model time",
|
||||
"provider": "Provider",
|
||||
"model": "Model",
|
||||
"inputTokens": "Input tokens",
|
||||
|
||||
@@ -501,6 +501,7 @@
|
||||
"fields": {
|
||||
"status": "Estado",
|
||||
"duration": "Duración",
|
||||
"modelTime": "Tiempo del modelo",
|
||||
"provider": "Proveedor",
|
||||
"model": "Modelo",
|
||||
"inputTokens": "Tokens de entrada",
|
||||
|
||||
@@ -500,6 +500,7 @@
|
||||
"fields": {
|
||||
"status": "ステータス",
|
||||
"duration": "所要時間",
|
||||
"modelTime": "モデル処理時間",
|
||||
"provider": "プロバイダー",
|
||||
"model": "モデル",
|
||||
"inputTokens": "入力トークン",
|
||||
|
||||
@@ -527,6 +527,7 @@
|
||||
"fields": {
|
||||
"status": "Статус",
|
||||
"duration": "Длительность",
|
||||
"modelTime": "Время модели",
|
||||
"provider": "Провайдер",
|
||||
"model": "Модель",
|
||||
"inputTokens": "Входные токены",
|
||||
|
||||
@@ -500,6 +500,7 @@
|
||||
"fields": {
|
||||
"status": "狀態",
|
||||
"duration": "耗時",
|
||||
"modelTime": "模型耗時",
|
||||
"provider": "供應商",
|
||||
"model": "模型",
|
||||
"inputTokens": "輸入詞元",
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in new issue
Block a user