[improvement][chat]Optimize robustness in semantic parsing.

chat/server/src/main/java/com/tencent/supersonic/chat/server/parser/NL2SQLParser.java

@@ -100,10 +100,17 @@ public void parse(ParseContext parseContext) {
                     queryNLReq.setMapModeEnum(mode);
                     doParse(queryNLReq, parseResp);
                 }
+
                 if (parseResp.getSelectedParses().isEmpty()) {
-                    queryNLReq.setMapModeEnum(MapModeEnum.LOOSE);
-                    doParse(queryNLReq, parseResp);
+                    for (MapModeEnum mode : Lists.newArrayList(MapModeEnum.LOOSE)) {
+                        queryNLReq.setMapModeEnum(mode);
+                        doParse(queryNLReq, parseResp);
+                        if (!parseResp.getSelectedParses().isEmpty()) {
+                            break;
+                        }
+                    }
                 }
+
                 if (parseResp.getSelectedParses().isEmpty()) {
                     errMsg.append(parseResp.getErrorMsg());
                     continue;
@@ -137,11 +144,18 @@ public void parse(ParseContext parseContext) {
             SemanticParseInfo userSelectParse = parseContext.getRequest().getSelectedParse();
             queryNLReq.setSelectedParseInfo(Objects.nonNull(userSelectParse) ? userSelectParse
                     : parseContext.getResponse().getSelectedParses().get(0));
-
             parseContext.setResponse(new ChatParseResp(parseContext.getResponse().getQueryId()));
+
             rewriteMultiTurn(parseContext, queryNLReq);
             addDynamicExemplars(parseContext, queryNLReq);
             doParse(queryNLReq, parseContext.getResponse());
+
+            // try again with all semantic fields passed to LLM
+            if (parseContext.getResponse().getState().equals(ParseResp.ParseState.FAILED)) {
+                queryNLReq.setSelectedParseInfo(null);
+                queryNLReq.setMapModeEnum(MapModeEnum.ALL);
+                doParse(queryNLReq, parseContext.getResponse());
+            }
         }
     }
 

common/src/main/java/com/tencent/supersonic/common/jsqlparser/AliasAcquireVisitor.java

@@ -0,0 +1,24 @@
+package com.tencent.supersonic.common.jsqlparser;
+
+import net.sf.jsqlparser.expression.Alias;
+import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
+import net.sf.jsqlparser.statement.select.SelectItem;
+
+import java.util.Set;
+
+public class AliasAcquireVisitor extends ExpressionVisitorAdapter {
+
+    private Set<String> aliases;
+
+    public AliasAcquireVisitor(Set<String> aliases) {
+        this.aliases = aliases;
+    }
+
+    @Override
+    public void visit(SelectItem selectItem) {
+        Alias alias = selectItem.getAlias();
+        if (alias != null) {
+            aliases.add(alias.getName());
+        }
+    }
+}

common/src/main/java/com/tencent/supersonic/common/jsqlparser/FieldAcquireVisitor.java

@@ -1,5 +1,6 @@
 package com.tencent.supersonic.common.jsqlparser;
 
+import com.google.common.collect.Sets;
 import net.sf.jsqlparser.expression.Alias;
 import net.sf.jsqlparser.expression.Expression;
 import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
@@ -11,6 +12,7 @@
 public class FieldAcquireVisitor extends ExpressionVisitorAdapter {
 
     private Set<String> fields;
+    private Set<String> aliases = Sets.newHashSet();
 
     public FieldAcquireVisitor(Set<String> fields) {
         this.fields = fields;
@@ -26,8 +28,9 @@ public void visit(Column column) {
     public void visit(SelectItem selectItem) {
         Alias alias = selectItem.getAlias();
         if (alias != null) {
-            fields.add(alias.getName());
+            aliases.add(alias.getName());
         }
+
         Expression expression = selectItem.getExpression();
         if (expression != null) {
             expression.accept(this);

common/src/main/java/com/tencent/supersonic/common/jsqlparser/SqlSelectHelper.java

@@ -133,6 +133,15 @@ public static Set<String> getSelectFields(List<PlainSelect> plainSelectList) {
         return result;
     }
 
+    public static Set<String> getAliasFields(PlainSelect plainSelect) {
+        Set<String> result = new HashSet<>();
+        List<SelectItem<?>> selectItems = plainSelect.getSelectItems();
+        for (SelectItem selectItem : selectItems) {
+            selectItem.accept(new AliasAcquireVisitor(result));
+        }
+        return result;
+    }
+
     public static List<PlainSelect> getPlainSelect(Select selectStatement) {
         if (selectStatement == null) {
             return null;
@@ -264,10 +273,14 @@ public void visit(PlainSelect plainSelect) {
     public static List<String> getAllSelectFields(String sql) {
         List<PlainSelect> plainSelects = getPlainSelects(getPlainSelect(sql));
         Set<String> results = new HashSet<>();
+        Set<String> aliases = new HashSet<>();
         for (PlainSelect plainSelect : plainSelects) {
             List<String> fields = getFieldsByPlainSelect(plainSelect);
             results.addAll(fields);
+            aliases.addAll(getAliasFields(plainSelect));
         }
+        // do not account in aliases
+        results.removeAll(aliases);
         return new ArrayList<>(results);
     }
 

headless/api/src/main/java/com/tencent/supersonic/headless/api/pojo/enums/MapModeEnum.java

@@ -1,7 +1,7 @@
 package com.tencent.supersonic.headless.api.pojo.enums;
 
 public enum MapModeEnum {
-    STRICT(0), MODERATE(2), LOOSE(4);
+    STRICT(0), MODERATE(2), LOOSE(4), ALL(6);
 
     public int threshold;
 

headless/chat/src/main/java/com/tencent/supersonic/headless/chat/mapper/AllFieldMapper.java

@@ -0,0 +1,40 @@
+package com.tencent.supersonic.headless.chat.mapper;
+
+import com.google.common.collect.Lists;
+import com.tencent.supersonic.headless.api.pojo.DataSetSchema;
+import com.tencent.supersonic.headless.api.pojo.SchemaElement;
+import com.tencent.supersonic.headless.api.pojo.SchemaElementMatch;
+import com.tencent.supersonic.headless.api.pojo.enums.MapModeEnum;
+import com.tencent.supersonic.headless.chat.ChatQueryContext;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.List;
+import java.util.Map;
+
+@Slf4j
+public class AllFieldMapper extends BaseMapper {
+
+    @Override
+    public void doMap(ChatQueryContext chatQueryContext) {
+        // Check if the map mode is ALL
+        if (!MapModeEnum.ALL.equals(chatQueryContext.getRequest().getMapModeEnum())) {
+            return;
+        }
+
+        Map<Long, DataSetSchema> schemaMap =
+                chatQueryContext.getSemanticSchema().getDataSetSchemaMap();
+        for (Map.Entry<Long, DataSetSchema> entry : schemaMap.entrySet()) {
+            List<SchemaElement> schemaElements = Lists.newArrayList();
+            schemaElements.addAll(entry.getValue().getDimensions());
+            schemaElements.addAll(entry.getValue().getMetrics());
+
+            for (SchemaElement schemaElement : schemaElements) {
+                chatQueryContext.getMapInfo().getMatchedElements(entry.getKey())
+                        .add(SchemaElementMatch.builder().word(schemaElement.getName())
+                                .element(schemaElement).detectWord(schemaElement.getName())
+                                .similarity(1.0).build());
+            }
+        }
+    }
+
+}

headless/chat/src/main/java/com/tencent/supersonic/headless/chat/mapper/TimeFieldMapper.java → headless/chat/src/main/java/com/tencent/supersonic/headless/chat/mapper/PartitionTimeMapper.java

@@ -12,19 +12,20 @@
 import java.util.stream.Collectors;
 
 @Slf4j
-public class TimeFieldMapper extends BaseMapper {
+public class PartitionTimeMapper extends BaseMapper {
 
     @Override
     public void doMap(ChatQueryContext chatQueryContext) {
-        if (chatQueryContext.getRequest().getText2SQLType().equals(Text2SQLType.ONLY_RULE)) {
+        if (chatQueryContext.getRequest().getText2SQLType().equals(Text2SQLType.ONLY_RULE)
+                || chatQueryContext.getMapInfo().isEmpty()) {
             return;
         }
 
         Map<Long, DataSetSchema> schemaMap =
                 chatQueryContext.getSemanticSchema().getDataSetSchemaMap();
         for (Map.Entry<Long, DataSetSchema> entry : schemaMap.entrySet()) {
             List<SchemaElement> timeDims = entry.getValue().getDimensions().stream()
-                    .filter(dim -> dim.getTimeFormat() != null).collect(Collectors.toList());
+                    .filter(SchemaElement::isPartitionTime).collect(Collectors.toList());
             for (SchemaElement schemaElement : timeDims) {
                 chatQueryContext.getMapInfo().getMatchedElements(entry.getKey())
                         .add(SchemaElementMatch.builder().word(schemaElement.getName())

headless/core/src/main/java/com/tencent/supersonic/headless/core/translator/DefaultSemanticTranslator.java

@@ -30,6 +30,9 @@ public void translate(QueryStatement queryStatement) {
                 if (parser.accept(queryStatement)) {
                     log.debug("QueryConverter accept [{}]", parser.getClass().getName());
                     parser.parse(queryStatement);
+                    if (queryStatement.getStatus() != 0) {
+                        break;
+                    }
                 }
             }
             if (!queryStatement.isOk()) {

headless/core/src/main/java/com/tencent/supersonic/headless/core/translator/parser/SqlQueryParser.java

@@ -63,6 +63,14 @@ public void parse(QueryStatement queryStatement) throws Exception {
         List<String> metrics =
                 metricSchemas.stream().map(SchemaItem::getBizName).collect(Collectors.toList());
         Set<String> dimensions = getDimensions(semanticSchemaResp, allFields);
+        // check if there are fields not matched with any metric or dimension
+        if (allFields.size() > metricSchemas.size() + dimensions.size()) {
+            queryStatement
+                    .setErrMsg("There are fields in the SQL not matched with any semantic column.");
+            queryStatement.setStatus(1);
+            return;
+        }
+
         OntologyQuery ontologyQuery = new OntologyQuery();
         ontologyQuery.getMetrics().addAll(metrics);
         ontologyQuery.getDimensions().addAll(dimensions);

headless/server/src/main/java/com/tencent/supersonic/headless/server/utils/ChatWorkflowEngine.java

@@ -76,7 +76,6 @@ public void start(ChatWorkflowState initialState, ChatQueryContext queryCtx,
                     long start = System.currentTimeMillis();
                     performTranslating(queryCtx, parseResult);
                     parseResult.getParseTimeCost().setSqlTime(System.currentTimeMillis() - start);
-                    parseResult.setState(ParseResp.ParseState.COMPLETED);
                     queryCtx.setChatWorkflowState(ChatWorkflowState.FINISHED);
                     break;
                 default:
@@ -137,7 +136,12 @@ private void performTranslating(ChatQueryContext queryCtx, ParseResp parseResult
                         ContextUtils.getBean(SemanticLayerService.class);
                 SemanticTranslateResp explain =
                         queryService.translate(semanticQueryReq, queryCtx.getRequest().getUser());
-                parseInfo.getSqlInfo().setQuerySQL(explain.getQuerySQL());
+                if (explain.isOk()) {
+                    parseInfo.getSqlInfo().setQuerySQL(explain.getQuerySQL());
+                    parseResult.setState(ParseResp.ParseState.COMPLETED);
+                } else {
+                    parseResult.setState(ParseResp.ParseState.FAILED);
+                }
                 if (StringUtils.isNotBlank(explain.getErrMsg())) {
                     errorMsg.add(explain.getErrMsg());
                 }

launchers/headless/src/main/resources/META-INF/spring.factories

@@ -4,7 +4,7 @@ com.tencent.supersonic.headless.chat.mapper.SchemaMapper=\
     com.tencent.supersonic.headless.chat.mapper.EmbeddingMapper, \
     com.tencent.supersonic.headless.chat.mapper.KeywordMapper, \
     com.tencent.supersonic.headless.chat.mapper.QueryFilterMapper, \
-    com.tencent.supersonic.headless.chat.mapper.TimeFieldMapper,\
+    com.tencent.supersonic.headless.chat.mapper.PartitionTimeMapper,\
     com.tencent.supersonic.headless.chat.mapper.TermDescMapper
 
 com.tencent.supersonic.headless.chat.parser.SemanticParser=\

launchers/standalone/src/main/resources/META-INF/spring.factories

@@ -4,8 +4,9 @@ com.tencent.supersonic.headless.chat.mapper.SchemaMapper=\
     com.tencent.supersonic.headless.chat.mapper.EmbeddingMapper, \
     com.tencent.supersonic.headless.chat.mapper.KeywordMapper, \
     com.tencent.supersonic.headless.chat.mapper.QueryFilterMapper, \
-    com.tencent.supersonic.headless.chat.mapper.TimeFieldMapper,\
-    com.tencent.supersonic.headless.chat.mapper.TermDescMapper
+    com.tencent.supersonic.headless.chat.mapper.PartitionTimeMapper,\
+    com.tencent.supersonic.headless.chat.mapper.TermDescMapper,\
+    com.tencent.supersonic.headless.chat.mapper.AllFieldMapper
 
 com.tencent.supersonic.headless.chat.parser.SemanticParser=\
     com.tencent.supersonic.headless.chat.parser.llm.LLMSqlParser,\

launchers/standalone/src/test/java/com/tencent/supersonic/evaluation/Text2SQLEval.java

@@ -24,6 +24,7 @@
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
+import org.junitpioneer.jupiter.SetSystemProperty;
 import org.springframework.test.context.TestPropertySource;
 
 import java.util.List;
@@ -95,6 +96,7 @@ public void test_drilldown() throws Exception {
     }
 
     @Test
+    @SetSystemProperty(key = "s2.test", value = "true")
     public void test_drilldown_and_topN() throws Exception {
         long start = System.currentTimeMillis();
         QueryResult result = submitNewChat("过去30天访问次数最高的部门top3", agent.getId());

launchers/standalone/src/test/java/com/tencent/supersonic/headless/QueryBySqlTest.java

@@ -7,6 +7,7 @@
 import com.tencent.supersonic.util.DataUtils;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.junitpioneer.jupiter.SetSystemProperty;
 
 import static java.time.LocalDate.now;
 import static org.junit.Assert.assertEquals;
@@ -29,6 +30,7 @@ public void testDetailQuery() throws Exception {
     }
 
     @Test
+    @SetSystemProperty(key = "s2.test", value = "true")
     public void testSumQuery() throws Exception {
         SemanticQueryResp semanticQueryResp =
                 queryBySql("SELECT SUM(访问次数) AS 总访问次数 FROM 超音数PVUV统计 ");