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(""); 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) => { const keyMap: Record = 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 (
{contextHolder}
AI创作
一句话生成思维导图~
setFocused(true)} onBlur={() => setFocused(false)} value={input} onChange={(e) => setInput(e.target.value)} disabled={loading} onPressEnter={onSubmit} />
{loading ? ( ) : ( )}
图形处理
{handleList.map((item) => (
handleAiFeature(item.key)} > {item.label}
))}
); }