springboot整合milvus MilvusClientV2

目录

1.依赖
2.配置
3.代码功能

1.依赖

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

2.配置

@Configuration
public class MilvusConfig {

    @Bean
    public MilvusClientV2 milvusClientV2(@Value("${milvusUrl}") String milvusUrl) {
        ConnectConfig connectConfig = ConnectConfig.builder()
                .uri(milvusUrl)
                .build();
        return new MilvusClientV2(connectConfig);
    }
}

3.代码功能

3.1创建映射mapping

#向量维度 dimension: 768

#创建映射mapping
public void createCollection() {
        CreateCollectionReq.CollectionSchema schema = client.createSchema();
      schema.addField(AddFieldReq.builder()
              .fieldName("id")
              .dataType(DataType.VarChar)
              .isPrimaryKey(true)
              .autoID(false)
              .build());

      schema.addField(AddFieldReq.builder()
              .fieldName("question")
              .dataType(DataType.VarChar)
              .maxLength(10000)
              .build());

      schema.addField(AddFieldReq.builder()
              .fieldName("question_vector")
              .dataType(DataType.FloatVector)
              .dimension(dimension)
              .build());

      IndexParam indexParam = IndexParam.builder()
              .fieldName("question_vector")
              .metricType(IndexParam.MetricType.COSINE)
              .build();

      CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
              .collectionName(embeddingConfig.getCollectionName())
              .collectionSchema(schema)
              .indexParams(Collections.singletonList(indexParam))
              .build();

      client.createCollection(createCollectionReq);
  }

#判断集合是否存在
public boolean existCollection() {
    return client.hasCollection(HasCollectionReq.builder().collectionName(embeddingConfig.getCollectionName()).build());
}

#新增向量数据
public void batchInsert(List<Embedding> list) {
      if (CollectionUtil.isEmpty(list)) {
         return;
      }

      Gson gson = new Gson();
      List<JsonObject> objectList = new ArrayList<>();
      list.forEach(l -> {
          JsonObject dict = new JsonObject();
          dict.addProperty("id", UUID.randomUUID().toString());
          dict.addProperty("question", l.getQuestion());
          dict.add("question_vector", gson.toJsonTree(l.getEmbedding()));
          objectList.add(dict);
      });

      InsertReq insertReq = InsertReq.builder()
              .collectionName(embeddingConfig.getCollectionName())
              .data(objectList).build();
      InsertResp insertResp = client.insert(insertReq);
      log.debug("batchInsert result:{}", insertResp);
  }

#检索数据topk
public List<ClassificationQuestionVO> search(Embedding embedding, Integer top) {
    FloatVec floatVec = new FloatVec(embedding.getEmbedding().stream()
          .map(Double::floatValue)
          .collect(Collectors.toList()));

    SearchReq searchReq = SearchReq.builder()
            .collectionName(embeddingConfig.getCollectionName())
            .outputFields(List.of("question"))
            .topK(top)
            .data(List.of(floatVec))
            .searchParams(Map.of("metric_type", IndexParam.MetricType.COSINE))
            .build();
    SearchResp searchResp = client.search(searchReq);
    log.debug("search text:{}, result:{}", embedding.getQuestion(), JSON.toJSONString(searchResp));
    List<List<SearchResp.SearchResult>> results = searchResp.getSearchResults();
}

#Embedding结构
@Data
public class Embedding {

    private String object;

    private List<Double> embedding;

    private Integer index;

    private String question;

}