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] 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"