Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -131,12 +131,12 @@ public StateGraph nl2sqlGraph(NodeBeanUtil nodeBeanUtil, CodeExecutorProperties
// QUERY_ENHANCE_NODE节点输出
keyStrategyHashMap.put(QUERY_ENHANCE_NODE_OUTPUT, KeyStrategy.REPLACE);
// Semantic model
keyStrategyHashMap.put(GENEGRATED_SEMANTIC_MODEL_PROMPT, KeyStrategy.REPLACE);
keyStrategyHashMap.put(GENERATED_SEMANTIC_MODEL_PROMPT, KeyStrategy.REPLACE);
// EVIDENCE节点输出
keyStrategyHashMap.put(EVIDENCE, KeyStrategy.REPLACE);
// schema recall节点输出
keyStrategyHashMap.put(TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, KeyStrategy.REPLACE);
keyStrategyHashMap.put(COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT, KeyStrategy.REPLACE);
keyStrategyHashMap.put(COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT, KeyStrategy.REPLACE);
// table relation节点输出
keyStrategyHashMap.put(TABLE_RELATION_OUTPUT, KeyStrategy.REPLACE);
keyStrategyHashMap.put(TABLE_RELATION_EXCEPTION_OUTPUT, KeyStrategy.REPLACE);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,15 +50,15 @@ private Constant() {

public static final String SCHEMA_RECALL_NODE_OUTPUT = "SCHEMA_RECALL_NODE_OUTPUT";

public static final String COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT = "COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT";
public static final String COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT = "COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT";

public static final String TABLE_RELATION_OUTPUT = "TABLE_RELATION_OUTPUT";

public static final String TABLE_RELATION_EXCEPTION_OUTPUT = "TABLE_RELATION_EXCEPTION_OUTPUT";

public static final String TABLE_RELATION_RETRY_COUNT = "TABLE_RELATION_RETRY_COUNT";

public static final String GENEGRATED_SEMANTIC_MODEL_PROMPT = "GENEGRATED_SEMANTIC_MODEL_PROMPT";
public static final String GENERATED_SEMANTIC_MODEL_PROMPT = "GENERATED_SEMANTIC_MODEL_PROMPT";

public static final String SQL_GENERATE_OUTPUT = "SQL_GENERATE_OUTPUT";

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ private Flux<ChatResponse> handlePlanGenerate(OverAllState state) {
}

// 构建提示参数
String semanticModel = (String) state.value(GENEGRATED_SEMANTIC_MODEL_PROMPT).orElse("");
String semanticModel = (String) state.value(GENERATED_SEMANTIC_MODEL_PROMPT).orElse("");
SchemaDTO schemaDTO = StateUtil.getObjectValue(state, TABLE_RELATION_OUTPUT, SchemaDTO.class);
String schemaStr = PromptHelper.buildMixMacSqlDbPrompt(schemaDTO, true);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ public Map<String, Object> apply(OverAllState state) throws Exception {
Flux<GraphResponse<StreamingOutput>> generator = FluxUtil
.createStreamingGeneratorWithMessages(this.getClass(), state, currentState -> {
return Map.of(TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, Collections.emptyList(),
COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT, Collections.emptyList());
COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT, Collections.emptyList());
}, displayFlux);

return Map.of(SCHEMA_RECALL_NODE_OUTPUT, generator);
Expand Down Expand Up @@ -128,8 +128,8 @@ public Map<String, Object> apply(OverAllState state) throws Exception {

Flux<GraphResponse<StreamingOutput>> generator = FluxUtil.createStreamingGeneratorWithMessages(this.getClass(),
state, currentState -> {
return Map.of(TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocuments,
COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT, columnDocuments);
return Map.of(TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocuments, COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT,
columnDocuments);
}, displayFlux);

// Return the processing result
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,10 @@
package com.alibaba.cloud.ai.dataagent.workflow.node;

import static com.alibaba.cloud.ai.dataagent.constant.Constant.AGENT_ID;
import static com.alibaba.cloud.ai.dataagent.constant.Constant.COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT;
import static com.alibaba.cloud.ai.dataagent.constant.Constant.COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT;
import static com.alibaba.cloud.ai.dataagent.constant.Constant.DB_DIALECT_TYPE;
import static com.alibaba.cloud.ai.dataagent.constant.Constant.EVIDENCE;
import static com.alibaba.cloud.ai.dataagent.constant.Constant.GENEGRATED_SEMANTIC_MODEL_PROMPT;
import static com.alibaba.cloud.ai.dataagent.constant.Constant.GENERATED_SEMANTIC_MODEL_PROMPT;
import static com.alibaba.cloud.ai.dataagent.constant.Constant.SQL_GENERATE_SCHEMA_MISSING_ADVICE;
import static com.alibaba.cloud.ai.dataagent.constant.Constant.TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT;
import static com.alibaba.cloud.ai.dataagent.constant.Constant.TABLE_RELATION_EXCEPTION_OUTPUT;
Expand Down Expand Up @@ -98,7 +98,7 @@ public Map<String, Object> apply(OverAllState state) throws Exception {

String evidence = StateUtil.getStringValue(state, EVIDENCE);
List<Document> tableDocuments = StateUtil.getDocumentList(state, TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT);
List<Document> columnDocuments = StateUtil.getDocumentList(state, COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT);
List<Document> columnDocuments = StateUtil.getDocumentList(state, COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT);
String agentIdStr = StateUtil.getStringValue(state, AGENT_ID);

// Execute business logic first - get final result immediately
Expand Down Expand Up @@ -130,7 +130,7 @@ public Map<String, Object> apply(OverAllState state) throws Exception {

// 构建语义模型提示并存储到resultMap中
String semanticModelPrompt = buildSemanticModelPrompt(semanticModels);
resultMap.put(GENEGRATED_SEMANTIC_MODEL_PROMPT, semanticModelPrompt);
resultMap.put(GENERATED_SEMANTIC_MODEL_PROMPT, semanticModelPrompt);
});

// Create display stream for user experience only
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ private OverAllState createTestState() {
state.registerKeyAndStrategy(AGENT_ID, new ReplaceStrategy());
state.registerKeyAndStrategy(SCHEMA_RECALL_NODE_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT, new ReplaceStrategy());
return state;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -84,14 +84,14 @@ private OverAllState createTestState() {
state.registerKeyAndStrategy(QUERY_ENHANCE_NODE_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(EVIDENCE, new ReplaceStrategy());
state.registerKeyAndStrategy(TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(AGENT_ID, new ReplaceStrategy());
state.registerKeyAndStrategy(TABLE_RELATION_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(DB_DIALECT_TYPE, new ReplaceStrategy());
state.registerKeyAndStrategy(TABLE_RELATION_RETRY_COUNT, new ReplaceStrategy());
state.registerKeyAndStrategy(TABLE_RELATION_EXCEPTION_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(SQL_GENERATE_SCHEMA_MISSING_ADVICE, new ReplaceStrategy());
state.registerKeyAndStrategy(GENEGRATED_SEMANTIC_MODEL_PROMPT, new ReplaceStrategy());
state.registerKeyAndStrategy(GENERATED_SEMANTIC_MODEL_PROMPT, new ReplaceStrategy());
return state;
}

Expand Down Expand Up @@ -132,7 +132,7 @@ void apply_validSchema_returnsTableRelationOutput() throws Exception {
List<Document> colDocs = List.of(new Document("column doc"));

state.updateState(Map.of(QUERY_ENHANCE_NODE_OUTPUT, dto, EVIDENCE, "evidence", AGENT_ID, "1",
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT, colDocs));
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT, colDocs));

setupCommonMocks("1", tableDocs);

Expand All @@ -154,7 +154,7 @@ void apply_emptyTableDocuments_returnsResultWithDefaults() throws Exception {
List<Document> emptyColDocs = Collections.emptyList();

state.updateState(Map.of(QUERY_ENHANCE_NODE_OUTPUT, dto, EVIDENCE, "evidence", AGENT_ID, "2",
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, emptyTableDocs, COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT, emptyColDocs));
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, emptyTableDocs, COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT, emptyColDocs));

setupCommonMocks("2", emptyTableDocs);

Expand All @@ -172,7 +172,7 @@ void apply_emptyColumnDocuments_returnsMinimalSchema() throws Exception {
List<Document> tableDocs = List.of(createTableDocument("users"));

state.updateState(Map.of(QUERY_ENHANCE_NODE_OUTPUT, dto, EVIDENCE, "evidence", AGENT_ID, "3",
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT,
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT,
Collections.emptyList()));

setupCommonMocks("3", tableDocs);
Expand All @@ -190,7 +190,7 @@ void apply_databaseConfigMissing_throwsException() {
List<Document> tableDocs = List.of(createTableDocument("users"));

state.updateState(Map.of(QUERY_ENHANCE_NODE_OUTPUT, dto, EVIDENCE, "evidence", AGENT_ID, "4",
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT,
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT,
Collections.emptyList()));

when(databaseUtil.getAgentDbConfig(4L)).thenThrow(new RuntimeException("DB config not found"));
Expand All @@ -205,7 +205,7 @@ void apply_noActiveDatasource_returnsEmptyForeignKeys() throws Exception {
List<Document> tableDocs = List.of(createTableDocument("users"));

state.updateState(Map.of(QUERY_ENHANCE_NODE_OUTPUT, dto, EVIDENCE, "evidence", AGENT_ID, "5",
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT,
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT,
Collections.emptyList()));

DbConfigBO dbConfig = DbConfigBO.builder().dialectType("mysql").schema("test_db").build();
Expand Down Expand Up @@ -239,7 +239,7 @@ void apply_withSchemaAdvice_passesAdviceToFineSelect() throws Exception {
List<Document> tableDocs = List.of(createTableDocument("users"));

state.updateState(Map.of(QUERY_ENHANCE_NODE_OUTPUT, dto, EVIDENCE, "evidence", AGENT_ID, "6",
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT,
TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocs, COLUMN_DOCUMENTS_FOR_SCHEMA_OUTPUT,
Collections.emptyList(), SQL_GENERATE_SCHEMA_MISSING_ADVICE, "add orders table"));

DbConfigBO dbConfig = DbConfigBO.builder().dialectType("mysql").schema("test_db").build();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -124,15 +124,15 @@ private OverAllState createTestState() {
state.registerKeyAndStrategy(QUERY_ENHANCE_NODE_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(TABLE_RELATION_OUTPUT, new ReplaceStrategy());
state.registerKeyAndStrategy(EVIDENCE, new ReplaceStrategy());
state.registerKeyAndStrategy(GENEGRATED_SEMANTIC_MODEL_PROMPT, new ReplaceStrategy());
state.registerKeyAndStrategy(GENERATED_SEMANTIC_MODEL_PROMPT, new ReplaceStrategy());
state.registerKeyAndStrategy(IS_ONLY_NL2SQL, new ReplaceStrategy());
state.registerKeyAndStrategy(PLAN_VALIDATION_ERROR, new ReplaceStrategy());
return state;
}

private void setupBasicState(OverAllState state) {
state.updateState(Map.of(QUERY_ENHANCE_NODE_OUTPUT, TEST_QUERY_ENHANCE, TABLE_RELATION_OUTPUT, TEST_SCHEMA,
EVIDENCE, "业务知识:用户表包含所有注册用户", GENEGRATED_SEMANTIC_MODEL_PROMPT, ""));
EVIDENCE, "业务知识:用户表包含所有注册用户", GENERATED_SEMANTIC_MODEL_PROMPT, ""));
}

@Test
Expand Down Expand Up @@ -218,7 +218,7 @@ void apply_withValidationError_regeneratesPlan() throws Exception {
void apply_withSemanticModel_includesInPrompt() throws Exception {
OverAllState state = createTestState();
setupBasicState(state);
state.updateState(Map.of(GENEGRATED_SEMANTIC_MODEL_PROMPT, "语义模型定义:PV表示页面浏览量"));
state.updateState(Map.of(GENERATED_SEMANTIC_MODEL_PROMPT, "语义模型定义:PV表示页面浏览量"));

when(llmService.callUser(anyString()))
.thenReturn(Flux.just(ChatResponseUtil.createPureResponse(VALID_PLAN_JSON)));
Expand All @@ -233,7 +233,7 @@ void apply_withSemanticModel_includesInPrompt() throws Exception {
void apply_emptyEvidence_generatesWithoutEvidence() throws Exception {
OverAllState state = createTestState();
state.updateState(Map.of(QUERY_ENHANCE_NODE_OUTPUT, TEST_QUERY_ENHANCE, TABLE_RELATION_OUTPUT, TEST_SCHEMA,
EVIDENCE, "", GENEGRATED_SEMANTIC_MODEL_PROMPT, ""));
EVIDENCE, "", GENERATED_SEMANTIC_MODEL_PROMPT, ""));

when(llmService.callUser(anyString()))
.thenReturn(Flux.just(ChatResponseUtil.createPureResponse(VALID_PLAN_JSON)));
Expand Down
Loading