mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 08:13:02 +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.
385 lines
13 KiB
Python
385 lines
13 KiB
Python
"""Tests for docsgpt/agents/workflow_agent.py graph loading and saving.
|
|
|
|
Tests _parse_embedded_workflow, _load_from_database, and _finalize_workflow_run
|
|
against the ephemeral ``pg_conn`` fixture. Agent construction is bypassed via
|
|
``__new__`` so we avoid the BaseAgent's LLM/tool wiring.
|
|
"""
|
|
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timezone
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
|
|
def _make_agent(*, workflow_id=None, workflow=None, workflow_owner=None,
|
|
decoded_token=None):
|
|
"""Construct a WorkflowAgent bypassing BaseAgent.__init__."""
|
|
from docsgpt.agents.workflow_agent import WorkflowAgent
|
|
agent = WorkflowAgent.__new__(WorkflowAgent)
|
|
agent.workflow_id = workflow_id
|
|
agent.workflow_owner = workflow_owner
|
|
agent._workflow_data = workflow
|
|
agent._engine = None
|
|
agent._run_persisted = False
|
|
agent._bridge_error = None
|
|
agent.decoded_token = decoded_token or {}
|
|
return agent
|
|
|
|
|
|
@contextmanager
|
|
def _patch_db(conn):
|
|
@contextmanager
|
|
def _yield():
|
|
yield conn
|
|
|
|
with patch(
|
|
"docsgpt.agents.workflow_agent.db_readonly", _yield
|
|
), patch(
|
|
"docsgpt.agents.workflow_agent.db_session", _yield
|
|
):
|
|
yield
|
|
|
|
|
|
class TestParseEmbeddedWorkflow:
|
|
def test_returns_none_for_malformed_workflow(self):
|
|
agent = _make_agent(workflow={"nodes": None}) # invalid shape
|
|
assert agent._parse_embedded_workflow() is None
|
|
|
|
def test_parses_valid_embedded_workflow(self):
|
|
agent = _make_agent(
|
|
workflow={
|
|
"name": "Embedded",
|
|
"description": "d",
|
|
"nodes": [
|
|
{
|
|
"id": "n1", "type": "start", "title": "Start",
|
|
"position": {"x": 0, "y": 0}, "data": {},
|
|
},
|
|
{
|
|
"id": "n2", "type": "end", "title": "End",
|
|
"position": {"x": 100, "y": 0}, "data": {},
|
|
},
|
|
],
|
|
"edges": [
|
|
{
|
|
"id": "e1", "source": "n1", "target": "n2",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
graph = agent._parse_embedded_workflow()
|
|
assert graph is not None
|
|
assert graph.workflow.name == "Embedded"
|
|
assert len(graph.nodes) == 2
|
|
assert len(graph.edges) == 1
|
|
|
|
def test_parses_edges_with_source_id_target_id(self):
|
|
"""Edges may arrive with ``source_id``/``target_id`` instead of
|
|
``source``/``target``."""
|
|
agent = _make_agent(
|
|
workflow={
|
|
"name": "v2",
|
|
"nodes": [
|
|
{"id": "a", "type": "start", "position": {"x": 0, "y": 0}, "data": {}},
|
|
{"id": "b", "type": "end", "position": {"x": 0, "y": 0}, "data": {}},
|
|
],
|
|
"edges": [
|
|
{
|
|
"id": "e1",
|
|
"source_id": "a", "target_id": "b",
|
|
"source_handle": "h", "target_handle": "h2",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
graph = agent._parse_embedded_workflow()
|
|
assert graph.edges[0].source_id == "a"
|
|
|
|
|
|
class TestLoadWorkflowGraph:
|
|
def test_returns_none_when_no_data_or_id(self):
|
|
agent = _make_agent()
|
|
assert agent._load_workflow_graph() is None
|
|
|
|
def test_uses_embedded_when_provided(self):
|
|
agent = _make_agent(workflow={"nodes": [], "edges": []})
|
|
# Empty nodes returns WorkflowGraph with no nodes, still valid object
|
|
got = agent._load_workflow_graph()
|
|
assert got is not None
|
|
|
|
def test_uses_database_when_id_set(self, pg_conn):
|
|
from docsgpt.storage.db.repositories.workflows import (
|
|
WorkflowsRepository,
|
|
)
|
|
user = "u-loadwf"
|
|
workflow = WorkflowsRepository(pg_conn).create(user, "wf")
|
|
agent = _make_agent(
|
|
workflow_id=str(workflow["id"]),
|
|
workflow_owner=user,
|
|
)
|
|
with _patch_db(pg_conn):
|
|
got = agent._load_workflow_graph()
|
|
assert got is not None
|
|
|
|
|
|
class TestLoadFromDatabase:
|
|
def test_returns_none_no_workflow_id(self):
|
|
agent = _make_agent()
|
|
assert agent._load_from_database() is None
|
|
|
|
def test_returns_none_no_owner(self):
|
|
agent = _make_agent(workflow_id="some-id")
|
|
assert agent._load_from_database() is None
|
|
|
|
def test_owner_from_decoded_token(self, pg_conn):
|
|
from docsgpt.storage.db.repositories.workflows import (
|
|
WorkflowsRepository,
|
|
)
|
|
user = "u-token-owner"
|
|
wf = WorkflowsRepository(pg_conn).create(user, "wf")
|
|
agent = _make_agent(
|
|
workflow_id=str(wf["id"]),
|
|
decoded_token={"sub": user},
|
|
)
|
|
with _patch_db(pg_conn):
|
|
got = agent._load_from_database()
|
|
assert got is not None
|
|
|
|
def test_returns_none_when_workflow_missing(self, pg_conn):
|
|
agent = _make_agent(
|
|
workflow_id="00000000-0000-0000-0000-000000000000",
|
|
workflow_owner="u",
|
|
)
|
|
with _patch_db(pg_conn):
|
|
got = agent._load_from_database()
|
|
assert got is None
|
|
|
|
def test_invalid_version_falls_back_to_1(self, pg_conn):
|
|
"""When current_graph_version is invalid it falls back to 1."""
|
|
from docsgpt.storage.db.repositories.workflows import (
|
|
WorkflowsRepository,
|
|
)
|
|
user = "u-bad-version"
|
|
wf = WorkflowsRepository(pg_conn).create(user, "wf")
|
|
agent = _make_agent(
|
|
workflow_id=str(wf["id"]), workflow_owner=user,
|
|
)
|
|
with _patch_db(pg_conn):
|
|
got = agent._load_from_database()
|
|
assert got is not None
|
|
|
|
def test_handles_exception(self):
|
|
@contextmanager
|
|
def _broken():
|
|
raise RuntimeError("boom")
|
|
yield
|
|
|
|
agent = _make_agent(workflow_id="x", workflow_owner="u")
|
|
with patch(
|
|
"docsgpt.agents.workflow_agent.db_readonly", _broken
|
|
):
|
|
got = agent._load_from_database()
|
|
assert got is None
|
|
|
|
|
|
class TestSaveWorkflowRun:
|
|
def test_returns_when_no_engine(self):
|
|
agent = _make_agent(workflow_id="x", workflow_owner="u")
|
|
# _engine is None
|
|
agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "query")
|
|
# should not raise
|
|
|
|
def test_returns_when_no_workflow_id(self):
|
|
agent = _make_agent(workflow_owner="u")
|
|
agent._engine = MagicMock()
|
|
agent._engine.execution_log = []
|
|
agent._engine.state = {}
|
|
agent._engine.get_execution_summary.return_value = []
|
|
agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "query")
|
|
|
|
def test_returns_when_workflow_missing_in_db(self, pg_conn):
|
|
agent = _make_agent(
|
|
workflow_id="00000000-0000-0000-0000-000000000000",
|
|
workflow_owner="u",
|
|
)
|
|
agent._engine = MagicMock()
|
|
agent._engine.execution_log = []
|
|
agent._engine.state = {}
|
|
agent._engine.get_execution_summary.return_value = []
|
|
with _patch_db(pg_conn):
|
|
# Should just return None since workflow not found in DB
|
|
agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "q")
|
|
|
|
def test_creates_run_row(self, pg_conn):
|
|
from docsgpt.storage.db.repositories.workflows import (
|
|
WorkflowsRepository,
|
|
)
|
|
from docsgpt.storage.db.repositories.workflow_runs import (
|
|
WorkflowRunsRepository,
|
|
)
|
|
|
|
import uuid as _uuid
|
|
|
|
user = "u-saverun"
|
|
wf = WorkflowsRepository(pg_conn).create(user, "wf")
|
|
agent = _make_agent(
|
|
workflow_id=str(wf["id"]), workflow_owner=user,
|
|
)
|
|
agent._engine = MagicMock()
|
|
agent._engine.execution_log = []
|
|
agent._engine.state = {"output": "hello"}
|
|
agent._engine.get_execution_summary.return_value = []
|
|
# The run row is persisted under the engine's run id (the parent for any
|
|
# run-scoped artifacts); a MagicMock id can't be adapted by the driver.
|
|
run_id = str(_uuid.uuid4())
|
|
agent._engine.workflow_run_id = run_id
|
|
|
|
with _patch_db(pg_conn):
|
|
agent._finalize_workflow_run(
|
|
agent.workflow_owner, agent.workflow_owner, None, "my query"
|
|
)
|
|
|
|
runs = WorkflowRunsRepository(pg_conn).list_for_workflow(str(wf["id"]))
|
|
assert len(runs) >= 1
|
|
assert str(runs[0]["id"]) == run_id
|
|
|
|
def test_exception_is_swallowed(self):
|
|
agent = _make_agent(workflow_id="x", workflow_owner="u")
|
|
agent._engine = MagicMock()
|
|
agent._engine.execution_log = []
|
|
agent._engine.state = {}
|
|
agent._engine.get_execution_summary.return_value = []
|
|
|
|
@contextmanager
|
|
def _broken():
|
|
raise RuntimeError("boom")
|
|
yield
|
|
|
|
with patch(
|
|
"docsgpt.agents.workflow_agent.db_session", _broken
|
|
):
|
|
# Should not raise
|
|
agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "q")
|
|
|
|
|
|
class TestDetermineRunStatus:
|
|
def test_completed_when_no_engine(self):
|
|
from docsgpt.agents.workflows.schemas import ExecutionStatus
|
|
|
|
agent = _make_agent()
|
|
assert agent._determine_run_status() == ExecutionStatus.COMPLETED
|
|
|
|
def test_completed_when_log_empty(self):
|
|
from docsgpt.agents.workflows.schemas import ExecutionStatus
|
|
|
|
agent = _make_agent()
|
|
agent._engine = MagicMock()
|
|
agent._engine.execution_log = []
|
|
assert agent._determine_run_status() == ExecutionStatus.COMPLETED
|
|
|
|
def test_failed_if_any_log_failed(self):
|
|
from docsgpt.agents.workflows.schemas import ExecutionStatus
|
|
|
|
agent = _make_agent()
|
|
agent._engine = MagicMock()
|
|
agent._engine.execution_log = [
|
|
{"status": ExecutionStatus.COMPLETED.value},
|
|
{"status": ExecutionStatus.FAILED.value},
|
|
]
|
|
assert agent._determine_run_status() == ExecutionStatus.FAILED
|
|
|
|
|
|
class TestSerializeState:
|
|
def test_primitive_passes_through(self):
|
|
agent = _make_agent()
|
|
assert agent._serialize_state_value(42) == 42
|
|
assert agent._serialize_state_value("x") == "x"
|
|
assert agent._serialize_state_value(None) is None
|
|
assert agent._serialize_state_value(True) is True
|
|
|
|
def test_datetime_becomes_iso(self):
|
|
agent = _make_agent()
|
|
now = datetime(2024, 1, 2, 3, 4, 5, tzinfo=timezone.utc)
|
|
assert agent._serialize_state_value(now) == now.isoformat()
|
|
|
|
def test_dict_becomes_dict_with_string_keys(self):
|
|
agent = _make_agent()
|
|
got = agent._serialize_state_value({1: "a", "b": 2})
|
|
assert got == {"1": "a", "b": 2}
|
|
|
|
def test_tuple_becomes_list(self):
|
|
agent = _make_agent()
|
|
assert agent._serialize_state_value((1, 2, 3)) == [1, 2, 3]
|
|
|
|
def test_unknown_becomes_string(self):
|
|
agent = _make_agent()
|
|
|
|
class Foo:
|
|
def __str__(self):
|
|
return "custom-foo"
|
|
|
|
assert agent._serialize_state_value(Foo()) == "custom-foo"
|
|
|
|
def test_serialize_state_dict(self):
|
|
agent = _make_agent()
|
|
got = agent._serialize_state({"x": 1, "y": [1, 2]})
|
|
assert got == {"x": 1, "y": [1, 2]}
|
|
|
|
|
|
class TestGen:
|
|
def test_yields_error_when_graph_fails(self):
|
|
agent = _make_agent() # no workflow data
|
|
# _load_workflow_graph returns None so gen yields error
|
|
results = list(agent._gen_inner("q", log_context=None))
|
|
assert results == [
|
|
{"type": "error", "error": "Failed to load workflow configuration."}
|
|
]
|
|
|
|
|
|
class TestAgentNodeApprovalPause:
|
|
"""A node agent whose tool pauses for approval must fail the node visibly, not emit empty output."""
|
|
|
|
def test_tool_calls_pending_raises_clear_error(self, monkeypatch):
|
|
from docsgpt.agents.workflows import workflow_engine as we
|
|
from docsgpt.agents.workflows.schemas import (
|
|
NodeType,
|
|
Workflow,
|
|
WorkflowGraph,
|
|
WorkflowNode,
|
|
)
|
|
from docsgpt.agents.workflows.workflow_engine import WorkflowEngine
|
|
|
|
# An ephemeral node agent whose LLM handler yields the pause signal and ends,
|
|
# emitting no "answer". Previously the engine dropped it and the node completed
|
|
# with empty output (or raised a confusing "Structured output was expected").
|
|
class _PendingAgent:
|
|
attachments = None
|
|
|
|
def gen(self, prompt):
|
|
yield {"type": "tool_calls_pending", "data": {"pending_tool_calls": [{"id": "x"}]}}
|
|
|
|
monkeypatch.setattr(
|
|
we.WorkflowNodeAgentFactory, "create", lambda **kw: _PendingAgent()
|
|
)
|
|
|
|
graph = WorkflowGraph(workflow=Workflow(name="Approval Pause"), nodes=[], edges=[])
|
|
agent = SimpleNamespace(
|
|
endpoint="stream", llm_name="openai", model_id="gpt-4o-mini", api_key="k",
|
|
chat_history=[], user="u", decoded_token={"sub": "u"},
|
|
)
|
|
engine = WorkflowEngine(
|
|
graph, agent, workflow_run_id="22222222-2222-2222-2222-222222222222"
|
|
)
|
|
|
|
node = WorkflowNode(
|
|
id="agent_1", workflow_id="wf-1", type=NodeType.AGENT, title="Extractor",
|
|
position={"x": 0, "y": 0},
|
|
config={"llm_name": "openai", "system_prompt": "s", "tools": ["t"]},
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="requires approval"):
|
|
list(engine._execute_agent_node(node))
|