feat(retrieval): 实现 Embedding/Vector/RRF/Reranker 混合检索引擎
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Claude
@@ -0,0 +1,8 @@
|
||||
"""跨层共享的常量。
|
||||
|
||||
放在这里是为了让 database（建向量表）与 retrieval（生成向量）共享同一个维度值，
|
||||
避免二者各自硬编码导致不一致。真实 BGE-M3 接入后把维度改为 1024 并重建向量索引即可。
|
||||
"""
|
||||
|
||||
# 轻量哈希向量的维度；vec0 虚拟表建表时按此维度固定列宽。
|
||||
EMBEDDING_DIM = 128
|
||||
@@ -0,0 +1 @@
|
||||
"""Retrieval Core：Embedding、VectorStore、RRF、Reranker 与混合检索引擎。"""
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Embedding 统一接口与轻量实现。
|
||||
|
||||
真实默认是本地 BGE-M3 类模型，但第一阶段先跑通链路，这里用确定性的特征哈希向量代替。
|
||||
后续接入真实模型时实现同样的 EmbeddingProvider 接口替换即可，上层检索逻辑不变。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import math
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
from app.constants import EMBEDDING_DIM
|
||||
from app.textutils import tokens
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class EmbeddingProvider(Protocol):
|
||||
"""统一 Embedding 接口（与文档一致）。"""
|
||||
|
||||
model_id: str
|
||||
dim: int
|
||||
|
||||
async def embed_documents(self, texts: list[str]) -> list[list[float]]: ...
|
||||
async def embed_query(self, query: str) -> list[float]: ...
|
||||
|
||||
|
||||
class HashEmbeddingProvider:
|
||||
"""轻量确定性向量：特征哈希 + 符号 + L2 归一化。
|
||||
|
||||
同一文本永远得到相同向量，可离线复现、无外部依赖。向量维度为 EMBEDDING_DIM，
|
||||
与 vec_blocks 建表维度一致。
|
||||
"""
|
||||
|
||||
model_id = "hash-v1"
|
||||
dim = EMBEDDING_DIM
|
||||
|
||||
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
return [self._embed(text) for text in texts]
|
||||
|
||||
async def embed_query(self, query: str) -> list[float]:
|
||||
return self._embed(query)
|
||||
|
||||
def _embed(self, text: str) -> list[float]:
|
||||
vec = [0.0] * self.dim
|
||||
for tok in tokens(text):
|
||||
digest = hashlib.sha256(tok.encode("utf-8")).digest()
|
||||
index = int.from_bytes(digest[:4], "little") % self.dim
|
||||
sign = 1.0 if digest[4] % 2 == 0 else -1.0
|
||||
vec[index] += sign
|
||||
norm = math.sqrt(sum(v * v for v in vec)) or 1.0
|
||||
return [v / norm for v in vec]
|
||||
@@ -0,0 +1,164 @@
|
||||
"""混合检索引擎：编排 FTS5 / Vector / RRF / Reranker / Metadata Filter / Citation。
|
||||
|
||||
对调用方（搜索页、RAG Engine、Agent Tool）暴露统一的 search(request) -> SearchResponse。
|
||||
引擎只依赖 VectorStore / EmbeddingProvider / RerankerProvider 抽象与 Repository，
|
||||
不直接拼接 vec0 内部 SQL，也不向前端输出聊天文本。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app import repository
|
||||
from app.contracts import (
|
||||
Citation,
|
||||
PageMeta,
|
||||
SearchMode,
|
||||
SearchRequest,
|
||||
SearchResponse,
|
||||
SearchResult,
|
||||
)
|
||||
from app.repository import BlockHit
|
||||
from app.retrieval.embedding import EmbeddingProvider, HashEmbeddingProvider
|
||||
from app.retrieval.hybrid import normalize_scores, rrf_fuse
|
||||
from app.retrieval.reranker import LexicalReranker, RankedCandidate, RerankerProvider
|
||||
from app.retrieval.vectorstore import SqliteVecStore, VectorStore
|
||||
from app.textutils import make_snippet, match_query
|
||||
|
||||
# 每个通道的候选池大小；真实规模上来后按 Retrieval Config 调整
|
||||
CANDIDATE_POOL = 50
|
||||
|
||||
|
||||
class RetrievalEngine:
|
||||
def __init__(
|
||||
self,
|
||||
embedding: EmbeddingProvider,
|
||||
reranker: RerankerProvider,
|
||||
vector_store: VectorStore,
|
||||
) -> None:
|
||||
self.embedding = embedding
|
||||
self.reranker = reranker
|
||||
self.vector_store = vector_store
|
||||
|
||||
async def search(self, request: SearchRequest) -> SearchResponse:
|
||||
# 1. 按模式收集候选（FTS 与 Vector 各产出「按相关性降序」的 block_id 列表）
|
||||
fts_ranked: list[str] = []
|
||||
vec_ranked: list[str] = []
|
||||
fts_scores: dict[str, float] = {}
|
||||
vec_scores: dict[str, float] = {}
|
||||
|
||||
if request.mode in (SearchMode.fts, SearchMode.hybrid):
|
||||
match = match_query(request.query)
|
||||
if match:
|
||||
fts_hits = repository.fts_search(match, CANDIDATE_POOL)
|
||||
fts_ranked = [h.block_id for h in fts_hits]
|
||||
# bm25 越小越相关，取反后统一为「越大越相关」
|
||||
fts_scores = {h.block_id: -h.bm25 for h in fts_hits}
|
||||
|
||||
if request.mode in (SearchMode.vector, SearchMode.hybrid):
|
||||
query_vec = await self.embedding.embed_query(request.query)
|
||||
vec_hits = await self.vector_store.search(query_vec, top_k=CANDIDATE_POOL)
|
||||
vec_ranked = [v.id for v in vec_hits]
|
||||
vec_scores = {v.id: v.score for v in vec_hits}
|
||||
|
||||
if request.mode == SearchMode.fts:
|
||||
candidate_scores = fts_scores
|
||||
elif request.mode == SearchMode.vector:
|
||||
candidate_scores = vec_scores
|
||||
else: # hybrid：RRF 融合
|
||||
candidate_scores = rrf_fuse([fts_ranked, vec_ranked])
|
||||
|
||||
if not candidate_scores:
|
||||
return self._empty(request)
|
||||
|
||||
# 2. 取完整 Block 上下文（用于过滤、摘要与 Citation 定位）
|
||||
hits = {h.block_id: h for h in repository.get_block_hits(list(candidate_scores.keys()))}
|
||||
|
||||
# 3. Metadata Filter
|
||||
filtered = [h for h in hits.values() if self._matches(h, request)]
|
||||
if not filtered:
|
||||
return self._empty(request)
|
||||
|
||||
# 4. 排序 / 精排
|
||||
if request.mode == SearchMode.hybrid:
|
||||
candidates = [
|
||||
RankedCandidate(block_id=h.block_id, score=candidate_scores[h.block_id], text=h.content)
|
||||
for h in filtered
|
||||
]
|
||||
ranked = await self.reranker.rerank(request.query, candidates)
|
||||
ordered = [(c.block_id, c.score) for c in ranked]
|
||||
else:
|
||||
ordered = sorted(
|
||||
((h.block_id, candidate_scores[h.block_id]) for h in filtered),
|
||||
key=lambda item: -item[1],
|
||||
)
|
||||
|
||||
ordered = normalize_scores(ordered)
|
||||
|
||||
# 5. 分页
|
||||
total = len(ordered)
|
||||
page = ordered[request.offset : request.offset + request.limit]
|
||||
items = [self._build_result(hits[block_id], request, score) for block_id, score in page]
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
mode=request.mode,
|
||||
items=items,
|
||||
page=PageMeta(total=total, limit=request.limit, offset=request.offset),
|
||||
)
|
||||
|
||||
def _matches(self, hit: BlockHit, request: SearchRequest) -> bool:
|
||||
if request.folders and hit.folder not in request.folders:
|
||||
return False
|
||||
if request.note_ids and hit.note_id not in request.note_ids:
|
||||
return False
|
||||
if request.tags and not (set(hit.tags) & set(request.tags)):
|
||||
return False
|
||||
if request.created_from and _utc(hit.created_at) < _utc(request.created_from):
|
||||
return False
|
||||
if request.created_to and _utc(hit.created_at) > _utc(request.created_to):
|
||||
return False
|
||||
if request.updated_from and _utc(hit.updated_at) < _utc(request.updated_from):
|
||||
return False
|
||||
if request.updated_to and _utc(hit.updated_at) > _utc(request.updated_to):
|
||||
return False
|
||||
return True
|
||||
|
||||
def _build_result(self, hit: BlockHit, request: SearchRequest, score: float) -> SearchResult:
|
||||
citation = Citation(
|
||||
citation_id=f"cit_{hit.block_id}",
|
||||
note_id=hit.note_id,
|
||||
block_id=hit.block_id,
|
||||
file_path=hit.file_path,
|
||||
heading_path=hit.heading_path,
|
||||
start_offset=hit.start_offset,
|
||||
end_offset=hit.end_offset,
|
||||
)
|
||||
snippet = make_snippet(hit.content, request.query) if request.include_snippet else None
|
||||
return SearchResult(
|
||||
note_id=hit.note_id,
|
||||
block_id=hit.block_id,
|
||||
title=hit.title,
|
||||
file_path=hit.file_path,
|
||||
heading_path=hit.heading_path,
|
||||
snippet=snippet,
|
||||
score=score,
|
||||
citation=citation,
|
||||
)
|
||||
|
||||
def _empty(self, request: SearchRequest) -> SearchResponse:
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
mode=request.mode,
|
||||
page=PageMeta(total=0, limit=request.limit, offset=request.offset),
|
||||
)
|
||||
|
||||
|
||||
def _utc(dt: datetime) -> datetime:
|
||||
"""把时间统一到 naive UTC 再比较，避免 aware/naive 混用报错。"""
|
||||
if dt.tzinfo is None:
|
||||
return dt
|
||||
return dt.astimezone(timezone.utc).replace(tzinfo=None)
|
||||
|
||||
|
||||
# 默认引擎实例：轻量实现跑通链路，后续可替换真实模型实现
|
||||
engine = RetrievalEngine(HashEmbeddingProvider(), LexicalReranker(), SqliteVecStore())
|
||||
@@ -0,0 +1,25 @@
|
||||
"""RRF 排名融合与分数归一化。"""
|
||||
|
||||
|
||||
def rrf_fuse(ranked_lists: list[list[str]], k: int = 60) -> dict[str, float]:
|
||||
"""Reciprocal Rank Fusion：对多个「按相关性降序」的 block_id 列表做排名融合。
|
||||
|
||||
每个 block 的融合分 = Σ 1/(k + rank)，rank 从 1 开始。返回 block_id -> 融合分。
|
||||
"""
|
||||
scores: dict[str, float] = {}
|
||||
for ids in ranked_lists:
|
||||
for rank, block_id in enumerate(ids, start=1):
|
||||
scores[block_id] = scores.get(block_id, 0.0) + 1.0 / (k + rank)
|
||||
return scores
|
||||
|
||||
|
||||
def normalize_scores(items: list[tuple[str, float]]) -> list[tuple[str, float]]:
|
||||
"""把 (block_id, score) 列表 min-max 归一化到 [0,1]，score 越大越相关。"""
|
||||
if not items:
|
||||
return []
|
||||
values = [score for _, score in items]
|
||||
lo, hi = min(values), max(values)
|
||||
span = hi - lo
|
||||
if span == 0:
|
||||
return [(block_id, 1.0) for block_id, _ in items]
|
||||
return [(block_id, round((score - lo) / span, 6)) for block_id, score in items]
|
||||
@@ -0,0 +1,60 @@
|
||||
"""Reranker 统一接口与轻量实现。
|
||||
|
||||
真实默认是 BGE reranker 类 Cross-Encoder，第一阶段先用词面重叠 + 原始分数加权的
|
||||
确定性精排跑通链路；后续替换实现即可。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
from app.textutils import tokens
|
||||
|
||||
|
||||
@dataclass
|
||||
class RankedCandidate:
|
||||
block_id: str
|
||||
score: float
|
||||
text: str = "" # 块正文，供轻量精排计算词面重叠
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class RerankerProvider(Protocol):
|
||||
"""统一 Reranker 接口：输入候选块，输出按相关性重排后的候选块。"""
|
||||
|
||||
model_id: str
|
||||
|
||||
async def rerank(self, query: str, candidates: list[RankedCandidate]) -> list[RankedCandidate]: ...
|
||||
|
||||
|
||||
class LexicalReranker:
|
||||
"""轻量精排：query 与块正文的词面重叠度，与归一化后的原始分数加权求和。"""
|
||||
|
||||
model_id = "lexical-v1"
|
||||
|
||||
def __init__(self, lexical_weight: float = 0.5) -> None:
|
||||
self.lexical_weight = lexical_weight
|
||||
|
||||
async def rerank(self, query: str, candidates: list[RankedCandidate]) -> list[RankedCandidate]:
|
||||
if not candidates:
|
||||
return []
|
||||
|
||||
# 把原始分数（RRF 等）归一化到 [0,1]，便于与重叠度同量纲加权
|
||||
scores = [c.score for c in candidates]
|
||||
lo, hi = min(scores), max(scores)
|
||||
span = (hi - lo) or 1.0
|
||||
|
||||
query_tokens = set(tokens(query))
|
||||
ranked: list[RankedCandidate] = []
|
||||
for c in candidates:
|
||||
norm = (c.score - lo) / span
|
||||
if query_tokens:
|
||||
overlap = len(query_tokens & set(tokens(c.text))) / len(query_tokens)
|
||||
else:
|
||||
overlap = 0.0
|
||||
final = self.lexical_weight * overlap + (1 - self.lexical_weight) * norm
|
||||
ranked.append(RankedCandidate(block_id=c.block_id, score=final, text=c.text))
|
||||
|
||||
ranked.sort(key=lambda c: c.score, reverse=True)
|
||||
return ranked
|
||||
@@ -0,0 +1,86 @@
|
||||
"""VectorStore 统一接口与 sqlite-vec 实现。
|
||||
|
||||
vec0 虚拟表返回的 distance 是欧氏距离（非平方）。入库前向量已做 L2 归一化，
|
||||
因此 distance² = 2(1-cos)，余弦相似度 = 1 - distance² / 2。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol, runtime_checkable
|
||||
|
||||
import sqlite_vec
|
||||
|
||||
from app.database.db import connect, transaction
|
||||
|
||||
|
||||
@dataclass
|
||||
class VectorRecord:
|
||||
id: str
|
||||
vector: list[float]
|
||||
|
||||
|
||||
@dataclass
|
||||
class VectorHit:
|
||||
id: str
|
||||
score: float # 余弦相似度 [0,1]
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class VectorStore(Protocol):
|
||||
"""统一向量存储接口（与文档一致）。上层只依赖此抽象，不读 vec0 内部表。"""
|
||||
|
||||
async def upsert(self, records: list[VectorRecord]) -> None: ...
|
||||
async def delete(self, ids: list[str]) -> None: ...
|
||||
async def search(self, vector: list[float], *, top_k: int) -> list[VectorHit]: ...
|
||||
|
||||
|
||||
class SqliteVecStore:
|
||||
"""sqlite-vec 默认实现。"""
|
||||
|
||||
async def upsert(self, records: list[VectorRecord]) -> None:
|
||||
if not records:
|
||||
return
|
||||
conn = connect()
|
||||
try:
|
||||
with transaction(conn):
|
||||
for record in records:
|
||||
conn.execute(
|
||||
"INSERT INTO vec_blocks (block_id, embedding) VALUES (?, ?)",
|
||||
(record.id, sqlite_vec.serialize_float32(record.vector)),
|
||||
)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
async def delete(self, ids: list[str]) -> None:
|
||||
if not ids:
|
||||
return
|
||||
conn = connect()
|
||||
try:
|
||||
with transaction(conn):
|
||||
for bid in ids:
|
||||
conn.execute("DELETE FROM vec_blocks WHERE block_id = ?", (bid,))
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
async def search(self, vector: list[float], *, top_k: int) -> list[VectorHit]:
|
||||
conn = connect()
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"SELECT block_id, distance FROM vec_blocks WHERE embedding MATCH ? AND k = ?",
|
||||
(sqlite_vec.serialize_float32(vector), top_k),
|
||||
).fetchall()
|
||||
return [
|
||||
VectorHit(id=row["block_id"], score=max(0.0, 1.0 - row["distance"] ** 2 / 2.0))
|
||||
for row in rows
|
||||
]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
async def clear(self) -> None:
|
||||
conn = connect()
|
||||
try:
|
||||
with transaction(conn):
|
||||
conn.execute("DELETE FROM vec_blocks")
|
||||
finally:
|
||||
conn.close()
|
||||
Reference in New Issue
Block a user