diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/AgentController.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/AgentController.java index 84a3a36a8..01d8e46bb 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/AgentController.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/AgentController.java @@ -19,9 +19,11 @@ import com.alibaba.cloud.ai.dataagent.service.agent.AgentService; import com.alibaba.cloud.ai.dataagent.vo.ApiKeyResponse; import com.alibaba.cloud.ai.dataagent.vo.ApiResponse; +import java.util.List; import lombok.AllArgsConstructor; import lombok.extern.slf4j.Slf4j; -import org.springframework.http.ResponseEntity; +import org.apache.commons.lang3.StringUtils; +import org.springframework.http.HttpStatus; import org.springframework.web.bind.annotation.CrossOrigin; import org.springframework.web.bind.annotation.DeleteMapping; import org.springframework.web.bind.annotation.GetMapping; @@ -32,12 +34,9 @@ import org.springframework.web.bind.annotation.RequestMapping; import org.springframework.web.bind.annotation.RequestParam; import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.server.ResponseStatusException; -import java.util.List; - -/** - * Agent Management Controller - */ +/** Agent Management Controller */ @Slf4j @RestController @RequestMapping("/api/agent") @@ -47,168 +46,117 @@ public class AgentController { private final AgentService agentService; - /** - * Get agent list - */ + /** Get agent list */ @GetMapping("/list") - public ResponseEntity> list(@RequestParam(value = "status", required = false) String status, + public List list(@RequestParam(value = "status", required = false) String status, @RequestParam(value = "keyword", required = false) String keyword) { List result; - if (keyword != null && !keyword.trim().isEmpty()) { + if (StringUtils.isNotBlank(keyword)) { result = agentService.search(keyword); } - else if (status != null && !status.trim().isEmpty()) { + else if (StringUtils.isNotBlank(status)) { result = agentService.findByStatus(status); } else { result = agentService.findAll(); } - return ResponseEntity.ok(result); + return result; } - /** - * Get agent details by ID - */ + /** Get agent details by ID */ @GetMapping("/{id}") - public ResponseEntity get(@PathVariable(value = "id") Long id) { - Agent agent = agentService.findById(id); - if (agent == null) { - return ResponseEntity.notFound().build(); - } - return ResponseEntity.ok(agent); + public Agent get(@PathVariable Long id) { + return checkAgentExists(id); } - /** - * Create agent - */ + /** Create agent */ @PostMapping - public ResponseEntity create(@RequestBody Agent agent) { + public Agent create(@RequestBody Agent agent) { // Set default status - if (agent.getStatus() == null || agent.getStatus().trim().isEmpty()) { + if (StringUtils.isBlank(agent.getStatus())) { agent.setStatus("draft"); } - Agent saved = agentService.save(agent); - return ResponseEntity.ok(saved); + return agentService.save(agent); } - /** - * Update agent - */ + /** Update agent */ @PutMapping("/{id}") - public ResponseEntity update(@PathVariable(value = "id") Long id, @RequestBody Agent agent) { - if (agentService.findById(id) == null) { - return ResponseEntity.notFound().build(); - } + public Agent update(@PathVariable Long id, @RequestBody Agent agent) { + checkAgentExists(id); agent.setId(id); - Agent updated = agentService.save(agent); - return ResponseEntity.ok(updated); + return agentService.save(agent); } - /** - * Delete agent - */ + /** Delete agent */ @DeleteMapping("/{id}") - public ResponseEntity delete(@PathVariable(value = "id") Long id) { - if (agentService.findById(id) == null) { - return ResponseEntity.notFound().build(); - } + public void delete(@PathVariable Long id) { + checkAgentExists(id); agentService.deleteById(id); - return ResponseEntity.ok().build(); } - /** - * Publish agent - */ + /** Publish agent */ @PostMapping("/{id}/publish") - public ResponseEntity publish(@PathVariable(value = "id") Long id) { - Agent agent = agentService.findById(id); - if (agent == null) { - return ResponseEntity.notFound().build(); - } + public Agent publish(@PathVariable Long id) { + Agent agent = checkAgentExists(id); agent.setStatus("published"); - Agent updated = agentService.save(agent); - return ResponseEntity.ok(updated); + return agentService.save(agent); } - /** - * Offline agent - */ + /** Offline agent */ @PostMapping("/{id}/offline") - public ResponseEntity offline(@PathVariable(value = "id") Long id) { - Agent agent = agentService.findById(id); - if (agent == null) { - return ResponseEntity.notFound().build(); - } + public Agent offline(@PathVariable Long id) { + Agent agent = checkAgentExists(id); agent.setStatus("offline"); - Agent updated = agentService.save(agent); - return ResponseEntity.ok(updated); + return agentService.save(agent); } - /** - * Get masked API Key status - */ + /** Get masked API Key status */ @GetMapping("/{id}/api-key") - public ResponseEntity> getApiKey(@PathVariable("id") Long id) { - Agent agent = agentService.findById(id); - if (agent == null) { - return ResponseEntity.notFound().build(); - } + public ApiResponse getApiKey(@PathVariable Long id) { + Agent agent = checkAgentExists(id); String masked = agentService.getApiKeyMasked(id); - return ResponseEntity.ok(buildApiKeyResponse(masked, agent.getApiKeyEnabled(), "获取 API Key 成功")); + return buildApiKeyResponse(masked, agent.getApiKeyEnabled(), "获取 API Key 成功"); } - /** - * Generate API Key - */ + /** Generate API Key */ @PostMapping("/{id}/api-key/generate") - public ResponseEntity> generateApiKey(@PathVariable("id") Long id) { - Agent existing = agentService.findById(id); - if (existing == null) { - return ResponseEntity.notFound().build(); - } + public ApiResponse generateApiKey(@PathVariable Long id) { + checkAgentExists(id); Agent agent = agentService.generateApiKey(id); - return ResponseEntity.ok(buildApiKeyResponse(agent.getApiKey(), agent.getApiKeyEnabled(), "生成 API Key 成功")); + return buildApiKeyResponse(agent.getApiKey(), agent.getApiKeyEnabled(), "生成 API Key 成功"); } - /** - * Reset API Key - */ + /** Reset API Key */ @PostMapping("/{id}/api-key/reset") - public ResponseEntity> resetApiKey(@PathVariable("id") Long id) { - Agent existing = agentService.findById(id); - if (existing == null) { - return ResponseEntity.notFound().build(); - } + public ApiResponse resetApiKey(@PathVariable Long id) { + checkAgentExists(id); Agent agent = agentService.resetApiKey(id); - return ResponseEntity.ok(buildApiKeyResponse(agent.getApiKey(), agent.getApiKeyEnabled(), "重置 API Key 成功")); + return buildApiKeyResponse(agent.getApiKey(), agent.getApiKeyEnabled(), "重置 API Key 成功"); } - /** - * Delete API Key - */ + /** Delete API Key */ @DeleteMapping("/{id}/api-key") - public ResponseEntity> deleteApiKey(@PathVariable("id") Long id) { - Agent existing = agentService.findById(id); - if (existing == null) { - return ResponseEntity.notFound().build(); - } + public ApiResponse deleteApiKey(@PathVariable Long id) { + checkAgentExists(id); Agent agent = agentService.deleteApiKey(id); - return ResponseEntity.ok(buildApiKeyResponse(agent.getApiKey(), agent.getApiKeyEnabled(), "删除 API Key 成功")); + return buildApiKeyResponse(agent.getApiKey(), agent.getApiKeyEnabled(), "删除 API Key 成功"); } - /** - * Toggle API Key enable flag - */ + /** Toggle API Key enable flag */ @PostMapping("/{id}/api-key/enable") - public ResponseEntity> toggleApiKey(@PathVariable("id") Long id, - @RequestParam("enabled") boolean enabled) { - Agent existing = agentService.findById(id); - if (existing == null) { - return ResponseEntity.notFound().build(); - } + public ApiResponse toggleApiKey(@PathVariable Long id, @RequestParam("enabled") boolean enabled) { + checkAgentExists(id); Agent agent = agentService.toggleApiKey(id, enabled); - return ResponseEntity.ok(buildApiKeyResponse(agent.getApiKey() == null ? null : "****", - agent.getApiKeyEnabled(), "更新 API Key 状态成功")); + return buildApiKeyResponse(agent.getApiKey() == null ? null : "****", agent.getApiKeyEnabled(), + "更新 API Key 状态成功"); + } + + private Agent checkAgentExists(Long id) { + Agent agent = agentService.findById(id); + if (agent == null) { + throw new ResponseStatusException(HttpStatus.NOT_FOUND, "agent with id: %d not found".formatted(id)); + } + return agent; } private ApiResponse buildApiKeyResponse(String apiKey, Integer apiKeyEnabled, String message) { diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/AgentDatasourceController.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/AgentDatasourceController.java index 11cdffdd8..f9962af1e 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/AgentDatasourceController.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/AgentDatasourceController.java @@ -18,17 +18,17 @@ import com.alibaba.cloud.ai.dataagent.dto.datasource.ToggleDatasourceDTO; import com.alibaba.cloud.ai.dataagent.dto.datasource.UpdateDatasourceTablesDTO; import com.alibaba.cloud.ai.dataagent.entity.AgentDatasource; +import com.alibaba.cloud.ai.dataagent.exception.InternalServerException; +import com.alibaba.cloud.ai.dataagent.exception.InvalidInputException; import com.alibaba.cloud.ai.dataagent.service.datasource.AgentDatasourceService; import com.alibaba.cloud.ai.dataagent.vo.ApiResponse; +import java.util.List; +import java.util.Optional; import lombok.AllArgsConstructor; import lombok.extern.slf4j.Slf4j; -import org.springframework.http.ResponseEntity; import org.springframework.validation.annotation.Validated; import org.springframework.web.bind.annotation.*; -import java.util.List; -import java.util.Optional; - /** * Agent Schema Initialization Controller Handles agent's database Schema initialization * to vector storage @@ -47,10 +47,10 @@ public class AgentDatasourceController { * Information Source" function on the frontend */ @PostMapping("/init") - public ResponseEntity initSchema(@PathVariable(value = "agentId") Long agentId) { + public ApiResponse initSchema(@PathVariable Long agentId) { // 防止前端恶意请求，dto数据应该在后端获取 try { - AgentDatasource agentDatasource = agentDatasourceService.getCurrentAgentDatasource(agentId.intValue()); + AgentDatasource agentDatasource = agentDatasourceService.getCurrentAgentDatasource(agentId); log.info("Initializing schema for agent: {}", agentId); // Extract data source ID and table list from request @@ -59,11 +59,11 @@ public ResponseEntity initSchema(@PathVariable(value = "agentId") L // Validate request parameters if (datasourceId == null) { - return ResponseEntity.badRequest().body(ApiResponse.error("数据源ID不能为空")); + throw new InvalidInputException("数据源ID不能为空"); } if (tables.isEmpty()) { - return ResponseEntity.badRequest().body(ApiResponse.error("表列表不能为空")); + throw new InvalidInputException("表列表不能为空"); } // Execute Schema initialization @@ -72,116 +72,106 @@ public ResponseEntity initSchema(@PathVariable(value = "agentId") L if (result) { log.info("Successfully initialized schema for agent: {}, tables: {}", agentId, tables.size()); - return ResponseEntity.ok(ApiResponse.success("Schema初始化成功")); + return ApiResponse.success("Schema初始化成功"); } else { - return ResponseEntity.internalServerError().body(ApiResponse.error("Schema初始化失败")); + throw new InternalServerException("Schema初始化失败"); } } catch (Exception e) { log.error("Failed to initialize schema for agent: {}", agentId, e); - return ResponseEntity.internalServerError().body(ApiResponse.error("Schema初始化失败：" + e.getMessage())); + throw new InternalServerException("Schema初始化失败：%s".formatted(e.getMessage())); } } - /** - * Get list of data sources configured for agent - */ + /** Get list of data sources configured for agent */ @GetMapping - public ResponseEntity getAgentDatasource(@PathVariable(value = "agentId") Long agentId) { + public ApiResponse> getAgentDatasource(@PathVariable Long agentId) { try { log.info("Getting datasources for agent: {}", agentId); - List datasources = agentDatasourceService.getAgentDatasource(Math.toIntExact(agentId)); + List datasources = agentDatasourceService.getAgentDatasource(agentId); log.info("Successfully retrieved {} datasources for agent: {}", datasources.size(), agentId); - return ResponseEntity.ok(ApiResponse.success("操作成功", datasources)); + return ApiResponse.success("操作成功", datasources); } catch (Exception e) { log.error("Failed to get datasources for agent: {}", agentId, e); - return ResponseEntity.badRequest().body(ApiResponse.error("获取数据源失败：" + e.getMessage(), List.of())); + throw new InvalidInputException("获取数据源失败：%s".formatted(e.getMessage()), List.of()); } } @GetMapping("/active") - public ResponseEntity getActiveAgentDatasource(@PathVariable(value = "agentId") Long agentId) { + public ApiResponse getActiveAgentDatasource(@PathVariable Long agentId) { try { log.info("Getting active datasource for agent: {}", agentId); - AgentDatasource datasource = agentDatasourceService.getCurrentAgentDatasource(Math.toIntExact(agentId)); - return ResponseEntity.ok(ApiResponse.success("操作成功", datasource)); + AgentDatasource datasource = agentDatasourceService.getCurrentAgentDatasource(agentId); + return ApiResponse.success("操作成功", datasource); } catch (Exception e) { log.error("Failed to get active datasource for agent: {}", agentId, e); - return ResponseEntity.badRequest().body(ApiResponse.error("获取数据源失败：" + e.getMessage())); + throw new InvalidInputException("获取数据源失败：%s".formatted(e.getMessage()), List.of()); } } - /** - * Add data source for agent - */ + /** Add data source for agent */ @PostMapping("/{datasourceId}") - public ResponseEntity addDatasourceToAgent(@PathVariable(value = "agentId") Integer agentId, - @PathVariable(value = "datasourceId") Integer datasourceId) { + public ApiResponse addDatasourceToAgent(@PathVariable Long agentId, + @PathVariable Integer datasourceId) { try { if (datasourceId == null) { - return ResponseEntity.badRequest().body(ApiResponse.error("数据源ID不能为空")); + throw new InvalidInputException("数据源ID不能为空"); } AgentDatasource agentDatasource = agentDatasourceService.addDatasourceToAgent(agentId, datasourceId); - return ResponseEntity.ok(ApiResponse.success("数据源添加成功", agentDatasource)); + return ApiResponse.success("数据源添加成功", agentDatasource); } catch (Exception e) { - return ResponseEntity.badRequest().body(ApiResponse.error("添加失败：" + e.getMessage())); + throw new InternalServerException("数据源添加失败：%s".formatted(e.getMessage())); } } // 更新选择的数据表 @PostMapping("/tables") - public ResponseEntity updateDatasourceTables(@PathVariable("agentId") String agentId, + public ApiResponse updateDatasourceTables(@PathVariable Long agentId, @RequestBody @Validated UpdateDatasourceTablesDTO dto) { try { dto.setTables(Optional.ofNullable(dto.getTables()).orElse(List.of())); - agentDatasourceService.updateDatasourceTables(Integer.parseInt(agentId), dto.getDatasourceId(), - dto.getTables()); - return ResponseEntity.ok(ApiResponse.success("更新成功")); + agentDatasourceService.updateDatasourceTables(agentId, dto.getDatasourceId(), dto.getTables()); + return ApiResponse.success("更新成功"); } catch (Exception e) { log.error("Error: ", e); - return ResponseEntity.badRequest().body(ApiResponse.error("更新失败：" + e.getMessage())); + throw new InternalServerException("更新失败：%s".formatted(e.getMessage())); } } - /** - * Remove data source association from agent - */ + /** Remove data source association from agent */ @DeleteMapping("/{datasourceId}") - public ResponseEntity removeDatasourceFromAgent(@PathVariable("agentId") Integer agentId, - @PathVariable("datasourceId") Integer datasourceId) { + public ApiResponse removeDatasourceFromAgent(@PathVariable Long agentId, @PathVariable Integer datasourceId) { try { agentDatasourceService.removeDatasourceFromAgent(agentId, datasourceId); - return ResponseEntity.ok(ApiResponse.success("数据源已移除")); + return ApiResponse.success("数据源已移除"); } catch (Exception e) { - return ResponseEntity.badRequest().body(ApiResponse.error("移除失败：" + e.getMessage())); + throw new InternalServerException("移除失败：%s".formatted(e.getMessage())); } } - /** - * 启用/禁用智能体的数据源 - */ + /** 启用/禁用智能体的数据源 */ @PutMapping("/toggle") - public ResponseEntity toggleDatasourceForAgent(@PathVariable("agentId") Integer agentId, + public ApiResponse toggleDatasourceForAgent(@PathVariable Long agentId, @RequestBody ToggleDatasourceDTO dto) { try { Boolean isActive = dto.getIsActive(); Integer datasourceId = dto.getDatasourceId(); if (isActive == null || datasourceId == null) { - return ResponseEntity.badRequest().body(ApiResponse.error("激活状态不能为空")); + throw new InvalidInputException("激活状态不能为空"); } AgentDatasource agentDatasource = agentDatasourceService.toggleDatasourceForAgent(agentId, dto.getDatasourceId(), isActive); - return ResponseEntity.ok(ApiResponse.success(isActive ? "数据源已启用" : "数据源已禁用", agentDatasource)); + return ApiResponse.success(isActive ? "数据源已启用" : "数据源已禁用", agentDatasource); } catch (Exception e) { - return ResponseEntity.badRequest().body(ApiResponse.error("操作失败：" + e.getMessage())); + throw new InternalServerException("操作失败：%s".formatted(e.getMessage())); } } diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/GlobalExceptionHandler.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/GlobalExceptionHandler.java new file mode 100644 index 000000000..0a236a609 --- /dev/null +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/controller/GlobalExceptionHandler.java @@ -0,0 +1,42 @@ +/* + * Copyright 2024-2026 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.dataagent.controller; + +import com.alibaba.cloud.ai.dataagent.exception.InternalServerException; +import com.alibaba.cloud.ai.dataagent.exception.InvalidInputException; +import com.alibaba.cloud.ai.dataagent.vo.ApiResponse; +import org.springframework.http.HttpStatus; +import org.springframework.web.bind.annotation.ExceptionHandler; +import org.springframework.web.bind.annotation.ResponseStatus; +import org.springframework.web.bind.annotation.RestControllerAdvice; +import org.springframework.web.servlet.mvc.method.annotation.ResponseEntityExceptionHandler; + +@RestControllerAdvice +public class GlobalExceptionHandler extends ResponseEntityExceptionHandler { + + @ExceptionHandler(InvalidInputException.class) + @ResponseStatus(HttpStatus.BAD_REQUEST) + public ApiResponse handleInvalidInputException(InvalidInputException e) { + return ApiResponse.error(e.getMessage(), e.getData()); + } + + @ExceptionHandler(InternalServerException.class) + @ResponseStatus(HttpStatus.INTERNAL_SERVER_ERROR) + public ApiResponse handleInvalidInputException(InternalServerException e) { + return ApiResponse.error(e.getMessage()); + } + +} diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/dto/schema/SemanticModelAddDTO.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/dto/schema/SemanticModelAddDTO.java index bb82c9468..44bfdfda9 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/dto/schema/SemanticModelAddDTO.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/dto/schema/SemanticModelAddDTO.java @@ -15,65 +15,46 @@ */ package com.alibaba.cloud.ai.dataagent.dto.schema; +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.NotNull; import lombok.AllArgsConstructor; import lombok.Builder; import lombok.Data; import lombok.NoArgsConstructor; -import jakarta.validation.constraints.NotBlank; -import jakarta.validation.constraints.NotNull; - -/** - * 添加语义模型的DTO类 - */ +/** 添加语义模型的DTO类 */ @Data @Builder @NoArgsConstructor @AllArgsConstructor public class SemanticModelAddDTO { - /** - * 关联的智能体ID - */ + /** 关联的智能体ID */ @NotNull(message = "智能体ID不能为空") - private Integer agentId; + private Long agentId; - /** - * 关联的表名 - */ + /** 关联的表名 */ @NotBlank(message = "表名不能为空") private String tableName; - /** - * 数据库中的物理字段名 (例如: csat_score) - */ + /** 数据库中的物理字段名 (例如: csat_score) */ @NotBlank(message = "数据库字段名不能为空") private String columnName; - /** - * 业务名/别名 (例如: 客户满意度分数) - */ + /** 业务名/别名 (例如: 客户满意度分数) */ @NotBlank(message = "业务名称不能为空") private String businessName; - /** - * 业务名的同义词 (例如: 满意度,客户评分) - */ + /** 业务名的同义词 (例如: 满意度,客户评分) */ private String synonyms; - /** - * 业务描述 (用于向LLM解释字段的业务含义) - */ + /** 业务描述 (用于向LLM解释字段的业务含义) */ private String businessDescription; - /** - * 数据库中的物理字段的原始注释 - */ + /** 数据库中的物理字段的原始注释 */ private String columnComment; - /** - * 物理数据类型 (例如: int, varchar(20)) - */ + /** 物理数据类型 (例如: int, varchar(20)) */ @NotBlank(message = "数据类型不能为空") private String dataType; diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/entity/AgentDatasource.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/entity/AgentDatasource.java index fa02956dc..fdd7ae0b5 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/entity/AgentDatasource.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/entity/AgentDatasource.java @@ -16,14 +16,12 @@ package com.alibaba.cloud.ai.dataagent.entity; import com.fasterxml.jackson.annotation.JsonFormat; -import org.springframework.format.annotation.DateTimeFormat; - import java.time.LocalDateTime; import java.util.List; - +import lombok.AllArgsConstructor; import lombok.Data; import lombok.NoArgsConstructor; -import lombok.AllArgsConstructor; +import org.springframework.format.annotation.DateTimeFormat; @Data @NoArgsConstructor @@ -32,7 +30,7 @@ public class AgentDatasource { private Integer id; - private Integer agentId; + private Long agentId; private Integer datasourceId; @@ -52,7 +50,7 @@ public class AgentDatasource { // 当前数据源选中的表 private List selectTables; - public AgentDatasource(Integer agentId, Integer datasourceId) { + public AgentDatasource(Long agentId, Integer datasourceId) { this.agentId = agentId; this.datasourceId = datasourceId; this.isActive = 1; diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/entity/SemanticModel.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/entity/SemanticModel.java index a3842441d..298819c23 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/entity/SemanticModel.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/entity/SemanticModel.java @@ -16,13 +16,12 @@ package com.alibaba.cloud.ai.dataagent.entity; import com.fasterxml.jackson.annotation.JsonFormat; +import java.time.LocalDateTime; import lombok.AllArgsConstructor; import lombok.Builder; import lombok.Data; import lombok.NoArgsConstructor; -import java.time.LocalDateTime; - /** * Semantic Model Configuration Entity Class */ @@ -40,7 +39,7 @@ public class SemanticModel { /** * 关联的智能体ID */ - private Integer agentId; + private Long agentId; /** * 关联的数据源ID diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/exception/InternalServerException.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/exception/InternalServerException.java new file mode 100644 index 000000000..801096012 --- /dev/null +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/exception/InternalServerException.java @@ -0,0 +1,24 @@ +/* + * Copyright 2024-2026 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.dataagent.exception; + +public class InternalServerException extends RuntimeException { + + public InternalServerException(String message) { + super(message); + } + +} diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/exception/InvalidInputException.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/exception/InvalidInputException.java new file mode 100644 index 000000000..de817fb13 --- /dev/null +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/exception/InvalidInputException.java @@ -0,0 +1,34 @@ +/* + * Copyright 2024-2026 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.dataagent.exception; + +import lombok.Getter; + +public class InvalidInputException extends RuntimeException { + + @Getter + private Object data; + + public InvalidInputException(String message) { + super(message); + } + + public InvalidInputException(String message, Object data) { + super(message); + this.data = data; + } + +} diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/mapper/AgentDatasourceMapper.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/mapper/AgentDatasourceMapper.java index 55657fd45..1f6d9d6b7 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/mapper/AgentDatasourceMapper.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/mapper/AgentDatasourceMapper.java @@ -16,6 +16,7 @@ package com.alibaba.cloud.ai.dataagent.mapper; import com.alibaba.cloud.ai.dataagent.entity.AgentDatasource; +import java.util.List; import org.apache.ibatis.annotations.Delete; import org.apache.ibatis.annotations.Insert; import org.apache.ibatis.annotations.Mapper; @@ -23,62 +24,52 @@ import org.apache.ibatis.annotations.Select; import org.apache.ibatis.annotations.Update; -import java.util.List; - @Mapper public interface AgentDatasourceMapper { - /** - * Query associated data sources by agent ID (including data source details) - */ + /** Query associated data sources by agent ID (including data source details) */ @Select("SELECT ad.*, d.name, d.type, d.host, d.port, d.database_name, " + "d.connection_url, d.username, d.password, d.status, d.test_status, d.description " + "FROM agent_datasource ad " + "LEFT JOIN datasource d ON ad.datasource_id = d.id " + "WHERE ad.agent_id = #{agentId} " + "ORDER BY ad.create_time DESC") - List selectByAgentIdWithDatasource(@Param("agentId") Integer agentId); + List selectByAgentIdWithDatasource(@Param("agentId") Long agentId); - /** - * Query associated data sources by agent ID - */ + /** Query associated data sources by agent ID */ @Select("SELECT * FROM agent_datasource WHERE agent_id = #{agentId} ORDER BY create_time DESC") - List selectByAgentId(@Param("agentId") Integer agentId); + List selectByAgentId(@Param("agentId") Long agentId); - /** - * Query association by agent ID and data source ID - */ + /** Query association by agent ID and data source ID */ @Select("SELECT * FROM agent_datasource WHERE agent_id = #{agentId} AND datasource_id = #{datasourceId}") - AgentDatasource selectByAgentIdAndDatasourceId(@Param("agentId") Integer agentId, + AgentDatasource selectByAgentIdAndDatasourceId(@Param("agentId") Long agentId, @Param("datasourceId") Integer datasourceId); - /** - * Disable all data sources for an agent - */ + /** Disable all data sources for an agent */ @Update("UPDATE agent_datasource SET is_active = 0 WHERE agent_id = #{agentId}") - int disableAllByAgentId(@Param("agentId") Integer agentId); + int disableAllByAgentId(@Param("agentId") Long agentId); /** * Count the number of enabled data sources for an agent (excluding the specified data * source) */ @Select("SELECT COUNT(*) FROM agent_datasource WHERE agent_id = #{agentId} AND is_active = 1 AND datasource_id != #{excludeDatasourceId}") - int countActiveByAgentIdExcluding(@Param("agentId") Integer agentId, + int countActiveByAgentIdExcluding(@Param("agentId") Long agentId, @Param("excludeDatasourceId") Integer excludeDatasourceId); @Delete("DELETE FROM agent_datasource WHERE datasource_id = #{datasourceId}") int deleteAllByDatasourceId(@Param("datasourceId") Integer datasourceId); @Insert("INSERT INTO agent_datasource (agent_id, datasource_id, is_active) VALUES (#{agentId}, #{datasourceId}, 1)") - int createNewRelationEnabled(@Param("agentId") Integer agentId, @Param("datasourceId") Integer datasourceId); + int createNewRelationEnabled(@Param("agentId") Long agentId, @Param("datasourceId") Integer datasourceId); @Update("UPDATE agent_datasource SET is_active = #{isActive} WHERE agent_id = #{agentId} AND datasource_id = #{datasourceId}") - int updateRelation(@Param("agentId") Integer agentId, @Param("datasourceId") Integer datasourceId, + int updateRelation(@Param("agentId") Long agentId, @Param("datasourceId") Integer datasourceId, @Param("isActive") Integer isActive); - default int enableRelation(Integer agentId, Integer datasourceId) { + default int enableRelation(Long agentId, Integer datasourceId) { return updateRelation(agentId, datasourceId, 1); } @Delete("DELETE FROM agent_datasource WHERE agent_id = #{agentId} AND datasource_id = #{datasourceId}") - int removeRelation(@Param("agentId") Integer agentId, @Param("datasourceId") Integer datasourceId); + int removeRelation(@Param("agentId") Long agentId, @Param("datasourceId") Integer datasourceId); } diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/agent/AgentStartupInitialization.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/agent/AgentStartupInitialization.java index ec56c7523..2df68800b 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/agent/AgentStartupInitialization.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/agent/AgentStartupInitialization.java @@ -19,6 +19,9 @@ import com.alibaba.cloud.ai.dataagent.entity.AgentDatasource; import com.alibaba.cloud.ai.dataagent.service.datasource.AgentDatasourceService; import com.alibaba.cloud.ai.dataagent.service.vectorstore.AgentVectorStoreService; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.beans.factory.DisposableBean; @@ -26,10 +29,6 @@ import org.springframework.boot.ApplicationRunner; import org.springframework.stereotype.Service; -import java.util.List; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutorService; - @Slf4j @Service @RequiredArgsConstructor @@ -60,9 +59,7 @@ public void run(ApplicationArguments args) { } } - /** - * Initialize all published agents - */ + /** Initialize all published agents */ private void initializePublishedAgents() { try { List publishedAgents = agentService.findByStatus("published"); @@ -130,8 +127,7 @@ private boolean initializeAgentDataSource(Agent agent) { return true; } - AgentDatasource activeDatasource = agentDatasourceService - .getCurrentAgentDatasource(Math.toIntExact(agentId)); + AgentDatasource activeDatasource = agentDatasourceService.getCurrentAgentDatasource(agentId); Integer datasourceId = activeDatasource.getDatasourceId(); diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/AgentDatasourceService.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/AgentDatasourceService.java index 6cbf7363f..742928589 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/AgentDatasourceService.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/AgentDatasourceService.java @@ -16,31 +16,28 @@ package com.alibaba.cloud.ai.dataagent.service.datasource; import com.alibaba.cloud.ai.dataagent.entity.AgentDatasource; - import java.util.List; public interface AgentDatasourceService { - /** - * Initialize agent's database schema using datasource - */ + /** Initialize agent's database schema using datasource */ Boolean initializeSchemaForAgentWithDatasource(Long agentId, Integer datasourceId, List tables); - List getAgentDatasource(Integer agentId); + List getAgentDatasource(Long agentId); - default AgentDatasource getCurrentAgentDatasource(Integer agentId) { + default AgentDatasource getCurrentAgentDatasource(Long agentId) { return getAgentDatasource(agentId).stream() .filter(a -> a.getIsActive() != 0) .findFirst() .orElseThrow(() -> new IllegalStateException("Agent " + agentId + " has no active datasource")); } - AgentDatasource addDatasourceToAgent(Integer agentId, Integer datasourceId); + AgentDatasource addDatasourceToAgent(Long agentId, Integer datasourceId); - void removeDatasourceFromAgent(Integer agentId, Integer datasourceId); + void removeDatasourceFromAgent(Long agentId, Integer datasourceId); - AgentDatasource toggleDatasourceForAgent(Integer agentId, Integer datasourceId, Boolean isActive); + AgentDatasource toggleDatasourceForAgent(Long agentId, Integer datasourceId, Boolean isActive); - void updateDatasourceTables(Integer agentId, Integer datasourceId, List tables); + void updateDatasourceTables(Long agentId, Integer datasourceId, List tables); } diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/DatasourceService.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/DatasourceService.java index c923e6e12..6af2dd570 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/DatasourceService.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/DatasourceService.java @@ -19,7 +19,6 @@ import com.alibaba.cloud.ai.dataagent.entity.AgentDatasource; import com.alibaba.cloud.ai.dataagent.entity.Datasource; import com.alibaba.cloud.ai.dataagent.entity.LogicalRelation; - import java.util.List; public interface DatasourceService { @@ -74,7 +73,7 @@ public interface DatasourceService { */ // 应该使用 AgentDatasourceService 中的方法 @Deprecated - List getAgentDatasource(Integer agentId); + List getAgentDatasource(Long agentId); List getDatasourceTables(Integer datasourceId) throws Exception; diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/impl/AgentDatasourceServiceImpl.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/impl/AgentDatasourceServiceImpl.java index 75f4fd7b7..c072bde55 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/impl/AgentDatasourceServiceImpl.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/impl/AgentDatasourceServiceImpl.java @@ -24,15 +24,14 @@ import com.alibaba.cloud.ai.dataagent.service.datasource.AgentDatasourceService; import com.alibaba.cloud.ai.dataagent.service.datasource.DatasourceService; import com.alibaba.cloud.ai.dataagent.service.schema.SchemaService; +import java.util.List; +import java.util.Optional; import lombok.AllArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.stereotype.Service; import org.springframework.transaction.annotation.Transactional; import org.springframework.util.Assert; -import java.util.List; -import java.util.Optional; - @Slf4j @Service @AllArgsConstructor @@ -83,7 +82,7 @@ public Boolean initializeSchemaForAgentWithDatasource(Long agentId, Integer data } @Override - public List getAgentDatasource(Integer agentId) { + public List getAgentDatasource(Long agentId) { Assert.notNull(agentId, "Agent ID cannot be null"); List adentDatasources = agentDatasourceMapper.selectByAgentIdWithDatasource(agentId); @@ -105,7 +104,7 @@ public List getAgentDatasource(Integer agentId) { @Override @Transactional - public AgentDatasource addDatasourceToAgent(Integer agentId, Integer datasourceId) { + public AgentDatasource addDatasourceToAgent(Long agentId, Integer datasourceId) { // First, disable other data sources for this agent (an agent can only have one // enabled data source) agentDatasourceMapper.disableAllByAgentId(agentId); @@ -136,12 +135,12 @@ public AgentDatasource addDatasourceToAgent(Integer agentId, Integer datasourceI } @Override - public void removeDatasourceFromAgent(Integer agentId, Integer datasourceId) { + public void removeDatasourceFromAgent(Long agentId, Integer datasourceId) { agentDatasourceMapper.removeRelation(agentId, datasourceId); } @Override - public AgentDatasource toggleDatasourceForAgent(Integer agentId, Integer datasourceId, Boolean isActive) { + public AgentDatasource toggleDatasourceForAgent(Long agentId, Integer datasourceId, Boolean isActive) { // If enabling data source, first check if there are other enabled data sources if (isActive) { int activeCount = agentDatasourceMapper.countActiveByAgentIdExcluding(agentId, datasourceId); @@ -163,7 +162,7 @@ public AgentDatasource toggleDatasourceForAgent(Integer agentId, Integer datasou @Override @Transactional - public void updateDatasourceTables(Integer agentId, Integer datasourceId, List tables) { + public void updateDatasourceTables(Long agentId, Integer datasourceId, List tables) { if (agentId == null || datasourceId == null || tables == null) { throw new IllegalArgumentException("参数不能为空"); } diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/impl/DatasourceServiceImpl.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/impl/DatasourceServiceImpl.java index 1e8d887de..513b0291d 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/impl/DatasourceServiceImpl.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/datasource/impl/DatasourceServiceImpl.java @@ -15,30 +15,24 @@ */ package com.alibaba.cloud.ai.dataagent.service.datasource.impl; -import com.alibaba.cloud.ai.dataagent.connector.accessor.Accessor; -import com.alibaba.cloud.ai.dataagent.connector.accessor.AccessorFactory; +import com.alibaba.cloud.ai.dataagent.bo.DbConfigBO; import com.alibaba.cloud.ai.dataagent.bo.schema.ColumnInfoBO; -import com.alibaba.cloud.ai.dataagent.connector.DbQueryParameter; import com.alibaba.cloud.ai.dataagent.bo.schema.TableInfoBO; +import com.alibaba.cloud.ai.dataagent.connector.DbQueryParameter; +import com.alibaba.cloud.ai.dataagent.connector.accessor.Accessor; +import com.alibaba.cloud.ai.dataagent.connector.accessor.AccessorFactory; import com.alibaba.cloud.ai.dataagent.connector.pool.DBConnectionPool; import com.alibaba.cloud.ai.dataagent.connector.pool.DBConnectionPoolFactory; -import com.alibaba.cloud.ai.dataagent.bo.DbConfigBO; -import com.alibaba.cloud.ai.dataagent.entity.Datasource; import com.alibaba.cloud.ai.dataagent.entity.AgentDatasource; +import com.alibaba.cloud.ai.dataagent.entity.Datasource; import com.alibaba.cloud.ai.dataagent.entity.LogicalRelation; import com.alibaba.cloud.ai.dataagent.enums.ErrorCodeEnum; -import com.alibaba.cloud.ai.dataagent.service.datasource.handler.DatasourceTypeHandler; -import com.alibaba.cloud.ai.dataagent.service.datasource.handler.registry.DatasourceTypeHandlerRegistry; -import com.alibaba.cloud.ai.dataagent.mapper.DatasourceMapper; import com.alibaba.cloud.ai.dataagent.mapper.AgentDatasourceMapper; +import com.alibaba.cloud.ai.dataagent.mapper.DatasourceMapper; import com.alibaba.cloud.ai.dataagent.mapper.LogicalRelationMapper; import com.alibaba.cloud.ai.dataagent.service.datasource.DatasourceService; -import lombok.AllArgsConstructor; -import lombok.extern.slf4j.Slf4j; -import org.apache.commons.lang3.StringUtils; -import org.springframework.stereotype.Service; -import org.springframework.transaction.annotation.Transactional; - +import com.alibaba.cloud.ai.dataagent.service.datasource.handler.DatasourceTypeHandler; +import com.alibaba.cloud.ai.dataagent.service.datasource.handler.registry.DatasourceTypeHandlerRegistry; import java.util.ArrayList; import java.util.HashSet; import java.util.List; @@ -46,6 +40,11 @@ import java.util.Objects; import java.util.Set; import java.util.stream.Collectors; +import lombok.AllArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; +import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Transactional; // todo: 检查Mapper的返回值，判断是否执行成功（或者对Mapper进行AOP） @Slf4j @@ -185,7 +184,7 @@ private boolean realConnectionTest(Datasource datasource) { @Override @Deprecated - public List getAgentDatasource(Integer agentId) { + public List getAgentDatasource(Long agentId) { List adentDatasources = agentDatasourceMapper.selectByAgentIdWithDatasource(agentId); // Manually fill in the data source information (since MyBatis Plus does not diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/semantic/SemanticModelServiceImpl.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/semantic/SemanticModelServiceImpl.java index fd9843d5a..41101383c 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/semantic/SemanticModelServiceImpl.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/semantic/SemanticModelServiceImpl.java @@ -23,13 +23,12 @@ import com.alibaba.cloud.ai.dataagent.mapper.AgentDatasourceMapper; import com.alibaba.cloud.ai.dataagent.mapper.SemanticModelMapper; import com.alibaba.cloud.ai.dataagent.vo.BatchImportResult; +import java.util.List; import lombok.AllArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.stereotype.Service; import org.springframework.web.multipart.MultipartFile; -import java.util.List; - @Service @AllArgsConstructor @Slf4j @@ -75,7 +74,7 @@ public void addSemanticModel(SemanticModel semanticModel) { @Override public boolean addSemanticModel(SemanticModelAddDTO dto) { // 根据agentId查询关联的datasourceId - Integer datasourceId = findDatasourceIdByAgentId(dto.getAgentId().longValue()); + Integer datasourceId = findDatasourceIdByAgentId(dto.getAgentId()); // 转换DTO为Entity SemanticModel semanticModel = SemanticModel.builder() @@ -97,11 +96,9 @@ public boolean addSemanticModel(SemanticModelAddDTO dto) { return true; } - /** - * 根据agentId查找关联的datasourceId 如果有多个数据源，返回第一个启用的数据源 - */ + /** 根据agentId查找关联的datasourceId 如果有多个数据源，返回第一个启用的数据源 */ private Integer findDatasourceIdByAgentId(Long agentId) { - List agentDatasources = agentDatasourceMapper.selectByAgentId(agentId.intValue()); + List agentDatasources = agentDatasourceMapper.selectByAgentId(agentId); if (agentDatasources.isEmpty()) { throw new RuntimeException("No datasource found for Agent ID " + agentId); @@ -143,12 +140,6 @@ public void deleteSemanticModel(Long id) { semanticModelMapper.deleteById(id); } - @Override - public void updateSemanticModel(Long id, SemanticModel semanticModel) { - semanticModel.setId(id); - semanticModelMapper.updateById(semanticModel); - } - @Override public BatchImportResult batchImport(SemanticModelBatchImportDTO dto) { BatchImportResult result = BatchImportResult.builder() @@ -191,7 +182,7 @@ public BatchImportResult batchImport(SemanticModelBatchImportDTO dto) { else { // 插入新记录 SemanticModel newModel = SemanticModel.builder() - .agentId(dto.getAgentId().intValue()) + .agentId(dto.getAgentId()) .datasourceId(datasourceId) .tableName(item.getTableName()) .columnName(item.getColumnName()) @@ -252,4 +243,10 @@ public BatchImportResult importFromExcel(MultipartFile file, Long agentId) { } } + @Override + public void updateSemanticModel(Long id, SemanticModel semanticModel) { + semanticModel.setId(id); + semanticModelMapper.updateById(semanticModel); + } + } diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/DatabaseUtil.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/DatabaseUtil.java index ed186d2a7..98481ea36 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/DatabaseUtil.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/DatabaseUtil.java @@ -15,9 +15,9 @@ */ package com.alibaba.cloud.ai.dataagent.util; +import com.alibaba.cloud.ai.dataagent.bo.DbConfigBO; import com.alibaba.cloud.ai.dataagent.connector.accessor.Accessor; import com.alibaba.cloud.ai.dataagent.connector.accessor.AccessorFactory; -import com.alibaba.cloud.ai.dataagent.bo.DbConfigBO; import com.alibaba.cloud.ai.dataagent.entity.AgentDatasource; import com.alibaba.cloud.ai.dataagent.service.datasource.AgentDatasourceService; import com.alibaba.cloud.ai.dataagent.service.datasource.DatasourceService; @@ -39,7 +39,7 @@ public class DatabaseUtil { private final DatasourceService datasourceService; - public DbConfigBO getAgentDbConfig(Integer agentId) { + public DbConfigBO getAgentDbConfig(Long agentId) { log.info("Getting datasource config for agent: {}", agentId); // Get the enabled data source for the agent @@ -52,7 +52,7 @@ public DbConfigBO getAgentDbConfig(Integer agentId) { return dbConfig; } - public Accessor getAgentAccessor(Integer agentId) { + public Accessor getAgentAccessor(Long agentId) { DbConfigBO dbConfig = getAgentDbConfig(agentId); return accessorFactory.getAccessorByDbConfig(dbConfig); } diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/FluxUtil.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/FluxUtil.java index 8f46cdc97..3b4138ebb 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/FluxUtil.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/FluxUtil.java @@ -20,14 +20,13 @@ import com.alibaba.cloud.ai.graph.action.NodeAction; import com.alibaba.cloud.ai.graph.streaming.OutputType; import com.alibaba.cloud.ai.graph.streaming.StreamingOutput; +import java.util.Map; +import java.util.function.Function; +import java.util.function.Supplier; import org.springframework.ai.chat.model.ChatResponse; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import java.util.Map; -import java.util.function.Supplier; -import java.util.function.Function; - /** * @author vlsmb * @since 2025/10/22 @@ -35,7 +34,6 @@ public final class FluxUtil { private FluxUtil() { - } /** diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/workflow/node/SqlExecuteNode.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/workflow/node/SqlExecuteNode.java index 58f1ae157..66821d4d5 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/workflow/node/SqlExecuteNode.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/workflow/node/SqlExecuteNode.java @@ -53,6 +53,7 @@ import java.util.Map; 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; @@ -60,6 +61,7 @@ /** * SQL execution node that executes SQL queries against the database. * + *

* This node is responsible for: - Executing SQL queries generated by previous nodes - * Handling query results and errors - Providing streaming feedback to users during * execution - Managing step-by-step result accumulation @@ -93,11 +95,11 @@ public Map apply(OverAllState state) throws Exception { // Get the agent ID from the state String agentIdStr = StateUtil.getStringValue(state, Constant.AGENT_ID); - if (agentIdStr == null || agentIdStr.trim().isEmpty()) { + if (StringUtils.isBlank(agentIdStr)) { throw new IllegalStateException("Agent ID cannot be empty."); } - Integer agentId = Integer.valueOf(agentIdStr); + Long agentId = Long.valueOf(agentIdStr); // Dynamically get the data source configuration for an agent DbConfigBO dbConfig = databaseUtil.getAgentDbConfig(agentId); @@ -108,6 +110,7 @@ public Map apply(OverAllState state) throws Exception { /** * Executes the SQL query against the database and handles the results. * + *

* This method follows the business-logic-first pattern: 1. Execute the actual SQL * query immediately 2. Process and store the results 3. Create streaming output for * user experience only @@ -120,7 +123,7 @@ public Map apply(OverAllState state) throws Exception { */ @SuppressWarnings("unchecked") private Map executeSqlQuery(OverAllState state, Integer currentStep, String sqlQuery, - DbConfigBO dbConfig, Integer agentId) { + DbConfigBO dbConfig, Long agentId) { // Execute business logic first - actual SQL execution DbQueryParameter dbQueryParameter = new DbQueryParameter(); dbQueryParameter.setSql(sqlQuery); @@ -192,7 +195,6 @@ private Map executeSqlQuery(OverAllState state, Integer currentS // result Flux> generator = FluxUtil.createStreamingGeneratorWithMessages(this.getClass(), state, v -> result, displayFlux); - return Map.of(SQL_EXECUTE_NODE_OUTPUT, generator); } diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/workflow/node/TableRelationNode.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/workflow/node/TableRelationNode.java index b80c77707..4d657d7de 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/workflow/node/TableRelationNode.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/workflow/node/TableRelationNode.java @@ -15,6 +15,18 @@ */ package com.alibaba.cloud.ai.dataagent.workflow.node; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.AGENT_ID; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.COLUMN_DOCUMENTS__FOR_SCHEMA_OUTPUT; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.DB_DIALECT_TYPE; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.EVIDENCE; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.GENEGRATED_SEMANTIC_MODEL_PROMPT; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.SQL_GENERATE_SCHEMA_MISSING_ADVICE; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.TABLE_DOCUMENTS_FOR_SCHEMA_OUTPUT; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.TABLE_RELATION_EXCEPTION_OUTPUT; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.TABLE_RELATION_OUTPUT; +import static com.alibaba.cloud.ai.dataagent.constant.Constant.TABLE_RELATION_RETRY_COUNT; +import static com.alibaba.cloud.ai.dataagent.prompt.PromptHelper.buildSemanticModelPrompt; + import com.alibaba.cloud.ai.dataagent.bo.DbConfigBO; import com.alibaba.cloud.ai.dataagent.dto.schema.SchemaDTO; import com.alibaba.cloud.ai.dataagent.dto.schema.TableDTO; @@ -24,17 +36,25 @@ import com.alibaba.cloud.ai.dataagent.enums.TextType; import com.alibaba.cloud.ai.dataagent.service.datasource.AgentDatasourceService; import com.alibaba.cloud.ai.dataagent.service.datasource.DatasourceService; -import com.alibaba.cloud.ai.dataagent.util.DatabaseUtil; -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.dataagent.service.nl2sql.Nl2SqlService; import com.alibaba.cloud.ai.dataagent.service.schema.SchemaService; import com.alibaba.cloud.ai.dataagent.service.semantic.SemanticModelService; import com.alibaba.cloud.ai.dataagent.util.ChatResponseUtil; +import com.alibaba.cloud.ai.dataagent.util.DatabaseUtil; import com.alibaba.cloud.ai.dataagent.util.FluxUtil; import com.alibaba.cloud.ai.dataagent.util.StateUtil; +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 java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Consumer; +import java.util.stream.Collectors; import lombok.AllArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.chat.model.ChatResponse; @@ -42,17 +62,11 @@ import org.springframework.stereotype.Component; import reactor.core.publisher.Flux; -import java.util.*; -import java.util.function.Consumer; -import java.util.stream.Collectors; - -import static com.alibaba.cloud.ai.dataagent.constant.Constant.*; -import static com.alibaba.cloud.ai.dataagent.prompt.PromptHelper.buildSemanticModelPrompt; - /** * Table relationship inference node that automatically completes complex structures like * JOINs and foreign keys. * + *

* This node is responsible for: - Inferring relationships between tables and fields - * Building initial schema from documents - Processing schema selection based on input and * evidence - Handling schema advice for missing information @@ -88,9 +102,9 @@ public Map apply(OverAllState state) throws Exception { String agentIdStr = StateUtil.getStringValue(state, AGENT_ID); // Execute business logic first - get final result immediately - DbConfigBO agentDbConfig = databaseUtil.getAgentDbConfig(Integer.valueOf(agentIdStr)); + DbConfigBO agentDbConfig = databaseUtil.getAgentDbConfig(Long.valueOf(agentIdStr)); - List logicalForeignKeys = getLogicalForeignKeys(Integer.valueOf(agentIdStr), tableDocuments); + List logicalForeignKeys = getLogicalForeignKeys(Long.valueOf(agentIdStr), tableDocuments); log.info("Found {} logical foreign keys for agent: {}", logicalForeignKeys.size(), agentIdStr); SchemaDTO initialSchema = buildInitialSchema(agentIdStr, columnDocuments, tableDocuments, agentDbConfig, @@ -142,12 +156,9 @@ public Map apply(OverAllState state) throws Exception { // subsequent nodes return Map.of(TABLE_RELATION_OUTPUT, generator, DB_DIALECT_TYPE, agentDbConfig.getDialectType(), TABLE_RELATION_RETRY_COUNT, 0, TABLE_RELATION_EXCEPTION_OUTPUT, ""); - } - /** - * Builds initial schema from column and table documents. - */ + /** Builds initial schema from column and table documents. */ private SchemaDTO buildInitialSchema(String agentId, List columnDocuments, List tableDocuments, DbConfigBO agentDbConfig, List logicalForeignKeys) { SchemaDTO schemaDTO = new SchemaDTO(); @@ -174,9 +185,7 @@ private SchemaDTO buildInitialSchema(String agentId, List columnDocume return schemaDTO; } - /** - * Processes schema selection based on input, evidence, and optional advice. - */ + /** Processes schema selection based on input, evidence, and optional advice. */ private Flux processSchemaSelection(SchemaDTO schemaDTO, String input, String evidence, OverAllState state, DbConfigBO agentDbConfig, Consumer dtoConsumer) { String schemaAdvice = StateUtil.getStringValue(state, SQL_GENERATE_SCHEMA_MISSING_ADVICE, null); @@ -199,10 +208,8 @@ private Flux processSchemaSelection(SchemaDTO schemaDTO, String in ChatResponseUtil.createResponse("\n\n选择数据表完成。"))); } - /** - * 获取逻辑外键信息，并过滤只保留与当前召回表相关的外键 - */ - private List getLogicalForeignKeys(Integer agentId, List tableDocuments) { + /** 获取逻辑外键信息，并过滤只保留与当前召回表相关的外键 */ + private List getLogicalForeignKeys(Long agentId, List tableDocuments) { try { // 获取当前 agent 激活的数据源 AgentDatasource agentDatasource = agentDatasourceService.getCurrentAgentDatasource(agentId); diff --git a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/mapper/AgentDatasourceMapperTest.java b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/mapper/AgentDatasourceMapperTest.java index 8de72af79..df8b7460f 100644 --- a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/mapper/AgentDatasourceMapperTest.java +++ b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/mapper/AgentDatasourceMapperTest.java @@ -15,9 +15,16 @@ */ package com.alibaba.cloud.ai.dataagent.mapper; +import static org.junit.jupiter.api.Assertions.*; +import static org.junit.jupiter.api.Assertions.assertNull; + import com.alibaba.cloud.ai.dataagent.entity.AgentDatasource; import com.alibaba.cloud.ai.dataagent.entity.Datasource; import com.alibaba.cloud.ai.dataagent.service.MySqlContainerConfiguration; +import java.sql.Statement; +import java.time.LocalDateTime; +import java.util.List; +import java.util.UUID; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mybatis.spring.boot.test.autoconfigure.MybatisTest; @@ -29,14 +36,6 @@ import org.springframework.jdbc.support.KeyHolder; import org.springframework.test.context.TestPropertySource; -import java.sql.Statement; -import java.time.LocalDateTime; -import java.util.List; -import java.util.UUID; - -import static org.junit.jupiter.api.Assertions.*; -import static org.junit.jupiter.api.Assertions.assertNull; - @MybatisTest @TestPropertySource( properties = { "spring.sql.init.mode=never", "mybatis.configuration.map-underscore-to-camel-case=true" }) @@ -53,7 +52,7 @@ public class AgentDatasourceMapperTest { @Autowired private JdbcTemplate jdbcTemplate; - private Integer agentId; + private Long agentId; @BeforeEach public void setUpAgent() { @@ -65,7 +64,7 @@ public void setUpAgent() { @Test public void testAgentDatasourceCreateNewRelationEnabled() { // Given - Integer agentId = this.agentId; + Long agentId = this.agentId; Integer datasourceId = createTestDatasourceAndGetId(); // When @@ -83,7 +82,7 @@ public void testAgentDatasourceCreateNewRelationEnabled() { @Test public void testAgentDatasourceSelectByAgentId() { // Given - Integer agentId = this.agentId; + Long agentId = this.agentId; Integer datasourceId1 = createTestDatasourceAndGetId(); Integer datasourceId2 = createTestDatasourceAndGetId(); @@ -103,7 +102,7 @@ public void testAgentDatasourceSelectByAgentId() { @Test public void testAgentDatasourceSelectByAgentIdAndDatasourceId() { // Given - Integer agentId = this.agentId; + Long agentId = this.agentId; Integer datasourceId = createTestDatasourceAndGetId(); agentDatasourceMapper.createNewRelationEnabled(agentId, datasourceId); @@ -119,7 +118,7 @@ public void testAgentDatasourceSelectByAgentIdAndDatasourceId() { @Test public void testAgentDatasourceUpdateRelation() { // Given - Integer agentId = this.agentId; + Long agentId = this.agentId; Integer datasourceId = createTestDatasourceAndGetId(); agentDatasourceMapper.createNewRelationEnabled(agentId, datasourceId); @@ -143,7 +142,7 @@ public void testAgentDatasourceUpdateRelation() { @Test public void testAgentDatasourceDisableAllByAgentId() { // Given - Integer agentId = this.agentId; + Long agentId = this.agentId; Integer datasourceId1 = createTestDatasourceAndGetId(); Integer datasourceId2 = createTestDatasourceAndGetId(); @@ -163,7 +162,7 @@ public void testAgentDatasourceDisableAllByAgentId() { @Test public void testAgentDatasourceCountActiveByAgentIdExcluding() { // Given - Integer agentId = this.agentId; + Long agentId = this.agentId; Integer datasourceId1 = createTestDatasourceAndGetId(); Integer datasourceId2 = createTestDatasourceAndGetId(); Integer datasourceId3 = createTestDatasourceAndGetId(); @@ -182,7 +181,7 @@ public void testAgentDatasourceCountActiveByAgentIdExcluding() { @Test public void testAgentDatasourceRemoveRelation() { // Given - Integer agentId = this.agentId; + Long agentId = this.agentId; Integer datasourceId = createTestDatasourceAndGetId(); agentDatasourceMapper.createNewRelationEnabled(agentId, datasourceId); @@ -198,8 +197,8 @@ public void testAgentDatasourceRemoveRelation() { @Test public void testAgentDatasourceDeleteAllByDatasourceId() { // Given - Integer agentId1 = createTestAgentAndGetId(); - Integer agentId2 = createTestAgentAndGetId(); + Long agentId1 = createTestAgentAndGetId(); + Long agentId2 = createTestAgentAndGetId(); Integer datasourceId = createTestDatasourceAndGetId(); agentDatasourceMapper.createNewRelationEnabled(agentId1, datasourceId); @@ -243,7 +242,7 @@ private Integer createTestDatasourceAndGetId() { return datasource.getId(); } - private Integer createTestAgentAndGetId() { + private Long createTestAgentAndGetId() { KeyHolder keyHolder = new GeneratedKeyHolder(); jdbcTemplate.update(con -> { var ps = con.prepareStatement( @@ -254,12 +253,12 @@ private Integer createTestAgentAndGetId() { return ps; }, keyHolder); Number key = keyHolder.getKey(); - Integer id = key == null ? null : key.intValue(); + Long id = key == null ? null : key.longValue(); // 立即校验插入是否可见且存在，避免外键引用失败 if (id == null) { throw new IllegalStateException("创建测试Agent失败：未获取到自增ID"); } - Integer exists = jdbcTemplate.queryForObject("SELECT COUNT(1) FROM agent WHERE id = ?", Integer.class, id); + Long exists = jdbcTemplate.queryForObject("SELECT COUNT(1) FROM agent WHERE id = ?", Long.class, id); if (exists == null || exists == 0) { throw new IllegalStateException("创建测试Agent失败：数据库中不存在id=" + id); } diff --git a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/mapper/MappersTest.java b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/mapper/MappersTest.java index 486304767..bd152708a 100644 --- a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/mapper/MappersTest.java +++ b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/mapper/MappersTest.java @@ -15,8 +15,10 @@ */ package com.alibaba.cloud.ai.dataagent.mapper; -import com.alibaba.cloud.ai.dataagent.service.MySqlContainerConfiguration; import com.alibaba.cloud.ai.dataagent.entity.*; +import com.alibaba.cloud.ai.dataagent.service.MySqlContainerConfiguration; +import java.time.LocalDateTime; +import java.util.List; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.mybatis.spring.boot.test.autoconfigure.MybatisTest; @@ -25,9 +27,6 @@ import org.springframework.boot.testcontainers.context.ImportTestcontainers; import org.springframework.test.context.TestPropertySource; -import java.time.LocalDateTime; -import java.util.List; - /** * Mappers 单元测试类 * @@ -161,7 +160,7 @@ public void testChatSessionAndMessageCrud() { public void testSemanticModelCrud() { Long agentId = createAgent("semantic-holder"); SemanticModel m = new SemanticModel(); - m.setAgentId(Math.toIntExact(agentId)); + m.setAgentId(agentId); m.setDatasourceId(1); // 添加数据源ID m.setTableName("test_table"); // 添加表名 m.setColumnName("origin_tc"); @@ -175,7 +174,7 @@ public void testSemanticModelCrud() { Assertions.assertEquals(1, ins); Assertions.assertNotNull(m.getId()); - List query = semanticModelMapper.selectByAgentId(Long.valueOf(m.getAgentId())); + List query = semanticModelMapper.selectByAgentId(m.getAgentId()); Assertions.assertTrue(query.stream().anyMatch(x -> x.getId().equals(m.getId()))); semanticModelMapper.disableById(m.getId()); diff --git a/pom.xml b/pom.xml index 9247f2e1b..758ede4a7 100644 --- a/pom.xml +++ b/pom.xml @@ -63,7 +63,7 @@ 4.37.0 2.3.232 8.1.3.140 - 1.19.3 + 1.21.4 8.18.0