diff --git a/app/search/prompt.py b/app/search/prompt.py new file mode 100644 index 0000000..4435dc0 --- /dev/null +++ b/app/search/prompt.py @@ -0,0 +1,124 @@ +"""搜索 Agent 提示词管理模块。 + +提供 PromptManager 类,统一管理循环级 prompt 的加载与组装。 +工具级 prompt(extract/verify)不在管理范围内。 + +与 TRM4 ``core/search/prompt.py`` 的差异: +- 工具描述从 ``app.search.tools.get_tool_descriptions`` 获取(路径变更); +- ``format_user_prompt`` 参数显式化(question/options/l1_node_ids/task_type)。 +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from app.search.tools import get_tool_descriptions + +if TYPE_CHECKING: + from pathlib import Path + + +class PromptManager: + """管理循环级 prompt 的加载与组装。 + + 构造时缓存 system.md 作为 inference 基础模板。 + 后续步骤(diagnose/evolve/question_gen)通过 load() 按文件名读取。 + + 参数: + prompts_dir: prompt 文件目录的绝对路径。 + """ + + def __init__(self, prompts_dir: Path) -> None: + self._prompts_dir = prompts_dir + system_path = prompts_dir / "system.md" + if not system_path.exists(): + raise FileNotFoundError(f"system.md 不存在: {system_path}") + self._system_base = system_path.read_text(encoding="utf-8") + + def build_inference_prompt( + self, + skill_mode: str, + task_type: str, + always_skills_text: str, + task_skill_map: dict[str, str], + catalog_text: str, + ) -> str: + """组装 inference 步骤的完整 system prompt。 + + 参数: + skill_mode: "auto" / "manual" / "none"。 + task_type: 当前 QA 的题型。 + always_skills_text: always 层 skill 正文(已拼接)。 + task_skill_map: {task_type: skill_body} 映射。 + catalog_text: manual 模式的 skill 目录文本。 + + 返回: + 拼装后的完整 system prompt。 + """ + include_read_skill = skill_mode == "manual" + parts = [ + self._system_base, + f"\n\n---\n\n{get_tool_descriptions(include_read_skill=include_read_skill)}", + ] + if always_skills_text: + parts.append(f"\n\n---\n\n# 通用搜索策略\n\n{always_skills_text}") + if skill_mode == "auto": + skill_text = task_skill_map.get(task_type) or task_skill_map.get("_default") + if skill_text: + parts.append(f"\n\n---\n\n# 当前题型搜索策略\n\n{skill_text}") + elif skill_mode == "manual": + if catalog_text: + parts.append( + "\n\n---\n\n# 可用搜索策略\n\n" + "以下技能扩展了你的导航能力。当问题匹配某技能的适用题型时," + "用 read_skill 工具加载该技能,然后按其指引操作。\n\n" + f"{catalog_text}" + ) + return "".join(parts) + + def format_user_prompt( + self, + question: str, + options: list[str], + l1_node_ids: list[str], + task_type: str | None = None, + ) -> str: + """格式化 inference 步骤的用户提示词。 + + 参数: + question: 问题文本。 + options: 选项列表(如 ["A. 历史", "B. 科学"])。 + l1_node_ids: L1 根节点 ID 列表(如 ["L1_000", "L1_001"])。 + task_type: 可选题型标签,非 None 时插入题型行(oracle 实验用)。 + + 返回: + 格式化后的用户提示词。 + """ + options_text = "\n".join(options) + roots_text = ", ".join(l1_node_ids) + task_type_line = f"**题型**: {task_type}\n" if task_type else "" + return ( + f"请回答以下关于这个视频的多选题:\n\n" + f"{task_type_line}" + f"**问题**: {question}\n" + f"**选项**:\n{options_text}\n\n" + f"**视频树 L1 根节点**: {roots_text}\n" + f"请从以上 L1 节点开始导航,收集证据后回答。" + ) + + def load(self, name: str) -> str: + """按文件名加载 prompt 内容。 + + 参数: + name: prompt 文件名(如 "diagnose_span.md")。 + + 返回: + 文件内容字符串。 + + 异常: + FileNotFoundError: 文件不存在。 + """ + path = self._prompts_dir / name + if not path.exists(): + raise FileNotFoundError(f"prompt 文件不存在: {path}") + return path.read_text(encoding="utf-8") diff --git a/tests/unit/test_search_prompt.py b/tests/unit/test_search_prompt.py new file mode 100644 index 0000000..6083428 --- /dev/null +++ b/tests/unit/test_search_prompt.py @@ -0,0 +1,259 @@ +"""app/search/prompt 模块的单元测试。 + +覆盖 PromptManager 的 __init__、build_inference_prompt、format_user_prompt、load。 +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from unittest.mock import patch + +import pytest + +from app.search.prompt import PromptManager + +if TYPE_CHECKING: + from pathlib import Path + + +# ── 辅助 fixture ───────────────────────────────────────────────────── + + +@pytest.fixture() +def prompts_dir(tmp_path: Path) -> Path: + """创建包含 system.md 的临时 prompt 目录。""" + system_md = tmp_path / "system.md" + system_md.write_text("你是搜索 Agent。", encoding="utf-8") + return tmp_path + + +@pytest.fixture() +def manager(prompts_dir: Path) -> PromptManager: + """构造一个 PromptManager 实例。""" + return PromptManager(prompts_dir) + + +# ── __init__ ────────────────────────────────────────────────────────── + + +class TestInit: + """PromptManager 构造函数测试集。""" + + def test_load_system_md(self, prompts_dir: Path) -> None: + """构造时应成功加载 system.md 内容。""" + mgr = PromptManager(prompts_dir) + assert mgr._system_base == "你是搜索 Agent。" + + def test_missing_system_md_raises(self, tmp_path: Path) -> None: + """system.md 不存在时应抛出 FileNotFoundError。""" + with pytest.raises(FileNotFoundError, match="system.md"): + PromptManager(tmp_path) + + +# ── build_inference_prompt ──────────────────────────────────────────── + + +class TestBuildInferencePrompt: + """build_inference_prompt 三种 skill_mode 的测试集。""" + + def _build( + self, + manager: PromptManager, + *, + skill_mode: str = "none", + task_type: str = "qa", + always_skills_text: str = "通用策略正文", + task_skill_map: dict[str, str] | None = None, + catalog_text: str = "", + ) -> str: + """build_inference_prompt 的便捷包装。""" + if task_skill_map is None: + task_skill_map = {} + with patch( + "app.search.prompt.get_tool_descriptions", + return_value="[工具描述]", + ): + return manager.build_inference_prompt( + skill_mode=skill_mode, + task_type=task_type, + always_skills_text=always_skills_text, + task_skill_map=task_skill_map, + catalog_text=catalog_text, + ) + + def test_auto_mode_appends_task_skill(self, manager: PromptManager) -> None: + """auto 模式应追加 always + 匹配 task_type 的 skill 正文。""" + result = self._build( + manager, + skill_mode="auto", + task_type="qa", + task_skill_map={"qa": "QA 策略正文"}, + ) + assert "你是搜索 Agent。" in result + assert "[工具描述]" in result + assert "通用搜索策略" in result + assert "通用策略正文" in result + assert "当前题型搜索策略" in result + assert "QA 策略正文" in result + + def test_auto_mode_falls_back_to_default(self, manager: PromptManager) -> None: + """auto 模式中 task_type 无匹配时回退到 _default。""" + result = self._build( + manager, + skill_mode="auto", + task_type="unknown_type", + task_skill_map={"_default": "默认策略"}, + ) + assert "当前题型搜索策略" in result + assert "默认策略" in result + + def test_auto_mode_no_match_no_default(self, manager: PromptManager) -> None: + """auto 模式中 task_type 无匹配且无 _default 时不追加题型策略段。""" + result = self._build( + manager, + skill_mode="auto", + task_type="unknown_type", + task_skill_map={}, + ) + assert "当前题型搜索策略" not in result + # 但 always 仍在 + assert "通用搜索策略" in result + + def test_manual_mode_appends_catalog(self, manager: PromptManager) -> None: + """manual 模式应追加 always + catalog 目录文本。""" + result = self._build( + manager, + skill_mode="manual", + catalog_text="- skill_a\n- skill_b", + ) + assert "通用搜索策略" in result + assert "可用搜索策略" in result + assert "read_skill" in result + assert "- skill_a" in result + + def test_manual_mode_include_read_skill(self, manager: PromptManager) -> None: + """manual 模式应传 include_read_skill=True 给 get_tool_descriptions。""" + with patch( + "app.search.prompt.get_tool_descriptions", + return_value="[工具描述]", + ) as mock_get: + manager.build_inference_prompt( + skill_mode="manual", + task_type="qa", + always_skills_text="", + task_skill_map={}, + catalog_text="目录", + ) + mock_get.assert_called_once_with(include_read_skill=True) + + def test_auto_mode_include_read_skill_false(self, manager: PromptManager) -> None: + """auto 模式应传 include_read_skill=False 给 get_tool_descriptions。""" + with patch( + "app.search.prompt.get_tool_descriptions", + return_value="[工具描述]", + ) as mock_get: + manager.build_inference_prompt( + skill_mode="auto", + task_type="qa", + always_skills_text="", + task_skill_map={}, + catalog_text="", + ) + mock_get.assert_called_once_with(include_read_skill=False) + + def test_none_mode_only_base_and_tools(self, manager: PromptManager) -> None: + """none 模式应仅包含 base + 工具描述 + always(若有)。""" + result = self._build( + manager, + skill_mode="none", + always_skills_text="通用策略正文", + catalog_text="不应出现", + ) + assert "你是搜索 Agent。" in result + assert "[工具描述]" in result + assert "通用搜索策略" in result + assert "可用搜索策略" not in result + assert "当前题型搜索策略" not in result + + def test_none_mode_empty_always(self, manager: PromptManager) -> None: + """none 模式下 always_skills_text 为空时不追加通用策略段。""" + result = self._build( + manager, + skill_mode="none", + always_skills_text="", + ) + assert "通用搜索策略" not in result + + +# ── format_user_prompt ──────────────────────────────────────────────── + + +class TestFormatUserPrompt: + """format_user_prompt 的测试集。""" + + def test_with_task_type(self, manager: PromptManager) -> None: + """指定 task_type 时应在输出中插入题型行。""" + result = manager.format_user_prompt( + question="这段视频讲了什么?", + options=["A. 历史", "B. 科学", "C. 艺术", "D. 体育"], + l1_node_ids=["L1_000", "L1_001"], + task_type="qa", + ) + assert "**题型**: qa" in result + assert "**问题**: 这段视频讲了什么?" in result + assert "A. 历史" in result + assert "D. 体育" in result + assert "L1_000, L1_001" in result + assert "请从以上 L1 节点开始导航" in result + + def test_without_task_type(self, manager: PromptManager) -> None: + """task_type 为 None 时输出不应包含题型行。""" + result = manager.format_user_prompt( + question="问题内容", + options=["A. 选项1", "B. 选项2"], + l1_node_ids=["L1_000"], + ) + assert "**题型**" not in result + assert "**问题**: 问题内容" in result + assert "L1_000" in result + + def test_options_each_on_own_line(self, manager: PromptManager) -> None: + """每个选项应独占一行。""" + result = manager.format_user_prompt( + question="Q", + options=["A. 一", "B. 二", "C. 三"], + l1_node_ids=["L1_000"], + ) + lines = result.split("\n") + # 选项应连续出现在各自行上 + option_lines = [line for line in lines if line.startswith(("A.", "B.", "C."))] + assert len(option_lines) == 3 + + def test_single_l1_node(self, manager: PromptManager) -> None: + """单个 L1 节点时根节点文本不含逗号。""" + result = manager.format_user_prompt( + question="Q", + options=["A. x"], + l1_node_ids=["L1_000"], + ) + assert "**视频树 L1 根节点**: L1_000" in result + assert "," not in result.split("根节点**: ")[1].split("\n")[0] + + +# ── load ────────────────────────────────────────────────────────────── + + +class TestLoad: + """load 方法的测试集。""" + + def test_load_existing_file(self, prompts_dir: Path) -> None: + """加载存在的 prompt 文件应返回其内容。""" + (prompts_dir / "diagnose_span.md").write_text("诊断模板内容", encoding="utf-8") + mgr = PromptManager(prompts_dir) + content = mgr.load("diagnose_span.md") + assert content == "诊断模板内容" + + def test_load_missing_file_raises(self, manager: PromptManager) -> None: + """加载不存在的 prompt 文件应抛出 FileNotFoundError。""" + with pytest.raises(FileNotFoundError, match="not_exist.md"): + manager.load("not_exist.md")