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(""); 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("流程图"); 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) => { const keyMap: Record = 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 (
{contextHolder}
AI创作
绘制图形
帮我绘制-{graphType}
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}
))}
); }