mirror of
https://github.com/arc53/DocsGPT.git
synced 2026-03-02 16:21:50 +00:00
fix: mini workflow fixes
This commit is contained in:
Alex
332
tests/agents/test_workflow_engine.py
Normal file
332
tests/agents/test_workflow_engine.py
Normal file
@@ -0,0 +1,332 @@
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from application.api.user.workflows import routes as workflow_routes
|
||||
from application.agents.workflows.node_agent import WorkflowNodeAgentFactory
|
||||
from application.agents.workflows.schemas import (
|
||||
NodeType,
|
||||
Workflow,
|
||||
WorkflowGraph,
|
||||
WorkflowNode,
|
||||
)
|
||||
from application.agents.workflows.workflow_engine import WorkflowEngine
|
||||
from application.api.user.workflows.routes import validate_workflow_structure
|
||||
|
||||
|
||||
class StubNodeAgent:
|
||||
def __init__(self, events):
|
||||
self.events = events
|
||||
|
||||
def gen(self, _prompt):
|
||||
yield from self.events
|
||||
|
||||
|
||||
def create_engine() -> WorkflowEngine:
|
||||
graph = WorkflowGraph(workflow=Workflow(name="Engine Test"), nodes=[], edges=[])
|
||||
agent = SimpleNamespace(
|
||||
endpoint="stream",
|
||||
llm_name="openai",
|
||||
model_id="gpt-4o-mini",
|
||||
api_key="test-key",
|
||||
chat_history=[],
|
||||
decoded_token={"sub": "user-1"},
|
||||
)
|
||||
return WorkflowEngine(graph, agent)
|
||||
|
||||
|
||||
def create_agent_node(
|
||||
node_id: str,
|
||||
output_variable: str = "",
|
||||
json_schema: Optional[Dict[str, Any]] = None,
|
||||
) -> WorkflowNode:
|
||||
config = {
|
||||
"agent_type": "classic",
|
||||
"system_prompt": "You are a helpful assistant.",
|
||||
"prompt_template": "",
|
||||
"stream_to_user": False,
|
||||
"tools": [],
|
||||
}
|
||||
if output_variable:
|
||||
config["output_variable"] = output_variable
|
||||
if json_schema is not None:
|
||||
config["json_schema"] = json_schema
|
||||
|
||||
return WorkflowNode(
|
||||
id=node_id,
|
||||
workflow_id="workflow-1",
|
||||
type=NodeType.AGENT,
|
||||
title="Agent",
|
||||
position={"x": 0, "y": 0},
|
||||
config=config,
|
||||
)
|
||||
|
||||
|
||||
def test_execute_agent_node_saves_structured_output_as_json(monkeypatch):
|
||||
engine = create_engine()
|
||||
node = create_agent_node(
|
||||
node_id="agent_1",
|
||||
output_variable="result",
|
||||
json_schema={"type": "object"},
|
||||
)
|
||||
node_events = [
|
||||
{"answer": '{"summary":"ok",', "structured": True},
|
||||
{"answer": '"score":2}', "structured": True},
|
||||
]
|
||||
|
||||
monkeypatch.setattr(
|
||||
WorkflowNodeAgentFactory,
|
||||
"create",
|
||||
staticmethod(lambda **kwargs: StubNodeAgent(node_events)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_api_key_for_provider",
|
||||
lambda _provider: None,
|
||||
)
|
||||
|
||||
list(engine._execute_agent_node(node))
|
||||
|
||||
expected_output = {"summary": "ok", "score": 2}
|
||||
assert engine.state["node_agent_1_output"] == expected_output
|
||||
assert engine.state["result"] == expected_output
|
||||
|
||||
|
||||
def test_execute_agent_node_normalizes_wrapped_schema_before_agent_create(monkeypatch):
|
||||
engine = create_engine()
|
||||
node = create_agent_node(
|
||||
node_id="agent_wrapped",
|
||||
json_schema={"schema": {"type": "object"}},
|
||||
)
|
||||
node_events = [{"answer": '{"summary":"ok"}', "structured": True}]
|
||||
captured: Dict[str, Any] = {}
|
||||
|
||||
def create_node_agent(**kwargs):
|
||||
captured["json_schema"] = kwargs.get("json_schema")
|
||||
return StubNodeAgent(node_events)
|
||||
|
||||
monkeypatch.setattr(
|
||||
WorkflowNodeAgentFactory,
|
||||
"create",
|
||||
staticmethod(create_node_agent),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_api_key_for_provider",
|
||||
lambda _provider: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_model_capabilities",
|
||||
lambda _model_id: {"supports_structured_output": True},
|
||||
)
|
||||
|
||||
list(engine._execute_agent_node(node))
|
||||
|
||||
assert captured["json_schema"] == {"type": "object"}
|
||||
assert engine.state["node_agent_wrapped_output"] == {"summary": "ok"}
|
||||
|
||||
|
||||
def test_execute_agent_node_falls_back_to_text_when_schema_not_configured(monkeypatch):
|
||||
engine = create_engine()
|
||||
node = create_agent_node(node_id="agent_2", output_variable="result")
|
||||
node_events = [{"answer": "plain text answer"}]
|
||||
|
||||
monkeypatch.setattr(
|
||||
WorkflowNodeAgentFactory,
|
||||
"create",
|
||||
staticmethod(lambda **kwargs: StubNodeAgent(node_events)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_api_key_for_provider",
|
||||
lambda _provider: None,
|
||||
)
|
||||
|
||||
list(engine._execute_agent_node(node))
|
||||
|
||||
assert engine.state["node_agent_2_output"] == "plain text answer"
|
||||
assert engine.state["result"] == "plain text answer"
|
||||
|
||||
|
||||
def test_validate_workflow_structure_rejects_invalid_agent_json_schema():
|
||||
nodes = [
|
||||
{"id": "start", "type": "start", "title": "Start", "data": {}},
|
||||
{
|
||||
"id": "agent",
|
||||
"type": "agent",
|
||||
"title": "Agent",
|
||||
"data": {"json_schema": "invalid"},
|
||||
},
|
||||
{"id": "end", "type": "end", "title": "End", "data": {}},
|
||||
]
|
||||
edges = [
|
||||
{"id": "edge_1", "source": "start", "target": "agent"},
|
||||
{"id": "edge_2", "source": "agent", "target": "end"},
|
||||
]
|
||||
|
||||
errors = validate_workflow_structure(nodes, edges)
|
||||
|
||||
assert any(
|
||||
"Agent node 'Agent' JSON schema must be a valid JSON object" in err
|
||||
for err in errors
|
||||
)
|
||||
|
||||
|
||||
def test_validate_workflow_structure_accepts_valid_agent_json_schema():
|
||||
nodes = [
|
||||
{"id": "start", "type": "start", "title": "Start", "data": {}},
|
||||
{
|
||||
"id": "agent",
|
||||
"type": "agent",
|
||||
"title": "Agent",
|
||||
"data": {"json_schema": {"type": "object"}},
|
||||
},
|
||||
{"id": "end", "type": "end", "title": "End", "data": {}},
|
||||
]
|
||||
edges = [
|
||||
{"id": "edge_1", "source": "start", "target": "agent"},
|
||||
{"id": "edge_2", "source": "agent", "target": "end"},
|
||||
]
|
||||
|
||||
errors = validate_workflow_structure(nodes, edges)
|
||||
|
||||
assert errors == []
|
||||
|
||||
|
||||
def test_validate_workflow_structure_accepts_wrapped_agent_json_schema():
|
||||
nodes = [
|
||||
{"id": "start", "type": "start", "title": "Start", "data": {}},
|
||||
{
|
||||
"id": "agent",
|
||||
"type": "agent",
|
||||
"title": "Agent",
|
||||
"data": {"json_schema": {"schema": {"type": "object"}}},
|
||||
},
|
||||
{"id": "end", "type": "end", "title": "End", "data": {}},
|
||||
]
|
||||
edges = [
|
||||
{"id": "edge_1", "source": "start", "target": "agent"},
|
||||
{"id": "edge_2", "source": "agent", "target": "end"},
|
||||
]
|
||||
|
||||
errors = validate_workflow_structure(nodes, edges)
|
||||
|
||||
assert errors == []
|
||||
|
||||
|
||||
def test_validate_workflow_structure_accepts_output_variable_and_schema_together():
|
||||
nodes = [
|
||||
{"id": "start", "type": "start", "title": "Start", "data": {}},
|
||||
{
|
||||
"id": "agent",
|
||||
"type": "agent",
|
||||
"title": "Agent",
|
||||
"data": {
|
||||
"output_variable": "answer",
|
||||
"json_schema": {"type": "object"},
|
||||
},
|
||||
},
|
||||
{"id": "end", "type": "end", "title": "End", "data": {}},
|
||||
]
|
||||
edges = [
|
||||
{"id": "edge_1", "source": "start", "target": "agent"},
|
||||
{"id": "edge_2", "source": "agent", "target": "end"},
|
||||
]
|
||||
|
||||
errors = validate_workflow_structure(nodes, edges)
|
||||
|
||||
assert errors == []
|
||||
|
||||
|
||||
def test_validate_workflow_structure_rejects_unsupported_structured_output_model(
|
||||
monkeypatch,
|
||||
):
|
||||
monkeypatch.setattr(
|
||||
workflow_routes,
|
||||
"get_model_capabilities",
|
||||
lambda _model_id: {"supports_structured_output": False},
|
||||
)
|
||||
|
||||
nodes = [
|
||||
{"id": "start", "type": "start", "title": "Start", "data": {}},
|
||||
{
|
||||
"id": "agent",
|
||||
"type": "agent",
|
||||
"title": "Agent",
|
||||
"data": {
|
||||
"model_id": "some-model",
|
||||
"json_schema": {"type": "object"},
|
||||
},
|
||||
},
|
||||
{"id": "end", "type": "end", "title": "End", "data": {}},
|
||||
]
|
||||
edges = [
|
||||
{"id": "edge_1", "source": "start", "target": "agent"},
|
||||
{"id": "edge_2", "source": "agent", "target": "end"},
|
||||
]
|
||||
|
||||
errors = validate_workflow_structure(nodes, edges)
|
||||
|
||||
assert any(
|
||||
"Agent node 'Agent' selected model does not support structured output"
|
||||
in err
|
||||
for err in errors
|
||||
)
|
||||
|
||||
|
||||
def test_execute_agent_node_raises_when_structured_output_violates_schema(monkeypatch):
|
||||
engine = create_engine()
|
||||
node = create_agent_node(
|
||||
node_id="agent_3",
|
||||
json_schema={
|
||||
"type": "object",
|
||||
"properties": {"summary": {"type": "string"}},
|
||||
"required": ["summary"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
)
|
||||
node_events = [{"answer": '{"score":2}', "structured": True}]
|
||||
|
||||
monkeypatch.setattr(
|
||||
WorkflowNodeAgentFactory,
|
||||
"create",
|
||||
staticmethod(lambda **kwargs: StubNodeAgent(node_events)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_api_key_for_provider",
|
||||
lambda _provider: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_model_capabilities",
|
||||
lambda _model_id: {"supports_structured_output": True},
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Structured output did not match schema"):
|
||||
list(engine._execute_agent_node(node))
|
||||
|
||||
|
||||
def test_execute_agent_node_raises_when_schema_set_and_response_not_json(monkeypatch):
|
||||
engine = create_engine()
|
||||
node = create_agent_node(
|
||||
node_id="agent_4",
|
||||
json_schema={"type": "object"},
|
||||
)
|
||||
node_events = [{"answer": "not-json"}]
|
||||
|
||||
monkeypatch.setattr(
|
||||
WorkflowNodeAgentFactory,
|
||||
"create",
|
||||
staticmethod(lambda **kwargs: StubNodeAgent(node_events)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_api_key_for_provider",
|
||||
lambda _provider: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_model_capabilities",
|
||||
lambda _model_id: {"supports_structured_output": True},
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Structured output was expected but response was not valid JSON",
|
||||
):
|
||||
list(engine._execute_agent_node(node))
|
||||
Reference in New Issue
Block a user