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.
This commit is contained in:
arc53-machine committed 2026-09-23 17:34:54 +01:00
1 parent da1bc16eb6
commit bae842d151
9 files changed
+415 -9

No files matched your search

+4 -1
View File
@@ -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)
+80 -2
View File
@@ -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
+4 -1
View File
@@ -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",
@@ -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":
+5 -1
View File
@@ -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)
+5
View File
@@ -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:
+248
View File
@@ -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
+4
View File
@@ -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
+14
View File
@@ -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"