545 lines
20 KiB
Python
545 lines
20 KiB
Python
"""app/search/summarizer 模块的单元测试。
|
||
|
||
覆盖 anchor 工具函数（纯函数）和 summarize_* 异步函数（FakeLLMProvider mock）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from typing import TYPE_CHECKING, Any
|
||
|
||
import pytest
|
||
|
||
if TYPE_CHECKING:
|
||
from pathlib import Path
|
||
|
||
from app.search.summarizer import (
|
||
_expand_anchor_ids,
|
||
assemble_anchored_output,
|
||
check_anchors,
|
||
summarize_children,
|
||
summarize_node,
|
||
summarize_nodes_batch,
|
||
)
|
||
|
||
# ── Fake LLM 基础设施 ──────────────────────────────────────────────
|
||
|
||
|
||
@dataclass
|
||
class FakeLLMResponse:
|
||
"""FakeLLMProvider 返回的响应对象。"""
|
||
|
||
content: str
|
||
thinking: str = ""
|
||
model: str = "fake"
|
||
provider: str = "fake"
|
||
prompt_tokens: int = 0
|
||
completion_tokens: int = 0
|
||
latency_ms: int = 0
|
||
ttft_ms: float | None = None
|
||
max_inter_token_ms: float | None = None
|
||
cache_hit: bool = False
|
||
call_id: str = "fake-call"
|
||
|
||
|
||
class FakeLLMProvider:
|
||
"""按顺序返回预设响应的 LLMProvider 假实现。"""
|
||
|
||
def __init__(self, responses: list[str]) -> None:
|
||
self._responses = iter(responses)
|
||
|
||
async def chat(
|
||
self,
|
||
messages: list[dict[str, Any]],
|
||
*,
|
||
session_id: str | None = None,
|
||
parent_call_id: str | None = None,
|
||
) -> FakeLLMResponse:
|
||
"""返回下一个预设响应。"""
|
||
return FakeLLMResponse(content=next(self._responses))
|
||
|
||
|
||
class FailingLLMProvider:
|
||
"""始终抛出异常的 LLMProvider 假实现。"""
|
||
|
||
def __init__(self, error_msg: str = "LLM 调用失败") -> None:
|
||
self._error_msg = error_msg
|
||
|
||
async def chat(
|
||
self,
|
||
messages: list[dict[str, Any]],
|
||
*,
|
||
session_id: str | None = None,
|
||
parent_call_id: str | None = None,
|
||
) -> FakeLLMResponse:
|
||
"""始终抛出异常。"""
|
||
raise RuntimeError(self._error_msg)
|
||
|
||
|
||
class FailOnNthLLMProvider:
|
||
"""第 N 次调用抛异常，其余正常返回的 LLMProvider。"""
|
||
|
||
def __init__(self, responses: list[str], fail_on: int) -> None:
|
||
self._responses = list(responses)
|
||
self._fail_on = fail_on
|
||
self._call_count = 0
|
||
|
||
async def chat(
|
||
self,
|
||
messages: list[dict[str, Any]],
|
||
*,
|
||
session_id: str | None = None,
|
||
parent_call_id: str | None = None,
|
||
) -> FakeLLMResponse:
|
||
"""第 fail_on 次调用抛异常。"""
|
||
self._call_count += 1
|
||
if self._call_count == self._fail_on:
|
||
raise RuntimeError(f"第 {self._fail_on} 次调用失败")
|
||
idx = self._call_count - 1
|
||
if self._call_count > self._fail_on:
|
||
idx -= 1
|
||
return FakeLLMResponse(content=self._responses[idx])
|
||
|
||
|
||
# ── Prompt 文件 fixture ──────────────────────────────────────────────
|
||
|
||
|
||
@pytest.fixture()
|
||
def prompts_dir(tmp_path: Path) -> Path:
|
||
"""创建包含最小化 prompt 文件的临时目录。"""
|
||
prompts = {
|
||
"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": "核实搜索结果摘要。",
|
||
}
|
||
for filename, content in prompts.items():
|
||
(tmp_path / filename).write_text(content, encoding="utf-8")
|
||
return tmp_path
|
||
|
||
|
||
# ══════════════════════════════════════════════════════════════════════
|
||
# Part A: Anchor 工具函数测试（纯函数，无需 mock）
|
||
# ══════════════════════════════════════════════════════════════════════
|
||
|
||
|
||
class TestExpandAnchorIds:
|
||
"""_expand_anchor_ids 展开范围语法。"""
|
||
|
||
def test_single_ids(self) -> None:
|
||
"""单个 id 不展开。"""
|
||
assert _expand_anchor_ids("s1") == ["s1"]
|
||
assert _expand_anchor_ids("c2") == ["c2"]
|
||
|
||
def test_comma_separated(self) -> None:
|
||
"""逗号分隔的多个 id。"""
|
||
assert _expand_anchor_ids("s1,c2,s5") == ["s1", "c2", "s5"]
|
||
|
||
def test_range_expansion(self) -> None:
|
||
"""范围语法 s3-s5 展开为 [s3, s4, s5]。"""
|
||
assert _expand_anchor_ids("s3-s5") == ["s3", "s4", "s5"]
|
||
|
||
def test_range_short_form(self) -> None:
|
||
"""短范围语法 s3-5（省略第二个前缀）也应展开。"""
|
||
assert _expand_anchor_ids("s3-5") == ["s3", "s4", "s5"]
|
||
|
||
def test_range_with_c_prefix(self) -> None:
|
||
"""c 前缀范围展开。"""
|
||
assert _expand_anchor_ids("c1-c3") == ["c1", "c2", "c3"]
|
||
|
||
def test_mixed_ids_and_ranges(self) -> None:
|
||
"""混合单 id 和范围。"""
|
||
result = _expand_anchor_ids("s1,c2-c4,s10")
|
||
assert result == ["s1", "c2", "c3", "c4", "s10"]
|
||
|
||
def test_cross_prefix_range_kept_as_token(self) -> None:
|
||
"""跨前缀范围（c3-s5）保留原 token。"""
|
||
result = _expand_anchor_ids("c3-s5")
|
||
assert result == ["c3-s5"]
|
||
|
||
def test_reversed_range_kept_as_token(self) -> None:
|
||
"""起点>终点的范围保留原 token。"""
|
||
result = _expand_anchor_ids("s5-s3")
|
||
assert result == ["s5-s3"]
|
||
|
||
def test_explosion_guard(self) -> None:
|
||
"""超过 50 条展开上限的范围保留原 token。"""
|
||
result = _expand_anchor_ids("s1-s100")
|
||
assert result == ["s1-s100"]
|
||
|
||
def test_fullwidth_comma(self) -> None:
|
||
"""全角逗号分隔。"""
|
||
result = _expand_anchor_ids("s1，s2")
|
||
assert result == ["s1", "s2"]
|
||
|
||
|
||
class TestCheckAnchors:
|
||
"""check_anchors 校验行号引注。"""
|
||
|
||
def test_legal_anchors_preserved(self) -> None:
|
||
"""合法锚保留不变。"""
|
||
anchor_map = {"s1": "第一行", "s2": "第二行", "c1": "字幕一"}
|
||
summary = "[相关信息]\n- 关键发现(s1)\n- 另一个发现(c1)"
|
||
cleaned, stats = check_anchors(summary, anchor_map)
|
||
assert "(s1)" in cleaned
|
||
assert "(c1)" in cleaned
|
||
assert stats["n_illegal"] == 0
|
||
assert stats["n_assertions"] == 2
|
||
assert stats["n_anchored"] == 2
|
||
|
||
def test_illegal_anchors_removed(self) -> None:
|
||
"""非法锚被删除，断言文本保留。"""
|
||
anchor_map = {"s1": "第一行"}
|
||
summary = "[相关信息]\n- 关键发现(s99)"
|
||
cleaned, stats = check_anchors(summary, anchor_map)
|
||
assert "(s99)" not in cleaned
|
||
assert "关键发现" in cleaned
|
||
assert stats["n_illegal"] == 1
|
||
assert stats["n_assertions"] == 1
|
||
assert stats["n_anchored"] == 0
|
||
|
||
def test_range_expansion_in_check(self) -> None:
|
||
"""范围语法在 check_anchors 中展开并校验。"""
|
||
anchor_map = {"s1": "行1", "s2": "行2", "s3": "行3"}
|
||
summary = "[相关信息]\n- 发现(s1-s3)"
|
||
cleaned, stats = check_anchors(summary, anchor_map)
|
||
assert "(s1,s2,s3)" in cleaned
|
||
assert stats["n_illegal"] == 0
|
||
|
||
def test_partial_legal_range(self) -> None:
|
||
"""范围中部分合法：仅保留合法子集。"""
|
||
anchor_map = {"s1": "行1", "s2": "行2"}
|
||
summary = "[相关信息]\n- 发现(s1-s4)"
|
||
cleaned, stats = check_anchors(summary, anchor_map)
|
||
assert "(s1,s2)" in cleaned
|
||
assert stats["n_illegal"] == 2 # s3, s4 非法
|
||
|
||
def test_no_info_statement_not_counted(self) -> None:
|
||
"""声明句"未包含…相关…信息"不计入 n_assertions。"""
|
||
anchor_map = {"s1": "行1"}
|
||
summary = (
|
||
"[相关信息]\n"
|
||
"- 该节点未包含与问题直接相关的信息\n"
|
||
"- 关键发现(s1)"
|
||
)
|
||
_, stats = check_anchors(summary, anchor_map)
|
||
assert stats["n_assertions"] == 1 # 声明句不计
|
||
assert stats["n_anchored"] == 1
|
||
|
||
def test_no_relevant_section(self) -> None:
|
||
"""无 [相关信息] 段落时，只清理锚，统计为零。"""
|
||
anchor_map = {"s1": "行1"}
|
||
summary = "一些分析文本(s1)(s99)"
|
||
cleaned, stats = check_anchors(summary, anchor_map)
|
||
assert "(s1)" in cleaned
|
||
assert "(s99)" not in cleaned
|
||
assert stats["n_assertions"] == 0
|
||
assert stats["n_anchored"] == 0
|
||
assert stats["n_illegal"] == 1
|
||
|
||
def test_fullwidth_brackets(self) -> None:
|
||
"""全角括号也应被识别。"""
|
||
anchor_map = {"s1": "行1"}
|
||
summary = "[相关信息]\n- 发现（s1）"
|
||
cleaned, stats = check_anchors(summary, anchor_map)
|
||
assert stats["n_anchored"] == 1
|
||
|
||
def test_all_illegal_group_removed(self) -> None:
|
||
"""组内全非法则整组删除。"""
|
||
anchor_map = {"s1": "行1"}
|
||
summary = "[相关信息]\n- 发现(s99,s100)"
|
||
cleaned, stats = check_anchors(summary, anchor_map)
|
||
assert "(s99" not in cleaned
|
||
assert "(s100" not in cleaned
|
||
assert stats["n_illegal"] == 2
|
||
|
||
|
||
class TestAssembleAnchoredOutput:
|
||
"""assemble_anchored_output 三种模式 + 封顶逻辑。"""
|
||
|
||
def test_ids_mode_no_expansion(self) -> None:
|
||
"""ids 模式：不展开引文，原样输出。"""
|
||
anchor_map = {"s1": "行1", "s2": "行2"}
|
||
summary = "关键发现(s1)"
|
||
result, stats = assemble_anchored_output(summary, anchor_map, "ids")
|
||
assert result == summary
|
||
assert stats["n_expanded"] == 0
|
||
|
||
def test_ids_expand_mode(self) -> None:
|
||
"""ids_expand 模式：保留行号 + 附加引文段。"""
|
||
anchor_map = {"s1": "第一行内容", "s2": "第二行内容"}
|
||
summary = "关键发现(s1,s2)"
|
||
result, stats = assemble_anchored_output(
|
||
summary, anchor_map, "ids_expand"
|
||
)
|
||
assert "(s1,s2)" in result
|
||
assert "[引文]" in result
|
||
assert 's1: "第一行内容"' in result
|
||
assert 's2: "第二行内容"' in result
|
||
assert stats["n_expanded"] == 2
|
||
|
||
def test_expand_only_mode_strips_anchors(self) -> None:
|
||
"""expand_only 模式：剥除行号 + 附加引文段。"""
|
||
anchor_map = {"s1": "第一行内容"}
|
||
summary = "关键发现(s1)"
|
||
result, stats = assemble_anchored_output(
|
||
summary, anchor_map, "expand_only"
|
||
)
|
||
assert "(s1)" not in result
|
||
assert "[引文]" in result
|
||
assert 's1: "第一行内容"' in result
|
||
assert stats["n_expanded"] == 1
|
||
|
||
def test_max_items_cap(self) -> None:
|
||
"""超过 5 条引文的封顶。"""
|
||
anchor_map = {f"s{i}": f"行{i}" for i in range(1, 10)}
|
||
refs = ",".join(f"s{i}" for i in range(1, 10))
|
||
summary = f"发现({refs})"
|
||
result, stats = assemble_anchored_output(
|
||
summary, anchor_map, "ids_expand"
|
||
)
|
||
assert stats["n_expanded"] == 5
|
||
|
||
def test_max_chars_cap(self) -> None:
|
||
"""总字符超过 800 时截断。"""
|
||
anchor_map = {
|
||
f"s{i}": "A" * 300 for i in range(1, 6)
|
||
}
|
||
refs = ",".join(f"s{i}" for i in range(1, 6))
|
||
summary = f"发现({refs})"
|
||
result, stats = assemble_anchored_output(
|
||
summary, anchor_map, "ids_expand"
|
||
)
|
||
# 300 字符原文 + 前缀 ≈ 310+ 每条，800 / 310 ≈ 2 条
|
||
assert stats["n_expanded"] < 5
|
||
|
||
def test_line_cap_truncation(self) -> None:
|
||
"""单行超 200 字符截断并标记 n_trunc。"""
|
||
anchor_map = {"s1": "A" * 250}
|
||
summary = "发现(s1)"
|
||
result, stats = assemble_anchored_output(
|
||
summary, anchor_map, "ids_expand"
|
||
)
|
||
assert stats["n_trunc"] == 1
|
||
assert "…" in result
|
||
|
||
def test_invalid_mode_raises(self) -> None:
|
||
"""无效模式应抛出 AssertionError。"""
|
||
with pytest.raises(AssertionError, match="未知装配形态"):
|
||
assemble_anchored_output("text", {}, "bad_mode")
|
||
|
||
|
||
# ══════════════════════════════════════════════════════════════════════
|
||
# Part B: summarize_* 异步函数测试（FakeLLMProvider mock）
|
||
# ══════════════════════════════════════════════════════════════════════
|
||
|
||
|
||
class TestSummarizeNode:
|
||
"""summarize_node 两轮摘要。"""
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_normal_two_round(self, prompts_dir: Path) -> None:
|
||
"""正常两轮：提取 + 核实。"""
|
||
llm = FakeLLMProvider(["提取结果摘要", "核实通过"])
|
||
result = await summarize_node(
|
||
llm,
|
||
"视频片段内容",
|
||
"这个视频讲了什么？",
|
||
prompts_dir,
|
||
anchor_map=None,
|
||
assemble_mode="ids",
|
||
)
|
||
assert "[内容摘要] 提取结果摘要" in result
|
||
assert "[核实] 核实通过" in result
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_extract_failure(self, prompts_dir: Path) -> None:
|
||
"""提取轮失败返回错误信息。"""
|
||
llm = FailingLLMProvider("网络超时")
|
||
result = await summarize_node(
|
||
llm,
|
||
"视频片段内容",
|
||
"问题",
|
||
prompts_dir,
|
||
anchor_map=None,
|
||
assemble_mode="ids",
|
||
)
|
||
assert "[摘要错误]" in result
|
||
assert "网络超时" in result
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_verify_failure_degrades(self, prompts_dir: Path) -> None:
|
||
"""核实轮失败降级为"跳过"。"""
|
||
llm = FailOnNthLLMProvider(["提取结果"], fail_on=2)
|
||
result = await summarize_node(
|
||
llm,
|
||
"视频片段内容",
|
||
"问题",
|
||
prompts_dir,
|
||
anchor_map=None,
|
||
assemble_mode="ids",
|
||
)
|
||
assert "[内容摘要] 提取结果" in result
|
||
assert "跳过（调用失败）" in result
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_anchor_mode(self, prompts_dir: Path) -> None:
|
||
"""锚模式：check_anchors + assemble。"""
|
||
anchor_map = {"s1": "第一行", "s2": "第二行"}
|
||
llm = FakeLLMProvider([
|
||
"[相关信息]\n- 关键发现(s1)\n- 补充(s2)",
|
||
"核实通过",
|
||
])
|
||
result = await summarize_node(
|
||
llm,
|
||
"带行号的内容",
|
||
"问题",
|
||
prompts_dir,
|
||
anchor_map=anchor_map,
|
||
assemble_mode="ids_expand",
|
||
)
|
||
assert "[内容摘要]" in result
|
||
assert "[核实] 核实通过" in result
|
||
assert "[引文]" in result
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_anchor_mode_with_stats_sink(self, prompts_dir: Path) -> None:
|
||
"""锚模式 stats_sink 回调接收完整统计。"""
|
||
anchor_map = {"s1": "第一行"}
|
||
collected: list[dict] = []
|
||
llm = FakeLLMProvider([
|
||
"[相关信息]\n- 关键发现(s1)",
|
||
"核实通过",
|
||
])
|
||
await summarize_node(
|
||
llm,
|
||
"内容",
|
||
"问题",
|
||
prompts_dir,
|
||
anchor_map=anchor_map,
|
||
assemble_mode="ids_expand",
|
||
stats_sink=collected.append,
|
||
)
|
||
assert len(collected) == 1
|
||
s = collected[0]
|
||
assert "n_assertions" in s
|
||
assert "n_anchored" in s
|
||
assert "n_expanded" in s
|
||
assert "output_chars" in s
|
||
assert "pre_assembly" in s
|
||
assert "anchor_map" in s
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_session_id_forwarded(self, prompts_dir: Path) -> None:
|
||
"""session_id 和 parent_call_id 应透传给 LLM。"""
|
||
received_kwargs: list[dict] = []
|
||
|
||
class CaptureLLM:
|
||
"""捕获 kwargs 的 LLM。"""
|
||
|
||
async def chat(self, messages: list, **kwargs: Any) -> FakeLLMResponse:
|
||
received_kwargs.append(kwargs)
|
||
return FakeLLMResponse(content="ok")
|
||
|
||
llm = CaptureLLM()
|
||
await summarize_node(
|
||
llm,
|
||
"内容",
|
||
"问题",
|
||
prompts_dir,
|
||
anchor_map=None,
|
||
assemble_mode="ids",
|
||
session_id="sess-1",
|
||
parent_call_id="call-0",
|
||
)
|
||
assert len(received_kwargs) == 2
|
||
for kw in received_kwargs:
|
||
assert kw["session_id"] == "sess-1"
|
||
assert kw["parent_call_id"] == "call-0"
|
||
|
||
|
||
class TestSummarizeChildren:
|
||
"""summarize_children 子节点标注。"""
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_normal(self, prompts_dir: Path) -> None:
|
||
"""正常两轮标注。"""
|
||
children_info = [
|
||
{"id": "n1", "time_range": (0.0, 30.0), "summary": "开头"},
|
||
{"id": "n2", "time_range": (30.0, 60.0), "summary": "中间"},
|
||
]
|
||
llm = FakeLLMProvider(["相关性标注结果", "核实通过"])
|
||
result = await summarize_children(
|
||
llm, children_info, "问题", prompts_dir
|
||
)
|
||
assert "相关性标注结果" in result
|
||
assert "[核实] 核实通过" in result
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_extract_failure_fallback(self, prompts_dir: Path) -> None:
|
||
"""提取失败回退到原始列表。"""
|
||
children_info = [
|
||
{"id": "n1", "time_range": (0.0, 30.0), "summary": "开头"},
|
||
]
|
||
llm = FailingLLMProvider("网络错误")
|
||
result = await summarize_children(
|
||
llm, children_info, "问题", prompts_dir
|
||
)
|
||
assert "n1" in result
|
||
assert "0-30s" in result
|
||
assert "开头" in result
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_verify_failure_returns_extract_only(
|
||
self, prompts_dir: Path
|
||
) -> None:
|
||
"""核实轮失败仍返回提取结果。"""
|
||
children_info = [
|
||
{"id": "n1", "time_range": (0.0, 30.0), "summary": "开头"},
|
||
]
|
||
llm = FailOnNthLLMProvider(["标注结果"], fail_on=2)
|
||
result = await summarize_children(
|
||
llm, children_info, "问题", prompts_dir
|
||
)
|
||
assert "标注结果" in result
|
||
|
||
|
||
class TestSummarizeNodesBatch:
|
||
"""summarize_nodes_batch 并发多节点。"""
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_batch_normal(self, prompts_dir: Path) -> None:
|
||
"""并发三个节点，结果顺序与输入一致。"""
|
||
# 每个节点需要 2 轮 LLM 调用（提取 + 核实）
|
||
llm = FakeLLMProvider([
|
||
"摘要A", "核实A",
|
||
"摘要B", "核实B",
|
||
"摘要C", "核实C",
|
||
])
|
||
items = [
|
||
("n1", "内容1", "extra1"),
|
||
("n2", "内容2", "extra2"),
|
||
("n3", "内容3", "extra3"),
|
||
]
|
||
results = await summarize_nodes_batch(
|
||
llm, items, "问题", prompts_dir
|
||
)
|
||
assert len(results) == 3
|
||
assert results[0][0] == "n1"
|
||
assert results[1][0] == "n2"
|
||
assert results[2][0] == "n3"
|
||
assert "[内容摘要]" in results[0][1]
|
||
assert "[内容摘要]" in results[1][1]
|
||
assert "[内容摘要]" in results[2][1]
|
||
|
||
@pytest.mark.asyncio()
|
||
async def test_batch_empty(self, prompts_dir: Path) -> None:
|
||
"""空列表返回空结果。"""
|
||
llm = FakeLLMProvider([])
|
||
results = await summarize_nodes_batch(
|
||
llm, [], "问题", prompts_dir
|
||
)
|
||
assert results == []
|