【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.xxport: 6380password: xxxai:openai:api-key: sk-xxxxxxxxxxxxxxxxxxxxembedding:base-url: https://api.siliconflow.cnoptions:model: BAAI/bge-m3vectorstore:redis:## 是否初始化所需的 schemainitialize-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 filteringMetadataField.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 {@Autowiredprivate 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);

预览效果