Video-Tree-TRM5/app/search/tools.py

"""搜索 Agent 工具调度器 — 工具描述与 dispatch 分发。

实现 ``core/agent/protocols.ToolDispatcher`` Protocol。
连接 TreeEnvironment（数据）、summarizer（LLM 摘要）、
vision（VLM 观察）和 skills（策略加载）。

与 TRM4 ``core/tree/tools.py`` 的差异:
- 自由函数 ``dispatch()`` → ``SearchToolDispatcher`` 类（依赖注入）;
- 同步 → 全异步;
- view_node / search_similar 内部拆分为 env 数据读取 + summarizer LLM 摘要。
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from app.search.summarizer import summarize_children, summarize_node, summarize_nodes_batch
from app.search.vision import observe_frame
from app.tree.environment import _LEVEL_LABEL, TreeEnvironment, _node_level

if TYPE_CHECKING:
    from collections.abc import Callable
    from pathlib import Path

    import numpy as np

    from app.ports import OCRProvider
    from app.search.skills import SkillRegistry
    from core.protocols import LLMProvider, VLMProvider

# ── 工具描述文本（与 TRM4 core/tree/tools.py 完全一致） ─────────────────

_BASE_DESCRIPTIONS = """\
## 可用工具

在 action 中指定 tool 和 args 来调用工具。

### view_node
查看节点信息，获取与问题相关的内容摘要和子节点概览。
- args: {"node_id": "节点 ID", "question": "当前关注的具体问题"}

### search_similar
语义检索最相关的节点，返回与问题相关的内容摘要。
- args: {"query": "搜索关键词（2-4 词）", "question": "当前关注的具体问题", "k": 返回数量（可选，默认 5）}

### observe_frame
调用视觉模型查看关键帧图像，回答针对性的视觉问题。
- args: {"node_ids": ["L3 节点 ID 列表（1-4 个），或单个 L2 节点 ID"], "question": "针对帧内容的具体视觉问题"}

### submit_answer
提交最终答案。
- args: {"answer": "选项字母 A/B/C/D", "evidence": "关键证据摘要", "reasoning": "每个选项的判断理由"}"""

_SKILL_DESCRIPTION = """

### read_skill
加载指定题型技能的详细搜索策略。
- args: {"name": "技能名称"}"""


def get_tool_descriptions(include_read_skill: bool = False) -> str:
    """返回工具描述文本，用于写入 system prompt。

    参数:
        include_read_skill: 是否包含 read_skill 工具（manual 模式用）。

    返回:
        Markdown 格式的工具描述文本。
    """
    text = _BASE_DESCRIPTIONS
    if include_read_skill:
        text += _SKILL_DESCRIPTION
    return text


# ── SearchToolDispatcher ──────────────────────────────────────────────


class SearchToolDispatcher:
    """搜索 Agent 工具调度器，实现 ToolDispatcher Protocol。

    按工具名路由到对应私有处理方法。未知工具抛 ValueError
    （AgentLoop 捕获后不计步数）；节点不存在等运行时错误
    捕获后返回错误文本。

    参数:
        env: 视频树运行时环境（纯数据访问）。
        tool_llm: 摘要用 LLM 端口。
        vlm: 视觉模型端口。
        ocr: 帧文字转录端口（None 不启用）。
        prompts_dir: prompt 文件目录。
        skills: 技能注册表（None 不启用 read_skill）。
        embed_fn: 文本嵌入函数（search_similar 用）。
        verify_vision: observe_frame 是否执行验证轮。
        anchor: view_node 是否启用行号锚模式。
        assemble_mode: 锚模式装配形态（"ids"/"ids_expand"/"expand_only"）。
        stats_sink: 统计回调（None 不收集）。
    """

    def __init__(
        self,
        env: TreeEnvironment,
        tool_llm: LLMProvider,
        vlm: VLMProvider,
        ocr: OCRProvider | None,
        prompts_dir: Path,
        skills: SkillRegistry | None,
        *,
        embed_fn: Callable[[str | list[str]], np.ndarray],
        verify_vision: bool,
        anchor: bool,
        assemble_mode: str,
        stats_sink: Callable[[dict[str, Any]], None] | None = None,
    ) -> None:
        self._env = env
        self._tool_llm = tool_llm
        self._vlm = vlm
        self._ocr = ocr
        self._prompts_dir = prompts_dir
        self._skills = skills
        self._embed_fn = embed_fn
        self._verify_vision = verify_vision
        self._anchor = anchor
        self._assemble_mode = assemble_mode
        self._stats_sink = stats_sink

    # ── ToolDispatcher Protocol 实现 ──────────────────────────────────

    async def dispatch(
        self, tool_name: str, args: dict[str, Any], *, context: dict[str, Any]
    ) -> str:
        """按工具名分发到对应处理方法。

        参数:
            tool_name: 工具名称。
            args: 工具参数字典。
            context: 调用上下文（含 session_id、parent_call_id 等遥测字段）。

        返回:
            工具执行结果文本。

        异常:
            ValueError: 未知工具名——上抛给 AgentLoop，不计步数。
        """
        try:
            if tool_name == "view_node":
                return await self._handle_view_node(args, context)
            if tool_name == "search_similar":
                return await self._handle_search_similar(args, context)
            if tool_name == "observe_frame":
                return await self._handle_observe_frame(args, context)
            if tool_name == "submit_answer":
                return f"[ok] 答案已提交: {args['answer']}"
            if tool_name == "read_skill":
                return self._handle_read_skill(args)
        except (KeyError, FileNotFoundError) as e:
            return f"工具执行错误: {e}"

        raise ValueError(f"未知工具: {tool_name}")

    # ── 私有处理方法 ──────────────────────────────────────────────────

    async def _handle_view_node(self, args: dict[str, Any], context: dict[str, Any]) -> str:
        """view_node：节点摘要 + 子节点概览。

        参数:
            args: {"node_id": str, "question": str}。
            context: 遥测上下文。

        返回:
            "[节点] {id} | {level} | {time}\\n\\n{summary}\\n\\n[子节点概览] ..."
        """
        node_id: str = args["node_id"]
        question: str = args["question"]
        session_id = context.get("session_id")
        parent_call_id = context.get("parent_call_id")

        # Phase 1: 节点元数据（头部格式化）
        node = self._env._id_to_node[node_id]
        level = _node_level(node)
        level_label = _LEVEL_LABEL[level]
        time_str = TreeEnvironment._format_time_range(node)

        # Phase 2: 节点内容摘要
        raw_text, anchor_map = self._env.get_node_text(node_id, anchor=self._anchor)
        summary = await summarize_node(
            self._tool_llm,
            raw_text,
            question,
            self._prompts_dir,
            anchor_map=anchor_map,
            assemble_mode=self._assemble_mode,
            stats_sink=self._stats_sink,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )

        parts: list[str] = [
            f"[节点] {node_id} | {level_label} | {time_str}",
            "",
            summary,
        ]

        # Phase 3: 子节点概览
        children_info = self._env.get_children_info(node_id)
        if children_info:
            children_text = await summarize_children(
                self._tool_llm,
                children_info,
                question,
                self._prompts_dir,
                session_id=session_id,
                parent_call_id=parent_call_id,
            )
            parts.append(f"\n[子节点概览] {len(children_info)} 个子节点\n{children_text}")

        return "\n".join(parts)

    async def _handle_search_similar(self, args: dict[str, Any], context: dict[str, Any]) -> str:
        """search_similar：语义检索 + 批量摘要。

        参数:
            args: {"query": str, "question": str, "k": int (可选)}。
            context: 遥测上下文。

        返回:
            "[搜索结果] 查询 \\"{query}\\" → N 个相关节点\\n\\n1. ..."
        """
        query: str = args["query"]
        question: str = args["question"]
        top_k: int = args.get("k", 5)
        session_id = context.get("session_id")
        parent_call_id = context.get("parent_call_id")

        # Phase 1: 语义检索
        results = self._env.search_similar(query, top_k=top_k, embed_fn=self._embed_fn)

        if not results:
            return f'[搜索结果] 查询 "{query}" → 0 个相关节点'

        # Phase 2: 构建摘要输入
        items: list[tuple[str, str, str]] = []
        for nid, score in results:
            node = self._env._id_to_node[nid]
            raw_text, _ = self._env.get_node_text(nid)
            level = _node_level(node)
            time_str = TreeEnvironment._format_time_range(node)
            extra = f"{level}  score={score:.4f}  [{time_str}]"
            items.append((nid, raw_text, extra))

        # Phase 3: 并发批量摘要
        summaries = await summarize_nodes_batch(
            self._tool_llm,
            items,
            question,
            self._prompts_dir,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )

        # Phase 4: 格式化输出
        lines: list[str] = []
        for i, (nid, summary_text) in enumerate(summaries):
            _, _, extra = items[i]
            lines.append(f"{i + 1}. {nid} | {extra}\n   {summary_text}")

        header = f'[搜索结果] 查询 "{query}" → {len(results)} 个相关节点'
        return header + "\n\n" + "\n\n".join(lines)

    async def _handle_observe_frame(self, args: dict[str, Any], context: dict[str, Any]) -> str:
        """observe_frame：VLM 帧观察 + 字幕前置。

        参数:
            args: {"node_ids": list[str], "question": str}。
            context: 遥测上下文。

        返回:
            "[字幕上下文] ...\\n[视觉观察] ..." 或 "[视觉观察] ..."
        """
        node_ids: list[str] = args["node_ids"]
        question: str = args.get("question", "")
        session_id = context.get("session_id")
        parent_call_id = context.get("parent_call_id")

        if not question.strip():
            return "工具执行错误: question 不能为空"

        # Phase 1: 解析帧路径和字幕
        frame_paths = self._env.resolve_frame_paths(node_ids)
        subtitle = self._env.get_subtitle(node_ids[0])

        # Phase 2: VLM 调用
        result = await observe_frame(
            self._vlm,
            frame_paths,
            question,
            self._prompts_dir,
            ocr=self._ocr,
            verify=self._verify_vision,
            stats_sink=self._stats_sink,
            session_id=session_id,
            parent_call_id=parent_call_id,
        )

        # Phase 3: 字幕前置拼接
        if subtitle:
            return f"[字幕上下文] {subtitle}\n{result}"
        return result

    def _handle_read_skill(self, args: dict[str, Any]) -> str:
        """read_skill：加载指定技能的搜索策略正文。

        参数:
            args: {"name": str}。

        返回:
            技能正文或错误提示。
        """
        if self._skills is None:
            return "错误: skills 未启用"
        return self._skills.read(args["name"])