content: /*
   *  Copyright (c) 2023-2025, Agents-Flex (fuhai999@gmail.com).
   *  <p>
   *  Licensed under the Apache License, Version 2.0 (the "License");
   *  you may not use this file except in compliance with the License.
   *  You may obtain a copy of the License at
   *  <p>
   *  http://www.apache.org/licenses/LICENSE-2.0
   *  <p>
   *  Unless required by applicable law or agreed to in writing, software
   *  distributed under the License is distributed on an "AS IS" BASIS,
   *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   *  See the License for the specific language governing permissions and
   *  limitations under the License.
   */
  package com.agentsflex.store.milvus;
  
  import com.agentsflex.core.document.Document;
  import com.agentsflex.core.store.DocumentStore;
  import com.agentsflex.core.store.SearchWrapper;
  import com.agentsflex.core.store.StoreOptions;
  import com.agentsflex.core.store.StoreResult;
  import com.agentsflex.core.util.Maps;
  import com.agentsflex.core.util.StringUtil;
  import com.agentsflex.core.util.VectorUtil;
  import com.alibaba.fastjson.JSON;
  import com.alibaba.fastjson.JSONObject;
  import io.milvus.v2.client.ConnectConfig;
  import io.milvus.v2.client.MilvusClientV2;
  import io.milvus.v2.common.ConsistencyLevel;
  import io.milvus.v2.common.DataType;
  import io.milvus.v2.common.IndexParam;
  import io.milvus.v2.exception.MilvusClientException;
  import io.milvus.v2.service.collection.request.CreateCollectionReq;
  import io.milvus.v2.service.collection.request.GetLoadStateReq;
  import io.milvus.v2.service.vector.request.*;
  import io.milvus.v2.service.vector.response.InsertResp;
  import io.milvus.v2.service.vector.response.QueryResp;
  import io.milvus.v2.service.vector.response.SearchResp;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.*;
  
  /**
   * MilvusVectorStore class provides an interface to interact with Milvus Vector Database.
   */
  public class MilvusVectorStore extends DocumentStore {
  
      private static final Logger logger = LoggerFactory.getLogger(MilvusVectorStore.class);
      private final MilvusClientV2 client;
      private final String defaultCollectionName;
      private final MilvusVectorStoreConfig config;
  
      public MilvusVectorStore(MilvusVectorStoreConfig config) {
          ConnectConfig connectConfig = ConnectConfig.builder()
              .uri(config.getUri())
              .dbName(config.getDatabaseName())
              .token(config.getToken())
              .username(config.getUsername())
              .password(config.getPassword())
              .build();
  
          this.client = new MilvusClientV2(connectConfig);
          this.defaultCollectionName = config.getDefaultCollectionName();
          this.config = config;
      }
  
      @Override
      public StoreResult storeInternal(List<Document> documents, StoreOptions options) {
          List<JSONObject> data = new ArrayList<>();
          for (Document doc : documents) {
              JSONObject dict = new JSONObject();
              dict.put("id", String.valueOf(doc.getId()));
              dict.put("content", doc.getContent());
              dict.put("vector", VectorUtil.toFloatList(doc.getVector()));
  
              Map<String, Object> metadatas = doc.getMetadataMap();
              JSONObject jsonObject = JSON.parseObject(JSON.toJSONBytes(metadatas == null ? Collections.EMPTY_MAP : metadatas));
              dict.put("metadata", jsonObject);
              data.add(dict);
          }
  
          String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
          InsertReq.InsertReqBuilder<?, ?> builder = InsertReq.builder();
          if (StringUtil.hasText(options.getPartitionName())) {
              builder.partitionName(options.getPartitionName());
          }
          InsertReq insertReq = builder
              .collectionName(collectionName)
              .data(data)
              .build();
          try {
              InsertResp insertResp = client.insert(insertReq);
          } catch (MilvusClientException e) {
              if (e.getMessage() != null && e.getMessage().contains("collection not found")
                  && config.isAutoCreateCollection()
                  && options.getMetadata("forInternal") == null) {
  
                  Boolean success = createCollection(collectionName);
                  if (success != null && success) {
                      //store
                      options.addMetadata("forInternal", true);
                      storeInternal(documents, options);
                  }
              } else {
                  return StoreResult.fail();
              }
          }
  
          return StoreResult.successWithIds(documents);
      }
  
  
      private Boolean createCollection(String collectionName) {
          List<CreateCollectionReq.FieldSchema> fieldSchemaList = new ArrayList<>();
  
          //id
          CreateCollectionReq.FieldSchema id = CreateCollectionReq.FieldSchema.builder()
              .name("id")
              .dataType(DataType.VarChar)
              .maxLength(36)
              .isPrimaryKey(true)
              .autoID(false)
              .build();
          fieldSchemaList.add(id);
  
          //content
          CreateCollectionReq.FieldSchema content = CreateCollectionReq.FieldSchema.builder()
              .name("content")
              .dataType(DataType.VarChar)
              .maxLength(65535)
              .build();
          fieldSchemaList.add(content);
  
          //metadata
          CreateCollectionReq.FieldSchema metadata = CreateCollectionReq.FieldSchema.builder()
              .name("metadata")
              .dataType(DataType.JSON)
              .build();
          fieldSchemaList.add(metadata);
  
          //vector
          CreateCollectionReq.FieldSchema vector = CreateCollectionReq.FieldSchema.builder()
              .name("vector")
              .dataType(DataType.FloatVector)
              .dimension(this.getEmbeddingModel().dimensions())
              .build();
          fieldSchemaList.add(vector);
  
          CreateCollectionReq.CollectionSchema collectionSchema = CreateCollectionReq.CollectionSchema
              .builder()
              .fieldSchemaList(fieldSchemaList)
              .build();
  
  
          List<IndexParam> indexParams = new ArrayList<>();
          IndexParam vectorIndex = IndexParam.builder().fieldName("vector")
              .indexType(IndexParam.IndexType.IVF_FLAT)
              .metricType(IndexParam.MetricType.COSINE)
              .indexName("vector")
              .extraParams(Maps.of("nlist", 1024))
              .build();
          indexParams.add(vectorIndex);
  
  
          CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
              .collectionName(collectionName)
              .collectionSchema(collectionSchema)
              .primaryFieldName("id")
              .vectorFieldName("vector")
              .description("Agents Flex Vector Store")
              .indexParams(indexParams)
              .build();
  
          client.createCollection(createCollectionReq);
  
          GetLoadStateReq quickSetupLoadStateReq = GetLoadStateReq.builder()
              .collectionName(collectionName)
              .build();
  
          return client.getLoadState(quickSetupLoadStateReq);
      }
  
      @Override
      public StoreResult deleteInternal(Collection<?> ids, StoreOptions options) {
  
          DeleteReq.DeleteReqBuilder<?, ?> builder = DeleteReq.builder();
          if (StringUtil.hasText(options.getPartitionName())) {
              builder.partitionName(options.getPartitionName());
          }
  
          DeleteReq deleteReq = builder
              .collectionName(options.getCollectionNameOrDefault(defaultCollectionName))
              .ids(new ArrayList<>(ids))
              .build();
  
          try {
              client.delete(deleteReq);
          } catch (Exception e) {
              logger.error("delete document error: " + e, e);
              return StoreResult.fail();
          }
  
          return StoreResult.success();
  
      }
  
      @Override
      public List<Document> searchInternal(SearchWrapper searchWrapper, StoreOptions options) {
          List<String> outputFields = searchWrapper.isOutputVector()
              ? Arrays.asList("id", "vector", "content", "metadata")
              : Arrays.asList("id", "content", "metadata");
  
          // 判断是否为向量查询
          if (searchWrapper.getVector() != null && searchWrapper.getVector().length > 0) {
              // 向量查询 - 使用SearchReq
              SearchReq.SearchReqBuilder<?, ?> builder = SearchReq.builder();
              if (StringUtil.hasText(options.getPartitionName())) {
                  builder.partitionNames(options.getPartitionNamesOrEmpty());
              }
  
              SearchReq searchReq = builder
                  .collectionName(options.getCollectionNameOrDefault(defaultCollectionName))
                  .consistencyLevel(ConsistencyLevel.STRONG)
                  .outputFields(outputFields)
                  .topK(searchWrapper.getMaxResults())
                  .annsField("vector")
                  .data(Collections.singletonList(VectorUtil.toFloatList(searchWrapper.getVector())))
                  .filter(searchWrapper.toFilterExpression(MilvusExpressionAdaptor.DEFAULT))
                  .build();
  
              try {
                  SearchResp resp = client.search(searchReq);
                  // Parse and convert search results to Document list
                  List<List<SearchResp.SearchResult>> results = resp.getSearchResults();
                  List<Document> documents = new ArrayList<>();
                  for (List<SearchResp.SearchResult> resultList : results) {
                      for (SearchResp.SearchResult result : resultList) {
                          Map<String, Object> entity = result.getEntity();
                          if (entity == null || entity.isEmpty()) {
                              continue;
                          }
  
                          Document doc = new Document();
                          doc.setId(result.getId());
  
                          Object vectorObj = entity.get("vector");
                          if (vectorObj instanceof List) {
                              //noinspection unchecked
                              doc.setVector(VectorUtil.convertToVector((List<Float>) vectorObj));
                          }
  
                          doc.setContent((String) entity.get("content"));
  
                          // 根据 metric 类型计算相似度
                          Float distance = result.getDistance();
                          if (distance != null) {
                              // 根据 https://milvus.io/docs/zh/single-vector-search.md#Single-Vector-Search
                              // 适用的度量类型和相应的距离范围表
                              // 当 metricType 类 COSINE 时，数值越大，表示相似度越高。相似度即为 distance 值;
  
                              // distance 的范围是 [-1, 1], 需要统一转换为 [0, 1]
                              double score = (distance + 1) / 2;
                              doc.setScore(score);
                          }
  
                          JSONObject object = (JSONObject) entity.get("metadata");
                          doc.addMetadata(object);
                          documents.add(doc);
                      }
                  }
                  return documents;
              } catch (Exception e) {
                  logger.error("Error searching in Milvus", e);
                  return Collections.emptyList();
              }
          } else {
              // 非向量查询 - 使用QueryReq
              QueryReq.QueryReqBuilder<?, ?> builder = QueryReq.builder();
              if (StringUtil.hasText(options.getPartitionName())) {
                  builder.partitionNames(options.getPartitionNamesOrEmpty());
              }
  
              QueryReq queryReq = builder
                  .collectionName(options.getCollectionNameOrDefault(defaultCollectionName))
                  .consistencyLevel(ConsistencyLevel.STRONG)
                  .outputFields(outputFields)
                  .filter(searchWrapper.toFilterExpression(MilvusExpressionAdaptor.DEFAULT))
                  .build();
  
              try {
                  QueryResp resp = client.query(queryReq);
                  List<QueryResp.QueryResult> results = resp.getQueryResults();
                  List<Document> documents = new ArrayList<>();
  
                  for (QueryResp.QueryResult result : results) {
                      Map<String, Object> entity = result.getEntity();
                      if (entity == null || entity.isEmpty()) {
                          continue;
                      }
  
                      Document doc = new Document();
                      doc.setId(result.getEntity().get("id"));
  
                      Object vectorObj = entity.get("vector");
                      if (vectorObj instanceof List) {
                          //noinspection unchecked
                          doc.setVector(VectorUtil.convertToVector((List<Float>) vectorObj));
                      }
  
                      doc.setContent((String) entity.get("content"));
  
                      JSONObject object = (JSONObject) entity.get("metadata");
                      doc.addMetadata(object);
                      documents.add(doc);
                  }
                  return documents;
              } catch (Exception e) {
                  logger.error("Error querying in Milvus", e);
                  return Collections.emptyList();
              }
          }
      }
  
      @Override
      public StoreResult updateInternal(List<Document> documents, StoreOptions options) {
          if (documents == null || documents.isEmpty()) {
              return StoreResult.success();
          }
          List<JSONObject> data = new ArrayList<>();
          for (Document doc : documents) {
              JSONObject dict = new JSONObject();
  
              dict.put("id", String.valueOf(doc.getId()));
              dict.put("content", doc.getContent());
              dict.put("vector", VectorUtil.toFloatList(doc.getVector()));
  
              Map<String, Object> metadatas = doc.getMetadataMap();
              JSONObject jsonObject = JSON.parseObject(JSON.toJSONBytes(metadatas == null ? Collections.EMPTY_MAP : metadatas));
              dict.put("metadata", jsonObject);
              data.add(dict);
          }
  
          UpsertReq upsertReq = UpsertReq.builder()
              .collectionName(options.getCollectionNameOrDefault(defaultCollectionName))
              .partitionName(options.getPartitionName())
              .data(data)
              .build();
          client.upsert(upsertReq);
          return StoreResult.successWithIds(documents);
      }
  
  
      public MilvusClientV2 getClient() {
          return client;
      }
  }
  