kplam
6ae91679e2
- .env.example 嵌入默认改为本地 BGE-M3（EMBEDDING_PROVIDER=local, dim=1024），NAS 全新部署即用 - 新增 app/services/reranker.py：基于 Ollama /api/rerank 的 cross-encoder 重排服务 - retriever 在 chunk 候选阶段接入语义精排，调用失败优雅降级为 RRF 顺序 - docker-compose 拉取 qwen3-reranker:0.6b，OLLAMA_MAX_LOADED_MODELS 提至 3 - .gitignore 排除 .workbuddy/ 与 .trae/（防止误提交项目记忆与 IDE spec） - 新增 reranker 单元与集成测试，全量测试 529 passed
80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
"""语义重排服务（cross-encoder reranker）
|
||
|
||
调用 Ollama 的 `/api/rerank` 端点（Qwen3-Reranker 等重排模型），对检索召回的
|
||
chunk 候选池做精排，提升最终 Top-K 的相关性。相比纯向量/RRF 融合，cross-encoder
|
||
以 (query, doc) 联合编码，能捕捉字词不匹配的语义关联，典型带来约 10% 的精度提升。
|
||
|
||
工厂 `create_reranker_service()` 按 `settings.reranker_enabled` 返回实例或 None：
|
||
关闭时上层检索链路退化为原有的 RRF 融合结果，行为完全不变。
|
||
"""
|
||
|
||
from typing import Protocol, runtime_checkable
|
||
|
||
import httpx
|
||
import structlog
|
||
|
||
from app.config import settings
|
||
|
||
logger = structlog.get_logger()
|
||
|
||
|
||
@runtime_checkable
|
||
class RerankerService(Protocol):
|
||
"""重排服务统一接口"""
|
||
|
||
async def rerank(self, query: str, documents: list[str]) -> list[float]:
|
||
"""对候选文档按与 query 的相关性打分
|
||
|
||
Args:
|
||
query: 查询文本
|
||
documents: 候选文档文本列表（按输入顺序对齐）
|
||
|
||
Returns:
|
||
与 documents 等长的相关性分数列表（分数越高越相关）
|
||
"""
|
||
...
|
||
|
||
|
||
class OllamaRerankerService:
|
||
"""基于 Ollama /api/rerank 的重排服务"""
|
||
|
||
def __init__(self, base_url: str, model: str, timeout: float = 30.0) -> None:
|
||
self.base_url = base_url.rstrip("/")
|
||
self.model = model
|
||
self.timeout = timeout
|
||
|
||
async def rerank(self, query: str, documents: list[str]) -> list[float]:
|
||
if not documents:
|
||
return []
|
||
url = f"{self.base_url}/api/rerank"
|
||
payload = {"model": self.model, "query": query, "documents": documents}
|
||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||
resp = await client.post(url, json=payload)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
|
||
results = data.get("results", [])
|
||
# 按 Ollama 返回的 index 映射 relevance_score；缺失项补 0，保证输出与输入等长
|
||
scores: list[float] = [0.0] * len(documents)
|
||
for item in results:
|
||
idx = item.get("index")
|
||
score = float(item.get("relevance_score", 0.0))
|
||
if idx is not None and 0 <= idx < len(documents):
|
||
scores[idx] = score
|
||
logger.debug("重排完成", model=self.model, candidates=len(documents))
|
||
return scores
|
||
|
||
|
||
def create_reranker_service() -> RerankerService | None:
|
||
"""按 settings.reranker_enabled 创建重排服务实例
|
||
|
||
关闭时返回 None，上层据此跳过精排、退化为 RRF 融合结果。
|
||
"""
|
||
if not settings.reranker_enabled:
|
||
return None
|
||
return OllamaRerankerService(
|
||
base_url=settings.ollama_base_url,
|
||
model=settings.reranker_model,
|
||
timeout=settings.reranker_timeout,
|
||
)
|