kplam
dce9e31bde
- 新增 JWT 认证模块，支持登录/注册/用户管理 - 新增文件上传接口，支持 .txt/.md/.html/.pdf/.docx 等格式解析入库 - 新增检索结果 AI 总结功能 - 新增文本去重缓存机制 - 新增全局认证夹具简化测试 - 新增配置项与环境变量支持 - 完善文档与测试覆盖
139 lines
5.0 KiB
Python
139 lines
5.0 KiB
Python
"""ResultSummarizer 单元测试（FakeOllama，不依赖真实 Ollama）
|
||
|
||
覆盖：
|
||
- summarize 正常：调用 ollama，prompt 含 query 与 hits 文本，返回 strip 后的 summary
|
||
- summarize 无 hits：返回空串且不调 ollama
|
||
- summarize ollama 异常：返回空串
|
||
- _build_context 拼接格式：编号 / title / section_path / 文本
|
||
- 取前 settings.result_summary_max_hits 条：超量 hits 只取前 N 条进入 prompt
|
||
"""
|
||
|
||
import pytest
|
||
|
||
from app.config import settings
|
||
from app.core.result_summarizer import ResultSummarizer
|
||
from app.models.search import SearchHit
|
||
|
||
|
||
class FakeOllama:
|
||
"""记录调用并返回固定响应的假 OllamaClient"""
|
||
|
||
def __init__(self, response: str = "这是总结") -> None:
|
||
self.response = response
|
||
self.calls: list[dict] = []
|
||
|
||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||
self.calls.append({"prompt": prompt, "json_mode": json_mode})
|
||
return self.response
|
||
|
||
|
||
class FailingOllama:
|
||
"""generate 抛异常的假 OllamaClient，用于测试容错降级"""
|
||
|
||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||
raise RuntimeError("ollama 不可用")
|
||
|
||
|
||
def _hit(idx: int, title: str = "", section_path: str = "") -> SearchHit:
|
||
"""构造测试用 SearchHit"""
|
||
return SearchHit(
|
||
text=f"文本内容-{idx}",
|
||
doc_id=f"doc-{idx}",
|
||
title=title or f"标题-{idx}",
|
||
section_path=section_path,
|
||
score=0.1 * idx,
|
||
)
|
||
|
||
|
||
class TestSummarize:
|
||
"""ResultSummarizer.summarize 主流程"""
|
||
|
||
async def test_summarize_normal(self):
|
||
"""正常总结：调用 ollama，prompt 含 query 与 hits 文本，返回 strip 后的 summary"""
|
||
ollama = FakeOllama(response=" 这是 AI 总结 ")
|
||
summarizer = ResultSummarizer(ollama=ollama) # type: ignore[arg-type]
|
||
hits = [_hit(1, "文档A", "章节A"), _hit(2, "文档B", "章节B")]
|
||
|
||
result = await summarizer.summarize("安装步骤是什么", hits)
|
||
|
||
assert result == "这是 AI 总结"
|
||
assert len(ollama.calls) == 1
|
||
prompt = ollama.calls[0]["prompt"]
|
||
assert "安装步骤是什么" in prompt
|
||
assert "文本内容-1" in prompt
|
||
assert "文本内容-2" in prompt
|
||
assert "文档A" in prompt
|
||
assert "文档B" in prompt
|
||
assert "章节A" in prompt
|
||
assert "章节B" in prompt
|
||
|
||
async def test_summarize_empty_hits_returns_empty(self):
|
||
"""无 hits：返回空串且不调用 ollama"""
|
||
ollama = FakeOllama()
|
||
summarizer = ResultSummarizer(ollama=ollama) # type: ignore[arg-type]
|
||
|
||
result = await summarizer.summarize("任何问题", [])
|
||
|
||
assert result == ""
|
||
assert ollama.calls == []
|
||
|
||
async def test_summarize_ollama_error_returns_empty(self):
|
||
"""ollama.generate 抛异常：返回空串（容错降级）"""
|
||
summarizer = ResultSummarizer(ollama=FailingOllama()) # type: ignore[arg-type]
|
||
hits = [_hit(1)]
|
||
|
||
result = await summarizer.summarize("问题", hits)
|
||
|
||
assert result == ""
|
||
|
||
|
||
class TestBuildContext:
|
||
"""_build_context 拼接格式：编号 / title / section_path / 文本"""
|
||
|
||
def test_build_context_format(self):
|
||
hits = [
|
||
_hit(1, "文档A", "第一章"),
|
||
_hit(2, "文档B", "第二章"),
|
||
]
|
||
context = ResultSummarizer._build_context(hits)
|
||
|
||
# 每条带编号、title、section_path
|
||
assert "[1] / 文档A / 第一章" in context
|
||
assert "[2] / 文档B / 第二章" in context
|
||
# 文本内容拼接
|
||
assert "文本内容-1" in context
|
||
assert "文本内容-2" in context
|
||
# 分隔符
|
||
assert "---" in context
|
||
|
||
def test_build_context_empty_title_and_section(self):
|
||
"""title 与 section_path 为空时头部仅保留编号"""
|
||
hit = SearchHit(text="纯文本", doc_id="d1", title="", section_path="", score=1.0)
|
||
context = ResultSummarizer._build_context([hit])
|
||
assert context.strip().startswith("[1]")
|
||
assert "纯文本" in context
|
||
|
||
|
||
class TestMaxHitsLimit:
|
||
"""取前 settings.result_summary_max_hits 条命中"""
|
||
|
||
async def test_only_first_n_hits_in_prompt(self, monkeypatch: pytest.MonkeyPatch):
|
||
"""构造超过 max_hits 的 hits，验证 prompt 只含前 N 条文本"""
|
||
max_hits = 3
|
||
monkeypatch.setattr(settings, "result_summary_max_hits", max_hits)
|
||
ollama = FakeOllama(response="总结")
|
||
summarizer = ResultSummarizer(ollama=ollama) # type: ignore[arg-type]
|
||
|
||
# 构造 max_hits + 5 条 hits
|
||
hits = [_hit(i) for i in range(max_hits + 5)]
|
||
await summarizer.summarize("问题", hits)
|
||
|
||
assert len(ollama.calls) == 1
|
||
prompt = ollama.calls[0]["prompt"]
|
||
# 前 max_hits 条文本出现在 prompt 中
|
||
for i in range(max_hits):
|
||
assert f"文本内容-{i}" in prompt
|
||
# 超出的文本不出现在 prompt 中
|
||
for i in range(max_hits, max_hits + 5):
|
||
assert f"文本内容-{i}" not in prompt
|