From f6269c483f55398d262d1daad5f1d38c77afc3b9 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:14:51 +0100 Subject: [PATCH 01/33] Add execution-trace settings and declare opentelemetry-api TRACES_* settings for the per-request trace timeline and its GenAI OTel export. opentelemetry-api was only transitive; the tracing package imports it directly. --- docs/content/Deploying/Settings-Reference.mdx | 41 +++++++++++++++++ docsgpt/core/settings/__init__.py | 2 + docsgpt/core/settings/tracing.py | 45 +++++++++++++++++++ docsgpt/requirements-docling.txt | 1 + docsgpt/requirements-milvus.txt | 1 + docsgpt/requirements.txt | 1 + pyproject.toml | 2 + uv.lock | 2 + 8 files changed, 95 insertions(+) create mode 100644 docsgpt/core/settings/tracing.py diff --git a/docs/content/Deploying/Settings-Reference.mdx b/docs/content/Deploying/Settings-Reference.mdx index d9f59ec7..b08af68a 100644 --- a/docs/content/Deploying/Settings-Reference.mdx +++ b/docs/content/Deploying/Settings-Reference.mdx @@ -1451,6 +1451,47 @@ Type `int`, default `30`, must be `>= 1`. Days guardrail events are kept before the cleanup task removes them. +## Execution traces + +Recording of agent, LLM, tool and retrieval steps per request. + +### `TRACES_ENABLED` + +Type `bool`, default `true`. + +Record an execution trace (agent runs, LLM calls, tool calls, retrieval, embeddings) for every request, store it in request_traces and show it in the Logs UI. False records nothing. + +### `TRACES_CAPTURE_CONTENT` + +Type `bool`, default `true`. + +Store short, secret-redacted previews (tool arguments and results, retrieved chunk titles, rephrased queries, answer excerpts) with each stored trace. Full prompts are never stored. OTel export follows OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT instead. + +### `TRACES_PREVIEW_CHARS` + +Type `int`, default `2000`, must be `>= 100`. + +Maximum characters kept per stored trace preview. + +### `TRACES_MAX_SPANS` + +Type `int`, default `500`, must be `>= 10`. + +Maximum spans recorded per trace; further spans are counted as dropped, not stored. + +### `TRACES_RETENTION_DAYS` + +Type `int`, default `30`, must be `>= 1`. + +Days stored traces are kept before the cleanup task removes them. + +### `TRACES_OTEL_EXPORT` + +Type `bool`, default `true`. + +Also emit each finished trace as OpenTelemetry GenAI spans (gen_ai.*) and metrics. Has no effect unless an OTel SDK is configured, e.g. by launching under opentelemetry-instrument. + + ## Quotas Quota window and the treatment of unpriced models. diff --git a/docsgpt/core/settings/__init__.py b/docsgpt/core/settings/__init__.py index 817281c5..aab95515 100644 --- a/docsgpt/core/settings/__init__.py +++ b/docsgpt/core/settings/__init__.py @@ -35,6 +35,7 @@ from docsgpt.core.settings.scheduler import SchedulerSettings from docsgpt.core.settings.server import ServerSettings from docsgpt.core.settings.speech import SpeechSettings from docsgpt.core.settings.storage import StorageSettings +from docsgpt.core.settings.tracing import TracingSettings from docsgpt.core.settings.vectorstores import VectorStoreSettings from docsgpt.core.settings.workers import WorkerSettings @@ -55,6 +56,7 @@ SETTINGS_GROUPS: tuple[tuple[str, type[SettingsGroup]], ...] = ( ("Events and devices", EventsSettings), ("Agents", AgentSettings), ("Guardrails", GuardrailSettings), + ("Execution traces", TracingSettings), ("Quotas", QuotaSettings), ("Scheduler", SchedulerSettings), ("Sandbox", SandboxSettings), diff --git a/docsgpt/core/settings/tracing.py b/docsgpt/core/settings/tracing.py new file mode 100644 index 00000000..19a26341 --- /dev/null +++ b/docsgpt/core/settings/tracing.py @@ -0,0 +1,45 @@ +"""Execution traces: the per-request timeline shown in the Logs UI and exported as OTel GenAI spans.""" + +from __future__ import annotations + +from pydantic import Field + +from docsgpt.core.settings._shared import SettingsGroup + + +class TracingSettings(SettingsGroup): + """Recording of agent, LLM, tool and retrieval steps per request.""" + + TRACES_ENABLED: bool = Field( + default=True, + description=( + "Record an execution trace (agent runs, LLM calls, tool calls, retrieval, embeddings) for every " + "request, store it in request_traces and show it in the Logs UI. False records nothing." + ), + ) + TRACES_CAPTURE_CONTENT: bool = Field( + default=True, + description=( + "Store short, secret-redacted previews (tool arguments and results, retrieved chunk titles, " + "rephrased queries, answer excerpts) with each stored trace. Full prompts are never stored. OTel " + "export follows OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT instead." + ), + ) + TRACES_PREVIEW_CHARS: int = Field( + default=2000, ge=100, description="Maximum characters kept per stored trace preview." + ) + TRACES_MAX_SPANS: int = Field( + default=500, + ge=10, + description="Maximum spans recorded per trace; further spans are counted as dropped, not stored.", + ) + TRACES_RETENTION_DAYS: int = Field( + default=30, ge=1, description="Days stored traces are kept before the cleanup task removes them." + ) + TRACES_OTEL_EXPORT: bool = Field( + default=True, + description=( + "Also emit each finished trace as OpenTelemetry GenAI spans (gen_ai.*) and metrics. Has no effect " + "unless an OTel SDK is configured, e.g. by launching under opentelemetry-instrument." + ), + ) diff --git a/docsgpt/requirements-docling.txt b/docsgpt/requirements-docling.txt index 2e272a2e..a22caa19 100644 --- a/docsgpt/requirements-docling.txt +++ b/docsgpt/requirements-docling.txt @@ -519,6 +519,7 @@ openpyxl==3.1.5 opentelemetry-api==1.44.0 # via # daytona + # docsgpt # fastmcp-slim # google-api-core # mcp diff --git a/docsgpt/requirements-milvus.txt b/docsgpt/requirements-milvus.txt index 46d3d2b9..713ded5c 100644 --- a/docsgpt/requirements-milvus.txt +++ b/docsgpt/requirements-milvus.txt @@ -438,6 +438,7 @@ openpyxl==3.1.5 opentelemetry-api==1.44.0 # via # daytona + # docsgpt # fastmcp-slim # google-api-core # mcp diff --git a/docsgpt/requirements.txt b/docsgpt/requirements.txt index 17db5b1d..b4b45c8f 100644 --- a/docsgpt/requirements.txt +++ b/docsgpt/requirements.txt @@ -429,6 +429,7 @@ openpyxl==3.1.5 opentelemetry-api==1.44.0 # via # daytona + # docsgpt # fastmcp-slim # google-api-core # mcp diff --git a/pyproject.toml b/pyproject.toml index 5e24bbaa..dbbfc373 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,6 +78,8 @@ dependencies = [ "openapi3-parser>=2.0.0,<3", # pandas reads .xlsx through openpyxl but does not depend on it. "openpyxl>=3.1.5,<4", + # Imported directly by docsgpt.tracing to replay traces as GenAI spans. + "opentelemetry-api>=1.29.0,<2", "opentelemetry-distro>=0.50b0,<1", "opentelemetry-exporter-otlp>=1.29.0,<2", "opentelemetry-instrumentation-celery>=0.50b0,<1", diff --git a/uv.lock b/uv.lock index dba25614..13780640 100644 --- a/uv.lock +++ b/uv.lock @@ -1375,6 +1375,7 @@ dependencies = [ { name = "openai" }, { name = "openapi3-parser" }, { name = "openpyxl" }, + { name = "opentelemetry-api" }, { name = "opentelemetry-distro" }, { name = "opentelemetry-exporter-otlp" }, { name = "opentelemetry-instrumentation-celery" }, @@ -1488,6 +1489,7 @@ requires-dist = [ { name = "openai", specifier = ">=3.13.0,<4" }, { name = "openapi3-parser", specifier = ">=2.0.0,<3" }, { name = "openpyxl", specifier = ">=3.1.5,<4" }, + { name = "opentelemetry-api", specifier = ">=1.29.0,<2" }, { name = "opentelemetry-distro", specifier = ">=0.50b0,<1" }, { name = "opentelemetry-exporter-otlp", specifier = ">=1.29.0,<2" }, { name = "opentelemetry-instrumentation-celery", specifier = ">=0.50b0,<1" }, From f3f817b1e8b45bd603e215dfe4f83d7429edfc7e Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:19:54 +0100 Subject: [PATCH 02/33] Store execution traces in a new request_traces table One row per execution holding its span tree as JSONB, linked to Logs rows by request, message, activity and workflow-run ids. message_id cascades so deleting or truncating a conversation drops its traces; a daily beat task enforces TRACES_RETENTION_DAYS. --- .../alembic/versions/0037_request_traces.py | 84 +++++++ docsgpt/api/user/tasks.py | 29 +++ docsgpt/storage/db/models.py | 36 +++ .../storage/db/repositories/request_traces.py | 219 ++++++++++++++++++ tests/api/user/test_tasks.py | 84 ++++++- .../db/repositories/test_request_traces.py | 179 ++++++++++++++ tests/storage/db/test_migration_0037.py | 65 ++++++ 7 files changed, 686 insertions(+), 10 deletions(-) create mode 100644 docsgpt/alembic/versions/0037_request_traces.py create mode 100644 docsgpt/storage/db/repositories/request_traces.py create mode 100644 tests/storage/db/repositories/test_request_traces.py create mode 100644 tests/storage/db/test_migration_0037.py diff --git a/docsgpt/alembic/versions/0037_request_traces.py b/docsgpt/alembic/versions/0037_request_traces.py new file mode 100644 index 00000000..4929af62 --- /dev/null +++ b/docsgpt/alembic/versions/0037_request_traces.py @@ -0,0 +1,84 @@ +"""0037 request_traces — one execution trace per request for the Logs UI. + +Each row is the full span tree of one execution (a chat turn, a tool-approval +continuation, a scheduled or webhook run, a search, a graph extraction): +agent runs, LLM calls, tool calls, retrieval and embeddings with their +timings, stored as a ``spans`` JSONB array. The Logs UI loads a trace whole, +so one row per trace keeps the write to a single INSERT and lets retention +and conversation deletion remove a trace in one step. + +``message_id`` cascades: deleting a conversation, or truncating it when a +turn is superseded, removes that turn's traces with it. The link ids +(``request_id``, ``activity_id``, ``workflow_run_id``) carry partial indexes +because only the Logs rows that have them look traces up by them. + +``status`` is the only CHECK; span kinds and sources are validated in code so +new ones need no migration. + +Idempotent both ways. + +Revision ID: 0037_request_traces +Revises: 0036_device_audit_created_idx +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0037_request_traces" +down_revision: Union[str, None] = "0036_device_audit_created_idx" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + """ + CREATE TABLE IF NOT EXISTS request_traces ( + id UUID PRIMARY KEY, + request_id TEXT, + message_id UUID REFERENCES conversation_messages(id) ON DELETE CASCADE, + conversation_id UUID, + activity_id TEXT, + workflow_run_id UUID, + user_id TEXT, + agent_id UUID, + source TEXT NOT NULL, + name TEXT, + status TEXT NOT NULL + CONSTRAINT request_traces_status_chk + CHECK (status IN ('ok', 'error', 'paused', 'cancelled')), + started_at TIMESTAMPTZ NOT NULL, + duration_ms INTEGER NOT NULL DEFAULT 0, + span_count INTEGER NOT NULL DEFAULT 0, + dropped_spans INTEGER NOT NULL DEFAULT 0, + summary JSONB NOT NULL DEFAULT '{}'::jsonb, + spans JSONB NOT NULL DEFAULT '[]'::jsonb, + otel_trace_id TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + """ + ) + op.execute( + """ + CREATE INDEX IF NOT EXISTS request_traces_request_idx + ON request_traces (request_id) WHERE request_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_message_idx + ON request_traces (message_id) WHERE message_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_activity_idx + ON request_traces (activity_id) WHERE activity_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_workflow_run_idx + ON request_traces (workflow_run_id) WHERE workflow_run_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_user_started_idx + ON request_traces (user_id, started_at DESC); + CREATE INDEX IF NOT EXISTS request_traces_agent_started_idx + ON request_traces (agent_id, started_at DESC) WHERE agent_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_created_idx + ON request_traces (created_at); + """ + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS request_traces;") diff --git a/docsgpt/api/user/tasks.py b/docsgpt/api/user/tasks.py index 82d247e7..b34945fe 100644 --- a/docsgpt/api/user/tasks.py +++ b/docsgpt/api/user/tasks.py @@ -608,6 +608,11 @@ def setup_periodic_tasks(sender, **kwargs): cleanup_guardrail_events.s(), name="cleanup-guardrail-events", ) + sender.add_periodic_task( + timedelta(hours=24), + cleanup_traces.s(), + name="cleanup-traces", + ) sender.add_periodic_task( timedelta(hours=24), cleanup_orphan_memories.s(), @@ -827,6 +832,30 @@ def cleanup_guardrail_events(self): return {"deleted": deleted, "ttl_days": ttl_days} +@celery.task(bind=True, acks_late=False) +def cleanup_traces(self): + """Delete ``request_traces`` rows older than ``TRACES_RETENTION_DAYS``. + + Every chat turn, scheduled run and search writes a trace, and each one + carries content previews, so the table is bounded by a retention window + like the other per-request journals. + """ + from docsgpt.core.settings import settings + if not settings.POSTGRES_URI: + return {"deleted": 0, "skipped": "POSTGRES_URI not set"} + + from docsgpt.storage.db.engine import get_engine + from docsgpt.storage.db.repositories.request_traces import ( + RequestTracesRepository, + ) + + ttl_days = settings.TRACES_RETENTION_DAYS + engine = get_engine() + with engine.begin() as conn: + deleted = RequestTracesRepository(conn).purge_older_than(ttl_days) + return {"deleted": deleted, "ttl_days": ttl_days} + + @celery.task(bind=True, acks_late=False) def cleanup_orphan_memories(self): """Sweep orphan memories left by the 0009 FK-to-trigger orphan window. diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index 79773b6e..067dfb4f 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -825,6 +825,42 @@ Index( Index("ix_guardrail_events_message", guardrail_events_table.c.message_id) Index("ix_guardrail_events_created", guardrail_events_table.c.created_at) +# One execution trace per request (chat turn, continuation, scheduled or +# webhook run, search, graph extraction): the span tree as a JSONB array, +# rendered as a waterfall in the Logs UI. ``message_id`` cascades so deleting +# or truncating a conversation removes its traces. Migration 0037. +request_traces_table = Table( + "request_traces", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("request_id", Text), + Column( + "message_id", + UUID(as_uuid=True), + ForeignKey("conversation_messages.id", ondelete="CASCADE"), + ), + Column("conversation_id", UUID(as_uuid=True)), + Column("activity_id", Text), + Column("workflow_run_id", UUID(as_uuid=True)), + Column("user_id", Text), + Column("agent_id", UUID(as_uuid=True)), + Column("source", Text, nullable=False), + Column("name", Text), + # ok | error | paused | cancelled + Column("status", Text, nullable=False), + Column("started_at", DateTime(timezone=True), nullable=False), + Column("duration_ms", Integer, nullable=False, server_default="0"), + Column("span_count", Integer, nullable=False, server_default="0"), + Column("dropped_spans", Integer, nullable=False, server_default="0"), + Column("summary", JSONB, nullable=False, server_default="{}"), + Column("spans", JSONB, nullable=False, server_default="[]"), + Column("otel_trace_id", Text), + Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), +) + +Index("request_traces_user_started_idx", request_traces_table.c.user_id, request_traces_table.c.started_at) +Index("request_traces_created_idx", request_traces_table.c.created_at) + tool_call_attempts_table = Table( "tool_call_attempts", metadata, diff --git a/docsgpt/storage/db/repositories/request_traces.py b/docsgpt/storage/db/repositories/request_traces.py new file mode 100644 index 00000000..1fb5c594 --- /dev/null +++ b/docsgpt/storage/db/repositories/request_traces.py @@ -0,0 +1,219 @@ +"""Repository for ``request_traces``: one stored execution trace per request.""" + +from __future__ import annotations + +import datetime +import json +import logging +from typing import Any, Dict, Iterable, List, Optional + +from sqlalchemy import Connection, text +from sqlalchemy.exc import IntegrityError + +from docsgpt.storage.db.base_repository import looks_like_uuid, row_to_dict +from docsgpt.storage.db.serialization import PGNativeJSONEncoder +from docsgpt.utils import strip_null_bytes + +logger = logging.getLogger(__name__) + +#: Columns a trace can be looked up by, mapped to their SQL type. +REF_FIELDS: Dict[str, str] = { + "id": "uuid", + "request_id": "text", + "message_id": "uuid", + "activity_id": "text", + "workflow_run_id": "uuid", +} + +_SUMMARY_COLUMNS = ( + "id, request_id, message_id, conversation_id, activity_id, workflow_run_id, " + "user_id, agent_id, source, name, status, started_at, duration_ms, " + "span_count, dropped_spans, summary, otel_trace_id" +) + + +def _dump_jsonb(value: Any) -> str: + return json.dumps(strip_null_bytes(value), cls=PGNativeJSONEncoder) + + +def _uuid_or_none(value: Optional[str]) -> Optional[str]: + return str(value) if value and looks_like_uuid(str(value)) else None + + +def _started_at(record: Dict[str, Any]) -> datetime.datetime: + started_ns = record.get("started_at_ns") + if started_ns: + return datetime.datetime.fromtimestamp(started_ns / 1e9, tz=datetime.timezone.utc) + return datetime.datetime.now(datetime.timezone.utc) + + +def _scope_clause(user_id: Optional[str], agent_id: Optional[str]) -> tuple[str, Dict[str, Any]]: + """Owner scoping: an owned agent's traces, else the caller's own traces.""" + if agent_id: + return "agent_id = CAST(:scope_agent AS uuid)", {"scope_agent": str(agent_id)} + return "user_id = :scope_user", {"scope_user": user_id} + + +class RequestTracesRepository: + def __init__(self, conn: Connection) -> None: + self._conn = conn + + def insert(self, record: Dict[str, Any]) -> bool: + """Insert one finished trace (the dict from ``Trace.to_record``). + + A foreign-key failure means the message was deleted before the trace + flushed (e.g. the turn was superseded); the trace describes a turn the + user discarded, so it is dropped rather than raised. The insert runs + in a savepoint so that drop leaves the caller's transaction usable. + + Args: + record: The trace record. + + Returns: + True when the row was written. + """ + params = { + "id": record["id"], + "request_id": record.get("request_id"), + "message_id": _uuid_or_none(record.get("message_id")), + "conversation_id": _uuid_or_none(record.get("conversation_id")), + "activity_id": record.get("activity_id"), + "workflow_run_id": _uuid_or_none(record.get("workflow_run_id")), + "user_id": record.get("user_id"), + "agent_id": _uuid_or_none(record.get("agent_id")), + "source": record.get("source") or "unknown", + "name": strip_null_bytes(record.get("name")), + "status": record.get("status") or "ok", + "started_at": _started_at(record), + "duration_ms": int(record.get("duration_ms") or 0), + "span_count": int(record.get("span_count") or 0), + "dropped_spans": int(record.get("dropped_spans") or 0), + "summary": _dump_jsonb(record.get("summary") or {}), + "spans": _dump_jsonb(record.get("spans") or []), + "otel_trace_id": record.get("otel_trace_id"), + } + statement = text( + """ + INSERT INTO request_traces + (id, request_id, message_id, conversation_id, activity_id, + workflow_run_id, user_id, agent_id, source, name, status, + started_at, duration_ms, span_count, dropped_spans, summary, + spans, otel_trace_id) + VALUES + (CAST(:id AS uuid), :request_id, CAST(:message_id AS uuid), + CAST(:conversation_id AS uuid), :activity_id, + CAST(:workflow_run_id AS uuid), :user_id, CAST(:agent_id AS uuid), + :source, :name, :status, :started_at, :duration_ms, :span_count, + :dropped_spans, CAST(:summary AS jsonb), CAST(:spans AS jsonb), + :otel_trace_id) + ON CONFLICT (id) DO NOTHING + """ + ) + try: + with self._conn.begin_nested(): + result = self._conn.execute(statement, params) + except IntegrityError as exc: + sqlstate = getattr(getattr(exc, "orig", None), "sqlstate", None) + if sqlstate != "23503": + raise + logger.info("Dropped trace %s: its message no longer exists", record["id"]) + return False + return (result.rowcount or 0) > 0 + + def list_by_ref( + self, + field: str, + value: str, + *, + user_id: Optional[str], + agent_id: Optional[str] = None, + limit: int = 20, + ) -> List[dict]: + """Full traces matching ``field = value``, oldest first, owner-scoped. + + Args: + field: One of :data:`REF_FIELDS`. + value: The id to match. + user_id: The caller; used when ``agent_id`` is not given. + agent_id: An agent the caller owns; scopes to that agent's traces. + limit: Maximum traces returned. + + Returns: + Trace dicts including ``spans``; empty for an unknown field or id. + """ + sql_type = REF_FIELDS.get(field) + if sql_type is None or not value: + return [] + if sql_type == "uuid" and not looks_like_uuid(str(value)): + return [] + if not agent_id and not user_id: + return [] + scope, scope_params = _scope_clause(user_id, agent_id) + result = self._conn.execute( + text( + f""" + SELECT {_SUMMARY_COLUMNS}, spans FROM request_traces + WHERE {field} = CAST(:value AS {sql_type}) AND {scope} + ORDER BY started_at + LIMIT :limit + """ + ), + {"value": str(value), "limit": max(1, min(int(limit), 100)), **scope_params}, + ) + return [row_to_dict(row) for row in result.fetchall()] + + def summaries_for_refs( + self, + refs: Dict[str, Iterable[str]], + *, + user_id: Optional[str], + agent_id: Optional[str] = None, + ) -> Dict[str, Dict[str, List[dict]]]: + """Trace summaries (no spans) for a page of Logs rows. + + Args: + refs: ``{field: [ids...]}`` for fields in :data:`REF_FIELDS`. + user_id: The caller; used when ``agent_id`` is not given. + agent_id: An agent the caller owns. + + Returns: + ``{field: {id: [summary, ...]}}`` with summaries oldest first. + """ + out: Dict[str, Dict[str, List[dict]]] = {} + if not agent_id and not user_id: + return out + scope, scope_params = _scope_clause(user_id, agent_id) + for field, values in refs.items(): + sql_type = REF_FIELDS.get(field) + if sql_type is None: + continue + ids = sorted({str(v) for v in values if v}) + if sql_type == "uuid": + ids = [v for v in ids if looks_like_uuid(v)] + if not ids: + continue + result = self._conn.execute( + text( + f""" + SELECT {_SUMMARY_COLUMNS} FROM request_traces + WHERE {field} = ANY(CAST(:ids AS {sql_type}[])) AND {scope} + ORDER BY started_at + """ + ), + {"ids": ids, **scope_params}, + ) + for row in result.fetchall(): + data = row_to_dict(row) + out.setdefault(field, {}).setdefault(str(data[field]), []).append(data) + return out + + def purge_older_than(self, days: int) -> int: + """Delete traces older than ``days`` (the retention window).""" + result = self._conn.execute( + text( + "DELETE FROM request_traces " + "WHERE created_at < NOW() - CAST(:days || ' days' AS interval)" + ), + {"days": str(max(1, days))}, + ) + return result.rowcount or 0 diff --git a/tests/api/user/test_tasks.py b/tests/api/user/test_tasks.py index a593bbb5..a3ff2a19 100644 --- a/tests/api/user/test_tasks.py +++ b/tests/api/user/test_tasks.py @@ -299,7 +299,7 @@ class TestSetupPeriodicTasks: setup_periodic_tasks(sender) - assert sender.add_periodic_task.call_count == 14 + assert sender.add_periodic_task.call_count == 15 calls = sender.add_periodic_task.call_args_list @@ -326,20 +326,23 @@ class TestSetupPeriodicTasks: # guardrail_events retention sweep (24h) assert calls[8][0][0] == timedelta(hours=24) assert calls[8][1].get("name") == "cleanup-guardrail-events" - # orphan memories sweep (24h) + # request_traces retention sweep (24h) assert calls[9][0][0] == timedelta(hours=24) - assert calls[9][1].get("name") == "cleanup-orphan-memories" + assert calls[9][1].get("name") == "cleanup-traces" + # orphan memories sweep (24h) + assert calls[10][0][0] == timedelta(hours=24) + assert calls[10][1].get("name") == "cleanup-orphan-memories" # scheduler dispatcher - assert calls[10][1].get("name") == "dispatch-scheduled-runs" + assert calls[11][1].get("name") == "dispatch-scheduled-runs" # schedule runs cleanup (24h) - assert calls[11][0][0] == timedelta(hours=24) - assert calls[11][1].get("name") == "cleanup-schedule-runs" + assert calls[12][0][0] == timedelta(hours=24) + assert calls[12][1].get("name") == "cleanup-schedule-runs" # sandbox session reaper (60s) - assert calls[12][0][0] == timedelta(seconds=60) - assert calls[12][1].get("name") == "reap-sandbox-sessions" + assert calls[13][0][0] == timedelta(seconds=60) + assert calls[13][1].get("name") == "reap-sandbox-sessions" # stale workflow-run reaper (5m) - assert calls[13][0][0] == timedelta(seconds=300) - assert calls[13][1].get("name") == "reap-stale-workflow-runs" + assert calls[14][0][0] == timedelta(seconds=300) + assert calls[14][1].get("name") == "reap-stale-workflow-runs" class TestMcpOauthTask: @@ -679,6 +682,67 @@ class TestCleanupMessageEventsTask: assert [r["sequence_no"] for r in rows] == [1] +class TestCleanupTracesTask: + """Retention janitor for ``request_traces``.""" + + @pytest.mark.unit + def test_skips_when_postgres_uri_missing(self, monkeypatch): + from docsgpt.api.user.tasks import cleanup_traces + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "POSTGRES_URI", None, raising=False) + + assert cleanup_traces.run() == {"deleted": 0, "skipped": "POSTGRES_URI not set"} + + @pytest.mark.unit + def test_deletes_traces_past_retention_window(self, pg_conn, monkeypatch): + import time + import uuid + + from sqlalchemy import text as _text + + from docsgpt.api.user.tasks import cleanup_traces + from docsgpt.core.settings import settings + from docsgpt.storage.db.repositories.request_traces import ( + RequestTracesRepository, + ) + + repo = RequestTracesRepository(pg_conn) + for request_id in ("stale", "fresh"): + repo.insert( + { + "id": str(uuid.uuid4()), + "request_id": request_id, + "user_id": "u1", + "source": "stream", + "status": "ok", + "started_at_ns": time.time_ns(), + } + ) + pg_conn.execute( + _text( + "UPDATE request_traces SET created_at = now() - interval '45 days' " + "WHERE request_id = 'stale'" + ) + ) + monkeypatch.setattr(settings, "POSTGRES_URI", "postgresql://stub", raising=False) + monkeypatch.setattr(settings, "TRACES_RETENTION_DAYS", 30) + + @contextmanager + def _fake_begin(): + yield pg_conn + + fake_engine = MagicMock() + fake_engine.begin = _fake_begin + + with patch("docsgpt.storage.db.engine.get_engine", return_value=fake_engine): + result = cleanup_traces.run() + + assert result == {"deleted": 1, "ttl_days": 30} + remaining = pg_conn.execute(_text("SELECT request_id FROM request_traces")).scalars().all() + assert remaining == ["fresh"] + + class TestCleanupOrphanMemoriesTask: """Sweeps orphan memories from the FK-to-trigger orphan window.""" diff --git a/tests/storage/db/repositories/test_request_traces.py b/tests/storage/db/repositories/test_request_traces.py new file mode 100644 index 00000000..0dc72b05 --- /dev/null +++ b/tests/storage/db/repositories/test_request_traces.py @@ -0,0 +1,179 @@ +"""Tests for RequestTracesRepository against a real Postgres instance.""" + +from __future__ import annotations + +import time +import uuid + +from sqlalchemy import text + +from docsgpt.storage.db.repositories.conversations import ConversationsRepository +from docsgpt.storage.db.repositories.request_traces import RequestTracesRepository + + +def _record(**overrides): + base = { + "id": str(uuid.uuid4()), + "request_id": "req-1", + "message_id": None, + "conversation_id": None, + "activity_id": None, + "workflow_run_id": None, + "user_id": "u1", + "agent_id": None, + "source": "stream", + "name": "stream", + "status": "ok", + "started_at_ns": time.time_ns(), + "duration_ms": 1234, + "span_count": 1, + "dropped_spans": 0, + "summary": {"llm_calls": 1}, + "spans": [ + { + "id": "s1", + "parent_id": None, + "kind": "llm", + "name": "chat m", + "status": "ok", + "offset_ms": 1.0, + "duration_ms": 10.0, + "attributes": {"gen_ai.request.model": "m"}, + } + ], + "otel_trace_id": None, + } + base.update(overrides) + return base + + +def _message(pg_conn, user_id="u1"): + convs = ConversationsRepository(pg_conn) + conv = convs.create(user_id, "t") + msg = convs.reserve_message( + str(conv["id"]), prompt="q", placeholder_response="..." + ) + return str(conv["id"]), str(msg["id"]) + + +class TestInsert: + def test_round_trip(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + record = _record() + assert repo.insert(record) is True + rows = repo.list_by_ref("request_id", "req-1", user_id="u1") + assert len(rows) == 1 + row = rows[0] + assert row["id"] == record["id"] + assert row["duration_ms"] == 1234 + assert row["summary"] == {"llm_calls": 1} + assert row["spans"][0]["attributes"]["gen_ai.request.model"] == "m" + + def test_strips_nul_bytes(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + record = _record( + spans=[{"id": "s", "name": "a\x00b", "preview": {"result": "x\x00y"}}] + ) + assert repo.insert(record) + row = repo.list_by_ref("id", record["id"], user_id="u1")[0] + assert row["spans"][0]["name"] == "ab" + assert row["spans"][0]["preview"]["result"] == "xy" + + def test_non_uuid_link_ids_become_null(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + record = _record(agent_id="not-a-uuid", conversation_id="nope") + assert repo.insert(record) + row = repo.list_by_ref("id", record["id"], user_id="u1")[0] + assert row["agent_id"] is None + assert row["conversation_id"] is None + + def test_missing_message_drops_row_and_keeps_transaction_usable(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + assert repo.insert(_record(message_id=str(uuid.uuid4()))) is False + # The savepoint rolled back; the outer transaction still works. + assert repo.insert(_record(request_id="req-after")) is True + + def test_duplicate_id_is_ignored(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + record = _record() + assert repo.insert(record) is True + assert repo.insert(record) is False + + +class TestScoping: + def test_other_users_cannot_read(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + repo.insert(_record()) + assert repo.list_by_ref("request_id", "req-1", user_id="u2") == [] + + def test_agent_scope_returns_agent_traces_from_any_user(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + agent_id = str(uuid.uuid4()) + repo.insert(_record(user_id="someone-else", agent_id=agent_id)) + rows = repo.list_by_ref("request_id", "req-1", user_id="owner", agent_id=agent_id) + assert len(rows) == 1 + + def test_unknown_field_and_bad_uuid_return_empty(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + repo.insert(_record()) + assert repo.list_by_ref("user_id", "u1", user_id="u1") == [] + assert repo.list_by_ref("message_id", "not-uuid", user_id="u1") == [] + assert repo.list_by_ref("request_id", "req-1", user_id=None) == [] + + +class TestMessageLink: + def test_rounds_for_a_message_are_ordered(self, pg_conn): + conv_id, msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + first = _record(message_id=msg_id, status="paused", started_at_ns=time.time_ns()) + second = _record(message_id=msg_id, started_at_ns=time.time_ns() + 1_000_000) + repo.insert(second) + repo.insert(first) + rows = repo.list_by_ref("message_id", msg_id, user_id="u1") + assert [r["id"] for r in rows] == [first["id"], second["id"]] + + def test_deleting_the_conversation_cascades(self, pg_conn): + conv_id, msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + repo.insert(_record(message_id=msg_id, conversation_id=conv_id)) + ConversationsRepository(pg_conn).delete(conv_id, "u1") + count = pg_conn.execute(text("SELECT count(*) FROM request_traces")).scalar() + assert count == 0 + + +class TestSummaries: + def test_summaries_grouped_by_field_and_id(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + a = _record(request_id="r-a") + b = _record(request_id="r-b", activity_id="act-1") + repo.insert(a) + repo.insert(b) + out = repo.summaries_for_refs( + {"request_id": ["r-a", "r-b", "r-missing"], "activity_id": ["act-1"]}, + user_id="u1", + ) + assert set(out["request_id"]) == {"r-a", "r-b"} + assert out["activity_id"]["act-1"][0]["id"] == b["id"] + assert "spans" not in out["request_id"]["r-a"][0] + + def test_summaries_are_scoped(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + repo.insert(_record(request_id="r-a")) + assert repo.summaries_for_refs({"request_id": ["r-a"]}, user_id="u2") == {} + + +class TestPurge: + def test_purge_older_than(self, pg_conn): + repo = RequestTracesRepository(pg_conn) + old = _record(request_id="old") + repo.insert(old) + repo.insert(_record(request_id="new")) + pg_conn.execute( + text( + "UPDATE request_traces SET created_at = now() - interval '40 days' " + "WHERE request_id = 'old'" + ) + ) + assert repo.purge_older_than(30) == 1 + assert repo.list_by_ref("request_id", "old", user_id="u1") == [] + assert len(repo.list_by_ref("request_id", "new", user_id="u1")) == 1 diff --git a/tests/storage/db/test_migration_0037.py b/tests/storage/db/test_migration_0037.py new file mode 100644 index 00000000..bb453df2 --- /dev/null +++ b/tests/storage/db/test_migration_0037.py @@ -0,0 +1,65 @@ +"""Migration round-trip test for 0037_request_traces.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from sqlalchemy import text + + +pytestmark = pytest.mark.integration + + +def _alembic_ini() -> Path: + return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini" + + +def _run_alembic(url: str, *args: str) -> None: + subprocess.check_call( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + ) + + +def _alembic_version(conn) -> str: + return conn.execute(text("SELECT version_num FROM alembic_version")).scalar() + + +def _table_exists(conn, table: str) -> bool: + return conn.execute(text("SELECT to_regclass(:t)"), {"t": f"public.{table}"}).scalar() is not None + + +_0037 = "0037_request_traces" +_0036 = "0036_device_audit_created_idx" + + +class TestMigration0037RoundTrip: + def test_head_has_request_traces(self, pg_engine): + with pg_engine.connect() as conn: + assert _alembic_version(conn) >= _0037 + assert _table_exists(conn, "request_traces") + + def test_downgrade_drops_then_upgrade_restores(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0036) + with pg_engine.connect() as conn: + assert _alembic_version(conn) == _0036 + assert not _table_exists(conn, "request_traces") + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + assert _table_exists(conn, "request_traces") + + def test_status_check_rejects_unknown_status(self, pg_engine): + with pytest.raises(Exception): + with pg_engine.begin() as conn: + conn.execute( + text( + "INSERT INTO request_traces (id, source, status, started_at) " + "VALUES (gen_random_uuid(), 'stream', 'weird', now())" + ) + ) From 3f45d019b05967fec905fec07e4f8893fb0452ff Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:19:54 +0100 Subject: [PATCH 03/33] Add the per-request trace recorder docsgpt.tracing records agent, LLM, tool, retrieval and embedding spans in memory and writes the finished trace once. Container spans nest by a per-thread stack so suspended generators cannot corrupt the tree; open spans are cancelled at flush and previews are bounded and redacted. --- docsgpt/tracing/__init__.py | 87 ++++++ docsgpt/tracing/core.py | 585 ++++++++++++++++++++++++++++++++++++ docsgpt/tracing/preview.py | 58 ++++ docsgpt/tracing/sink.py | 53 ++++ tests/tracing/__init__.py | 0 tests/tracing/test_trace.py | 341 +++++++++++++++++++++ 6 files changed, 1124 insertions(+) create mode 100644 docsgpt/tracing/__init__.py create mode 100644 docsgpt/tracing/core.py create mode 100644 docsgpt/tracing/preview.py create mode 100644 docsgpt/tracing/sink.py create mode 100644 tests/tracing/__init__.py create mode 100644 tests/tracing/test_trace.py diff --git a/docsgpt/tracing/__init__.py b/docsgpt/tracing/__init__.py new file mode 100644 index 00000000..0cf018d2 --- /dev/null +++ b/docsgpt/tracing/__init__.py @@ -0,0 +1,87 @@ +"""Per-request execution traces. + +Records agent runs, LLM calls, tool calls, retrieval and embeddings as a tree +of timed spans. A finished trace is stored in ``request_traces`` (rendered as +a waterfall in the Logs UI) and replayed as OpenTelemetry GenAI spans when an +OTel SDK is configured. See ``docs/content/Deploying/Observability.mdx``. + +Typical use:: + + trace = tracing.start_trace(source="stream", request_id=request_id) + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "execute_tool search") as s: + s.preview("arguments", args) + ... + tracing.flush(trace) + +Every call is a no-op when no trace is active. +""" + +from docsgpt.tracing.core import ( + BINDABLE_IDS, + CONTAINER_KINDS, + KIND_AGENT, + KIND_EMBEDDING, + KIND_GUARDRAIL, + KIND_LLM, + KIND_RERANK, + KIND_RETRIEVAL, + KIND_SEARCH, + KIND_STEP, + KIND_TOOL, + NOOP_SPAN, + STATUS_CANCELLED, + STATUS_DENIED, + STATUS_ERROR, + STATUS_OK, + STATUS_PAUSED, + STATUS_PENDING, + STATUS_SKIPPED, + Span, + Trace, + activate, + bind, + bind_if_unset, + current_trace, + mark_content_blocked, + span, + start_span, + start_trace, + wrap, +) +from docsgpt.tracing.sink import discard, flush + +__all__ = [ + "BINDABLE_IDS", + "CONTAINER_KINDS", + "KIND_AGENT", + "KIND_EMBEDDING", + "KIND_GUARDRAIL", + "KIND_LLM", + "KIND_RERANK", + "KIND_RETRIEVAL", + "KIND_SEARCH", + "KIND_STEP", + "KIND_TOOL", + "NOOP_SPAN", + "STATUS_CANCELLED", + "STATUS_DENIED", + "STATUS_ERROR", + "STATUS_OK", + "STATUS_PAUSED", + "STATUS_PENDING", + "STATUS_SKIPPED", + "Span", + "Trace", + "activate", + "bind", + "bind_if_unset", + "current_trace", + "discard", + "flush", + "mark_content_blocked", + "span", + "start_span", + "start_trace", + "wrap", +] diff --git a/docsgpt/tracing/core.py b/docsgpt/tracing/core.py new file mode 100644 index 00000000..b0289152 --- /dev/null +++ b/docsgpt/tracing/core.py @@ -0,0 +1,585 @@ +"""In-memory recording of one execution trace. + +A :class:`Trace` collects spans for a single request (a chat turn, a +scheduled run, a search). Spans are recorded while the request runs and the +whole trace is written once, when its owner calls :func:`flush`. + +Nesting is tracked with one span stack per thread. Only *container* spans +(agent, tool, retrieval, step) are pushed, so a leaf such as an LLM call can +never become the parent of a sibling that starts while it is still open -- +which matters because DocsGPT's agent loop is a chain of suspended +generators. Ending a span pops anything left above it (marked +``cancelled``), and :meth:`Trace.finish` closes whatever is still open, so a +generator finalized late can never corrupt the tree or the stored record. +""" + +from __future__ import annotations + +import contextlib +import functools +import logging +import threading +import time +import uuid +from contextvars import ContextVar +from typing import Any, Callable, Dict, Iterator, List, Optional + +from docsgpt.core.settings import settings +from docsgpt.tracing.preview import make_preview + +logger = logging.getLogger(__name__) + +KIND_AGENT = "agent" +KIND_LLM = "llm" +KIND_TOOL = "tool" +KIND_RETRIEVAL = "retrieval" +KIND_SEARCH = "search" +KIND_EMBEDDING = "embedding" +KIND_RERANK = "rerank" +KIND_GUARDRAIL = "guardrail" +KIND_STEP = "step" + +#: Kinds that become the implicit parent of spans started while they are open. +CONTAINER_KINDS = frozenset({KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_STEP}) + +STATUS_OK = "ok" +STATUS_ERROR = "error" +STATUS_CANCELLED = "cancelled" +STATUS_PAUSED = "paused" +STATUS_PENDING = "pending" +STATUS_DENIED = "denied" +STATUS_SKIPPED = "skipped" + +#: Trace ids ``bind`` may set; anything else is ignored. +BINDABLE_IDS = frozenset( + { + "request_id", + "message_id", + "conversation_id", + "activity_id", + "workflow_run_id", + "user_id", + "agent_id", + "name", + } +) + +_current: ContextVar[Optional["Trace"]] = ContextVar("docsgpt_trace", default=None) + + +def _new_id() -> str: + return uuid.uuid4().hex[:16] + + +class Span: + """One timed step inside a trace. + + Attribute keys follow the OTel GenAI conventions (``gen_ai.*``) where one + exists and ``docsgpt.*`` otherwise, so the stored trace and the exported + spans share a vocabulary. + """ + + __slots__ = ( + "trace", + "id", + "parent_id", + "kind", + "name", + "attributes", + "previews", + "status", + "error", + "start_perf_ns", + "end_perf_ns", + "_thread", + "_pushed", + ) + + def __init__( + self, + trace: "Trace", + kind: str, + name: str, + parent_id: Optional[str], + attributes: Optional[Dict[str, Any]] = None, + ) -> None: + self.trace = trace + self.id = _new_id() + self.parent_id = parent_id + self.kind = kind + self.name = name + self.attributes: Dict[str, Any] = dict(attributes or {}) + self.previews: Dict[str, Any] = {} + self.status: Optional[str] = None + self.error: Optional[str] = None + self.start_perf_ns = time.perf_counter_ns() + self.end_perf_ns: Optional[int] = None + self._thread: Optional[int] = None + self._pushed = False + + @property + def ended(self) -> bool: + return self.end_perf_ns is not None + + @property + def duration_ms(self) -> Optional[float]: + if self.end_perf_ns is None: + return None + return (self.end_perf_ns - self.start_perf_ns) / 1e6 + + def set(self, **attributes: Any) -> "Span": + """Merge attributes; ``None`` values are ignored. Keys may contain dots via ``**{...}``.""" + if not self.ended: + self.attributes.update({k: v for k, v in attributes.items() if v is not None}) + return self + + def preview(self, key: str, value: Any) -> "Span": + """Attach a bounded, redacted content preview (skipped when capture is off).""" + if self.ended or value is None or not settings.TRACES_CAPTURE_CONTENT: + return self + try: + self.previews[key] = make_preview(value) + except Exception: # noqa: BLE001 - a preview must never break a request + logger.debug("trace preview failed for %s", key, exc_info=True) + return self + + def fail(self, exc: BaseException) -> "Span": + """Record ``exc`` on the span without ending it.""" + self.status = STATUS_ERROR + self.error = str(exc)[:500] or type(exc).__name__ + self.attributes["error.type"] = type(exc).__name__ + return self + + def end( + self, + status: Optional[str] = None, + *, + error: Optional[BaseException] = None, + attributes: Optional[Dict[str, Any]] = None, + ) -> None: + """Close the span. A second end, or an end after the trace finished, is ignored.""" + if self.ended: + return + if attributes: + self.set(**attributes) + if error is not None: + self.fail(error) + if status is not None: + self.status = status + elif self.status is None: + self.status = STATUS_OK + self.end_perf_ns = time.perf_counter_ns() + self.trace._on_end(self) + + def __enter__(self) -> "Span": + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + if exc is None: + self.end() + elif isinstance(exc, GeneratorExit): + self.end(STATUS_CANCELLED) + else: + self.end(error=exc) + return False + + +class _NoopSpan: + """Stand-in returned when no trace is active or the span cap is reached.""" + + id = None + parent_id = None + kind = None + name = "" + status = None + ended = True + duration_ms = None + + @property + def attributes(self) -> Dict[str, Any]: + return {} + + @property + def previews(self) -> Dict[str, Any]: + return {} + + def set(self, **attributes: Any) -> "_NoopSpan": + return self + + def preview(self, key: str, value: Any) -> "_NoopSpan": + return self + + def fail(self, exc: BaseException) -> "_NoopSpan": + return self + + def end(self, status=None, *, error=None, attributes=None) -> None: + return None + + def __enter__(self) -> "_NoopSpan": + return self + + def __exit__(self, exc_type, exc, tb) -> bool: + return False + + def __bool__(self) -> bool: + return False + + +NOOP_SPAN = _NoopSpan() + + +class Trace: + """All spans recorded for one execution, plus the ids that link it to logs.""" + + def __init__( + self, + *, + source: str, + name: Optional[str] = None, + request_id: Optional[str] = None, + message_id: Optional[str] = None, + conversation_id: Optional[str] = None, + activity_id: Optional[str] = None, + workflow_run_id: Optional[str] = None, + user_id: Optional[str] = None, + agent_id: Optional[str] = None, + otel_context: Any = None, + ) -> None: + self.id = str(uuid.uuid4()) + self.source = source + self.name = name or source + self.request_id = request_id + self.message_id = message_id + self.conversation_id = conversation_id + self.activity_id = activity_id + self.workflow_run_id = workflow_run_id + self.user_id = user_id + self.agent_id = agent_id + self.otel_context = otel_context + self.otel_trace_id: Optional[str] = None + self.start_ns = time.time_ns() + self.start_perf_ns = time.perf_counter_ns() + self.end_perf_ns: Optional[int] = None + self.status: Optional[str] = None + self.spans: List[Span] = [] + self.dropped_spans = 0 + self.content_blocked = False + self.attributes: Dict[str, Any] = {} + self.finished = False + self.flushed = False + self._lock = threading.Lock() + self._stacks: Dict[int, List[Any]] = {} + + # -- span lifecycle ------------------------------------------------- + + def _parent_for_current_thread(self) -> Optional[str]: + stack = self._stacks.get(threading.get_ident()) + return stack[-1].id if stack else None + + def start_span( + self, + kind: str, + name: str, + *, + parent: Any = None, + attributes: Optional[Dict[str, Any]] = None, + ): + """Record a new span; returns :data:`NOOP_SPAN` once finished or over the cap.""" + with self._lock: + if self.finished: + return NOOP_SPAN + if len(self.spans) >= settings.TRACES_MAX_SPANS: + self.dropped_spans += 1 + return NOOP_SPAN + if parent is not None: + parent_id = getattr(parent, "id", None) + else: + parent_id = self._parent_for_current_thread() + span = Span(self, kind, name, parent_id, attributes) + self.spans.append(span) + if kind in CONTAINER_KINDS: + ident = threading.get_ident() + self._stacks.setdefault(ident, []).append(span) + span._thread = ident + span._pushed = True + return span + + def _on_end(self, span: Span) -> None: + if not span._pushed: + return + with self._lock: + stack = self._stacks.get(span._thread) + if not stack or span not in stack: + return + index = stack.index(span) + abandoned = stack[index + 1:] + del stack[index:] + if not stack: + self._stacks.pop(span._thread, None) + for child in reversed(abandoned): + if isinstance(child, Span) and not child.ended: + child.end(STATUS_CANCELLED) + + def _seed_thread(self, parent: Optional[Span]) -> Callable[[], None]: + """Make ``parent`` the implicit parent in this thread; returns an undo callable.""" + if parent is None: + return lambda: None + ident = threading.get_ident() + marker = _Seed(parent.id) + with self._lock: + self._stacks.setdefault(ident, []).append(marker) + + def _undo() -> None: + with self._lock: + stack = self._stacks.get(ident) + if stack and marker in stack: + del stack[stack.index(marker):] + if not stack: + self._stacks.pop(ident, None) + + return _undo + + def current_parent(self) -> Optional["_Seed | Span"]: + stack = self._stacks.get(threading.get_ident()) + return stack[-1] if stack else None + + # -- ids -------------------------------------------------------------- + + def bind(self, *, only_if_unset: bool = False, **ids: Any) -> None: + for key, value in ids.items(): + if key not in BINDABLE_IDS or value is None: + continue + if only_if_unset and getattr(self, key, None): + continue + setattr(self, key, str(value)) + + # -- completion --------------------------------------------------------- + + def finish(self, status: Optional[str] = None) -> None: + """Freeze the trace: close open spans as ``cancelled`` and set the status.""" + with self._lock: + if self.finished: + return + self.finished = True + open_spans = [s for s in self.spans if not s.ended] + self._stacks.clear() + now = time.perf_counter_ns() + for span in open_spans: + span.status = STATUS_CANCELLED + span.end_perf_ns = now + self.end_perf_ns = now + if status is not None: + self.status = status + else: + failed = any( + s.parent_id is None and s.status == STATUS_ERROR for s in self.spans + ) + self.status = STATUS_ERROR if failed else STATUS_OK + + @property + def duration_ms(self) -> Optional[float]: + if self.end_perf_ns is None: + return None + return (self.end_perf_ns - self.start_perf_ns) / 1e6 + + def span_start_ns(self, span: Span) -> int: + """Wall-clock start of ``span`` in ns, derived from the monotonic offset.""" + return self.start_ns + (span.start_perf_ns - self.start_perf_ns) + + def span_end_ns(self, span: Span) -> int: + end = span.end_perf_ns if span.end_perf_ns is not None else self.end_perf_ns + return self.start_ns + ((end or span.start_perf_ns) - self.start_perf_ns) + + def summary(self) -> Dict[str, Any]: + """Aggregate counts shown as chips in the Logs UI.""" + by_id = {s.id: s for s in self.spans} + llm = [s for s in self.spans if s.kind == KIND_LLM] + # Outermost retrieval spans only, so nested dispatcher/retriever + # spans are not double-counted. + retrieval = [ + s + for s in self.spans + if s.kind == KIND_RETRIEVAL + and not (s.parent_id in by_id and by_id[s.parent_id].kind == KIND_RETRIEVAL) + ] + tools = [s for s in self.spans if s.kind == KIND_TOOL] + + def _tokens(key: str) -> int: + total = 0 + for s in llm: + value = s.attributes.get(key) + if isinstance(value, (int, float)): + total += int(value) + return total + + return { + "llm_calls": len(llm), + "tool_calls": len(tools), + "retrieval_calls": len(retrieval), + "retrieval_ms": round(sum(s.duration_ms or 0 for s in retrieval), 1), + "input_tokens": _tokens("gen_ai.usage.input_tokens"), + "output_tokens": _tokens("gen_ai.usage.output_tokens"), + "errors": sum(1 for s in self.spans if s.status == STATUS_ERROR), + } + + def to_record(self) -> Dict[str, Any]: + """The ``request_traces`` row for this (finished) trace.""" + spans = [] + for s in self.spans: + entry: Dict[str, Any] = { + "id": s.id, + "parent_id": s.parent_id, + "kind": s.kind, + "name": s.name, + "status": s.status or STATUS_CANCELLED, + "offset_ms": round((s.start_perf_ns - self.start_perf_ns) / 1e6, 2), + "duration_ms": round(s.duration_ms or 0.0, 2), + "attributes": s.attributes, + } + if s.error: + entry["error"] = s.error + if s.previews and not self.content_blocked: + entry["preview"] = s.previews + spans.append(entry) + return { + "id": self.id, + "request_id": self.request_id, + "message_id": self.message_id, + "conversation_id": self.conversation_id, + "activity_id": self.activity_id, + "workflow_run_id": self.workflow_run_id, + "user_id": self.user_id, + "agent_id": self.agent_id, + "source": self.source, + "name": self.name, + "status": self.status or STATUS_OK, + "started_at_ns": self.start_ns, + "duration_ms": int(round(self.duration_ms or 0)), + "span_count": len(spans), + "dropped_spans": self.dropped_spans, + "summary": self.summary(), + "spans": spans, + "otel_trace_id": self.otel_trace_id, + } + + +class _Seed: + """Placeholder stack entry naming a parent span owned by another thread.""" + + __slots__ = ("id",) + + def __init__(self, span_id: Optional[str]) -> None: + self.id = span_id + + +# -- module-level API -------------------------------------------------------- + + +def current_trace() -> Optional[Trace]: + """The trace active in this context, or ``None``.""" + return _current.get() + + +def start_trace(*, source: str, capture_otel_context: bool = True, **ids: Any) -> Optional[Trace]: + """Create a trace, or return ``None`` when ``TRACES_ENABLED`` is off. + + The OTel context current at this point (normally the HTTP server span) is + captured so the exported GenAI spans hang off the request's own trace. + """ + if not settings.TRACES_ENABLED: + return None + otel_context = None + if capture_otel_context: + try: + from opentelemetry import context as otel_ctx + + otel_context = otel_ctx.get_current() + except Exception: # noqa: BLE001 + otel_context = None + known = {k: v for k, v in ids.items() if k in BINDABLE_IDS} + known = {k: (str(v) if v is not None else None) for k, v in known.items()} + return Trace(source=source, otel_context=otel_context, **known) + + +@contextlib.contextmanager +def activate(trace: Optional[Trace]) -> Iterator[Optional[Trace]]: + """Make ``trace`` current for the enclosed block (a no-op for ``None``).""" + if trace is None: + yield None + return + token = _current.set(trace) + try: + yield trace + finally: + try: + _current.reset(token) + except ValueError: + # Reset from a different context (a generator finalized + # elsewhere); clearing is the closest safe equivalent. + _current.set(None) + + +def start_span(kind: str, name: str, *, parent: Any = None, attributes: Optional[Dict[str, Any]] = None): + """Start a span in the current trace; returns :data:`NOOP_SPAN` without one.""" + trace = _current.get() + if trace is None: + return NOOP_SPAN + try: + return trace.start_span(kind, name, parent=parent, attributes=attributes) + except Exception: # noqa: BLE001 - tracing must never break a request + logger.debug("trace start_span failed", exc_info=True) + return NOOP_SPAN + + +def span(kind: str, name: str, *, parent: Any = None, attributes: Optional[Dict[str, Any]] = None): + """Context-manager form of :func:`start_span` (errors mark the span and re-raise).""" + return start_span(kind, name, parent=parent, attributes=attributes) + + +def bind(**ids: Any) -> None: + """Set link ids (``message_id``, ``conversation_id``, ...) on the current trace.""" + trace = _current.get() + if trace is not None: + trace.bind(**ids) + + +def bind_if_unset(**ids: Any) -> None: + """Like :func:`bind` but keeps any value already set.""" + trace = _current.get() + if trace is not None: + trace.bind(only_if_unset=True, **ids) + + +def mark_content_blocked() -> None: + """A guardrail blocked or retracted content: drop every preview from the stored trace.""" + trace = _current.get() + if trace is not None: + trace.content_blocked = True + + +def wrap(fn: Callable[..., Any]) -> Callable[..., Any]: + """Bind the current trace and parent span into ``fn`` for another thread. + + Use around work handed to a thread pool or ``threading.Thread``; those do + not inherit context variables. Without an active trace ``fn`` is returned + unchanged. + """ + trace = _current.get() + if trace is None: + return fn + parent = trace.current_parent() + + @functools.wraps(fn) + def _run(*args: Any, **kwargs: Any) -> Any: + token = _current.set(trace) + undo = trace._seed_thread(parent) + try: + return fn(*args, **kwargs) + finally: + undo() + try: + _current.reset(token) + except ValueError: + _current.set(None) + + return _run diff --git a/docsgpt/tracing/preview.py b/docsgpt/tracing/preview.py new file mode 100644 index 00000000..892e00f2 --- /dev/null +++ b/docsgpt/tracing/preview.py @@ -0,0 +1,58 @@ +"""Bounded, secret-redacted content previews stored with a trace span.""" + +from __future__ import annotations + +import json +from typing import Any + +from docsgpt.core.settings import settings +from docsgpt.storage.db.redaction import redact_secrets +from docsgpt.utils import strip_null_bytes + +_ELLIPSIS = "…" +# A list longer than this keeps its head only; the preview is for reading, +# not for reconstructing the payload. +_MAX_LIST_ITEMS = 50 + + +def _bound(value: Any, limit: int) -> Any: + """Truncate strings and long lists inside ``value``; stringify unknown objects.""" + if isinstance(value, str): + return value if len(value) <= limit else value[:limit] + _ELLIPSIS + if isinstance(value, dict): + return {str(k): _bound(v, limit) for k, v in value.items()} + if isinstance(value, (list, tuple)): + items = [_bound(v, limit) for v in list(value)[:_MAX_LIST_ITEMS]] + if len(value) > _MAX_LIST_ITEMS: + items.append(f"{_ELLIPSIS} {len(value) - _MAX_LIST_ITEMS} more") + return items + if value is None or isinstance(value, (bool, int, float)): + return value + return _bound(str(value), limit) + + +def make_preview(value: Any, limit: int | None = None) -> Any: + """Return a storable preview of ``value``. + + Secret-keyed fields are redacted, NUL bytes stripped and every string + truncated to ``limit`` (``TRACES_PREVIEW_CHARS`` by default). A structure + whose JSON form is still far larger than ``limit`` collapses to a + truncated JSON string so one huge tool result cannot bloat the row. + + Args: + value: Any JSON-like value (tool arguments, a result, a query string). + limit: Maximum characters per string; defaults to the setting. + + Returns: + A JSON-serializable preview. + """ + limit = limit or settings.TRACES_PREVIEW_CHARS + bounded = strip_null_bytes(_bound(redact_secrets(value), limit)) + if isinstance(bounded, (dict, list)): + try: + encoded = json.dumps(bounded, ensure_ascii=False, default=str) + except (TypeError, ValueError): + encoded = str(bounded) + if len(encoded) > limit * 4: + return encoded[:limit] + _ELLIPSIS + return bounded diff --git a/docsgpt/tracing/sink.py b/docsgpt/tracing/sink.py new file mode 100644 index 00000000..c5f9ad6c --- /dev/null +++ b/docsgpt/tracing/sink.py @@ -0,0 +1,53 @@ +"""Write a finished trace to its two sinks: Postgres and OpenTelemetry.""" + +from __future__ import annotations + +import logging +from typing import Optional + +from docsgpt.tracing.core import Trace + +logger = logging.getLogger(__name__) + + +def flush(trace: Optional[Trace], status: Optional[str] = None) -> None: + """Finish ``trace`` and persist it once; later calls are no-ops. + + OTel replay runs first so the exported trace id can be stored with the + row. Both sinks swallow their own failures: a trace is diagnostic data + and must never fail the request that produced it. + + Args: + trace: The trace to write; ``None`` is accepted and ignored. + status: Final status; defaults to ``error`` when a top-level span + failed, else ``ok``. + """ + if trace is None or trace.flushed: + return + trace.flushed = True + trace.finish(status) + try: + from docsgpt.tracing.otel import export_trace + + trace.otel_trace_id = export_trace(trace) + except Exception: # noqa: BLE001 + logger.warning("Failed to export trace %s to OpenTelemetry", trace.id, exc_info=True) + if not trace.spans: + # Nothing ran worth showing (e.g. an early validation error). + return + try: + from docsgpt.storage.db.repositories.request_traces import RequestTracesRepository + from docsgpt.storage.db.session import db_session + + with db_session() as conn: + RequestTracesRepository(conn).insert(trace.to_record()) + except Exception: # noqa: BLE001 + logger.warning("Failed to store trace %s", trace.id, exc_info=True) + + +def discard(trace: Optional[Trace]) -> None: + """Drop ``trace`` without writing it (the request was rejected or superseded).""" + if trace is None or trace.flushed: + return + trace.flushed = True + trace.finish() diff --git a/tests/tracing/__init__.py b/tests/tracing/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/tracing/test_trace.py b/tests/tracing/test_trace.py new file mode 100644 index 00000000..685a5595 --- /dev/null +++ b/tests/tracing/test_trace.py @@ -0,0 +1,341 @@ +"""Tests for the execution-trace recorder in ``docsgpt.tracing``.""" + +from __future__ import annotations + +import threading +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings + + +@pytest.fixture(autouse=True) +def _tracing_on(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + monkeypatch.setattr(settings, "TRACES_MAX_SPANS", 500) + monkeypatch.setattr(settings, "TRACES_PREVIEW_CHARS", 2000) + + +def _by_name(trace): + return {s.name: s for s in trace.spans} + + +class TestNoActiveTrace: + def test_span_calls_are_noops_without_a_trace(self): + assert tracing.current_trace() is None + with tracing.span(tracing.KIND_LLM, "chat gpt") as s: + s.set(foo=1) + s.preview("output", "hi") + handle = tracing.start_span(tracing.KIND_TOOL, "tool") + handle.end(status="error") + tracing.bind(message_id="m") + + def test_start_trace_returns_none_when_disabled(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", False) + assert tracing.start_trace(source="stream") is None + with tracing.activate(None): + assert tracing.current_trace() is None + + +class TestNesting: + def test_containers_push_and_leaves_record_parent(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_AGENT, "agent"): + with tracing.span(tracing.KIND_LLM, "llm-1"): + pass + with tracing.span(tracing.KIND_TOOL, "tool"): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + with tracing.span(tracing.KIND_EMBEDDING, "embed"): + pass + with tracing.span(tracing.KIND_LLM, "llm-2"): + pass + spans = _by_name(trace) + assert spans["agent"].parent_id is None + assert spans["llm-1"].parent_id == spans["agent"].id + assert spans["tool"].parent_id == spans["agent"].id + assert spans["retrieval"].parent_id == spans["tool"].id + assert spans["embed"].parent_id == spans["retrieval"].id + assert spans["llm-2"].parent_id == spans["agent"].id + assert all(s.status == "ok" for s in trace.spans) + + def test_leaf_span_does_not_become_a_parent(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + llm = tracing.start_span(tracing.KIND_LLM, "llm") + with tracing.span(tracing.KIND_TOOL, "tool"): + pass + llm.end() + spans = _by_name(trace) + assert spans["tool"].parent_id is None + + def test_explicit_parent_wins(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + outer = tracing.start_span(tracing.KIND_RETRIEVAL, "outer") + with tracing.span(tracing.KIND_AGENT, "agent"): + child = tracing.start_span(tracing.KIND_SEARCH, "child", parent=outer) + child.end() + outer.end() + assert _by_name(trace)["child"].parent_id == _by_name(trace)["outer"].id + + def test_interleaved_generators_nest_by_start_order(self): + """An agent generator suspended at yield keeps its children nested.""" + trace = tracing.start_trace(source="stream") + + def llm_stream(): + handle = tracing.start_span(tracing.KIND_LLM, "llm") + try: + yield "a" + yield "b" + finally: + handle.end() + + def agent_gen(): + handle = tracing.start_span(tracing.KIND_AGENT, "agent") + try: + yield from llm_stream() + with tracing.span(tracing.KIND_TOOL, "tool"): + pass + yield "c" + finally: + handle.end() + + with tracing.activate(trace): + assert list(agent_gen()) == ["a", "b", "c"] + spans = _by_name(trace) + assert spans["llm"].parent_id == spans["agent"].id + assert spans["tool"].parent_id == spans["agent"].id + + def test_out_of_order_end_cancels_abandoned_children(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + agent = tracing.start_span(tracing.KIND_AGENT, "agent") + tool = tracing.start_span(tracing.KIND_TOOL, "tool") + agent.end() + after = tracing.start_span(tracing.KIND_LLM, "after") + after.end() + tool.end() # late end is ignored + spans = _by_name(trace) + assert spans["tool"].status == "cancelled" + assert spans["agent"].status == "ok" + assert spans["after"].parent_id is None + + def test_exception_marks_span_error_and_propagates(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with pytest.raises(ValueError): + with tracing.span(tracing.KIND_TOOL, "tool"): + raise ValueError("boom") + span = trace.spans[0] + assert span.status == "error" + assert span.attributes["error.type"] == "ValueError" + assert span.error == "boom" + + def test_generator_exit_marks_span_cancelled(self): + trace = tracing.start_trace(source="stream") + + def gen(): + with tracing.span(tracing.KIND_AGENT, "agent"): + yield 1 + yield 2 + + with tracing.activate(trace): + g = gen() + next(g) + g.close() + assert trace.spans[0].status == "cancelled" + + +class TestFinish: + def test_finish_cancels_open_spans_and_ignores_late_ends(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + open_span = tracing.start_span(tracing.KIND_AGENT, "agent") + trace.finish() + assert open_span.status == "cancelled" + duration = open_span.duration_ms + open_span.end(status="ok") + assert open_span.status == "cancelled" + assert open_span.duration_ms == duration + + def test_finish_status_defaults_to_error_when_a_top_level_span_failed(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + tracing.start_span(tracing.KIND_AGENT, "agent").end(status="error") + trace.finish() + assert trace.status == "error" + + def test_explicit_finish_status(self): + trace = tracing.start_trace(source="stream") + trace.finish(status="paused") + assert trace.status == "paused" + + def test_spans_after_finish_are_dropped(self): + trace = tracing.start_trace(source="stream") + trace.finish() + with tracing.activate(trace): + tracing.start_span(tracing.KIND_LLM, "late").end() + assert trace.spans == [] + + def test_summary_rolls_up_llm_tool_retrieval(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_AGENT, "agent"): + with tracing.span( + tracing.KIND_LLM, + "chat", + attributes={ + "gen_ai.usage.input_tokens": 10, + "gen_ai.usage.output_tokens": 4, + }, + ): + pass + with tracing.span(tracing.KIND_RETRIEVAL, "r"): + with tracing.span(tracing.KIND_RETRIEVAL, "inner"): + pass + with tracing.span(tracing.KIND_TOOL, "t") as t: + t.end(status="error") + trace.finish() + summary = trace.summary() + assert summary["llm_calls"] == 1 + assert summary["tool_calls"] == 1 + assert summary["input_tokens"] == 10 + assert summary["output_tokens"] == 4 + assert summary["retrieval_calls"] == 1 + assert summary["errors"] == 1 + assert summary["retrieval_ms"] >= 0 + + +class TestCap: + def test_span_cap_counts_dropped(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_MAX_SPANS", 10) + trace = tracing.start_trace(source="graph_extraction") + with tracing.activate(trace): + for i in range(15): + with tracing.span(tracing.KIND_LLM, f"llm-{i}"): + pass + assert len(trace.spans) == 10 + assert trace.dropped_spans == 5 + + +class TestThreads: + def test_wrap_carries_trace_and_parent_into_pool_threads(self): + trace = tracing.start_trace(source="stream") + + def work(i): + with tracing.span(tracing.KIND_SEARCH, f"search-{i}"): + with tracing.span(tracing.KIND_EMBEDDING, f"embed-{i}"): + pass + return threading.get_ident() + + with tracing.activate(trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + with ThreadPoolExecutor(max_workers=3) as pool: + list(pool.map(tracing.wrap(work), range(3))) + spans = _by_name(trace) + for i in range(3): + assert spans[f"search-{i}"].parent_id == spans["retrieval"].id + assert spans[f"embed-{i}"].parent_id == spans["retrieval"].id + + def test_pool_threads_without_wrap_see_no_trace(self): + trace = tracing.start_trace(source="stream") + seen = [] + with tracing.activate(trace): + t = threading.Thread(target=lambda: seen.append(tracing.current_trace())) + t.start() + t.join() + assert seen == [None] + + def test_wrap_without_trace_is_passthrough(self): + fn = tracing.wrap(lambda x: x + 1) + assert fn(1) == 2 + + +class TestBind: + def test_bind_sets_ids_on_active_trace(self): + trace = tracing.start_trace(source="stream", request_id="r1") + with tracing.activate(trace): + tracing.bind(message_id="m1", conversation_id="c1", unknown="x") + assert trace.request_id == "r1" + assert trace.message_id == "m1" + assert trace.conversation_id == "c1" + + def test_bind_if_unset_keeps_first_value(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + tracing.bind_if_unset(activity_id="a1") + tracing.bind_if_unset(activity_id="a2") + assert trace.activity_id == "a1" + + +class TestPreviews: + def test_preview_redacts_secrets_and_bounds_strings(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_PREVIEW_CHARS", 100) + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "tool") as s: + s.preview("arguments", {"api_key": "sk-123", "q": "x" * 500}) + s.preview("result", "y" * 500 + "\x00") + preview = trace.spans[0].previews + assert preview["arguments"]["api_key"] == "[REDACTED]" + assert len(preview["arguments"]["q"]) < 200 + assert preview["result"].endswith("…") + assert "\x00" not in preview["result"] + + def test_preview_skipped_when_capture_disabled(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", False) + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "tool") as s: + s.preview("result", "secret stuff") + assert trace.spans[0].previews == {} + + def test_guardrail_trigger_strips_all_previews(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "tool") as s: + s.preview("result", "leaked") + tracing.mark_content_blocked() + trace.finish() + record = trace.to_record() + assert all(not s.get("preview") for s in record["spans"]) + assert trace.content_blocked is True + + def test_huge_structures_collapse_to_a_string(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_PREVIEW_CHARS", 100) + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with tracing.span(tracing.KIND_TOOL, "tool") as s: + s.preview("result", [{"k": "v" * 50} for _ in range(100)]) + value = trace.spans[0].previews["result"] + assert isinstance(value, str) + assert len(value) <= 101 + + +class TestRecord: + def test_to_record_shape(self): + trace = tracing.start_trace( + source="stream", request_id="r", user_id="u", agent_id="a" + ) + with tracing.activate(trace): + with tracing.span( + tracing.KIND_LLM, "chat m", attributes={"gen_ai.request.model": "m"} + ): + pass + trace.finish() + record = trace.to_record() + assert record["source"] == "stream" + assert record["request_id"] == "r" + assert record["status"] == "ok" + assert record["span_count"] == 1 + span = record["spans"][0] + assert set(span) >= { + "id", "parent_id", "kind", "name", "status", "offset_ms", + "duration_ms", "attributes", + } + assert span["attributes"]["gen_ai.request.model"] == "m" + assert span["offset_ms"] >= 0 From b5e92576590f5e6988f5a1187a89f91ffdde5026 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:21:25 +0100 Subject: [PATCH 04/33] Export traces as OpenTelemetry GenAI spans and metrics Finished traces are replayed with their recorded timestamps and explicit parents under the request's server span, using gen_ai.* attribute names. Content is exported only when OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT opts in. gen_ai.client token-usage and operation-duration histograms are recorded per model call. --- docsgpt/tracing/otel.py | 274 ++++++++++++++++++++++++++++++ tests/tracing/test_otel_export.py | 235 +++++++++++++++++++++++++ tests/tracing/test_sink.py | 100 +++++++++++ 3 files changed, 609 insertions(+) create mode 100644 docsgpt/tracing/otel.py create mode 100644 tests/tracing/test_otel_export.py create mode 100644 tests/tracing/test_sink.py diff --git a/docsgpt/tracing/otel.py b/docsgpt/tracing/otel.py new file mode 100644 index 00000000..40a8b5ef --- /dev/null +++ b/docsgpt/tracing/otel.py @@ -0,0 +1,274 @@ +"""OpenTelemetry export for execution traces (GenAI semantic conventions). + +A finished :class:`~docsgpt.tracing.core.Trace` is *replayed*: every recorded +span is re-created with its original start and end timestamps and explicit +parents, under the OTel context captured when the trace started (normally the +HTTP server span). Replaying after the fact means no OTel context is ever +attached across the agent loop's generator yields. The cost is that GenAI +spans reach the backend only when the request ends, and auto-instrumented +HTTP/DB spans made during a step do not nest under it. + +Attribute names follow the (still incubating) ``gen_ai.*`` conventions and are +written as literals rather than imported from the semconv package. Prompt and +tool content is exported only when +``OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT`` opts in. + +Nothing here runs unless an OTel SDK is installed as the global provider (for +example by launching under ``opentelemetry-instrument``). +""" + +from __future__ import annotations + +import json +import logging +import os +import threading +from typing import Any, Dict, Optional + +from docsgpt.core.settings import settings +from docsgpt.tracing.core import ( + KIND_EMBEDDING, + KIND_LLM, + KIND_TOOL, + STATUS_ERROR, + STATUS_OK, + Trace, +) + +logger = logging.getLogger(__name__) + +_INSTRUMENTATION_NAME = "docsgpt.tracing" + +#: DocsGPT provider names that differ from the ``gen_ai.provider.name`` values. +_PROVIDER_NAMES = { + "google": "gcp.gen_ai", + "azure_openai": "azure.ai.openai", + "aws_bedrock": "aws.bedrock", + "mistral": "mistral_ai", +} + +#: Preview keys with a dedicated GenAI attribute; others export as ``docsgpt.preview.``. +_TOOL_PREVIEW_ATTRIBUTES = { + "arguments": "gen_ai.tool.call.arguments", + "result": "gen_ai.tool.call.result", +} + +_CAPTURE_VALUES = {"span_only", "span_and_event", "true"} + +# Advisory buckets from the GenAI metrics conventions. +_DURATION_BUCKETS = [0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28, 2.56, 5.12, 10.24, 20.48, 40.96, 81.92] +_TOKEN_BUCKETS = [1, 4, 16, 64, 256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864] + + +def provider_name(internal: Optional[str]) -> str: + """Map a DocsGPT ``provider_name`` to the ``gen_ai.provider.name`` value.""" + if not internal: + return "unknown" + return _PROVIDER_NAMES.get(internal, internal) + + +def content_capture_enabled() -> bool: + """Whether the standard GenAI content-capture variable opts into span content.""" + value = os.getenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "") + return value.strip().lower() in _CAPTURE_VALUES + + +def _is_noop(provider: Any) -> bool: + from opentelemetry import trace as ot_trace + + return isinstance(provider, (ot_trace.ProxyTracerProvider, ot_trace.NoOpTracerProvider)) + + +def _attr_value(value: Any) -> Any: + """Coerce a recorded attribute into an OTel-legal value (or ``None`` to skip).""" + if value is None: + return None + if isinstance(value, (str, bool, int, float)): + return value + if isinstance(value, (list, tuple)): + items = [v for v in value if v is not None] + if items and all(isinstance(v, str) for v in items): + return list(items) + if items and all(isinstance(v, bool) for v in items): + return list(items) + if items and all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in items): + return list(items) + if not items: + return None + try: + return json.dumps(value, ensure_ascii=False, default=str) + except (TypeError, ValueError): + return str(value) + + +def _attributes(raw: Dict[str, Any]) -> Dict[str, Any]: + out = {} + for key, value in raw.items(): + coerced = _attr_value(value) + if coerced is not None: + out[key] = coerced + return out + + +def _preview_attributes(kind: str, previews: Dict[str, Any]) -> Dict[str, Any]: + out = {} + for key, value in previews.items(): + name = _TOOL_PREVIEW_ATTRIBUTES.get(key) if kind == KIND_TOOL else None + out[name or f"docsgpt.preview.{key}"] = value if isinstance(value, str) else _attr_value(value) + return out + + +def _root_attributes(trace: Trace) -> Dict[str, Any]: + return _attributes( + { + "docsgpt.trace.id": trace.id, + "docsgpt.source": trace.source, + "docsgpt.request_id": trace.request_id, + "docsgpt.message_id": trace.message_id, + "docsgpt.activity_id": trace.activity_id, + "docsgpt.workflow_run_id": trace.workflow_run_id, + "gen_ai.conversation.id": trace.conversation_id, + "gen_ai.agent.id": trace.agent_id, + "docsgpt.status": trace.status, + "docsgpt.dropped_spans": trace.dropped_spans or None, + } + ) + + +def export_trace(trace: Trace, tracer_provider: Any = None) -> Optional[str]: + """Replay ``trace`` as OTel spans; returns the OTel trace id (hex) or ``None``. + + Args: + trace: A finished trace. + tracer_provider: Override the global provider (tests). + + Returns: + The 32-char hex OTel trace id, or ``None`` when export is disabled, + no SDK is configured, or the trace is empty. + """ + if not settings.TRACES_OTEL_EXPORT or not trace.spans: + return None + from opentelemetry import trace as ot_trace + from opentelemetry.trace import SpanKind, Status, StatusCode + + provider = tracer_provider or ot_trace.get_tracer_provider() + if _is_noop(provider): + return None + tracer = provider.get_tracer(_INSTRUMENTATION_NAME) + capture = content_capture_enabled() and not trace.content_blocked + + root = tracer.start_span( + f"docsgpt {trace.source}", + context=trace.otel_context, + kind=SpanKind.INTERNAL, + attributes=_root_attributes(trace), + start_time=trace.start_ns, + ) + if trace.status == STATUS_ERROR: + root.set_status(Status(StatusCode.ERROR)) + exported: Dict[str, Any] = {} + for span in trace.spans: + parent = exported.get(span.parent_id, root) + attributes = dict(span.attributes) + if trace.conversation_id and span.kind in (KIND_LLM, KIND_TOOL): + attributes.setdefault("gen_ai.conversation.id", trace.conversation_id) + if span.status and span.status not in (STATUS_OK, STATUS_ERROR): + attributes["docsgpt.status"] = span.status + attributes = _attributes(attributes) + if capture and span.previews: + attributes.update(_preview_attributes(span.kind, span.previews)) + otel_span = tracer.start_span( + span.name, + context=ot_trace.set_span_in_context(parent), + kind=SpanKind.CLIENT if span.kind in (KIND_LLM, KIND_EMBEDDING) else SpanKind.INTERNAL, + attributes=attributes, + start_time=trace.span_start_ns(span), + ) + if span.status == STATUS_ERROR: + otel_span.set_status(Status(StatusCode.ERROR, span.error)) + exported[span.id] = otel_span + for span in trace.spans: + exported[span.id].end(end_time=trace.span_end_ns(span)) + root.end(end_time=trace.start_ns + int((trace.duration_ms or 0) * 1e6)) + context = root.get_span_context() + return format(context.trace_id, "032x") if context.trace_id else None + + +# -- metrics -------------------------------------------------------------------- + +_instruments_lock = threading.Lock() +_global_instruments: Optional[tuple] = None + + +def _make_instruments(meter_provider: Any = None) -> tuple: + from opentelemetry import metrics + + meter = ( + meter_provider.get_meter(_INSTRUMENTATION_NAME) + if meter_provider is not None + else metrics.get_meter(_INSTRUMENTATION_NAME) + ) + + def _histogram(name: str, unit: str, description: str, buckets: list): + try: + return meter.create_histogram( + name, unit=unit, description=description, + explicit_bucket_boundaries_advisory=buckets, + ) + except TypeError: # API older than bucket advisories + return meter.create_histogram(name, unit=unit, description=description) + + duration = _histogram( + "gen_ai.client.operation.duration", "s", "GenAI operation duration.", _DURATION_BUCKETS + ) + tokens = _histogram( + "gen_ai.client.token.usage", "{token}", "Number of input and output tokens used.", _TOKEN_BUCKETS + ) + return duration, tokens + + +def _instruments(meter_provider: Any = None) -> tuple: + global _global_instruments + if meter_provider is not None: + return _make_instruments(meter_provider) + with _instruments_lock: + if _global_instruments is None: + _global_instruments = _make_instruments() + return _global_instruments + + +def record_llm_metrics( + *, + provider: Optional[str], + model: Optional[str], + input_tokens: int, + output_tokens: int, + duration_s: float, + error_type: Optional[str], + operation: str = "chat", + meter_provider: Any = None, +) -> None: + """Record the two GenAI client metrics for one model call. + + Token usage is recorded only for successful calls, as the conventions + require; duration is recorded for every call, with ``error.type`` on + failures. Without a configured SDK the meter is a no-op. + """ + if not settings.TRACES_OTEL_EXPORT: + return + try: + duration, tokens = _instruments(meter_provider) + attributes = { + "gen_ai.operation.name": operation, + "gen_ai.provider.name": provider_name(provider), + } + if model: + attributes["gen_ai.request.model"] = str(model) + if error_type: + duration.record(duration_s, {**attributes, "error.type": error_type}) + return + duration.record(duration_s, attributes) + tokens.record(int(input_tokens or 0), {**attributes, "gen_ai.token.type": "input"}) + tokens.record(int(output_tokens or 0), {**attributes, "gen_ai.token.type": "output"}) + except Exception: # noqa: BLE001 - metrics must never break a call + logger.debug("Failed to record GenAI metrics", exc_info=True) diff --git a/tests/tracing/test_otel_export.py b/tests/tracing/test_otel_export.py new file mode 100644 index 00000000..c93ca935 --- /dev/null +++ b/tests/tracing/test_otel_export.py @@ -0,0 +1,235 @@ +"""Tests for replaying a finished trace as OpenTelemetry GenAI spans.""" + +from __future__ import annotations + +import pytest +from opentelemetry import trace as ot_trace +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanKind, StatusCode + +from docsgpt import tracing +from docsgpt.core.settings import settings +from docsgpt.tracing import otel as trace_otel + + +@pytest.fixture(autouse=True) +def _settings(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + monkeypatch.delenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", raising=False) + + +@pytest.fixture() +def provider(): + exporter = InMemorySpanExporter() + tp = TracerProvider() + tp.add_span_processor(SimpleSpanProcessor(exporter)) + tp.exporter = exporter + return tp + + +def _sample_trace(otel_context=None): + trace = tracing.start_trace( + source="stream", + request_id="req-1", + message_id="msg-1", + conversation_id="conv-1", + agent_id="agent-1", + capture_otel_context=False, + ) + trace.otel_context = otel_context + with tracing.activate(trace): + with tracing.span( + tracing.KIND_AGENT, + "invoke_agent Support", + attributes={"gen_ai.operation.name": "invoke_agent", "gen_ai.agent.name": "Support"}, + ): + with tracing.span( + tracing.KIND_LLM, + "chat gpt-4o", + attributes={ + "gen_ai.operation.name": "chat", + "gen_ai.provider.name": "openai", + "gen_ai.request.model": "gpt-4o", + "gen_ai.usage.input_tokens": 12, + "gen_ai.usage.output_tokens": 3, + "docsgpt.sources": ["a", "b"], + "docsgpt.meta": {"nested": True}, + }, + ) as llm: + llm.preview("output", "hello") + with tracing.span( + tracing.KIND_TOOL, + "execute_tool search", + attributes={"gen_ai.tool.name": "search"}, + ) as tool: + tool.preview("arguments", {"q": "x"}) + tool.preview("result", "found") + tool.end(error=RuntimeError("tool broke")) + trace.finish() + return trace + + +class TestReplay: + def test_spans_are_parented_and_timed(self, provider): + trace = _sample_trace() + otel_trace_id = trace_otel.export_trace(trace, tracer_provider=provider) + spans = {s.name: s for s in provider.exporter.get_finished_spans()} + root = spans["docsgpt stream"] + agent = spans["invoke_agent Support"] + llm = spans["chat gpt-4o"] + tool = spans["execute_tool search"] + assert agent.parent.span_id == root.context.span_id + assert llm.parent.span_id == agent.context.span_id + assert tool.parent.span_id == agent.context.span_id + assert otel_trace_id == format(root.context.trace_id, "032x") + recorded = {s.name: s for s in trace.spans} + assert llm.start_time == trace.span_start_ns(recorded["chat gpt-4o"]) + assert llm.end_time == trace.span_end_ns(recorded["chat gpt-4o"]) + assert root.start_time == trace.start_ns + assert root.end_time >= agent.end_time + + def test_semconv_attributes_and_kinds(self, provider): + trace_otel.export_trace(_sample_trace(), tracer_provider=provider) + spans = {s.name: s for s in provider.exporter.get_finished_spans()} + llm = spans["chat gpt-4o"] + assert llm.kind == SpanKind.CLIENT + assert llm.attributes["gen_ai.usage.input_tokens"] == 12 + assert llm.attributes["gen_ai.conversation.id"] == "conv-1" + assert tuple(llm.attributes["docsgpt.sources"]) == ("a", "b") + assert llm.attributes["docsgpt.meta"] == '{"nested": true}' + root = spans["docsgpt stream"] + assert root.attributes["docsgpt.request_id"] == "req-1" + assert root.attributes["docsgpt.message_id"] == "msg-1" + + def test_error_status(self, provider): + trace_otel.export_trace(_sample_trace(), tracer_provider=provider) + tool = {s.name: s for s in provider.exporter.get_finished_spans()}["execute_tool search"] + assert tool.status.status_code == StatusCode.ERROR + assert tool.attributes["error.type"] == "RuntimeError" + + def test_no_content_by_default(self, provider): + trace_otel.export_trace(_sample_trace(), tracer_provider=provider) + for span in provider.exporter.get_finished_spans(): + assert not any("preview" in k or "call.arguments" in k for k in span.attributes) + + def test_content_when_opted_in(self, provider, monkeypatch): + monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "span_only") + trace_otel.export_trace(_sample_trace(), tracer_provider=provider) + spans = {s.name: s for s in provider.exporter.get_finished_spans()} + tool = spans["execute_tool search"] + assert tool.attributes["gen_ai.tool.call.arguments"] == '{"q": "x"}' + assert tool.attributes["gen_ai.tool.call.result"] == "found" + assert spans["chat gpt-4o"].attributes["docsgpt.preview.output"] == "hello" + + def test_content_blocked_trace_never_exports_content(self, provider, monkeypatch): + monkeypatch.setenv("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "SPAN_ONLY") + trace = _sample_trace() + trace.content_blocked = True + trace_otel.export_trace(trace, tracer_provider=provider) + for span in provider.exporter.get_finished_spans(): + assert not any("preview" in k or "call.result" in k for k in span.attributes) + + def test_root_is_parented_to_captured_context(self, provider): + server = provider.get_tracer("test").start_span("GET /stream") + context = ot_trace.set_span_in_context(server) + server.end() + trace_otel.export_trace(_sample_trace(otel_context=context), tracer_provider=provider) + root = {s.name: s for s in provider.exporter.get_finished_spans()}["docsgpt stream"] + assert root.parent.span_id == server.get_span_context().span_id + assert root.context.trace_id == server.get_span_context().trace_id + + def test_skipped_without_sdk_provider(self): + assert trace_otel.export_trace(_sample_trace(), tracer_provider=ot_trace.NoOpTracerProvider()) is None + + def test_skipped_when_disabled(self, provider, monkeypatch): + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + assert trace_otel.export_trace(_sample_trace(), tracer_provider=provider) is None + assert provider.exporter.get_finished_spans() == () + + +class TestProviderName: + @pytest.mark.parametrize( + "internal, expected", + [ + ("openai", "openai"), + ("google", "gcp.gen_ai"), + ("anthropic", "anthropic"), + ("azure_openai", "azure.ai.openai"), + ("groq", "groq"), + (None, "unknown"), + ], + ) + def test_mapping(self, internal, expected): + assert trace_otel.provider_name(internal) == expected + + +class TestMetrics: + def test_llm_metrics_recorded(self): + reader = InMemoryMetricReader() + mp = MeterProvider(metric_readers=[reader]) + trace_otel.record_llm_metrics( + provider="openai", + model="gpt-4o", + input_tokens=10, + output_tokens=5, + duration_s=0.25, + error_type=None, + meter_provider=mp, + ) + data = reader.get_metrics_data() + metrics = { + m.name: m + for rm in data.resource_metrics + for sm in rm.scope_metrics + for m in sm.metrics + } + usage = metrics["gen_ai.client.token.usage"] + by_type = { + p.attributes["gen_ai.token.type"]: p.sum for p in usage.data.data_points + } + assert by_type == {"input": 10, "output": 5} + point = usage.data.data_points[0] + assert point.attributes["gen_ai.provider.name"] == "openai" + assert point.attributes["gen_ai.request.model"] == "gpt-4o" + assert point.attributes["gen_ai.operation.name"] == "chat" + duration = metrics["gen_ai.client.operation.duration"].data.data_points[0] + assert duration.sum == pytest.approx(0.25) + + def test_error_type_on_duration_only(self): + reader = InMemoryMetricReader() + mp = MeterProvider(metric_readers=[reader]) + trace_otel.record_llm_metrics( + provider="openai", + model="m", + input_tokens=0, + output_tokens=0, + duration_s=0.1, + error_type="Timeout", + meter_provider=mp, + ) + metrics = { + m.name: m + for rm in reader.get_metrics_data().resource_metrics + for sm in rm.scope_metrics + for m in sm.metrics + } + assert "gen_ai.client.token.usage" not in metrics + point = metrics["gen_ai.client.operation.duration"].data.data_points[0] + assert point.attributes["error.type"] == "Timeout" + + def test_metrics_disabled(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + reader = InMemoryMetricReader() + mp = MeterProvider(metric_readers=[reader]) + trace_otel.record_llm_metrics( + provider="openai", model="m", input_tokens=1, output_tokens=1, + duration_s=0.1, error_type=None, meter_provider=mp, + ) + data = reader.get_metrics_data() + assert data is None or not data.resource_metrics diff --git a/tests/tracing/test_sink.py b/tests/tracing/test_sink.py new file mode 100644 index 00000000..7fc7b188 --- /dev/null +++ b/tests/tracing/test_sink.py @@ -0,0 +1,100 @@ +"""Tests for writing a finished trace to Postgres and OTel.""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings + + +@pytest.fixture(autouse=True) +def _settings(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + + +def _trace_with_span(): + trace = tracing.start_trace(source="stream", request_id="r1", user_id="u1") + with tracing.activate(trace): + with tracing.span(tracing.KIND_LLM, "chat m"): + pass + return trace + + +@contextmanager +def _patched_store(): + repo = MagicMock() + repo.insert.return_value = True + + @contextmanager + def _session(): + yield MagicMock() + + with patch("docsgpt.storage.db.session.db_session", _session), patch( + "docsgpt.storage.db.repositories.request_traces.RequestTracesRepository", + return_value=repo, + ): + yield repo + + +class TestFlush: + def test_flush_writes_once(self): + trace = _trace_with_span() + with _patched_store() as repo: + tracing.flush(trace) + tracing.flush(trace) + assert repo.insert.call_count == 1 + record = repo.insert.call_args[0][0] + assert record["request_id"] == "r1" + assert record["status"] == "ok" + assert trace.finished and trace.flushed + + def test_flush_status_override(self): + trace = _trace_with_span() + with _patched_store() as repo: + tracing.flush(trace, status="paused") + assert repo.insert.call_args[0][0]["status"] == "paused" + + def test_empty_trace_is_not_stored(self): + trace = tracing.start_trace(source="stream") + with _patched_store() as repo: + tracing.flush(trace) + repo.insert.assert_not_called() + + def test_store_failure_is_swallowed(self): + trace = _trace_with_span() + with _patched_store() as repo: + repo.insert.side_effect = RuntimeError("db down") + tracing.flush(trace) # must not raise + + def test_otel_failure_does_not_block_store(self): + trace = _trace_with_span() + with _patched_store() as repo, patch( + "docsgpt.tracing.otel.export_trace", side_effect=RuntimeError("otel") + ): + tracing.flush(trace) + repo.insert.assert_called_once() + + def test_otel_trace_id_is_stored(self): + trace = _trace_with_span() + with _patched_store() as repo, patch( + "docsgpt.tracing.otel.export_trace", return_value="ab" * 16 + ): + tracing.flush(trace) + assert repo.insert.call_args[0][0]["otel_trace_id"] == "ab" * 16 + + def test_none_is_ignored(self): + tracing.flush(None) + tracing.discard(None) + + +class TestDiscard: + def test_discard_never_writes(self): + trace = _trace_with_span() + with _patched_store() as repo: + tracing.discard(trace) + tracing.flush(trace) + repo.insert.assert_not_called() From e8166e6c29db22535732a6e83f3bcdff1cafda30 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:23:06 +0100 Subject: [PATCH 05/33] Record a chat span for every LLM call The token-usage wrappers open one span per decorated invocation, so a primary attempt, its retry and a fallback appear as siblings with the provider that ran. Stream spans start on the first pull, carry tokens, cost, time to first token and cache hits, and feed the GenAI metrics. --- docsgpt/cache.py | 8 ++ docsgpt/tracing/llm.py | 119 ++++++++++++++++++ docsgpt/usage.py | 46 ++++++- tests/tracing/test_llm_spans.py | 216 ++++++++++++++++++++++++++++++++ 4 files changed, 385 insertions(+), 4 deletions(-) create mode 100644 docsgpt/tracing/llm.py create mode 100644 tests/tracing/test_llm_spans.py 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 From a15940c28121ab638a3f5ccce744b0504d1727c7 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:25:41 +0100 Subject: [PATCH 06/33] Trace agent runs and tool calls @log_activity opens the invoke_agent span for every agent run and names the trace's activity; gen_continuation opens its own. ToolExecutor.execute records each executed call with redacted argument and result previews, and paused, denied, skipped and client-executed calls are recorded where they are decided. --- docsgpt/agents/base.py | 81 +++++++--- docsgpt/agents/tool_executor.py | 82 ++++++++++- docsgpt/llm/handlers/base.py | 45 ++++-- docsgpt/logging.py | 48 ++++++ tests/test_logging.py | 92 ++++++++++++ tests/tracing/test_tool_spans.py | 244 +++++++++++++++++++++++++++++++ 6 files changed, 557 insertions(+), 35 deletions(-) create mode 100644 tests/tracing/test_tool_spans.py diff --git a/docsgpt/agents/base.py b/docsgpt/agents/base.py index 030d76fc..002e7fa6 100644 --- a/docsgpt/agents/base.py +++ b/docsgpt/agents/base.py @@ -6,8 +6,11 @@ from abc import ABC, abstractmethod from datetime import datetime, timezone from typing import Any, Dict, Generator, List, Optional +from docsgpt import tracing from docsgpt.agents.tool_executor import ( ToolExecutor, + finish_tool_span, + record_tool_span_start, result_status, truncate_tool_result, ) @@ -492,12 +495,51 @@ class BaseAgent(ABC): or client-side results), appends the resulting messages, then hands back to the LLM to continue the conversation. + Unlike :meth:`gen` this is not wrapped by ``@log_activity``, so the + continuation's ``invoke_agent`` trace span is opened here. + Args: messages: The saved messages array from the pause point. tools_dict: The saved tools dictionary. 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 + }, + ) + completed = False + try: + yield from self._gen_continuation_inner( + messages, tools_dict, pending_tool_calls, tool_actions, reasoning_content + ) + completed = True + except Exception as exc: + span.end(error=exc) + raise + finally: + span.end(None if completed else tracing.STATUS_CANCELLED) + + def _gen_continuation_inner( + self, + messages: List[Dict], + tools_dict: Dict, + pending_tool_calls: List[Dict], + tool_actions: List[Dict], + reasoning_content: str = "", + ) -> Generator[Dict, None, None]: self._prepare_tools(tools_dict) actions_by_id = {a["call_id"]: a for a in tool_actions} @@ -583,16 +625,15 @@ class BaseAgent(ABC): messages.append( self.llm_handler.create_tool_message(tc, denial) ) - yield { - "type": "tool_call", - "data": { - "tool_name": pending.get("tool_name", "unknown"), - "call_id": call_id, - "action_name": pending.get("llm_name", pending["name"]), - "arguments": args, - "status": "denied", - }, + denied_data = { + "tool_name": pending.get("tool_name", "unknown"), + "call_id": call_id, + "action_name": pending.get("llm_name", pending["name"]), + "arguments": args, + "status": "denied", } + finish_tool_span(record_tool_span_start(tc), {**denied_data, "error": comment or None}) + yield {"type": "tool_call", "data": denied_data} elif "result" in action: result = action["result"] @@ -616,17 +657,19 @@ class BaseAgent(ABC): tc, _bound_tool_response_for_llm(result_str) ) ) - yield { - "type": "tool_call", - "data": { - "tool_name": pending.get("tool_name", "unknown"), - "call_id": call_id, - "action_name": pending.get("llm_name", pending["name"]), - "arguments": args, - "result": truncate_tool_result(result_str), - "status": result_status(result), - }, + client_data = { + "tool_name": pending.get("tool_name", "unknown"), + "call_id": call_id, + "action_name": pending.get("llm_name", pending["name"]), + "arguments": args, + "result": truncate_tool_result(result_str), + "status": result_status(result), } + finish_tool_span( + record_tool_span_start(tc, **{"docsgpt.client_executed": True}), + client_data, + ) + yield {"type": "tool_call", "data": client_data} # Resume the LLM loop with the updated messages llm_response = self._llm_gen(messages, preserve_responses_state=True) diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 9952e3d2..155ea925 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -13,6 +13,7 @@ from docsgpt.agents.default_tools import ( resolve_tool_by_id, synthesized_default_tools, ) +from docsgpt import tracing from docsgpt.agents.tools.tool_action_parser import ToolActionParser from docsgpt.agents.tools.tool_manager import ToolManager from docsgpt.guardrails.types import Stage as GuardrailStage, resolve_tool_result @@ -29,6 +30,65 @@ from docsgpt.storage.db.session import db_readonly, db_session logger = logging.getLogger(__name__) +def record_tool_span_start(call: Any, **attributes: Any): + """Open an ``execute_tool`` span for ``call`` (no-op without an active trace).""" + llm_name = getattr(call, "name", None) or "unknown" + base = { + "gen_ai.operation.name": "execute_tool", + "gen_ai.tool.name": llm_name, + "gen_ai.tool.call.id": getattr(call, "id", None), + "gen_ai.tool.type": "function", + } + base.update(attributes) + return tracing.start_span( + tracing.KIND_TOOL, + f"execute_tool {llm_name}", + attributes={k: v for k, v in base.items() if v is not None}, + ) + + +#: ``tool_calls`` statuses mapped to trace span statuses. +_TOOL_SPAN_STATUS = { + "error": tracing.STATUS_ERROR, + "denied": tracing.STATUS_DENIED, + "skipped": tracing.STATUS_SKIPPED, + "awaiting_approval": tracing.STATUS_PENDING, + "requires_client_execution": tracing.STATUS_PENDING, +} + + +def finish_tool_span(span: Any, data: Dict[str, Any]) -> None: + """Close a tool span from the ``tool_calls`` entry the call produced. + + Args: + span: The span from :func:`record_tool_span_start`. + data: The tool-call record (``tool_name``, ``arguments``, ``result``, + ``status``, ...); results in it are already guardrail-scanned. + """ + if not span: + return + status = str(data.get("status") or "completed") + span.set( + **{ + "docsgpt.tool": data.get("tool_name"), + "docsgpt.action": data.get("action_name"), + "docsgpt.tool_status": status, + "docsgpt.artifact_id": data.get("artifact_id"), + } + ) + if data.get("arguments") is not None: + span.preview("arguments", data.get("arguments")) + if data.get("result") is not None: + span.preview("result", data.get("result")) + span_status = _TOOL_SPAN_STATUS.get(status, tracing.STATUS_OK) + if span_status == tracing.STATUS_ERROR: + span.error = str(data.get("error") or data.get("result") or "tool error")[:500] + span.set(**{"error.type": str(data.get("error_type") or "ToolError")}) + elif data.get("error"): + span.error = str(data["error"])[:500] + span.end(span_status) + + def _is_foreign_key_violation(exc: BaseException) -> bool: """Whether ``exc`` is a Postgres FK violation (SQLSTATE 23503).""" if not isinstance(exc, IntegrityError): @@ -950,7 +1010,27 @@ class ToolExecutor: MAX_ADVERTISED_TOOL_NAMES = 30 def execute(self, tools_dict: Dict, call, llm_class_name: str): - """Execute a tool call. Yields status events, returns (result, call_id).""" + """Execute a tool call. Yields status events, returns (result, call_id). + + Every call is recorded as an ``execute_tool`` trace span; its outcome + is read from the ``tool_calls`` entry each branch of :meth:`_execute` + appends, so the span matches what the conversation stores. + """ + span = record_tool_span_start(call) + recorded = len(self.tool_calls) + try: + outcome = yield from self._execute(tools_dict, call, llm_class_name) + except Exception as exc: + span.end(error=exc) + raise + except GeneratorExit: + span.end(tracing.STATUS_CANCELLED) + raise + data = self.tool_calls[-1] if len(self.tool_calls) > recorded else {} + finish_tool_span(span, data) + return outcome + + def _execute(self, tools_dict: Dict, call, llm_class_name: str): parser = ToolActionParser(llm_class_name, name_mapping=self._name_to_tool) tool_id, action_name, call_args = parser.parse_args(call) llm_name = getattr(call, "name", "unknown") diff --git a/docsgpt/llm/handlers/base.py b/docsgpt/llm/handlers/base.py index 039d7007..c52f1823 100644 --- a/docsgpt/llm/handlers/base.py +++ b/docsgpt/llm/handlers/base.py @@ -97,6 +97,17 @@ class LLMResponse: return bool(self.tool_calls) and self.finish_reason == "tool_calls" +def _trace_unexecuted_tool(call: Any, data: Dict[str, Any]) -> None: + """Record a tool call that was paused, denied or skipped instead of run. + + Those calls never reach ``ToolExecutor.execute`` (which traces executed + calls), so without this the trace would not show them at all. + """ + from docsgpt.agents.tool_executor import finish_tool_span, record_tool_span_start + + finish_tool_span(record_tool_span_start(call), data) + + def _is_restated_payload(existing: str, incoming: str) -> bool: """True when ``incoming`` restates an already-complete ``existing`` payload. @@ -1127,6 +1138,10 @@ class LLMHandler(ABC): # Mark remaining tools as skipped for remaining_call in tool_calls[i:]: + _trace_unexecuted_tool( + remaining_call, + {"tool_name": "system", "action_name": remaining_call.name, "status": "skipped"}, + ) skip_message = { "type": "tool_call", "data": { @@ -1207,22 +1222,21 @@ class LLMHandler(ABC): message_id=agent.tool_executor.message_id, user_id=agent.tool_executor.user, ) - yield { - "type": "tool_call", - "data": { - "tool_name": pause_info["tool_name"], - "call_id": pause_info["call_id"], - "action_name": pause_info.get( - "llm_name", pause_info["name"] - ), - "arguments": pause_info["arguments"], - "status": "denied", - "error": deny_reason, - "error_type": pause_info.get( - "error_type", "tool_not_allowed" - ), - }, + denied_data = { + "tool_name": pause_info["tool_name"], + "call_id": pause_info["call_id"], + "action_name": pause_info.get( + "llm_name", pause_info["name"] + ), + "arguments": pause_info["arguments"], + "status": "denied", + "error": deny_reason, + "error_type": pause_info.get( + "error_type", "tool_not_allowed" + ), } + _trace_unexecuted_tool(call, denied_data) + yield {"type": "tool_call", "data": denied_data} continue # Yield pause event so the client knows this tool is waiting pause_data = { @@ -1236,6 +1250,7 @@ class LLMHandler(ABC): # can wire the sticky "don't ask again" button. if pause_info.get("device_id"): pause_data["device_id"] = pause_info["device_id"] + _trace_unexecuted_tool(call, pause_data) yield {"type": "tool_call", "data": pause_data} pending_actions.append(pause_info) # Do NOT add messages for pending tools here. diff --git a/docsgpt/logging.py b/docsgpt/logging.py index c05265e8..9689b1b2 100644 --- a/docsgpt/logging.py +++ b/docsgpt/logging.py @@ -7,6 +7,7 @@ import logging import uuid from typing import Any, Callable, Dict, Generator, List, Optional +from docsgpt import tracing from docsgpt.core import log_context from docsgpt.storage.db.repositories.stack_logs import StackLogsRepository from docsgpt.storage.db.session import db_session @@ -136,10 +137,17 @@ 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) + error: BaseException | None = None + completed = False try: generator = func(*args, **kwargs) yield from _consume_and_log(generator, context) + completed = True except Exception as exc: # Only ``Exception`` counts as an activity error; ``GeneratorExit`` # (consumer disconnected mid-stream) and ``KeyboardInterrupt`` @@ -154,6 +162,7 @@ def log_activity() -> Callable: started_at=started_at, error=error, ) + _end_agent_span(span, context, error=error, completed=completed) log_context.reset(ctx_token) return wrapper @@ -161,6 +170,45 @@ 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).""" + label = type(agent).__name__ + return 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(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, + }.items() + if v is not None + }, + ) + + +def _end_agent_span(span: Any, context: "LogContext", *, error: BaseException | None, completed: bool) -> None: + """Close the agent span with the run's response aggregates and outcome.""" + if not span: + return + attributes = { + "docsgpt.answer_chars": context.answer_length, + "docsgpt.source_count": context.source_count, + "docsgpt.tool_call_count": context.tool_call_count, + } + if error is not None: + span.end(error=error, attributes=attributes) + elif context.stream_error: + span.set(**attributes, **{"error.type": "StreamError"}) + span.error = context.stream_error + span.end(tracing.STATUS_ERROR) + else: + span.end(None if completed else tracing.STATUS_CANCELLED, attributes=attributes) + + def _emit_activity_finished( *, context: "LogContext", diff --git a/tests/test_logging.py b/tests/test_logging.py index 2a18573a..784b9234 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -398,3 +398,95 @@ class TestAccumulateResponseSummary: assert ctx.tool_call_count == 2 + + +class TestLogActivityTraceSpan: + """``@log_activity`` opens the ``invoke_agent`` span for every agent run.""" + + @pytest.fixture(autouse=True) + def _tracing_on(self, monkeypatch): + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + + class _Agent: + endpoint = "stream" + user = "user1" + user_api_key = "" + agent_id = "agent-1" + model_id = "gpt-4o" + + def test_agent_span_wraps_the_run_and_binds_activity_id(self): + from docsgpt import tracing + from docsgpt.logging import log_activity + + @log_activity() + def gen(agent, log_context=None): + with tracing.span(tracing.KIND_LLM, "chat m"): + pass + yield {"answer": "hi"} + yield {"sources": [{"title": "a"}]} + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with patch("docsgpt.logging._log_activity_to_db"), tracing.activate(trace): + list(gen(self._Agent())) + agent_span, llm_span = trace.spans + assert agent_span.kind == tracing.KIND_AGENT + assert agent_span.name == "invoke_agent _Agent" + assert agent_span.status == "ok" + 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 + assert agent_span.attributes["docsgpt.source_count"] == 1 + assert llm_span.parent_id == agent_span.id + assert trace.activity_id is not None + + def test_nested_agent_keeps_outer_activity_id(self): + from docsgpt import tracing + from docsgpt.logging import log_activity + + @log_activity() + def inner(agent, log_context=None): + yield "x" + + @log_activity() + def outer(agent, log_context=None): + yield log_context.activity_id + yield from inner(agent) + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with patch("docsgpt.logging._log_activity_to_db"), tracing.activate(trace): + outer_activity = list(outer(self._Agent()))[0] + assert trace.activity_id == outer_activity + outer_span, inner_span = trace.spans + assert inner_span.parent_id == outer_span.id + + def test_yielded_error_marks_span_error(self): + from docsgpt import tracing + from docsgpt.logging import log_activity + + @log_activity() + def gen(agent, log_context=None): + yield {"type": "error", "error": "node failed"} + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with patch("docsgpt.logging._log_activity_to_db"), tracing.activate(trace): + list(gen(self._Agent())) + assert trace.spans[0].status == "error" + assert trace.spans[0].error == "node failed" + + def test_raised_error_marks_span_error(self): + from docsgpt import tracing + from docsgpt.logging import log_activity + + @log_activity() + def gen(agent, log_context=None): + yield "a" + raise RuntimeError("boom") + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with patch("docsgpt.logging._log_activity_to_db"), tracing.activate(trace), \ + pytest.raises(RuntimeError): + list(gen(self._Agent())) + assert trace.spans[0].status == "error" + assert trace.spans[0].attributes["error.type"] == "RuntimeError" diff --git a/tests/tracing/test_tool_spans.py b/tests/tracing/test_tool_spans.py new file mode 100644 index 00000000..6f2447db --- /dev/null +++ b/tests/tracing/test_tool_spans.py @@ -0,0 +1,244 @@ +"""Tool calls become ``execute_tool`` spans: executed, paused, denied, skipped.""" + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest + +from docsgpt import tracing +from docsgpt.agents.tool_executor import ToolExecutor +from docsgpt.core.settings import settings +from docsgpt.llm.handlers.base import LLMHandler, ToolCall + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + + +@pytest.fixture() +def trace(): + t = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(t): + yield t + + +def _executor_with(inner): + executor = ToolExecutor.__new__(ToolExecutor) + executor.tool_calls = [] + executor._execute = lambda tools_dict, call, llm_class_name: inner(executor, call) + return executor + + +def _drain(gen): + events = [] + while True: + try: + events.append(next(gen)) + except StopIteration as stop: + return events, stop.value + + +class TestExecute: + def test_completed_call(self, trace): + def inner(executor, call): + yield {"type": "tool_call", "data": {"status": "pending"}} + executor.tool_calls.append( + { + "tool_name": "brave", + "action_name": "search", + "arguments": {"q": "x", "api_key": "k"}, + "result": "3 results", + "status": "completed", + } + ) + return "3 results", call.id + + executor = _executor_with(inner) + call = ToolCall(id="c1", name="search_1", arguments="{}") + events, outcome = _drain(executor.execute({}, call, "OpenAILLM")) + assert outcome == ("3 results", "c1") + assert len(events) == 1 + (span,) = trace.spans + assert span.kind == tracing.KIND_TOOL + assert span.name == "execute_tool search_1" + assert span.status == "ok" + assert span.attributes["gen_ai.tool.call.id"] == "c1" + assert span.attributes["docsgpt.tool"] == "brave" + assert span.previews["arguments"]["api_key"] == "[REDACTED]" + assert span.previews["result"] == "3 results" + + def test_in_band_error(self, trace): + def inner(executor, call): + executor.tool_calls.append( + {"tool_name": "unknown", "result": "no such tool", "status": "error"} + ) + return "no such tool", call.id + yield # pragma: no cover + + executor = _executor_with(inner) + _drain(executor.execute({}, ToolCall(id="c", name="x", arguments="{}"), "L")) + span = trace.spans[0] + assert span.status == "error" + assert span.error == "no such tool" + + def test_raised_error(self, trace): + def inner(executor, call): + yield {"type": "tool_call", "data": {"status": "pending"}} + raise RuntimeError("api down") + + executor = _executor_with(inner) + with pytest.raises(RuntimeError): + _drain(executor.execute({}, ToolCall(id="c", name="x", arguments="{}"), "L")) + assert trace.spans[0].status == "error" + assert trace.spans[0].attributes["error.type"] == "RuntimeError" + + def test_nested_retrieval_parents_to_tool(self, trace): + def inner(executor, call): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + executor.tool_calls.append({"status": "completed"}) + return "", call.id + yield # pragma: no cover + + executor = _executor_with(inner) + _drain(executor.execute({}, ToolCall(id="c", name="internal_search", arguments="{}"), "L")) + tool, retrieval = trace.spans + assert retrieval.parent_id == tool.id + + +class _Handler(LLMHandler): + def parse_response(self, response): # pragma: no cover - unused + raise NotImplementedError + + def create_tool_message(self, tool_call, result): + return {"role": "tool", "content": str(result)} + + def _iterate_stream(self, response): # pragma: no cover - unused + return iter(()) + + +class TestPausedCalls: + def _agent(self, pause): + agent = Mock() + agent._check_context_limit = Mock(return_value=False) + agent.context_limit_reached = False + agent.llm.__class__.__name__ = "MockLLM" + agent.tool_executor.check_pause = Mock(return_value=pause) + return agent + + def test_awaiting_approval_is_pending(self, trace): + agent = self._agent( + { + "call_id": "c1", + "name": "send_0", + "tool_name": "telegram", + "tool_id": "0", + "action_name": "send", + "arguments": {"text": "hi"}, + "pause_type": "awaiting_approval", + } + ) + call = ToolCall(id="c1", name="send_0", arguments='{"text": "hi"}') + _drain(_Handler().handle_tool_calls(agent, [call], {"0": {"name": "telegram"}}, [])) + (span,) = trace.spans + assert span.status == "pending" + assert span.attributes["docsgpt.tool_status"] == "awaiting_approval" + + def test_headless_denial_is_denied(self, trace): + agent = self._agent( + { + "call_id": "c1", + "name": "send_0", + "tool_name": "telegram", + "tool_id": "0", + "action_name": "send", + "arguments": {}, + "pause_type": "headless_denied", + "deny_reason": "not allowed", + } + ) + agent.tool_executor.message_id = None + agent.tool_executor.user = "u" + agent.tool_executor.agent_id = None + call = ToolCall(id="c1", name="send_0", arguments="{}") + with pytest.MonkeyPatch.context() as mp: + mp.setattr("docsgpt.agents.tool_executor._record_proposed", lambda *a, **k: False) + _drain(_Handler().handle_tool_calls(agent, [call], {"0": {"name": "telegram"}}, [])) + (span,) = trace.spans + assert span.status == "denied" + assert span.error == "not allowed" + + +class TestContinuation: + def _agent(self): + from docsgpt.agents.classic_agent import ClassicAgent + + llm = Mock() + llm._supports_tools = True + llm.gen_stream = Mock(return_value=iter(["Answer"])) + llm._supports_structured_output = Mock(return_value=False) + handler = Mock() + handler.process_message_flow = Mock(return_value=iter([])) + handler.create_tool_message = Mock(return_value={"role": "tool", "content": "x"}) + executor = Mock() + executor.tool_calls = [] + executor.prepare_tools_for_llm = Mock(return_value=[]) + executor.get_truncated_tool_calls = Mock(return_value=[]) + return ClassicAgent( + endpoint="stream", + llm_name="openai", + model_id="gpt-4", + api_key="test", + llm=llm, + llm_handler=handler, + tool_executor=executor, + ) + + def _pending(self): + return [ + { + "call_id": "c1", + "name": "danger_0", + "tool_name": "danger", + "tool_id": "0", + "action_name": "danger", + "arguments": {}, + "pause_type": "awaiting_approval", + } + ] + + def test_continuation_opens_agent_span_with_denied_tool(self, trace): + agent = self._agent() + list( + agent.gen_continuation( + [{"role": "system", "content": "s"}], + {"0": {"name": "danger"}}, + self._pending(), + [{"call_id": "c1", "decision": "denied", "comment": "too risky"}], + ) + ) + agent_span = trace.spans[0] + tool_span = trace.spans[1] + assert agent_span.kind == tracing.KIND_AGENT + assert agent_span.attributes["docsgpt.continuation"] is True + assert agent_span.status == "ok" + assert tool_span.parent_id == agent_span.id + assert tool_span.status == "denied" + assert tool_span.error == "too risky" + + def test_client_result_is_recorded(self, trace): + agent = self._agent() + list( + agent.gen_continuation( + [{"role": "system", "content": "s"}], + {"0": {"name": "danger"}}, + self._pending(), + [{"call_id": "c1", "result": {"ok": True}}], + ) + ) + tool_span = trace.spans[1] + assert tool_span.status == "ok" + assert tool_span.attributes["docsgpt.client_executed"] is True From f9605e7794741f1729b6d61abda644ca7a2564c2 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:27:43 +0100 Subject: [PATCH 07/33] Trace RAG retrieval, query embeddings and prescreen The dispatcher and each retriever's search open retrieval spans recording sources, chunk counts, top scores and a chunk preview. The shared query embedding, every per-source vector search and the prescreen rerank get their own spans; pool workers carry the trace so they nest correctly. --- docsgpt/retriever/classic_rag.py | 45 ++++++-- docsgpt/retriever/dispatcher.py | 12 ++ docsgpt/retriever/fanout.py | 21 +++- docsgpt/retriever/graph_rag.py | 34 ++++-- docsgpt/retriever/stages/prescreen.py | 37 ++++--- docsgpt/tracing/core.py | 2 +- docsgpt/tracing/retrieval.py | 94 ++++++++++++++++ tests/tracing/test_retrieval_spans.py | 153 ++++++++++++++++++++++++++ 8 files changed, 362 insertions(+), 36 deletions(-) create mode 100644 docsgpt/tracing/retrieval.py create mode 100644 tests/tracing/test_retrieval_spans.py diff --git a/docsgpt/retriever/classic_rag.py b/docsgpt/retriever/classic_rag.py index db5e09b9..bcf49038 100644 --- a/docsgpt/retriever/classic_rag.py +++ b/docsgpt/retriever/classic_rag.py @@ -2,6 +2,11 @@ import logging from typing import Any, Dict, List, Optional, Tuple from docsgpt.core.settings import settings +from docsgpt.tracing.retrieval import ( + describe_documents, + start_retrieval_span, + start_source_search_span, +) from docsgpt.llm.llm_creator import LLMCreator from docsgpt.retriever.base import BaseRetriever from docsgpt.retriever.fanout import fetch_per_source, max_parallel_sources @@ -266,6 +271,15 @@ class ClassicRAG(BaseRetriever): logged and reported as ``None`` so one bad source cannot take the rest of the retrieval down with it. """ + span = start_source_search_span( + plan["id"], + top_k=plan["src_k"], + **{ + "docsgpt.vector_store": settings.VECTOR_STORE, + "docsgpt.retriever": type(self).__name__, + "docsgpt.shared_query_vector": query_vector is not None, + }, + ) try: if docsearch is None: docsearch = VectorCreator.create_vectorstore( @@ -279,8 +293,14 @@ class ClassicRAG(BaseRetriever): query_vector=query_vector, ) score_kind = self._score_kind(docsearch) if self.include_scores else None + try: + candidates = len(docs_temp) + except TypeError: + candidates = None + span.end(attributes={"docsgpt.candidate_count": candidates}) return docs_temp, score_kind except Exception as e: + span.end(error=e) logger.error( f"Error searching vectorstore {plan['id']}: {e}", exc_info=True ) @@ -409,11 +429,20 @@ class ClassicRAG(BaseRetriever): def search(self, query: str = ""): """Search for documents using optional query override""" - if query: - self.original_question = query - # Invalidate the cached rephrase so a per-source path that opts in - # rephrases against the new query, not a stale one. - self._rephrased_question = None - self.question = self._rephrase_query() - self._rephrased_question = self.question - return self._get_data() + with start_retrieval_span( + f"retrieval {type(self).__name__}", + sources=self.vectorstores, + **{"docsgpt.retriever": type(self).__name__, "docsgpt.top_k": self.chunks}, + ) as span: + if query: + self.original_question = query + # Invalidate the cached rephrase so a per-source path that opts in + # rephrases against the new query, not a stale one. + self._rephrased_question = None + self.question = self._rephrase_query() + self._rephrased_question = self.question + docs = self._get_data() + rephrased = self._rephrased_question + span.set(**{"docsgpt.rephrased": bool(rephrased and rephrased != self.original_question)}) + describe_documents(span, docs, query=rephrased or self.original_question) + return docs diff --git a/docsgpt/retriever/dispatcher.py b/docsgpt/retriever/dispatcher.py index 31cc75e0..7ec23c5c 100644 --- a/docsgpt/retriever/dispatcher.py +++ b/docsgpt/retriever/dispatcher.py @@ -23,6 +23,7 @@ from docsgpt.retriever.stages.prescreen import ( max_candidate_k, ) from docsgpt.storage.db.source_config import RetrievalConfig +from docsgpt.tracing.retrieval import describe_documents, start_retrieval_span from docsgpt.utils import num_tokens_from_string logger = logging.getLogger(__name__) @@ -294,6 +295,17 @@ class Dispatcher(BaseRetriever): def search(self, query: str = "") -> List[Dict[str, Any]]: """Run every group under the shared budget and merge the results.""" + sources = [s for group in self._groups for s in group.get("doc_ids", [])] + with start_retrieval_span( + "retrieval", + sources=sources, + **{"docsgpt.retriever": "Dispatcher", "docsgpt.group_count": len(self._groups)}, + ) as span: + docs = self._search_groups(query) + describe_documents(span, docs) + return docs + + def _search_groups(self, query: str) -> List[Dict[str, Any]]: groups = self._groups n_groups = len(groups) diff --git a/docsgpt/retriever/fanout.py b/docsgpt/retriever/fanout.py index cb1ad3a6..fcca623d 100644 --- a/docsgpt/retriever/fanout.py +++ b/docsgpt/retriever/fanout.py @@ -12,7 +12,9 @@ import logging from concurrent.futures import ThreadPoolExecutor from typing import Any, Callable, Dict, Iterable, List, Optional, TypeVar +from docsgpt import tracing from docsgpt.core.settings import settings +from docsgpt.tracing.retrieval import start_embedding_span logger = logging.getLogger(__name__) @@ -85,16 +87,22 @@ def embed_questions(docsearch, questions: Iterable[str]) -> Dict[str, List[float embedder = store_embeddings(docsearch) if embedder is None: return {} + distinct = list(dict.fromkeys(questions)) + span = start_embedding_span( + settings.EMBEDDINGS_NAME, + inputs=len(distinct), + **{"docsgpt.embeddings_backend": type(embedder).__name__}, + ) try: - return { - question: embedder.embed_query(question) - for question in dict.fromkeys(questions) - } + vectors = {question: embedder.embed_query(question) for question in distinct} except Exception as e: + span.end(error=e) logger.warning( "Query embedding failed (%s); each store will embed its own.", e ) return {} + span.end() + return vectors def run_source_jobs( @@ -121,10 +129,13 @@ def run_source_jobs( workers = max_parallel_sources(len(jobs)) if workers == 1: return [fn(job) for job in jobs] + # Pool threads don't inherit context: carry the trace and its current + # retrieval span in so per-source spans nest under it. + traced_fn = tracing.wrap(fn) with ThreadPoolExecutor( max_workers=workers, thread_name_prefix="rag-source" ) as pool: - return list(pool.map(fn, jobs)) + return list(pool.map(traced_fn, jobs)) def fetch_per_source( diff --git a/docsgpt/retriever/graph_rag.py b/docsgpt/retriever/graph_rag.py index 8d191f7b..59ee6163 100644 --- a/docsgpt/retriever/graph_rag.py +++ b/docsgpt/retriever/graph_rag.py @@ -33,6 +33,11 @@ from docsgpt.retriever.base import BaseRetriever from docsgpt.retriever.classic_rag import ClassicRAG from docsgpt.retriever.labels import labels_from_metadata from docsgpt.storage.db.source_config import GraphRetrievalConfig +from docsgpt.tracing.retrieval import ( + describe_documents, + start_embedding_span, + start_retrieval_span, +) from docsgpt.utils import num_tokens_from_string from docsgpt.vectorstore.base import get_embeddings @@ -231,7 +236,11 @@ class GraphRAGRetriever(BaseRetriever): def _embed_query(self, question: str) -> List[float]: embedding = get_embeddings() - return embedding.embed_query(question) + with start_embedding_span( + settings.EMBEDDINGS_NAME, + **{"docsgpt.embeddings_backend": type(embedding).__name__}, + ): + return embedding.embed_query(question) def _ppr_scores(self, subgraph, seeds) -> Dict[str, float]: """Run Personalized PageRank, then down-weight hub nodes by IDF. @@ -730,10 +739,19 @@ class GraphRAGRetriever(BaseRetriever): logging.debug("Error closing GraphRAG store: %s", e) def search(self, query: str = "") -> List[Dict[str, Any]]: - if query: - self.original_question = query - self._classic.original_question = query - self._classic._rephrased_question = None - self._classic.question = self._classic._rephrase_query() - self._classic._rephrased_question = self._classic.question - return self._get_data() + with start_retrieval_span( + f"retrieval {type(self).__name__}", + sources=self.vectorstores, + **{"docsgpt.retriever": type(self).__name__, "docsgpt.top_k": self.chunks}, + ) as span: + if query: + self.original_question = query + self._classic.original_question = query + self._classic._rephrased_question = None + self._classic.question = self._classic._rephrase_query() + self._classic._rephrased_question = self._classic.question + docs = self._get_data() + describe_documents( + span, docs, query=self._classic._rephrased_question or self.original_question + ) + return docs diff --git a/docsgpt/retriever/stages/prescreen.py b/docsgpt/retriever/stages/prescreen.py index cae50dea..2dfa5140 100644 --- a/docsgpt/retriever/stages/prescreen.py +++ b/docsgpt/retriever/stages/prescreen.py @@ -18,6 +18,7 @@ import re from concurrent.futures import ThreadPoolExecutor from typing import Any, Callable, Dict, List, Optional +from docsgpt import tracing from docsgpt.llm.llm_creator import LLMCreator from docsgpt.storage.db.source_config import PreScreenConfig @@ -164,20 +165,28 @@ class PreScreenStage: ] llm = self._build_llm() - max_workers = min(_MAX_WORKERS, len(batches)) - kept: List[Dict[str, Any]] = [] - if max_workers <= 1: - for batch in batches: - kept.extend(self._screen_batch(llm, query, batch)) - else: - with ThreadPoolExecutor(max_workers=max_workers) as pool: - results = pool.map( - lambda b: self._screen_batch(llm, query, b), batches - ) - for batch_result in results: - kept.extend(batch_result) - - return kept[: self.config.max_keep] + with tracing.span( + tracing.KIND_RERANK, + "rerank prescreen", + attributes={ + "gen_ai.operation.name": "rerank", + "docsgpt.candidate_count": len(docs), + "docsgpt.batch_count": len(batches), + }, + ) as span: + max_workers = min(_MAX_WORKERS, len(batches)) + kept: List[Dict[str, Any]] = [] + if max_workers <= 1: + for batch in batches: + kept.extend(self._screen_batch(llm, query, batch)) + else: + screen = tracing.wrap(lambda b: self._screen_batch(llm, query, b)) + with ThreadPoolExecutor(max_workers=max_workers) as pool: + for batch_result in pool.map(screen, batches): + kept.extend(batch_result) + survivors = kept[: self.config.max_keep] + span.set(**{"docsgpt.kept_count": len(survivors)}) + return survivors def build_prescreen_stages( diff --git a/docsgpt/tracing/core.py b/docsgpt/tracing/core.py index b0289152..46e9c4a7 100644 --- a/docsgpt/tracing/core.py +++ b/docsgpt/tracing/core.py @@ -40,7 +40,7 @@ KIND_GUARDRAIL = "guardrail" KIND_STEP = "step" #: Kinds that become the implicit parent of spans started while they are open. -CONTAINER_KINDS = frozenset({KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_STEP}) +CONTAINER_KINDS = frozenset({KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_RERANK, KIND_STEP}) STATUS_OK = "ok" STATUS_ERROR = "error" diff --git a/docsgpt/tracing/retrieval.py b/docsgpt/tracing/retrieval.py new file mode 100644 index 00000000..f9ce7103 --- /dev/null +++ b/docsgpt/tracing/retrieval.py @@ -0,0 +1,94 @@ +"""Spans for RAG retrieval: the search itself, the query embedding, each source.""" + +from __future__ import annotations + +from typing import Any, Dict, Iterable, List, Optional + +from docsgpt.tracing import core + +#: Characters of each retrieved chunk kept in the stored preview. +_SNIPPET_CHARS = 240 +#: Chunks listed in the preview; the count attribute still reports all of them. +_MAX_PREVIEW_CHUNKS = 20 + + +def start_retrieval_span(name: str, *, sources: Optional[Iterable[Any]] = None, **attributes: 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 = { + "gen_ai.operation.name": "retrieval", + "docsgpt.source_ids": source_ids or None, + "gen_ai.data_source.id": source_ids[0] if len(source_ids) == 1 else None, + } + base.update(attributes) + return core.start_span( + core.KIND_RETRIEVAL, + name, + attributes={k: v for k, v in base.items() if v is not None}, + ) + + +def describe_documents(span: Any, docs: Optional[List[Dict[str, Any]]], *, query: Optional[str] = None) -> None: + """Record what a retrieval returned: counts, sources, top scores and a chunk preview.""" + if not span: + return + docs = docs or [] + scores = [d.get("score") for d in docs if isinstance(d, dict) and isinstance(d.get("score"), (int, float))] + score_kinds = {d.get("score_kind") for d in docs if isinstance(d, dict) and d.get("score_kind")} + span.set( + **{ + "docsgpt.chunk_count": len(docs), + "docsgpt.top_score": max(scores) if scores else None, + "docsgpt.score_kind": score_kinds.pop() if len(score_kinds) == 1 else None, + } + ) + if query: + span.preview("query", query) + if docs: + span.preview( + "chunks", + [ + { + k: v + for k, v in { + "title": d.get("title") or d.get("filename"), + "source": d.get("source"), + "score": d.get("score"), + "text": str(d.get("text") or "")[:_SNIPPET_CHARS], + }.items() + if v not in (None, "") + } + for d in docs[:_MAX_PREVIEW_CHUNKS] + if isinstance(d, dict) + ], + ) + + +def start_embedding_span(model: Optional[str], *, inputs: int = 1, **attributes: Any): + """Open an ``embeddings`` span for a query embedding (leaf).""" + base = { + "gen_ai.operation.name": "embeddings", + "gen_ai.request.model": model, + "docsgpt.input_count": inputs, + } + base.update(attributes) + return core.start_span( + core.KIND_EMBEDDING, + f"embeddings {model}" if model else "embeddings", + attributes={k: v for k, v in base.items() if v is not None}, + ) + + +def start_source_search_span(source_id: Any, *, top_k: Optional[int] = None, **attributes: Any): + """Open a per-source vector ``search`` span (leaf).""" + base = { + "gen_ai.operation.name": "search", + "gen_ai.data_source.id": str(source_id) if source_id else None, + "docsgpt.top_k": top_k, + } + base.update(attributes) + return core.start_span( + core.KIND_SEARCH, + f"search {source_id}" if source_id else "search", + attributes={k: v for k, v in base.items() if v is not None}, + ) diff --git a/tests/tracing/test_retrieval_spans.py b/tests/tracing/test_retrieval_spans.py new file mode 100644 index 00000000..be45d98c --- /dev/null +++ b/tests/tracing/test_retrieval_spans.py @@ -0,0 +1,153 @@ +"""RAG retrieval becomes retrieval / embeddings / search / rerank spans.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings +from docsgpt.retriever import fanout + + +@pytest.fixture(autouse=True) +def _env(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + + +@pytest.fixture() +def trace(): + t = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(t): + yield t + + +class TestFanout: + def test_embedding_span(self, trace): + embedder = MagicMock() + embedder.embed_query.side_effect = lambda q: [0.1, 0.2] + store = MagicMock(_embedding=embedder) + vectors = fanout.embed_questions(store, ["q", "q", "r"]) + assert set(vectors) == {"q", "r"} + (span,) = trace.spans + assert span.kind == tracing.KIND_EMBEDDING + assert span.attributes["gen_ai.operation.name"] == "embeddings" + assert span.attributes["docsgpt.input_count"] == 2 + + def test_embedding_failure_is_recorded(self, trace): + embedder = MagicMock() + embedder.embed_query.side_effect = RuntimeError("model down") + assert fanout.embed_questions(MagicMock(_embedding=embedder), ["q"]) == {} + assert trace.spans[0].status == "error" + + def test_pool_jobs_nest_under_the_open_retrieval(self, trace, monkeypatch): + monkeypatch.setattr(settings, "RETRIEVAL_MAX_PARALLEL_SOURCES", 4) + + def job(i): + with tracing.span(tracing.KIND_SEARCH, f"search {i}"): + return i + + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval") as outer: + assert fanout.run_source_jobs(job, [1, 2, 3]) == [1, 2, 3] + searches = [s for s in trace.spans if s.kind == tracing.KIND_SEARCH] + assert len(searches) == 3 + assert all(s.parent_id == outer.id for s in searches) + + +class TestClassicRAG: + def _rag(self, stores): + from docsgpt.retriever.classic_rag import ClassicRAG + + with patch("docsgpt.retriever.classic_rag.LLMCreator.create_llm", return_value=MagicMock()): + rag = ClassicRAG( + source={"question": "what is x", "active_docs": list(stores)}, + chunks=4, + decoded_token={"sub": "u1"}, + ) + return rag + + def test_search_records_retrieval_tree(self, trace, monkeypatch): + monkeypatch.setattr(settings, "RETRIEVAL_MAX_PARALLEL_SOURCES", 4) + embedder = MagicMock() + embedder.embed_query.return_value = [0.1] + + def make_store(*_args, **_kwargs): + store = MagicMock(_embedding=embedder) + store.search.return_value = [ + {"text": "chunk body", "metadata": {"title": "Doc", "source": "doc.md"}} + ] + return store + + rag = self._rag(["s1", "s2"]) + with patch( + "docsgpt.retriever.classic_rag.VectorCreator.create_vectorstore", + side_effect=make_store, + ): + docs = rag.search() + assert len(docs) == 2 + retrieval = trace.spans[0] + assert retrieval.kind == tracing.KIND_RETRIEVAL + assert retrieval.attributes["docsgpt.source_ids"] == ["s1", "s2"] + assert retrieval.attributes["docsgpt.chunk_count"] == 2 + assert retrieval.previews["query"] == "what is x" + assert retrieval.previews["chunks"][0]["text"] == "chunk body" + children = [s for s in trace.spans[1:]] + assert {s.kind for s in children} == {tracing.KIND_EMBEDDING, tracing.KIND_SEARCH} + assert all(s.parent_id == retrieval.id for s in children) + searches = [s for s in children if s.kind == tracing.KIND_SEARCH] + assert sorted(s.attributes["gen_ai.data_source.id"] for s in searches) == ["s1", "s2"] + assert all(s.attributes["docsgpt.candidate_count"] == 1 for s in searches) + + def test_failed_source_search_is_an_error_span(self, trace, monkeypatch): + monkeypatch.setattr(settings, "RETRIEVAL_MAX_PARALLEL_SOURCES", 1) + rag = self._rag(["s1", "s2"]) + + def make_store(_kind, source_id, _key): + if source_id == "s2": + raise RuntimeError("no store") + store = MagicMock(_embedding=None, _embeddings=None, embeddings=None) + store._get_embeddings.side_effect = RuntimeError("no embedder") + store.search.return_value = [] + return store + + with patch( + "docsgpt.retriever.classic_rag.VectorCreator.create_vectorstore", + side_effect=make_store, + ): + assert rag.search() == [] + search = {s.attributes["gen_ai.data_source.id"]: s for s in trace.spans if s.kind == tracing.KIND_SEARCH} + assert search["s1"].status == "ok" + assert search["s2"].status == "error" + assert trace.spans[0].status == "ok" + + +class TestPrescreen: + def test_rerank_span_contains_screening_calls(self, trace): + from docsgpt.retriever.stages.prescreen import PreScreenStage + from docsgpt.storage.db.source_config import PreScreenConfig + + stage = PreScreenStage( + PreScreenConfig(candidate_k=10, max_keep=2, batch_size=1), + llm_name="openai", + api_key=None, + model_id="m", + ) + llm = MagicMock() + + def fake_gen(**_kwargs): + with tracing.span(tracing.KIND_LLM, "chat m"): + return '{"keep": [0]}' + + llm.gen.side_effect = fake_gen + stage._build_llm = lambda: llm + docs = [{"text": f"d{i}"} for i in range(3)] + kept = stage(docs, {"query": "q"}) + assert len(kept) == 2 + rerank = trace.spans[0] + assert rerank.kind == tracing.KIND_RERANK + assert rerank.attributes["docsgpt.kept_count"] == 2 + llm_spans = [s for s in trace.spans if s.kind == tracing.KIND_LLM] + assert len(llm_spans) == 3 + assert all(s.parent_id == rerank.id for s in llm_spans) From da1bc16eb6a877b027ba817e88cce5045b11caf0 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:29:21 +0100 Subject: [PATCH 08/33] Trace guardrails, workflow steps and research phases Guardrail evaluations are traced when they call a remote check or fire; a firing guardrail drops every content preview from the stored trace. Each workflow node and each research phase is a step span, and the workflow run id links the trace to its Logs row. --- docsgpt/agents/research_agent.py | 38 ++++++++++-- docsgpt/agents/workflows/workflow_engine.py | 22 ++++++- docsgpt/guardrails/engine.py | 46 ++++++++++++++- docsgpt/tracing/core.py | 6 +- tests/agents/test_research_agent.py | 50 ++++++++++++++++ tests/agents/test_workflow_engine_coverage.py | 48 +++++++++++++++ tests/guardrails/test_engine.py | 58 +++++++++++++++++++ 7 files changed, 256 insertions(+), 12 deletions(-) diff --git a/docsgpt/agents/research_agent.py b/docsgpt/agents/research_agent.py index de96cf35..2abb6cdb 100644 --- a/docsgpt/agents/research_agent.py +++ b/docsgpt/agents/research_agent.py @@ -4,6 +4,7 @@ import os import time from typing import Dict, Generator, List, Optional +from docsgpt import tracing from docsgpt.agents.base import BaseAgent from docsgpt.agents.tool_executor import ToolExecutor from docsgpt.agents.tools.graph_search import add_graph_search_tool @@ -51,6 +52,15 @@ SYNTHESIS_PROMPT = _load_prompt("synthesis.txt") # --------------------------------------------------------------------------- + +def _phase_span(phase: str, **attributes): + """Open a trace span for one research phase (clarify, plan, a step, synthesis).""" + return tracing.start_span( + tracing.KIND_STEP, + f"research {phase}", + attributes={"docsgpt.research.phase": phase.split(" ")[0], **attributes}, + ) + class CitationManager: """Tracks and deduplicates citations across research steps.""" @@ -169,7 +179,8 @@ class ResearchAgent(BaseAgent): # Phase 0: Clarification (skip if user is responding to a prior clarification) if not self._is_follow_up(): - clarification = self._clarification_phase(query) + with _phase_span("clarify"): + clarification = self._clarification_phase(query) if clarification: yield {"metadata": {"is_clarification": True}} yield {"answer": clarification} @@ -182,7 +193,11 @@ class ResearchAgent(BaseAgent): # Phase 1: Planning (with adaptive depth) yield {"type": "research_progress", "data": {"status": "planning"}} - plan, complexity = self._planning_phase(query) + with _phase_span("plan") as plan_span: + plan, complexity = self._planning_phase(query) + plan_span.set( + **{"docsgpt.research.steps": len(plan or []), "docsgpt.research.complexity": complexity} + ) if not plan: logger.warning("ResearchAgent: Planning produced no steps, falling back") @@ -222,7 +237,9 @@ class ResearchAgent(BaseAgent): }, } - report = self._research_step(step_query, tools_dict) + with _phase_span(f"step {step_num}", **{"docsgpt.research.step": step_num}) as step_span: + step_span.preview("query", step_query) + report = self._research_step(step_query, tools_dict) intermediate_reports.append({"step": step, "content": report}) yield { @@ -249,9 +266,18 @@ class ResearchAgent(BaseAgent): "tokens_used": self._tokens_used, }, } - yield from self._synthesis_phase( - query, plan, intermediate_reports, tools_dict, log_context - ) + synthesis_span = _phase_span("synthesis") + try: + yield from self._synthesis_phase( + query, plan, intermediate_reports, tools_dict, log_context + ) + except GeneratorExit: + synthesis_span.end(tracing.STATUS_CANCELLED) + raise + except Exception as exc: + synthesis_span.end(error=exc) + raise + synthesis_span.end() # Sources and tool calls self.retrieved_docs = self.citations.get_all_docs() diff --git a/docsgpt/agents/workflows/workflow_engine.py b/docsgpt/agents/workflows/workflow_engine.py index 15876c82..f296fc05 100644 --- a/docsgpt/agents/workflows/workflow_engine.py +++ b/docsgpt/agents/workflows/workflow_engine.py @@ -6,6 +6,7 @@ import uuid from datetime import datetime, timezone from typing import Any, Dict, Generator, List, Optional, TYPE_CHECKING +from docsgpt import tracing from docsgpt.agents.workflows.cel_evaluator import CelEvaluationError, evaluate_cel from docsgpt.agents.workflows.node_agent import WorkflowNodeAgentFactory from docsgpt.agents.workflows.schemas import ( @@ -114,6 +115,9 @@ class WorkflowEngine: self, initial_inputs: WorkflowState, query: str ) -> Generator[Dict[str, str], None, None]: self._initialize_state(initial_inputs, query) + # A workflow run's Logs row finds its trace by this id; a nested + # engine (a workflow inside a workflow) keeps the outermost run's. + tracing.bind_if_unset(workflow_run_id=self.workflow_run_id) # Surface the run id up front so the client can list this run's # artifacts (GET /api/artifacts?workflow_run_id=) once it has been @@ -140,6 +144,16 @@ class WorkflowEngine: break log_entry = self._create_log_entry(node) self._last_node_tool_calls = [] + step_span = tracing.start_span( + tracing.KIND_STEP, + f"workflow_step {node.title or node.type.value}", + attributes={ + "docsgpt.workflow.node_id": node.id, + "docsgpt.workflow.node_type": node.type.value, + "docsgpt.workflow.node_title": node.title or None, + "docsgpt.workflow_run_id": self.workflow_run_id, + }, + ) yield { "type": "workflow_step", @@ -150,7 +164,12 @@ class WorkflowEngine: } try: - yield from self._execute_node(node) + try: + yield from self._execute_node(node) + except GeneratorExit: + step_span.end(tracing.STATUS_CANCELLED) + raise + step_span.end() log_entry["status"] = ExecutionStatus.COMPLETED.value self._finalize_log_entry(log_entry, pre_state) @@ -167,6 +186,7 @@ class WorkflowEngine: "output": node_output, } except Exception as e: + step_span.end(error=e) logger.error(f"Error executing node {node.id}: {e}", exc_info=True) log_entry["status"] = ExecutionStatus.FAILED.value log_entry["error"] = str(e) diff --git a/docsgpt/guardrails/engine.py b/docsgpt/guardrails/engine.py index 97a359b7..be489385 100644 --- a/docsgpt/guardrails/engine.py +++ b/docsgpt/guardrails/engine.py @@ -7,6 +7,7 @@ from threading import Thread from time import monotonic from typing import List, Optional +from docsgpt import tracing from docsgpt.guardrails.base import ScanContext from docsgpt.guardrails.config import GuardrailsConfig from docsgpt.guardrails.guardrail_creator import GuardrailCreator @@ -26,6 +27,35 @@ logger = logging.getLogger(__name__) _MAX_WORKERS = 8 +def _start_guardrail_span(stage: Stage, controls) -> "tracing.Span": + return tracing.start_span( + tracing.KIND_GUARDRAIL, + f"guardrail {stage.value}", + attributes={ + "docsgpt.guardrail.stage": stage.value, + "docsgpt.guardrail.checks": [c.check for c in controls], + }, + ) + + +def _describe_decision(span, decision: StageDecision) -> None: + """Record a stage decision on its span; a firing guardrail drops trace previews.""" + triggered = [v.check for v in decision.triggered] + span.set( + **{ + "docsgpt.guardrail.triggered": triggered or None, + "docsgpt.guardrail.blocked": decision.blocked, + "docsgpt.guardrail.redacted": decision.redacted, + "docsgpt.guardrail.categories": decision.categories() or None, + "docsgpt.guardrail.unevaluated": [v.check for v in decision.unevaluated] or None, + } + ) + if triggered or decision.blocked or decision.redacted: + # The scanned text (or text near it) sits in other spans' previews: + # the retrieval query, tool results, the answer. Keep none of it. + tracing.mark_content_blocked() + + class GuardrailEngine: """Stateless evaluator bound to one agent's resolved config.""" @@ -93,14 +123,24 @@ class GuardrailEngine: return decision if any(self._needs_deadline(c) for c in controls): - decision.verdicts = self._run_concurrent(controls, text, stage) + # Remote checks (e.g. an LLM judge) cost real time, so they are + # always traced; the judge's own LLM call nests under this span. + with _start_guardrail_span(stage, controls) as span: + decision.verdicts = self._run_concurrent(controls, text, stage) + self._reduce(decision) + _describe_decision(span, decision) else: # Bounded local checks are pattern matches measured in # microseconds. Running them inline keeps the streaming hot loop # free of thread churn. decision.verdicts = [self._run_control(c, text, stage) for c in controls] + self._reduce(decision) + # The output guard evaluates every streamed segment; tracing each + # clean local scan would bury the trace, so only a firing is kept. + if not decision.clean or decision.unevaluated: + with _start_guardrail_span(stage, controls) as span: + _describe_decision(span, decision) - self._reduce(decision) self._record(decision) return decision @@ -117,7 +157,7 @@ class GuardrailEngine: for control in controls[:_MAX_WORKERS]: slot: dict = {"control": control, "verdict": None} thread = Thread( - target=self._fill_slot, + target=tracing.wrap(self._fill_slot), args=(slot, control, text, stage), daemon=True, name=f"guardrail-{control.check}", diff --git a/docsgpt/tracing/core.py b/docsgpt/tracing/core.py index 46e9c4a7..d9042f42 100644 --- a/docsgpt/tracing/core.py +++ b/docsgpt/tracing/core.py @@ -5,7 +5,7 @@ scheduled run, a search). Spans are recorded while the request runs and the whole trace is written once, when its owner calls :func:`flush`. Nesting is tracked with one span stack per thread. Only *container* spans -(agent, tool, retrieval, step) are pushed, so a leaf such as an LLM call can +(agent, tool, retrieval, rerank, guardrail, step) are pushed, so a leaf such as an LLM call can never become the parent of a sibling that starts while it is still open -- which matters because DocsGPT's agent loop is a chain of suspended generators. Ending a span pops anything left above it (marked @@ -40,7 +40,9 @@ KIND_GUARDRAIL = "guardrail" KIND_STEP = "step" #: Kinds that become the implicit parent of spans started while they are open. -CONTAINER_KINDS = frozenset({KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_RERANK, KIND_STEP}) +CONTAINER_KINDS = frozenset( + {KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_RERANK, KIND_GUARDRAIL, KIND_STEP} +) STATUS_OK = "ok" STATUS_ERROR = "error" diff --git a/tests/agents/test_research_agent.py b/tests/agents/test_research_agent.py index cd9aabd9..79b899e7 100644 --- a/tests/agents/test_research_agent.py +++ b/tests/agents/test_research_agent.py @@ -1626,3 +1626,53 @@ class TestResearchAgentClarificationCoverage: text = 'Here is my response: {"needs_clarification": false} end.' result = agent._parse_clarification_json(text) assert result == {"needs_clarification": False} + + +@pytest.mark.unit +class TestGenInnerTraceSpans: + """Each research phase is a step span in the execution trace.""" + + def test_phases_become_step_spans( + self, + agent_base_params, + mock_llm, + mock_llm_creator, + mock_llm_handler_creator, + log_context, + monkeypatch, + ): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + agent = ResearchAgent(**agent_base_params) + plan_steps = [{"query": "s1", "rationale": "r"}, {"query": "s2", "rationale": "r"}] + + def fake_step(query, tools_dict): + with tracing.span(tracing.KIND_LLM, "chat m"): + return f"report {query}" + + trace = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(trace), \ + patch.object(agent, "_setup_tools", return_value={}), \ + patch.object(agent, "_is_follow_up", return_value=False), \ + patch.object(agent, "_clarification_phase", return_value=None), \ + patch.object(agent, "_planning_phase", return_value=(plan_steps, "moderate")), \ + patch.object(agent, "_research_step", side_effect=fake_step), \ + patch.object(agent, "_synthesis_phase", return_value=iter([{"answer": "final"}])), \ + patch.object(agent, "_get_truncated_tool_calls", return_value=[]): + list(agent._gen_inner("Compare A and B", log_context)) + + names = [s.name for s in trace.spans if s.kind == tracing.KIND_STEP] + assert names == [ + "research clarify", + "research plan", + "research step 1", + "research step 2", + "research synthesis", + ] + by_name = {s.name: s for s in trace.spans} + assert by_name["research plan"].attributes["docsgpt.research.steps"] == 2 + llm_spans = [s for s in trace.spans if s.kind == tracing.KIND_LLM] + assert llm_spans[0].parent_id == by_name["research step 1"].id + assert by_name["research step 1"].previews["query"] == "s1" diff --git a/tests/agents/test_workflow_engine_coverage.py b/tests/agents/test_workflow_engine_coverage.py index 5382523c..1b8c6707 100644 --- a/tests/agents/test_workflow_engine_coverage.py +++ b/tests/agents/test_workflow_engine_coverage.py @@ -1081,3 +1081,51 @@ class TestNodeDocumentManifest: monkeypatch.setattr("docsgpt.storage.db.session.db_readonly", broken_readonly) config = AgentNodeConfig(agent_type="classic", input_documents=["*"]) assert engine._node_document_manifest(config) == "" + + +class TestTraceSpans: + """Every executed node is a ``workflow_step`` span; the run id links the trace.""" + + @pytest.fixture(autouse=True) + def _trace(self, monkeypatch): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + self.trace = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(self.trace): + yield + + @pytest.mark.unit + def test_step_spans_and_run_binding(self): + nodes = [ + _make_node("n1", NodeType.START, "Start"), + _make_node("n2", NodeType.END, "End", config={"config": {}}), + ] + engine = WorkflowEngine( + _make_graph(nodes, [_make_edge("e1", "n1", "n2")]), _make_agent() + ) + list(engine.execute({}, "hello")) + steps = [s for s in self.trace.spans if s.kind == "step"] + assert [s.name for s in steps] == ["workflow_step Start", "workflow_step End"] + assert all(s.status == "ok" for s in steps) + assert steps[0].attributes["docsgpt.workflow.node_type"] == "start" + assert self.trace.workflow_run_id == str(engine.workflow_run_id) + + @pytest.mark.unit + def test_failed_node_is_an_error_span(self): + nodes = [ + _make_node("n1", NodeType.START), + _make_node( + "n2", + NodeType.STATE, + "State", + config={"config": {"operations": [{"expression": "bad!!!", "target_variable": "x"}]}}, + ), + ] + engine = WorkflowEngine( + _make_graph(nodes, [_make_edge("e1", "n1", "n2")]), _make_agent() + ) + list(engine.execute({}, "q")) + state_span = [s for s in self.trace.spans if s.name == "workflow_step State"][0] + assert state_span.status == "error" diff --git a/tests/guardrails/test_engine.py b/tests/guardrails/test_engine.py index b3316804..cf7b704a 100644 --- a/tests/guardrails/test_engine.py +++ b/tests/guardrails/test_engine.py @@ -340,3 +340,61 @@ def test_action_enum_serialises_as_value(): dumped = config.model_dump(mode="json") assert dumped["controls"][0]["action"] == Action.REDACT.value assert dumped["controls"][0]["stage"] == Stage.INPUT.value + + +class QuickRemoteCheck(GuardrailCheck): + name = "_test_quick_remote" + label = "Quick remote" + supported_stages = {Stage.INPUT, Stage.OUTPUT} + remote = True + + def scan(self, text, stage, context): + from docsgpt import tracing + + with tracing.span(tracing.KIND_LLM, "chat judge"): + pass + return CheckOutcome.clean() + + +class TestTraceSpans: + """Guardrail evaluations appear in the execution trace only when they matter.""" + + @pytest.fixture(autouse=True) + def _trace(self, monkeypatch): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + GuardrailCreator.register(QuickRemoteCheck.name, QuickRemoteCheck) + self.trace = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(self.trace): + yield + GuardrailCreator.checks.pop(QuickRemoteCheck.name, None) + + def test_clean_local_scan_is_not_traced(self): + engine = GuardrailEngine( + _config(controls=[{"check": "_test_never", "stage": "output", "action": "block"}]) + ) + engine.evaluate("text", Stage.OUTPUT) + assert self.trace.spans == [] + assert self.trace.content_blocked is False + + def test_firing_local_scan_is_traced_and_blocks_previews(self): + engine = GuardrailEngine(_config()) + engine.evaluate("some text", Stage.INPUT) + (span,) = self.trace.spans + assert span.kind == "guardrail" + assert span.attributes["docsgpt.guardrail.blocked"] is True + assert span.attributes["docsgpt.guardrail.triggered"] == ["_test_always"] + assert self.trace.content_blocked is True + + def test_remote_scan_is_traced_with_judge_nested(self): + engine = GuardrailEngine( + _config(controls=[{"check": "_test_quick_remote", "stage": "input", "action": "block"}]) + ) + decision = engine.evaluate("text", Stage.INPUT) + assert decision.clean + guard, judge = self.trace.spans + assert guard.kind == "guardrail" + assert judge.parent_id == guard.id + assert self.trace.content_blocked is False From bae842d151770ddf7acedab91e01f1037813082a Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:34:54 +0100 Subject: [PATCH 09/33] Trace every chat turn from agent setup to the last event StreamProcessor starts the trace and mints the request id before the agent is built, so pre-fetch retrieval and compression are inside it and side-channel LLM calls share the id. complete_stream activates it in the SSE pump thread, binds the message and conversation, and writes it once however the stream ends: paused, failed, abandoned or superseded (dropped). user_logs rows now carry request_id and message_id. --- docsgpt/api/answer/routes/answer.py | 5 +- docsgpt/api/answer/routes/base.py | 82 +++++- docsgpt/api/answer/routes/stream.py | 5 +- .../api/answer/services/stream_processor.py | 55 +++- docsgpt/api/v1/routes.py | 6 +- docsgpt/tracing/core.py | 5 + tests/api/answer/test_trace_wiring.py | 248 ++++++++++++++++++ tests/conftest.py | 4 + tests/tracing/test_trace.py | 14 + 9 files changed, 415 insertions(+), 9 deletions(-) create mode 100644 tests/api/answer/test_trace_wiring.py diff --git a/docsgpt/api/answer/routes/answer.py b/docsgpt/api/answer/routes/answer.py index 7cdcd722..bc2eb726 100644 --- a/docsgpt/api/answer/routes/answer.py +++ b/docsgpt/api/answer/routes/answer.py @@ -87,7 +87,7 @@ class AnswerResource(Resource, BaseAnswerResource): if error := self.validate_request(data): return error decoded_token = getattr(request, "decoded_token", None) - processor = StreamProcessor(data, decoded_token) + processor = StreamProcessor(data, decoded_token, trace_source="answer") try: # ---- Continuation mode ---- if data.get("tool_actions"): @@ -115,6 +115,7 @@ class AnswerResource(Resource, BaseAnswerResource): decoded_token=processor.decoded_token, agent_id=processor.agent_id, model_id=processor.model_id, + trace=processor.trace, _continuation={ "messages": messages, "tools_dict": tools_dict, @@ -156,6 +157,8 @@ class AnswerResource(Resource, BaseAnswerResource): is_shared_usage=processor.is_shared_usage, shared_token=processor.shared_token, model_id=processor.model_id, + request_id=processor.request_id, + trace=processor.trace, ) stream_result = self.process_response_stream(stream) diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index e0cfe4cd..019ab261 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -1,4 +1,5 @@ import datetime +import inspect import json import logging import threading @@ -9,6 +10,7 @@ from typing import Any, Dict, Generator, List, Optional from flask import jsonify, make_response, Response from flask_restx import Namespace +from docsgpt import tracing from docsgpt.api.answer.services.continuation_service import ContinuationService from docsgpt.api.answer.services.conversation_service import ( ConversationService, @@ -245,6 +247,59 @@ class BaseAnswerResource: return error def complete_stream( + self, + *args: Any, + request_id: Optional[str] = None, + trace: Optional["tracing.Trace"] = None, + **kwargs: Any, + ) -> Generator[str, None, None]: + """Stream the conversation response and record its execution trace. + + Accepts every argument of :meth:`_complete_stream_body`. The trace is + activated for the whole stream -- which runs in the SSE pump thread, + not the request thread -- and written exactly once when the stream + ends, however it ends. + + Args: + request_id: The request's id, generated by the route before the + agent was built so pre-fetch retrieval shares it. A + continuation's saved id takes precedence; absent both, a new + one is minted. + trace: The trace the route started (it already holds the + pre-fetch retrieval). A new one is started when omitted. + + Yields: + Server-sent event strings. + """ + bound = inspect.signature(self._complete_stream_body).bind(*args, **kwargs) + continuation = bound.arguments.get("_continuation") + request_id = ( + continuation.get("request_id") if continuation else None + ) or request_id or str(uuid.uuid4()) + decoded_token = bound.arguments.get("decoded_token") or {} + if not isinstance(trace, tracing.Trace): + trace = None + if trace is None: + agent = bound.arguments.get("agent") + trace = tracing.start_trace( + source=str(getattr(agent, "endpoint", None) or "stream"), + capture_otel_context=False, + ) + if trace is not None: + trace.bind( + request_id=request_id, + user_id=decoded_token.get("sub") if isinstance(decoded_token, dict) else None, + agent_id=bound.arguments.get("agent_id"), + ) + with tracing.activate(trace): + try: + yield from self._complete_stream_body( + *bound.args, request_id=request_id, **bound.kwargs + ) + finally: + tracing.flush(trace) + + def _complete_stream_body( self, question: str, agent: Any, @@ -263,6 +318,7 @@ class BaseAnswerResource: model_user_id: Optional[str] = None, _continuation: Optional[Dict] = None, finalize_tool_pause_as_complete: bool = False, + request_id: Optional[str] = None, ) -> Generator[str, None, None]: """ Generator function that streams the complete conversation response. @@ -336,10 +392,10 @@ class BaseAnswerResource: pending_pause_event: Optional[dict] = None # One id shared across the WAL row, primary LLM (token_usage - # attribution), the SSE event, and resumed continuations. + # attribution), the SSE event, resumed continuations and the trace. request_id = ( _continuation.get("request_id") if _continuation else None - ) or str(uuid.uuid4()) + ) or request_id or str(uuid.uuid4()) # Reserve the placeholder row before the LLM call so a crash # mid-stream still leaves the question queryable. Continuations @@ -405,6 +461,17 @@ class BaseAnswerResource: primary_llm = getattr(agent, "llm", None) if primary_llm is not None: primary_llm._request_id = request_id + # Side-channel LLMs built later (guardrail judge, retrievers the agent + # creates for its search tools) read the id off the agent. + if getattr(agent, "request_id", None) is None: + try: + agent.request_id = request_id + except Exception: + logger.debug("Could not stamp request_id on the agent") + tracing.bind( + message_id=reserved_message_id, + conversation_id=str(conversation_id) if conversation_id else None, + ) # Flipped to ``streaming`` on the first ``answer``/``sources`` chunk; # the reconciler reads ``status`` to tell "never started" from "in @@ -795,6 +862,9 @@ class BaseAnswerResource: # ---- Paused: save continuation state and end stream early ---- if paused: continuation = getattr(agent, "_pending_continuation", None) + trace = tracing.current_trace() + if trace is not None: + trace.outcome = tracing.STATUS_PAUSED # ---- Stateless-tool-round mode (OpenAI-compatible /v1) ---- # OpenAI clients resume by re-POSTing the whole message @@ -1232,6 +1302,9 @@ class BaseAnswerResource: "tool_calls": tool_calls_for_logging, "attachments": attachment_ids, "timestamp": datetime.datetime.now(datetime.timezone.utc), + # Links the Logs row to this turn's execution trace. + "request_id": request_id, + "message_id": reserved_message_id, } if is_structured: log_data["structured_output"] = True @@ -1469,9 +1542,14 @@ class BaseAnswerResource: ) if journal_writer is not None: journal_writer.close() + # The user replaced this turn; its trace describes nothing kept. + tracing.discard(tracing.current_trace()) return except Exception as e: logger.error(f"Error in stream: {str(e)}", exc_info=True) + trace = tracing.current_trace() + if trace is not None: + trace.outcome = tracing.STATUS_ERROR # This process took the resume claim, so it owns releasing it. The # only other way back is ``revert_stale_resuming``'s 600 s grace, # which leaves the user locked out of their own conversation for diff --git a/docsgpt/api/answer/routes/stream.py b/docsgpt/api/answer/routes/stream.py index cdee274b..36cb0472 100644 --- a/docsgpt/api/answer/routes/stream.py +++ b/docsgpt/api/answer/routes/stream.py @@ -94,7 +94,7 @@ class StreamResource(Resource, BaseAnswerResource): if error := self.validate_request(data, "index" in data): return error decoded_token = getattr(request, "decoded_token", None) - processor = StreamProcessor(data, decoded_token) + processor = StreamProcessor(data, decoded_token, trace_source="stream") try: # ---- Continuation mode ---- @@ -130,6 +130,7 @@ class StreamResource(Resource, BaseAnswerResource): agent_id=processor.agent_id, model_id=processor.model_id, model_user_id=processor.model_user_id, + trace=processor.trace, _continuation={ "messages": messages, "tools_dict": tools_dict, @@ -181,6 +182,8 @@ class StreamResource(Resource, BaseAnswerResource): shared_token=processor.shared_token, model_id=processor.model_id, model_user_id=processor.model_user_id, + request_id=processor.request_id, + trace=processor.trace, ), ), mimetype="text/event-stream", diff --git a/docsgpt/api/answer/services/stream_processor.py b/docsgpt/api/answer/services/stream_processor.py index 004b616c..bb7f47a3 100644 --- a/docsgpt/api/answer/services/stream_processor.py +++ b/docsgpt/api/answer/services/stream_processor.py @@ -1,9 +1,12 @@ import datetime +import functools import json import logging +import uuid from pathlib import Path from typing import Any, Dict, Optional, Set +from docsgpt import tracing from docsgpt.agents.agent_creator import AgentCreator from docsgpt.agents.default_tools import synthesized_default_tools from docsgpt.api.answer.services.compression import CompressionOrchestrator @@ -104,9 +107,41 @@ def get_prompt(prompt_id: str, prompts_collection=None) -> str: raise ValueError(f"Invalid prompt ID: {prompt_id}") from e +def _traced_setup(method): + """Run a request-setup method inside the request's execution trace. + + Agent setup does real work worth seeing in the trace -- pre-fetch + retrieval, history compression -- before ``complete_stream`` runs, so + the trace is started here, in the request thread, and handed on. + """ + + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + trace = getattr(self, "trace", None) + if trace is None: + trace = tracing.start_trace(source=getattr(self, "trace_source", "stream")) + self.trace = trace + with tracing.activate(trace): + try: + return method(self, *args, **kwargs) + finally: + if trace is not None: + decoded = getattr(self, "decoded_token", None) + trace.bind( + request_id=getattr(self, "request_id", None), + user_id=decoded.get("sub") if isinstance(decoded, dict) else None, + agent_id=getattr(self, "agent_id", None), + ) + + return wrapper + + class StreamProcessor: def __init__( - self, request_data: Dict[str, Any], decoded_token: Optional[Dict[str, Any]] + self, + request_data: Dict[str, Any], + decoded_token: Optional[Dict[str, Any]], + trace_source: str = "stream", ): # Legacy attribute retained as None for any external callers that # introspect the processor; all DB access uses per-op connections. @@ -137,6 +172,11 @@ class StreamProcessor: self.reserved_message_id: Optional[str] = None # Carried through resumes so multi-pause runs keep one request_id. self.request_id: Optional[str] = None + # The request's execution trace, started by the first traced setup + # step and handed to ``complete_stream``; ``trace_source`` names the + # entry point it is stored under. + self.trace: Optional[tracing.Trace] = None + self.trace_source = trace_source self.conversation_service = ConversationService() self.compression_orchestrator = CompressionOrchestrator( self.conversation_service @@ -161,12 +201,17 @@ class StreamProcessor: self._load_conversation_history() self._process_attachments() + @_traced_setup def build_agent(self, question: str): """One call to go from request data to a ready-to-run agent. Combines initialize(), pre_fetch_docs(), pre_fetch_tools(), and - create_agent() into a single convenience method. + create_agent() into a single convenience method. The request id is + minted first so pre-fetch retrieval and its side-channel LLM calls + share it with the rest of the turn. """ + if not getattr(self, "request_id", None): + self.request_id = self.data.get("request_id") or str(uuid.uuid4()) self.initialize() agent_type = self.agent_config.get("agent_type", "classic") @@ -221,6 +266,7 @@ class StreamProcessor: tools_data=tools_data, ) + @_traced_setup def build_continuation_from_messages(self, messages, tool_actions): """Rebuild a tool continuation from the request messages (STATELESS). @@ -1118,7 +1164,7 @@ class StreamProcessor: user_api_key=self.agent_config["user_api_key"], agent_id=self.agent_id, decoded_token=self.decoded_token, - request_id=self.data.get("request_id"), + request_id=self.request_id or self.data.get("request_id"), ) def _legacy_classic(): @@ -1440,6 +1486,7 @@ class StreamProcessor: logger.warning(f"Failed to fetch memory tool data: {str(e)}") return None + @_traced_setup def resume_from_tool_actions( self, tool_actions: list, @@ -1813,7 +1860,7 @@ class StreamProcessor: "llm_name": provider or settings.LLM_PROVIDER, "api_key": system_api_key, "decoded_token": self.decoded_token, - "request_id": self.data.get("request_id"), + "request_id": self.request_id or self.data.get("request_id"), } elif agent_type == "workflow": diff --git a/docsgpt/api/v1/routes.py b/docsgpt/api/v1/routes.py index fa915d24..9294c23a 100644 --- a/docsgpt/api/v1/routes.py +++ b/docsgpt/api/v1/routes.py @@ -257,7 +257,7 @@ def chat_completions(): internal_data["persist"] = True try: - processor = StreamProcessor(internal_data, decoded_token) + processor = StreamProcessor(internal_data, decoded_token, trace_source="v1") # Set when this request took the resume claim, so a refusal can release it. claimed_conversation_id = None @@ -495,6 +495,8 @@ def _stream_response( visibility=visibility, _continuation=continuation, finalize_tool_pause_as_complete=finalize_stateless_tool_pause, + request_id=processor.request_id, + trace=getattr(processor, "trace", None), ) translation_state = StreamTranslationState() @@ -572,6 +574,8 @@ def _non_stream_response( visibility=visibility, _continuation=continuation, finalize_tool_pause_as_complete=finalize_stateless_tool_pause, + request_id=processor.request_id, + trace=getattr(processor, "trace", None), ) result = helper.process_response_stream(stream) diff --git a/docsgpt/tracing/core.py b/docsgpt/tracing/core.py index d9042f42..47a2c49a 100644 --- a/docsgpt/tracing/core.py +++ b/docsgpt/tracing/core.py @@ -263,6 +263,9 @@ class Trace: self.start_perf_ns = time.perf_counter_ns() self.end_perf_ns: Optional[int] = None self.status: Optional[str] = None + # Status the owner determined mid-run (e.g. ``paused`` at a tool + # approval); used by :meth:`finish` when no explicit status is given. + self.outcome: Optional[str] = None self.spans: List[Span] = [] self.dropped_spans = 0 self.content_blocked = False @@ -370,6 +373,8 @@ class Trace: span.status = STATUS_CANCELLED span.end_perf_ns = now self.end_perf_ns = now + if status is None: + status = self.outcome if status is not None: self.status = status else: diff --git a/tests/api/answer/test_trace_wiring.py b/tests/api/answer/test_trace_wiring.py new file mode 100644 index 00000000..49fcd6e0 --- /dev/null +++ b/tests/api/answer/test_trace_wiring.py @@ -0,0 +1,248 @@ +"""``complete_stream`` owns the request's execution trace and writes it once.""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings + + +@pytest.fixture(autouse=True) +def _tracing_on(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + + +@contextmanager +def _captured_flushes(): + """Record every flushed trace instead of writing it.""" + flushed = [] + + def _fake_flush(trace, status=None): + if trace is None or trace.flushed: + return + trace.flushed = True + trace.finish(status) + flushed.append(trace) + + with patch("docsgpt.tracing.flush", side_effect=_fake_flush): + yield flushed + + +def _agent(events): + agent = MagicMock() + + def _gen(query): + with tracing.span(tracing.KIND_AGENT, "invoke_agent Fake"): + with tracing.span(tracing.KIND_LLM, "chat m"): + pass + yield from events + + agent.gen.side_effect = _gen + agent.tool_calls = [] + agent.compression_metadata = None + agent.compression_saved = False + return agent + + +def _run(resource, agent, **kwargs): + base = dict( + question="q", + agent=agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u-trace"}, + should_persist=False, + ) + base.update(kwargs) + return list(resource.complete_stream(**base)) + + +@pytest.mark.unit +class TestTraceLifecycle: + def test_normal_turn_flushes_once_with_ids(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(), _captured_flushes() as flushed: + _run(BaseAnswerResource(), _agent([{"answer": "hi"}]), request_id="req-1") + (trace,) = flushed + assert trace.request_id == "req-1" + assert trace.user_id == "u-trace" + assert trace.status == "ok" + assert [s.name for s in trace.spans] == ["invoke_agent Fake", "chat m"] + + def test_route_trace_is_reused(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + route_trace = tracing.start_trace(source="answer", capture_otel_context=False) + with tracing.activate(route_trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + with flask_app.app_context(), _captured_flushes() as flushed: + _run( + BaseAnswerResource(), + _agent([{"answer": "hi"}]), + request_id="req-2", + trace=route_trace, + ) + assert flushed == [route_trace] + assert route_trace.source == "answer" + assert [s.kind for s in route_trace.spans] == ["retrieval", "agent", "llm"] + + def test_mock_trace_is_ignored(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(), _captured_flushes() as flushed: + _run(BaseAnswerResource(), _agent([{"answer": "hi"}]), trace=MagicMock()) + assert len(flushed) == 1 + assert isinstance(flushed[0], tracing.Trace) + + def test_continuation_keeps_saved_request_id(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + agent = _agent([]) + agent.gen_continuation.side_effect = lambda **_kw: iter([{"answer": "done"}]) + with flask_app.app_context(), _captured_flushes() as flushed: + _run( + BaseAnswerResource(), + agent, + question="", + request_id="fresh", + _continuation={ + "messages": [], + "tools_dict": {}, + "pending_tool_calls": [], + "tool_actions": [], + "request_id": "saved-req", + }, + ) + assert flushed[0].request_id == "saved-req" + + def test_agent_error_marks_trace_error(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + agent = MagicMock() + agent.gen.side_effect = RuntimeError("upstream down") + with flask_app.app_context(), _captured_flushes() as flushed: + stream = _run(BaseAnswerResource(), agent) + assert any('"type": "error"' in s for s in stream) + assert flushed[0].status == "error" + + def test_abandoned_stream_still_flushes(self, flask_app, mock_mongo_db): + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(), _captured_flushes() as flushed: + gen = BaseAnswerResource().complete_stream( + question="q", + agent=_agent([{"answer": "a"}, {"answer": "b"}]), + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "u"}, + should_persist=False, + ) + next(gen) + gen.close() + assert len(flushed) == 1 + + +@pytest.mark.unit +class TestTraceWithPersistence: + def test_message_and_conversation_ids_bound(self, pg_conn, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + from tests.api.answer.test_base_routes import _patch_db_session + + with flask_app.app_context(), _patch_db_session(pg_conn), _captured_flushes() as flushed: + _run( + BaseAnswerResource(), + _agent([{"answer": "persisted"}]), + should_persist=True, + model_id="gpt-4", + request_id="req-p", + ) + trace = flushed[0] + assert trace.message_id + assert trace.conversation_id + from sqlalchemy import text as sql_text + + row = pg_conn.execute( + sql_text("SELECT data FROM user_logs WHERE user_id = 'u-trace'") + ).fetchone() + assert row[0]["request_id"] == "req-p" + assert row[0]["message_id"] == trace.message_id + + def test_paused_turn_is_flushed_paused(self, pg_conn, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + from tests.api.answer.test_base_routes import _patch_db_session + + agent = _agent( + [ + { + "type": "tool_calls_pending", + "data": {"pending_tool_calls": [{"call_id": "c1"}]}, + } + ] + ) + agent._pending_continuation = { + "messages": [], + "tools_dict": {}, + "pending_tool_calls": [{"call_id": "c1"}], + } + with flask_app.app_context(), _patch_db_session(pg_conn), patch( + "docsgpt.api.answer.services.continuation_service.ContinuationService.save_state", + return_value=True, + ), _captured_flushes() as flushed: + _run(BaseAnswerResource(), agent, should_persist=True, model_id="gpt-4") + assert flushed[0].status == "paused" + + +@pytest.mark.unit +class TestProcessorTraceSetup: + def test_build_agent_mints_request_id_inside_the_trace(self): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + seen = {} + + class _Stop(Exception): + pass + + def _initialize(): + seen["trace"] = tracing.current_trace() + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + raise _Stop() + + processor = StreamProcessor({"question": "q"}, {"sub": "u1"}, trace_source="answer") + with patch.object(processor, "initialize", side_effect=_initialize): + with pytest.raises(_Stop): + processor.build_agent("q") + trace = processor.trace + assert seen["trace"] is trace + assert trace.source == "answer" + assert processor.request_id and trace.request_id == processor.request_id + assert trace.user_id == "u1" + assert [s.kind for s in trace.spans] == ["retrieval"] + assert tracing.current_trace() is None + + def test_client_supplied_request_id_is_kept(self): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + processor = StreamProcessor({"request_id": "client-rid"}, {"sub": "u1"}) + with patch.object(processor, "initialize", side_effect=RuntimeError("stop")): + with pytest.raises(RuntimeError): + processor.build_agent("q") + assert processor.request_id == "client-rid" + + def test_tracing_disabled_leaves_no_trace(self, monkeypatch): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + monkeypatch.setattr(settings, "TRACES_ENABLED", False) + processor = StreamProcessor({}, {"sub": "u1"}) + with patch.object(processor, "initialize", side_effect=RuntimeError("stop")): + with pytest.raises(RuntimeError): + processor.build_agent("q") + assert processor.trace is None + assert processor.request_id diff --git a/tests/conftest.py b/tests/conftest.py index 876e6d89..f5bee6a6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -44,6 +44,10 @@ import os os.environ.setdefault("AUTO_MIGRATE", "false") os.environ.setdefault("AUTO_CREATE_DB", "false") os.environ.setdefault("AUTO_VECTOR_SCHEMA", "false") +# Execution traces write to ``request_traces`` through their own DB session +# when a request finishes. Off by default so unrelated tests never touch a +# database through that path; tracing tests switch it on explicitly. +os.environ.setdefault("TRACES_ENABLED", "false") import subprocess import sys diff --git a/tests/tracing/test_trace.py b/tests/tracing/test_trace.py index 685a5595..03a4caca 100644 --- a/tests/tracing/test_trace.py +++ b/tests/tracing/test_trace.py @@ -339,3 +339,17 @@ class TestRecord: } assert span["attributes"]["gen_ai.request.model"] == "m" assert span["offset_ms"] >= 0 + + +class TestOutcome: + def test_outcome_is_used_when_no_status_given(self): + trace = tracing.start_trace(source="stream") + trace.outcome = "paused" + trace.finish() + assert trace.status == "paused" + + def test_explicit_status_beats_outcome(self): + trace = tracing.start_trace(source="stream") + trace.outcome = "paused" + trace.finish(status="error") + assert trace.status == "error" From 5d0992eef85fdf13d1d9ccdddf72b48d70d79bc4 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:37:24 +0100 Subject: [PATCH 10/33] Trace scheduled, webhook, search, MCP and graph-extraction runs run_agent_headless records each unattended run under its endpoint; the scheduler passes its run id and the webhook worker its task id so Logs rows can find their trace, while the LLM's own request id stays untouched for quota counts. /api/search and MCP search_docs record their retrieval, and a graph build records every extraction call under one step. --- docsgpt/agents/headless_runner.py | 46 +++++++ docsgpt/api/user/scheduler_worker.py | 2 + docsgpt/graphrag/extraction.py | 5 +- docsgpt/mcp_server.py | 2 +- docsgpt/services/search_service.py | 29 ++++- docsgpt/worker.py | 39 ++++-- tests/services/test_mcp_server.py | 6 +- tests/tracing/test_entry_points.py | 180 +++++++++++++++++++++++++++ tests/worker/test_extract_graph.py | 50 ++++++++ 9 files changed, 344 insertions(+), 15 deletions(-) create mode 100644 tests/tracing/test_entry_points.py diff --git a/docsgpt/agents/headless_runner.py b/docsgpt/agents/headless_runner.py index c2a4d713..e50e2478 100644 --- a/docsgpt/agents/headless_runner.py +++ b/docsgpt/agents/headless_runner.py @@ -5,6 +5,7 @@ from __future__ import annotations import logging from typing import Any, Dict, Iterable, List, Optional +from docsgpt import tracing from docsgpt.agents.agent_creator import AgentCreator from docsgpt.agents.tool_executor import ToolExecutor from docsgpt.api.answer.services.prompt_renderer import ( @@ -69,12 +70,57 @@ def run_agent_headless( endpoint: str = "headless", chat_history: Optional[List[Dict[str, Any]]] = None, conversation_id: Optional[str] = None, + request_id: Optional[str] = None, ) -> Dict[str, Any]: """Run an agent with no live client; returns a structured outcome dict. + The run is recorded as one execution trace under ``endpoint`` as its + source. ``request_id`` links that trace to the caller's own record (the + scheduler passes its run id, the webhook worker its task id); it is kept + off the LLM's token-usage rows, whose request ids drive request counts. + Raises: QuotaExceededError: If the agent owner's usage quota is exhausted. """ + trace = tracing.start_trace( + source=endpoint, + request_id=request_id, + user_id=_resolve_owner(agent_config), + agent_id=_resolve_agent_id(agent_config), + conversation_id=conversation_id, + ) + status = None + with tracing.activate(trace): + try: + outcome = _run_agent_headless( + agent_config, + query, + tool_allowlist=tool_allowlist, + model_id_override=model_id_override, + endpoint=endpoint, + chat_history=chat_history, + conversation_id=conversation_id, + ) + if outcome.get("error"): + status = tracing.STATUS_ERROR + return outcome + except BaseException: + status = tracing.STATUS_ERROR + raise + finally: + tracing.flush(trace, status) + + +def _run_agent_headless( + agent_config: Dict[str, Any], + query: str, + *, + tool_allowlist: Optional[Iterable[str]] = None, + model_id_override: Optional[str] = None, + endpoint: str = "headless", + chat_history: Optional[List[Dict[str, Any]]] = None, + conversation_id: Optional[str] = None, +) -> Dict[str, Any]: from docsgpt.core.model_utils import ( get_api_key_for_provider, get_default_model_id, diff --git a/docsgpt/api/user/scheduler_worker.py b/docsgpt/api/user/scheduler_worker.py index 3e085f0e..15c73f76 100644 --- a/docsgpt/api/user/scheduler_worker.py +++ b/docsgpt/api/user/scheduler_worker.py @@ -277,6 +277,8 @@ def execute_scheduled_run_body(run_id: str, celery_task_id: Optional[str]) -> Di endpoint="schedule", conversation_id=schedule.get("origin_conversation_id"), chat_history=chat_history, + # Links the run's execution trace to its Logs row. + request_id=str(run_id), ) except SoftTimeLimitExceeded: timed_out = True diff --git a/docsgpt/graphrag/extraction.py b/docsgpt/graphrag/extraction.py index daa8f910..80efd1fa 100644 --- a/docsgpt/graphrag/extraction.py +++ b/docsgpt/graphrag/extraction.py @@ -23,6 +23,7 @@ import logging import re from typing import Any, Callable, Dict, List, Optional +from docsgpt import tracing from docsgpt.core.model_utils import ( get_api_key_for_provider, get_provider_from_model_id, @@ -353,7 +354,9 @@ def extract_graph_for_source( if pool is not None: # ``map`` yields in submission order, so chunks are still applied in # the order they were given and a run stays reproducible. - prepared = pool.map(_prepare, items) + # Pool threads don't inherit context; carry the trace in so each + # chunk's extraction LLM call is recorded. + prepared = pool.map(tracing.wrap(_prepare), items) else: prepared = (_prepare(item) for item in items) missed = [] diff --git a/docsgpt/mcp_server.py b/docsgpt/mcp_server.py index f67120b2..9e8c7fbf 100644 --- a/docsgpt/mcp_server.py +++ b/docsgpt/mcp_server.py @@ -53,7 +53,7 @@ async def search_docs(query: str, chunks: int = 5) -> list[dict]: if not api_key: raise PermissionError("Missing Bearer token") try: - return await asyncio.to_thread(search, api_key, query, chunks) + return await asyncio.to_thread(search, api_key, query, chunks, source="mcp") except InvalidAPIKey as exc: raise PermissionError("Invalid API key") from exc except SearchFailed: diff --git a/docsgpt/services/search_service.py b/docsgpt/services/search_service.py index b67a1aca..68756c91 100644 --- a/docsgpt/services/search_service.py +++ b/docsgpt/services/search_service.py @@ -10,10 +10,12 @@ from __future__ import annotations import logging from typing import Any, Dict, List, Optional +from docsgpt import tracing from docsgpt.core.settings import settings from docsgpt.retriever.fanout import fetch_per_source from docsgpt.storage.db.repositories.agents import AgentsRepository from docsgpt.storage.db.session import db_readonly +from docsgpt.tracing.retrieval import describe_documents, start_retrieval_span from docsgpt.vectorstore.vector_creator import VectorCreator logger = logging.getLogger(__name__) @@ -219,14 +221,21 @@ def _search_sources( return results[:chunks] -def search(api_key: str, query: str, chunks: int = 5) -> List[Dict[str, Any]]: +def search( + api_key: str, query: str, chunks: int = 5, *, source: str = "search" +) -> List[Dict[str, Any]]: """Resolve an agent by API key and search its sources. + Every search that reaches the sources is recorded as an execution trace + owned by the agent's owner, under ``source``. + Args: api_key: Agent API key (the opaque string stored on ``agents.key`` in Postgres). query: Free-text search query. chunks: Max number of hits to return. + source: Trace source name: ``search`` for ``/api/search``, ``mcp`` + for the MCP ``search_docs`` tool. Returns: List of hit dicts with ``text``, ``title``, ``source`` keys. @@ -256,4 +265,20 @@ def search(api_key: str, query: str, chunks: int = 5) -> List[Dict[str, Any]]: if not source_ids: return [] - return _search_sources(query, source_ids, chunks) + trace = tracing.start_trace( + source=source, + user_id=agent.get("user_id"), + agent_id=str(agent.get("id")) if agent.get("id") else None, + ) + with tracing.activate(trace): + try: + with start_retrieval_span( + f"retrieval {source}", + sources=source_ids, + **{"docsgpt.top_k": chunks}, + ) as span: + results = _search_sources(query, source_ids, chunks) + describe_documents(span, results, query=query) + return results + finally: + tracing.flush(trace) diff --git a/docsgpt/worker.py b/docsgpt/worker.py index 9fbf57f3..f3311859 100755 --- a/docsgpt/worker.py +++ b/docsgpt/worker.py @@ -16,6 +16,7 @@ from urllib.parse import urljoin, urlsplit import requests +from docsgpt import tracing from docsgpt.core.settings import settings from docsgpt.events.publisher import publish_user_event from docsgpt.parser.chunking_creator import ChunkerCreator @@ -2129,6 +2130,7 @@ def agent_webhook_worker(self, agent_id, payload): input_data, tool_allowlist=_webhook_tool_allowlist(agent_config), endpoint="webhook", + request_id=getattr(getattr(self, "request", None), "id", None), ) result = { "answer": outcome.get("answer", ""), @@ -2942,16 +2944,36 @@ def extract_graph_worker(self, source_id, user): }, ) + trace = tracing.start_trace( + source="graph_extraction", + name=f"graph_extraction {source.get('name') or source_id}", + request_id=getattr(self.request, "id", None), + user_id=user, + ) try: - summary = extract_graph_for_source( - source_id, - user, - chunks, - config=cfg, - request_id=getattr(self.request, "id", None), - progress_cb=_progress, - ) + with tracing.activate(trace), tracing.span( + tracing.KIND_STEP, + "graph_extraction", + attributes={"docsgpt.source_id": source_id, "docsgpt.chunk_count": total}, + ) as span: + summary = extract_graph_for_source( + source_id, + user, + chunks, + config=cfg, + request_id=getattr(self.request, "id", None), + progress_cb=_progress, + ) + if isinstance(summary, dict): + span.set( + **{ + "docsgpt.graph.nodes": summary.get("nodes"), + "docsgpt.graph.edges": summary.get("edges"), + "docsgpt.graph.chunks_processed": summary.get("chunks_processed"), + } + ) except Exception as e: + tracing.flush(trace, tracing.STATUS_ERROR) _publish_graph_event( user, source_id, @@ -2960,6 +2982,7 @@ def extract_graph_worker(self, source_id, user): ) raise + tracing.flush(trace) _publish_graph_event( user, source_id, diff --git a/tests/services/test_mcp_server.py b/tests/services/test_mcp_server.py index d8441db5..b048eb77 100644 --- a/tests/services/test_mcp_server.py +++ b/tests/services/test_mcp_server.py @@ -99,7 +99,7 @@ class TestSearchDocsTool: ): out = await search_docs(query="q", chunks=7) assert out == hits - mock_search.assert_called_once_with("the-key", "q", 7) + mock_search.assert_called_once_with("the-key", "q", 7, source="mcp") @pytest.mark.asyncio async def test_default_chunks_is_5(self): @@ -115,7 +115,7 @@ class TestSearchDocsTool: ) as mock_search, ): await search_docs(query="q") - mock_search.assert_called_once_with("k", "q", 5) + mock_search.assert_called_once_with("k", "q", 5, source="mcp") @pytest.mark.asyncio async def test_bearer_scheme_case_insensitive(self): @@ -131,4 +131,4 @@ class TestSearchDocsTool: ) as mock_search, ): await search_docs(query="q") - mock_search.assert_called_once_with("lowercase-scheme", "q", 5) + mock_search.assert_called_once_with("lowercase-scheme", "q", 5, source="mcp") diff --git a/tests/tracing/test_entry_points.py b/tests/tracing/test_entry_points.py new file mode 100644 index 00000000..221b0c29 --- /dev/null +++ b/tests/tracing/test_entry_points.py @@ -0,0 +1,180 @@ +"""Non-chat entry points each record and write one execution trace.""" + +from __future__ import annotations + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest + +from docsgpt import tracing +from docsgpt.core.settings import settings + + +@pytest.fixture(autouse=True) +def _tracing_on(monkeypatch): + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + + +@pytest.fixture() +def flushed(): + """Capture flushed traces instead of writing them.""" + captured = [] + + def _fake_flush(trace, status=None): + if trace is None or trace.flushed: + return + trace.flushed = True + trace.finish(status) + captured.append(trace) + + with patch("docsgpt.tracing.flush", side_effect=_fake_flush): + yield captured + + +def _headless(events, monkeypatch, **kwargs): + from docsgpt.agents import headless_runner as hr + + agent = MagicMock(name="agent") + + def _gen(query): + with tracing.span(tracing.KIND_AGENT, "invoke_agent Fake"): + yield from events + + agent.gen.side_effect = _gen + agent.llm.token_usage = {"prompt_tokens": 1, "generated_tokens": 1} + retriever = MagicMock(name="retriever") + + def _search(query): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + return [] + + retriever.search.side_effect = _search + tool_executor = MagicMock(name="tool_executor") + tool_executor.headless_denials = [] + monkeypatch.setattr(hr, "get_prompt", lambda _pid: "system prompt") + monkeypatch.setattr( + hr.RetrieverCreator, "create_retriever", classmethod(lambda cls, *a, **kw: retriever) + ) + monkeypatch.setattr(hr, "ToolExecutor", lambda *a, **kw: tool_executor) + monkeypatch.setattr( + hr.AgentCreator, "create_agent", classmethod(lambda cls, *a, **kw: agent) + ) + with patch("docsgpt.core.model_utils.validate_model_id", return_value=True), \ + patch("docsgpt.core.model_utils.get_default_model_id", return_value="m"), \ + patch("docsgpt.core.model_utils.get_provider_from_model_id", return_value="openai"), \ + patch("docsgpt.core.model_utils.get_api_key_for_provider", return_value="k"), \ + patch("docsgpt.utils.calculate_doc_token_budget", return_value=1000): + return hr.run_agent_headless( + {"user_id": "owner-1", "id": "11111111-1111-1111-1111-111111111111"}, + "do the thing", + **kwargs, + ) + + +@pytest.mark.unit +class TestHeadless: + def test_scheduled_run_is_traced_with_its_run_id(self, monkeypatch, flushed): + _headless([{"answer": "done"}], monkeypatch, endpoint="schedule", request_id="run-1") + (trace,) = flushed + assert trace.source == "schedule" + assert trace.request_id == "run-1" + assert trace.user_id == "owner-1" + assert trace.status == "ok" + assert [s.kind for s in trace.spans] == ["retrieval", "agent"] + + def test_stream_error_marks_trace_error(self, monkeypatch, flushed): + outcome = _headless([{"type": "error", "error": "boom"}], monkeypatch, endpoint="webhook") + assert outcome["error_type"] == "stream_error" + assert flushed[0].status == "error" + + def test_raised_error_still_flushes(self, monkeypatch, flushed): + from docsgpt.agents import headless_runner as hr + + monkeypatch.setattr(hr, "_run_agent_headless", MagicMock(side_effect=RuntimeError("x"))) + with pytest.raises(RuntimeError): + hr.run_agent_headless({"user_id": "u"}, "q") + assert flushed[0].status == "error" + + def test_trace_request_id_stays_off_llm_usage_rows(self, monkeypatch, flushed): + """The headless LLM's own request id is untouched (quota counts depend on it).""" + from docsgpt.agents import headless_runner as hr + + created = {} + + def _create(cls, *a, **kw): + agent = MagicMock() + agent.gen.return_value = iter([{"answer": "x"}]) + agent.llm.token_usage = {} + agent.llm._request_id = None + created["agent"] = agent + return agent + + monkeypatch.setattr(hr.AgentCreator, "create_agent", classmethod(_create)) + monkeypatch.setattr(hr, "get_prompt", lambda _pid: "p") + monkeypatch.setattr( + hr.RetrieverCreator, "create_retriever", + classmethod(lambda cls, *a, **kw: MagicMock(search=MagicMock(return_value=[]))), + ) + monkeypatch.setattr(hr, "ToolExecutor", lambda *a, **kw: MagicMock(headless_denials=[])) + with patch("docsgpt.core.model_utils.validate_model_id", return_value=True), \ + patch("docsgpt.core.model_utils.get_default_model_id", return_value="m"), \ + patch("docsgpt.core.model_utils.get_provider_from_model_id", return_value="openai"), \ + patch("docsgpt.core.model_utils.get_api_key_for_provider", return_value="k"), \ + patch("docsgpt.utils.calculate_doc_token_budget", return_value=1000): + hr.run_agent_headless({"user_id": "u"}, "q", endpoint="schedule", request_id="run-9") + assert created["agent"].llm._request_id is None + + +@contextmanager +def _search_env(agent): + repo = MagicMock() + repo.find_by_key.return_value = agent + + @contextmanager + def _conn(): + yield MagicMock() + + store = MagicMock() + store.search.return_value = [{"text": "hit", "metadata": {"title": "T", "source": "s"}}] + with patch("docsgpt.api.user.team_sharing.can_access", return_value=True), \ + patch("docsgpt.services.search_service.db_readonly", _conn), \ + patch("docsgpt.services.search_service.AgentsRepository", return_value=repo), \ + patch( + "docsgpt.services.search_service.VectorCreator.create_vectorstore", + return_value=store, + ): + yield + + +@pytest.mark.unit +class TestSearch: + def test_search_is_traced(self, flushed): + from docsgpt.services.search_service import search + + agent = {"id": "a-1", "source_id": "src-1", "extra_source_ids": [], "user_id": "owner"} + with _search_env(agent): + results = search("k", "what is x", 3) + assert results + (trace,) = flushed + assert trace.source == "search" + assert trace.user_id == "owner" + retrieval = trace.spans[0] + assert retrieval.kind == tracing.KIND_RETRIEVAL + assert retrieval.attributes["docsgpt.chunk_count"] == 1 + + def test_mcp_source_name(self, flushed): + from docsgpt.services.search_service import search + + agent = {"id": "a-1", "source_id": "src-1", "extra_source_ids": [], "user_id": "owner"} + with _search_env(agent): + search("k", "q", 3, source="mcp") + assert flushed[0].source == "mcp" + + def test_no_sources_records_nothing(self, flushed): + from docsgpt.services.search_service import search + + with _search_env({"id": "a-1", "source_id": None, "user_id": "owner"}): + assert search("k", "q", 3) == [] + assert flushed == [] diff --git a/tests/worker/test_extract_graph.py b/tests/worker/test_extract_graph.py index 7e040fbe..d6dd5701 100644 --- a/tests/worker/test_extract_graph.py +++ b/tests/worker/test_extract_graph.py @@ -170,3 +170,53 @@ class TestExtractGraphWorker: worker.extract_graph_worker(task_self, source_id, "alice") assert "graph.extract.failed" in [e[0] for e in events] + + +@pytest.mark.unit +class TestExtractGraphTrace: + """A graph build is one execution trace holding every extraction call.""" + + def _run(self, pg_conn, monkeypatch, task_self, extract): + from docsgpt import tracing, worker + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + source_id = _seed_source(pg_conn) + _patch_store(monkeypatch, [{"doc_id": "c1", "text": "alpha"}]) + monkeypatch.setattr("docsgpt.graphrag.graphrag_available", lambda: True) + monkeypatch.setattr("docsgpt.graphrag.extraction.extract_graph_for_source", extract) + flushed = [] + + def _flush(trace, status=None): + trace.flushed = True + trace.finish(status) + flushed.append(trace) + + monkeypatch.setattr(tracing, "flush", _flush) + return worker, source_id, flushed + + def test_successful_build_is_traced(self, pg_conn, patch_worker_db, task_self, monkeypatch): + from docsgpt import tracing + + def _extract(*_a, **_kw): + with tracing.span(tracing.KIND_LLM, "chat m"): + pass + return {"nodes": 3, "edges": 2, "chunks_processed": 1} + + worker, source_id, flushed = self._run(pg_conn, monkeypatch, task_self, _extract) + worker.extract_graph_worker(task_self, source_id, "alice") + (trace,) = flushed + assert trace.source == "graph_extraction" + assert trace.user_id == "alice" + step, llm = trace.spans + assert step.attributes["docsgpt.graph.nodes"] == 3 + assert llm.parent_id == step.id + + def test_failed_build_is_an_error_trace(self, pg_conn, patch_worker_db, task_self, monkeypatch): + worker, source_id, flushed = self._run( + pg_conn, monkeypatch, task_self, MagicMock(side_effect=RuntimeError("llm down")) + ) + with pytest.raises(RuntimeError): + worker.extract_graph_worker(task_self, source_id, "alice") + assert flushed[0].status == "error" From 26881ba66b594f849c714f076d99e5a0cbd4dd23 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:39:09 +0100 Subject: [PATCH 11/33] Serve traces to the Logs UI GET /api/traces returns a Logs row's stored traces, scoped to the caller or to an agent they own. get_user_logs rows now carry the ids that find their trace and a merged summary (duration, LLM and tool calls, tokens, retrieval time) from one batched lookup per page; searches, MCP searches and graph builds, which have no log row of their own, are listed from their traces. --- docsgpt/api/user/analytics/routes.py | 190 ++++++++++++++++++++++++++- tests/api/user/test_traces_routes.py | 188 ++++++++++++++++++++++++++ 2 files changed, 375 insertions(+), 3 deletions(-) create mode 100644 tests/api/user/test_traces_routes.py diff --git a/docsgpt/api/user/analytics/routes.py b/docsgpt/api/user/analytics/routes.py index 129f7e97..cb4486e5 100644 --- a/docsgpt/api/user/analytics/routes.py +++ b/docsgpt/api/user/analytics/routes.py @@ -14,6 +14,10 @@ from docsgpt.api.user.base import ( ) from docsgpt.storage.db.redaction import redact_secrets from docsgpt.storage.db.repositories.agents import AgentsRepository +from docsgpt.storage.db.repositories.request_traces import ( + REF_FIELDS as TRACE_REF_FIELDS, + RequestTracesRepository, +) from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository from docsgpt.storage.db.session import db_readonly @@ -90,6 +94,154 @@ def _resolve_agent(conn, api_key_id, user_id): return agent, api_key, agent_pg_id +def _trace_branch(name: str, sources_sql: str, scope: str) -> dict: + """A ``get_user_logs`` branch listing stored traces of the given sources.""" + return { + "name": name, + "level": "CASE WHEN t.status = 'error' THEN 'error' ELSE 'info' END", + # A search's query sits in its retrieval span's preview; graph builds + # are named after their source. + "summary": "COALESCE(t.spans->0->'preview'->>'query', t.name, t.source)", + "where": [f"t.source IN {sources_sql}", scope], + "sql": f""" + SELECT '{name}' AS event_type, + CAST(t.id AS text) AS id, + t.user_id AS user_id, + t.started_at AS timestamp, + {{level}} AS level, + t.source AS action, + {{summary}} AS summary, + jsonb_build_object( + 'status', t.status, + 'source', t.source, + 'duration_ms', t.duration_ms + ) AS payload + FROM request_traces t + WHERE {{where}} + """, + } + + +def _trace_ref(item: dict): + """The ``(field, value)`` that finds a Logs row's trace, if it has one.""" + event_type = item["event_type"] + row_id = item["id"].split("-", 1)[1] + if event_type == "chat": + return ("request_id", item.get("request_id")) if item.get("request_id") else None + if event_type in ("system", "webhook"): + return ("activity_id", item.get("activity_id")) if item.get("activity_id") else None + if event_type == "workflow": + return ("workflow_run_id", row_id) + if event_type == "schedule": + # The scheduler records the run under its run id. + return ("request_id", row_id) + if event_type in ("search", "graph"): + return ("id", row_id) + return None + + +def _merge_trace_summaries(traces: list) -> dict: + """One Logs-row summary over every trace for it (a turn plus its resumes).""" + totals: dict = {} + for trace in traces: + for key, value in (trace.get("summary") or {}).items(): + if isinstance(value, (int, float)) and not isinstance(value, bool): + totals[key] = totals.get(key, 0) + value + if "retrieval_ms" in totals: + totals["retrieval_ms"] = round(totals["retrieval_ms"], 1) + return { + "count": len(traces), + "duration_ms": sum(int(t.get("duration_ms") or 0) for t in traces), + "status": traces[-1].get("status"), + "started_at": traces[0].get("started_at"), + "summary": totals, + } + + +def _attach_trace_summaries(conn, items: list, *, user_id, agent_id) -> None: + """Add a ``trace`` summary to each Logs row that has a stored trace. + + One batched lookup per link field for the whole page, rather than a join + inside the UNION, so the timeline query is unchanged. Rows without a + trace (older than the feature, or tracing disabled) get no key. + """ + refs: dict = {} + for item in items: + ref = _trace_ref(item) + if ref: + refs.setdefault(ref[0], set()).add(str(ref[1])) + if not refs: + return + found = RequestTracesRepository(conn).summaries_for_refs( + refs, user_id=user_id, agent_id=agent_id + ) + for item in items: + ref = _trace_ref(item) + if not ref: + continue + traces = found.get(ref[0], {}).get(str(ref[1])) + if traces: + item["trace"] = { + "ref": {"field": ref[0], "value": str(ref[1])}, + **_merge_trace_summaries(traces), + } + + +@analytics_ns.route("/traces") +class GetTraces(Resource): + @api.doc( + description=( + "Stored execution traces (span timelines) for one Logs row. Pass exactly " + "one of message_id, request_id, activity_id, workflow_run_id or id; " + "api_key_id scopes to an agent you own." + ), + params={ + "message_id": "Assistant message id", + "request_id": "Request id (chat turn, scheduled run)", + "activity_id": "Agent activity id (webhook and system rows)", + "workflow_run_id": "Workflow run id", + "id": "Trace id", + "api_key_id": "Agent id to scope to", + }, + ) + def get(self): + decoded_token = request.decoded_token + if not decoded_token: + return make_response(jsonify({"success": False}), 401) + user = decoded_token.get("sub") + refs = [ + (field, request.args.get(field)) + for field in TRACE_REF_FIELDS + if request.args.get(field) + ] + if len(refs) != 1: + return make_response( + jsonify( + { + "success": False, + "message": "Pass exactly one of: " + ", ".join(TRACE_REF_FIELDS), + } + ), + 400, + ) + field, value = refs[0] + api_key_id = request.args.get("api_key_id") + try: + with db_readonly() as conn: + agent, _api_key, agent_pg_id = _resolve_agent(conn, api_key_id, user) + if api_key_id and agent is None: + return make_response(jsonify({"success": True, "traces": []}), 200) + traces = RequestTracesRepository(conn).list_by_ref( + field, value, user_id=user, agent_id=agent_pg_id + ) + except Exception as err: + current_app.logger.error(f"Error getting traces: {err}", exc_info=True) + return make_response(jsonify({"success": False}), 400) + for trace in traces: + trace.pop("_id", None) + return make_response(jsonify({"success": True, "traces": traces}), 200) + + @analytics_ns.route("/get_message_analytics") class GetMessageAnalytics(Resource): get_message_analytics_model = api.model( @@ -701,7 +853,9 @@ class GetUserLogs(Resource): "event_type": fields.String( required=False, description="Filter by event source", - enum=["chat", "schedule", "webhook", "workflow", "system"], + enum=[ + "chat", "schedule", "webhook", "workflow", "system", "search", "graph", + ], ), "search": fields.String( required=False, description="Substring filter on the summary" @@ -741,6 +895,8 @@ class GetUserLogs(Resource): "webhook", "workflow", "system", + "search", + "graph", ): return make_response( jsonify({"success": False, "message": "Invalid option"}), 400 @@ -831,6 +987,7 @@ class GetUserLogs(Resource): "wr.user_id = :user_id", "wr.workflow_id = CAST(:agent_workflow_id AS uuid)", ] + trace_scope = "t.agent_id = CAST(:agent_pg_id AS uuid)" else: chat_where = ["l.user_id = :user_id"] webhook_where = [ @@ -852,6 +1009,7 @@ class GetUserLogs(Resource): "r.status IN ('success', 'completed', 'failed', 'timeout', 'skipped')", ] workflow_where = ["wr.user_id = :user_id"] + trace_scope = "t.user_id = :user_id" # One normalized timeline over five event sources. # ``payload`` carries the per-type detail; the outer query @@ -892,7 +1050,8 @@ class GetUserLogs(Resource): {summary} AS summary, jsonb_build_object( 'endpoint', s.endpoint, - 'stacks', s.stacks + 'stacks', s.stacks, + 'activity_id', s.activity_id ) AS payload FROM stack_logs s WHERE {where} @@ -913,7 +1072,8 @@ class GetUserLogs(Resource): {summary} AS summary, jsonb_build_object( 'endpoint', s.endpoint, - 'stacks', s.stacks + 'stacks', s.stacks, + 'activity_id', s.activity_id ) AS payload FROM stack_logs s WHERE {where} @@ -989,6 +1149,11 @@ class GetUserLogs(Resource): WHERE {where} """, }, + # Runs with no log row of their own are listed from their + # stored trace: searches (/api/search and MCP) and graph + # builds. + _trace_branch("search", "('search', 'mcp')", trace_scope), + _trace_branch("graph", "('graph_extraction')", trace_scope), ] if level: @@ -1062,6 +1227,8 @@ class GetUserLogs(Resource): "tool_calls": payload.get("tool_calls"), "agent_id": payload.get("agent_id"), "attachments": payload.get("attachments"), + "request_id": payload.get("request_id"), + "message_id": payload.get("message_id"), } ) elif m["event_type"] in ("system", "webhook"): @@ -1072,6 +1239,15 @@ class GetUserLogs(Resource): # before write-time redaction still carry the # reflected provider/user secrets in ``stacks``. "stacks": redact_secrets(payload.get("stacks")), + "activity_id": payload.get("activity_id"), + } + ) + elif m["event_type"] in ("search", "graph"): + item.update( + { + "status": payload.get("status"), + "source": payload.get("source"), + "duration_ms": payload.get("duration_ms"), } ) elif m["event_type"] == "workflow": @@ -1104,6 +1280,14 @@ class GetUserLogs(Resource): } ) results.append(item) + if results: + with db_readonly() as conn: + _attach_trace_summaries( + conn, + results, + user_id=user, + agent_id=agent_pg_id if api_key_id else None, + ) except Exception as err: current_app.logger.error( f"Error getting user logs: {err}", exc_info=True diff --git a/tests/api/user/test_traces_routes.py b/tests/api/user/test_traces_routes.py new file mode 100644 index 00000000..cb918493 --- /dev/null +++ b/tests/api/user/test_traces_routes.py @@ -0,0 +1,188 @@ +"""Tests for ``GET /api/traces`` and the trace summaries on ``get_user_logs`` rows.""" + +from __future__ import annotations + +import time +import uuid +from contextlib import contextmanager +from unittest.mock import patch + +import pytest +from flask import Flask + +from docsgpt.storage.db.repositories.request_traces import RequestTracesRepository + + +@pytest.fixture +def app(): + return Flask(__name__) + + +@contextmanager +def _patch_db(conn): + @contextmanager + def _yield_conn(): + yield conn + + with patch("docsgpt.api.user.analytics.routes.db_readonly", _yield_conn): + yield + + +def _trace(pg_conn, **overrides): + record = { + "id": str(uuid.uuid4()), + "request_id": "req-1", + "user_id": "owner", + "source": "stream", + "name": "stream", + "status": "ok", + "started_at_ns": time.time_ns(), + "duration_ms": 900, + "span_count": 2, + "summary": {"llm_calls": 1, "tool_calls": 1, "input_tokens": 50, "retrieval_ms": 12.5}, + "spans": [ + { + "id": "s1", + "parent_id": None, + "kind": "retrieval", + "name": "retrieval", + "status": "ok", + "offset_ms": 0, + "duration_ms": 12.5, + "attributes": {}, + "preview": {"query": "how do I deploy"}, + } + ], + } + record.update(overrides) + RequestTracesRepository(pg_conn).insert(record) + return record + + +def _get(app, pg_conn, user, query): + from docsgpt.api.user.analytics.routes import GetTraces + + with _patch_db(pg_conn), app.test_request_context("/api/traces", query_string=query): + from flask import request + + request.decoded_token = {"sub": user} if user else None + return GetTraces().get() + + +def _logs(app, pg_conn, user, body): + from docsgpt.api.user.analytics.routes import GetUserLogs + + with _patch_db(pg_conn), app.test_request_context( + "/api/get_user_logs", method="POST", json=body + ): + from flask import request + + request.decoded_token = {"sub": user} + return GetUserLogs().post() + + +class TestGetTraces: + def test_requires_auth(self, app, pg_conn): + assert _get(app, pg_conn, None, {"request_id": "x"}).status_code == 401 + + def test_requires_exactly_one_ref(self, app, pg_conn): + assert _get(app, pg_conn, "owner", {}).status_code == 400 + assert _get(app, pg_conn, "owner", {"request_id": "a", "id": "b"}).status_code == 400 + + def test_returns_own_trace_with_spans(self, app, pg_conn): + record = _trace(pg_conn) + response = _get(app, pg_conn, "owner", {"request_id": "req-1"}) + assert response.status_code == 200 + (trace,) = response.json["traces"] + assert trace["id"] == record["id"] + assert trace["spans"][0]["preview"]["query"] == "how do I deploy" + assert "_id" not in trace + + def test_other_users_get_nothing(self, app, pg_conn): + _trace(pg_conn) + response = _get(app, pg_conn, "intruder", {"request_id": "req-1"}) + assert response.json["traces"] == [] + + def test_agent_owner_sees_shared_callers_traces(self, app, pg_conn): + from docsgpt.storage.db.repositories.agents import AgentsRepository + + agent = AgentsRepository(pg_conn).create("owner", "a", "published", key="k1") + _trace(pg_conn, user_id="caller", agent_id=str(agent["id"])) + response = _get( + app, pg_conn, "owner", {"request_id": "req-1", "api_key_id": str(agent["id"])} + ) + assert len(response.json["traces"]) == 1 + + def test_unowned_agent_scope_is_empty(self, app, pg_conn): + from docsgpt.storage.db.repositories.agents import AgentsRepository + + agent = AgentsRepository(pg_conn).create("someone", "a", "published", key="k2") + _trace(pg_conn, user_id="someone", agent_id=str(agent["id"])) + response = _get( + app, pg_conn, "owner", {"request_id": "req-1", "api_key_id": str(agent["id"])} + ) + assert response.json["traces"] == [] + + +class TestLogsTraceSummaries: + def test_chat_row_gets_merged_summary_across_rounds(self, app, pg_conn): + from docsgpt.storage.db.repositories.user_logs import UserLogsRepository + + UserLogsRepository(pg_conn).insert( + user_id="owner", + endpoint="stream_answer", + data={"action": "stream_answer", "question": "q", "request_id": "req-1"}, + ) + _trace(pg_conn, status="paused") + _trace(pg_conn, started_at_ns=time.time_ns() + 1_000_000) + response = _logs(app, pg_conn, "owner", {"event_type": "chat"}) + (row,) = response.json["logs"] + assert row["request_id"] == "req-1" + trace = row["trace"] + assert trace["ref"] == {"field": "request_id", "value": "req-1"} + assert trace["count"] == 2 + assert trace["duration_ms"] == 1800 + assert trace["status"] == "ok" + assert trace["summary"]["llm_calls"] == 2 + assert trace["summary"]["retrieval_ms"] == 25.0 + + def test_row_without_trace_has_no_summary(self, app, pg_conn): + from docsgpt.storage.db.repositories.user_logs import UserLogsRepository + + UserLogsRepository(pg_conn).insert( + user_id="owner", endpoint="stream_answer", data={"question": "old row"} + ) + (row,) = _logs(app, pg_conn, "owner", {}).json["logs"] + assert "trace" not in row + + def test_webhook_row_links_by_activity_id(self, app, pg_conn): + from docsgpt.storage.db.repositories.stack_logs import StackLogsRepository + + StackLogsRepository(pg_conn).insert( + activity_id="act-1", endpoint="webhook", level="info", user_id="owner", query="{}" + ) + _trace(pg_conn, request_id="task-1", source="webhook", activity_id="act-1") + (row,) = _logs(app, pg_conn, "owner", {"event_type": "webhook"}).json["logs"] + assert row["activity_id"] == "act-1" + assert row["trace"]["ref"] == {"field": "activity_id", "value": "act-1"} + + def test_search_traces_are_listed(self, app, pg_conn): + record = _trace(pg_conn, request_id=None, source="mcp", name="mcp", status="error") + _trace(pg_conn, request_id=None, source="graph_extraction", name="graph_extraction Docs") + search_rows = _logs(app, pg_conn, "owner", {"event_type": "search"}).json["logs"] + assert len(search_rows) == 1 + row = search_rows[0] + assert row["id"] == f"search-{record['id']}" + assert row["question"] == "how do I deploy" + assert row["level"] == "error" + assert row["action"] == "mcp" + assert row["trace"]["ref"] == {"field": "id", "value": record["id"]} + graph_rows = _logs(app, pg_conn, "owner", {"event_type": "graph"}).json["logs"] + assert len(graph_rows) == 1 + + def test_search_traces_of_other_users_are_hidden(self, app, pg_conn): + _trace(pg_conn, request_id=None, source="search", user_id="someone") + assert _logs(app, pg_conn, "owner", {"event_type": "search"}).json["logs"] == [] + + def test_unknown_event_type_is_rejected(self, app, pg_conn): + assert _logs(app, pg_conn, "owner", {"event_type": "bogus"}).status_code == 400 From 3310b7fa1bfc9a0e4a760ae6dae2d6de9e66fa06 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:45:19 +0100 Subject: [PATCH 12/33] Show execution traces in the Logs UI Expanded Logs rows show the trace's duration, LLM calls, tokens, tool calls and retrieval time, and open a side panel with a waterfall of every step: retrieval, embeddings, per-source searches, model calls and tool calls, with each step's details, previews and raw attributes. Chat turns paused for approval show one waterfall per round. Searches and graph builds get their own log types. Strings are translated in every locale. --- frontend/src/api/endpoints.ts | 1 + frontend/src/api/services/userService.ts | 11 + frontend/src/conversation/AnswerFlow.tsx | 2 +- frontend/src/locale/de.json | 111 +++++++++- frontend/src/locale/en.json | 111 +++++++++- frontend/src/locale/es.json | 111 +++++++++- frontend/src/locale/jp.json | 111 +++++++++- frontend/src/locale/logsTrace.test.ts | 63 ++++++ frontend/src/locale/ru.json | 125 ++++++++++- frontend/src/locale/zh-TW.json | 111 +++++++++- frontend/src/locale/zh.json | 111 +++++++++- frontend/src/settings/Logs.tsx | 66 +++++- frontend/src/settings/traces/TraceChips.tsx | 93 ++++++++ frontend/src/settings/traces/TraceSheet.tsx | 148 +++++++++++++ .../src/settings/traces/TraceSpanDetails.tsx | 202 ++++++++++++++++++ .../src/settings/traces/TraceWaterfall.tsx | 187 ++++++++++++++++ .../src/settings/traces/traceUtils.test.ts | 159 ++++++++++++++ frontend/src/settings/traces/traceUtils.ts | 140 ++++++++++++ frontend/src/settings/types/index.ts | 77 ++++++- 19 files changed, 1921 insertions(+), 19 deletions(-) create mode 100644 frontend/src/locale/logsTrace.test.ts create mode 100644 frontend/src/settings/traces/TraceChips.tsx create mode 100644 frontend/src/settings/traces/TraceSheet.tsx create mode 100644 frontend/src/settings/traces/TraceSpanDetails.tsx create mode 100644 frontend/src/settings/traces/TraceWaterfall.tsx create mode 100644 frontend/src/settings/traces/traceUtils.test.ts create mode 100644 frontend/src/settings/traces/traceUtils.ts diff --git a/frontend/src/api/endpoints.ts b/frontend/src/api/endpoints.ts index 28a95dbc..7989be60 100644 --- a/frontend/src/api/endpoints.ts +++ b/frontend/src/api/endpoints.ts @@ -57,6 +57,7 @@ const endpoints = { TOOL_ANALYTICS: '/api/get_tool_analytics', SCHEDULE_ANALYTICS: '/api/get_schedule_analytics', LOGS: `/api/get_user_logs`, + TRACES: (params: string) => `/api/traces?${params}`, MANAGE_SYNC: '/api/manage_sync', SYNC_SOURCE: '/api/sync_source', REINGEST_SOURCE: '/api/sources/reingest', diff --git a/frontend/src/api/services/userService.ts b/frontend/src/api/services/userService.ts index fd9d81f6..85253b01 100644 --- a/frontend/src/api/services/userService.ts +++ b/frontend/src/api/services/userService.ts @@ -109,6 +109,17 @@ const userService = { apiClient.post(endpoints.USER.SCHEDULE_ANALYTICS, data, token), getLogs: (data: any, token: string | null): Promise => apiClient.post(endpoints.USER.LOGS, data, token), + getTraces: ( + params: Record, + token: string | null, + signal?: AbortSignal, + ): Promise => + apiClient.get( + endpoints.USER.TRACES(new URLSearchParams(params).toString()), + token, + {}, + signal, + ), manageSync: (data: any, token: string | null): Promise => apiClient.post(endpoints.USER.MANAGE_SYNC, data, token), syncSource: (data: any, token: string | null): Promise => diff --git a/frontend/src/conversation/AnswerFlow.tsx b/frontend/src/conversation/AnswerFlow.tsx index c9f32892..6fa071fa 100644 --- a/frontend/src/conversation/AnswerFlow.tsx +++ b/frontend/src/conversation/AnswerFlow.tsx @@ -333,7 +333,7 @@ function InlineToolCallChip({ ); } -function ToolCallPanel({ +export function ToolCallPanel({ title, copyText, children, diff --git a/frontend/src/locale/de.json b/frontend/src/locale/de.json index 061b5480..90cdca21 100644 --- a/frontend/src/locale/de.json +++ b/frontend/src/locale/de.json @@ -411,7 +411,9 @@ "schedule": "Geplant", "webhook": "Webhook", "workflow": "Workflow", - "system": "System" + "system": "System", + "search": "Suche", + "graph": "Graph-Aufbau" }, "detail": { "agent": "Agent", @@ -431,7 +433,112 @@ "workflow": "Workflow", "activity": "Aktivität", "steps": "Schritte", - "result": "Ergebnis" + "result": "Ergebnis", + "source": "Quelle" + }, + "trace": { + "view": "Trace anzeigen", + "viewRounds_one": "Trace anzeigen ({{count}} Runde)", + "viewRounds_other": "Trace anzeigen ({{count}} Runden)", + "title": "Ausführungs-Trace", + "subtitle": "Jeder Schritt hinter diesem Eintrag – Abruf, Modellaufrufe und Tool-Aufrufe – mit seiner Dauer.", + "empty": "Für diesen Eintrag wurde kein Trace aufgezeichnet.", + "failed": "Der Trace konnte nicht geladen werden.", + "noSpans": "Für diese Ausführung wurde nichts aufgezeichnet.", + "round": "Runde {{n}}", + "step": "Schritt", + "toggleChildren": "Verschachtelte Schritte ein- oder ausblenden", + "droppedSpans_one": "{{count}} weiterer Schritt wurde nicht aufgezeichnet (Größenlimit des Traces).", + "droppedSpans_other": "{{count}} weitere Schritte wurden nicht aufgezeichnet (Größenlimit des Traces).", + "chips": { + "llmCalls_one": "{{count}} LLM-Aufruf", + "llmCalls_other": "{{count}} LLM-Aufrufe", + "tokens": "{{input}} ein / {{output}} aus", + "toolCalls_one": "{{count}} Tool-Aufruf", + "toolCalls_other": "{{count}} Tool-Aufrufe", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} Fehler", + "errors_other": "{{count}} Fehler" + }, + "headline": { + "tokens": "{{input}} → {{output}} Tok.", + "chunks_one": "{{count}} Abschnitt", + "chunks_other": "{{count}} Abschnitte", + "hits_one": "{{count}} Treffer", + "hits_other": "{{count}} Treffer", + "cached": "aus Cache" + }, + "status": { + "ok": "OK", + "error": "Fehler", + "cancelled": "Abgebrochen", + "paused": "Pausiert", + "pending": "Wartet auf Freigabe", + "denied": "Abgelehnt", + "skipped": "Übersprungen" + }, + "kinds": { + "agent": "Agent", + "llm": "LLM-Aufruf", + "tool": "Tool-Aufruf", + "retrieval": "Abruf", + "search": "Vektorsuche", + "embedding": "Embedding", + "rerank": "Neusortierung", + "guardrail": "Guardrail", + "step": "Schritt" + }, + "sources": { + "stream": "Chat", + "answer": "Antwort-API", + "v1": "OpenAI-kompatible API", + "webhook": "Webhook", + "schedule": "Geplante Ausführung", + "search": "Such-API", + "mcp": "MCP-Suche", + "graph_extraction": "Graph-Aufbau", + "headless": "Hintergrundausführung" + }, + "fields": { + "status": "Status", + "duration": "Dauer", + "provider": "Anbieter", + "model": "Modell", + "inputTokens": "Eingabe-Tokens", + "outputTokens": "Ausgabe-Tokens", + "cachedTokens": "Gecachte Tokens", + "timeToFirstToken": "Zeit bis zum ersten Token", + "cost": "Kosten", + "tokenSource": "Zweck", + "cacheHit": "Aus dem Cache", + "estimatedUsage": "Geschätzte Tokens", + "yes": "Ja", + "tool": "Tool", + "action": "Aktion", + "callId": "Aufruf-ID", + "retriever": "Retriever", + "sources": "Quellen", + "topK": "Top K", + "chunks": "Abschnitte", + "candidates": "Kandidaten", + "topScore": "Beste Bewertung", + "vectorStore": "Vektorspeicher", + "agentType": "Agententyp", + "toolCalls": "Tool-Aufrufe", + "stage": "Phase", + "checks": "Prüfungen", + "triggered": "Ausgelöst", + "blocked": "Blockiert", + "redacted": "Geschwärzt", + "nodeType": "Knotentyp", + "phase": "Abschnitt", + "query": "Anfrage", + "arguments": "Argumente", + "result": "Ergebnis", + "output": "Ausgabe", + "retrievedChunks": "Abgerufene Abschnitte", + "allAttributes": "Alle Attribute" + } } }, "teams": { diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index 665e746c..d02926ce 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -416,7 +416,9 @@ "schedule": "Scheduled", "webhook": "Webhook", "workflow": "Workflow", - "system": "System" + "system": "System", + "search": "Search", + "graph": "Graph build" }, "detail": { "agent": "Agent", @@ -436,7 +438,112 @@ "workflow": "Workflow", "activity": "Activity", "steps": "Steps", - "result": "Result" + "result": "Result", + "source": "Source" + }, + "trace": { + "view": "View trace", + "viewRounds_one": "View trace ({{count}} round)", + "viewRounds_other": "View trace ({{count}} rounds)", + "title": "Execution trace", + "subtitle": "Every step behind this entry — retrieval, model calls and tool calls — with its timing.", + "empty": "No trace was recorded for this entry.", + "failed": "Could not load the trace.", + "noSpans": "Nothing was recorded for this run.", + "round": "Round {{n}}", + "step": "Step", + "toggleChildren": "Show or hide nested steps", + "droppedSpans_one": "{{count}} more step was not recorded (trace size limit).", + "droppedSpans_other": "{{count}} more steps were not recorded (trace size limit).", + "chips": { + "llmCalls_one": "{{count}} LLM call", + "llmCalls_other": "{{count}} LLM calls", + "tokens": "{{input}} in / {{output}} out", + "toolCalls_one": "{{count}} tool call", + "toolCalls_other": "{{count}} tool calls", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} error", + "errors_other": "{{count}} errors" + }, + "headline": { + "tokens": "{{input}} → {{output}} tok", + "chunks_one": "{{count}} chunk", + "chunks_other": "{{count}} chunks", + "hits_one": "{{count}} hit", + "hits_other": "{{count}} hits", + "cached": "cached" + }, + "status": { + "ok": "OK", + "error": "Error", + "cancelled": "Cancelled", + "paused": "Paused", + "pending": "Awaiting approval", + "denied": "Denied", + "skipped": "Skipped" + }, + "kinds": { + "agent": "Agent", + "llm": "LLM call", + "tool": "Tool call", + "retrieval": "Retrieval", + "search": "Vector search", + "embedding": "Embedding", + "rerank": "Rerank", + "guardrail": "Guardrail", + "step": "Step" + }, + "sources": { + "stream": "Chat", + "answer": "Answer API", + "v1": "OpenAI-compatible API", + "webhook": "Webhook", + "schedule": "Scheduled run", + "search": "Search API", + "mcp": "MCP search", + "graph_extraction": "Graph build", + "headless": "Background run" + }, + "fields": { + "status": "Status", + "duration": "Duration", + "provider": "Provider", + "model": "Model", + "inputTokens": "Input tokens", + "outputTokens": "Output tokens", + "cachedTokens": "Cached tokens", + "timeToFirstToken": "Time to first token", + "cost": "Cost", + "tokenSource": "Purpose", + "cacheHit": "Served from cache", + "estimatedUsage": "Estimated tokens", + "yes": "Yes", + "tool": "Tool", + "action": "Action", + "callId": "Call ID", + "retriever": "Retriever", + "sources": "Sources", + "topK": "Top K", + "chunks": "Chunks", + "candidates": "Candidates", + "topScore": "Top score", + "vectorStore": "Vector store", + "agentType": "Agent type", + "toolCalls": "Tool calls", + "stage": "Stage", + "checks": "Checks", + "triggered": "Triggered", + "blocked": "Blocked", + "redacted": "Redacted", + "nodeType": "Node type", + "phase": "Phase", + "query": "Query", + "arguments": "Arguments", + "result": "Result", + "output": "Output", + "retrievedChunks": "Retrieved chunks", + "allAttributes": "All attributes" + } } }, "teams": { diff --git a/frontend/src/locale/es.json b/frontend/src/locale/es.json index 1c27a945..e1a3c307 100644 --- a/frontend/src/locale/es.json +++ b/frontend/src/locale/es.json @@ -411,7 +411,9 @@ "schedule": "Programado", "webhook": "Webhook", "workflow": "Flujo de trabajo", - "system": "Sistema" + "system": "Sistema", + "search": "Búsqueda", + "graph": "Creación de grafo" }, "detail": { "agent": "Agente", @@ -431,7 +433,112 @@ "workflow": "Flujo de trabajo", "activity": "Actividad", "steps": "Pasos", - "result": "Resultado" + "result": "Resultado", + "source": "Origen" + }, + "trace": { + "view": "Ver traza", + "viewRounds_one": "Ver traza ({{count}} ronda)", + "viewRounds_other": "Ver traza ({{count}} rondas)", + "title": "Traza de ejecución", + "subtitle": "Cada paso detrás de esta entrada (recuperación, llamadas al modelo y llamadas a herramientas) con su duración.", + "empty": "No se registró ninguna traza para esta entrada.", + "failed": "No se pudo cargar la traza.", + "noSpans": "No se registró nada en esta ejecución.", + "round": "Ronda {{n}}", + "step": "Paso", + "toggleChildren": "Mostrar u ocultar pasos anidados", + "droppedSpans_one": "{{count}} paso más no se registró (límite de tamaño de la traza).", + "droppedSpans_other": "{{count}} pasos más no se registraron (límite de tamaño de la traza).", + "chips": { + "llmCalls_one": "{{count}} llamada al LLM", + "llmCalls_other": "{{count}} llamadas al LLM", + "tokens": "{{input}} entrada / {{output}} salida", + "toolCalls_one": "{{count}} llamada a herramienta", + "toolCalls_other": "{{count}} llamadas a herramientas", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} error", + "errors_other": "{{count}} errores" + }, + "headline": { + "tokens": "{{input}} → {{output}} tok", + "chunks_one": "{{count}} fragmento", + "chunks_other": "{{count}} fragmentos", + "hits_one": "{{count}} resultado", + "hits_other": "{{count}} resultados", + "cached": "en caché" + }, + "status": { + "ok": "OK", + "error": "Error", + "cancelled": "Cancelado", + "paused": "En pausa", + "pending": "Esperando aprobación", + "denied": "Denegado", + "skipped": "Omitido" + }, + "kinds": { + "agent": "Agente", + "llm": "Llamada al LLM", + "tool": "Llamada a herramienta", + "retrieval": "Recuperación", + "search": "Búsqueda vectorial", + "embedding": "Embedding", + "rerank": "Reordenación", + "guardrail": "Guardrail", + "step": "Paso" + }, + "sources": { + "stream": "Chat", + "answer": "API de respuestas", + "v1": "API compatible con OpenAI", + "webhook": "Webhook", + "schedule": "Ejecución programada", + "search": "API de búsqueda", + "mcp": "Búsqueda MCP", + "graph_extraction": "Creación de grafo", + "headless": "Ejecución en segundo plano" + }, + "fields": { + "status": "Estado", + "duration": "Duración", + "provider": "Proveedor", + "model": "Modelo", + "inputTokens": "Tokens de entrada", + "outputTokens": "Tokens de salida", + "cachedTokens": "Tokens en caché", + "timeToFirstToken": "Tiempo hasta el primer token", + "cost": "Coste", + "tokenSource": "Propósito", + "cacheHit": "Servido desde caché", + "estimatedUsage": "Tokens estimados", + "yes": "Sí", + "tool": "Herramienta", + "action": "Acción", + "callId": "ID de llamada", + "retriever": "Recuperador", + "sources": "Fuentes", + "topK": "Top K", + "chunks": "Fragmentos", + "candidates": "Candidatos", + "topScore": "Mejor puntuación", + "vectorStore": "Almacén vectorial", + "agentType": "Tipo de agente", + "toolCalls": "Llamadas a herramientas", + "stage": "Etapa", + "checks": "Comprobaciones", + "triggered": "Activadas", + "blocked": "Bloqueado", + "redacted": "Censurado", + "nodeType": "Tipo de nodo", + "phase": "Fase", + "query": "Consulta", + "arguments": "Argumentos", + "result": "Resultado", + "output": "Salida", + "retrievedChunks": "Fragmentos recuperados", + "allAttributes": "Todos los atributos" + } } }, "teams": { diff --git a/frontend/src/locale/jp.json b/frontend/src/locale/jp.json index 537a6373..ef316ca3 100644 --- a/frontend/src/locale/jp.json +++ b/frontend/src/locale/jp.json @@ -410,7 +410,9 @@ "schedule": "スケジュール", "webhook": "Webhook", "workflow": "ワークフロー", - "system": "システム" + "system": "システム", + "search": "検索", + "graph": "グラフ構築" }, "detail": { "agent": "エージェント", @@ -430,7 +432,112 @@ "workflow": "ワークフロー", "activity": "アクティビティ", "steps": "ステップ", - "result": "結果" + "result": "結果", + "source": "ソース" + }, + "trace": { + "view": "トレースを表示", + "viewRounds_one": "トレースを表示（{{count}} ラウンド）", + "viewRounds_other": "トレースを表示（{{count}} ラウンド）", + "title": "実行トレース", + "subtitle": "このエントリの各ステップ（検索、モデル呼び出し、ツール呼び出し）とその所要時間。", + "empty": "このエントリのトレースは記録されていません。", + "failed": "トレースを読み込めませんでした。", + "noSpans": "この実行では何も記録されませんでした。", + "round": "ラウンド {{n}}", + "step": "ステップ", + "toggleChildren": "ネストされたステップの表示を切り替え", + "droppedSpans_one": "さらに {{count}} 個のステップは記録されていません（トレースのサイズ上限）。", + "droppedSpans_other": "さらに {{count}} 個のステップは記録されていません（トレースのサイズ上限）。", + "chips": { + "llmCalls_one": "LLM 呼び出し {{count}} 回", + "llmCalls_other": "LLM 呼び出し {{count}} 回", + "tokens": "入力 {{input}} / 出力 {{output}}", + "toolCalls_one": "ツール呼び出し {{count}} 回", + "toolCalls_other": "ツール呼び出し {{count}} 回", + "retrieval": "RAG {{duration}}", + "errors_one": "エラー {{count}} 件", + "errors_other": "エラー {{count}} 件" + }, + "headline": { + "tokens": "{{input}} → {{output}} トークン", + "chunks_one": "{{count}} チャンク", + "chunks_other": "{{count}} チャンク", + "hits_one": "{{count}} 件ヒット", + "hits_other": "{{count}} 件ヒット", + "cached": "キャッシュ" + }, + "status": { + "ok": "OK", + "error": "エラー", + "cancelled": "キャンセル", + "paused": "一時停止", + "pending": "承認待ち", + "denied": "拒否", + "skipped": "スキップ" + }, + "kinds": { + "agent": "エージェント", + "llm": "LLM 呼び出し", + "tool": "ツール呼び出し", + "retrieval": "検索", + "search": "ベクトル検索", + "embedding": "埋め込み", + "rerank": "再ランク付け", + "guardrail": "ガードレール", + "step": "ステップ" + }, + "sources": { + "stream": "チャット", + "answer": "回答 API", + "v1": "OpenAI 互換 API", + "webhook": "Webhook", + "schedule": "スケジュール実行", + "search": "検索 API", + "mcp": "MCP 検索", + "graph_extraction": "グラフ構築", + "headless": "バックグラウンド実行" + }, + "fields": { + "status": "ステータス", + "duration": "所要時間", + "provider": "プロバイダー", + "model": "モデル", + "inputTokens": "入力トークン", + "outputTokens": "出力トークン", + "cachedTokens": "キャッシュされたトークン", + "timeToFirstToken": "最初のトークンまでの時間", + "cost": "コスト", + "tokenSource": "用途", + "cacheHit": "キャッシュから提供", + "estimatedUsage": "推定トークン", + "yes": "はい", + "tool": "ツール", + "action": "アクション", + "callId": "呼び出し ID", + "retriever": "リトリーバー", + "sources": "ソース", + "topK": "Top K", + "chunks": "チャンク", + "candidates": "候補", + "topScore": "最高スコア", + "vectorStore": "ベクトルストア", + "agentType": "エージェントの種類", + "toolCalls": "ツール呼び出し", + "stage": "ステージ", + "checks": "チェック", + "triggered": "発動", + "blocked": "ブロック", + "redacted": "マスク済み", + "nodeType": "ノードの種類", + "phase": "フェーズ", + "query": "クエリ", + "arguments": "引数", + "result": "結果", + "output": "出力", + "retrievedChunks": "取得したチャンク", + "allAttributes": "すべての属性" + } } }, "teams": { diff --git a/frontend/src/locale/logsTrace.test.ts b/frontend/src/locale/logsTrace.test.ts new file mode 100644 index 00000000..98f36f1a --- /dev/null +++ b/frontend/src/locale/logsTrace.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from 'vitest'; + +import de from './de.json'; +import en from './en.json'; +import es from './es.json'; +import jp from './jp.json'; +import ru from './ru.json'; +import zhTW from './zh-TW.json'; +import zh from './zh.json'; + +type Tree = { [key: string]: string | Tree }; + +const PLURAL_SUFFIX = /_(zero|one|two|few|many|other)$/; + +const flatten = (tree: Tree, prefix = ''): string[] => + Object.entries(tree).flatMap(([key, value]) => + typeof value === 'string' + ? [prefix + key.replace(PLURAL_SUFFIX, '')] + : flatten(value, `${prefix}${key}.`), + ); + +const traceBlock = (locale: { settings: object }): Tree => + ((locale.settings as Tree).logs as Tree).trace as Tree; + +const keysOf = (locale: { settings: object }): string[] => + Array.from(new Set(flatten(traceBlock(locale)))).sort(); + +const LOCALES = { es, de, jp, ru, zh, zhTW }; + +describe('settings.logs.trace locale block', () => { + it.each(Object.entries(LOCALES))( + '%s has the same keys as en', + (_name, locale) => { + expect(keysOf(locale)).toEqual(keysOf(en)); + }, + ); + + it.each(Object.entries(LOCALES))( + '%s is translated, not an English copy', + (_name, locale) => { + expect(traceBlock(locale).title).not.toBe(traceBlock(en).title); + expect((traceBlock(locale).fields as Tree).duration).not.toBe( + (traceBlock(en).fields as Tree).duration, + ); + }, + ); + + it('every locale provides both plural forms i18next falls back through', () => { + [en, ...Object.values(LOCALES)].forEach((locale) => { + const chips = traceBlock(locale).chips as Tree; + expect(chips.llmCalls_one).toBeTruthy(); + expect(chips.llmCalls_other).toBeTruthy(); + }); + }); + + it('every locale labels the new log event types', () => { + [en, ...Object.values(LOCALES)].forEach((locale) => { + const types = ((locale.settings as Tree).logs as Tree).types as Tree; + expect(types.search).toBeTruthy(); + expect(types.graph).toBeTruthy(); + }); + }); +}); diff --git a/frontend/src/locale/ru.json b/frontend/src/locale/ru.json index 243954cb..bf090d6b 100644 --- a/frontend/src/locale/ru.json +++ b/frontend/src/locale/ru.json @@ -423,7 +423,9 @@ "schedule": "По расписанию", "webhook": "Вебхук", "workflow": "Рабочий процесс", - "system": "Система" + "system": "Система", + "search": "Поиск", + "graph": "Построение графа" }, "detail": { "agent": "Агент", @@ -443,7 +445,126 @@ "workflow": "Рабочий процесс", "activity": "Активность", "steps": "Шаги", - "result": "Результат" + "result": "Результат", + "source": "Источник" + }, + "trace": { + "view": "Открыть трассировку", + "viewRounds_one": "Открыть трассировку ({{count}} раунд)", + "viewRounds_few": "Открыть трассировку ({{count}} раунда)", + "viewRounds_many": "Открыть трассировку ({{count}} раундов)", + "viewRounds_other": "Открыть трассировку ({{count}} раунда)", + "title": "Трассировка выполнения", + "subtitle": "Каждый шаг за этой записью — поиск, вызовы модели и инструментов — с длительностью.", + "empty": "Для этой записи трассировка не сохранена.", + "failed": "Не удалось загрузить трассировку.", + "noSpans": "Для этого запуска ничего не записано.", + "round": "Раунд {{n}}", + "step": "Шаг", + "toggleChildren": "Показать или скрыть вложенные шаги", + "droppedSpans_one": "Ещё {{count}} шаг не записан (лимит размера трассировки).", + "droppedSpans_few": "Ещё {{count}} шага не записаны (лимит размера трассировки).", + "droppedSpans_many": "Ещё {{count}} шагов не записаны (лимит размера трассировки).", + "droppedSpans_other": "Ещё {{count}} шага не записаны (лимит размера трассировки).", + "chips": { + "llmCalls_one": "{{count}} вызов LLM", + "llmCalls_few": "{{count}} вызова LLM", + "llmCalls_many": "{{count}} вызовов LLM", + "llmCalls_other": "{{count}} вызова LLM", + "tokens": "{{input}} вход / {{output}} выход", + "toolCalls_one": "{{count}} вызов инструмента", + "toolCalls_few": "{{count}} вызова инструментов", + "toolCalls_many": "{{count}} вызовов инструментов", + "toolCalls_other": "{{count}} вызова инструментов", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} ошибка", + "errors_few": "{{count}} ошибки", + "errors_many": "{{count}} ошибок", + "errors_other": "{{count}} ошибки" + }, + "headline": { + "tokens": "{{input}} → {{output}} ток.", + "chunks_one": "{{count}} фрагмент", + "chunks_few": "{{count}} фрагмента", + "chunks_many": "{{count}} фрагментов", + "chunks_other": "{{count}} фрагмента", + "hits_one": "{{count}} совпадение", + "hits_few": "{{count}} совпадения", + "hits_many": "{{count}} совпадений", + "hits_other": "{{count}} совпадения", + "cached": "из кэша" + }, + "status": { + "ok": "ОК", + "error": "Ошибка", + "cancelled": "Отменено", + "paused": "Приостановлено", + "pending": "Ожидает подтверждения", + "denied": "Отклонено", + "skipped": "Пропущено" + }, + "kinds": { + "agent": "Агент", + "llm": "Вызов LLM", + "tool": "Вызов инструмента", + "retrieval": "Поиск", + "search": "Векторный поиск", + "embedding": "Эмбеддинг", + "rerank": "Переранжирование", + "guardrail": "Ограничитель", + "step": "Шаг" + }, + "sources": { + "stream": "Чат", + "answer": "API ответов", + "v1": "OpenAI-совместимый API", + "webhook": "Вебхук", + "schedule": "Запуск по расписанию", + "search": "API поиска", + "mcp": "Поиск MCP", + "graph_extraction": "Построение графа", + "headless": "Фоновый запуск" + }, + "fields": { + "status": "Статус", + "duration": "Длительность", + "provider": "Провайдер", + "model": "Модель", + "inputTokens": "Входные токены", + "outputTokens": "Выходные токены", + "cachedTokens": "Токены из кэша", + "timeToFirstToken": "Время до первого токена", + "cost": "Стоимость", + "tokenSource": "Назначение", + "cacheHit": "Из кэша", + "estimatedUsage": "Оценочные токены", + "yes": "Да", + "tool": "Инструмент", + "action": "Действие", + "callId": "ID вызова", + "retriever": "Ретривер", + "sources": "Источники", + "topK": "Top K", + "chunks": "Фрагменты", + "candidates": "Кандидаты", + "topScore": "Лучшая оценка", + "vectorStore": "Векторное хранилище", + "agentType": "Тип агента", + "toolCalls": "Вызовы инструментов", + "stage": "Этап", + "checks": "Проверки", + "triggered": "Сработали", + "blocked": "Заблокировано", + "redacted": "Скрыто", + "nodeType": "Тип узла", + "phase": "Фаза", + "query": "Запрос", + "arguments": "Аргументы", + "result": "Результат", + "output": "Вывод", + "retrievedChunks": "Найденные фрагменты", + "allAttributes": "Все атрибуты" + } } }, "teams": { diff --git a/frontend/src/locale/zh-TW.json b/frontend/src/locale/zh-TW.json index 3e10c564..c7184ad8 100644 --- a/frontend/src/locale/zh-TW.json +++ b/frontend/src/locale/zh-TW.json @@ -410,7 +410,9 @@ "schedule": "排程", "webhook": "Webhook", "workflow": "工作流程", - "system": "系統" + "system": "系統", + "search": "搜尋", + "graph": "圖譜建置" }, "detail": { "agent": "代理", @@ -430,7 +432,112 @@ "workflow": "工作流程", "activity": "活動", "steps": "步驟", - "result": "結果" + "result": "結果", + "source": "來源" + }, + "trace": { + "view": "檢視追蹤", + "viewRounds_one": "檢視追蹤（{{count}} 輪）", + "viewRounds_other": "檢視追蹤（{{count}} 輪）", + "title": "執行追蹤", + "subtitle": "此項目背後的每個步驟（檢索、模型呼叫和工具呼叫）及其耗時。", + "empty": "此項目沒有記錄追蹤。", + "failed": "無法載入追蹤。", + "noSpans": "此次執行沒有記錄任何內容。", + "round": "第 {{n}} 輪", + "step": "步驟", + "toggleChildren": "顯示或隱藏巢狀步驟", + "droppedSpans_one": "另有 {{count}} 個步驟未被記錄（追蹤大小限制）。", + "droppedSpans_other": "另有 {{count}} 個步驟未被記錄（追蹤大小限制）。", + "chips": { + "llmCalls_one": "{{count}} 次 LLM 呼叫", + "llmCalls_other": "{{count}} 次 LLM 呼叫", + "tokens": "輸入 {{input}} / 輸出 {{output}}", + "toolCalls_one": "{{count}} 次工具呼叫", + "toolCalls_other": "{{count}} 次工具呼叫", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} 個錯誤", + "errors_other": "{{count}} 個錯誤" + }, + "headline": { + "tokens": "{{input}} → {{output}} 詞元", + "chunks_one": "{{count}} 個片段", + "chunks_other": "{{count}} 個片段", + "hits_one": "{{count}} 個命中", + "hits_other": "{{count}} 個命中", + "cached": "快取" + }, + "status": { + "ok": "正常", + "error": "錯誤", + "cancelled": "已取消", + "paused": "已暫停", + "pending": "等待核准", + "denied": "已拒絕", + "skipped": "已略過" + }, + "kinds": { + "agent": "代理", + "llm": "LLM 呼叫", + "tool": "工具呼叫", + "retrieval": "檢索", + "search": "向量搜尋", + "embedding": "嵌入", + "rerank": "重新排序", + "guardrail": "防護機制", + "step": "步驟" + }, + "sources": { + "stream": "聊天", + "answer": "回答 API", + "v1": "OpenAI 相容 API", + "webhook": "Webhook", + "schedule": "排程執行", + "search": "搜尋 API", + "mcp": "MCP 搜尋", + "graph_extraction": "圖譜建置", + "headless": "背景執行" + }, + "fields": { + "status": "狀態", + "duration": "耗時", + "provider": "供應商", + "model": "模型", + "inputTokens": "輸入詞元", + "outputTokens": "輸出詞元", + "cachedTokens": "快取詞元", + "timeToFirstToken": "首個詞元耗時", + "cost": "費用", + "tokenSource": "用途", + "cacheHit": "來自快取", + "estimatedUsage": "估算詞元", + "yes": "是", + "tool": "工具", + "action": "動作", + "callId": "呼叫 ID", + "retriever": "檢索器", + "sources": "來源", + "topK": "Top K", + "chunks": "片段", + "candidates": "候選", + "topScore": "最高分", + "vectorStore": "向量儲存", + "agentType": "代理類型", + "toolCalls": "工具呼叫", + "stage": "階段", + "checks": "檢查項目", + "triggered": "已觸發", + "blocked": "已封鎖", + "redacted": "已遮蔽", + "nodeType": "節點類型", + "phase": "階段", + "query": "查詢", + "arguments": "參數", + "result": "結果", + "output": "輸出", + "retrievedChunks": "檢索到的片段", + "allAttributes": "所有屬性" + } } }, "teams": { diff --git a/frontend/src/locale/zh.json b/frontend/src/locale/zh.json index 201b6c4c..9db7d12f 100644 --- a/frontend/src/locale/zh.json +++ b/frontend/src/locale/zh.json @@ -410,7 +410,9 @@ "schedule": "计划", "webhook": "Webhook", "workflow": "工作流", - "system": "系统" + "system": "系统", + "search": "搜索", + "graph": "图谱构建" }, "detail": { "agent": "代理", @@ -430,7 +432,112 @@ "workflow": "工作流", "activity": "活动", "steps": "步骤", - "result": "结果" + "result": "结果", + "source": "来源" + }, + "trace": { + "view": "查看追踪", + "viewRounds_one": "查看追踪（{{count}} 轮）", + "viewRounds_other": "查看追踪（{{count}} 轮）", + "title": "执行追踪", + "subtitle": "此条目背后的每一步（检索、模型调用和工具调用）及其耗时。", + "empty": "此条目没有记录追踪。", + "failed": "无法加载追踪。", + "noSpans": "此次运行没有记录任何内容。", + "round": "第 {{n}} 轮", + "step": "步骤", + "toggleChildren": "显示或隐藏嵌套步骤", + "droppedSpans_one": "另有 {{count}} 个步骤未被记录（追踪大小限制）。", + "droppedSpans_other": "另有 {{count}} 个步骤未被记录（追踪大小限制）。", + "chips": { + "llmCalls_one": "{{count}} 次 LLM 调用", + "llmCalls_other": "{{count}} 次 LLM 调用", + "tokens": "输入 {{input}} / 输出 {{output}}", + "toolCalls_one": "{{count}} 次工具调用", + "toolCalls_other": "{{count}} 次工具调用", + "retrieval": "RAG {{duration}}", + "errors_one": "{{count}} 个错误", + "errors_other": "{{count}} 个错误" + }, + "headline": { + "tokens": "{{input}} → {{output}} 词元", + "chunks_one": "{{count}} 个片段", + "chunks_other": "{{count}} 个片段", + "hits_one": "{{count}} 个命中", + "hits_other": "{{count}} 个命中", + "cached": "缓存" + }, + "status": { + "ok": "正常", + "error": "错误", + "cancelled": "已取消", + "paused": "已暂停", + "pending": "等待批准", + "denied": "已拒绝", + "skipped": "已跳过" + }, + "kinds": { + "agent": "智能体", + "llm": "LLM 调用", + "tool": "工具调用", + "retrieval": "检索", + "search": "向量搜索", + "embedding": "嵌入", + "rerank": "重排序", + "guardrail": "防护栏", + "step": "步骤" + }, + "sources": { + "stream": "聊天", + "answer": "回答 API", + "v1": "OpenAI 兼容 API", + "webhook": "Webhook", + "schedule": "定时运行", + "search": "搜索 API", + "mcp": "MCP 搜索", + "graph_extraction": "图谱构建", + "headless": "后台运行" + }, + "fields": { + "status": "状态", + "duration": "耗时", + "provider": "提供商", + "model": "模型", + "inputTokens": "输入词元", + "outputTokens": "输出词元", + "cachedTokens": "缓存词元", + "timeToFirstToken": "首个词元耗时", + "cost": "费用", + "tokenSource": "用途", + "cacheHit": "来自缓存", + "estimatedUsage": "估算词元", + "yes": "是", + "tool": "工具", + "action": "操作", + "callId": "调用 ID", + "retriever": "检索器", + "sources": "来源", + "topK": "Top K", + "chunks": "片段", + "candidates": "候选", + "topScore": "最高分", + "vectorStore": "向量存储", + "agentType": "智能体类型", + "toolCalls": "工具调用", + "stage": "阶段", + "checks": "检查项", + "triggered": "已触发", + "blocked": "已拦截", + "redacted": "已脱敏", + "nodeType": "节点类型", + "phase": "阶段", + "query": "查询", + "arguments": "参数", + "result": "结果", + "output": "输出", + "retrievedChunks": "检索到的片段", + "allAttributes": "全部属性" + } } }, "teams": { diff --git a/frontend/src/settings/Logs.tsx b/frontend/src/settings/Logs.tsx index b5b80a9b..8f38af82 100644 --- a/frontend/src/settings/Logs.tsx +++ b/frontend/src/settings/Logs.tsx @@ -1,3 +1,4 @@ +import { Activity } from 'lucide-react'; import React, { useCallback, useEffect, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { useSelector } from 'react-redux'; @@ -16,7 +17,10 @@ import { } from '../components/ui/select'; import { useLoaderState } from '../hooks'; import { selectToken } from '../preferences/preferenceSlice'; -import { LogData } from './types'; +import TraceChips from './traces/TraceChips'; +import TraceSheet from './traces/TraceSheet'; +import { formatDurationMs } from './traces/traceUtils'; +import { LogData, TraceRef } from './types'; type LogsProps = { agentId?: string; @@ -35,6 +39,7 @@ export default function Logs({ agentId, tableHeader }: LogsProps) { const [typeFilter, setTypeFilter] = useState('all'); const [searchInput, setSearchInput] = useState(''); const [search, setSearch] = useState(''); + const [openTrace, setOpenTrace] = useState(null); const logs = Object.values(logsByPage).flat(); @@ -127,6 +132,8 @@ export default function Logs({ agentId, tableHeader }: LogsProps) { { label: t('settings.logs.types.webhook'), value: 'webhook' }, { label: t('settings.logs.types.workflow'), value: 'workflow' }, { label: t('settings.logs.types.system'), value: 'system' }, + { label: t('settings.logs.types.search'), value: 'search' }, + { label: t('settings.logs.types.graph'), value: 'graph' }, ]; return ( @@ -178,8 +185,14 @@ export default function Logs({ agentId, tableHeader }: LogsProps) { setPage={setPage} loading={loadingLogs} tableHeader={tableHeader} + onViewTrace={setOpenTrace} /> + setOpenTrace(null)} + /> ); } @@ -189,8 +202,15 @@ type LogsTableProps = { setPage: React.Dispatch>; loading: boolean; tableHeader?: string; + onViewTrace: (ref: TraceRef) => void; }; -function LogsTable({ logs, setPage, loading, tableHeader }: LogsTableProps) { +function LogsTable({ + logs, + setPage, + loading, + tableHeader, + onViewTrace, +}: LogsTableProps) { const { t } = useTranslation(); const observerRef = useRef(null); const [openLogId, setOpenLogId] = useState(null); @@ -248,6 +268,7 @@ function LogsTable({ logs, setPage, loading, tableHeader }: LogsTableProps) { log={log} isOpen={openLogId === log.id} onToggle={handleLogToggle} + onViewTrace={onViewTrace} /> ); @@ -258,6 +279,7 @@ function LogsTable({ logs, setPage, loading, tableHeader }: LogsTableProps) { log={log} isOpen={openLogId === log.id} onToggle={handleLogToggle} + onViewTrace={onViewTrace} /> ); })} @@ -271,17 +293,19 @@ function formatDuration(start?: string, end?: string): string | null { if (!start || !end) return null; const ms = new Date(end).getTime() - new Date(start).getTime(); if (isNaN(ms) || ms < 0) return null; - return ms < 1000 ? `${ms}ms` : `${(ms / 1000).toFixed(1)}s`; + return formatDurationMs(ms); } function Log({ log, isOpen, onToggle, + onViewTrace, }: { log: LogData; isOpen: boolean; onToggle: (id: string) => void; + onViewTrace: (ref: TraceRef) => void; }) { const { t } = useTranslation(); const logLevelColor = { @@ -336,6 +360,17 @@ function Log({ } else if (log.event_type === 'system' || log.event_type === 'webhook') { if (log.endpoint) detailRows.push([t('settings.logs.detail.endpoint'), log.endpoint]); + } else if (log.event_type === 'search' || log.event_type === 'graph') { + if (log.source) + detailRows.push([ + t('settings.logs.detail.source'), + t(`settings.logs.trace.sources.${log.source}`, log.source), + ]); + if (log.status) + detailRows.push([ + t('settings.logs.detail.status'), + t(`settings.logs.trace.status.${log.status}`, log.status), + ]); } const textBlocks: { label: string; text: string; isError?: boolean }[] = []; @@ -412,6 +447,11 @@ function Log({ )}

{`[${log.action}]`}

+ {log.trace && ( +

+ {formatDurationMs(log.trace.duration_ms)} +

+ )}

@@ -423,6 +463,26 @@ function Log({ {isOpen && (
+ {log.trace && ( +
+ + +
+ )} {detailRows.length > 0 && (
{detailRows.map(([label, value]) => ( diff --git a/frontend/src/settings/traces/TraceChips.tsx b/frontend/src/settings/traces/TraceChips.tsx new file mode 100644 index 00000000..e5aa32c7 --- /dev/null +++ b/frontend/src/settings/traces/TraceChips.tsx @@ -0,0 +1,93 @@ +import { + Brain, + Coins, + Search, + Timer, + TriangleAlert, + Wrench, +} from 'lucide-react'; +import React from 'react'; +import { useTranslation } from 'react-i18next'; + +import { TraceCounts } from '../types'; +import { formatDurationMs, formatTokens } from './traceUtils'; + +type TraceChipsProps = { + durationMs?: number; + counts: TraceCounts; +}; + +/** Compact stat chips for a trace: duration, LLM calls, tokens, tools, RAG time, errors. */ +export default function TraceChips({ durationMs, counts }: TraceChipsProps) { + const { t } = useTranslation(); + const chips: { + key: string; + icon: React.ReactNode; + label: string; + tone?: 'danger'; + }[] = []; + if (durationMs !== undefined) + chips.push({ + key: 'duration', + icon: , + label: formatDurationMs(durationMs), + }); + if (counts.llm_calls) + chips.push({ + key: 'llm', + icon: , + label: t('settings.logs.trace.chips.llmCalls', { + count: counts.llm_calls, + }), + }); + if (counts.input_tokens || counts.output_tokens) + chips.push({ + key: 'tokens', + icon: , + label: t('settings.logs.trace.chips.tokens', { + input: formatTokens(counts.input_tokens), + output: formatTokens(counts.output_tokens), + }), + }); + if (counts.tool_calls) + chips.push({ + key: 'tools', + icon: , + label: t('settings.logs.trace.chips.toolCalls', { + count: counts.tool_calls, + }), + }); + if (counts.retrieval_calls) + chips.push({ + key: 'retrieval', + icon: , + label: t('settings.logs.trace.chips.retrieval', { + duration: formatDurationMs(counts.retrieval_ms), + }), + }); + if (counts.errors) + chips.push({ + key: 'errors', + icon: , + label: t('settings.logs.trace.chips.errors', { count: counts.errors }), + tone: 'danger', + }); + if (!chips.length) return null; + return ( +
+ {chips.map((chip) => ( + + {chip.icon} + {chip.label} + + ))} +
+ ); +} diff --git a/frontend/src/settings/traces/TraceSheet.tsx b/frontend/src/settings/traces/TraceSheet.tsx new file mode 100644 index 00000000..027c4172 --- /dev/null +++ b/frontend/src/settings/traces/TraceSheet.tsx @@ -0,0 +1,148 @@ +import React, { useEffect, useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { useSelector } from 'react-redux'; + +import { Pill } from '../../admin/AdminUI'; +import userService from '../../api/services/userService'; +import Spinner from '../../components/Spinner'; +import { + Sheet, + SheetContent, + SheetDescription, + SheetHeader, + SheetTitle, +} from '../../components/ui/sheet'; +import { selectToken } from '../../preferences/preferenceSlice'; +import { formatDateTime } from '../../utils/dateTimeUtils'; +import { Trace, TraceRef } from '../types'; +import TraceChips from './TraceChips'; +import TraceWaterfall from './TraceWaterfall'; + +type TraceSheetProps = { + traceRef: TraceRef | null; + agentId?: string; + onClose: () => void; +}; + +const STATUS_TONE: Record = + { + ok: 'success', + error: 'danger', + paused: 'warning', + cancelled: 'muted', + }; + +/** + * Side panel showing the execution trace(s) behind one Logs row. A chat turn + * paused for tool approval has one trace per round, shown in order. + */ +export default function TraceSheet({ + traceRef, + agentId, + onClose, +}: TraceSheetProps) { + const { t } = useTranslation(); + const token = useSelector(selectToken); + const [traces, setTraces] = useState([]); + const [loading, setLoading] = useState(false); + const [failed, setFailed] = useState(false); + + useEffect(() => { + if (!traceRef) return; + const controller = new AbortController(); + setLoading(true); + setFailed(false); + setTraces([]); + const params: Record = { + [traceRef.field]: traceRef.value, + }; + if (agentId) params.api_key_id = agentId; + userService + .getTraces(params, token, controller.signal) + .then(async (response: Response) => { + if (!response.ok) throw new Error('Failed to load traces'); + const data = await response.json(); + setTraces(data.traces ?? []); + }) + .catch((error: unknown) => { + if ((error as Error)?.name === 'AbortError') return; + console.error(error); + setFailed(true); + }) + .finally(() => { + if (!controller.signal.aborted) setLoading(false); + }); + return () => controller.abort(); + }, [traceRef?.field, traceRef?.value, agentId, token]); + + return ( + !open && onClose()}> + + + {t('settings.logs.trace.title')} + + {t('settings.logs.trace.subtitle')} + + +
+ {loading && ( +
+ +
+ )} + {!loading && failed && ( +

+ {t('settings.logs.trace.failed')} +

+ )} + {!loading && !failed && traces.length === 0 && ( +

+ {t('settings.logs.trace.empty')} +

+ )} + {traces.map((trace, index) => ( +
+
+ {traces.length > 1 && ( + + {t('settings.logs.trace.round', { n: index + 1 })} + + )} + + {t( + `settings.logs.trace.sources.${trace.source}`, + trace.source, + )} + + + {formatDateTime(trace.started_at)} + + + {t( + `settings.logs.trace.status.${trace.status}`, + trace.status, + )} + +
+ + {trace.dropped_spans > 0 && ( +

+ {t('settings.logs.trace.droppedSpans', { + count: trace.dropped_spans, + })} +

+ )} + +
+ ))} +
+ + + ); +} diff --git a/frontend/src/settings/traces/TraceSpanDetails.tsx b/frontend/src/settings/traces/TraceSpanDetails.tsx new file mode 100644 index 00000000..576c42a7 --- /dev/null +++ b/frontend/src/settings/traces/TraceSpanDetails.tsx @@ -0,0 +1,202 @@ +import React from 'react'; +import { useTranslation } from 'react-i18next'; + +import { ToolCallPanel } from '../../conversation/AnswerFlow'; +import { TraceSpan } from '../types'; +import { formatDurationMs, formatTokens } from './traceUtils'; + +type Row = [string, string]; + +function str(value: unknown): string | undefined { + if (value === null || value === undefined || value === '') return undefined; + if (Array.isArray(value)) return value.map(String).join(', '); + return String(value); +} + +function jsonText(value: unknown): string { + return typeof value === 'string' ? value : JSON.stringify(value, null, 2); +} + +type ChunkPreview = { + title?: string; + source?: string; + score?: number; + text?: string; +}; + +/** Everything recorded for one span: its key facts, previews and raw attributes. */ +export default function TraceSpanDetails({ span }: { span: TraceSpan }) { + const { t } = useTranslation(); + const a = span.attributes || {}; + const f = (key: string) => t(`settings.logs.trace.fields.${key}`); + const yes = t('settings.logs.trace.fields.yes'); + const rows: Row[] = []; + const push = (label: string, value: unknown) => { + const text = str(value); + if (text !== undefined) rows.push([label, text]); + }; + + push( + f('status'), + t(`settings.logs.trace.status.${span.status}`, span.status), + ); + push(f('duration'), formatDurationMs(span.duration_ms)); + switch (span.kind) { + case 'llm': { + push(f('provider'), a['gen_ai.provider.name']); + push(f('model'), a['gen_ai.request.model']); + const input = a['gen_ai.usage.input_tokens']; + const output = a['gen_ai.usage.output_tokens']; + if (typeof input === 'number') + push(f('inputTokens'), formatTokens(input)); + if (typeof output === 'number') + push(f('outputTokens'), formatTokens(output)); + push(f('cachedTokens'), a['gen_ai.usage.cache_read.input_tokens']); + if (typeof a['docsgpt.ttft_ms'] === 'number') + push(f('timeToFirstToken'), formatDurationMs(a['docsgpt.ttft_ms'])); + if ( + typeof a['docsgpt.cost_usd'] === 'number' && + a['docsgpt.cost_usd'] > 0 + ) + push(f('cost'), `$${a['docsgpt.cost_usd'].toFixed(6)}`); + push(f('tokenSource'), a['docsgpt.token_source']); + if (a['docsgpt.cache_hit']) push(f('cacheHit'), yes); + if (a['docsgpt.usage_estimated']) push(f('estimatedUsage'), yes); + break; + } + case 'tool': + push(f('tool'), a['docsgpt.tool']); + push(f('action'), a['docsgpt.action'] ?? a['gen_ai.tool.name']); + push(f('callId'), a['gen_ai.tool.call.id']); + break; + case 'retrieval': + case 'search': + case 'embedding': + case 'rerank': + push(f('retriever'), a['docsgpt.retriever']); + push(f('sources'), a['docsgpt.source_ids'] ?? a['gen_ai.data_source.id']); + push(f('model'), a['gen_ai.request.model']); + push(f('topK'), a['docsgpt.top_k']); + push(f('chunks'), a['docsgpt.chunk_count'] ?? a['docsgpt.kept_count']); + push(f('candidates'), a['docsgpt.candidate_count']); + push(f('topScore'), a['docsgpt.top_score']); + push(f('vectorStore'), a['docsgpt.vector_store']); + break; + case 'agent': + push(f('agentType'), a['docsgpt.agent_type']); + push(f('model'), a['gen_ai.request.model']); + push(f('sources'), a['docsgpt.source_count']); + push(f('toolCalls'), a['docsgpt.tool_call_count']); + break; + case 'guardrail': + push(f('stage'), a['docsgpt.guardrail.stage']); + push(f('checks'), a['docsgpt.guardrail.checks']); + push(f('triggered'), a['docsgpt.guardrail.triggered']); + if (a['docsgpt.guardrail.blocked']) push(f('blocked'), yes); + if (a['docsgpt.guardrail.redacted']) push(f('redacted'), yes); + break; + case 'step': + push(f('nodeType'), a['docsgpt.workflow.node_type']); + push(f('phase'), a['docsgpt.research.phase']); + break; + } + + const preview = span.preview || {}; + const chunks = Array.isArray(preview.chunks) + ? (preview.chunks as ChunkPreview[]) + : []; + const otherPreviews = Object.entries(preview).filter( + ([key]) => + !['arguments', 'result', 'query', 'chunks', 'output'].includes(key), + ); + + return ( +
+
+ {rows.map(([label, value]) => ( + + {label} + {value} + + ))} +
+ {span.error && ( +
+ {span.error} +
+ )} + {preview.query !== undefined && ( + +

+ {jsonText(preview.query)} +

+ + )} + {preview.arguments !== undefined && ( + +

+ {jsonText(preview.arguments)} +

+ + )} + {preview.result !== undefined && ( + +

+ {jsonText(preview.result)} +

+ + )} + {preview.output !== undefined && ( + +

+ {jsonText(preview.output)} +

+ + )} + {chunks.length > 0 && ( + +
    + {chunks.map((chunk, index) => ( +
  1. + + {index + 1}. {chunk.title || chunk.source || '—'} + {typeof chunk.score === 'number' && ( + + {chunk.score.toFixed(3)} + + )} + + {chunk.text && ( + + {chunk.text} + + )} +
    2. + ))} +
+ + )} + {otherPreviews.map(([key, value]) => ( + +

+ {jsonText(value)} +

+ + ))} +
+ + {f('allAttributes')} + + 
+          {JSON.stringify(a, null, 2)}
+
+
+
+ ); +} diff --git a/frontend/src/settings/traces/TraceWaterfall.tsx b/frontend/src/settings/traces/TraceWaterfall.tsx new file mode 100644 index 00000000..f7f70f05 --- /dev/null +++ b/frontend/src/settings/traces/TraceWaterfall.tsx @@ -0,0 +1,187 @@ +import { + Binary, + Bot, + Brain, + ChevronRight, + Database, + ListFilter, + ListTree, + Search, + ShieldCheck, + Wrench, +} from 'lucide-react'; +import React, { useMemo, useState } from 'react'; +import { useTranslation } from 'react-i18next'; + +import { Trace, TraceSpan } from '../types'; +import TraceSpanDetails from './TraceSpanDetails'; +import { + barGeometry, + buildSpanRows, + formatDurationMs, + spanHeadline, + traceTotalMs, +} from './traceUtils'; + +const KIND_STYLE: Record = { + agent: { icon: Bot, bar: 'bg-violet-500' }, + llm: { icon: Brain, bar: 'bg-sky-500' }, + tool: { icon: Wrench, bar: 'bg-amber-500' }, + retrieval: { icon: Search, bar: 'bg-emerald-500' }, + search: { icon: Database, bar: 'bg-emerald-400' }, + embedding: { icon: Binary, bar: 'bg-teal-500' }, + rerank: { icon: ListFilter, bar: 'bg-lime-500' }, + guardrail: { icon: ShieldCheck, bar: 'bg-rose-400' }, + step: { icon: ListTree, bar: 'bg-indigo-400' }, +}; + +function barClass(span: TraceSpan): string { + if (span.status === 'error') return 'bg-red-500'; + if (['cancelled', 'pending', 'denied', 'skipped'].includes(span.status)) + return 'bg-gray-400 dark:bg-gray-500'; + return (KIND_STYLE[span.kind] ?? KIND_STYLE.step).bar; +} + +const SCALE_STEPS = [0, 0.25, 0.5, 0.75, 1]; +const INDENT_PX = 14; + +/** A waterfall of one trace's spans; click a row for its details. */ +export default function TraceWaterfall({ trace }: { trace: Trace }) { + const { t } = useTranslation(); + const rows = useMemo(() => buildSpanRows(trace.spans), [trace.spans]); + const totalMs = traceTotalMs(trace.duration_ms, trace.spans); + const [selectedId, setSelectedId] = useState(null); + const [collapsed, setCollapsed] = useState>(new Set()); + + const visibleRows = useMemo(() => { + const out: typeof rows = []; + let hiddenBelowDepth: number | null = null; + for (const row of rows) { + if (hiddenBelowDepth !== null && row.depth > hiddenBelowDepth) continue; + hiddenBelowDepth = collapsed.has(row.span.id) ? row.depth : null; + out.push(row); + } + return out; + }, [rows, collapsed]); + + const toggleCollapsed = (id: string) => + setCollapsed((prev) => { + const next = new Set(prev); + if (next.has(id)) next.delete(id); + else next.add(id); + return next; + }); + + if (!rows.length) { + return ( +

+ {t('settings.logs.trace.noSpans')} +

+ ); + } + + return ( +
+
+ + {t('settings.logs.trace.step')} + +
+ {SCALE_STEPS.map((step) => ( + + {formatDurationMs(totalMs * step)} + + ))} +
+
+ {visibleRows.map(({ span, depth, hasChildren }) => { + const Icon = (KIND_STYLE[span.kind] ?? KIND_STYLE.step).icon; + const { left, width } = barGeometry(span, totalMs); + const selected = selectedId === span.id; + const headline = spanHeadline(span, t); + return ( +
+
setSelectedId(selected ? null : span.id)} + onKeyDown={(e) => { + if (e.key === 'Enter' || e.key === ' ') { + e.preventDefault(); + setSelectedId(selected ? null : span.id); + } + }} + className={`grid w-full cursor-pointer grid-cols-[minmax(0,5fr)_minmax(0,4fr)] items-center gap-3 rounded-md py-1 text-left ${ + selected + ? 'bg-muted dark:bg-white/10' + : 'hover:bg-muted/60 dark:hover:bg-white/5' + }`} + > + + {hasChildren ? ( + + ) : ( + + )} + + + {span.name} + + {headline && ( + + {headline} + + )} + + {formatDurationMs(span.duration_ms)} + + + + + +
+ {selected && ( +
+ +
+ )} +
+ ); + })} +
+ ); +} diff --git a/frontend/src/settings/traces/traceUtils.test.ts b/frontend/src/settings/traces/traceUtils.test.ts new file mode 100644 index 00000000..1a07167e --- /dev/null +++ b/frontend/src/settings/traces/traceUtils.test.ts @@ -0,0 +1,159 @@ +import { describe, expect, it } from 'vitest'; + +import { TraceSpan } from '../types'; +import { + barGeometry, + buildSpanRows, + formatDurationMs, + formatTokens, + spanHeadline, + traceTotalMs, +} from './traceUtils'; + +const span = (overrides: Partial): TraceSpan => ({ + id: 's', + parent_id: null, + kind: 'llm', + name: 'chat m', + status: 'ok', + offset_ms: 0, + duration_ms: 10, + attributes: {}, + ...overrides, +}); + +describe('buildSpanRows', () => { + it('orders depth-first by start offset with depths', () => { + const rows = buildSpanRows([ + span({ id: 'llm2', parent_id: 'agent', offset_ms: 50 }), + span({ id: 'agent', kind: 'agent', offset_ms: 5, duration_ms: 100 }), + span({ id: 'retrieval', kind: 'retrieval', offset_ms: 0 }), + span({ id: 'tool', kind: 'tool', parent_id: 'agent', offset_ms: 20 }), + span({ id: 'inner', parent_id: 'tool', offset_ms: 21 }), + ]); + expect(rows.map((r) => [r.span.id, r.depth])).toEqual([ + ['retrieval', 0], + ['agent', 0], + ['tool', 1], + ['inner', 2], + ['llm2', 1], + ]); + }); + + it('treats spans whose parent is missing as roots', () => { + const rows = buildSpanRows([span({ id: 'a', parent_id: 'gone' })]); + expect(rows).toHaveLength(1); + expect(rows[0].depth).toBe(0); + }); + + it('marks rows that have children', () => { + const rows = buildSpanRows([ + span({ id: 'p', kind: 'agent' }), + span({ id: 'c', parent_id: 'p' }), + ]); + expect(rows[0].hasChildren).toBe(true); + expect(rows[1].hasChildren).toBe(false); + }); +}); + +describe('traceTotalMs', () => { + it('uses the larger of the stored duration and the last span end', () => { + expect(traceTotalMs(100, [span({ offset_ms: 90, duration_ms: 30 })])).toBe( + 120, + ); + expect(traceTotalMs(500, [span({ offset_ms: 0, duration_ms: 30 })])).toBe( + 500, + ); + }); + + it('never returns zero', () => { + expect(traceTotalMs(0, [])).toBeGreaterThan(0); + }); +}); + +describe('barGeometry', () => { + it('positions a span as percentages of the total', () => { + expect(barGeometry(span({ offset_ms: 25, duration_ms: 50 }), 100)).toEqual({ + left: 25, + width: 50, + }); + }); + + it('keeps tiny spans visible and inside the track', () => { + const tiny = barGeometry(span({ offset_ms: 100, duration_ms: 0 }), 100); + expect(tiny.width).toBeGreaterThan(0); + expect(tiny.left + tiny.width).toBeLessThanOrEqual(100); + }); +}); + +describe('formatDurationMs', () => { + it.each([ + [0.4, '<1 ms'], + [12.4, '12 ms'], + [999, '999 ms'], + [1234, '1.23 s'], + [15432, '15.4 s'], + [65000, '1m 05s'], + ])('%s -> %s', (ms, expected) => { + expect(formatDurationMs(ms)).toBe(expected); + }); + + it('shows a dash for missing values', () => { + expect(formatDurationMs(undefined)).toBe('—'); + }); +}); + +describe('formatTokens', () => { + it('abbreviates large counts', () => { + expect(formatTokens(950)).toBe('950'); + expect(formatTokens(12345)).toBe('12.3k'); + expect(formatTokens(2_500_000)).toBe('2.5M'); + }); +}); + +const HEADLINES: Record = { + 'settings.logs.trace.headline.tokens': '{{input}} → {{output}} tok', + 'settings.logs.trace.headline.chunks': '{{count}} chunks', + 'settings.logs.trace.headline.hits': '{{count}} hits', + 'settings.logs.trace.headline.cached': 'cached', +}; + +const t = (key: string, options: Record = {}) => + (HEADLINES[key] ?? key).replace(/\{\{(\w+)\}\}/g, (_m, name) => + String(options[name]), + ); + +describe('spanHeadline', () => { + it('summarises an llm span by tokens', () => { + expect( + spanHeadline( + span({ + attributes: { + 'gen_ai.usage.input_tokens': 1200, + 'gen_ai.usage.output_tokens': 80, + }, + }), + t, + ), + ).toBe('1.2k → 80 tok'); + }); + + it('summarises a retrieval span by chunk count', () => { + expect( + spanHeadline( + span({ kind: 'retrieval', attributes: { 'docsgpt.chunk_count': 4 } }), + t, + ), + ).toBe('4 chunks'); + }); + + it('is empty when nothing is known', () => { + expect(spanHeadline(span({ kind: 'step' }), t)).toBe(''); + }); + + it('labels a cached llm call', () => { + expect( + spanHeadline(span({ attributes: { 'docsgpt.cache_hit': true } }), t), + ).toBe('cached'); + }); +}); diff --git a/frontend/src/settings/traces/traceUtils.ts b/frontend/src/settings/traces/traceUtils.ts new file mode 100644 index 00000000..3f748328 --- /dev/null +++ b/frontend/src/settings/traces/traceUtils.ts @@ -0,0 +1,140 @@ +import { TraceSpan } from '../types'; + +export type SpanRow = { + span: TraceSpan; + depth: number; + hasChildren: boolean; +}; + +/** + * Flatten a trace's spans into display rows: depth-first, children in start + * order under their parent. A span whose parent is missing (dropped by the + * span cap) is shown at the top level rather than hidden. + */ +export function buildSpanRows(spans: TraceSpan[]): SpanRow[] { + const ids = new Set(spans.map((s) => s.id)); + const children = new Map(); + for (const span of spans) { + const parent = + span.parent_id && ids.has(span.parent_id) ? span.parent_id : null; + const list = children.get(parent) ?? []; + list.push(span); + children.set(parent, list); + } + for (const list of children.values()) { + list.sort((a, b) => a.offset_ms - b.offset_ms); + } + const rows: SpanRow[] = []; + const visit = (parent: string | null, depth: number) => { + for (const span of children.get(parent) ?? []) { + const kids = children.get(span.id); + rows.push({ span, depth, hasChildren: Boolean(kids?.length) }); + visit(span.id, depth + 1); + } + }; + visit(null, 0); + return rows; +} + +/** Width of the timeline: the stored duration, or the last span end if later. */ +export function traceTotalMs( + durationMs: number | undefined, + spans: TraceSpan[], +): number { + const lastEnd = spans.reduce( + (max, s) => Math.max(max, (s.offset_ms || 0) + (s.duration_ms || 0)), + 0, + ); + return Math.max(durationMs || 0, lastEnd, 1); +} + +const MIN_BAR_PERCENT = 0.6; + +/** A span's bar as left offset and width, in percent of the timeline. */ +export function barGeometry( + span: TraceSpan, + totalMs: number, +): { left: number; width: number } { + const total = totalMs > 0 ? totalMs : 1; + const width = Math.max( + ((span.duration_ms || 0) / total) * 100, + MIN_BAR_PERCENT, + ); + const left = Math.min( + Math.max(((span.offset_ms || 0) / total) * 100, 0), + 100 - width, + ); + return { left: round(left), width: round(width) }; +} + +function round(n: number): number { + return Math.round(n * 100) / 100; +} + +/** Human duration: `<1 ms`, `12 ms`, `1.23 s`, `15.4 s`, `1m 05s`. */ +export function formatDurationMs(ms: number | null | undefined): string { + if (ms === null || ms === undefined || Number.isNaN(ms)) return '—'; + if (ms < 1) return '<1 ms'; + if (ms < 1000) return `${Math.round(ms)} ms`; + if (ms < 10_000) return `${(ms / 1000).toFixed(2)} s`; + if (ms < 60_000) return `${(ms / 1000).toFixed(1)} s`; + const minutes = Math.floor(ms / 60_000); + const seconds = Math.round((ms % 60_000) / 1000); + return `${minutes}m ${String(seconds).padStart(2, '0')}s`; +} + +/** Token counts abbreviated: `950`, `12.3k`, `2.5M`. */ +export function formatTokens(n: number | null | undefined): string { + const value = n ?? 0; + if (value >= 1_000_000) return `${trimZero((value / 1_000_000).toFixed(1))}M`; + if (value >= 1000) return `${trimZero((value / 1000).toFixed(1))}k`; + return String(Math.round(value)); +} + +function trimZero(s: string): string { + return s.endsWith('.0') ? s.slice(0, -2) : s; +} + +function num(value: unknown): number | undefined { + return typeof value === 'number' && Number.isFinite(value) + ? value + : undefined; +} + +type Translate = (key: string, options?: Record) => string; + +/** + * A short, kind-specific summary shown next to a span's name, e.g. + * `1.2k → 80 tok` for an LLM call or `4 chunks` for a retrieval. + */ +export function spanHeadline(span: TraceSpan, t: Translate): string { + const a = span.attributes || {}; + switch (span.kind) { + case 'llm': { + if (a['docsgpt.cache_hit']) + return t('settings.logs.trace.headline.cached'); + const input = num(a['gen_ai.usage.input_tokens']); + const output = num(a['gen_ai.usage.output_tokens']); + if (input === undefined && output === undefined) return ''; + return t('settings.logs.trace.headline.tokens', { + input: formatTokens(input), + output: formatTokens(output), + }); + } + case 'retrieval': + case 'rerank': { + const chunks = num(a['docsgpt.chunk_count'] ?? a['docsgpt.kept_count']); + return chunks === undefined + ? '' + : t('settings.logs.trace.headline.chunks', { count: chunks }); + } + case 'search': { + const hits = num(a['docsgpt.candidate_count']); + return hits === undefined + ? '' + : t('settings.logs.trace.headline.hits', { count: hits }); + } + default: + return ''; + } +} diff --git a/frontend/src/settings/types/index.ts b/frontend/src/settings/types/index.ts index 81e557a8..76006b84 100644 --- a/frontend/src/settings/types/index.ts +++ b/frontend/src/settings/types/index.ts @@ -17,7 +17,74 @@ export type ChunkType = { }; export type LogEventType = - 'chat' | 'schedule' | 'webhook' | 'workflow' | 'system'; + 'chat' | 'schedule' | 'webhook' | 'workflow' | 'system' | 'search' | 'graph'; + +/** Counts rolled up from a trace's spans; every field may be absent. */ +export type TraceCounts = { + llm_calls?: number; + tool_calls?: number; + retrieval_calls?: number; + retrieval_ms?: number; + input_tokens?: number; + output_tokens?: number; + errors?: number; +}; + +/** The id a Logs row's traces are looked up by. */ +export type TraceRef = { + field: 'id' | 'request_id' | 'message_id' | 'activity_id' | 'workflow_run_id'; + value: string; +}; + +/** Trace summary attached to a Logs row (merged over its rounds). */ +export type LogTraceSummary = { + ref: TraceRef; + count: number; + duration_ms: number; + status?: string; + started_at?: string; + summary: TraceCounts; +}; + +export type TraceSpanKind = + | 'agent' + | 'llm' + | 'tool' + | 'retrieval' + | 'search' + | 'embedding' + | 'rerank' + | 'guardrail' + | 'step'; + +export type TraceSpan = { + id: string; + parent_id: string | null; + kind: TraceSpanKind | string; + name: string; + status: string; + offset_ms: number; + duration_ms: number; + attributes: Record; + preview?: Record; + error?: string; +}; + +export type Trace = { + id: string; + request_id?: string | null; + message_id?: string | null; + source: string; + name?: string | null; + status: string; + started_at: string; + duration_ms: number; + span_count: number; + dropped_spans: number; + summary: TraceCounts; + spans: TraceSpan[]; + otel_trace_id?: string | null; +}; export type LogData = { id: string; @@ -33,9 +100,17 @@ export type LogData = { tool_calls?: Record[]; agent_id?: string; attachments?: string[]; + request_id?: string; + message_id?: string; // system + webhook events (stack_logs) endpoint?: string; stacks?: Record[]; + activity_id?: string; + // search + graph events (request_traces) + source?: string; + duration_ms?: number; + // Present when the row has a stored execution trace. + trace?: LogTraceSummary; // workflow events (workflow_runs) workflow_name?: string; result?: Record; From 4e1002f5e0e00fa48b90cbec67c08ff121f4c98d Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:46:18 +0100 Subject: [PATCH 13/33] Document execution traces and the GenAI OpenTelemetry export Covers what is recorded, where it is shown, the TRACES_* settings, the gen_ai.* spans and metrics, content capture, and the limits of exporting spans when the request ends. --- docs/content/Deploying/Observability.mdx | 87 +++++++++++++++++++++++- 1 file changed, 85 insertions(+), 2 deletions(-) diff --git a/docs/content/Deploying/Observability.mdx b/docs/content/Deploying/Observability.mdx index a0d8afaf..4643fcd0 100644 --- a/docs/content/Deploying/Observability.mdx +++ b/docs/content/Deploying/Observability.mdx @@ -14,8 +14,9 @@ launch command with `opentelemetry-instrument` and setting OTLP env vars. Auto-instrumentation covers Flask, Starlette, Celery, SQLAlchemy, -psycopg, Redis, requests, and Python logging. LLM/retriever calls are -not captured at this layer — see *Going further* below. +psycopg, Redis, requests, and Python logging. Agent runs, LLM calls, +tool calls and retrieval are recorded by DocsGPT itself and exported as +OpenTelemetry GenAI spans — see [Execution traces](#execution-traces). ## Enabling @@ -67,6 +68,88 @@ dotenv run -- opentelemetry-instrument celery -A docsgpt.app.celery worker -l IN the OTEL log handler. Without it, `logging` writes only to stdout. +## Execution traces + +Every request records an **execution trace**: a timed tree of the steps +behind it. Traces are recorded for chat turns (`/stream`, `/api/answer`, +`/v1/chat/completions`, including each round of a tool-approval pause), +scheduled and webhook runs, workflows, the research agent, `/api/search`, +the MCP `search_docs` tool, and graph builds. + +| Step | Recorded when | +| --- | --- | +| `invoke_agent` | An agent (or a workflow node's agent) runs | +| `chat` | An LLM call made during the request, including retries, fallbacks, query rephrasing, prescreening, history compression and guardrail judges | +| `execute_tool` | A tool is executed, paused for approval, denied or skipped | +| `retrieval` | A retriever or the multi-source dispatcher searches | +| `embeddings` | The query is embedded | +| `search` | One source is searched | +| `rerank` | Prescreening filters retrieved chunks | +| `guardrail` | A guardrail calls a remote check or fires | +| `step` | A workflow node or research phase runs | + +Traces are stored in the `request_traces` table and shown in the app: open +**Settings → Logs** (or an agent's **Logs** tab), expand an entry and choose +**View trace** to see a waterfall of every step with its timing, tokens, +cost and details. + +Stored traces keep short previews — tool arguments and results, retrieved +chunk titles and snippets, rephrased queries, answer excerpts — truncated +and with secret-named fields redacted. Full prompts are never stored. When a +guardrail fires during a request, every preview is dropped from its trace. + +```bash +TRACES_ENABLED=true # record traces at all +TRACES_CAPTURE_CONTENT=true # keep previews in stored traces +TRACES_PREVIEW_CHARS=2000 # characters kept per preview +TRACES_MAX_SPANS=500 # steps kept per trace; the rest are counted +TRACES_RETENTION_DAYS=30 # a daily task deletes older traces +TRACES_OTEL_EXPORT=true # also export traces as OTel GenAI spans +``` + +### GenAI spans and metrics + +When DocsGPT runs under `opentelemetry-instrument`, each finished trace is +also exported as spans that follow the +[OpenTelemetry GenAI semantic conventions](https://opentelemetry.io/docs/specs/semconv/gen-ai/): +`invoke_agent {agent}`, `chat {model}`, `execute_tool {tool}`, +`embeddings {model}` and `retrieval`, with attributes such as +`gen_ai.provider.name`, `gen_ai.request.model`, +`gen_ai.usage.input_tokens`, `gen_ai.usage.output_tokens`, +`gen_ai.usage.cache_read.input_tokens`, `gen_ai.conversation.id`, +`gen_ai.agent.id` and `gen_ai.tool.name`. DocsGPT-specific details use the +`docsgpt.*` prefix (`docsgpt.request_id`, `docsgpt.token_source`, +`docsgpt.cache_hit`, `docsgpt.ttft_ms`, ...). The trace's root span is a +child of the request's HTTP server span, and the stored trace keeps the +OTel trace id so you can move between the two. + +Two metrics are recorded for every model call: +`gen_ai.client.token.usage` and `gen_ai.client.operation.duration`. + +Backends that understand the GenAI conventions — Langfuse +(`/api/public/otel`), Arize Phoenix, Datadog LLM Observability, Grafana — +render these as LLM traces with token and cost views. + +Prompt and tool content is **not** exported by default, because the OTLP +backend may be a third party. Opt in with the standard variable: + +```bash +OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT=SPAN_ONLY +``` + +This adds the same redacted previews the app stores (for example +`gen_ai.tool.call.arguments` and `gen_ai.tool.call.result`). + + + GenAI spans are exported when the request finishes, with their original + timestamps. Consequences: a long research run appears only when it ends; + HTTP and database spans made during a step sit beside the step rather than + under it; and log records carry the request's span ids, not the step's. + + +The GenAI conventions are still in development upstream, so attribute names +may change in later releases. + ## Backend examples ### Axiom From 35126857281fd609f8463f6bdb9152af706edc7a Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:50:19 +0100 Subject: [PATCH 14/33] Keep complete_stream's signature and classify /api/traces for tokens The trace lifecycle moves into a decorator so complete_stream keeps its real parameters for introspection. GET /api/traces is readable with the analytics:read scope, like the Logs endpoint it complements. --- docsgpt/api/answer/routes/base.py | 108 +++++++++++++++--------------- docsgpt/api/pat/rules.py | 1 + 2 files changed, 56 insertions(+), 53 deletions(-) diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index 019ab261..469271b7 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -1,4 +1,5 @@ import datetime +import functools import inspect import json import logging @@ -68,6 +69,50 @@ class StreamSuperseded(Exception): answer_ns = Namespace("answer", description="Answer related operations", path="/") +def _traced_stream(method): + """Record ``complete_stream`` as the request's execution trace. + + The stream runs in the SSE pump thread, not the request thread, so the + trace the route started is activated here for the whole stream and written + exactly once when the stream ends -- normally, paused, failed or + abandoned. The request id is resolved here too so the trace and the + stream agree on it. + """ + signature = inspect.signature(method) + + @functools.wraps(method) + def wrapper(*args, **kwargs): + bound = signature.bind(*args, **kwargs) + arguments = bound.arguments + continuation = arguments.get("_continuation") + request_id = ( + continuation.get("request_id") if continuation else None + ) or arguments.get("request_id") or str(uuid.uuid4()) + arguments["request_id"] = request_id + trace = arguments.get("trace") + if not isinstance(trace, tracing.Trace): + agent = arguments.get("agent") + trace = tracing.start_trace( + source=str(getattr(agent, "endpoint", None) or "stream"), + capture_otel_context=False, + ) + arguments["trace"] = trace + if trace is not None: + decoded_token = arguments.get("decoded_token") + trace.bind( + request_id=request_id, + user_id=decoded_token.get("sub") if isinstance(decoded_token, dict) else None, + agent_id=arguments.get("agent_id"), + ) + with tracing.activate(trace): + try: + yield from method(*bound.args, **bound.kwargs) + finally: + tracing.flush(trace) + + return wrapper + + class BaseAnswerResource: """Shared base class for answer endpoints""" @@ -246,60 +291,8 @@ class BaseAnswerResource: logger.exception("Failed to release resume claim after a usage refusal") return error + @_traced_stream def complete_stream( - self, - *args: Any, - request_id: Optional[str] = None, - trace: Optional["tracing.Trace"] = None, - **kwargs: Any, - ) -> Generator[str, None, None]: - """Stream the conversation response and record its execution trace. - - Accepts every argument of :meth:`_complete_stream_body`. The trace is - activated for the whole stream -- which runs in the SSE pump thread, - not the request thread -- and written exactly once when the stream - ends, however it ends. - - Args: - request_id: The request's id, generated by the route before the - agent was built so pre-fetch retrieval shares it. A - continuation's saved id takes precedence; absent both, a new - one is minted. - trace: The trace the route started (it already holds the - pre-fetch retrieval). A new one is started when omitted. - - Yields: - Server-sent event strings. - """ - bound = inspect.signature(self._complete_stream_body).bind(*args, **kwargs) - continuation = bound.arguments.get("_continuation") - request_id = ( - continuation.get("request_id") if continuation else None - ) or request_id or str(uuid.uuid4()) - decoded_token = bound.arguments.get("decoded_token") or {} - if not isinstance(trace, tracing.Trace): - trace = None - if trace is None: - agent = bound.arguments.get("agent") - trace = tracing.start_trace( - source=str(getattr(agent, "endpoint", None) or "stream"), - capture_otel_context=False, - ) - if trace is not None: - trace.bind( - request_id=request_id, - user_id=decoded_token.get("sub") if isinstance(decoded_token, dict) else None, - agent_id=bound.arguments.get("agent_id"), - ) - with tracing.activate(trace): - try: - yield from self._complete_stream_body( - *bound.args, request_id=request_id, **bound.kwargs - ) - finally: - tracing.flush(trace) - - def _complete_stream_body( self, question: str, agent: Any, @@ -319,6 +312,7 @@ class BaseAnswerResource: _continuation: Optional[Dict] = None, finalize_tool_pause_as_complete: bool = False, request_id: Optional[str] = None, + trace: Optional["tracing.Trace"] = None, ) -> Generator[str, None, None]: """ Generator function that streams the complete conversation response. @@ -342,6 +336,14 @@ class BaseAnswerResource: shared_token: Token for shared agent model_id: Model ID used for the request retrieved_docs: Pre-fetched documents for sources (optional) + request_id: The request's id, minted by the route before the + agent was built so pre-fetch retrieval shares it. A + continuation's saved id takes precedence; absent both, a new + one is minted. + trace: The execution trace the route started (it already holds + the pre-fetch retrieval); :func:`_traced_stream` activates it + and writes it when the stream ends. A new one is started when + omitted. finalize_tool_pause_as_complete: Stateless-tool-round mode for the OpenAI-compatible ``/v1/chat/completions`` endpoint. OpenAI clients resume a tool call by re-POSTing the full diff --git a/docsgpt/api/pat/rules.py b/docsgpt/api/pat/rules.py index 227fe11a..c1d7683d 100644 --- a/docsgpt/api/pat/rules.py +++ b/docsgpt/api/pat/rules.py @@ -298,6 +298,7 @@ RULES: dict[tuple[str, str], Rule] = { ("/api/get_tool_analytics", "POST"): _rule("analytics:read", blocked_by=_ALL_FAMILIES), ("/api/get_schedule_analytics", "POST"): _rule("analytics:read", blocked_by=_ALL_FAMILIES), ("/api/get_user_logs", "POST"): _rule("analytics:read", blocked_by=_ALL_FAMILIES), + ("/api/traces", "GET"): _rule("analytics:read", blocked_by=_ALL_FAMILIES), # Teams (read only) ("/api/teams", "GET"): _rule("teams:read"), ("/api/teams/", "GET"): _rule("teams:read"), From 285fdc914c3322c96bef85c7fa5fdf4124b7f4ca Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:55:30 +0100 Subject: [PATCH 15/33] Add an e2e spec for execution traces; tidy the timeline scale A tool-calling chat turn and a RAG turn are traced, linked from their Logs rows and opened in the trace panel. The timeline scale starts at 0 and its labels no longer wrap. --- .../src/settings/traces/TraceWaterfall.tsx | 4 +- .../e2e/specs/tier-a/execution-traces.spec.ts | 181 ++++++++++++++++++ 2 files changed, 183 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/specs/tier-a/execution-traces.spec.ts diff --git a/frontend/src/settings/traces/TraceWaterfall.tsx b/frontend/src/settings/traces/TraceWaterfall.tsx index f7f70f05..fba28239 100644 --- a/frontend/src/settings/traces/TraceWaterfall.tsx +++ b/frontend/src/settings/traces/TraceWaterfall.tsx @@ -90,10 +90,10 @@ export default function TraceWaterfall({ trace }: { trace: Trace }) { {SCALE_STEPS.map((step) => ( - {formatDurationMs(totalMs * step)} + {step === 0 ? '0' : formatDurationMs(totalMs * step)} ))}
diff --git a/tests/e2e/specs/tier-a/execution-traces.spec.ts b/tests/e2e/specs/tier-a/execution-traces.spec.ts new file mode 100644 index 00000000..2f1f4653 --- /dev/null +++ b/tests/e2e/specs/tier-a/execution-traces.spec.ts @@ -0,0 +1,181 @@ +/** + * Execution traces — a chat turn's steps are stored and shown in the Logs UI. + * + * One `/stream` turn with a tool call writes a `request_traces` row holding + * the agent run, every LLM call and the tool call. The Logs row links to it + * through `request_id` (stamped into `user_logs.data`), carries a trace + * summary, and "View trace" opens the waterfall. + * + * // Silent-break covered: the Logs row loses its trace link. If the + * request id stops reaching `user_logs` (or the stored trace), the row + * silently renders without the "View trace" button and nothing errors. + */ + +import * as playwright from '@playwright/test'; +const { expect, test } = playwright; + +import { dirname, resolve } from 'node:path'; +import { fileURLToPath } from 'node:url'; + +import { authedRequest } from '../../helpers/api.js'; +import { newUserContext } from '../../helpers/auth.js'; +import { pg } from '../../helpers/db.js'; +import { resetDb } from '../../helpers/reset.js'; +import { + multipartContext, + postUpload, + waitForTask, +} from '../../helpers/uploads.js'; + +const HERE = dirname(fileURLToPath(import.meta.url)); +const README_MD = resolve(HERE, '..', '..', 'fixtures', 'docs', 'readme.md'); + +/** `memory_view` ships in DEFAULT_CHAT_TOOLS, so every user can run it. */ +const TOOL_DIRECTIVE = '[[MOCK_LLM_TOOLCALL:memory_view:once]]'; + +interface TraceRow { + source: string; + status: string; + request_id: string | null; + message_id: string | null; + // pg parses JSONB to JS values + // eslint-disable-next-line @typescript-eslint/no-explicit-any + spans: any[]; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + summary: any; +} + +async function tracesFor(userId: string): Promise { + const { rows } = await pg.query( + `SELECT source, status, request_id, message_id::text AS message_id, spans, summary + FROM request_traces WHERE user_id = $1 ORDER BY started_at`, + [userId], + ); + return rows; +} + +test.describe('tier-a · execution traces', () => { + test.beforeEach(async () => { + await resetDb(); + }); + + test('a chat turn with a tool call is traced and opens from the Logs page', async ({ + browser, + }) => { + const { context, sub, token } = await newUserContext(browser); + const api = await authedRequest(playwright, token); + try { + const res = await api.post('/stream', { + data: { + question: `what do you remember about me? ${TOOL_DIRECTIVE}`, + save_conversation: true, + }, + timeout: 180_000, + }); + expect(res.ok()).toBeTruthy(); + expect(await res.text()).toContain('"type": "end"'); + + // Stored trace: one row for the turn, linked to its message. + const traces = await tracesFor(sub); + expect(traces).toHaveLength(1); + const [trace] = traces; + expect(trace.source).toBe('stream'); + expect(trace.request_id).toBeTruthy(); + expect(trace.message_id).toBeTruthy(); + const kinds = new Set(trace.spans.map((s) => s.kind)); + expect(kinds.has('agent')).toBe(true); + expect(kinds.has('llm')).toBe(true); + expect(kinds.has('tool')).toBe(true); + expect(trace.summary.llm_calls).toBeGreaterThan(0); + const tool = trace.spans.find((s) => s.kind === 'tool'); + expect(tool.attributes['gen_ai.tool.name']).toBe('memory_view'); + + // The Logs API links the row to the trace. + const logsRes = await api.post('/api/get_user_logs', { + data: { page: 1, page_size: 10, event_type: 'chat' }, + }); + const logs = (await logsRes.json()).logs; + expect(logs[0].request_id).toBe(trace.request_id); + expect(logs[0].trace.ref).toEqual({ + field: 'request_id', + value: trace.request_id, + }); + + // And the UI opens it. + const page = await context.newPage(); + await page.goto('/settings/logs'); + await page.getByText('what do you remember about me?').first().click(); + await page.getByRole('button', { name: 'View trace' }).click(); + const sheet = page.getByRole('dialog'); + await expect(sheet.getByText('Execution trace')).toBeVisible(); + await expect(sheet.getByText('execute_tool memory_view').first()).toBeVisible(); + await sheet.getByText('execute_tool memory_view').first().click(); + await expect(sheet.getByText('Arguments')).toBeVisible(); + if (process.env.TRACE_SCREENSHOT) { + await page.screenshot({ path: process.env.TRACE_SCREENSHOT, fullPage: false }); + } + } finally { + await api.dispose(); + await context.close(); + } + }); + + test('a RAG turn records the retrieval, its query embedding and each source search', async ({ + browser, + }) => { + const { context, sub, token } = await newUserContext(browser); + const api = await authedRequest(playwright, token); + const multi = await multipartContext(token); + try { + const taskId = await postUpload(multi, README_MD, { + user: sub, + name: 'trace-docs', + mimeType: 'text/markdown', + }); + expect((await waitForTask(api, taskId, 120_000)).status).toBe('SUCCESS'); + const { rows } = await pg.query<{ id: string }>( + `SELECT id::text AS id FROM sources WHERE user_id = $1 ORDER BY date DESC LIMIT 1`, + [sub], + ); + const sourceId = rows[0].id; + + const res = await api.post('/stream', { + data: { + question: 'how do I get started?', + active_docs: sourceId, + save_conversation: true, + }, + timeout: 180_000, + }); + expect(await res.text()).toContain('"type": "end"'); + + const [trace] = await tracesFor(sub); + const byKind = (kind: string) => trace.spans.filter((s) => s.kind === kind); + const retrieval = byKind('retrieval'); + expect(retrieval.length).toBeGreaterThan(0); + expect(byKind('embedding').length).toBeGreaterThan(0); + const searches = byKind('search'); + expect(searches.map((s) => s.attributes['gen_ai.data_source.id'])).toContain(sourceId); + // Retrieval ran before the agent answered, inside the same trace. + const agent = byKind('agent')[0]; + expect(retrieval[0].offset_ms).toBeLessThan(agent.offset_ms); + expect(trace.summary.retrieval_calls).toBeGreaterThan(0); + + const page = await context.newPage(); + await page.goto('/settings/logs'); + await page.getByText('how do I get started?').first().click(); + await page.getByRole('button', { name: 'View trace' }).click(); + const sheet = page.getByRole('dialog'); + await expect(sheet.getByText('Execution trace')).toBeVisible(); + await sheet.getByText(/^retrieval/).first().click(); + await expect(sheet.getByText('Retrieved chunks')).toBeVisible(); + if (process.env.TRACE_SCREENSHOT_RAG) { + await page.screenshot({ path: process.env.TRACE_SCREENSHOT_RAG, fullPage: false }); + } + } finally { + await multi.dispose(); + await api.dispose(); + await context.close(); + } + }); +}); From efec1a2f831eadf9db18c2c7366af9ddd3b74ae7 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 18:11:19 +0100 Subject: [PATCH 16/33] Keep tool output out of span errors and fix trace status and listing Tool results, denial comments and tool exception text now reach a span only as a capture-gated preview; span.error is a fixed message, since it is stored and exported regardless of content settings. A turn whose stream yields an error event is recorded as failed. Search traces are listed by a query copied into the small summary column, so the Logs timeline never reads the spans JSONB. Adds GraphRAG span tests. --- docsgpt/agents/tool_executor.py | 18 +++++++++---- docsgpt/api/answer/routes/base.py | 7 +++++ docsgpt/api/user/analytics/routes.py | 7 ++--- docsgpt/tracing/core.py | 10 ++++++- tests/api/answer/test_trace_wiring.py | 11 ++++++++ tests/api/user/test_traces_routes.py | 9 ++++++- tests/retriever/test_graph_rag.py | 38 +++++++++++++++++++++++++++ tests/tracing/test_tool_spans.py | 28 +++++++++++++++++--- tests/tracing/test_trace.py | 19 ++++++++++++++ 9 files changed, 134 insertions(+), 13 deletions(-) diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 155ea925..2c7b446d 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -81,11 +81,15 @@ def finish_tool_span(span: Any, data: Dict[str, Any]) -> None: if data.get("result") is not None: span.preview("result", data.get("result")) span_status = _TOOL_SPAN_STATUS.get(status, tracing.STATUS_OK) + # ``error`` and ``result`` are tool output or user text (a denial + # comment), so they travel only as previews, which honour the content + # capture settings; ``span.error`` is exported and stored unconditionally. + if data.get("error"): + span.preview("error", data["error"]) if span_status == tracing.STATUS_ERROR: - span.error = str(data.get("error") or data.get("result") or "tool error")[:500] - span.set(**{"error.type": str(data.get("error_type") or "ToolError")}) - elif data.get("error"): - span.error = str(data["error"])[:500] + error_type = str(data.get("error_type") or "ToolError") + span.error = f"Tool call failed ({error_type})" + span.set(**{"error.type": error_type}) span.end(span_status) @@ -1021,7 +1025,11 @@ class ToolExecutor: try: outcome = yield from self._execute(tools_dict, call, llm_class_name) except Exception as exc: - span.end(error=exc) + # A tool's exception text can quote its response; like tool + # output it goes only into the capture-gated preview. + span.preview("error", str(exc)) + span.error = f"Tool call failed ({type(exc).__name__})" + span.end(tracing.STATUS_ERROR, attributes={"error.type": type(exc).__name__}) raise except GeneratorExit: span.end(tracing.STATUS_CANCELLED) diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index 469271b7..f23b82f3 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -860,6 +860,13 @@ class BaseAnswerResource: # error silently — the exact shape of the bug being fixed here. if stream_error: query_metadata.setdefault("error", stream_error) + # A yielded error (e.g. a failed workflow node) ends the + # generator normally, so no span raised; the user still saw + # the turn fail, and its trace should say so. A pause below + # overrides this, as it does for the message row. + trace = tracing.current_trace() + if trace is not None: + trace.outcome = tracing.STATUS_ERROR # ---- Paused: save continuation state and end stream early ---- if paused: diff --git a/docsgpt/api/user/analytics/routes.py b/docsgpt/api/user/analytics/routes.py index cb4486e5..570c141e 100644 --- a/docsgpt/api/user/analytics/routes.py +++ b/docsgpt/api/user/analytics/routes.py @@ -99,9 +99,10 @@ def _trace_branch(name: str, sources_sql: str, scope: str) -> dict: return { "name": name, "level": "CASE WHEN t.status = 'error' THEN 'error' ELSE 'info' END", - # A search's query sits in its retrieval span's preview; graph builds - # are named after their source. - "summary": "COALESCE(t.spans->0->'preview'->>'query', t.name, t.source)", + # A search is listed by its query (copied into the small ``summary`` + # at flush so this never detoasts ``spans``); graph builds are named + # after their source. + "summary": "COALESCE(t.summary->>'query', t.name, t.source)", "where": [f"t.source IN {sources_sql}", scope], "sql": f""" SELECT '{name}' AS event_type, diff --git a/docsgpt/tracing/core.py b/docsgpt/tracing/core.py index 47a2c49a..aa0f3b07 100644 --- a/docsgpt/tracing/core.py +++ b/docsgpt/tracing/core.py @@ -431,6 +431,14 @@ class Trace: def to_record(self) -> Dict[str, Any]: """The ``request_traces`` row for this (finished) trace.""" + summary = self.summary() + # Traces without a log row of their own (searches) are listed by + # their query; keep it in the small summary so listing never reads + # ``spans``. Dropped with the other previews when content is blocked. + if not self.content_blocked and self.spans: + query = self.spans[0].previews.get("query") + if isinstance(query, str): + summary["query"] = query spans = [] for s in self.spans: entry: Dict[str, Any] = { @@ -464,7 +472,7 @@ class Trace: "duration_ms": int(round(self.duration_ms or 0)), "span_count": len(spans), "dropped_spans": self.dropped_spans, - "summary": self.summary(), + "summary": summary, "spans": spans, "otel_trace_id": self.otel_trace_id, } diff --git a/tests/api/answer/test_trace_wiring.py b/tests/api/answer/test_trace_wiring.py index 49fcd6e0..9153c424 100644 --- a/tests/api/answer/test_trace_wiring.py +++ b/tests/api/answer/test_trace_wiring.py @@ -132,6 +132,17 @@ class TestTraceLifecycle: assert any('"type": "error"' in s for s in stream) assert flushed[0].status == "error" + def test_yielded_error_marks_trace_error(self, flask_app, mock_mongo_db): + """A failed workflow node yields an error event instead of raising.""" + from docsgpt.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(), _captured_flushes() as flushed: + _run( + BaseAnswerResource(), + _agent([{"type": "error", "error": "node failed"}]), + ) + assert flushed[0].status == "error" + def test_abandoned_stream_still_flushes(self, flask_app, mock_mongo_db): from docsgpt.api.answer.routes.base import BaseAnswerResource diff --git a/tests/api/user/test_traces_routes.py b/tests/api/user/test_traces_routes.py index cb918493..a0b2bbbe 100644 --- a/tests/api/user/test_traces_routes.py +++ b/tests/api/user/test_traces_routes.py @@ -167,7 +167,14 @@ class TestLogsTraceSummaries: assert row["trace"]["ref"] == {"field": "activity_id", "value": "act-1"} def test_search_traces_are_listed(self, app, pg_conn): - record = _trace(pg_conn, request_id=None, source="mcp", name="mcp", status="error") + record = _trace( + pg_conn, + request_id=None, + source="mcp", + name="mcp", + status="error", + summary={"llm_calls": 0, "query": "how do I deploy"}, + ) _trace(pg_conn, request_id=None, source="graph_extraction", name="graph_extraction Docs") search_rows = _logs(app, pg_conn, "owner", {"event_type": "search"}).json["logs"] assert len(search_rows) == 1 diff --git a/tests/retriever/test_graph_rag.py b/tests/retriever/test_graph_rag.py index 5203eee1..802bfc99 100644 --- a/tests/retriever/test_graph_rag.py +++ b/tests/retriever/test_graph_rag.py @@ -1021,3 +1021,41 @@ class TestPersonalizedPageRankWithoutScipy: docs = rag._get_data() assert [doc["text"] for doc in docs] == ["near", "far"] + + +class TestTraceSpans: + """GraphRAG searches and query embeddings are recorded in the execution trace.""" + + @pytest.fixture(autouse=True) + def _trace(self, monkeypatch): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True) + self.trace = tracing.start_trace(source="stream", capture_otel_context=False) + with tracing.activate(self.trace): + yield + + def test_search_is_a_retrieval_span(self, _patch_llm_creator): + rag = _make_retriever() + docs = [{"text": "alpha", "title": "Doc A", "source": "a.md"}] + with patch.object(rag, "_get_data", return_value=docs): + assert rag.search("new question") == docs + (span,) = [s for s in self.trace.spans if s.kind == "retrieval"] + assert span.name == "retrieval GraphRAGRetriever" + assert span.attributes["docsgpt.retriever"] == "GraphRAGRetriever" + assert span.attributes["gen_ai.data_source.id"] == "src1" + assert span.attributes["docsgpt.chunk_count"] == 1 + assert span.previews["chunks"][0]["title"] == "Doc A" + + def test_embed_query_is_an_embedding_span(self): + embedder = Mock() + embedder.embed_query.return_value = [0.5] + with patch( + "docsgpt.retriever.graph_rag.get_embeddings", return_value=embedder + ): + assert GraphRAGRetriever._embed_query(object(), "q") == [0.5] + (span,) = self.trace.spans + assert span.kind == "embedding" + assert span.attributes["gen_ai.operation.name"] == "embeddings" diff --git a/tests/tracing/test_tool_spans.py b/tests/tracing/test_tool_spans.py index 6f2447db..899b6a8b 100644 --- a/tests/tracing/test_tool_spans.py +++ b/tests/tracing/test_tool_spans.py @@ -82,7 +82,9 @@ class TestExecute: _drain(executor.execute({}, ToolCall(id="c", name="x", arguments="{}"), "L")) span = trace.spans[0] assert span.status == "error" - assert span.error == "no such tool" + # Tool output stays in the capture-gated preview, never in ``error``. + assert span.error == "Tool call failed (ToolError)" + assert span.previews["result"] == "no such tool" def test_raised_error(self, trace): def inner(executor, call): @@ -94,6 +96,8 @@ class TestExecute: _drain(executor.execute({}, ToolCall(id="c", name="x", arguments="{}"), "L")) assert trace.spans[0].status == "error" assert trace.spans[0].attributes["error.type"] == "RuntimeError" + assert trace.spans[0].error == "Tool call failed (RuntimeError)" + assert trace.spans[0].previews["error"] == "api down" def test_nested_retrieval_parents_to_tool(self, trace): def inner(executor, call): @@ -109,6 +113,22 @@ class TestExecute: assert retrieval.parent_id == tool.id +def test_tool_content_is_not_kept_when_capture_is_off(trace, monkeypatch): + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", False) + + def inner(executor, call): + executor.tool_calls.append( + {"tool_name": "t", "result": "secret tool output", "status": "error"} + ) + return "", call.id + yield # pragma: no cover + + _drain(_executor_with(inner).execute({}, ToolCall(id="c", name="x", arguments="{}"), "L")) + trace.finish() + record = trace.to_record() + assert "secret tool output" not in str(record) + + class _Handler(LLMHandler): def parse_response(self, response): # pragma: no cover - unused raise NotImplementedError @@ -169,7 +189,8 @@ class TestPausedCalls: _drain(_Handler().handle_tool_calls(agent, [call], {"0": {"name": "telegram"}}, [])) (span,) = trace.spans assert span.status == "denied" - assert span.error == "not allowed" + assert span.error is None + assert span.previews["error"] == "not allowed" class TestContinuation: @@ -227,7 +248,8 @@ class TestContinuation: assert agent_span.status == "ok" assert tool_span.parent_id == agent_span.id assert tool_span.status == "denied" - assert tool_span.error == "too risky" + assert tool_span.error is None + assert tool_span.previews["error"] == "too risky" def test_client_result_is_recorded(self, trace): agent = self._agent() diff --git a/tests/tracing/test_trace.py b/tests/tracing/test_trace.py index 03a4caca..437f1b53 100644 --- a/tests/tracing/test_trace.py +++ b/tests/tracing/test_trace.py @@ -353,3 +353,22 @@ class TestOutcome: trace.outcome = "paused" trace.finish(status="error") assert trace.status == "error" + + +class TestRecordQuery: + def test_first_span_query_is_copied_into_summary(self): + trace = tracing.start_trace(source="search") + with tracing.activate(trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval") as s: + s.preview("query", "how do I deploy") + trace.finish() + assert trace.to_record()["summary"]["query"] == "how do I deploy" + + def test_query_is_dropped_when_content_is_blocked(self): + trace = tracing.start_trace(source="search") + with tracing.activate(trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval") as s: + s.preview("query", "leaky") + tracing.mark_content_blocked() + trace.finish() + assert "query" not in trace.to_record()["summary"] From 8157bcf3bfdeb24ba62ab49c70caee01aef85644 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 18:11:19 +0100 Subject: [PATCH 17/33] Make Logs rows keyboard-operable and fix minute rounding Rows open with Enter or Space and expose aria-expanded, so the View trace action is reachable without a mouse. Durations round to whole seconds before splitting, so 119.6 s reads 2m 00s instead of 1m 60s. --- frontend/src/settings/Logs.tsx | 9 +++++++++ frontend/src/settings/traces/traceUtils.test.ts | 1 + frontend/src/settings/traces/traceUtils.ts | 6 ++++-- 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/frontend/src/settings/Logs.tsx b/frontend/src/settings/Logs.tsx index 8f38af82..aec4c42f 100644 --- a/frontend/src/settings/Logs.tsx +++ b/frontend/src/settings/Logs.tsx @@ -429,7 +429,16 @@ function Log({ return (
onToggle(log.id)} + onKeyDown={(e) => { + if (e.key === 'Enter' || e.key === ' ') { + e.preventDefault(); + onToggle(log.id); + } + }} className={`text-foreground flex cursor-pointer flex-row items-start gap-2 p-2 px-4 py-3 ${ isOpen ? 'dark:bg-background rounded-t-xl bg-[#F1F1F1]' : '' }`} diff --git a/frontend/src/settings/traces/traceUtils.test.ts b/frontend/src/settings/traces/traceUtils.test.ts index 1a07167e..106c3dc8 100644 --- a/frontend/src/settings/traces/traceUtils.test.ts +++ b/frontend/src/settings/traces/traceUtils.test.ts @@ -94,6 +94,7 @@ describe('formatDurationMs', () => { [1234, '1.23 s'], [15432, '15.4 s'], [65000, '1m 05s'], + [119_600, '2m 00s'], ])('%s -> %s', (ms, expected) => { expect(formatDurationMs(ms)).toBe(expected); }); diff --git a/frontend/src/settings/traces/traceUtils.ts b/frontend/src/settings/traces/traceUtils.ts index 3f748328..e077299c 100644 --- a/frontend/src/settings/traces/traceUtils.ts +++ b/frontend/src/settings/traces/traceUtils.ts @@ -78,8 +78,10 @@ export function formatDurationMs(ms: number | null | undefined): string { if (ms < 1000) return `${Math.round(ms)} ms`; if (ms < 10_000) return `${(ms / 1000).toFixed(2)} s`; if (ms < 60_000) return `${(ms / 1000).toFixed(1)} s`; - const minutes = Math.floor(ms / 60_000); - const seconds = Math.round((ms % 60_000) / 1000); + // Round to whole seconds before splitting, so 119.6 s reads 2m 00s, not 1m 60s. + const totalSeconds = Math.round(ms / 1000); + const minutes = Math.floor(totalSeconds / 60); + const seconds = totalSeconds % 60; return `${minutes}m ${String(seconds).padStart(2, '0')}s`; } From d7319928d18354eaf70ab7ea451e79496f9afd5a Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:53:18 +0100 Subject: [PATCH 18/33] Build the trace UI on theme tokens and existing components Colours come from theme tokens only: status pills and chips use muted, primary and destructive; waterfall bars use chart-1..5, one per group of span kinds, with destructive for failures and muted-foreground for cancelled or pending steps. Rows select through ghost Buttons, the error box is a destructive Alert and View trace is an outline Button. Bar geometry and indentation move to CSS custom properties, scale ticks to static positions, and the sheet takes layout classes only. --- frontend/src/settings/Logs.tsx | 13 +- frontend/src/settings/traces/TraceChips.tsx | 4 +- frontend/src/settings/traces/TraceSheet.tsx | 33 +-- .../src/settings/traces/TraceSpanDetails.tsx | 9 +- .../src/settings/traces/TraceWaterfall.tsx | 195 ++++++++++-------- 5 files changed, 141 insertions(+), 113 deletions(-) diff --git a/frontend/src/settings/Logs.tsx b/frontend/src/settings/Logs.tsx index aec4c42f..52d7e08e 100644 --- a/frontend/src/settings/Logs.tsx +++ b/frontend/src/settings/Logs.tsx @@ -7,6 +7,7 @@ import userService from '../api/services/userService'; import ChevronRight from '../assets/chevron-right.svg'; import CopyButton from '../components/CopyButton'; import SkeletonLoader from '../components/SkeletonLoader'; +import { Button } from '../components/ui/button'; import { Input } from '../components/ui/input'; import { Select, @@ -439,7 +440,7 @@ function Log({ onToggle(log.id); } }} - className={`text-foreground flex cursor-pointer flex-row items-start gap-2 p-2 px-4 py-3 ${ + className={`text-foreground focus-visible:ring-ring/50 flex cursor-pointer flex-row items-start gap-2 p-2 px-4 py-3 outline-none focus-visible:ring-3 focus-visible:ring-inset ${ isOpen ? 'dark:bg-background rounded-t-xl bg-[#F1F1F1]' : '' }`} > @@ -478,18 +479,18 @@ function Log({ durationMs={log.trace.duration_ms} counts={log.trace.summary} /> - +
)} {detailRows.length > 0 && ( diff --git a/frontend/src/settings/traces/TraceChips.tsx b/frontend/src/settings/traces/TraceChips.tsx index e5aa32c7..d1240a57 100644 --- a/frontend/src/settings/traces/TraceChips.tsx +++ b/frontend/src/settings/traces/TraceChips.tsx @@ -80,8 +80,8 @@ export default function TraceChips({ durationMs, counts }: TraceChipsProps) { key={chip.key} className={`inline-flex items-center gap-1 rounded-full px-2 py-0.5 text-xs tabular-nums ${ chip.tone === 'danger' - ? 'bg-red-100 text-red-700 dark:bg-red-900/40 dark:text-red-300' - : 'bg-muted text-foreground dark:bg-white/10' + ? 'bg-destructive/10 text-destructive' + : 'bg-muted text-foreground' }`} > {chip.icon} diff --git a/frontend/src/settings/traces/TraceSheet.tsx b/frontend/src/settings/traces/TraceSheet.tsx index 027c4172..5a847c96 100644 --- a/frontend/src/settings/traces/TraceSheet.tsx +++ b/frontend/src/settings/traces/TraceSheet.tsx @@ -2,7 +2,6 @@ import React, { useEffect, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { useSelector } from 'react-redux'; -import { Pill } from '../../admin/AdminUI'; import userService from '../../api/services/userService'; import Spinner from '../../components/Spinner'; import { @@ -24,13 +23,13 @@ type TraceSheetProps = { onClose: () => void; }; -const STATUS_TONE: Record = - { - ok: 'success', - error: 'danger', - paused: 'warning', - cancelled: 'muted', - }; +/** Status pill tones. The theme has no success or warning token yet. */ +const STATUS_TONE: Record = { + ok: 'bg-muted text-foreground', + error: 'bg-destructive/10 text-destructive', + paused: 'bg-primary/10 text-primary', + cancelled: 'bg-muted text-muted-foreground', +}; /** * Side panel showing the execution trace(s) behind one Logs row. A chat turn @@ -79,22 +78,22 @@ export default function TraceSheet({ !open && onClose()}> - + {t('settings.logs.trace.title')} {t('settings.logs.trace.subtitle')} -
+
{loading && (
)} {!loading && failed && ( -

+

{t('settings.logs.trace.failed')}

)} @@ -120,19 +119,23 @@ export default function TraceSheet({ {formatDateTime(trace.started_at)} - + {t( `settings.logs.trace.status.${trace.status}`, trace.status, )} - +
{trace.dropped_spans > 0 && ( -

+

{t('settings.logs.trace.droppedSpans', { count: trace.dropped_spans, })} diff --git a/frontend/src/settings/traces/TraceSpanDetails.tsx b/frontend/src/settings/traces/TraceSpanDetails.tsx index 576c42a7..09554057 100644 --- a/frontend/src/settings/traces/TraceSpanDetails.tsx +++ b/frontend/src/settings/traces/TraceSpanDetails.tsx @@ -1,6 +1,7 @@ import React from 'react'; import { useTranslation } from 'react-i18next'; +import { Alert, AlertDescription } from '../../components/ui/alert'; import { ToolCallPanel } from '../../conversation/AnswerFlow'; import { TraceSpan } from '../types'; import { formatDurationMs, formatTokens } from './traceUtils'; @@ -121,9 +122,11 @@ export default function TraceSpanDetails({ span }: { span: TraceSpan }) { ))}

{span.error && ( -
- {span.error} -
+ + + {span.error} + + )} {preview.query !== undefined && ( diff --git a/frontend/src/settings/traces/TraceWaterfall.tsx b/frontend/src/settings/traces/TraceWaterfall.tsx index fba28239..30957719 100644 --- a/frontend/src/settings/traces/TraceWaterfall.tsx +++ b/frontend/src/settings/traces/TraceWaterfall.tsx @@ -13,6 +13,7 @@ import { import React, { useMemo, useState } from 'react'; import { useTranslation } from 'react-i18next'; +import { Button } from '../../components/ui/button'; import { Trace, TraceSpan } from '../types'; import TraceSpanDetails from './TraceSpanDetails'; import { @@ -23,29 +24,48 @@ import { traceTotalMs, } from './traceUtils'; +/** + * Icon per span kind, and one of the five chart series per group of kinds: + * the agent and its steps, model calls, tool calls, retrieval (with its + * embedding, per-source search and rerank), and guardrails. + */ const KIND_STYLE: Record = { - agent: { icon: Bot, bar: 'bg-violet-500' }, - llm: { icon: Brain, bar: 'bg-sky-500' }, - tool: { icon: Wrench, bar: 'bg-amber-500' }, - retrieval: { icon: Search, bar: 'bg-emerald-500' }, - search: { icon: Database, bar: 'bg-emerald-400' }, - embedding: { icon: Binary, bar: 'bg-teal-500' }, - rerank: { icon: ListFilter, bar: 'bg-lime-500' }, - guardrail: { icon: ShieldCheck, bar: 'bg-rose-400' }, - step: { icon: ListTree, bar: 'bg-indigo-400' }, + agent: { icon: Bot, bar: 'bg-chart-1' }, + step: { icon: ListTree, bar: 'bg-chart-1' }, + llm: { icon: Brain, bar: 'bg-chart-2' }, + tool: { icon: Wrench, bar: 'bg-chart-3' }, + retrieval: { icon: Search, bar: 'bg-chart-4' }, + search: { icon: Database, bar: 'bg-chart-4' }, + embedding: { icon: Binary, bar: 'bg-chart-4' }, + rerank: { icon: ListFilter, bar: 'bg-chart-4' }, + guardrail: { icon: ShieldCheck, bar: 'bg-chart-5' }, }; -function barClass(span: TraceSpan): string { - if (span.status === 'error') return 'bg-red-500'; - if (['cancelled', 'pending', 'denied', 'skipped'].includes(span.status)) - return 'bg-gray-400 dark:bg-gray-500'; - return (KIND_STYLE[span.kind] ?? KIND_STYLE.step).bar; +const INACTIVE_STATUSES = ['cancelled', 'pending', 'denied', 'skipped']; + +function kindStyle(span: TraceSpan) { + return KIND_STYLE[span.kind] ?? KIND_STYLE.step; } -const SCALE_STEPS = [0, 0.25, 0.5, 0.75, 1]; -const INDENT_PX = 14; +function barClass(span: TraceSpan): string { + if (span.status === 'error') return 'bg-destructive'; + if (INACTIVE_STATUSES.includes(span.status)) return 'bg-muted-foreground/40'; + return kindStyle(span).bar; +} -/** A waterfall of one trace's spans; click a row for its details. */ +/** Scale ticks: 0, a quarter, half, three quarters and the full duration. */ +const SCALE_TICKS = [ + { at: 0, className: 'left-0' }, + { at: 0.25, className: 'left-1/4 -translate-x-1/2' }, + { at: 0.5, className: 'left-1/2 -translate-x-1/2' }, + { at: 0.75, className: 'left-3/4 -translate-x-1/2' }, + { at: 1, className: 'right-0' }, +]; + +/** Indent per nesting level, in rem. */ +const INDENT_REM = 0.875; + +/** A waterfall of one trace's spans; select a row for its details. */ export default function TraceWaterfall({ trace }: { trace: Trace }) { const { t } = useTranslation(); const rows = useMemo(() => buildSpanRows(trace.spans), [trace.spans]); @@ -72,6 +92,9 @@ export default function TraceWaterfall({ trace }: { trace: Trace }) { return next; }); + const toggleSelected = (id: string) => + setSelectedId((current) => (current === id ? null : id)); + if (!rows.length) { return (

@@ -81,101 +104,99 @@ export default function TraceWaterfall({ trace }: { trace: Trace }) { } return ( -

-
- - {t('settings.logs.trace.step')} - -
- {SCALE_STEPS.map((step) => ( +
+
+ {t('settings.logs.trace.step')} +
+ {SCALE_TICKS.map((tick) => ( - {step === 0 ? '0' : formatDurationMs(totalMs * step)} + {tick.at === 0 ? '0' : formatDurationMs(totalMs * tick.at)} ))}
{visibleRows.map(({ span, depth, hasChildren }) => { - const Icon = (KIND_STYLE[span.kind] ?? KIND_STYLE.step).icon; + const Icon = kindStyle(span).icon; const { left, width } = barGeometry(span, totalMs); const selected = selectedId === span.id; const headline = spanHeadline(span, t); + const isCollapsed = collapsed.has(span.id); + const geometry = { + '--trace-indent': `${depth * INDENT_REM}rem`, + '--trace-bar-left': `${left}%`, + '--trace-bar-width': `${width}%`, + } as React.CSSProperties; return ( -
+
setSelectedId(selected ? null : span.id)} - onKeyDown={(e) => { - if (e.key === 'Enter' || e.key === ' ') { - e.preventDefault(); - setSelectedId(selected ? null : span.id); - } - }} - className={`grid w-full cursor-pointer grid-cols-[minmax(0,5fr)_minmax(0,4fr)] items-center gap-3 rounded-md py-1 text-left ${ - selected - ? 'bg-muted dark:bg-white/10' - : 'hover:bg-muted/60 dark:hover:bg-white/5' - }`} + className={`grid grid-cols-9 items-center gap-3 rounded-md ${selected ? 'bg-accent' : ''}`} > - +
{hasChildren ? ( - + ) : ( - + )} - - - {span.name} - - {headline && ( - - {headline} + +
+ {/* Mouse shortcut to the row's button; keyboard users use the button. */} +
{selected && ( -
+
)} From 56e8437a83d33f77ab8eca605eb8f84a4127cf97 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:06:49 +0100 Subject: [PATCH 19/33] Drop the visible heading from the trace drawer The drawer opens straight onto the trace. Its title stays as the dialog's screen-reader name, and the unused subtitle string is removed from every locale. --- frontend/src/locale/de.json | 1 - frontend/src/locale/en.json | 1 - frontend/src/locale/es.json | 1 - frontend/src/locale/jp.json | 1 - frontend/src/locale/ru.json | 1 - frontend/src/locale/zh-TW.json | 1 - frontend/src/locale/zh.json | 1 - frontend/src/settings/traces/TraceSheet.tsx | 23 +++++++------------ .../e2e/specs/tier-a/execution-traces.spec.ts | 8 +++---- 9 files changed, 12 insertions(+), 26 deletions(-) diff --git a/frontend/src/locale/de.json b/frontend/src/locale/de.json index 90cdca21..2c1bceab 100644 --- a/frontend/src/locale/de.json +++ b/frontend/src/locale/de.json @@ -441,7 +441,6 @@ "viewRounds_one": "Trace anzeigen ({{count}} Runde)", "viewRounds_other": "Trace anzeigen ({{count}} Runden)", "title": "Ausführungs-Trace", - "subtitle": "Jeder Schritt hinter diesem Eintrag – Abruf, Modellaufrufe und Tool-Aufrufe – mit seiner Dauer.", "empty": "Für diesen Eintrag wurde kein Trace aufgezeichnet.", "failed": "Der Trace konnte nicht geladen werden.", "noSpans": "Für diese Ausführung wurde nichts aufgezeichnet.", diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index d02926ce..6ee9dfaf 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -446,7 +446,6 @@ "viewRounds_one": "View trace ({{count}} round)", "viewRounds_other": "View trace ({{count}} rounds)", "title": "Execution trace", - "subtitle": "Every step behind this entry — retrieval, model calls and tool calls — with its timing.", "empty": "No trace was recorded for this entry.", "failed": "Could not load the trace.", "noSpans": "Nothing was recorded for this run.", diff --git a/frontend/src/locale/es.json b/frontend/src/locale/es.json index e1a3c307..d0aa427c 100644 --- a/frontend/src/locale/es.json +++ b/frontend/src/locale/es.json @@ -441,7 +441,6 @@ "viewRounds_one": "Ver traza ({{count}} ronda)", "viewRounds_other": "Ver traza ({{count}} rondas)", "title": "Traza de ejecución", - "subtitle": "Cada paso detrás de esta entrada (recuperación, llamadas al modelo y llamadas a herramientas) con su duración.", "empty": "No se registró ninguna traza para esta entrada.", "failed": "No se pudo cargar la traza.", "noSpans": "No se registró nada en esta ejecución.", diff --git a/frontend/src/locale/jp.json b/frontend/src/locale/jp.json index ef316ca3..33a05be7 100644 --- a/frontend/src/locale/jp.json +++ b/frontend/src/locale/jp.json @@ -440,7 +440,6 @@ "viewRounds_one": "トレースを表示（{{count}} ラウンド）", "viewRounds_other": "トレースを表示（{{count}} ラウンド）", "title": "実行トレース", - "subtitle": "このエントリの各ステップ（検索、モデル呼び出し、ツール呼び出し）とその所要時間。", "empty": "このエントリのトレースは記録されていません。", "failed": "トレースを読み込めませんでした。", "noSpans": "この実行では何も記録されませんでした。", diff --git a/frontend/src/locale/ru.json b/frontend/src/locale/ru.json index bf090d6b..5e724a0b 100644 --- a/frontend/src/locale/ru.json +++ b/frontend/src/locale/ru.json @@ -455,7 +455,6 @@ "viewRounds_many": "Открыть трассировку ({{count}} раундов)", "viewRounds_other": "Открыть трассировку ({{count}} раунда)", "title": "Трассировка выполнения", - "subtitle": "Каждый шаг за этой записью — поиск, вызовы модели и инструментов — с длительностью.", "empty": "Для этой записи трассировка не сохранена.", "failed": "Не удалось загрузить трассировку.", "noSpans": "Для этого запуска ничего не записано.", diff --git a/frontend/src/locale/zh-TW.json b/frontend/src/locale/zh-TW.json index c7184ad8..fbc8b567 100644 --- a/frontend/src/locale/zh-TW.json +++ b/frontend/src/locale/zh-TW.json @@ -440,7 +440,6 @@ "viewRounds_one": "檢視追蹤（{{count}} 輪）", "viewRounds_other": "檢視追蹤（{{count}} 輪）", "title": "執行追蹤", - "subtitle": "此項目背後的每個步驟（檢索、模型呼叫和工具呼叫）及其耗時。", "empty": "此項目沒有記錄追蹤。", "failed": "無法載入追蹤。", "noSpans": "此次執行沒有記錄任何內容。", diff --git a/frontend/src/locale/zh.json b/frontend/src/locale/zh.json index 9db7d12f..c1db1fee 100644 --- a/frontend/src/locale/zh.json +++ b/frontend/src/locale/zh.json @@ -440,7 +440,6 @@ "viewRounds_one": "查看追踪（{{count}} 轮）", "viewRounds_other": "查看追踪（{{count}} 轮）", "title": "执行追踪", - "subtitle": "此条目背后的每一步（检索、模型调用和工具调用）及其耗时。", "empty": "此条目没有记录追踪。", "failed": "无法加载追踪。", "noSpans": "此次运行没有记录任何内容。", diff --git a/frontend/src/settings/traces/TraceSheet.tsx b/frontend/src/settings/traces/TraceSheet.tsx index 5a847c96..f459527f 100644 --- a/frontend/src/settings/traces/TraceSheet.tsx +++ b/frontend/src/settings/traces/TraceSheet.tsx @@ -4,13 +4,7 @@ import { useSelector } from 'react-redux'; import userService from '../../api/services/userService'; import Spinner from '../../components/Spinner'; -import { - Sheet, - SheetContent, - SheetDescription, - SheetHeader, - SheetTitle, -} from '../../components/ui/sheet'; +import { Sheet, SheetContent } from '../../components/ui/sheet'; import { selectToken } from '../../preferences/preferenceSlice'; import { formatDateTime } from '../../utils/dateTimeUtils'; import { Trace, TraceRef } from '../types'; @@ -76,17 +70,15 @@ export default function TraceSheet({ return ( !open && onClose()}> + {/* The waterfall speaks for itself: the title is for screen readers + only, and there is no description to announce. */} - - {t('settings.logs.trace.title')} - - {t('settings.logs.trace.subtitle')} - - -
+
{loading && (
@@ -104,7 +96,8 @@ export default function TraceSheet({ )} {traces.map((trace, index) => (
-
+ {/* Right padding keeps the first row clear of the close button. */} +
{traces.length > 1 && ( {t('settings.logs.trace.round', { n: index + 1 })} diff --git a/tests/e2e/specs/tier-a/execution-traces.spec.ts b/tests/e2e/specs/tier-a/execution-traces.spec.ts index 2f1f4653..e355b43d 100644 --- a/tests/e2e/specs/tier-a/execution-traces.spec.ts +++ b/tests/e2e/specs/tier-a/execution-traces.spec.ts @@ -106,8 +106,8 @@ test.describe('tier-a · execution traces', () => { await page.goto('/settings/logs'); await page.getByText('what do you remember about me?').first().click(); await page.getByRole('button', { name: 'View trace' }).click(); - const sheet = page.getByRole('dialog'); - await expect(sheet.getByText('Execution trace')).toBeVisible(); + const sheet = page.getByRole('dialog', { name: 'Execution trace' }); + await expect(sheet).toBeVisible(); await expect(sheet.getByText('execute_tool memory_view').first()).toBeVisible(); await sheet.getByText('execute_tool memory_view').first().click(); await expect(sheet.getByText('Arguments')).toBeVisible(); @@ -165,8 +165,8 @@ test.describe('tier-a · execution traces', () => { await page.goto('/settings/logs'); await page.getByText('how do I get started?').first().click(); await page.getByRole('button', { name: 'View trace' }).click(); - const sheet = page.getByRole('dialog'); - await expect(sheet.getByText('Execution trace')).toBeVisible(); + const sheet = page.getByRole('dialog', { name: 'Execution trace' }); + await expect(sheet).toBeVisible(); await sheet.getByText(/^retrieval/).first().click(); await expect(sheet.getByText('Retrieved chunks')).toBeVisible(); if (process.env.TRACE_SCREENSHOT_RAG) { From 5272440a218546a1441436044d64ca17e6caf400 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:08:27 +0100 Subject: [PATCH 20/33] Return focus to View trace when the drawer closes The drawer is opened from a Logs row rather than a SheetTrigger, so Radix had nothing to restore focus to. It now remembers what had focus when it opened and returns there on close; the e2e spec checks it. --- frontend/src/settings/traces/TraceSheet.tsx | 19 ++++++++++++++++++- .../e2e/specs/tier-a/execution-traces.spec.ts | 5 +++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/frontend/src/settings/traces/TraceSheet.tsx b/frontend/src/settings/traces/TraceSheet.tsx index f459527f..eb7e2385 100644 --- a/frontend/src/settings/traces/TraceSheet.tsx +++ b/frontend/src/settings/traces/TraceSheet.tsx @@ -1,4 +1,4 @@ -import React, { useEffect, useState } from 'react'; +import React, { useEffect, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { useSelector } from 'react-redux'; @@ -39,6 +39,9 @@ export default function TraceSheet({ const [traces, setTraces] = useState([]); const [loading, setLoading] = useState(false); const [failed, setFailed] = useState(false); + // The sheet is opened from a Logs row, not a SheetTrigger, so Radix has no + // trigger to return focus to on close; remember what had focus instead. + const returnFocusRef = useRef(null); useEffect(() => { if (!traceRef) return; @@ -76,6 +79,20 @@ export default function TraceSheet({ side="right" title={t('settings.logs.trace.title')} aria-describedby={undefined} + onOpenAutoFocus={() => { + returnFocusRef.current = + document.activeElement instanceof HTMLElement + ? document.activeElement + : null; + }} + onCloseAutoFocus={(event) => { + const target = returnFocusRef.current; + returnFocusRef.current = null; + if (target?.isConnected) { + event.preventDefault(); + target.focus(); + } + }} className="w-full overflow-y-auto sm:max-w-3xl" >
diff --git a/tests/e2e/specs/tier-a/execution-traces.spec.ts b/tests/e2e/specs/tier-a/execution-traces.spec.ts index e355b43d..709585bb 100644 --- a/tests/e2e/specs/tier-a/execution-traces.spec.ts +++ b/tests/e2e/specs/tier-a/execution-traces.spec.ts @@ -114,6 +114,11 @@ test.describe('tier-a · execution traces', () => { if (process.env.TRACE_SCREENSHOT) { await page.screenshot({ path: process.env.TRACE_SCREENSHOT, fullPage: false }); } + + // Closing returns focus to the button that opened the drawer. + await page.keyboard.press('Escape'); + await expect(sheet).toBeHidden(); + await expect(page.getByRole('button', { name: 'View trace' })).toBeFocused(); } finally { await api.dispose(); await context.close(); From 82b31760eefffd6f88556a54702887ae728445f9 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:08:27 +0100 Subject: [PATCH 21/33] Put View trace before the trace chips in a Logs row --- frontend/src/settings/Logs.tsx | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/frontend/src/settings/Logs.tsx b/frontend/src/settings/Logs.tsx index 52d7e08e..16a0b247 100644 --- a/frontend/src/settings/Logs.tsx +++ b/frontend/src/settings/Logs.tsx @@ -475,10 +475,6 @@ function Log({
{log.trace && (
-
)} {detailRows.length > 0 && ( From 24796ae61fe9f6d4184132f4a07495810c7cbc15 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:09:58 +0100 Subject: [PATCH 22/33] Mint request ids on the server and trace refused requests build_agent no longer takes request_id from the request body: it becomes the primary LLM's usage request id, and quotas count distinct request ids, so a client could make every call count as one. Requests refused after setup started (unauthorized, over quota, resume conflict, setup error) now write their trace, marked error, through an after-request hook; streaming routes hand the trace to complete_stream instead. --- docsgpt/api/answer/routes/answer.py | 8 ++- docsgpt/api/answer/routes/base.py | 8 ++- docsgpt/api/answer/routes/stream.py | 10 ++- .../api/answer/services/stream_processor.py | 61 +++++++++++++++-- docsgpt/api/v1/routes.py | 8 ++- tests/api/answer/test_trace_wiring.py | 67 ++++++++++++++++++- 6 files changed, 147 insertions(+), 15 deletions(-) diff --git a/docsgpt/api/answer/routes/answer.py b/docsgpt/api/answer/routes/answer.py index bc2eb726..4d7f8dd2 100644 --- a/docsgpt/api/answer/routes/answer.py +++ b/docsgpt/api/answer/routes/answer.py @@ -13,7 +13,10 @@ from docsgpt.api.answer.services.continuation_service import ( ResumeInProgressError, ) from docsgpt.api.answer.services.persistence_policy import resolve_persistence -from docsgpt.api.answer.services.stream_processor import StreamProcessor +from docsgpt.api.answer.services.stream_processor import ( + StreamProcessor, + flush_trace_after_request, +) logger = logging.getLogger(__name__) @@ -88,6 +91,9 @@ class AnswerResource(Resource, BaseAnswerResource): return error decoded_token = getattr(request, "decoded_token", None) processor = StreamProcessor(data, decoded_token, trace_source="answer") + # ``complete_stream`` is consumed below and writes the trace itself; + # this covers requests refused before it runs. + flush_trace_after_request(processor) try: # ---- Continuation mode ---- if data.get("tool_actions"): diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index f23b82f3..943efedc 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -6,7 +6,7 @@ import logging import threading import time import uuid -from typing import Any, Dict, Generator, List, Optional +from typing import Any, Callable, Dict, Generator, List, Optional from flask import jsonify, make_response, Response from flask_restx import Namespace @@ -69,7 +69,9 @@ class StreamSuperseded(Exception): answer_ns = Namespace("answer", description="Answer related operations", path="/") -def _traced_stream(method): +def _traced_stream( + method: Callable[..., Generator[str, None, None]], +) -> Callable[..., Generator[str, None, None]]: """Record ``complete_stream`` as the request's execution trace. The stream runs in the SSE pump thread, not the request thread, so the @@ -81,7 +83,7 @@ def _traced_stream(method): signature = inspect.signature(method) @functools.wraps(method) - def wrapper(*args, **kwargs): + def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]: bound = signature.bind(*args, **kwargs) arguments = bound.arguments continuation = arguments.get("_continuation") diff --git a/docsgpt/api/answer/routes/stream.py b/docsgpt/api/answer/routes/stream.py index 36cb0472..21dbcbff 100644 --- a/docsgpt/api/answer/routes/stream.py +++ b/docsgpt/api/answer/routes/stream.py @@ -13,7 +13,10 @@ from docsgpt.api.answer.services.continuation_service import ( ResumeInProgressError, ) from docsgpt.api.answer.services.persistence_policy import resolve_persistence -from docsgpt.api.answer.services.stream_processor import StreamProcessor +from docsgpt.api.answer.services.stream_processor import ( + StreamProcessor, + flush_trace_after_request, +) from docsgpt.streaming.sse_keepalive import with_sse_keepalive logger = logging.getLogger(__name__) @@ -95,6 +98,7 @@ class StreamResource(Resource, BaseAnswerResource): return error decoded_token = getattr(request, "decoded_token", None) processor = StreamProcessor(data, decoded_token, trace_source="stream") + flush_trace_after_request(processor) try: # ---- Continuation mode ---- @@ -130,7 +134,7 @@ class StreamResource(Resource, BaseAnswerResource): agent_id=processor.agent_id, model_id=processor.model_id, model_user_id=processor.model_user_id, - trace=processor.trace, + trace=processor.handoff_trace(), _continuation={ "messages": messages, "tools_dict": tools_dict, @@ -183,7 +187,7 @@ class StreamResource(Resource, BaseAnswerResource): model_id=processor.model_id, model_user_id=processor.model_user_id, request_id=processor.request_id, - trace=processor.trace, + trace=processor.handoff_trace(), ), ), mimetype="text/event-stream", diff --git a/docsgpt/api/answer/services/stream_processor.py b/docsgpt/api/answer/services/stream_processor.py index bb7f47a3..27be1028 100644 --- a/docsgpt/api/answer/services/stream_processor.py +++ b/docsgpt/api/answer/services/stream_processor.py @@ -4,7 +4,9 @@ import json import logging import uuid from pathlib import Path -from typing import Any, Dict, Optional, Set +from typing import Any, Callable, Dict, Optional, Set, TypeVar + +from flask import after_this_request from docsgpt import tracing from docsgpt.agents.agent_creator import AgentCreator @@ -107,7 +109,10 @@ def get_prompt(prompt_id: str, prompts_collection=None) -> str: raise ValueError(f"Invalid prompt ID: {prompt_id}") from e -def _traced_setup(method): +T = TypeVar("T") + + +def _traced_setup(method: Callable[..., T]) -> Callable[..., T]: """Run a request-setup method inside the request's execution trace. Agent setup does real work worth seeing in the trace -- pre-fetch @@ -116,7 +121,7 @@ def _traced_setup(method): """ @functools.wraps(method) - def wrapper(self, *args, **kwargs): + def wrapper(self: "StreamProcessor", *args: Any, **kwargs: Any) -> T: trace = getattr(self, "trace", None) if trace is None: trace = tracing.start_trace(source=getattr(self, "trace_source", "stream")) @@ -136,6 +141,25 @@ def _traced_setup(method): return wrapper +def flush_trace_after_request(processor: "StreamProcessor") -> None: + """Write ``processor``'s setup trace when the request ends, unless it was claimed. + + Registered on the current request with ``after_this_request``; the hook + always hands the response back unchanged and never raises. + + Args: + processor: The request's processor. + """ + + @after_this_request + def _flush(response: Any) -> Any: + try: + processor.flush_unclaimed_trace() + except Exception: + logger.warning("Could not write an unclaimed request trace", exc_info=True) + return response + + class StreamProcessor: def __init__( self, @@ -201,6 +225,31 @@ class StreamProcessor: self._load_conversation_history() self._process_attachments() + def handoff_trace(self) -> Optional[tracing.Trace]: + """Hand the setup trace to a streaming ``complete_stream``. + + The stream writes the trace when it ends, after the view has + returned; marking the hand-off stops :meth:`flush_unclaimed_trace` + from writing it first. + + Returns: + The trace to pass as ``complete_stream(trace=...)``. + """ + self._trace_handed_off = True + return getattr(self, "trace", None) + + def flush_unclaimed_trace(self) -> None: + """Write the setup trace of a request that ended before streaming. + + A request refused after setup started (unauthorized, over its usage + limit, a resume conflict, a setup error) still records what ran, + marked ``error``. A trace the request already wrote, or handed to a + stream, is left alone. Routes arrange this with + :func:`flush_trace_after_request`. + """ + if not getattr(self, "_trace_handed_off", False): + tracing.flush(getattr(self, "trace", None), tracing.STATUS_ERROR) + @_traced_setup def build_agent(self, question: str): """One call to go from request data to a ready-to-run agent. @@ -208,10 +257,12 @@ class StreamProcessor: Combines initialize(), pre_fetch_docs(), pre_fetch_tools(), and create_agent() into a single convenience method. The request id is minted first so pre-fetch retrieval and its side-channel LLM calls - share it with the rest of the turn. + share it with the rest of the turn. It is always generated here, never + taken from the request body: request quotas count distinct request + ids, so a client-chosen id would let every call count as one. """ if not getattr(self, "request_id", None): - self.request_id = self.data.get("request_id") or str(uuid.uuid4()) + self.request_id = str(uuid.uuid4()) self.initialize() agent_type = self.agent_config.get("agent_type", "classic") diff --git a/docsgpt/api/v1/routes.py b/docsgpt/api/v1/routes.py index 9294c23a..c177ad6c 100644 --- a/docsgpt/api/v1/routes.py +++ b/docsgpt/api/v1/routes.py @@ -21,7 +21,10 @@ from docsgpt.api.answer.services.continuation_service import ( RESUME_IN_PROGRESS_MESSAGE, ResumeInProgressError, ) -from docsgpt.api.answer.services.stream_processor import StreamProcessor +from docsgpt.api.answer.services.stream_processor import ( + StreamProcessor, + flush_trace_after_request, +) from docsgpt.api.v1 import idempotency as v1_idempotency from docsgpt.api.v1.session_store import ( V1Session, @@ -258,6 +261,7 @@ def chat_completions(): try: processor = StreamProcessor(internal_data, decoded_token, trace_source="v1") + flush_trace_after_request(processor) # Set when this request took the resume claim, so a refusal can release it. claimed_conversation_id = None @@ -371,6 +375,8 @@ def chat_completions(): # safe way to re-emit a recorded SSE stream (and the regression / # b2b client is non-streaming), so a streaming request never # claims a key. This is a known, accepted limitation. + # The stream writes the trace once it runs, after this returns. + processor.handoff_trace() return Response( with_sse_keepalive( _stream_response( diff --git a/tests/api/answer/test_trace_wiring.py b/tests/api/answer/test_trace_wiring.py index 9153c424..7cf75f7a 100644 --- a/tests/api/answer/test_trace_wiring.py +++ b/tests/api/answer/test_trace_wiring.py @@ -238,14 +238,46 @@ class TestProcessorTraceSetup: assert [s.kind for s in trace.spans] == ["retrieval"] assert tracing.current_trace() is None - def test_client_supplied_request_id_is_kept(self): + def test_client_supplied_request_id_is_ignored(self): + """Quotas count distinct request ids; a client must not choose its own.""" from docsgpt.api.answer.services.stream_processor import StreamProcessor processor = StreamProcessor({"request_id": "client-rid"}, {"sub": "u1"}) with patch.object(processor, "initialize", side_effect=RuntimeError("stop")): with pytest.raises(RuntimeError): processor.build_agent("q") - assert processor.request_id == "client-rid" + assert processor.request_id and processor.request_id != "client-rid" + + def test_refused_request_still_writes_its_trace(self): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + processor = StreamProcessor({}, {"sub": "u1"}) + + def _initialize(): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + + with patch.object(processor, "initialize", side_effect=_initialize), patch.object( + processor, "pre_fetch_docs", return_value=(None, None) + ), patch.object(processor, "pre_fetch_tools", return_value=None), patch.object( + processor, "create_agent", return_value=MagicMock() + ), patch.object(processor, "_exposure_partition", return_value=([], [])): + processor.build_agent("q") + with _captured_flushes() as flushed: + processor.flush_unclaimed_trace() + (trace,) = flushed + assert trace.status == "error" + assert [s.kind for s in trace.spans] == ["retrieval"] + + def test_handed_off_trace_is_left_to_the_stream(self): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + processor = StreamProcessor({}, {"sub": "u1"}) + processor.trace = tracing.start_trace(source="stream", capture_otel_context=False) + assert processor.handoff_trace() is processor.trace + with _captured_flushes() as flushed: + processor.flush_unclaimed_trace() + assert flushed == [] def test_tracing_disabled_leaves_no_trace(self, monkeypatch): from docsgpt.api.answer.services.stream_processor import StreamProcessor @@ -257,3 +289,34 @@ class TestProcessorTraceSetup: processor.build_agent("q") assert processor.trace is None assert processor.request_id + + +@pytest.mark.unit +class TestRouteFlushesRefusedRequests: + def test_unauthorized_answer_request_writes_its_trace(self, mock_mongo_db, flask_app): + """The route registers the flush, and the hook never replaces the response.""" + import json + + from flask_restx import Api + + from docsgpt.api.answer.routes.answer import answer_ns + + api = Api(flask_app) + api.add_namespace(answer_ns) + client = flask_app.test_client() + processor = MagicMock() + processor.decoded_token = None + processor.flush_unclaimed_trace.return_value = "not a response" + with patch( + "docsgpt.api.answer.routes.answer.StreamProcessor", return_value=processor + ), patch( + "docsgpt.api.answer.routes.answer.AnswerResource.validate_request", + return_value=None, + ): + resp = client.post( + "/api/answer", + data=json.dumps({"question": "q"}), + content_type="application/json", + ) + assert resp.status_code == 401 + processor.flush_unclaimed_trace.assert_called_once_with() From 9f7f0b2f1e0f71fc52e3636ed40febff36253bfe Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:09:58 +0100 Subject: [PATCH 23/33] Keep trace summaries from breaking Logs and fix their counts A failed trace-summary lookup now leaves the Logs page intact without chips. Tool-call counts include only calls that ran, not their paused, denied or skipped records. A local guardrail that fires unchanged on every streamed segment is recorded once, so it cannot use up the span cap. --- docsgpt/api/user/analytics/routes.py | 29 ++++++++++++++------- docsgpt/guardrails/engine.py | 31 ++++++++++++++++++---- docsgpt/tracing/__init__.py | 2 ++ docsgpt/tracing/core.py | 39 ++++++++++++++++++++++++---- tests/api/user/test_traces_routes.py | 19 ++++++++++++++ tests/guardrails/test_engine.py | 10 +++++++ tests/tracing/test_trace.py | 23 ++++++++++++++++ 7 files changed, 134 insertions(+), 19 deletions(-) diff --git a/docsgpt/api/user/analytics/routes.py b/docsgpt/api/user/analytics/routes.py index 570c141e..c77d9eb1 100644 --- a/docsgpt/api/user/analytics/routes.py +++ b/docsgpt/api/user/analytics/routes.py @@ -1,10 +1,11 @@ """Analytics and reporting routes.""" import datetime +from typing import Optional, Tuple from flask import current_app, jsonify, make_response, request from flask_restx import fields, Namespace, Resource -from sqlalchemy import text as _sql_text +from sqlalchemy import Connection, text as _sql_text from docsgpt.api import api from docsgpt.api.user.base import ( @@ -123,7 +124,7 @@ def _trace_branch(name: str, sources_sql: str, scope: str) -> dict: } -def _trace_ref(item: dict): +def _trace_ref(item: dict) -> Optional[Tuple[str, str]]: """The ``(field, value)`` that finds a Logs row's trace, if it has one.""" event_type = item["event_type"] row_id = item["id"].split("-", 1)[1] @@ -159,7 +160,9 @@ def _merge_trace_summaries(traces: list) -> dict: } -def _attach_trace_summaries(conn, items: list, *, user_id, agent_id) -> None: +def _attach_trace_summaries( + conn: Connection, items: list, *, user_id: Optional[str], agent_id: Optional[str] +) -> None: """Add a ``trace`` summary to each Logs row that has a stored trace. One batched lookup per link field for the whole page, rather than a join @@ -1282,12 +1285,20 @@ class GetUserLogs(Resource): ) results.append(item) if results: - with db_readonly() as conn: - _attach_trace_summaries( - conn, - results, - user_id=user, - agent_id=agent_pg_id if api_key_id else None, + # Trace chips are an extra: a failed lookup must leave the + # page intact, just without them. + try: + with db_readonly() as conn: + _attach_trace_summaries( + conn, + results, + user_id=user, + agent_id=agent_pg_id if api_key_id else None, + ) + except Exception: + current_app.logger.warning( + "Could not attach trace summaries to the logs page", + exc_info=True, ) except Exception as err: current_app.logger.error( diff --git a/docsgpt/guardrails/engine.py b/docsgpt/guardrails/engine.py index be489385..834edc36 100644 --- a/docsgpt/guardrails/engine.py +++ b/docsgpt/guardrails/engine.py @@ -38,7 +38,23 @@ def _start_guardrail_span(stage: Stage, controls) -> "tracing.Span": ) -def _describe_decision(span, decision: StageDecision) -> None: +def _content_fired(decision: StageDecision) -> bool: + return bool(decision.triggered or decision.blocked or decision.redacted) + + +def _decision_key(decision: StageDecision) -> tuple: + """Identity of a guardrail outcome, for recording each distinct one once.""" + return ( + "guardrail", + decision.stage.value, + tuple(sorted(v.check for v in decision.triggered)), + tuple(sorted(v.check for v in decision.unevaluated)), + decision.blocked, + decision.redacted, + ) + + +def _describe_decision(span: "tracing.Span", decision: StageDecision) -> None: """Record a stage decision on its span; a firing guardrail drops trace previews.""" triggered = [v.check for v in decision.triggered] span.set( @@ -50,7 +66,7 @@ def _describe_decision(span, decision: StageDecision) -> None: "docsgpt.guardrail.unevaluated": [v.check for v in decision.unevaluated] or None, } ) - if triggered or decision.blocked or decision.redacted: + if _content_fired(decision): # The scanned text (or text near it) sits in other spans' previews: # the retrieval query, tool results, the answer. Keep none of it. tracing.mark_content_blocked() @@ -136,10 +152,15 @@ class GuardrailEngine: decision.verdicts = [self._run_control(c, text, stage) for c in controls] self._reduce(decision) # The output guard evaluates every streamed segment; tracing each - # clean local scan would bury the trace, so only a firing is kept. + # clean local scan would bury the trace, so only a firing is kept, + # and a firing that repeats unchanged segment after segment is + # recorded once so it cannot use up the trace's span cap. if not decision.clean or decision.unevaluated: - with _start_guardrail_span(stage, controls) as span: - _describe_decision(span, decision) + if tracing.first_occurrence(_decision_key(decision)): + with _start_guardrail_span(stage, controls) as span: + _describe_decision(span, decision) + elif _content_fired(decision): + tracing.mark_content_blocked() self._record(decision) return decision diff --git a/docsgpt/tracing/__init__.py b/docsgpt/tracing/__init__.py index 0cf018d2..d31ba387 100644 --- a/docsgpt/tracing/__init__.py +++ b/docsgpt/tracing/__init__.py @@ -43,6 +43,7 @@ from docsgpt.tracing.core import ( bind, bind_if_unset, current_trace, + first_occurrence, mark_content_blocked, span, start_span, @@ -78,6 +79,7 @@ __all__ = [ "bind_if_unset", "current_trace", "discard", + "first_occurrence", "flush", "mark_content_blocked", "span", diff --git a/docsgpt/tracing/core.py b/docsgpt/tracing/core.py index aa0f3b07..47c5dbba 100644 --- a/docsgpt/tracing/core.py +++ b/docsgpt/tracing/core.py @@ -272,6 +272,7 @@ class Trace: self.attributes: Dict[str, Any] = {} self.finished = False self.flushed = False + self._seen_keys: set = set() self._lock = threading.Lock() self._stacks: Dict[int, List[Any]] = {} @@ -288,7 +289,7 @@ class Trace: *, parent: Any = None, attributes: Optional[Dict[str, Any]] = None, - ): + ) -> Any: """Record a new span; returns :data:`NOOP_SPAN` once finished or over the cap.""" with self._lock: if self.finished: @@ -344,10 +345,18 @@ class Trace: return _undo - def current_parent(self) -> Optional["_Seed | Span"]: + def current_parent(self) -> Optional[Any]: stack = self._stacks.get(threading.get_ident()) return stack[-1] if stack else None + def first_occurrence(self, key: Any) -> bool: + """True the first time ``key`` is seen in this trace, then False.""" + with self._lock: + if key in self._seen_keys: + return False + self._seen_keys.add(key) + return True + # -- ids -------------------------------------------------------------- def bind(self, *, only_if_unset: bool = False, **ids: Any) -> None: @@ -409,7 +418,12 @@ class Trace: if s.kind == KIND_RETRIEVAL and not (s.parent_id in by_id and by_id[s.parent_id].kind == KIND_RETRIEVAL) ] - tools = [s for s in self.spans if s.kind == KIND_TOOL] + # Only calls that ran count: a call paused for approval is recorded + # again when it runs in the next round, and denied or skipped calls + # never ran at all. + tools = [ + s for s in self.spans if s.kind == KIND_TOOL and s.status in (STATUS_OK, STATUS_ERROR) + ] def _tokens(key: str) -> int: total = 0 @@ -534,7 +548,9 @@ def activate(trace: Optional[Trace]) -> Iterator[Optional[Trace]]: _current.set(None) -def start_span(kind: str, name: str, *, parent: Any = None, attributes: Optional[Dict[str, Any]] = None): +def start_span( + kind: str, name: str, *, parent: Any = None, attributes: Optional[Dict[str, Any]] = None +) -> Any: """Start a span in the current trace; returns :data:`NOOP_SPAN` without one.""" trace = _current.get() if trace is None: @@ -546,7 +562,9 @@ def start_span(kind: str, name: str, *, parent: Any = None, attributes: Optional return NOOP_SPAN -def span(kind: str, name: str, *, parent: Any = None, attributes: Optional[Dict[str, Any]] = None): +def span( + kind: str, name: str, *, parent: Any = None, attributes: Optional[Dict[str, Any]] = None +) -> Any: """Context-manager form of :func:`start_span` (errors mark the span and re-raise).""" return start_span(kind, name, parent=parent, attributes=attributes) @@ -565,6 +583,17 @@ def bind_if_unset(**ids: Any) -> None: trace.bind(only_if_unset=True, **ids) +def first_occurrence(key: Any) -> bool: + """True the first time ``key`` is seen in the current trace. + + For steps that repeat identically many times in one request (a guardrail + re-firing on every streamed segment): record the first, skip the rest, so + they cannot use up the span cap. False without an active trace. + """ + trace = _current.get() + return trace.first_occurrence(key) if trace is not None else False + + def mark_content_blocked() -> None: """A guardrail blocked or retracted content: drop every preview from the stored trace.""" trace = _current.get() diff --git a/tests/api/user/test_traces_routes.py b/tests/api/user/test_traces_routes.py index a0b2bbbe..7c89d864 100644 --- a/tests/api/user/test_traces_routes.py +++ b/tests/api/user/test_traces_routes.py @@ -193,3 +193,22 @@ class TestLogsTraceSummaries: def test_unknown_event_type_is_rejected(self, app, pg_conn): assert _logs(app, pg_conn, "owner", {"event_type": "bogus"}).status_code == 400 + + +class TestSummaryFailure: + def test_page_survives_a_failed_trace_lookup(self, app, pg_conn): + from docsgpt.storage.db.repositories.user_logs import UserLogsRepository + + UserLogsRepository(pg_conn).insert( + user_id="owner", + endpoint="stream_answer", + data={"question": "q", "request_id": "req-1"}, + ) + with patch( + "docsgpt.api.user.analytics.routes.RequestTracesRepository.summaries_for_refs", + side_effect=RuntimeError("statement timeout"), + ): + response = _logs(app, pg_conn, "owner", {}) + assert response.status_code == 200 + (row,) = response.json["logs"] + assert "trace" not in row diff --git a/tests/guardrails/test_engine.py b/tests/guardrails/test_engine.py index cf7b704a..420192de 100644 --- a/tests/guardrails/test_engine.py +++ b/tests/guardrails/test_engine.py @@ -388,6 +388,16 @@ class TestTraceSpans: assert span.attributes["docsgpt.guardrail.triggered"] == ["_test_always"] assert self.trace.content_blocked is True + def test_a_repeating_firing_is_recorded_once(self): + """The output guard re-scans every segment; one firing must not fill the span cap.""" + engine = GuardrailEngine( + _config(controls=[{"check": "_test_always", "stage": "output", "action": "flag"}]) + ) + for _ in range(50): + engine.evaluate("segment", Stage.OUTPUT) + assert len(self.trace.spans) == 1 + assert self.trace.content_blocked is True + def test_remote_scan_is_traced_with_judge_nested(self): engine = GuardrailEngine( _config(controls=[{"check": "_test_quick_remote", "stage": "input", "action": "block"}]) diff --git a/tests/tracing/test_trace.py b/tests/tracing/test_trace.py index 437f1b53..e728e0d3 100644 --- a/tests/tracing/test_trace.py +++ b/tests/tracing/test_trace.py @@ -372,3 +372,26 @@ class TestRecordQuery: tracing.mark_content_blocked() trace.finish() assert "query" not in trace.to_record()["summary"] + + +class TestToolCount: + def test_only_tool_calls_that_ran_are_counted(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + tracing.start_span(tracing.KIND_TOOL, "paused").end(tracing.STATUS_PENDING) + tracing.start_span(tracing.KIND_TOOL, "denied").end(tracing.STATUS_DENIED) + tracing.start_span(tracing.KIND_TOOL, "skipped").end(tracing.STATUS_SKIPPED) + tracing.start_span(tracing.KIND_TOOL, "ran").end() + tracing.start_span(tracing.KIND_TOOL, "failed").end(tracing.STATUS_ERROR) + trace.finish() + assert trace.summary()["tool_calls"] == 2 + + +class TestFirstOccurrence: + def test_key_is_new_once_per_trace(self): + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + assert tracing.first_occurrence(("a", 1)) is True + assert tracing.first_occurrence(("a", 1)) is False + assert tracing.first_occurrence(("a", 2)) is True + assert tracing.first_occurrence(("a", 3)) is False # no active trace From 699906a0154af0afdc8261eb328e4bc7bd54a002 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:09:58 +0100 Subject: [PATCH 24/33] Give a scheduled run's trace to the user who scheduled it On a shared agent the run's Logs row belongs to the scheduling user, not the agent owner, so the trace is stored under that user. --- docsgpt/agents/headless_runner.py | 6 +++++- docsgpt/api/user/scheduler_worker.py | 4 +++- tests/tracing/test_entry_points.py | 11 +++++++++++ 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/docsgpt/agents/headless_runner.py b/docsgpt/agents/headless_runner.py index e50e2478..d97e6407 100644 --- a/docsgpt/agents/headless_runner.py +++ b/docsgpt/agents/headless_runner.py @@ -71,6 +71,7 @@ def run_agent_headless( chat_history: Optional[List[Dict[str, Any]]] = None, conversation_id: Optional[str] = None, request_id: Optional[str] = None, + trace_user_id: Optional[str] = None, ) -> Dict[str, Any]: """Run an agent with no live client; returns a structured outcome dict. @@ -78,6 +79,9 @@ def run_agent_headless( source. ``request_id`` links that trace to the caller's own record (the scheduler passes its run id, the webhook worker its task id); it is kept off the LLM's token-usage rows, whose request ids drive request counts. + ``trace_user_id`` owns the trace when the run belongs to someone other + than the agent's owner (a schedule a user set on a shared agent), so the + trace is visible wherever that user sees the run; it defaults to the owner. Raises: QuotaExceededError: If the agent owner's usage quota is exhausted. @@ -85,7 +89,7 @@ def run_agent_headless( trace = tracing.start_trace( source=endpoint, request_id=request_id, - user_id=_resolve_owner(agent_config), + user_id=trace_user_id or _resolve_owner(agent_config), agent_id=_resolve_agent_id(agent_config), conversation_id=conversation_id, ) diff --git a/docsgpt/api/user/scheduler_worker.py b/docsgpt/api/user/scheduler_worker.py index 15c73f76..be05c101 100644 --- a/docsgpt/api/user/scheduler_worker.py +++ b/docsgpt/api/user/scheduler_worker.py @@ -277,8 +277,10 @@ def execute_scheduled_run_body(run_id: str, celery_task_id: Optional[str]) -> Di endpoint="schedule", conversation_id=schedule.get("origin_conversation_id"), chat_history=chat_history, - # Links the run's execution trace to its Logs row. + # Links the run's execution trace to its Logs row, which belongs + # to the user who scheduled it (not always the agent's owner). request_id=str(run_id), + trace_user_id=run.get("user_id") or schedule.get("user_id"), ) except SoftTimeLimitExceeded: timed_out = True diff --git a/tests/tracing/test_entry_points.py b/tests/tracing/test_entry_points.py index 221b0c29..95ba3809 100644 --- a/tests/tracing/test_entry_points.py +++ b/tests/tracing/test_entry_points.py @@ -84,6 +84,17 @@ class TestHeadless: assert trace.status == "ok" assert [s.kind for s in trace.spans] == ["retrieval", "agent"] + def test_trace_can_belong_to_the_scheduling_user(self, monkeypatch, flushed): + """A schedule on a shared agent is the scheduler's run, not the owner's.""" + _headless( + [{"answer": "done"}], + monkeypatch, + endpoint="schedule", + request_id="run-2", + trace_user_id="scheduler-user", + ) + assert flushed[0].user_id == "scheduler-user" + def test_stream_error_marks_trace_error(self, monkeypatch, flushed): outcome = _headless([{"type": "error", "error": "boom"}], monkeypatch, endpoint="webhook") assert outcome["error_type"] == "stream_error" From 0bbbac0eb22fd01344e61faf635acf1c87ae4e5d Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:09:58 +0100 Subject: [PATCH 25/33] 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 | 19 +-------- docsgpt/agents/research_agent.py | 4 +- docsgpt/agents/tool_executor.py | 2 +- docsgpt/logging.py | 39 +++++++++++++++++-- docsgpt/tracing/llm.py | 18 ++++++--- docsgpt/tracing/retrieval.py | 10 +++-- docsgpt/usage.py | 4 +- frontend/src/locale/de.json | 1 + frontend/src/locale/en.json | 1 + frontend/src/locale/es.json | 1 + frontend/src/locale/jp.json | 1 + frontend/src/locale/ru.json | 1 + frontend/src/locale/zh-TW.json | 1 + frontend/src/locale/zh.json | 1 + .../src/settings/traces/TraceSpanDetails.tsx | 4 ++ tests/test_logging.py | 2 + tests/tracing/test_llm_spans.py | 12 ++++++ tests/tracing/test_tool_spans.py | 2 + 18 files changed, 88 insertions(+), 35 deletions(-) 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" From 0591580014db2283d797d670a50d73b04a458e18 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:42:05 +0100 Subject: [PATCH 26/33] Keep exception text out of stored and exported span errors Span.fail stored the exception message in span.error, which is stored and exported whatever the content settings; a provider's content-filter error can quote the prompt. span.error now holds the exception type, and the message is a capture-gated preview, like tool errors already were. A yielded stream error on the agent span is treated the same way. --- docsgpt/logging.py | 4 +++- docsgpt/tracing/core.py | 15 +++++++++++++-- tests/test_logging.py | 5 ++--- tests/tracing/test_trace.py | 15 ++++++++++++++- 4 files changed, 32 insertions(+), 7 deletions(-) diff --git a/docsgpt/logging.py b/docsgpt/logging.py index 38ddf6da..5440631f 100644 --- a/docsgpt/logging.py +++ b/docsgpt/logging.py @@ -234,7 +234,9 @@ def _end_agent_span(span: Any, context: "LogContext", *, error: BaseException | span.end(error=error, attributes=attributes) elif context.stream_error: span.set(**attributes, **{"error.type": "StreamError"}) - span.error = context.stream_error + # The error event's text is gated like any other content. + span.preview("error", context.stream_error) + span.error = "StreamError" span.end(tracing.STATUS_ERROR) else: span.end(None if completed else tracing.STATUS_CANCELLED, attributes=attributes) diff --git a/docsgpt/tracing/core.py b/docsgpt/tracing/core.py index 47c5dbba..6e3aef52 100644 --- a/docsgpt/tracing/core.py +++ b/docsgpt/tracing/core.py @@ -146,10 +146,21 @@ class Span: return self def fail(self, exc: BaseException) -> "Span": - """Record ``exc`` on the span without ending it.""" + """Record ``exc`` on the span without ending it. + + ``error`` and ``error.type`` hold only the exception type: they are + stored and exported whatever the content settings, and exception text + can quote the prompt (a provider's content-filter error) or a tool + result. The message itself becomes the capture-gated ``error`` + preview, dropped with every other preview when capture is off or a + guardrail fired. + """ self.status = STATUS_ERROR - self.error = str(exc)[:500] or type(exc).__name__ + self.error = type(exc).__name__ self.attributes["error.type"] = type(exc).__name__ + message = str(exc) + if message: + self.preview("error", message) return self def end( diff --git a/tests/test_logging.py b/tests/test_logging.py index 69642c82..07fba72a 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -398,8 +398,6 @@ class TestAccumulateResponseSummary: assert ctx.tool_call_count == 2 - - class TestLogActivityTraceSpan: """``@log_activity`` opens the ``invoke_agent`` span for every agent run.""" @@ -475,7 +473,8 @@ class TestLogActivityTraceSpan: with patch("docsgpt.logging._log_activity_to_db"), tracing.activate(trace): list(gen(self._Agent())) assert trace.spans[0].status == "error" - assert trace.spans[0].error == "node failed" + assert trace.spans[0].error == "StreamError" + assert trace.spans[0].previews["error"] == "node failed" def test_raised_error_marks_span_error(self): from docsgpt import tracing diff --git a/tests/tracing/test_trace.py b/tests/tracing/test_trace.py index e728e0d3..b6e8c558 100644 --- a/tests/tracing/test_trace.py +++ b/tests/tracing/test_trace.py @@ -133,7 +133,20 @@ class TestNesting: span = trace.spans[0] assert span.status == "error" assert span.attributes["error.type"] == "ValueError" - assert span.error == "boom" + # Exception text can quote a prompt: it is a capture-gated preview, + # never the always-stored/exported ``error``. + assert span.error == "ValueError" + assert span.previews["error"] == "boom" + + def test_exception_text_is_dropped_when_capture_is_off(self, monkeypatch): + monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", False) + trace = tracing.start_trace(source="stream") + with tracing.activate(trace): + with pytest.raises(ValueError): + with tracing.span(tracing.KIND_LLM, "chat"): + raise ValueError("content filter: ") + trace.finish() + assert "the user's prompt" not in str(trace.to_record()) def test_generator_exit_marks_span_cancelled(self): trace = tracing.start_trace(source="stream") From 778830208bfc203ed9f60f4e2db5afdaf582b58f Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:42:05 +0100 Subject: [PATCH 27/33] Write chat traces off the stream's thread The OTel replay and the trace INSERT ran in the stream's finally, so a slow database held the SSE connection open after the last event. The trace is still frozen when the stream ends, but written on a small writer pool. --- docsgpt/api/answer/routes/base.py | 4 +- docsgpt/tracing/sink.py | 66 +++++++++++++++---- tests/api/answer/test_trace_wiring.py | 2 +- .../e2e/specs/tier-a/execution-traces.spec.ts | 6 +- tests/tracing/test_entry_points.py | 2 +- tests/tracing/test_sink.py | 28 ++++++++ tests/worker/test_extract_graph.py | 2 +- 7 files changed, 91 insertions(+), 19 deletions(-) diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index 943efedc..782e11b5 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -110,7 +110,9 @@ def _traced_stream( try: yield from method(*bound.args, **bound.kwargs) finally: - tracing.flush(trace) + # Written on a writer thread so the stream's connection closes + # without waiting on the OTel replay and the INSERT. + tracing.flush(trace, background=True) return wrapper diff --git a/docsgpt/tracing/sink.py b/docsgpt/tracing/sink.py index c5f9ad6c..49a516a8 100644 --- a/docsgpt/tracing/sink.py +++ b/docsgpt/tracing/sink.py @@ -3,29 +3,36 @@ from __future__ import annotations import logging +import threading +from concurrent.futures import Future, ThreadPoolExecutor from typing import Optional from docsgpt.tracing.core import Trace logger = logging.getLogger(__name__) +# Writes handed off by ``flush(background=True)``. Two workers keep up with +# chat traffic (each write is one INSERT); the pool's threads are joined at +# interpreter exit, so queued traces are still written on a clean shutdown. +_writer: Optional[ThreadPoolExecutor] = None +_writer_lock = threading.Lock() -def flush(trace: Optional[Trace], status: Optional[str] = None) -> None: - """Finish ``trace`` and persist it once; later calls are no-ops. + +def _executor() -> ThreadPoolExecutor: + global _writer + with _writer_lock: + if _writer is None: + _writer = ThreadPoolExecutor(max_workers=2, thread_name_prefix="trace-writer") + return _writer + + +def _write(trace: Trace) -> None: + """Export ``trace`` to OTel, then store it; each sink swallows its own failure. OTel replay runs first so the exported trace id can be stored with the - row. Both sinks swallow their own failures: a trace is diagnostic data - and must never fail the request that produced it. - - Args: - trace: The trace to write; ``None`` is accepted and ignored. - status: Final status; defaults to ``error`` when a top-level span - failed, else ``ok``. + row. A trace is diagnostic data and must never fail the request that + produced it. """ - if trace is None or trace.flushed: - return - trace.flushed = True - trace.finish(status) try: from docsgpt.tracing.otel import export_trace @@ -45,6 +52,39 @@ def flush(trace: Optional[Trace], status: Optional[str] = None) -> None: logger.warning("Failed to store trace %s", trace.id, exc_info=True) +def flush( + trace: Optional[Trace], status: Optional[str] = None, *, background: bool = False +) -> Optional[Future]: + """Finish ``trace`` and persist it once; later calls are no-ops. + + The trace is frozen immediately, so nothing recorded afterwards is + included, whichever way it is written. + + Args: + trace: The trace to write; ``None`` is accepted and ignored. + status: Final status; defaults to ``error`` when a top-level span + failed, else ``ok``. + background: Write on a writer thread instead of the caller's. A chat + stream flushes this way so its connection closes without waiting + on the OTel replay and the INSERT. + + Returns: + The pending write when ``background`` is set, else ``None``. + """ + if trace is None or trace.flushed: + return None + trace.flushed = True + trace.finish(status) + if background: + try: + return _executor().submit(_write, trace) + except RuntimeError: + # The pool is shut down (interpreter exit): write inline instead. + pass + _write(trace) + return None + + def discard(trace: Optional[Trace]) -> None: """Drop ``trace`` without writing it (the request was rejected or superseded).""" if trace is None or trace.flushed: diff --git a/tests/api/answer/test_trace_wiring.py b/tests/api/answer/test_trace_wiring.py index 7cf75f7a..d03dec01 100644 --- a/tests/api/answer/test_trace_wiring.py +++ b/tests/api/answer/test_trace_wiring.py @@ -22,7 +22,7 @@ def _captured_flushes(): """Record every flushed trace instead of writing it.""" flushed = [] - def _fake_flush(trace, status=None): + def _fake_flush(trace, status=None, **_kwargs): if trace is None or trace.flushed: return trace.flushed = True diff --git a/tests/e2e/specs/tier-a/execution-traces.spec.ts b/tests/e2e/specs/tier-a/execution-traces.spec.ts index 709585bb..d0474961 100644 --- a/tests/e2e/specs/tier-a/execution-traces.spec.ts +++ b/tests/e2e/specs/tier-a/execution-traces.spec.ts @@ -75,9 +75,10 @@ test.describe('tier-a · execution traces', () => { expect(res.ok()).toBeTruthy(); expect(await res.text()).toContain('"type": "end"'); - // Stored trace: one row for the turn, linked to its message. + // Stored trace: one row for the turn, linked to its message. It is + // written just after the stream closes, so wait for it. + await expect.poll(async () => (await tracesFor(sub)).length).toBe(1); const traces = await tracesFor(sub); - expect(traces).toHaveLength(1); const [trace] = traces; expect(trace.source).toBe('stream'); expect(trace.request_id).toBeTruthy(); @@ -154,6 +155,7 @@ test.describe('tier-a · execution traces', () => { }); expect(await res.text()).toContain('"type": "end"'); + await expect.poll(async () => (await tracesFor(sub)).length).toBe(1); const [trace] = await tracesFor(sub); const byKind = (kind: string) => trace.spans.filter((s) => s.kind === kind); const retrieval = byKind('retrieval'); diff --git a/tests/tracing/test_entry_points.py b/tests/tracing/test_entry_points.py index 95ba3809..5a613892 100644 --- a/tests/tracing/test_entry_points.py +++ b/tests/tracing/test_entry_points.py @@ -22,7 +22,7 @@ def flushed(): """Capture flushed traces instead of writing them.""" captured = [] - def _fake_flush(trace, status=None): + def _fake_flush(trace, status=None, **_kwargs): if trace is None or trace.flushed: return trace.flushed = True diff --git a/tests/tracing/test_sink.py b/tests/tracing/test_sink.py index 7fc7b188..4581434e 100644 --- a/tests/tracing/test_sink.py +++ b/tests/tracing/test_sink.py @@ -98,3 +98,31 @@ class TestDiscard: tracing.discard(trace) tracing.flush(trace) repo.insert.assert_not_called() + + +class TestBackgroundFlush: + def test_background_flush_writes_off_the_calling_thread(self): + import threading + + trace = _trace_with_span() + seen = {} + + def _insert(record): + seen["thread"] = threading.current_thread().name + return True + + with _patched_store() as repo: + repo.insert.side_effect = _insert + future = tracing.flush(trace, background=True) + assert trace.finished # frozen at once, written later + future.result(timeout=5) + assert seen["thread"].startswith("trace-writer") + assert repo.insert.call_count == 1 + + def test_background_flush_is_still_once(self): + trace = _trace_with_span() + with _patched_store() as repo: + future = tracing.flush(trace, background=True) + assert tracing.flush(trace, background=True) is None + future.result(timeout=5) + assert repo.insert.call_count == 1 diff --git a/tests/worker/test_extract_graph.py b/tests/worker/test_extract_graph.py index d6dd5701..10283ad5 100644 --- a/tests/worker/test_extract_graph.py +++ b/tests/worker/test_extract_graph.py @@ -188,7 +188,7 @@ class TestExtractGraphTrace: monkeypatch.setattr("docsgpt.graphrag.extraction.extract_graph_for_source", extract) flushed = [] - def _flush(trace, status=None): + def _flush(trace, status=None, **_kwargs): trace.flushed = True trace.finish(status) flushed.append(trace) From 248ebc80505c2b545eb4f0b3907d276e204bac27 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:42:05 +0100 Subject: [PATCH 28/33] Discard the setup trace of a replayed /v1 request An Idempotency-Key retry returns the cached response after setup already ran; its trace was then written as a failure no Logs row points to. --- docsgpt/api/v1/routes.py | 4 +++ tests/api/v1/test_v1_idempotency.py | 43 +++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/docsgpt/api/v1/routes.py b/docsgpt/api/v1/routes.py index c177ad6c..2b943088 100644 --- a/docsgpt/api/v1/routes.py +++ b/docsgpt/api/v1/routes.py @@ -21,6 +21,7 @@ from docsgpt.api.answer.services.continuation_service import ( RESUME_IN_PROGRESS_MESSAGE, ResumeInProgressError, ) +from docsgpt import tracing from docsgpt.api.answer.services.stream_processor import ( StreamProcessor, flush_trace_after_request, @@ -410,6 +411,9 @@ def chat_completions(): if not claimed: # ``completed`` cache hit, or a 409 for an in-flight same-key # request — either way return without re-running the agent. + # The original request already has its trace; this retry's + # setup (pre-fetch retrieval) is not a failed run to record. + tracing.discard(processor.trace) return replay # An exception from the agent run propagates to the ``except`` handlers diff --git a/tests/api/v1/test_v1_idempotency.py b/tests/api/v1/test_v1_idempotency.py index 08897b24..eda35b4d 100644 --- a/tests/api/v1/test_v1_idempotency.py +++ b/tests/api/v1/test_v1_idempotency.py @@ -437,3 +437,46 @@ class TestV1IdempotencyRoute: {"n": TASK_NAME}, ).scalar() assert completed == 1 + + +@pytest.mark.integration +class TestV1ReplayTrace: + """A replayed request must not leave a failed trace behind.""" + + QUESTION = TestV1IdempotencyRoute.QUESTION + + def test_replay_discards_its_setup_trace(self, pg_engine, monkeypatch): + from docsgpt import tracing + from docsgpt.core.settings import settings + + monkeypatch.setattr(settings, "TRACES_ENABLED", True) + monkeypatch.setattr(settings, "TRACES_OTEL_EXPORT", False) + user_id = f"user-{uuid.uuid4().hex[:8]}" + api_key = f"key-{uuid.uuid4().hex[:8]}" + with pg_engine.begin() as conn: + _seed_user(conn, user_id) + _seed_agent(conn, user_id, api_key) + + def _build_agent(processor, question): # noqa: ARG001 + # What the real setup does: a trace holding the pre-fetch retrieval. + processor.trace = tracing.start_trace(source="v1", capture_otel_context=False) + with tracing.activate(processor.trace): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + return _CountingAnswerAgent() + + written = [] + app = _build_app() + idem_key = f"idem-{uuid.uuid4().hex}" + with _wire_v1_route_db(pg_engine, monkeypatch), patch( + "docsgpt.api.answer.services.stream_processor.StreamProcessor.build_agent", + _build_agent, + ), patch("docsgpt.tracing.sink._write", side_effect=written.append): + _wire_idem_db(pg_engine, monkeypatch) + with app.test_client() as c: + resp1 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key) + resp2 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key) + + assert resp1.status_code == resp2.status_code == 200 + # Only the first request's trace is written, and it is not a failure. + assert [t.status for t in written] == ["ok"] From b3aa509d3a32b987bad1123130d06b4974d601b6 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:42:05 +0100 Subject: [PATCH 29/33] Delete all of a conversation's traces and index trace listing Only message-linked traces cascaded, so scheduled runs, stateless /v1 rounds and unreserved turns kept their previews after the conversation was deleted; they are now deleted by conversation_id. The Logs search and graph branches filter by source, so the listing indexes lead with it. --- .../alembic/versions/0037_request_traces.py | 23 ++++++++++++------- docsgpt/storage/db/models.py | 11 ++++++--- .../storage/db/repositories/conversations.py | 20 ++++++++++++++++ .../db/repositories/test_request_traces.py | 21 +++++++++++++++++ tests/storage/db/test_migration_0037.py | 13 +++++++++++ 5 files changed, 77 insertions(+), 11 deletions(-) diff --git a/docsgpt/alembic/versions/0037_request_traces.py b/docsgpt/alembic/versions/0037_request_traces.py index 4929af62..0f57f343 100644 --- a/docsgpt/alembic/versions/0037_request_traces.py +++ b/docsgpt/alembic/versions/0037_request_traces.py @@ -7,10 +7,15 @@ timings, stored as a ``spans`` JSONB array. The Logs UI loads a trace whole, so one row per trace keeps the write to a single INSERT and lets retention and conversation deletion remove a trace in one step. -``message_id`` cascades: deleting a conversation, or truncating it when a -turn is superseded, removes that turn's traces with it. The link ids -(``request_id``, ``activity_id``, ``workflow_run_id``) carry partial indexes -because only the Logs rows that have them look traces up by them. +``message_id`` cascades: truncating a conversation when a turn is +superseded removes that turn's traces with it. Deleting a conversation also +deletes, by ``conversation_id``, the traces that have no message (scheduled +runs, stateless ``/v1`` rounds, turns whose message was never reserved); the +conversation index serves that. The link ids (``request_id``, +``activity_id``, ``workflow_run_id``) carry partial indexes because only the +Logs rows that have them look traces up by them. Traces with no log row of +their own (searches, graph builds) are listed per user or agent and source, +hence the source-leading indexes. ``status`` is the only CHECK; span kinds and sources are validated in code so new ones need no migration. @@ -70,10 +75,12 @@ def upgrade() -> None: ON request_traces (activity_id) WHERE activity_id IS NOT NULL; CREATE INDEX IF NOT EXISTS request_traces_workflow_run_idx ON request_traces (workflow_run_id) WHERE workflow_run_id IS NOT NULL; - CREATE INDEX IF NOT EXISTS request_traces_user_started_idx - ON request_traces (user_id, started_at DESC); - CREATE INDEX IF NOT EXISTS request_traces_agent_started_idx - ON request_traces (agent_id, started_at DESC) WHERE agent_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_conversation_idx + ON request_traces (conversation_id) WHERE conversation_id IS NOT NULL; + CREATE INDEX IF NOT EXISTS request_traces_user_source_started_idx + ON request_traces (user_id, source, started_at DESC); + CREATE INDEX IF NOT EXISTS request_traces_agent_source_started_idx + ON request_traces (agent_id, source, started_at DESC) WHERE agent_id IS NOT NULL; CREATE INDEX IF NOT EXISTS request_traces_created_idx ON request_traces (created_at); """ diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index 067dfb4f..ebd25567 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -827,8 +827,8 @@ Index("ix_guardrail_events_created", guardrail_events_table.c.created_at) # One execution trace per request (chat turn, continuation, scheduled or # webhook run, search, graph extraction): the span tree as a JSONB array, -# rendered as a waterfall in the Logs UI. ``message_id`` cascades so deleting -# or truncating a conversation removes its traces. Migration 0037. +# rendered as a waterfall in the Logs UI. ``message_id`` cascades; deleting a +# conversation also deletes its traces by ``conversation_id``. Migration 0037. request_traces_table = Table( "request_traces", metadata, @@ -858,7 +858,12 @@ request_traces_table = Table( Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), ) -Index("request_traces_user_started_idx", request_traces_table.c.user_id, request_traces_table.c.started_at) +Index( + "request_traces_user_source_started_idx", + request_traces_table.c.user_id, + request_traces_table.c.source, + request_traces_table.c.started_at, +) Index("request_traces_created_idx", request_traces_table.c.created_at) tool_call_attempts_table = Table( diff --git a/docsgpt/storage/db/repositories/conversations.py b/docsgpt/storage/db/repositories/conversations.py index 95205350..7d27a3f5 100644 --- a/docsgpt/storage/db/repositories/conversations.py +++ b/docsgpt/storage/db/repositories/conversations.py @@ -576,6 +576,18 @@ class ConversationsRepository: artifacts = ArtifactsRepository(self._conn) paths = artifacts.storage_paths_for_conversation(conversation_id) + # Traces linked to a message go with it (ON DELETE CASCADE). This + # takes the rest -- scheduled runs, stateless /v1 rounds, turns whose + # message was never reserved -- whose tool results and chunk previews + # the user expects gone with the conversation. + self._conn.execute( + text( + "DELETE FROM request_traces WHERE conversation_id = CAST(:id AS uuid) " + "AND EXISTS (SELECT 1 FROM conversations " + "WHERE id = CAST(:id AS uuid) AND user_id = :user_id)" + ), + {"id": conversation_id, "user_id": user_id}, + ) result = self._conn.execute( text( "DELETE FROM conversations " @@ -597,6 +609,14 @@ class ConversationsRepository: artifacts = ArtifactsRepository(self._conn) paths = artifacts.storage_paths_for_user_conversations(user_id) artifacts.delete_for_user_conversations(user_id) + # See ``delete``: traces not linked to a message are removed explicitly. + self._conn.execute( + text( + "DELETE FROM request_traces WHERE conversation_id IN " + "(SELECT id FROM conversations WHERE user_id = :user_id)" + ), + {"user_id": user_id}, + ) result = self._conn.execute( text("DELETE FROM conversations WHERE user_id = :user_id"), {"user_id": user_id}, diff --git a/tests/storage/db/repositories/test_request_traces.py b/tests/storage/db/repositories/test_request_traces.py index 0dc72b05..7d380188 100644 --- a/tests/storage/db/repositories/test_request_traces.py +++ b/tests/storage/db/repositories/test_request_traces.py @@ -141,6 +141,27 @@ class TestMessageLink: assert count == 0 +class TestConversationDeletion: + """Every trace of a conversation goes with it, not only message-linked ones.""" + + def test_trace_without_message_is_deleted_with_its_conversation(self, pg_conn): + conv_id, _msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + # e.g. a scheduled run in this conversation, or a stateless /v1 round. + repo.insert(_record(conversation_id=conv_id, source="schedule")) + repo.insert(_record(request_id="other")) # another conversation's trace + ConversationsRepository(pg_conn).delete(conv_id, "u1") + remaining = pg_conn.execute(text("SELECT request_id FROM request_traces")).scalars().all() + assert remaining == ["other"] + + def test_delete_all_for_user_removes_their_conversation_traces(self, pg_conn): + conv_id, _msg_id = _message(pg_conn) + repo = RequestTracesRepository(pg_conn) + repo.insert(_record(conversation_id=conv_id)) + ConversationsRepository(pg_conn).delete_all_for_user("u1") + assert pg_conn.execute(text("SELECT count(*) FROM request_traces")).scalar() == 0 + + class TestSummaries: def test_summaries_grouped_by_field_and_id(self, pg_conn): repo = RequestTracesRepository(pg_conn) diff --git a/tests/storage/db/test_migration_0037.py b/tests/storage/db/test_migration_0037.py index bb453df2..e8838a74 100644 --- a/tests/storage/db/test_migration_0037.py +++ b/tests/storage/db/test_migration_0037.py @@ -63,3 +63,16 @@ class TestMigration0037RoundTrip: "VALUES (gen_random_uuid(), 'stream', 'weird', now())" ) ) + + def test_listing_and_deletion_indexes_exist(self, pg_engine): + with pg_engine.connect() as conn: + names = set( + conn.execute( + text("SELECT indexname FROM pg_indexes WHERE tablename = 'request_traces'") + ).scalars() + ) + assert { + "request_traces_conversation_idx", + "request_traces_user_source_started_idx", + "request_traces_agent_source_started_idx", + } <= names From a74420ed2c1a3fe95e1ccd3cebdbef1f9757da66 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:42:05 +0100 Subject: [PATCH 30/33] Trace unexecuted tool calls in one place; tidy span lifecycles Paused, denied, skipped and client-run tool calls all go through trace_unexecuted_tool_call. The research synthesis, continuation and workflow step spans use with-blocks instead of hand-rolled GeneratorExit handling. Fixes blank lines and adds a missing hint and docstrings. --- docsgpt/agents/base.py | 21 ++++-------------- docsgpt/agents/research_agent.py | 22 ++++++------------- docsgpt/agents/tool_executor.py | 16 ++++++++++++++ docsgpt/agents/workflows/workflow_engine.py | 7 +----- docsgpt/guardrails/engine.py | 6 +++-- docsgpt/llm/handlers/base.py | 18 ++++----------- .../storage/db/repositories/request_traces.py | 2 ++ 7 files changed, 38 insertions(+), 54 deletions(-) diff --git a/docsgpt/agents/base.py b/docsgpt/agents/base.py index 3dea7254..4f620236 100644 --- a/docsgpt/agents/base.py +++ b/docsgpt/agents/base.py @@ -6,11 +6,9 @@ from abc import ABC, abstractmethod from datetime import datetime, timezone from typing import Any, Dict, Generator, List, Optional -from docsgpt import tracing from docsgpt.agents.tool_executor import ( ToolExecutor, - finish_tool_span, - record_tool_span_start, + trace_unexecuted_tool_call, result_status, truncate_tool_result, ) @@ -504,18 +502,10 @@ class BaseAgent(ABC): pending_tool_calls: The pending tool call descriptors from the pause. tool_actions: Client-provided actions resolving the pending calls. """ - span = start_agent_span(self, continuation=True) - completed = False - try: + with start_agent_span(self, continuation=True): yield from self._gen_continuation_inner( messages, tools_dict, pending_tool_calls, tool_actions, reasoning_content ) - completed = True - except Exception as exc: - span.end(error=exc) - raise - finally: - span.end(None if completed else tracing.STATUS_CANCELLED) def _gen_continuation_inner( self, @@ -617,7 +607,7 @@ class BaseAgent(ABC): "arguments": args, "status": "denied", } - finish_tool_span(record_tool_span_start(tc), {**denied_data, "error": comment or None}) + trace_unexecuted_tool_call(tc, {**denied_data, "error": comment or None}) yield {"type": "tool_call", "data": denied_data} elif "result" in action: @@ -650,10 +640,7 @@ class BaseAgent(ABC): "result": truncate_tool_result(result_str), "status": result_status(result), } - finish_tool_span( - record_tool_span_start(tc, **{"docsgpt.client_executed": True}), - client_data, - ) + trace_unexecuted_tool_call(tc, client_data, **{"docsgpt.client_executed": True}) yield {"type": "tool_call", "data": client_data} # Resume the LLM loop with the updated messages diff --git a/docsgpt/agents/research_agent.py b/docsgpt/agents/research_agent.py index c4858ec0..859bc872 100644 --- a/docsgpt/agents/research_agent.py +++ b/docsgpt/agents/research_agent.py @@ -47,12 +47,6 @@ STEP_PROMPT = _load_prompt("step.txt") SYNTHESIS_PROMPT = _load_prompt("synthesis.txt") -# --------------------------------------------------------------------------- -# CitationManager -# --------------------------------------------------------------------------- - - - 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( @@ -61,6 +55,12 @@ def _phase_span(phase: str, **attributes: Any) -> Any: attributes={"docsgpt.research.phase": phase.split(" ")[0], **attributes}, ) + +# --------------------------------------------------------------------------- +# CitationManager +# --------------------------------------------------------------------------- + + class CitationManager: """Tracks and deduplicates citations across research steps.""" @@ -266,18 +266,10 @@ class ResearchAgent(BaseAgent): "tokens_used": self._tokens_used, }, } - synthesis_span = _phase_span("synthesis") - try: + with _phase_span("synthesis"): yield from self._synthesis_phase( query, plan, intermediate_reports, tools_dict, log_context ) - except GeneratorExit: - synthesis_span.end(tracing.STATUS_CANCELLED) - raise - except Exception as exc: - synthesis_span.end(error=exc) - raise - synthesis_span.end() # Sources and tool calls self.retrieved_docs = self.citations.get_all_docs() diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 8e2b65e1..bc7f809a 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -57,6 +57,22 @@ _TOOL_SPAN_STATUS = { } +def trace_unexecuted_tool_call(call: Any, data: Dict[str, Any], **attributes: Any) -> None: + """Record a tool call that this process did not execute. + + The single entry point for calls that never reach :meth:`ToolExecutor.execute` + (which traces executed calls itself): calls paused for approval or for the + client, denied ones, ones skipped at the context limit, and results a client + sends back on resume. A new pause or refusal path records its call here. + + Args: + call: The tool call (anything with ``name`` and ``id``). + data: The tool-call record emitted for it (``status``, ``arguments``, ...). + **attributes: Extra span attributes. + """ + finish_tool_span(record_tool_span_start(call, **attributes), data) + + def finish_tool_span(span: Any, data: Dict[str, Any]) -> None: """Close a tool span from the ``tool_calls`` entry the call produced. diff --git a/docsgpt/agents/workflows/workflow_engine.py b/docsgpt/agents/workflows/workflow_engine.py index f296fc05..3b80dc81 100644 --- a/docsgpt/agents/workflows/workflow_engine.py +++ b/docsgpt/agents/workflows/workflow_engine.py @@ -164,12 +164,8 @@ class WorkflowEngine: } try: - try: + with step_span: yield from self._execute_node(node) - except GeneratorExit: - step_span.end(tracing.STATUS_CANCELLED) - raise - step_span.end() log_entry["status"] = ExecutionStatus.COMPLETED.value self._finalize_log_entry(log_entry, pre_state) @@ -186,7 +182,6 @@ class WorkflowEngine: "output": node_output, } except Exception as e: - step_span.end(error=e) logger.error(f"Error executing node {node.id}: {e}", exc_info=True) log_entry["status"] = ExecutionStatus.FAILED.value log_entry["error"] = str(e) diff --git a/docsgpt/guardrails/engine.py b/docsgpt/guardrails/engine.py index 834edc36..5f42cf02 100644 --- a/docsgpt/guardrails/engine.py +++ b/docsgpt/guardrails/engine.py @@ -5,7 +5,7 @@ from __future__ import annotations import logging from threading import Thread from time import monotonic -from typing import List, Optional +from typing import Any, List, Optional, Sequence from docsgpt import tracing from docsgpt.guardrails.base import ScanContext @@ -27,7 +27,8 @@ logger = logging.getLogger(__name__) _MAX_WORKERS = 8 -def _start_guardrail_span(stage: Stage, controls) -> "tracing.Span": +def _start_guardrail_span(stage: Stage, controls: Sequence[Any]) -> "tracing.Span": + """Open a ``guardrail`` span for evaluating ``controls`` at ``stage``.""" return tracing.start_span( tracing.KIND_GUARDRAIL, f"guardrail {stage.value}", @@ -39,6 +40,7 @@ def _start_guardrail_span(stage: Stage, controls) -> "tracing.Span": def _content_fired(decision: StageDecision) -> bool: + """True when a check matched, or the text was blocked or redacted.""" return bool(decision.triggered or decision.blocked or decision.redacted) diff --git a/docsgpt/llm/handlers/base.py b/docsgpt/llm/handlers/base.py index c52f1823..0954dd53 100644 --- a/docsgpt/llm/handlers/base.py +++ b/docsgpt/llm/handlers/base.py @@ -5,6 +5,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Any, Dict, Generator, List, Optional, Union +from docsgpt.agents.tool_executor import trace_unexecuted_tool_call from docsgpt.logging import build_stack_data logger = logging.getLogger(__name__) @@ -97,17 +98,6 @@ class LLMResponse: return bool(self.tool_calls) and self.finish_reason == "tool_calls" -def _trace_unexecuted_tool(call: Any, data: Dict[str, Any]) -> None: - """Record a tool call that was paused, denied or skipped instead of run. - - Those calls never reach ``ToolExecutor.execute`` (which traces executed - calls), so without this the trace would not show them at all. - """ - from docsgpt.agents.tool_executor import finish_tool_span, record_tool_span_start - - finish_tool_span(record_tool_span_start(call), data) - - def _is_restated_payload(existing: str, incoming: str) -> bool: """True when ``incoming`` restates an already-complete ``existing`` payload. @@ -1138,7 +1128,7 @@ class LLMHandler(ABC): # Mark remaining tools as skipped for remaining_call in tool_calls[i:]: - _trace_unexecuted_tool( + trace_unexecuted_tool_call( remaining_call, {"tool_name": "system", "action_name": remaining_call.name, "status": "skipped"}, ) @@ -1235,7 +1225,7 @@ class LLMHandler(ABC): "error_type", "tool_not_allowed" ), } - _trace_unexecuted_tool(call, denied_data) + trace_unexecuted_tool_call(call, denied_data) yield {"type": "tool_call", "data": denied_data} continue # Yield pause event so the client knows this tool is waiting @@ -1250,7 +1240,7 @@ class LLMHandler(ABC): # can wire the sticky "don't ask again" button. if pause_info.get("device_id"): pause_data["device_id"] = pause_info["device_id"] - _trace_unexecuted_tool(call, pause_data) + trace_unexecuted_tool_call(call, pause_data) yield {"type": "tool_call", "data": pause_data} pending_actions.append(pause_info) # Do NOT add messages for pending tools here. diff --git a/docsgpt/storage/db/repositories/request_traces.py b/docsgpt/storage/db/repositories/request_traces.py index 1fb5c594..79a7e3cb 100644 --- a/docsgpt/storage/db/repositories/request_traces.py +++ b/docsgpt/storage/db/repositories/request_traces.py @@ -33,10 +33,12 @@ _SUMMARY_COLUMNS = ( def _dump_jsonb(value: Any) -> str: + """Serialize ``value`` for a JSONB parameter, without NUL bytes Postgres rejects.""" return json.dumps(strip_null_bytes(value), cls=PGNativeJSONEncoder) def _uuid_or_none(value: Optional[str]) -> Optional[str]: + """``value`` as a string when it is a UUID, else ``None`` (the column is typed).""" return str(value) if value and looks_like_uuid(str(value)) else None From 6b145d7aad276b27232b9b890a10031615a013e4 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:23:41 +0100 Subject: [PATCH 31/33] Wait for the background trace write in the /v1 replay test The first request's trace is written on the trace-writer pool after the response returns; assert once it has landed rather than racing it. --- tests/api/v1/test_v1_idempotency.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/api/v1/test_v1_idempotency.py b/tests/api/v1/test_v1_idempotency.py index eda35b4d..7715a886 100644 --- a/tests/api/v1/test_v1_idempotency.py +++ b/tests/api/v1/test_v1_idempotency.py @@ -19,6 +19,7 @@ These tests pin both layers against an ephemeral Postgres (``pg_engine``): from __future__ import annotations +import time import uuid from contextlib import contextmanager from typing import Any, Dict, List @@ -478,5 +479,12 @@ class TestV1ReplayTrace: resp2 = _post_chat(c, self.QUESTION, api_key, idem_key=idem_key) assert resp1.status_code == resp2.status_code == 200 + # The first request's trace is written on the background writer, so + # wait for it. A replay's trace is dropped synchronously, and an + # unclaimed-trace flush would be synchronous too, so any stray write + # is already in ``written`` by the time the responses return. + deadline = time.monotonic() + 5 + while not written and time.monotonic() < deadline: + time.sleep(0.02) # Only the first request's trace is written, and it is not a failure. assert [t.status for t in written] == ["ok"] From 3156ee37a194acebdf6649c3374544a0c1b22672 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Thu, 24 Sep 2026 00:08:53 +0100 Subject: [PATCH 32/33] Name the provider a model call actually reached DeepSeek and every other OpenAI-compatible API run through OpenAILLM, so spans and metrics reported them as openai. The provider now comes from the endpoint: known API hosts map to their gen_ai.provider.name (deepseek, azure.ai.openai, mistral_ai, ...), and an OpenAI client pointed anywhere else is openai_compatible. LLM spans also record server.address. --- docsgpt/llm/llm_creator.py | 4 ++ docsgpt/tracing/llm.py | 66 ++++++++++++++++++++++++++++++++- tests/tracing/test_llm_spans.py | 51 +++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 2 deletions(-) diff --git a/docsgpt/llm/llm_creator.py b/docsgpt/llm/llm_creator.py index 995600cb..d64a37a2 100644 --- a/docsgpt/llm/llm_creator.py +++ b/docsgpt/llm/llm_creator.py @@ -128,6 +128,10 @@ class LLMCreator: # llm.model_id is the upstream name (BYOM resolves it above); stamp # the canonical id (UUID for BYOM) separately for token_usage. llm._canonical_model_id = model_id + # The provider plugin that built it: ``openai_compatible`` endpoints + # run through the OpenAI client, so the class alone cannot name them + # in traces (see ``docsgpt.tracing.llm.llm_provider``). + llm._provider_plugin = plugin.name # Calls to a user's own model are recorded at $0 (see ``docsgpt/usage.py``). llm._is_byom = model is not None and getattr(model, "source", "builtin") == "user" return llm diff --git a/docsgpt/tracing/llm.py b/docsgpt/tracing/llm.py index 556fc373..ef169186 100644 --- a/docsgpt/tracing/llm.py +++ b/docsgpt/tracing/llm.py @@ -8,6 +8,7 @@ fallback each get their own span with the provider that actually ran. from __future__ import annotations from typing import Any, Dict, Iterable, Optional, Union +from urllib.parse import urlparse from docsgpt.core.settings import settings from docsgpt.tracing import core @@ -16,6 +17,66 @@ 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" +#: API hosts whose provider differs from the client class that calls them +#: (every OpenAI-compatible API runs through ``OpenAILLM``), mapped to their +#: ``gen_ai.provider.name``. Matched on the host or any subdomain of it. +_PROVIDER_HOSTS = ( + ("api.openai.com", "openai"), + ("openai.azure.com", "azure.ai.openai"), + ("deepseek.com", "deepseek"), + ("mistral.ai", "mistral_ai"), + ("x.ai", "x_ai"), + ("perplexity.ai", "perplexity"), + ("groq.com", "groq"), + ("openrouter.ai", "openrouter"), + ("novita.ai", "novita"), + ("cohere.com", "cohere"), + ("cohere.ai", "cohere"), + ("together.xyz", "together_ai"), + ("together.ai", "together_ai"), + ("fireworks.ai", "fireworks"), + ("anthropic.com", "anthropic"), + ("generativelanguage.googleapis.com", "gcp.gemini"), +) + + +def _endpoint_host(llm: Any) -> Optional[str]: + url = getattr(llm, "_effective_base_url", None) or getattr(llm, "base_url", None) + if not isinstance(url, str) or not url: + return None + try: + return urlparse(url).hostname + except ValueError: + return None + + +def llm_provider(llm: Any) -> str: + """The ``gen_ai.provider.name`` of the API ``llm`` actually calls. + + The client class alone is not enough: DeepSeek, a self-hosted vLLM or any + other OpenAI-compatible endpoint runs through ``OpenAILLM``, whose + ``provider_name`` is ``openai``. A known API host wins; an OpenAI client + pointed at any other host is reported as ``openai_compatible``. + + Args: + llm: The LLM instance making the call. + + Returns: + The provider name for spans and metrics. + """ + base = provider_name(getattr(llm, "provider_name", None)) + host = _endpoint_host(llm) + if host: + for suffix, name in _PROVIDER_HOSTS: + if host == suffix or host.endswith("." + suffix): + return name + if getattr(llm, "_provider_plugin", None) == "openai_compatible": + return "openai_compatible" + if base == "openai" and host: + # An OpenAI client aimed somewhere other than OpenAI's API. + return "openai_compatible" + return base + def start_llm_span( llm: Any, model: Optional[str], *, stream: bool, tools: Any = None @@ -25,7 +86,8 @@ def start_llm_span( return core.NOOP_SPAN attributes = { "gen_ai.operation.name": "chat", - "gen_ai.provider.name": provider_name(getattr(llm, "provider_name", None)), + "gen_ai.provider.name": llm_provider(llm), + "server.address": _endpoint_host(llm), "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), @@ -83,7 +145,7 @@ def finish_llm_call( if not settings.TRACES_ENABLED: return record_llm_metrics( - provider=getattr(llm, "provider_name", None), + provider=llm_provider(llm), model=str(model) if model else None, input_tokens=call_usage.get("prompt_tokens", 0), output_tokens=call_usage.get("generated_tokens", 0), diff --git a/tests/tracing/test_llm_spans.py b/tests/tracing/test_llm_spans.py index 897f0a89..966719de 100644 --- a/tests/tracing/test_llm_spans.py +++ b/tests/tracing/test_llm_spans.py @@ -226,3 +226,54 @@ class TestCacheHits: first, second = trace.spans assert first.attributes.get("docsgpt.cache_hit") is None assert second.attributes["docsgpt.cache_hit"] is True + + +class TestProviderResolution: + """The span names the provider actually called, not the client class used.""" + + @staticmethod + def _llm(provider="openai", base_url=None, plugin=None): + llm = _LLM() + llm.provider_name = provider + if base_url is not None: + llm._effective_base_url = base_url + if plugin is not None: + llm._provider_plugin = plugin + return llm + + @pytest.mark.parametrize( + "base_url, plugin, expected", + [ + ("https://api.deepseek.com/v1", "openai_compatible", "deepseek"), + ("https://my-res.openai.azure.com/openai", "openai", "azure.ai.openai"), + ("https://api.mistral.ai/v1", "openai_compatible", "mistral_ai"), + ("https://api.x.ai/v1", "openai_compatible", "x_ai"), + ("https://api.openai.com/v1", "openai", "openai"), + ("http://10.0.0.5:8000/v1", "openai_compatible", "openai_compatible"), + ("http://127.0.0.1:7899/v1", "openai", "openai_compatible"), + (None, "openai", "openai"), + ], + ) + def test_provider_from_endpoint(self, base_url, plugin, expected): + from docsgpt.tracing.llm import llm_provider + + assert llm_provider(self._llm(base_url=base_url, plugin=plugin)) == expected + + def test_native_providers_are_unchanged(self): + from docsgpt.tracing.llm import llm_provider + + assert llm_provider(self._llm(provider="anthropic")) == "anthropic" + assert llm_provider(self._llm(provider="google")) == "gcp.gen_ai" + + def test_span_and_metrics_use_it_and_record_server_address(self, trace, metrics): + llm = self._llm(base_url="https://api.deepseek.com/v1", plugin="openai_compatible") + + @gen_token_usage + def _gen(self, model, messages, stream, tools, **kwargs): + return "x" + + _gen(llm, "deepseek-chat", [], False, None) + attrs = trace.spans[0].attributes + assert attrs["gen_ai.provider.name"] == "deepseek" + assert attrs["server.address"] == "api.deepseek.com" + assert metrics.call_args.kwargs["provider"] == "deepseek" From f2d92bf8df882075aa41aacb60b3532b465c7ec3 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Thu, 24 Sep 2026 00:08:53 +0100 Subject: [PATCH 33/33] Log failed chat turns as chat entries A turn whose agent raised wrote no user_logs row, so it only surfaced as the agent's system error row. Every finished turn now writes its chat row, at level error with the error when it failed, and linked to its trace; the system row for the same traced activity is no longer listed twice. --- docsgpt/api/answer/routes/base.py | 147 +++++++++++++++++++------- docsgpt/api/user/analytics/routes.py | 13 +++ tests/api/answer/test_trace_wiring.py | 46 ++++++++ tests/api/user/test_traces_routes.py | 36 +++++++ 4 files changed, 203 insertions(+), 39 deletions(-) diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index 782e11b5..ca92d3cf 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -144,6 +144,79 @@ class BaseAnswerResource: return missing_fields return None + def _persist_turn_log( + self, + *, + decoded_token: Dict[str, Any], + user_api_key: Optional[str], + agent_id: Optional[str], + question: str, + response: str, + sources: List[Dict[str, Any]], + tool_calls: Any, + attachment_ids: Optional[List[str]], + request_id: Optional[str], + message_id: Optional[str], + error: Optional[str] = None, + extra: Optional[Dict[str, Any]] = None, + ) -> None: + """Write the turn's ``user_logs`` row: its Logs entry, linked to its trace. + + Written for every finished turn, failed ones included (``level`` is + ``error`` then), so a failed chat shows up as a chat entry. A failure + to write is logged, never raised. + + Args: + decoded_token: The caller's token. + user_api_key: The agent API key the request used, if any. + agent_id: The agent that answered. + question: The question as stored (after input redaction). + response: The answer, or what streamed before a failure. + sources: Retrieved sources. + tool_calls: The turn's tool calls, before log truncation. + attachment_ids: Attached file ids. + request_id: The turn's request id (links its trace). + message_id: The reserved message id. + error: What failed the turn, if it failed. + extra: More fields for the row (structured-output details). + """ + log_data: Dict[str, Any] = { + "action": "stream_answer", + "level": "error" if error else "info", + "user": decoded_token.get("sub"), + "api_key": user_api_key, + "agent_id": agent_id, + "question": question, + "response": response, + "sources": sources, + "tool_calls": self._prepare_tool_calls_for_logging(tool_calls), + "attachments": attachment_ids, + "timestamp": datetime.datetime.now(datetime.timezone.utc), + # Links the Logs row to this turn's execution trace. + "request_id": request_id, + "message_id": message_id, + } + if error: + log_data["error"] = error + if extra: + log_data.update(extra) + # Clean up text fields to be no longer than 10000 characters. + for key, value in log_data.items(): + if isinstance(value, str) and len(value) > 10000: + log_data[key] = value[:10000] + try: + with db_session() as conn: + UserLogsRepository(conn).insert( + user_id=log_data.get("user"), + endpoint="stream_answer", + data=log_data, + ) + except Exception as log_err: + logger.error( + f"Failed to persist stream_answer user log: {log_err}", + exc_info=True, + ) + @staticmethod def _prepare_tool_calls_for_logging( tool_calls: Optional[List[Dict[str, Any]]], max_chars: int = 10000 @@ -1299,47 +1372,27 @@ class BaseAnswerResource: ) yield _emit({"type": "id", "id": str(conversation_id)}) - tool_calls_for_logging = self._prepare_tool_calls_for_logging( - getattr(agent, "tool_calls", tool_calls) or tool_calls - ) - - log_data = { - "action": "stream_answer", - "level": "info", - "user": decoded_token.get("sub"), - "api_key": user_api_key, - "agent_id": agent_id, - "question": question, - "response": response_full, - "sources": source_log_docs, - "tool_calls": tool_calls_for_logging, - "attachments": attachment_ids, - "timestamp": datetime.datetime.now(datetime.timezone.utc), - # Links the Logs row to this turn's execution trace. - "request_id": request_id, - "message_id": reserved_message_id, - } + extra: Dict[str, Any] = {} if is_structured: - log_data["structured_output"] = True + extra["structured_output"] = True if schema_info: - log_data["schema"] = schema_info - # Clean up text fields to be no longer than 10000 characters - - for key, value in log_data.items(): - if isinstance(value, str) and len(value) > 10000: - log_data[key] = value[:10000] - try: - with db_session() as conn: - UserLogsRepository(conn).insert( - user_id=log_data.get("user"), - endpoint="stream_answer", - data=log_data, - ) - except Exception as log_err: - logger.error( - f"Failed to persist stream_answer user log: {log_err}", - exc_info=True, - ) + extra["schema"] = schema_info + self._persist_turn_log( + decoded_token=decoded_token, + user_api_key=user_api_key, + agent_id=agent_id, + question=question, + response=response_full, + sources=source_log_docs, + tool_calls=getattr(agent, "tool_calls", tool_calls) or tool_calls, + attachment_ids=attachment_ids, + request_id=request_id, + message_id=reserved_message_id, + # A yielded error (a failed workflow node) ends the turn + # normally but still failed it. + error=stream_error, + extra=extra, + ) yield _emit({"type": "end"}) # Drain the journal buffer so the terminal ``end`` event is @@ -1617,6 +1670,22 @@ class BaseAnswerResource: f"Failed to finalize errored message: {fin_err}", exc_info=True, ) + # A failed turn is still a chat turn: log it as one (level + # ``error``), with its trace link, instead of leaving only the + # agent's system error row. + self._persist_turn_log( + decoded_token=decoded_token, + user_api_key=user_api_key, + agent_id=agent_id, + question=question, + response=response_full, + sources=source_log_docs, + tool_calls=getattr(agent, "tool_calls", tool_calls) or tool_calls, + attachment_ids=attachment_ids, + request_id=request_id, + message_id=reserved_message_id, + error=f"{type(e).__name__}: {e}", + ) yield _emit( { "type": "error", diff --git a/docsgpt/api/user/analytics/routes.py b/docsgpt/api/user/analytics/routes.py index c77d9eb1..3642ed3a 100644 --- a/docsgpt/api/user/analytics/routes.py +++ b/docsgpt/api/user/analytics/routes.py @@ -943,6 +943,16 @@ class GetUserLogs(Resource): "WHERE e.activity_id = s.activity_id " "AND e.level = 'error'))" ) + # A failed chat turn now writes its own chat row (level + # ``error``) linked to its trace; the agent's error row for + # the same activity would list the failure twice. Rows from + # before execution traces, or with tracing off, have no such + # trace and still show here. + chat_failure_dedupe = ( + "NOT EXISTS (SELECT 1 FROM request_traces t " + "WHERE t.activity_id = s.activity_id " + "AND t.source IN ('stream', 'answer', 'v1'))" + ) if api_key_id: # The owner-scoped lookup gates access, so the # chat/webhook/system branches match on the agent @@ -978,6 +988,7 @@ class GetUserLogs(Resource): "s.level = 'error'", "COALESCE(s.endpoint, '') NOT IN ('webhook', 'schedule')", stack_agent_match, + chat_failure_dedupe, ] # Owner-gated agent match: drop the user clause so a # shared agent's runs (stamped with the caller's @@ -1003,6 +1014,7 @@ class GetUserLogs(Resource): "s.user_id = :user_id", "s.level = 'error'", "COALESCE(s.endpoint, '') NOT IN ('webhook', 'schedule')", + chat_failure_dedupe, ] # Terminal statuses only (worker writes ``success`` / # ``failed`` / ``timeout`` / ``skipped``; ``completed`` @@ -1233,6 +1245,7 @@ class GetUserLogs(Resource): "attachments": payload.get("attachments"), "request_id": payload.get("request_id"), "message_id": payload.get("message_id"), + "error": payload.get("error"), } ) elif m["event_type"] in ("system", "webhook"): diff --git a/tests/api/answer/test_trace_wiring.py b/tests/api/answer/test_trace_wiring.py index d03dec01..01be1475 100644 --- a/tests/api/answer/test_trace_wiring.py +++ b/tests/api/answer/test_trace_wiring.py @@ -185,6 +185,52 @@ class TestTraceWithPersistence: assert row[0]["request_id"] == "req-p" assert row[0]["message_id"] == trace.message_id + def test_failed_turn_is_logged_as_a_chat_row(self, pg_conn, flask_app): + """A raised failure still writes the turn's chat row, at level error.""" + from docsgpt.api.answer.routes.base import BaseAnswerResource + from tests.api.answer.test_base_routes import _patch_db_session + + agent = MagicMock() + agent.gen.side_effect = RuntimeError("upstream down") + agent.tool_calls = [] + with flask_app.app_context(), _patch_db_session(pg_conn), _captured_flushes(): + _run( + BaseAnswerResource(), + agent, + should_persist=True, + model_id="gpt-4", + request_id="req-failed", + ) + from sqlalchemy import text as sql_text + + rows = pg_conn.execute( + sql_text("SELECT data FROM user_logs WHERE user_id = 'u-trace'") + ).fetchall() + assert len(rows) == 1 + data = rows[0][0] + assert data["level"] == "error" + assert data["request_id"] == "req-failed" + assert data["error"] == "RuntimeError: upstream down" + + def test_yielded_error_logs_the_chat_row_at_error_level(self, pg_conn, flask_app): + from docsgpt.api.answer.routes.base import BaseAnswerResource + from tests.api.answer.test_base_routes import _patch_db_session + + with flask_app.app_context(), _patch_db_session(pg_conn), _captured_flushes(): + _run( + BaseAnswerResource(), + _agent([{"type": "error", "error": "node failed"}]), + should_persist=True, + model_id="gpt-4", + ) + from sqlalchemy import text as sql_text + + data = pg_conn.execute( + sql_text("SELECT data FROM user_logs WHERE user_id = 'u-trace'") + ).fetchone()[0] + assert data["level"] == "error" + assert data["error"] == "node failed" + def test_paused_turn_is_flushed_paused(self, pg_conn, flask_app): from docsgpt.api.answer.routes.base import BaseAnswerResource from tests.api.answer.test_base_routes import _patch_db_session diff --git a/tests/api/user/test_traces_routes.py b/tests/api/user/test_traces_routes.py index 7c89d864..bbafdd79 100644 --- a/tests/api/user/test_traces_routes.py +++ b/tests/api/user/test_traces_routes.py @@ -212,3 +212,39 @@ class TestSummaryFailure: assert response.status_code == 200 (row,) = response.json["logs"] assert "trace" not in row + + +class TestFailedChatTurns: + def _failed_turn(self, pg_conn, *, with_trace): + from docsgpt.storage.db.repositories.stack_logs import StackLogsRepository + from docsgpt.storage.db.repositories.user_logs import UserLogsRepository + + StackLogsRepository(pg_conn).insert( + activity_id="act-chat", endpoint="stream", level="error", user_id="owner", query="q" + ) + UserLogsRepository(pg_conn).insert( + user_id="owner", + endpoint="stream_answer", + data={ + "question": "q", + "level": "error", + "error": "RuntimeError: upstream down", + "request_id": "req-f", + }, + ) + if with_trace: + _trace(pg_conn, request_id="req-f", activity_id="act-chat", status="error") + + def test_failed_turn_is_one_chat_row_not_also_a_system_row(self, app, pg_conn): + self._failed_turn(pg_conn, with_trace=True) + rows = _logs(app, pg_conn, "owner", {}).json["logs"] + assert [r["event_type"] for r in rows] == ["chat"] + assert rows[0]["level"] == "error" + assert rows[0]["error"] == "RuntimeError: upstream down" + assert rows[0]["trace"]["status"] == "error" + + def test_untraced_failures_keep_their_system_row(self, app, pg_conn): + """Rows from before traces (or with tracing off) have nothing to dedupe against.""" + self._failed_turn(pg_conn, with_trace=False) + rows = _logs(app, pg_conn, "owner", {}).json["logs"] + assert sorted(r["event_type"] for r in rows) == ["chat", "system"]