向量检索

实现一个基于内存 + numpy 的向量数据库，演示 RAG 中的索引与检索流程： - 文本向量化（确定性哈希向量器） - 余弦相似度检索 - 元数据过滤 生产环境可替换为 Qdrant、Milvus、Chroma 等专业向量数据库。 对应教程：第6卷《RAG》第2章 向量数据库。

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

---

源代码

"""向量检索示例。

实现一个基于内存 + ``numpy`` 的向量数据库，演示 RAG 中的索引与检索流程：
    - 文本向量化（确定性哈希向量器）
    - 余弦相似度检索
    - 元数据过滤

生产环境可替换为 Qdrant、Milvus、Chroma 等专业向量数据库。

对应教程：第6卷《RAG》第2章 向量数据库。
"""

from __future__ import annotations

import hashlib
import math
import re
from dataclasses import dataclass, field
from typing import Any

import numpy as np


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


class HashEmbedder:
    """基于哈希的确定性向量器（教学用，归一化后可直接用点积算余弦）。"""

    def __init__(self, dim: int = 128, seed: int = 0) -> None:
        if dim <= 0:
            raise ValueError("dim 必须为正")
        self.dim = dim
        self.seed = seed

    def embed(self, text: str) -> np.ndarray:
        """将文本嵌入为 L2 归一化的 numpy 向量。

        Args:
            text: 原始文本。

        Returns:
            形状 ``(dim,)`` 的 float32 归一化向量。
        """
        vec = np.zeros(self.dim, dtype=np.float32)
        for tok in tokenize(text):
            digest = hashlib.md5(f"{self.seed}:{tok}".encode()).hexdigest()  # noqa: S324
            index = int(digest[:8], 16) % self.dim
            sign = 1.0 if int(digest[8:16], 16) % 2 == 0 else -1.0
            vec[index] += sign
        norm = float(np.linalg.norm(vec))
        if norm > 0:
            vec /= norm
        return vec

    def embed_batch(self, texts: list[str]) -> np.ndarray:
        """批量向量化。

        Returns:
            形状 ``(n, dim)`` 的矩阵。
        """
        return np.vstack([self.embed(t) for t in texts]) if texts else np.zeros((0, self.dim))


@dataclass
class VectorRecord:
    """向量记录。

    Attributes:
        id: 记录 ID。
        text: 原始文本。
        vector: 向量。
        metadata: 附加元数据。
    """

    id: str
    text: str
    vector: np.ndarray
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class SearchHit:
    """检索命中结果。

    Attributes:
        id: 记录 ID。
        text: 文本。
        score: 相似度分数。
        metadata: 元数据。
    """

    id: str
    text: str
    score: float
    metadata: dict[str, Any]


class InMemoryVectorStore:
    """内存向量库。

    支持：
        - ``add`` 批量写入
        - ``search`` 余弦相似度 top-k 检索
        - 元数据过滤（精确匹配）
    """

    def __init__(self, embedder: HashEmbedder | None = None) -> None:
        self.embedder = embedder or HashEmbedder()
        self._records: list[VectorRecord] = []
        self._matrix: np.ndarray | None = None  # 缓存的 (n, dim) 矩阵

    def add(self, texts: list[str], metadatas: list[dict[str, Any]] | None = None) -> list[str]:
        """批量写入文本。

        Args:
            texts: 文本列表。
            metadatas: 对应元数据列表，可选。

        Returns:
            生成的 ID 列表。
        """
        metadatas = metadatas or [{} for _ in texts]
        if len(metadatas) != len(texts):
            raise ValueError("metadatas 长度必须与 texts 一致")
        ids: list[str] = []
        for i, (text, meta) in enumerate(zip(texts, metadatas, strict=True)):
            rec_id = meta.get("id") or f"rec-{len(self._records) + i}"
            ids.append(rec_id)
            self._records.append(
                VectorRecord(
                    id=rec_id, text=text, vector=self.embedder.embed(text), metadata=dict(meta)
                )
            )
        self._matrix = None  # 失效缓存
        return ids

    def _build_matrix(self) -> np.ndarray:
        if not self._records:
            return np.zeros((0, self.embedder.dim), dtype=np.float32)
        self._matrix = np.vstack([r.vector for r in self._records])
        return self._matrix

    def search(
        self,
        query: str,
        top_k: int = 3,
        filter: dict[str, Any] | None = None,
    ) -> list[SearchHit]:
        """检索与 query 最相似的 top_k 条记录。

        Args:
            query: 查询文本。
            top_k: 返回数量。
            filter: 元数据过滤条件（键值精确匹配）。

        Returns:
            按相似度降序排列的 ``SearchHit`` 列表。
        """
        if not self._records:
            return []
        q_vec = self.embedder.embed(query)
        matrix = self._build_matrix()
        # 向量已归一化，点积即余弦
        scores = matrix @ q_vec
        # 应用元数据过滤
        candidates: list[tuple[int, float]] = []
        for i, rec in enumerate(self._records):
            if filter and not all(rec.metadata.get(k) == v for k, v in filter.items()):
                continue
            candidates.append((i, float(scores[i])))
        candidates.sort(key=lambda x: x[1], reverse=True)
        hits: list[SearchHit] = []
        for i, score in candidates[:top_k]:
            rec = self._records[i]
            hits.append(SearchHit(id=rec.id, text=rec.text, score=score, metadata=rec.metadata))
        return hits

    def __len__(self) -> int:
        return len(self._records)


def main() -> None:
    """演示向量检索流程。"""
    store = InMemoryVectorStore(HashEmbedder(dim=128))
    texts = [
        "AI Agent 是能自主决策的智能体",
        "RAG 通过检索增强大模型生成",
        "今天天气晴朗适合户外活动",
        "MCP 是连接大模型与工具的协议",
    ]
    metadatas = [
        {"topic": "agent"},
        {"topic": "rag"},
        {"topic": "weather"},
        {"topic": "mcp"},
    ]
    store.add(texts, metadatas)
    query = "智能体自己做决定"
    print(f"查询: {query}\n")
    for hit in store.search(query, top_k=2):
        print(f"  {hit.score:.4f}  [{hit.metadata['topic']}]  {hit.text}")
    print("\n带过滤 (topic=rag):")
    for hit in store.search("检索", top_k=2, filter={"topic": "rag"}):
        print(f"  {hit.score:.4f}  {hit.text}")


if __name__ == "__main__":
    main()