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