Video-Tree-TRM5/tests/unit/test_search_tools.py

"""SearchToolDispatcher 与 get_tool_descriptions 单元测试。

验证工具描述生成和五种工具的 dispatch 路由：
view_node、search_similar、observe_frame、submit_answer、read_skill，
以及未知工具 ValueError 和节点不存在错误文本。
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from pathlib import Path

import numpy as np
import pytest

from app.search.skills import SkillRegistry
from app.search.tools import SearchToolDispatcher, get_tool_descriptions
from app.tree.environment import TreeEnvironment
from app.tree.index import (
    IndexMeta,
    L1Card,
    L1Node,
    L2Card,
    L2Node,
    L3Card,
    L3Node,
    TreeIndex,
)
from core.types import LLMResponse

# ── 假实现 ────────────────────────────────────────────────────────────


def _make_llm_response(content: str = "fake summary") -> LLMResponse:
    """构造固定的 LLMResponse 实例。"""
    return LLMResponse(
        content=content,
        thinking="",
        model="fake-model",
        provider="fake",
        prompt_tokens=10,
        completion_tokens=5,
        latency_ms=50,
        ttft_ms=None,
        max_inter_token_ms=None,
        cache_hit=False,
        call_id="fake-call-id",
    )


class FakeLLM:
    """最小 LLMProvider 假实现。"""

    async def chat(
        self,
        messages: list[dict[str, Any]],
        *,
        session_id: str | None = None,
        parent_call_id: str | None = None,
    ) -> LLMResponse:
        """返回固定摘要内容。"""
        return _make_llm_response("fake summary")


class FakeVLM:
    """最小 VLMProvider 假实现。"""

    async def chat_with_images(
        self,
        messages: list[dict[str, Any]],
        images: list[str | Path],
        *,
        session_id: str | None = None,
        parent_call_id: str | None = None,
    ) -> LLMResponse:
        """返回固定视觉观察内容。"""
        return _make_llm_response("fake visual observation")


class FakeOCR:
    """最小 OCRProvider 假实现。"""

    async def transcribe_frames(self, frame_paths: list[Path]) -> str:
        """返回固定 OCR 文本。"""
        return "OCR text"


def _fake_embed_fn(texts: str | list[str]) -> np.ndarray:
    """返回固定维度的 L2 归一化嵌入向量。"""
    if isinstance(texts, str):
        vec = np.ones((1, 4), dtype=np.float32)
    else:
        vec = np.ones((len(texts), 4), dtype=np.float32)
    norms = np.linalg.norm(vec, axis=1, keepdims=True)
    return vec / norms


# ── Fixtures ──────────────────────────────────────────────────────────


def _make_test_tree() -> TreeIndex:
    """构建包含 L1→L2→L3 的最小测试树。"""
    l3 = L3Node(
        id="vid_L1_000_L2_000_L3_000",
        card=L3Card(
            frame_summary="test frame summary",
            visible_entities=["person"],
            ongoing_actions=["walking"],
            visible_text=[],
            spatial_layout="center",
            visual_attributes={},
        ),
        timestamp=10.0,
        frame_path="frames/L1_000_L2_000_L3_000.jpg",
        subtitle="test subtitle text",
    )
    l2 = L2Node(
        id="vid_L1_000_L2_000",
        card=L2Card(
            event_description="test event description",
            entities=["person"],
            actions=["walking"],
            action_subjects=["person"],
            visible_text=[],
            spatial_relations="none",
            state_changes=None,
        ),
        time_range=(5.0, 15.0),
        children=[l3],
    )
    l1 = L1Node(
        id="vid_L1_000",
        card=L1Card(
            scene_summary="test scene summary",
            main_setting="outdoor",
            key_entities=["person"],
            main_actions=["walking"],
            topic_keywords=["outdoor"],
            visible_text=[],
            temporal_flow="linear",
        ),
        time_range=(0.0, 30.0),
        children=[l2],
    )
    return TreeIndex(
        metadata=IndexMeta(source_path="test.mp4", modality="video"),
        roots=[l1],
    )


@pytest.fixture()
def env() -> TreeEnvironment:
    """带最小树的 TreeEnvironment。"""
    return TreeEnvironment(_make_test_tree())


@pytest.fixture()
def prompts_dir(tmp_path: Path) -> Path:
    """在 tmp 目录中创建必需的 prompt 文件。"""
    prompt_files = [
        "view_node_extract.md",
        "view_node_verify.md",
        "view_node_children_extract.md",
        "view_node_children_verify.md",
        "search_similar_extract.md",
        "search_similar_verify.md",
        "observe_frame_extract.md",
        "observe_frame_verify.md",
    ]
    for name in prompt_files:
        (tmp_path / name).write_text(f"fake prompt for {name}", encoding="utf-8")
    return tmp_path


@pytest.fixture()
def skills_registry(tmp_path: Path) -> SkillRegistry:
    """带一个预注册技能的 SkillRegistry。"""
    skill_path = tmp_path / "test_skill.md"
    skill_path.write_text(
        "---\nname: test_skill\ndescription: test\n---\nskill body content",
        encoding="utf-8",
    )
    registry = SkillRegistry()
    registry.set_paths({"test_skill": skill_path})
    return registry


@pytest.fixture()
def dispatcher(
    env: TreeEnvironment,
    prompts_dir: Path,
    skills_registry: SkillRegistry,
) -> SearchToolDispatcher:
    """标准配置的 SearchToolDispatcher 实例。"""
    return SearchToolDispatcher(
        env=env,
        tool_llm=FakeLLM(),
        vlm=FakeVLM(),
        ocr=FakeOCR(),
        prompts_dir=prompts_dir,
        skills=skills_registry,
        embed_fn=_fake_embed_fn,
        verify_vision=False,
        anchor=False,
        assemble_mode="ids",
    )


@pytest.fixture()
def dispatcher_no_skills(
    env: TreeEnvironment,
    prompts_dir: Path,
) -> SearchToolDispatcher:
    """skills=None 的 SearchToolDispatcher 实例。"""
    return SearchToolDispatcher(
        env=env,
        tool_llm=FakeLLM(),
        vlm=FakeVLM(),
        ocr=None,
        prompts_dir=prompts_dir,
        skills=None,
        embed_fn=_fake_embed_fn,
        verify_vision=False,
        anchor=False,
        assemble_mode="ids",
    )


# ── get_tool_descriptions 测试 ───────────────────────────────────────


class TestGetToolDescriptions:
    """get_tool_descriptions 工具描述生成测试。"""

    def test_without_read_skill(self) -> None:
        """不含 read_skill 时应包含四个基础工具。"""
        text = get_tool_descriptions(include_read_skill=False)
        assert "view_node" in text
        assert "search_similar" in text
        assert "observe_frame" in text
        assert "submit_answer" in text
        assert "read_skill" not in text

    def test_with_read_skill(self) -> None:
        """含 read_skill 时应额外包含 read_skill 工具描述。"""
        text = get_tool_descriptions(include_read_skill=True)
        assert "view_node" in text
        assert "read_skill" in text
        assert "加载指定题型技能" in text


# ── dispatch 路由测试 ─────────────────────────────────────────────────


class TestDispatchViewNode:
    """dispatch view_node 工具测试。"""

    @pytest.mark.asyncio()
    async def test_view_node_returns_header_and_summary(
        self, dispatcher: SearchToolDispatcher
    ) -> None:
        """view_node 应返回含节点头部、摘要和子节点概览的文本。"""
        result = await dispatcher.dispatch(
            "view_node",
            {"node_id": "vid_L1_000", "question": "what happens?"},
            context={},
        )
        # 头部格式
        assert "[节点] vid_L1_000 | 场景层 |" in result
        assert "0.0-30.0s" in result
        # 摘要内容（来自 FakeLLM）
        assert "fake summary" in result
        # 子节点概览（L1 有 L2 子节点）
        assert "[子节点概览]" in result
        assert "1 个子节点" in result

    @pytest.mark.asyncio()
    async def test_view_node_l3_no_children(self, dispatcher: SearchToolDispatcher) -> None:
        """L3 叶子节点应无子节点概览段。"""
        result = await dispatcher.dispatch(
            "view_node",
            {"node_id": "vid_L1_000_L2_000_L3_000", "question": "test"},
            context={},
        )
        assert "[节点] vid_L1_000_L2_000_L3_000 | 关键帧层 |" in result
        assert "[子节点概览]" not in result


class TestDispatchSearchSimilar:
    """dispatch search_similar 工具测试。"""

    @pytest.mark.asyncio()
    async def test_search_similar_returns_results(self, dispatcher: SearchToolDispatcher) -> None:
        """search_similar 应返回搜索头部和编号结果列表。"""
        result = await dispatcher.dispatch(
            "search_similar",
            {"query": "walking", "question": "what is the person doing?"},
            context={},
        )
        assert '[搜索结果] 查询 "walking"' in result
        assert "个相关节点" in result
        # 至少有一个编号结果
        assert "1." in result
        # 包含分数信息
        assert "score=" in result

    @pytest.mark.asyncio()
    async def test_search_similar_custom_k(self, dispatcher: SearchToolDispatcher) -> None:
        """search_similar 的 k 参数应限制返回数量。"""
        result = await dispatcher.dispatch(
            "search_similar",
            {"query": "test", "question": "test", "k": 1},
            context={},
        )
        assert "1 个相关节点" in result


class TestDispatchObserveFrame:
    """dispatch observe_frame 工具测试。"""

    @pytest.mark.asyncio()
    async def test_observe_frame_with_subtitle(
        self, dispatcher: SearchToolDispatcher, tmp_path: Path
    ) -> None:
        """有字幕的 L3 节点应在输出前添加字幕上下文。"""
        # 创建帧文件使路径存在检查通过
        frame_file = tmp_path / "L1_000_L2_000_L3_000.jpg"
        frame_file.write_bytes(b"\xff\xd8\xff\xe0")

        # 重建 dispatcher 指定 frames_dir
        tree = _make_test_tree()
        env_with_frames = TreeEnvironment(tree, frames_dir=tmp_path)

        d = SearchToolDispatcher(
            env=env_with_frames,
            tool_llm=FakeLLM(),
            vlm=FakeVLM(),
            ocr=FakeOCR(),
            prompts_dir=dispatcher._prompts_dir,
            skills=None,
            embed_fn=_fake_embed_fn,
            verify_vision=False,
            anchor=False,
            assemble_mode="ids",
        )

        result = await d.dispatch(
            "observe_frame",
            {
                "node_ids": ["vid_L1_000_L2_000_L3_000"],
                "question": "what is visible?",
            },
            context={},
        )
        assert "[字幕上下文] test subtitle text" in result
        assert "fake visual observation" in result

    @pytest.mark.asyncio()
    async def test_observe_frame_empty_question(self, dispatcher: SearchToolDispatcher) -> None:
        """空 question 应返回错误文本。"""
        result = await dispatcher.dispatch(
            "observe_frame",
            {"node_ids": ["vid_L1_000_L2_000_L3_000"], "question": "  "},
            context={},
        )
        assert "question 不能为空" in result


class TestDispatchSubmitAnswer:
    """dispatch submit_answer 工具测试。"""

    @pytest.mark.asyncio()
    async def test_submit_answer_returns_confirmation(
        self, dispatcher: SearchToolDispatcher
    ) -> None:
        """submit_answer 应返回确认文本。"""
        result = await dispatcher.dispatch(
            "submit_answer",
            {"answer": "B", "evidence": "seen in frame", "reasoning": "clear visual"},
            context={},
        )
        assert result == "[ok] 答案已提交: B"


class TestDispatchReadSkill:
    """dispatch read_skill 工具测试。"""

    @pytest.mark.asyncio()
    async def test_read_skill_returns_body(self, dispatcher: SearchToolDispatcher) -> None:
        """read_skill 应返回去除 frontmatter 后的技能正文。"""
        result = await dispatcher.dispatch(
            "read_skill",
            {"name": "test_skill"},
            context={},
        )
        assert "skill body content" in result

    @pytest.mark.asyncio()
    async def test_read_skill_disabled(self, dispatcher_no_skills: SearchToolDispatcher) -> None:
        """skills=None 时 read_skill 应返回未启用提示。"""
        result = await dispatcher_no_skills.dispatch(
            "read_skill",
            {"name": "anything"},
            context={},
        )
        assert result == "错误: skills 未启用"


# ── 错误处理测试 ──────────────────────────────────────────────────────


class TestDispatchErrors:
    """dispatch 错误处理测试。"""

    @pytest.mark.asyncio()
    async def test_unknown_tool_raises_value_error(self, dispatcher: SearchToolDispatcher) -> None:
        """未知工具应抛出 ValueError。"""
        with pytest.raises(ValueError, match="未知工具: nonexistent_tool"):
            await dispatcher.dispatch("nonexistent_tool", {}, context={})

    @pytest.mark.asyncio()
    async def test_node_not_found_returns_error_text(
        self, dispatcher: SearchToolDispatcher
    ) -> None:
        """节点不存在时应返回错误文本（非异常）。"""
        result = await dispatcher.dispatch(
            "view_node",
            {"node_id": "nonexistent_node", "question": "test"},
            context={},
        )
        assert "工具执行错误" in result
        assert "nonexistent_node" in result

    @pytest.mark.asyncio()
    async def test_read_skill_not_found_returns_error_text(
        self, dispatcher: SearchToolDispatcher
    ) -> None:
        """未注册的技能名应返回错误文本。"""
        result = await dispatcher.dispatch(
            "read_skill",
            {"name": "nonexistent_skill"},
            context={},
        )
        assert "工具执行错误" in result