diff --git a/backend/tests/test_deepagents_e2e.py b/backend/tests/test_deepagents_e2e.py new file mode 100644 index 0000000..fbb7447 --- /dev/null +++ b/backend/tests/test_deepagents_e2e.py @@ -0,0 +1,79 @@ +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