SpringBoot 中操作 Milvus 向量数据库的工具类分享

本文将介绍如何在 SpringBoot 项目中集成 Milvus 向量数据库，并提供一个封装了常用 CRUD 和检索操作的工具类，方便快速开发具备向量检索能力的应用。

一、环境准备与依赖导入

首先，在 pom.xml 中添加 Milvus 的 Java SDK 依赖：

<dependency>
    <groupId>io.milvus</groupId>
    <artifactId>milvus-sdk-java</artifactId>
    <version>2.6.3</version>
</dependency>

二、配置 Milvus 连接

创建一个配置类 MilvusConfig，用于建立与 Milvus 服务的连接：

import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
​
@Configuration
public class MilvusConfig {
​
    @Value("${langchain4j.milvus.database-name}")
    private String dbName;
    
    @Value("${milvus.uri}")
    private String uri;
​
    @Bean
    public MilvusClientV2 milvusClientV2() {
        ConnectConfig connectConfig = ConnectConfig.builder()
                .uri(uri)
                .dbName(dbName)
                .build();
        return new MilvusClientV2(connectConfig);
    }
}

请在 application.properties 或 application.yml 中配置相应的连接参数。

三、工具类实现：CRUD 与向量检索功能

以下工具类 MilvusTools 封装了常见的数据库操作，包括插入、删除、更新、按主键查询和基于文本的向量检索：

import com.google.gson.JsonObject;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.ConsistencyLevel;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.vector.request.*;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
​
import java.util.*;
​
import static com.alibaba.dashscope.utils.JsonUtils.gson;
​
/**
 * Milvus 工具类，提供与 Milvus 向量数据库交互的常用操作
 */
@Component
public class MilvusTools {
​
    @Autowired
    private EmbeddingModel embeddingModel;
​
    @Autowired
    private MilvusClientV2 client;
​
    /**
     * 向指定集合插入文本及对应向量
     */
    public Map<String, Long> insert(String text, float[] vector, String collectionName) {
        JsonObject entity = new JsonObject();
        entity.addProperty("text", text);
        Long id = System.currentTimeMillis();
        entity.add("id", gson.toJsonTree(id));
        entity.add("vector", gson.toJsonTree(vector));
​
        List<JsonObject> entities = Collections.singletonList(entity);
​
        InsertReq insertReq = InsertReq.builder()
                .collectionName(collectionName)
                .data(entities)
                .build();
​
        InsertResp insertResp = client.insert(insertReq);
        Map<String, Long> result = new HashMap<>();
        result.put("insertCnt", insertResp.getInsertCnt());
        result.put("primaryKey", (Long) insertResp.getPrimaryKeys().get(0));
        
        client.close();
        return result;
    }
​
    /**
     * 根据主键列表删除记录
     */
    public Long delete(List<Long> ids, String collectionName) {
        DeleteReq deleteReq = DeleteReq.builder()
                .collectionName(collectionName)
                .ids(ids)
                .build();
​
        DeleteResp deleteResp = client.delete(deleteReq);
        client.close();
        return deleteResp.getDeleteCnt();
    }
​
    /**
     * 更新指定主键对应的记录
     */
    public UpsertResp update(String text, float[] vector, String collectionName, Long primaryKey) {
        JsonObject entity = new JsonObject();
        entity.addProperty("text", text);
        entity.addProperty("id", primaryKey);
        entity.add("vector", gson.toJsonTree(vector));
​
        UpsertReq upsertReq = UpsertReq.builder()
                .collectionName(collectionName)
                .data(Collections.singletonList(entity))
                .build();
​
        UpsertResp resp = client.upsert(upsertReq);
        client.close();
        return resp;
    }
​
    /**
     * 根据主键查询记录
     */
    public List<QueryResp.QueryResult> queryByPrimaryKey(String primaryKeyValue, String collectionName) {
        try {
            String expr = isNumeric(primaryKeyValue) ? 
                "id == " + primaryKeyValue : 
                "id == '" + primaryKeyValue + "'";
​
            QueryReq queryReq = QueryReq.builder()
                    .collectionName(collectionName)
                    .filter(expr)
                    .build();
​
            QueryResp response = client.query(queryReq);
            return response.getQueryResults();
        } catch (Exception e) {
            System.err.println("查询异常: " + e.getMessage());
            return Collections.emptyList();
        }
    }
​
    /**
     * 基于文本内容进行向量相似性搜索
     */
    public SearchResp searchByText(String collectionName, String text) {
        Response<Embedding> embeddings = embeddingModel.embed(text);
        float[] floatVector = embeddings.content().vector();
        BaseVector queryVector = new FloatVec(floatVector);
​
        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(Collections.singletonList(queryVector))
                .topK(3)
                .outputFields(Collections.singletonList("*"))
                .metricType(IndexParam.MetricType.COSINE)
                .consistencyLevel(ConsistencyLevel.EVENTUALLY)
                .build();
​
        return client.search(searchReq);
    }
​
    private boolean isNumeric(String str) {
        try {
            Double.parseDouble(str);
            return true;
        } catch (NumberFormatException e) {
            return false;
        }
    }
}

四、功能测试用例

提供以下测试方法，用于验证工具类的各项功能：

@Autowired
private EmbeddingModel embeddingModel;
​
@Autowired
private MilvusTools milvusTools;
​
@Test
public void insertData() {
    String text = "这是一个测试插入语句";
    Response<Embedding> embed = embeddingModel.embed(text);
    float[] vector = embed.content().vector();
    String collectionName = "example_collection";
​
    Map<String, Long> result = milvusTools.insert(text, vector, collectionName);
    System.out.println(result);
}
​
@Test
public void deleteData() {
    String collectionName = "example_collection";
    Long count = milvusTools.delete(Collections.singletonList(1757427143710L), collectionName);
    System.out.println("删除记录数: " + count);
}
​
@Test
public void searchData() {
    String collectionName = "example_collection";
    String text = "这是一个测试插入语句";
    SearchResp searchResp = milvusTools.searchByText(collectionName, text);
    System.out.println(searchResp);
}
​
@Test
public void searchById() {
    String collectionName = "example_collection";
    List<QueryResp.QueryResult> results = milvusTools.queryByPrimaryKey("1", collectionName);
    System.out.println(results);
}
​
@Test
public void updateData() {
    String collectionName = "example_collection";
    Long id = 1757427358507L;
    String text = "这是一个测试修改语句";
    Response<Embedding> embed = embeddingModel.embed(text);
    float[] vector = embed.content().vector();
    
    UpsertResp resp = milvusTools.update(text, vector, collectionName, id);
    System.out.println(resp);
}

五、注意事项

1. 本工具类中使用的 EmbeddingModel 是为文本生成向量表示的组件，可根据实际需求替换为其他支持的模型（如 OpenAI、Hugging Face 等）。

2. 使用前请确保已正确配置 Milvus 服务地址和数据库名称。

3. 在实际生产环境中，建议对连接进行池化管理，避免频繁创建和关闭连接。

通过上述工具类，可以方便地在 SpringBoot 项目中集成 Milvus 实现向量数据的存储与检索，为构建 AI 应用提供基础设施支持。