From afa82bf98dcd000f60185ea9ca5b6dc32ef5f321 Mon Sep 17 00:00:00 2001
From: lexluo09 <39718951+lexluo09@users.noreply.github.com>
Date: Wed, 21 Aug 2024 16:43:51 +0800
Subject: [PATCH] (improvement)(chat) Support connecting to Milvus service in
 the cloud. (#1588)

---
 .../config/EmbeddingStoreParameterConfig.java | 19 ++++++++++++++++---
 .../common/pojo/EmbeddingStoreConfig.java     |  1 +
 .../spring/MilvusEmbeddingStoreFactory.java   |  1 +
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java b/common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java
index cd4937491..9ab1d74d5 100644
--- a/common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java
+++ b/common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java
@@ -45,12 +45,17 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
             new Parameter("s2.embedding.store.dimension", "",
                     "纬度", "", "number",
                     "向量库配置", null, getDimensionDependency());
+    public static final Parameter EMBEDDING_STORE_DATABASE_NAME =
+            new Parameter("s2.embedding.store.databaseName", "",
+                    "DatabaseName", "", "string",
+                    "向量库配置", null, getDatabaseNameDependency());
 
     @Override
     public List<Parameter> getSysParameters() {
         return Lists.newArrayList(
                 EMBEDDING_STORE_PROVIDER, EMBEDDING_STORE_BASE_URL, EMBEDDING_STORE_API_KEY,
-                EMBEDDING_STORE_PERSIST_PATH, EMBEDDING_STORE_TIMEOUT, EMBEDDING_STORE_DIMENSION
+                EMBEDDING_STORE_DATABASE_NAME, EMBEDDING_STORE_PERSIST_PATH,
+                EMBEDDING_STORE_TIMEOUT, EMBEDDING_STORE_DIMENSION
         );
     }
 
@@ -60,12 +65,13 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
         String apiKey = getParameterValue(EMBEDDING_STORE_API_KEY);
         String persistPath = getParameterValue(EMBEDDING_STORE_PERSIST_PATH);
         String timeOut = getParameterValue(EMBEDDING_STORE_TIMEOUT);
+        String databaseName = getParameterValue(EMBEDDING_STORE_DATABASE_NAME);
         Integer dimension = null;
         if (StringUtils.isNumeric(getParameterValue(EMBEDDING_STORE_DIMENSION))) {
             dimension = Integer.valueOf(getParameterValue(EMBEDDING_STORE_DIMENSION));
         }
-        return EmbeddingStoreConfig.builder().provider(provider)
-                .baseUrl(baseUrl).apiKey(apiKey).persistPath(persistPath)
+        return EmbeddingStoreConfig.builder().provider(provider).baseUrl(baseUrl)
+                .apiKey(apiKey).persistPath(persistPath).databaseName(databaseName)
                 .timeOut(Long.valueOf(timeOut)).dimension(dimension).build();
     }
 
@@ -106,4 +112,11 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
                 ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "384")
         );
     }
+
+    private static List<Parameter.Dependency> getDatabaseNameDependency() {
+        return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
+                Lists.newArrayList(EmbeddingStoreType.MILVUS.name()),
+                ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "")
+        );
+    }
 }
\ No newline at end of file
diff --git a/common/src/main/java/com/tencent/supersonic/common/pojo/EmbeddingStoreConfig.java b/common/src/main/java/com/tencent/supersonic/common/pojo/EmbeddingStoreConfig.java
index cd9ad5275..aa6dfff5d 100644
--- a/common/src/main/java/com/tencent/supersonic/common/pojo/EmbeddingStoreConfig.java
+++ b/common/src/main/java/com/tencent/supersonic/common/pojo/EmbeddingStoreConfig.java
@@ -21,4 +21,5 @@ public class EmbeddingStoreConfig implements Serializable {
     private String apiKey;
     private Long timeOut = 60L;
     private Integer dimension;
+    private String databaseName;
 }
\ No newline at end of file
diff --git a/common/src/main/java/dev/langchain4j/milvus/spring/MilvusEmbeddingStoreFactory.java b/common/src/main/java/dev/langchain4j/milvus/spring/MilvusEmbeddingStoreFactory.java
index 2ff3733fa..3e3aac552 100644
--- a/common/src/main/java/dev/langchain4j/milvus/spring/MilvusEmbeddingStoreFactory.java
+++ b/common/src/main/java/dev/langchain4j/milvus/spring/MilvusEmbeddingStoreFactory.java
@@ -22,6 +22,7 @@ public class MilvusEmbeddingStoreFactory extends BaseEmbeddingStoreFactory {
         EmbeddingStoreProperties embeddingStore = new EmbeddingStoreProperties();
         BeanUtils.copyProperties(storeConfig, embeddingStore);
         embeddingStore.setUri(storeConfig.getBaseUrl());
+        embeddingStore.setToken(storeConfig.getApiKey());
         return embeddingStore;
     }