知识库-向量化功能-混合查询

造轮子的猪
23 阅读8分钟

知识库-向量化功能-混合查询

一、功能目标

基于 Elasticsearch 8.8+ 实现 BM25（关键词检索）+ 向量相似度（kNN） 混合检索，通过 Rank Fusion（RRF）算法融合两路结果，解决单一检索方式的局限性：

  • 纯关键词检索（BM25）：易遗漏语义相关但关键词不匹配的文档；
  • 纯向量检索（kNN）：对关键词精准匹配场景效果弱； 最终提升召回文档的准确率，优化 RAG（检索增强生成）输出质量。

二、核心原理

检索方式算法核心适用场景局限性
BM25基于词频、文档长度的关键词相关性计算（改进版 TF-IDF）关键词精准匹配、结构化文本检索无法理解语义，易受同义词/近义词影响
kNN（向量检索）基于向量空间距离（余弦/欧氏）的语义相似度计算语义相似检索、模糊匹配对关键词精准匹配不敏感，计算成本较高
RRF 融合无需调参的排序融合算法，公式：1/(k + rank)（k 为常数）平衡关键词与语义检索结果自动加权，两路结果排名越靠前，融合得分越高

三、核心实现代码

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.search.Hit;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.stream.Collectors;

/**
 * ES 混合检索服务（BM25 + kNN + RRF 融合）
 * 适配 Elasticsearch 8.8+，需提前确保向量字段已创建（dimension=768，similarity=cosine）
 */
@Service
public class HybridSearchService {
    // 依赖注入：ES 客户端、文本嵌入模型
    private final ElasticsearchClient client;
    private final EmbeddingModel embeddingModel;

    // 配置常量（可封装为配置类，支持外部配置）
    private static final int K = 20;                // 单路召回数量（建议10-50）
    private static final int FINAL_K = 5;           // 最终返回数量（适配LLM上下文窗口）
    private static final int RRF_RANK_CONSTANT = 60;// RRF 排名常数（经验值，建议50-100）
    private static final int KNN_NUM_CANDIDATES = 100; // kNN 候选集数量（需大于K）

    public HybridSearchService(ElasticsearchClient client, EmbeddingModel embeddingModel) {
        this.client = client;
        this.embeddingModel = embeddingModel;
    }

    /**
     * 文档分片实体类（需与ES索引字段映射）
     */
    public static class DocFragment {
        private String text;   // ES 文本字段（对应text-field-name）
        private String docId;  // 原始文档ID（非分片ID，用于溯源）

        // 必须提供getter/setter（Jackson反序列化需要）
        public String getText() { return text; }
        public void setText(String text) { this.text = text; }
        public String getDocId() { return docId; }
        public void setDocId(String docId) { this.docId = docId; }
    }

    /**
     * 辅助类：封装ES检索结果（ID + 分数 + 文档内容）
     */
    private static class HitWithScore {
        String id;          // ES文档ID（分片ID）
        Double score;       // 原始检索分数（BM25/kNN）
        DocFragment doc;    // 文档内容

        HitWithScore(String id, Double score, DocFragment doc) {
            this.id = id;
            this.score = score;
            this.doc = doc;
        }
    }

    /**
     * 核心混合检索方法
     * @param index ES索引名（需与vectorstore配置的index-name一致）
     * @param queryText 检索文本（用户提问/关键词）
     * @return 融合后的Top-N文档列表
     * @throws Exception 检索异常（ES连接/模型调用失败）
     */
    public List<DocFragment> hybridSearch(String index, String queryText) throws Exception {
        // ========== 步骤1：BM25关键词检索 ==========
        // 基于ES内置BM25算法，检索关键词匹配的文档
        SearchResponse<DocFragment> bm25Resp = client.search(
                s -> s.index(index)
                        .query(q -> q.match(m -> m
                                .field("text")       // 文本字段名（需与vectorstore配置一致）
                                .query(queryText)    // 检索文本
                                .fuzziness("AUTO")   // 可选：模糊匹配（容错拼写错误）
                        ))
                        .size(K),                  // 召回K个结果
                DocFragment.class
        );

        // ========== 步骤2：kNN向量语义检索 ==========
        // 1. 生成检索文本的向量（调用嵌入模型）
        float[] queryVector = embeddingModel.embed(queryText);
        // 2. 转换为ES支持的Float列表（优化：提前初始化容量）
        List<Float> floatList = new ArrayList<>(queryVector.length);
        for (float f : queryVector) {
            floatList.add(f);
        }
        // 3. 执行kNN检索（近似最近邻）

        SearchResponse<DocFragment> knnResp = client.search(
                s -> s.index(index)
                        .knn(k -> k
                                .field("vector")
                                .queryVector(floatList)
                                .k(K)
                                .numCandidates(100)
                        )
                        .size(K),
                DocFragment.class
        );

        // ========== 步骤3：RRF（Reciprocal Rank Fusion）融合 ==========
        Map<String, Double> rrfScores = new HashMap<>(); // 文档ID → 融合分数
        Map<String, DocFragment> allDocs = new HashMap<>(); // 文档ID → 文档内容

        // 处理BM25结果，计算RRF分数
        List<HitWithScore> bm25Hits = convertToHitWithScore(bm25Resp);
        for (int rank = 0; rank < bm25Hits.size(); rank++) {
            HitWithScore hit = bm25Hits.get(rank);
            double rrfScore = 1.0 / (RRF_RANK_CONSTANT + rank + 1); // RRF核心公式
            rrfScores.merge(hit.id, rrfScore, Double::sum); // 累加分数
            allDocs.putIfAbsent(hit.id, hit.doc); // 缓存文档内容
        }

        // 处理kNN结果，计算RRF分数
        List<HitWithScore> knnHits = convertToHitWithScore(knnResp);
        for (int rank = 0; rank < knnHits.size(); rank++) {
            HitWithScore hit = knnHits.get(rank);
            double rrfScore = 1.0 / (RRF_RANK_CONSTANT + rank + 1);
            rrfScores.merge(hit.id, rrfScore, Double::sum);
            allDocs.putIfAbsent(hit.id, hit.doc);
        }

        // ========== 步骤4：排序并返回Top-FINAL_K ==========
        return rrfScores.entrySet().stream()
                .sorted(Map.Entry.<String, Double>comparingByValue().reversed()) // 按融合分数降序
                .limit(FINAL_K)                                                  // 取前FINAL_K个
                .map(entry -> allDocs.get(entry.getKey()))                       // 映射为文档对象
                .collect(Collectors.toList());
    }

    /**
     * 辅助方法：将ES检索响应转换为HitWithScore列表
     */
    private List<HitWithScore> convertToHitWithScore(SearchResponse<DocFragment> response) {
        if (response == null || response.hits() == null || response.hits().hits() == null) {
            return Collections.emptyList();
        }
        return response.hits().hits().stream()
                .map(hit -> new HitWithScore(hit.id(), hit.score(), hit.source()))
                .collect(Collectors.toList());
    }
}

四、混合查询 vs vectorStore.similaritySearch()

4.1 核心差异对比表

维度混合查询（BM25 + kNN + RRF）vectorStore.similaritySearch()（Spring AI 内置）
检索方式多策略融合（关键词+语义）单一向量检索（仅kNN）
算法核心BM25 + kNN + RRF 融合排序仅向量相似度排序（cosine/l2_norm/dot_product）
适用场景通用检索、复杂语义+关键词混合需求、高召回率场景纯语义检索、简单相似性匹配、快速开发场景
召回精度高（兼顾关键词和语义，减少漏检/误检）中（易遗漏关键词精准匹配的文档）
性能开销较高（两次检索+融合计算）较低（单次kNN检索）
灵活性高（可调整K值、RRF常数、是否开启模糊匹配）低（封装度高，仅支持少量参数配置）
ES版本要求8.8+（需支持RRF/knn查询）8.x均可（仅依赖向量字段）
代码复杂度较高（需手动处理两路检索+融合）极低（一行代码调用，无需关注底层）

4.2 详细说明

（1）vectorStore.similaritySearch() 特点

Spring AI 封装的向量检索方法，核心逻辑：

// 内置方法调用示例
List<Document> results = vectorStore.similaritySearch(
    Query.query(queryText)
         .withTopK(5) // 仅支持设置返回数量
);
  • 优势：开箱即用，无需关注ES底层语法，适合快速开发；
  • 局限：
    • 仅基于向量相似度检索，无法利用关键词匹配；
    • 对“语义相似但关键词不匹配”的场景友好，但对“关键词精准匹配但语义模糊”的场景（如专业术语检索）效果差；
    • 无结果融合能力，召回结果单一。
（2）混合查询特点
  • 优势：
    • 兼顾“关键词精准匹配”和“语义相似”，召回结果更全面；
    • RRF融合无需手动调参（如加权系数），适配性强；
    • 可灵活扩展（如增加模糊匹配、过滤条件、自定义排序）；
  • 局限：
    • 需手动编写ES检索逻辑，代码量较大；
    • 两次检索增加ES查询开销，需合理设置K值（避免过大）；
    • 依赖ES 8.8+版本的kNN和RRF特性。

4.3 选型建议

场景推荐方案
快速原型开发、纯语义检索需求vectorStore.similaritySearch()
生产环境、通用检索、高召回率要求混合查询（BM25 + kNN + RRF）
专业术语检索、关键词精准匹配为主混合查询（可加大BM25权重）
模糊语义检索、无明确关键词vectorStore.similaritySearch()（或混合查询加大kNN权重）

五、核心配置说明

配置项取值建议作用
K（单路召回数量）10-50单路召回越多，融合结果越全面，但性能开销越大
FINAL_K（最终返回数量）3-10适配LLM上下文窗口大小（避免输入过长）
RRF_RANK_CONSTANT50-100经验值，数值越小，排名靠前的文档得分权重越高
KNN_NUM_CANDIDATES100-200kNN候选集数量，越大召回精度越高，但查询耗时越长
模糊匹配（fuzziness）AUTO可选开启，容错拼写错误（如“检索”→“搜素”）

六、使用示例

6.1 基础调用

@Autowired
private HybridSearchService hybridSearchService;

public void testHybridSearch() {
    try {
        String indexName = "ai_vector_index"; // 与vectorstore配置一致
        String queryText = "Spring AI 向量检索使用方法";
        // 执行混合检索
        List<HybridSearchService.DocFragment> results = hybridSearchService.hybridSearch(indexName, queryText);
        // 处理结果
        for (HybridSearchService.DocFragment doc : results) {
            System.out.printf("文档ID：%s，内容：%s%n", doc.getDocId(), doc.getText());
        }
    } catch (Exception e) {
        System.err.println("混合检索失败：" + e.getMessage());
    }
}

6.2 结合RAG场景

@Autowired
private HybridSearchService hybridSearchService;
@Autowired
private ChatClient chatClient;

public String ragAnswer(String userQuestion) {
    try {
        // 1. 混合检索获取相关文档
        List<HybridSearchService.DocFragment> docs = hybridSearchService.hybridSearch("ai_vector_index", userQuestion);
        // 2. 构建上下文
        String context = docs.stream()
                .map(HybridSearchService.DocFragment::getText)
                .collect(Collectors.joining("\n"));
        // 3. 调用LLM生成回答
        String prompt = String.format("基于以下上下文回答问题：%n%s%n问题：%s", context, userQuestion);
        return chatClient.call(prompt);
    } catch (Exception e) {
        return "检索失败：" + e.getMessage();
    }
}

七、注意事项

  1. ES索引配置：需确保向量字段vector的dimension与嵌入模型一致（如768），similarity设置为cosine；
  2. 性能优化：
    • 合理设置K值（建议20），避免单路召回过多；
    • 为ES索引的text字段配置分词器（如IK分词器），提升BM25检索效果；
    • 开启ES缓存，减少重复检索开销；
  3. 异常处理：增加ES连接超时、模型调用失败的容错逻辑；
  4. 监控指标：统计BM25/kNN召回文档的重合率、融合后Top-N的准确率，便于调参；
  5. 扩展能力：可增加过滤条件（如按文档ID/时间筛选）、自定义加权（如对重要文档提升分数）。

八、扩展建议

  1. 多字段检索：BM25检索可扩展为多字段（如text+docId+title），提升关键词匹配维度；
  2. 动态权重：根据检索场景动态调整RRF常数（如专业术语检索减小RRF常数，强化BM25权重）；
  3. 缓存优化：对高频检索词的混合结果进行缓存，减少重复计算；
  4. 异步检索：BM25和kNN检索异步执行，减少总耗时；
  5. 结果去重：融合前对BM25/kNN结果去重（如按文档内容相似度），避免重复文档。
avatar
文章
阅读
粉丝