topdog/supersonic
1
0
Fork 0
mirror of https://github.com/tencentmusic/supersonic.git synced 2025-12-13 13:07:32 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

[improvement][headless&chat]当LLM生成SQL包含日期类型字段时，Correcter不再额外增加日期 (#1473)

Browse Source
This commit is contained in:
yudong
2024-08-02 14:14:29 +08:00
committed by GitHub
parent e26263d229
commit 53a9f7c451
4 changed files with 32 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
headless/chat/src/main/java/com/tencent/supersonic/headless/chat/corrector/TimeCorrector.java
View File

@@ -8,6 +8,8 @@ import com.tencent.supersonic.common.jsqlparser.SqlDateSelectHelper;
import com.tencent.supersonic.common.jsqlparser.SqlSelectHelper;
import com.tencent.supersonic.common.jsqlparser.SqlRemoveHelper;
import com.tencent.supersonic.common.jsqlparser.DateVisitor.DateBoundInfo;
import com.tencent.supersonic.headless.api.pojo.DataSetSchema;
import com.tencent.supersonic.headless.api.pojo.SchemaElement;
import com.tencent.supersonic.headless.api.pojo.SemanticParseInfo;
import com.tencent.supersonic.headless.chat.ChatQueryContext;
import lombok.extern.slf4j.Slf4j;
@@ -56,6 +58,15 @@ public class TimeCorrector extends BaseSemanticCorrector {
}
}
private boolean checkIfNameInWhereFields(Set<SchemaElement> dims, List<String> whereFields) {
for (SchemaElement element : dims) {
if (whereFields.contains(element.getName())) {
return true;
}
}
return false;
}
private void addDateIfNotExist(ChatQueryContext chatQueryContext, SemanticParseInfo semanticParseInfo) {
String correctS2SQL = semanticParseInfo.getSqlInfo().getCorrectedS2SQL();
List<String> whereFields = SqlSelectHelper.getWhereFields(correctS2SQL);
@@ -66,6 +77,12 @@ public class TimeCorrector extends BaseSemanticCorrector {
if (StringUtils.isNotBlank(correctorDate) && !Boolean.parseBoolean(correctorDate)) {
return;
}
Long dataSetId = semanticParseInfo.getDataSetId();
DataSetSchema dataSetSchema = chatQueryContext.getSemanticSchema().getDataSetSchemaMap().get(dataSetId);
boolean isDateInWhere = checkIfNameInWhereFields(dataSetSchema.getDimensions(), whereFields);
if (isDateInWhere) {
return;
}
if (CollectionUtils.isEmpty(whereFields) || !TimeDimensionEnum.containsZhTimeDimension(whereFields)) {
Pair<String, String> startEndDate = S2SqlDateHelper.getStartEndDate(chatQueryContext,