lambdagent/agentexample/agent67v2/tools/tool_search.py

"""
ToolSearch — 工具懒加载元工具

借鉴 Claude Code 的 deferred tools 模式:
  不在 system prompt 中注入全部工具描述,
  而是提供一个 ToolSearch 元工具, LLM 需要时才查找具体工具。

Lambda 语义:
  ToolSearch = λquery. filter(Γ_tools, query)
  效果: system prompt 从 42 个工具描述缩减为 1 个元工具描述
  → token 节省 ~60%, 首次响应快 2-3x

用法:
  LLM 输出: {"action": "ToolSearch", "input": {"query": "文件操作"}}
  返回: 匹配的工具列表 + 简要描述, LLM 在下一步选择具体工具
"""
from __future__ import annotations

import json
from typing import Any, Dict, List, Optional

from lambdagent.core import Term, Context


# ════════════════════════════════════════════════════════════
# 工具目录: 所有可用工具的元数据 (不含实际实现)
# ════════════════════════════════════════════════════════════

TOOL_CATALOG: Dict[str, Dict[str, Any]] = {
    # ── 文件操作 (5) ──
    "ReadFile": {
        "category": "file",
        "description": "读取文件内容",
        "input_schema": {"path": "str, 文件路径", "offset": "int, 起始行(可选)", "limit": "int, 行数(可选)"},
        "tags": ["文件", "读取", "查看", "file", "read"],
    },
    "EditFile": {
        "category": "file",
        "description": "编辑文件: 查找并替换指定文本片段",
        "input_schema": {"path": "str", "old_text": "str", "new_text": "str"},
        "tags": ["文件", "编辑", "修改", "替换", "file", "edit"],
    },
    "WriteFile": {
        "category": "file",
        "description": "写入/创建文件 (覆盖写入)",
        "input_schema": {"path": "str", "content": "str"},
        "tags": ["文件", "写入", "创建", "保存", "file", "write", "create"],
    },
    "ListFiles": {
        "category": "file",
        "description": "列出目录下的文件和子目录",
        "input_schema": {"path": "str, 目录路径", "pattern": "str, glob模式(可选)"},
        "tags": ["文件", "列表", "目录", "ls", "file", "list", "directory"],
    },
    "SearchContent": {
        "category": "file",
        "description": "在文件内容中搜索关键词 (grep)",
        "input_schema": {"pattern": "str, 正则表达式", "path": "str, 搜索目录", "glob": "str, 文件过滤(可选)"},
        "tags": ["搜索", "grep", "内容", "查找", "search", "find"],
    },

    # ── 代码分析 (3) ──
    "CodeSearch": {
        "category": "code",
        "description": "在代码库中搜索函数/类/变量定义",
        "input_schema": {"query": "str, 搜索词", "language": "str, 语言(可选)"},
        "tags": ["代码", "搜索", "函数", "类", "code", "search", "definition"],
    },
    "ProjectMap": {
        "category": "code",
        "description": "生成项目结构图和代码地图",
        "input_schema": {"path": "str, 项目根目录"},
        "tags": ["项目", "结构", "地图", "架构", "project", "structure", "map"],
    },
    "RunTests": {
        "category": "code",
        "description": "运行测试用例 (pytest/unittest/maven test)",
        "input_schema": {"path": "str, 测试文件或目录", "framework": "str, 框架(可选)"},
        "tags": ["测试", "运行", "pytest", "test", "run"],
    },

    # ── Shell (1) ──
    "Bash": {
        "category": "shell",
        "description": "执行 Shell 命令 (支持持久 CWD 和后台执行)",
        "input_schema": {"command": "str, 命令", "timeout": "int, 超时秒数(可选)", "background": "bool(可选)"},
        "tags": ["终端", "命令", "shell", "bash", "terminal", "命令行"],
    },

    # ── Git (5) ──
    "GitStatus": {
        "category": "git",
        "description": "查看 Git 工作区状态 (git status)",
        "input_schema": {"path": "str, 仓库路径(可选)"},
        "tags": ["git", "状态", "status"],
    },
    "GitDiff": {
        "category": "git",
        "description": "查看 Git 差异 (git diff)",
        "input_schema": {"path": "str(可选)", "staged": "bool(可选)"},
        "tags": ["git", "差异", "diff", "变更"],
    },
    "GitLog": {
        "category": "git",
        "description": "查看 Git 提交历史 (git log)",
        "input_schema": {"path": "str(可选)", "limit": "int(可选)"},
        "tags": ["git", "历史", "log", "提交"],
    },
    "GitCommit": {
        "category": "git",
        "description": "创建 Git 提交 (git add + commit)",
        "input_schema": {"message": "str", "files": "list[str](可选)"},
        "tags": ["git", "提交", "commit"],
    },
    "GitBranch": {
        "category": "git",
        "description": "Git 分支操作 (create/switch/list/delete)",
        "input_schema": {"action": "str, create|switch|list|delete", "name": "str(可选)"},
        "tags": ["git", "分支", "branch"],
    },

    # ── Web (2) ──
    "WebSearch": {
        "category": "web",
        "description": "搜索引擎查询 (返回摘要和链接)",
        "input_schema": {"query": "str, 搜索词"},
        "tags": ["网络", "搜索", "web", "search", "google", "查询"],
    },
    "WebFetch": {
        "category": "web",
        "description": "获取网页内容 (HTML → Markdown)",
        "input_schema": {"url": "str, 网址"},
        "tags": ["网络", "网页", "抓取", "web", "fetch", "url"],
    },

    # ── Notebook & 文档 (4) ──
    "NotebookEdit": {
        "category": "doc",
        "description": "编辑 Jupyter Notebook 单元格",
        "input_schema": {"path": "str", "cell_index": "int", "content": "str", "cell_type": "str(可选)"},
        "tags": ["notebook", "jupyter", "编辑", "cell"],
    },
    "DocGen": {
        "category": "doc",
        "description": "文档生成 (Markdown → HTML/PDF/Word)",
        "input_schema": {"content": "str", "format": "str, html|pdf|docx", "output_path": "str"},
        "tags": ["文档", "生成", "导出", "doc", "pdf", "word"],
    },
    "ChunkSplit": {
        "category": "doc",
        "description": "文本分块 (用于 RAG 或长文处理)",
        "input_schema": {"text": "str", "chunk_size": "int(可选)", "overlap": "int(可选)"},
        "tags": ["分块", "切分", "chunk", "split", "RAG"],
    },
    "OCR": {
        "category": "doc",
        "description": "图片文字识别 (OCR)",
        "input_schema": {"image_path": "str"},
        "tags": ["OCR", "图片", "文字", "识别", "image"],
    },

    # ── 知识库 (4) ──
    "KBCreate": {
        "category": "knowledge",
        "description": "创建知识库",
        "input_schema": {"name": "str", "description": "str(可选)"},
        "tags": ["知识库", "创建", "knowledge", "create", "KB"],
    },
    "KBAdd": {
        "category": "knowledge",
        "description": "向知识库添加文档",
        "input_schema": {"kb_name": "str", "content": "str", "metadata": "dict(可选)"},
        "tags": ["知识库", "添加", "导入", "knowledge", "add"],
    },
    "KBSearch": {
        "category": "knowledge",
        "description": "在知识库中检索",
        "input_schema": {"kb_name": "str", "query": "str", "top_k": "int(可选)"},
        "tags": ["知识库", "检索", "搜索", "knowledge", "search", "query"],
    },
    "KBList": {
        "category": "knowledge",
        "description": "列出所有知识库",
        "input_schema": {},
        "tags": ["知识库", "列表", "knowledge", "list"],
    },

    # ── 任务/记忆/调度 (10) ──
    "TaskCreate": {
        "category": "task",
        "description": "创建任务",
        "input_schema": {"title": "str", "description": "str(可选)", "priority": "str(可选)"},
        "tags": ["任务", "创建", "todo", "task", "create"],
    },
    "TaskUpdate": {
        "category": "task",
        "description": "更新任务状态",
        "input_schema": {"task_id": "str", "status": "str", "note": "str(可选)"},
        "tags": ["任务", "更新", "task", "update"],
    },
    "TaskList": {
        "category": "task",
        "description": "列出任务",
        "input_schema": {"status": "str(可选)", "limit": "int(可选)"},
        "tags": ["任务", "列表", "task", "list"],
    },
    "MemoryStore": {
        "category": "memory",
        "description": "保存记忆 (持久化存储)",
        "input_schema": {"key": "str", "value": "str", "tags": "list[str](可选)"},
        "tags": ["记忆", "保存", "存储", "memory", "store", "save"],
    },
    "MemoryRecall": {
        "category": "memory",
        "description": "回忆 (按 key 或语义检索)",
        "input_schema": {"query": "str", "limit": "int(可选)"},
        "tags": ["记忆", "回忆", "检索", "memory", "recall", "retrieve"],
    },
    "MemoryList": {
        "category": "memory",
        "description": "列出所有记忆",
        "input_schema": {},
        "tags": ["记忆", "列表", "memory", "list"],
    },
    "MemoryForget": {
        "category": "memory",
        "description": "删除记忆",
        "input_schema": {"key": "str"},
        "tags": ["记忆", "删除", "遗忘", "memory", "forget", "delete"],
    },
    "ScheduleCreate": {
        "category": "schedule",
        "description": "创建定时任务",
        "input_schema": {"cron": "str, cron表达式", "task": "str, 要执行的任务"},
        "tags": ["调度", "定时", "schedule", "cron", "定时任务"],
    },
    "ScheduleList": {
        "category": "schedule",
        "description": "列出定时任务",
        "input_schema": {},
        "tags": ["调度", "列表", "schedule", "list"],
    },
    "ScheduleDelete": {
        "category": "schedule",
        "description": "删除定时任务",
        "input_schema": {"schedule_id": "str"},
        "tags": ["调度", "删除", "schedule", "delete"],
    },

    # ── 系统控制 (4, macOS) ──
    "browser": {
        "category": "system",
        "description": "浏览器控制 (打开URL, 获取标签页, 执行JS)",
        "input_schema": {"action": "str, open|tabs|execute_js", "url": "str(可选)", "script": "str(可选)"},
        "tags": ["浏览器", "browser", "网页", "URL", "系统"],
    },
    "app": {
        "category": "system",
        "description": "应用控制 (启动/切换/关闭 macOS 应用)",
        "input_schema": {"action": "str, launch|switch|quit|list", "name": "str(可选)"},
        "tags": ["应用", "app", "启动", "切换", "系统"],
    },
    "system": {
        "category": "system",
        "description": "系统信息查询 (CPU/内存/磁盘/网络)",
        "input_schema": {"query": "str, cpu|memory|disk|network|all"},
        "tags": ["系统", "信息", "system", "info", "CPU", "内存"],
    },
    "screenshot": {
        "category": "system",
        "description": "屏幕截图",
        "input_schema": {"region": "str, full|window|area(可选)", "save_path": "str(可选)"},
        "tags": ["截图", "screenshot", "屏幕", "系统"],
    },

    # ── 画像 & 学习 (4) ──
    "ProfileGet": {
        "category": "profile",
        "description": "获取用户画像",
        "input_schema": {},
        "tags": ["画像", "profile", "用户", "偏好"],
    },
    "ProfileUpdate": {
        "category": "profile",
        "description": "更新用户画像",
        "input_schema": {"key": "str", "value": "str"},
        "tags": ["画像", "更新", "profile", "update"],
    },
    "LearningFeedback": {
        "category": "profile",
        "description": "提交学习反馈 (帮助agent自我改进)",
        "input_schema": {"feedback": "str", "rating": "int(可选)"},
        "tags": ["学习", "反馈", "learning", "feedback"],
    },
    "LearningStrategies": {
        "category": "profile",
        "description": "查看当前学习策略",
        "input_schema": {},
        "tags": ["学习", "策略", "learning", "strategies"],
    },

    # ── 通知 & 事件 (3) ──
    "Notify": {
        "category": "notify",
        "description": "发送通知 (终端/macOS 通知中心)",
        "input_schema": {"message": "str", "title": "str(可选)", "sound": "bool(可选)"},
        "tags": ["通知", "notify", "提醒", "alert"],
    },
    "EventSubscribe": {
        "category": "notify",
        "description": "订阅事件 (文件变化/定时/webhook)",
        "input_schema": {"event_type": "str", "config": "dict"},
        "tags": ["事件", "订阅", "event", "subscribe", "watch"],
    },
    "EventList": {
        "category": "notify",
        "description": "列出已订阅的事件",
        "input_schema": {},
        "tags": ["事件", "列表", "event", "list"],
    },
}

# 按类别索引
CATEGORIES = {}
for tool_name, meta in TOOL_CATALOG.items():
    cat = meta["category"]
    if cat not in CATEGORIES:
        CATEGORIES[cat] = []
    CATEGORIES[cat].append(tool_name)


# ════════════════════════════════════════════════════════════
# ToolSearch: 懒加载元工具
# ════════════════════════════════════════════════════════════

class ToolSearch(Term):
    """
    工具懒加载元工具。

    Lambda 语义:
        ToolSearch = λquery. filter(TOOL_CATALOG, query)

    不执行任何工具, 只返回匹配的工具描述,
    让 LLM 在下一步选择具体工具。

    搜索策略:
        1. 按类别名精确匹配 (如 "file", "git", "web")
        2. 按工具名精确匹配 (如 "ReadFile", "Bash")
        3. 按 tags 模糊匹配 (如 "搜索", "编辑")
        4. 按 description 模糊匹配
    """

    def __init__(self):
        super().__init__("ToolSearch")

    def apply(self, input: Any, ctx: Context | None = None) -> str:
        ctx = ctx or Context()

        # 解析输入
        if isinstance(input, str):
            try:
                data = json.loads(input)
                query = data.get("query", data.get("keyword", input))
                category = data.get("category", "")
            except (json.JSONDecodeError, AttributeError):
                query = input.strip()
                category = ""
        elif isinstance(input, dict):
            query = input.get("query", input.get("keyword", ""))
            category = input.get("category", "")
        else:
            query = str(input)
            category = ""

        results = self.search(query, category)

        if not results:
            return json.dumps({
                "matched": 0,
                "hint": f"没有找到匹配 '{query}' 的工具。可用类别: {', '.join(sorted(CATEGORIES.keys()))}",
                "categories": {cat: tools for cat, tools in sorted(CATEGORIES.items())},
            }, ensure_ascii=False, indent=2)

        return json.dumps({
            "matched": len(results),
            "tools": results,
            "hint": "请从以上工具中选择一个, 用 JSON 代码块调用它。",
        }, ensure_ascii=False, indent=2)

    def search(self, query: str, category: str = "") -> List[dict]:
        """搜索工具目录"""
        results = []
        query_lower = query.lower()
        query_words = query_lower.split()

        for tool_name, meta in TOOL_CATALOG.items():
            score = 0

            # 按类别精确匹配
            if category and meta["category"] == category:
                score += 10

            # 按工具名匹配
            if query_lower in tool_name.lower():
                score += 8

            # 按类别名匹配
            if query_lower == meta["category"]:
                score += 10

            # 按 tags 匹配
            tags_text = " ".join(meta["tags"]).lower()
            for word in query_words:
                if word in tags_text:
                    score += 3

            # 按 description 匹配
            desc_lower = meta["description"].lower()
            for word in query_words:
                if word in desc_lower:
                    score += 2

            if score > 0:
                results.append({
                    "name": tool_name,
                    "category": meta["category"],
                    "description": meta["description"],
                    "input_schema": meta["input_schema"],
                    "_score": score,
                })

        # 按 score 排序, 返回前 10 个
        results.sort(key=lambda x: -x["_score"])
        for r in results:
            del r["_score"]
        return results[:10]

    @staticmethod
    def list_categories() -> dict:
        """列出所有工具类别"""
        return {
            cat: {
                "count": len(tools),
                "tools": tools,
            }
            for cat, tools in sorted(CATEGORIES.items())
        }

    @staticmethod
    def get_tool_description() -> str:
        """返回 ToolSearch 自身的工具描述 (注入 system prompt)"""
        category_summary = ", ".join(
            f"{cat}({len(tools)})" for cat, tools in sorted(CATEGORIES.items())
        )
        return (
            f"ToolSearch: 搜索可用工具。共 {len(TOOL_CATALOG)} 个工具, "
            f"分 {len(CATEGORIES)} 类: {category_summary}。\n"
            f"输入: {{\"query\": \"关键词\"}} 或 {{\"category\": \"类别名\"}}\n"
            f"返回: 匹配的工具列表 + 调用格式。\n"
            f"你应该先用 ToolSearch 找到需要的工具, 再调用具体工具。"
        )


# 单例
tool_search = ToolSearch()