YunxiaoXu/ea-chatbot-lg
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore(graph): Relocate QueryAnalysis schema and update existing tests for Orchestrator architecture

Browse Source
This commit is contained in:
Yunxiao Xu
2026-02-23 05:58:58 -08:00
parent ad7845cc6a
commit f4d09c07c4
7 changed files with 199 additions and 227 deletions
Download Patch File Download Diff File Expand all files Collapse all files

139
backend/tests/test_workflow_e2e.py
View File

@@ -1,60 +1,51 @@
import pytest
from unittest.mock import MagicMock, patch
from langchain_core.messages import AIMessage
from ea_chatbot.graph.workflow import app
from ea_chatbot.graph.nodes.query_analyzer import QueryAnalysis
from ea_chatbot.schemas import TaskPlanResponse, TaskPlanContext, CodeGenerationResponse
from ea_chatbot.schemas import QueryAnalysis, ChecklistResponse, ChecklistTask, CodeGenerationResponse
from ea_chatbot.graph.state import AgentState
from langchain_core.messages import AIMessage
@pytest.fixture
def mock_llms():
with patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model") as mock_qa_llm, \
patch("ea_chatbot.graph.nodes.planner.get_llm_model") as mock_planner_llm, \
patch("ea_chatbot.graph.nodes.coder.get_llm_model") as mock_coder_llm, \
patch("ea_chatbot.graph.nodes.summarizer.get_llm_model") as mock_summarizer_llm, \
patch("ea_chatbot.graph.nodes.researcher.get_llm_model") as mock_researcher_llm, \
patch("ea_chatbot.graph.nodes.summarize_conversation.get_llm_model") as mock_summary_llm, \
patch("ea_chatbot.utils.database_inspection.get_data_summary") as mock_get_summary:
mock_get_summary.return_value = "Data summary"
# Mock summary LLM to return a simple response
mock_summary_instance = MagicMock()
mock_summary_llm.return_value = mock_summary_instance
mock_summary_instance.invoke.return_value = AIMessage(content="Turn summary")
with patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model") as mock_qa, \
patch("ea_chatbot.graph.nodes.planner.get_llm_model") as mock_planner, \
patch("ea_chatbot.graph.workers.data_analyst.nodes.coder.get_llm_model") as mock_coder, \
patch("ea_chatbot.graph.workers.data_analyst.nodes.summarizer.get_llm_model") as mock_worker_summarizer, \
patch("ea_chatbot.graph.nodes.synthesizer.get_llm_model") as mock_synthesizer, \
patch("ea_chatbot.graph.nodes.researcher.get_llm_model") as mock_researcher:
yield {
"qa": mock_qa_llm,
"planner": mock_planner_llm,
"coder": mock_coder_llm,
"summarizer": mock_summarizer_llm,
"researcher": mock_researcher_llm,
"summary": mock_summary_llm
"qa": mock_qa,
"planner": mock_planner,
"coder": mock_coder,
"worker_summarizer": mock_worker_summarizer,
"synthesizer": mock_synthesizer,
"researcher": mock_researcher
}
def test_workflow_data_analysis_flow(mock_llms):
"""Test full flow: QueryAnalyzer -> Planner -> Coder -> Executor -> Summarizer."""
"""Test full flow: QueryAnalyzer -> Planner -> Delegate -> DataAnalyst -> Reflector -> Synthesizer."""
# 1. Mock Query Analyzer (routes to plan)
# 1. Mock Query Analyzer
mock_qa_instance = MagicMock()
mock_llms["qa"].return_value = mock_qa_instance
mock_qa_instance.with_structured_output.return_value.invoke.return_value = QueryAnalysis(
data_required=["2024 results"],
unknowns=[],
data_required=["2024 results"],
unknowns=[],
ambiguities=[],
conditions=[],
conditions=[],
next_action="plan"
)
# 2. Mock Planner
mock_planner_instance = MagicMock()
mock_llms["planner"].return_value = mock_planner_instance
mock_planner_instance.with_structured_output.return_value.invoke.return_value = TaskPlanResponse(
mock_planner_instance.with_structured_output.return_value.invoke.return_value = ChecklistResponse(
goal="Get results",
reflection="Reflect",
context=TaskPlanContext(initial_context="Ctx", assumptions=[], constraints=[]),
steps=["Step 1"]
checklist=[ChecklistTask(task="Query Data", worker="data_analyst")]
)
# 3. Mock Coder
# 3. Mock Coder (Worker)
mock_coder_instance = MagicMock()
mock_llms["coder"].return_value = mock_coder_instance
mock_coder_instance.with_structured_output.return_value.invoke.return_value = CodeGenerationResponse(
@@ -62,10 +53,15 @@ def test_workflow_data_analysis_flow(mock_llms):
explanation="Explain"
)
# 4. Mock Summarizer
mock_summarizer_instance = MagicMock()
mock_llms["summarizer"].return_value = mock_summarizer_instance
mock_summarizer_instance.invoke.return_value = AIMessage(content="Final Summary: Success")
# 4. Mock Worker Summarizer
mock_ws_instance = MagicMock()
mock_llms["worker_summarizer"].return_value = mock_ws_instance
mock_ws_instance.invoke.return_value = AIMessage(content="Worker Summary")
# 5. Mock Synthesizer
mock_syn_instance = MagicMock()
mock_llms["synthesizer"].return_value = mock_syn_instance
mock_syn_instance.invoke.return_value = AIMessage(content="Final Summary: Success")
# Initial state
initial_state = {
@@ -73,66 +69,67 @@ def test_workflow_data_analysis_flow(mock_llms):
"question": "Show me 2024 results",
"analysis": None,
"next_action": "",
"plan": None,
"code": None,
"error": None,
"iterations": 0,
"checklist": [],
"current_step": 0,
"vfs": {},
"plots": [],
"dfs": {}
}
# Run the graph
result = app.invoke(initial_state, config={"recursion_limit": 15})
result = app.invoke(initial_state, config={"recursion_limit": 20})
assert result["next_action"] == "plan"
assert "Execution Success" in result["code_output"]
assert "Final Summary: Success" in result["messages"][-1].content
assert "Final Summary: Success" in [m.content for m in result["messages"]]
assert result["current_step"] == 1
def test_workflow_research_flow(mock_llms):
"""Test flow: QueryAnalyzer -> Researcher -> Summarizer."""
"""Test flow with research task."""
# 1. Mock Query Analyzer (routes to research)
# 1. Mock Query Analyzer
mock_qa_instance = MagicMock()
mock_llms["qa"].return_value = mock_qa_instance
mock_qa_instance.with_structured_output.return_value.invoke.return_value = QueryAnalysis(
data_required=[],
unknowns=[],
data_required=[],
unknowns=[],
ambiguities=[],
conditions=[],
conditions=[],
next_action="research"
)
# 2. Mock Researcher
mock_researcher_instance = MagicMock()
mock_llms["researcher"].return_value = mock_researcher_instance
# Researcher node uses bind_tools if it's ChatOpenAI/ChatGoogleGenerativeAI
# Since it's a MagicMock, it will fallback to using the base instance
mock_researcher_instance.invoke.return_value = AIMessage(content="Research Results")
# 2. Mock Planner
mock_planner_instance = MagicMock()
mock_llms["planner"].return_value = mock_planner_instance
mock_planner_instance.with_structured_output.return_value.invoke.return_value = ChecklistResponse(
goal="Search",
reflection="Reflect",
checklist=[ChecklistTask(task="Search Web", worker="researcher")]
)
# Also mock bind_tools just in case we ever use spec
mock_llm_with_tools = MagicMock()
mock_researcher_instance.bind_tools.return_value = mock_llm_with_tools
mock_llm_with_tools.invoke.return_value = AIMessage(content="Research Results")
# 3. Mock Researcher
mock_res_instance = MagicMock()
mock_llms["researcher"].return_value = mock_res_instance
mock_res_instance.invoke.return_value = AIMessage(content="Research Result")
# 3. Mock Summarizer (not used in this flow, but kept for completeness)
mock_summarizer_instance = MagicMock()
mock_llms["summarizer"].return_value = mock_summarizer_instance
mock_summarizer_instance.invoke.return_value = AIMessage(content="Final Summary: Research Success")
# 4. Mock Synthesizer
mock_syn_instance = MagicMock()
mock_llms["synthesizer"].return_value = mock_syn_instance
mock_syn_instance.invoke.return_value = AIMessage(content="Final Research Summary")
# Initial state
initial_state = {
"messages": [],
"question": "Who is the governor of Florida?",
"question": "Who is the governor?",
"analysis": None,
"next_action": "",
"plan": None,
"code": None,
"error": None,
"iterations": 0,
"checklist": [],
"current_step": 0,
"vfs": {},
"plots": [],
"dfs": {}
}
# Run the graph
result = app.invoke(initial_state, config={"recursion_limit": 10})
result = app.invoke(initial_state, config={"recursion_limit": 20})
assert result["next_action"] == "research"
assert "Research Results" in result["messages"][-1].content
assert "Final Research Summary" in [m.content for m in result["messages"]]
assert result["current_step"] == 1