diff --git a/backend/app/constants.py b/backend/app/constants.py new file mode 100644 index 0000000..1fb5a7a --- /dev/null +++ b/backend/app/constants.py @@ -0,0 +1,8 @@ +"""跨层共享的常量。 + +放在这里是为了让 database(建向量表)与 retrieval(生成向量)共享同一个维度值, +避免二者各自硬编码导致不一致。真实 BGE-M3 接入后把维度改为 1024 并重建向量索引即可。 +""" + +# 轻量哈希向量的维度;vec0 虚拟表建表时按此维度固定列宽。 +EMBEDDING_DIM = 128 diff --git a/backend/app/retrieval/__init__.py b/backend/app/retrieval/__init__.py new file mode 100644 index 0000000..9e24cc2 --- /dev/null +++ b/backend/app/retrieval/__init__.py @@ -0,0 +1 @@ +"""Retrieval Core:Embedding、VectorStore、RRF、Reranker 与混合检索引擎。""" diff --git a/backend/app/retrieval/embedding.py b/backend/app/retrieval/embedding.py new file mode 100644 index 0000000..6449d4e --- /dev/null +++ b/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] diff --git a/backend/app/retrieval/engine.py b/backend/app/retrieval/engine.py new file mode 100644 index 0000000..8f28e47 --- /dev/null +++ b/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()) diff --git a/backend/app/retrieval/hybrid.py b/backend/app/retrieval/hybrid.py new file mode 100644 index 0000000..62eb38c --- /dev/null +++ b/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] diff --git a/backend/app/retrieval/reranker.py b/backend/app/retrieval/reranker.py new file mode 100644 index 0000000..15548d6 --- /dev/null +++ b/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 diff --git a/backend/app/retrieval/vectorstore.py b/backend/app/retrieval/vectorstore.py new file mode 100644 index 0000000..2dd50e4 --- /dev/null +++ b/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()