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

chapter4/active-tool-discovery/agent.py

@@ -0,0 +1,281 @@
+"""
+三种工具发现策略的 Agent 循环（文本/ReAct 协议）。
+
+为什么用"文本注入 + 文本解析工具调用"而不是 OpenAI 原生 function calling？
+—— 本实验要复现的正是书中所述：把 120+ 工具 schema **一次性注入 system prompt（几万 token）**，
+   模型在超长上下文下"指令遵循退化"。OpenAI 原生 function-calling 接口对工具选择做了很强的
+   约束/优化，即使上百个工具也很少选错，无法体现该退化；而把 schema 当作纯文本塞进 prompt、
+   让模型自己以 JSON 形式输出工具调用，才是书中控制组的真实机制，也才能观察到退化。
+
+协议：模型每一步只输出一个 JSON：
+    {"thought": "...", "tool": "工具名", "arguments": {...}}
+任务完成时输出：
+    {"thought": "...", "tool": "finish", "arguments": {"answer": "..."}}
+
+1) run_full_injection —— 对照组（全量注入）
+   system prompt 里以文本列出全部 126 个工具。injected_tokens = 该工具清单文本的 token 数。
+
+2) run_retrieval_prefilter —— 对照组之二（检索预筛选）
+   按用户初始查询做**一次性**语义检索，只把 top-n 个候选工具注入 system prompt。
+   token 已大幅下降，但一次性匹配无法预见执行中才浮现的跨领域需求（书中所述局限）。
+
+3) run_active_discovery —— 实验组（主动发现）
+   system prompt 只列出少量基础工具 + discover_tools 元工具。
+   模型调用 discover_tools(need) 时，用嵌入相似度返回 3-5 个候选工具，其文本清单作为
+   **user message** 追加进对话（保护 system 前缀 KV Cache），并更新状态栏可用工具列表。
+   injected_tokens = 基础工具 + discover_tools + 实际发现加载的工具清单的 token 数。
+"""
+
+import json
+import re
+from typing import Dict, List
+
+import tiktoken
+
+from discovery import ToolIndex  # noqa: F401  (类型提示用)
+from tools_library import (ALL_TOOLS, BASE_TOOL_NAMES, TOOL_IMPLS,
+                           TOOLS_BY_NAME)
+
+try:
+    _ENC = tiktoken.get_encoding("o200k_base")  # gpt-4o 系列编码
+except Exception:
+    _ENC = tiktoken.get_encoding("cl100k_base")
+
+
+# ---------------------------------------------------------------------------
+# 工具清单文本渲染 & token 统计
+# ---------------------------------------------------------------------------
+
+def render_tool(tool: Dict) -> str:
+    """把单个工具渲染成完整 JSON schema 文本（与真实注入到 prompt 的形式一致）。"""
+    return json.dumps(tool["function"], ensure_ascii=False, indent=2)
+
+
+def render_tools(tools: List[Dict]) -> str:
+    return "\n".join(render_tool(t) for t in tools)
+
+
+def count_tokens(text: str) -> int:
+    return len(_ENC.encode(text)) if text else 0
+
+
+# discover_tools 元工具（也用文本形式呈现给模型）
+DISCOVER_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "discover_tools",
+        "description": ("发现新工具：当缺少合适的专用工具时调用它，用一句自然语言描述你需要的"
+                        "『能力』(need)，系统会用语义检索返回最匹配的若干专用工具及其定义，之后即可调用它们。"),
+        "parameters": {"type": "object",
+                       "properties": {"need": {"type": "string"}}, "required": ["need"]},
+    },
+}
+
+FINISH_TOOL_DESC = "- finish(answer: string): 所有子任务都完成后调用，给出最终回答。"
+
+
+_PROTOCOL = (
+    "你每一步都必须、且只能输出一个 JSON 对象，不要输出任何多余文字，格式为：\n"
+    '{"thought": "简要思考", "tool": "工具名", "arguments": {参数键值}}\n'
+    "系统会执行该工具并把结果返回给你，然后你再输出下一步。\n"
+    "当且仅当任务的所有子任务都已用合适的工具完成后，输出："
+    '{"thought": "...", "tool": "finish", "arguments": {"answer": "最终回答"}}\n'
+    "注意：请为每个子任务选择最匹配的『专用工具』，而不是笼统的通用搜索工具。"
+)
+
+
+def _extract_json(text: str):
+    """从模型回复里抽取第一个 JSON 对象。"""
+    text = text.strip()
+    text = re.sub(r"^```(?:json)?|```$", "", text, flags=re.MULTILINE).strip()
+    # 找到第一个 { 到匹配的 }
+    start = text.find("{")
+    if start == -1:
+        return None
+    depth = 0
+    for i in range(start, len(text)):
+        if text[i] == "{":
+            depth += 1
+        elif text[i] == "}":
+            depth -= 1
+            if depth == 0:
+                try:
+                    return json.loads(text[start:i + 1])
+                except json.JSONDecodeError:
+                    return None
+    return None
+
+
+def _run_loop(client, model, system_prompt, task_prompt, available_names,
+              on_discover=None, max_steps=10):
+    """
+    文本 ReAct 循环。
+    available_names: set，当前允许调用的工具名（不含 discover_tools/finish）。
+      —— 主动发现模式下会随 discover_tools 动态增长。
+    返回 (called_tools, trace, finished)。
+    """
+    messages = [{"role": "system", "content": system_prompt},
+                {"role": "user", "content": task_prompt}]
+    called: List[str] = []
+    trace: List[str] = []
+    finished = False
+
+    for _ in range(max_steps):
+        try:
+            resp = client.chat.completions.create(
+                model=model, messages=messages, temperature=0)
+        except Exception as e:
+            # 部分推理型模型（如 gpt-5.x）只支持默认 temperature=1，此时退回默认值重试。
+            if "temperature" in str(e):
+                resp = client.chat.completions.create(
+                    model=model, messages=messages)
+            else:
+                raise
+        content = resp.choices[0].message.content or ""
+        messages.append({"role": "assistant", "content": content})
+
+        action = _extract_json(content)
+        if action is None or "tool" not in action:
+            trace.append(f"[格式错误] 模型未输出合法 JSON: {content[:80]!r}")
+            messages.append({"role": "user",
+                             "content": "你的回复不是合法的 JSON，请只输出规定格式的 JSON 对象。"})
+            continue
+
+        name = action.get("tool")
+        args = action.get("arguments") or {}
+
+        if name == "finish":
+            trace.append(f"[finish] {str(args.get('answer',''))[:100]}")
+            finished = True
+            break
+
+        if name == "discover_tools" and on_discover is not None:
+            need = args.get("need", "")
+            result_text, new_names = on_discover(need)
+            called.append(name)
+            trace.append(f"[discover_tools] need='{need}' -> {new_names}")
+            available_names.update(new_names)
+            messages.append({"role": "user", "content": result_text})
+            continue
+
+        # 普通工具调用
+        if name not in available_names:
+            # 该工具当前不可用（主动发现里还没发现 / 预筛选没选中 / 或纯属幻觉）——
+            # 不计入 called（未真正执行），判分因此能体现该子任务失败。
+            trace.append(f"[不可用] {name}")
+            hint = ("该工具当前不可用。"
+                    + ("请先用 discover_tools 发现所需能力的工具。" if on_discover else
+                       "请从工具清单中选择一个存在的工具。"))
+            messages.append({"role": "user", "content": hint})
+            continue
+
+        called.append(name)
+        impl = TOOL_IMPLS.get(name)
+        result = impl(args) if impl else json.dumps({"error": f"unknown tool {name}"})
+        trace.append(f"[call] {name}({json.dumps(args, ensure_ascii=False)})")
+        messages.append({"role": "user", "content": f"工具 {name} 返回：{result}"})
+
+    return called, trace, finished
+
+
+# ---------------------------------------------------------------------------
+# 对照组：全量注入
+# ---------------------------------------------------------------------------
+
+def run_full_injection(client, model, task_prompt: str, tools: List[Dict] = None,
+                       max_steps: int = 10) -> Dict:
+    tools = tools if tools is not None else ALL_TOOLS
+    tools_text = render_tools(tools) + "\n" + FINISH_TOOL_DESC
+    injected = count_tokens(tools_text)
+    system = (
+        f"你是一个智能助手。下面是你可以使用的全部工具清单（共 {len(tools)} 个），"
+        "请根据任务选择最合适的工具来完成。若任务包含多个子任务，请确保每个子任务都被处理。\n\n"
+        "【工具清单】\n" + tools_text + "\n\n" + _PROTOCOL
+    )
+    available = {t["function"]["name"] for t in tools}
+    called, trace, finished = _run_loop(client, model, system, task_prompt, available,
+                                        max_steps=max_steps)
+    return {"mode": "full_injection", "injected_tokens": injected,
+            "num_tools_exposed": len(tools), "called": called,
+            "trace": trace, "finished": finished}
+
+
+# ---------------------------------------------------------------------------
+# 对照组之二：检索预筛选（书中"检索式预筛选"）
+#   —— 按用户初始查询做**一次性**语义检索，只把 top-n 个候选工具注入 system prompt。
+#      它介于"全量注入"与"主动发现"之间：token 已大幅下降，但只匹配一次，无法预见
+#      任务执行中才浮现的跨领域需求（书中所述的内在局限）——若第二个子任务所需的
+#      专用工具没被这一次检索选中，模型就无从调用它，导致该子任务失败。
+# ---------------------------------------------------------------------------
+
+def run_retrieval_prefilter(client, model, task_prompt: str, index, top_n: int = 10,
+                            tools: List[Dict] = None, max_steps: int = 10) -> Dict:
+    tools = tools if tools is not None else ALL_TOOLS
+    tbn = {t["function"]["name"]: t for t in tools}
+    hits = index.search(task_prompt, top_k=top_n)
+    picked = [name for name, _ in hits if name in tbn]
+    picked_tools = [tbn[n] for n in picked]
+    tools_text = render_tools(picked_tools) + "\n" + FINISH_TOOL_DESC
+    injected = count_tokens(tools_text)
+    system = (
+        f"你是一个智能助手。系统已根据你的任务预先检索出下列可能相关的工具（共 {len(picked_tools)} 个），"
+        "请从中选择合适的工具完成任务。若某个子任务在清单中找不到合适的工具，请如实说明。\n\n"
+        "【工具清单】\n" + tools_text + "\n\n" + _PROTOCOL
+    )
+    available = set(picked)
+    called, trace, finished = _run_loop(client, model, system, task_prompt, available,
+                                        max_steps=max_steps)
+    return {"mode": "retrieval_prefilter", "injected_tokens": injected,
+            "num_tools_exposed": len(picked_tools), "prefiltered": picked,
+            "called": called, "trace": trace, "finished": finished}
+
+
+# ---------------------------------------------------------------------------
+# 实验组：主动发现
+# ---------------------------------------------------------------------------
+
+def run_active_discovery(client, model, task_prompt: str, index, top_k=4,
+                         tools: List[Dict] = None, max_steps: int = 10) -> Dict:
+    tools = tools if tools is not None else ALL_TOOLS
+    tbn = {t["function"]["name"]: t for t in tools}
+    base_tools = [tbn[n] for n in BASE_TOOL_NAMES]
+    base_text = (render_tools(base_tools) + "\n"
+                 + render_tool(DISCOVER_TOOL) + "\n" + FINISH_TOOL_DESC)
+
+    discovered_names = set()          # 本轮实际发现加载的专用工具
+    discovered_texts: List[str] = []  # 对应的文本清单（用于统计按需注入 token）
+    available = set(BASE_TOOL_NAMES)
+
+    def on_discover(need: str):
+        hits = index.search(need, top_k=top_k)
+        names, lines = [], []
+        for name, score in hits:
+            if name in BASE_TOOL_NAMES:
+                continue
+            names.append(name)
+            lines.append(render_tool(tbn[name]) + f"   (相似度 {score:.3f})")
+            if name not in discovered_names:
+                discovered_names.add(name)
+                discovered_texts.append(render_tool(tbn[name]))
+        status = f"\n\n【状态栏｜当前可用工具】{sorted(available | set(names))}"
+        body = ("discover_tools 匹配到以下专用工具，已加载，可直接调用：\n"
+                + "\n".join(lines) + status)
+        return body, names
+
+    system = (
+        "你是一个智能助手。你当前只掌握少量基础工具（见下）。"
+        "当任务需要你没有的能力时，先调用 discover_tools，用自然语言描述你需要的能力，"
+        "系统会返回并加载匹配的专用工具，然后你再调用它们。"
+        "若任务包含多个子任务（如既要查询又要下载），请针对每一项能力分别调用 discover_tools，"
+        "并在结束前确认每个子任务都已用合适的工具完成。\n\n"
+        "【基础工具】\n" + base_text + "\n\n" + _PROTOCOL
+    )
+    called, trace, finished = _run_loop(client, model, system, task_prompt,
+                                        available, on_discover=on_discover,
+                                        max_steps=max_steps)
+
+    injected = count_tokens(base_text) + count_tokens("\n".join(discovered_texts))
+    return {"mode": "active_discovery", "injected_tokens": injected,
+            "num_tools_exposed": len(BASE_TOOL_NAMES) + 1 + len(discovered_names),
+            "discovered": sorted(discovered_names),
+            "called": called, "trace": trace, "finished": finished}