jupyter-agent 2025.6.104__py3-none-any.whl → 2025.7.100__py3-none-any.whl

jupyter_agent/bot_evaluators/flow_task_executor.py

@@ -0,0 +1,152 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import time
+
+from enum import Enum
+from typing import Optional, List
+from pydantic import BaseModel, Field
+from IPython.display import Markdown
+from .base import BaseEvaluator
+from ..bot_outputs import _D, _I, _W, _E, _F, _A, _O, _C, _M, _B
+from ..bot_evaluation import FlowEvaluationRecord
+
+
+FLOW_TASK_EXEC_EVAL_PROMPT = """\
+**角色定义**：
+
+你是一个任务规划评估专家，负责对任务规划的结果进行评估。
+
+**任务要求**：
+
+请你根据任务规划的结果，评估任务规划的质量和准确性，并给出相应的评分和反馈。
+
+{% include "TASK_OUTPUT_FORMAT" %}
+
+---
+
+{% include "TASK_CONTEXTS" %}
+
+---
+
+{% include "CODE_CONTEXTS" %}
+
+---
+
+**当前子任务规划信息**:
+
+### 当前子任务规划目标：
+
+{{ task.subject }}
+
+{% if task.coding_prompt %}
+### 当前子任务代码需求：
+
+{{ task.coding_prompt }}
+
+### 当前子任务生成的代码：
+
+```python
+{{ task.source }}
+```
+
+### 当前代码执行的输出与结果：
+
+{{ task.output }}
+{% endif %}
+
+### 当前子任务总结要求：
+
+{{ task.summary_prompt }}
+
+
+### 当前子任务输出的分析总结后的最终结果：
+
+```markdown
+{{ task.result }}
+```
+
+{% if task.important_infos %}
+### 当前子任务输出的重要信息：
+
+```json
+{{ task.important_infos | json }}
+```
+{% endif %}
+
+{% if task.request_below_supply_infos %}
+### 当前子任务输出的请求用户补充确认的信息：
+
+```json
+{{ task.request_below_supply_infos | json }}
+```
+{% endif %}
+
+---
+
+请按要求给出当前子任务规划的评估结果：
+"""
+
+
+class FlowTaskExecEvalResult(BaseModel):
+    """
+    任务规划评估结果
+    """
+
+    is_correct: bool = Field(description="最终结果是否符合当前子任务的目标", examples=[True, False])
+    correct_score: float = Field(
+        description="最终结果符合当前子任务目标的分数，范围0-1，>=0.5表示符合目标，<0.5表示不符合目标",
+        examples=[0.95, 0.3],
+    )
+    planning_score: float = Field(
+        description="当前子任务的目标规划、代码生成、总结是否符合全局目标规划要求，范围0-1，>=0.5表示符合要求，<0.5表示不符合要求",
+        examples=[0.85, 0.25],
+    )
+    reasoning_score: float = Field(
+        description="当前子任务的推理过程是否合理，是否存在逻辑错误，是否存在与前置子任务相冲突的情况，"
+        "范围0-1，>=0.5表示合理、正确、无冲突，<0.5表示不合理",
+        examples=[0.9, 0.4],
+    )
+    coding_score: float = Field(
+        description="代码生成的质量评分，代码逻辑是否符合规划要求，是否存在逻辑错误，是否存在冗余、抽象不合理等情况，"
+        "范围0-1，>=0.5表示代码质量较高，<0.5表示代码质量较低",
+        examples=[0.75, 0.2],
+    )
+    important_info_score: float = Field(
+        description="重要信息分数，当前子任务的规划、代码生成、总结是否充分考虑了前置任务生成的重要信息，"
+        "以及当前子任务的重要信息是否完整、准确、无误导、无冲突，"
+        "范围0-1，>=0.5表示重要信息完整、准确，<0.5表示重要信息不完整或不准确",
+        examples=[0.9, 0.4],
+    )
+    user_supply_info_score: float = Field(
+        description="用户补充信息分数，当前子任务的规划、代码生成、总结是否充分考虑了用户补充的信息，"
+        "范围0-1，>=0.5表示充分考虑，<0.5表示未充分考虑",
+        examples=[0.8, 0.3],
+    )
+    feedback: Optional[str] = Field(default=None, description="评估反馈")
+
+
+class FlowTaskExecEvaluator(BaseEvaluator):
+    """
+    任务规划评估器
+    """
+
+    PROMPT = FLOW_TASK_EXEC_EVAL_PROMPT
+    OUTPUT_JSON_SCHEMA = FlowTaskExecEvalResult
+
+    def on_reply(self, reply):
+        reply = super().on_reply(reply)
+        return FlowEvaluationRecord(
+            timestamp=time.time(),
+            evaluator="flow_task_executor",
+            correct_score=reply.correct_score,
+            planning_score=reply.planning_score,
+            reasoning_score=reply.reasoning_score,
+            coding_score=reply.coding_score,
+            important_score=reply.important_info_score,
+            user_supply_score=reply.user_supply_info_score,
+        )

jupyter_agent/bot_flows/__init__.py

@@ -7,14 +7,10 @@ https://opensource.org/licenses/MIT
 
 from .base import BaseTaskFlow
 from .master_planner import MasterPlannerFlow
-from .task_executor_v1 import TaskExecutorFlowV1
-from .task_executor_v2 import TaskExecutorFlowV2
 from .task_executor_v3 import TaskExecutorFlowV3
 
 __all__ = [
     "BaseTaskFlow",
     "MasterPlannerFlow",
-    "TaskExecutorFlowV1",
-    "TaskExecutorFlowV2",
     "TaskExecutorFlowV3",
 ]

jupyter_agent/bot_flows/base.py

@@ -13,9 +13,12 @@ from enum import Enum
 from typing import List, Dict, Optional, Type
 from IPython.display import Markdown
 from ..bot_agents.base import BaseAgent
+from ..bot_evaluators.base import BaseEvaluator
+from ..bot_evaluators.dummy_global import DummyGlobalEvaluator
+from ..bot_evaluators.flow_task_executor import FlowTaskExecEvaluator
 from ..bot_outputs import _D, _I, _W, _E, _F, _M, _B
 from ..bot_outputs import set_stage, flush_output, output_evaluation
-from ..bot_outputs import FlowEvalutionRecord, StageEvalutionRecord, NotebookEvalutionRecord
+from ..bot_evaluation import FlowEvaluationRecord, StageEvaluationRecord, NotebookEvaluationRecord
 
 TASK_AGENT_STATE_ERROR = "_AGENT_STATE_ERROR_32534526_"
 TASK_STAGE_START = "start"
@@ -37,9 +40,10 @@ class StageNext[ST](BaseModel):
     message: str = ""
 
 
-class StageTransition[ST, AS](BaseModel):
+class StageNode[ST, AS](BaseModel):
     stage: ST | str
-    agent: Type[BaseAgent] | str
+    agents: Type[BaseAgent] | List[Type[BaseAgent]]
+    evaluators: Optional[Type[BaseEvaluator] | List[Type[BaseEvaluator]]] = None
     states: Dict[AS | str, StageNext[ST] | List[StageNext[ST]] | Dict[TaskAction, StageNext[ST]] | ST | str] = {}
     next_stage: Optional[StageNext[ST] | List[StageNext[ST]] | Dict[TaskAction, StageNext[ST]] | ST | str] = None
 
@@ -49,15 +53,18 @@ class BaseTaskFlow:
     基础任务流程
     """
 
-    STAGE_TRANSITIONS: List[StageTransition] = []
+    STAGE_NODES: List[StageNode] = []
     START_STAGE = TASK_STAGE_START
     STOP_STAGES = [TASK_STAGE_COMPLETED, TASK_STAGE_GLOBAL_FINISHED]
+    FLOW_EVALUATOR = FlowTaskExecEvaluator
+    GLOBAL_EVALUATOR = DummyGlobalEvaluator
 
-    def __init__(self, notebook_context, agent_factory):
+    def __init__(self, notebook_context, agent_factory, evaluator_factory=None):
         self.notebook_context = notebook_context
         self.agent_factory = agent_factory
-        self.stage_transitions = {}
-        self.prepare_stage_transitions()
+        self.evaluator_factory = evaluator_factory
+        self.stage_nodes = {}
+        self.prepare_stage_nodes()
 
     @property
     def task(self):
@@ -67,10 +74,10 @@ class BaseTaskFlow:
     def cells(self):
         return self.notebook_context.cells
 
-    def prepare_stage_transitions(self):
-        for st in self.STAGE_TRANSITIONS:
+    def prepare_stage_nodes(self):
+        for st in self.STAGE_NODES:
             assert not (st.next_stage and st.states), "next_stage and states are mutually exclusive"
-            self.stage_transitions[st.stage] = st
+            self.stage_nodes[st.stage] = st
             if st.next_stage:
                 st.states[TaskAction.DEFAULT] = st.next_stage
                 st.next_stage = None
@@ -99,15 +106,29 @@ class BaseTaskFlow:
             if TASK_AGENT_STATE_ERROR not in st.states:
                 st.states[TASK_AGENT_STATE_ERROR] = {"*": StageNext(stage=st.stage)}
 
-    def get_stage_agent(self, stage):
-        for t in self.STAGE_TRANSITIONS:
+    def get_stage_agents(self, stage) -> List[BaseAgent]:
+        for t in self.STAGE_NODES:
             if t.stage == stage:
-                return self.agent_factory(t.agent)
+                if isinstance(t.agents, list):
+                    return [self.agent_factory(a) for a in t.agents]
+                else:
+                    return [self.agent_factory(t.agents)]
         raise ValueError(f"No agent for stage `{stage}`")
 
+    def get_stage_evaluators(self, stage) -> List[BaseEvaluator]:
+        if self.evaluator_factory is None:
+            return []
+        for t in self.STAGE_NODES:
+            if t.stage == stage:
+                if isinstance(t.evaluators, list):
+                    return [self.evaluator_factory(e) for e in t.evaluators]
+                else:
+                    return [self.evaluator_factory(t.evaluators)]
+        return []
+
     def _get_next_stage_trans(self, stage, state, action=TaskAction.CONTINUE):
 
-        st = self.stage_transitions.get(stage)
+        st = self.stage_nodes.get(stage)
         if st:
             state_ns = st.states.get(state) or st.states.get("*")
             assert state_ns, f"No next stage for stage `{stage}` and state `{state}`"
@@ -157,16 +178,19 @@ class BaseTaskFlow:
         stage_count = 0
         # Initialize the task stage
         stage = stage or self.START_STAGE
+        agent = None
         while n_tries <= max_tries:
             stage_st = time.time()
             try:
                 stage_name = stage.value if isinstance(stage, Enum) else stage
                 stage_name = stage_name.replace(".", "-").capitalize()
                 set_stage(stage_name)
-                agent = self.get_stage_agent(stage)
-                _M(f"**Executing** stage `{stage}` with agent `{type(agent).__name__}` ...")
-                failed, state = agent()
+                agents = self.get_stage_agents(stage)
+                for agent in agents:
+                    _I(f"Executing stage `{stage}` with agent `{type(agent).__name__}` ...")
+                    failed, state = agent()
             except Exception as e:
+                _W(f"Error during task execution stage `{stage}`: `{type(e)}`: `{e}`")
                 _M(f"**Error** during task execution stage `{stage}`: `{type(e)}`: `{e}`")
                 _M(f"```python\n{traceback.format_exc()}\n```")
                 state = TASK_AGENT_STATE_ERROR
@@ -174,16 +198,39 @@ class BaseTaskFlow:
             stage_count += 1
             stage_duration = time.time() - stage_st
             flow_duration += stage_duration
-            _M(f"Stage `{stage}` completed in {stage_duration:.2f} seconds with state `{state}` and failed `{failed}`")
-            output_evaluation(
-                StageEvalutionRecord(
-                    cell_index=self.task.cell_idx,
-                    flow=type(self).__name__,
-                    stage=str(stage),
-                    execution_duration=stage_duration,
-                    is_success=not failed,
+            _I(f"Stage `{stage}` completed in {stage_duration:.2f} seconds with state `{state}` and failed `{failed}`")
+            if evaluators := self.get_stage_evaluators(stage):
+                for evaluator in evaluators:
+                    # If the agent has evaluators, run them
+                    try:
+                        _I(f"Evaluating stage `{stage}` with evaluator `{type(evaluator).__name__}` ...")
+                        evaluation_result = evaluator()
+                        evaluation_result.timestamp = evaluation_result.timestamp or time.time()
+                        evaluation_result.evaluator = evaluation_result.evaluator or type(evaluator).__name__
+                        evaluation_result.cell_index = self.task.cell_idx
+                        evaluation_result.flow = type(self).__name__
+                        evaluation_result.stage = str(stage)
+                        evaluation_result.agent = type(agent).__name__
+                        evaluation_result.execution_duration = stage_duration
+                        evaluation_result.is_success = not failed
+                        output_evaluation(evaluation_result)
+                    except Exception as e:
+                        _W(f"Error during task evaluation stage `{stage}`: `{type(e)}`: `{e}`")
+                        _M(f"**Error** during task evaluation stage `{stage}`: `{type(e)}`: `{e}`")
+                        _M(f"```python\n{traceback.format_exc()}\n```")
+            else:
+                output_evaluation(
+                    StageEvaluationRecord(
+                        timestamp=time.time(),
+                        evaluator="default",
+                        cell_index=self.task.cell_idx,
+                        flow=type(self).__name__,
+                        stage=str(stage),
+                        agent=type(agent).__name__,
+                        execution_duration=stage_duration,
+                        is_success=not failed,
+                    )
                 )
-            )
 
             if state != TASK_AGENT_STATE_ERROR:
                 # Agent did not fail, check if we have reached the final stage
@@ -191,7 +238,7 @@ class BaseTaskFlow:
                 self.task.agent_stage = next_stage
                 self.task.update_cell()
                 if next_stage in self.STOP_STAGES:
-                    _M(f"Task execution **Stopped** at stage `{next_stage}`")
+                    _I(f"Task execution **Stopped** at stage `{next_stage}`")
                     stage = next_stage
                     break
 
@@ -212,37 +259,69 @@ class BaseTaskFlow:
                 self.task.agent_stage = next_stage
                 self.task.update_cell()
                 if action == TaskAction.STOP:
-                    _M(f"Task execution **Stopped**, and set next stage to `{next_stage}`")
+                    _I(f"Task execution **Stopped**, and set next stage to `{next_stage}`")
                     stage = next_stage
                     break
                 else:
-                    _M(f"**Action**: `{action}` transits stage to `{next_stage}`")
+                    _I(f"Action: `{action}` transits stage to `{next_stage}`")
                     stage = next_stage
             else:
                 # transit to the next stage without confirmation
                 next_stage = self.get_next_stage(stage, state, TaskAction.CONTINUE)
                 self.task.agent_stage = next_stage
                 self.task.update_cell()
-                _M(f"**Transits** stage to `{next_stage}`")
+                _I(f"Transits stage to `{next_stage}`")
                 stage = next_stage
-
             if not stage_continue:
                 break
         # Finalize the task execution
         stage_name = stage.value if isinstance(stage, Enum) else stage
         if stage_name == TASK_STAGE_GLOBAL_FINISHED:
             _M("Task execution **finished** globally.")
-            output_evaluation(NotebookEvalutionRecord(cell_index=self.task.cell_idx, is_success=True))
+            if self.evaluator_factory is not None and hasattr(self, "GLOBAL_EVALUATOR") and self.GLOBAL_EVALUATOR:
+                evaluator = self.evaluator_factory(self.GLOBAL_EVALUATOR)
+                _I(f"Evaluating notebook with evaluator `{type(evaluator).__name__}` ...")
+                evaluation_result = evaluator()
+                evaluation_result.timestamp = evaluation_result.timestamp or time.time()
+                evaluation_result.evaluator = evaluation_result.evaluator or type(evaluator).__name__
+                evaluation_result.cell_index = self.task.cell_idx
+                evaluation_result.is_success = True
+                output_evaluation(evaluation_result)
+            else:
+                output_evaluation(
+                    NotebookEvaluationRecord(
+                        timestamp=time.time(),
+                        evaluator="default",
+                        cell_index=self.task.cell_idx,
+                        is_success=True,
+                    )
+                )
         elif stage_name == TASK_STAGE_COMPLETED:
-            _M(f"Task execution **completed** in {flow_duration:.2f} seconds with {stage_count} stages.")
-            output_evaluation(
-                FlowEvalutionRecord(
-                    cell_index=self.task.cell_idx,
-                    flow=type(self).__name__,
-                    stage_count=stage_count,
-                    execution_duration=flow_duration,
-                    is_success=True,
+            _I(f"Task execution **completed** in {flow_duration:.2f} seconds with {stage_count} stages.")
+            if self.evaluator_factory is not None and hasattr(self, "FLOW_EVALUATOR") and self.FLOW_EVALUATOR:
+                evaluator = self.evaluator_factory(self.FLOW_EVALUATOR)
+                _I(f"Evaluating flow `{type(self).__name__}` with evaluator `{type(evaluator).__name__}` ...")
+                evaluation_result = evaluator()
+                evaluation_result.timestamp = evaluation_result.timestamp or time.time()
+                evaluation_result.evaluator = evaluation_result.evaluator or type(evaluator).__name__
+                evaluation_result.cell_index = self.task.cell_idx
+                evaluation_result.flow = type(self).__name__
+                evaluation_result.stage_count = stage_count
+                evaluation_result.execution_duration = flow_duration
+                evaluation_result.is_success = True
+                output_evaluation(evaluation_result)
+            else:
+                # If no evaluator, just output the evaluation record
+                output_evaluation(
+                    FlowEvaluationRecord(
+                        timestamp=time.time(),
+                        evaluator="default",
+                        cell_index=self.task.cell_idx,
+                        flow=type(self).__name__,
+                        stage_count=stage_count,
+                        execution_duration=flow_duration,
+                        is_success=True,
+                    )
                 )
-            )
         flush_output()
         return stage

jupyter_agent/bot_flows/master_planner.py

@@ -5,13 +5,24 @@ This software is released under the MIT License.
 https://opensource.org/licenses/MIT
 """
 
-from .base import BaseTaskFlow, StageTransition, TASK_STAGE_START, TASK_STAGE_COMPLETED
+from .base import BaseTaskFlow, StageNode, TASK_STAGE_START, TASK_STAGE_COMPLETED
+from ..bot_evaluators.flow_global_planning import FlowGlobalPlanningEvaluator
+from ..bot_agents.master_planner import MasterPlannerAgent
+from ..bot_agents.output_task_result import OutputTaskResult
+from ..bot_agents.prepare_next_cell import PrepareNextCell
+from ..bot_evaluators.dummy_task import DummyTaskEvaluator
 
 
 class MasterPlannerFlow(BaseTaskFlow):
 
-    STAGE_TRANSITIONS = [
-        StageTransition(stage=TASK_STAGE_START, agent="MasterPlannerAgent", next_stage=TASK_STAGE_COMPLETED),
-        StageTransition(stage=TASK_STAGE_COMPLETED, agent="OutputTaskResult", next_stage=TASK_STAGE_COMPLETED),
+    STAGE_NODES = [
+        StageNode(
+            stage=TASK_STAGE_START,
+            agents=MasterPlannerAgent,
+            evaluators=DummyTaskEvaluator,
+            next_stage=TASK_STAGE_COMPLETED,
+        ),
+        StageNode(stage=TASK_STAGE_COMPLETED, agents=OutputTaskResult, next_stage=TASK_STAGE_COMPLETED),
     ]
     STOP_STAGES = [TASK_STAGE_COMPLETED]
+    FLOW_EVALUATOR = FlowGlobalPlanningEvaluator

jupyter_agent/bot_flows/task_executor_v3.py

@@ -8,32 +8,34 @@ https://opensource.org/licenses/MIT
 from enum import Enum
 from .base import (
     BaseTaskFlow,
-    StageTransition,
+    StageNode,
     StageNext,
     TaskAction,
     TASK_STAGE_COMPLETED,
     TASK_STAGE_GLOBAL_FINISHED,
 )
-from ..bot_agents import (
-    TaskPlannerAgentV3,
-    TaskCodingAgent,
-    CodeDebugerAgent,
-    CodeExecutor,
-    TaskStructureSummaryAgent,
-    TaskStructureReasoningAgent,
-    OutputTaskResult,
-)
-from ..bot_agents.task_planner_v3 import TaskPlannerState
+from ..bot_agents.task_planner_v3 import TaskPlannerAgentV3, TaskPlannerState
+from ..bot_agents.task_coder import TaskCodingAgent
+from ..bot_agents.task_debuger import CodeDebugerAgent
+from ..bot_agents.task_code_executor import CodeExecutor
+from ..bot_agents.task_structrue_summarier import TaskStructureSummaryAgent, TaskStructureSummaryState
+from ..bot_agents.task_structrue_reasoner import TaskStructureReasoningAgent, TaskStructureReasonState
+from ..bot_agents.output_task_result import OutputTaskResult
+from ..bot_agents.request_user_supply import RequestAboveUserSupplyAgent, RequestBelowUserSupplyAgent
+from ..bot_agents.prepare_next_cell import PrepareNextCell
 
 
 class TaskStage(str, Enum):
     PLANNING = "planning"
     PLANNING_PAUSED = "planning_paused"
+    REQUEST_INFO_ABOVE = "request_info_above"
+    REQUEST_INFO_BELOW = "request_info_below"
     CODING = "coding"
     EXECUTING = "executing"
     DEBUGGING = "debugging"
     REASONING = "reasoning"
     SUMMARY = "summary"
+    PREPARE_NEXT = "prepare_next"
     OUTPUT_RESULT = "output_result"
     COMPLETED = TASK_STAGE_COMPLETED
     GLOBAL_FINISHED = TASK_STAGE_GLOBAL_FINISHED
@@ -43,20 +45,25 @@ class TaskExecutorFlowV3(BaseTaskFlow):
 
     START_STAGE = TaskStage.PLANNING
     STOP_STAGES = [TaskStage.COMPLETED, TaskStage.PLANNING_PAUSED, TaskStage.GLOBAL_FINISHED]
-    STAGE_TRANSITIONS = [
-        StageTransition[TaskStage, TaskPlannerState](
+    STAGE_NODES = [
+        StageNode[TaskStage, TaskPlannerState](
             stage=TaskStage.PLANNING,
-            agent=TaskPlannerAgentV3,
+            agents=TaskPlannerAgentV3,
             states={
                 TaskPlannerState.CODING_PLANNED: TaskStage.CODING,
                 TaskPlannerState.REASONING_PLANNED: TaskStage.REASONING,
-                TaskPlannerState.REQUEST_INFO: TaskStage.PLANNING_PAUSED,
+                TaskPlannerState.REQUEST_INFO: TaskStage.REQUEST_INFO_ABOVE,
                 TaskPlannerState.GLOBAL_FINISHED: TaskStage.GLOBAL_FINISHED,
             },
         ),
-        StageTransition[TaskStage, TaskPlannerState](
+        StageNode[TaskStage, None](
+            stage=TaskStage.REQUEST_INFO_ABOVE,
+            agents=RequestAboveUserSupplyAgent,
+            next_stage=TaskStage.PLANNING_PAUSED,
+        ),
+        StageNode[TaskStage, TaskPlannerState](
             stage=TaskStage.PLANNING_PAUSED,
-            agent=TaskPlannerAgentV3,
+            agents=TaskPlannerAgentV3,
             states={
                 TaskPlannerState.CODING_PLANNED: TaskStage.CODING,
                 TaskPlannerState.REASONING_PLANNED: TaskStage.REASONING,
@@ -64,37 +71,49 @@ class TaskExecutorFlowV3(BaseTaskFlow):
                 TaskPlannerState.GLOBAL_FINISHED: TaskStage.COMPLETED,
             },
         ),
-        StageTransition[TaskStage, None](
-            stage=TaskStage.CODING, agent=TaskCodingAgent, next_stage=TaskStage.EXECUTING
-        ),
-        StageTransition[TaskStage, bool](
+        StageNode[TaskStage, None](stage=TaskStage.CODING, agents=TaskCodingAgent, next_stage=TaskStage.EXECUTING),
+        StageNode[TaskStage, bool](
             stage=TaskStage.EXECUTING,
-            agent=CodeExecutor,
+            agents=CodeExecutor,
             states={True: TaskStage.SUMMARY, False: TaskStage.DEBUGGING},
         ),
-        StageTransition[TaskStage, None](
-            stage=TaskStage.DEBUGGING, agent=CodeDebugerAgent, next_stage=TaskStage.EXECUTING
-        ),
-        StageTransition[TaskStage, None](
-            stage=TaskStage.REASONING, agent=TaskStructureReasoningAgent, next_stage=TaskStage.COMPLETED
+        StageNode[TaskStage, None](stage=TaskStage.DEBUGGING, agents=CodeDebugerAgent, next_stage=TaskStage.EXECUTING),
+        StageNode[TaskStage, TaskStructureReasonState](
+            stage=TaskStage.REASONING,
+            agents=TaskStructureReasoningAgent,
+            states={
+                TaskStructureReasonState.DONE: TaskStage.COMPLETED,
+                TaskStructureReasonState.REQUEST_INFO: TaskStage.REQUEST_INFO_BELOW,
+            },
         ),
-        StageTransition[TaskStage, None](
+        StageNode[TaskStage, TaskStructureSummaryState](
             stage=TaskStage.SUMMARY,
-            agent=TaskStructureSummaryAgent,
-            next_stage={
-                TaskAction.DEFAULT: StageNext(stage=TaskStage.COMPLETED),
-                TaskAction.STOP: StageNext(stage=TaskStage.EXECUTING),
+            agents=TaskStructureSummaryAgent,
+            states={
+                TaskStructureSummaryState.DONE: {
+                    TaskAction.DEFAULT: StageNext(stage=TaskStage.PREPARE_NEXT),
+                    TaskAction.STOP: StageNext(stage=TaskStage.EXECUTING),
+                },
+                TaskStructureSummaryState.REQUEST_INFO: TaskStage.REQUEST_INFO_BELOW,
             },
         ),
-        StageTransition[TaskStage, bool](
+        StageNode[TaskStage, None](
+            stage=TaskStage.PREPARE_NEXT, agents=PrepareNextCell, next_stage=TaskStage.COMPLETED
+        ),
+        StageNode[TaskStage, None](
+            stage=TaskStage.REQUEST_INFO_BELOW,
+            agents=[PrepareNextCell, RequestBelowUserSupplyAgent],
+            next_stage=TaskStage.COMPLETED,
+        ),
+        StageNode[TaskStage, bool](
             stage=TaskStage.COMPLETED,
-            agent=CodeExecutor,
+            agents=CodeExecutor,
             states={True: TaskStage.OUTPUT_RESULT, False: TaskStage.DEBUGGING},
         ),
-        StageTransition[TaskStage, None](
-            stage=TaskStage.OUTPUT_RESULT, agent=OutputTaskResult, next_stage=TaskStage.COMPLETED
+        StageNode[TaskStage, None](
+            stage=TaskStage.OUTPUT_RESULT, agents=OutputTaskResult, next_stage=TaskStage.COMPLETED
         ),
-        StageTransition[TaskStage, None](
-            stage=TaskStage.GLOBAL_FINISHED, agent=OutputTaskResult, next_stage=TaskStage.GLOBAL_FINISHED
+        StageNode[TaskStage, None](
+            stage=TaskStage.GLOBAL_FINISHED, agents=OutputTaskResult, next_stage=TaskStage.GLOBAL_FINISHED
         ),
     ]