mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 12:13:38 +00:00
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.
This commit is contained in:
1 parent
f9605e7794
commit
da1bc16eb6
7 files changed
+256
-12
No files matched your search
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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
|
||||
Reference in new issue
Block a user