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

[improvement][headless]Add databaseType into the Schema part of the Text2SQL prompt. #1621

Browse Source
This commit is contained in:
jerryjzhang
2024-09-12 23:20:49 +08:00
parent 47cc933aec
commit 37f12391b0
9 changed files with 44 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
headless/chat/src/main/java/com/tencent/supersonic/headless/chat/parser/llm/LLMRequestService.java
View File

@@ -61,6 +61,7 @@ public class LLMRequestService {
llmReq.setQueryText(queryText);
LLMReq.LLMSchema llmSchema = new LLMReq.LLMSchema();
llmReq.setSchema(llmSchema);
llmSchema.setDatabaseType(getDatabaseType(queryCtx, dataSetId));
llmSchema.setDataSetId(dataSetId);
llmSchema.setDataSetName(dataSetIdToName.get(dataSetId));
llmSchema.setMetrics(getMappedMetrics(queryCtx, dataSetId));
@@ -205,4 +206,14 @@ public class LLMRequestService {
DataSetSchema dataSetSchema = dataSetSchemaMap.get(dataSetId);
return dataSetSchema.getPrimaryKey();
}
protected String getDatabaseType(@NotNull ChatQueryContext queryCtx, Long dataSetId) {
SemanticSchema semanticSchema = queryCtx.getSemanticSchema();
if (semanticSchema == null || semanticSchema.getDataSetSchemaMap() == null) {
return null;
}
Map<Long, DataSetSchema> dataSetSchemaMap = semanticSchema.getDataSetSchemaMap();
DataSetSchema dataSetSchema = dataSetSchemaMap.get(dataSetId);
return dataSetSchema.getDatabaseType();
}
}

8
headless/chat/src/main/java/com/tencent/supersonic/headless/chat/parser/llm/PromptHelper.java
View File

@@ -162,11 +162,17 @@ public class PromptHelper {
primaryKeyStr = String.format("%s", llmReq.getSchema().getPrimaryKey().getName());
}
String databaseTypeStr = "";
if (llmReq.getSchema().getDatabaseType() != null) {
databaseTypeStr = llmReq.getSchema().getDatabaseType();
}
String template =
"Table=[%s], PartitionTimeField=[%s], PrimaryKeyField=[%s], "
"DatabaseType=[%s], Table=[%s], PartitionTimeField=[%s], PrimaryKeyField=[%s], "
+ "Metrics=[%s], Dimensions=[%s], Values=[%s]";
return String.format(
template,
databaseTypeStr,
tableStr,
partitionTimeStr,
primaryKeyStr,

1
headless/chat/src/main/java/com/tencent/supersonic/headless/chat/query/llm/s2sql/LLMReq.java
View File

@@ -33,6 +33,7 @@ public class LLMReq {
@Data
public static class LLMSchema {
private String databaseType;
private Long dataSetId;
private String dataSetName;
private List<SchemaElement> metrics;