混合检索

演示 RAG 中向量检索与关键词检索（BM25 风格）的融合： - 向量检索捕获语义相似； - 关键词检索捕获字面命中； - 用倒数排名融合（Reciprocal Rank Fusion, RRF）合并两路结果。 全部基于标准库实现，便于离线测试。 对应教程：第6卷《RAG》第3章 检索策略。

文件路径: examples/volume-06-rag\hybrid_search.py

---

源代码

"""混合检索示例。

演示 RAG 中向量检索与关键词检索（BM25 风格）的融合：
    - 向量检索捕获语义相似；
    - 关键词检索捕获字面命中；
    - 用倒数排名融合（Reciprocal Rank Fusion, RRF）合并两路结果。

全部基于标准库实现，便于离线测试。

对应教程：第6卷《RAG》第3章 检索策略。
"""

from __future__ import annotations

import math
import re
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any


def tokenize(text: str) -> list[str]:
    """简易分词：英文单词 + 中文字符。"""
    text = text.lower()
    return re.findall(r"[a-z]+", text) + re.findall(r"[\u4e00-\u9fff]", text)


@dataclass
class Document:
    """检索文档。

    Attributes:
        id: 文档 ID。
        text: 文本。
        metadata: 元数据。
    """

    id: str
    text: str
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class KeywordIndex:
    """基于词频的简易倒排索引（BM25 简化版）。"""

    _docs: list[Document] = field(default_factory=list)
    _doc_tokens: list[list[str]] = field(default_factory=list)
    _df: dict[str, int] = field(default_factory=dict)  # 文档频率

    def add(self, docs: list[Document]) -> None:
        """批量添加文档。"""
        for doc in docs:
            tokens = tokenize(doc.text)
            self._docs.append(doc)
            self._doc_tokens.append(tokens)
            for term in set(tokens):
                self._df[term] = self._df.get(term, 0) + 1

    def search(self, query: str, top_k: int = 5) -> list[tuple[str, float]]:
        """BM25 风格检索。

        Args:
            query: 查询文本。
            top_k: 返回数量。

        Returns:
            ``(doc_id, score)`` 列表，按分数降序。
        """
        q_terms = tokenize(query)
        if not q_terms or not self._docs:
            return []
        n = len(self._docs)
        avg_len = sum(len(t) for t in self._doc_tokens) / n
        scores: dict[str, float] = {doc.id: 0.0 for doc in self._docs}
        k1, b = 1.5, 0.75
        for term in q_terms:
            df = self._df.get(term, 0)
            if df == 0:
                continue
            idf = math.log((n - df + 0.5) / (df + 0.5) + 1)
            for i, doc_tokens in enumerate(self._doc_tokens):
                tf = doc_tokens.count(term)
                if tf == 0:
                    continue
                denom = tf + k1 * (1 - b + b * len(doc_tokens) / max(avg_len, 1))
                scores[self._docs[i].id] += idf * tf * (k1 + 1) / denom
        ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        return [(doc_id, s) for doc_id, s in ranked[:top_k] if s > 0]


@dataclass
class VectorRetriever:
    """向量检索器（复用哈希向量器思路，独立实现以避免跨卷依赖）。"""

    embedder: Callable[[str], list[float]]
    _docs: list[Document] = field(default_factory=list)
    _vectors: list[list[float]] = field(default_factory=list)

    def add(self, docs: list[Document]) -> None:
        """批量添加文档。"""
        for doc in docs:
            self._docs.append(doc)
            self._vectors.append(self.embedder(doc.text))

    def search(self, query: str, top_k: int = 5) -> list[tuple[str, float]]:
        """余弦相似度检索。

        Returns:
            ``(doc_id, score)`` 列表，按分数降序。
        """
        if not self._docs:
            return []
        q_vec = self.embedder(query)
        scored: list[tuple[str, float]] = []
        for doc, vec in zip(self._docs, self._vectors, strict=True):
            scored.append((doc.id, _cosine(q_vec, vec)))
        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:top_k]


def _cosine(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b, strict=True))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    if na == 0 or nb == 0:
        return 0.0
    return dot / (na * nb)


def reciprocal_rank_fusion(
    rankings: list[list[tuple[str, float]]],
    k: int = 60,
) -> list[tuple[str, float]]:
    """倒数排名融合（RRF）。

    对每个文档 ID，累加 ``1 / (k + rank)`` 作为融合分数，``rank`` 从 1 开始。

    Args:
        rankings: 多路检索结果，每路为 ``(doc_id, score)`` 列表（按分数降序）。
        k: RRF 平滑常数，默认 60。

    Returns:
        融合后的 ``(doc_id, rrf_score)`` 列表，按融合分数降序。
    """
    fused: dict[str, float] = {}
    for ranking in rankings:
        for rank, (doc_id, _score) in enumerate(ranking, start=1):
            fused[doc_id] = fused.get(doc_id, 0.0) + 1.0 / (k + rank)
    return sorted(fused.items(), key=lambda x: x[1], reverse=True)


def hybrid_search(
    query: str,
    docs: list[Document],
    embedder: Callable[[str], list[float]],
    top_k: int = 3,
) -> list[tuple[str, float]]:
    """执行混合检索：向量 + 关键词，并用 RRF 融合。

    Args:
        query: 查询文本。
        docs: 文档列表。
        embedder: 向量化函数。
        top_k: 最终返回数量。

    Returns:
        融合后的 ``(doc_id, score)`` 列表。
    """
    keyword_index = KeywordIndex()
    keyword_index.add(docs)
    vector_retriever = VectorRetriever(embedder=embedder)
    vector_retriever.add(docs)

    keyword_hits = keyword_index.search(query, top_k=max(top_k * 2, 5))
    vector_hits = vector_retriever.search(query, top_k=max(top_k * 2, 5))

    fused = reciprocal_rank_fusion([keyword_hits, vector_hits])
    return fused[:top_k]


def main() -> None:
    """演示混合检索。"""
    docs = [
        Document("d1", "AI Agent 能自主决策并执行任务"),
        Document("d2", "RAG 通过检索增强大模型生成"),
        Document("d3", "今天天气晴朗适合户外活动"),
        Document("d4", "智能体可以自己做决定"),
    ]

    # 复用一个简单的哈希向量器
    import hashlib

    def embed(text: str, dim: int = 64) -> list[float]:
        vec = [0.0] * dim
        for tok in tokenize(text):
            h = int(hashlib.md5(tok.encode()).hexdigest(), 16) % dim  # noqa: S324
            vec[h] += 1.0
        norm = math.sqrt(sum(v * v for v in vec))
        return [v / norm for v in vec] if norm > 0 else vec

    query = "智能体自己做决定"
    print(f"查询: {query}\n")
    for doc_id, score in hybrid_search(query, docs, embed, top_k=3):
        text = next(d.text for d in docs if d.id == doc_id)
        print(f"  {doc_id}  rrf={score:.4f}  {text}")


if __name__ == "__main__":
    main()