From 899047dbd1845cf8e5f0db05f44af537565eba4b Mon Sep 17 00:00:00 2001
From: lexluo09 <39718951+lexluo09@users.noreply.github.com>
Date: Mon, 11 Dec 2023 10:24:22 +0800
Subject: [PATCH] [improvement](chat) Adjust related packages, add other
 SqlGeneration implementations, and rename EmbedLLMProxy to JavaLLMProxy.
 (#489)

---
 .../{EmbedLLMProxy.java => JavaLLMProxy.java} | 13 ++++++----
 .../chat/parser/PythonLLMProxy.java           |  2 +-
 ...ator.java => FunctionPromptGenerator.java} |  4 ++--
 .../sql/llm/{prompt => }/InputFormat.java     |  2 +-
 .../parser/sql/llm/OneStepsSqlGeneration.java | 24 +++++++++++++++++++
 .../sql/llm/{prompt => }/OutputFormat.java    |  2 +-
 .../sql/llm/{prompt => }/SqlExample.java      |  2 +-
 .../llm/{prompt => }/SqlExampleLoader.java    |  2 +-
 .../chat/parser/sql/llm/SqlGeneration.java    |  3 +++
 .../llm/{prompt => }/SqlPromptGenerator.java  |  2 +-
 .../sql/llm/TwoStepCSSqlGeneration.java       | 24 +++++++++++++++++++
 ...eration.java => TwoStepSqlGeneration.java} |  7 ++----
 .../chat/query/llm/s2sql/LLMReq.java          |  6 ++---
 .../supersonic/EmbeddingInitListener.java     |  8 +++----
 .../main/resources/META-INF/spring.factories  |  2 +-
 15 files changed, 77 insertions(+), 26 deletions(-)
 rename chat/core/src/main/java/com/tencent/supersonic/chat/parser/{EmbedLLMProxy.java => JavaLLMProxy.java} (84%)
 rename chat/core/src/main/java/com/tencent/supersonic/chat/parser/plugin/function/{FunctionCallPromptGenerator.java => FunctionPromptGenerator.java} (94%)
 rename chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/{prompt => }/InputFormat.java (95%)
 create mode 100644 chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/OneStepsSqlGeneration.java
 rename chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/{prompt => }/OutputFormat.java (96%)
 rename chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/{prompt => }/SqlExample.java (91%)
 rename chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/{prompt => }/SqlExampleLoader.java (98%)
 rename chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/{prompt => }/SqlPromptGenerator.java (98%)
 create mode 100644 chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepCSSqlGeneration.java
 rename chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/{TwoStepsSqlGeneration.java => TwoStepSqlGeneration.java} (89%)

diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/EmbedLLMProxy.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/JavaLLMProxy.java
similarity index 84%
rename from chat/core/src/main/java/com/tencent/supersonic/chat/parser/EmbedLLMProxy.java
rename to chat/core/src/main/java/com/tencent/supersonic/chat/parser/JavaLLMProxy.java
index 4be84f15d..7f6ef1a2d 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/EmbedLLMProxy.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/JavaLLMProxy.java
@@ -1,12 +1,12 @@
 package com.tencent.supersonic.chat.parser;
 
 import com.tencent.supersonic.chat.api.pojo.QueryContext;
-import com.tencent.supersonic.chat.parser.plugin.function.FunctionCallPromptGenerator;
+import com.tencent.supersonic.chat.parser.plugin.function.FunctionPromptGenerator;
 import com.tencent.supersonic.chat.parser.plugin.function.FunctionReq;
 import com.tencent.supersonic.chat.parser.plugin.function.FunctionResp;
 import com.tencent.supersonic.chat.parser.sql.llm.SqlGeneration;
 import com.tencent.supersonic.chat.parser.sql.llm.SqlGenerationFactory;
-import com.tencent.supersonic.chat.parser.sql.llm.prompt.OutputFormat;
+import com.tencent.supersonic.chat.parser.sql.llm.OutputFormat;
 import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq;
 import com.tencent.supersonic.chat.query.llm.s2sql.LLMResp;
 import com.tencent.supersonic.common.util.ContextUtils;
@@ -14,14 +14,17 @@ import dev.langchain4j.model.chat.ChatLanguageModel;
 import java.util.Objects;
 import lombok.extern.slf4j.Slf4j;
 
+/**
+ * LLMProxy based on langchain4j Java version.
+ */
 @Slf4j
-public class EmbedLLMProxy implements LLMProxy {
+public class JavaLLMProxy implements LLMProxy {
 
     @Override
     public boolean isSkip(QueryContext queryContext) {
         ChatLanguageModel chatLanguageModel = ContextUtils.getBean(ChatLanguageModel.class);
         if (Objects.isNull(chatLanguageModel)) {
-            log.warn("chatLanguageModel is null, skip EmbedLLMProxy");
+            log.warn("chatLanguageModel is null, skip :{}", JavaLLMProxy.class.getName());
             return true;
         }
         return false;
@@ -43,7 +46,7 @@ public class EmbedLLMProxy implements LLMProxy {
     @Override
     public FunctionResp requestFunction(FunctionReq functionReq) {
 
-        FunctionCallPromptGenerator promptGenerator = ContextUtils.getBean(FunctionCallPromptGenerator.class);
+        FunctionPromptGenerator promptGenerator = ContextUtils.getBean(FunctionPromptGenerator.class);
 
         String functionCallPrompt = promptGenerator.generateFunctionCallPrompt(functionReq.getQueryText(),
                 functionReq.getPluginConfigs());
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/PythonLLMProxy.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/PythonLLMProxy.java
index 78d573b97..e36b51a5b 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/PythonLLMProxy.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/PythonLLMProxy.java
@@ -32,7 +32,7 @@ public class PythonLLMProxy implements LLMProxy {
     public boolean isSkip(QueryContext queryContext) {
         LLMParserConfig llmParserConfig = ContextUtils.getBean(LLMParserConfig.class);
         if (StringUtils.isEmpty(llmParserConfig.getUrl())) {
-            log.warn("llmParserUrl is empty, skip PythonLLMProxy, config:{}", llmParserConfig);
+            log.warn("llmParserUrl is empty, skip :{}", PythonLLMProxy.class.getName());
             return true;
         }
         return false;
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/plugin/function/FunctionCallPromptGenerator.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/plugin/function/FunctionPromptGenerator.java
similarity index 94%
rename from chat/core/src/main/java/com/tencent/supersonic/chat/parser/plugin/function/FunctionCallPromptGenerator.java
rename to chat/core/src/main/java/com/tencent/supersonic/chat/parser/plugin/function/FunctionPromptGenerator.java
index a1f21ae30..c1f9f5ed6 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/plugin/function/FunctionCallPromptGenerator.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/plugin/function/FunctionPromptGenerator.java
@@ -1,6 +1,6 @@
 package com.tencent.supersonic.chat.parser.plugin.function;
 
-import com.tencent.supersonic.chat.parser.sql.llm.prompt.InputFormat;
+import com.tencent.supersonic.chat.parser.sql.llm.InputFormat;
 import com.tencent.supersonic.chat.plugin.PluginParseConfig;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -9,7 +9,7 @@ import org.springframework.stereotype.Component;
 
 @Component
 @Slf4j
-public class FunctionCallPromptGenerator {
+public class FunctionPromptGenerator {
 
     public String generateFunctionCallPrompt(String queryText, List<PluginParseConfig> toolConfigList) {
         List<String> toolExplainList = toolConfigList.stream()
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/InputFormat.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/InputFormat.java
similarity index 95%
rename from chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/InputFormat.java
rename to chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/InputFormat.java
index 146e126d3..014186078 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/InputFormat.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/InputFormat.java
@@ -1,4 +1,4 @@
-package com.tencent.supersonic.chat.parser.sql.llm.prompt;
+package com.tencent.supersonic.chat.parser.sql.llm;
 
 import java.util.ArrayList;
 import java.util.HashMap;
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/OneStepsSqlGeneration.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/OneStepsSqlGeneration.java
new file mode 100644
index 000000000..bee7ac4a4
--- /dev/null
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/OneStepsSqlGeneration.java
@@ -0,0 +1,24 @@
+package com.tencent.supersonic.chat.parser.sql.llm;
+
+
+import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq;
+import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq.SqlGenerationMode;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.stereotype.Service;
+
+@Service
+@Slf4j
+public class OneStepsSqlGeneration implements SqlGeneration, InitializingBean {
+
+    @Override
+    public String generation(LLMReq llmReq, String modelClusterKey) {
+        //TODO
+        return "";
+    }
+
+    @Override
+    public void afterPropertiesSet() {
+        SqlGenerationFactory.addSqlGenerationForFactory(SqlGenerationMode.ONE_STEP, this);
+    }
+}
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/OutputFormat.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/OutputFormat.java
similarity index 96%
rename from chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/OutputFormat.java
rename to chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/OutputFormat.java
index ff09b79fb..fda0031a1 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/OutputFormat.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/OutputFormat.java
@@ -1,4 +1,4 @@
-package com.tencent.supersonic.chat.parser.sql.llm.prompt;
+package com.tencent.supersonic.chat.parser.sql.llm;
 
 import com.tencent.supersonic.chat.parser.plugin.function.FunctionResp;
 import com.tencent.supersonic.common.util.JsonUtil;
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/SqlExample.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlExample.java
similarity index 91%
rename from chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/SqlExample.java
rename to chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlExample.java
index c2f8df143..504266e83 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/SqlExample.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlExample.java
@@ -1,4 +1,4 @@
-package com.tencent.supersonic.chat.parser.sql.llm.prompt;
+package com.tencent.supersonic.chat.parser.sql.llm;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import lombok.Data;
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/SqlExampleLoader.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlExampleLoader.java
similarity index 98%
rename from chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/SqlExampleLoader.java
rename to chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlExampleLoader.java
index a8ec09ecb..4f14c45b0 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/SqlExampleLoader.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlExampleLoader.java
@@ -1,4 +1,4 @@
-package com.tencent.supersonic.chat.parser.sql.llm.prompt;
+package com.tencent.supersonic.chat.parser.sql.llm;
 
 
 import com.fasterxml.jackson.core.type.TypeReference;
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlGeneration.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlGeneration.java
index 7970828d0..8342282e6 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlGeneration.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlGeneration.java
@@ -3,6 +3,9 @@ package com.tencent.supersonic.chat.parser.sql.llm;
 
 import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq;
 
+/**
+ * Sql Generation
+ */
 public interface SqlGeneration {
 
     String generation(LLMReq llmReq, String modelClusterKey);
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/SqlPromptGenerator.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlPromptGenerator.java
similarity index 98%
rename from chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/SqlPromptGenerator.java
rename to chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlPromptGenerator.java
index 3552f06f4..5768c3d4c 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/prompt/SqlPromptGenerator.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/SqlPromptGenerator.java
@@ -1,4 +1,4 @@
-package com.tencent.supersonic.chat.parser.sql.llm.prompt;
+package com.tencent.supersonic.chat.parser.sql.llm;
 
 import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq.ElementValue;
 import java.util.Arrays;
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepCSSqlGeneration.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepCSSqlGeneration.java
new file mode 100644
index 000000000..bf013a2a3
--- /dev/null
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepCSSqlGeneration.java
@@ -0,0 +1,24 @@
+package com.tencent.supersonic.chat.parser.sql.llm;
+
+
+import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq;
+import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq.SqlGenerationMode;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.stereotype.Service;
+
+@Service
+@Slf4j
+public class TwoStepCSSqlGeneration implements SqlGeneration, InitializingBean {
+
+    @Override
+    public String generation(LLMReq llmReq, String modelClusterKey) {
+        //TODO
+        return "";
+    }
+
+    @Override
+    public void afterPropertiesSet() {
+        SqlGenerationFactory.addSqlGenerationForFactory(SqlGenerationMode.TWO_STEP_CS, this);
+    }
+}
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepsSqlGeneration.java b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepSqlGeneration.java
similarity index 89%
rename from chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepsSqlGeneration.java
rename to chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepSqlGeneration.java
index 54839bbd5..43285d883 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepsSqlGeneration.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/parser/sql/llm/TwoStepSqlGeneration.java
@@ -2,9 +2,6 @@ package com.tencent.supersonic.chat.parser.sql.llm;
 
 
 import com.tencent.supersonic.chat.config.OptimizationConfig;
-import com.tencent.supersonic.chat.parser.sql.llm.prompt.OutputFormat;
-import com.tencent.supersonic.chat.parser.sql.llm.prompt.SqlExampleLoader;
-import com.tencent.supersonic.chat.parser.sql.llm.prompt.SqlPromptGenerator;
 import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq;
 import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq.ElementValue;
 import com.tencent.supersonic.chat.query.llm.s2sql.LLMReq.SqlGenerationMode;
@@ -25,7 +22,7 @@ import org.springframework.stereotype.Service;
 
 @Service
 @Slf4j
-public class TwoStepsSqlGeneration implements SqlGeneration, InitializingBean {
+public class TwoStepSqlGeneration implements SqlGeneration, InitializingBean {
 
     @Autowired
     private ChatLanguageModel chatLanguageModel;
@@ -68,6 +65,6 @@ public class TwoStepsSqlGeneration implements SqlGeneration, InitializingBean {
 
     @Override
     public void afterPropertiesSet() {
-        SqlGenerationFactory.addSqlGenerationForFactory(SqlGenerationMode.TWO_STEPS, this);
+        SqlGenerationFactory.addSqlGenerationForFactory(SqlGenerationMode.TWO_STEP, this);
     }
 }
diff --git a/chat/core/src/main/java/com/tencent/supersonic/chat/query/llm/s2sql/LLMReq.java b/chat/core/src/main/java/com/tencent/supersonic/chat/query/llm/s2sql/LLMReq.java
index cb5c215ce..2784c4ec0 100644
--- a/chat/core/src/main/java/com/tencent/supersonic/chat/query/llm/s2sql/LLMReq.java
+++ b/chat/core/src/main/java/com/tencent/supersonic/chat/query/llm/s2sql/LLMReq.java
@@ -18,7 +18,7 @@ public class LLMReq {
 
     private String priorExts;
 
-    private SqlGenerationMode sqlGenerationMode = SqlGenerationMode.TWO_STEPS;
+    private SqlGenerationMode sqlGenerationMode = SqlGenerationMode.TWO_STEP;
 
     @Data
     public static class ElementValue {
@@ -50,9 +50,9 @@ public class LLMReq {
 
         ONE_STEP("ONE_STEP"),
 
-        TWO_STEPS("TWO_STEPS"),
+        TWO_STEP("TWO_STEP"),
 
-        TWO_STEPS_WITH_CS("TWO_STEPS_WITH_CS");
+        TWO_STEP_CS("TWO_STEP_CS");
 
 
         private String name;
diff --git a/launchers/standalone/src/main/java/com/tencent/supersonic/EmbeddingInitListener.java b/launchers/standalone/src/main/java/com/tencent/supersonic/EmbeddingInitListener.java
index 1a46e2261..1b2c6e117 100644
--- a/launchers/standalone/src/main/java/com/tencent/supersonic/EmbeddingInitListener.java
+++ b/launchers/standalone/src/main/java/com/tencent/supersonic/EmbeddingInitListener.java
@@ -1,9 +1,9 @@
 package com.tencent.supersonic;
 
 import com.tencent.supersonic.chat.config.OptimizationConfig;
-import com.tencent.supersonic.chat.parser.sql.llm.prompt.SqlExample;
-import com.tencent.supersonic.chat.parser.sql.llm.prompt.SqlExampleLoader;
-import com.tencent.supersonic.chat.parser.EmbedLLMProxy;
+import com.tencent.supersonic.chat.parser.sql.llm.SqlExample;
+import com.tencent.supersonic.chat.parser.sql.llm.SqlExampleLoader;
+import com.tencent.supersonic.chat.parser.JavaLLMProxy;
 import com.tencent.supersonic.chat.parser.LLMProxy;
 import com.tencent.supersonic.chat.utils.ComponentFactory;
 import java.util.List;
@@ -31,7 +31,7 @@ public class EmbeddingInitListener implements CommandLineRunner {
 
     public void initSqlExamples() {
         try {
-            if (llmProxy instanceof EmbedLLMProxy) {
+            if (llmProxy instanceof JavaLLMProxy) {
                 List<SqlExample> sqlExamples = sqlExampleLoader.getSqlExamples();
                 String collectionName = optimizationConfig.getText2sqlCollectionName();
                 sqlExampleLoader.addEmbeddingStore(sqlExamples, collectionName);
diff --git a/launchers/standalone/src/main/resources/META-INF/spring.factories b/launchers/standalone/src/main/resources/META-INF/spring.factories
index 35545c639..3a94267ba 100644
--- a/launchers/standalone/src/main/resources/META-INF/spring.factories
+++ b/launchers/standalone/src/main/resources/META-INF/spring.factories
@@ -31,7 +31,7 @@ com.tencent.supersonic.chat.processor.ParseResultProcessor=\
     com.tencent.supersonic.chat.processor.RespBuildProcessor
 
 com.tencent.supersonic.chat.parser.LLMProxy=\
-  com.tencent.supersonic.chat.parser.EmbedLLMProxy
+  com.tencent.supersonic.chat.parser.JavaLLMProxy
 
 com.tencent.supersonic.chat.api.component.SemanticInterpreter=\
     com.tencent.supersonic.knowledge.semantic.LocalSemanticInterpreter