211 lines
7.4 KiB
Python
211 lines
7.4 KiB
Python
"""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
|
||
|
||
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: object | None = None # 延迟初始化，避免 import 失败
|
||
|
||
def _get_reranker(self) -> object | None:
|
||
"""延迟加载 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) -> object:
|
||
"""构建 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, object] = {
|
||
"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) -> object:
|
||
"""构建 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, object] = {
|
||
"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"]
|