跳转至

active-tool-discovery

第4章 · 工具 · 配套项目 chapter4/active-tool-discovery

项目说明

实验 4-6：主动工具发现（Active Tool Discovery）

《深入理解 AI Agent》配套代码 · ★★★

在 126 个跨领域工具的工具库上，一次运行对比三种"工具发现"策略——全量注入、 检索预筛选（一次性语义预选 top-n）、主动发现（执行中按需 discover_tools）—— 并输出统一的对比表（准确率 / 注入 token / 延迟）。量化全量注入的 token 浪费，展示主动发现 如何用嵌入向量相似度把上百个工具收敛成几条精准候选，同时揭示"一次性预筛选"在多步跨领域 任务上的内在局限。无 API key 也能跑：python demo.py --offline 用本地嵌入 + mock 模型自检机制。

目的

当一个 Agent 拥有上百个工具时，常见做法是把全部工具的 JSON schema 一次性塞进 system prompt。 这会带来两个问题：

  1. token 浪费：126 个工具的完整 schema 约 1.16 万 token，每一步推理都要重复计费。
  2. 指令遵循退化：措辞稍泛的任务下，模型会"广撒网"地把通用兜底工具（web_search / google_search / universal_search）和专用工具一起调用，甚至用通用搜索替代专用工具 —— 即书中所说的"查股价却选了通用 web_search"。

主动发现只在 system 里保留少量基础工具 + 一个 discover_tools(need) 元工具。模型遇到能力缺口时， 用自然语言描述需求，系统用嵌入相似度从工具库检索 3-5 个最相关的专用工具，把它们的 schema 作为 user message 追加进对话（保护 system 前缀的 KV Cache），并更新状态栏可用工具列表。

机制

tools_library.py   126 个跨领域工具（finance/web/arxiv/github/geo/weather/media/... 共 17 个领域）
                   每个工具有真实 name/description/parameters；执行为轻量 mock（重点是"选对工具"）
                   其中故意混入 8 个"通用/近义"工具（web_search 等），它们的描述夸大自己无所不能
                   select_tools(size)：按 --tool-set-size 截取子集，演示"工具集越大全量注入越吃亏"
discovery.py       可插拔嵌入后端 + 工具向量索引；OpenAIEmbedder 用 text-embedding-3-small 生成向量
                   并缓存到 .cache/；search(need) = 把 need 向量化后与工具向量做余弦相似度返回 top-k
agent.py           三种策略的 ReAct 循环（文本协议：模型每步输出一个 JSON 工具调用）
                   - run_full_injection：126 个工具 schema 全部写进 system prompt
                   - run_retrieval_prefilter：按初始查询一次性检索 top-n 工具注入（书中"检索式预筛选"）
                   - run_active_discovery：基础工具 + discover_tools，执行中按需检索加载
offline_backend.py 离线后端：LocalEmbedder（本地哈希词袋嵌入）+ MockChatClient（脚本化 mock 模型），
                   让 --offline 无需任何 API key 即可跑通全流程（token/延迟真实，准确率仅反映启发式路由）
demo.py            对同一组任务分别跑所选策略，打印 token / 延迟 / 调用轨迹 / 是否精确选对，并汇总对比表

为什么用"文本注入 + 文本解析"而不是 OpenAI 原生 function calling？ 原生 function-calling 接口对工具选择做了很强的约束优化，即使上百个工具也极少选错，无法体现书中 所述的"超长上下文指令遵循退化"。把 schema 当作纯文本塞进 prompt、让模型自己以 JSON 输出工具调用， 才是控制组的真实机制，也才能观察到退化。这也正是书中"把 schema 注入 system prompt（几万 token）"的写法。

为什么嵌入检索能避免错选？ 通用工具 web_search 的描述"什么都能做"，语义被稀释；而专用工具 （如 search_news）描述聚焦。对一个聚焦的 need（"获取特斯拉最近的新闻"），聚焦的专用工具余弦相似度 更高、排在前面，通用工具往往进不了 top-k，于是根本不会被加载 —— 检索层天然起到了"精度过滤"作用。

检索预筛选为什么不够？ 检索预筛选（run_retrieval_prefilter）只按初始查询做一次语义匹配、 一次性注入 top-n 工具。对"查股价 + 搜新闻"这类多步跨领域任务，初始查询的向量往往偏向第一个领域， 第二个子任务需要的专用工具可能挤不进 top-n，模型执行到一半才发现"想调用的工具根本没在清单里"—— 这正是书中指出的一次性匹配的内在局限。主动发现把"发现"延后到执行中、按每个真实浮现的 need 分别检索， 从而补齐这一缺口（离线自检里可直接观察到：检索预筛选在半数多步任务上漏掉了第二个工具，见下表）。

运行

pip install -r requirements.txt

# 方式 A：离线机制自检（无需任何 key；token/延迟真实，准确率仅反映启发式路由）
python demo.py --offline

# 方式 B：真实模型（体现小模型"指令遵循退化"需要真实 LLM）
cp env.example .env    # 填入 OPENAI_API_KEY（chat 与 embeddings 都用 OpenAI）
# 兜底：若无 OPENAI_API_KEY 但设置了 OPENROUTER_API_KEY，chat 会自动改走 OpenRouter
#（模型映射到 openai/gpt-5.6-luna 等），工具检索退用本地哈希嵌入（OpenRouter 无 embeddings 接口）。
python demo.py                                   # 全部 8 任务 × 三种策略
python demo.py --strategies full,discovery       # 只跑其中两种策略对比
python demo.py --tasks finance+news,crypto+news  # 只跑指定任务（逗号分隔）
python demo.py --tasks 'opinion(诱导)'            # 含括号的任务 id 记得加引号
python demo.py --tool-set-size 20                # 缩小工具集，看全量注入的劣势如何随规模放大
python demo.py --query '查英伟达股价再搜点相关新闻' --offline   # 临时单条自然语言任务
python demo.py --offline --output results/offline.json         # 导出结构化结果

默认模型 gpt-5.6-luna，可用 --model 或 env 覆盖：python demo.py --model gpt-5.6-luna。 首次运行会为工具生成嵌入向量并缓存到 .cache/，之后复用。python demo.py --help 查看全部参数 （--query / --tasks / --strategies / --tool-set-size / --top-k / --prefilter-n / --model / --embed-model / --max-steps / --offline / --output）。

如何适配 / 扩展

  • 换模型：MODEL=gpt-4.1-mini python demo.py（chat 模型）；EMBED_MODEL=text-embedding-3-large 换嵌入模型 （换嵌入模型会因签名变化自动重建 .cache/ 索引）。
  • 换供应商 / 网关：chat 与 embeddings 都走 OpenAI SDK，OpenAI() 会自动读取环境变量 OPENAI_BASE_URL， 因此指向任意 OpenAI 兼容的网关/代理只需 OPENAI_BASE_URL=https://your-gateway/v1（该端点需同时提供 chat 与 embeddings）。
  • 换任务 / 输入：编辑 tools_library.py 里的 TASKS（每条含 prompt 与判分用的能力槽位），或用 --tasks 只跑其中几条，或用 --query 传一句临时需求；想扩充工具库同样在 tools_library.py 的 ALL_TOOLS 中增删。
  • 离线自检：--offline 用 offline_backend.py 的本地哈希嵌入 + 脚本化 mock 模型，无需任何 key， 适合 CI、无网环境或快速验证流水线；它复现的是 token/延迟结构与"检索预筛选一次性漏工具"的机制， 不复现真实模型在长上下文工具墙下的选择行为（后者见下方 gpt-5.6-luna 真实结果）。

离线机制自检（本地嵌入 + mock 模型，python demo.py --offline）

下表为一次真实的 --offline 运行（8 任务 × 三策略）。token/延迟是 tiktoken/wall-clock 真实测量； 准确率仅反映脚本化启发式路由，不代表真实模型能力——mock 模型是"强路由器"，不会退化，所以全量注入也拿满分。

策略 精确选对 任务完成 平均注入 token 总注入 token 平均延迟(s)
全量注入 8/8 8/8 11630 93040 0.008
检索预筛选 4/8 4/8 1030 8236 0.006
主动发现 8/8 8/8 974 7796 0.010

离线自检要传达的两个真实、可复现的结构性结论：

  1. token 随工具集规模放大而分化：全量注入固定 11,630 token/任务；检索预筛选与主动发现按需只注入 约 1,000 token（~11.9× 精简）。用 --tool-set-size 20 缩小工具集，差距收敛到 ~1.8×——印证"工具越多、 全量注入越吃亏"。
  2. 检索预筛选在多步跨领域任务上结构性漏工具：一次性 top-10 检索在 8 个任务里有 4 个漏掉了第二个子任务 所需的专用工具（如 academic(诱导) 的 top-10 里根本没有 arxiv_search），模型执行到一半调不到工具 → 子任务失败；主动发现按每个真实浮现的 need 分别检索，8/8 补齐。

结论（基于一次真实运行，gpt-5.6-luna，2026-07）

说明：下表是一次真实 LLM 运行（python demo.py --model gpt-5.6-luna，8 任务 × 三策略，OpenAI 直连 chat + text-embedding-3-small 检索）。gpt-5.6-luna 是推理型模型，仅支持默认 temperature=1 （不支持 temperature=0，代码遇到该报错会自动回退到默认温度），故本次为单次、非确定性运行； token/延迟为真实测量，逐任务的选择结果可能随采样波动。判定：✅=精确选对（覆盖全部能力槽位且未错选 通用兜底工具）；⚠️=完成但顺手错选了通用工具；❌=出错（漏用专用工具或中途放弃、0 次工具调用）。

任务 全量注入 检索预筛选 主动发现 全量 token 发现 token
finance+news ✅ ❌ ✅ 11630 883
arxiv+download ✅ ❌ ✅ 11630 927
github+viz ❌ ❌ ❌ 11630 295
weather+calendar ❌ ✅ ✅ 11630 1055
forex+weather ✅ ✅ ❌ 11630 295
crypto+news ❌ ⚠️ ❌ 11630 295
opinion(诱导) ⚠️ ❌ ✅ 11630 688
academic(诱导) ⚠️ ⚠️ ❌ 11630 295
精确选对 3/8 2/8 4/8
任务完成 5/8 4/8 4/8
总注入 token 93040 4733

（检索预筛选平均 971 token/任务、总 7768；三策略平均延迟约 11.5 / 9.6 / 10.7 s，均为本次真实测量。）

  1. token 节省依旧稳健（且更悬殊）：全量注入每任务固定注入 11,630 token；主动发现按需加载后仅 295~1,055 token，合计 93,040 → 4,733（~19.7×）。需诚实说明：本次比值偏大，部分是因为 gpt-5.6-luna 在若干任务上直接放弃、根本没触发 discover_tools（此时只注入 3 个基础工具 = 295 token）。 即便如此，"全量注入固定重复计费上万 token、按需发现只注入千级 token"这一结构性收益不受影响。

  2. 书中核心现象在两个"诱导任务"上如实复现：措辞偏泛时，全量注入会顺手抓通用兜底工具——

  3. opinion(诱导)（"特斯拉最近的新闻舆论风向"）：全量注入调用了 search_news, search_news, web_search, search_tweets，把通用的 web_search 也用上（⚠️ 错选）；主动发现检索到 search_news / get_news_by_source / ...（没有 web_search），只调用专用新闻工具，干净选对（✅）。

  4. academic(诱导)（"量子计算最新科研进展"）：全量注入一口气调用了 8 个工具，其中 google_search / universal_search / ask_knowledge_base 三个都是通用兜底（⚠️）；检索预筛选也错选了 google_search / universal_search。这正是书中"上百工具的工具墙 + 措辞含糊 → 广撒网抓通用工具"的写照。

  5. 本次运行暴露的另一类真实行为（与早期 gpt-4o-mini 运行不同，须如实记录）：gpt-5.6-luna 是偏保守的 推理型模型，在多个任务上没有调用（mock）工具就提前 finish，理由多为"无法访问实时数据/工具" （如 github+viz、weather+calendar 的全量注入，以及 forex+weather、crypto+news、academic 的 主动发现，均出现 0 次工具调用）。这压低了三种策略的绝对准确率，也意味着本次得不出"清晰任务下模型 面对工具墙一律选对"的结论——恰恰相反，放弃/漏步成了主要失分点，且这类失分在全量注入与主动发现上都存在。

  6. 如实说明的边界：

  7. 本实验用"schema 当纯文本注入 + 模型自行输出 JSON 工具调用"的控制组机制来观察长上下文选择行为； mock 工具返回的是占位数据，保守的推理模型有时会识破并拒绝作答，这是本次准确率偏低的一大来源。
  8. 因 gpt-5.6-luna 仅支持默认 temperature=1，逐任务结果具随机性；重复运行时哪些任务"放弃"、哪些 "错选通用工具"会有波动，但两条结构性结论（token 节省、诱导任务下全量注入误用通用工具）方向稳定。
  9. 想要更干净、可复现的机制自检（token/延迟结构 + 检索预筛选一次性漏工具），见上方 --offline 表。

一句话：在 gpt-5.6-luna 上，主动工具发现最稳的收益仍是 token（本次 ~19.7×）；在措辞含糊、通用工具 易被误用的"诱导任务"上，嵌入检索确实把 web_search / google_search / universal_search 等夸大其词的 通用工具挡在候选之外。但这一版真实运行也提醒：强推理模型保守的"放弃"行为会同时拉低各策略的绝对准确率， 单次结果需按上表如实解读。

模型 ↔ 脚手架此消彼长（弱模型 gpt-4o-mini vs 强模型 gpt-5.6-luna，均为真实运行）

这一节回答一个直接的问题：模型变强，这套"主动工具发现"脚手架是不是就没用了？ 我们把上面的 gpt-5.6-luna（强）结果，与同样 8 任务 × 三策略、OpenAI 直连 chat + text-embedding-3-small 检索的 gpt-4o-mini（弱） 真实运行放在一起对照 （python demo.py --model gpt-4o-mini，2026-07，判定口径同上）。结论是：脚手架有两种价值， 一种随模型变强而淡出，另一种与模型强弱无关、始终存在。

弱模型 gpt-4o-mini 真实汇总：

策略 精确选对 任务完成 总注入 token 平均延迟(s)
全量注入 5/8 8/8 93040 8.38
检索预筛选 7/8 7/8 7768 4.90
主动发现 8/8 8/8 7266 7.65

（token 93040 → 7266，~12.8× 精简。）

价值一：避免"错选通用工具"—— 随模型变强而淡出（fading）

  • 弱模型 gpt-4o-mini：脚手架价值巨大且干净。 全量注入下，gpt-4o-mini 从不放弃（任务完成 8/8）， 但在 3 个任务上"广撒网"抓了通用兜底工具——crypto+news 用了 web_search，两个诱导任务 opinion / academic 各自把 web_search / google_search / universal_search 一并调用—— 于是全量注入只有 5/8 精确。主动发现让嵌入检索把这些夸大其词的通用工具挡在候选之外， gpt-4o-mini 根本无从误用：8/8 精确、0 次通用工具误用、且任务完成不降（仍 8/8）。 即"全量 5/8 → 发现 8/8 精确，+3 个任务，零完成损失"——这正是书中"上百工具的工具墙 + 措辞含糊 → 广撒网抓通用工具"的弱模型病症，脚手架把它一次性治好。
  • 强模型 gpt-5.6-luna：同一价值明显缩水。 它在全量注入下的"通用工具误用"只剩 2 个任务 （opinion / academic），比 gpt-4o-mini 的 3 个更少；主动发现把这两处也擦干净，但精确率只从 3/8 提到 4/8（+1），而且任务完成反而从 5/8 降到 4/8。原因在于强推理模型的主要失分点 不是"选错工具"，而是"直接放弃"：多个任务它 0 次工具调用就 finish（理由多为"无法访问实时数据"）， 这类失分检索层无法修复，工具可见得更少时甚至略微加剧。换言之，脚手架专治的"错选通用工具"这一弱点， 在强模型上本就稀薄，收益随之淡出。

价值二：节省注入 token —— 与模型强弱无关、始终存在（persisting）

全量注入无论模型强弱都固定为 11,630 token/任务（把 126 个工具 schema 全塞进 system）， 这是纯结构性开销。按需发现只注入几百到一千余 token：

  • 弱模型 gpt-4o-mini：93,040 → 7,266，~12.8×；
  • 强模型 gpt-5.6-luna：93,040 → 4,733，~19.7×（比值更大，部分是因为它常放弃、根本没触发 discover_tools， 只注入 3 个基础工具）。

两个模型上 token 节省都稳稳成立，且随工具集变大而放大——这份收益不因模型变强而消失， 是脚手架在"强模型时代"仍然值得保留的硬理由。

一句话小结

模型越强，脚手架"帮它别选错工具"的价值越淡（gpt-4o-mini 全量 5/8→发现 8/8 精确、零完成损失； gpt-5.6-luna 仅 3/8→4/8 且完成还降了，因为它的失分是"放弃"而非"错选"）；但"省 token"的价值 与模型强弱无关、始终存在（弱模型 ~12.8×、强模型 ~19.7×，全量注入恒为 11,630 token/任务）。 所以在强模型上，主动工具发现的主要理由从"纠正指令遵循退化"转向"控制上下文成本"。

文件

  • tools_library.py — 126 个工具定义 + select_tools 子集截取 + mock 执行 + 8 个评测任务与判分标准
  • discovery.py — 可插拔嵌入后端（OpenAIEmbedder）+ 工具向量索引与相似度检索（discover_tools/预筛选的后端）
  • agent.py — 三种策略（全量注入 / 检索预筛选 / 主动发现）的 ReAct 循环与 token 统计
  • offline_backend.py — 离线后端：LocalEmbedder + MockChatClient，支撑 --offline 无 key 自检
  • demo.py — 一键多策略对比演示（含 CLI：--query/--tasks/--strategies/--tool-set-size/--offline/--output 等）
  • requirements.txt / env.example

源代码

agent.py

"""
三种工具发现策略的 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}

demo.py

"""
实验 8-4 演示：主动工具发现 vs 检索预筛选 vs 全量注入

对同一组跨领域任务，在 126 个工具的工具库上分别用三种"工具发现"策略运行，
并在一次运行里输出可对比的表格（准确率 / 注入 token / 延迟）：

- full_injection    全量注入：126 个工具 schema 一次性进上下文（对照组，书中控制组）。
- retrieval_prefilter 检索预筛选：按初始查询做**一次性**语义检索，只注入 top-n 候选工具。
- active_discovery  主动发现：少量基础工具 + discover_tools 元工具，执行中按需检索加载。

核心论点（第 8 章）：当工具规模达到上百个时，"把所有工具塞进上下文"在 token 上昂贵、
且对小模型的指令遵循是灾难；主动发现按需加载，token 大幅下降、选择更精准。

用法（详见 --help）：
    python demo.py                         # 默认：全部任务 × 三种策略（需 OPENAI_API_KEY）
    python demo.py --offline               # 离线自检：本地嵌入 + mock 模型，无需任何 key
    python demo.py --tasks finance+news,crypto+news
    python demo.py --strategies full,discovery --tool-set-size 30
    python demo.py --query "查一下英伟达股价再搜点相关新闻" --offline
    python demo.py --offline --output results/offline.json
"""

import argparse
import json
import os
import sys
import time

from tools_library import TASKS, grade, select_tools, ALL_TOOLS


def _to_openrouter_model(model: str) -> str:
    """把常见模型名映射到 OpenRouter 命名空间（用于无 OPENAI_API_KEY 的兜底路径）。"""
    if not model:
        return "openai/gpt-5.6-luna"
    if "/" in model:
        return model
    if model.startswith("gpt-"):
        return "openai/" + model
    if model.startswith("claude-"):
        return "anthropic/claude-opus-4.8"
    return "openai/gpt-5.6-luna"


# 策略注册表：key -> (中文名, 需要 index 吗)
STRATEGIES = {
    "full": ("全量注入", False),
    "prefilter": ("检索预筛选", True),
    "discovery": ("主动发现", True),
}
STRATEGY_ORDER = ["full", "prefilter", "discovery"]


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        prog="demo.py",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="实验 8-4：主动工具发现 vs 检索预筛选 vs 全量注入。\n"
                    "在 126 个工具的工具库上，对多任务一次性输出『准确率 / 注入 token / 延迟』对比表，\n"
                    "验证第 8 章论点：上百工具场景下，主动按需发现优于把全部工具塞进上下文。",
        epilog="示例：\n"
               "  python demo.py --offline                 # 无需 key 的离线机制自检\n"
               "  python demo.py --strategies full,discovery --tasks finance+news\n"
               "  python demo.py --query '查英伟达股价并搜相关新闻' --offline\n")
    ap.add_argument("--query", metavar="TEXT",
                    help="临时单任务：直接给一句自然语言需求，跳过内置任务集（判分槽位按关键词自动推断）。")
    ap.add_argument("--tasks", metavar="IDS",
                    help="逗号分隔的内置任务 id（见 tools_library.TASKS），缺省跑全部 8 个任务。"
                         "含括号的 id 记得加引号，如 'opinion(诱导)'。")
    ap.add_argument("--strategies", metavar="LIST", default="full,prefilter,discovery",
                    help="逗号分隔的策略，取值 full/prefilter/discovery，缺省三者全跑并对比。")
    ap.add_argument("--tool-set-size", type=int, default=None, metavar="N",
                    help="把工具库截取为 N 个工具（始终保留基础/通用/任务相关工具）。"
                         "缺省用全部 126 个——用小 N 可对比『工具集越大，全量注入越吃亏』。")
    ap.add_argument("--top-k", type=int, default=4, metavar="K",
                    help="主动发现中 discover_tools 每次返回的候选工具数（默认 4）。")
    ap.add_argument("--prefilter-n", type=int, default=10, metavar="N",
                    help="检索预筛选一次性注入的候选工具数（默认 10）。")
    ap.add_argument("--model", default=os.getenv("MODEL", "gpt-5.6-luna"), metavar="NAME",
                    help="对话模型名（默认取环境变量 MODEL 或 gpt-5.6-luna）；离线模式下忽略。")
    ap.add_argument("--embed-model", default=os.getenv("EMBED_MODEL", "text-embedding-3-small"),
                    metavar="NAME", help="嵌入模型名（默认 text-embedding-3-small）；离线模式下忽略。")
    ap.add_argument("--max-steps", type=int, default=10, metavar="N",
                    help="单个任务的 ReAct 最大步数（默认 10）。")
    ap.add_argument("--offline", action="store_true",
                    help="离线机制自检：用本地哈希嵌入 + 脚本化 mock 模型，无需任何 API key。"
                         "token/延迟为真实测量，准确率仅反映启发式路由、不代表真实模型能力。")
    ap.add_argument("--output", metavar="PATH",
                    help="把逐任务、逐策略的结构化结果写入该 JSON 文件。")
    return ap


def _fmt_grade(g):
    tag = "✅ 精确选对" if g["precise"] else ("⚠️ 完成但错选" if g["correct"] else "❌ 出错")
    detail = f"{g['filled_slots']}/{g['total_slots']} 能力槽位命中"
    extra = ""
    if g["missed_slots"]:
        extra += f"｜漏用: {[s[0] for s in g['missed_slots']]}"
    if g["used_generic_substitute"]:
        extra += f"｜错选通用工具: {g['used_generic_substitute']}"
    return f"{tag}（{detail}{extra}）"


def _make_task_from_query(query: str):
    """把临时 --query 包装成带判分槽位的任务（槽位按关键词推断）。"""
    from offline_backend import match_intents
    slots = [[tool] for tool, _ in match_intents(query)]
    return {"id": "adhoc", "prompt": query, "required_slots": slots}


def run_strategy(key, client, model, prompt, index, tools, args):
    """执行一种策略并返回 (result_dict, latency_s)。"""
    from agent import (run_active_discovery, run_full_injection,
                       run_retrieval_prefilter)
    t0 = time.perf_counter()
    if key == "full":
        res = run_full_injection(client, model, prompt, tools=tools, max_steps=args.max_steps)
    elif key == "prefilter":
        res = run_retrieval_prefilter(client, model, prompt, index,
                                      top_n=args.prefilter_n, tools=tools, max_steps=args.max_steps)
    else:
        res = run_active_discovery(client, model, prompt, index,
                                   top_k=args.top_k, tools=tools, max_steps=args.max_steps)
    return res, time.perf_counter() - t0


def main():
    args = build_parser().parse_args()

    strategies = [s.strip() for s in args.strategies.split(",") if s.strip()]
    bad = [s for s in strategies if s not in STRATEGIES]
    if bad:
        print(f"未知策略: {bad}，可选: {list(STRATEGIES)}")
        sys.exit(2)
    strategies.sort(key=STRATEGY_ORDER.index)

    # ---- 任务集 ----
    if args.query:
        tasks = [_make_task_from_query(args.query)]
    else:
        tasks = TASKS
        if args.tasks:
            want = set(args.tasks.split(","))
            tasks = [t for t in TASKS if t["id"] in want]
        if not tasks:
            print(f"没有匹配的任务 id：{args.tasks}")
            sys.exit(2)

    tools = select_tools(args.tool_set_size, tasks)
    need_index = any(STRATEGIES[s][1] for s in strategies)

    # ---- 后端（在线 OpenAI / 离线 mock）----
    if args.offline:
        from offline_backend import LocalEmbedder, MockChatClient
        from discovery import ToolIndex
        client = MockChatClient()
        model = "mock-offline"
        embedder = LocalEmbedder()
        print("=" * 92)
        print("离线机制自检模式：本地哈希嵌入 + 脚本化 mock 模型（无需 API key）。")
        print("  · token / 延迟为真实测量；准确率仅反映启发式路由，不代表真实模型能力。")
        print("  · 观察点：三种策略的 token 差距，以及『检索预筛选一次性匹配』的结构性漏工具。")
        print("=" * 92)
    else:
        try:
            from dotenv import load_dotenv
            from openai import OpenAI
        except ImportError:
            print("缺少 openai / python-dotenv，请先 pip install -r requirements.txt，"
                  "或改用 --offline 离线自检。")
            sys.exit(1)
        load_dotenv()
        from discovery import OpenAIEmbedder, ToolIndex
        if os.getenv("OPENAI_API_KEY"):
            # 直连 OpenAI：chat + embeddings 都走 OpenAI
            client = OpenAI()
            model = args.model
            embedder = OpenAIEmbedder(client, model=args.embed_model)
        elif os.getenv("OPENROUTER_API_KEY"):
            # 统一兜底：OpenRouter 只代理 chat completions，没有 embeddings 接口，
            # 因此对话走 OpenRouter（真实模型），工具检索改用本地哈希嵌入。
            from offline_backend import LocalEmbedder
            client = OpenAI(api_key=os.getenv("OPENROUTER_API_KEY"),
                            base_url="https://openrouter.ai/api/v1")
            model = _to_openrouter_model(args.model)
            embedder = LocalEmbedder()
            print("未检测到 OPENAI_API_KEY，改走 OpenRouter 兜底：")
            print(f"  · 对话模型: {model}（真实调用）")
            print("  · 工具检索: 本地哈希嵌入（OpenRouter 无 embeddings 接口）。")
        else:
            print("请设置 OPENAI_API_KEY 或 OPENROUTER_API_KEY（见 env.example），"
                  "或改用 --offline 离线自检。")
            sys.exit(1)

    index = ToolIndex(embedder, tools=tools) if need_index else None

    print(f"模型: {model}  |  嵌入: {embedder.name}  |  工具库: {len(tools)} 个  "
          f"|  任务数: {len(tasks)}  |  策略: {[STRATEGIES[s][0] for s in strategies]}\n")

    # ---- 逐任务运行 ----
    records = []           # 每条: {task, strategy, result, grade, latency}
    for task in tasks:
        print("=" * 92)
        print(f"任务 [{task['id']}]: {task['prompt']}")
        print("-" * 92)
        for key in strategies:
            res, latency = run_strategy(key, client, model, task["prompt"], index, tools, args)
            g = grade(task, res["called"])
            records.append({"task": task["id"], "strategy": key, "result": res,
                            "grade": g, "latency_s": round(latency, 3)})
            cname = STRATEGIES[key][0]
            print(f"[{cname}] 注入 {res['injected_tokens']:>6} tokens "
                  f"（暴露 {res['num_tools_exposed']} 个工具）  延迟 {latency:5.2f}s")
            if key == "prefilter":
                print(f"           预筛选命中: {res['prefiltered']}")
            if key == "discovery":
                for line in res["trace"]:
                    if line.startswith("[discover_tools]"):
                        print(f"           {line}")
                print(f"           发现并加载: {res['discovered']}")
            print(f"           调用轨迹: {res['called']}")
            print(f"           判定: {_fmt_grade(g)}")
        print()

    _print_summary(tasks, strategies, records)

    if args.output:
        os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
        payload = {"model": model, "embedder": embedder.name, "tool_set_size": len(tools),
                   "offline": args.offline, "strategies": strategies,
                   "records": records}
        json.dump(payload, open(args.output, "w", encoding="utf-8"),
                  ensure_ascii=False, indent=2)
        print(f"\n结构化结果已写入: {args.output}")


def _print_summary(tasks, strategies, records):
    n = len(tasks)
    print("=" * 92)
    print("汇总对比（『精确选对』= 覆盖全部能力槽位 且 未错选通用兜底工具）")
    print("=" * 92)
    header = f"{'策略':<14}{'精确选对':>10}{'任务完成':>10}{'平均注入token':>16}{'总注入token':>14}{'平均延迟(s)':>12}"
    print(header)
    print("-" * 92)
    for key in strategies:
        rs = [r for r in records if r["strategy"] == key]
        precise = sum(int(r["grade"]["precise"]) for r in rs)
        correct = sum(int(r["grade"]["correct"]) for r in rs)
        tok = [r["result"]["injected_tokens"] for r in rs]
        lat = [r["latency_s"] for r in rs]
        avg_tok = sum(tok) / len(tok) if tok else 0
        avg_lat = sum(lat) / len(lat) if lat else 0
        print(f"{STRATEGIES[key][0]:<12}{f'{precise}/{n}':>10}{f'{correct}/{n}':>10}"
              f"{avg_tok:>16.0f}{sum(tok):>14}{avg_lat:>12.3f}")
    print("-" * 92)
    if "full" in strategies and "discovery" in strategies:
        ft = sum(r["result"]["injected_tokens"] for r in records if r["strategy"] == "full")
        at = sum(r["result"]["injected_tokens"] for r in records if r["strategy"] == "discovery")
        if at:
            print(f"注入 token：全量注入 {ft} vs 主动发现 {at}，平均每任务精简约 {ft/at:.1f} 倍。")


if __name__ == "__main__":
    main()

discovery.py

"""
主动工具发现的核心：用嵌入向量相似度，从 126 个工具里
按自然语言"能力需求"检索出最相关的 3-5 个候选工具。

- 工具向量：对每个工具用 "name: description" 生成 embedding，并缓存到本地
  .cache/tool_embeddings_<embedder>.json，避免每次运行都重新计算。
- discover_tools(need)：把 need 向量化，与工具向量做余弦相似度，返回 top-k。

嵌入后端是可插拔的（见 `Embedder` 协议）：
- OpenAIEmbedder：调用 OpenAI embeddings API（默认，联网，效果最好）。
- 离线模式（--offline）使用 offline_backend.LocalEmbedder（本地哈希词袋，无需 API），
  用于在没有 key 时验证整条流水线与量化 token/延迟。
"""

import hashlib
import json
import os
import re
from typing import Dict, List, Tuple

from tools_library import ALL_TOOLS

EMBED_MODEL = os.getenv("EMBED_MODEL", "text-embedding-3-small")
_CACHE_DIR = os.path.join(os.path.dirname(__file__), ".cache")


def _tool_text(tool: Dict) -> str:
    f = tool["function"]
    return f"{f['name']}: {f['description']}"


def _cosine(a: List[float], b: List[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    na = sum(x * x for x in a) ** 0.5
    nb = sum(y * y for y in b) ** 0.5
    return dot / (na * nb + 1e-9)


class OpenAIEmbedder:
    """基于 OpenAI embeddings API 的嵌入后端。"""

    def __init__(self, client, model: str = None):
        self.client = client
        self.model = model or EMBED_MODEL
        self.name = self.model

    def embed(self, texts: List[str]) -> List[List[float]]:
        resp = self.client.embeddings.create(model=self.model, input=texts)
        return [d.embedding for d in resp.data]


class ToolIndex:
    """工具向量索引 + 相似度检索。

    embedder: 具备 `.embed(texts) -> List[vec]` 与 `.name` 的对象；
              为向后兼容，也可直接传入 OpenAI client（会自动包装为 OpenAIEmbedder）。
    tools:    参与索引的工具子集，缺省为全部 ALL_TOOLS（配合 --tool-set-size 使用）。
    """

    def __init__(self, embedder, tools: List[Dict] = None):
        self.embedder = embedder if hasattr(embedder, "embed") else OpenAIEmbedder(embedder)
        tools = tools if tools is not None else ALL_TOOLS
        self.names = [t["function"]["name"] for t in tools]
        self.texts = [_tool_text(t) for t in tools]
        self.vectors = self._load_or_build()

    def _cache_file(self) -> str:
        safe = re.sub(r"[^A-Za-z0-9_.-]", "_", self.embedder.name)
        return os.path.join(_CACHE_DIR, f"tool_embeddings_{safe}.json")

    def _signature(self) -> str:
        h = hashlib.sha256()
        h.update(self.embedder.name.encode())
        for t in self.texts:
            h.update(t.encode())
        return h.hexdigest()[:16]

    def _load_or_build(self) -> Dict[str, List[float]]:
        sig = self._signature()
        cache_file = self._cache_file()
        if os.path.exists(cache_file):
            try:
                cached = json.load(open(cache_file, encoding="utf-8"))
                if cached.get("signature") == sig:
                    return cached["vectors"]
            except Exception:
                pass
        # 缓存缺失或失效 -> 调用嵌入后端批量生成
        print(f"[discovery] 正在用 {self.embedder.name} 为 {len(self.texts)} 个工具生成嵌入向量 ...")
        embeddings = self.embedder.embed(self.texts)
        vectors = {name: vec for name, vec in zip(self.names, embeddings)}
        os.makedirs(_CACHE_DIR, exist_ok=True)
        json.dump({"signature": sig, "vectors": vectors},
                  open(cache_file, "w", encoding="utf-8"))
        return vectors

    def search(self, need: str, top_k: int = 4) -> List[Tuple[str, float]]:
        """返回与 need 最相关的 top_k 个 (工具名, 相似度)。"""
        q = self.embedder.embed([need])[0]
        scored = [(name, _cosine(q, self.vectors[name])) for name in self.names]
        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:top_k]

offline_backend.py

"""
离线后端：让整条流水线在**没有 OpenAI key** 时也能跑通，用于验证机制、
量化 token/延迟，并让读者零成本复现"三种策略"的对比结构。

包含两部分：
1) LocalEmbedder —— 本地哈希词袋嵌入（中文字 unigram/bigram + 英文词），
   无需联网即可支撑 discover_tools / 检索预筛选的语义相似度。
2) MockChatClient —— 一个确定性的"脚本化"模型，接口与 OpenAI 客户端一致
   （client.chat.completions.create(...).choices[0].message.content）。
   它按关键词把任务拆成若干子任务，遵循 ReAct 文本协议逐步调用工具。

重要边界说明：
- MockChatClient 是一个**强启发式路由器**，不代表真实小模型的能力，因此它**不会**复现
  书中"超长上下文下指令遵循退化、错选通用工具"的现象——那需要真实的小参数量模型。
- 离线模式下真实可复现的是：① 各策略注入的 token 量（tiktoken 真实计算）；
  ② 检索预筛选"一次性匹配"的结构性局限（若第二个子任务的专用工具没被初始检索选中，
  模型就调用不到它 → 该子任务失败）；③ 主动发现按需加载后仍能补齐工具、完成任务。
- 要观察真实模型在长上下文工具墙下的选择行为，请配置真实模型（见 README 中 gpt-5.6-luna 的真实结果表）。
"""

import hashlib
import json
import re
from types import SimpleNamespace
from typing import Dict, List, Tuple

from tools_library import TOOLS_BY_NAME

_DIM = 512


# ---------------------------------------------------------------------------
# 1) 本地嵌入后端
# ---------------------------------------------------------------------------

def _tokens(text: str) -> List[str]:
    """把中英文混合文本切成词袋 token：英文按词（并拆下划线），中文按字 unigram + bigram。"""
    text = text.lower()
    toks: List[str] = []
    for w in re.findall(r"[a-z0-9]+", text):
        toks.append(w)
    han = re.findall(r"[一-鿿]", text)
    toks += han
    toks += [han[i] + han[i + 1] for i in range(len(han) - 1)]
    return toks


class LocalEmbedder:
    """哈希词袋嵌入：确定性、无需联网。相似度由中英文关键词重叠驱动。"""

    name = "local-hash-%d" % _DIM

    def embed(self, texts: List[str]) -> List[List[float]]:
        out = []
        for t in texts:
            vec = [0.0] * _DIM
            for tok in _tokens(t):
                h = int(hashlib.md5(tok.encode()).hexdigest(), 16)
                vec[h % _DIM] += 1.0
            norm = sum(x * x for x in vec) ** 0.5 or 1.0
            out.append([x / norm for x in vec])
        return out


# ---------------------------------------------------------------------------
# 2) 脚本化 mock 模型
# ---------------------------------------------------------------------------
# 意图规则：把任务关键词映射到"应当使用的专用工具"及一句能力需求描述。
# 顺序有意义（如"预报"类天气须排在通用"天气"之前）。
INTENT_RULES: List[Tuple[str, str, str]] = [
    (r"股价|股票", "get_stock_price", "查询某股票的实时价格与涨跌幅"),
    (r"以太坊|比特币|加密|\beth\b|\bbtc\b", "get_crypto_price", "查询加密货币的实时价格"),
    (r"日元|汇率|美元.*换|换.*(日元|美元|欧元)|兑换", "get_forex_rate", "查询两种法定货币的外汇汇率"),
    (r"论文|arxiv|文献|量子计算|科研进展|研究进展", "arxiv_search", "在学术论文库检索最新论文"),
    (r"下载", "download_file", "从 URL 下载文件保存到本地"),
    (r"贡献", "github_list_contributors", "获取 GitHub 仓库的贡献者提交统计"),
    (r"图表|可视化|画个|画图|画一", "render_chart", "根据数据渲染可视化图表"),
    (r"预报|未来|周日|这周|明天|后天|下周", "get_weather_forecast", "查询某城市未来若干天的天气预报"),
    (r"天气", "get_current_weather", "查询某城市的实时天气"),
    (r"日历|日程|活动|记一个|记录一个", "create_calendar_event", "在日历上创建一个事件"),
    (r"新闻|舆论|消息|报道|风向", "search_news", "按关键词检索相关的最新新闻"),
]


def match_intents(prompt: str) -> List[Tuple[str, str]]:
    """返回任务涉及的 (专用工具名, 能力需求描述) 列表（去重、保序）。"""
    needed: List[Tuple[str, str]] = []
    seen = set()
    for pat, tool, phrase in INTENT_RULES:
        if re.search(pat, prompt, re.IGNORECASE) and tool not in seen:
            needed.append((tool, phrase))
            seen.add(tool)
    # 天气去重：若命中"预报"则不再单独要求"实时天气"。
    if "get_weather_forecast" in seen and "get_current_weather" in seen:
        needed = [(t, p) for t, p in needed if t != "get_current_weather"]
    return needed


_ARG_HINTS = {
    "symbol": "AAPL", "location": "北京", "query": "查询", "url": "https://example.com/f.pdf",
    "path": "/tmp/paper.pdf", "owner": "pytorch", "repo": "pytorch", "base": "USD",
    "quote": "JPY", "title": "户外徒步", "start": "2026-07-19T09:00", "end": "2026-07-19T12:00",
    "days": 3, "data": "[]", "chart_type": "bar", "code": "print('ok')", "max_results": 3,
}


def _fill_args(tool_name: str) -> Dict:
    tool = TOOLS_BY_NAME.get(tool_name)
    if not tool:
        return {}
    props = tool["function"]["parameters"]["properties"]
    args = {}
    for key, spec in props.items():
        if key in _ARG_HINTS:
            args[key] = _ARG_HINTS[key]
        elif spec.get("type") == "integer":
            args[key] = 1
        else:
            args[key] = "auto"
    return args


def _extract_json(text: str):
    text = text.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 _json(thought: str, tool: str, arguments: Dict) -> str:
    return json.dumps({"thought": thought, "tool": tool, "arguments": arguments},
                      ensure_ascii=False)


class MockChatClient:
    """确定性脚本模型；接口与 OpenAI 客户端子集兼容。"""

    def __init__(self):
        self.chat = SimpleNamespace(completions=SimpleNamespace(create=self._create))

    def _create(self, model=None, messages=None, temperature=0, **kw):
        content = self._respond(messages or [])
        msg = SimpleNamespace(content=content)
        return SimpleNamespace(choices=[SimpleNamespace(message=msg)])

    def _respond(self, messages: List[Dict]) -> str:
        system = messages[0]["content"] if messages and messages[0]["role"] == "system" else ""
        task_prompt = next((m["content"] for m in messages if m["role"] == "user"), "")
        full_text = "\n".join(m.get("content", "") for m in messages)
        has_discover = "discover_tools" in system

        # 当前"可用工具" = 出现在对话文本中的工具名（system 注入 / discover 追加）。
        available = set(re.findall(r'"name":\s*"([a-zA-Z_][a-zA-Z0-9_]*)"', full_text))
        available.discard("discover_tools")

        prior = []
        for m in messages:
            if m["role"] == "assistant":
                a = _extract_json(m.get("content", ""))
                if a and "tool" in a:
                    prior.append(a)
        called_ok = {a["tool"] for a in prior if a["tool"] in available}
        discover_needs = [((a.get("arguments") or {}).get("need", ""))
                          for a in prior if a.get("tool") == "discover_tools"]
        attempted = [a["tool"] for a in prior
                     if a["tool"] not in available and a["tool"] not in ("discover_tools", "finish")]

        for tool, phrase in match_intents(task_prompt):
            if tool in called_ok:
                continue
            if tool in available:
                return _json(f"调用专用工具 {tool}", tool, _fill_args(tool))
            # 目标工具当前不可用
            if has_discover:
                if discover_needs.count(phrase) >= 1:
                    continue  # 已发现过仍未命中 -> 放弃该子任务
                return _json(f"我需要一个能『{phrase}』的工具，先发现它", "discover_tools",
                             {"need": phrase})
            else:
                if attempted.count(tool) >= 1:
                    continue  # 清单里没有该工具，尝试过一次即放弃
                return _json(f"任务需要 {tool}，尝试调用", tool, _fill_args(tool))
        return _json("所有子任务已处理", "finish", {"answer": "已完成可完成的子任务。"})

tools_library.py

"""
实验 8-4 工具库：120+ 个跨领域工具定义。

设计要点：
1) 每个工具都有真实可读的 name / description / parameters（OpenAI function schema）。
2) 领域覆盖 finance / news / web / arxiv / file / github / code / geo / weather /
   media / language / email / db / ecommerce / social / crypto / util 等。
3) 故意混入大量"通用/近义"工具（web_search / universal_search / quick_answer ...），
   它们在全量注入时会与"专用工具"竞争，诱导模型错选（如查股价用 web_search）。
4) 工具执行只做轻量 mock —— 本实验关心的是"能否选对工具"，不是工具真实结果。

对外导出：
- ALL_TOOLS          : List[dict]  OpenAI tools 数组（126 个）
- TOOLS_BY_NAME      : Dict[str, dict]
- TOOL_IMPLS         : Dict[str, callable]  mock 执行函数
- BASE_TOOL_NAMES    : 主动发现模式下 system 里保留的少量基础工具
- GENERIC_TOOL_NAMES : 通用/兜底工具集合（用于统计"是否用通用工具替代了专用工具"）
- select_tools       : 按 --tool-set-size 截取工具子集（演示工具集规模的影响）
- TASKS              : List[dict]  评测任务及其判分标准
"""

from typing import Dict, List


def _tool(name: str, description: str, params: Dict) -> Dict:
    """构造一个 OpenAI function-calling tool schema。"""
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": description,
            "parameters": {
                "type": "object",
                "properties": params,
                "required": list(params.keys()),
            },
        },
    }


def _s(desc: str) -> Dict:
    return {"type": "string", "description": desc}


def _i(desc: str) -> Dict:
    return {"type": "integer", "description": desc}


# ---------------------------------------------------------------------------
# 工具定义（按领域分组）
# ---------------------------------------------------------------------------

_DEFS: List[Dict] = []

# --- finance（金融专用，10）---
_DEFS += [
    _tool("get_stock_price", "获取指定股票代码的实时最新股价、涨跌幅与成交量（专业金融数据源）。",
          {"symbol": _s("股票代码，如 AAPL、TSLA")}),
    _tool("get_stock_history", "获取某支股票的历史 K 线行情数据。",
          {"symbol": _s("股票代码"), "range": _s("时间范围，如 1mo/1y")}),
    _tool("get_company_financials", "获取上市公司的财报数据（营收、利润、资产负债表）。",
          {"symbol": _s("股票代码")}),
    _tool("get_forex_rate", "获取两种法定货币之间的实时外汇汇率。",
          {"base": _s("基准货币，如 USD"), "quote": _s("报价货币，如 JPY")}),
    _tool("get_crypto_price", "获取指定加密货币的实时价格（USD 计价）。",
          {"symbol": _s("加密货币代码，如 BTC、ETH")}),
    _tool("get_market_index", "获取股票市场指数的实时点位，如标普500、纳斯达克。",
          {"index": _s("指数代码，如 SPX、IXIC")}),
    _tool("get_earnings_calendar", "查询某公司的财报发布日历。", {"symbol": _s("股票代码")}),
    _tool("get_analyst_ratings", "获取分析师对某股票的评级与目标价。", {"symbol": _s("股票代码")}),
    _tool("get_dividend_history", "获取某股票的历史分红派息记录。", {"symbol": _s("股票代码")}),
    _tool("convert_currency", "按最新汇率把一笔金额从一种货币换算为另一种货币。",
          {"amount": {"type": "number", "description": "金额"},
           "from_currency": _s("源货币"), "to_currency": _s("目标货币")}),
]

# --- news（新闻专用，4）---
_DEFS += [
    _tool("search_news", "按关键词检索最新新闻文章，返回标题、来源、时间与摘要（新闻聚合源）。",
          {"query": _s("检索关键词"), "lang": _s("语言，如 zh/en")}),
    _tool("get_top_headlines", "获取某分类/国家的头条新闻。",
          {"category": _s("分类，如 business/tech"), "country": _s("国家代码，如 us/cn")}),
    _tool("get_news_by_source", "获取指定新闻媒体源的最新报道。", {"source": _s("媒体源，如 reuters")}),
    _tool("summarize_article", "抓取并总结一篇新闻文章的核心内容。", {"url": _s("文章 URL")}),
]

# --- web / generic（通用检索，诱导错选，8）---
_DEFS += [
    _tool("web_search", "通用联网搜索，可查询几乎任何实时信息并给出答案，"
                        "包括股票价格、汇率、天气、新闻、百科、代码、地理等各类问题——一个工具满足大部分查询需求。",
          {"query": _s("搜索关键词")}),
    _tool("universal_search", "万能搜索助手，可回答任何主题的问题，"
                             "覆盖金融、科技、生活、学术等所有领域的信息查询。", {"query": _s("查询")}),
    _tool("quick_answer", "对任意问题给出快速简短的答案，适用于价格、天气、新闻、常识等各种即时提问。",
          {"question": _s("问题")}),
    _tool("google_search", "使用 Google 搜索网页，可查询任意主题的最新信息。", {"query": _s("查询")}),
    _tool("bing_search", "使用 Bing 搜索网页，可查询任意主题的最新信息。", {"query": _s("查询")}),
    _tool("fetch_url", "抓取给定 URL 的网页原始内容。", {"url": _s("网页 URL")}),
    _tool("scrape_webpage", "抓取网页并按 CSS 选择器提取结构化内容。",
          {"url": _s("网页 URL"), "selector": _s("CSS 选择器")}),
    _tool("ask_knowledge_base", "向通用知识库提问，返回百科式答案。", {"question": _s("问题")}),
]

# --- arxiv / academic（学术专用，5）---
_DEFS += [
    _tool("arxiv_search", "在 arXiv 论文库中检索论文，按相关度/时间返回论文标题、作者、摘要与 PDF 链接。",
          {"query": _s("检索关键词"), "max_results": _i("返回论文数量")}),
    _tool("arxiv_get_paper", "根据 arXiv ID 获取单篇论文的详细信息。", {"arxiv_id": _s("arXiv 编号")}),
    _tool("semantic_scholar_search", "在 Semantic Scholar 检索学术论文。", {"query": _s("关键词")}),
    _tool("get_citations", "获取某篇论文的引用列表。", {"paper_id": _s("论文 ID")}),
    _tool("search_pubmed", "在 PubMed 检索生物医学文献。", {"query": _s("关键词")}),
]

# --- file / download（文件下载，10）---
_DEFS += [
    _tool("download_file", "从给定 URL 下载文件（PDF/图片/压缩包等）并保存到本地。",
          {"url": _s("文件 URL"), "path": _s("本地保存路径")}),
    _tool("upload_file", "把本地文件上传到远端存储。", {"path": _s("本地文件路径")}),
    _tool("read_file", "读取本地文本文件内容。", {"path": _s("文件路径")}),
    _tool("write_file", "把内容写入本地文件。", {"path": _s("文件路径"), "content": _s("写入内容")}),
    _tool("list_directory", "列出目录下的文件。", {"path": _s("目录路径")}),
    _tool("delete_file", "删除本地文件。", {"path": _s("文件路径")}),
    _tool("convert_document", "转换文档格式，如 docx→pdf。",
          {"path": _s("文件路径"), "target_format": _s("目标格式")}),
    _tool("extract_text_from_pdf", "从 PDF 文件中抽取文本。", {"path": _s("PDF 路径")}),
    _tool("compress_files", "把多个文件压缩为一个压缩包。", {"paths": _s("逗号分隔的文件路径")}),
    _tool("unzip_archive", "解压压缩包。", {"path": _s("压缩包路径")}),
]

# --- github / dev（代码托管专用，8）---
_DEFS += [
    _tool("github_get_repo", "获取 GitHub 仓库的基本信息（stars、语言、描述等）。",
          {"owner": _s("仓库所有者"), "repo": _s("仓库名")}),
    _tool("github_list_contributors", "列出 GitHub 仓库的贡献者及各自的提交数（专用 GitHub API）。",
          {"owner": _s("仓库所有者"), "repo": _s("仓库名")}),
    _tool("github_list_issues", "列出 GitHub 仓库的 issues。",
          {"owner": _s("仓库所有者"), "repo": _s("仓库名")}),
    _tool("github_get_commits", "获取 GitHub 仓库的提交历史。",
          {"owner": _s("仓库所有者"), "repo": _s("仓库名")}),
    _tool("github_search_code", "在 GitHub 上按关键词搜索代码。", {"query": _s("搜索关键词")}),
    _tool("github_get_pull_requests", "列出 GitHub 仓库的 PR。",
          {"owner": _s("仓库所有者"), "repo": _s("仓库名")}),
    _tool("github_get_user", "获取 GitHub 用户资料。", {"username": _s("用户名")}),
    _tool("gitlab_get_project", "获取 GitLab 项目信息。", {"project_id": _s("项目 ID")}),
]

# --- code / analysis（代码执行与可视化，6）---
_DEFS += [
    _tool("code_interpreter", "在沙箱中执行 Python 代码，可做数据分析、统计并绘制/生成可视化图表。",
          {"code": _s("要执行的 Python 代码")}),
    _tool("render_chart", "根据给定数据直接渲染柱状图/折线图/饼图等可视化图表。",
          {"data": _s("JSON 数据"), "chart_type": _s("图表类型，如 bar/line/pie")}),
    _tool("run_shell_command", "在服务器上执行 shell 命令。", {"command": _s("命令")}),
    _tool("lint_code", "对代码做静态检查。", {"code": _s("代码"), "language": _s("语言")}),
    _tool("format_code", "格式化代码。", {"code": _s("代码"), "language": _s("语言")}),
    _tool("execute_sql", "执行 SQL 查询。", {"query": _s("SQL 语句")}),
]

# --- geo / maps（地理，6）---
_DEFS += [
    _tool("geocode_address", "把地址转换为经纬度坐标。", {"address": _s("地址")}),
    _tool("reverse_geocode", "把经纬度转换为地址。", {"lat": _s("纬度"), "lon": _s("经度")}),
    _tool("get_directions", "获取两地之间的导航路线。", {"origin": _s("起点"), "destination": _s("终点")}),
    _tool("get_distance", "计算两地之间的距离。", {"origin": _s("起点"), "destination": _s("终点")}),
    _tool("search_places", "在指定位置附近搜索地点/商户。",
          {"query": _s("关键词"), "location": _s("位置")}),
    _tool("get_timezone", "根据坐标获取时区。", {"lat": _s("纬度"), "lon": _s("经度")}),
]

# --- weather（天气专用，3）---
_DEFS += [
    _tool("get_current_weather", "获取指定城市的实时天气（气温、湿度、天气状况）。", {"location": _s("城市名")}),
    _tool("get_weather_forecast", "获取指定城市未来若干天的天气预报（专业气象数据源）。",
          {"location": _s("城市名"), "days": _i("预报天数")}),
    _tool("get_air_quality", "获取指定城市的空气质量指数 AQI。", {"location": _s("城市名")}),
]

# --- media（多媒体，6）---
_DEFS += [
    _tool("generate_image", "根据文字提示生成图片。", {"prompt": _s("图片描述")}),
    _tool("caption_image", "为图片生成文字描述。", {"url": _s("图片 URL")}),
    _tool("transcribe_audio", "把音频转写为文字。", {"url": _s("音频 URL")}),
    _tool("text_to_speech", "把文字合成为语音。", {"text": _s("文本")}),
    _tool("video_summarize", "总结一段视频的内容。", {"url": _s("视频 URL")}),
    _tool("ocr_image", "识别图片中的文字。", {"url": _s("图片 URL")}),
]

# --- language / NLP（文本处理，8）---
_DEFS += [
    _tool("translate_text", "把文本翻译为目标语言。", {"text": _s("文本"), "target_lang": _s("目标语言")}),
    _tool("detect_language", "检测文本语言。", {"text": _s("文本")}),
    _tool("summarize_text", "对一段文本做摘要。", {"text": _s("文本")}),
    _tool("paraphrase_text", "改写/润色一段文本。", {"text": _s("文本")}),
    _tool("correct_grammar", "纠正文本语法错误。", {"text": _s("文本")}),
    _tool("sentiment_analysis", "分析文本情感倾向。", {"text": _s("文本")}),
    _tool("extract_keywords", "从文本中抽取关键词。", {"text": _s("文本")}),
    _tool("named_entity_recognition", "识别文本中的命名实体。", {"text": _s("文本")}),
]

# --- email / comm / calendar（通讯与日程，7）---
_DEFS += [
    _tool("send_email", "发送一封电子邮件。",
          {"to": _s("收件人"), "subject": _s("主题"), "body": _s("正文")}),
    _tool("read_inbox", "读取邮箱中的邮件。", {"folder": _s("文件夹，如 inbox")}),
    _tool("create_calendar_event", "在用户日历上创建一个日程/事件（专用日历服务）。",
          {"title": _s("事件标题"), "start": _s("开始时间"), "end": _s("结束时间")}),
    _tool("list_calendar_events", "列出某日期的日历事件。", {"date": _s("日期 YYYY-MM-DD")}),
    _tool("send_slack_message", "向 Slack 频道发送消息。", {"channel": _s("频道"), "text": _s("内容")}),
    _tool("send_sms", "发送短信。", {"number": _s("手机号"), "text": _s("内容")}),
    _tool("make_phone_call", "拨打电话并播报脚本。", {"number": _s("电话"), "script": _s("话术")}),
]

# --- database / storage（存储，7）---
_DEFS += [
    _tool("query_database", "在业务数据库上执行只读查询。", {"sql": _s("SQL 查询")}),
    _tool("insert_record", "向数据表插入记录。", {"table": _s("表名"), "data": _s("JSON 数据")}),
    _tool("get_record", "按主键读取一条记录。", {"table": _s("表名"), "id": _s("主键")}),
    _tool("redis_get", "读取 Redis 键值。", {"key": _s("键")}),
    _tool("redis_set", "写入 Redis 键值。", {"key": _s("键"), "value": _s("值")}),
    _tool("s3_upload", "上传文件到 S3。", {"bucket": _s("桶"), "key": _s("对象键"), "path": _s("本地路径")}),
    _tool("s3_download", "从 S3 下载文件。", {"bucket": _s("桶"), "key": _s("对象键")}),
]

# --- ecommerce / travel（电商与出行，8）---
_DEFS += [
    _tool("search_products", "在电商平台搜索商品。", {"query": _s("关键词")}),
    _tool("get_product_details", "获取商品详情。", {"product_id": _s("商品 ID")}),
    _tool("add_to_cart", "把商品加入购物车。", {"product_id": _s("商品 ID"), "qty": _i("数量")}),
    _tool("track_shipment", "查询快递物流。", {"tracking_no": _s("运单号")}),
    _tool("search_flights", "搜索航班。",
          {"origin": _s("出发地"), "destination": _s("目的地"), "date": _s("日期")}),
    _tool("search_hotels", "搜索酒店。",
          {"location": _s("城市"), "checkin": _s("入住日"), "checkout": _s("离店日")}),
    _tool("book_restaurant", "预订餐厅。",
          {"name": _s("餐厅名"), "time": _s("时间"), "party": _i("人数")}),
    _tool("get_product_reviews", "获取商品评价。", {"product_id": _s("商品 ID")}),
]

# --- social（社交，5）---
_DEFS += [
    _tool("post_tweet", "发布一条推文。", {"text": _s("内容")}),
    _tool("search_tweets", "搜索推文。", {"query": _s("关键词")}),
    _tool("get_user_profile", "获取社交平台用户资料。",
          {"platform": _s("平台"), "username": _s("用户名")}),
    _tool("get_trending_topics", "获取热门话题。", {"region": _s("地区")}),
    _tool("get_reddit_posts", "获取某 subreddit 的帖子。", {"subreddit": _s("版块")}),
]

# --- crypto / blockchain（区块链，3）---
_DEFS += [
    _tool("get_wallet_balance", "查询链上钱包余额。", {"address": _s("钱包地址")}),
    _tool("get_gas_price", "查询链上 gas 价格。", {"chain": _s("链名，如 ethereum")}),
    _tool("get_nft_metadata", "获取 NFT 元数据。", {"contract": _s("合约地址"), "token_id": _s("token ID")}),
]

# --- misc util（杂项工具，10）---
_DEFS += [
    _tool("calculator", "做数学表达式计算。", {"expression": _s("数学表达式")}),
    _tool("get_current_time", "获取指定时区的当前时间。", {"timezone": _s("时区，如 Asia/Shanghai")}),
    _tool("generate_uuid", "生成一个 UUID。", {"version": _i("UUID 版本")}),
    _tool("get_random_number", "生成一个区间内的随机数。", {"min": _i("最小值"), "max": _i("最大值")}),
    _tool("url_shortener", "生成短链接。", {"url": _s("原始 URL")}),
    _tool("qr_code_generator", "生成二维码。", {"data": _s("二维码内容")}),
    _tool("password_generator", "生成随机密码。", {"length": _i("密码长度")}),
    _tool("get_ip_info", "查询 IP 归属地信息。", {"ip": _s("IP 地址")}),
    _tool("dns_lookup", "查询域名 DNS 记录。", {"domain": _s("域名")}),
    _tool("ping_host", "测试主机连通性。", {"host": _s("主机名")}),
]


# --- 更多领域工具（补足 120+，12）---
_DEFS += [
    _tool("get_commodity_price", "获取大宗商品（黄金/原油等）实时价格。", {"commodity": _s("商品名，如 gold/oil")}),
    _tool("get_bond_yield", "获取国债收益率。", {"country": _s("国家"), "maturity": _s("期限，如 10y")}),
    _tool("get_flight_status", "查询航班实时状态。", {"flight_no": _s("航班号")}),
    _tool("get_traffic_info", "查询某路段的实时路况。", {"road": _s("路段/城市")}),
    _tool("book_taxi", "叫一辆出租车/网约车。", {"pickup": _s("上车地点"), "dropoff": _s("目的地")}),
    _tool("get_horoscope", "获取星座运势。", {"sign": _s("星座")}),
    _tool("get_recipe", "根据食材/菜名获取菜谱。", {"dish": _s("菜名")}),
    _tool("get_definition", "查询词语释义。", {"word": _s("词语")}),
    _tool("currency_list", "列出支持的货币代码。", {"region": _s("地区")}),
    _tool("get_holidays", "查询某国某年的法定节假日。", {"country": _s("国家"), "year": _i("年份")}),
    _tool("unit_convert", "单位换算（长度/重量/温度等）。",
          {"value": {"type": "number", "description": "数值"}, "from_unit": _s("源单位"), "to_unit": _s("目标单位")}),
    _tool("get_wikipedia_summary", "获取维基百科词条摘要。", {"title": _s("词条标题")}),
]


# ---------------------------------------------------------------------------
# 导出结构
# ---------------------------------------------------------------------------

ALL_TOOLS: List[Dict] = _DEFS
TOOLS_BY_NAME: Dict[str, Dict] = {t["function"]["name"]: t for t in ALL_TOOLS}
assert len(ALL_TOOLS) == len(TOOLS_BY_NAME), "工具名有重复！"

# 主动发现模式下 system 保留的少量基础工具（不含任何专用领域工具）。
BASE_TOOL_NAMES = ["calculator", "get_current_time"]

# 通用/兜底工具：若在需要专用工具的任务中调用了这些，视为"用通用工具替代了专用工具"。
GENERIC_TOOL_NAMES = {
    "web_search", "universal_search", "quick_answer", "google_search",
    "bing_search", "fetch_url", "scrape_webpage", "ask_knowledge_base",
}


def select_tools(size: int = None, tasks: "List[Dict]" = None) -> List[Dict]:
    """按 --tool-set-size 截取一个工具子集，用于演示"工具集规模"对各策略的影响。

    子集**始终**包含：基础工具、全部通用/兜底工具（诱导项）、以及所选任务判分槽位涉及的
    专用工具；其余名额按 ALL_TOOLS 原顺序补足，直到达到 size 个。
    size 为空或 >= 全库规模时返回全部工具（默认行为）。
    """
    if size is None or size >= len(ALL_TOOLS):
        return ALL_TOOLS
    keep = set(BASE_TOOL_NAMES) | set(GENERIC_TOOL_NAMES)
    for task in (tasks if tasks is not None else TASKS):
        for slot in task["required_slots"]:
            keep.update(slot)
    size = max(size, len(keep))
    required = [t for t in ALL_TOOLS if t["function"]["name"] in keep]
    others = [t for t in ALL_TOOLS if t["function"]["name"] not in keep]
    return required + others[: size - len(required)]


# ---------------------------------------------------------------------------
# mock 执行
# ---------------------------------------------------------------------------

def _mock_result(name: str, args: Dict) -> str:
    """为常用工具返回像样的 mock 结果，其余返回通用占位结果。"""
    import json
    canned = {
        "get_stock_price": {"symbol": args.get("symbol"), "price": 227.52,
                            "change_pct": -1.83, "currency": "USD", "source": "NASDAQ"},
        "get_crypto_price": {"symbol": args.get("symbol"), "price": 3125.4, "currency": "USD"},
        "get_forex_rate": {"base": args.get("base"), "quote": args.get("quote"), "rate": 156.7},
        "convert_currency": {"amount": args.get("amount"), "from": args.get("from_currency"),
                             "to": args.get("to_currency"), "result": 15670.0, "rate": 156.7},
        "search_news": {"results": [
            {"title": "Apple shares slip on iPhone demand concerns", "source": "Reuters"},
            {"title": "Analysts weigh in on AAPL pullback", "source": "Bloomberg"}]},
        "arxiv_search": {"results": [
            {"id": "2406.00001", "title": "Efficient Transformers Revisited",
             "pdf": "https://arxiv.org/pdf/2406.00001"},
            {"id": "2406.00002", "title": "Sparse Attention Transformers",
             "pdf": "https://arxiv.org/pdf/2406.00002"},
            {"id": "2406.00003", "title": "Transformer Scaling Laws 2024",
             "pdf": "https://arxiv.org/pdf/2406.00003"}]},
        "download_file": {"saved": args.get("path"), "bytes": 482113, "status": "ok"},
        "github_list_contributors": {"contributors": [
            {"login": "alice", "commits": 1240}, {"login": "bob", "commits": 830},
            {"login": "carol", "commits": 617}]},
        "code_interpreter": {"stdout": "chart saved to /tmp/contrib.png", "status": "ok"},
        "render_chart": {"chart": "/tmp/contrib.png", "status": "ok"},
        "get_weather_forecast": {"location": args.get("location"),
                                 "forecast": [{"day": "Sun", "cond": "Sunny", "high": 31}]},
        "get_current_weather": {"location": args.get("location"), "cond": "Clear", "temp": 28},
        "create_calendar_event": {"event": args.get("title"), "status": "created"},
    }
    if name in canned:
        return json.dumps(canned[name], ensure_ascii=False)
    return json.dumps({"tool": name, "args": args, "status": "ok",
                       "result": f"<{name} 的 mock 结果>"}, ensure_ascii=False)


# 所有工具共用一个 mock 分发器
TOOL_IMPLS: Dict[str, callable] = {
    name: (lambda args, n=name: _mock_result(n, args)) for name in TOOLS_BY_NAME
}


# ---------------------------------------------------------------------------
# 评测任务及判分标准
# ---------------------------------------------------------------------------
# required_slots: List[List[str]]
#   每个内层 list 是"一个能力槽位"的可接受工具集合（任一命中即算填上该槽位）。
#   一个任务判为"选对"，当且仅当所有槽位都被填上。
# 这些任务都需要跨领域协作，且都存在"通用工具易被误选"的陷阱。

TASKS: List[Dict] = [
    {
        "id": "finance+news",
        "prompt": "苹果公司最近股价怎么样？帮我看看有没有相关新闻能解释一下原因。",
        "required_slots": [
            ["get_stock_price"],
            ["search_news", "get_top_headlines", "get_news_by_source"],
        ],
    },
    {
        "id": "arxiv+download",
        "prompt": "我想看看 transformer 领域最新的研究论文，帮我找几篇最新的，并把排在前三的下载下来。",
        "required_slots": [
            ["arxiv_search"],
            ["download_file"],
        ],
    },
    {
        "id": "github+viz",
        "prompt": "帮我看看 pytorch/pytorch 这个仓库都有谁贡献最多，最好能画个各人提交量的图表。",
        "required_slots": [
            ["github_list_contributors"],
            ["code_interpreter", "render_chart"],
        ],
    },
    {
        "id": "weather+calendar",
        "prompt": "这周日北京天气怎么样？要是晴天的话，帮我在日历里记一个'户外徒步'的活动。",
        "required_slots": [
            ["get_weather_forecast"],
            ["create_calendar_event"],
        ],
    },
    {
        "id": "forex+weather",
        "prompt": "100 美元现在能换多少日元？顺便告诉我东京现在的天气怎么样。",
        "required_slots": [
            ["get_forex_rate", "convert_currency"],
            ["get_current_weather"],
        ],
    },
    {
        "id": "crypto+news",
        "prompt": "以太坊现在多少钱一个？另外有什么最新的相关消息吗？",
        "required_slots": [
            ["get_crypto_price"],
            ["search_news", "get_top_headlines", "get_news_by_source"],
        ],
    },
    # 下面两个是"通用工具诱导"任务：措辞偏泛，容易让模型误用 web_search 等通用兜底工具，
    # 而其实存在更合适的专用工具。用来体现"全量注入错选通用工具、主动发现选对专用工具"。
    {
        "id": "opinion(诱导)",
        "prompt": "帮我了解一下特斯拉这家公司最近的新闻舆论风向。",
        "required_slots": [
            ["search_news", "get_news_by_source", "get_top_headlines"],
        ],
    },
    {
        "id": "academic(诱导)",
        "prompt": "帮我了解一下最近'量子计算'方面有什么新的科研进展。",
        "required_slots": [
            ["arxiv_search", "semantic_scholar_search", "search_pubmed"],
        ],
    },
]


def grade(task: Dict, called_tools: List[str]) -> Dict:
    """根据实际调用的工具给某个任务打分。"""
    called = set(called_tools)
    filled = []
    missed = []
    for slot in task["required_slots"]:
        if any(t in called for t in slot):
            filled.append(slot)
        else:
            missed.append(slot)
    used_generic = sorted(called & GENERIC_TOOL_NAMES)
    correct = len(missed) == 0
    return {
        "correct": correct,                       # 是否覆盖了全部能力槽位
        # 精确选对 = 覆盖全部能力槽位 且 没有误用通用兜底工具（web_search 等）
        "precise": correct and not used_generic,
        "filled_slots": len(filled),
        "total_slots": len(task["required_slots"]),
        "missed_slots": missed,
        "used_generic_substitute": used_generic,
    }


if __name__ == "__main__":
    print(f"工具总数: {len(ALL_TOOLS)}")
    print(f"基础工具: {BASE_TOOL_NAMES}")
    print(f"任务数: {len(TASKS)}")