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 {
@Bean
public 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")
@Autowired
private 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);
//指定向量数据的知识库Id
document.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 {
@Autowired
private 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);