mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-03 22:13:01 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
1084 lines
39 KiB
Python
1084 lines
39 KiB
Python
"""Tests covering gaps in WorkflowEngine: execute loop, state/condition/end nodes,
|
|
template context, source data, structured output parsing, get_execution_summary."""
|
|
|
|
from datetime import datetime, timezone
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from docsgpt.agents.workflows.schemas import (
|
|
ExecutionStatus,
|
|
NodeType,
|
|
WorkflowEdge,
|
|
WorkflowGraph,
|
|
WorkflowNode,
|
|
Workflow,
|
|
)
|
|
from docsgpt.agents.workflows.workflow_engine import WorkflowEngine
|
|
|
|
|
|
def _make_graph(nodes, edges):
|
|
wf = Workflow(name="Test", description="test workflow")
|
|
return WorkflowGraph(workflow=wf, nodes=nodes, edges=edges)
|
|
|
|
|
|
def _make_node(id, type, title="Node", config=None, position=None):
|
|
return WorkflowNode(
|
|
id=id,
|
|
workflow_id="wf1",
|
|
type=type,
|
|
title=title,
|
|
position=position or {"x": 0, "y": 0},
|
|
config=config or {},
|
|
)
|
|
|
|
|
|
def _make_edge(id, source, target, source_handle=None, target_handle=None):
|
|
return WorkflowEdge(
|
|
id=id,
|
|
workflow_id="wf1",
|
|
source=source,
|
|
target=target,
|
|
sourceHandle=source_handle,
|
|
targetHandle=target_handle,
|
|
)
|
|
|
|
|
|
def _make_agent():
|
|
agent = MagicMock()
|
|
agent.chat_history = []
|
|
agent.endpoint = "https://api.example.com"
|
|
agent.llm_name = "openai"
|
|
agent.model_id = "gpt-4"
|
|
agent.api_key = "key"
|
|
agent.decoded_token = {"sub": "user1"}
|
|
agent.retrieved_docs = None
|
|
return agent
|
|
|
|
|
|
class TestExecuteLoop:
|
|
|
|
@pytest.mark.unit
|
|
def test_no_start_node_yields_error(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
events = list(engine.execute({}, "query"))
|
|
assert any(e.get("type") == "error" and "start node" in e.get("error", "") for e in events)
|
|
|
|
@pytest.mark.unit
|
|
def test_emits_workflow_run_id_first(self):
|
|
"""A ``workflow_run`` event carrying the run id precedes any step event."""
|
|
nodes = [
|
|
_make_node("n1", NodeType.START, "Start"),
|
|
_make_node("n2", NodeType.END, "End", config={"config": {}}),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2")]
|
|
engine = WorkflowEngine(_make_graph(nodes, edges), _make_agent())
|
|
events = list(engine.execute({}, "hello"))
|
|
run_events = [e for e in events if e.get("type") == "workflow_run"]
|
|
assert len(run_events) == 1
|
|
assert run_events[0]["workflow_run_id"] == engine.workflow_run_id
|
|
assert events[0]["type"] == "workflow_run"
|
|
|
|
@pytest.mark.unit
|
|
def test_start_to_end(self):
|
|
nodes = [
|
|
_make_node("n1", NodeType.START, "Start"),
|
|
_make_node("n2", NodeType.END, "End", config={"config": {}}),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2")]
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
events = list(engine.execute({}, "hello"))
|
|
step_events = [e for e in events if e.get("type") == "workflow_step"]
|
|
assert len(step_events) >= 2 # At least start + end
|
|
|
|
@pytest.mark.unit
|
|
def test_node_not_found_yields_error(self):
|
|
nodes = [_make_node("n1", NodeType.START)]
|
|
edges = [_make_edge("e1", "n1", "nonexistent")]
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
events = list(engine.execute({}, "q"))
|
|
assert any("not found" in e.get("error", "") for e in events)
|
|
|
|
@pytest.mark.unit
|
|
def test_node_execution_error_yields_error(self):
|
|
nodes = [
|
|
_make_node("n1", NodeType.START),
|
|
_make_node("n2", NodeType.STATE, "State", config={"config": {"operations": [{"expression": "bad!!!", "target_variable": "x"}]}}),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2")]
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
events = list(engine.execute({}, "q"))
|
|
failed_events = [e for e in events if e.get("status") == "failed"]
|
|
assert len(failed_events) >= 1
|
|
|
|
@pytest.mark.unit
|
|
def test_max_steps_limit(self):
|
|
# Create a cycle: start -> state -> state (loop)
|
|
nodes = [
|
|
_make_node("n1", NodeType.START),
|
|
_make_node("n2", NodeType.NOTE, "Note"),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2"), _make_edge("e2", "n2", "n2")]
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.MAX_EXECUTION_STEPS = 5
|
|
events = list(engine.execute({}, "q"))
|
|
# Should not run forever
|
|
assert len(events) > 0
|
|
|
|
@pytest.mark.unit
|
|
def test_branch_ends_without_end_node(self):
|
|
nodes = [
|
|
_make_node("n1", NodeType.START),
|
|
_make_node("n2", NodeType.NOTE, "Note"),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2")] # n2 has no outgoing edges
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
events = list(engine.execute({}, "q"))
|
|
assert len(events) > 0
|
|
|
|
|
|
class TestInitializeState:
|
|
|
|
@pytest.mark.unit
|
|
def test_sets_query_and_history(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.chat_history = [{"prompt": "hi", "response": "hey"}]
|
|
engine = WorkflowEngine(graph, agent)
|
|
engine._initialize_state({"custom": "value"}, "test query")
|
|
assert engine.state["query"] == "test query"
|
|
assert "custom" in engine.state
|
|
assert engine.state["chat_history"] is not None
|
|
|
|
|
|
class TestGetNextNodeId:
|
|
|
|
@pytest.mark.unit
|
|
def test_no_edges_returns_none(self):
|
|
nodes = [_make_node("n1", NodeType.START)]
|
|
graph = _make_graph(nodes, [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
assert engine._get_next_node_id("n1") is None
|
|
|
|
@pytest.mark.unit
|
|
def test_returns_first_edge_target(self):
|
|
nodes = [_make_node("n1", NodeType.START), _make_node("n2", NodeType.END)]
|
|
edges = [_make_edge("e1", "n1", "n2")]
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
assert engine._get_next_node_id("n1") == "n2"
|
|
|
|
@pytest.mark.unit
|
|
def test_condition_uses_matched_handle(self):
|
|
nodes = [
|
|
_make_node("n1", NodeType.CONDITION),
|
|
_make_node("n2", NodeType.END, "Yes End"),
|
|
_make_node("n3", NodeType.END, "No End"),
|
|
]
|
|
edges = [
|
|
_make_edge("e1", "n1", "n2", source_handle="yes"),
|
|
_make_edge("e2", "n1", "n3", source_handle="no"),
|
|
]
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine._condition_result = "no"
|
|
assert engine._get_next_node_id("n1") == "n3"
|
|
assert engine._condition_result is None # Cleared after use
|
|
|
|
@pytest.mark.unit
|
|
def test_condition_no_matching_handle_returns_none(self):
|
|
nodes = [_make_node("n1", NodeType.CONDITION)]
|
|
edges = [_make_edge("e1", "n1", "n2", source_handle="yes")]
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine._condition_result = "nonexistent"
|
|
assert engine._get_next_node_id("n1") is None
|
|
|
|
|
|
class TestExecuteStateNode:
|
|
|
|
@pytest.mark.unit
|
|
def test_evaluates_operations(self):
|
|
node = _make_node("n1", NodeType.STATE, config={
|
|
"config": {
|
|
"operations": [
|
|
{"expression": "x + 1", "target_variable": "result"},
|
|
]
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {"x": 5}
|
|
list(engine._execute_state_node(node))
|
|
assert engine.state["result"] == 6
|
|
|
|
@pytest.mark.unit
|
|
def test_skips_empty_expression(self):
|
|
node = _make_node("n1", NodeType.STATE, config={
|
|
"config": {
|
|
"operations": [
|
|
{"expression": "", "target_variable": "result"},
|
|
]
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {}
|
|
list(engine._execute_state_node(node))
|
|
assert "result" not in engine.state
|
|
|
|
|
|
class TestEndNodeSeparator:
|
|
|
|
@pytest.mark.unit
|
|
def test_separates_from_prior_streamed_output(self):
|
|
# A prior streaming node ran, so the end template must be preceded by a
|
|
# paragraph break instead of mashing ("...growth.Sales analysis...").
|
|
node = _make_node("n1", NodeType.END, config={"config": {"output_template": "Final line."}})
|
|
engine = WorkflowEngine(_make_graph([node], []), _make_agent())
|
|
engine.state = {}
|
|
engine._has_streamed = True
|
|
answers = [e["answer"] for e in engine._execute_end_node(node)]
|
|
assert answers[0] == "\n\n"
|
|
assert answers[-1].strip() == "Final line."
|
|
|
|
@pytest.mark.unit
|
|
def test_no_separator_when_nothing_streamed(self):
|
|
node = _make_node("n1", NodeType.END, config={"config": {"output_template": "Final line."}})
|
|
engine = WorkflowEngine(_make_graph([node], []), _make_agent())
|
|
engine.state = {}
|
|
answers = [e["answer"] for e in engine._execute_end_node(node)]
|
|
assert answers == ["Final line."]
|
|
|
|
|
|
class TestExecuteConditionNode:
|
|
|
|
@pytest.mark.unit
|
|
def test_matches_first_true_case(self):
|
|
node = _make_node("n1", NodeType.CONDITION, config={
|
|
"config": {
|
|
"cases": [
|
|
{"expression": "x > 10", "source_handle": "high"},
|
|
{"expression": "x > 5", "source_handle": "medium"},
|
|
]
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {"x": 7}
|
|
list(engine._execute_condition_node(node))
|
|
assert engine._condition_result == "medium"
|
|
|
|
@pytest.mark.unit
|
|
def test_falls_through_to_else(self):
|
|
node = _make_node("n1", NodeType.CONDITION, config={
|
|
"config": {
|
|
"cases": [
|
|
{"expression": "x > 100", "source_handle": "high"},
|
|
]
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {"x": 1}
|
|
list(engine._execute_condition_node(node))
|
|
assert engine._condition_result == "else"
|
|
|
|
@pytest.mark.unit
|
|
def test_skips_empty_expression(self):
|
|
node = _make_node("n1", NodeType.CONDITION, config={
|
|
"config": {
|
|
"cases": [
|
|
{"expression": " ", "source_handle": "a"},
|
|
{"expression": "true", "source_handle": "b"},
|
|
]
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {}
|
|
list(engine._execute_condition_node(node))
|
|
assert engine._condition_result == "b"
|
|
|
|
@pytest.mark.unit
|
|
def test_cel_error_continues(self):
|
|
node = _make_node("n1", NodeType.CONDITION, config={
|
|
"config": {
|
|
"cases": [
|
|
{"expression": "bad!!!", "source_handle": "a"},
|
|
{"expression": "true", "source_handle": "b"},
|
|
]
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {}
|
|
list(engine._execute_condition_node(node))
|
|
assert engine._condition_result == "b"
|
|
|
|
|
|
class TestExecuteEndNode:
|
|
|
|
@pytest.mark.unit
|
|
def test_with_output_template(self):
|
|
node = _make_node("n1", NodeType.END, config={
|
|
"config": {"output_template": "Result: {{ query }}"}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {"query": "hello"}
|
|
engine._format_template = MagicMock(return_value="Result: hello")
|
|
events = list(engine._execute_end_node(node))
|
|
assert len(events) == 1
|
|
assert events[0]["answer"] == "Result: hello"
|
|
|
|
@pytest.mark.unit
|
|
def test_without_output_template(self):
|
|
node = _make_node("n1", NodeType.END, config={"config": {}})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
events = list(engine._execute_end_node(node))
|
|
assert len(events) == 0
|
|
|
|
|
|
class TestParseStructuredOutput:
|
|
|
|
@pytest.mark.unit
|
|
def test_valid_json(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
success, data = engine._parse_structured_output('{"key": "value"}')
|
|
assert success is True
|
|
assert data == {"key": "value"}
|
|
|
|
@pytest.mark.unit
|
|
def test_invalid_json(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
success, data = engine._parse_structured_output("not json")
|
|
assert success is False
|
|
assert data is None
|
|
|
|
@pytest.mark.unit
|
|
def test_empty_string(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
success, data = engine._parse_structured_output("")
|
|
assert success is False
|
|
assert data is None
|
|
|
|
@pytest.mark.unit
|
|
def test_whitespace_only(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
success, data = engine._parse_structured_output(" ")
|
|
assert success is False
|
|
assert data is None
|
|
|
|
|
|
class TestNormalizeNodeJsonSchema:
|
|
|
|
@pytest.mark.unit
|
|
def test_none_returns_none(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
assert engine._normalize_node_json_schema(None, "Node") is None
|
|
|
|
@pytest.mark.unit
|
|
def test_valid_schema(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
schema = {"type": "object", "properties": {"name": {"type": "string"}}}
|
|
result = engine._normalize_node_json_schema(schema, "Node")
|
|
assert result is not None
|
|
|
|
@pytest.mark.unit
|
|
def test_invalid_schema_raises(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
with patch("docsgpt.agents.workflows.workflow_engine.normalize_json_schema_payload") as mock_norm:
|
|
from docsgpt.core.json_schema_utils import JsonSchemaValidationError
|
|
mock_norm.side_effect = JsonSchemaValidationError("bad schema")
|
|
with pytest.raises(ValueError, match="Invalid JSON schema"):
|
|
engine._normalize_node_json_schema({"bad": True}, "TestNode")
|
|
|
|
|
|
class TestValidateStructuredOutput:
|
|
|
|
@pytest.mark.unit
|
|
def test_valid_output_passes(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
schema = {"type": "object", "properties": {"name": {"type": "string"}}}
|
|
engine._validate_structured_output(schema, {"name": "Alice"}) # Should not raise
|
|
|
|
@pytest.mark.unit
|
|
def test_invalid_output_raises(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
schema = {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}
|
|
with pytest.raises(ValueError, match="did not match schema"):
|
|
engine._validate_structured_output(schema, {})
|
|
|
|
@pytest.mark.unit
|
|
def test_no_jsonschema_module(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
with patch("docsgpt.agents.workflows.workflow_engine.jsonschema", None):
|
|
engine._validate_structured_output({"type": "object"}, {}) # Should not raise
|
|
|
|
|
|
class TestFormatTemplate:
|
|
|
|
@pytest.mark.unit
|
|
def test_renders_template(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {"query": "hello"}
|
|
engine._build_template_context = MagicMock(return_value={"query": "hello"})
|
|
engine._template_engine = MagicMock()
|
|
engine._template_engine.render.return_value = "hello world"
|
|
result = engine._format_template("{{ query }} world")
|
|
assert result == "hello world"
|
|
|
|
@pytest.mark.unit
|
|
def test_render_error_returns_raw(self):
|
|
from docsgpt.templates.template_engine import TemplateRenderError
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine._build_template_context = MagicMock(return_value={})
|
|
engine._template_engine = MagicMock()
|
|
engine._template_engine.render.side_effect = TemplateRenderError("fail")
|
|
result = engine._format_template("{{ bad }}")
|
|
assert result == "{{ bad }}"
|
|
|
|
|
|
class TestBuildTemplateContext:
|
|
|
|
@pytest.mark.unit
|
|
def test_includes_state_variables(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = None
|
|
engine = WorkflowEngine(graph, agent)
|
|
engine.state = {"query": "hello", "custom_var": "value"}
|
|
context = engine._build_template_context()
|
|
assert context["agent"]["query"] == "hello"
|
|
assert "custom_var" in context
|
|
|
|
@pytest.mark.unit
|
|
def test_reserved_namespace_gets_prefixed(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = None
|
|
engine = WorkflowEngine(graph, agent)
|
|
engine.state = {"source": "my_source_val"}
|
|
context = engine._build_template_context()
|
|
assert context.get("agent_source") == "my_source_val"
|
|
|
|
@pytest.mark.unit
|
|
def test_passthrough_data(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = None
|
|
engine = WorkflowEngine(graph, agent)
|
|
engine.state = {"passthrough": {"key": "val"}}
|
|
context = engine._build_template_context()
|
|
assert "passthrough" in context or "agent_passthrough" in context
|
|
|
|
@pytest.mark.unit
|
|
def test_tools_data(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = None
|
|
engine = WorkflowEngine(graph, agent)
|
|
engine.state = {"tools": {"tool1": "result"}}
|
|
context = engine._build_template_context()
|
|
assert "agent" in context
|
|
|
|
|
|
class TestGetSourceTemplateData:
|
|
|
|
@pytest.mark.unit
|
|
def test_no_docs_returns_none(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = None
|
|
engine = WorkflowEngine(graph, agent)
|
|
docs, together = engine._get_source_template_data()
|
|
assert docs is None
|
|
assert together is None
|
|
|
|
@pytest.mark.unit
|
|
def test_empty_docs_returns_none(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = []
|
|
engine = WorkflowEngine(graph, agent)
|
|
docs, together = engine._get_source_template_data()
|
|
assert docs is None
|
|
|
|
@pytest.mark.unit
|
|
def test_docs_with_filename(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = [{"text": "content", "filename": "doc.txt"}]
|
|
engine = WorkflowEngine(graph, agent)
|
|
docs, together = engine._get_source_template_data()
|
|
assert docs is not None
|
|
assert "doc.txt" in together
|
|
assert "content" in together
|
|
|
|
@pytest.mark.unit
|
|
def test_docs_without_filename(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = [{"text": "content only"}]
|
|
engine = WorkflowEngine(graph, agent)
|
|
docs, together = engine._get_source_template_data()
|
|
assert together == "content only"
|
|
|
|
@pytest.mark.unit
|
|
def test_skips_non_dict_docs(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = ["not a dict", {"text": "ok"}]
|
|
engine = WorkflowEngine(graph, agent)
|
|
docs, together = engine._get_source_template_data()
|
|
assert together == "ok"
|
|
|
|
@pytest.mark.unit
|
|
def test_skips_non_string_text(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = [{"text": 123}]
|
|
engine = WorkflowEngine(graph, agent)
|
|
docs, together = engine._get_source_template_data()
|
|
assert together is None
|
|
|
|
@pytest.mark.unit
|
|
def test_doc_with_title_fallback(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = [{"text": "content", "title": "doc_title"}]
|
|
engine = WorkflowEngine(graph, agent)
|
|
docs, together = engine._get_source_template_data()
|
|
assert "doc_title" in together
|
|
|
|
@pytest.mark.unit
|
|
def test_doc_with_source_fallback(self):
|
|
graph = _make_graph([], [])
|
|
agent = _make_agent()
|
|
agent.retrieved_docs = [{"text": "content", "source": "src"}]
|
|
engine = WorkflowEngine(graph, agent)
|
|
docs, together = engine._get_source_template_data()
|
|
assert "src" in together
|
|
|
|
|
|
class TestGetExecutionSummary:
|
|
|
|
@pytest.mark.unit
|
|
def test_returns_log_entries(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
now = datetime.now(timezone.utc)
|
|
engine.execution_log = [
|
|
{
|
|
"node_id": "n1",
|
|
"node_type": "start",
|
|
"status": "completed",
|
|
"started_at": now,
|
|
"completed_at": now,
|
|
"error": None,
|
|
"state_delta": {},
|
|
}
|
|
]
|
|
summary = engine.get_execution_summary()
|
|
assert len(summary) == 1
|
|
assert summary[0].node_id == "n1"
|
|
assert summary[0].status == ExecutionStatus.COMPLETED
|
|
|
|
@pytest.mark.unit
|
|
def test_empty_log(self):
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
assert engine.get_execution_summary() == []
|
|
|
|
|
|
class TestAgentNodeExecution:
|
|
"""Cover lines 204, 213-215, 223, 232-233, 283-284, 289, 355, 375."""
|
|
|
|
@pytest.mark.unit
|
|
def test_agent_node_without_prompt_template(self):
|
|
"""Cover line 204/206: agent node without prompt_template uses query."""
|
|
node = _make_node("n1", NodeType.AGENT, "Agent", config={
|
|
"config": {
|
|
"agent_type": "classic",
|
|
"stream_to_user": False,
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {"query": "test question"}
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = [{"answer": "response"}]
|
|
|
|
with patch(
|
|
"docsgpt.agents.workflows.workflow_engine.WorkflowNodeAgentFactory"
|
|
) as mock_factory, \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_provider_from_model_id",
|
|
return_value="openai",
|
|
), \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_api_key_for_provider",
|
|
return_value="key",
|
|
), \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_model_capabilities",
|
|
return_value=None,
|
|
):
|
|
mock_factory.create.return_value = mock_agent
|
|
list(engine._execute_agent_node(node))
|
|
|
|
output_key = f"node_{node.id}_output"
|
|
assert output_key in engine.state
|
|
|
|
@pytest.mark.unit
|
|
def test_agent_node_with_structured_output(self):
|
|
"""Cover lines 283-284, 289: structured output parsing."""
|
|
node = _make_node("n1", NodeType.AGENT, "Agent", config={
|
|
"config": {
|
|
"agent_type": "classic",
|
|
"stream_to_user": False,
|
|
"json_schema": {
|
|
"type": "object",
|
|
"properties": {"name": {"type": "string"}},
|
|
},
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {"query": "test"}
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = [
|
|
{"answer": '{"name": "Alice"}', "structured": True}
|
|
]
|
|
|
|
with patch(
|
|
"docsgpt.agents.workflows.workflow_engine.WorkflowNodeAgentFactory"
|
|
) as mock_factory, \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_provider_from_model_id",
|
|
return_value="openai",
|
|
), \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_api_key_for_provider",
|
|
return_value="key",
|
|
), \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_model_capabilities",
|
|
return_value={"supports_structured_output": True},
|
|
):
|
|
mock_factory.create.return_value = mock_agent
|
|
list(engine._execute_agent_node(node))
|
|
|
|
output_key = f"node_{node.id}_output"
|
|
assert engine.state[output_key] == {"name": "Alice"}
|
|
|
|
@pytest.mark.unit
|
|
def test_agent_node_model_no_structured_support_raises(self):
|
|
"""Cover lines 223: model without structured output raises ValueError."""
|
|
node = _make_node("n1", NodeType.AGENT, "Agent", config={
|
|
"config": {
|
|
"agent_type": "classic",
|
|
"json_schema": {
|
|
"type": "object",
|
|
"properties": {"x": {"type": "string"}},
|
|
},
|
|
"model_id": "test-model",
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {"query": "test"}
|
|
|
|
with patch(
|
|
"docsgpt.core.model_utils.get_provider_from_model_id",
|
|
return_value="openai",
|
|
), \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_api_key_for_provider",
|
|
return_value="key",
|
|
), \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_model_capabilities",
|
|
return_value={"supports_structured_output": False},
|
|
):
|
|
with pytest.raises(ValueError, match="does not support structured output"):
|
|
list(engine._execute_agent_node(node))
|
|
|
|
@pytest.mark.unit
|
|
def test_agent_node_output_variable(self):
|
|
"""Cover line 300: output_variable stores result."""
|
|
node = _make_node("n1", NodeType.AGENT, "Agent", config={
|
|
"config": {
|
|
"agent_type": "classic",
|
|
"stream_to_user": False,
|
|
"output_variable": "my_result",
|
|
}
|
|
})
|
|
graph = _make_graph([node], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.state = {"query": "test"}
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.gen.return_value = [{"answer": "output text"}]
|
|
|
|
with patch(
|
|
"docsgpt.agents.workflows.workflow_engine.WorkflowNodeAgentFactory"
|
|
) as mock_factory, \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_provider_from_model_id",
|
|
return_value="openai",
|
|
), \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_api_key_for_provider",
|
|
return_value="key",
|
|
), \
|
|
patch(
|
|
"docsgpt.core.model_utils.get_model_capabilities",
|
|
return_value=None,
|
|
):
|
|
mock_factory.create.return_value = mock_agent
|
|
list(engine._execute_agent_node(node))
|
|
|
|
assert engine.state["my_result"] == "output text"
|
|
|
|
@pytest.mark.unit
|
|
def test_validate_structured_output_schema_error(self):
|
|
"""Cover line 375/382-383: invalid schema raises ValueError."""
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
import jsonschema as js
|
|
|
|
with patch(
|
|
"docsgpt.agents.workflows.workflow_engine.normalize_json_schema_payload",
|
|
return_value={"type": "invalid_schema_type"},
|
|
), \
|
|
patch(
|
|
"docsgpt.agents.workflows.workflow_engine.jsonschema"
|
|
) as mock_js:
|
|
mock_js.validate.side_effect = js.exceptions.SchemaError("bad schema")
|
|
mock_js.exceptions = js.exceptions
|
|
with pytest.raises(ValueError, match="Invalid JSON schema"):
|
|
engine._validate_structured_output(
|
|
{"type": "object"}, {"name": "test"}
|
|
)
|
|
|
|
@pytest.mark.unit
|
|
def test_parse_structured_output_invalid_json(self):
|
|
"""Cover lines 349-352: invalid JSON returns False."""
|
|
graph = _make_graph([], [])
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
success, data = engine._parse_structured_output("not json {")
|
|
assert success is False
|
|
assert data is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional coverage for workflow_engine.py
|
|
# Lines: 96-114 (exception in node execution), 122-130 (branch/max steps)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestWorkflowNodeExecutionException:
|
|
"""Cover lines 96-114: exception during _execute_node yields error events."""
|
|
|
|
def test_node_raises_exception_yields_error(self):
|
|
"""Force _execute_node to raise, covering lines 96-114."""
|
|
nodes = [
|
|
_make_node("n1", NodeType.START),
|
|
_make_node("n2", NodeType.AGENT, "Agent"),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2")]
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
|
|
# Patch _execute_node to raise on agent node
|
|
original_execute = engine._execute_node
|
|
|
|
def patched_execute(node):
|
|
if node.type == NodeType.AGENT:
|
|
raise RuntimeError("Agent exploded")
|
|
yield from original_execute(node)
|
|
|
|
engine._execute_node = patched_execute
|
|
events = list(engine.execute({}, "test query"))
|
|
|
|
error_events = [e for e in events if e.get("type") == "error"]
|
|
assert len(error_events) >= 1
|
|
failed_steps = [e for e in events if e.get("status") == "failed"]
|
|
assert len(failed_steps) >= 1
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestWorkflowMaxStepsReached:
|
|
"""Cover lines 127-130: max steps limit warning."""
|
|
|
|
def test_max_steps_exactly_reached(self):
|
|
nodes = [
|
|
_make_node("n1", NodeType.START),
|
|
_make_node("n2", NodeType.NOTE, "Note"),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2"), _make_edge("e2", "n2", "n2")]
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
engine.MAX_EXECUTION_STEPS = 3
|
|
events = list(engine.execute({}, "q"))
|
|
# The while loop runs 3 times then exits, steps >= MAX
|
|
assert len(events) >= 3
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestWorkflowBranchEndsNonEndNode:
|
|
"""Cover lines 122-125: branch ends at non-end node without outgoing edges."""
|
|
|
|
def test_branch_ends_at_state_node(self):
|
|
nodes = [
|
|
_make_node("n1", NodeType.START),
|
|
_make_node(
|
|
"n2",
|
|
NodeType.STATE,
|
|
"State",
|
|
config={"config": {"operations": []}},
|
|
),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2")] # n2 has no outgoing
|
|
graph = _make_graph(nodes, edges)
|
|
engine = WorkflowEngine(graph, _make_agent())
|
|
events = list(engine.execute({}, "q"))
|
|
# Should complete without crash, branch ended warning logged
|
|
assert len(events) > 0
|
|
|
|
|
|
class TestStepTelemetry:
|
|
"""Steps persist state DELTAS plus duration_ms and a per-node tool summary."""
|
|
|
|
def _run_state_workflow(self):
|
|
nodes = [
|
|
_make_node("n1", NodeType.START),
|
|
_make_node(
|
|
"n2",
|
|
NodeType.STATE,
|
|
"Set",
|
|
config={
|
|
"config": {
|
|
"operations": [
|
|
{"expression": "1 + 1", "target_variable": "x"}
|
|
]
|
|
}
|
|
},
|
|
),
|
|
_make_node("n3", NodeType.END, "End", config={"config": {}}),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2"), _make_edge("e2", "n2", "n3")]
|
|
engine = WorkflowEngine(_make_graph(nodes, edges), _make_agent())
|
|
events = list(engine.execute({}, "q"))
|
|
return engine, events
|
|
|
|
@pytest.mark.unit
|
|
def test_step_snapshots_are_deltas_not_full_state(self):
|
|
engine, events = self._run_state_workflow()
|
|
by_node = {log["node_id"]: log for log in engine.execution_log}
|
|
|
|
# The start node's delta carries the initial state it set...
|
|
assert "query" in by_node["n1"]["state_delta"]
|
|
# ...and later nodes repeat none of it, only what they changed.
|
|
assert by_node["n2"]["state_delta"] == {"x": 2}
|
|
assert by_node["n3"]["state_delta"] == {}
|
|
|
|
# The streamed step events carry the same deltas.
|
|
streamed = {
|
|
e["node_id"]: e
|
|
for e in events
|
|
if e.get("type") == "workflow_step" and e.get("status") == "completed"
|
|
}
|
|
assert streamed["n2"]["state_delta"] == {"x": 2}
|
|
assert "query" not in streamed["n2"]["state_delta"]
|
|
|
|
@pytest.mark.unit
|
|
def test_steps_carry_duration_ms(self):
|
|
engine, _ = self._run_state_workflow()
|
|
for log in engine.execution_log:
|
|
assert isinstance(log["duration_ms"], int)
|
|
assert log["duration_ms"] >= 0
|
|
|
|
summary = engine.get_execution_summary()
|
|
assert all(step.duration_ms is not None for step in summary)
|
|
|
|
@pytest.mark.unit
|
|
def test_failed_step_snapshot_is_a_delta_too(self):
|
|
nodes = [
|
|
_make_node("n1", NodeType.START),
|
|
_make_node(
|
|
"n2",
|
|
NodeType.STATE,
|
|
"Bad",
|
|
config={
|
|
"config": {
|
|
"operations": [
|
|
{"expression": "bad!!!", "target_variable": "x"}
|
|
]
|
|
}
|
|
},
|
|
),
|
|
]
|
|
edges = [_make_edge("e1", "n1", "n2")]
|
|
engine = WorkflowEngine(_make_graph(nodes, edges), _make_agent())
|
|
events = list(engine.execute({}, "q"))
|
|
|
|
failed = [e for e in events if e.get("status") == "failed"]
|
|
assert failed and "query" not in failed[0]["state_delta"]
|
|
failed_log = engine.execution_log[-1]
|
|
assert failed_log["status"] == "failed"
|
|
assert failed_log["duration_ms"] >= 0
|
|
|
|
@pytest.mark.unit
|
|
def test_summarize_tool_calls_compacts_executor_calls(self):
|
|
engine = WorkflowEngine(_make_graph([], []), _make_agent())
|
|
node_agent = MagicMock()
|
|
node_agent.tool_executor.tool_calls = [
|
|
{
|
|
"tool_name": "artifact_generator",
|
|
"action_name": "create_artifact",
|
|
"arguments": {"kind": "html"},
|
|
"result": "x" * 5000,
|
|
"status": "error",
|
|
},
|
|
{
|
|
"tool_name": "artifact_generator",
|
|
"action_name": "create_artifact",
|
|
"arguments": {"kind": "html"},
|
|
"result": "ok",
|
|
},
|
|
]
|
|
summary = engine._summarize_tool_calls(node_agent)
|
|
assert summary == [
|
|
{
|
|
"tool_name": "artifact_generator",
|
|
"action_name": "create_artifact",
|
|
"status": "error",
|
|
},
|
|
{
|
|
"tool_name": "artifact_generator",
|
|
"action_name": "create_artifact",
|
|
"status": "completed",
|
|
},
|
|
]
|
|
|
|
|
|
class TestNodeDocumentManifest:
|
|
"""The engine appends a per-node manifest of staged input documents to the
|
|
node system prompt, scoped to that node's input_documents selection."""
|
|
|
|
_ART_1 = "11111111-1111-4111-8111-111111111111"
|
|
_ART_2 = "22222222-2222-4222-8222-222222222222"
|
|
|
|
def _engine_with_docs(self, monkeypatch):
|
|
from contextlib import contextmanager
|
|
|
|
engine = WorkflowEngine(_make_graph([], []), _make_agent())
|
|
engine.state["input_documents"] = [
|
|
{"artifact_id": self._ART_1, "ref": "A1", "filename": "portfolio_v1.csv"},
|
|
{"artifact_id": self._ART_2, "ref": "A2", "filename": "notes.txt"},
|
|
]
|
|
|
|
rows = {
|
|
self._ART_1: (
|
|
{"current_version": 3, "metadata": {"ref_seq": 1}, "title": "Portfolio"},
|
|
{"filename": "portfolio_v1.csv", "mime_type": "text/csv"},
|
|
),
|
|
self._ART_2: (
|
|
{"current_version": 1, "metadata": {}, "title": "Notes"},
|
|
{"filename": "notes.txt", "mime_type": "text/plain"},
|
|
),
|
|
}
|
|
|
|
class FakeRepo:
|
|
def __init__(self, conn):
|
|
pass
|
|
|
|
def resolve_id_by_ref_seq(self, position, conversation_id=None, workflow_run_id=None):
|
|
for artifact_id, (artifact, _version) in rows.items():
|
|
if artifact["metadata"].get("ref_seq") == position:
|
|
return artifact_id
|
|
return None
|
|
|
|
def artifact_id_at_position(self, position, conversation_id=None, workflow_run_id=None):
|
|
return None
|
|
|
|
def get_artifact_in_parent(self, artifact_id, conversation_id=None, workflow_run_id=None):
|
|
entry = rows.get(artifact_id)
|
|
return entry[0] if entry else None
|
|
|
|
def get_version(self, artifact_id, version):
|
|
return rows[artifact_id][1]
|
|
|
|
@contextmanager
|
|
def fake_readonly():
|
|
yield object()
|
|
|
|
monkeypatch.setattr("docsgpt.storage.db.session.db_readonly", fake_readonly)
|
|
monkeypatch.setattr(
|
|
"docsgpt.storage.db.repositories.artifacts.ArtifactsRepository", FakeRepo
|
|
)
|
|
return engine
|
|
|
|
@pytest.mark.unit
|
|
def test_manifest_lists_refs_filenames_and_mimes(self, monkeypatch):
|
|
from docsgpt.agents.workflows.schemas import AgentNodeConfig
|
|
|
|
engine = self._engine_with_docs(monkeypatch)
|
|
config = AgentNodeConfig(agent_type="classic", input_documents=["*"])
|
|
manifest = engine._node_document_manifest(config)
|
|
assert "- A1: portfolio_v1.csv (text/csv)" in manifest
|
|
# Legacy artifact without ref_seq falls back to the full id, which tools accept.
|
|
assert f"- {self._ART_2}: notes.txt (text/plain)" in manifest
|
|
assert "Pass a ref" in manifest
|
|
|
|
@pytest.mark.unit
|
|
def test_manifest_empty_without_selection(self, monkeypatch):
|
|
from docsgpt.agents.workflows.schemas import AgentNodeConfig
|
|
|
|
engine = self._engine_with_docs(monkeypatch)
|
|
config = AgentNodeConfig(agent_type="classic", input_documents=[])
|
|
assert engine._node_document_manifest(config) == ""
|
|
|
|
@pytest.mark.unit
|
|
def test_manifest_failure_never_breaks_the_node(self, monkeypatch):
|
|
from contextlib import contextmanager
|
|
|
|
from docsgpt.agents.workflows.schemas import AgentNodeConfig
|
|
|
|
engine = WorkflowEngine(_make_graph([], []), _make_agent())
|
|
engine.state["input_documents"] = [{"artifact_id": self._ART_1}]
|
|
|
|
@contextmanager
|
|
def broken_readonly():
|
|
raise RuntimeError("db down")
|
|
yield
|
|
|
|
monkeypatch.setattr("docsgpt.storage.db.session.db_readonly", broken_readonly)
|
|
config = AgentNodeConfig(agent_type="classic", input_documents=["*"])
|
|
assert engine._node_document_manifest(config) == ""
|