0606

admin

2025-06-07 3732bde983f34f5cb729ac1724b9795fb97d8941

 aiflowy-modules/aiflowy-module-ai/src/main/java/tech/aiflowy/ai/controller/AiBotController.java

@@ -15,25 +15,36 @@
import com.agentsflex.core.prompt.ToolPrompt;
import com.agentsflex.core.util.CollectionUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializeConfig;
import com.jfinal.template.stat.ast.Break;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonObject;
import com.mybatisflex.core.query.QueryWrapper;
import com.mybatisflex.core.table.TableInfo;
import com.mybatisflex.core.table.TableInfoFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.*;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import tech.aiflowy.ai.config.DifyStreamClient;
import tech.aiflowy.ai.config.FileReference;
import tech.aiflowy.ai.entity.*;
import tech.aiflowy.ai.mapper.AiBotConversationMessageMapper;
import tech.aiflowy.ai.service.*;
import tech.aiflowy.common.ai.ChatManager;
import tech.aiflowy.common.ai.MySseEmitter;
import tech.aiflowy.common.domain.Result;
import tech.aiflowy.common.entity.LoginAccount;
import tech.aiflowy.common.satoken.util.SaTokenUtil;
import tech.aiflowy.common.util.StringUtil;
import tech.aiflowy.common.web.controller.BaseCurdController;
@@ -46,10 +57,6 @@
import javax.servlet.http.HttpServletResponse;
import java.math.BigInteger;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 控制层。
@@ -73,6 +80,8 @@
    private AiBotConversationMessageMapper aiBotConversationMessageMapper;

    private static final Logger logger = LoggerFactory.getLogger(AiBotController.class);
    @Autowired
    private RestTemplate restTemplate;

    public AiBotController(AiBotService service, AiLlmService aiLlmService, AiBotWorkflowService aiBotWorkflowService, AiBotKnowledgeService aiBotKnowledgeService, AiBotMessageService aiBotMessageService) {
        super(service);
@@ -128,104 +137,292 @@
     * @param response
     * @return
     */

    @Autowired
    private ObjectMapper objectMapper;


    @PostMapping("chat")
    public SseEmitter chat(@JsonBody(value = "prompt", required = true) String prompt,
                           @JsonBody(value = "botId", required = true) BigInteger botId,
                           @JsonBody(value = "sessionId", required = true) String sessionId,
                           @JsonBody(value = "isExternalMsg") int isExternalMsg,
                           @JsonBody(value = "file") String file,
                           HttpServletResponse response) {
        response.setContentType("text/event-stream");
        AiBot aiBot = service.getById(botId);
        if (aiBot == null) {
            return ChatManager.getInstance().sseEmitterForContent("机器人不存在");
        }

        Map<String, Object> llmOptions = aiBot.getLlmOptions();
        String systemPrompt = llmOptions != null ? (String) llmOptions.get("systemPrompt") : null;
        AiLlm aiLlm = aiLlmService.getById(aiBot.getLlmId());

        if (aiLlm == null) {
            return ChatManager.getInstance().sseEmitterForContent("LLM不存在");
        }
        if (StringUtil.hasText(aiBot.getModelAPI())){
            if (aiBot.getBotTypeId() == 2) {

        Llm llm = aiLlm.toLlm();
                String apiUrl = aiBot.getModelAPI()+"/workflows/run"; // 替换为实际API URL
                String apiKey = aiBot.getModelKEY(); // 替换为实际API Key

        if (llm == null) {
            return ChatManager.getInstance().sseEmitterForContent("LLM获取为空");
        }
                DifyStreamClient client = new DifyStreamClient(apiUrl, apiKey, aiBotMessageService);
                DifyStreamClient uploadClient = new DifyStreamClient(aiBot.getModelAPI()+"/files/upload", apiKey, aiBotMessageService);
                String test = uploadClient.fileUpload(SaTokenUtil.getLoginAccount().getId() + "", file);
                System.out.println(test);
                Gson gson = new GsonBuilder().setPrettyPrinting().create();
                JsonObject fileJson = gson.fromJson(test, JsonObject.class);
                String fileId = fileJson.get("id").getAsString();

        AiBotMessageMemory memory = new AiBotMessageMemory(botId, SaTokenUtil.getLoginAccount().getId(),
                sessionId, isExternalMsg, aiBotMessageService, aiBotConversationMessageMapper,
                aiBotConversationMessageService);
                // 2. 构建文件参数对象
                Map<String, Object> fileParam = new HashMap<>();
                fileParam.put("transfer_method", "local_file");
                fileParam.put("upload_file_id", fileId);
                String[] split = file.split("\\.");
//                fileParam.put("type", fileJson.get("extension").getAsString()); // 例如 "excel"、"pdf" 等
                fileParam.put("type", "document"); // 例如 "excel"、"pdf" 等

        final HistoriesPrompt historiesPrompt = new HistoriesPrompt();
        if (systemPrompt != null) {
            historiesPrompt.setSystemMessage(SystemMessage.of(systemPrompt));
        }
                // 3. 组装 inputs 参数
                Map<String, Object> inputs = new HashMap<>();
                inputs.put("w", fileParam); // 添加文件参数，variableName 如 "document"

        historiesPrompt.setMemory(memory);

        HumanMessage humanMessage = new HumanMessage(prompt);

        // 添加插件相关的function calling
        appendPluginToolFunction(botId, humanMessage);
                AiBotMessageMemory memory = new AiBotMessageMemory(botId, SaTokenUtil.getLoginAccount().getId(),
                        sessionId, isExternalMsg, aiBotMessageService, aiBotConversationMessageMapper,
                        aiBotConversationMessageService);

        //添加工作流相关的 Function Calling
        appendWorkflowFunctions(botId, humanMessage);
                final HistoriesPrompt historiesPrompt = new HistoriesPrompt();
                if (systemPrompt != null) {
                    historiesPrompt.setSystemMessage(SystemMessage.of(systemPrompt));
                }

        //添加知识库相关的 Function Calling
        appendKnowledgeFunctions(botId, humanMessage);
                historiesPrompt.setMemory(memory);

        historiesPrompt.addMessage(humanMessage);
                HumanMessage humanMessage = new HumanMessage(prompt);

        MySseEmitter emitter = new MySseEmitter((long) (1000 * 60 * 2));
                // 添加插件相关的function calling
                appendPluginToolFunction(botId, humanMessage);

        final Boolean[] needClose = {true};
                //添加工作流相关的 Function Calling
                appendWorkflowFunctions(botId, humanMessage);

        ServletRequestAttributes sra = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        // 统一使用流式处理，无论是否有 Function Calling
        llm.chatStream(historiesPrompt, new StreamResponseListener() {
            @Override
            public void onMessage(ChatContext context, AiMessageResponse response) {
                //添加知识库相关的 Function Calling
                appendKnowledgeFunctions(botId, humanMessage);

                historiesPrompt.addMessage(humanMessage);


                final Boolean[] needClose = {true};

                ServletRequestAttributes sra = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();


                MySseEmitter emitter = new MySseEmitter(1000L * 60 * 2); // 2分钟超时

                try {
                    RequestContextHolder.setRequestAttributes(sra, true);
                    if (response != null) {
                        // 检查是否需要触发 Function Calling
                        if (response.getFunctionCallers() != null && CollectionUtil.hasItems(response.getFunctionCallers())) {
                            needClose[0] = false;
                            function_call(response, emitter, needClose, historiesPrompt, llm, prompt, false);
                        } else {
                            // 强制流式返回，即使有 Function Calling 也先返回部分结果
                            if (response.getMessage() != null) {
                                String content = response.getMessage().getContent();
                                if (StringUtil.hasText(content)) {
                                    emitter.send(JSON.toJSONString(response.getMessage()));
                                }
                            }
                        }
                    }


                    String userId = SaTokenUtil.getLoginAccount().getId() + "";
                    client.runWorkflow(inputs, prompt, userId, emitter, sessionId, botId);
                } catch (Exception e) {
                    emitter.completeWithError(e);
                }
//                System.out.println(emitter.toString());

                return emitter;
            }
            aiBot.setModelAPI(aiBot.getModelAPI()+"/chat-messages");
            String apiUrl = aiBot.getModelAPI(); // 替换为实际API URL
            String apiKey = aiBot.getModelKEY(); // 替换为实际API Key

            DifyStreamClient client = new DifyStreamClient(apiUrl, apiKey, aiBotMessageService);
            AiBotMessageMemory memory = new AiBotMessageMemory(botId, SaTokenUtil.getLoginAccount().getId(),
                    sessionId, isExternalMsg, aiBotMessageService, aiBotConversationMessageMapper,
                    aiBotConversationMessageService);

            final HistoriesPrompt historiesPrompt = new HistoriesPrompt();
            if (systemPrompt != null) {
                historiesPrompt.setSystemMessage(SystemMessage.of(systemPrompt));
            }

            @Override
            public void onStop(ChatContext context) {
                if (needClose[0]) {
                    emitter.complete();
            historiesPrompt.setMemory(memory);

            HumanMessage humanMessage = new HumanMessage(prompt);

            // 添加插件相关的function calling
            appendPluginToolFunction(botId, humanMessage);

            //添加工作流相关的 Function Calling
            appendWorkflowFunctions(botId, humanMessage);

            //添加知识库相关的 Function Calling
            appendKnowledgeFunctions(botId, humanMessage);

            historiesPrompt.addMessage(humanMessage);


            final Boolean[] needClose = {true};

            ServletRequestAttributes sra = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();


            MySseEmitter emitter = new MySseEmitter(1000L * 60 * 2); // 2分钟超时

            try {
                String userId = SaTokenUtil.getLoginAccount().getId() + "";
                client.chatStream(prompt, userId, emitter, sessionId, botId);
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
            System.out.println(emitter.toString());
            return emitter;
        }
        else{
            AiLlm aiLlm = aiLlmService.getById(aiBot.getLlmId());

            if (aiLlm == null) {
                return ChatManager.getInstance().sseEmitterForContent("LLM不存在");
            }

            Llm llm = aiLlm.toLlm();
            if (llm == null) {
                return ChatManager.getInstance().sseEmitterForContent("LLM获取为空");
            }

            AiBotMessageMemory memory = new AiBotMessageMemory(botId, SaTokenUtil.getLoginAccount().getId(),
                    sessionId, isExternalMsg, aiBotMessageService, aiBotConversationMessageMapper,
                    aiBotConversationMessageService);

            final HistoriesPrompt historiesPrompt = new HistoriesPrompt();
            if (systemPrompt != null) {
                historiesPrompt.setSystemMessage(SystemMessage.of(systemPrompt));
            }

            historiesPrompt.setMemory(memory);

            HumanMessage humanMessage = new HumanMessage(prompt);

            // 添加插件相关的function calling
            appendPluginToolFunction(botId, humanMessage);

            //添加工作流相关的 Function Calling
            appendWorkflowFunctions(botId, humanMessage);

            //添加知识库相关的 Function Calling
            appendKnowledgeFunctions(botId, humanMessage);

            historiesPrompt.addMessage(humanMessage);

            MySseEmitter emitter = new MySseEmitter((long) (1000 * 60 * 2));

            final Boolean[] needClose = {true};

            ServletRequestAttributes sra = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
            // 统一使用流式处理，无论是否有 Function Calling
            llm.chatStream(historiesPrompt, new StreamResponseListener() {
                @Override
                public void onMessage(ChatContext context, AiMessageResponse response) {
                    try {
                        RequestContextHolder.setRequestAttributes(sra, true);
                        if (response != null) {
                            // 检查是否需要触发 Function Calling
                            if (response.getFunctionCallers() != null && CollectionUtil.hasItems(response.getFunctionCallers())) {
                                needClose[0] = false;
                                function_call(response, emitter, needClose, historiesPrompt, llm, prompt, false);
                            } else {
                                // 强制流式返回，即使有 Function Calling 也先返回部分结果
                                if (response.getMessage() != null) {
                                    String content = response.getMessage().getContent();
                                    if (StringUtil.hasText(content)) {
                                        System.out.println(response);
                                        emitter.send(JSON.toJSONString(response.getMessage()));
                                    }
                                }
                            }
                        }


                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                }

                @Override
                public void onStop(ChatContext context) {
                    if (needClose[0]) {
                        emitter.complete();
                    }
                }

                @Override
                public void onFailure(ChatContext context, Throwable throwable) {
                    emitter.completeWithError(throwable);
                }
            });

            System.out.println(emitter.toString());
            return emitter;
        }

    }

    @PostMapping("files/upload")
    public Result filesUpload(@JsonBody(value = "botId", required = true) BigInteger botId,
                           String file,
                           HttpServletResponse response){
        try{
            String userId = SaTokenUtil.getLoginAccount().getId() + "";
            response.setContentType("text/event-stream");
            AiBot aiBot = service.getById(botId);
            aiBot.setModelAPI(aiBot.getModelAPI()+"/files/upload");
            String apiUrl = aiBot.getModelAPI(); // 替换为实际API URL
            String apiKey = aiBot.getModelKEY(); // 替换为实际API Key
            DifyStreamClient client = new DifyStreamClient(apiUrl, apiKey, aiBotMessageService);
            client.fileUpload(userId,file);
            return Result.success("成功！");
        }catch (Exception e){
            return Result.fail(400, String.valueOf(e));
        }
    }



    public Result save(@RequestBody String jsonStr) {
        // 解析JSON
        JSONObject json = JSONObject.parseObject(jsonStr);

        // 合并所有secondMenuId*字段
        List<Integer> menuIds = new ArrayList<>();
        for (String key : json.keySet()) {
            if (key.startsWith("secondMenuId")) {
                Object value = json.get(key);
                if (value instanceof Integer) {
                    menuIds.add((Integer) value);
                }
            }
        }

            @Override
            public void onFailure(ChatContext context, Throwable throwable) {
                emitter.completeWithError(throwable);
            }
        });
        // 保留第一个ID（根据需要调整）
        if (!menuIds.isEmpty()) {
            json.put("secondMenuId", menuIds.get(0));
        }

        return emitter;
        // 转换为实体类
        AiBot entity = json.toJavaObject(AiBot.class);

        // 后续处理保持不变
        Result result = onSaveOrUpdateBefore(entity, true);
        if (result != null) return result;

        if (entity == null) {
            throw new NullPointerException("entity is null");
        }
        LoginAccount loginAccount = SaTokenUtil.getLoginAccount();
        commonFiled(entity,loginAccount.getId(),loginAccount.getTenantId(),loginAccount.getDeptId());
        boolean success = service.save(entity);
        onSaveOrUpdateAfter(entity, true);
        TableInfo tableInfo = TableInfoFactory.ofEntityClass(entity.getClass());
        Object[] pkArgs = tableInfo.buildPkSqlArgs(entity);
        return Result.create(success).set("id", pkArgs);
    }



    /**
     * 外部用户调用智能体进行对话
@@ -548,4 +745,4 @@


    }
}
}