SpringAI 基于 milvus 的向量检索

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

1. 启动milvus

为了尽可能快速上手springai的vectordb功能，我们推荐使用云上的milvus，注册就能创建免费的milvus实例，测试完全够了。从控制台复制域名和token

docker安装一个attu可视化工具去连接公网的milvus

docker run -p 8000:3000  -e MILVUS_URL=0.0.0.0:19530 zilliz/attu:latest

2. pom添加依赖

因为考虑到后期可能随时更换向量模型，所以不推荐以yml方式让springboot自动配置VectorStore，这里以手动方式配置

    <dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-milvus-store</artifactId></dependency>

3. 注册Bean

@Configuration
public class SpringAIConfig {@Beanpublic MilvusServiceClient milvusClient() {return new MilvusServiceClient(ConnectParam.newBuilder().withToken("9b06645c438b57b982585fc9b4bd678e6d74d3ae62771exxxxxxxxxxxxxxxxxxxxxxxx").withUri("https://in03-d7f9d7fd8895405.serverless.ali-cn-hangzhou.cloud.zilliz.com.cn").withDatabaseName("db_d7f9d7fxxxxxxx").build());}//如果同时定义多个向量库这里需要起个名字@Bean("milvusVectorStore")public VectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingModel embeddingModel) {return MilvusVectorStore.builder(milvusClient, embeddingModel).collectionName("test_vector_store").databaseName("db_d7f9d7fxxxxxxx").indexType(IndexType.IVF_FLAT).metricType(MetricType.COSINE).batchingStrategy(new TokenCountBatchingStrategy()).initializeSchema(true).build();}
}

4. 创建向量数据的创建和检索方法

保存时在Metadata字段保存知识库的id，用于过滤向量数据，实现知识库的知识隔离

@Component
@RequiredArgsConstructor
public class VectorService {@Qualifier("milvusVectorStore")@Autowiredprivate VectorStore milvusVectorStore;public void embedFileToMilvus(MultipartFile file,String knowledgeId) {try {// 读取上传文件内容String content = new String(file.getBytes(), StandardCharsets.UTF_8);// 切分为小块List<Document> docs = splitTextToDocuments(content,knowledgeId); // 每500字符为一块// 写入向量库milvusVectorStore.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);//指定向量数据的知识库Iddocument.getMetadata().put("knowledgeId",knowledgeId);docs.add(document);}return docs;}public void store(List<Document> documents) {if (documents == null || documents.isEmpty()) {return;}milvusVectorStore.add(documents);}public List<Document> search(String query,String knowledgeId,Double threshold) {FilterExpressionBuilder b = new FilterExpressionBuilder();return milvusVectorStore.similaritySearch(SearchRequest.builder().query(query).topK(5)   //返回条数.similarityThreshold(threshold)   //相似度，阈值范围0~1，值越大匹配越严格‌.filterExpression(b.eq("knowledgeId", knowledgeId).build()).build());}public void delete(Set<String> ids) {milvusVectorStore.delete(new ArrayList<>(ids));}}

5. 测试接口

@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.embedFileToMilvus(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);}
}

数据库插入内容预览

检索效果

6. 将检索结果作为上下文

  // 系统提示词private final static String SYSTEM_PROMPT = """你需要使用文档内容对用户提出的问题进行回复，同时你需要表现得天生就知道这些内容，不能在回复中体现出你是根据给出的文档内容进行回复的，这点非常重要。当用户提出的问题无法根据文档内容进行回复或者你也不清楚时，回复不知道即可。文档内容如下:{documents}""";...String systemPrompt;// 判断是否需要检索知识库if (body.getKnowledgeId() != null) {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(""));// 用文档内容填充SYSTEM_PROMPT模板String filledSystemPrompt = SYSTEM_PROMPT.replace("{documents}", context);// 判断用户是否指定自定义系统提示词if (body.getSystemPrompt() != null && !body.getSystemPrompt().trim().isEmpty()) {systemPrompt = body.getSystemPrompt() + "\n" + filledSystemPrompt;} else {systemPrompt = filledSystemPrompt;}} else {// 没有检索到内容，判断用户是否指定自定义系统提示词systemPrompt = (body.getSystemPrompt() != null && !body.getSystemPrompt().trim().isEmpty())? body.getSystemPrompt(): "中文回答";}}...//将系统提示词添加到提示词消息中messages.add(new SystemMessage(systemPrompt));Prompt prompt = new Prompt(messages);