feat(search): PromptManager — 搜索 Agent 提示词加载与组装
从 TRM4 core/search/prompt.py 迁移。有意变更: - 工具描述从 app.search.tools.get_tool_descriptions 获取 - format_user_prompt 参数显式化(question/options/l1_node_ids/task_type) 16 个单元测试覆盖 __init__、build_inference_prompt(auto/manual/none 三种 skill_mode)、format_user_prompt、load。 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -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")
|
||||||
@@ -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")
|
||||||
Reference in New Issue
Block a user