From ff0833ddfa5e8006651b1756d0365787c7bfe479 Mon Sep 17 00:00:00 2001
From: lexluo09 <39718951+lexluo09@users.noreply.github.com>
Date: Wed, 11 Sep 2024 22:52:46 +0800
Subject: [PATCH] [improvement](headless) Optimize the reload embedding
 interface (#1655)

---
 .../common/service/ExemplarService.java       |  2 ++
 .../service/impl/ExemplarServiceImpl.java     | 19 +++++++++----------
 .../server/rest/KnowledgeController.java      | 14 +++++++++-----
 3 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/common/src/main/java/com/tencent/supersonic/common/service/ExemplarService.java b/common/src/main/java/com/tencent/supersonic/common/service/ExemplarService.java
index 2706274b7..f0e5bfdad 100644
--- a/common/src/main/java/com/tencent/supersonic/common/service/ExemplarService.java
+++ b/common/src/main/java/com/tencent/supersonic/common/service/ExemplarService.java
@@ -12,4 +12,6 @@ public interface ExemplarService {
     List<Text2SQLExemplar> recallExemplars(String collection, String query, int num);
 
     List<Text2SQLExemplar> recallExemplars(String query, int num);
+
+    void loadSysExemplars();
 }
diff --git a/common/src/main/java/com/tencent/supersonic/common/service/impl/ExemplarServiceImpl.java b/common/src/main/java/com/tencent/supersonic/common/service/impl/ExemplarServiceImpl.java
index 7b4eb92a7..d52e03d08 100644
--- a/common/src/main/java/com/tencent/supersonic/common/service/impl/ExemplarServiceImpl.java
+++ b/common/src/main/java/com/tencent/supersonic/common/service/impl/ExemplarServiceImpl.java
@@ -20,7 +20,6 @@ import org.springframework.core.annotation.Order;
 import org.springframework.core.io.ClassPathResource;
 import org.springframework.stereotype.Service;
 
-import java.io.IOException;
 import java.io.InputStream;
 import java.util.List;
 
@@ -88,18 +87,18 @@ public class ExemplarServiceImpl implements ExemplarService, CommandLineRunner {
 
     @Override
     public void run(String... args) {
+        loadSysExemplars();
+    }
+
+    public void loadSysExemplars() {
         try {
-            loadSysExemplars();
+            ClassPathResource resource = new ClassPathResource(SYS_EXEMPLAR_FILE);
+            InputStream inputStream = resource.getInputStream();
+            List<Text2SQLExemplar> exemplars = objectMapper.readValue(inputStream, valueTypeRef);
+            String collection = embeddingConfig.getText2sqlCollectionName();
+            exemplars.stream().forEach(e -> storeExemplar(collection, e));
         } catch (Exception e) {
             log.error("Failed to load system exemplars", e);
         }
     }
-
-    private void loadSysExemplars() throws IOException {
-        ClassPathResource resource = new ClassPathResource(SYS_EXEMPLAR_FILE);
-        InputStream inputStream = resource.getInputStream();
-        List<Text2SQLExemplar> exemplars = objectMapper.readValue(inputStream, valueTypeRef);
-        String collection = embeddingConfig.getText2sqlCollectionName();
-        exemplars.stream().forEach(e -> storeExemplar(collection, e));
-    }
 }
diff --git a/headless/server/src/main/java/com/tencent/supersonic/headless/server/rest/KnowledgeController.java b/headless/server/src/main/java/com/tencent/supersonic/headless/server/rest/KnowledgeController.java
index 64dcc64ab..0c782571c 100644
--- a/headless/server/src/main/java/com/tencent/supersonic/headless/server/rest/KnowledgeController.java
+++ b/headless/server/src/main/java/com/tencent/supersonic/headless/server/rest/KnowledgeController.java
@@ -7,6 +7,7 @@ import javax.validation.Valid;
 import com.github.pagehelper.PageInfo;
 import com.tencent.supersonic.auth.api.authentication.pojo.User;
 import com.tencent.supersonic.auth.api.authentication.utils.UserHolder;
+import com.tencent.supersonic.common.service.ExemplarService;
 import com.tencent.supersonic.headless.api.pojo.request.DictItemFilter;
 import com.tencent.supersonic.headless.api.pojo.request.DictItemReq;
 import com.tencent.supersonic.headless.api.pojo.request.DictSingleTaskReq;
@@ -36,10 +37,12 @@ public class KnowledgeController {
 
     @Autowired private DictConfService confService;
 
-    @Autowired private MetaEmbeddingTask embeddingTask;
+    @Autowired private MetaEmbeddingTask metaEmbeddingTask;
 
     @Autowired private DictionaryReloadTask dictionaryReloadTask;
 
+    @Autowired private ExemplarService exemplarService;
+
     /**
      * addDictConf-新增item的字典配置 Add configuration information for dictionary entries
      *
@@ -130,15 +133,16 @@ public class KnowledgeController {
         return taskService.queryLatestDictTask(taskReq, user);
     }
 
-    @GetMapping("/meta/embedding/reload")
-    public Object reloadMetaEmbedding() {
-        embeddingTask.reloadMetaEmbedding();
+    @GetMapping("/embedding/reload")
+    public Object reloadEmbedding() {
+        metaEmbeddingTask.reloadMetaEmbedding();
+        exemplarService.loadSysExemplars();
         return true;
     }
 
     @GetMapping("/embedding/persistFile")
     public Object executePersistFileTask() {
-        embeddingTask.executePersistFileTask();
+        metaEmbeddingTask.executePersistFileTask();
         return true;
     }