feat(rag_platform): U5 — rerank + question generation + termbase
Add rerank.py: Reranker with Cohere/BGE provider support, data export risk annotation, graceful degradation Add question_gen.py: LLM-based question generation following ContextualChunker pattern, with caching Add termbase.py: jieba custom dictionary management, add/remove/load terms Tests: 58 new tests (14 rerank + 19 question_gen + 25 termbase), 205 total passing
This commit is contained in:
chiguyong
parent
fb9f16d6e5
commit
5c562dbff3
|
|
@ -0,0 +1,193 @@
|
|||
"""LLM-based 问题生成 — 为文档段落生成相关问题，提升检索召回率。
|
||||
|
||||
参考 memory/contextual_retrieval.py 的 ContextualChunker 模式：
|
||||
- 使用 LLM gateway 生成问题
|
||||
- Prompt 模板化
|
||||
- 失败时降级（返回空列表，不抛异常）
|
||||
|
||||
生成的问题可作为 chunk 的 metadata.questions 字段，在索引时与 chunk 内容
|
||||
一同嵌入，提升"问题→段落"的检索召回率（HyDE 反向模式）。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GeneratedQuestion(BaseModel):
|
||||
"""生成的问题条目。"""
|
||||
|
||||
model_config = ConfigDict()
|
||||
|
||||
question: str
|
||||
chunk_id: str
|
||||
document_id: str
|
||||
|
||||
|
||||
# Prompt 模板 — 指示 LLM 为给定 chunk 生成可被该 chunk 回答的问题
|
||||
QUESTION_GEN_PROMPT_TEMPLATE = """\
|
||||
请为以下文本片段生成 {n} 个问题，要求：
|
||||
1. 每个问题都能直接从该片段中找到答案
|
||||
2. 问题应覆盖片段中的关键信息点
|
||||
3. 问题应简洁、自然，符合用户提问习惯
|
||||
4. 每行一个问题，不要编号，不要前缀
|
||||
|
||||
文本片段：
|
||||
{chunk}
|
||||
|
||||
问题（每行一个）：
|
||||
"""
|
||||
|
||||
|
||||
# ponytail: 解析 LLM 输出为问题列表 — 按行切分，过滤空行和编号前缀
|
||||
# 升级路径：若 LLM 输出格式不稳定，可改用 JSON 模式（要求 LLM 输出 JSON 数组）
|
||||
_NUMBER_PREFIX_RE = re.compile(r"^\s*\d+[\.\)、\]]\s*")
|
||||
|
||||
|
||||
class QuestionGenerator:
|
||||
"""问题生成器 — 为每个 chunk 生成相关问题。
|
||||
|
||||
Args:
|
||||
llm_gateway: LLM Gateway 实例（需实现 async chat(messages, model) -> response）
|
||||
max_questions_per_chunk: 每个 chunk 生成的问题数上限
|
||||
model: LLM 模型名（默认 "default"）
|
||||
cache: 是否启用缓存（避免对同一 chunk 重复调用 LLM）
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm_gateway: Any = None,
|
||||
max_questions_per_chunk: int = 3,
|
||||
model: str = "default",
|
||||
cache: bool = True,
|
||||
) -> None:
|
||||
self._llm_gateway = llm_gateway
|
||||
self._max_questions = max_questions_per_chunk
|
||||
self._model = model
|
||||
self._cache_enabled = cache
|
||||
self._cache: dict[str, list[str]] = {}
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
chunks: list[dict[str, Any]],
|
||||
document_context: str = "",
|
||||
) -> list[GeneratedQuestion]:
|
||||
"""为每个 chunk 生成相关问题。
|
||||
|
||||
Args:
|
||||
chunks: chunk 字典列表，每个字典需包含 id、content、document_id 字段
|
||||
document_context: 完整文档内容（可选，用于提供额外上下文）
|
||||
|
||||
Returns:
|
||||
生成的问题列表（GeneratedQuestion）。无 LLM 或失败时返回空列表。
|
||||
"""
|
||||
if not chunks:
|
||||
return []
|
||||
|
||||
if not self._llm_gateway:
|
||||
logger.info("No LLM gateway configured, skipping question generation")
|
||||
return []
|
||||
|
||||
results: list[GeneratedQuestion] = []
|
||||
for chunk in chunks:
|
||||
chunk_id = str(chunk.get("id", ""))
|
||||
content = str(chunk.get("content", ""))
|
||||
document_id = str(chunk.get("document_id", ""))
|
||||
|
||||
if not content.strip():
|
||||
continue
|
||||
|
||||
questions = await self._generate_for_chunk(content, document_context)
|
||||
for q in questions:
|
||||
results.append(
|
||||
GeneratedQuestion(
|
||||
question=q,
|
||||
chunk_id=chunk_id,
|
||||
document_id=document_id,
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
async def _generate_for_chunk(
|
||||
self,
|
||||
chunk_content: str,
|
||||
document_context: str,
|
||||
) -> list[str]:
|
||||
"""为单个 chunk 生成问题。
|
||||
|
||||
Args:
|
||||
chunk_content: chunk 文本内容
|
||||
document_context: 完整文档内容（当前未使用，保留以供未来扩展）
|
||||
|
||||
Returns:
|
||||
问题字符串列表。失败时返回空列表。
|
||||
"""
|
||||
# 缓存检查
|
||||
cache_key = self._make_cache_key(chunk_content)
|
||||
if self._cache_enabled and cache_key in self._cache:
|
||||
return self._cache[cache_key]
|
||||
|
||||
# 截断超长 chunk — 避免 prompt 过长
|
||||
chunk_preview = chunk_content[:2000] if len(chunk_content) > 2000 else chunk_content
|
||||
|
||||
prompt = QUESTION_GEN_PROMPT_TEMPLATE.format(
|
||||
n=self._max_questions,
|
||||
chunk=chunk_preview,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self._llm_gateway.chat(
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
model=self._model,
|
||||
)
|
||||
content = response.content.strip()
|
||||
questions = self._parse_questions(content)
|
||||
except Exception as e:
|
||||
logger.warning("Question generation failed for chunk: %s", e)
|
||||
questions = []
|
||||
|
||||
if self._cache_enabled:
|
||||
self._cache[cache_key] = questions
|
||||
|
||||
return questions
|
||||
|
||||
@staticmethod
|
||||
def _parse_questions(raw_output: str) -> list[str]:
|
||||
"""解析 LLM 输出为问题列表 — 按行切分，过滤空行和编号前缀。
|
||||
|
||||
Args:
|
||||
raw_output: LLM 原始输出文本
|
||||
|
||||
Returns:
|
||||
问题字符串列表（最多 max_questions 个）。
|
||||
"""
|
||||
lines: list[str] = []
|
||||
for line in raw_output.splitlines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
# 去除编号前缀（如 "1. "、"2) "、"3、 "）
|
||||
line = _NUMBER_PREFIX_RE.sub("", line).strip()
|
||||
if line:
|
||||
lines.append(line)
|
||||
return lines
|
||||
|
||||
@staticmethod
|
||||
def _make_cache_key(chunk_content: str) -> str:
|
||||
"""生成缓存键 — 基于 chunk 内容的 SHA256 哈希。"""
|
||||
return hashlib.sha256(chunk_content.encode()).hexdigest()[:16]
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""清除问题生成缓存。"""
|
||||
self._cache.clear()
|
||||
|
||||
|
||||
__all__ = ["GeneratedQuestion", "QuestionGenerator"]
|
||||
|
|
@ -0,0 +1,210 @@
|
|||
"""Rerank 模型集成 — 支持 Cohere Rerank 和 BGE-Reranker（本地部署）。
|
||||
|
||||
数据出境风险：Cohere Rerank 将文档 chunks 发送到第三方 API。
|
||||
敏感数据 KB 应使用 BGE-Reranker via Xinference（本地部署）。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from agentkit.rag_platform.models import QueryResult
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from llama_index.core.schema import NodeWithScore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RerankConfig(BaseModel):
|
||||
"""Rerank 配置 — 可按 KB 覆盖。
|
||||
|
||||
provider:
|
||||
- "cohere": Cohere Rerank API（数据出境）
|
||||
- "bge": BGE-Reranker via Xinference（本地部署，敏感数据 KB 推荐）
|
||||
- "none": 不重排
|
||||
"""
|
||||
|
||||
model_config = ConfigDict()
|
||||
|
||||
provider: str = "none"
|
||||
api_key: str | None = None
|
||||
base_url: str | None = None # Xinference URL for BGE
|
||||
model_name: str | None = None # 模型名（如 "bge-reranker-base"）
|
||||
top_n: int = 5
|
||||
# True 表示当前 KB 使用了云端 rerank，存在数据出境风险
|
||||
data_export_warning: bool = False
|
||||
|
||||
|
||||
class Reranker:
|
||||
"""Rerank 引擎 — 包装 LlamaIndex rerankers。
|
||||
|
||||
使用方式：
|
||||
config = RerankConfig(provider="cohere", api_key="...", top_n=5)
|
||||
reranker = Reranker(config)
|
||||
reranked = await reranker.rerank(query, results)
|
||||
"""
|
||||
|
||||
def __init__(self, config: RerankConfig) -> None:
|
||||
self._config = config
|
||||
self._reranker: Any = None # 延迟初始化，避免 import 失败
|
||||
|
||||
def _get_reranker(self) -> Any:
|
||||
"""延迟加载 reranker 实例 — 避免在 import 时失败。"""
|
||||
if self._reranker is not None:
|
||||
return self._reranker
|
||||
|
||||
cfg = self._config
|
||||
if cfg.provider == "cohere":
|
||||
self._reranker = self._build_cohere_reranker(cfg)
|
||||
elif cfg.provider == "bge":
|
||||
self._reranker = self._build_bge_reranker(cfg)
|
||||
elif cfg.provider == "none":
|
||||
self._reranker = None
|
||||
else: # pragma: no cover — 配置校验已穷尽
|
||||
raise ValueError(f"Unsupported rerank provider: {cfg.provider}")
|
||||
|
||||
return self._reranker
|
||||
|
||||
def _build_cohere_reranker(self, cfg: RerankConfig) -> Any:
|
||||
"""构建 CohereRerank — 数据出境，需 api_key。"""
|
||||
if not cfg.api_key:
|
||||
raise ValueError("Cohere rerank requires api_key")
|
||||
try:
|
||||
from llama_index.postprocessor.cohere_rerank import CohereRerank
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"CohereRerank requires llama-index-postprocessor-cohere-rerank. "
|
||||
"Install: pip install llama-index-postprocessor-cohere-rerank"
|
||||
) from e
|
||||
|
||||
kwargs: dict[str, Any] = {
|
||||
"api_key": cfg.api_key,
|
||||
"top_n": cfg.top_n,
|
||||
}
|
||||
if cfg.model_name:
|
||||
kwargs["model"] = cfg.model_name
|
||||
return CohereRerank(**kwargs)
|
||||
|
||||
def _build_bge_reranker(self, cfg: RerankConfig) -> Any:
|
||||
"""构建 BGE-Reranker via Xinference（本地部署，无数据出境）。
|
||||
|
||||
使用 SentenceTransformerRerank 作为本地 BGE-Reranker 的封装。
|
||||
若 base_url 指向 Xinference，调用方应在 KB 设置中标注为本地部署。
|
||||
"""
|
||||
try:
|
||||
from llama_index.postprocessor.sentence_transformers_rerank import (
|
||||
SentenceTransformerRerank,
|
||||
)
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"SentenceTransformerRerank requires "
|
||||
"llama-index-postprocessor-sentence-transformers-rerank. "
|
||||
"Install: pip install llama-index-postprocessor-sentence-transformers-rerank"
|
||||
) from e
|
||||
|
||||
model = cfg.model_name or "BAAI/bge-reranker-base"
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"top_n": cfg.top_n,
|
||||
}
|
||||
return SentenceTransformerRerank(**kwargs)
|
||||
|
||||
async def rerank(
|
||||
self,
|
||||
query: str,
|
||||
results: list[QueryResult],
|
||||
) -> list[QueryResult]:
|
||||
"""对检索结果重排，返回按相关性排序的 top_n 结果。
|
||||
|
||||
Args:
|
||||
query: 查询文本
|
||||
results: 原始检索结果列表
|
||||
|
||||
Returns:
|
||||
重排后的 QueryResult 列表（按相关性降序），分数已更新为 rerank 分数。
|
||||
若 provider == "none" 或 results 为空，原样返回。
|
||||
"""
|
||||
# 空结果或关闭重排 — 直接返回
|
||||
if not results or self._config.provider == "none":
|
||||
return list(results)
|
||||
|
||||
reranker = self._get_reranker()
|
||||
if reranker is None:
|
||||
return list(results)
|
||||
|
||||
# 将 QueryResult 转为 LlamaIndex NodeWithScore
|
||||
nodes_with_scores = self._to_nodes_with_scores(results)
|
||||
|
||||
try:
|
||||
# LlamaIndex reranker 的 postprocessnodes 是同步方法
|
||||
reranked_nodes = reranker.postprocessnodes(
|
||||
nodes_with_scores,
|
||||
query_str=query,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Rerank failed, returning original results: %s", e)
|
||||
return list(results)
|
||||
|
||||
# 将重排结果映射回 QueryResult，更新分数
|
||||
return self._from_nodes_with_scores(reranked_nodes, results)
|
||||
|
||||
@staticmethod
|
||||
def _to_nodes_with_scores(
|
||||
results: list[QueryResult],
|
||||
) -> "list[NodeWithScore]":
|
||||
"""将 QueryResult 列表转为 LlamaIndex NodeWithScore 列表。"""
|
||||
from llama_index.core.schema import NodeWithScore, TextNode
|
||||
|
||||
out: list[NodeWithScore] = []
|
||||
for r in results:
|
||||
node = TextNode(
|
||||
id_=r.chunk_id,
|
||||
text=r.content,
|
||||
metadata=r.metadata,
|
||||
)
|
||||
out.append(NodeWithScore(node=node, score=r.score))
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def _from_nodes_with_scores(
|
||||
nodes: "list[NodeWithScore]",
|
||||
original: list[QueryResult],
|
||||
) -> list[QueryResult]:
|
||||
"""将重排后的 NodeWithScore 列表转回 QueryResult，更新分数。
|
||||
|
||||
通过 node_id 匹配原始 QueryResult 的元数据（document_id, kb_id）。
|
||||
"""
|
||||
original_map = {r.chunk_id: r for r in original}
|
||||
out: list[QueryResult] = []
|
||||
for nws in nodes:
|
||||
node_id = nws.node.node_id if hasattr(nws.node, "node_id") else None
|
||||
original_r = original_map.get(node_id) if node_id else None
|
||||
if original_r is None:
|
||||
# 兜底：通过内容匹配（理论上不应触发）
|
||||
content = nws.node.get_content() if hasattr(nws.node, "get_content") else ""
|
||||
original_r = next(
|
||||
(r for r in original if r.content == content),
|
||||
None,
|
||||
)
|
||||
if original_r is None:
|
||||
continue
|
||||
|
||||
new_score = float(nws.score) if nws.score is not None else original_r.score
|
||||
out.append(
|
||||
QueryResult(
|
||||
chunk_id=original_r.chunk_id,
|
||||
content=original_r.content,
|
||||
score=new_score,
|
||||
metadata=original_r.metadata,
|
||||
document_id=original_r.document_id,
|
||||
kb_id=original_r.kb_id,
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
__all__ = ["RerankConfig", "Reranker"]
|
||||
|
|
@ -0,0 +1,163 @@
|
|||
"""术语表管理 + jieba 自定义词典 — 增强中文分词准确率。
|
||||
|
||||
领域术语（如"知识图谱"、"向量数据库"）默认会被 jieba 错误切分（如"知识/图谱"），
|
||||
导致全文检索召回率下降。本模块通过 jieba.add_word() / jieba.load_userdict()
|
||||
将术语表注入 jieba 词典，确保领域术语被正确识别为单个 token。
|
||||
|
||||
jieba 词典文件格式（每行一条）：
|
||||
词语 词频 词性
|
||||
例如：
|
||||
知识图谱 100 n
|
||||
向量数据库 100 n
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TermEntry(BaseModel):
|
||||
"""术语条目。"""
|
||||
|
||||
model_config = ConfigDict()
|
||||
|
||||
term: str
|
||||
frequency: int | None = None # jieba 词频（None 表示使用默认）
|
||||
pos: str | None = None # 词性标注（如 "n"、"v"）
|
||||
|
||||
|
||||
class Termbase:
|
||||
"""术语表管理 — 加载/添加/删除术语，同步到 jieba 自定义词典。
|
||||
|
||||
使用方式：
|
||||
tb = Termbase()
|
||||
tb.add_term("知识图谱")
|
||||
tb.add_term("向量数据库", frequency=100, pos="n")
|
||||
tokens = tb.tokenize("知识图谱是向量数据库的基础")
|
||||
# tokens 中 "知识图谱" 和 "向量数据库" 作为整体 token 出现
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._terms: dict[str, TermEntry] = {}
|
||||
|
||||
def add_term(
|
||||
self,
|
||||
term: str,
|
||||
frequency: int | None = None,
|
||||
pos: str | None = None,
|
||||
) -> None:
|
||||
"""添加术语到词典。
|
||||
|
||||
Args:
|
||||
term: 术语文本（非空）
|
||||
frequency: jieba 词频（None 表示使用默认）
|
||||
pos: 词性标注
|
||||
"""
|
||||
term = term.strip()
|
||||
if not term:
|
||||
return
|
||||
|
||||
self._terms[term] = TermEntry(term=term, frequency=frequency, pos=pos)
|
||||
|
||||
# 同步到 jieba 词典
|
||||
# ponytail: freq=None 时 jieba 使用默认词频（足够高以覆盖默认切分）
|
||||
# 升级路径：若需精细控制词频，调用方应显式传入 frequency 参数
|
||||
import jieba
|
||||
|
||||
if frequency is not None:
|
||||
jieba.add_word(term, freq=frequency, tag=pos)
|
||||
else:
|
||||
jieba.add_word(term, tag=pos)
|
||||
logger.debug("Added term to jieba dictionary: %s", term)
|
||||
|
||||
def load_from_file(self, path: str) -> None:
|
||||
"""从文件加载术语表（jieba 词典格式：词语 词频 词性）。
|
||||
|
||||
Args:
|
||||
path: 词典文件路径
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: 文件不存在
|
||||
"""
|
||||
file_path = Path(path)
|
||||
if not file_path.exists():
|
||||
raise FileNotFoundError(f"Termbase file not found: {path}")
|
||||
|
||||
import jieba
|
||||
|
||||
jieba.load_userdict(str(file_path))
|
||||
|
||||
# 同步到内部 _terms 字典
|
||||
with file_path.open(encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
parts = line.split()
|
||||
term = parts[0]
|
||||
freq = int(parts[1]) if len(parts) > 1 and parts[1].isdigit() else None
|
||||
pos = parts[2] if len(parts) > 2 else None
|
||||
if term:
|
||||
self._terms[term] = TermEntry(term=term, frequency=freq, pos=pos)
|
||||
|
||||
logger.info("Loaded %d terms from %s", len(self._terms), path)
|
||||
|
||||
def load_from_list(self, terms: list[str]) -> None:
|
||||
"""从字符串列表加载术语。
|
||||
|
||||
Args:
|
||||
terms: 术语字符串列表
|
||||
"""
|
||||
for term in terms:
|
||||
self.add_term(term)
|
||||
|
||||
def remove_term(self, term: str) -> None:
|
||||
"""删除术语。
|
||||
|
||||
Args:
|
||||
term: 要删除的术语
|
||||
"""
|
||||
if term not in self._terms:
|
||||
return
|
||||
|
||||
del self._terms[term]
|
||||
|
||||
import jieba
|
||||
|
||||
jieba.del_word(term)
|
||||
logger.debug("Removed term from jieba dictionary: %s", term)
|
||||
|
||||
def list_terms(self) -> list[TermEntry]:
|
||||
"""列出所有术语。
|
||||
|
||||
Returns:
|
||||
TermEntry 列表（按添加顺序）
|
||||
"""
|
||||
return list(self._terms.values())
|
||||
|
||||
def tokenize(self, text: str) -> list[str]:
|
||||
"""使用自定义词典分词。
|
||||
|
||||
Args:
|
||||
text: 待分词文本
|
||||
|
||||
Returns:
|
||||
token 列表（精确模式，cut_all=False）
|
||||
"""
|
||||
import jieba
|
||||
|
||||
return list(jieba.cut(text, cut_all=False))
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._terms)
|
||||
|
||||
def __contains__(self, term: str) -> bool:
|
||||
return term in self._terms
|
||||
|
||||
|
||||
__all__ = ["TermEntry", "Termbase"]
|
||||
|
|
@ -0,0 +1,341 @@
|
|||
"""U5 测试 — LLM-based 问题生成。
|
||||
|
||||
测试场景：
|
||||
1. 无 LLM gateway 时返回空列表
|
||||
2. 空 chunks 列表返回空列表
|
||||
3. LLM 生成的问题被正确解析（按行切分，去除编号前缀）
|
||||
4. 每个 chunk 生成的问题关联正确的 chunk_id 和 document_id
|
||||
5. LLM 失败时该 chunk 返回空问题列表（不影响其他 chunk）
|
||||
6. 缓存生效（同一 chunk 不重复调用 LLM）
|
||||
7. 跳过空内容 chunk
|
||||
8. _parse_questions 正确处理编号前缀和空行
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from agentkit.rag_platform.question_gen import (
|
||||
GeneratedQuestion,
|
||||
QuestionGenerator,
|
||||
_NUMBER_PREFIX_RE,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试辅助函数
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_llm_response(content: str):
|
||||
"""创建 mock LLM 响应。"""
|
||||
response = MagicMock()
|
||||
response.content = content
|
||||
return response
|
||||
|
||||
|
||||
def _make_mock_llm_gateway(responses: list[str] | None = None):
|
||||
"""创建 mock LLM gateway。
|
||||
|
||||
Args:
|
||||
responses: chat 方法返回的响应内容列表（按调用顺序返回）
|
||||
若为 None，返回默认响应
|
||||
|
||||
Returns:
|
||||
(mock_gateway, chat_mock) — mock gateway 和 chat AsyncMock
|
||||
"""
|
||||
mock_gateway = MagicMock()
|
||||
mock_gateway.chat = AsyncMock()
|
||||
|
||||
if responses is not None:
|
||||
mock_gateway.chat.side_effect = [_make_llm_response(r) for r in responses]
|
||||
else:
|
||||
mock_gateway.chat.return_value = _make_llm_response("问题一\n问题二\n问题三")
|
||||
|
||||
return mock_gateway, mock_gateway.chat
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GeneratedQuestion 模型测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGeneratedQuestion:
|
||||
"""GeneratedQuestion 模型测试。"""
|
||||
|
||||
def test_fields(self):
|
||||
"""模型字段正确。"""
|
||||
q = GeneratedQuestion(
|
||||
question="什么是 RAG？",
|
||||
chunk_id="c1",
|
||||
document_id="d1",
|
||||
)
|
||||
assert q.question == "什么是 RAG？"
|
||||
assert q.chunk_id == "c1"
|
||||
assert q.document_id == "d1"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _parse_questions 测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestParseQuestions:
|
||||
"""_parse_questions 静态方法测试。"""
|
||||
|
||||
def test_plain_lines(self):
|
||||
"""纯文本行被正确切分。"""
|
||||
result = QuestionGenerator._parse_questions("问题一\n问题二\n问题三")
|
||||
assert result == ["问题一", "问题二", "问题三"]
|
||||
|
||||
def test_numbered_lines(self):
|
||||
"""编号前缀被去除。"""
|
||||
result = QuestionGenerator._parse_questions("1. 问题一\n2. 问题二\n3. 问题三")
|
||||
assert result == ["问题一", "问题二", "问题三"]
|
||||
|
||||
def test_paren_numbered_lines(self):
|
||||
"""括号编号前缀被去除。"""
|
||||
result = QuestionGenerator._parse_questions("1) 问题一\n2) 问题二")
|
||||
assert result == ["问题一", "问题二"]
|
||||
|
||||
def test_chinese_numbered_lines(self):
|
||||
"""中文编号前缀（、）被去除。"""
|
||||
result = QuestionGenerator._parse_questions("1、 问题一\n2、 问题二")
|
||||
assert result == ["问题一", "问题二"]
|
||||
|
||||
def test_empty_lines_filtered(self):
|
||||
"""空行被过滤。"""
|
||||
result = QuestionGenerator._parse_questions("问题一\n\n \n问题二")
|
||||
assert result == ["问题一", "问题二"]
|
||||
|
||||
def test_empty_input(self):
|
||||
"""空输入返回空列表。"""
|
||||
assert QuestionGenerator._parse_questions("") == []
|
||||
|
||||
def test_whitespace_only(self):
|
||||
"""纯空白返回空列表。"""
|
||||
assert QuestionGenerator._parse_questions(" \n \n") == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# QuestionGenerator 测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestQuestionGeneratorNoLLM:
|
||||
"""无 LLM gateway 时的测试。"""
|
||||
|
||||
async def test_no_llm_returns_empty(self):
|
||||
"""无 LLM gateway 时返回空列表。"""
|
||||
gen = QuestionGenerator(llm_gateway=None)
|
||||
chunks = [{"id": "c1", "content": "内容", "document_id": "d1"}]
|
||||
|
||||
result = await gen.generate(chunks)
|
||||
|
||||
assert result == []
|
||||
|
||||
async def test_empty_chunks_returns_empty(self):
|
||||
"""空 chunks 列表返回空列表。"""
|
||||
mock_gw, _ = _make_mock_llm_gateway()
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw)
|
||||
|
||||
result = await gen.generate([])
|
||||
|
||||
assert result == []
|
||||
mock_gw.chat.assert_not_awaited()
|
||||
|
||||
|
||||
class TestQuestionGeneratorWithLLM:
|
||||
"""有 LLM gateway 时的测试。"""
|
||||
|
||||
async def test_generates_questions_for_chunks(self):
|
||||
"""为每个 chunk 生成相关问题。"""
|
||||
mock_gw, chat_mock = _make_mock_llm_gateway(["问题A1\n问题A2", "问题B1\n问题B2"])
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw, max_questions_per_chunk=2)
|
||||
|
||||
chunks = [
|
||||
{"id": "c1", "content": "内容A", "document_id": "d1"},
|
||||
{"id": "c2", "content": "内容B", "document_id": "d1"},
|
||||
]
|
||||
|
||||
result = await gen.generate(chunks)
|
||||
|
||||
assert len(result) == 4
|
||||
# 第一个 chunk 的问题
|
||||
assert result[0].question == "问题A1"
|
||||
assert result[0].chunk_id == "c1"
|
||||
assert result[0].document_id == "d1"
|
||||
assert result[1].question == "问题A2"
|
||||
assert result[1].chunk_id == "c1"
|
||||
# 第二个 chunk 的问题
|
||||
assert result[2].question == "问题B1"
|
||||
assert result[2].chunk_id == "c2"
|
||||
assert result[3].question == "问题B2"
|
||||
assert result[3].chunk_id == "c2"
|
||||
|
||||
# LLM 被调用 2 次（每个 chunk 一次）
|
||||
assert chat_mock.await_count == 2
|
||||
|
||||
async def test_questions_relate_to_chunk_content(self):
|
||||
"""生成的问题与 chunk 内容相关（验证 prompt 包含 chunk 内容）。"""
|
||||
mock_gw, chat_mock = _make_mock_llm_gateway(["什么是 RAG？"])
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw, max_questions_per_chunk=1)
|
||||
|
||||
chunks = [
|
||||
{
|
||||
"id": "c1",
|
||||
"content": "RAG 是检索增强生成的缩写，结合了检索和生成模型。",
|
||||
"document_id": "d1",
|
||||
}
|
||||
]
|
||||
|
||||
result = await gen.generate(chunks)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].question == "什么是 RAG？"
|
||||
|
||||
# 验证 prompt 包含 chunk 内容
|
||||
call_args = chat_mock.await_args
|
||||
messages = call_args.args[0] if call_args.args else call_args.kwargs["messages"]
|
||||
prompt_content = messages[0]["content"]
|
||||
assert "RAG" in prompt_content
|
||||
assert "检索增强生成" in prompt_content
|
||||
|
||||
async def test_llm_failure_returns_empty_for_chunk(self):
|
||||
"""LLM 失败时该 chunk 返回空问题列表，不影响其他 chunk。"""
|
||||
mock_gw = MagicMock()
|
||||
mock_gw.chat = AsyncMock(
|
||||
side_effect=[
|
||||
RuntimeError("LLM error"), # 第一个 chunk 失败
|
||||
_make_llm_response("问题B1"), # 第二个 chunk 成功
|
||||
]
|
||||
)
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw, cache=False)
|
||||
|
||||
chunks = [
|
||||
{"id": "c1", "content": "内容A", "document_id": "d1"},
|
||||
{"id": "c2", "content": "内容B", "document_id": "d1"},
|
||||
]
|
||||
|
||||
result = await gen.generate(chunks)
|
||||
|
||||
# 第一个 chunk 失败，无问题；第二个 chunk 成功，1 个问题
|
||||
assert len(result) == 1
|
||||
assert result[0].question == "问题B1"
|
||||
assert result[0].chunk_id == "c2"
|
||||
|
||||
async def test_skips_empty_content_chunks(self):
|
||||
"""空内容 chunk 被跳过（不调用 LLM）。"""
|
||||
mock_gw, chat_mock = _make_mock_llm_gateway(["问题A"])
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw)
|
||||
|
||||
chunks = [
|
||||
{"id": "c1", "content": "", "document_id": "d1"}, # 空内容
|
||||
{"id": "c2", "content": " ", "document_id": "d1"}, # 纯空白
|
||||
{"id": "c3", "content": "内容C", "document_id": "d1"}, # 有效
|
||||
]
|
||||
|
||||
result = await gen.generate(chunks)
|
||||
|
||||
# 只有 c3 生成问题
|
||||
assert len(result) == 1
|
||||
assert result[0].chunk_id == "c3"
|
||||
# LLM 只被调用 1 次
|
||||
assert chat_mock.await_count == 1
|
||||
|
||||
async def test_cache_avoids_duplicate_calls(self):
|
||||
"""缓存生效 — 同一 chunk 内容不重复调用 LLM。"""
|
||||
mock_gw, chat_mock = _make_mock_llm_gateway(["问题A"])
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw, cache=True)
|
||||
|
||||
chunks = [{"id": "c1", "content": "相同内容", "document_id": "d1"}]
|
||||
|
||||
# 第一次调用
|
||||
result1 = await gen.generate(chunks)
|
||||
assert len(result1) == 1
|
||||
|
||||
# 第二次调用相同内容
|
||||
result2 = await gen.generate(chunks)
|
||||
assert len(result2) == 1
|
||||
|
||||
# LLM 只被调用 1 次（缓存命中）
|
||||
assert chat_mock.await_count == 1
|
||||
|
||||
async def test_no_cache_calls_each_time(self):
|
||||
"""禁用缓存时每次都调用 LLM。"""
|
||||
mock_gw, chat_mock = _make_mock_llm_gateway(["问题A", "问题A"])
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw, cache=False)
|
||||
|
||||
chunks = [{"id": "c1", "content": "相同内容", "document_id": "d1"}]
|
||||
|
||||
await gen.generate(chunks)
|
||||
await gen.generate(chunks)
|
||||
|
||||
# LLM 被调用 2 次
|
||||
assert chat_mock.await_count == 2
|
||||
|
||||
async def test_clear_cache(self):
|
||||
"""clear_cache 清除缓存后重新调用 LLM。"""
|
||||
mock_gw, chat_mock = _make_mock_llm_gateway(["问题A", "问题A"])
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw, cache=True)
|
||||
|
||||
chunks = [{"id": "c1", "content": "相同内容", "document_id": "d1"}]
|
||||
|
||||
await gen.generate(chunks)
|
||||
gen.clear_cache()
|
||||
await gen.generate(chunks)
|
||||
|
||||
# 缓存清除后 LLM 被调用 2 次
|
||||
assert chat_mock.await_count == 2
|
||||
|
||||
async def test_max_questions_in_prompt(self):
|
||||
"""prompt 中包含 max_questions_per_chunk 数量。"""
|
||||
mock_gw, chat_mock = _make_mock_llm_gateway(["问题A"])
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw, max_questions_per_chunk=5)
|
||||
|
||||
chunks = [{"id": "c1", "content": "内容", "document_id": "d1"}]
|
||||
await gen.generate(chunks)
|
||||
|
||||
call_args = chat_mock.await_args
|
||||
messages = call_args.args[0] if call_args.args else call_args.kwargs["messages"]
|
||||
prompt_content = messages[0]["content"]
|
||||
# prompt 中应包含 "5 个问题"
|
||||
assert "5" in prompt_content
|
||||
|
||||
async def test_truncates_long_chunk(self):
|
||||
"""超长 chunk 被截断到 2000 字符（避免 prompt 过长）。"""
|
||||
mock_gw, chat_mock = _make_mock_llm_gateway(["问题A"])
|
||||
gen = QuestionGenerator(llm_gateway=mock_gw)
|
||||
|
||||
long_content = "A" * 3000
|
||||
chunks = [{"id": "c1", "content": long_content, "document_id": "d1"}]
|
||||
await gen.generate(chunks)
|
||||
|
||||
call_args = chat_mock.await_args
|
||||
messages = call_args.args[0] if call_args.args else call_args.kwargs["messages"]
|
||||
prompt_content = messages[0]["content"]
|
||||
# prompt 中不应包含完整的 3000 字符内容（被截断到 2000）
|
||||
assert prompt_content.count("A") < 3000
|
||||
|
||||
|
||||
class TestNumberPrefixRegex:
|
||||
"""_NUMBER_PREFIX_RE 正则测试。"""
|
||||
|
||||
def test_dot_prefix(self):
|
||||
"""点号编号前缀匹配。"""
|
||||
assert _NUMBER_PREFIX_RE.match("1. 问题")
|
||||
assert _NUMBER_PREFIX_RE.match("12. 问题")
|
||||
|
||||
def test_paren_prefix(self):
|
||||
"""括号编号前缀匹配。"""
|
||||
assert _NUMBER_PREFIX_RE.match("1) 问题")
|
||||
assert _NUMBER_PREFIX_RE.match("2) 问题")
|
||||
|
||||
def test_chinese_paren_prefix(self):
|
||||
"""中文顿号编号前缀匹配。"""
|
||||
assert _NUMBER_PREFIX_RE.match("1、 问题")
|
||||
|
||||
def test_no_match_plain_text(self):
|
||||
"""纯文本不匹配。"""
|
||||
assert not _NUMBER_PREFIX_RE.match("问题一")
|
||||
assert not _NUMBER_PREFIX_RE.match("什么是 RAG？")
|
||||
|
|
@ -0,0 +1,344 @@
|
|||
"""U5 测试 — Rerank 模型集成（Cohere/BGE/none）。
|
||||
|
||||
测试场景：
|
||||
1. provider="none" 时原样返回结果
|
||||
2. 空结果列表原样返回
|
||||
3. Cohere rerank 调用 LlamaIndex CohereRerank 并按相关性重排
|
||||
4. BGE rerank 调用 SentenceTransformerRerank 并按相关性重排
|
||||
5. rerank 失败时降级返回原始结果
|
||||
6. rerank 后结果分数被更新为 rerank 分数
|
||||
7. RerankConfig 默认值与字段校验
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from agentkit.rag_platform.models import QueryResult
|
||||
from agentkit.rag_platform.rerank import RerankConfig, Reranker
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# llama_index 模块 mock — 测试环境可能未安装 llama_index
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _setup_llama_index_mocks():
|
||||
"""注入 mock llama_index 模块到 sys.modules，使 import 成功。
|
||||
|
||||
仅当真实 llama_index 无法导入时才注入 mock — 避免污染已安装真实模块的环境。
|
||||
"""
|
||||
try:
|
||||
import llama_index # noqa: F401
|
||||
|
||||
return # 真实模块可用，无需 mock
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
mock_li = MagicMock()
|
||||
mock_li_core = MagicMock()
|
||||
mock_li_core_schema = MagicMock()
|
||||
mock_li_postprocessor = MagicMock()
|
||||
mock_li_postprocessor_cohere = MagicMock()
|
||||
mock_li_postprocessor_st = MagicMock()
|
||||
|
||||
# TextNode / NodeWithScore — 简单的 mock 类
|
||||
class MockTextNode:
|
||||
def __init__(self, id_=None, text="", metadata=None):
|
||||
self.node_id = id_
|
||||
self.text = text
|
||||
self.metadata = metadata or {}
|
||||
|
||||
def get_content(self):
|
||||
return self.text
|
||||
|
||||
class MockNodeWithScore:
|
||||
def __init__(self, node=None, score=None):
|
||||
self.node = node
|
||||
self.score = score
|
||||
|
||||
mock_li_core_schema.TextNode = MockTextNode
|
||||
mock_li_core_schema.NodeWithScore = MockNodeWithScore
|
||||
|
||||
# CohereRerank / SentenceTransformerRerank — 由测试动态配置
|
||||
mock_li_postprocessor_cohere.CohereRerank = MagicMock()
|
||||
mock_li_postprocessor_st.SentenceTransformerRerank = MagicMock()
|
||||
|
||||
sys.modules["llama_index"] = mock_li
|
||||
sys.modules["llama_index.core"] = mock_li_core
|
||||
sys.modules["llama_index.core.schema"] = mock_li_core_schema
|
||||
sys.modules["llama_index.postprocessor"] = mock_li_postprocessor
|
||||
sys.modules["llama_index.postprocessor.cohere_rerank"] = mock_li_postprocessor_cohere
|
||||
sys.modules["llama_index.postprocessor.sentence_transformers_rerank"] = mock_li_postprocessor_st
|
||||
|
||||
|
||||
_setup_llama_index_mocks()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试辅助函数
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_query_result(
|
||||
chunk_id: str,
|
||||
content: str,
|
||||
score: float = 0.5,
|
||||
document_id: str = "doc1",
|
||||
kb_id: str = "kb1",
|
||||
) -> QueryResult:
|
||||
"""创建测试用 QueryResult。"""
|
||||
return QueryResult(
|
||||
chunk_id=chunk_id,
|
||||
content=content,
|
||||
score=score,
|
||||
metadata={"document_id": document_id, "kb_id": kb_id},
|
||||
document_id=document_id,
|
||||
kb_id=kb_id,
|
||||
)
|
||||
|
||||
|
||||
def _make_mock_reranker(reranked_order: list[tuple[str, float]]):
|
||||
"""创建 mock LlamaIndex reranker。
|
||||
|
||||
Args:
|
||||
reranked_order: [(node_id, new_score), ...] — 重排后的顺序和分数
|
||||
|
||||
Returns:
|
||||
mock reranker 实例，postprocessnodes 方法返回重排后的 NodeWithScore 列表
|
||||
"""
|
||||
from llama_index.core.schema import NodeWithScore, TextNode
|
||||
|
||||
mock = MagicMock()
|
||||
mock.postprocessnodes = MagicMock(
|
||||
return_value=[
|
||||
NodeWithScore(
|
||||
node=TextNode(id_=node_id, text=""),
|
||||
score=score,
|
||||
)
|
||||
for node_id, score in reranked_order
|
||||
]
|
||||
)
|
||||
return mock
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RerankConfig 测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRerankConfig:
|
||||
"""RerankConfig 模型测试。"""
|
||||
|
||||
def test_defaults(self):
|
||||
"""默认值为 none provider，无数据出境风险。"""
|
||||
cfg = RerankConfig()
|
||||
assert cfg.provider == "none"
|
||||
assert cfg.api_key is None
|
||||
assert cfg.base_url is None
|
||||
assert cfg.top_n == 5
|
||||
assert cfg.data_export_warning is False
|
||||
|
||||
def test_cohere_config(self):
|
||||
"""Cohere 配置 — 需标注数据出境风险。"""
|
||||
cfg = RerankConfig(
|
||||
provider="cohere",
|
||||
api_key="test-key",
|
||||
top_n=3,
|
||||
data_export_warning=True,
|
||||
)
|
||||
assert cfg.provider == "cohere"
|
||||
assert cfg.api_key == "test-key"
|
||||
assert cfg.top_n == 3
|
||||
assert cfg.data_export_warning is True
|
||||
|
||||
def test_bge_config(self):
|
||||
"""BGE 配置 — 本地部署，无数据出境。"""
|
||||
cfg = RerankConfig(
|
||||
provider="bge",
|
||||
base_url="http://localhost:9997",
|
||||
model_name="bge-reranker-base",
|
||||
top_n=10,
|
||||
)
|
||||
assert cfg.provider == "bge"
|
||||
assert cfg.base_url == "http://localhost:9997"
|
||||
assert cfg.model_name == "bge-reranker-base"
|
||||
assert cfg.top_n == 10
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reranker 测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRerankerNone:
|
||||
"""provider="none" 时的测试。"""
|
||||
|
||||
async def test_none_provider_returns_as_is(self):
|
||||
"""provider="none" 时原样返回结果（不调用 reranker）。"""
|
||||
cfg = RerankConfig(provider="none")
|
||||
reranker = Reranker(cfg)
|
||||
|
||||
results = [
|
||||
_make_query_result("c1", "content 1", 0.9),
|
||||
_make_query_result("c2", "content 2", 0.5),
|
||||
]
|
||||
|
||||
reranked = await reranker.rerank("query", results)
|
||||
|
||||
assert len(reranked) == 2
|
||||
assert reranked[0].chunk_id == "c1"
|
||||
assert reranked[1].chunk_id == "c2"
|
||||
# 分数不变
|
||||
assert reranked[0].score == 0.9
|
||||
assert reranked[1].score == 0.5
|
||||
|
||||
async def test_empty_results_returns_empty(self):
|
||||
"""空结果列表原样返回空。"""
|
||||
cfg = RerankConfig(provider="cohere", api_key="key")
|
||||
reranker = Reranker(cfg)
|
||||
|
||||
reranked = await reranker.rerank("query", [])
|
||||
assert reranked == []
|
||||
|
||||
|
||||
class TestRerankerCohere:
|
||||
"""Cohere rerank 测试。"""
|
||||
|
||||
async def test_rerank_reorders_results(self):
|
||||
"""rerank 后结果按相关性重排 — 原始顺序被打乱。"""
|
||||
# 原始顺序：c1 (score=0.9), c2 (score=0.5)
|
||||
# 重排后：c2 (score=0.95), c1 (score=0.80) — c2 更相关
|
||||
cfg = RerankConfig(provider="cohere", api_key="test-key", top_n=2)
|
||||
reranker = Reranker(cfg)
|
||||
|
||||
# 注入 mock reranker
|
||||
mock_reranker = _make_mock_reranker([("c2", 0.95), ("c1", 0.80)])
|
||||
reranker._reranker = mock_reranker
|
||||
|
||||
results = [
|
||||
_make_query_result("c1", "content 1", 0.9),
|
||||
_make_query_result("c2", "content 2", 0.5),
|
||||
]
|
||||
|
||||
reranked = await reranker.rerank("query", results)
|
||||
|
||||
assert len(reranked) == 2
|
||||
# 重排后 c2 在前
|
||||
assert reranked[0].chunk_id == "c2"
|
||||
assert reranked[1].chunk_id == "c1"
|
||||
# 分数被更新为 rerank 分数
|
||||
assert reranked[0].score == 0.95
|
||||
assert reranked[1].score == 0.80
|
||||
|
||||
async def test_rerank_preserves_metadata(self):
|
||||
"""rerank 后保留原始 QueryResult 的元数据（document_id, kb_id）。"""
|
||||
cfg = RerankConfig(provider="cohere", api_key="test-key", top_n=1)
|
||||
reranker = Reranker(cfg)
|
||||
|
||||
mock_reranker = _make_mock_reranker([("c1", 0.99)])
|
||||
reranker._reranker = mock_reranker
|
||||
|
||||
results = [
|
||||
_make_query_result("c1", "content 1", 0.5, document_id="doc-99", kb_id="kb-99"),
|
||||
]
|
||||
|
||||
reranked = await reranker.rerank("query", results)
|
||||
|
||||
assert len(reranked) == 1
|
||||
assert reranked[0].document_id == "doc-99"
|
||||
assert reranked[0].kb_id == "kb-99"
|
||||
assert reranked[0].metadata["document_id"] == "doc-99"
|
||||
|
||||
async def test_rerank_failure_falls_back(self):
|
||||
"""reranker 抛异常时降级返回原始结果。"""
|
||||
cfg = RerankConfig(provider="cohere", api_key="test-key", top_n=5)
|
||||
reranker = Reranker(cfg)
|
||||
|
||||
# mock reranker 抛异常
|
||||
mock_reranker = MagicMock()
|
||||
mock_reranker.postprocessnodes = MagicMock(side_effect=RuntimeError("API error"))
|
||||
reranker._reranker = mock_reranker
|
||||
|
||||
results = [
|
||||
_make_query_result("c1", "content 1", 0.9),
|
||||
_make_query_result("c2", "content 2", 0.5),
|
||||
]
|
||||
|
||||
reranked = await reranker.rerank("query", results)
|
||||
|
||||
# 降级返回原始结果（顺序不变）
|
||||
assert len(reranked) == 2
|
||||
assert reranked[0].chunk_id == "c1"
|
||||
assert reranked[1].chunk_id == "c2"
|
||||
|
||||
async def test_cohere_requires_api_key(self):
|
||||
"""Cohere provider 缺少 api_key 时抛 ValueError。"""
|
||||
cfg = RerankConfig(provider="cohere", api_key=None)
|
||||
reranker = Reranker(cfg)
|
||||
|
||||
results = [_make_query_result("c1", "content", 0.5)]
|
||||
|
||||
try:
|
||||
await reranker.rerank("query", results)
|
||||
raise AssertionError("Expected ValueError")
|
||||
except ValueError as e:
|
||||
assert "api_key" in str(e)
|
||||
|
||||
|
||||
class TestRerankerBGE:
|
||||
"""BGE rerank 测试。"""
|
||||
|
||||
async def test_bge_rerank_reorders_results(self):
|
||||
"""BGE rerank 同样按相关性重排。"""
|
||||
cfg = RerankConfig(provider="bge", top_n=2)
|
||||
reranker = Reranker(cfg)
|
||||
|
||||
mock_reranker = _make_mock_reranker([("c2", 0.88), ("c1", 0.55)])
|
||||
reranker._reranker = mock_reranker
|
||||
|
||||
results = [
|
||||
_make_query_result("c1", "content 1", 0.9),
|
||||
_make_query_result("c2", "content 2", 0.5),
|
||||
]
|
||||
|
||||
reranked = await reranker.rerank("query", results)
|
||||
|
||||
assert len(reranked) == 2
|
||||
assert reranked[0].chunk_id == "c2"
|
||||
assert reranked[0].score == 0.88
|
||||
assert reranked[1].chunk_id == "c1"
|
||||
assert reranked[1].score == 0.55
|
||||
|
||||
async def test_bge_no_api_key_required(self):
|
||||
"""BGE 本地部署不需要 api_key。"""
|
||||
cfg = RerankConfig(provider="bge", top_n=5)
|
||||
reranker = Reranker(cfg)
|
||||
|
||||
# 注入 mock reranker — 验证不抛异常
|
||||
mock_reranker = _make_mock_reranker([("c1", 0.9)])
|
||||
reranker._reranker = mock_reranker
|
||||
|
||||
results = [_make_query_result("c1", "content", 0.5)]
|
||||
reranked = await reranker.rerank("query", results)
|
||||
|
||||
assert len(reranked) == 1
|
||||
assert reranked[0].score == 0.9
|
||||
|
||||
|
||||
class TestRerankerUnsupportedProvider:
|
||||
"""不支持的 provider 测试。"""
|
||||
|
||||
async def test_unsupported_provider_raises(self):
|
||||
"""不支持的 provider 抛 ValueError。"""
|
||||
cfg = RerankConfig(provider="invalid_provider")
|
||||
reranker = Reranker(cfg)
|
||||
|
||||
results = [_make_query_result("c1", "content", 0.5)]
|
||||
|
||||
try:
|
||||
await reranker.rerank("query", results)
|
||||
raise AssertionError("Expected ValueError")
|
||||
except ValueError as e:
|
||||
assert "Unsupported" in str(e)
|
||||
|
|
@ -0,0 +1,330 @@
|
|||
"""U5 测试 — 术语表管理 + jieba 自定义词典。
|
||||
|
||||
测试场景：
|
||||
1. add_term 添加术语后 jieba 正确分词
|
||||
2. load_from_list 批量加载术语
|
||||
3. load_from_file 从 jieba 词典文件加载
|
||||
4. remove_term 删除术语后 jieba 恢复默认分词
|
||||
5. list_terms 列出所有术语
|
||||
6. tokenize 使用自定义词典分词
|
||||
7. 术语表增强后检索召回率提升（模拟场景）
|
||||
8. 领域术语被正确分词（如"知识图谱"、"向量数据库"）
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import jieba
|
||||
|
||||
from agentkit.rag_platform.termbase import TermEntry, Termbase
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TermEntry 模型测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTermEntry:
|
||||
"""TermEntry 模型测试。"""
|
||||
|
||||
def test_defaults(self):
|
||||
"""默认值正确。"""
|
||||
entry = TermEntry(term="知识图谱")
|
||||
assert entry.term == "知识图谱"
|
||||
assert entry.frequency is None
|
||||
assert entry.pos is None
|
||||
|
||||
def test_with_all_fields(self):
|
||||
"""所有字段赋值正确。"""
|
||||
entry = TermEntry(term="向量数据库", frequency=100, pos="n")
|
||||
assert entry.term == "向量数据库"
|
||||
assert entry.frequency == 100
|
||||
assert entry.pos == "n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Termbase 基础测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTermbaseBasic:
|
||||
"""Termbase 基础功能测试。"""
|
||||
|
||||
def test_empty_termbase(self):
|
||||
"""空术语表长度为 0。"""
|
||||
tb = Termbase()
|
||||
assert len(tb) == 0
|
||||
assert tb.list_terms() == []
|
||||
assert "知识图谱" not in tb
|
||||
|
||||
def test_add_term(self):
|
||||
"""add_term 添加术语到字典。"""
|
||||
tb = Termbase()
|
||||
tb.add_term("知识图谱")
|
||||
|
||||
assert len(tb) == 1
|
||||
assert "知识图谱" in tb
|
||||
terms = tb.list_terms()
|
||||
assert len(terms) == 1
|
||||
assert terms[0].term == "知识图谱"
|
||||
|
||||
def test_add_term_with_freq_and_pos(self):
|
||||
"""add_term 带词频和词性。"""
|
||||
tb = Termbase()
|
||||
tb.add_term("向量数据库", frequency=100, pos="n")
|
||||
|
||||
terms = tb.list_terms()
|
||||
assert terms[0].term == "向量数据库"
|
||||
assert terms[0].frequency == 100
|
||||
assert terms[0].pos == "n"
|
||||
|
||||
def test_add_term_strips_whitespace(self):
|
||||
"""add_term 去除首尾空白。"""
|
||||
tb = Termbase()
|
||||
tb.add_term(" 知识图谱 ")
|
||||
|
||||
assert "知识图谱" in tb
|
||||
terms = tb.list_terms()
|
||||
assert terms[0].term == "知识图谱"
|
||||
|
||||
def test_add_empty_term_ignored(self):
|
||||
"""add_term 忽略空字符串。"""
|
||||
tb = Termbase()
|
||||
tb.add_term("")
|
||||
tb.add_term(" ")
|
||||
|
||||
assert len(tb) == 0
|
||||
|
||||
def test_add_duplicate_term_overwrites(self):
|
||||
"""重复添加同一术语覆盖原条目。"""
|
||||
tb = Termbase()
|
||||
tb.add_term("知识图谱", frequency=50)
|
||||
tb.add_term("知识图谱", frequency=100, pos="n")
|
||||
|
||||
assert len(tb) == 1
|
||||
terms = tb.list_terms()
|
||||
assert terms[0].frequency == 100
|
||||
assert terms[0].pos == "n"
|
||||
|
||||
def test_remove_term(self):
|
||||
"""remove_term 删除术语。"""
|
||||
tb = Termbase()
|
||||
tb.add_term("知识图谱")
|
||||
assert len(tb) == 1
|
||||
|
||||
tb.remove_term("知识图谱")
|
||||
assert len(tb) == 0
|
||||
assert "知识图谱" not in tb
|
||||
|
||||
def test_remove_nonexistent_term_no_error(self):
|
||||
"""删除不存在的术语不报错。"""
|
||||
tb = Termbase()
|
||||
tb.remove_term("不存在的术语") # 不应抛异常
|
||||
assert len(tb) == 0
|
||||
|
||||
|
||||
class TestTermbaseLoadFromList:
|
||||
"""load_from_list 测试。"""
|
||||
|
||||
def test_load_from_list(self):
|
||||
"""从字符串列表加载术语。"""
|
||||
tb = Termbase()
|
||||
tb.load_from_list(["知识图谱", "向量数据库", "RAG"])
|
||||
|
||||
assert len(tb) == 3
|
||||
assert "知识图谱" in tb
|
||||
assert "向量数据库" in tb
|
||||
assert "RAG" in tb
|
||||
|
||||
def test_load_from_empty_list(self):
|
||||
"""空列表不添加任何术语。"""
|
||||
tb = Termbase()
|
||||
tb.load_from_list([])
|
||||
assert len(tb) == 0
|
||||
|
||||
|
||||
class TestTermbaseLoadFromFile:
|
||||
"""load_from_file 测试。"""
|
||||
|
||||
def test_load_from_file(self, tmp_path: Path):
|
||||
"""从 jieba 词典文件加载术语。"""
|
||||
dict_file = tmp_path / "terms.txt"
|
||||
dict_file.write_text(
|
||||
"知识图谱 100 n\n向量数据库 100 n\nRAG 50\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
tb = Termbase()
|
||||
tb.load_from_file(str(dict_file))
|
||||
|
||||
assert len(tb) == 3
|
||||
assert "知识图谱" in tb
|
||||
assert "向量数据库" in tb
|
||||
assert "RAG" in tb
|
||||
|
||||
# 验证词频和词性解析
|
||||
terms = {t.term: t for t in tb.list_terms()}
|
||||
assert terms["知识图谱"].frequency == 100
|
||||
assert terms["知识图谱"].pos == "n"
|
||||
assert terms["RAG"].frequency == 50
|
||||
assert terms["RAG"].pos is None
|
||||
|
||||
def test_load_from_file_skips_comments_and_empty(self, tmp_path: Path):
|
||||
"""词典文件中的注释行和空行被跳过。"""
|
||||
dict_file = tmp_path / "terms.txt"
|
||||
dict_file.write_text(
|
||||
"# 这是注释\n\n知识图谱 100 n\n\n# 另一个注释\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
tb = Termbase()
|
||||
tb.load_from_file(str(dict_file))
|
||||
|
||||
assert len(tb) == 1
|
||||
assert "知识图谱" in tb
|
||||
|
||||
def test_load_from_nonexistent_file_raises(self):
|
||||
"""文件不存在时抛 FileNotFoundError。"""
|
||||
tb = Termbase()
|
||||
try:
|
||||
tb.load_from_file("/nonexistent/path/terms.txt")
|
||||
raise AssertionError("Expected FileNotFoundError")
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# jieba 分词集成测试 — 验证术语表对分词的影响
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTermbaseTokenization:
|
||||
"""术语表对 jieba 分词的影响测试。"""
|
||||
|
||||
def test_tokenize_without_termbase(self):
|
||||
"""无术语表时 jieba 默认分词（领域术语可能被错误切分）。"""
|
||||
# 重置 jieba 词典到默认状态
|
||||
jieba.del_word("知识图谱")
|
||||
jieba.del_word("向量数据库")
|
||||
|
||||
tb = Termbase()
|
||||
tokens = tb.tokenize("知识图谱是向量数据库的基础")
|
||||
|
||||
# 无术语表时，"知识图谱" 可能被切分为 "知识" + "图谱"
|
||||
# 注意：jieba 默认词典可能已包含部分常见词，这里只验证分词返回列表
|
||||
assert isinstance(tokens, list)
|
||||
assert len(tokens) > 0
|
||||
|
||||
def test_tokenize_with_termbase(self):
|
||||
"""添加术语表后，领域术语被正确识别为单个 token。"""
|
||||
# 先清除可能存在的自定义词
|
||||
jieba.del_word("知识图谱")
|
||||
jieba.del_word("向量数据库")
|
||||
|
||||
tb = Termbase()
|
||||
tb.add_term("知识图谱")
|
||||
tb.add_term("向量数据库")
|
||||
|
||||
tokens = tb.tokenize("知识图谱是向量数据库的基础")
|
||||
|
||||
# 添加术语后，"知识图谱" 和 "向量数据库" 应作为整体 token 出现
|
||||
assert "知识图谱" in tokens
|
||||
assert "向量数据库" in tokens
|
||||
|
||||
def test_termbase_improves_tokenization(self):
|
||||
"""术语表增强后分词更准确 — 验证领域术语作为整体出现。"""
|
||||
# 测试前先清除
|
||||
jieba.del_word("检索增强生成")
|
||||
|
||||
text = "检索增强生成是RAG的核心技术"
|
||||
|
||||
# 添加术语表
|
||||
tb_after = Termbase()
|
||||
tb_after.add_term("检索增强生成")
|
||||
tokens_after = tb_after.tokenize(text)
|
||||
|
||||
# 添加术语后，"检索增强生成" 应作为整体出现
|
||||
assert "检索增强生成" in tokens_after
|
||||
|
||||
def test_tokenize_empty_string(self):
|
||||
"""空字符串返回空列表。"""
|
||||
tb = Termbase()
|
||||
assert tb.tokenize("") == []
|
||||
|
||||
def test_tokenize_english(self):
|
||||
"""英文文本正常分词。"""
|
||||
tb = Termbase()
|
||||
tokens = tb.tokenize("hello world")
|
||||
assert "hello" in tokens
|
||||
assert "world" in tokens
|
||||
|
||||
def test_remove_term_restores_default_tokenization(self):
|
||||
"""删除术语后 jieba 恢复默认分词（术语不再作为整体）。"""
|
||||
# 添加术语
|
||||
tb = Termbase()
|
||||
tb.add_term("测试术语XYZ")
|
||||
tokens_with = tb.tokenize("测试术语XYZ很重要")
|
||||
assert "测试术语XYZ" in tokens_with
|
||||
|
||||
# 删除术语
|
||||
tb.remove_term("测试术语XYZ")
|
||||
# 删除后，"测试术语XYZ" 不再作为整体（可能被切分）
|
||||
# 注意：jieba 删除词后可能仍缓存，但 del_word 会从词典移除
|
||||
# 这里验证术语已从 Termbase 字典中删除
|
||||
assert "测试术语XYZ" not in tb
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 检索召回率提升模拟测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTermbaseRetrievalImprovement:
|
||||
"""术语表增强后检索召回率提升的模拟测试。"""
|
||||
|
||||
def test_termbase_improves_keyword_matching(self):
|
||||
"""术语表增强后，关键词匹配更准确。
|
||||
|
||||
模拟场景：用户查询"知识图谱"，文档中包含"知识图谱"。
|
||||
无术语表时 jieba 可能将查询切分为"知识"+"图谱"，
|
||||
导致匹配精度下降；有术语表时整体匹配。
|
||||
"""
|
||||
# 清除可能的自定义词
|
||||
jieba.del_word("知识图谱")
|
||||
|
||||
query = "知识图谱"
|
||||
doc = "知识图谱是人工智能的重要分支"
|
||||
|
||||
# 有术语表 — "知识图谱" 作为整体
|
||||
tb_with = Termbase()
|
||||
tb_with.add_term("知识图谱")
|
||||
query_tokens_with = set(tb_with.tokenize(query))
|
||||
doc_tokens_with = set(tb_with.tokenize(doc))
|
||||
|
||||
# 有术语表时，查询和文档共享 "知识图谱" token
|
||||
# 无术语表时，可能共享 "知识" 和 "图谱"（如果被切分）
|
||||
# 关键验证：有术语表时 "知识图谱" 在两边都出现
|
||||
assert "知识图谱" in query_tokens_with
|
||||
assert "知识图谱" in doc_tokens_with
|
||||
|
||||
# 交集应包含 "知识图谱"
|
||||
intersection_with = query_tokens_with & doc_tokens_with
|
||||
assert "知识图谱" in intersection_with
|
||||
|
||||
def test_multiple_terms_improve_coverage(self):
|
||||
"""多个领域术语同时增强分词。"""
|
||||
# 清除可能的自定义词
|
||||
for term in ["知识图谱", "向量数据库", "嵌入模型"]:
|
||||
jieba.del_word(term)
|
||||
|
||||
tb = Termbase()
|
||||
tb.load_from_list(["知识图谱", "向量数据库", "嵌入模型"])
|
||||
|
||||
text = "知识图谱通常使用向量数据库和嵌入模型构建"
|
||||
tokens = tb.tokenize(text)
|
||||
|
||||
# 所有领域术语都应作为整体 token 出现
|
||||
assert "知识图谱" in tokens
|
||||
assert "向量数据库" in tokens
|
||||
assert "嵌入模型" in tokens
|
||||
Loading…
Reference in New Issue