fix(orchestrator): Apply refinements from code review
This commit is contained in:
@@ -12,14 +12,16 @@ def mock_llms():
|
||||
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:
|
||||
patch("ea_chatbot.graph.nodes.researcher.get_llm_model") as mock_researcher, \
|
||||
patch("ea_chatbot.graph.nodes.reflector.get_llm_model") as mock_reflector:
|
||||
yield {
|
||||
"qa": mock_qa,
|
||||
"planner": mock_planner,
|
||||
"coder": mock_coder,
|
||||
"worker_summarizer": mock_worker_summarizer,
|
||||
"synthesizer": mock_synthesizer,
|
||||
"researcher": mock_researcher
|
||||
"researcher": mock_researcher,
|
||||
"reflector": mock_reflector
|
||||
}
|
||||
|
||||
def test_workflow_data_analysis_flow(mock_llms):
|
||||
@@ -58,7 +60,12 @@ def test_workflow_data_analysis_flow(mock_llms):
|
||||
mock_llms["worker_summarizer"].return_value = mock_ws_instance
|
||||
mock_ws_instance.invoke.return_value = AIMessage(content="Worker Summary")
|
||||
|
||||
# 5. Mock Synthesizer
|
||||
# 5. Mock Reflector
|
||||
mock_reflector_instance = MagicMock()
|
||||
mock_llms["reflector"].return_value = mock_reflector_instance
|
||||
mock_reflector_instance.with_structured_output.return_value.invoke.return_value = MagicMock(satisfied=True, reasoning="Good.")
|
||||
|
||||
# 6. 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")
|
||||
@@ -111,7 +118,12 @@ def test_workflow_research_flow(mock_llms):
|
||||
mock_llms["researcher"].return_value = mock_res_instance
|
||||
mock_res_instance.invoke.return_value = AIMessage(content="Research Result")
|
||||
|
||||
# 4. Mock Synthesizer
|
||||
# 4. Mock Reflector
|
||||
mock_reflector_instance = MagicMock()
|
||||
mock_llms["reflector"].return_value = mock_reflector_instance
|
||||
mock_reflector_instance.with_structured_output.return_value.invoke.return_value = MagicMock(satisfied=True, reasoning="Good.")
|
||||
|
||||
# 5. 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")
|
||||
|
||||
Reference in New Issue
Block a user