jupyter-agent 2025.6.100__py3-none-any.whl → 2025.6.101__py3-none-any.whl

jupyter_agent/bot_flows/base.py

@@ -0,0 +1,209 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import traceback
+
+from pydantic import BaseModel
+from enum import Enum
+from typing import List, Dict, Optional, Type
+from IPython.display import Markdown
+from ..bot_agents.base import BaseAgent
+from ..bot_outputs import _D, _I, _W, _E, _F, _M, _B, set_stage, flush_output
+
+TASK_AGENT_STATE_ERROR = "_AGENT_STATE_ERROR_32534526_"
+TASK_STAGE_START = "start"
+TASK_STAGE_COMPLETED = "completed"
+
+
+class TaskAction(str, Enum):
+    DEFAULT = "*"
+    CONTINUE = "continue"
+    RETRY = "retry"
+    SKIP = "skip"
+    STOP = "stop"
+
+
+class StageNext[ST](BaseModel):
+    action: Optional[TaskAction] = None
+    stage: ST | str
+    message: str = ""
+
+
+class StageTransition[ST, AS](BaseModel):
+    stage: ST | str
+    agent: Type[BaseAgent] | str
+    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
+
+
+class BaseTaskFlow:
+    """
+    基础任务流程
+    """
+
+    STAGE_TRANSITIONS: List[StageTransition] = []
+    START_STAGE = TASK_STAGE_START
+    STOP_STAGES = [TASK_STAGE_COMPLETED]
+
+    def __init__(self, notebook_context, agent_factory):
+        self.notebook_context = notebook_context
+        self.agent_factory = agent_factory
+        self.stage_transitions = {}
+        self.prepare_stage_transitions()
+
+    @property
+    def task(self):
+        return self.notebook_context.cur_task
+
+    @property
+    def cells(self):
+        return self.notebook_context.cells
+
+    def prepare_stage_transitions(self):
+        for st in self.STAGE_TRANSITIONS:
+            assert not (st.next_stage and st.states), "next_stage and states are mutually exclusive"
+            self.stage_transitions[st.stage] = st
+            if st.next_stage:
+                st.states[TaskAction.DEFAULT] = st.next_stage
+                st.next_stage = None
+            for state, ns in st.states.items():
+                st.states[state] = {}
+                if isinstance(ns, str):
+                    st.states[state] = {TaskAction.DEFAULT: StageNext(stage=ns)}
+                elif isinstance(ns, StageNext):
+                    action = ns.action or TaskAction.DEFAULT
+                    st.states[state] = {action: ns}
+                elif isinstance(ns, list):
+                    state_dict = {}
+                    for n in ns:
+                        action = n.action or TaskAction.DEFAULT
+                        state_dict[action] = n
+                    st.states[state] = state_dict
+                elif isinstance(ns, dict):
+                    st.states[state] = ns
+                else:
+                    raise ValueError(f"Unknown next stage: {ns}")
+                state_ns: Dict = st.states[state]  # type: ignore
+                state_ns[TaskAction.CONTINUE] = state_ns.get(TaskAction.CONTINUE) or state_ns.get(TaskAction.DEFAULT)
+                state_ns[TaskAction.RETRY] = state_ns.get(TaskAction.RETRY) or StageNext(stage=st.stage)
+                state_ns[TaskAction.STOP] = state_ns.get(TaskAction.STOP) or StageNext(stage=st.stage)
+                state_ns[TaskAction.SKIP] = state_ns.get(TaskAction.SKIP) or state_ns.get(TaskAction.CONTINUE)
+            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:
+            if t.stage == stage:
+                return self.agent_factory(t.agent)
+        raise ValueError(f"No agent for stage `{stage}`")
+
+    def _get_next_stage_trans(self, stage, state, action=TaskAction.CONTINUE):
+
+        st = self.stage_transitions.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}`"
+            act_ns = state_ns.get(action) or state_ns.get("*")
+            assert act_ns, f"No next stage for stage `{stage}`, state `{state}`, action `{action}`"
+            return act_ns
+        else:
+            raise ValueError(f"No next stage for stage `{stage}`")
+
+    def get_prompt_message(self, stage, state, failed):
+
+        ns = self._get_next_stage_trans(stage, state)
+        if failed:
+            msg = ns.message or f"Staget `{stage}` FAILED!"
+            return (
+                f"{msg}\n Continue from stage `{ns.stage}`? \n"
+                f"(C)ontinue, (R)etry, s(K)ip, (S)top, default `continue`"
+            )
+        else:
+            return (
+                f"{ns.message}\n Continue to stage `{ns.stage}`? \n"
+                f"(C)ontinue, (R)etry, s(K)ip, (S)top, default `continue`"
+            )
+
+    def match_action(self, input):
+        input = input.lower().strip()
+        if input == "" or input == "c" or (len(input) > 1 and TaskAction.CONTINUE.value.startswith(input)):
+            return TaskAction.CONTINUE
+        elif input == "r" or (len(input) > 1 and TaskAction.RETRY.value.startswith(input)):
+            return TaskAction.RETRY
+        elif input == "k" or (len(input) > 1 and TaskAction.SKIP.value.startswith(input)):
+            return TaskAction.SKIP
+        elif input == "s" or (len(input) > 1 and TaskAction.STOP.value.startswith(input)):
+            return TaskAction.STOP
+        else:
+            raise ValueError(f"Unknown action: {input}")
+
+    def get_next_stage(self, stage, state, action):
+
+        ns = self._get_next_stage_trans(stage, state, action)
+        return ns.stage
+
+    def __call__(self, stage, max_tries=3, stage_continue=True, stage_confirm=True):
+
+        n_tries = 0
+        stage = stage or self.START_STAGE
+        while n_tries <= max_tries:
+            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()
+            except Exception as 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
+                failed = True
+
+            if state != TASK_AGENT_STATE_ERROR:
+                # Agent did not fail, check if we have reached the final stage
+                next_stage = self.get_next_stage(stage, state, TaskAction.CONTINUE)
+                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}`")
+                    break
+
+            if failed:
+                # Agent failed
+                n_tries += 1
+
+            if failed or stage_confirm:
+                # Agent failed or we need to confirm
+                message = self.get_prompt_message(stage, state, failed)
+                _M("**Confirm**: " + message)
+                flush_output()
+                action = self.match_action(input(message))
+                next_stage = self.get_next_stage(stage, state, action)
+                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}`")
+                    break
+                elif n_tries > max_tries:
+                    _M(f"**Max tries reached** during task execution stage `{stage}`, **Stop!**")
+                    break
+                else:
+                    _M(f"**Action**: `{action}` transits stage to `{next_stage}`")
+                    stage = next_stage
+            else:
+                # Agent succeeded, 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}`")
+                stage = next_stage
+
+            if not stage_continue:
+                break
+        flush_output()
+        return stage

jupyter_agent/bot_flows/master_planner.py

@@ -0,0 +1,16 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+from .base import BaseTaskFlow, StageTransition, TASK_STAGE_START, TASK_STAGE_COMPLETED
+
+
+class MasterPlannerFlow(BaseTaskFlow):
+
+    STAGE_TRANSITIONS = [
+        StageTransition(stage=TASK_STAGE_START, agent="MasterPlannerAgent", next_stage=TASK_STAGE_START)
+    ]
+    STOP_STAGES = [TASK_STAGE_START]

jupyter_agent/bot_flows/task_executor_v1.py

@@ -0,0 +1,86 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+from enum import Enum
+from .base import BaseTaskFlow, StageTransition, StageNext, TaskAction
+from ..bot_agents import (
+    TaskPlannerAgentV1,
+    TaskCodingAgent,
+    CodeDebugerAgent,
+    CodeExecutor,
+    TaskVerifyAgent,
+    TaskSummaryAgent,
+    TaskVerifyState,
+)
+from ..bot_agents.task_planner_v1 import TaskPlannerState
+
+
+class TaskStage(str, Enum):
+    PLANNING = "planning"
+    PLANNING_PAUSED = "planning_paused"
+    CODING = "coding"
+    EXECUTING = "executing"
+    DEBUGGING = "debugging"
+    VERIFYING = "verifying"
+    SUMMARY = "summary"
+    COMPLETED = "completed"
+
+
+class TaskExecutorFlowV1(BaseTaskFlow):
+
+    START_STAGE = TaskStage.PLANNING
+    STOP_STAGES = [TaskStage.COMPLETED, TaskStage.PLANNING_PAUSED]
+    STAGE_TRANSITIONS = [
+        StageTransition[TaskStage, TaskPlannerState](
+            stage=TaskStage.PLANNING,
+            agent=TaskPlannerAgentV1,
+            states={
+                TaskPlannerState.PLANNED: TaskStage.CODING,
+                TaskPlannerState.REQUEST_INFO: TaskStage.PLANNING_PAUSED,
+                TaskPlannerState.GLOBAL_FINISHED: TaskStage.COMPLETED,
+            },
+        ),
+        StageTransition[TaskStage, TaskPlannerState](
+            stage=TaskStage.PLANNING_PAUSED,
+            agent=TaskPlannerAgentV1,
+            states={
+                TaskPlannerState.PLANNED: TaskStage.CODING,
+                TaskPlannerState.REQUEST_INFO: TaskStage.PLANNING_PAUSED,
+                TaskPlannerState.GLOBAL_FINISHED: TaskStage.COMPLETED,
+            },
+        ),
+        StageTransition[TaskStage, None](
+            stage=TaskStage.CODING, agent=TaskCodingAgent, next_stage=TaskStage.EXECUTING
+        ),
+        StageTransition[TaskStage, bool](
+            stage=TaskStage.EXECUTING,
+            agent=CodeExecutor,
+            states={True: TaskStage.VERIFYING, False: TaskStage.DEBUGGING},
+        ),
+        StageTransition[TaskStage, None](
+            stage=TaskStage.DEBUGGING, agent=CodeDebugerAgent, next_stage=TaskStage.EXECUTING
+        ),
+        StageTransition[TaskStage, TaskVerifyState](
+            stage=TaskStage.VERIFYING,
+            agent=TaskVerifyAgent,
+            states={
+                TaskVerifyState.PASSED: TaskStage.SUMMARY,
+                TaskVerifyState.FAILED: [
+                    StageNext[TaskStage](action=TaskAction.CONTINUE, stage=TaskStage.PLANNING),
+                    StageNext[TaskStage](action=TaskAction.SKIP, stage=TaskStage.SUMMARY),
+                ],
+            },
+        ),
+        StageTransition[TaskStage, None](
+            stage=TaskStage.SUMMARY, agent=TaskSummaryAgent, next_stage=TaskStage.COMPLETED
+        ),
+        StageTransition[TaskStage, bool](
+            stage=TaskStage.COMPLETED,
+            agent=CodeExecutor,
+            states={True: TaskStage.COMPLETED, False: TaskStage.DEBUGGING},
+        ),
+    ]

jupyter_agent/bot_flows/task_executor_v2.py

@@ -0,0 +1,84 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+from enum import Enum
+from .base import BaseTaskFlow, StageTransition, StageNext, TaskAction
+from ..bot_agents import (
+    TaskPlannerAgentV2,
+    TaskCodingAgent,
+    CodeDebugerAgent,
+    CodeExecutor,
+    TaskSummaryAgent,
+    TaskReasoningAgent,
+)
+from ..bot_agents.task_planner_v2 import TaskPlannerState
+
+
+class TaskStage(str, Enum):
+    PLANNING = "planning"
+    PLANNING_PAUSED = "planning_paused"
+    CODING = "coding"
+    EXECUTING = "executing"
+    DEBUGGING = "debugging"
+    REASONING = "reasoning"
+    SUMMARY = "summary"
+    COMPLETED = "completed"
+
+
+class TaskExecutorFlowV2(BaseTaskFlow):
+
+    START_STAGE = TaskStage.PLANNING
+    STOP_STAGES = [TaskStage.COMPLETED, TaskStage.PLANNING_PAUSED]
+    STAGE_TRANSITIONS = [
+        StageTransition[TaskStage, TaskPlannerState](
+            stage=TaskStage.PLANNING,
+            agent=TaskPlannerAgentV2,
+            states={
+                TaskPlannerState.CODING_PLANNED: TaskStage.CODING,
+                TaskPlannerState.REASONING_PLANNED: TaskStage.REASONING,
+                TaskPlannerState.REQUEST_INFO: TaskStage.PLANNING_PAUSED,
+                TaskPlannerState.GLOBAL_FINISHED: TaskStage.COMPLETED,
+            },
+        ),
+        StageTransition[TaskStage, TaskPlannerState](
+            stage=TaskStage.PLANNING_PAUSED,
+            agent=TaskPlannerAgentV2,
+            states={
+                TaskPlannerState.CODING_PLANNED: TaskStage.CODING,
+                TaskPlannerState.REASONING_PLANNED: TaskStage.REASONING,
+                TaskPlannerState.REQUEST_INFO: TaskStage.PLANNING_PAUSED,
+                TaskPlannerState.GLOBAL_FINISHED: TaskStage.COMPLETED,
+            },
+        ),
+        StageTransition[TaskStage, None](
+            stage=TaskStage.CODING, agent=TaskCodingAgent, next_stage=TaskStage.EXECUTING
+        ),
+        StageTransition[TaskStage, bool](
+            stage=TaskStage.EXECUTING,
+            agent=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=TaskReasoningAgent, next_stage=TaskStage.COMPLETED
+        ),
+        StageTransition[TaskStage, None](
+            stage=TaskStage.SUMMARY,
+            agent=TaskSummaryAgent,
+            next_stage={
+                TaskAction.DEFAULT: StageNext(stage=TaskStage.COMPLETED),
+                TaskAction.STOP: StageNext(stage=TaskStage.EXECUTING),
+            },
+        ),
+        StageTransition[TaskStage, bool](
+            stage=TaskStage.COMPLETED,
+            agent=CodeExecutor,
+            states={True: TaskStage.COMPLETED, False: TaskStage.DEBUGGING},
+        ),
+    ]

jupyter_agent/bot_flows/task_executor_v3.py

@@ -0,0 +1,89 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+from enum import Enum
+from .base import BaseTaskFlow, StageTransition, StageNext, TaskAction
+from ..bot_agents import (
+    TaskPlannerAgentV3,
+    TaskCodingAgent,
+    CodeDebugerAgent,
+    CodeExecutor,
+    TaskStructureSummaryAgent,
+    TaskStructureReasoningAgent,
+    OutputTaskResult,
+)
+from ..bot_agents.task_planner_v3 import TaskPlannerState
+
+
+class TaskStage(str, Enum):
+    PLANNING = "planning"
+    PLANNING_PAUSED = "planning_paused"
+    CODING = "coding"
+    EXECUTING = "executing"
+    DEBUGGING = "debugging"
+    REASONING = "reasoning"
+    SUMMARY = "summary"
+    COMPLETED = "completed"
+    OUTPUT_RESULT = "output_result"
+
+
+class TaskExecutorFlowV3(BaseTaskFlow):
+
+    START_STAGE = TaskStage.PLANNING
+    STOP_STAGES = [TaskStage.COMPLETED, TaskStage.PLANNING_PAUSED]
+    STAGE_TRANSITIONS = [
+        StageTransition[TaskStage, TaskPlannerState](
+            stage=TaskStage.PLANNING,
+            agent=TaskPlannerAgentV3,
+            states={
+                TaskPlannerState.CODING_PLANNED: TaskStage.CODING,
+                TaskPlannerState.REASONING_PLANNED: TaskStage.REASONING,
+                TaskPlannerState.REQUEST_INFO: TaskStage.PLANNING_PAUSED,
+                TaskPlannerState.GLOBAL_FINISHED: TaskStage.COMPLETED,
+            },
+        ),
+        StageTransition[TaskStage, TaskPlannerState](
+            stage=TaskStage.PLANNING_PAUSED,
+            agent=TaskPlannerAgentV3,
+            states={
+                TaskPlannerState.CODING_PLANNED: TaskStage.CODING,
+                TaskPlannerState.REASONING_PLANNED: TaskStage.REASONING,
+                TaskPlannerState.REQUEST_INFO: TaskStage.PLANNING_PAUSED,
+                TaskPlannerState.GLOBAL_FINISHED: TaskStage.COMPLETED,
+            },
+        ),
+        StageTransition[TaskStage, None](
+            stage=TaskStage.CODING, agent=TaskCodingAgent, next_stage=TaskStage.EXECUTING
+        ),
+        StageTransition[TaskStage, bool](
+            stage=TaskStage.EXECUTING,
+            agent=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
+        ),
+        StageTransition[TaskStage, None](
+            stage=TaskStage.SUMMARY,
+            agent=TaskStructureSummaryAgent,
+            next_stage={
+                TaskAction.DEFAULT: StageNext(stage=TaskStage.COMPLETED),
+                TaskAction.STOP: StageNext(stage=TaskStage.EXECUTING),
+            },
+        ),
+        StageTransition[TaskStage, bool](
+            stage=TaskStage.COMPLETED,
+            agent=CodeExecutor,
+            states={True: TaskStage.OUTPUT_RESULT, False: TaskStage.DEBUGGING},
+        ),
+        StageTransition[TaskStage, None](
+            stage=TaskStage.OUTPUT_RESULT, agent=OutputTaskResult, next_stage=TaskStage.COMPLETED
+        ),
+    ]

jupyter_agent/bot_magics.py

@@ -0,0 +1,127 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import time
+import shlex
+import argparse
+import ipynbname
+import traceback
+
+from IPython.display import Markdown
+from IPython.core.magic import Magics, magics_class, cell_magic
+from traitlets import Unicode, Int, Bool
+from traitlets.config.configurable import Configurable
+from .bot_contexts import NotebookContext, AgentCellContext
+from .bot_agents import AgentFactory
+from .bot_agents.base import AgentModelType
+from .bot_flows import MasterPlannerFlow, TaskExecutorFlowV1, TaskExecutorFlowV2, TaskExecutorFlowV3
+from .bot_outputs import _D, _I, _W, _E, _F, _M, _B, _O, reset_output, set_logging_level
+
+
+@magics_class
+class BotMagics(Magics, Configurable):
+    """Jupyter cell magic for OpenAI bot integration"""
+
+    # 配置项
+    logging_level = Unicode("INFO", help="Debug level for logging").tag(config=True)
+    default_api_url = Unicode(None, allow_none=True, help="Default API URL").tag(config=True)
+    default_api_key = Unicode("API_KEY", help="Default API Key").tag(config=True)
+    default_model_name = Unicode("", help="Default Model Name").tag(config=True)
+    planner_api_url = Unicode(None, allow_none=True, help="Planner API URL").tag(config=True)
+    planner_api_key = Unicode("API_KEY", help="Planner API Key").tag(config=True)
+    planner_model_name = Unicode("", help="Planner Model Name").tag(config=True)
+    coding_api_url = Unicode(None, allow_none=True, help="Coding API URL").tag(config=True)
+    coding_api_key = Unicode("API_KEY", help="Coding API Key").tag(config=True)
+    coding_model_name = Unicode("", help="Coding Model Name").tag(config=True)
+    reasoning_api_url = Unicode(None, allow_none=True, help="Reasoning API URL").tag(config=True)
+    reasoning_api_key = Unicode("API_KEY", help="Reasoning API Key").tag(config=True)
+    reasoning_model_name = Unicode("", help="Reasoning Model Name").tag(config=True)
+    display_message = Bool(False, help="Display chat message").tag(config=True)
+    display_think = Bool(True, help="Display chatthink response").tag(config=True)
+    display_response = Bool(False, help="Display chat full response").tag(config=True)
+    notebook_path = Unicode(None, allow_none=True, help="Path to Notebook file").tag(config=True)
+    default_task_flow = Unicode("v3", allow_none=True, help="Default task flow").tag(config=True)
+    support_save_meta = Bool(False, help="Support save metadata to cell").tag(config=True)
+
+    def parse_args(self, line):
+        """解析命令行参数"""
+        parser = argparse.ArgumentParser()
+        parser.add_argument("-l", "--logging-level", type=str, default=self.logging_level, help="level for logging")
+        parser.add_argument("-P", "--planning", action="store_true", default=False, help="Run in planning mode")
+        parser.add_argument("-s", "--stage", type=str, default=None, help="Task stage")
+        parser.add_argument("-f", "--flow", type=str, default=self.default_task_flow, help="Flow name")
+        parser.add_argument("-m", "--max-tries", type=int, default=3, help="Max tries")
+        parser.add_argument("-S", "--step-mode", action="store_true", default=False, help="Run in single step mode")
+        parser.add_argument("-Y", "--auto-confirm", action="store_true", default=False, help="Run without confirm")
+        options, _ = parser.parse_known_args(shlex.split(line.strip()))
+
+        return options
+
+    @cell_magic
+    def bot(self, line, cell):
+        """Jupyter cell magic: %%bot"""
+        try:
+            AgentCellContext.SUPPORT_SAVE_META = self.support_save_meta
+            reset_output(stage="Logging", logging_level=self.logging_level)
+            _I("Cell magic %%bot executing ...")
+            _D(f"Cell magic called with line: {line}")
+            _D(f"Cell magic called with cell: {repr(cell)[:50]} ...")
+            if not cell.strip():
+                _O(
+                    Markdown(
+                        "The cell is **empty**, we can't do anything.\n\n"
+                        "We will fill it with some random characters, please **RERUN** the cell again."
+                    )
+                )
+                if self.shell is not None:
+                    self.shell.set_next_input(
+                        "%%bot {}\n\n# {}".format(line.strip(), time.strftime("%Y-%m-%d %H:%M:%S")), replace=True
+                    )
+                return
+            options = self.parse_args(line)
+            _D(f"Cell magic called with options: {options}")
+            set_logging_level(options.logging_level)
+            self.notebook_path = self.notebook_path or ipynbname.path()
+            _D(f"Cell magic called with notebook path: {self.notebook_path}")
+            nb_context = NotebookContext(line, cell, notebook_path=self.notebook_path)
+            agent_factory = AgentFactory(
+                nb_context,
+                display_think=self.display_think,
+                display_message=self.display_message,
+                display_response=self.display_response,
+            )
+            agent_factory.config_model(
+                AgentModelType.DEFAULT, self.default_api_url, self.default_api_key, self.default_model_name
+            )
+            agent_factory.config_model(
+                AgentModelType.PLANNER, self.planner_api_url, self.planner_api_key, self.planner_model_name
+            )
+            agent_factory.config_model(
+                AgentModelType.CODING, self.coding_api_url, self.coding_api_key, self.coding_model_name
+            )
+            agent_factory.config_model(
+                AgentModelType.REASONING, self.reasoning_api_url, self.reasoning_api_key, self.reasoning_model_name
+            )
+            if options.planning:
+                flow = MasterPlannerFlow(nb_context, agent_factory)
+            else:
+                if options.flow == "v1":
+                    flow = TaskExecutorFlowV1(nb_context, agent_factory)
+                elif options.flow == "v2":
+                    flow = TaskExecutorFlowV2(nb_context, agent_factory)
+                elif options.flow == "v3":
+                    flow = TaskExecutorFlowV3(nb_context, agent_factory)
+                else:
+                    raise ValueError(f"Unknown flow: {options.flow}")
+            flow(options.stage, options.max_tries, not options.step_mode, not options.auto_confirm)
+        except Exception as e:
+            traceback.print_exc()
+
+
+def load_ipython_extension(ipython):
+    """Load the bot magic extension."""
+    ipython.register_magics(BotMagics)