shalu-system-designer/apps/designer/src/pages/mindmap/components/Config/AiCreator.tsx

import React, { useRef, useState } from "react";
import { CloseOutlined, LoadingOutlined, SendOutlined } from "@ant-design/icons";
import { Button, Tooltip, Input, Form, message } from "antd";
import { useChat } from "@/hooks/useChat";
import { Cell } from "@antv/x6";
import { handleParseFowlAIData } from "@/utils";
import { useModel } from "umi";
import { traverseNode } from "@/utils/mindmapHander";

export default function AICreator(props: {
  type: "mindmap" | "flow";
  onClose?: () => void;
  onChange?: (data: any) => void;
  onError?: (err: Error) => void;
}) {
  const [focused, setFocused] = React.useState(false);
  const [input, setInput] = useState("");
  const [messageApi, contextHolder] = message.useMessage();
  const messageKey = "ailoading";
  const msgContent = useRef<string>("");
  const { mindProjectInfo, setMindProjectInfo, graph } = useModel("mindMapModel");

  const { loading, onRequest, cancel } = useChat({
    app_name: "system_design",
    onUpdate: (msg) => {
      setInput("");
      msgContent.current += msg.answer;
    },
    onSuccess: (msg) => {
      console.log("加载完毕！", msgContent.current);
      messageApi.open({
        key: messageKey,
        type: "success",
        content: "AI创作完成",
        duration: 2,
        style: {
          marginTop: 300,
        },
      });
      handleParseFowlAIData({
        content: msgContent.current,
        message: {
          key: messageKey,
          instance: messageApi,
        },
        onSuccess: props.onChange,
        onError: props.onError,
      });
    },
    onError: (err) => {
      messageApi.open({
        key: messageKey,
        type: "error",
        content: err.message || "AI创作失败",
        duration: 2,
        style: {
          marginTop: 300,
        },
      });
    },
  });
  const handleStop = () => {
    cancel();
    messageApi.open({
      key: messageKey,
      type: "error",
      content: "AI创作已取消",
      duration: 2,
      style: {
        marginTop: 300,
      },
    });
  };

  // 处理提交
  const onSubmit = () => {
    if (input.trim()) {
      onRequest(`#角色：你是一个思维导图设计助手，根据需求部分的描述为用户使用设计思维导图工具生成对应的json数据
        #需求：${input}`, undefined, input);

      messageApi.open({
        key: messageKey,
        type: "loading",
        content: "AI创作中...",
        duration: 0,
        style: {
          marginTop: 300,
        },
      });
    }
  };

  React.useEffect(() => {
    return () => {
      // 取消所有进行中的请求
      const controller = new AbortController();
      controller.abort();
    };
  }, []);

  const handleList = [
    {
      key: "style",
      label: "风格美化",
      icon: "icon-yijianmeihua",
      color: "#a171f2",
    },
    {
      key: "grammar",
      label: "语法修复",
      icon: "icon-tubiao_yufajiucuo",
      color: "#00c4ad",
    },
    {
      key: "translation_en",
      label: "翻译为英文",
      icon: "icon-fanyiweiyingwen",
      color: "#8c4ff0",
    },
    {
      key: "translation_zh",
      label: "翻译为中文",
      icon: "icon-fanyiweizhongwen",
      color: "#3d72fb",
    },
  ];

  type LabelMap = Record<
    string,
    { cell: string; key: string; cellType: string }[]
  >;

  // 从元素中提取文本 当前无选中元素时，提取所有元素
  const getLabels = () => {
    const labelMap: LabelMap = {};
    let cells: Cell[] | undefined;

    cells = graph?.getSelectedCells();
    if (!cells || cells.length === 0) {
      cells = graph?.getCells();
    }

    if (!cells || !cells.length) return;

    cells.forEach((cell) => {
      const data = cell.getData();
      // 从label中提取文本
      if (data?.label?.trim()) {
        if (!labelMap[data.label.trim()]) {
          labelMap[data.label.trim()] = [
            { cell: cell.id, key: "label", cellType: cell.shape },
          ];
        } else {
          labelMap[data.label.trim()].push({
            cell: cell.id,
            key: "label",
            cellType: cell.shape,
          });
        }
      }
    });
    return labelMap;
  };

  // 替换节点文本内容
  const handleReplace = (labelMap: LabelMap, data: Record<string, string>) => {
    const keyMap: Record<string, string> = data;
    if(Array.isArray(data)) {
      data.forEach(item => {
        keyMap[item.original] = item.value
      })
    }
    mindProjectInfo && setMindProjectInfo({
      ...mindProjectInfo,
      topics: traverseNode(mindProjectInfo?.topics || [], (topic) => {
        // 判断当前节点是否需要替换
        const ids = labelMap[topic.label?.trim()]?.map(item => item.cell);
        if (keyMap[topic.label?.trim()] && ids?.includes(topic.id)) {
          topic.label = keyMap[topic.label?.trim()];
        }
        return topic
      })
    });
  };

  // 风格美化
  const handleStyle = () => {
    onRequest("生成一个美化风格的配置", {
      onUpdate: (data) => {
        console.log("style update:", data);
      },
      onSuccess: (data) => {
        console.log(data);
      },
      onError: (err) => {
        console.error(err);
      },
    });
  };

  // 语法修复
  const handleGrammar = () => {
    const labelMap = getLabels();
    if (!labelMap) {
      messageApi.open({
        key: messageKey,
        type: "info",
        content: "无可修复的数据",
      });
      return;
    }

    messageApi.open({
      key: messageKey,
      type: "loading",
      content: "AI正在修复语法...",
      duration: 0,
      style: {
        marginTop: 300,
      },
    });

    const data = JSON.stringify(Object.keys(labelMap));
    let result = "";
    onRequest(
      `修复语法错误，需要修复的数据：${data}`,
      {
        onUpdate: (data) => {
          result += data.answer;
        },
        onSuccess: (data) => {
          messageApi.open({
            key: messageKey,
            type: "success",
            content: "AI修复语法完成",
            duration: 2,
            style: {
              marginTop: 300,
            },
          });
          handleParseFowlAIData({
            content: result,
            message: {
              key: messageKey,
              instance: messageApi,
            },
            onSuccess: (data) => {
              handleReplace(labelMap, data);
            },
          });
        },
        onError: (err) => {
          console.error(err);
          messageApi.open({
            key: messageKey,
            type: "error",
            content: err.message || "AI创作失败",
            duration: 2,
            style: {
              marginTop: 300,
            },
          });
        },
      },
      "语法修复"
    );
  };

  // 翻译
  const handleTranslation = (lang: "en" | "zh") => {
    const labelMap = getLabels();
    if (!labelMap) {
      messageApi.open({
        key: messageKey,
        type: "info",
        content: "无可翻译的数据",
      });
      return;
    }

    messageApi.open({
      key: messageKey,
      type: "loading",
      content: "AI正在翻译...",
      duration: 0,
      style: {
        marginTop: 300,
      },
    });

    const data = JSON.stringify(Object.keys(labelMap));
    let result = "";
    onRequest(
      `翻译成${lang === "en" ? "英文" : "中文"}，需要翻译的数据：${data}`,
      {
        onUpdate: (data) => {
          result += data.answer;
        },
        onSuccess: (data) => {
          messageApi.open({
            key: messageKey,
            type: "success",
            content: "AI翻译完成",
            duration: 2,
            style: {
              marginTop: 300,
            },
          });
          handleParseFowlAIData({
            content: result,
            message: {
              key: messageKey,
              instance: messageApi,
            },
            onSuccess: (data) => {
              handleReplace(labelMap, data);
            },
          });
        },
        onError: (err) => {
          console.error(err);
          messageApi.open({
            key: messageKey,
            type: "error",
            content: err.message || "AI创作失败",
            duration: 2,
            style: {
              marginTop: 300,
            },
          });
        },
      },
      "翻译内容"
    );
  };

  const handleAiFeature = (feature: string) => {
    switch (feature) {
      case "style":
        handleStyle();
        break;
      case "grammar":
        handleGrammar();
        break;
      case "translation_en":
        handleTranslation("en");
        break;
      case "translation_zh":
        handleTranslation("zh");
        break;
      default:
        break;
    }
  };

  return (
    <div
      className="flex-1 h-full"
      style={{
        backgroundImage: "linear-gradient(137deg, #e5f4ff 0%, #efe7ff 100%)",
      }}
    >
      {contextHolder}
      <div className="chat-head w-full h-40px px-10px color-#333 flex items-center justify-between">
        <span className="text-14px">
          <svg className="icon h-32px w-32px" aria-hidden="true">
            <use xlinkHref="#icon-AI1"></use>
          </svg>
          <span>AI创作</span>
        </span>
        <span>
          <Button
            type="text"
            size="small"
            icon={<CloseOutlined />}
            onClick={() => props.onClose?.()}
          ></Button>
        </span>
      </div>

      <div className="text-14px pl-12px text-#333">一句话生成思维导图~</div>
      <div className="chat-content px-10px overflow-y-auto mt-12px">
        <div
          style={{
            borderColor: focused ? "#1890ff" : "#ddd",
          }}
          className="chat-foot bg-#fff rounded-10px border border-solid border-1px shadow-sm"
        >
          <Form onFinish={onSubmit}>
            <Input.TextArea
              rows={3}
              autoSize={{ maxRows: 3, minRows: 3 }}
              placeholder="输入你的主题或问题"
              variant="borderless"
              onFocus={() => setFocused(true)}
              onBlur={() => setFocused(false)}
              value={input}
              onChange={(e) => setInput(e.target.value)}
              disabled={loading}
              onPressEnter={onSubmit}
            />
            <div className="text-right p-10px">
              {loading ? (
                <Tooltip title="停止生成">
                  <Button
                    type="primary"
                    shape="circle"
                    icon={<LoadingOutlined />}
                    onClick={handleStop}
                  ></Button>
                </Tooltip>
              ) : (
                <Button
                  type="text"
                  icon={<SendOutlined />}
                  disabled={!input.trim()}
                  htmlType="submit"
                ></Button>
              )}
            </div>
          </Form>
        </div>
      </div>

      <div className="text-14px pl-12px text-#333 mt-32px">图形处理</div>
      <div className="flex flex-wrap gap-10px p-10px">
        {handleList.map((item) => (
          <div
            key={item.key}
            className="flex-[40%] h-50px bg-#fff rounded-10px shadow-sm flex items-center pl-10px text-12px cursor-pointer"
            onClick={() => handleAiFeature(item.key)}
          >
            <i
              className={`iconfont ${item.icon} text-16px`}
              style={{ color: item.color }}
            ></i>
            <span className="ml-10px">{item.label}</span>
          </div>
        ))}
      </div>
    </div>
  );
}