diff --git a/README.md b/README.md index 33bc02fe9..e566146e4 100644 --- a/README.md +++ b/README.md @@ -117,7 +117,13 @@ spring-ai-alibaba-nl2sql/ 将表和数据导入到你的MySQL数据库中。 -### 2. 配置management数据库 + + +### 2. 配置 + +#### 2.1 配置management数据库 + + 在`spring-ai-alibaba-nl2sql-management/src/main/resources/application.yml`中配置你的MySQL数据库连接信息。 @@ -133,7 +139,9 @@ spring: type: com.alibaba.druid.pool.DruidDataSource ``` -### 2.1 可选：启用/关闭自动初始化（schema.sql + data.sql） +#### 2.2 可选：启用/关闭自动初始化（schema.sql + data.sql） + + - 默认配置：`application.yml` 中已设置为开启 @@ -162,7 +170,7 @@ spring: ![Snipaste_2025-09-18_14-35-29.jpg](img/Snipaste_2025-09-18_14-35-29.jpg) -### 3. 配置 API Key +#### 2.3 配置 API Key ```yaml spring: @@ -178,15 +186,52 @@ spring: ``` -### 4. 启动管理端 + +#### 2.4 嵌入模型批处理策略配置 + + + +| 属性 | 说明 | 默认值 | +| ----------------------------------------------------------- | ------------------------------------------------------------ | ----------- | +| spring.ai.alibaba.nl2sql.embedding-batch.encoding-type | 文本编码类型，可参考com.knuddels.jtokkit.api.EncodingType | cl100k_base | +| spring.ai.alibaba.nl2sql.embedding-batch.max-token-count | 每批次最大令牌数 值越小，每批次文档越少，但更安全 值越大，处理效率越高，但可能超出API限制 建议值：2000-8000，根据实际API限制调整 | 2000 | +| spring.ai.alibaba.nl2sql.embedding-batch.reserve-percentage | 预留百分比 用于预留缓冲空间，避免超出限制 建议值：0.1-0.2（10%-20%） | 0.2 | +| spring.ai.alibaba.nl2sql.embedding-batch.max-text-count | 每批次最大文本数量 适用于DashScope等有文本数量限制的API DashScope限制为10 | 10 | + + + +#### 2.5 向量库配置 + + + +| 属性 | 说明 | 默认值 | +| ------------------------------------------------------------ | ------------------------------------------------------------ | --------- | +| spring.ai.alibaba.nl2sql.vector-store.similarity-threshold | 相似度阈值配置，用于过滤相似度分数大于等于此阈值的文档 | 0.2 | +| spring.ai.alibaba.nl2sql.vector-store.batch-del-topk-limit | 一次删除操作中，最多删除的文档数量 | 5000 | +| spring.ai.alibaba.nl2sql.vector-store.topk-limit | 查询返回最大文档数 | 30 | +| **spring.ai.alibaba.nl2sql.vector-store.enable-hybrid-search** | 是否启用混合搜索。**注意**：**项目目前默认只提供ES的混合检索能力，
如需要扩展其他向量库可自行继承重写 com.alibaba.cloud.ai.service.hybrid.retrieval
.AbstractHybridRetrievalStrategy#retrieve 该方法**
**并且修改com.alibaba.cloud.ai.service.hybrid.
factory.HybridRetrievalStrategyFactory#getObject
注册相应的bean** | **false** | +| spring.ai.alibaba.nl2sql.vector-store.elasticsearch-min-score | Elasticsearch最小分数阈值，用于es执行关键词搜索时过滤相关性较低的文档。
**开发时使用的es服务端版本 8.15.0** | 0.5 | + + + +#### 2.6 检索融合策略 + +| 属性 | 说明 | 默认值 | +| ---------------------------------------- | -------------------- | ------ | +| spring.ai.alibaba.nl2sql.fusion-strategy | 多路召回结果融合策略 | rrf | +| | | | + + + +### 3. 启动管理端 在`spring-ai-alibaba-nl2sql-management`目录下，运行 `spring-ai-alibaba-nl2sql/spring-ai-alibaba-nl2sql-management/src/main/java/com/alibaba/cloud/ai/DataAgentApplication.java` 类。 -### 5. 启动WEB页面 +### 4. 启动WEB页面 进入 `spring-ai-alibaba-nl2sql/spring-ai-alibaba-nl2sql-web-ui` 目录 -#### 安装依赖 +#### 4.1 安装依赖 ```bash @@ -197,7 +242,7 @@ npm install yarn install ``` -### 启动服务 +#### 4.2 启动服务 ```bash # 使用 npm diff --git a/spring-ai-alibaba-data-agent-chat/pom.xml b/spring-ai-alibaba-data-agent-chat/pom.xml index 802411eda..ea9620c7f 100644 --- a/spring-ai-alibaba-data-agent-chat/pom.xml +++ b/spring-ai-alibaba-data-agent-chat/pom.xml @@ -39,6 +39,12 @@ spring-boot-starter-web + + + org.springframework.boot + spring-boot-starter-aop + + org.springframework.boot spring-boot-starter-test @@ -199,7 +205,10 @@ elasticsearch-rest-client ${elasticsearch-client.version} - + + jakarta.validation + jakarta.validation-api + diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/aop/NodeEntryLoggingAspect.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/aop/NodeEntryLoggingAspect.java new file mode 100644 index 000000000..7653652fa --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/aop/NodeEntryLoggingAspect.java @@ -0,0 +1,57 @@ +/* + * Copyright 2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.alibaba.cloud.ai.aop; + +import com.alibaba.cloud.ai.graph.OverAllState; +import lombok.extern.slf4j.Slf4j; +import org.aspectj.lang.JoinPoint; +import org.aspectj.lang.annotation.Aspect; +import org.aspectj.lang.annotation.Before; +import org.aspectj.lang.annotation.Pointcut; +import org.springframework.stereotype.Component; + +/** + * AOP切面类，用于记录所有Node类的入口日志 + * + * + */ +@Aspect +@Component +@Slf4j +public class NodeEntryLoggingAspect { + + @Pointcut("execution(* com.alibaba.cloud.ai.node..*.apply(com.alibaba.cloud.ai.graph.OverAllState))") + public void nodeEntry() { + } + + /** + * 在所有实现NodeAction接口的类的apply方法执行前记录日志 + * @param joinPoint 连接点 + */ + @Before("nodeEntry()") + public void logNodeEntry(JoinPoint joinPoint) { + String className = joinPoint.getTarget().getClass().getSimpleName(); + log.info("Entering {} node", className); + + // 获取方法参数并打印状态信息 + Object[] args = joinPoint.getArgs(); + if (args != null && args.length > 0 && args[0] instanceof OverAllState state) { + log.debug("State: {}", state); + } + } + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/config/DataAgentConfiguration.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/config/DataAgentConfiguration.java index 4d5c39481..9af906f83 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/config/DataAgentConfiguration.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/config/DataAgentConfiguration.java @@ -108,22 +108,23 @@ public StateGraph nl2sqlGraph(NodeBeanUtil nodeBeanUtil) throws GraphStateExcept keyStrategyHashMap.put(INPUT_KEY, new ReplaceStrategy()); // Agent ID keyStrategyHashMap.put(AGENT_ID, new ReplaceStrategy()); - // Business knowledge - keyStrategyHashMap.put(BUSINESS_KNOWLEDGE, new ReplaceStrategy()); + // Intent recognition + keyStrategyHashMap.put(INTENT_RECOGNITION_NODE_OUTPUT, new ReplaceStrategy()); + // QUERY_ENHANCE_NODE节点输出 + keyStrategyHashMap.put(QUERY_ENHANCE_NODE_OUTPUT, new ReplaceStrategy()); // Semantic model - keyStrategyHashMap.put(SEMANTIC_MODEL, new ReplaceStrategy()); - // queryWrite节点输出 - keyStrategyHashMap.put(QUERY_REWRITE_NODE_OUTPUT, new ReplaceStrategy()); - // keyword extract节点输出 - keyStrategyHashMap.put(KEYWORD_EXTRACT_NODE_OUTPUT, new ReplaceStrategy()); - keyStrategyHashMap.put(EVIDENCES, new ReplaceStrategy()); + keyStrategyHashMap.put(GENEGRATED_SEMANTIC_MODEL_PROMPT, new ReplaceStrategy()); + // EVIDENCE节点输出 + keyStrategyHashMap.put(EVIDENCE, new ReplaceStrategy()); // schema recall节点输出 keyStrategyHashMap.put(TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, new ReplaceStrategy()); - keyStrategyHashMap.put(COLUMN_DOCUMENTS_BY_KEYWORDS_OUTPUT, new ReplaceStrategy()); + keyStrategyHashMap.put(COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT, new ReplaceStrategy()); // table relation节点输出 keyStrategyHashMap.put(TABLE_RELATION_OUTPUT, new ReplaceStrategy()); keyStrategyHashMap.put(TABLE_RELATION_EXCEPTION_OUTPUT, new ReplaceStrategy()); keyStrategyHashMap.put(TABLE_RELATION_RETRY_COUNT, new ReplaceStrategy()); + // Feasibility Assessment 节点输出 + keyStrategyHashMap.put(FEASIBILITY_ASSESSMENT_NODE_OUTPUT, new ReplaceStrategy()); // sql generate节点输出 keyStrategyHashMap.put(SQL_GENERATE_SCHEMA_MISSING_ADVICE, new ReplaceStrategy()); keyStrategyHashMap.put(SQL_GENERATE_OUTPUT, new ReplaceStrategy()); @@ -160,10 +161,12 @@ public StateGraph nl2sqlGraph(NodeBeanUtil nodeBeanUtil) throws GraphStateExcept }; StateGraph stateGraph = new StateGraph(NL2SQL_GRAPH_NAME, keyStrategyFactory) - .addNode(QUERY_REWRITE_NODE, nodeBeanUtil.getNodeBeanAsync(QueryRewriteNode.class)) - .addNode(KEYWORD_EXTRACT_NODE, nodeBeanUtil.getNodeBeanAsync(KeywordExtractNode.class)) + .addNode(INTENT_RECOGNITION_NODE, nodeBeanUtil.getNodeBeanAsync(IntentRecognitionNode.class)) + .addNode(EVIDENCE_RECALL_NODE, nodeBeanUtil.getNodeBeanAsync(EvidenceRecallNode.class)) + .addNode(QUERY_ENHANCE_NODE, nodeBeanUtil.getNodeBeanAsync(QueryEnhanceNode.class)) .addNode(SCHEMA_RECALL_NODE, nodeBeanUtil.getNodeBeanAsync(SchemaRecallNode.class)) .addNode(TABLE_RELATION_NODE, nodeBeanUtil.getNodeBeanAsync(TableRelationNode.class)) + .addNode(FEASIBILITY_ASSESSMENT_NODE, nodeBeanUtil.getNodeBeanAsync(FeasibilityAssessmentNode.class)) .addNode(SQL_GENERATE_NODE, nodeBeanUtil.getNodeBeanAsync(SqlGenerateNode.class)) .addNode(PLANNER_NODE, nodeBeanUtil.getNodeBeanAsync(PlannerNode.class)) .addNode(PLAN_EXECUTOR_NODE, nodeBeanUtil.getNodeBeanAsync(PlanExecutorNode.class)) @@ -175,13 +178,19 @@ public StateGraph nl2sqlGraph(NodeBeanUtil nodeBeanUtil) throws GraphStateExcept .addNode(SEMANTIC_CONSISTENCY_NODE, nodeBeanUtil.getNodeBeanAsync(SemanticConsistencyNode.class)) .addNode(HUMAN_FEEDBACK_NODE, nodeBeanUtil.getNodeBeanAsync(HumanFeedbackNode.class)); - stateGraph.addEdge(START, QUERY_REWRITE_NODE) - .addConditionalEdges(QUERY_REWRITE_NODE, edge_async(new QueryRewriteDispatcher()), - Map.of(KEYWORD_EXTRACT_NODE, KEYWORD_EXTRACT_NODE, END, END)) - .addEdge(KEYWORD_EXTRACT_NODE, SCHEMA_RECALL_NODE) + stateGraph.addEdge(START, INTENT_RECOGNITION_NODE) + .addConditionalEdges(INTENT_RECOGNITION_NODE, edge_async(new IntentRecognitionDispatcher()), + Map.of(EVIDENCE_RECALL_NODE, EVIDENCE_RECALL_NODE, END, END)) + .addEdge(EVIDENCE_RECALL_NODE, QUERY_ENHANCE_NODE) + .addConditionalEdges(QUERY_ENHANCE_NODE, edge_async(new QueryEnhanceDispatcher()), + Map.of(SCHEMA_RECALL_NODE, SCHEMA_RECALL_NODE, END, END)) .addEdge(SCHEMA_RECALL_NODE, TABLE_RELATION_NODE) .addConditionalEdges(TABLE_RELATION_NODE, edge_async(new TableRelationDispatcher()), - Map.of(PLANNER_NODE, PLANNER_NODE, END, END, TABLE_RELATION_NODE, TABLE_RELATION_NODE)) // retry + Map.of(FEASIBILITY_ASSESSMENT_NODE, FEASIBILITY_ASSESSMENT_NODE, END, END, TABLE_RELATION_NODE, + TABLE_RELATION_NODE)) // retry + .addConditionalEdges(FEASIBILITY_ASSESSMENT_NODE, edge_async(new FeasibilityAssessmentDispatcher()), + Map.of(PLANNER_NODE, PLANNER_NODE, END, END)) + // The edge from PlannerNode now goes to PlanExecutorNode for validation and // execution .addEdge(PLANNER_NODE, PLAN_EXECUTOR_NODE) @@ -214,16 +223,15 @@ public StateGraph nl2sqlGraph(NodeBeanUtil nodeBeanUtil) throws GraphStateExcept .addConditionalEdges(SQL_EXECUTE_NODE, edge_async(new SQLExecutorDispatcher()), Map.of(SQL_GENERATE_NODE, SQL_GENERATE_NODE, SEMANTIC_CONSISTENCY_NODE, SEMANTIC_CONSISTENCY_NODE)) .addConditionalEdges(SQL_GENERATE_NODE, edge_async(new SqlGenerateDispatcher()), - Map.of(KEYWORD_EXTRACT_NODE, KEYWORD_EXTRACT_NODE, END, END, SQL_EXECUTE_NODE, SQL_EXECUTE_NODE)) + Map.of(FEASIBILITY_ASSESSMENT_NODE, FEASIBILITY_ASSESSMENT_NODE, END, END, SQL_EXECUTE_NODE, + SQL_EXECUTE_NODE)) .addConditionalEdges(SEMANTIC_CONSISTENCY_NODE, edge_async(new SemanticConsistenceDispatcher()), Map.of(SQL_GENERATE_NODE, SQL_GENERATE_NODE, PLAN_EXECUTOR_NODE, PLAN_EXECUTOR_NODE)); GraphRepresentation graphRepresentation = stateGraph.getGraph(GraphRepresentation.Type.PLANTUML, "workflow graph"); - log.info("\n\n"); - log.info(graphRepresentation.content()); - log.info("\n\n"); + log.info("workflow in PlantUML format as follows \n\n" + graphRepresentation.content() + "\n\n"); return stateGraph; } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/config/DataAgentProperties.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/config/DataAgentProperties.java index 51e7f5200..51abd8d8f 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/config/DataAgentProperties.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/config/DataAgentProperties.java @@ -37,6 +37,8 @@ public class DataAgentProperties { */ private EmbeddingBatch embeddingBatch = new EmbeddingBatch(); + private VectorStoreProperties vectorStore = new VectorStoreProperties(); + @Getter @Setter public static class EmbeddingBatch { @@ -63,4 +65,35 @@ public static class EmbeddingBatch { } + @Getter + @Setter + public static class VectorStoreProperties { + + /** + * 相似度阈值配置，用于过滤相似度分数大于等于此阈值的文档 + */ + private double similarityThreshold = 0.2; + + /** + * 一次删除操作中，最多删除的文档数量 + */ + private int batchDelTopkLimit = 5000; + + /** + * 查询时返回的最大文档数量 + */ + private int topkLimit = 30; + + /** + * 是否启用混合搜索 + */ + private boolean enableHybridSearch = false; + + /** + * Elasticsearch最小分数阈值，用于es执行关键词搜索时过滤相关性较低的文档 + */ + private double elasticsearchMinScore = 0.5; + + } + } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/constant/Constant.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/constant/Constant.java index 35bd698b7..758dad4fe 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/constant/Constant.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/constant/Constant.java @@ -35,17 +35,19 @@ private Constant() { public static final String NL2SQL_GRAPH_NAME = "nl2sqlGraph"; - public static final String QUERY_REWRITE_NODE_OUTPUT = "QUERY_REWRITE_NODE_OUTPUT"; + public static final String INTENT_RECOGNITION_NODE_OUTPUT = "INTENT_RECOGNITION_NODE_OUTPUT"; - public static final String KEYWORD_EXTRACT_NODE_OUTPUT = "KEYWORD_EXTRACT_NODE_OUTPUT"; + public static final String QUERY_ENHANCE_NODE_OUTPUT = "QUERY_ENHANCE_NODE_OUTPUT"; - public static final String EVIDENCES = "EVIDENCES"; + public static final String FEASIBILITY_ASSESSMENT_NODE_OUTPUT = "FEASIBILITY_ASSESSMENT_NODE_OUTPUT"; + + public static final String EVIDENCE = "EVIDENCE"; public static final String TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT = "TABLE_DOCUMENTS_FOR_SCHEMA"; public static final String SCHEMA_RECALL_NODE_OUTPUT = "SCHEMA_RECALL_NODE_OUTPUT"; - public static final String COLUMN_DOCUMENTS_BY_KEYWORDS_OUTPUT = "COLUMN_DOCUMENTS_BY_KEYWORDS_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"; @@ -53,9 +55,7 @@ private Constant() { public static final String TABLE_RELATION_RETRY_COUNT = "TABLE_RELATION_RETRY_COUNT"; - public static final String BUSINESS_KNOWLEDGE = "BUSINESS_KNOWLEDGE"; - - public static final String SEMANTIC_MODEL = "SEMANTIC_MODEL"; + public static final String GENEGRATED_SEMANTIC_MODEL_PROMPT = "GENEGRATED_SEMANTIC_MODEL_PROMPT"; public static final String SQL_GENERATE_OUTPUT = "SQL_GENERATE_OUTPUT"; @@ -93,6 +93,14 @@ private Constant() { public static final String PLAN_EXECUTOR_NODE = "PLAN_EXECUTOR_NODE"; + public static final String INTENT_RECOGNITION_NODE = "INTENT_RECOGNITION_NODE"; + + public static final String EVIDENCE_RECALL_NODE = "EVIDENCE_RECALL_NODE"; + + public static final String QUERY_ENHANCE_NODE = "QUERY_ENHANCE_NODE"; + + public static final String FEASIBILITY_ASSESSMENT_NODE = "FEASIBILITY_ASSESSMENT_NODE"; + public static final String QUERY_REWRITE_NODE = "QUERY_REWRITE_NODE"; public static final String REPORT_GENERATOR_NODE = "REPORT_GENERATOR_NODE"; @@ -144,16 +152,4 @@ private Constant() { // 人类复核相关 public static final String HUMAN_REVIEW_ENABLED = "HUMAN_REVIEW_ENABLED"; - // column - public static final String COLUMN = "column"; - - // table - public static final String TABLE = "table"; - - // vectorType - public static final String VECTOR_TYPE = "vectorType"; - - // knowledgeId - public static final String KNOWLEDGE_ID = "knowledgeId"; - } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/constant/DocumentMetadataConstant.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/constant/DocumentMetadataConstant.java new file mode 100644 index 000000000..3379e2303 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/constant/DocumentMetadataConstant.java @@ -0,0 +1,49 @@ +/* + * Copyright 2024-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.cloud.ai.constant; + +// Document中metadata中存储的key +public final class DocumentMetadataConstant { + + private DocumentMetadataConstant() { + + } + + // column + public static final String COLUMN = "column"; + + // table + public static final String TABLE = "table"; + + // vectorType + public static final String VECTOR_TYPE = "vectorType"; + + // knowledgeId + public static final String KNOWLEDGE_ID = "knowledgeId"; + + // businessTerm + public static final String BUSINESS_TERM = "businessTerm"; + + // 智能体的知识 + public static final String AGENT_KNOWLEDGE = "agentKnowledge"; + + // 关联该文档在数据库中的记录id + public static final String DB_RECORD_ID = "dbRecordId"; + + // isRecall + public static final String IS_RECALL = "isRecall"; + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/QueryRewriteDispatcher.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/FeasibilityAssessmentDispatcher.java similarity index 54% rename from spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/QueryRewriteDispatcher.java rename to spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/FeasibilityAssessmentDispatcher.java index 5f901210d..b35ca74b7 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/QueryRewriteDispatcher.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/FeasibilityAssessmentDispatcher.java @@ -23,30 +23,25 @@ import static com.alibaba.cloud.ai.constant.Constant.*; import static com.alibaba.cloud.ai.graph.StateGraph.END; -/** - * @author zhangshenghang - */ @Slf4j -public class QueryRewriteDispatcher implements EdgeAction { +public class FeasibilityAssessmentDispatcher implements EdgeAction { @Override - public String apply(OverAllState state) { - // value的值是和 resources/init-rewrite.txt的输出一致，例如 - // 需求类型：《自由闲聊》......... - // 语种类型：《中文》 - // 需求内容：2025-10-17天气如何 - String value = state.value(QUERY_REWRITE_NODE_OUTPUT, END); - log.info("[QueryRewriteDispatcher]apply方法被调用，QUERY_REWRITE_NODE_OUTPUT的value如下\n {}", value); - - if (value != null && value.contains("需求类型：《数据分析》")) { - log.info("[QueryRewriteDispatcher]需求类型为数据分析，进入KEYWORD_EXTRACT_NODE节点"); - return KEYWORD_EXTRACT_NODE; + public String apply(OverAllState state) throws Exception { + // value的值是和 resources/feasibility-assessment.txt的输出一致，例如 + // 【需求类型】：《数据分析》 + // 【语种类型】：《中文》 + // 【需求内容】：查询所有“核心用户”的数量 + String value = state.value(FEASIBILITY_ASSESSMENT_NODE_OUTPUT, END); + + if (value != null && value.contains("【需求类型】：《数据分析》")) { + log.info("[FeasibilityAssessmentNodeDispatcher]需求类型为数据分析，进入PlannerNode节点"); + return PLANNER_NODE; } else { - log.info("[QueryRewriteDispatcher]需求类型非数据分析，返回END节点"); + log.info("[FeasibilityAssessmentNodeDispatcher]需求类型非数据分析，返回END节点"); return END; } - } } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/IntentRecognitionDispatcher.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/IntentRecognitionDispatcher.java new file mode 100644 index 000000000..ceeece468 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/IntentRecognitionDispatcher.java @@ -0,0 +1,69 @@ +/* + * Copyright 2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.alibaba.cloud.ai.dispatcher; + +import com.alibaba.cloud.ai.graph.OverAllState; +import com.alibaba.cloud.ai.graph.action.EdgeAction; +import com.alibaba.cloud.ai.util.JsonUtil; +import com.alibaba.cloud.ai.util.MarkdownParserUtil; +import com.alibaba.cloud.ai.util.StateUtil; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import lombok.extern.slf4j.Slf4j; + +import static com.alibaba.cloud.ai.constant.Constant.*; +import static com.alibaba.cloud.ai.graph.StateGraph.END; + +/** + * 根据意图识别结果决定下一个节点的分发器 + */ +@Slf4j +public class IntentRecognitionDispatcher implements EdgeAction { + + @Override + public String apply(OverAllState state) throws Exception { + // 获取意图识别结果 + String intentResult = StateUtil.getStringValue(state, INTENT_RECOGNITION_NODE_OUTPUT); + + if (intentResult == null || intentResult.trim().isEmpty()) { + log.warn("Intent recognition result is null or empty, defaulting to END"); + return END; + } + + try { + // 解析JSON格式的意图识别结果 + JsonNode jsonNode = JsonUtil.getObjectMapper() + .readTree(MarkdownParserUtil.extractText(intentResult.trim())); + String classification = jsonNode.path("classification").asText(); + + // 根据分类结果决定下一个节点 + if ("《闲聊或无关指令》".equals(classification)) { + log.warn("Intent classified as chat or irrelevant, ending conversation"); + return END; + } + else { + log.info("Intent classified as potential data analysis request, proceeding to evidence recall"); + return EVIDENCE_RECALL_NODE; + } + } + catch (JsonProcessingException e) { + log.error("Failed to parse intent recognition result as JSON: {}", intentResult, e); + return END; + } + } + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/QueryEnhanceDispatcher.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/QueryEnhanceDispatcher.java new file mode 100644 index 000000000..c7c73aacc --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/QueryEnhanceDispatcher.java @@ -0,0 +1,66 @@ +/* + * Copyright 2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.cloud.ai.dispatcher; + +import com.alibaba.cloud.ai.dto.QueryEnhanceOutputDTO; +import com.alibaba.cloud.ai.graph.OverAllState; +import com.alibaba.cloud.ai.graph.action.EdgeAction; +import com.alibaba.cloud.ai.util.StateUtil; +import lombok.extern.slf4j.Slf4j; + +import static com.alibaba.cloud.ai.constant.Constant.QUERY_ENHANCE_NODE_OUTPUT; +import static com.alibaba.cloud.ai.constant.Constant.SCHEMA_RECALL_NODE; +import static com.alibaba.cloud.ai.graph.StateGraph.END; + +/** + * 根据查询增强结果决定下一个节点的分发器 + */ +@Slf4j +public class QueryEnhanceDispatcher implements EdgeAction { + + @Override + public String apply(OverAllState state) throws Exception { + // 获取查询处理结果 + QueryEnhanceOutputDTO queryProcessOutput = StateUtil.getObjectValue(state, QUERY_ENHANCE_NODE_OUTPUT, + QueryEnhanceOutputDTO.class); + + // 检查查询处理结果是否为空 + if (queryProcessOutput == null) { + log.warn("Query process output is null, ending conversation"); + return END; + } + + // 检查各个字段是否为空 + boolean isCanonicalQueryEmpty = queryProcessOutput.getCanonicalQuery() == null + || queryProcessOutput.getCanonicalQuery().trim().isEmpty(); + boolean isExpandedQueriesEmpty = queryProcessOutput.getExpandedQueries() == null + || queryProcessOutput.getExpandedQueries().isEmpty(); + boolean isConsolidatedKeywordsEmpty = queryProcessOutput.getConsolidatedKeywords() == null + || queryProcessOutput.getConsolidatedKeywords().isEmpty(); + + if (isCanonicalQueryEmpty || isExpandedQueriesEmpty || isConsolidatedKeywordsEmpty) { + log.warn( + "Query process output contains empty fields - canonicalQuery: {}, expandedQueries: {}, consolidatedKeywords: {}", + isCanonicalQueryEmpty, isExpandedQueriesEmpty, isConsolidatedKeywordsEmpty); + return END; + } + else { + log.info("Query process output is valid, proceeding to schema recall"); + return SCHEMA_RECALL_NODE; + } + } + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/SqlGenerateDispatcher.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/SqlGenerateDispatcher.java index 530f1f9aa..64a600bf0 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/SqlGenerateDispatcher.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/SqlGenerateDispatcher.java @@ -38,9 +38,10 @@ public String apply(OverAllState state) { log.info("检测到流程结束标志: {}", END); yield END; } + // TODO 需要优化，不能简单跳转 case SQL_GENERATE_SCHEMA_MISSING -> { - log.warn("SQL生成缺少Schema，跳转到{}节点", KEYWORD_EXTRACT_NODE); - yield KEYWORD_EXTRACT_NODE; + log.warn("SQL生成缺少Schema，跳转到{}节点", FEASIBILITY_ASSESSMENT_NODE); + yield FEASIBILITY_ASSESSMENT_NODE; } default -> { log.info("SQL生成成功，进入SQL执行节点: {}", SQL_EXECUTE_NODE); diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/TableRelationDispatcher.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/TableRelationDispatcher.java index 7ab10b354..9b3094817 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/TableRelationDispatcher.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dispatcher/TableRelationDispatcher.java @@ -46,7 +46,7 @@ public String apply(OverAllState state) throws Exception { Optional tableRelationOutput = state.value(TABLE_RELATION_OUTPUT); if (tableRelationOutput.isPresent()) { - return PLANNER_NODE; // next node is planner + return FEASIBILITY_ASSESSMENT_NODE; } // no output, end diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/BusinessKnowledgeDTO.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/BusinessKnowledgeDTO.java index 60b659055..7b5ebf775 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/BusinessKnowledgeDTO.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/BusinessKnowledgeDTO.java @@ -16,6 +16,8 @@ package com.alibaba.cloud.ai.dto; import com.alibaba.cloud.ai.entity.BusinessKnowledge; +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.NotNull; import lombok.AllArgsConstructor; import lombok.Builder; import lombok.Data; @@ -30,8 +32,10 @@ @Builder public class BusinessKnowledgeDTO { + @NotBlank(message = "Business term cannot be empty") private String businessTerm; // Business term + @NotBlank(message = "Description cannot be empty") private String description; // Description private String synonyms; // Synonyms, comma separated @@ -39,6 +43,7 @@ public class BusinessKnowledgeDTO { @Builder.Default private Boolean isRecall = true; // is recall + @NotNull(message = "Agent ID cannot be Null") private Long agentId; // Associated agent ID public BusinessKnowledge toEntity() { diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/QueryEnhanceOutputDTO.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/QueryEnhanceOutputDTO.java new file mode 100644 index 000000000..0ea499a7c --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/QueryEnhanceOutputDTO.java @@ -0,0 +1,41 @@ +/* + * Copyright 2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.cloud.ai.dto; + +import com.fasterxml.jackson.annotation.JsonProperty; +import lombok.Data; +import lombok.NoArgsConstructor; + +import java.util.List; + +// 对应 模板query-enhancement.txt的输出 +@Data +@NoArgsConstructor +public class QueryEnhanceOutputDTO { + + // 经LLM重写后的 规范化查询 + @JsonProperty("canonical_query") + private String canonicalQuery; + + // 基于canonicalQuery的扩展查询 + @JsonProperty("expanded_queries") + private List expandedQueries; + + // 从"规范化查询"、所有"扩展问题"以及 "业务知识涉及的关键概念和数值"中，提取出的所有核心关键词。 + @JsonProperty("consolidated_keywords") + private List consolidatedKeywords; + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/SemanticModelAddDTO.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/SemanticModelAddDTO.java new file mode 100644 index 000000000..17e3e6430 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/SemanticModelAddDTO.java @@ -0,0 +1,80 @@ +/* + * Copyright 2024-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.cloud.ai.dto; + +import lombok.AllArgsConstructor; +import lombok.Builder; +import lombok.Data; +import lombok.NoArgsConstructor; + +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.NotNull; + +/** + * 添加语义模型的DTO类 + */ +@Data +@Builder +@NoArgsConstructor +@AllArgsConstructor +public class SemanticModelAddDTO { + + /** + * 关联的智能体ID + */ + @NotNull(message = "智能体ID不能为空") + private Integer agentId; + + /** + * 关联的表名 + */ + @NotBlank(message = "表名不能为空") + private String tableName; + + /** + * 数据库中的物理字段名 (例如: csat_score) + */ + @NotBlank(message = "数据库字段名不能为空") + private String columnName; + + /** + * 业务名/别名 (例如: 客户满意度分数) + */ + @NotBlank(message = "业务名称不能为空") + private String businessName; + + /** + * 业务名的同义词 (例如: 满意度,客户评分) + */ + private String synonyms; + + /** + * 业务描述 (用于向LLM解释字段的业务含义) + */ + private String businessDescription; + + /** + * 数据库中的物理字段的原始注释 + */ + private String columnComment; + + /** + * 物理数据类型 (例如: int, varchar(20)) + */ + @NotBlank(message = "数据类型不能为空") + private String dataType; + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/schema/ColumnDTO.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/schema/ColumnDTO.java index 6b4a6b14c..adc05b14f 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/schema/ColumnDTO.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/dto/schema/ColumnDTO.java @@ -38,8 +38,6 @@ public class ColumnDTO { private String type; - private List samples; - private List data; private Map mapping; diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/entity/SemanticModel.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/entity/SemanticModel.java index 8298a4c85..9cce5d0f5 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/entity/SemanticModel.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/entity/SemanticModel.java @@ -31,29 +31,77 @@ @AllArgsConstructor public class SemanticModel { + /** + * 主键ID + */ private Long id; - private Long agentId; + /** + * 关联的智能体ID + */ + private Integer agentId; - private String fieldName; + /** + * 关联的数据源ID + */ + private Integer datasourceId; - private String conversationName; + /** + * 关联的表名 + */ + private String tableName; + /** + * 数据库中的物理字段名 (例如: csat_score) + */ + private String columnName; + + /** + * 业务名/别名 (例如: 客户满意度分数) + */ + private String businessName; + + /** + * 业务名的同义词 (例如: 满意度,客户评分) + */ private String synonyms; - private String description; + /** + * 业务描述 (用于向LLM解释字段的业务含义) + */ + private String businessDescription; - private String type; + /** + * 数据库中的物理字段的原始注释 + */ + private String columnComment; + /** + * 物理数据类型 (例如: int, varchar(20)) + */ + private String dataType; + + /** + * 状态: 0 停用 1 启用 + */ + private Integer status; + + /** + * 创建时间 + */ private LocalDateTime createdTime; + /** + * 更新时间 + */ private LocalDateTime updateTime; - private Integer status; - + /** + * 获取用于提示的信息 + */ public String getPromptInfo() { - return String.format("对话字段名称: %s, 数据库字段名: %s, 字段同义词: %s, 字段描述: %s, 字段类型: %s", conversationName, fieldName, - synonyms, description, type); + return String.format("业务名称: %s, 表名: %s, 数据库字段名: %s, 字段同义词: %s, 业务描述: %s, 数据类型: %s", businessName, tableName, + columnName, synonyms, businessDescription, dataType); } } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/mapper/SemanticModelMapper.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/mapper/SemanticModelMapper.java index 49afb1052..c88af1a53 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/mapper/SemanticModelMapper.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/mapper/SemanticModelMapper.java @@ -51,9 +51,9 @@ public interface SemanticModelMapper { */ @Select(""" SELECT * FROM semantic_model - WHERE field_name LIKE CONCAT('%', #{keyword}, '%') - OR conversation_name LIKE CONCAT('%', #{keyword}, '%') - OR description LIKE CONCAT('%', #{keyword}, '%') + WHERE column_name LIKE CONCAT('%', #{keyword}, '%') + OR business_name LIKE CONCAT('%', #{keyword}, '%') + OR business_description LIKE CONCAT('%', #{keyword}, '%') OR synonyms LIKE CONCAT('%', #{keyword}, '%') ORDER BY created_time DESC """) @@ -92,9 +92,9 @@ OR synonyms LIKE CONCAT('%', #{keyword}, '%') @Insert(""" INSERT INTO semantic_model - (agent_id, field_name, conversation_name, synonyms, description, type, created_time, updated_time, status) + (agent_id, datasource_id, table_name, column_name, business_name, synonyms, business_description, column_comment, data_type, created_time, updated_time, status) VALUES - (#{agentId}, #{fieldName}, #{conversationName}, #{synonyms}, #{description}, #{type}, NOW(), NOW(), #{status}) + (#{agentId}, #{datasourceId}, #{tableName}, #{columnName}, #{businessName}, #{synonyms}, #{businessDescription}, #{columnComment}, #{dataType}, NOW(), NOW(), #{status}) """) @Options(useGeneratedKeys = true, keyProperty = "id", keyColumn = "id") int insert(SemanticModel model); @@ -104,11 +104,14 @@ OR synonyms LIKE CONCAT('%', #{keyword}, '%') UPDATE semantic_model agent_id = #{agentId}, - field_name = #{fieldName}, - conversation_name = #{conversationName}, + datasource_id = #{datasourceId}, + table_name = #{tableName}, + column_name = #{columnName}, + business_name = #{businessName}, synonyms = #{synonyms}, - description = #{description}, - type = #{type}, + business_description = #{businessDescription}, + column_comment = #{columnComment}, + data_type = #{dataType}, status = #{status}, updated_time = NOW() @@ -123,4 +126,22 @@ OR synonyms LIKE CONCAT('%', #{keyword}, '%') """) int deleteById(@Param("id") Long id); + /** + * Query semantic models by datasource ID, status and table names + */ + @Select(""" + + """) + List selectByDatasourceIdAndTableNames(@Param("datasourceId") Integer datasourceId, + @Param("tableNames") List tableNames); + } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/EvidenceRecallNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/EvidenceRecallNode.java new file mode 100644 index 000000000..ee6807f35 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/EvidenceRecallNode.java @@ -0,0 +1,171 @@ +/* + * Copyright 2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.cloud.ai.node; + +import com.alibaba.cloud.ai.constant.DocumentMetadataConstant; +import com.alibaba.cloud.ai.enums.TextType; +import com.alibaba.cloud.ai.graph.GraphResponse; +import com.alibaba.cloud.ai.graph.OverAllState; +import com.alibaba.cloud.ai.graph.action.NodeAction; +import com.alibaba.cloud.ai.graph.streaming.FluxConverter; +import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; +import com.alibaba.cloud.ai.prompt.PromptConstant; +import com.alibaba.cloud.ai.prompt.PromptHelper; +import com.alibaba.cloud.ai.service.llm.LlmService; +import com.alibaba.cloud.ai.service.vectorstore.AgentVectorStoreService; +import com.alibaba.cloud.ai.util.*; +import com.fasterxml.jackson.core.type.TypeReference; +import lombok.AllArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.jetbrains.annotations.Nullable; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.document.Document; +import org.springframework.stereotype.Component; +import org.springframework.util.Assert; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Sinks; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static com.alibaba.cloud.ai.constant.Constant.*; + +@Slf4j +@Component +@AllArgsConstructor +public class EvidenceRecallNode implements NodeAction { + + private final LlmService llmService; + + private final AgentVectorStoreService vectorStoreService; + + @Override + public Map apply(OverAllState state) throws Exception { + + // 从state中提取question和agentId + String question = StateUtil.getStringValue(state, INPUT_KEY); + String agentId = StateUtil.getStringValue(state, AGENT_ID); + Assert.hasText(agentId, "Agent ID cannot be empty."); + + log.info("Extracting keywords before getting evidence in question: {}", question); + log.debug("Agent ID: {}", agentId); + + // 构建关键词提取提示词 + String prompt = PromptConstant.getQuestionToKeywordsPromptTemplate().render(Map.of("question", question)); + log.debug("Built evidence keyword extraction prompt as follows \n {} \n", prompt); + + // 调用LLM提取关键词 + Flux responseFlux = llmService.callUser(prompt); + Sinks.Many evidenceDisplaySink = Sinks.many().multicast().onBackpressureBuffer(); + + final Map resultMap = new HashMap<>(); + Flux> generator = FluxUtil.createStreamingGenerator(this.getClass(), state, + responseFlux, + Flux.just(ChatResponseUtil.createResponse("正在获取关键词..."), + ChatResponseUtil.createPureResponse(TextType.JSON.getStartSign())), + Flux.just(ChatResponseUtil.createPureResponse(TextType.JSON.getEndSign()), + ChatResponseUtil.createResponse("\n关键词获取完成！")), + result -> { + resultMap.putAll(getEvidences(result, agentId, evidenceDisplaySink)); + return resultMap; + }); + + Flux> evidenceFlux = FluxConverter.builder() + .startingNode(this.getClass().getSimpleName()) + .startingState(state) + .mapResult(r -> resultMap) + .build(evidenceDisplaySink.asFlux().map(ChatResponseUtil::createPureResponse)); + return Map.of(EVIDENCE, generator.concatWith(evidenceFlux)); + } + + private Map getEvidences(String llmOutput, String agentId, Sinks.Many sink) { + try { + List keywords = extractKeywords(llmOutput); + + if (null == keywords || keywords.isEmpty()) { + log.debug("No keywords extracted from LLM output"); + sink.tryEmitNext("未找到关键词！\n"); + return Map.of(EVIDENCE, "无"); + } + + // 将关键词列表用空格拼接成字符串 + sink.tryEmitNext("关键词：\n"); + keywords.forEach(keyword -> sink.tryEmitNext(keyword + " ")); + sink.tryEmitNext("\n"); + String keywordsString = String.join(" ", keywords); + log.debug("Joined keywords string: {}", keywordsString); + sink.tryEmitNext("正在获取证据..."); + + // 获取业务知识和智能体的知识 + List businessTermDocuments = vectorStoreService + .getDocumentsForAgent(agentId, keywordsString, DocumentMetadataConstant.BUSINESS_TERM) + .stream() + .filter(doc -> Boolean.TRUE.equals(doc.getMetadata().get(DocumentMetadataConstant.IS_RECALL))) + .toList(); + + // 检查是否有证据文档 + if (businessTermDocuments.isEmpty()) { + log.debug("No evidence documents found for agent: {} with keywords: {}", agentId, keywordsString); + sink.tryEmitNext("未找到证据！\n"); + return Map.of(EVIDENCE, "无"); + } + + // 构建业务知识提示 + String businessKnowledgePrompt = PromptHelper.buildBusinessKnowledgePrompt( + businessTermDocuments.stream().map(Document::getText).collect(Collectors.joining(";\n"))); + // TODO 根据知识库模板渲染智能体的知识，然后拼接成EVIDENCE。 businessKnowledgePrompt + "\n\n" + + // agentKnowledgePrompt; + String evidence = businessKnowledgePrompt + "\n\n"; + + // 输出证据内容 + sink.tryEmitNext("证据内容：\n"); + businessTermDocuments.forEach(e -> sink.tryEmitNext(e.getText() + "\n")); + // TODO agentKnowledge.forEach + + // 返回结果 + return Map.of(EVIDENCE, evidence); + } + catch (Exception e) { + log.error("Error occurred while getting evidences", e); + sink.tryEmitError(e); + return Map.of(EVIDENCE, ""); + } + finally { + sink.tryEmitComplete(); + } + } + + @Nullable + private static List extractKeywords(String llmOutput) { + // 解析关键词列表 + List keywords; + try { + String content = MarkdownParserUtil.extractText(llmOutput.trim()); + keywords = JsonUtil.getObjectMapper().readValue(content, new TypeReference>() { + }); + log.info("For getting evidence keyword,extracted {} keywords: {}", keywords != null ? keywords.size() : 0, + keywords); + } + catch (Exception e) { + log.error("Failed to parse keywords from LLM response", e); + keywords = List.of(); // 使用空列表作为默认值 + } + return keywords; + } + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/FeasibilityAssessmentNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/FeasibilityAssessmentNode.java new file mode 100644 index 000000000..03dad7cb0 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/FeasibilityAssessmentNode.java @@ -0,0 +1,74 @@ +/* + * Copyright 2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.cloud.ai.node; + +import com.alibaba.cloud.ai.dto.schema.SchemaDTO; +import com.alibaba.cloud.ai.graph.GraphResponse; +import com.alibaba.cloud.ai.graph.OverAllState; +import com.alibaba.cloud.ai.graph.action.NodeAction; +import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; +import com.alibaba.cloud.ai.prompt.PromptHelper; +import com.alibaba.cloud.ai.service.llm.LlmService; +import com.alibaba.cloud.ai.util.FluxUtil; +import com.alibaba.cloud.ai.util.StateUtil; +import lombok.AllArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.stereotype.Component; +import reactor.core.publisher.Flux; + +import java.util.Map; + +import static com.alibaba.cloud.ai.constant.Constant.*; + +// 可行性评估节点，看需求是 数据分析/需要澄清 或者最终确认为自由闲聊 +@Slf4j +@Component +@AllArgsConstructor +public class FeasibilityAssessmentNode implements NodeAction { + + private final LlmService llmService; + + @Override + public Map apply(OverAllState state) throws Exception { + // 获取canonical_query + String canonicalQuery = StateUtil.getCanonicalQuery(state); + + // 获取召回的Schema + SchemaDTO recalledSchema = StateUtil.getObjectValue(state, TABLE_RELATION_OUTPUT, SchemaDTO.class); + + // 获取证据信息 + String evidence = StateUtil.getStringValue(state, EVIDENCE); + + // 构建可行性评估提示词，多轮对话暂时为空 + String prompt = PromptHelper.buildFeasibilityAssessmentPrompt(canonicalQuery, recalledSchema, evidence, null); + log.debug("Built feasibility assessment prompt as follows \n {} \n", prompt); + + // 调用LLM进行可行性评估 + Flux responseFlux = llmService.callUser(prompt); + + Flux> generator = FluxUtil.createStreamingGeneratorWithMessages(this.getClass(), + state, "正在进行可行性评估...", "可行性评估完成！", llmOutput -> { + // 获取评估结果 + String assessmentResult = llmOutput.trim(); + log.info("Feasibility assessment result: {}", assessmentResult); + // 返回评估结果 + return Map.of(FEASIBILITY_ASSESSMENT_NODE_OUTPUT, assessmentResult); + }, responseFlux); + return Map.of(FEASIBILITY_ASSESSMENT_NODE_OUTPUT, generator); + } + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/HumanFeedbackNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/HumanFeedbackNode.java index 61d63317a..2f401ded2 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/HumanFeedbackNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/HumanFeedbackNode.java @@ -39,7 +39,6 @@ public class HumanFeedbackNode implements NodeAction { @Override public Map apply(OverAllState state) throws Exception { - log.info("Processing human feedback"); Map updated = new HashMap<>(); // 检查最大修复次数 diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/IntentRecognitionNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/IntentRecognitionNode.java new file mode 100644 index 000000000..2d6054891 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/IntentRecognitionNode.java @@ -0,0 +1,73 @@ +/* + * Copyright 2024-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.cloud.ai.node; + +import com.alibaba.cloud.ai.enums.TextType; +import com.alibaba.cloud.ai.graph.GraphResponse; +import com.alibaba.cloud.ai.graph.OverAllState; +import com.alibaba.cloud.ai.graph.action.NodeAction; +import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; +import com.alibaba.cloud.ai.prompt.PromptHelper; +import com.alibaba.cloud.ai.service.llm.LlmService; +import com.alibaba.cloud.ai.util.ChatResponseUtil; +import com.alibaba.cloud.ai.util.FluxUtil; +import com.alibaba.cloud.ai.util.StateUtil; +import lombok.AllArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.stereotype.Component; +import reactor.core.publisher.Flux; + +import java.util.Map; + +import static com.alibaba.cloud.ai.constant.Constant.INPUT_KEY; +import static com.alibaba.cloud.ai.constant.Constant.INTENT_RECOGNITION_NODE_OUTPUT; + +/** + * 意图识别节点，用于识别用户输入是闲聊还是数据分析请求 + */ +@Slf4j +@Component +@AllArgsConstructor +public class IntentRecognitionNode implements NodeAction { + + private final LlmService llmService; + + @Override + public Map apply(OverAllState state) throws Exception { + + // 获取用户输入 + String userInput = StateUtil.getStringValue(state, INPUT_KEY); + log.info("User input for intent recognition: {}", userInput); + + // 构建意图识别提示，多轮对话暂时为空 + String prompt = PromptHelper.buildIntentRecognitionPrompt(null, userInput); + log.debug("Built intent recognition prompt as follows \n {} \n", prompt); + + // 调用LLM进行意图识别 + Flux responseFlux = llmService.callUser(prompt); + + Flux> generator = FluxUtil.createStreamingGenerator(this.getClass(), state, + responseFlux, + Flux.just(ChatResponseUtil.createResponse("正在进行意图识别..."), + ChatResponseUtil.createPureResponse(TextType.JSON.getStartSign())), + Flux.just(ChatResponseUtil.createPureResponse(TextType.JSON.getEndSign()), + ChatResponseUtil.createResponse("\n意图识别完成！")), + result -> Map.of(INTENT_RECOGNITION_NODE_OUTPUT, result)); + return Map.of(INTENT_RECOGNITION_NODE_OUTPUT, generator); + } + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/KeywordExtractNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/KeywordExtractNode.java deleted file mode 100644 index 1e2c6fdeb..000000000 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/KeywordExtractNode.java +++ /dev/null @@ -1,187 +0,0 @@ -/* - * Copyright 2025 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.alibaba.cloud.ai.node; - -import com.alibaba.cloud.ai.enums.TextType; -import com.alibaba.cloud.ai.graph.GraphResponse; -import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; -import com.alibaba.cloud.ai.pojo.KeywordExtractionResult; -import com.alibaba.cloud.ai.graph.OverAllState; -import com.alibaba.cloud.ai.graph.action.NodeAction; -import com.alibaba.cloud.ai.service.processing.QueryProcessingService; -import com.alibaba.cloud.ai.util.ChatResponseUtil; -import com.alibaba.cloud.ai.util.FluxUtil; -import com.alibaba.cloud.ai.util.StateUtil; -import lombok.AllArgsConstructor; -import lombok.extern.slf4j.Slf4j; -import org.springframework.ai.chat.model.ChatResponse; -import org.springframework.stereotype.Component; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; -import reactor.core.publisher.Sinks; - -import java.util.ArrayList; -import java.util.HashMap; -import java.util.HashSet; -import java.util.LinkedHashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; - -import static com.alibaba.cloud.ai.constant.Constant.*; - -/** - * Keyword, entity, and temporal information extraction node to prepare for subsequent - * schema recall. - * - * This node is responsible for: - Extracting evidences from user input - Extracting - * keywords based on evidences - Preparing structured information for schema recall - - * Providing streaming feedback during extraction process - * - * @author zhangshenghang - */ -@Slf4j -@Component -@AllArgsConstructor -public class KeywordExtractNode implements NodeAction { - - private final QueryProcessingService queryProcessingService; - - /** - * Process multiple question variants, extract keywords and merge results Use parallel - * stream processing to improve multi-question processing efficiency - * @param questions list of question variants - * @return list of extraction results - */ - private Flux processMultipleQuestions(List questions, String agentId, - Consumer> resultConsumer) { - List resultList = new ArrayList<>(questions.size()); - AtomicReference> fluxRef = new AtomicReference<>(Flux.empty()); - questions.forEach(question -> { - List evidences = queryProcessingService.extractEvidences(question, agentId); - fluxRef - .set(fluxRef.get().concatWith(queryProcessingService.extractKeywords(question, evidences, keywords -> { - log.info("成功从问题变体提取关键词: 问题=\"{}\", 关键词={}", question, keywords); - resultList.add(new KeywordExtractionResult(question, evidences, keywords)); - }).doOnError(e -> { - log.warn("从问题变体提取关键词失败: 问题={}", question, e); - resultList.add(new KeywordExtractionResult(question, false)); - }))); - }); - return fluxRef.get().doOnComplete(() -> resultConsumer.accept(resultList)); - } - - /** - * Merge keywords from multiple question variants, deduplicate and keep original - * question keywords priority - * @param extractionResults list of extraction results - * @param originalQuestion original question - * @return merged keyword list - */ - private List mergeKeywords(List extractionResults, String originalQuestion) { - if (extractionResults.isEmpty()) { - return List.of(); - } - - Set mergedKeywords = new LinkedHashSet<>(); - - extractionResults.stream() - .filter(result -> result.isSuccessful() && result.getQuestion().equals(originalQuestion)) - .findFirst() - .ifPresent(result -> mergedKeywords.addAll(result.getKeywords())); - - extractionResults.stream() - .filter(result -> result.isSuccessful() && !result.getQuestion().equals(originalQuestion)) - .forEach(result -> mergedKeywords.addAll(result.getKeywords())); - - return new ArrayList<>(mergedKeywords); - } - - /** - * Merge evidences from multiple question variants, deduplicate - * @param extractionResults list of extraction results - * @return merged evidence list - */ - private List mergeEvidences(List extractionResults) { - Set mergedEvidences = new HashSet<>(); - - extractionResults.stream() - .filter(KeywordExtractionResult::isSuccessful) - .forEach(result -> mergedEvidences.addAll(result.getEvidences())); - - return new ArrayList<>(mergedEvidences); - } - - @Override - public Map apply(OverAllState state) throws Exception { - log.info("Entering {} node", this.getClass().getSimpleName()); - - String input = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT, - StateUtil.getStringValue(state, INPUT_KEY)); - - try { - log.info("开始增强关键词提取处理..."); - - Map resultMap = new HashMap<>(); - - AtomicReference> expandedQuestionsRef = new AtomicReference<>(null); - Sinks.Many sink = Sinks.many().multicast().onBackpressureBuffer(); - Flux displayFlux = FluxUtil - .cascadeFlux(queryProcessingService.expandQuestion(input, expandedQuestions -> { - log.info("问题扩展结果: {}", expandedQuestions); - expandedQuestionsRef.set(expandedQuestions); - }), r -> processMultipleQuestions(expandedQuestionsRef.get(), StateUtil.getStringValue(state, AGENT_ID), - extractionResults -> { - List mergedKeywords = mergeKeywords(extractionResults, input); - List mergedEvidences = mergeEvidences(extractionResults); - - log.info("[{}] 增强提取结果 - 证据: {}, 关键词: {}", this.getClass().getSimpleName(), mergedEvidences, - mergedKeywords); - resultMap.putAll(Map.of(KEYWORD_EXTRACT_NODE_OUTPUT, mergedKeywords, EVIDENCES, - mergedEvidences, RESULT, mergedKeywords)); - sink.tryEmitNext(ChatResponseUtil.createPureResponse(TextType.JSON.getEndSign())); - sink.tryEmitNext(ChatResponseUtil - .createResponse("\n\n合并后的证据: " + String.join(", ", mergedEvidences))); - sink.tryEmitNext( - ChatResponseUtil.createResponse("合并后的关键词: " + String.join(", ", mergedKeywords))); - sink.tryEmitNext(ChatResponseUtil.createResponse("关键词提取完成.")); - sink.tryEmitComplete(); - }), flux -> Mono.just(""), - Flux.just(ChatResponseUtil.createResponse("开始增强关键词提取..."), - ChatResponseUtil.createResponse("正在扩展问题理解...\n"), - ChatResponseUtil.createPureResponse(TextType.JSON.getStartSign())), - Flux.defer(() -> Flux.just(ChatResponseUtil.createPureResponse(TextType.JSON.getEndSign()), - ChatResponseUtil.createResponse("\n\n问题扩展结束！"), - ChatResponseUtil.createResponse("合并多个问题变体的结果...\n"), - ChatResponseUtil.createPureResponse(TextType.JSON.getStartSign()))), - sink.asFlux()); - - Flux> generator = FluxUtil - .createStreamingGeneratorWithMessages(this.getClass(), state, v -> resultMap, displayFlux); - - return Map.of(KEYWORD_EXTRACT_NODE_OUTPUT, generator); - - } - catch (Exception e) { - log.error("增强关键词提取失败{}", e.getMessage()); - throw new RuntimeException(e); - } - } - -} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PlanExecutorNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PlanExecutorNode.java index f1a684583..0799ad8db 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PlanExecutorNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PlanExecutorNode.java @@ -18,12 +18,12 @@ import com.alibaba.cloud.ai.graph.OverAllState; import com.alibaba.cloud.ai.graph.StateGraph; +import com.alibaba.cloud.ai.graph.action.NodeAction; import com.alibaba.cloud.ai.pojo.ExecutionStep; import com.alibaba.cloud.ai.pojo.Plan; +import com.alibaba.cloud.ai.util.PlanProcessUtil; import com.alibaba.cloud.ai.util.StateUtil; import lombok.extern.slf4j.Slf4j; -import org.springframework.ai.converter.BeanOutputConverter; -import org.springframework.core.ParameterizedTypeReference; import org.springframework.stereotype.Component; import java.util.List; @@ -34,7 +34,6 @@ import static com.alibaba.cloud.ai.constant.Constant.IS_ONLY_NL2SQL; import static com.alibaba.cloud.ai.constant.Constant.HUMAN_REVIEW_ENABLED; import static com.alibaba.cloud.ai.constant.Constant.ONLY_NL2SQL_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.PLANNER_NODE_OUTPUT; import static com.alibaba.cloud.ai.constant.Constant.PLAN_CURRENT_STEP; import static com.alibaba.cloud.ai.constant.Constant.PLAN_NEXT_NODE; import static com.alibaba.cloud.ai.constant.Constant.PLAN_REPAIR_COUNT; @@ -52,29 +51,19 @@ */ @Slf4j @Component -public class PlanExecutorNode extends AbstractPlanBasedNode { +public class PlanExecutorNode implements NodeAction { // Supported node types private static final Set SUPPORTED_NODES = Set.of(SQL_EXECUTE_NODE, PYTHON_GENERATE_NODE, REPORT_GENERATOR_NODE); - private final BeanOutputConverter converter; - - public PlanExecutorNode() { - this.converter = new BeanOutputConverter<>(new ParameterizedTypeReference<>() { - }); - } - @Override public Map apply(OverAllState state) throws Exception { - logNodeEntry(); - // 1. Validate the Plan - String plannerOutput = StateUtil.getStringValue(state, PLANNER_NODE_OUTPUT); try { - Plan plan = converter.convert(plannerOutput); + Plan plan = PlanProcessUtil.getPlan(state); - if (plan == null || plan.getExecutionPlan() == null || plan.getExecutionPlan().isEmpty()) { + if (plan.getExecutionPlan() == null || plan.getExecutionPlan().isEmpty()) { return buildValidationResult(state, false, "Validation failed: The generated plan is empty or has no execution steps."); } @@ -115,8 +104,8 @@ public Map apply(OverAllState state) throws Exception { return Map.of(PLAN_VALIDATION_STATUS, true, PLAN_NEXT_NODE, HUMAN_FEEDBACK_NODE); } - Plan plan = getPlan(state); - Integer currentStep = getCurrentStepNumber(state); + Plan plan = PlanProcessUtil.getPlan(state); + int currentStep = PlanProcessUtil.getCurrentStepNumber(state); List executionPlan = plan.getExecutionPlan(); // Check if the plan is completed diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PlannerNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PlannerNode.java index 5349cabb9..c9dbd17b9 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PlannerNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PlannerNode.java @@ -36,14 +36,7 @@ import java.util.Map; -import static com.alibaba.cloud.ai.constant.Constant.BUSINESS_KNOWLEDGE; -import static com.alibaba.cloud.ai.constant.Constant.INPUT_KEY; -import static com.alibaba.cloud.ai.constant.Constant.IS_ONLY_NL2SQL; -import static com.alibaba.cloud.ai.constant.Constant.PLANNER_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.PLAN_VALIDATION_ERROR; -import static com.alibaba.cloud.ai.constant.Constant.QUERY_REWRITE_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.SEMANTIC_MODEL; -import static com.alibaba.cloud.ai.constant.Constant.TABLE_RELATION_OUTPUT; +import static com.alibaba.cloud.ai.constant.Constant.*; /** * @author zhangshenghang @@ -57,11 +50,9 @@ public class PlannerNode implements NodeAction { @Override public Map apply(OverAllState state) throws Exception { - String input = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT, - StateUtil.getStringValue(state, INPUT_KEY)); - // 使用经过时间表达式处理的重写查询，如果没有则回退到原始输入 - String processedQuery = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT, input); - log.info("Using processed query for planning: {}", processedQuery); + // 获取查询增强节点的输出 + String canonicalQuery = StateUtil.getCanonicalQuery(state); + log.info("Using processed query for planning: {}", canonicalQuery); // 是否为NL2SQL模式 Boolean onlyNl2sql = state.value(IS_ONLY_NL2SQL, false); @@ -76,24 +67,23 @@ public Map apply(OverAllState state) throws Exception { } // 构建提示参数 - String businessKnowledge = (String) state.value(BUSINESS_KNOWLEDGE).orElse(""); - String semanticModel = (String) state.value(SEMANTIC_MODEL).orElse(""); + String semanticModel = (String) state.value(GENEGRATED_SEMANTIC_MODEL_PROMPT).orElse(""); SchemaDTO schemaDTO = StateUtil.getObjectValue(state, TABLE_RELATION_OUTPUT, SchemaDTO.class); String schemaStr = PromptHelper.buildMixMacSqlDbPrompt(schemaDTO, true); // 构建用户提示 - String userPrompt = buildUserPrompt(processedQuery, validationError, state); + String userPrompt = buildUserPrompt(canonicalQuery, validationError, state); + String evidence = StateUtil.getStringValue(state, EVIDENCE); // 构建模板参数 - Map params = Map.of("user_question", userPrompt, "schema", schemaStr, "business_knowledge", - businessKnowledge, "semantic_model", semanticModel, "plan_validation_error", - formatValidationError(validationError)); + Map params = Map.of("user_question", userPrompt, "schema", schemaStr, "evidence", evidence, + "semantic_model", semanticModel, "plan_validation_error", formatValidationError(validationError)); // 生成计划 String plannerPrompt = (onlyNl2sql ? PromptConstant.getPlannerNl2sqlOnlyTemplate() : PromptConstant.getPlannerPromptTemplate()) .render(params); - + log.debug("Planner prompt: as follows \n{}\n", plannerPrompt); Flux chatResponseFlux = Flux.concat( Flux.just(ChatResponseUtil.createPureResponse(TextType.JSON.getStartSign())), llmService.callUser(plannerPrompt), diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonAnalyzeNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonAnalyzeNode.java index fb916b405..8fd436443 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonAnalyzeNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonAnalyzeNode.java @@ -23,8 +23,10 @@ import com.alibaba.cloud.ai.prompt.PromptConstant; import com.alibaba.cloud.ai.service.llm.LlmService; import com.alibaba.cloud.ai.util.FluxUtil; +import com.alibaba.cloud.ai.util.PlanProcessUtil; import com.alibaba.cloud.ai.util.StateUtil; -import com.alibaba.cloud.ai.util.StepResultUtil; +import lombok.AllArgsConstructor; +import lombok.extern.slf4j.Slf4j; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.stereotype.Component; import reactor.core.publisher.Flux; @@ -32,11 +34,7 @@ import java.util.HashMap; import java.util.Map; -import static com.alibaba.cloud.ai.constant.Constant.PLAN_CURRENT_STEP; -import static com.alibaba.cloud.ai.constant.Constant.PYTHON_ANALYSIS_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.PYTHON_EXECUTE_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.QUERY_REWRITE_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.SQL_EXECUTE_NODE_OUTPUT; +import static com.alibaba.cloud.ai.constant.Constant.*; /** * 根据Python代码的运行结果做总结分析 @@ -44,24 +42,20 @@ * @author vlsmb * @since 2025/7/30 */ +@Slf4j @Component -public class PythonAnalyzeNode extends AbstractPlanBasedNode implements NodeAction { +@AllArgsConstructor +public class PythonAnalyzeNode implements NodeAction { private final LlmService llmService; - public PythonAnalyzeNode(LlmService llmService) { - super(); - this.llmService = llmService; - } - @Override public Map apply(OverAllState state) throws Exception { - this.logNodeEntry(); // Get context - String userQuery = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT); + String userQuery = StateUtil.getCanonicalQuery(state); String pythonOutput = StateUtil.getStringValue(state, PYTHON_EXECUTE_NODE_OUTPUT); - int currentStep = this.getCurrentStepNumber(state); + int currentStep = PlanProcessUtil.getCurrentStepNumber(state); @SuppressWarnings("unchecked") Map sqlExecuteResult = StateUtil.getObjectValue(state, SQL_EXECUTE_NODE_OUTPUT, Map.class, new HashMap<>()); @@ -74,9 +68,9 @@ public Map apply(OverAllState state) throws Exception { Flux> generator = FluxUtil.createStreamingGeneratorWithMessages(this.getClass(), state, "正在分析代码运行结果...\n", "\n结果分析完成。", aiResponse -> { - Map updatedSqlResult = StepResultUtil.addStepResult(sqlExecuteResult, currentStep, + Map updatedSqlResult = PlanProcessUtil.addStepResult(sqlExecuteResult, currentStep, aiResponse); - this.logNodeOutput("python_analysis_result", aiResponse); + log.info("python analyze result: {}", aiResponse); return Map.of(SQL_EXECUTE_NODE_OUTPUT, updatedSqlResult, PLAN_CURRENT_STEP, currentStep + 1); }, pythonAnalyzeFlux); diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonExecuteNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonExecuteNode.java index 5a17c91d8..ddb0fcc27 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonExecuteNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonExecuteNode.java @@ -24,6 +24,7 @@ import com.alibaba.cloud.ai.service.code.CodePoolExecutorService; import com.alibaba.cloud.ai.util.ChatResponseUtil; import com.alibaba.cloud.ai.util.FluxUtil; +import com.alibaba.cloud.ai.util.JsonUtil; import com.alibaba.cloud.ai.util.StateUtil; import com.fasterxml.jackson.databind.ObjectMapper; import lombok.extern.slf4j.Slf4j; @@ -47,21 +48,19 @@ */ @Slf4j @Component -public class PythonExecuteNode extends AbstractPlanBasedNode implements NodeAction { +public class PythonExecuteNode implements NodeAction { private final CodePoolExecutorService codePoolExecutor; private final ObjectMapper objectMapper; public PythonExecuteNode(CodePoolExecutorService codePoolExecutor) { - super(); this.codePoolExecutor = codePoolExecutor; - this.objectMapper = new ObjectMapper(); + this.objectMapper = JsonUtil.getObjectMapper(); } @Override public Map apply(OverAllState state) throws Exception { - this.logNodeEntry(); try { // Get context @@ -98,7 +97,7 @@ public Map apply(OverAllState state) throws Exception { emitter.next(ChatResponseUtil.createResponse("标准输出：")); emitter.next(ChatResponseUtil.createPureResponse(TextType.JSON.getStartSign())); emitter.next(ChatResponseUtil.createResponse(finalStdout)); - emitter.next(ChatResponseUtil.createResponse(TextType.JSON.getEndSign())); + emitter.next(ChatResponseUtil.createPureResponse(TextType.JSON.getEndSign())); emitter.next(ChatResponseUtil.createResponse("Python代码执行成功！")); emitter.complete(); }); diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonGenerateNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonGenerateNode.java index 0210ef27a..ae32c6b01 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonGenerateNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/PythonGenerateNode.java @@ -29,6 +29,7 @@ import com.alibaba.cloud.ai.util.ChatResponseUtil; import com.alibaba.cloud.ai.util.FluxUtil; import com.alibaba.cloud.ai.util.MarkdownParserUtil; +import com.alibaba.cloud.ai.util.PlanProcessUtil; import com.alibaba.cloud.ai.util.StateUtil; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.databind.ObjectMapper; @@ -40,13 +41,7 @@ import java.util.List; import java.util.Map; -import static com.alibaba.cloud.ai.constant.Constant.PYTHON_EXECUTE_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.PYTHON_GENERATE_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.PYTHON_IS_SUCCESS; -import static com.alibaba.cloud.ai.constant.Constant.PYTHON_TRIES_COUNT; -import static com.alibaba.cloud.ai.constant.Constant.QUERY_REWRITE_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.SQL_RESULT_LIST_MEMORY; -import static com.alibaba.cloud.ai.constant.Constant.TABLE_RELATION_OUTPUT; +import static com.alibaba.cloud.ai.constant.Constant.*; /** * 生成Python代码的节点 @@ -56,7 +51,7 @@ */ @Slf4j @Component -public class PythonGenerateNode extends AbstractPlanBasedNode implements NodeAction { +public class PythonGenerateNode implements NodeAction { private static final int SAMPLE_DATA_NUMBER = 5; @@ -69,7 +64,6 @@ public class PythonGenerateNode extends AbstractPlanBasedNode implements NodeAct private final LlmService llmService; public PythonGenerateNode(CodeExecutorProperties codeExecutorProperties, LlmService llmService) { - super(); this.codeExecutorProperties = codeExecutorProperties; this.llmService = llmService; this.objectMapper = new ObjectMapper().setSerializationInclusion(JsonInclude.Include.NON_NULL); @@ -77,7 +71,6 @@ public PythonGenerateNode(CodeExecutorProperties codeExecutorProperties, LlmServ @Override public Map apply(OverAllState state) throws Exception { - this.logNodeEntry(); // Get context SchemaDTO schemaDTO = StateUtil.getObjectValue(state, TABLE_RELATION_OUTPUT, SchemaDTO.class); @@ -85,7 +78,7 @@ public Map apply(OverAllState state) throws Exception { boolean codeRunSuccess = StateUtil.getObjectValue(state, PYTHON_IS_SUCCESS, Boolean.class, true); int triesCount = StateUtil.getObjectValue(state, PYTHON_TRIES_COUNT, Integer.class, MAX_TRIES_COUNT); - String userPrompt = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT); + String userPrompt = StateUtil.getCanonicalQuery(state); if (!codeRunSuccess) { // Last generated Python code failed to run, inform AI model of this // information @@ -104,7 +97,7 @@ public Map apply(OverAllState state) throws Exception { """, lastCode, lastError); } - ExecutionStep executionStep = this.getCurrentExecutionStep(state); + ExecutionStep executionStep = PlanProcessUtil.getCurrentExecutionStep(state); ExecutionStep.ToolParameters toolParameters = executionStep.getToolParameters(); diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/QueryEnhanceNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/QueryEnhanceNode.java new file mode 100644 index 000000000..e52a79107 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/QueryEnhanceNode.java @@ -0,0 +1,99 @@ +/* + * Copyright 2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.alibaba.cloud.ai.node; + +import com.alibaba.cloud.ai.dto.QueryEnhanceOutputDTO; +import com.alibaba.cloud.ai.enums.TextType; +import com.alibaba.cloud.ai.graph.GraphResponse; +import com.alibaba.cloud.ai.graph.OverAllState; +import com.alibaba.cloud.ai.graph.action.NodeAction; +import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; +import com.alibaba.cloud.ai.prompt.PromptHelper; +import com.alibaba.cloud.ai.service.llm.LlmService; +import com.alibaba.cloud.ai.util.ChatResponseUtil; +import com.alibaba.cloud.ai.util.FluxUtil; +import com.alibaba.cloud.ai.util.JsonUtil; +import com.alibaba.cloud.ai.util.StateUtil; +import lombok.AllArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.stereotype.Component; +import reactor.core.publisher.Flux; + +import java.util.Map; + +import static com.alibaba.cloud.ai.constant.Constant.*; + +/** + * 查询处理节点，用于查询改写、分解和关键词提取 + */ +@Slf4j +@Component +@AllArgsConstructor +public class QueryEnhanceNode implements NodeAction { + + private final LlmService llmService; + + @Override + public Map apply(OverAllState state) throws Exception { + + // 获取用户输入 + String userInput = StateUtil.getStringValue(state, INPUT_KEY); + log.info("User input for query enhance: {}", userInput); + + String evidence = StateUtil.getStringValue(state, EVIDENCE); + + // 构建查询处理提示，多轮对话暂时为空 + String prompt = PromptHelper.buildQueryEnhancePrompt(null, userInput, evidence); + log.debug("Built query enhance prompt as follows \n {} \n", prompt); + + // 调用LLM进行查询处理 + Flux responseFlux = llmService.callUser(prompt); + + Flux> generator = FluxUtil.createStreamingGenerator(this.getClass(), state, + responseFlux, + Flux.just(ChatResponseUtil.createResponse("正在进行问题增强..."), + ChatResponseUtil.createPureResponse(TextType.JSON.getStartSign())), + Flux.just(ChatResponseUtil.createPureResponse(TextType.JSON.getEndSign()), + ChatResponseUtil.createResponse("\n问题增强完成！")), + this::handleQueryEnhance); + + return Map.of(QUERY_ENHANCE_NODE_OUTPUT, generator); + } + + private Map handleQueryEnhance(String llmOutput) { + // 获取处理结果 + String enhanceResult = llmOutput.trim(); + log.info("Query enhance result: {}", enhanceResult); + + // 解析处理结果，转成 QueryProcessOutputDTO + QueryEnhanceOutputDTO queryEnhanceOutputDTO = null; + try { + queryEnhanceOutputDTO = JsonUtil.getObjectMapper().readValue(enhanceResult, QueryEnhanceOutputDTO.class); + log.info("Successfully parsed query enhance result: {}", queryEnhanceOutputDTO); + } + catch (Exception e) { + log.error("Failed to parse query enhance result: {}", enhanceResult, e); + } + + if (queryEnhanceOutputDTO == null) + return Map.of(); + // 返回处理结果 + return Map.of(QUERY_ENHANCE_NODE_OUTPUT, queryEnhanceOutputDTO); + } + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/QueryRewriteNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/QueryRewriteNode.java deleted file mode 100644 index 3e6ef7bf7..000000000 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/QueryRewriteNode.java +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright 2025 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.alibaba.cloud.ai.node; - -import com.alibaba.cloud.ai.graph.GraphResponse; -import com.alibaba.cloud.ai.graph.OverAllState; -import com.alibaba.cloud.ai.graph.action.NodeAction; -import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; -import com.alibaba.cloud.ai.service.processing.QueryProcessingService; -import com.alibaba.cloud.ai.util.FluxUtil; -import com.alibaba.cloud.ai.util.StateUtil; -import lombok.AllArgsConstructor; -import lombok.extern.slf4j.Slf4j; -import org.springframework.stereotype.Component; -import reactor.core.publisher.Flux; - -import java.util.Map; - -import static com.alibaba.cloud.ai.constant.Constant.AGENT_ID; -import static com.alibaba.cloud.ai.constant.Constant.INPUT_KEY; -import static com.alibaba.cloud.ai.constant.Constant.QUERY_REWRITE_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.RESULT; - -/** - * Query rewriting and intent clarification node to improve intent understanding accuracy. - * - * This node is responsible for: - Rewriting user queries to clarify intent - Improving - * understanding accuracy through query transformation - Providing streaming feedback - * during rewriting process - * - * @author zhangshenghang - */ -@Slf4j -@Component -@AllArgsConstructor -public class QueryRewriteNode implements NodeAction { - - private final QueryProcessingService queryProcessingService; - - @Override - public Map apply(OverAllState state) throws Exception { - log.info("Entering {} node", this.getClass().getSimpleName()); - - String input = StateUtil.getStringValue(state, INPUT_KEY); - String agentId = StateUtil.getStringValue(state, AGENT_ID); // Get agent ID - log.info("[{}] Processing user input: {} for agentId: {}", this.getClass().getSimpleName(), input, agentId); - - final StringBuilder queryResultCollector = new StringBuilder(); - // Use streaming utility class for content collection and result mapping - Flux> generator = FluxUtil.createStreamingGeneratorWithMessages(this.getClass(), - state, "开始进行问题重写...", "\n\n问题重写完成！", finalResult -> { - finalResult = queryResultCollector.toString(); - return Map.of(QUERY_REWRITE_NODE_OUTPUT, finalResult, RESULT, finalResult); - }, queryProcessingService.rewriteStream(input, agentId, queryResultCollector)); - - return Map.of(QUERY_REWRITE_NODE_OUTPUT, generator); - } - -} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/ReportGeneratorNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/ReportGeneratorNode.java index f2050b522..0f978f4e5 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/ReportGeneratorNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/ReportGeneratorNode.java @@ -41,12 +41,7 @@ import java.util.List; import java.util.Map; -import static com.alibaba.cloud.ai.constant.Constant.INPUT_KEY; -import static com.alibaba.cloud.ai.constant.Constant.PLANNER_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.PLAN_CURRENT_STEP; -import static com.alibaba.cloud.ai.constant.Constant.QUERY_REWRITE_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.RESULT; -import static com.alibaba.cloud.ai.constant.Constant.SQL_EXECUTE_NODE_OUTPUT; +import static com.alibaba.cloud.ai.constant.Constant.*; /** * Report generation node that creates comprehensive analysis reports based on execution @@ -77,12 +72,10 @@ public ReportGeneratorNode(LlmService llmService, UserPromptService promptConfig @Override public Map apply(OverAllState state) throws Exception { - log.info("Entering {} node", this.getClass().getSimpleName()); // Get necessary input parameters String plannerNodeOutput = StateUtil.getStringValue(state, PLANNER_NODE_OUTPUT); - String userInput = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT, - StateUtil.getStringValue(state, INPUT_KEY)); + String userInput = StateUtil.getCanonicalQuery(state); Integer currentStep = StateUtil.getObjectValue(state, PLAN_CURRENT_STEP, Integer.class, 1); @SuppressWarnings("unchecked") HashMap executionResults = StateUtil.getObjectValue(state, SQL_EXECUTE_NODE_OUTPUT, diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SchemaRecallNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SchemaRecallNode.java index 46eff0760..ce865e3a1 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SchemaRecallNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SchemaRecallNode.java @@ -16,6 +16,7 @@ package com.alibaba.cloud.ai.node; +import com.alibaba.cloud.ai.dto.QueryEnhanceOutputDTO; import com.alibaba.cloud.ai.graph.GraphResponse; import com.alibaba.cloud.ai.graph.OverAllState; import com.alibaba.cloud.ai.graph.action.NodeAction; @@ -35,13 +36,7 @@ import java.util.List; import java.util.Map; -import static com.alibaba.cloud.ai.constant.Constant.COLUMN_DOCUMENTS_BY_KEYWORDS_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.AGENT_ID; -import static com.alibaba.cloud.ai.constant.Constant.INPUT_KEY; -import static com.alibaba.cloud.ai.constant.Constant.KEYWORD_EXTRACT_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.QUERY_REWRITE_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.SCHEMA_RECALL_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT; +import static com.alibaba.cloud.ai.constant.Constant.*; /** * Schema recall node that retrieves relevant database schema information based on @@ -62,50 +57,68 @@ public class SchemaRecallNode implements NodeAction { @Override public Map apply(OverAllState state) throws Exception { - log.info("Entering {} node", this.getClass().getSimpleName()); - String input = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT, - StateUtil.getStringValue(state, INPUT_KEY)); - List keywords = StateUtil.getListValue(state, KEYWORD_EXTRACT_NODE_OUTPUT); + // get input information + QueryEnhanceOutputDTO queryEnhanceOutputDTO = StateUtil.getObjectValue(state, QUERY_ENHANCE_NODE_OUTPUT, + QueryEnhanceOutputDTO.class); + String input = queryEnhanceOutputDTO.getCanonicalQuery(); + List keywords = queryEnhanceOutputDTO.getConsolidatedKeywords(); String agentId = StateUtil.getStringValue(state, AGENT_ID); // Execute business logic first - recall schema information immediately - List tableDocuments; - List> columnDocumentsByKeywords; - - // If agentId exists, use agent-specific search, otherwise use global search - if (agentId != null && !agentId.trim().isEmpty()) { - log.info("Using agent-specific schema recall for agent: {}", agentId); - tableDocuments = new ArrayList<>(schemaService.getTableDocumentsForAgent(agentId, input)); - columnDocumentsByKeywords = schemaService.getColumnDocumentsByKeywordsForAgent(agentId, keywords); - } - else { - log.info("Using global schema recall (no agentId provided)"); - tableDocuments = schemaService.getTableDocuments(input); - columnDocumentsByKeywords = schemaService.getColumnDocumentsByKeywords(keywords); - } + List tableDocuments = new ArrayList<>(schemaService.getTableDocumentsForAgent(agentId, input)); + List columnDocumentsByKeywords = schemaService.getColumnDocumentsByKeywordsForAgent(agentId, keywords, + tableDocuments); - log.info("[{}] Schema recall results - table documents count: {}, keyword-related column document groups: {}", - this.getClass().getSimpleName(), tableDocuments.size(), columnDocumentsByKeywords.size()); + // extract table names + List recallTableNames = extractTableName(tableDocuments); + // extract column names + List recallColumnNames = extractColumnName(columnDocumentsByKeywords); Flux displayFlux = Flux.create(emitter -> { - emitter.next(ChatResponseUtil.createResponse("开始召回Schema信息...")); - emitter.next(ChatResponseUtil.createResponse("表信息召回完成，数量: " + tableDocuments.size())); - emitter.next(ChatResponseUtil.createResponse("列信息召回完成，数量: " + columnDocumentsByKeywords.size())); - emitter.next(ChatResponseUtil.createResponse("Schema信息召回完成.")); + emitter.next(ChatResponseUtil.createResponse("开始初步召回Schema信息...")); + emitter.next(ChatResponseUtil.createResponse( + "初步表信息召回完成，数量: " + tableDocuments.size() + "，表名: " + String.join(", ", recallTableNames))); + emitter.next(ChatResponseUtil.createResponse("初步列信息召回完成，数量: " + columnDocumentsByKeywords.size() + "，列名: " + + String.join(", ", recallColumnNames))); + emitter.next(ChatResponseUtil.createResponse("初步Schema信息召回完成.")); emitter.complete(); }); Flux> generator = FluxUtil.createStreamingGeneratorWithMessages(this.getClass(), state, currentState -> { - log.info("Table document details: {}", tableDocuments); - log.info("Keyword-related column document details: {}", columnDocumentsByKeywords); return Map.of(TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT, tableDocuments, - COLUMN_DOCUMENTS_BY_KEYWORDS_OUTPUT, columnDocumentsByKeywords); + COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT, columnDocumentsByKeywords); }, displayFlux); // Return the processing result return Map.of(SCHEMA_RECALL_NODE_OUTPUT, generator); } + private static List extractTableName(List tableDocuments) { + List tableNames = new ArrayList<>(); + // metadata中的name字段 + for (Document document : tableDocuments) { + String name = (String) document.getMetadata().get("name"); + if (name != null && !name.isEmpty()) + tableNames.add(name); + } + log.info("At this SchemaRecallNode, Recall tables are: {}", tableNames); + return tableNames; + + } + + private static List extractColumnName(List columnDocuments) { + // metadata中的 tableName + . + name字段 + List columnNames = new ArrayList<>(); + for (Document document : columnDocuments) { + String name = (String) document.getMetadata().get("name"); + String tableName = (String) document.getMetadata().get("tableName"); + if (name != null && !name.isEmpty() && tableName != null && !tableName.isEmpty()) + columnNames.add(tableName + "." + name); + } + log.info("At this SchemaRecallNode, Recall columns are: {}", columnNames); + return columnNames; + } + } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SemanticConsistencyNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SemanticConsistencyNode.java index 4ca972bfc..e7662b5d2 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SemanticConsistencyNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SemanticConsistencyNode.java @@ -19,26 +19,23 @@ import com.alibaba.cloud.ai.dto.schema.SchemaDTO; import com.alibaba.cloud.ai.graph.GraphResponse; import com.alibaba.cloud.ai.graph.OverAllState; +import com.alibaba.cloud.ai.graph.action.NodeAction; import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; import com.alibaba.cloud.ai.pojo.ExecutionStep; import com.alibaba.cloud.ai.prompt.PromptHelper; import com.alibaba.cloud.ai.service.nl2sql.Nl2SqlService; import com.alibaba.cloud.ai.util.FluxUtil; +import com.alibaba.cloud.ai.util.PlanProcessUtil; import com.alibaba.cloud.ai.util.StateUtil; +import lombok.AllArgsConstructor; import lombok.extern.slf4j.Slf4j; -import org.apache.commons.lang3.StringUtils; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.stereotype.Component; import reactor.core.publisher.Flux; -import java.util.List; import java.util.Map; -import static com.alibaba.cloud.ai.constant.Constant.EVIDENCES; -import static com.alibaba.cloud.ai.constant.Constant.PLAN_CURRENT_STEP; -import static com.alibaba.cloud.ai.constant.Constant.SEMANTIC_CONSISTENCY_NODE_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.SEMANTIC_CONSISTENCY_NODE_RECOMMEND_OUTPUT; -import static com.alibaba.cloud.ai.constant.Constant.TABLE_RELATION_OUTPUT; +import static com.alibaba.cloud.ai.constant.Constant.*; /** * Semantic consistency validation node that checks SQL query semantic consistency. @@ -51,34 +48,29 @@ */ @Slf4j @Component -public class SemanticConsistencyNode extends AbstractPlanBasedNode { +@AllArgsConstructor +public class SemanticConsistencyNode implements NodeAction { private final Nl2SqlService nl2SqlService; - public SemanticConsistencyNode(Nl2SqlService nl2SqlService) { - super(); - this.nl2SqlService = nl2SqlService; - } - @Override public Map apply(OverAllState state) throws Exception { - logNodeEntry(); // Get necessary input parameters - List evidenceList = StateUtil.getListValue(state, EVIDENCES); + String evidence = StateUtil.getStringValue(state, EVIDENCE); SchemaDTO schemaDTO = StateUtil.getObjectValue(state, TABLE_RELATION_OUTPUT, SchemaDTO.class); // Get current execution step and SQL query - ExecutionStep executionStep = getCurrentExecutionStep(state); - Integer currentStep = getCurrentStepNumber(state); + ExecutionStep executionStep = PlanProcessUtil.getCurrentExecutionStep(state); + Integer currentStep = PlanProcessUtil.getCurrentStepNumber(state); ExecutionStep.ToolParameters toolParameters = executionStep.getToolParameters(); String sqlQuery = toolParameters.getSqlQuery(); log.info("Starting semantic consistency validation - SQL: {}", sqlQuery); log.info("Step description: {}", toolParameters.getDescription()); - Flux validationResultFlux = performSemanticValidationStream(schemaDTO, evidenceList, - toolParameters, sqlQuery); + Flux validationResultFlux = performSemanticValidationStream(schemaDTO, evidence, toolParameters, + sqlQuery); Flux> generator = FluxUtil.createStreamingGeneratorWithMessages(this.getClass(), state, "开始语义一致性校验", "语义一致性校验完成", validationResult -> { @@ -95,11 +87,10 @@ public Map apply(OverAllState state) throws Exception { /** * Perform streaming semantic consistency validation */ - private Flux performSemanticValidationStream(SchemaDTO schemaDTO, List evidenceList, + private Flux performSemanticValidationStream(SchemaDTO schemaDTO, String evidence, ExecutionStep.ToolParameters toolParameters, String sqlQuery) throws Exception { // Build validation context String schema = PromptHelper.buildMixMacSqlDbPrompt(schemaDTO, true); - String evidence = StringUtils.join(evidenceList, ";\n"); String context = String.join("\n", schema, evidence, toolParameters.getDescription()); // Execute semantic consistency check diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SqlExecuteNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SqlExecuteNode.java index f08c40841..840a7893b 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SqlExecuteNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SqlExecuteNode.java @@ -26,6 +26,7 @@ import com.alibaba.cloud.ai.enums.TextType; import com.alibaba.cloud.ai.graph.GraphResponse; import com.alibaba.cloud.ai.graph.OverAllState; +import com.alibaba.cloud.ai.graph.action.NodeAction; import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; import com.alibaba.cloud.ai.pojo.ExecutionStep; import com.alibaba.cloud.ai.service.datasource.DatasourceService; @@ -33,8 +34,8 @@ import com.alibaba.cloud.ai.entity.Datasource; import com.alibaba.cloud.ai.util.ChatResponseUtil; import com.alibaba.cloud.ai.util.FluxUtil; +import com.alibaba.cloud.ai.util.PlanProcessUtil; import com.alibaba.cloud.ai.util.StateUtil; -import com.alibaba.cloud.ai.util.StepResultUtil; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.stereotype.Component; @@ -58,7 +59,7 @@ */ @Slf4j @Component -public class SqlExecuteNode extends AbstractPlanBasedNode { +public class SqlExecuteNode implements NodeAction { // todo: 根据数据库配置动态获取Accessor private final Accessor dbAccessor; @@ -68,7 +69,6 @@ public class SqlExecuteNode extends AbstractPlanBasedNode { private final DbConfig dbConfig; public SqlExecuteNode(AccessorFactory accessorFactory, DatasourceService datasourceService, DbConfig dbConfig) { - super(); this.dbAccessor = accessorFactory.getAccessorByDbConfig(dbConfig); this.datasourceService = datasourceService; this.dbConfig = dbConfig; @@ -76,10 +76,9 @@ public SqlExecuteNode(AccessorFactory accessorFactory, DatasourceService datasou @Override public Map apply(OverAllState state) throws Exception { - logNodeEntry(); - ExecutionStep executionStep = getCurrentExecutionStep(state); - Integer currentStep = getCurrentStepNumber(state); + ExecutionStep executionStep = PlanProcessUtil.getCurrentExecutionStep(state); + Integer currentStep = PlanProcessUtil.getCurrentStepNumber(state); ExecutionStep.ToolParameters toolParameters = executionStep.getToolParameters(); String sqlQuery = toolParameters.getSqlQuery(); @@ -199,7 +198,7 @@ private Map executeSqlQuery(OverAllState state, Integer currentS // Update step results with the query output Map existingResults = StateUtil.getObjectValue(state, SQL_EXECUTE_NODE_OUTPUT, Map.class, new HashMap<>()); - Map updatedResults = StepResultUtil.addStepResult(existingResults, currentStep, jsonStr); + Map updatedResults = PlanProcessUtil.addStepResult(existingResults, currentStep, jsonStr); log.info("SQL execution successful, result count: {}", resultSetBO.getData() != null ? resultSetBO.getData().size() : 0); diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SqlGenerateNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SqlGenerateNode.java index 0e5141a2c..ebf0286e5 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SqlGenerateNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/SqlGenerateNode.java @@ -27,18 +27,17 @@ import com.alibaba.cloud.ai.service.nl2sql.Nl2SqlService; import com.alibaba.cloud.ai.util.ChatResponseUtil; import com.alibaba.cloud.ai.util.FluxUtil; +import com.alibaba.cloud.ai.util.PlanProcessUtil; import com.alibaba.cloud.ai.util.StateUtil; import com.google.common.util.concurrent.AtomicDouble; +import lombok.AllArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.chat.model.ChatResponse; -import org.springframework.ai.converter.BeanOutputConverter; -import org.springframework.core.ParameterizedTypeReference; import org.springframework.stereotype.Component; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import java.util.HashMap; -import java.util.List; import java.util.Map; import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; @@ -61,31 +60,19 @@ */ @Slf4j @Component +@AllArgsConstructor public class SqlGenerateNode implements NodeAction { private static final int MAX_OPTIMIZATION_ROUNDS = 3; private final Nl2SqlService nl2SqlService; - private final BeanOutputConverter converter; - - public SqlGenerateNode(Nl2SqlService nl2SqlService) { - this.nl2SqlService = nl2SqlService; - this.converter = new BeanOutputConverter<>(new ParameterizedTypeReference<>() { - }); - } - @Override public Map apply(OverAllState state) throws Exception { - log.info("Entering {} node", this.getClass().getSimpleName()); // Get necessary input parameters - String plannerNodeOutput = StateUtil.getStringValue(state, PLANNER_NODE_OUTPUT); - Plan plan = converter.convert(plannerNodeOutput); - Integer currentStep = StateUtil.getObjectValue(state, PLAN_CURRENT_STEP, Integer.class, 1); - - List executionPlan = Optional.ofNullable(plan).orElseThrow().getExecutionPlan(); - ExecutionStep executionStep = executionPlan.get(currentStep - 1); + Plan plan = PlanProcessUtil.getPlan(state); + ExecutionStep executionStep = PlanProcessUtil.getCurrentExecutionStep(state); ExecutionStep.ToolParameters toolParameters = executionStep.getToolParameters(); // Execute business logic first - determine what needs to be regenerated @@ -133,11 +120,11 @@ private Flux handleSqlExecutionException(OverAllState state, ExecutionSt String sqlException = StateUtil.getStringValue(state, SQL_EXECUTE_NODE_EXCEPTION_OUTPUT); log.info("Detected SQL execution exception, starting to regenerate SQL: {}", sqlException); - List evidenceList = StateUtil.getListValue(state, EVIDENCES); + String evidence = StateUtil.getStringValue(state, EVIDENCE); SchemaDTO schemaDTO = StateUtil.getObjectValue(state, TABLE_RELATION_OUTPUT, SchemaDTO.class); - return regenerateSql(state, toolParameters.toJsonStr(), evidenceList, schemaDTO, - SQL_EXECUTE_NODE_EXCEPTION_OUTPUT, toolParameters.getSqlQuery(), finalSqlConsumer); + return regenerateSql(state, toolParameters.toJsonStr(), evidence, schemaDTO, SQL_EXECUTE_NODE_EXCEPTION_OUTPUT, + toolParameters.getSqlQuery(), finalSqlConsumer); } /** @@ -147,10 +134,10 @@ private Flux handleSemanticConsistencyFailure(OverAllState state, ExecutionStep.ToolParameters toolParameters, Consumer finalSqlConsumer) { log.info("Semantic consistency validation failed, starting to regenerate SQL"); - List evidenceList = StateUtil.getListValue(state, EVIDENCES); + String evidence = StateUtil.getStringValue(state, EVIDENCE); SchemaDTO schemaDTO = StateUtil.getObjectValue(state, TABLE_RELATION_OUTPUT, SchemaDTO.class); - return regenerateSql(state, toolParameters.toJsonStr(), evidenceList, schemaDTO, + return regenerateSql(state, toolParameters.toJsonStr(), evidence, schemaDTO, SEMANTIC_CONSISTENCY_NODE_RECOMMEND_OUTPUT, toolParameters.getSqlQuery(), finalSqlConsumer); } @@ -167,7 +154,7 @@ private boolean isSemanticConsistencyFailed(OverAllState state) { * @param finalSqlConsumer 处理最终生成SQL的消费者，如果失败，则sql为null * @return AI中间输出过程的Flux */ - private Flux regenerateSql(OverAllState state, String input, List evidenceList, SchemaDTO schemaDTO, + private Flux regenerateSql(OverAllState state, String input, String evidence, SchemaDTO schemaDTO, String exceptionOutputKey, String originalSql, Consumer finalSqlConsumer) { String exceptionMessage = StateUtil.getStringValue(state, exceptionOutputKey); log.info("开始增强SQL生成流程 - 原始SQL: {}, 异常信息: {}", originalSql, exceptionMessage); @@ -239,7 +226,7 @@ public Flux get() { } }; - Flux sqlFlux = nl2SqlService.generateSql(evidenceList, input, schemaDTO, originalSql, exceptionMessage); + Flux sqlFlux = nl2SqlService.generateSql(evidence, input, schemaDTO, originalSql, exceptionMessage); Mono sqlMono = sqlFlux.collect(StringBuilder::new, StringBuilder::append).map(StringBuilder::toString); return Flux.just("正在生成SQL...\n", TextType.SQL.getStartSign()) .concatWith(sqlMono.flatMapMany(sql -> Flux.just(nl2SqlService.sqlTrim(sql)).expand(newSql -> { diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/TableRelationNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/TableRelationNode.java index 0b6f8bc15..2629ee783 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/TableRelationNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/TableRelationNode.java @@ -17,17 +17,17 @@ package com.alibaba.cloud.ai.node; import com.alibaba.cloud.ai.connector.config.DbConfig; -import com.alibaba.cloud.ai.dto.BusinessKnowledgeDTO; +import com.alibaba.cloud.ai.dto.schema.SchemaDTO; +import com.alibaba.cloud.ai.dto.schema.TableDTO; import com.alibaba.cloud.ai.entity.Datasource; import com.alibaba.cloud.ai.entity.SemanticModel; import com.alibaba.cloud.ai.enums.TextType; import com.alibaba.cloud.ai.graph.GraphResponse; import com.alibaba.cloud.ai.graph.OverAllState; import com.alibaba.cloud.ai.graph.action.NodeAction; -import com.alibaba.cloud.ai.dto.schema.SchemaDTO; import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; -import com.alibaba.cloud.ai.service.datasource.DatasourceService; import com.alibaba.cloud.ai.service.business.BusinessKnowledgeService; +import com.alibaba.cloud.ai.service.datasource.DatasourceService; import com.alibaba.cloud.ai.service.nl2sql.Nl2SqlService; import com.alibaba.cloud.ai.service.schema.SchemaService; import com.alibaba.cloud.ai.service.semantic.SemanticModelService; @@ -41,7 +41,6 @@ import org.springframework.ai.document.Document; import org.springframework.dao.DataAccessException; import org.springframework.stereotype.Component; -import org.springframework.util.StringUtils; import reactor.core.publisher.Flux; import java.util.HashMap; @@ -50,7 +49,6 @@ import java.util.function.Consumer; import static com.alibaba.cloud.ai.constant.Constant.*; -import static com.alibaba.cloud.ai.prompt.PromptHelper.buildBusinessKnowledgePrompt; import static com.alibaba.cloud.ai.prompt.PromptHelper.buildSemanticModelPrompt; /** @@ -80,52 +78,37 @@ public class TableRelationNode implements NodeAction { @Override public Map apply(OverAllState state) throws Exception { - log.info("Entering {} node", this.getClass().getSimpleName()); - - int retryCount = StateUtil.getObjectValue(state, TABLE_RELATION_RETRY_COUNT, Integer.class, 0); // Get necessary input parameters - String input = StateUtil.getStringValue(state, QUERY_REWRITE_NODE_OUTPUT, - StateUtil.getStringValue(state, INPUT_KEY)); - List evidenceList = StateUtil.getListValue(state, EVIDENCES); + String canonicalQuery = StateUtil.getCanonicalQuery(state); + + String evidence = StateUtil.getStringValue(state, EVIDENCE); List tableDocuments = StateUtil.getDocumentList(state, TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT); - List> columnDocumentsByKeywords = StateUtil.getDocumentListList(state, - COLUMN_DOCUMENTS_BY_KEYWORDS_OUTPUT); + List columnDocuments = StateUtil.getDocumentList(state, COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT); String agentIdStr = StateUtil.getStringValue(state, AGENT_ID); - if (!StringUtils.hasText(agentIdStr)) - throw new RuntimeException("Agent ID is empty."); // Execute business logic first - get final result immediately DbConfig agentDbConfig = getAgentDbConfig(Integer.valueOf(agentIdStr)); - SchemaDTO schemaDTO = buildInitialSchema(agentIdStr, columnDocumentsByKeywords, tableDocuments, agentDbConfig); - - List businessKnowledges; - List semanticModel = List.of(); - try { - // Extract business knowledge and semantic model - businessKnowledges = businessKnowledgeService.getKnowledgeDtoRecalled(Long.parseLong(agentIdStr)); - // todo: 等待SemanticModelService写完再打开 - // semanticModel = - // semanticModelService.getEnabledByAgentId(Long.parseLong(agentIdStr)); - } - catch (DataAccessException e) { - log.warn("Database query failed (attempt {}): {}", retryCount + 1, e.getMessage()); - - String errorType = classifyDatabaseError(e); - return Map.of(TABLE_RELATION_EXCEPTION_OUTPUT, errorType + ": " + e.getMessage(), - TABLE_RELATION_RETRY_COUNT, retryCount + 1); - } - // load prompt template - String businessKnowledgePrompt = buildBusinessKnowledgePrompt(businessKnowledges); - String semanticModelPrompt = buildSemanticModelPrompt(semanticModel); + SchemaDTO initialSchema = buildInitialSchema(agentIdStr, columnDocuments, tableDocuments, agentDbConfig); - Map resultMap = new HashMap<>( - Map.of(BUSINESS_KNOWLEDGE, businessKnowledgePrompt, SEMANTIC_MODEL, semanticModelPrompt)); + Map resultMap = new HashMap<>(); - Flux schemaFlux = processSchemaSelection(schemaDTO, input, evidenceList, state, agentDbConfig, - result -> { + Flux schemaFlux = processSchemaSelection(initialSchema, canonicalQuery, evidence, state, + agentDbConfig, result -> { log.info("[{}] Schema processing result: {}", this.getClass().getSimpleName(), result); + // 将处理后的SchemaDTO存储到resultMap中 resultMap.put(TABLE_RELATION_OUTPUT, result); + + // 从最终的SchemaDTO中获取表名列表 + List tableNames = result.getTable().stream().map(TableDTO::getName).toList(); + + // 根据agentId和表名列表获取语义模型 + List semanticModels = semanticModelService + .getByAgentIdAndTableNames(Integer.valueOf(agentIdStr), tableNames); + + // 构建语义模型提示并存储到resultMap中 + String semanticModelPrompt = buildSemanticModelPrompt(semanticModels); + resultMap.put(GENEGRATED_SEMANTIC_MODEL_PROMPT, semanticModelPrompt); }); // Create display stream for user experience only @@ -146,8 +129,8 @@ public Map apply(OverAllState state) throws Exception { state, v -> resultMap, displayFlux); // need to reset retry count and exception - return Map.of(TABLE_RELATION_OUTPUT, generator, BUSINESS_KNOWLEDGE, businessKnowledgePrompt, SEMANTIC_MODEL, - semanticModelPrompt, TABLE_RELATION_RETRY_COUNT, 0, TABLE_RELATION_EXCEPTION_OUTPUT, ""); + return Map.of(TABLE_RELATION_OUTPUT, generator, TABLE_RELATION_RETRY_COUNT, 0, TABLE_RELATION_EXCEPTION_OUTPUT, + ""); } @@ -165,12 +148,12 @@ private String classifyDatabaseError(DataAccessException e) { /** * Builds initial schema from column and table documents. */ - private SchemaDTO buildInitialSchema(String agentId, List> columnDocumentsByKeywords, - List tableDocuments, DbConfig agentDbConfig) { + private SchemaDTO buildInitialSchema(String agentId, List columnDocuments, List tableDocuments, + DbConfig agentDbConfig) { SchemaDTO schemaDTO = new SchemaDTO(); schemaService.extractDatabaseName(schemaDTO, agentDbConfig); - schemaService.buildSchemaFromDocuments(agentId, columnDocumentsByKeywords, tableDocuments, schemaDTO); + schemaService.buildSchemaFromDocuments(agentId, columnDocuments, tableDocuments, schemaDTO); return schemaDTO; } @@ -191,7 +174,7 @@ private DbConfig getAgentDbConfig(Integer agentId) { /** * Processes schema selection based on input, evidence, and optional advice. */ - private Flux processSchemaSelection(SchemaDTO schemaDTO, String input, List evidenceList, + private Flux processSchemaSelection(SchemaDTO schemaDTO, String input, String evidence, OverAllState state, DbConfig agentDbConfig, Consumer dtoConsumer) { String schemaAdvice = StateUtil.getStringValue(state, SQL_GENERATE_SCHEMA_MISSING_ADVICE, null); @@ -199,12 +182,11 @@ private Flux processSchemaSelection(SchemaDTO schemaDTO, String in if (schemaAdvice != null) { log.info("[{}] Processing with schema supplement advice: {}", this.getClass().getSimpleName(), schemaAdvice); - schemaFlux = nl2SqlService.fineSelect(schemaDTO, input, evidenceList, schemaAdvice, agentDbConfig, - dtoConsumer); + schemaFlux = nl2SqlService.fineSelect(schemaDTO, input, evidence, schemaAdvice, agentDbConfig, dtoConsumer); } else { log.info("[{}] Executing regular schema selection", this.getClass().getSimpleName()); - schemaFlux = nl2SqlService.fineSelect(schemaDTO, input, evidenceList, null, agentDbConfig, dtoConsumer); + schemaFlux = nl2SqlService.fineSelect(schemaDTO, input, evidence, null, agentDbConfig, dtoConsumer); } return Flux .just(ChatResponseUtil.createResponse("正在选择合适的数据表...\n"), diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/prompt/PromptConstant.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/prompt/PromptConstant.java index 72ed89fdf..ad4b3750f 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/prompt/PromptConstant.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/prompt/PromptConstant.java @@ -30,6 +30,20 @@ public static PromptTemplate getInitRewritePromptTemplate() { return new PromptTemplate(PromptLoader.loadPrompt("init-rewrite")); } + // intent-recognition + public static PromptTemplate getIntentRecognitionPromptTemplate() { + return new PromptTemplate(PromptLoader.loadPrompt("intent-recognition")); + } + + public static PromptTemplate getQueryEnhancementPromptTemplate() { + return new PromptTemplate(PromptLoader.loadPrompt("query-enhancement")); + } + + // feasibility-assessment + public static PromptTemplate getFeasibilityAssessmentPromptTemplate() { + return new PromptTemplate(PromptLoader.loadPrompt("feasibility-assessment")); + } + public static PromptTemplate getQuestionToKeywordsPromptTemplate() { return new PromptTemplate(PromptLoader.loadPrompt("question-to-keywords")); } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/prompt/PromptHelper.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/prompt/PromptHelper.java index a685f989f..e447f9ddf 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/prompt/PromptHelper.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/prompt/PromptHelper.java @@ -15,19 +15,17 @@ */ package com.alibaba.cloud.ai.prompt; -import com.alibaba.cloud.ai.entity.SemanticModel; -import com.alibaba.cloud.ai.enums.BizDataSourceTypeEnum; import com.alibaba.cloud.ai.connector.config.DbConfig; -import com.alibaba.cloud.ai.dto.BusinessKnowledgeDTO; import com.alibaba.cloud.ai.dto.schema.ColumnDTO; import com.alibaba.cloud.ai.dto.schema.SchemaDTO; import com.alibaba.cloud.ai.dto.schema.TableDTO; +import com.alibaba.cloud.ai.entity.SemanticModel; import com.alibaba.cloud.ai.entity.UserPromptConfig; - +import com.alibaba.cloud.ai.enums.BizDataSourceTypeEnum; +import org.apache.commons.collections.CollectionUtils; import org.apache.commons.lang3.BooleanUtils; import org.apache.commons.lang3.StringUtils; import org.springframework.ai.chat.prompt.PromptTemplate; -import org.apache.commons.collections.CollectionUtils; import java.time.LocalDateTime; import java.time.format.DateTimeFormatter; @@ -44,6 +42,7 @@ public static String buildRewritePrompt(String query, SchemaDTO schemaDTO, List< for (TableDTO tableDTO : schemaDTO.getTable()) { dbContent.append(buildMacSqlTablePrompt(tableDTO)).append("\n"); } + // TODO 待完善多轮输入 StringBuilder multiTurn = new StringBuilder(); multiTurn.append("<最新>").append("用户: ").append(query); @@ -88,13 +87,15 @@ public static String buildQueryToKeywordsPrompt(String question) { return PromptConstant.getQuestionToKeywordsPromptTemplate().render(params); } - public static String buildMixSelectorPrompt(List evidences, String question, SchemaDTO schemaDTO) { + public static String buildMixSelectorPrompt(String evidence, String question, SchemaDTO schemaDTO) { String schemaInfo = buildMixMacSqlDbPrompt(schemaDTO, true); Map params = new HashMap<>(); params.put("schema_info", schemaInfo); params.put("question", question); - String evidence = CollectionUtils.isEmpty(evidences) ? "" : StringUtils.join(evidences, ";\n"); - params.put("evidence", evidence); + if (StringUtils.isBlank(evidence)) + params.put("evidence", "无"); + else + params.put("evidence", evidence); return PromptConstant.getMixSelectorPromptTemplate().render(params); } @@ -170,15 +171,7 @@ public static String buildMixMacSqlTablePrompt(TableDTO tableDTO, Boolean withCo List data = new ArrayList<>(enumData.subList(0, Math.min(3, enumData.size()))); line.append(StringUtils.join(data, ",")).append("]"); } - else if (CollectionUtils.isNotEmpty(columnDTO.getSamples())) { - List data = columnDTO.getSamples().subList(0, Math.min(3, columnDTO.getSamples().size())); - data = data.stream().filter(item -> StringUtils.isNotBlank(item)).collect(Collectors.toList()); - if (CollectionUtils.isNotEmpty(data)) { - line.append(", Examples: ["); - data = processSamples(data, columnDTO); - line.append(StringUtils.join(data, ",")).append("]"); - } - } + line.append(")"); columnLines.add(line.toString()); } @@ -212,8 +205,7 @@ private static List processSamples(List samples, ColumnDTO colum } public static List buildMixSqlGeneratorPrompt(String question, DbConfig dbConfig, SchemaDTO schemaDTO, - List evidenceList) { - String evidence = StringUtils.join(evidenceList, ";\n"); + String evidence) { String schemaInfo = buildMixMacSqlDbPrompt(schemaDTO, true); String dialect = BizDataSourceTypeEnum.fromTypeName(dbConfig.getDialectType()).getDialect(); Map params = new HashMap<>(); @@ -271,8 +263,7 @@ public static String buildReportGeneratorPromptWithOptimization(String userRequi } public static String buildSqlErrorFixerPrompt(String question, DbConfig dbConfig, SchemaDTO schemaDTO, - List evidenceList, String errorSql, String errorMessage) { - String evidence = StringUtils.join(evidenceList, ";\n"); + String evidence, String errorSql, String errorMessage) { String schemaInfo = buildMixMacSqlDbPrompt(schemaDTO, true); String dialect = BizDataSourceTypeEnum.fromTypeName(dbConfig.getDialectType()).getDialect(); @@ -287,11 +278,12 @@ public static String buildSqlErrorFixerPrompt(String question, DbConfig dbConfig return PromptConstant.getSqlErrorFixerPromptTemplate().render(params); } - public static String buildBusinessKnowledgePrompt(List businessKnowledgeDTOS) { + public static String buildBusinessKnowledgePrompt(String businessTerms) { Map params = new HashMap<>(); - String businessKnowledge = CollectionUtils.isEmpty(businessKnowledgeDTOS) ? "" - : StringUtils.join(businessKnowledgeDTOS, ";\n"); - params.put("businessKnowledge", businessKnowledge); + if (StringUtils.isNotBlank(businessTerms)) + params.put("businessKnowledge", businessTerms); + else + params.put("businessKnowledge", "无"); return PromptConstant.getBusinessKnowledgePromptTemplate().render(params); } @@ -329,6 +321,56 @@ private static String buildOptimizationSection(List optimizati return result.toString().trim(); } + /** + * 构建意图识别提示词 + * @param multiTurn 多轮对话历史 + * @param latestQuery 最新用户输入 + * @return 意图识别提示词 + */ + public static String buildIntentRecognitionPrompt(String multiTurn, String latestQuery) { + Map params = new HashMap<>(); + params.put("multi_turn", multiTurn != null ? multiTurn : "(无)"); + params.put("latest_query", latestQuery); + return PromptConstant.getIntentRecognitionPromptTemplate().render(params); + } + + /** + * 构建查询处理提示词 + * @param multiTurn 多轮对话历史 + * @param latestQuery 最新用户输入 + * @return 查询处理提示词 + */ + public static String buildQueryEnhancePrompt(String multiTurn, String latestQuery, String evidence) { + Map params = new HashMap<>(); + params.put("multi_turn", multiTurn != null ? multiTurn : "(无)"); + params.put("latest_query", latestQuery); + if (StringUtils.isEmpty(evidence)) + params.put("evidence", "无"); + else + params.put("evidence", evidence); + params.put("current_time_info", LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"))); + return PromptConstant.getQueryEnhancementPromptTemplate().render(params); + } + + /** + * 构建可行性评估提示词 + * @param canonicalQuery 规范化查询 + * @param recalledSchema 召回的数据库Schema + * @param evidence 参考信息 + * @param multiTurn 多轮对话历史 + * @return 可行性评估提示词 + */ + public static String buildFeasibilityAssessmentPrompt(String canonicalQuery, SchemaDTO recalledSchema, + String evidence, String multiTurn) { + Map params = new HashMap<>(); + String schemaInfo = buildMixMacSqlDbPrompt(recalledSchema, true); + params.put("canonical_query", canonicalQuery != null ? canonicalQuery : ""); + params.put("recalled_schema", schemaInfo); + params.put("evidence", evidence != null ? evidence : ""); + params.put("multi_turn", multiTurn != null ? multiTurn : "(无)"); + return PromptConstant.getFeasibilityAssessmentPromptTemplate().render(params); + } + /** * 渲染优化提示词模板 * @param optimizationPrompt 优化提示词模板 diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/business/BusinessKnowledgeServiceImpl.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/business/BusinessKnowledgeServiceImpl.java index 7019c404c..3be1e4415 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/business/BusinessKnowledgeServiceImpl.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/business/BusinessKnowledgeServiceImpl.java @@ -16,20 +16,38 @@ package com.alibaba.cloud.ai.service.business; +import com.alibaba.cloud.ai.constant.Constant; +import com.alibaba.cloud.ai.constant.DocumentMetadataConstant; import com.alibaba.cloud.ai.dto.BusinessKnowledgeDTO; import com.alibaba.cloud.ai.entity.BusinessKnowledge; import com.alibaba.cloud.ai.mapper.BusinessKnowledgeMapper; +import com.alibaba.cloud.ai.service.vectorstore.AgentVectorStoreService; +import com.alibaba.cloud.ai.util.DocumentConverterUtil; +import lombok.extern.slf4j.Slf4j; +import org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.VectorStore; import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Transactional; +import java.util.HashMap; import java.util.List; +import java.util.Map; +@Slf4j @Service public class BusinessKnowledgeServiceImpl implements BusinessKnowledgeService { private final BusinessKnowledgeMapper businessKnowledgeMapper; - public BusinessKnowledgeServiceImpl(BusinessKnowledgeMapper businessKnowledgeMapper) { + private final VectorStore vectorStore; + + private final AgentVectorStoreService agentVectorStoreService; + + public BusinessKnowledgeServiceImpl(BusinessKnowledgeMapper businessKnowledgeMapper, VectorStore vectorStore, + AgentVectorStoreService agentVectorStoreService) { this.businessKnowledgeMapper = businessKnowledgeMapper; + this.vectorStore = vectorStore; + this.agentVectorStoreService = agentVectorStoreService; } @Override @@ -58,38 +76,125 @@ public BusinessKnowledge getKnowledgeById(Long id) { } @Override + @Transactional(rollbackFor = Exception.class) public Long addKnowledge(BusinessKnowledgeDTO knowledgeDTO) { BusinessKnowledge entity = knowledgeDTO.toEntity(); - if (businessKnowledgeMapper.insert(entity) > 0) { + + // 插入数据库 + if (businessKnowledgeMapper.insert(entity) <= 0) { + throw new RuntimeException("Failed to add knowledge to database"); + } + + try { + // 转换为文档并插入向量库 + Document document = DocumentConverterUtil.convertBusinessTermToDocument(knowledgeDTO, + entity.getId().toString()); + vectorStore.add(List.of(document)); return entity.getId(); } - throw new RuntimeException("Failed to add knowledge"); + catch (Exception e) { + // 向量库插入失败，清理数据库记录 + try { + businessKnowledgeMapper.deleteById(entity.getId()); + } + catch (Exception deleteException) { + log.error("Failed to clean up database record after vector store error: {}", + deleteException.getMessage()); + } + throw new RuntimeException("Failed to add knowledge: " + e.getMessage(), e); + } } @Override + @Transactional(rollbackFor = Exception.class) public void updateKnowledge(Long id, BusinessKnowledgeDTO knowledgeDTO) { BusinessKnowledge entity = knowledgeDTO.toEntity(); entity.setId(id); - if (businessKnowledgeMapper.updateById(entity) > 0) { - return; + + // 更新数据库 + if (businessKnowledgeMapper.updateById(entity) <= 0) { + throw new RuntimeException("Failed to update knowledge in database"); + } + + try { + // 更新向量库中的文档 + Map metadata = new HashMap<>(); + metadata.put(Constant.AGENT_ID, entity.getAgentId()); + metadata.put(DocumentMetadataConstant.VECTOR_TYPE, DocumentMetadataConstant.BUSINESS_TERM); + metadata.put(DocumentMetadataConstant.DB_RECORD_ID, id.toString()); + agentVectorStoreService.deleteDocumentsByMetedata(knowledgeDTO.getAgentId().toString(), metadata); + + Document document = DocumentConverterUtil.convertBusinessTermToDocument(knowledgeDTO, id.toString()); + vectorStore.add(List.of(document)); + } + catch (Exception e) { + throw new RuntimeException("Failed to update knowledge in vector store: " + e.getMessage(), e); } - throw new RuntimeException("Failed to update knowledge"); } @Override + @Transactional(rollbackFor = Exception.class) public void deleteKnowledge(Long id) { - if (businessKnowledgeMapper.deleteById(id) > 0) { + // 从数据库获取原始数据 + BusinessKnowledge knowledge = businessKnowledgeMapper.selectById(id); + if (knowledge == null) { + log.warn("Knowledge not found with id: " + id); return; } - throw new RuntimeException("Failed to delete knowledge"); + + try { + // 从向量库删除文档 + Map metadata = new HashMap<>(); + metadata.put(Constant.AGENT_ID, knowledge.getAgentId()); + metadata.put(DocumentMetadataConstant.VECTOR_TYPE, DocumentMetadataConstant.BUSINESS_TERM); + metadata.put(DocumentMetadataConstant.DB_RECORD_ID, id.toString()); + agentVectorStoreService.deleteDocumentsByMetedata(knowledge.getAgentId().toString(), metadata); + } + catch (Exception e) { + log.warn("Failed to delete document from vector store: {}", e.getMessage()); + } + + // 从数据库删除记录 + if (businessKnowledgeMapper.deleteById(id) <= 0) { + throw new RuntimeException("Failed to delete knowledge from database"); + } } @Override + @Transactional(rollbackFor = Exception.class) public void recallKnowledge(Long id, boolean isRecall) { - if (businessKnowledgeMapper.changeRecall(id, isRecall) > 0) { - return; + // 从数据库获取原始数据 + BusinessKnowledge knowledge = businessKnowledgeMapper.selectById(id); + if (knowledge == null) { + throw new RuntimeException("Knowledge not found with id: " + id); + } + + // 更新数据库中的召回状态 + if (businessKnowledgeMapper.changeRecall(id, isRecall) <= 0) { + throw new RuntimeException("Failed to change recall status in database"); + } + + try { + // 更新向量库中的文档 + Map metadata = new HashMap<>(); + metadata.put(Constant.AGENT_ID, knowledge.getAgentId()); + metadata.put(DocumentMetadataConstant.VECTOR_TYPE, DocumentMetadataConstant.BUSINESS_TERM); + metadata.put(DocumentMetadataConstant.DB_RECORD_ID, id.toString()); + agentVectorStoreService.deleteDocumentsByMetedata(knowledge.getAgentId().toString(), metadata); + + BusinessKnowledgeDTO dto = new BusinessKnowledgeDTO(); + dto.setBusinessTerm(knowledge.getBusinessTerm()); + dto.setDescription(knowledge.getDescription()); + dto.setSynonyms(knowledge.getSynonyms()); + dto.setAgentId(knowledge.getAgentId()); + dto.setIsRecall(isRecall); + + Document document = DocumentConverterUtil.convertBusinessTermToDocument(dto, id.toString()); + vectorStore.add(List.of(document)); + } + catch (Exception e) { + throw new RuntimeException("Failed to update recall status in vector store: " + e.getMessage(), e); } - throw new RuntimeException("Failed to change recall status"); } } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/factory/FusionStrategyFactory.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/factory/FusionStrategyFactory.java index f44f95df2..cc90ef449 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/factory/FusionStrategyFactory.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/factory/FusionStrategyFactory.java @@ -20,28 +20,26 @@ import com.alibaba.cloud.ai.service.hybrid.fusion.impl.RrfFusionStrategy; import com.alibaba.cloud.ai.service.hybrid.fusion.impl.WeightedAverageStrategy; import lombok.extern.slf4j.Slf4j; +import org.springframework.beans.factory.FactoryBean; import org.springframework.beans.factory.annotation.Value; -import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; -import org.springframework.context.annotation.Bean; -import org.springframework.context.annotation.Configuration; +import org.springframework.stereotype.Component; /** * FusionStrategy工厂类，根据配置创建相应的FusionStrategy实现类 */ @Slf4j -@Configuration -public class FusionStrategyFactory { +@Component +public class FusionStrategyFactory implements FactoryBean { @Value("${" + Constant.PROJECT_PROPERTIES_PREFIX + ".fusion-strategy:rrf}") private String fusionStrategyType; /** - * 根据配置创建FusionStrategy实例 + * 创建FusionStrategy实例 * @return FusionStrategy实例 */ - @Bean - @ConditionalOnMissingBean(FusionStrategy.class) - public FusionStrategy fusionStrategy() { + @Override + public FusionStrategy getObject() throws Exception { log.info("Creating FusionStrategy with type: {}", fusionStrategyType); return switch (fusionStrategyType.toLowerCase()) { @@ -60,4 +58,22 @@ public FusionStrategy fusionStrategy() { }; } + /** + * 返回FusionStrategy的类型 + * @return FusionStrategy类对象 + */ + @Override + public Class getObjectType() { + return FusionStrategy.class; + } + + /** + * 设置为单例模式 + * @return true + */ + @Override + public boolean isSingleton() { + return true; + } + } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/factory/HybridRetrievalStrategyFactory.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/factory/HybridRetrievalStrategyFactory.java index 759f43ec4..aa767ed59 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/factory/HybridRetrievalStrategyFactory.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/factory/HybridRetrievalStrategyFactory.java @@ -16,6 +16,7 @@ package com.alibaba.cloud.ai.service.hybrid.factory; +import com.alibaba.cloud.ai.config.DataAgentProperties; import com.alibaba.cloud.ai.constant.Constant; import com.alibaba.cloud.ai.service.hybrid.fusion.FusionStrategy; import com.alibaba.cloud.ai.service.hybrid.retrieval.HybridRetrievalStrategy; @@ -23,10 +24,10 @@ import com.alibaba.cloud.ai.service.hybrid.retrieval.impl.ElasticsearchHybridRetrievalStrategy; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.beans.factory.FactoryBean; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Value; -import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; -import org.springframework.context.annotation.Bean; -import org.springframework.context.annotation.Configuration; +import org.springframework.stereotype.Component; import java.util.concurrent.ExecutorService; @@ -34,44 +35,64 @@ * 混合检索策略工厂类 根据配置条件创建并注册相应的 HybridRetrievalStrategy 实现类 */ @Slf4j -@Configuration -public class HybridRetrievalStrategyFactory { +@Component +public class HybridRetrievalStrategyFactory implements FactoryBean { + // spring ai 的官方属性 @Value("${spring.ai.vectorstore.type:simple}") private String vectorStoreType; + // spring ai Elasticsearch VectorStore的官方属性 @Value("${spring.ai.vectorstore.elasticsearch.index-name:spring-ai-document-index}") private String elasticsearchIndexName; @Value("${" + Constant.PROJECT_PROPERTIES_PREFIX + ".elasticsearch-min-score:0.5}") private Double minScore; - /** - * 创建并注册 HybridRetrievalStrategy Bean - * @param executorService 线程池执行器 - * @param vectorStore 向量存储 - * @param fusionStrategy 融合策略 - * @return HybridRetrievalStrategy 实例 - */ - @Bean - @ConditionalOnMissingBean(HybridRetrievalStrategy.class) - public HybridRetrievalStrategy hybridRetrievalStrategy(ExecutorService executorService, VectorStore vectorStore, - FusionStrategy fusionStrategy) { - log.info("Creating HybridRetrievalStrategy with vectorStore type: {}", vectorStoreType); + @Autowired + private ExecutorService executorService; + @Autowired + private VectorStore vectorStore; + + @Autowired + private FusionStrategy fusionStrategy; + + @Autowired + private DataAgentProperties dataAgentProperties; + + @Override + public HybridRetrievalStrategy getObject() throws Exception { + + if (!dataAgentProperties.getVectorStore().isEnableHybridSearch()) { + log.warn("Hybrid search is disabled. Returning null HybridRetrievalStrategy."); + return null; + } if ("elasticsearch".equalsIgnoreCase(vectorStoreType)) { log.info("Creating ElasticsearchHybridRetrievalStrategy with index: {}", elasticsearchIndexName); ElasticsearchHybridRetrievalStrategy strategy = new ElasticsearchHybridRetrievalStrategy(executorService, vectorStore, fusionStrategy); // 设置索引名称 strategy.setIndexName(elasticsearchIndexName); - strategy.setMinScore(minScore); + // 从DataAgentProperties获取最小分数 + strategy.setMinScore(dataAgentProperties.getVectorStore().getElasticsearchMinScore()); return strategy; } else { - log.warn("Creating DefaultHybridRetrievalStrategy (default) without keyword-search ability"); + log.warn( + "Creating DefaultHybridRetrievalStrategy (default) without keyword-search ability,maybe you should implement interface -> HybridRetrievalStrategy "); return new DefaultHybridRetrievalStrategy(executorService, vectorStore, fusionStrategy); } } + @Override + public Class getObjectType() { + return HybridRetrievalStrategy.class; + } + + @Override + public boolean isSingleton() { + return true; + } + } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/FusionStrategy.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/FusionStrategy.java index 5638cb58d..a7f6495e9 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/FusionStrategy.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/FusionStrategy.java @@ -22,6 +22,7 @@ /*融合策略接口*/ public interface FusionStrategy { - List fuseResults(List vectorResults, List keywordResults, int topK); + @SuppressWarnings("unchecked") + List fuseResults(int topK, List... resultLists); } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/impl/RrfFusionStrategy.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/impl/RrfFusionStrategy.java index 6d0076de0..7da32d853 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/impl/RrfFusionStrategy.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/impl/RrfFusionStrategy.java @@ -28,36 +28,31 @@ @Slf4j public class RrfFusionStrategy implements FusionStrategy { + @SuppressWarnings("unchecked") @Override - public List fuseResults(List vectorResults, List keywordResults, int topK) { - if (vectorResults == null) { - vectorResults = List.of(); - } - if (keywordResults == null) { - keywordResults = List.of(); - } - - if (vectorResults.isEmpty() && keywordResults.isEmpty()) { + public List fuseResults(int topK, List... resultLists) { + if (resultLists == null || resultLists.length == 0) { return List.of(); } // RRF参数配置 int k = 60; - // 将两个结果列表组合 - List> results = List.of(vectorResults, keywordResults); - // 使用Map存储每个文档的RRF分数 Map rrfScores = new HashMap<>(); // 使用Map存储文档ID到Document对象的映射 Map documentMap = new HashMap<>(); - // 遍历每个结果列表 - for (List resultList : results) { + for (List resultList : resultLists) { + if (resultList == null) { + continue; + } + for (int i = 0; i < resultList.size(); i++) { Document doc = resultList.get(i); // 排名从1开始 int rank = i + 1; + // 假设 getDocumentId 方法能稳定获取唯一ID String docId = getDocumentId(doc); // 计算并累加RRF分数: score = 1 / (k + rank) @@ -67,6 +62,11 @@ public List fuseResults(List vectorResults, List k } } + // 如果处理完所有列表后没有任何文档，提前返回 + if (rrfScores.isEmpty()) { + return List.of(); + } + // 按RRF分数降序排序，取topK个 return rrfScores.entrySet() .stream() @@ -74,6 +74,7 @@ public List fuseResults(List vectorResults, List k .limit(topK) .map(entry -> documentMap.get(entry.getKey())) .collect(Collectors.toList()); + } private String getDocumentId(Document document) { diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/impl/WeightedAverageStrategy.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/impl/WeightedAverageStrategy.java index 45fce9163..8f948f33b 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/impl/WeightedAverageStrategy.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/fusion/impl/WeightedAverageStrategy.java @@ -23,8 +23,9 @@ public class WeightedAverageStrategy implements FusionStrategy { + @SuppressWarnings("unchecked") @Override - public List fuseResults(List vectorResults, List keywordResults, int topK) { + public List fuseResults(int topK, List... resultLists) { throw new UnsupportedOperationException("Not implemented"); } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/AbstractHybridRetrievalStrategy.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/AbstractHybridRetrievalStrategy.java index 6ffca80e2..9d907e854 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/AbstractHybridRetrievalStrategy.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/AbstractHybridRetrievalStrategy.java @@ -16,16 +16,12 @@ package com.alibaba.cloud.ai.service.hybrid.retrieval; -import com.alibaba.cloud.ai.constant.Constant; import com.alibaba.cloud.ai.request.AgentSearchRequest; import com.alibaba.cloud.ai.service.hybrid.fusion.FusionStrategy; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.document.Document; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; -import org.springframework.ai.vectorstore.filter.Filter; -import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder; -import org.springframework.util.CollectionUtils; import org.springframework.util.StringUtils; import java.util.List; @@ -34,6 +30,8 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; +import static com.alibaba.cloud.ai.util.SearchUtil.buildVectorSearchRequest; + @Slf4j public abstract class AbstractHybridRetrievalStrategy implements HybridRetrievalStrategy { @@ -67,40 +65,34 @@ public List retrieve(AgentSearchRequest request) { if (Objects.isNull(request.getTopK())) throw new IllegalArgumentException("topK cannot be null"); - SearchRequest vectorSearchRequest = buildVectorSearchRequest(request); + SearchRequest vectorSearchRequest = buildVectorSearchRequest(request, true); // 异步执行向量搜索 CompletableFuture> vectorSearchFuture = CompletableFuture.supplyAsync(() -> { List vectorResults = vectorStore.similaritySearch(vectorSearchRequest); - log.info("Vector Search completed. Found {} documents for SearchRequest: {}", vectorResults.size(), + log.debug("Vector Search completed. Found {} documents for SearchRequest: {}", vectorResults.size(), vectorSearchRequest); return vectorResults; }, executorService); // 异步执行关键词搜索 - CompletableFuture> keywordSearchFuture = null; - if (!CollectionUtils.isEmpty(request.getKeywords())) { - keywordSearchFuture = CompletableFuture.supplyAsync(() -> { - List results = getDocumentsByKeywords(request); - log.info("Keyword Search completed. Found {} documents, using keywords: {}", results.size(), - request.getKeywords()); - return results; - }, executorService); - } + CompletableFuture> keywordSearchFuture = CompletableFuture.supplyAsync(() -> { + List results = getDocumentsByKeywords(request); + log.debug("Keyword Search completed. Found {} documents, with query: {}", results.size(), + request.getQuery()); + return results; + }, executorService); try { List vectorResults = vectorSearchFuture.get(); - if (keywordSearchFuture == null) { - return vectorResults; - } // 等待关键词搜索完成 List keywordResults = keywordSearchFuture.get(); // 融合结果 - List finalDocuments = fusionStrategy.fuseResults(vectorResults, keywordResults, - request.getTopK()); - log.info("Fusion completed. Found {} documents", finalDocuments.size()); + List finalDocuments = fusionStrategy.fuseResults(request.getTopK(), vectorResults, + keywordResults); + log.debug("Fusion completed. Found {} documents", finalDocuments.size()); return finalDocuments; } catch (InterruptedException e) { @@ -113,24 +105,6 @@ public List retrieve(AgentSearchRequest request) { } - private SearchRequest buildVectorSearchRequest(AgentSearchRequest request) { - SearchRequest.Builder builder = SearchRequest.builder(); - if (StringUtils.hasText(request.getQuery())) - builder.query(request.getQuery()); - - // 和keyword查找一样设置2倍topK去查找，因为后续需要和keyword的查找融合返回topK条 - builder.topK(request.getTopK() * 2); - - FilterExpressionBuilder b = new FilterExpressionBuilder(); - Filter.Expression expression = b - .and(b.eq(Constant.VECTOR_TYPE, request.getDocVectorType()), b.eq(Constant.AGENT_ID, request.getAgentId())) - .build(); - builder.filterExpression(expression); - - builder.similarityThreshold(request.getSimilarityThreshold()); - return builder.build(); - } - public abstract List getDocumentsByKeywords(AgentSearchRequest agentSearchRequest); } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/impl/DefaultHybridRetrievalStrategy.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/impl/DefaultHybridRetrievalStrategy.java index 3bf05641b..ce978a9b4 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/impl/DefaultHybridRetrievalStrategy.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/impl/DefaultHybridRetrievalStrategy.java @@ -28,7 +28,7 @@ import java.util.concurrent.ExecutorService; /** - * 适合测试使用，无关键词搜索能力 + * 适合测试以及没有继承实现AbstractHybridRetrievalStrategy的向量库（如Pg,milvus）使用，无关键词搜索能力 */ @Slf4j public class DefaultHybridRetrievalStrategy extends AbstractHybridRetrievalStrategy { diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/impl/ElasticsearchHybridRetrievalStrategy.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/impl/ElasticsearchHybridRetrievalStrategy.java index a26978cde..a33599b4c 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/impl/ElasticsearchHybridRetrievalStrategy.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/hybrid/retrieval/impl/ElasticsearchHybridRetrievalStrategy.java @@ -32,7 +32,7 @@ import org.springframework.ai.document.Document; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore; -import org.springframework.util.CollectionUtils; +import org.springframework.util.StringUtils; import java.io.IOException; import java.util.Collections; @@ -74,7 +74,7 @@ public List getDocumentsByKeywords(AgentSearchRequest agentSearchReque * "term": { "metadata.vectorType": "table" } } ] } }, "size": 20, "_source": true * } */ - if (CollectionUtils.isEmpty(agentSearchRequest.getKeywords())) + if (!StringUtils.hasText(agentSearchRequest.getQuery())) return Collections.emptyList(); ElasticsearchVectorStore vectorStore = (ElasticsearchVectorStore) this.vectorStore; @@ -99,10 +99,9 @@ public List getDocumentsByKeywords(AgentSearchRequest agentSearchReque private SearchRequest buildSearchRequest(AgentSearchRequest agentSearchRequest) { // 拼接keywords 通过空格连接 - String keywords = String.join(" ", agentSearchRequest.getKeywords()); - log.debug("ElasticsearchClient search keywords: {}", keywords); + log.debug("ElasticsearchClient search with query: {}", agentSearchRequest.getQuery()); - Query matchQuery = MatchQuery.of(m -> m.field("content").query(keywords))._toQuery(); + Query matchQuery = MatchQuery.of(m -> m.field("content").query(agentSearchRequest.getQuery()))._toQuery(); // 创建元数据过滤条件 Query agentIdFilter = TermQuery.of(t -> t.field("metadata.agentId").value(agentSearchRequest.getAgentId())) diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/knowledge/AgentKnowledgeServiceImpl.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/knowledge/AgentKnowledgeServiceImpl.java index 66e55207f..9b80ad1af 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/knowledge/AgentKnowledgeServiceImpl.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/knowledge/AgentKnowledgeServiceImpl.java @@ -17,6 +17,7 @@ package com.alibaba.cloud.ai.service.knowledge; import com.alibaba.cloud.ai.constant.Constant; +import com.alibaba.cloud.ai.constant.DocumentMetadataConstant; import com.alibaba.cloud.ai.entity.AgentKnowledge; import com.alibaba.cloud.ai.mapper.AgentKnowledgeMapper; import com.alibaba.cloud.ai.service.vectorstore.AgentVectorStoreService; @@ -159,7 +160,7 @@ public void deleteKnowledgeFromVectorStore(Long agentId, Integer knowledgeId) { try { String agentIdStr = String.valueOf(agentId); Map metadata = new HashMap<>(Map.ofEntries(Map.entry(Constant.AGENT_ID, agentIdStr), - Map.entry(Constant.KNOWLEDGE_ID, knowledgeId))); + Map.entry(DocumentMetadataConstant.KNOWLEDGE_ID, knowledgeId))); vectorStoreService.deleteDocumentsByMetedata(agentIdStr, metadata); diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/nl2sql/Nl2SqlService.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/nl2sql/Nl2SqlService.java index 01260537a..506b5ed5d 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/nl2sql/Nl2SqlService.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/nl2sql/Nl2SqlService.java @@ -22,19 +22,17 @@ import org.springframework.ai.chat.model.ChatResponse; import reactor.core.publisher.Flux; -import java.util.List; import java.util.function.Consumer; public interface Nl2SqlService { Flux semanticConsistencyStream(String sql, String queryPrompt); - Flux generateSql(List evidenceList, String query, SchemaDTO schemaDTO, String sql, - String exceptionMessage); + Flux generateSql(String evidence, String query, SchemaDTO schemaDTO, String sql, String exceptionMessage); Flux generateOptimizedSql(String previousSql, String exceptionMessage, int round); - Flux fineSelect(SchemaDTO schemaDTO, String query, List evidenceList, + Flux fineSelect(SchemaDTO schemaDTO, String query, String evidence, String sqlGenerateSchemaMissingAdvice, DbConfig specificDbConfig, Consumer dtoConsumer); default String sqlTrim(String sql) { diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/nl2sql/Nl2SqlServiceImpl.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/nl2sql/Nl2SqlServiceImpl.java index ac9e64cd3..03ae6d241 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/nl2sql/Nl2SqlServiceImpl.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/nl2sql/Nl2SqlServiceImpl.java @@ -17,29 +17,21 @@ import com.alibaba.cloud.ai.connector.accessor.Accessor; import com.alibaba.cloud.ai.connector.accessor.AccessorFactory; -import com.alibaba.cloud.ai.connector.bo.DbQueryParameter; -import com.alibaba.cloud.ai.connector.bo.ResultSetBO; import com.alibaba.cloud.ai.connector.config.DbConfig; -import com.alibaba.cloud.ai.dto.schema.ColumnDTO; import com.alibaba.cloud.ai.dto.schema.SchemaDTO; -import com.alibaba.cloud.ai.dto.schema.TableDTO; import com.alibaba.cloud.ai.prompt.PromptHelper; import com.alibaba.cloud.ai.service.llm.LlmService; - import com.alibaba.cloud.ai.util.FluxUtil; +import com.alibaba.cloud.ai.util.JsonUtil; import com.alibaba.cloud.ai.util.MarkdownParserUtil; import com.fasterxml.jackson.core.type.TypeReference; -import com.alibaba.cloud.ai.util.JsonUtil; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.stereotype.Service; import reactor.core.publisher.Flux; -import java.util.ArrayList; -import java.util.HashMap; import java.util.HashSet; import java.util.List; -import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.function.Consumer; @@ -73,7 +65,7 @@ public Flux semanticConsistencyStream(String sql, String queryProm } @Override - public Flux generateSql(List evidenceList, String query, SchemaDTO schemaDTO, String sql, + public Flux generateSql(String evidence, String query, SchemaDTO schemaDTO, String sql, String exceptionMessage) { log.info("Generating SQL for query: {}, hasExistingSql: {}", query, sql != null && !sql.isEmpty()); @@ -81,15 +73,15 @@ public Flux generateSql(List evidenceList, String query, SchemaD if (sql != null && !sql.isEmpty()) { // Use professional SQL error repair prompt log.debug("Using SQL error fixer for existing SQL: {}", sql); - String errorFixerPrompt = PromptHelper.buildSqlErrorFixerPrompt(query, dbConfig, schemaDTO, evidenceList, - sql, exceptionMessage); + String errorFixerPrompt = PromptHelper.buildSqlErrorFixerPrompt(query, dbConfig, schemaDTO, evidence, sql, + exceptionMessage); newSqlFlux = llmService.toStringFlux(llmService.callUser(errorFixerPrompt)); log.info("SQL error fixing completed"); } else { // Normal SQL generation process log.debug("Generating new SQL from scratch"); - List prompts = PromptHelper.buildMixSqlGeneratorPrompt(query, dbConfig, schemaDTO, evidenceList); + List prompts = PromptHelper.buildMixSqlGeneratorPrompt(query, dbConfig, schemaDTO, evidence); newSqlFlux = llmService.toStringFlux(llmService.call(prompts.get(0), prompts.get(1))); log.info("New SQL generation completed"); } @@ -164,18 +156,13 @@ private Flux fineSelect(SchemaDTO schemaDTO, String sqlGenerateSch } @Override - public Flux fineSelect(SchemaDTO schemaDTO, String query, List evidenceList, + public Flux fineSelect(SchemaDTO schemaDTO, String query, String evidence, String sqlGenerateSchemaMissingAdvice, DbConfig specificDbConfig, Consumer dtoConsumer) { - log.debug("Fine selecting schema for query: {} with {} evidences and specificDbConfig: {}", query, - evidenceList.size(), specificDbConfig != null ? specificDbConfig.getUrl() : "default"); + log.debug("Fine selecting schema for query: {} with evidences and specificDbConfig: {}", query, + specificDbConfig != null ? specificDbConfig.getUrl() : "default"); - // 增加具体的样例数据，让模型根据样例数据进行选择 - SchemaDTO enrichedSchema = enrichSchemaWithSampleData(schemaDTO, specificDbConfig); - log.debug("Schema enriched with sample data for {} tables", - enrichedSchema.getTable() != null ? enrichedSchema.getTable().size() : 0); - - String prompt = buildMixSelectorPrompt(evidenceList, query, enrichedSchema); - log.debug("Calling LLM for schema fine selection"); + String prompt = buildMixSelectorPrompt(evidence, query, schemaDTO); + log.debug("Built schema fine selection prompt as follows \n {} \n", prompt); Set selectedTables = new HashSet<>(); @@ -209,11 +196,14 @@ public Flux fineSelect(SchemaDTO schemaDTO, String query, List !selectedTables.contains(table.getName().toLowerCase())); - int finalTableCount = schemaDTO.getTable() != null ? schemaDTO.getTable().size() : 0; - log.debug("Fine selection completed: {} -> {} tables, selected tables: {}", originalTableCount, - finalTableCount, selectedTables); + if (schemaDTO.getTable() != null) { + int originalTableCount = schemaDTO.getTable().size(); + schemaDTO.getTable() + .removeIf(table -> !selectedTables.contains(table.getName().toLowerCase())); + int finalTableCount = schemaDTO.getTable().size(); + log.debug("Fine selection completed: {} -> {} tables, selected tables: {}", + originalTableCount, finalTableCount, selectedTables); + } } } dtoConsumer.accept(schemaDTO); @@ -223,233 +213,4 @@ public Flux fineSelect(SchemaDTO schemaDTO, String query, List columnHeaders = tableData.getColumn(); - List> rows = tableData.getData(); - - // 为每个列创建样例数据映射 - Map> columnSamples = new HashMap<>(); - - // 遍历每一行数据 - for (Map row : rows) { - for (String columnName : columnHeaders) { - String value = row.get(columnName); - - if (value != null && !value.trim().isEmpty()) { - columnSamples.computeIfAbsent(columnName, k -> new ArrayList<>()).add(value); - } - } - } - - // 将样例数据分配给对应的列 - for (ColumnDTO columnDTO : tableDTO.getColumn()) { - String columnName = columnDTO.getName(); - List samples = columnSamples.get(columnName); - - if (samples != null && !samples.isEmpty()) { - // 去重并限制样例数量 - List filteredSamples = samples.stream() - .filter(sample -> sample != null && !sample.trim().isEmpty()) - .distinct() - .limit(5) // 最多保留5个样例值 - .collect(Collectors.toList()); - - if (!filteredSamples.isEmpty()) { - columnDTO.setSamples(filteredSamples); - log.debug("Added {} sample values for column '{}.{}': {}", filteredSamples.size(), - tableDTO.getName(), columnName, filteredSamples); - } - } - } - } - - /** - * 检查数据库配置是否有效 - * @param dbConfig 数据库配置 - * @return true如果配置有效，false otherwise - */ - private boolean isDatabaseConfigValid(DbConfig dbConfig) { - if (dbConfig == null) { - log.debug("dbConfig is null"); - return false; - } - - if (dbAccessor == null) { - log.debug("dbAccessor is null"); - return false; - } - - // 检查基本的连接信息 - boolean hasBasicInfo = dbConfig.getUrl() != null && !dbConfig.getUrl().trim().isEmpty() - && dbConfig.getUsername() != null && !dbConfig.getUsername().trim().isEmpty(); - - if (!hasBasicInfo) { - log.debug("dbConfig missing basic connection info - url: {}, username: {}", - dbConfig.getUrl() != null ? "present" : "null", - dbConfig.getUsername() != null ? "present" : "null"); - return false; - } - - return true; - } - - /** - * 创建SchemaDTO的深拷贝 - * @param originalSchema 原始Schema - * @return Schema的深拷贝 - */ - private SchemaDTO copySchemaDTO(SchemaDTO originalSchema) { - SchemaDTO copy = new SchemaDTO(); - copy.setName(originalSchema.getName()); - copy.setDescription(originalSchema.getDescription()); - copy.setTableCount(originalSchema.getTableCount()); - copy.setForeignKeys(originalSchema.getForeignKeys()); - - if (originalSchema.getTable() != null) { - List copiedTables = originalSchema.getTable() - .stream() - .map(this::copyTableDTO) - .collect(Collectors.toList()); - copy.setTable(copiedTables); - } - - return copy; - } - - /** - * 创建TableDTO的深拷贝 - * @param originalTable 原始表 - * @return 表的深拷贝 - */ - private TableDTO copyTableDTO(TableDTO originalTable) { - TableDTO copy = new TableDTO(); - copy.setName(originalTable.getName()); - copy.setDescription(originalTable.getDescription()); - copy.setPrimaryKeys(originalTable.getPrimaryKeys()); - - if (originalTable.getColumn() != null) { - List copiedColumns = originalTable.getColumn() - .stream() - .map(this::copyColumnDTO) - .collect(Collectors.toList()); - copy.setColumn(copiedColumns); - } - - return copy; - } - - /** - * 创建ColumnDTO的深拷贝 - * @param originalColumn 原始列 - * @return 列的深拷贝 - */ - private ColumnDTO copyColumnDTO(ColumnDTO originalColumn) { - ColumnDTO copy = new ColumnDTO(); - copy.setName(originalColumn.getName()); - copy.setDescription(originalColumn.getDescription()); - copy.setEnumeration(originalColumn.getEnumeration()); - copy.setRange(originalColumn.getRange()); - copy.setType(originalColumn.getType()); - copy.setMapping(originalColumn.getMapping()); - - // 复制现有的样例数据 - if (originalColumn.getSamples() != null) { - copy.setSamples(new ArrayList<>(originalColumn.getSamples())); - } - if (originalColumn.getData() != null) { - copy.setData(new ArrayList<>(originalColumn.getData())); - } - - return copy; - } - } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/processing/AbstractQueryProcessingService.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/processing/AbstractQueryProcessingService.java deleted file mode 100644 index b58e8ebf3..000000000 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/processing/AbstractQueryProcessingService.java +++ /dev/null @@ -1,235 +0,0 @@ -/* - * Copyright 2024-2025 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.alibaba.cloud.ai.service.processing; - -import com.alibaba.cloud.ai.dto.schema.SchemaDTO; -import com.alibaba.cloud.ai.entity.Datasource; -import com.alibaba.cloud.ai.enums.TextType; -import com.alibaba.cloud.ai.prompt.PromptConstant; -import com.alibaba.cloud.ai.prompt.PromptHelper; -import com.alibaba.cloud.ai.service.datasource.DatasourceService; -import com.alibaba.cloud.ai.service.llm.LlmService; -import com.alibaba.cloud.ai.service.nl2sql.Nl2SqlService; -import com.alibaba.cloud.ai.service.schema.SchemaService; -import com.alibaba.cloud.ai.service.vectorstore.AgentVectorStoreService; -import com.alibaba.cloud.ai.util.ChatResponseUtil; -import com.alibaba.cloud.ai.util.FluxUtil; -import com.alibaba.cloud.ai.util.JsonUtil; -import com.alibaba.cloud.ai.util.SchemaProcessorUtil; -import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.core.type.TypeReference; -import lombok.AllArgsConstructor; -import lombok.extern.slf4j.Slf4j; -import org.springframework.ai.chat.model.ChatResponse; -import org.springframework.ai.document.Document; -import org.springframework.util.Assert; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; -import java.util.stream.Collectors; - -@Slf4j -@AllArgsConstructor -public abstract class AbstractQueryProcessingService implements QueryProcessingService { - - private final LlmService llmService; - - private final DatasourceService datasourceService; - - protected abstract AgentVectorStoreService getVectorStoreService(); - - protected abstract SchemaService getSchemaService(); - - protected abstract Nl2SqlService getNl2SqlService(); - - @Override - public List extractEvidences(String query, String agentId) { - log.debug("Extracting evidences for query: {} with agentId: {}", query, agentId); - Assert.notNull(agentId, "AgentId cannot be null"); - List evidenceDocuments = getVectorStoreService().getDocumentsForAgent(agentId, query, "evidence"); - - List evidences = evidenceDocuments.stream().map(Document::getText).collect(Collectors.toList()); - log.debug("Extracted {} evidences: {}", evidences.size(), evidences); - return evidences; - } - - @Override - public Flux extractKeywords(String query, List evidenceList, - Consumer> resultConsumer) { - log.debug("Extracting keywords from query: {} with {} evidences", query, evidenceList.size()); - StringBuilder queryBuilder = new StringBuilder(query); - for (String evidence : evidenceList) { - queryBuilder.append(evidence).append("。"); - } - query = queryBuilder.toString(); - - String prompt = PromptHelper.buildQueryToKeywordsPrompt(query); - log.debug("Calling LLM for keyword extraction"); - StringBuilder sb = new StringBuilder(); - return llmService.callUser(prompt).doOnNext(response -> { - String text = response.getResult().getOutput().getText(); - sb.append(text); - }).doOnComplete(() -> { - String content = sb.toString(); - List keywords; - try { - keywords = JsonUtil.getObjectMapper().readValue(content, new TypeReference<>() { - }); - } - catch (JsonProcessingException e) { - throw new RuntimeException(e); - } - log.debug("Extracted {} keywords: {}", keywords != null ? keywords.size() : 0, keywords); - resultConsumer.accept(keywords); - }); - } - - @Override - public Flux expandQuestion(String query, Consumer> resultConsumer) { - log.info("Starting question expansion for query: {}", query); - try { - // Build question expansion prompt - Map params = new HashMap<>(); - params.put("question", query); - String prompt = PromptConstant.getQuestionExpansionPromptTemplate().render(params); - - // Call LLM to get expanded questions - log.debug("Calling LLM for question expansion"); - StringBuilder sb = new StringBuilder(); - Flux responseFlux = llmService.callUser(prompt); - return responseFlux.doOnNext(response -> { - String text = response.getResult().getOutput().getText(); - sb.append(text); - }).doOnComplete(() -> { - String content = sb.toString(); - try { - // Parse JSON response - List expandedQuestions = JsonUtil.getObjectMapper() - .readValue(content, new TypeReference<>() { - }); - - if (expandedQuestions == null || expandedQuestions.isEmpty()) { - log.warn("No expanded questions generated, returning original query"); - expandedQuestions = Collections.singletonList(query); - } - - log.info("Question expansion completed successfully: {} questions generated", - expandedQuestions.size()); - log.debug("Expanded questions: {}", expandedQuestions); - resultConsumer.accept(expandedQuestions); - } - catch (Exception e) { - throw new RuntimeException(e); - } - }); - } - catch (Exception e) { - log.warn("Question expansion failed, returning original query: {}", e.getMessage()); - return Flux.error(e); - } - } - - @Override - public Flux rewriteStream(String query, String agentId, StringBuilder queryResultCollector) { - return FluxUtil.cascadeFlux(processTimeExpressions(query), timeRewrittenQuery -> { - - log.debug("Time rewritten query: {} -> {}", query, timeRewrittenQuery); - - List evidences = extractEvidences(timeRewrittenQuery, agentId); - log.debug("Extracted {} evidences for rewriteStream, they are {}", evidences.size(), evidences); - - AtomicReference schemaDTORef = new AtomicReference<>(null); - return FluxUtil.cascadeFlux( - select(timeRewrittenQuery, evidences, agentId, schemaDTORef::set), r -> { - SchemaDTO schemaDTO = schemaDTORef.get(); - log.debug("SchemaDTO is {}", schemaDTO); - String prompt = PromptHelper.buildRewritePrompt(timeRewrittenQuery, schemaDTO, evidences); - log.debug("Built rewrite prompt for streaming, prompt is as follows \n {}", prompt); - return llmService.callUser(prompt) - .doOnNext(chatResponse -> queryResultCollector - .append(ChatResponseUtil.getText(chatResponse))); - }, flux -> Mono.just(""), - Flux.just(ChatResponseUtil.createResponse("正在选择合适的数据表...\n"), - ChatResponseUtil.createPureResponse(TextType.JSON.getStartSign())), - Flux.just(ChatResponseUtil.createPureResponse(TextType.JSON.getEndSign()), - ChatResponseUtil.createResponse("\n\n选择数据表完成！")), - Flux.empty()); - }, flux -> flux.map(r -> Optional.ofNullable(r.getResult().getOutput().getText()).orElse("")) - .collect(StringBuilder::new, StringBuilder::append) - .map(StringBuilder::toString), Flux.just(ChatResponseUtil.createResponse("正在替换问题中的时间表达式...\n")), - Flux.just(ChatResponseUtil.createResponse("\n\n重写时间表达式完成！\n正在提取用户问题关键词...")), Flux.empty()) - .doOnSubscribe(s -> log.info("Starting rewriteStream for query: {} with agentId: {}", query, agentId)) - .doOnError(e -> log.error("RewriteStream failed for query: {}", e.getMessage())); - } - - /** - * 处理查询中的时间表达式，将相对时间转换为具体时间 - * @param query 原始查询 - * @return 处理后的查询 - */ - private Flux processTimeExpressions(String query) { - log.debug("Processing time expressions in query: {}", query); - - // 使用统一管理的提示词构建时间转换提示 - String timeConversionPrompt = PromptHelper.buildTimeConversionPrompt(query); - - // 调用模型进行时间转换 - return llmService.callUser(timeConversionPrompt); - } - - private Flux select(String query, List evidenceList, String agentId, - Consumer dtoConsumer) { - Assert.notNull(agentId, "AgentId cannot be null"); - log.debug("Starting schema selection for query: {} with {} evidences and agentId: {}", query, - evidenceList.size(), agentId); - AtomicReference oldSchemaDTORef = new AtomicReference<>(null); - return FluxUtil.cascadeFlux(extractKeywords(query, evidenceList, keywords -> { - log.debug("Using {} keywords for schema selection", keywords != null ? keywords.size() : 0); - - Datasource datasource = datasourceService.getActiveDatasourceByAgentId(Integer.valueOf(agentId)); - if (datasource == null) { - throw new RuntimeException("No active datasource found for agentId: " + agentId); - } - SchemaDTO schemaDTO = getSchemaService().mixRagForAgent(agentId, query, keywords, - SchemaProcessorUtil.createDbConfigFromDatasource(datasource)); - oldSchemaDTORef.set(schemaDTO); - - log.debug("Retrieved schema with {} tables", - schemaDTO.getTable() != null ? schemaDTO.getTable().size() : 0); - }), r -> { - Flux flux = fineSelect(oldSchemaDTORef.get(), query, evidenceList, result -> { - log.debug("Fine selection completed, final schema has {} tables", - result.getTable() != null ? result.getTable().size() : 0); - dtoConsumer.accept(result); - }); - return Flux.concat(Flux.just(ChatResponseUtil.createResponse("")), flux); - }, flux -> Mono.just("")); - } - - private Flux fineSelect(SchemaDTO schemaDTO, String query, List evidenceList, - Consumer dtoConsumer) { - return getNl2SqlService().fineSelect(schemaDTO, query, evidenceList, null, null, dtoConsumer); - } - -} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/processing/QueryProcessingService.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/processing/QueryProcessingService.java deleted file mode 100644 index b48cf9845..000000000 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/processing/QueryProcessingService.java +++ /dev/null @@ -1,44 +0,0 @@ -/* - * Copyright 2024-2025 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.alibaba.cloud.ai.service.processing; - -import org.springframework.ai.chat.model.ChatResponse; -import reactor.core.publisher.Flux; - -import java.util.List; -import java.util.function.Consumer; - -/** - * 问题加工服务 - */ -public interface QueryProcessingService { - - List extractEvidences(String query, String agentId); - - Flux extractKeywords(String query, List evidenceList, Consumer> resultConsumer); - - /** - * 扩展用户问题为一个List列表 - * @param query 用户问题 - * @param resultConsumer 处理扩展结果的消费者 - * @return AI模型的流式结果 - */ - Flux expandQuestion(String query, Consumer> resultConsumer); - - Flux rewriteStream(String query, String agentId, StringBuilder queryResultCollector) throws Exception; - -} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/processing/QueryProcessingServiceImpl.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/processing/QueryProcessingServiceImpl.java deleted file mode 100644 index fc37c87a4..000000000 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/processing/QueryProcessingServiceImpl.java +++ /dev/null @@ -1,58 +0,0 @@ -/* - * Copyright 2024-2025 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.alibaba.cloud.ai.service.processing; - -import com.alibaba.cloud.ai.service.datasource.DatasourceService; -import com.alibaba.cloud.ai.service.llm.LlmService; -import com.alibaba.cloud.ai.service.nl2sql.Nl2SqlService; -import com.alibaba.cloud.ai.service.schema.SchemaService; -import com.alibaba.cloud.ai.service.vectorstore.AgentVectorStoreService; -import org.springframework.stereotype.Service; - -@Service -public class QueryProcessingServiceImpl extends AbstractQueryProcessingService { - - private final AgentVectorStoreService vectorStoreService; - - private final SchemaService schemaService; - - private final Nl2SqlService nl2SqlService; - - public QueryProcessingServiceImpl(LlmService aiService, AgentVectorStoreService vectorStoreService, - SchemaService schemaService, Nl2SqlService nl2SqlService, DatasourceService datasourceService) { - super(aiService, datasourceService); - this.vectorStoreService = vectorStoreService; - this.schemaService = schemaService; - this.nl2SqlService = nl2SqlService; - } - - @Override - protected AgentVectorStoreService getVectorStoreService() { - return this.vectorStoreService; - } - - @Override - protected SchemaService getSchemaService() { - return this.schemaService; - } - - @Override - protected Nl2SqlService getNl2SqlService() { - return this.nl2SqlService; - } - -} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/schema/SchemaService.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/schema/SchemaService.java index 334c78a12..ced55e198 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/schema/SchemaService.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/schema/SchemaService.java @@ -26,21 +26,12 @@ public interface SchemaService { List getTableDocumentsForAgent(String agentId, String query); - List> getColumnDocumentsByKeywordsForAgent(String agentId, List keywords); + List getColumnDocumentsByKeywordsForAgent(String agentId, List keywords, + List tableDocuments); void extractDatabaseName(SchemaDTO schemaDTO, DbConfig dbConfig); - void buildSchemaFromDocuments(String agentId, List> columnDocumentList, - List tableDocuments, SchemaDTO schemaDTO); - - SchemaDTO mixRagForAgent(String agentId, String query, List keywords, DbConfig dbConfig); - - default List getTableDocuments(String query) { - return getTableDocumentsForAgent(null, query); - } - - default List> getColumnDocumentsByKeywords(List keywords) { - return getColumnDocumentsByKeywordsForAgent(null, keywords); - } + void buildSchemaFromDocuments(String agentId, List columnDocumentList, List tableDocuments, + SchemaDTO schemaDTO); } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/schema/SchemaServiceImpl.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/schema/SchemaServiceImpl.java index 8a0610f22..06b6df70a 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/schema/SchemaServiceImpl.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/schema/SchemaServiceImpl.java @@ -15,12 +15,14 @@ */ package com.alibaba.cloud.ai.service.schema; -import com.alibaba.cloud.ai.constant.Constant; -import com.alibaba.cloud.ai.enums.BizDataSourceTypeEnum; +import com.alibaba.cloud.ai.config.DataAgentProperties; import com.alibaba.cloud.ai.connector.config.DbConfig; +import com.alibaba.cloud.ai.constant.DocumentMetadataConstant; import com.alibaba.cloud.ai.dto.schema.ColumnDTO; import com.alibaba.cloud.ai.dto.schema.SchemaDTO; import com.alibaba.cloud.ai.dto.schema.TableDTO; +import com.alibaba.cloud.ai.enums.BizDataSourceTypeEnum; +import com.alibaba.cloud.ai.service.hybrid.fusion.FusionStrategy; import com.alibaba.cloud.ai.service.vectorstore.AgentVectorStoreService; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; @@ -32,13 +34,7 @@ import org.springframework.stereotype.Service; import org.springframework.util.Assert; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; +import java.util.*; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; @@ -58,38 +54,13 @@ public class SchemaServiceImpl implements SchemaService { */ private final AgentVectorStoreService vectorStoreService; - /** - * Build schema based on RAG - supports agent isolation - * @param agentId agent ID - * @param query query - * @param keywords keyword list - * @param dbConfig Database configuration - * @return SchemaDTO - */ - @Override - public SchemaDTO mixRagForAgent(String agentId, String query, List keywords, DbConfig dbConfig) { - SchemaDTO schemaDTO = new SchemaDTO(); - extractDatabaseName(schemaDTO, dbConfig); // Set database name or schema name - - // Get table documents - List tableDocuments = new ArrayList<>(getTableDocumentsForAgent(agentId, query)); - - // Get column documents - List> columnDocumentList = getColumnDocumentsByKeywordsForAgent(agentId, keywords); - - buildSchemaFromDocuments(agentId, columnDocumentList, tableDocuments, schemaDTO); + private final FusionStrategy fusionStrategy; - return schemaDTO; - } + private final DataAgentProperties dataAgentProperties; @Override - public void buildSchemaFromDocuments(String agentId, List> columnDocumentList, + public void buildSchemaFromDocuments(String agentId, List currentColumnDocuments, List tableDocuments, SchemaDTO schemaDTO) { - // Process column weights and sort by table association - updateAndSortColumnScoresByTableWeights(columnDocumentList, tableDocuments); - - // Initialize column selector, TODO upper limit 100 has issues - Map weightedColumns = selectColumnsByRoundRobin(columnDocumentList, 100); // 如果外键关系是"订单表.订单ID=订单详情表.订单ID"，那么 relatedNamesFromForeignKeys // 将包含"订单表.订单ID"和"订单详情表.订单ID" @@ -100,10 +71,10 @@ public void buildSchemaFromDocuments(String agentId, List> column // Supplement missing foreign key corresponding tables expandTableDocumentsWithForeignKeys(agentId, tableDocuments, relatedNamesFromForeignKeys); - expandColumnDocumentsWithForeignKeys(agentId, weightedColumns, relatedNamesFromForeignKeys); + expandColumnDocumentsWithForeignKeys(agentId, currentColumnDocuments, relatedNamesFromForeignKeys); - // Attach weighted columns to corresponding tables - attachColumnsToTables(weightedColumns, tableList); + // Attach columns to corresponding tables + attachColumnsToTables(currentColumnDocuments, tableList); // Finally assemble SchemaDTO schemaDTO.setTable(tableList); @@ -122,40 +93,94 @@ public void buildSchemaFromDocuments(String agentId, List> column @Override public List getTableDocumentsForAgent(String agentId, String query) { Assert.notNull(agentId, "agentId cannot be null"); - return vectorStoreService.getDocumentsForAgent(agentId, query, Constant.TABLE); + return vectorStoreService.getDocumentsForAgent(agentId, query, DocumentMetadataConstant.TABLE); } /** * Get all column documents by keywords for specified agent */ @Override - public List> getColumnDocumentsByKeywordsForAgent(String agentId, List keywords) { + public List getColumnDocumentsByKeywordsForAgent(String agentId, List keywords, + List tableDocuments) { - Assert.notNull(agentId, "agentId cannot be null"); + Assert.hasText(agentId, "agentId cannot be empty"); + + List> allResults = new ArrayList<>(); + for (String kw : keywords) { + List docs = vectorStoreService.getDocumentsForAgent(agentId, kw, DocumentMetadataConstant.COLUMN); + if (CollectionUtils.isEmpty(docs)) { + continue; + } + List filterDocs = filterColumnsWithMatchingTables(docs, tableDocuments); + if (CollectionUtils.isNotEmpty(filterDocs)) + allResults.add(filterDocs); + + } + + List[] resultArray = allResults.toArray(new List[0]); + return fusionStrategy.fuseResults(dataAgentProperties.getVectorStore().getTopkLimit(), resultArray); - return keywords.stream() - .map(kw -> vectorStoreService.getDocumentsForAgent(agentId, kw, Constant.COLUMN)) - .collect(Collectors.toList()); } /** * Expand column documents (supplement missing columns through foreign keys) */ - private void expandColumnDocumentsWithForeignKeys(String agentId, Map weightedColumns, + private void expandColumnDocumentsWithForeignKeys(String agentId, List currentColumns, Set foreignKeySet) { + // 提取已存在的列名，格式为 tableName.columnName + Set existingColumnNames = currentColumns.stream() + .map(doc -> (String) doc.getMetadata().get("tableName") + "." + (String) doc.getMetadata().get("name")) + .collect(Collectors.toSet()); - Set existingColumnNames = weightedColumns.keySet(); - Set missingColumns = new HashSet<>(); + // 找出缺少的列 + Set missingColumnNames = new HashSet<>(); for (String key : foreignKeySet) { if (!existingColumnNames.contains(key)) { - missingColumns.add(key); + missingColumnNames.add(key); + } + } + + if (missingColumnNames.isEmpty()) { + return; + } + + // 按表名分组缺少的列 + Map> missingColumnsByTable = new HashMap<>(); + for (String columnFullName : missingColumnNames) { + String[] parts = columnFullName.split("\\."); + if (parts.length == 2) { + String tableName = parts[0]; + String columnName = parts[1]; + missingColumnsByTable.computeIfAbsent(tableName, k -> new ArrayList<>()).add(columnName); } } - for (String columnName : missingColumns) { - addColumnsDocument(agentId, weightedColumns, columnName); + // 为每个表批量获取缺少的列文档 + for (Map.Entry> entry : missingColumnsByTable.entrySet()) { + String tableName = entry.getKey(); + List columnNames = entry.getValue(); + + List foundColumnDocs = vectorStoreService.getColumnDocuments(agentId, tableName, columnNames); + + // 添加到当前列列表中 + if (!foundColumnDocs.isEmpty()) + currentColumns.addAll(foundColumnDocs); + + // 记录日志 + if (foundColumnDocs.size() < columnNames.size()) { + Set foundColumnNames = foundColumnDocs.stream() + .map(doc -> (String) doc.getMetadata().get("name")) + .collect(Collectors.toSet()); + + Set notFoundColumns = new HashSet<>(columnNames); + notFoundColumns.removeAll(foundColumnNames); + + log.warn("When we search from vector store,agentId: {} - columns not found in table {}: {}", agentId, + tableName, notFoundColumns); + } } + log.debug("Finish expanding column documents with foreign keys: {}", missingColumnNames); } /** @@ -178,68 +203,21 @@ private void expandTableDocumentsWithForeignKeys(String agentId, List } } - for (String tableName : missingTables) { - addTableDocument(agentId, tableDocuments, tableName); + if (!CollectionUtils.isEmpty(missingTables)) { + log.debug("Expand table documents with foreign keys: {}", missingTables); + loadMissingTableDocument(agentId, tableDocuments, new ArrayList<>(missingTables)); } - } - - protected void addTableDocument(String agentId, List tableDocuments, String tableName) { - List documentsForAgent = vectorStoreService.getDocumentsForAgent(agentId, tableName, Constant.TABLE); - if (documentsForAgent != null && !documentsForAgent.isEmpty()) - tableDocuments.addAll(documentsForAgent); - } - protected void addColumnsDocument(String agentId, Map weightedColumns, String columnName) { - List documentsForAgent = vectorStoreService.getDocumentsForAgent(agentId, columnName, - Constant.COLUMN); - if (documentsForAgent != null && !documentsForAgent.isEmpty()) { - for (Document document : documentsForAgent) - weightedColumns.putIfAbsent(document.getId(), document); - } } - /** - * Select up to maxCount columns by weight using a round-robin approach to ensure - * balanced selection across different tables - */ - protected Map selectColumnsByRoundRobin(List> columnDocumentList, int maxCount) { - Map selectedColumns = new HashMap<>(); - int currentRound = 0; - - // Continue selecting columns until we reach maxCount or exhaust all columns - while (selectedColumns.size() < maxCount) { - boolean hasMoreColumnsInAnyList = false; - - // Process each table's column list in the current round - for (List tableColumns : columnDocumentList) { - if (currentRound < tableColumns.size()) { - // Get the column at current position (already sorted by weight) - Document column = tableColumns.get(currentRound); - String columnId = column.getId(); - - // Add to selection if not already selected - if (!selectedColumns.containsKey(columnId)) { - selectedColumns.put(columnId, column); - - // Stop if we've reached the maximum count - if (selectedColumns.size() >= maxCount) { - break; - } - } - - hasMoreColumnsInAnyList = true; - } - } - - // If no more columns in any list, exit the loop - if (!hasMoreColumnsInAnyList) { - break; - } - - currentRound++; - } - - return selectedColumns; + private void loadMissingTableDocument(String agentId, List tableDocuments, + List missingtableNames) { + List foundTableDocs = vectorStoreService.getTableDocuments(agentId, missingtableNames); + if (foundTableDocs.size() > missingtableNames.size()) + log.error("When we search missing tables:{}, more than expected tables for agent: {}", missingtableNames, + agentId); + if (!foundTableDocs.isEmpty()) + tableDocuments.addAll(foundTableDocs); } /** @@ -438,12 +416,12 @@ protected Set extractRelatedNamesFromForeignKeys(List tableDoc /** * Attach column documents to corresponding tables */ - protected void attachColumnsToTables(Map weightedColumns, List tableList) { - if (CollectionUtils.isEmpty(weightedColumns.values())) { + private void attachColumnsToTables(List columns, List tableList) { + if (CollectionUtils.isEmpty(columns)) { return; } - for (Document columnDoc : weightedColumns.values()) { + for (Document columnDoc : columns) { Map meta = columnDoc.getMetadata(); ColumnDTO columnDTO = new ColumnDTO(); columnDTO.setName((String) meta.get("name")); diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/semantic/SemanticModelService.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/semantic/SemanticModelService.java index bb688abc9..a7867e890 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/semantic/SemanticModelService.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/semantic/SemanticModelService.java @@ -16,6 +16,7 @@ package com.alibaba.cloud.ai.service.semantic; +import com.alibaba.cloud.ai.dto.SemanticModelAddDTO; import com.alibaba.cloud.ai.entity.SemanticModel; import java.util.List; @@ -26,10 +27,14 @@ public interface SemanticModelService { List getEnabledByAgentId(Long agentId); + List getByAgentIdAndTableNames(Integer agentId, List tableNames); + SemanticModel getById(Long id); void addSemanticModel(SemanticModel semanticModel); + boolean addSemanticModel(SemanticModelAddDTO dto); + void enableSemanticModel(Long id); void disableSemanticModel(Long id); diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/semantic/SemanticModelServiceImpl.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/semantic/SemanticModelServiceImpl.java index eb2e27abf..dac58212a 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/semantic/SemanticModelServiceImpl.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/semantic/SemanticModelServiceImpl.java @@ -16,7 +16,10 @@ package com.alibaba.cloud.ai.service.semantic; +import com.alibaba.cloud.ai.dto.SemanticModelAddDTO; +import com.alibaba.cloud.ai.entity.AgentDatasource; import com.alibaba.cloud.ai.entity.SemanticModel; +import com.alibaba.cloud.ai.mapper.AgentDatasourceMapper; import com.alibaba.cloud.ai.mapper.SemanticModelMapper; import lombok.AllArgsConstructor; import org.springframework.stereotype.Service; @@ -29,6 +32,8 @@ public class SemanticModelServiceImpl implements SemanticModelService { private final SemanticModelMapper semanticModelMapper; + private final AgentDatasourceMapper agentDatasourceMapper; + @Override public List getAll() { return semanticModelMapper.selectAll(); @@ -39,6 +44,17 @@ public List getEnabledByAgentId(Long agentId) { return semanticModelMapper.selectEnabledByAgentId(agentId); } + @Override + public List getByAgentIdAndTableNames(Integer agentId, List tableNames) { + Integer datasourceId = findDatasourceIdByAgentId(agentId); + + if (datasourceId == null || tableNames == null || tableNames.isEmpty()) { + return List.of(); + } + + return semanticModelMapper.selectByDatasourceIdAndTableNames(datasourceId, tableNames); + } + @Override public SemanticModel getById(Long id) { return semanticModelMapper.selectById(id); @@ -49,6 +65,52 @@ public void addSemanticModel(SemanticModel semanticModel) { semanticModelMapper.insert(semanticModel); } + @Override + public boolean addSemanticModel(SemanticModelAddDTO dto) { + // 根据agentId查询关联的datasourceId + Integer datasourceId = findDatasourceIdByAgentId(dto.getAgentId()); + + // 转换DTO为Entity + SemanticModel semanticModel = SemanticModel.builder() + .agentId(dto.getAgentId()) + .datasourceId(datasourceId) + .tableName(dto.getTableName()) + .columnName(dto.getColumnName()) + .businessName(dto.getBusinessName()) + .synonyms(dto.getSynonyms()) + .businessDescription(dto.getBusinessDescription()) + .columnComment(dto.getColumnComment()) + .dataType(dto.getDataType()) + .status(1) // 默认启用状态 + .build(); + + // 保存到数据库 + semanticModelMapper.insert(semanticModel); + + return true; + } + + /** + * 根据agentId查找关联的datasourceId 如果有多个数据源，返回第一个启用的数据源 + */ + private Integer findDatasourceIdByAgentId(Integer agentId) { + List agentDatasources = agentDatasourceMapper.selectByAgentId(agentId); + + if (agentDatasources.isEmpty()) { + throw new RuntimeException("No datasource found for Agent ID " + agentId); + } + + // 优先返回启用的数据源 + for (AgentDatasource ad : agentDatasources) { + if (ad.getIsActive() != null && ad.getIsActive() == 1) { + return ad.getDatasourceId(); + } + } + + // 如果没有启用的数据源，返回第一个 + return agentDatasources.get(0).getDatasourceId(); + } + @Override public void enableSemanticModel(Long id) { semanticModelMapper.enableById(id); diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/vectorstore/AgentVectorStoreService.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/vectorstore/AgentVectorStoreService.java index 5485890d4..9dc2b5abe 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/vectorstore/AgentVectorStoreService.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/vectorstore/AgentVectorStoreService.java @@ -16,13 +16,12 @@ package com.alibaba.cloud.ai.service.vectorstore; -import java.util.List; -import java.util.Map; - -import org.springframework.ai.document.Document; - import com.alibaba.cloud.ai.request.AgentSearchRequest; import com.alibaba.cloud.ai.request.SchemaInitRequest; +import org.springframework.ai.document.Document; + +import java.util.List; +import java.util.Map; public interface AgentVectorStoreService { @@ -31,6 +30,7 @@ public interface AgentVectorStoreService { */ List search(AgentSearchRequest searchRequest); + // TODO 2025-11-10 后续应该移动到 schemaservice，本service 只负责数据处理 Boolean schema(String agentId, SchemaInitRequest schemaInitRequest) throws Exception; Boolean deleteDocumentsByVectorType(String agentId, String vectorType) throws Exception; @@ -42,10 +42,14 @@ public interface AgentVectorStoreService { */ List getDocumentsForAgent(String agentId, String query, String vectorType); + List getDocumentsOnlyByFilter(String filterExpression, int topK); + + List getTableDocuments(String agentId, List tableNames); + + List getColumnDocuments(String agentId, String upstreamTableName, List columnNames); + boolean hasDocuments(String agentId); void addDocuments(String agentId, List documents); - int estimateDocuments(String agentId); - } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/vectorstore/AgentVectorStoreServiceImpl.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/vectorstore/AgentVectorStoreServiceImpl.java index 1a6e162ad..364f3afc5 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/vectorstore/AgentVectorStoreServiceImpl.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/service/vectorstore/AgentVectorStoreServiceImpl.java @@ -15,6 +15,7 @@ */ package com.alibaba.cloud.ai.service.vectorstore; +import com.alibaba.cloud.ai.config.DataAgentProperties; import com.alibaba.cloud.ai.connector.accessor.Accessor; import com.alibaba.cloud.ai.connector.accessor.AccessorFactory; import com.alibaba.cloud.ai.connector.bo.DbQueryParameter; @@ -22,27 +23,29 @@ import com.alibaba.cloud.ai.connector.bo.TableInfoBO; import com.alibaba.cloud.ai.connector.config.DbConfig; import com.alibaba.cloud.ai.constant.Constant; +import com.alibaba.cloud.ai.constant.DocumentMetadataConstant; import com.alibaba.cloud.ai.request.AgentSearchRequest; import com.alibaba.cloud.ai.request.SchemaInitRequest; import com.alibaba.cloud.ai.service.TableMetadataService; -import jakarta.annotation.PreDestroy; +import com.alibaba.cloud.ai.service.hybrid.retrieval.HybridRetrievalStrategy; +import com.alibaba.cloud.ai.util.SearchUtil; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.BatchingStrategy; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.SimpleVectorStore; import org.springframework.ai.vectorstore.VectorStore; -import org.springframework.beans.factory.annotation.Value; import org.springframework.stereotype.Service; import org.springframework.util.Assert; -import org.springframework.util.StringUtils; import java.util.*; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; import java.util.concurrent.ExecutorService; -import java.util.stream.Collectors; import static com.alibaba.cloud.ai.util.DocumentConverterUtil.convertColumnsToDocuments; import static com.alibaba.cloud.ai.util.DocumentConverterUtil.convertTablesToDocuments; +import static com.alibaba.cloud.ai.util.SearchUtil.buildFilterExpressionString; @Slf4j @Service @@ -50,21 +53,13 @@ public class AgentVectorStoreServiceImpl implements AgentVectorStoreService { private static final String DEFAULT = "default"; - private final VectorStore vectorStore; // 由Spring Boot自动配置 + private final VectorStore vectorStore; private final ExecutorService dbOperationExecutor; - /** - * 相似度阈值配置，用于过滤相似度分数大于等于此阈值的文档 - */ - @Value("${spring.ai.vectorstore.similarityThreshold:0.2}") - protected double similarityThreshold; - - @Value("${spring.ai.vectorstore.batch-topk-limit:5000}") - protected int batchTopkLimit; + private final Optional hybridRetrievalStrategy; - @Value("${spring.ai.vectorstore.agent-query-topk-limit:25}") - protected int agentQueryTopkLimit; + private final DataAgentProperties dataAgentProperties; protected final AccessorFactory accessorFactory; @@ -73,10 +68,13 @@ public class AgentVectorStoreServiceImpl implements AgentVectorStoreService { protected final TableMetadataService tableMetadataService; public AgentVectorStoreServiceImpl(VectorStore vectorStore, ExecutorService dbOperationExecutor, + Optional hybridRetrievalStrategy, DataAgentProperties dataAgentProperties, BatchingStrategy batchingStrategy, AccessorFactory accessorFactory, TableMetadataService tableMetadataService) { this.vectorStore = vectorStore; this.dbOperationExecutor = dbOperationExecutor; + this.hybridRetrievalStrategy = hybridRetrievalStrategy; + this.dataAgentProperties = dataAgentProperties; this.batchingStrategy = batchingStrategy; this.accessorFactory = accessorFactory; this.tableMetadataService = tableMetadataService; @@ -85,23 +83,16 @@ public AgentVectorStoreServiceImpl(VectorStore vectorStore, ExecutorService dbOp @Override public List search(AgentSearchRequest searchRequest) { - Assert.notNull(searchRequest, "SearchRequest cannot be null"); - Assert.notNull(searchRequest.getAgentId(), "agentId cannot be null"); - org.springframework.ai.vectorstore.SearchRequest.Builder builder = org.springframework.ai.vectorstore.SearchRequest - .builder(); - - if (StringUtils.hasText(searchRequest.getQuery())) - builder.query(searchRequest.getQuery()); - - if (Objects.nonNull(searchRequest.getTopK())) - builder.topK(searchRequest.getTopK()); - - String filterFormatted = buildFilterExpressionString(searchRequest.getMetadataFilter()); - if (StringUtils.hasText(filterFormatted)) - builder.filterExpression(filterFormatted); - builder.similarityThreshold(similarityThreshold); - List results = vectorStore.similaritySearch(builder.build()); - log.info("Search completed. Found {} documents for SearchRequest: {}", results.size(), searchRequest); + Assert.hasText(searchRequest.getAgentId(), "AgentId cannot be empty"); + Assert.hasText(searchRequest.getDocVectorType(), "DocVectorType cannot be empty"); + if (dataAgentProperties.getVectorStore().isEnableHybridSearch() && hybridRetrievalStrategy.isPresent()) { + return hybridRetrievalStrategy.get().retrieve(searchRequest); + } + log.debug("Hybrid search is not enabled. use vector-search only"); + SearchRequest vectorSearchRequest = SearchUtil.buildVectorSearchRequest(searchRequest); + List results = vectorStore.similaritySearch(vectorSearchRequest); + log.debug("Search completed with vectorType: {}, found {} documents for SearchRequest: {}", + searchRequest.getDocVectorType(), results.size(), searchRequest); return results; } @@ -248,8 +239,8 @@ protected Map> buildForeignKeyMap(List fo } protected void clearSchemaDataForAgent(String agentId) throws Exception { - deleteDocumentsByVectorType(agentId, Constant.COLUMN); - deleteDocumentsByVectorType(agentId, Constant.TABLE); + deleteDocumentsByVectorType(agentId, DocumentMetadataConstant.COLUMN); + deleteDocumentsByVectorType(agentId, DocumentMetadataConstant.TABLE); } @Override @@ -257,8 +248,8 @@ public Boolean deleteDocumentsByVectorType(String agentId, String vectorType) th Assert.notNull(agentId, "AgentId cannot be null."); Assert.notNull(vectorType, "VectorType cannot be null."); - Map metadata = new HashMap<>( - Map.ofEntries(Map.entry(Constant.AGENT_ID, agentId), Map.entry(Constant.VECTOR_TYPE, vectorType))); + Map metadata = new HashMap<>(Map.ofEntries(Map.entry(Constant.AGENT_ID, agentId), + Map.entry(DocumentMetadataConstant.VECTOR_TYPE, vectorType))); return this.deleteDocumentsByMetedata(agentId, metadata); } @@ -278,75 +269,22 @@ public void addDocuments(String agentId, List documents) { vectorStore.add(documents); } - @Override - public int estimateDocuments(String agentId) { - Assert.notNull(agentId, "AgentId cannot be null."); - - // 使用相似度搜索来估计文档数量 - int totalCount = 0; - Set seenDocumentIds = new HashSet<>(); - int newDocumentsCount; - int maxIterations = 100; // 设置最大迭代次数，防止死循环 - int currentIteration = 0; - - do { - currentIteration++; - if (currentIteration > maxIterations) { - log.warn("Reached maximum iterations ({}) while estimating documents for agent: {}", maxIterations, - agentId); - break; - } - - // 使用不同的查询字符串来获取不同的文档集合 - // 使用迭代次数作为查询的一部分，确保每次查询略有不同 - String query = currentIteration == 1 ? "" : "iteration_" + currentIteration; - - List batch = vectorStore - .similaritySearch(org.springframework.ai.vectorstore.SearchRequest.builder() - .query(query) - .filterExpression(buildFilterExpressionString(Map.of(Constant.AGENT_ID, agentId))) - .topK(batchTopkLimit) - .similarityThreshold(0.0) // 设置最低相似度阈值以获取所有文档 - .build()); - - // 计算新文档数量 - newDocumentsCount = 0; - for (Document doc : batch) { - if (seenDocumentIds.add(doc.getId())) { - // 如果add返回true，表示这是一个新的文档ID - newDocumentsCount++; - } - } - - totalCount += newDocumentsCount; - - // 如果这批文档数量小于batchSize，说明已经获取了所有文档 - // 或者没有新文档，也说明已经获取了所有文档 - if (batch.size() < batchTopkLimit || newDocumentsCount == 0) { - break; - } - - } - while (newDocumentsCount > 0); // 只有当获取到新文档时才继续循环 - - log.info("Estimated {} documents for agent: {} after {} iterations", totalCount, agentId, currentIteration); - return totalCount; - } - @Override public Boolean deleteDocumentsByMetedata(String agentId, Map metadata) throws Exception { - Assert.notNull(agentId, "AgentId cannot be null."); + Assert.hasText(agentId, "AgentId cannot be empty."); Assert.notNull(metadata, "Metadata cannot be null."); // 添加agentId元数据过滤条件, 用于删除指定agentId下的所有数据，因为metadata中用户调用可能忘记添加agentId metadata.put(Constant.AGENT_ID, agentId); String filterExpression = buildFilterExpressionString(metadata); - // TODO 后续改成getVectorStore().delete(filterExpression); - // TODO 目前不支持通过元数据删除，使用会抛出UnsupportedOperationException，后续spring - // TODO ai发布1.1.0正式版本后再修改，现在是通过id删除 - - // 搜索和删除文档 - batchDelDocumentsWithFilter(filterExpression); + // es的可以直接元数据删除 + if (vectorStore instanceof SimpleVectorStore) { + // 目前SimpleVectorStore不支持通过元数据删除，使用会抛出UnsupportedOperationException,现在是通过id删除 + batchDelDocumentsWithFilter(filterExpression); + } + else { + vectorStore.delete(filterExpression); + } return true; } @@ -363,7 +301,7 @@ private void batchDelDocumentsWithFilter(String filterExpression) { .query(DEFAULT)// 使用默认的查询字符串，因为有的嵌入模型不支持空字符串 .filterExpression(filterExpression) .similarityThreshold(0.0)// 设置最低相似度阈值以获取元数据匹配的所有文档 - .topK(batchTopkLimit) + .topK(dataAgentProperties.getVectorStore().getBatchDelTopkLimit()) .build()); // 过滤掉已经处理过的文档，只删除未处理的文档 @@ -390,91 +328,49 @@ private void batchDelDocumentsWithFilter(String filterExpression) { log.info("Deleted {} documents with filter expression: {}", totalDeleted, filterExpression); } - /** - * 构建过滤表达式字符串，目前FilterExpressionBuilder 不支持链式拼接元数据过滤，所以只能使用字符串拼接 - * @param filterMap - * @return - */ - protected final String buildFilterExpressionString(Map filterMap) { - if (filterMap == null || filterMap.isEmpty()) { - return null; - } - - // 验证键名是否合法（只包含字母、数字和下划线） - for (String key : filterMap.keySet()) { - if (!key.matches("[a-zA-Z_][a-zA-Z0-9_]*")) { - throw new IllegalArgumentException("Invalid key name: " + key - + ". Keys must start with a letter or underscore and contain only alphanumeric characters and underscores."); - } - } - - return filterMap.entrySet().stream().map(entry -> { - String key = entry.getKey(); - Object value = entry.getValue(); - - // 处理空值 - if (value == null) { - return key + " == null"; - } - - // 根据值的类型决定如何格式化 - if (value instanceof String) { - // 转义字符串中的特殊字符 - String escapedValue = escapeStringLiteral((String) value); - return key + " == '" + escapedValue + "'"; - } - else if (value instanceof Number) { - // 数字类型直接使用 - return key + " == " + value; - } - else if (value instanceof Boolean) { - // 布尔值使用小写形式 - return key + " == " + ((Boolean) value).toString().toLowerCase(); - } - else if (value instanceof Enum) { - // 枚举类型，转换为字符串并转义 - String enumValue = ((Enum) value).name(); - String escapedValue = escapeStringLiteral(enumValue); - return key + " == '" + escapedValue + "'"; - } - else { - // 其他类型尝试转换为字符串并转义 - String stringValue = value.toString(); - String escapedValue = escapeStringLiteral(stringValue); - return key + " == '" + escapedValue + "'"; - } - }).collect(Collectors.joining(" && ")); + @Override + public List getDocumentsForAgent(String agentId, String query, String vectorType) { + AgentSearchRequest searchRequest = AgentSearchRequest.builder() + .agentId(agentId) + .docVectorType(vectorType) + .query(query) + .topK(dataAgentProperties.getVectorStore().getTopkLimit()) + .similarityThreshold(dataAgentProperties.getVectorStore().getSimilarityThreshold()) + .build(); + return search(searchRequest); } - /** - * 转义字符串字面量中的特殊字符 - */ - private String escapeStringLiteral(String input) { - if (input == null) { - return ""; - } - - // 转义反斜杠和单引号 - String escaped = input.replace("\\", "\\\\").replace("'", "\\'"); - - // 转义其他特殊字符 - escaped = escaped.replace("\n", "\\n") - .replace("\r", "\\r") - .replace("\t", "\\t") - .replace("\b", "\\b") - .replace("\f", "\\f"); - - return escaped; + @Override + public List getDocumentsOnlyByFilter(String filterExpression, int topK) { + Assert.hasText(filterExpression, "filterExpression cannot be empty."); + SearchRequest searchRequest = SearchRequest.builder() + .query(DEFAULT) + .topK(topK) + .filterExpression(filterExpression) + .similarityThreshold(0.0) + .build(); + return vectorStore.similaritySearch(searchRequest); } @Override - public List getDocumentsForAgent(String agentId, String query, String vectorType) { - AgentSearchRequest searchRequest = AgentSearchRequest.getInstance(agentId); - searchRequest.setQuery(query); - searchRequest.setTopK(agentQueryTopkLimit); - searchRequest.setMetadataFilter(Map.of(Constant.VECTOR_TYPE, vectorType)); + public List getTableDocuments(String agentId, List tableNames) { + Assert.hasText(agentId, "AgentId cannot be empty."); + if (tableNames.isEmpty()) + return Collections.emptyList(); + // 通过元数据过滤查找目标表 + String filterExpression = SearchUtil.buildFilterExpressionForSearchTables(agentId, tableNames); + return this.getDocumentsOnlyByFilter(filterExpression, tableNames.size() + 5); + } - return search(searchRequest); + @Override + public List getColumnDocuments(String agentId, String upstreamTableName, List columnNames) { + Assert.hasText(agentId, "AgentId cannot be empty."); + Assert.hasText(upstreamTableName, "UpstreamTableName cannot be empty."); + if (columnNames.isEmpty()) + return Collections.emptyList(); + String filterExpression = SearchUtil.buildFilterExpressionForSearchColumns(agentId, upstreamTableName, + columnNames); + return this.getDocumentsOnlyByFilter(filterExpression, columnNames.size() + 5); } @Override @@ -489,27 +385,4 @@ public boolean hasDocuments(String agentId) { return !docs.isEmpty(); } - /** - * 销毁方法，在Spring容器销毁bean时调用，用于关闭线程池 - */ - @PreDestroy - public void destroy() { - if (dbOperationExecutor != null && !dbOperationExecutor.isShutdown()) { - log.info("Shutting down database operation executor"); - dbOperationExecutor.shutdown(); - try { - // 等待正在执行的任务完成，最多等待30秒 - if (!dbOperationExecutor.awaitTermination(30, java.util.concurrent.TimeUnit.SECONDS)) { - log.warn("Database operation executor did not terminate gracefully, forcing shutdown"); - dbOperationExecutor.shutdownNow(); - } - } - catch (InterruptedException e) { - log.warn("Interrupted while waiting for database operation executor to terminate"); - dbOperationExecutor.shutdownNow(); - Thread.currentThread().interrupt(); - } - } - } - } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/DocumentConverterUtil.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/DocumentConverterUtil.java index f8b480eeb..8324503cd 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/DocumentConverterUtil.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/DocumentConverterUtil.java @@ -15,20 +15,23 @@ */ package com.alibaba.cloud.ai.util; -import java.util.*; -import java.util.stream.Collectors; - +import com.alibaba.cloud.ai.connector.bo.ColumnInfoBO; +import com.alibaba.cloud.ai.connector.bo.TableInfoBO; +import com.alibaba.cloud.ai.constant.Constant; +import com.alibaba.cloud.ai.constant.DocumentMetadataConstant; +import com.alibaba.cloud.ai.dto.BusinessKnowledgeDTO; import com.alibaba.cloud.ai.entity.AgentKnowledge; +import lombok.extern.slf4j.Slf4j; import org.springframework.ai.document.Document; -import com.alibaba.cloud.ai.connector.bo.ColumnInfoBO; -import com.alibaba.cloud.ai.connector.bo.TableInfoBO; -import com.alibaba.cloud.ai.request.EvidenceRequest; +import java.util.*; +import java.util.stream.Collectors; /** * Utility class for converting business objects to Document objects. Provides common * document conversion functionality for vector store operations. */ +@Slf4j public class DocumentConverterUtil { public static List convertColumnsToDocuments(String agentId, List tables) { @@ -68,7 +71,7 @@ public static Document convertColumnToDocumentForAgent(String agentId, TableInfo metadata.put("samples", columnInfoBO.getSamples()); } - return new Document(columnInfoBO.getName(), text, metadata); + return new Document(text, metadata); } /** @@ -86,7 +89,7 @@ public static Document convertTableToDocumentForAgent(String agentId, TableInfoB metadata.put("primaryKey", Optional.ofNullable(tableInfoBO.getPrimaryKeys()).orElse(new ArrayList<>())); metadata.put("vectorType", "table"); metadata.put("agentId", agentId); - return new Document(tableInfoBO.getName(), text, metadata); + return new Document(text, metadata); } public static List convertTablesToDocuments(String agentId, List tables) { @@ -95,28 +98,28 @@ public static List convertTablesToDocuments(String agentId, List convertEvidenceToDocumentsForAgent(String agentId, - List evidenceRequests) { - return evidenceRequests.stream().map(evidenceRequest -> { - Map metadata = new HashMap<>(); - metadata.put("evidenceType", evidenceRequest.getType()); - metadata.put("vectorType", "evidence"); - - metadata.put("agentId", agentId); - return new Document(UUID.randomUUID().toString(), evidenceRequest.getContent(), metadata); - }).toList(); + public static Document convertBusinessTermToDocument(BusinessKnowledgeDTO dto, String dbRecoredId) { + // 构建文档内容，包含业务名词、说明和同义词 + String businessTerm = Optional.ofNullable(dto.getBusinessTerm()).orElse("无"); + String description = Optional.ofNullable(dto.getDescription()).orElse("无"); + String synonyms = Optional.ofNullable(dto.getSynonyms()).orElse("无"); + + String content = String.format("业务名词: %s, 说明: %s, 同义词: %s", businessTerm, description, synonyms); + + // 构建元数据 + Map metadata = new HashMap<>(); + metadata.put(DocumentMetadataConstant.VECTOR_TYPE, DocumentMetadataConstant.BUSINESS_TERM); + metadata.put(Constant.AGENT_ID, dto.getAgentId().toString()); + metadata.put(DocumentMetadataConstant.DB_RECORD_ID, dbRecoredId); + metadata.put(DocumentMetadataConstant.IS_RECALL, dto.getIsRecall()); + + return new Document(content, metadata); } /** * Create Document from AgentKnowledge */ public static Document createDocumentFromKnowledge(String agentId, AgentKnowledge knowledge) { - String documentId = agentId + ":knowledge:" + knowledge.getId(); String content = knowledge.getContent(); if (content == null || content.trim().isEmpty()) { content = knowledge.getTitle(); // If content is empty, use title @@ -136,7 +139,7 @@ public static Document createDocumentFromKnowledge(String agentId, AgentKnowledg metadata.put("embeddingStatus", knowledge.getEmbeddingStatus()); metadata.put("createTime", knowledge.getCreateTime()); - return new Document(documentId, content, metadata); + return new Document(content, metadata); } /** diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/FluxUtil.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/FluxUtil.java index 4e8e075f4..8efdb7bc8 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/FluxUtil.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/FluxUtil.java @@ -80,7 +80,7 @@ public static Flux cascadeFlux(Flux originFlux, Function * @param completionMessage completion message * @param resultMapper result mapping function * @param sourceFlux source data stream - * @return AsyncGenerator instance + * @return Flux instance */ public static Flux> createStreamingGeneratorWithMessages( Class nodeClass, OverAllState state, String startMessage, String completionMessage, @@ -113,4 +113,28 @@ public static Flux> createStreamingGeneratorWithM return createStreamingGeneratorWithMessages(nodeClass, state, null, null, resultMapper, sourceFlux); } + /** + * create streaming generator with start and end flux + * @param nodeClass node class + * @param state state + * @param sourceFlux source data stream + * @param preFlux preFlux + * @param sufFlux sufFlux + * @param sourceMapper result of sourceFlux mapping function + * @return Flux instance + */ + public static Flux> createStreamingGenerator(Class nodeClass, + OverAllState state, Flux sourceFlux, Flux preFlux, Flux sufFlux, + Function> sourceMapper) { + String nodeName = nodeClass.getSimpleName(); + // Used to collect actual processing results + final StringBuilder collectedResult = new StringBuilder(); + sourceFlux = sourceFlux.doOnNext(r -> collectedResult.append(ChatResponseUtil.getText(r))); + return FluxConverter.builder() + .startingNode(nodeName) + .startingState(state) + .mapResult(r -> sourceMapper.apply(collectedResult.toString())) + .build(Flux.concat(preFlux, sourceFlux, sufFlux)); + } + } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/NodeBeanUtil.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/NodeBeanUtil.java index 2975a6699..1d9118941 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/NodeBeanUtil.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/NodeBeanUtil.java @@ -18,13 +18,10 @@ import com.alibaba.cloud.ai.graph.action.AsyncNodeAction; import com.alibaba.cloud.ai.graph.action.NodeAction; +import lombok.AllArgsConstructor; +import org.springframework.context.ApplicationContext; import org.springframework.stereotype.Component; -import java.util.List; -import java.util.Map; -import java.util.function.Function; -import java.util.stream.Collectors; - /** * 管理Node Bean * @@ -32,17 +29,13 @@ * @since 2025/9/28 */ @Component +@AllArgsConstructor public class NodeBeanUtil { - private final Map, ? extends NodeAction> nodeMap; - - public NodeBeanUtil(List actions) { - this.nodeMap = actions.stream() - .collect(Collectors.toUnmodifiableMap(NodeAction::getClass, Function.identity())); - } + private final ApplicationContext context; public NodeAction getNodeBean(Class clazz) { - return nodeMap.get(clazz); + return context.getBean(clazz); } public AsyncNodeAction getNodeBeanAsync(Class clazz) { diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/AbstractPlanBasedNode.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/PlanProcessUtil.java similarity index 64% rename from spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/AbstractPlanBasedNode.java rename to spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/PlanProcessUtil.java index 2648e1ea4..9addfdbeb 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/node/AbstractPlanBasedNode.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/PlanProcessUtil.java @@ -14,38 +14,42 @@ * limitations under the License. */ -package com.alibaba.cloud.ai.node; +package com.alibaba.cloud.ai.util; import com.alibaba.cloud.ai.graph.OverAllState; -import com.alibaba.cloud.ai.graph.action.NodeAction; import com.alibaba.cloud.ai.pojo.ExecutionStep; import com.alibaba.cloud.ai.pojo.Plan; -import lombok.extern.slf4j.Slf4j; import org.springframework.ai.converter.BeanOutputConverter; import org.springframework.core.ParameterizedTypeReference; +import java.util.HashMap; import java.util.List; +import java.util.Map; -import static com.alibaba.cloud.ai.constant.Constant.PLAN_CURRENT_STEP; import static com.alibaba.cloud.ai.constant.Constant.PLANNER_NODE_OUTPUT; +import static com.alibaba.cloud.ai.constant.Constant.PLAN_CURRENT_STEP; /** - * Abstract base class for plan-based execution nodes Provides common functionality for - * nodes that execute based on predefined plans + * util class for plan-based execution nodes Provides common functionality for nodes that + * execute based on predefined plans * * @author zhangshenghang */ -// todo: 改为工具类，注入到Node实现类中 -@Slf4j -public abstract class AbstractPlanBasedNode implements NodeAction { +public final class PlanProcessUtil { + + private static final BeanOutputConverter converter; - private final BeanOutputConverter converter; + private static final String STEP_PREFIX = "step_"; - protected AbstractPlanBasedNode() { - this.converter = new BeanOutputConverter<>(new ParameterizedTypeReference<>() { + static { + converter = new BeanOutputConverter<>(new ParameterizedTypeReference<>() { }); } + private PlanProcessUtil() { + + } + /** * Get the current execution step from the plan * @param state the overall state containing plan information @@ -53,16 +57,9 @@ protected AbstractPlanBasedNode() { * @throws IllegalStateException if plan output is empty, plan parsing fails, or step * index is out of range */ - protected ExecutionStep getCurrentExecutionStep(OverAllState state) { - String plannerNodeOutput = (String) state.value(PLANNER_NODE_OUTPUT) - .orElseThrow(() -> new IllegalStateException("计划节点输出为空")); - - Plan plan = converter.convert(plannerNodeOutput); - if (plan == null) { - throw new IllegalStateException("计划解析失败"); - } - - Integer currentStep = state.value(PLAN_CURRENT_STEP, 1); + public static ExecutionStep getCurrentExecutionStep(OverAllState state) { + Plan plan = getPlan(state); + int currentStep = getCurrentStepNumber(state); List executionPlan = plan.getExecutionPlan(); if (executionPlan == null || executionPlan.isEmpty()) { @@ -83,7 +80,7 @@ protected ExecutionStep getCurrentExecutionStep(OverAllState state) { * @return the parsed plan object * @throws IllegalStateException if plan output is empty or plan parsing fails */ - protected Plan getPlan(OverAllState state) { + public static Plan getPlan(OverAllState state) { String plannerNodeOutput = (String) state.value(PLANNER_NODE_OUTPUT) .orElseThrow(() -> new IllegalStateException("计划节点输出为空")); Plan plan = converter.convert(plannerNodeOutput); @@ -98,22 +95,22 @@ protected Plan getPlan(OverAllState state) { * @param state the overall state * @return the current step number (defaults to 1 if not set) */ - protected Integer getCurrentStepNumber(OverAllState state) { + public static int getCurrentStepNumber(OverAllState state) { return state.value(PLAN_CURRENT_STEP, 1); } /** - * Log node entry - */ - protected void logNodeEntry() { - log.info("Entering {} node", this.getClass().getSimpleName()); - } - - /** - * Log node output + * Add step result + * @param existingResults existing result collection + * @param stepNumber step number + * @param result result content + * @return updated result collection */ - protected void logNodeOutput(String outputKey, Object output) { - log.info("{} node output {}: {}", this.getClass().getSimpleName(), outputKey, output); + public static Map addStepResult(Map existingResults, Integer stepNumber, + String result) { + Map updatedResults = new HashMap<>(existingResults); + updatedResults.put(STEP_PREFIX + stepNumber, result); + return updatedResults; } } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/SearchUtil.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/SearchUtil.java new file mode 100644 index 000000000..fea9381b9 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/SearchUtil.java @@ -0,0 +1,213 @@ +/* + * Copyright 2024-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.alibaba.cloud.ai.util; + +import com.alibaba.cloud.ai.constant.Constant; +import com.alibaba.cloud.ai.constant.DocumentMetadataConstant; +import com.alibaba.cloud.ai.request.AgentSearchRequest; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder; +import org.springframework.util.StringUtils; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +public class SearchUtil { + + /** + * 构建向量搜索请求（普通搜索） + * @param request 搜索请求参数 + * @return 构建的搜索请求 + */ + public static SearchRequest buildVectorSearchRequest(AgentSearchRequest request) { + return buildVectorSearchRequest(request, false); + } + + /** + * 构建向量搜索请求 + * @param request 搜索请求参数 + * @param isHybridSearch 是否为混合搜索（需要RRF融合） + * @return 构建的搜索请求 + */ + public static SearchRequest buildVectorSearchRequest(AgentSearchRequest request, boolean isHybridSearch) { + SearchRequest.Builder builder = SearchRequest.builder(); + if (StringUtils.hasText(request.getQuery())) + builder.query(request.getQuery()); + + // 如果是混合搜索，提高topK为 RRF 等融合策略提供更丰富的候选集 + if (request.getTopK() != null) { + if (isHybridSearch) { + builder.topK(request.getTopK() * 2); + } + else { + builder.topK(request.getTopK()); + } + } + else { + builder.topK(isHybridSearch ? 40 : 20); + } + + builder.filterExpression(buildFilter(request.getAgentId(), request.getDocVectorType())); + builder.similarityThreshold(request.getSimilarityThreshold()); + return builder.build(); + } + + public static Filter.Expression buildFilter(String agentId, String docVectorType) { + FilterExpressionBuilder b = new FilterExpressionBuilder(); + return b.and(b.eq(DocumentMetadataConstant.VECTOR_TYPE, docVectorType), b.eq(Constant.AGENT_ID, agentId)) + .build(); + } + + /** + * 构建过滤表达式字符串，目前FilterExpressionBuilder 不支持链式拼接元数据过滤，所以只能使用字符串拼接 + * @param filterMap + * @return + */ + public static String buildFilterExpressionString(Map filterMap) { + if (filterMap == null || filterMap.isEmpty()) { + return null; + } + + // 验证键名是否合法（只包含字母、数字和下划线） + for (String key : filterMap.keySet()) { + if (!key.matches("[a-zA-Z_][a-zA-Z0-9_]*")) { + throw new IllegalArgumentException("Invalid key name: " + key + + ". Keys must start with a letter or underscore and contain only alphanumeric characters and underscores."); + } + } + + return filterMap.entrySet().stream().map(entry -> { + String key = entry.getKey(); + Object value = entry.getValue(); + + // 处理空值 + if (value == null) { + return key + " == null"; + } + + // 根据值的类型决定如何格式化 + if (value instanceof String) { + // 转义字符串中的特殊字符 + String escapedValue = escapeStringLiteral((String) value); + return key + " == '" + escapedValue + "'"; + } + else if (value instanceof Number) { + // 数字类型直接使用 + return key + " == " + value; + } + else if (value instanceof Boolean) { + // 布尔值使用小写形式 + return key + " == " + ((Boolean) value).toString().toLowerCase(); + } + else if (value instanceof Enum) { + // 枚举类型，转换为字符串并转义 + String enumValue = ((Enum) value).name(); + String escapedValue = escapeStringLiteral(enumValue); + return key + " == '" + escapedValue + "'"; + } + else { + // 其他类型尝试转换为字符串并转义 + String stringValue = value.toString(); + String escapedValue = escapeStringLiteral(stringValue); + return key + " == '" + escapedValue + "'"; + } + }).collect(Collectors.joining(" && ")); + } + + /** + * 转义字符串字面量中的特殊字符 + */ + public static String escapeStringLiteral(String input) { + if (input == null) { + return ""; + } + + // 转义反斜杠和单引号 + String escaped = input.replace("\\", "\\\\").replace("'", "\\'"); + + // 转义其他特殊字符 + escaped = escaped.replace("\n", "\\n") + .replace("\r", "\\r") + .replace("\t", "\\t") + .replace("\b", "\\b") + .replace("\f", "\\f"); + + return escaped; + } + + /** + * 构造元数据过滤表达式 + * @param agentId 代理ID + * @param tableNames 表名列表 + * @return SQL风格的过滤表达式字符串 + */ + public static String buildFilterExpressionForSearchTables(String agentId, List tableNames) { + // 输出:agentId == 'agent123' && vectorType == 'text' && name in ['users', 'orders', + // 'products'] + StringBuilder expression = new StringBuilder(); + expression.append("agentId == '").append(escapeStringLiteral(agentId)).append("'"); + expression.append(" && vectorType == '") + .append(escapeStringLiteral(DocumentMetadataConstant.TABLE)) + .append("'"); + + // 添加tableNames的IN条件（如果列表不为空） + if (tableNames != null && !tableNames.isEmpty()) { + String tableNamesList = tableNames.stream() + .map(SearchUtil::escapeStringLiteral) + .collect(Collectors.joining("', '", "'", "'")); + + expression.append(" && name in [").append(tableNamesList).append("]"); + } + + return expression.toString(); + } + + public static String buildFilterExpressionForSearchColumns(String agentId, String upstreamTableName, + List columnNames) { + // agentId == 'agent123' && tableName == 'users' && vectorType == 'column' && name + // in ['id', 'name', 'email'] + StringBuilder expression = new StringBuilder(); + + // 添加 agentId 条件 + expression.append("agentId == '").append(escapeStringLiteral(agentId)).append("'"); + + // 添加 tableName 条件 + expression.append(" && tableName == '").append(escapeStringLiteral(upstreamTableName)).append("'"); + + // 添加 vectorType 条件（固定为 column） + expression.append(" && vectorType == '") + .append(escapeStringLiteral(DocumentMetadataConstant.COLUMN)) + .append("'"); + + // 添加 columnNames 的 IN 条件（如果列表不为空） + if (columnNames != null && !columnNames.isEmpty()) { + String columnNamesList = columnNames.stream() + .map(SearchUtil::escapeStringLiteral) + .collect(Collectors.joining("', '", "'", "'")); + + expression.append(" && name in [").append(columnNamesList).append("]"); + } + + return expression.toString(); + } + + private SearchUtil() { + throw new AssertionError("Cannot instantiate utility class"); + } + +} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/StateUtil.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/StateUtil.java index 2eb208764..202cd35d5 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/StateUtil.java +++ b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/StateUtil.java @@ -16,6 +16,7 @@ package com.alibaba.cloud.ai.util; +import com.alibaba.cloud.ai.dto.QueryEnhanceOutputDTO; import com.alibaba.cloud.ai.graph.OverAllState; import com.fasterxml.jackson.databind.ObjectMapper; import org.springframework.ai.document.Document; @@ -24,6 +25,8 @@ import java.util.List; import java.util.function.Supplier; +import static com.alibaba.cloud.ai.constant.Constant.QUERY_ENHANCE_NODE_OUTPUT; + /** * State management utility class, providing type-safe state getting methods * @@ -114,10 +117,13 @@ public static List getDocumentList(OverAllState state, String key) { } /** - * Get list of Document lists + * Get canonical query */ - public static List> getDocumentListList(OverAllState state, String key) { - return getListValue(state, key); + public static String getCanonicalQuery(OverAllState state) { + QueryEnhanceOutputDTO queryEnhanceOutputDTO = getObjectValue(state, QUERY_ENHANCE_NODE_OUTPUT, + QueryEnhanceOutputDTO.class); + // 获取canonical_query + return queryEnhanceOutputDTO.getCanonicalQuery(); } } diff --git a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/StepResultUtil.java b/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/StepResultUtil.java deleted file mode 100644 index 1e78825be..000000000 --- a/spring-ai-alibaba-data-agent-chat/src/main/java/com/alibaba/cloud/ai/util/StepResultUtil.java +++ /dev/null @@ -1,77 +0,0 @@ -/* - * Copyright 2025 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.alibaba.cloud.ai.util; - -import java.util.HashMap; -import java.util.Map; - -/** - * Step result management utility class - * - * @author zhangshenghang - */ -public class StepResultUtil { - - private static final String STEP_PREFIX = "step_"; - - /** - * Add step result - * @param existingResults existing result collection - * @param stepNumber step number - * @param result result content - * @return updated result collection - */ - public static Map addStepResult(Map existingResults, Integer stepNumber, - String result) { - Map updatedResults = new HashMap<>(existingResults); - updatedResults.put(STEP_PREFIX + stepNumber, result); - return updatedResults; - } - - /** - * Get step result - * @param results result collection - * @param stepNumber step number - * @return step result, returns null if not exists - */ - public static String getStepResult(Map results, Integer stepNumber) { - return results.get(STEP_PREFIX + stepNumber); - } - - /** - * Check if step result exists - * @param results result collection - * @param stepNumber step number - * @return whether exists - */ - public static boolean hasStepResult(Map results, Integer stepNumber) { - return results.containsKey(STEP_PREFIX + stepNumber); - } - - /** - * Clear specific step result - * @param results result collection - * @param stepNumber step number - * @return result collection - */ - public static Map clearStepResult(Map results, Integer stepNumber) { - Map updatedResults = new HashMap<>(results); - updatedResults.remove(STEP_PREFIX + stepNumber); - return updatedResults; - } - -} diff --git a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/agent-knowledge.txt b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/agent-knowledge.txt new file mode 100644 index 000000000..e06bf3958 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/agent-knowledge.txt @@ -0,0 +1,20 @@ +### 领域知识库说明 +领域知识库包含从各类文档（如分析报告、产品手册、复盘总结）中提取的非结构化信息。 +这些知识旨在为智能体提供深入的背景信息、事实依据和过程分析，帮助其理解“为什么”和“怎么样”，而不仅仅是“是什么”。 + +#### 数据结构详解 +每个领域知识片段包含以下核心属性： +- `来源文档`：知识所源自的原始文档名称（如"2025年Q3用户增长分析报告.pdf"） +- `相关片段`：从文档中抽取的、与当前查询最相关的具体文本段落 + +#### 使用指南 +1. 将领域知识作为决策的背景和事实依据： + - 当需要解释原因、进行对比分析或总结情况时，应重点参考`相关片段`中的内容。 + - `来源文档`字段有助于追溯信息源头，判断信息的可信度和时效性。 + +2. 特别注意以下场景： + - 片段中可能包含观点、分析或假设，而不仅仅是客观事实。 + - 不同的文档片段之间可能存在时间上的先后关系或内容的补充。 + +#### 当前领域知识库 +{domainKnowledge} \ No newline at end of file diff --git a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/business-knowledge.txt b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/business-knowledge.txt index 605f9f0b5..ede78d08a 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/business-knowledge.txt +++ b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/business-knowledge.txt @@ -7,7 +7,6 @@ - `业务名词`：行业认可的标准术语（如"客户生命周期价值"） - `说明`：术语的详细业务定义和技术解释 - `同义词`：常见变体表达（逗号分隔，为空时表示无别名） -- `默认召回`：标识核心术语（仅true的条目参与匹配） #### 使用指南 1. 将用户原始查询与业务知识库结合分析： diff --git a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/feasibility-assessment.txt b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/feasibility-assessment.txt new file mode 100644 index 000000000..2ac1ff7c4 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/feasibility-assessment.txt @@ -0,0 +1,92 @@ +# 角色 +你是一名经验丰富且务实的数据分析专家。你的任务是结合用户的查询、相关的业务知识以及数据库结构，判断用户的请求是否有可能被回答，并决定是直接分析，还是需要澄清关键点。 + +# 核心原则 +你的判断应倾向于乐观和包容。只要用户的核心意图能够在召回的数据库Schema和业务参考信息中找到相关线索，就应该优先尝试数据分析。只在遇到无法绕过的关键信息缺失时，才请求澄清。 + +# 输入信息 +1. **【规范化查询】**: 经过预处理的用户最终需求。这是你的分析目标。 +2. **【召回的数据库Schema】**: 系统找到的相关表和字段。这是你的“数据结构”。 +3. **【参考信息 (Evidence)】**: 相关的业务规则、指标定义或术语解释。这是连接用户语言和数据库的“知识库”。 +4. **【多轮输入】**: 完整的对话历史上下文。 + +# 任务与决策逻辑 + +1. **综合分析可行性**: + * 首先，检查【规范化查询】中的核心概念是否能在【召回的数据库Schema】中直接找到对应字段。 + * 然后，如果找不到直接对应，立即查阅【参考信息 (Evidence)】，看是否有业务规则能将用户的术语（如“高价值客户”）映射到Schema中的字段和条件（如 `客户表` 和 `消费额 > 10000`）。 + * 综合判断，只要通过Schema和Evidence能覆盖查询的核心要求，就认为是可行的。 + +2. **确定需求类型**: + * **《数据分析》**: 只要【规范化查询】中的绝大多数关键信息，能通过**直接映射Schema**或**借助Evidence解释**后找到合理的对应，就选择此项。 + * **《需要澄清》**: 当且仅当遇到以下**严重阻碍**时，才选择此项： + * 核心实体或指标在Schema和Evidence中都**找不到任何解释或对应**。 + * 查询中的**决定性概念**（如“最好”）非常模糊，且Evidence中没有给出定义，同时Schema中存在多个可能的衡量标准。 + * **《自由闲聊》**: 当【召回的数据库Schema】为空，且查询与业务无关时，选择此项。 + +3. **生成需求内容**: + * 如果类型是《数据分析》，【需求内容】**直接使用【规范化查询】的文本**。 + * 如果类型是《需要澄清》，【需求内容】必须是一个简洁、聚焦的反问。 + * 如果类型是《自由闲聊》，【需求内容】应礼貌地告知无法回答。 + +# 输出格式 +【需求类型】：... +【语种类型】：... +【需求内容】：... + +--- +# 示例 + +## 示例1: 借助Evidence成功解析 -> 数据分析 +【规范化查询】: 查询所有“核心用户”的数量 +【召回的数据库Schema】: +# 表名: user, 包含字段: [(user_id), (name), (registration_date)] +# 表名: orders, 包含字段: [(order_id), (user_id), (order_amount), (order_date)] +【参考信息 (Evidence)】: “核心用户”被定义为最近30天内消费总额超过5000元的用户。 +【多轮输入】: ... +【输出】 +【需求类型】：《数据分析》 +【语种类型】：《中文》 +【需求内容】：查询所有“核心用户”的数量 +// (决策逻辑: Schema中没有“核心用户”字段，但Evidence提供了清晰的定义，使其可以被转换成对orders表的查询。因此可行。) + +## 示例2: 核心指标完全缺失 -> 需要澄清 +【规范化查询】: 统计所有部门的总毛利是多少 +【召回的数据库Schema】: +# 表名: 部门信息, 包含字段: [(部门编号), (部门名称), (员工人数)] +【参考信息 (Evidence)】: (无相关信息) +【多轮输入】: ... +【输出】 +【需求类型】：《需要澄清》 +【语种类型】：《中文》 +【需求内容】：抱歉，我没有找到关于“毛利”的直接数据或计算方式。我这里有关于部门的“员工人数”等信息，请问您还需要查询其他内容吗？ +// (决策逻辑: Schema和Evidence中都没有任何关于“毛利”或其组成部分（如营收、成本）的信息，障碍无法绕过。) + +## 示例3: 概念模糊且Evidence未定义 -> 需要澄清 +【规范化查询】: 找出最受欢迎的网络小说 +【召回的数据库Schema】: +# 表名: 网络小说, 包含字段: [(书名), (评分), (点击数)] +【参考信息 (Evidence)】: (无相关信息) +【多轮输入】: ... +【输出】 +【需求类型】：《需要澄清》 +【语种类型】：《中文》 +【需求内容】：请问您指的“最受欢迎”是按“评分”还是“点击数”来衡量呢？ +// (决策逻辑: “受欢迎”概念模糊，Evidence没有提供标准，而Schema中有多个可能选项，必须澄清。) + +--- +# 正式任务 + +【规范化查询】 +{canonical_query} + +【召回的数据库Schema】 +{recalled_schema} + +【参考信息 (Evidence)】 +{evidence} + +【多轮输入】 +{multi_turn} + +【输出】 \ No newline at end of file diff --git a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/intent-recognition.txt b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/intent-recognition.txt new file mode 100644 index 000000000..e3677fcc1 --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/intent-recognition.txt @@ -0,0 +1,68 @@ +# 角色 +你是一个轻量级、高效率的前置意图分类器，在整个数据分析工作流的最前端。 + +# 任务 +你的唯一任务是快速判断用户的最新输入是 **《闲聊或无关指令》** 还是 **《可能的数据分析请求》**，以便过滤掉明显无效的请求，节约后续昂贵的计算资源。 + +# 核心原则 +你的决策必须极端保守，遵循“宁放过，不杀错”的原则。也就是说，只要用户的输入有**一丝一毫**的可能性是想查询或分析数据，你就**必须**将其划分为《可能的数据分析请求》。只有当输入**明确、毫无疑问**是闲聊或与数据无关时，才进行拦截。 + +# 分类标准 + +## 《闲聊或无关指令》 +// 当且仅当输入完全符合以下特征之一时，才选择此项。 +- **纯情感或礼貌性用语**: "哈哈哈", "谢谢你！", "你好呀", "没问题" +- **关于AI自身的元问题**: "你是谁？", "你叫什么名字？", "你能干什么？" +- **与业务数据完全无关的指令或常识问题**: "帮我写一首诗", "今天天气怎么样？", "苹果公司的CEO是谁？" +- **无意义的乱码**: "asdfasdf", "123123" + +## 《可能的数据分析请求》 +// 只要输入沾上一点边，就选这个。 +- **包含明确的数据分析关键词**: "查询", "分析", "统计", "列表", "排名", "总和", "平均值", "对比一下" +- **包含业务名词**: 即使用户没有使用分析关键词，但提到了数据库里可能存在的实体或指标，如 "销售额", "xx部门", "那个员工" +- **模糊的询问**: 在多轮对话中，出现 "那个呢？", "他们的情况怎么样？", "具体一点" 这类指代性的、需要联系上下文的模糊问题，一律视为可能的数据分析。 +- **看似闲聊但可能涉及业务**: 例如 "我们公司哪个产品卖得最好？" 这句话口语化，但本质是数据查询。 + +# 输出格式 +请严格按照以下JSON格式输出，不要有任何额外的解释和说明，不要有json块(```json```)的输出标识。 +\{"classification": "分类名称"\} // 分类名称必须是《闲聊或无关指令》或《可能的数据分析请求》 + +--- +# 示例 + +## 示例1 +【多轮输入】 +用户: 查一下所有员工的工资 +<最新>用户输入: 哈哈哈哈，太棒了！ +# 输出 +\{"classification": "《闲聊或无关指令》"\} + +## 示例2 +【多轮输入】 +用户: 查一下所有员工的工资 +<最新>用户输入: 他们呢？ +# 输出 +\{"classification": "《可能的数据分析请求》"\} + +## 示例3 +【多轮输入】 +(无) +<最新>用户输入: 你好 +# 输出 +\{"classification": "《闲聊或无关指令》"\} + +## 示例4 +【多轮输入】 +(无) +<最新>用户输入: 给我看看销售部 +# 输出 +\{"classification": "《可能的数据分析请求》"\} + +--- +# 正式任务 + +【多轮输入】 +{multi_turn} +<最新>用户输入: {latest_query} + +# 输出 \ No newline at end of file diff --git a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/planner-nl2sql-only.txt b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/planner-nl2sql-only.txt index e12dd4175..9306ddba4 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/planner-nl2sql-only.txt +++ b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/planner-nl2sql-only.txt @@ -28,6 +28,12 @@ Based on the user's question, the following relevant database schemas have been ```sql {schema} ``` +【Reference information】 +{evidence} + +# SEMANTIC MODEL + +{semantic_model} # DATA-CENTRIC CHAIN OF THOUGHT (Internal Monologue) diff --git a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/planner.txt b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/planner.txt index c3992f758..cc562542a 100644 --- a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/planner.txt +++ b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/planner.txt @@ -42,6 +42,13 @@ Based on the user's question, the following relevant database schemas have been {schema} ``` +【Reference information】 +{evidence} + +# SEMANTIC MODEL + +{semantic_model} + # DATA-CENTRIC CHAIN OF THOUGHT (Internal Monologue) 1. **Understand Goal**: What is the user's ultimate business question? (e.g., "Analyze lead conversion quality"). diff --git a/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/query-enhancement.txt b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/query-enhancement.txt new file mode 100644 index 000000000..15ae3d05f --- /dev/null +++ b/spring-ai-alibaba-data-agent-chat/src/main/resources/prompts/query-enhancement.txt @@ -0,0 +1,78 @@ +# 角色 +你是一位顶级的自然语言理解专家，擅长融合业务知识，对用户查询进行澄清、转换、扩展和关键词提取。 + +# 上下文信息 +- 当前时间: {current_time_info} +- 参考信息 (Evidence): {evidence} + +# 任务 +你的任务是接收用户输入，并将其处理成一个结构化的JSON对象。你需要严格遵循以下三个步骤，充分利用所有上下文信息来完成任务： + +1. **查询澄清与规范化 (Clarify, Normalize & Translate)**: + * **理解与澄清**: 结合多轮对话历史，理解用户的完整意图，进行指代消解。 + * **时间转换**: 识别相对时间，并根据“当前时间”转换为绝对日期或范围。 + * **业务术语解析**: 查阅【参考信息 (Evidence)】，将查询中的业务术语（如“核心用户”）替换或补充为其具体的、可用数据描述的定义（如“消费总额超过5000元的用户”）。 + * **生成基准**: 综合以上所有信息，生成一个独立的、无歧义的、时间明确且业务术语已解析的“规范化查询”（canonical_query）。 + +2. **查询扩展 (Expand)**: + * 基于第一步生成的、信息极其丰富的“规范化查询”，生成2-3个语义相同但表达不同的“扩展问题”（expanded_queries）。 + +3. **关键词提取与合并 (Extract & Consolidate)**: + * 从“规范化查询”、所有“扩展问题”以及**解析业务术语过程中涉及的关键概念和数值**中，提取出所有核心关键词。 + * 对所有关键词进行去重、合并，形成一个最终的、最全面的“关键词列表”（consolidated_keywords）。 + +# 输出格式 +你必须严格按照下面的JSON格式输出，不能有任何多余的解释。 +\{ + "canonical_query": "对用户最终意图的单一、清晰的重写，包含绝对时间和解析后的业务术语", + "expanded_queries": [ + "基于完整信息的扩展问题表述 1", + "..." + ], + "consolidated_keywords": [ + "去重合并后的关键词1（包含日期、业务术语、解析后的指标等）", + "..." + ] +\} + +--- +# 示例 +[当前时间: 2025-11-08 11:11:12] +【参考信息 (Evidence)】: “核心用户”被定义为最近30天内消费总额超过5000元的用户。 +【多轮输入】 +<最新>用户输入: 帮我看看上个月的核心用户有多少 +# 输出 +\{ + "canonical_query": "查询上个月（2025-10-01至2025-10-31）期间，消费总额超过5000元的用户数量", + "expanded_queries": [ + "统计在2025年10月份，累计消费金额大于5000的客户总数是多少？", + "找出上个月消费超过5000元的核心用户有多少人" + ], + "consolidated_keywords": [ + "查询", + "上个月", + "2025-10-01", + "2025-10-31", + "期间", + "消费总额", + "超过", + "5000元", + "用户", + "客户", + "数量", + "统计", + "累计消费", + "金额", + "核心用户" + ] +\} +--- +# 正式任务 + +[当前时间: {current_time_info}] +[参考信息 (Evidence): {evidence}] +【多轮输入】 +{multi_turn} +<最新>用户输入: {latest_query} + +# 输出 \ No newline at end of file diff --git a/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/dispatcher/TableRelationDispatcherTest.java b/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/dispatcher/TableRelationDispatcherTest.java index 9d752346a..9b79e2be5 100644 --- a/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/dispatcher/TableRelationDispatcherTest.java +++ b/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/dispatcher/TableRelationDispatcherTest.java @@ -53,7 +53,7 @@ void testSuccessfulExecution() throws Exception { String result = dispatcher.apply(state); // Verify routing decision - assertEquals(PLANNER_NODE, result); + assertEquals(FEASIBILITY_ASSESSMENT_NODE, result); } @Test @@ -181,7 +181,7 @@ void testSuccessWithPreviousError() throws Exception { TABLE_RELATION_RETRY_COUNT, 0)); String result = dispatcher.apply(state); - assertEquals(PLANNER_NODE, result); + assertEquals(FEASIBILITY_ASSESSMENT_NODE, result); } } diff --git a/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/code/AgentVectorStoreServiceImplTest.java b/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/code/AgentVectorStoreServiceImplTest.java index 8603d1ef9..c827a1509 100644 --- a/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/code/AgentVectorStoreServiceImplTest.java +++ b/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/code/AgentVectorStoreServiceImplTest.java @@ -13,21 +13,21 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.alibaba.cloud.ai.service.code; -import com.alibaba.cloud.ai.connector.config.DbConfig; -import com.alibaba.cloud.ai.constant.Constant; +import com.alibaba.cloud.ai.config.DataAgentProperties; import com.alibaba.cloud.ai.connector.accessor.Accessor; import com.alibaba.cloud.ai.connector.accessor.AccessorFactory; import com.alibaba.cloud.ai.connector.bo.ColumnInfoBO; import com.alibaba.cloud.ai.connector.bo.ForeignKeyInfoBO; import com.alibaba.cloud.ai.connector.bo.TableInfoBO; +import com.alibaba.cloud.ai.connector.config.DbConfig; +import com.alibaba.cloud.ai.constant.Constant; +import com.alibaba.cloud.ai.constant.DocumentMetadataConstant; import com.alibaba.cloud.ai.request.AgentSearchRequest; import com.alibaba.cloud.ai.request.SchemaInitRequest; -import com.alibaba.cloud.ai.service.TableMetadataService; +import com.alibaba.cloud.ai.service.hybrid.retrieval.HybridRetrievalStrategy; import com.alibaba.cloud.ai.service.vectorstore.AgentVectorStoreServiceImpl; -import com.alibaba.cloud.ai.util.DocumentConverterUtil; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; @@ -40,28 +40,37 @@ import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.test.util.ReflectionTestUtils; +import org.testcontainers.shaded.com.google.common.util.concurrent.MoreExecutors; import java.util.*; +import java.util.concurrent.ExecutorService; import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.ArgumentMatchers.*; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.*; -/** - * Unit tests for {@link AgentVectorStoreServiceImpl} - */ @ExtendWith(MockitoExtension.class) -@DisplayName("AgentVectorStoreService Tests") -public class AgentVectorStoreServiceImplTest { +class AgentVectorStoreServiceImplTest { + + private static final String TEST_AGENT_ID = "test-agent-id"; + + private static final String TEST_VECTOR_TYPE = "test-vector-type"; + + private static final String TEST_QUERY = "test query"; @Mock private VectorStore vectorStore; + private ExecutorService directExecutor; + @Mock - private BatchingStrategy batchingStrategy; + private Optional hybridRetrievalStrategy; + + @Mock(lenient = true) + private DataAgentProperties dataAgentProperties; @Mock - private TableMetadataService tableMetadataService; + private BatchingStrategy batchingStrategy; @Mock private AccessorFactory accessorFactory; @@ -72,37 +81,36 @@ public class AgentVectorStoreServiceImplTest { @InjectMocks private AgentVectorStoreServiceImpl agentVectorStoreService; - private final String TEST_AGENT_ID = "test-agent-id"; - - private final String TEST_QUERY = "test query"; - - private final String TEST_VECTOR_TYPE = "column"; - @BeforeEach void setUp() { - // Set similarity threshold for testing - ReflectionTestUtils.setField(agentVectorStoreService, "similarityThreshold", 0.2); + // 使用 MoreExecutors.newDirectExecutorService() 创建单线程线程池 + directExecutor = MoreExecutors.newDirectExecutorService(); + + // Setup dataAgentProperties mock + DataAgentProperties.VectorStoreProperties vectorStoreProperties = new DataAgentProperties.VectorStoreProperties(); + vectorStoreProperties.setEnableHybridSearch(false); + vectorStoreProperties.setTopkLimit(10); + vectorStoreProperties.setSimilarityThreshold(0.7); + vectorStoreProperties.setBatchDelTopkLimit(100); + when(dataAgentProperties.getVectorStore()).thenReturn(vectorStoreProperties); } @Test @DisplayName("Test search with valid request") void testSearchWithValidRequest() { // Arrange - AgentSearchRequest searchRequest = AgentSearchRequest.getInstance(TEST_AGENT_ID); - searchRequest.setQuery(TEST_QUERY); - searchRequest.setTopK(5); - searchRequest.setMetadataFilter(Map.of(Constant.VECTOR_TYPE, TEST_VECTOR_TYPE)); + AgentSearchRequest searchRequest = AgentSearchRequest.builder() + .agentId(TEST_AGENT_ID) + .docVectorType(TEST_VECTOR_TYPE) + .query(TEST_QUERY) + .topK(5) + .similarityThreshold(0.8) + .build(); - List expectedDocuments = Arrays.asList( - new Document("Test content 1", - Map.of(Constant.AGENT_ID, TEST_AGENT_ID, Constant.VECTOR_TYPE, TEST_VECTOR_TYPE)), - new Document("Test content 2", - Map.of(Constant.AGENT_ID, TEST_AGENT_ID, Constant.VECTOR_TYPE, TEST_VECTOR_TYPE))); + List expectedDocuments = Arrays.asList(new Document("Document 1 content", Map.of("key1", "value1")), + new Document("Document 2 content", Map.of("key2", "value2"))); - // Mock the vectorStore.similaritySearch to return documents on first call and - // empty list on subsequent calls - when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(expectedDocuments) - .thenReturn(Collections.emptyList()); + when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(expectedDocuments); // Act List result = agentVectorStoreService.search(searchRequest); @@ -114,14 +122,50 @@ void testSearchWithValidRequest() { verify(vectorStore).similaritySearch(any(SearchRequest.class)); } + @Test + @DisplayName("Test search with hybrid search enabled") + void testSearchWithHybridSearchEnabled() { + // Arrange + DataAgentProperties.VectorStoreProperties vectorStoreProperties = new DataAgentProperties.VectorStoreProperties(); + vectorStoreProperties.setEnableHybridSearch(true); + when(dataAgentProperties.getVectorStore()).thenReturn(vectorStoreProperties); + + HybridRetrievalStrategy mockStrategy = mock(HybridRetrievalStrategy.class); + // 正确地mock Optional + when(hybridRetrievalStrategy.isPresent()).thenReturn(true); + when(hybridRetrievalStrategy.get()).thenReturn(mockStrategy); + + AgentSearchRequest searchRequest = AgentSearchRequest.builder() + .agentId(TEST_AGENT_ID) + .docVectorType(TEST_VECTOR_TYPE) + .query(TEST_QUERY) + .topK(5) + .similarityThreshold(0.8) + .build(); + + List expectedDocuments = Arrays.asList(new Document("Hybrid result 1", Map.of("key1", "value1")), + new Document("Hybrid result 2", Map.of("key2", "value2"))); + + when(mockStrategy.retrieve(any(AgentSearchRequest.class))).thenReturn(expectedDocuments); + + // Act + List result = agentVectorStoreService.search(searchRequest); + + // Assert + assertNotNull(result); + assertEquals(2, result.size()); + assertEquals(expectedDocuments, result); + verify(mockStrategy).retrieve(searchRequest); + verify(vectorStore, never()).similaritySearch(any(SearchRequest.class)); + } + @Test @DisplayName("Test addDocuments with valid input") void testAddDocuments() { - Map metadata = new HashMap<>(); - metadata.put(Constant.AGENT_ID, TEST_AGENT_ID); // Arrange - List documents = Arrays.asList(new Document("Test content 1", metadata), - new Document("Test content 2", metadata)); + List documents = Arrays.asList( + new Document("Document 1", Map.of(Constant.AGENT_ID, TEST_AGENT_ID, "key1", "value1")), + new Document("Document 2", Map.of(Constant.AGENT_ID, TEST_AGENT_ID, "key2", "value2"))); // Act agentVectorStoreService.addDocuments(TEST_AGENT_ID, documents); @@ -134,7 +178,7 @@ void testAddDocuments() { @DisplayName("Test addDocuments with null agentId should throw exception") void testAddDocumentsWithNullAgentId() { // Arrange - List documents = Arrays.asList(new Document("Test content")); + List documents = Arrays.asList(new Document("Document 1", Map.of(Constant.AGENT_ID, TEST_AGENT_ID))); // Act & Assert IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, @@ -152,90 +196,49 @@ void testAddDocumentsWithEmptyDocuments() { } @Test - @DisplayName("Test estimateDocuments with existing documents") - void testEstimateDocuments() { - // Arrange - List mockDocuments = Arrays.asList( - new Document("Test content 1", Map.of(Constant.AGENT_ID, TEST_AGENT_ID)), - new Document("Test content 2", Map.of(Constant.AGENT_ID, TEST_AGENT_ID))); - // Mock the vectorStore.similaritySearch to return documents on first call and - // empty list on subsequent calls using doReturn style for better mocking - doReturn(mockDocuments).doReturn(Collections.emptyList()) - .when(vectorStore) - .similaritySearch(any(SearchRequest.class)); - - // Act - int result = agentVectorStoreService.estimateDocuments(TEST_AGENT_ID); - - // Assert - assertEquals(2, result); - // Since the estimated Documents method is called multiple times until there are - // no new documents, we should verify that it has been called at least once - verify(vectorStore, atLeastOnce()).similaritySearch(any(SearchRequest.class)); - } - - @Test - @DisplayName("Test hasDocuments with existing documents returns true") - void testHasDocumentsWithExistingDocuments() { - // Arrange - List mockDocuments = Arrays - .asList(new Document("Test content", Map.of(Constant.AGENT_ID, TEST_AGENT_ID))); - when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(mockDocuments); - - // Act - boolean result = agentVectorStoreService.hasDocuments(TEST_AGENT_ID); - - // Assert - assertTrue(result); - verify(vectorStore).similaritySearch(any(SearchRequest.class)); - } - - @Test - @DisplayName("Test hasDocuments with no documents returns false") - void testHasDocumentsWithNoDocuments() { + @DisplayName("Test addDocuments with document missing agentId in metadata should throw exception") + void testAddDocumentsWithMissingAgentIdInMetadata() { // Arrange - when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(Collections.emptyList()); + List documents = Arrays.asList(new Document("Document 1", Map.of("key1", "value1"))); - // Act - boolean result = agentVectorStoreService.hasDocuments(TEST_AGENT_ID); - - // Assert - assertFalse(result); - verify(vectorStore).similaritySearch(any(SearchRequest.class)); + // Act & Assert + IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> agentVectorStoreService.addDocuments(TEST_AGENT_ID, documents)); + assertEquals("Document metadata must contain agentId.", exception.getMessage()); } @Test - @DisplayName("Test getDocumentsForAgent") - void testGetDocumentsForAgent() { + @DisplayName("Test addDocuments with document having different agentId in metadata should throw exception") + void testAddDocumentsWithDifferentAgentIdInMetadata() { // Arrange - List expectedDocuments = Arrays.asList(new Document("Test content", - Map.of(Constant.AGENT_ID, TEST_AGENT_ID, Constant.VECTOR_TYPE, TEST_VECTOR_TYPE))); - - when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(expectedDocuments); - - // Act - List result = agentVectorStoreService.getDocumentsForAgent(TEST_AGENT_ID, TEST_QUERY, - TEST_VECTOR_TYPE); + List documents = Arrays + .asList(new Document("Document 1", Map.of(Constant.AGENT_ID, "different-agent-id"))); - // Assert - assertNotNull(result); - assertEquals(1, result.size()); - assertEquals(expectedDocuments, result); - verify(vectorStore).similaritySearch(any(SearchRequest.class)); + // Act & Assert + IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, + () -> agentVectorStoreService.addDocuments(TEST_AGENT_ID, documents)); + assertEquals("Document metadata agentId does not match.", exception.getMessage()); } @Test @DisplayName("Test deleteDocumentsByVectorType") void testDeleteDocumentsByVectorType() throws Exception { // Arrange + // 创建一个新的测试实例，使用SimpleVectorStore类型的mock + VectorStore simpleVectorStore = mock(org.springframework.ai.vectorstore.SimpleVectorStore.class); + + // 使用反射设置vectorStore字段为SimpleVectorStore类型的mock + ReflectionTestUtils.setField(agentVectorStoreService, "vectorStore", simpleVectorStore); + List mockDocuments = Arrays.asList( new Document("Test content 1", - Map.of(Constant.AGENT_ID, TEST_AGENT_ID, Constant.VECTOR_TYPE, TEST_VECTOR_TYPE)), - new Document("Test content 2", - Map.of(Constant.AGENT_ID, TEST_AGENT_ID, Constant.VECTOR_TYPE, TEST_VECTOR_TYPE))); + Map.of(Constant.AGENT_ID, TEST_AGENT_ID, DocumentMetadataConstant.VECTOR_TYPE, + TEST_VECTOR_TYPE)), + new Document("Test content 2", Map.of(Constant.AGENT_ID, TEST_AGENT_ID, + DocumentMetadataConstant.VECTOR_TYPE, TEST_VECTOR_TYPE))); // Mock the vectorStore.similaritySearch to return documents on first call and // empty list on subsequent calls - when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(mockDocuments) + when(simpleVectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(mockDocuments) .thenReturn(Collections.emptyList()); // Act @@ -243,7 +246,7 @@ void testDeleteDocumentsByVectorType() throws Exception { // Assert assertTrue(result); - verify(vectorStore, times(2)).similaritySearch(any(SearchRequest.class)); + verify(simpleVectorStore, times(2)).similaritySearch(any(SearchRequest.class)); } @Test @@ -272,13 +275,19 @@ void testDeleteDocumentsByMetedata() throws Exception { metadata.put("key1", "value1"); metadata.put("key2", 123); + // 创建一个新的测试实例，使用SimpleVectorStore类型的mock + VectorStore simpleVectorStore = mock(org.springframework.ai.vectorstore.SimpleVectorStore.class); + + // 使用反射设置vectorStore字段为SimpleVectorStore类型的mock + ReflectionTestUtils.setField(agentVectorStoreService, "vectorStore", simpleVectorStore); + List mockDocuments = Arrays.asList( new Document("Test content 1", Map.of(Constant.AGENT_ID, TEST_AGENT_ID, "key1", "value1", "key2", 123)), new Document("Test content 2", Map.of(Constant.AGENT_ID, TEST_AGENT_ID, "key1", "value1", "key2", 123))); // Mock the vectorStore.similaritySearch to return documents on first call and // empty list on subsequent calls - when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(mockDocuments) + when(simpleVectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(mockDocuments) .thenReturn(Collections.emptyList()); // Act @@ -286,7 +295,7 @@ void testDeleteDocumentsByMetedata() throws Exception { // Assert assertTrue(result); - verify(vectorStore, times(2)).similaritySearch(any(SearchRequest.class)); + verify(simpleVectorStore, times(2)).similaritySearch(any(SearchRequest.class)); } @Test @@ -298,7 +307,7 @@ void testDeleteDocumentsByMetedataWithNullAgentId() { // Act & Assert IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> agentVectorStoreService.deleteDocumentsByMetedata(null, metadata)); - assertEquals("AgentId cannot be null.", exception.getMessage()); + assertEquals("AgentId cannot be empty.", exception.getMessage()); } @Test @@ -311,180 +320,232 @@ void testDeleteDocumentsByMetedataWithNullMetadata() { } @Test - @DisplayName("Test schema with valid input") - void testSchemaWithValidInput() throws Exception { + @DisplayName("Test schema with exception should return false") + void testSchemaWithException() throws Exception { // Arrange SchemaInitRequest schemaInitRequest = new SchemaInitRequest(); DbConfig dbConfig = new DbConfig(); - dbConfig.setSchema("test-schema"); schemaInitRequest.setDbConfig(dbConfig); - schemaInitRequest.setTables(Arrays.asList("table1", "table2")); - - List foreignKeys = Arrays.asList(new ForeignKeyInfoBO("table1", "col1", "table2", "col2")); - - // 创建带有列信息的表数据 - List columns1 = Arrays.asList(ColumnInfoBO.builder() - .name("id") - .type("BIGINT") - .description("Primary key") - .primary(true) - .notnull(true) - .build(), ColumnInfoBO.builder().name("name").type("VARCHAR").description("Name column").build()); - List columns2 = Arrays.asList(ColumnInfoBO.builder() - .name("id") - .type("BIGINT") - .description("Primary key") - .primary(true) - .notnull(true) - .build(), ColumnInfoBO.builder().name("value").type("INTEGER").description("Value column").build()); - - List tables = Arrays.asList(createTableInfoBO("table1", "Description 1", columns1), - createTableInfoBO("table2", "Description 2", columns2)); - - when(accessorFactory.getAccessorByDbConfig(any(DbConfig.class))).thenReturn(accessor); - when(accessor.showForeignKeys(any(DbConfig.class), any())).thenReturn(foreignKeys); - when(accessor.fetchTables(any(DbConfig.class), any())).thenReturn(tables); - - // 模拟DocumentConverterUtil.convertColumnsToDocuments和convertTablesToDocuments的返回值 - List columnDocs = DocumentConverterUtil.convertColumnsToDocuments(TEST_AGENT_ID, tables); - List tableDocs = DocumentConverterUtil.convertTablesToDocuments(TEST_AGENT_ID, tables); - - // 模拟批处理策略返回文档列表 - when(batchingStrategy.batch(any())).thenReturn(Arrays.asList(columnDocs, tableDocs)); + + when(accessorFactory.getAccessorByDbConfig(any(DbConfig.class))) + .thenThrow(new RuntimeException("Database connection failed")); // Act Boolean result = agentVectorStoreService.schema(TEST_AGENT_ID, schemaInitRequest); // Assert - assertTrue(result); + assertFalse(result); verify(accessorFactory).getAccessorByDbConfig(dbConfig); - verify(accessor).showForeignKeys(eq(dbConfig), any()); - verify(accessor).fetchTables(eq(dbConfig), any()); - verify(vectorStore, atLeastOnce()).add(any()); } @Test - @DisplayName("Test schema with exception should return false") - void testSchemaWithException() throws Exception { + @DisplayName("Test buildForeignKeyMap") + void testBuildForeignKeyMap() { // Arrange - SchemaInitRequest schemaInitRequest = new SchemaInitRequest(); - DbConfig dbConfig = new DbConfig(); - schemaInitRequest.setDbConfig(dbConfig); + List foreignKeys = Arrays.asList(new ForeignKeyInfoBO("table1", "col1", "table2", "col2"), + new ForeignKeyInfoBO("table3", "col3", "table4", "col4")); - when(accessorFactory.getAccessorByDbConfig(any(DbConfig.class))) - .thenThrow(new RuntimeException("Database connection failed")); + // Act + Map> result = (Map>) ReflectionTestUtils + .invokeMethod(agentVectorStoreService, "buildForeignKeyMap", foreignKeys); + + // Assert + assertNotNull(result); + assertEquals(4, result.size()); + assertTrue(result.containsKey("table1")); + assertTrue(result.containsKey("table2")); + assertTrue(result.containsKey("table3")); + assertTrue(result.containsKey("table4")); + assertEquals(1, result.get("table1").size()); + assertEquals(1, result.get("table2").size()); + assertEquals(1, result.get("table3").size()); + assertEquals(1, result.get("table4").size()); + assertTrue(result.get("table1").get(0).contains("table1.col1=table2.col2")); + assertTrue(result.get("table2").get(0).contains("table1.col1=table2.col2")); + assertTrue(result.get("table3").get(0).contains("table3.col3=table4.col4")); + assertTrue(result.get("table4").get(0).contains("table3.col3=table4.col4")); + } + + @Test + @DisplayName("Test partitionList") + void testPartitionList() { + // Arrange + List list = Arrays.asList("item1", "item2", "item3", "item4", "item5"); + int batchSize = 2; // Act - Boolean result = agentVectorStoreService.schema(TEST_AGENT_ID, schemaInitRequest); + @SuppressWarnings("unchecked") + List> result = (List>) ReflectionTestUtils.invokeMethod(agentVectorStoreService, + "partitionList", list, batchSize); // Assert - assertFalse(result); - verify(accessorFactory).getAccessorByDbConfig(dbConfig); + assertNotNull(result); + assertEquals(3, result.size()); + assertEquals(2, result.get(0).size()); + assertEquals(2, result.get(1).size()); + assertEquals(1, result.get(2).size()); + assertEquals("item1", result.get(0).get(0)); + assertEquals("item2", result.get(0).get(1)); + assertEquals("item3", result.get(1).get(0)); + assertEquals("item4", result.get(1).get(1)); + assertEquals("item5", result.get(2).get(0)); } @Test - @DisplayName("Test buildFilterExpressionString with null map") - void testBuildFilterExpressionStringWithNullMap() { + @DisplayName("Test getDocumentsForAgent") + void testGetDocumentsForAgent() { + // Arrange + AgentSearchRequest searchRequest = AgentSearchRequest.builder() + .agentId(TEST_AGENT_ID) + .docVectorType(TEST_VECTOR_TYPE) + .query(TEST_QUERY) + .topK(10) + .similarityThreshold(0.7) + .build(); + + List expectedDocuments = Arrays.asList(new Document("Document 1 content", Map.of("key1", "value1")), + new Document("Document 2 content", Map.of("key2", "value2"))); + + when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(expectedDocuments); + // Act - String result = (String) ReflectionTestUtils.invokeMethod(agentVectorStoreService, - "buildFilterExpressionString", (Map) null); + List result = agentVectorStoreService.getDocumentsForAgent(TEST_AGENT_ID, TEST_QUERY, + TEST_VECTOR_TYPE); // Assert - assertNull(result); + assertNotNull(result); + assertEquals(2, result.size()); + assertEquals(expectedDocuments, result); + verify(vectorStore).similaritySearch(any(SearchRequest.class)); } @Test - @DisplayName("Test buildFilterExpressionString with empty map") - void testBuildFilterExpressionStringWithEmptyMap() { + @DisplayName("Test getDocumentsOnlyByFilter") + void testGetDocumentsOnlyByFilter() { + // Arrange + String filterExpression = "agentId == 'test-agent-id' && vectorType == 'test-vector-type'"; + int topK = 5; + + List expectedDocuments = Arrays.asList(new Document("Document 1 content", Map.of("key1", "value1")), + new Document("Document 2 content", Map.of("key2", "value2"))); + + when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(expectedDocuments); + // Act - String result = (String) ReflectionTestUtils.invokeMethod(agentVectorStoreService, - "buildFilterExpressionString", new HashMap<>()); + List result = agentVectorStoreService.getDocumentsOnlyByFilter(filterExpression, topK); // Assert - assertNull(result); + assertNotNull(result); + assertEquals(2, result.size()); + assertEquals(expectedDocuments, result); + verify(vectorStore).similaritySearch(any(SearchRequest.class)); } @Test - @DisplayName("Test buildFilterExpressionString with valid map") - void testBuildFilterExpressionStringWithValidMap() { + @DisplayName("Test getTableDocuments") + void testGetTableDocuments() { // Arrange - Map filterMap = new HashMap<>(); - filterMap.put("key1", "value1"); - filterMap.put("key2", 123); - filterMap.put("key3", true); - filterMap.put("key4", null); + List tableNames = Arrays.asList("table1", "table2"); + + List expectedDocuments = Arrays.asList(new Document("Table 1 content", Map.of("key1", "value1")), + new Document("Table 2 content", Map.of("key2", "value2"))); + + when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(expectedDocuments); // Act - String result = (String) ReflectionTestUtils.invokeMethod(agentVectorStoreService, - "buildFilterExpressionString", filterMap); + List result = agentVectorStoreService.getTableDocuments(TEST_AGENT_ID, tableNames); // Assert assertNotNull(result); - assertTrue(result.contains("key1 == 'value1'")); - assertTrue(result.contains("key2 == 123")); - assertTrue(result.contains("key3 == true")); - assertTrue(result.contains("key4 == null")); - assertTrue(result.contains(" && ")); + assertEquals(2, result.size()); + assertEquals(expectedDocuments, result); + verify(vectorStore).similaritySearch(any(SearchRequest.class)); } @Test - @DisplayName("Test buildFilterExpressionString with invalid key should throw exception") - void testBuildFilterExpressionStringWithInvalidKey() { + @DisplayName("Test getTableDocuments with empty table names") + void testGetTableDocumentsWithEmptyTableNames() { // Arrange - Map filterMap = new HashMap<>(); - filterMap.put("invalid-key", "value1"); + List tableNames = Collections.emptyList(); - // Act & Assert - IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> ReflectionTestUtils - .invokeMethod(agentVectorStoreService, "buildFilterExpressionString", filterMap)); - assertTrue(exception.getMessage().contains("Invalid key name: invalid-key")); + // Act + List result = agentVectorStoreService.getTableDocuments(TEST_AGENT_ID, tableNames); + + // Assert + assertNotNull(result); + assertTrue(result.isEmpty()); + verify(vectorStore, never()).similaritySearch(any(SearchRequest.class)); } @Test - @DisplayName("Test buildForeignKeyMap") - void testBuildForeignKeyMap() { + @DisplayName("Test getColumnDocuments") + void testGetColumnDocuments() { // Arrange - List foreignKeys = Arrays.asList(new ForeignKeyInfoBO("table1", "col1", "table2", "col2"), - new ForeignKeyInfoBO("table3", "col3", "table4", "col4")); + String upstreamTableName = "table1"; + List columnNames = Arrays.asList("col1", "col2"); + + List expectedDocuments = Arrays.asList(new Document("Column 1 content", Map.of("key1", "value1")), + new Document("Column 2 content", Map.of("key2", "value2"))); + + when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(expectedDocuments); // Act - Map> result = (Map>) ReflectionTestUtils - .invokeMethod(agentVectorStoreService, "buildForeignKeyMap", foreignKeys); + List result = agentVectorStoreService.getColumnDocuments(TEST_AGENT_ID, upstreamTableName, + columnNames); // Assert assertNotNull(result); - assertEquals(4, result.size()); - assertTrue(result.containsKey("table1")); - assertTrue(result.containsKey("table2")); - assertTrue(result.containsKey("table3")); - assertTrue(result.containsKey("table4")); - assertEquals(1, result.get("table1").size()); - assertEquals(1, result.get("table2").size()); - assertEquals(1, result.get("table3").size()); - assertEquals(1, result.get("table4").size()); - assertTrue(result.get("table1").get(0).contains("table1.col1=table2.col2")); - assertTrue(result.get("table2").get(0).contains("table1.col1=table2.col2")); - assertTrue(result.get("table3").get(0).contains("table3.col3=table4.col4")); - assertTrue(result.get("table4").get(0).contains("table3.col3=table4.col4")); + assertEquals(2, result.size()); + assertEquals(expectedDocuments, result); + verify(vectorStore).similaritySearch(any(SearchRequest.class)); + } + + @Test + @DisplayName("Test getColumnDocuments with empty column names") + void testGetColumnDocumentsWithEmptyColumnNames() { + // Arrange + String upstreamTableName = "table1"; + List columnNames = Collections.emptyList(); + + // Act + List result = agentVectorStoreService.getColumnDocuments(TEST_AGENT_ID, upstreamTableName, + columnNames); + + // Assert + assertNotNull(result); + assertTrue(result.isEmpty()); + verify(vectorStore, never()).similaritySearch(any(SearchRequest.class)); } @Test - @DisplayName("Test escapeStringLiteral") - void testEscapeStringLiteral() { - // Test null input - String result1 = (String) ReflectionTestUtils.invokeMethod(agentVectorStoreService, "escapeStringLiteral", - (String) null); - assertEquals("", result1); - - // Test normal string - String result2 = (String) ReflectionTestUtils.invokeMethod(agentVectorStoreService, "escapeStringLiteral", - "normal string"); - assertEquals("normal string", result2); - - // Test string with special characters - String result3 = (String) ReflectionTestUtils.invokeMethod(agentVectorStoreService, "escapeStringLiteral", - "string'with\\special\nchars"); - assertEquals("string\\'with\\\\special\\nchars", result3); + @DisplayName("Test hasDocuments with existing documents") + void testHasDocumentsWithExistingDocuments() { + // Arrange + List mockDocuments = Arrays + .asList(new Document("Test content", Map.of(Constant.AGENT_ID, TEST_AGENT_ID))); + + when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(mockDocuments); + + // Act + boolean result = agentVectorStoreService.hasDocuments(TEST_AGENT_ID); + + // Assert + assertTrue(result); + verify(vectorStore).similaritySearch(any(SearchRequest.class)); + } + + @Test + @DisplayName("Test hasDocuments with no documents") + void testHasDocumentsWithNoDocuments() { + // Arrange + when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(Collections.emptyList()); + + // Act + boolean result = agentVectorStoreService.hasDocuments(TEST_AGENT_ID); + + // Assert + assertFalse(result); + verify(vectorStore).similaritySearch(any(SearchRequest.class)); } /** diff --git a/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/hybrid/fusion/RrfFusionStrategyTest.java b/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/hybrid/fusion/RrfFusionStrategyTest.java index 2cb724b67..cf01407c5 100644 --- a/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/hybrid/fusion/RrfFusionStrategyTest.java +++ b/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/hybrid/fusion/RrfFusionStrategyTest.java @@ -51,7 +51,7 @@ void testFuseResultsWithValidInput() { int topK = 5; // 执行测试 - List result = rrfFusionStrategy.fuseResults(vectorResults, keywordResults, topK); + List result = rrfFusionStrategy.fuseResults(topK, vectorResults, keywordResults); // 验证结果 assertNotNull(result); @@ -79,7 +79,7 @@ void testFuseResultsWithEmptyVectorResults() { int topK = 3; // 执行测试 - List result = rrfFusionStrategy.fuseResults(vectorResults, keywordResults, topK); + List result = rrfFusionStrategy.fuseResults(topK, vectorResults, keywordResults); // 验证结果 assertNotNull(result); @@ -98,7 +98,7 @@ void testFuseResultsWithEmptyKeywordResults() { int topK = 3; // 执行测试 - List result = rrfFusionStrategy.fuseResults(vectorResults, keywordResults, topK); + List result = rrfFusionStrategy.fuseResults(topK, vectorResults, keywordResults); // 验证结果 assertNotNull(result); @@ -116,7 +116,7 @@ void testFuseResultsWithBothEmptyResults() { int topK = 3; // 执行测试 - List result = rrfFusionStrategy.fuseResults(vectorResults, keywordResults, topK); + List result = rrfFusionStrategy.fuseResults(topK, vectorResults, keywordResults); // 验证结果 assertNotNull(result); @@ -134,7 +134,7 @@ void testFuseResultsWithZeroTopK() { int topK = 0; // 执行测试 - List result = rrfFusionStrategy.fuseResults(vectorResults, keywordResults, topK); + List result = rrfFusionStrategy.fuseResults(topK, vectorResults, keywordResults); // 验证结果 assertNotNull(result); @@ -152,7 +152,7 @@ void testFuseResultsWithDuplicateDocuments() { int topK = 3; // 执行测试 - List result = rrfFusionStrategy.fuseResults(vectorResults, keywordResults, topK); + List result = rrfFusionStrategy.fuseResults(topK, vectorResults, keywordResults); // 验证结果 assertNotNull(result); @@ -170,13 +170,13 @@ void testFuseResultsWithNullInput() { int topK = 3; // 执行测试 - 应该处理null输入而不抛出异常 - List result1 = rrfFusionStrategy.fuseResults(null, List.of(), topK); + List result1 = rrfFusionStrategy.fuseResults(topK, null, List.of()); assertNotNull(result1); - List result2 = rrfFusionStrategy.fuseResults(List.of(), null, topK); + List result2 = rrfFusionStrategy.fuseResults(topK, List.of(), null); assertNotNull(result2); - List result3 = rrfFusionStrategy.fuseResults(null, null, topK); + List result3 = rrfFusionStrategy.fuseResults(topK, null, null); assertNotNull(result3); } @@ -195,7 +195,7 @@ void testFuseResultsWithTopKLessThanTotalDocuments() { int topK = 4; // 小于总文档数(6) // 执行测试 - List result = rrfFusionStrategy.fuseResults(vectorResults, keywordResults, topK); + List result = rrfFusionStrategy.fuseResults(topK, vectorResults, keywordResults); // 验证结果 assertNotNull(result); diff --git a/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/hybrid/retrieval/AbstractHybridRetrievalStrategyTest.java b/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/hybrid/retrieval/AbstractHybridRetrievalStrategyTest.java index 772722370..acb50e2ca 100644 --- a/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/hybrid/retrieval/AbstractHybridRetrievalStrategyTest.java +++ b/spring-ai-alibaba-data-agent-chat/src/test/java/com/alibaba/cloud/ai/service/hybrid/retrieval/AbstractHybridRetrievalStrategyTest.java @@ -27,7 +27,6 @@ import org.springframework.ai.vectorstore.VectorStore; import org.testcontainers.shaded.com.google.common.util.concurrent.MoreExecutors; -import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.concurrent.ExecutorService; @@ -77,12 +76,14 @@ void setUp() { } @Test - void retrieve_ShouldFuseVectorAndKeywordResults_WhenKeywordsArePresent() { + void retrieve_ShouldFuseVectorAndKeywordResults_WhenQueryIsPresent() { // 1. 准备 (Arrange) - AgentSearchRequest request = AgentSearchRequest.getInstance("agent1"); - request.setDocVectorType("test-type"); - request.setKeywords(Arrays.asList("a", "b")); - request.setTopK(10); + AgentSearchRequest request = AgentSearchRequest.builder() + .agentId("agent1") + .docVectorType("test-type") + .query("test query") + .topK(10) + .build(); List vectorResults = List.of(new Document("vec_doc1"), new Document("vec_doc2")); List keywordResults = List.of(new Document("key_doc1")); @@ -92,7 +93,7 @@ void retrieve_ShouldFuseVectorAndKeywordResults_WhenKeywordsArePresent() { when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(vectorResults); // 使用 spy 来 mock getDocumentsByKeywords 方法 org.mockito.Mockito.doReturn(keywordResults).when(retrievalStrategy).getDocumentsByKeywords(request); - when(fusionStrategy.fuseResults(vectorResults, keywordResults, 10)).thenReturn(fusedResults); + when(fusionStrategy.fuseResults(10, vectorResults, keywordResults)).thenReturn(fusedResults); // 2. 执行 (Act) // 因为我们用了 directExecutor，这里的调用会同步执行所有逻辑 @@ -104,21 +105,32 @@ void retrieve_ShouldFuseVectorAndKeywordResults_WhenKeywordsArePresent() { } @Test - void retrieve_ShouldReturnOnlyVectorResults_WhenKeywordsAreEmpty() { + void retrieve_ShouldReturnOnlyVectorResults_WhenQueryIsEmpty() { // 1. 准备 (Arrange) - AgentSearchRequest request = AgentSearchRequest.getInstance("agent1"); - request.setDocVectorType("test-type"); - request.setTopK(10); - - List vectorResults = List.of(new Document("vec_doc1"), new Document("vec_doc2")); + AgentSearchRequest request = AgentSearchRequest.builder() + .agentId("agent1") + .docVectorType("test-type") + .query("") // 空查询字符串 + .topK(10) + .build(); + + Document doc1 = new Document("vec_doc1"); + Document doc2 = new Document("vec_doc2"); + List vectorResults = List.of(doc1, doc2); + List keywordResults = Collections.emptyList(); // 空查询应该返回空结果 + List fusedResults = List.of(doc1, doc2); // 融合后应该只有向量结果 + // 定义 mock 对象的行为 when(vectorStore.similaritySearch(any(SearchRequest.class))).thenReturn(vectorResults); + // 使用 spy 来 mock getDocumentsByKeywords 方法 + org.mockito.Mockito.doReturn(keywordResults).when(retrievalStrategy).getDocumentsByKeywords(request); + when(fusionStrategy.fuseResults(10, vectorResults, keywordResults)).thenReturn(fusedResults); // 2. 执行 (Act) List finalDocuments = retrievalStrategy.retrieve(request); // 3. 断言 (Assert) - assertEquals(vectorResults.size(), finalDocuments.size()); + assertEquals(fusedResults.size(), finalDocuments.size()); assertEquals(vectorResults, finalDocuments); } diff --git a/spring-ai-alibaba-data-agent-chat/src/test/resources/sql/schema.sql b/spring-ai-alibaba-data-agent-chat/src/test/resources/sql/schema.sql index bf8b005f5..89bc330c5 100644 --- a/spring-ai-alibaba-data-agent-chat/src/test/resources/sql/schema.sql +++ b/spring-ai-alibaba-data-agent-chat/src/test/resources/sql/schema.sql @@ -28,39 +28,37 @@ CREATE TABLE IF NOT EXISTS business_knowledge ( description TEXT COMMENT '描述', synonyms TEXT COMMENT '同义词', is_recall INT DEFAULT 0 COMMENT '是否召回', - data_set_id VARCHAR(255) COMMENT '数据集id', - agent_id INT COMMENT '关联的智能体ID', + agent_id INT NOT NULL COMMENT '关联的智能体ID', created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间', updated_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间', PRIMARY KEY (id), INDEX idx_business_term (business_term), - INDEX idx_data_set_id (data_set_id), INDEX idx_agent_id (agent_id), INDEX idx_is_recall (is_recall), - FOREIGN KEY (agent_id) REFERENCES agent(id) ON DELETE SET NULL + FOREIGN KEY (agent_id) REFERENCES agent(id) ON DELETE CASCADE ) ENGINE = InnoDB COMMENT = '业务知识表'; -- 语义模型表 -CREATE TABLE IF NOT EXISTS semantic_model ( - id INT NOT NULL AUTO_INCREMENT, - agent_id INT COMMENT '关联的智能体ID', - field_name VARCHAR(255) NOT NULL DEFAULT '' COMMENT '智能体字段名称', - synonyms TEXT COMMENT '字段名称同义词', - origin_name VARCHAR(255) DEFAULT '' COMMENT '原始字段名', - description TEXT COMMENT '字段描述', - origin_description VARCHAR(255) COMMENT '原始字段描述', - type VARCHAR(255) DEFAULT '' COMMENT '字段类型 (integer, varchar....)', - created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间', - updated_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间', - is_recall TINYINT DEFAULT 0 COMMENT '0 停用 1 启用', - status TINYINT DEFAULT 0 COMMENT '0 停用 1 启用', - PRIMARY KEY (id), - INDEX idx_agent_id (agent_id), - INDEX idx_field_name (field_name), - INDEX idx_status (status), - INDEX idx_is_recall (is_recall), - FOREIGN KEY (agent_id) REFERENCES agent(id) ON DELETE SET NULL -) ENGINE = InnoDB COMMENT = '语义模型表'; +CREATE TABLE IF NOT EXISTS `semantic_model` ( + `id` int(11) NOT NULL AUTO_INCREMENT, + `agent_id` int(11) NOT NULL COMMENT '关联的智能体ID', + `datasource_id` int(11) NOT NULL COMMENT '关联的数据源ID', + `table_name` varchar(255) COLLATE utf8mb4_bin NOT NULL COMMENT '关联的表名', + `column_name` varchar(255) COLLATE utf8mb4_bin NOT NULL DEFAULT '' COMMENT '数据库中的物理字段名 (例如: csat_score)', + `business_name` varchar(255) COLLATE utf8mb4_bin NOT NULL DEFAULT '' COMMENT '业务名/别名 (例如: 客户满意度分数)', + `synonyms` text COLLATE utf8mb4_bin COMMENT '业务名的同义词 (例如: 满意度,客户评分)', + `business_description` text COLLATE utf8mb4_bin COMMENT '业务描述 (用于向LLM解释字段的业务含义)', + `column_comment` varchar(255) COLLATE utf8mb4_bin DEFAULT NULL COMMENT '数据库中的物理字段的原始注释 ', + `data_type` varchar(255) COLLATE utf8mb4_bin NOT NULL DEFAULT '' COMMENT '物理数据类型 (例如: int, varchar(20))', + `status` tinyint(4) NOT NULL DEFAULT '1' COMMENT '0 停用 1 启用', + `created_time` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间', + `updated_time` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间', + PRIMARY KEY (`id`) USING BTREE, + KEY `idx_agent_id` (`agent_id`) USING BTREE, + KEY `idx_field_name` (`business_name`) USING BTREE, + KEY `idx_status` (`status`) USING BTREE, + CONSTRAINT `fk_semantic_model_agent` FOREIGN KEY (`agent_id`) REFERENCES `agent` (`id`) ON DELETE CASCADE +) ENGINE=InnoDB AUTO_INCREMENT=4 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin ROW_FORMAT=DYNAMIC COMMENT='语义模型表'; -- 智能体知识表 diff --git a/spring-ai-alibaba-data-agent-common/pom.xml b/spring-ai-alibaba-data-agent-common/pom.xml index 3f86ad6d3..84612d027 100644 --- a/spring-ai-alibaba-data-agent-common/pom.xml +++ b/spring-ai-alibaba-data-agent-common/pom.xml @@ -63,6 +63,10 @@ ${spring-ai-alibaba.version} compile + + org.projectlombok + lombok + diff --git a/spring-ai-alibaba-data-agent-common/src/main/java/com/alibaba/cloud/ai/request/AgentSearchRequest.java b/spring-ai-alibaba-data-agent-common/src/main/java/com/alibaba/cloud/ai/request/AgentSearchRequest.java index fa5cbb33e..8e591add5 100644 --- a/spring-ai-alibaba-data-agent-common/src/main/java/com/alibaba/cloud/ai/request/AgentSearchRequest.java +++ b/spring-ai-alibaba-data-agent-common/src/main/java/com/alibaba/cloud/ai/request/AgentSearchRequest.java @@ -16,145 +16,25 @@ package com.alibaba.cloud.ai.request; -import java.io.Serial; -import java.io.Serializable; -import java.util.List; -import java.util.Map; -import java.util.Objects; +import lombok.Builder; +import lombok.Data; -public class AgentSearchRequest implements Serializable { +@Data +@Builder +public class AgentSearchRequest implements java.io.Serializable { - @Serial + @java.io.Serial private static final long serialVersionUID = 1L; - private final String agentId; + private String agentId; private String docVectorType; - private List keywords; - + @Builder.Default private Double similarityThreshold = 0.2; private String query; private Integer topK; - // TODO 待废弃，因为元数据的过滤只需要考虑agentId,vectorType - private Map metadataFilter; - - /** - * 私有构造函数，禁止直接实例化 此构造函数仅内部使用，外部代码必须通过getInstance()方法创建实例 - */ - private AgentSearchRequest(String agentId) { - Objects.requireNonNull(agentId, "Agent ID cannot be null"); - this.agentId = agentId; - // 初始化metadataFilter，确保始终包含agentId - this.metadataFilter = Map.of("agentId", agentId); - } - - /** - * 创建AgentSearchRequest实例的工厂方法 - * @param agentId 代理ID，不能为空 - * @return AgentSearchRequest实例 - * @throws IllegalArgumentException 如果agentId为空 - */ - public static AgentSearchRequest getInstance(String agentId) { - return new AgentSearchRequest(agentId); - } - - /** - * 创建AgentSearchRequest实例的工厂方法 - * @param agentId 代理ID，不能为空 - * @param query 查询内容 - * @return AgentSearchRequest实例 - * @throws IllegalArgumentException 如果agentId为空 - */ - public static AgentSearchRequest getInstance(String agentId, String query) { - AgentSearchRequest request = new AgentSearchRequest(agentId); - request.setQuery(query); - return request; - } - - /** - * 创建AgentSearchRequest实例的工厂方法 - * @param agentId 代理ID，不能为空 - * @param query 查询内容 - * @param topK 返回结果数量 - * @return AgentSearchRequest实例 - * @throws IllegalArgumentException 如果agentId为空 - */ - public static AgentSearchRequest getInstance(String agentId, String query, Integer topK) { - AgentSearchRequest request = new AgentSearchRequest(agentId); - request.setQuery(query); - request.setTopK(topK); - return request; - } - - public String getAgentId() { - return agentId; - } - - public Map getMetadataFilter() { - return metadataFilter; - } - - public void setMetadataFilter(Map metadataFilter) { - if (metadataFilter == null) { - // 如果传入null，则创建只包含agentId的map - this.metadataFilter = Map.of("agentId", agentId); - } - else { - // 创建新的map，包含传入的所有参数和agentId - Map newFilter = new java.util.HashMap<>(metadataFilter); - newFilter.put("agentId", agentId); - this.metadataFilter = Map.copyOf(newFilter); // 创建不可变副本 - } - } - - public String getQuery() { - return query; - } - - public void setQuery(String query) { - this.query = query; - } - - public Integer getTopK() { - return topK; - } - - public void setTopK(Integer topK) { - this.topK = topK; - } - - @Override - public String toString() { - return "AgentSearchRequest{" + "agentId='" + agentId + '\'' + ", query='" + query + '\'' + ", topK=" + topK - + ", metadataFilter=" + metadataFilter + '}'; - } - - public String getDocVectorType() { - return docVectorType; - } - - public void setDocVectorType(String docVectorType) { - this.docVectorType = docVectorType; - } - - public List getKeywords() { - return keywords; - } - - public void setKeywords(List keywords) { - this.keywords = keywords; - } - - public Double getSimilarityThreshold() { - return similarityThreshold; - } - - public void setSimilarityThreshold(Double similarityThreshold) { - this.similarityThreshold = similarityThreshold; - } - } diff --git a/spring-ai-alibaba-data-agent-frontend/src/components/agent/BusinessKnowledgeConfig.vue b/spring-ai-alibaba-data-agent-frontend/src/components/agent/BusinessKnowledgeConfig.vue index 32246a401..af977128f 100644 --- a/spring-ai-alibaba-data-agent-frontend/src/components/agent/BusinessKnowledgeConfig.vue +++ b/spring-ai-alibaba-data-agent-frontend/src/components/agent/BusinessKnowledgeConfig.vue @@ -15,7 +15,7 @@ -->