diff --git a/src/graph/builder.py b/src/graph/builder.py index 9da3a81..defe879 100644 --- a/src/graph/builder.py +++ b/src/graph/builder.py @@ -22,14 +22,23 @@ def continue_to_running_research_team(state: State): current_plan = state.get("current_plan") if not current_plan or not current_plan.steps: return "planner" + if all(step.execution_res for step in current_plan.steps): return "planner" + + # Find first incomplete step + incomplete_step = None for step in current_plan.steps: if not step.execution_res: + incomplete_step = step break - if step.step_type and step.step_type == StepType.RESEARCH: + + if not incomplete_step: + return "planner" + + if incomplete_step.step_type == StepType.RESEARCH: return "researcher" - if step.step_type and step.step_type == StepType.PROCESSING: + if incomplete_step.step_type == StepType.PROCESSING: return "coder" return "planner" diff --git a/tests/unit/graph/test_builder.py b/tests/unit/graph/test_builder.py index fbbdaa7..079cac0 100644 --- a/tests/unit/graph/test_builder.py +++ b/tests/unit/graph/test_builder.py @@ -66,6 +66,17 @@ def test_continue_to_running_research_team_next_coder(mock_state): assert builder_mod.continue_to_running_research_team(state) == "coder" +def test_continue_to_running_research_team_next_coder_withresult(mock_state): + Step = mock_state["Step"] + Plan = mock_state["Plan"] + steps = [ + Step(execution_res=True), + Step(execution_res=True, step_type=builder_mod.StepType.PROCESSING), + ] + state = {"current_plan": Plan(steps=steps)} + assert builder_mod.continue_to_running_research_team(state) == "planner" + + def test_continue_to_running_research_team_default_planner(mock_state): Step = mock_state["Step"] Plan = mock_state["Plan"]