retriever.py

"""
检索模块：BM25 关键词检索 + ChromaDB 向量检索 + RRF 融合

混合检索的核心思想：
- 向量检索擅长语义匹配("汽车" 能匹配 "轿车")
- BM25 擅长精确匹配(错误码、专有名词)
- RRF 把两路结果按排名融合，互补短板

【面试重点】混合检索 vs 单路检索：
- 单路向量检索：精确关键词容易丢(搜 "ERR_0451" 可能匹配不到)
- 单路 BM25：不懂语义(搜 "怎么部署" 匹配不到 "安装步骤")
- 混合检索 + RRF：两路互补，且不需要对不同量纲的分数做归一化
"""

import jieba
import numpy as np
from rank_bm25 import BM25Okapi
import chromadb

from document_loader import Document
from embedder import embed_texts
from vector_store import get_collection
from reranker import rerank
from models import SearchResult


# ============ 中英文分词 ============

# 停用词：出现频率极高但没有检索价值的词
# BM25 的 IDF 会自动降低高频词的权重，但提前过滤能减少索引大小、加快速度
# 这里只列最常见的，不需要太全(IDF 兜底)
STOP_WORDS = {
    # 中文停用词
    "的", "了", "在", "是", "我", "有", "和", "就",
    "不", "人", "都", "一", "一个", "上", "也", "很",
    "到", "说", "要", "去", "你", "会", "着", "没有",
    "看", "好", "自己", "这", "他", "她", "它",
    # 英文停用词
    "the", "a", "an", "is", "are", "was", "were",
    "in", "on", "at", "to", "for", "of", "with",
    "and", "or", "but", "not", "this", "that",
    "it", "be", "as", "by", "from",
}


def tokenize(text: str) -> list[str]:
    """
    中英文混合分词。

    处理流程：
    1. jieba 分词(中文按词切，英文按连续字母保留)
    2. 英文统一小写("MCP" 和 "mcp" 视为同一个词)
    3. 过滤停用词和纯标点

    参数：
        text: 原始文本

    返回：
        分词后的词列表，如 ["chromadb", "嵌入式", "向量", "数据库"]
    """
    # jieba.lcut() 返回分词列表
    # "ChromaDB 是向量数据库" → ["ChromaDB", " ", "是", "向量", "数据库"]
    raw_tokens = jieba.lcut(text)

    tokens = []
    for token in raw_tokens:
        # strip() 去掉首尾空白
        token = token.strip()
        if not token:
            continue

        # 英文统一小写，让 "MCP" == "mcp"
        token = token.lower()

        # 过滤停用词
        if token in STOP_WORDS:
            continue

        # 过滤纯标点符号(长度 1 且不是字母/数字/中文的)
        # 保留单个中文字(如 "图")和单个字母/数字(如 "a" 已被停用词过滤)
        if len(token) == 1 and not token.isalnum() and not '\u4e00' <= token <= '\u9fff':
            continue

        tokens.append(token)

    return tokens


# ============ BM25 关键词检索 ============
class BM25Index:
    """
    BM25 关键词检索索引。

    使用方式：
        index = BM25Index()
        index.build(chunks)           # 启动时建一次索引
        results = index.search(query) # 每次搜索直接用

    【面试重点】BM25 基于 TF-IDF 思想：
    - TF(词频)：某个词在当前 chunk 出现越多，分数越高
    - IDF(逆文档频率)：某个词在越少的 chunk 里出现，说明越稀有，权重越高
    - 还会用文档长度做归一化，避免长 chunk 天然占优
    """

    def __init__(self):
        self._bm25 = None           # BM25Okapi 实例(rank_bm25 库)
        self._chunks = []            # 原始 chunk 列表，搜索时要返回原文和 metadata
        self._chunk_ids = []         # 每个 chunk 的 ID，跟 ChromaDB 里的对应

    def build(self, chunks: list[Document]) -> None:
        """
        用分块文档列表建 BM25 索引。

        流程：
        1. 保存 chunks(搜索时要取原文)
        2. 对每个 chunk 分词，得到二维列表
        3. 传给 BM25Okapi 建索引(内部自动建倒排索引、算 IDF)

        参数：
            chunks: split_documents() 输出的 Document 列表
        """
        if not chunks:
            print("BM25: 没有文档，跳过建索引")
            return

        self._chunks = chunks

        # 生成每个 chunk 的 ID(跟 vector_store.py 里的 make_chunk_id 逻辑一致)
        import os
        self._chunk_ids = []
        for chunk in chunks:
            filepath = chunk.metadata.get("filepath", "unknown")
            chunk_index = chunk.metadata.get("chunk_index", 0)
            filename = os.path.basename(filepath)
            self._chunk_ids.append(f"{filename}_chunk_{chunk_index}")

        # 对每个 chunk 的文本做分词
        # corpus 是二维列表：[[词1, 词2, ...], [词1, 词2, ...], ...]
        # 每个内层列表是一个 chunk 的分词结果
        corpus = []
        for chunk in chunks:
            tokens = tokenize(chunk.content)
            corpus.append(tokens)

        # BM25Okapi 拿到分词后的 corpus，内部自动建倒排索引
        # 之后调 get_scores() 就能算出每个 chunk 跟 query 的相关度
        self._bm25 = BM25Okapi(corpus)

        print(f"BM25: 索引建完，共 {len(chunks)} 个分块")

    def search(self, query: str, k: int = 5) -> list[SearchResult]:
        """
        BM25 关键词检索，返回 top-k 结果。

        流程：
        1. 对 query 分词
        2. get_scores() 算每个 chunk 的 BM25 分数
        3. 按分数从高到低排序，取前 k 个
        4. 包装成 SearchResult 返回

        参数：
            query: 用户的搜索文本
            k: 返回结果数

        返回：
            SearchResult 列表，按分数从高到低，rank 从 1 开始
        """
        if self._bm25 is None:
            print("BM25: 索引未建立，请先调用 build()")
            return []

        # 第一步：query 分词(跟建索引时用同一个 tokenize 函数)
        query_tokens = tokenize(query)
        if not query_tokens:
            return []

        # 第二步：算每个 chunk 的 BM25 分数
        # 返回一个数组，长度 = chunk 总数
        # scores[i] = query 对第 i 个 chunk 的 BM25 相关度分数
        scores = self._bm25.get_scores(query_tokens)

        # 第三步：按分数从高到低排序
        # np.argsort() 返回排序后的下标数组(默认升序)
        # [::-1] 反转变成降序
        # 例：scores = [0.0, 5.1, 2.3] → argsort 升序 [0,2,1] → 反转 [1,2,0]
        sorted_indices = np.argsort(scores)[::-1]

        # 取前 k 个(如果不足 k 个就全取)
        top_k_indices = sorted_indices[:k]

        # 第四步：包装成 SearchResult
        results = []
        for rank, idx in enumerate(top_k_indices, start=1):
            # 跳过分数为 0 的结果(完全不匹配的)
            if scores[idx] <= 0:
                break

            results.append(SearchResult(
                chunk_id=self._chunk_ids[idx],
                content=self._chunks[idx].content,
                metadata=self._chunks[idx].metadata,
                score=float(scores[idx]),
                rank=rank,
            ))

        return results


# ============ 向量检索 ============
def vector_search(
    query: str,
    collection: chromadb.Collection = None,
    k: int = 5,
) -> list[SearchResult]:
    """
    ChromaDB 向量检索：把 query 做 embedding，在向量空间找最近邻。

    流程：
    1. 把 query 文本做 embedding，得到一个 1024 维向量
    2. 调 collection.query() 在 HNSW 索引里找最近的 k 个向量
    3. 包装成 SearchResult 返回

    参数：
        query: 用户的搜索文本
        collection: ChromaDB Collection，不传则用默认的
        k: 返回结果数

    返回：
        SearchResult 列表，按相似度从高到低，rank 从 1 开始

    【面试重点】向量检索 vs BM25：
    - 向量检索用的是 HNSW(多层图索引)，O(log n) 近似最近邻
    - BM25 用的是倒排索引，精确匹配关键词
    - 两者互补：向量懂语义，BM25 精确匹配
    """
    if collection is None:
        collection = get_collection()

    # 第一步：query 做 embedding
    # embed_texts 接受列表，返回二维数组，[0] 取第一个(也是唯一一个)向量
    query_vector = embed_texts([query])[0]

    # 第二步：ChromaDB 向量检索
    # query_embeddings: 要搜的向量(二维列表格式)
    # n_results: 返回几条结果
    # include: 同时返回原文、元数据和距离
    #
    # ChromaDB 返回的 distances 是距离(越小越相似)
    # 因为我们用的 cosine space，距离 = 1 - 余弦相似度
    # 所以相似度 = 1 - distance
    results = collection.query(
        query_embeddings=[query_vector.tolist()],
        n_results=k,
        include=["documents", "metadatas", "distances"],
    )

    # 第三步：包装成 SearchResult
    # collection.query() 返回的是嵌套列表(支持批量查询)
    # 我们只查了一个 query，所以取 [0]
    search_results = []
    ids = results["ids"][0]
    documents = results["documents"][0]
    metadatas = results["metadatas"][0]
    distances = results["distances"][0]

    for rank, (chunk_id, content, metadata, distance) in enumerate(
        zip(ids, documents, metadatas, distances), start=1
    ):
        # 余弦相似度 = 1 - cosine distance
        similarity = 1.0 - distance

        search_results.append(SearchResult(
            chunk_id=chunk_id,
            content=content,
            metadata=metadata,
            score=similarity,
            rank=rank,
        ))

    return search_results


# ============ RRF 融合 ============
def rrf_fuse(
    results_list: list[list[SearchResult]],
    k: int = 5,
    rrf_k: int = 60,
    weights: list[float] = None,
) -> list[SearchResult]:
    """
    Reciprocal Rank Fusion：把多路检索结果按排名融合。

    公式：RRF_score(doc) = Σ weight_i / (rrf_k + rank_i)

    - 两路都靠前的 chunk → 总分最高
    - 只出现在一路的 chunk → 只算那一路的分
    - rrf_k=60 起平滑作用，防止排名第 1 的权重过大
    - weights 可以调整两路的偏好(默认平等)

    参数：
        results_list: 多路检索结果，如 [bm25_results, vector_results]
        k: 最终返回几条结果
        rrf_k: RRF 常数(默认 60)
        weights: 每路的权重，如 [0.5, 0.5]，默认平等

    返回：
        融合后的 top-k 结果，按 RRF 分数从高到低

    【面试重点】为什么用 RRF 而不直接加分数？
    - BM25 分数和余弦相似度量纲不同，直接加没有意义
    - RRF 只看排名不看分数，天然解决了跨检索器的分数不可比问题
    """
    # 默认权重：每路平等
    if weights is None:
        weights = [1.0] * len(results_list)

    # 用字典收集每个 chunk 的 RRF 总分
    # key = chunk_id, value = {"score": 总分, "result": SearchResult对象}
    # 同一个 chunk 可能同时出现在多路结果里，分数要累加
    fused_scores = {}

    for route_idx, results in enumerate(results_list):
        weight = weights[route_idx]

        for result in results:
            # RRF 公式：weight / (rrf_k + rank)
            rrf_score = weight / (rrf_k + result.rank)

            if result.chunk_id in fused_scores:
                # 这个 chunk 已经在别的路出现过，累加分数
                fused_scores[result.chunk_id]["score"] += rrf_score
            else:
                # 第一次见到这个 chunk，记录下来
                fused_scores[result.chunk_id] = {
                    "score": rrf_score,
                    "result": result,
                }

    # 按 RRF 总分从高到低排序
    sorted_items = sorted(
        fused_scores.values(),
        key=lambda x: x["score"],
        reverse=True,
    )

    # 取前 k 个，重新编排名
    final_results = []
    for rank, item in enumerate(sorted_items[:k], start=1):
        result = item["result"]
        final_results.append(SearchResult(
            chunk_id=result.chunk_id,
            content=result.content,
            metadata=result.metadata,
            score=item["score"],    # 这里的 score 是 RRF 融合后的分数
            rank=rank,
        ))

    return final_results


# ============ 入口函数：混合检索 ============
def hybrid_search(
    query: str,
    bm25_index: BM25Index,
    collection: chromadb.Collection = None,
    k: int = 5,
    bm25_weight: float = 0.5,
    vector_weight: float = 0.5,
    rrf_k: int = 60,
    use_rerank: bool = True,
    rerank_candidates: int = 20,
) -> list[SearchResult]:
    """
    混合检索入口：BM25 + 向量检索 + RRF 融合 + Cross-Encoder 精排。

    完整流程:
    1. BM25 关键词检索 -> top-N 候选
    2. ChromaDB 向量检索 -> top-N 候选
    3. RRF 加权融合 -> top-N 粗排结果
    4. Cross-Encoder Rerank -> 最终 top-k (可选)

    参数:
        query: 用户的搜索文本
        bm25_index: 已建好的 BM25 索引
        collection: ChromaDB Collection
        k: 最终返回结果数
        bm25_weight: BM25 路的权重(默认 0.5)
        vector_weight: 向量路的权重(默认 0.5)
        rrf_k: RRF 平滑常数(默认 60)
        use_rerank: 是否启用 Cross-Encoder 精排(默认 True)
        rerank_candidates: 粗排阶段取多少个候选给精排(默认 20)

    返回:
        最终 top-k SearchResult 列表
    """
    # 粗排阶段: 多取候选, 给精排更大的筛选空间
    # 如果不用 Rerank, 粗排直接取 k 个就够了
    coarse_k = rerank_candidates if use_rerank else k

    # 两路检索
    bm25_results = bm25_index.search(query, k=coarse_k)
    vector_results = vector_search(query, collection, k=coarse_k)

    # RRF 融合
    coarse_results = rrf_fuse(
        results_list=[bm25_results, vector_results],
        k=coarse_k,
        rrf_k=rrf_k,
        weights=[bm25_weight, vector_weight],
    )

    # 精排(可选)
    if use_rerank and coarse_results:
        final_results = rerank(query, coarse_results, top_k=k)
    else:
        final_results = coarse_results[:k]

    return final_results


# ============ 测试代码 ============
if __name__ == "__main__":
    from document_loader import load_documents
    from text_splitter import split_documents

    # --- 第一步：准备数据(加载 → 分块 → 存入 ChromaDB)---
    print("=" * 50)
    print("第一步：加载文档并分块")
    print("=" * 50)
    docs = load_documents("F:/AI_Program/enterprise-kb")
    chunks = split_documents(docs, chunk_size=500, chunk_overlap=100)

    # --- 第二步：建 BM25 索引 ---
    print("\n" + "=" * 50)
    print("第二步：建 BM25 索引")
    print("=" * 50)
    bm25_index = BM25Index()
    bm25_index.build(chunks)

    # --- 第三步：获取 ChromaDB Collection(之前已经存过数据)---
    print("\n" + "=" * 50)
    print("第三步：连接 ChromaDB")
    print("=" * 50)
    collection = get_collection()
    print(f"ChromaDB 中已有 {collection.count()} 个分块")

    # --- 第四步：测试混合检索(粗排 vs 精排对比) ---
    test_queries = [
        "MCP协议是什么",
        "BM25 search algorithm",
        "ChromaDB 向量数据库怎么用",
    ]

    for query in test_queries:
        print("\n" + "=" * 50)
        print(f"搜索: {query}")
        print("=" * 50)

        # 粗排(不用 Rerank)
        coarse = hybrid_search(
            query=query,
            bm25_index=bm25_index,
            collection=collection,
            k=3,
            use_rerank=False,
        )
        print("\n  --- 粗排(RRF) ---")
        for r in coarse:
            preview = r.content[:80].encode("gbk", errors="replace").decode("gbk")
            print(f"  #{r.rank} [RRF={r.score:.5f}] {r.chunk_id}")
            print(f"     {preview}...")

        # 精排(加 Rerank)
        refined = hybrid_search(
            query=query,
            bm25_index=bm25_index,
            collection=collection,
            k=3,
            use_rerank=True,
            rerank_candidates=10,
        )
        print("\n  --- 精排(Cross-Encoder) ---")
        for r in refined:
            preview = r.content[:80].encode("gbk", errors="replace").decode("gbk")
            print(f"  #{r.rank} [CE={r.score:.4f}] {r.chunk_id}")
            print(f"     {preview}...")