mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 16:14:10 +00:00
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:
1 parent
da1bc16eb6
commit
bae842d151
9 files changed
+415
-9
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
Reference in new issue
Block a user