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);