【RAG技术】Springboot整合Elasticsearch8 向量检索快速实战

本篇文章旨在快速完成Springboot整合ES8进行向量检索的实战，相关配置都是最少配置，需要阅读者具备一定的前置知识。主要技术用到了ES的查询模板。

废话不多说，直接开整！

第一步：配置pom.xml，添加springboot整合es的依赖

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-data-elasticsearch</artifactId>
        </dependency>

第二步：配置application.properties

# elastic
spring.elasticsearch.uris=localhost:9200

第三步：编写ES模板查询的工具类：

import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import org.springframework.data.elasticsearch.client.elc.ElasticsearchTemplate;
import org.springframework.data.elasticsearch.core.SearchHit;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.SearchTemplateQuery;
import org.springframework.data.elasticsearch.core.script.Script;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;
import java.util.Map;


@Component
public class ESHelper {

    @Resource
    private ElasticsearchTemplate elasticsearchTemplate;

    private Map<String,String> templateMap = new HashMap<>();

    private String vectorSearchTemplate;

    public static final String TEMPLATE_QA_ID="template_qa";

    @PostConstruct
    public void init() {
        vectorSearchTemplate = """
                {
                  "query": {
                    "function_score": {
                      "query": {
                        "match_all": {}
                      },
                      "functions": [
                        {
                          "script_score": {
                            "script": {
                              "source": "cosineSimilarity({{#toJson}}embedding{{/toJson}}, 'embedding') + 1.0"
                            }
                          }
                        }
                      ]
                    }
                  }
                }
                """;
        createTemplate(TEMPLATE_QA_ID);
        templateMap.put(TEMPLATE_QA_ID,vectorSearchTemplate);
    }

    public void createTemplate(String templateId) {
        elasticsearchTemplate.putScript(Script.builder()
                .withId(templateId)
                .withLanguage("mustache")
                .withSource(vectorSearchTemplate).build());
    }

    public <T> List<SearchHit<T>> templateQuery(String templateId, Map<String, Object> params,Class<T> clazz,String... indexNames) {
        SearchTemplateQuery query = SearchTemplateQuery.builder().withId(templateId).withParams(params).build();
        IndexCoordinates index = IndexCoordinates.of(indexNames);
        List<SearchHit<T>> searchHits = elasticsearchTemplate.search(query, clazz,index).getSearchHits();
        return searchHits;

    }


}

代码解释：

templateMap：用于存放查询模板的ID和查询模板。

vectorSearchTemplate：向量查询模板。

TEMPLATE_QA_ID：向量查询模板ID

public void init()：初始化Bean时，向ES中创建查询模板，然后缓存向量查询模版。

public <T> List<SearchHit<T>> templateQuery 方法：根据查询模板ID和其他参数根据模板查询。

特别注意：

在上面代码中的查询模板中有下面这段代码

"script": {
  "source": "cosineSimilarity({{#toJson}}embedding{{/toJson}}, 'embedding') + 1.0"
}

这段代码主要是通过计算传入的 embedding 向量和 ES 中向量字段embedding 进行余弦相似度计算。这里两个字段都叫 embedding，只是名字相同，可以修改为不同的字段。

第四步：以上演示了如何进行向量查询，下面补充下如何在ES中保存这些向量，这个操作和ES正常操作没啥区别。

# 单个文档构建
IndexQuery query= new IndexQueryBuilder().withIndex(索引名称)
.withObject(文档映射对象).build();

elasticsearchTemplate.doIndex(indexQueries, 文档映射实体.class);

也可以使用 elasticsearchTemplate.bulkIndex 方法 进行批量构建文档。

总结

本文演示了在RAG系统中，如何将ES作为向量数据库的使用介绍，包括如何进行索引文档的构建，如何构建查询模型进行相似度查询。

再顺嘴简单说下RAG，一般是先通过嵌入模型计算文档向量，然后在保存到向量数据库中，这里的ES就可以作为一个向量数据库。当用户提出问题时，首先通过嵌入模型拿到问题的向量表示，然后根据向量到向量数据库中进行向量相似度查询。根据拿到的结果，进行重排序等操作拿到最相关的文档，最终一起构建提示词，喂给大模型。

---

关于我：

资深程序员，曾在北京某AI公司从事智能对话问答平台研发，也曾在大厂历练过，对AI场景应用充满热情。