80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock
|
|
from ea_chatbot.graph.workflow import create_workflow
|
|
from ea_chatbot.graph.state import AgentState
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
|
|
def test_deepagents_multi_worker_sequential_flow():
|
|
"""Verify that the Orchestrator can handle a sequence of different workers."""
|
|
|
|
mock_analyzer = MagicMock()
|
|
mock_planner = MagicMock()
|
|
mock_delegate = MagicMock()
|
|
mock_analyst = MagicMock()
|
|
mock_researcher = MagicMock()
|
|
mock_reflector = MagicMock()
|
|
mock_synthesizer = MagicMock()
|
|
|
|
# 1. Analyzer: Plan
|
|
mock_analyzer.return_value = {"next_action": "plan"}
|
|
|
|
# 2. Planner: Two tasks
|
|
mock_planner.return_value = {
|
|
"checklist": [
|
|
{"task": "Get Numbers", "worker": "data_analyst"},
|
|
{"task": "Get Facts", "worker": "researcher"}
|
|
],
|
|
"current_step": 0
|
|
}
|
|
|
|
# 3. Delegate & Reflector: Loop through tasks
|
|
mock_delegate.side_effect = [
|
|
{"next_action": "data_analyst"}, # Step 0
|
|
{"next_action": "researcher"}, # Step 1
|
|
{"next_action": "summarize"} # Step 2 (done)
|
|
]
|
|
|
|
mock_analyst.return_value = {"messages": [AIMessage(content="Analyst summary")], "vfs": {"data.csv": "..."}}
|
|
mock_researcher.return_value = {"messages": [AIMessage(content="Researcher summary")]}
|
|
|
|
mock_reflector.side_effect = [
|
|
{"current_step": 1, "next_action": "delegate"}, # Done with Analyst
|
|
{"current_step": 2, "next_action": "delegate"} # Done with Researcher
|
|
]
|
|
|
|
mock_synthesizer.return_value = {
|
|
"messages": [AIMessage(content="Final multi-agent response")],
|
|
"next_action": "end"
|
|
}
|
|
|
|
app = create_workflow(
|
|
query_analyzer=mock_analyzer,
|
|
planner=mock_planner,
|
|
delegate=mock_delegate,
|
|
data_analyst_worker=mock_analyst,
|
|
researcher_worker=mock_researcher,
|
|
reflector=mock_reflector,
|
|
synthesizer=mock_synthesizer
|
|
)
|
|
|
|
initial_state = AgentState(
|
|
messages=[HumanMessage(content="Numbers and facts please")],
|
|
question="Numbers and facts please",
|
|
analysis={},
|
|
next_action="",
|
|
iterations=0,
|
|
checklist=[],
|
|
current_step=0,
|
|
vfs={},
|
|
plots=[],
|
|
dfs={}
|
|
)
|
|
|
|
final_state = app.invoke(initial_state, config={"recursion_limit": 30})
|
|
|
|
assert mock_analyst.called
|
|
assert mock_researcher.called
|
|
assert mock_reflector.call_count == 2
|
|
assert "Final multi-agent response" in [m.content for m in final_state["messages"]]
|
|
assert final_state["current_step"] == 2
|