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

jupyter_agent/bot_agents/task_structrue_summarier.py

@@ -0,0 +1,123 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import json
+
+from enum import Enum
+from typing import List, Optional, Dict, Any
+from pydantic import BaseModel, Field
+from IPython.display import Markdown
+from .base import BaseChatAgent, AgentOutputFormat
+from ..bot_outputs import ReplyType, _D, _I, _W, _E, _F, _M, _B, _C, _O, markdown_block
+from ..utils import RequestUserPrompt, format_user_prompts
+
+TASK_SUMMARY_PROMPT = """\
+**角色定义**：
+
+你是一个信息提炼专家，能够从分析结果中提取关键结论。
+
+**任务要求**：
+
+- 将代码执行的输出与结果转化为**人类可读的总结**
+- 包含以下内容：
+  1. 代码执行结果总结
+  2. 核心发现（如"Electronics类别月均增长12%"）
+  3. 数据支撑（引用关键数值或图表）
+  4. 其它建议（如新子任务Prompt等）
+- 在引用其它已完成的子任务的结果时，特别是其important_infos中的信息，要保证准确、清晰、完整，不要出现任何误导信息
+
+注：任务代码执行的结果不会记录在全局上下文中，只有任务总结的结果会记录在全局上下文中，
+因此任务总结中应包含对代码执行结果的简要说明，以便后续子任务使用。
+
+{% include "TASK_OUTPUT_FORMAT" %}
+
+---
+
+{% include "TASK_CONTEXTS" %}
+
+---
+
+{% include "CODE_CONTEXTS" %}
+
+---
+
+**当前子任务信息**:
+
+### 当前子任务目标：
+{{ task.subject }}
+
+### 当前子任务代码需求：
+{{ task.coding_prompt }}
+
+### 当前代码：
+```python
+{{ task.source }}
+```
+
+### 当前代码执行的输出与结果：
+{{ task.output }}
+
+### 当前任务总结要求：
+{{ task.summary_prompt }}
+
+---
+
+请按要求输出任务总结：
+"""
+
+
+class RequestInfo(BaseModel):
+    prompt: str = Field(description="需要用户补充更详细的信息的 Prompt", examples=["请补充与...相关的详细的信息"])
+    example: Optional[str] = Field(None, description="示例", examples=["..."])
+
+
+class TaskStructureSumaryOutput(BaseModel):
+
+    summary: str = Field(description=f"任务总结的详细描述", examples=["..."])
+    important_infos: Optional[Dict[str, Any]] = Field(
+        None,
+        description="任务总结中的重要信息，特别是需要后续子任务重点关注的信息。"
+        "注意：该字段仅支持结构化信息，不能使用代码、长文本等非结构化信息",
+        examples=[
+            {
+                "..._constraint": "...",
+                "..._expression": "...",
+                "..._patterns": ["...", "..."],
+                "..._execution_strategies": ["...", "..."],
+                "..._features": {"...": "...", "...": "..."},
+                "..._mapping_rules": {"...": "...", "...": "..."},
+                "...": "...",
+            }
+        ],
+    )
+    request_confirm_infos: Optional[List[RequestUserPrompt]] = Field(
+        None, description="需要用户补充确认的信息，问题应尽量简单，只需要用户回答是/否或在备选项中选择即可"
+    )
+
+
+class TaskStructureSummaryAgent(BaseChatAgent):
+
+    PROMPT = TASK_SUMMARY_PROMPT
+    OUTPUT_FORMAT = AgentOutputFormat.JSON
+    OUTPUT_JSON_SCHEMA = TaskStructureSumaryOutput
+    DISPLAY_REPLY = True
+
+    def on_reply(self, reply: TaskStructureSumaryOutput):
+        assert reply.summary, "Reply is empty"
+        _C(Markdown("### 任务总结\n\n" + reply.summary), reply_type=ReplyType.TASK_RESULT)
+        self.task.set_data("result", reply.summary)
+        if reply.important_infos:
+            self.task.set_data("important_infos", reply.important_infos)
+            _O(
+                markdown_block(
+                    f"```json\n{json.dumps(reply.important_infos, indent=4, ensure_ascii=False)}\n```",
+                    title="重要信息",
+                )
+            )
+        if reply.request_confirm_infos:
+            _O(Markdown(f"### 需要补充确认的信息\n"))
+            _O(Markdown(format_user_prompts(reply.request_confirm_infos, title="用户补充确认信息")))

jupyter_agent/bot_agents/task_summarier.py

@@ -0,0 +1,76 @@
+"""
+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, AgentOutputFormat
+from ..bot_outputs import _D, _I, _W, _E, _F, _M, _B, _C
+from ..bot_outputs import ReplyType
+
+TASK_SUMMARY_PROMPT = """\
+**角色定义**：
+
+你是一个信息提炼专家，能够从分析结果中提取关键结论。
+
+**任务要求**：
+
+- 将代码执行的输出与结果转化为**人类可读的总结**
+- 包含以下内容：
+  1. 代码执行结果总结
+  2. 核心发现（如"Electronics类别月均增长12%"）
+  3. 数据支撑（引用关键数值或图表）
+  4. 其它建议（如新子任务Prompt等）
+
+注：任务代码执行的结果不会记录在全局上下文中，只有任务总结的结果会记录在全局上下文中，
+因此任务总结中应包含对代码执行结果的简要说明，以便后续子任务使用。
+
+{% include "TASK_OUTPUT_FORMAT" %}
+
+---
+
+{% include "TASK_CONTEXTS" %}
+
+---
+
+{% include "CODE_CONTEXTS" %}
+
+---
+
+**当前子任务信息**:
+
+### 当前子任务目标：
+{{ task.subject }}
+
+### 当前子任务代码需求：
+{{ task.coding_prompt }}
+
+### 当前代码：
+```python
+{{ task.source }}
+```
+
+### 当前代码执行的输出与结果：
+{{ task.output }}
+
+### 当前任务总结要求：
+{{ task.summary_prompt }}
+
+---
+
+请按要求输出任务总结：
+"""
+
+
+class TaskSummaryAgent(BaseChatAgent):
+
+    PROMPT = TASK_SUMMARY_PROMPT
+    OUTPUT_FORMAT = AgentOutputFormat.TEXT
+    DISPLAY_REPLY = False
+
+    def on_reply(self, reply: str):
+        _C(Markdown("### 任务总结\n" + reply), reply_type=ReplyType.TASK_RESULT)
+        assert reply, "Reply is empty"
+        self.task.set_data("result", reply)

jupyter_agent/bot_agents/task_verifier.py

@@ -0,0 +1,99 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+from enum import Enum
+from typing import List
+from pydantic import BaseModel, Field
+from IPython.display import Markdown
+from .base import BaseChatAgent, AgentOutputFormat
+from ..bot_outputs import _D, _I, _W, _E, _F, _M, _B, _C, _O
+from ..bot_outputs import ReplyType
+
+
+TASK_VERIFY_PROMPT = """\
+**角色定义**：
+
+你是一个数据质量检查员，负责验证子任务的输出与结果的正确性。
+
+**任务要求**：
+
+- 对比子任务Prompt的预期输出和实际结果，验证以下维度：  
+  1. 数据完整性（如无缺失值、数据量合理）  
+  2. 逻辑一致性（如增长率计算正确）  
+- 输出验证结果和改进建议（如需要重新运行子任务则标记为失败）  
+
+{% include "TASK_OUTPUT_FORMAT" %}
+
+---
+
+{% include "TASK_CONTEXTS" %}
+
+---
+
+{% include "CODE_CONTEXTS" %}
+
+---
+
+**当前子任务信息**:
+
+### 当前子任务目标：
+{{ task.subject }}
+
+### 当前子任务代码需求：
+{{ task.coding_prompt }}
+
+### 当前代码：
+```python
+{{ task.source }}
+```
+
+### 当前输出：
+{{ task.output }}
+
+### 当前任务验证条件：
+{{ task.verify_prompt }}
+
+---
+
+请按要求输出验证结果：
+"""
+
+
+class TaskVerifyState(Enum):
+    FAILED = "failed"
+    PASSED = "passed"
+
+
+class TaskVerifyOutput(BaseModel):
+    state: TaskVerifyState = Field(description="任务验证结果", examples=[TaskVerifyState.PASSED.value])
+    issues: List[str] = Field(
+        [],
+        description="任务验证失败问题清单, 任务验证失败时必填, 任务验证通过时返回空列表",
+        examples=[["...未包含...字段...", "..字段值缺失...", "...字段值未在合理范围内...", "..."]],
+    )
+
+
+class TaskVerifyAgent(BaseChatAgent):
+
+    PROMPT = TASK_VERIFY_PROMPT
+    OUTPUT_FORMAT = AgentOutputFormat.JSON
+    OUTPUT_JSON_SCHEMA = TaskVerifyOutput
+
+    def on_reply(self, reply: TaskVerifyOutput):
+
+        if reply.state == TaskVerifyState.PASSED:
+            _M("### 任务验证通过！")
+            return False, reply.state
+        else:
+            _M("### 任务验证不通过！\n")
+            task_issue = ""
+            if reply.issues:
+                for issue in reply.issues:
+                    task_issue += "- {}\n".format(issue)
+            self.task.set_data("issue", task_issue)
+            _M(task_issue)
+            return True, reply.state

jupyter_agent/bot_agents/task_verify_summarier.py

@@ -0,0 +1,134 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+from enum import Enum
+from typing import List, Optional
+from pydantic import BaseModel, Field
+from IPython.display import Markdown
+from .base import BaseChatAgent, AgentOutputFormat
+from ..bot_outputs import _D, _I, _W, _E, _F, _M, _B, _C, _O
+from ..bot_outputs import ReplyType
+
+
+TASK_SUMMARY_PROMPT = """\
+**角色定义**：
+
+你是一个任务总结规划专家，能够从分析结果中提取关键结论，并对任务结果进行总结分析，给出优化建议。
+
+**任务要求**：
+
+- 对任务代码的执行结果进行进一步的推理分析总结，并输出**人类可读的总结**，包含以下内容：
+  1. 代码执行结果总结
+  1. 核心发现（如"Electronics类别月均增长12%"）
+  2. 数据支撑（引用关键数值或图表）
+  3. 其它建议（如新子任务Prompt等）
+- 若代码的执行结果不满足当前子任务的要求，则输出**人类可读的修改建议**，包含以下内容：
+  1. 当前结果不满足子任务目标的具体原因
+  2. 修改后的代码生成Prompt，包括：
+    - 需生成的代码类型（如数据处理、建模、可视化等）
+    - 具体输入（数据、变量、参数等）
+    - 预期输出形式（变量名、图表、文本等）
+    - 代码执行的结果仅在当前子任务中可见，不会记录在全局上下文中
+  3. 修改后的分析总结Prompt，包括：
+    - 说明本子任务结果总结的要点和输出要素，以便后续子任务使用
+    - 验证总结的结果会记录在全局上下文中
+
+注：任务代码执行的结果不会记录在全局上下文中，只有任务总结的结果会记录在全局上下文中，
+因此任务总结中应包含对代码执行结果的简要说明，以便后续子任务使用。
+
+{% include "TASK_OUTPUT_FORMAT" %}
+
+---
+
+{% include "TASK_CONTEXTS" %}
+
+---
+
+{% include "CODE_CONTEXTS" %}
+
+---
+
+**当前子任务信息**:
+
+### 当前子任务目标：
+{{ task.subject }}
+
+### 当前子任务代码需求：
+{{ task.coding_prompt }}
+
+### 当前代码：
+```python
+{{ task.source }}
+```
+
+### 当前输出：
+{{ task.output }}
+
+### 当前任务总结要求：
+{{ task.summary_prompt }}
+
+---
+
+请按要求输出验证结果：
+"""
+
+
+class TaskSummaryState(str, Enum):
+    SUCCESS = "success"
+    NOT_SATISFY = "not_satisfy"
+
+
+class TaskEnhancement(BaseModel):
+    issues: List[str] = Field([], description="当前子任务不满足要求的问题清单", examples=[["...", "..."]])
+    code_prompt: str = Field("", description="修改后的代码生成Prompt", examples=["..."])
+    summary_prompt: str = Field("", description="修改后的分析总结Prompt", examples=["..."])
+
+
+class TaskSummaryOutput(BaseModel):
+    state: TaskSummaryState = Field(description="是否完成总结", examples=[TaskSummaryState.SUCCESS.value])
+    summary: str = Field(
+        "", description=f'任务总结的详细描述，在 state="{TaskSummaryState.SUCCESS}" 时必填', examples=["..."]
+    )
+    enhancement: Optional[TaskEnhancement] = Field(
+        None,
+        description=f"任务不满足要求时的修改建议，在 state='{TaskSummaryState.NOT_SATISFY}' 时必填",
+        examples=[{"issues": ["...", "..."], "code_prompt": "...", "verify_prompt": "...", "summary_prompt": "..."}],
+    )
+
+
+class TaskVerifySummaryAgent(BaseChatAgent):
+
+    PROMPT = TASK_SUMMARY_PROMPT
+    OUTPUT_FORMAT = AgentOutputFormat.JSON
+    OUTPUT_JSON_SCHEMA = TaskSummaryOutput
+
+    def on_reply(self, reply: TaskSummaryOutput):
+
+        if reply.state == TaskSummaryState.SUCCESS:
+            _O(Markdown("### 任务总结"))
+            assert reply.summary, "Summary is empty"
+            _C(Markdown(reply.summary), reply_type=ReplyType.TASK_RESULT)
+            self.task.set_data("result", reply.summary)
+            return False, reply.state
+        else:
+            _M("### 任务验证不通过！\n")
+            assert reply.enhancement, "Enhancement is empty"
+            assert reply.enhancement.issues, "Issues is empty"
+            assert reply.enhancement.code_prompt, "Code prompt is empty"
+            assert reply.enhancement.summary_prompt, "Summary prompt is empty"
+            task_issue = ""
+            if reply.enhancement.issues:
+                for issue in reply.enhancement.issues:
+                    task_issue += "- {}\n".format(issue)
+            self.task.set_data("issue", task_issue)
+            self.task.set_data("coding_prompt", reply.enhancement.code_prompt)
+            self.task.set_data("summary_prompt", reply.enhancement.summary_prompt)
+            _M(task_issue)
+            _M("### 修改后的子任务信息\n")
+            _M(f"### 当前子任务代码需求：\n\n{reply.enhancement.code_prompt}")
+            _M(f"### 当前子任务总结要求：\n\n{reply.enhancement.summary_prompt}")
+            return True, reply.state

jupyter_agent/bot_chat.py

@@ -0,0 +1,218 @@
+"""
+Copyright (c) 2025 viewstar000
+
+This software is released under the MIT License.
+https://opensource.org/licenses/MIT
+"""
+
+import re
+import json
+import jinja2
+import openai
+
+
+from .bot_outputs import _D, _I, _W, _E, _F, _B, _M
+
+
+class ChatMessages:
+    def __init__(self, contexts=None, templates=None, display_message=True):
+        self.messages = []
+        self.contexts = contexts
+        self.templates = templates
+        self.display_message = display_message
+        if self.templates is not None:
+            self.jinja_env = jinja2.Environment(
+                loader=jinja2.DictLoader(self.templates), trim_blocks=True, lstrip_blocks=True
+            )
+        else:
+            self.jinja_env = jinja2.Environment(trim_blocks=True, lstrip_blocks=True)
+        self.jinja_env.filters["json"] = lambda x: json.dumps(x, indent=2, ensure_ascii=False)
+
+    def add(self, content, role="user", content_type="text", tpl_context=None):
+        tpl_context = tpl_context or self.contexts
+        if content_type == "text" and tpl_context is not None:
+            content = self.jinja_env.from_string(content).render(**tpl_context)
+        if content_type == "text":
+            content_key = "text"
+        else:
+            raise NotImplementedError
+        _D("Adding message: role={}, content_type={}".format(role, content_type))
+        if self.display_message:
+            _B(content, title="Chat Message {}: {}".format(role, content_type))
+        if len(self.messages) == 0 or self.messages[-1]["role"] != role:
+            self.messages.append({"role": role, "content": [{"type": content_type, content_key: content}]})
+        else:
+            self.messages[-1]["content"].append({"type": content_type, content_key: content})
+
+    def get(self):
+        return self.messages
+
+    def clear(self):
+        self.messages = []
+
+
+class BotChat:
+    """聊天混合类，提供聊天相关功能"""
+
+    display_think = True
+    display_message = True
+    display_response = False
+
+    def __init__(self, base_url, api_key, model_name, **chat_kwargs):
+        """初始化聊天混合类"""
+        self.base_url = base_url
+        self.api_key = api_key
+        self.model_name = model_name
+        self.dispaly_think = chat_kwargs.get("dispaly_think", self.display_think)
+        self.display_message = chat_kwargs.get("display_message", self.display_message)
+        self.display_response = chat_kwargs.get("display_response", self.display_response)
+
+    def parse_reply(self, reply, ret_think_block=False, ret_empty_block=False, display_reply=True):
+        """解析聊天回复"""
+
+        def _read_think_block(tokens):
+            text = ""
+            while True:
+                try:
+                    token = next(iter_tokens)
+                except StopIteration:
+                    break
+                if token is None:
+                    continue
+                if token == "</think>":
+                    break
+                elif token == "<think>":
+                    text += _read_think_block(tokens)
+                    text += "</think>"
+                # elif token.startswith("```") and len(token) > 3:
+                #     text += _read_code_block(tokens)
+                # elif token.startswith("```") and len(token) == 3:
+                #     text += _read_fence_block(tokens)
+                else:
+                    text += token
+            return text
+
+        def _read_code_block(tokens):
+            text = ""
+            while True:
+                try:
+                    token = next(iter_tokens)
+                except StopIteration:
+                    break
+                if token is None:
+                    continue
+                if token == "```":
+                    break
+                # elif token == "<think>":
+                #     text += _read_think_block(tokens)
+                elif token.startswith("```") and len(token) > 3:
+                    text += _read_code_block(tokens)
+                    text += "```"
+                else:
+                    text += token
+            return text
+
+        def _read_fence_block(tokens):
+            text = ""
+            while True:
+                try:
+                    token = next(iter_tokens)
+                except StopIteration:
+                    break
+                if token is None:
+                    continue
+                if token == "```":
+                    break
+                # elif token == "<think>":
+                #     text += _read_think_block(tokens)
+                elif token.startswith("```") and len(token) > 3:
+                    text += _read_code_block(tokens)
+                    text += "```"
+                else:
+                    text += token
+            return text
+
+        tokens = re.split(r"(<think>)|(</think>)|(```[a-zA-Z_0-9]+)|(```)", reply)
+        iter_tokens = iter(tokens)
+        while True:
+            try:
+                token = next(iter_tokens)
+            except StopIteration:
+                break
+            if token:
+                if token == "<think>":
+                    think_block = _read_think_block(iter_tokens)
+                    raw_think_block = token + think_block + "</think>"
+                    if (self.dispaly_think or display_reply) and think_block and think_block.strip():
+                        _B(think_block, title="Thought Block")
+                    if ret_think_block and (ret_empty_block or think_block and think_block.strip()):
+                        yield {"type": "think", "content": think_block, "raw": raw_think_block}
+                elif token.startswith("```") and len(token) > 3:
+                    content = _read_code_block(iter_tokens)
+                    raw_content = token + content + "```"
+                    lang = token[3:].lower()
+                    if display_reply and content and content.strip():
+                        _B(content, title="Code Block", format="code", code_language=lang)
+                    if ret_empty_block or content and content.strip():
+                        yield {"type": "code", "lang": lang, "content": content, "raw": raw_content}
+                elif token.startswith("```") and len(token) == 3:
+                    content = _read_fence_block(iter_tokens)
+                    raw_content = token + content + "```"
+                    if display_reply and content and content.strip():
+                        _B(content, title="Fence Block", format="code", code_language="text")
+                    if ret_empty_block or content and content.strip():
+                        yield {"type": "fence", "content": content, "raw": raw_content}
+                else:
+                    if display_reply and token and token.strip():
+                        _M(token)
+                    if ret_empty_block or token and token.strip():
+                        yield {"type": "text", "content": token, "raw": token}
+
+    def create_messages(self, contexts=None, templates=None):
+        return ChatMessages(contexts=contexts, templates=templates, display_message=self.display_message)
+
+    def chat(
+        self,
+        messages,
+        ret_think_block=False,
+        ret_empty_block=False,
+        display_reply=True,
+        max_tokens=32 * 1024,
+        max_completion_tokens=4 * 1024,
+        temperature=0.8,
+        n=1,
+        **kwargs,
+    ):
+        """发送聊天请求"""
+        sizes = [len(content["text"]) for content in messages[0]["content"]]
+        total_size = sum(sizes)
+        _D("Total message size: {} chars, {}".format(total_size, sizes))
+        _I("Connecting to OpenAI API: {}".format(self.base_url or "default"))
+        openai_client = openai.OpenAI(api_key=self.api_key, base_url=self.base_url)
+        _I("Sending request to OpenAI API, model: {}".format(self.model_name))
+        response = openai_client.chat.completions.create(
+            model=self.model_name,
+            messages=messages,
+            max_tokens=max_tokens,
+            max_completion_tokens=max_completion_tokens,
+            temperature=temperature,
+            n=n,
+            **kwargs,
+        )
+        if not response.choices or not response.choices[0].message:
+            _E("No valid response from OpenAI API")
+            return []
+        else:
+            _I("Received response from OpenAI API")
+            _D("Response content: " + repr(response.choices[0].message.content)[:50])
+            if self.display_response:
+                _B(response.choices[0].message.content, title="Chat Response")
+            reply = response.choices[0].message.content
+            return list(
+                self.parse_reply(
+                    reply,
+                    ret_think_block=ret_think_block,
+                    ret_empty_block=ret_empty_block,
+                    display_reply=display_reply,
+                )
+            )