ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter10/staged-system-prompt/tools.py

@@ -0,0 +1,329 @@
+"""
+工具实现与工具集定义。
+
+本文件包含两部分：
+1. Workspace：一个进程内的“虚拟工作区”，负责保存需求、文件内容，
+   并提供真实的代码执行 / 语法检查 / 复杂度分析能力。
+2. 三个阶段各自的工具 JSON Schema（供 OpenAI function calling 使用）。
+
+关键点：不同阶段暴露给模型的工具集是不同的，这是“阶段化系统提示词”实验
+的核心之一——提示词换了角色，工具也随之切换。
+"""
+
+from __future__ import annotations
+
+import ast
+import hashlib
+import os
+import subprocess
+import sys
+import tempfile
+from typing import Dict, List
+
+
+# ----------------------------------------------------------------------------
+# 触发阶段转换的“信号工具”名字。Agent 主循环看到这些工具被调用就切换阶段。
+# ----------------------------------------------------------------------------
+COMPLETE_REQUIREMENTS = "complete_requirements_analysis"  # 阶段1 -> 阶段2
+SUBMIT_FOR_REVIEW = "submit_for_review"                   # 阶段2 -> 阶段3
+REQUEST_REVISION = "request_revision"                     # 阶段3 -> 阶段2（回退）
+APPROVE_CODE = "approve_code"                             # 阶段3 -> 完成
+
+
+class Workspace:
+    """跨阶段共享的任务状态（需求、文件、审查意见）。"""
+
+    def __init__(self) -> None:
+        # 阶段1 收集到的、已确认的需求（key -> value）
+        self.requirements: Dict[str, str] = {}
+        # 阶段2 写出的“文件系统”（path -> content）
+        self.files: Dict[str, str] = {}
+        # 阶段3 退回时记录的问题清单，供阶段2 修复时参考
+        self.review_issues: List[str] = []
+
+    # --- 阶段1：需求分析师的工具实现 -------------------------------------
+    def save_requirement(self, key: str, value: str) -> str:
+        self.requirements[key] = value
+        return f"已记录需求 [{key}] = {value}"
+
+    # --- 阶段2：软件工程师的工具实现 -------------------------------------
+    def write_file(self, path: str, content: str) -> str:
+        self.files[path] = content
+        return f"已写入文件 {path}（{len(content)} 字符，{content.count(chr(10)) + 1} 行）"
+
+    def read_file(self, path: str) -> str:
+        if path not in self.files:
+            return f"错误：文件 {path} 不存在。当前文件列表：{list(self.files) or '空'}"
+        return self.files[path]
+
+    def execute_code(self, code: str) -> str:
+        """在含当前工作区文件的临时目录里真实执行代码（带超时）。"""
+        return _run_python_source(code, files=self.files)
+
+    # --- 阶段3：代码审查员的工具实现 -------------------------------------
+    def run_linter(self, path: str) -> str:
+        """轻量静态检查：语法编译 + 常见坏味道，不引入额外依赖。"""
+        if path not in self.files:
+            return f"错误：文件 {path} 不存在。"
+        source = self.files[path]
+        problems: List[str] = []
+
+        # 1) 语法能否编译
+        try:
+            tree = ast.parse(source)
+        except SyntaxError as exc:
+            return f"[linter] 语法错误：第 {exc.lineno} 行 {exc.msg}"
+
+        # 2) 逐行的风格问题（阈值定得“严格但可达标”，方便演示先退回再通过）
+        for i, line in enumerate(source.splitlines(), start=1):
+            if len(line) > 120:
+                problems.append(f"L{i}: 行超过 120 字符（{len(line)}），请折行或精简")
+            if line.rstrip() != line:
+                problems.append(f"L{i}: 行尾有多余空白")
+            if "\t" in line:
+                problems.append(f"L{i}: 使用了 Tab 缩进，建议用空格")
+
+        # 3) 基于 AST 的问题：缺少模块 docstring、裸 except
+        if not ast.get_docstring(tree):
+            problems.append("模块缺少文件级 docstring（请在文件开头加一段三引号说明）")
+        for node in ast.walk(tree):
+            if isinstance(node, ast.ExceptHandler) and node.type is None:
+                problems.append(f"L{node.lineno}: 使用了裸 except，建议捕获具体异常")
+
+        if not problems:
+            return "[linter] 通过：未发现问题。"
+        return "[linter] 发现 %d 个问题：\n- %s" % (len(problems), "\n- ".join(problems))
+
+    def run_tests(self, path: str) -> str:
+        """冒烟测试：把文件跑起来，验证 import / 主流程不崩溃。"""
+        if path not in self.files:
+            return f"错误：文件 {path} 不存在。"
+        # 造一个假的“下载文件夹”，让整理脚本有东西可整理
+        harness = (
+            "import os, tempfile, runpy, sys\n"
+            "d = tempfile.mkdtemp()\n"
+            "for name in ['a.jpg','b.pdf','c.txt','d.mp3','readme']:\n"
+            "    open(os.path.join(d, name), 'w').close()\n"
+            f"sys.argv = [{path!r}, d]\n"
+            "print('SMOKE_TEST target dir:', d)\n"
+            f"runpy.run_path({path!r}, run_name='__main__')\n"
+        )
+        # Keep the harness and submitted module as separate files.  Concatenating
+        # the harness before the module used to make a valid `from __future__`
+        # import fail solely because it was no longer at the start of its file.
+        result, returncode = _run_python(harness, files={path: self.files[path]})
+        # 以退出码为准：超时（returncode 为 None）和非零退出都不能算 PASS；
+        # Traceback/Error 子串检查作为额外防线保留。
+        ok = returncode == 0 and "Traceback" not in result and "Error" not in result
+        verdict = "PASS" if ok else "FAIL"
+        return f"[tests] 冒烟测试结果：{verdict}\n{result}"
+
+    def analyze_complexity(self, path: str) -> str:
+        """用 AST 估算复杂度：函数数量、最大分支数、最大嵌套深度。"""
+        if path not in self.files:
+            return f"错误：文件 {path} 不存在。"
+        try:
+            tree = ast.parse(self.files[path])
+        except SyntaxError as exc:
+            return f"[complexity] 无法解析：{exc.msg}"
+
+        funcs = [n for n in ast.walk(tree)
+                 if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))]
+        branch_types = (ast.If, ast.For, ast.While, ast.Try, ast.With)
+        total_branches = sum(1 for n in ast.walk(tree) if isinstance(n, branch_types))
+
+        def depth(node: ast.AST, level: int = 0) -> int:
+            best = level
+            for child in ast.iter_child_nodes(node):
+                inc = 1 if isinstance(child, branch_types) else 0
+                best = max(best, depth(child, level + inc))
+            return best
+
+        return (
+            "[complexity] 函数数量=%d，分支/循环语句=%d，最大嵌套深度=%d"
+            % (len(funcs), total_branches, depth(tree))
+        )
+
+    def inject_review_fault(self, path: str) -> dict:
+        """Insert one real, auditable lint defect for the rollback campaign.
+
+        The normal demo never calls this method.  The strict acceptance runner
+        uses it once, after the first implementation submission, so that the
+        reviewer must observe a genuine tool failure and the implementation
+        Agent must remove the defect before approval.  This is fault injection,
+        not a fabricated linter response: ``run_linter`` inspects the mutated
+        source normally.
+        """
+        if path not in self.files:
+            raise KeyError(f"cannot inject review fault: {path!r} is not in the workspace")
+        marker = "EXPERIMENT_10_1_REVIEW_CANARY"
+        before = self.files[path]
+        if marker in before:
+            raise ValueError(f"review fault marker already exists in {path!r}")
+        # A tab plus trailing spaces is valid Python inside a comment, but it
+        # deterministically violates two independent checks in run_linter.
+        after = before.rstrip("\n") + f"\n\t# {marker}  \n"
+        self.files[path] = after
+        return {
+            "path": path,
+            "kind": "tab_and_trailing_whitespace_comment",
+            "marker": marker,
+            "before_sha256": hashlib.sha256(before.encode()).hexdigest(),
+            "after_sha256": hashlib.sha256(after.encode()).hexdigest(),
+            "source_changed": before != after,
+        }
+
+
+def _run_python(source: str, timeout: int = 10, files: Dict[str, str] | None = None) -> tuple:
+    """执行源码，返回 (合并后的输出, 退出码)；超时时退出码为 None。"""
+    with tempfile.TemporaryDirectory() as tmp:
+        for name, content in (files or {}).items():
+            relative = os.path.normpath(name)
+            if os.path.isabs(relative) or relative == ".." or relative.startswith(".." + os.sep):
+                return f"工作区文件路径越界：{name}", 1
+            destination = os.path.join(tmp, relative)
+            os.makedirs(os.path.dirname(destination), exist_ok=True)
+            with open(destination, "w", encoding="utf-8") as fh:
+                fh.write(content)
+        script = os.path.join(tmp, "_workspace_runner.py")
+        with open(script, "w", encoding="utf-8") as fh:
+            fh.write(source)
+        try:
+            proc = subprocess.run(
+                [sys.executable, script],
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+                cwd=tmp,
+            )
+        except subprocess.TimeoutExpired:
+            return f"执行超时（>{timeout}s）", None
+        out = (proc.stdout or "").strip()
+        err = (proc.stderr or "").strip()
+        parts = []
+        if out:
+            parts.append("stdout:\n" + out)
+        if err:
+            parts.append("stderr:\n" + err)
+        parts.append(f"退出码: {proc.returncode}")
+        return "\n".join(parts), proc.returncode
+
+
+def _run_python_source(
+    source: str, timeout: int = 10, files: Dict[str, str] | None = None
+) -> str:
+    """把源码写到临时文件并用子进程执行，返回合并后的输出。"""
+    return _run_python(source, timeout, files=files)[0]
+
+
+# ----------------------------------------------------------------------------
+# 各阶段的工具 Schema（OpenAI tools 格式）。每个阶段只暴露自己那套工具。
+# ----------------------------------------------------------------------------
+
+def _tool(name: str, description: str, properties: dict, required: list) -> dict:
+    return {
+        "type": "function",
+        "function": {
+            "name": name,
+            "description": description,
+            "parameters": {
+                "type": "object",
+                "properties": properties,
+                "required": required,
+            },
+        },
+    }
+
+
+STAGE1_TOOLS = [
+    _tool(
+        "ask_clarifying_question",
+        "向用户提出一个澄清需求的问题，用户会回答。需求不明确时必须先问清楚。",
+        {"question": {"type": "string", "description": "要问用户的问题"}},
+        ["question"],
+    ),
+    _tool(
+        "save_requirement",
+        "把一条已经确认的需求记录到需求文档中，供后续实现阶段使用。",
+        {
+            "key": {"type": "string", "description": "需求项名称，如 file_types"},
+            "value": {"type": "string", "description": "需求项取值/描述"},
+        },
+        ["key", "value"],
+    ),
+    _tool(
+        COMPLETE_REQUIREMENTS,
+        "当所有关键需求都已澄清并记录后调用，结束需求分析阶段，进入代码实现阶段。",
+        {"summary": {"type": "string", "description": "对已确认需求的一句话总结"}},
+        ["summary"],
+    ),
+]
+
+STAGE2_TOOLS = [
+    _tool(
+        "write_file",
+        "写入（或覆盖）一个文件的完整内容。",
+        {
+            "path": {"type": "string", "description": "文件路径，如 organize_downloads.py"},
+            "content": {"type": "string", "description": "文件的完整内容"},
+        },
+        ["path", "content"],
+    ),
+    _tool(
+        "read_file",
+        "读取一个已写入文件的内容。",
+        {"path": {"type": "string", "description": "文件路径"}},
+        ["path"],
+    ),
+    _tool(
+        "execute_code",
+        "执行一段 Python 代码用于自测/验证，返回标准输出与错误。",
+        {"code": {"type": "string", "description": "要执行的 Python 代码"}},
+        ["code"],
+    ),
+    _tool(
+        SUBMIT_FOR_REVIEW,
+        "当代码实现完成且自测通过后调用，提交给代码审查员，进入审查阶段。",
+        {"file": {"type": "string", "description": "要提交审查的主文件路径"}},
+        ["file"],
+    ),
+]
+
+STAGE3_TOOLS = [
+    _tool(
+        "run_linter",
+        "对文件运行静态检查，返回代码风格/规范问题。",
+        {"file": {"type": "string", "description": "文件路径"}},
+        ["file"],
+    ),
+    _tool(
+        "run_tests",
+        "对文件运行冒烟测试，验证能否正常运行。",
+        {"file": {"type": "string", "description": "文件路径"}},
+        ["file"],
+    ),
+    _tool(
+        "analyze_complexity",
+        "分析文件的代码复杂度（函数数、分支数、嵌套深度）。",
+        {"file": {"type": "string", "description": "文件路径"}},
+        ["file"],
+    ),
+    _tool(
+        REQUEST_REVISION,
+        "当审查发现必须修复的问题时调用，把代码退回实现阶段并附上问题清单。",
+        {
+            "issues": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "需要修复的问题列表",
+            }
+        },
+        ["issues"],
+    ),
+    _tool(
+        APPROVE_CODE,
+        "当代码通过所有审查、质量达标时调用，批准代码，任务完成。",
+        {"comment": {"type": "string", "description": "审查通过的简短评语"}},
+        ["comment"],
+    ),
+]