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

import React, { useRef, useState } from "react";
import {
  CloseOutlined,
  SendOutlined,
  CaretDownOutlined,
} from "@ant-design/icons";
import { Button, Tooltip, Input, Form, Dropdown, message } from "antd";
import type { DropDownProps } from "antd";
import { useChat } from "@/hooks/useChat";
import { useModel } from "umi";
import { Cell } from "@antv/x6";
import { handleParseAIData } from "@/utils";

const items = [
  { key: "1", label: "流程图" },
  { key: "2", label: "泳道图" },
  { key: "3", label: "ER实体图" },
  { key: "4", label: "组织图" },
  { key: "5", label: "时序图" },
];
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 { graph } = useModel("flowchartModel");

  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,
        },
      });
      handleParseAIData({
        content: msgContent.current,
        onSuccess: props.onChange,
        onError: props.onError,
        message: {
          key: messageKey,
          instance: messageApi,
        },
      });
    },
    onError: (err) => {
      messageApi.open({
        key: messageKey,
        type: "error",
        content: err.message || "AI创作失败",
        duration: 2,
        style: {
          marginTop: 300,
        },
      });
    },
  });

  const [graphType, setGraphType] = useState<string>("流程图");
  const dropDownMenu: DropDownProps["menu"] = {
    items,
    onClick: (info) => {
      console.log(info);
      const type = items.find((item) => item.key === info.key);
      setGraphType(type?.label || "流程图");
    },
  };

  const handleStop = () => {
    cancel();
    messageApi.open({
      key: messageKey,
      type: "error",
      content: "AI创作已取消",
      duration: 2,
      style: {
        marginTop: 300,
      },
    });
  };

  // 处理提交
  const onSubmit = () => {
    if (input.trim()) {
      onRequest(
        `设计一个${graphType}, 返回图形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,
          });
        }
      }
      // 从name中提取文本
      if (data?.name?.trim()) {
        if (!labelMap[data.name.trim()]) {
          labelMap[data.name.trim()] = [
            { cell: cell.id, key: "name", cellType: cell.shape },
          ];
        } else {
          labelMap[data.name.trim()].push({
            cell: cell.id,
            key: "name",
            cellType: cell.shape,
          });
        }
      }
      // 从边线中提取文本
      if(cell.isEdge()) {
        (cell.labels || []).forEach((label) => {
          const labelText = (label?.attrs?.label?.text as string)?.trim();
          if(labelText) {
            if (!labelMap[labelText]) {
              labelMap[labelText] = [
                { cell: cell.id, key: "label", cellType: cell.shape },
              ];
            } else {
              labelMap[labelText].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
      })
    }
    Object.keys(keyMap).forEach((key) => {
      if (labelMap[key]) {
        labelMap[key].forEach((item) => {
          const cell = graph?.getCellById(item.cell);
          if (cell && cell.shape !== "edge") {
            cell.setData({
              [item.key]: keyMap[key],
            });
            graph?.select(cell);
          } else if (cell?.isEdge()) {
            // 设置边线文本
            const labels = cell.getLabels();
            cell.setLabels(labels.map(item => {
              return {
                ...item,
                attrs: {
                  label: {
                    text: keyMap?.[(item.attrs?.label?.text as string)?.trim()] || item.attrs?.label?.text
                  }
                }
              }
            }))
          }
        });
      }
    });
  };

  // 风格美化
  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,
            },
          });
          handleParseAIData({
            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,
            },
          });
          handleParseAIData({
            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"
        >
          <Dropdown menu={dropDownMenu} placement="bottomLeft">
            <div className="text-12px pl-10px pt-10px cursor-pointer">
              帮我绘制-{graphType}
              <CaretDownOutlined />
            </div>
          </Dropdown>
          <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={<i className="iconfont icon-stopcircle" />}
                    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>
  );
}