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