add new chat corrector

1985312383 · 1985312383 · commit f899d23b6320 · 2025-06-21T04:57:04.000+08:00
在助理最终执行物理SQL前，加入一步LLM优化性能功能
diff --git a/chat/server/src/main/java/com/tencent/supersonic/chat/server/executor/SqlExecutor.java b/chat/server/src/main/java/com/tencent/supersonic/chat/server/executor/SqlExecutor.java
@@ -75,8 +75,12 @@ private QueryResult doExecute(ExecuteContext executeContext) {
             return null;
         }
 
-        QuerySqlReq sqlReq =
-                QuerySqlReq.builder().sql(parseInfo.getSqlInfo().getCorrectedS2SQL()).build();
+        // 使用querySQL，它已经包含了所有修正（包括物理SQL修正）
+        String finalSql = StringUtils.isNotBlank(parseInfo.getSqlInfo().getQuerySQL())
+                ? parseInfo.getSqlInfo().getQuerySQL()
+                : parseInfo.getSqlInfo().getCorrectedS2SQL();
+
+        QuerySqlReq sqlReq = QuerySqlReq.builder().sql(finalSql).build();
         sqlReq.setSqlInfo(parseInfo.getSqlInfo());
         sqlReq.setDataSetId(parseInfo.getDataSetId());
 
@@ -90,7 +94,7 @@ private QueryResult doExecute(ExecuteContext executeContext) {
         queryResult.setQueryTimeCost(System.currentTimeMillis() - startTime);
         if (queryResp != null) {
             queryResult.setQueryAuthorization(queryResp.getQueryAuthorization());
-            queryResult.setQuerySql(queryResp.getSql());
+            queryResult.setQuerySql(finalSql);
             queryResult.setQueryResults(queryResp.getResultList());
             queryResult.setQueryColumns(queryResp.getColumns());
             queryResult.setQueryState(QueryState.SUCCESS);
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
@@ -22,6 +22,7 @@
 import org.springframework.beans.BeanUtils;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.context.annotation.Lazy;
 import org.springframework.stereotype.Service;
 import org.springframework.util.CollectionUtils;
 
@@ -39,6 +40,7 @@ public class AgentServiceImpl extends ServiceImpl<AgentDOMapper, AgentDO> implem
     private MemoryService memoryService;
 
     @Autowired
+    @Lazy
     private ChatQueryService chatQueryService;
 
     @Autowired
diff --git a/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/ChatQueryServiceImpl.java b/chat/server/src/main/java/com/tencent/supersonic/chat/server/service/impl/ChatQueryServiceImpl.java
@@ -49,6 +49,7 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Lazy;
 import org.springframework.stereotype.Service;
 import org.springframework.util.CollectionUtils;
 
@@ -66,6 +67,7 @@ public class ChatQueryServiceImpl implements ChatQueryService {
     @Autowired
     private SemanticLayerService semanticLayerService;
     @Autowired
+    @Lazy
     private AgentService agentService;
 
     private final List<ChatQueryParser> chatQueryParsers = ComponentFactory.getChatParsers();
diff --git a/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/SqlInfo.java b/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/SqlInfo.java
@@ -16,4 +16,7 @@ public class SqlInfo implements Serializable {
 
     // SQL to be executed finally
     private String querySQL;
+
+    // Physical SQL corrected by LLM for performance optimization
+    private String correctedQuerySQL;
 }
diff --git a/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/enums/ChatWorkflowState.java b/headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/enums/ChatWorkflowState.java
@@ -8,5 +8,6 @@ public enum ChatWorkflowState {
     VALIDATING,
     SQL_CORRECTING,
     PROCESSING,
+    PHYSICAL_SQL_CORRECTING,
     FINISHED
 }
diff --git a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/corrector/LLMPhysicalSqlCorrector.java b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/corrector/LLMPhysicalSqlCorrector.java
@@ -0,0 +1,97 @@
+package com.tencent.supersonic.headless.chat.corrector;
+
+import com.tencent.supersonic.common.pojo.ChatApp;
+import com.tencent.supersonic.common.pojo.enums.AppModule;
+import com.tencent.supersonic.common.util.ChatAppManager;
+import com.tencent.supersonic.headless.api.pojo.SemanticParseInfo;
+import com.tencent.supersonic.headless.chat.ChatQueryContext;
+import dev.langchain4j.model.chat.ChatLanguageModel;
+import dev.langchain4j.model.input.Prompt;
+import dev.langchain4j.model.input.PromptTemplate;
+import dev.langchain4j.model.output.structured.Description;
+import dev.langchain4j.provider.ModelProvider;
+import dev.langchain4j.service.AiServices;
+import lombok.Data;
+import lombok.ToString;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * 物理SQL修正器 - 使用LLM优化物理SQL性能
+ */
+@Slf4j
+public class LLMPhysicalSqlCorrector extends BaseSemanticCorrector {
+
+    private static final Logger keyPipelineLog = LoggerFactory.getLogger("keyPipeline");
+
+    public static final String APP_KEY = "PHYSICAL_SQL_CORRECTOR";
+    private static final String INSTRUCTION = ""
+            + "#Role: You are a senior database performance optimization expert experienced in SQL tuning."
+            + "\n\n#Task: You will be provided with a user question and the corresponding physical SQL query,"
+            + " please analyze and optimize this SQL to improve query performance." + "\n\n#Rules:"
+            + "\n1. ALWAYS add appropriate index hints if the database supports them."
+            + "\n2. Optimize JOIN order by placing smaller tables first."
+            + "\n3. Add reasonable query limits to prevent large result sets if no LIMIT exists."
+            + "\n4. Optimize WHERE condition order by placing high-selectivity conditions first."
+            + "\n5. Ensure the optimized SQL is syntactically correct and logically equivalent."
+            + "\n6. If the SQL is already well-optimized, return 'positive'."
+            + "\n\n#Question: {{question}}" + "\n\n#OriginalSQL: {{sql}}" + "\n\n#Response:";
+
+    public LLMPhysicalSqlCorrector() {
+        ChatAppManager.register(APP_KEY, ChatApp.builder().prompt(INSTRUCTION).name("物理SQL修正")
+                .appModule(AppModule.CHAT).description("通过大模型对物理SQL做性能优化").enable(false).build());
+    }
+
+    @Data
+    @ToString
+    static class PhysicalSql {
+        @Description("either positive or negative")
+        private String opinion;
+
+        @Description("optimized sql if negative")
+        private String sql;
+    }
+
+    interface PhysicalSqlExtractor {
+        PhysicalSql generatePhysicalSql(String text);
+    }
+
+    @Override
+    public void doCorrect(ChatQueryContext chatQueryContext, SemanticParseInfo semanticParseInfo) {
+        ChatApp chatApp = chatQueryContext.getRequest().getChatAppConfig().get(APP_KEY);
+        if (!chatQueryContext.getRequest().getText2SQLType().enableLLM() || Objects.isNull(chatApp)
+                || !chatApp.isEnable()) {
+            return;
+        }
+
+        ChatLanguageModel chatLanguageModel =
+                ModelProvider.getChatModel(chatApp.getChatModelConfig());
+        PhysicalSqlExtractor extractor =
+                AiServices.create(PhysicalSqlExtractor.class, chatLanguageModel);
+        Prompt prompt = generatePrompt(chatQueryContext.getRequest().getQueryText(),
+                semanticParseInfo, chatApp.getPrompt());
+        PhysicalSql physicalSql =
+                extractor.generatePhysicalSql(prompt.toUserMessage().singleText());
+        keyPipelineLog.info("LLMPhysicalSqlCorrector modelReq:\n{} \nmodelResp:\n{}", prompt.text(),
+                physicalSql);
+        if ("NEGATIVE".equalsIgnoreCase(physicalSql.getOpinion())
+                && StringUtils.isNotBlank(physicalSql.getSql())) {
+            semanticParseInfo.getSqlInfo().setCorrectedQuerySQL(physicalSql.getSql());
+        }
+    }
+
+    private Prompt generatePrompt(String queryText, SemanticParseInfo semanticParseInfo,
+            String promptTemplate) {
+        Map<String, Object> variable = new HashMap<>();
+        variable.put("question", queryText);
+        variable.put("sql", semanticParseInfo.getSqlInfo().getQuerySQL());
+
+        return PromptTemplate.from(promptTemplate).apply(variable);
+    }
+}
diff --git a/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/parser/llm/PromptHelper.java b/headless/chat/src/main/java/com/tencent/supersonic/headless/chat/parser/llm/PromptHelper.java
@@ -52,7 +52,8 @@ public List<List<Text2SQLExemplar>> getFewShotExemplars(LLMReq llmReq) {
         for (int i = 0; i < selfConsistencyNumber; i++) {
             List<Text2SQLExemplar> shuffledList = new ArrayList<>(exemplars);
             // only shuffle the exemplars from config
-            List<Text2SQLExemplar> subList=shuffledList.subList(llmReq.getDynamicExemplars().size(),shuffledList.size());
+            List<Text2SQLExemplar> subList =
+                    shuffledList.subList(llmReq.getDynamicExemplars().size(), shuffledList.size());
             Collections.shuffle(subList);
             results.add(shuffledList.subList(0, Math.min(shuffledList.size(), fewShotNumber)));
         }
diff --git a/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/ChatWorkflowEngine.java b/headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/ChatWorkflowEngine.java
@@ -8,6 +8,7 @@
 import com.tencent.supersonic.headless.api.pojo.response.ParseResp;
 import com.tencent.supersonic.headless.api.pojo.response.SemanticTranslateResp;
 import com.tencent.supersonic.headless.chat.ChatQueryContext;
+import com.tencent.supersonic.headless.chat.corrector.LLMPhysicalSqlCorrector;
 import com.tencent.supersonic.headless.chat.corrector.SemanticCorrector;
 import com.tencent.supersonic.headless.chat.mapper.SchemaMapper;
 import com.tencent.supersonic.headless.chat.parser.SemanticParser;
@@ -76,6 +77,10 @@ public void start(ChatWorkflowState initialState, ChatQueryContext queryCtx) {
                     long start = System.currentTimeMillis();
                     performTranslating(queryCtx, parseResult);
                     parseResult.getParseTimeCost().setSqlTime(System.currentTimeMillis() - start);
+                    queryCtx.setChatWorkflowState(ChatWorkflowState.PHYSICAL_SQL_CORRECTING);
+                    break;
+                case PHYSICAL_SQL_CORRECTING:
+                    performPhysicalSqlCorrecting(queryCtx);
                     queryCtx.setChatWorkflowState(ChatWorkflowState.FINISHED);
                     break;
                 default:
@@ -162,4 +167,25 @@ private void performTranslating(ChatQueryContext queryCtx, ParseResp parseResult
             parseResult.setErrorMsg(String.join("\n", errorMsg));
         }
     }
+
+    private void performPhysicalSqlCorrecting(ChatQueryContext queryCtx) {
+        List<SemanticQuery> candidateQueries = queryCtx.getCandidateQueries();
+        if (CollectionUtils.isNotEmpty(candidateQueries)) {
+            for (SemanticQuery semanticQuery : candidateQueries) {
+                for (SemanticCorrector corrector : semanticCorrectors) {
+                    if (corrector instanceof LLMPhysicalSqlCorrector) {
+                        corrector.correct(queryCtx, semanticQuery.getParseInfo());
+                        // 如果物理SQL被修正了，更新querySQL为修正后的版本
+                        SemanticParseInfo parseInfo = semanticQuery.getParseInfo();
+                        if (StringUtils.isNotBlank(parseInfo.getSqlInfo().getCorrectedQuerySQL())) {
+                            parseInfo.getSqlInfo().setQuerySQL(parseInfo.getSqlInfo().getCorrectedQuerySQL());
+                            log.info("Physical SQL corrected and updated querySQL: {}", 
+                                    parseInfo.getSqlInfo().getQuerySQL());
+                        }
+                        break;
+                    }
+                }
+            }
+        }
+    }
 }
diff --git a/launchers/headless/src/main/resources/META-INF/spring.factories b/launchers/headless/src/main/resources/META-INF/spring.factories
@@ -14,7 +14,8 @@ com.tencent.supersonic.headless.chat.parser.SemanticParser=\
 
 com.tencent.supersonic.headless.chat.corrector.SemanticCorrector=\
     com.tencent.supersonic.headless.chat.corrector.RuleSqlCorrector,\
-    com.tencent.supersonic.headless.chat.corrector.LLMSqlCorrector
+    com.tencent.supersonic.headless.chat.corrector.LLMSqlCorrector,\
+    com.tencent.supersonic.headless.chat.corrector.LLMPhysicalSqlCorrector
 
 com.tencent.supersonic.headless.chat.knowledge.file.FileHandler=\
     com.tencent.supersonic.headless.chat.knowledge.file.FileHandlerImpl
diff --git a/launchers/standalone/src/main/resources/META-INF/spring.factories b/launchers/standalone/src/main/resources/META-INF/spring.factories
@@ -15,7 +15,8 @@ com.tencent.supersonic.headless.chat.parser.SemanticParser=\
 
 com.tencent.supersonic.headless.chat.corrector.SemanticCorrector=\
     com.tencent.supersonic.headless.chat.corrector.RuleSqlCorrector,\
-    com.tencent.supersonic.headless.chat.corrector.LLMSqlCorrector
+    com.tencent.supersonic.headless.chat.corrector.LLMSqlCorrector,\
+    com.tencent.supersonic.headless.chat.corrector.LLMPhysicalSqlCorrector
 
 com.tencent.supersonic.headless.chat.knowledge.file.FileHandler=\
     com.tencent.supersonic.headless.chat.knowledge.file.FileHandlerImpl
diff --git a/pom.xml b/pom.xml
@@ -10,9 +10,9 @@
     <version>${revision}</version>
 
     <modules>
+        <module>common</module>
         <module>auth</module>
         <module>chat</module>
-        <module>common</module>
         <module>launchers</module>
         <module>headless</module>
     </modules>
@@ -31,6 +31,7 @@
         <java.target.version>21</java.target.version>
         <maven.compiler.source>21</maven.compiler.source>
         <maven.compiler.target>21</maven.compiler.target>
+        <maven.compiler.release>21</maven.compiler.release>
         <file.encoding>UTF-8</file.encoding>
         <jsqlparser.version>4.9</jsqlparser.version>
         <pagehelper.version>6.1.0</pagehelper.version>
@@ -254,6 +255,13 @@
                     <source>${java.source.version}</source>
                     <target>${java.target.version}</target>
                     <encoding>${file.encoding}</encoding>
+                    <annotationProcessorPaths>
+                        <path>
+                            <groupId>org.projectlombok</groupId>
+                            <artifactId>lombok</artifactId>
+                            <version>${lombok.version}</version>
+                        </path>
+                    </annotationProcessorPaths>
                 </configuration>
             </plugin>
             <plugin>
diff --git a/webapp/package.json b/webapp/package.json
@@ -0,0 +1,5 @@
+{
+	"dependencies": {
+		"supersonic-chat-sdk": "link:packages/chat-sdk"
+	}
+}
diff --git a/webapp/packages/chat-sdk/src/common/type.ts b/webapp/packages/chat-sdk/src/common/type.ts
@@ -79,6 +79,7 @@ export type SqlInfoType = {
   parsedS2SQL: string;
   correctedS2SQL: string;
   querySQL: string;
+  correctedQuerySQL?: string;
 };
 
 export type ChatContextType = {
diff --git a/webapp/packages/chat-sdk/src/components/ChatItem/SqlItem.tsx b/webapp/packages/chat-sdk/src/components/ChatItem/SqlItem.tsx

Original file line number	Diff line number	Diff line change
`@@ -16,4 +16,7 @@ public class SqlInfo implements Serializable {`
`16`	`16`
`17`	`17`	`// SQL to be executed finally`
`18`	`18`	`private String querySQL;`
	`19`	`+`
	`20`	`+ // Physical SQL corrected by LLM for performance optimization`
	`21`	`+ private String correctedQuerySQL;`
`19`	`22`	`}`
Original file line number	Diff line number	Diff line change
`@@ -8,5 +8,6 @@ public enum ChatWorkflowState {`
`8`	`8`	`VALIDATING,`
`9`	`9`	`SQL_CORRECTING,`
`10`	`10`	`PROCESSING,`
	`11`	`+ PHYSICAL_SQL_CORRECTING,`
`11`	`12`	`FINISHED`
`12`	`13`	`}`