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

jupyter_agent/bot_agents/__init__.py

@@ -0,0 +1,42 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+from .base import BaseChatAgent, AgentFactory
+from .master_planner import MasterPlannerAgent
+from .output_task_result import OutputTaskResult
+from .task_code_executor import CodeExecutor
+from .task_planner_v1 import TaskPlannerAgentV1
+from .task_planner_v2 import TaskPlannerAgentV2
+from .task_planner_v3 import TaskPlannerAgentV3
+from .task_coder import TaskCodingAgent
+from .task_debuger import CodeDebugerAgent
+from .task_verifier import TaskVerifyAgent, TaskVerifyState
+from .task_summarier import TaskSummaryAgent
+from .task_verify_summarier import TaskVerifySummaryAgent
+from .task_structrue_summarier import TaskStructureSummaryAgent
+from .task_reasoner import TaskReasoningAgent
+from .task_structrue_reasoner import TaskStructureReasoningAgent
+
+__all__ = [
+    "AgentFactory",
+    "BaseChatAgent",
+    "CodeDebugerAgent",
+    "CodeExecutor",
+    "MasterPlannerAgent",
+    "TaskCodingAgent",
+    "TaskPlannerAgentV1",
+    "TaskPlannerAgentV2",
+    "TaskPlannerAgentV3",
+    "TaskReasoningAgent",
+    "TaskStructureReasoningAgent",
+    "TaskStructureSummaryAgent",
+    "TaskSummaryAgent",
+    "TaskVerifyAgent",
+    "TaskVerifyState",
+    "TaskVerifySummaryAgent",
+    "OutputTaskResult",
+]

jupyter_agent/bot_agents/base.py

@@ -0,0 +1,324 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import json
+import importlib
+
+from typing import Tuple, Any
+from enum import Enum, unique
+from IPython.display import Markdown
+from ..bot_outputs import _C, flush_output
+from ..bot_chat import BotChat
+from ..utils import no_indent
+
+_TASK_CONTEXTS = no_indent(
+    """
+**全局任务规划及子任务完成情况**：
+
+{% for cell in cells %}
+    {% if cell.type == "planning" and cell.source.strip() %}
+        {{ cell.source }}
+        {{ cell.result }}
+    {% elif cell.type == "task" and cell.subject.strip() %}
+        ## 子任务[{{ cell.task_id }}] - {{ '已完成' if cell.result else '未完成' }}
+
+        ### 任务目标
+
+        {{ cell.subject }}
+
+        ### 任务结果
+
+        {{ cell.result }}
+
+        {% if cell.important_infos %}
+            ### 任务结论中的重要信息(Important Infos)
+
+            ```json
+            {{ cell.important_infos | json }}
+            ```
+        {%+ endif %}
+    {% elif cell.is_task_context and cell.source.strip() %}
+        {{ cell.source }}
+    {% endif %}
+{% endfor %}
+"""
+)
+
+
+_CODE_CONTEXTS = no_indent(
+    """
+**已执行的代码**：
+
+```python
+{% for cell in cells +%}
+    {% if cell.type == "task" and cell.source.strip() %}
+        ######## Cell[{{ cell.cell_idx }}] for Task[{{ cell.task_id }}] ########
+
+        {{ cell.source }}
+    {% elif cell.is_code_context and cell.source.strip() %}
+        ######## Cell[{{ cell.cell_idx }}] ########
+
+        {{ cell.source }}
+    {% endif %}
+{%+ endfor %}
+```
+"""
+)
+
+_TASK_OUTPUT_FORMAT = """
+{% if OUTPUT_FORMAT == "code" %}
+**输出格式**：
+
+输出{{ OUTPUT_CODE_LANG }}代码块，以Markdown格式显示，使用```{{ OUTPUT_CODE_LANG }}...```包裹。
+
+示例代码：
+
+```python
+def xxx(xxx):
+    ...
+    return xxx
+
+xxx(...)
+```
+
+{% elif OUTPUT_FORMAT == "json" %}
+**输出格式**：
+
+输出JSON格式数据，以Markdown格式显示，使用```json...```包裹。
+
+数据符合JSON Schema：
+
+```json
+{{ OUTPUT_JSON_SCHEMA }}
+```
+
+数据示例:
+
+```json
+{{ OUTPUT_JSON_EXAMPLE }}
+```
+
+{% endif %}
+"""
+
+PREDEFINE_PROMPT_BLOCKS = {
+    "TASK_CONTEXTS": _TASK_CONTEXTS,
+    "CODE_CONTEXTS": _CODE_CONTEXTS,
+    "TASK_OUTPUT_FORMAT": _TASK_OUTPUT_FORMAT,
+}
+
+
+@unique
+class AgentOutputFormat(str, Enum):
+    RAW = "raw"
+    TEXT = "text"
+    CODE = "code"
+    JSON = "json"
+
+
+@unique
+class AgentCombineReply(str, Enum):
+    FIRST = "first"
+    LAST = "last"
+    LIST = "list"
+    MERGE = "merge"
+
+
+@unique
+class AgentModelType(str, Enum):
+    DEFAULT = "default"
+    PLANNER = "planner"
+    CODING = "coding"
+    REASONING = "reasoning"
+
+
+class BaseAgent:
+    """基础代理类"""
+
+    def __init__(self, notebook_context):
+        self.notebook_context = notebook_context
+
+    @property
+    def task(self):
+        return self.notebook_context.cur_task
+
+    @property
+    def cells(self):
+        return self.notebook_context.cells
+
+
+class BaseChatAgent(BotChat, BaseAgent):
+    """基础聊天代理类"""
+
+    PROMPT = "You are a helpful assistant. {{ prompt }}\n\nAnswer:"
+    OUTPUT_FORMAT = AgentOutputFormat.RAW
+    OUTPUT_CODE_LANG = "python"
+    OUTPUT_JSON_SCHEMA = None  # Pydantic Model
+    DISPLAY_REPLY = True
+    COMBINE_REPLY = AgentCombineReply.MERGE
+    ACCEPT_EMPYT_REPLY = False
+    MODEL_TYPE = AgentModelType.REASONING
+
+    def __init__(self, notebook_context, base_url, api_key, model_name, **chat_kwargs):
+        """初始化基础任务代理"""
+        BaseAgent.__init__(self, notebook_context)
+        BotChat.__init__(self, base_url, api_key, model_name, **chat_kwargs)
+
+    def prepare_contexts(self, **kwargs):
+        contexts = {
+            "cells": self.cells,
+            "task": self.task,
+            "OUTPUT_FORMAT": self.OUTPUT_FORMAT,
+            "OUTPUT_CODE_LANG": self.OUTPUT_CODE_LANG,
+        }
+        if self.OUTPUT_JSON_SCHEMA:
+            json_schema = self.OUTPUT_JSON_SCHEMA.model_json_schema()
+            model_fields = getattr(self.OUTPUT_JSON_SCHEMA, "model_fields", None)
+            if model_fields and hasattr(model_fields, "items"):
+                json_example = {
+                    name: field.examples[0] if getattr(field, "examples", None) else getattr(field, "default", None)
+                    for name, field in model_fields.items()
+                }
+            else:
+                json_example = {}
+            contexts["OUTPUT_JSON_SCHEMA"] = json.dumps(json_schema, indent=2, ensure_ascii=False)
+            contexts["OUTPUT_JSON_EXAMPLE"] = json.dumps(json_example, indent=2, ensure_ascii=False)
+        contexts.update(kwargs)
+        return contexts
+
+    def create_messages(self, contexts):
+        messages = super().create_messages(contexts, templates=PREDEFINE_PROMPT_BLOCKS)
+        messages.add(self.PROMPT)
+        return messages
+
+    def combine_raw_replies(self, replies):
+        if self.COMBINE_REPLY == AgentCombineReply.FIRST:
+            return replies[0]["raw"]
+        elif self.COMBINE_REPLY == AgentCombineReply.LAST:
+            return replies[-1]["raw"]
+        elif self.COMBINE_REPLY == AgentCombineReply.MERGE:
+            return "".join([reply["raw"] for reply in replies])
+        else:
+            raise ValueError("Unsupported combine_reply: {} for raw output".format(self.COMBINE_REPLY))
+
+    def combine_code_replies(self, replies):
+        code_replies = [
+            reply for reply in replies if reply["type"] == "code" and reply["lang"] == self.OUTPUT_CODE_LANG
+        ]
+        if self.COMBINE_REPLY == AgentCombineReply.FIRST:
+            return code_replies[0]["content"]
+        elif self.COMBINE_REPLY == AgentCombineReply.LAST:
+            return code_replies[-1]["content"]
+        elif self.COMBINE_REPLY == AgentCombineReply.MERGE:
+            return "\n".join([reply["content"] for reply in code_replies])
+        else:
+            raise ValueError("Unsupported combine_reply: {} for code output".format(self.COMBINE_REPLY))
+
+    def combine_json_replies(self, replies):
+        json_replies = [reply for reply in replies if reply["type"] == "code" and reply["lang"] == "json"]
+        if self.COMBINE_REPLY == AgentCombineReply.FIRST:
+            json_obj = json.loads(json_replies[0]["content"])
+            if self.OUTPUT_JSON_SCHEMA:
+                json_obj = self.OUTPUT_JSON_SCHEMA(**json_obj)
+            return json_obj
+        elif self.COMBINE_REPLY == AgentCombineReply.LAST:
+            json_obj = json.loads(json_replies[-1]["content"])
+            if self.OUTPUT_JSON_SCHEMA:
+                json_obj = self.OUTPUT_JSON_SCHEMA(**json_obj)
+            return json_obj
+        elif self.COMBINE_REPLY == AgentCombineReply.LIST:
+            json_objs = [json.loads(reply["content"]) for reply in json_replies]
+            if self.OUTPUT_JSON_SCHEMA:
+                json_objs = [self.OUTPUT_JSON_SCHEMA(**json_obj) for json_obj in json_objs]
+            return json_objs
+        elif self.COMBINE_REPLY == AgentCombineReply.MERGE:
+            json_obj = {}
+            for json_reply in json_replies:
+                json_obj.update(json.loads(json_reply["content"]))
+            if self.OUTPUT_JSON_SCHEMA:
+                json_obj = self.OUTPUT_JSON_SCHEMA(**json_obj)
+            return json_obj
+        else:
+            raise ValueError("Unsupported combine_reply: {} for json output".format(self.COMBINE_REPLY))
+
+    def combine_text_replies(self, replies):
+        text_replies = [reply for reply in replies if reply["type"] == "text"]
+        if self.COMBINE_REPLY == AgentCombineReply.FIRST:
+            return text_replies[0]["content"]
+        elif self.COMBINE_REPLY == AgentCombineReply.LAST:
+            return text_replies[-1]["content"]
+        elif self.COMBINE_REPLY == AgentCombineReply.MERGE:
+            return "".join([reply["content"] for reply in text_replies])
+        else:
+            raise ValueError("Unsupported combine_reply: {} for text output".format(self.COMBINE_REPLY))
+
+    def combine_replies(self, replies):
+        if self.OUTPUT_FORMAT == AgentOutputFormat.RAW:
+            return self.combine_raw_replies(replies).strip()
+        elif self.OUTPUT_FORMAT == AgentOutputFormat.TEXT:
+            return self.combine_text_replies(replies).strip()
+        elif self.OUTPUT_FORMAT == AgentOutputFormat.CODE:
+            return self.combine_code_replies(replies).strip()
+        elif self.OUTPUT_FORMAT == AgentOutputFormat.JSON:
+            return self.combine_json_replies(replies)
+        else:
+            raise ValueError("Unsupported output format: {}".format(self.OUTPUT_FORMAT))
+
+    def on_reply(self, reply) -> Tuple[bool, Any] | Any:
+        _C(Markdown(reply))
+
+    def __call__(self, **kwargs) -> Tuple[bool, Any]:
+        contexts = self.prepare_contexts(**kwargs)
+        messages = self.create_messages(contexts)
+        replies = self.chat(messages.get(), display_reply=self.DISPLAY_REPLY)
+        reply = self.combine_replies(replies)
+        if not self.ACCEPT_EMPYT_REPLY and not reply:
+            raise ValueError("Reply is empty")
+        result = self.on_reply(reply)
+        flush_output()
+        if not isinstance(result, tuple):
+            return False, result
+        else:
+            return result
+
+
+class AgentFactory:
+
+    def __init__(self, notebook_context, **chat_kwargs):
+        self.notebook_context = notebook_context
+        self.chat_kwargs = chat_kwargs
+        self.models = {AgentModelType.DEFAULT: {"api_url": None, "api_key": None, "model": None}}
+
+    def config_model(self, agent_model, api_url, api_key, model_name):
+        self.models[agent_model] = {
+            "api_url": api_url,
+            "api_key": api_key,
+            "model": model_name,
+        }
+
+    def __call__(self, agent_class):
+
+        if isinstance(agent_class, str):
+            bot_agents = importlib.import_module("..bot_agents", __package__)
+            agent_class = getattr(bot_agents, agent_class)
+
+        if issubclass(agent_class, BaseChatAgent):
+            agent_model = agent_class.MODEL_TYPE if hasattr(agent_class, "MODEL_TYPE") else AgentModelType.DEFAULT
+            return agent_class(
+                notebook_context=self.notebook_context,
+                base_url=self.models.get(agent_model, {}).get("api_url")
+                or self.models[AgentModelType.DEFAULT]["api_url"],
+                api_key=self.models.get(agent_model, {}).get("api_key")
+                or self.models[AgentModelType.DEFAULT]["api_key"],
+                model_name=self.models.get(agent_model, {}).get("model")
+                or self.models[AgentModelType.DEFAULT]["model"],
+                **self.chat_kwargs,
+            )
+        elif issubclass(agent_class, BaseAgent):
+            return agent_class(notebook_context=self.notebook_context)
+        else:
+            raise ValueError("Unsupported agent class: {}".format(agent_class))

jupyter_agent/bot_agents/master_planner.py

@@ -0,0 +1,45 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+from IPython.display import Markdown
+from .base import BaseChatAgent, AgentModelType
+from ..bot_outputs import _C, ReplyType
+
+MASTER_PLANNER_PROMPT = """\
+**角色定义**：
+
+你是一个高级分析规划专家，擅长将自然语言任务拆解为可执行的Jupyter Notebook分析流程。
+
+**任务要求**：
+
+- 解析用户输入的自然语言指令（目标Prompt），提取核心需求（如数据来源、分析目标、输出格式等）
+- 将核心需求转化为可执行的Jupyter Notebook分析流程
+- 流程的每个环节应尽量简单明确
+- 生成**目标规划说明**，包含：
+  1. 需要执行的子任务列表（按逻辑顺序排列）
+  3. 每个子任务的输入/输出依赖关系
+  4. 预期的最终输出形式（文字/图表/表格等）
+
+**用户需求**：
+
+{{ task.source }}
+
+---
+
+请根据上述信息为用户规划全局执行计划：
+"""
+
+
+class MasterPlannerAgent(BaseChatAgent):
+    """全局规划器代理类"""
+
+    PROMPT = MASTER_PLANNER_PROMPT
+    DISPLAY_REPLY = False
+    MODEL_TYPE = AgentModelType.PLANNER
+
+    def on_reply(self, reply):
+        _C(Markdown(reply), reply_type=ReplyType.TASK_RESULT)

jupyter_agent/bot_agents/output_task_result.py

@@ -0,0 +1,29 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import json
+
+from IPython.display import Markdown
+from .base import BaseAgent
+from ..bot_outputs import _D, _I, _W, _E, _F, _M, _B, _C, _O, ReplyType, markdown_block
+
+
+class OutputTaskResult(BaseAgent):
+
+    def __call__(self):
+        """执行代码逻辑"""
+        if self.task.result:
+            _O(Markdown("### 任务结果\n\n"))
+            _C(Markdown(self.task.result), reply_type=ReplyType.TASK_RESULT)
+        if self.task.important_infos:
+            _O(
+                markdown_block(
+                    f"```json\n{json.dumps(self.task.important_infos, indent=4, ensure_ascii=False)}\n```",
+                    title="重要信息",
+                )
+            )
+        return False, None

jupyter_agent/bot_agents/task_code_executor.py

@@ -0,0 +1,53 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import re
+
+from IPython.core.getipython import get_ipython
+from IPython.display import Markdown, clear_output
+from .base import BaseAgent
+from ..utils import TeeOutputCapture
+from ..bot_outputs import _D, _I, _W, _E, _F, _M, _B, _C, ReplyType
+
+
+class CodeExecutor(BaseAgent):
+
+    def __call__(self):
+        """执行代码逻辑"""
+        _D(f"执行代码: {repr(self.task.source)[:80]}")
+        ipython = get_ipython()
+        exec_failed = False
+        with TeeOutputCapture() as captured:
+            if ipython is None:
+                exec_failed = True
+                self.task.cell_error = "IPython environment not found."
+                _E("执行失败: IPython environment not found.")
+                result = None
+            else:
+                result = ipython.run_cell(self.task.source)
+                if result.success:
+                    self.task.cell_result = "{}".format(result.result)
+                    _D(f"执行结果: {repr(self.task.cell_result)[:80]}")
+                else:
+                    exec_failed = True
+                    exc_info = ipython._format_exception_for_storage(result.error_before_exec or result.error_in_exec)
+                    ansi_escape = re.compile(r"\x1b\[[0-9;]*m")
+                    clean_traceback = "\n".join(ansi_escape.sub("", line) for line in exc_info["traceback"])
+                    self.task.cell_error = clean_traceback
+                    _E(f"执行失败: {clean_traceback}")
+        self.task.cell_output = ""
+        if captured.stdout:
+            self.task.cell_output += "Stdout:\n" + captured.stdout + "\n"
+        if captured.stderr:
+            self.task.cell_output += "Stderr:\n" + captured.stderr + "\n"
+        if captured.outputs:
+            self.task.cell_output += "Outputs:\n"
+            for output in captured.outputs:
+                output_content = output.data.get("text/markdown", "") or output.data.get("text/plain", "")
+                self.task.cell_output += output_content
+                self.task.cell_output += "\n"
+        return exec_failed, not exec_failed

jupyter_agent/bot_agents/task_coder.py

@@ -0,0 +1,71 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import time
+
+from .base import BaseChatAgent, AgentOutputFormat, AgentModelType
+
+
+TASK_CODING_PROMPT = """\
+**角色定义**：
+
+你是一个代码生成器，能够将自然语言描述转化为可执行的Jupyter Notebook代码。
+
+**任务要求**：
+
+- 根据子任务Prompt生成**Python代码**（需包含必要的依赖导入）
+- 代码需严格遵守以下原则：
+  - 使用已存在的变量（如`cleaned_sales_df`）
+  - 新生成的变量需命名清晰（如`growth_rates_df`）
+  - 发生错误时，直接抛出异常，不需要进行任何错误处理，以便于执行器发现问题
+  - 代码在Jupyter环境中执行，但也应考虑通用性，尽量封装为函数或类
+  - 执行结果应保存在变量中，同时打印输出（或做为cell的返回值输出）
+  - 代码中应保含必要的文档注释
+  - 避免生成重复的代码
+
+{% include "TASK_OUTPUT_FORMAT" %}
+
+---
+
+{% include "TASK_CONTEXTS" %}
+
+---
+
+{% include "CODE_CONTEXTS" %}
+
+---
+
+**当前子任务信息**:
+
+### 当前子任务目标：
+{{ task.subject }}
+
+### 当前子任务代码需求：
+{{ task.coding_prompt }}
+
+{% if task.issue %}
+### 尽量避免的问题：  
+{{ task.issue }}
+{% endif %}
+
+---
+
+请按要求生成任务代码：
+"""
+
+
+class TaskCodingAgent(BaseChatAgent):
+
+    PROMPT = TASK_CODING_PROMPT
+    OUTPUT_FORMAT = AgentOutputFormat.CODE
+    OUTPUT_CODE_LANG = "python"
+    MODEL_TYPE = AgentModelType.CODING
+
+    def on_reply(self, reply: str):
+        generated_code = "# Generated by Jupyter Agent (Coder) {}\n".format(time.strftime("%Y-%m-%d %H:%M:%S"))
+        generated_code += reply
+        self.task.source = generated_code

jupyter_agent/bot_agents/task_debuger.py

@@ -0,0 +1,69 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import time
+
+from .base import BaseChatAgent, AgentOutputFormat, AgentModelType
+
+TASK_DEBUGGER_PROMPT = """\
+**角色定义**：
+
+你是一个代码调试专家，擅长定位并修复Jupyter Notebook中的错误。
+
+**任务要求**：
+
+- 分析执行失败的代码和错误信息，提出修复方案
+- 修复方案需包含错误原因分析（语法错误/逻辑错误/依赖缺失等）
+
+{% include "TASK_OUTPUT_FORMAT" %}
+
+---
+
+{% include "TASK_CONTEXTS" %}
+
+---
+
+{% include "CODE_CONTEXTS" %}
+
+---
+
+**当前子任务信息**:
+
+### 当前子任务目标：
+{{ task.subject }}
+
+### 当前子任务代码需求：
+{{ task.coding_prompt }}
+
+### 当前代码：
+```python
+{{ task.source }}
+```
+
+### 当前输出：
+{{ task.cell_output }}
+
+### 当前错误：
+{{ task.cell_error }}
+
+---
+
+请输出修复后的代码：
+"""
+
+
+class CodeDebugerAgent(BaseChatAgent):
+
+    PROMPT = TASK_DEBUGGER_PROMPT
+    OUTPUT_FORMAT = AgentOutputFormat.CODE
+    OUTPUT_CODE_LANG = "python"
+    MODEL_TYPE = AgentModelType.CODING
+
+    def on_reply(self, reply: str):
+        generated_code = "# Generated by Jupyter Agent (Debugger) {}\n".format(time.strftime("%Y-%m-%d %H:%M:%S"))
+        generated_code += reply
+        self.task.source = generated_code