feat(retrieval): 实现 Embedding/Vector/RRF/Reranker 混合检索引擎

Co-Authored-By: Claude <noreply@anthropic.com>

backend/app/constants.py

@@ -0,0 +1,8 @@
+"""跨层共享的常量。
+
+放在这里是为了让 database（建向量表）与 retrieval（生成向量）共享同一个维度值，
+避免二者各自硬编码导致不一致。真实 BGE-M3 接入后把维度改为 1024 并重建向量索引即可。
+"""
+
+# 轻量哈希向量的维度；vec0 虚拟表建表时按此维度固定列宽。
+EMBEDDING_DIM = 128

backend/app/retrieval/__init__.py

@@ -0,0 +1 @@
+"""Retrieval Core：Embedding、VectorStore、RRF、Reranker 与混合检索引擎。"""

backend/app/retrieval/embedding.py

@@ -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]

backend/app/retrieval/engine.py

@@ -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())

backend/app/retrieval/hybrid.py

@@ -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]

backend/app/retrieval/reranker.py

@@ -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

backend/app/retrieval/vectorstore.py

@@ -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()