📷 [图片 token=H9UUbjpbyoOKQUx6KOIcsBzGnVf(未能下载,见飞书原文)]

📷 [图片 token=Vmz8bkS8Po57qvxcfF5cezr6nRh(未能下载,见飞书原文)]

前言

上一节我们实现了知识库Agent的上半部分,这一节我们来实现知识库的召回功能。

核心代码分布在:app/services/vector_search_service.py — 使用原生 PyMilvus 执行向量检索

📷 [图片 token=E8TWbe4qpoDMMXxKZdNcVvlJnRd(未能下载,见飞书原文)]

召回

我们之前已经将文档向量化存储到了 Milvus,所以召回时也从这个数据库去查询。流程分两步:

  1. 将查询文本向量化

  2. 相似度查询

search_similar_documents 是底层召回的完整实现:

def search_similar_documents(self, query: str, top_k: int = 3) -> List[SearchResult]:
    """
    搜索相似文档

    Args:
        query: 查询文本
        top_k: 返回最相似的K个结果

    Returns:
        List[SearchResult]: 搜索结果列表
    """
    logger.info(f"开始搜索相似文档, 查询: {query}, topK: {top_k}")

    # 1. 将查询文本向量化
    query_vector = vector_embedding_service.embed_query(query)
    logger.debug(f"查询向量生成成功, 维度: {len(query_vector)}")

    # 2. 获取 collection
    collection: Collection = milvus_manager.get_collection()

    # 3. 构建搜索参数
    search_params = {
        "metric_type": "L2",   # 欧氏距离,与入库时的索引类型保持一致
        "params": {"nprobe": 10},
    }

    # 4. 执行搜索
    results = collection.search(
        data=[query_vector],
        anns_field="vector",
        param=search_params,
        limit=top_k,
        output_fields=["id", "content", "metadata"],
    )

    # 5. 解析搜索结果
    search_results = []
    for hits in results:
        for hit in hits:
            result = SearchResult(
                id=hit.entity.get("id"),
                content=hit.entity.get("content"),
                score=hit.distance,   # L2 距离,越小越相似
                metadata=hit.entity.get("metadata", {}),
            )
            search_results.append(result)

    logger.info(f"搜索完成, 找到 {len(search_results)} 个相似文档")
    return search_results

查询文本向量化

首先对用户问题进行向量化,调用 DashScopeEmbeddings.embed_query,通过 DashScope OpenAI 兼容接口获取 1024 维向量:

# 1. 将查询文本向量化
query_vector = vector_embedding_service.embed_query(query)

embed_query 的实现(复用入库时相同的 API,无额外开销):

def embed_query(self, text: str) -> List[float]:
    """嵌入单个查询文本,返回单条向量"""
    response = self.client.embeddings.create(
        model=self.model,          # text-embedding-v4
        input=text,
        dimensions=self.dimensions,  # 1024
        encoding_format="float"
    )
    return response.data[0].embedding

构建搜索参数并执行向量检索

然后使用 PyMilvus 的 collection.search 进行相似度查询,获取距离最近的向量数据:

# 构建搜索参数
search_params = {
    "metric_type": "L2",       # 欧氏距离
    "params": {"nprobe": 10},  # 搜索时探测的 cluster 数量,越大越精确但越慢
}

# 执行搜索
results = collection.search(
    data=[query_vector],                       # 查询向量(批量,这里只有一条)
    anns_field="vector",                       # 向量字段名,与入库时一致
    param=search_params,
    limit=top_k,                               # 返回最相似的前 K 条
    output_fields=["id", "content", "metadata"],  # 需要返回的字段
)

搜索结果封装到 SearchResult 对象中,score 为 L2 欧氏距离,越小表示越相似

class SearchResult:
    def __init__(self, id: str, content: str, score: float, metadata: Dict[str, Any]):
        self.id = id
        self.content = content
        self.score = score     # L2 距离,越小越相似
        self.metadata = metadata

总结

至此,RAG 的分片、索引、召回功能都已实现完毕。后续会介绍 RAG Agent 是怎么使用知识库、怎么结合召回来与大模型进行交互的。