From c045b34328683800a0bde48597789cd06ff6d328 Mon Sep 17 00:00:00 2001
From: jerryjzhang <zhangjun2915@163.com>
Date: Fri, 20 Sep 2024 16:46:49 +0800
Subject: [PATCH] [improvement][chat]Expose default prompt template to
 facilitate customization.

---
 .../chat/server/rest/AgentController.java          | 10 +++++-----
 .../chat/server/service/impl/AgentServiceImpl.java | 14 ++++++++++++++
 .../chat/parser/llm/OnePassSCSqlGenStrategy.java   |  2 +-
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/chat/server/src/main/java/com/tencent/supersonic/chat/server/rest/AgentController.java b/chat/server/src/main/java/com/tencent/supersonic/chat/server/rest/AgentController.java
index cedc23462..62d8020ad 100644
--- a/chat/server/src/main/java/com/tencent/supersonic/chat/server/rest/AgentController.java
+++ b/chat/server/src/main/java/com/tencent/supersonic/chat/server/rest/AgentController.java
@@ -52,11 +52,6 @@ public class AgentController {
         return true;
     }
 
-    @PostMapping("/testLLMConn")
-    public boolean testLLMConn(@RequestBody ChatModelConfig modelConfig) {
-        return LLMConnHelper.testConnection(modelConfig);
-    }
-
     @RequestMapping("/getAgentList")
     public List<Agent> getAgentList() {
         return agentService.getAgents();
@@ -66,4 +61,9 @@ public class AgentController {
     public Map<AgentToolType, String> getToolTypes() {
         return AgentToolType.getToolTypes();
     }
+
+    @PostMapping("/testLLMConn")
+    public boolean testLLMConn(@RequestBody ChatModelConfig modelConfig) {
+        return LLMConnHelper.testConnection(modelConfig);
+    }
 }
diff --git a/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/AgentServiceImpl.java b/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/AgentServiceImpl.java
index 566d8d7b6..a6ef3e7ec 100644
--- a/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/AgentServiceImpl.java
+++ b/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/AgentServiceImpl.java
@@ -16,6 +16,7 @@ import com.tencent.supersonic.common.config.PromptConfig;
 import com.tencent.supersonic.common.config.VisualConfig;
 import com.tencent.supersonic.common.pojo.ChatModelConfig;
 import com.tencent.supersonic.common.util.JsonUtil;
+import com.tencent.supersonic.headless.chat.parser.llm.OnePassSCSqlGenStrategy;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.BeanUtils;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -23,6 +24,7 @@ import org.springframework.stereotype.Service;
 import org.springframework.util.CollectionUtils;
 
 import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.stream.Collectors;
@@ -44,6 +46,12 @@ public class AgentServiceImpl extends ServiceImpl<AgentDOMapper, AgentDO> implem
 
     @Override
     public Agent createAgent(Agent agent, User user) {
+        if (Objects.isNull(agent.getPromptConfig())
+                || Objects.isNull(agent.getPromptConfig().getPromptTemplate())) {
+            PromptConfig promptConfig = new PromptConfig();
+            promptConfig.setPromptTemplate(OnePassSCSqlGenStrategy.INSTRUCTION.trim());
+            agent.setPromptConfig(promptConfig);
+        }
         agent.createdBy(user.getName());
         AgentDO agentDO = convert(agent);
         save(agentDO);
@@ -54,6 +62,12 @@ public class AgentServiceImpl extends ServiceImpl<AgentDOMapper, AgentDO> implem
 
     @Override
     public Agent updateAgent(Agent agent, User user) {
+        if (Objects.isNull(agent.getPromptConfig())
+                || Objects.isNull(agent.getPromptConfig().getPromptTemplate())) {
+            PromptConfig promptConfig = new PromptConfig();
+            promptConfig.setPromptTemplate(OnePassSCSqlGenStrategy.INSTRUCTION.trim());
+            agent.setPromptConfig(promptConfig);
+        }
         agent.updatedBy(user.getName());
         updateById(convert(agent));
         executeAgentExamplesAsync(agent);
diff --git a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/parser/llm/OnePassSCSqlGenStrategy.java b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/parser/llm/OnePassSCSqlGenStrategy.java
index 3852f7120..8145da458 100644
--- a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/parser/llm/OnePassSCSqlGenStrategy.java
+++ b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/parser/llm/OnePassSCSqlGenStrategy.java
@@ -25,7 +25,7 @@ import java.util.concurrent.ConcurrentHashMap;
 @Slf4j
 public class OnePassSCSqlGenStrategy extends SqlGenStrategy {
 
-    private static final String INSTRUCTION =
+    public static final String INSTRUCTION =
             ""
                     + "\n#Role: You are a data analyst experienced in SQL languages."
                     + "\n#Task: You will be provided with a natural language question asked by users,"