【SpringAI】6.向量检索（redis）

基于redis的向量检索

向量数据库可以使用milvus，redis,Elasticsearch等，本文以redis为例：

docker启动redis-stack-server

docker run -d --name redis-stack-server -p 6380:6379  -v /home/redis/data:/data -e REDIS_ARGS="--requirepass mima --bind 0.0.0.0 --protected-mode no"  redis/redis-stack-server:latest

1,pom引入依赖

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-starter-vector-store-redis</artifactId>
</dependency> 

2,yml配置

这里以硅基流动的免费量化模型测试

spring:
  data:
    redis:
      host: xx.xx.xx.xx
      port: 6380
      password: xxx
  ai:
    openai:
      api-key: sk-xxxxxxxxxxxxxxxxxxxx
      embedding:
        base-url: https://api.siliconflow.cn
        options:
          model: BAAI/bge-m3

    vectorstore:
      redis:
        ## 是否初始化所需的 schema
        initialize-schema: true
        ## 用于存储向量的索引的名称
        index-name: knowledgeId
        ## Redis 键的前缀
        prefix: glmapper_

3,文本向量化和文本检索

import lombok.RequiredArgsConstructor;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.stereotype.Component;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.ai.document.Document;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.ArrayList;
import java.util.Set;

@Component
@RequiredArgsConstructor
public class VectorService {

    private final VectorStore vectorStore;


    // 将上传的文本文件向量化并保存到向量数据库
    public void embedFile(MultipartFile file,String knowledgeId) {
        try {
            // 读取上传文件内容
            String content = new String(file.getBytes(), StandardCharsets.UTF_8);
            // 切分为小块
            List<Document> docs = splitTextToDocuments(content,knowledgeId); // 每500字符为一块
            // 写入向量库
            vectorStore.add(docs);
        } catch (Exception e) {
            throw new RuntimeException("文件向量化失败: " + e.getMessage(), e);
        }
    }

    // 按固定长度分割文本为 Document 列表
    private List<Document> splitTextToDocuments(String text,String knowledgeId) {
        List<Document> docs = new ArrayList<>();
        int length = text.length();
        for (int i = 0; i < length; i += 500) {
            int end = Math.min(length, i + 500);
            String chunk = text.substring(i, end);
            Document document = new Document(chunk);
            //指定向量数据的知识库Id
//            document.getMetadata().put("knowledgeId",knowledgeId);
            docs.add(document);
        }
        return docs;
    }

    public void store(List<Document> documents) {
        if (documents == null || documents.isEmpty()) {
            return;
        }
        vectorStore.add(documents);
    }

    public List<Document> search(String query,String knowledgeId,Double threshold) {
        FilterExpressionBuilder b = new FilterExpressionBuilder();
        return vectorStore.similaritySearch(SearchRequest.builder()
                .query(query)
                .topK(5)   //返回条数
                .similarityThreshold(threshold)   //相似度，阈值范围0~1，值越大匹配越严格‌
//                .filterExpression(b.eq("knowledgeId", knowledgeId).build())
                .build());
    }

    public void delete(Set<String> ids) {
        vectorStore.delete(new ArrayList<>(ids));
    }

}

原本想将量化数据按照知识库分组和过滤，实际并没有效果，即使按照官方文档手动定义Bean指定knowledgeId的tag也无效，官方提供的手动定义如下：

@Bean
public VectorStore vectorStore(JedisPooled jedisPooled, EmbeddingModel embeddingModel) {
    return RedisVectorStore.builder(jedisPooled, embeddingModel)
        .indexName("custom-index")                // Optional: defaults to "spring-ai-index"
        .prefix("custom-prefix")                  // Optional: defaults to "embedding:"
        .metadataFields(                         // Optional: define metadata fields for filtering
            MetadataField.tag("country"),
            MetadataField.numeric("year"))
        .initializeSchema(true)                   // Optional: defaults to false
        .batchingStrategy(new TokenCountBatchingStrategy()) // Optional: defaults to TokenCountBatchingStrategy
        .build();
}

// This can be any EmbeddingModel implementation
@Bean
public EmbeddingModel embeddingModel() {
    return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
}

搜索时代码如下：

vectorStore.similaritySearch(SearchRequest.builder()
        .query("The World")
        .topK(TOP_K)
        .similarityThreshold(SIMILARITY_THRESHOLD)
        .filterExpression(b.and(
                b.in("country", "UK", "NL"),
                b.gte("year", 2020)).build()).build());

4,测试接口

@Tag(name = "向量检索", description = "向量检索")
@RestController
@RequestMapping("/vector")
public class VectorController {

    @Autowired
    private VectorService vectorService;

    @Operation(summary = "文本文件向量化", description = "文本文件向量化")
    @PostMapping("/uploadFile")
    public RestVO<Map<String, Object>> uploadFile(@RequestPart MultipartFile file, @RequestParam String knowledgeId) {
        vectorService.embedFile(file, knowledgeId);
        return RestVO.success(Map.of("success", true, "message", "文件已向量化"));
    }

    @Operation(summary = "向量检索", description = "向量检索")
    @GetMapping("/query")
    public RestVO<List<Document>> uploadFile(@RequestParam String query, @RequestParam Double threshold, @RequestParam(required = false) String knowledgeId) {
        List<Document> documentList = vectorService.search(query, knowledgeId,threshold);
        return RestVO.success(documentList);
    }
}

查询结果

5, 检索结果运用

        List<Message> messages = new ArrayList<>();
        //省略历史会话
        UserMessage userMessage;
        //TODO 可以先对用户问题做关键词提取再去检索
        List<Document> documentList = vectorStore.similaritySearch(body.getMessage());
        System.out.println("检索结果" + documentList.size());
        if (documentList != null && !documentList.isEmpty()) {
                String context = documentList.stream()
                        .map(Document::getText)
                        .collect(Collectors.joining(""));
                userMessage = new UserMessage("参考内容：【\n" + context + "】\n\n回答：【" + body.getMessage() + "】");
        } else {
                userMessage = new UserMessage(body.getMessage());
        }
        messages.add(userMessage);
        Prompt prompt = new Prompt(messages);

预览效果