hudi-service/service-web/client/src/pages/ai/flow/FlowEditor.tsx

import {PlusCircleFilled, SaveFilled} from '@ant-design/icons'
import {
  Background,
  BackgroundVariant,
  Controls,
  type Edge,
  MiniMap,
  type Node,
  type NodeProps,
  ReactFlow,
} from '@xyflow/react'
import {useMount} from 'ahooks'
import type {Schema} from 'amis'
import {Button, Drawer, Dropdown, message, Space} from 'antd'
import {arrToMap, find, findIdx, isEqual, isNil, randomId} from 'licia'
import {type JSX, useState} from 'react'
import styled from 'styled-components'
import '@xyflow/react/dist/style.css'
import {amisRender, commonInfo} from '../../../util/amis.tsx'
import StartNode from './node/StartNode.tsx'
import EndNode from './node/EndNode.tsx'
import LlmNode from './node/LlmNode.tsx'
import {useDataStore} from './store/DataStore.ts'
import {useFlowStore} from './store/FlowStore.ts'

const FlowableDiv = styled.div`
    height: 93vh;

    .toolbar {
        z-index: 999;
        position: absolute;
    }

    .node-card {
        cursor: grab;

        .card-container {
            cursor: default;
        }
    }
`

const initialNodes: Node[] = [
  {
    id: 'BMFP3Eov94',
    type: 'start-amis-node',
    position: {x: 10, y: 100},
    data: {},
  },
  {
    id: 'PYK8LjduQ1',
    type: 'end-amis-node',
    position: {x: 500, y: 100},
    data: {},
  },
]
const initialEdges: Edge[] = []

function FlowEditor() {
  const [messageApi, contextHolder] = message.useMessage()
  const [nodeDef] = useState<{
    key: string,
    name: string,
    component: (props: NodeProps) => JSX.Element
  }[]>([
    {
      key: 'start-amis-node',
      name: '开始',
      component: StartNode,
    },
    {
      key: 'end-amis-node',
      name: '结束',
      component: EndNode,
    },
    {
      key: 'llm-amis-node',
      name: '大模型',
      component: LlmNode,
    },
  ])
  const [open, setOpen] = useState(false)

  const {getData, getDataById, setDataById} = useDataStore()
  const {
    nodes,
    addNode,
    removeNode,
    setNodes,
    onNodesChange,
    edges,
    setEdges,
    onEdgesChange,
    onConnect,
  } = useFlowStore()

  const [currentNodeForm, setCurrentNodeForm] = useState<JSX.Element>()
  const editNode = (id: string, name: string, description: string, columnSchema?: Schema[]) => {
    if (!isNil(columnSchema)) {
      setCurrentNodeForm(
        amisRender(
          {
            type: 'wrapper',
            size: 'none',
            body: [
              {
                type: 'tpl',
                className: 'text-secondary',
                tpl: description,
              },
              {
                debug: commonInfo.debug,
                title: name,
                type: 'form',
                wrapWithPanel: false,
                onEvent: {
                  submitSucc: {
                    actions: [
                      {
                        actionType: 'custom',
                        // @ts-ignore
                        script: (context, action, event) => {
                          setDataById(id, context.props.data)
                          setOpen(false)
                        },
                      },
                    ],
                  },
                },
                body: [
                  ...(columnSchema ?? []),
                  {
                    type: 'wrapper',
                    size: 'none',
                    className: 'space-x-1 float-right',
                    body: [
                      {
                        type: 'action',
                        label: '取消',
                        onEvent: {
                          click: {
                            actions: [
                              {
                                actionType: 'custom',
                                // @ts-ignore
                                script: (context, action, event) => {
                                  setOpen(false)
                                },
                              },
                            ],
                          },
                        },
                      },
                      {
                        type: 'submit',
                        label: '保存',
                        level: 'primary',
                      },
                    ],
                  },
                ],
              },
            ]
          },
          getDataById(id),
        ),
      )
      setOpen(true)
    }
  }

  useMount(() => {
    for (let node of initialNodes) {
      node.data = {
        getDataById,
        setDataById,
        removeNode,
        editNode,
      }
    }
    setNodes(initialNodes)
    setEdges(initialEdges)
  })

  return (
    <FlowableDiv>
      {contextHolder}
      <Space className="toolbar">
        <Button type="primary" onClick={() => console.log(JSON.stringify(getData()))}>
          <SaveFilled/>
          保存
        </Button>
        <Dropdown
          menu={{
            items: nodeDef.map(def => ({key: def.key, label: def.name})),
            onClick: ({key}) => {
              if (isEqual(key, 'start-amis-node') && findIdx(nodes, (node: Node) => isEqual(key, node.type)) > -1) {
                messageApi.error('只能存在1个开始节点')
                return
              }
              if (isEqual(key, 'end-amis-node') && findIdx(nodes, (node: Node) => isEqual(key, node.type)) > -1) {
                messageApi.error('只能存在1个结束节点')
                return
              }

              addNode({
                id: randomId(10),
                type: key,
                position: {x: 100, y: 100},
                data: {
                  getDataById,
                  setDataById,
                  removeNode,
                  editNode,
                },
              })
            },
          }}
        >
          <Button type="dashed">
            <PlusCircleFilled/>
            新增节点
          </Button>
        </Dropdown>
      </Space>
      <Drawer
        title="节点编辑"
        open={open}
        closeIcon={false}
        maskClosable={false}
        destroyOnHidden
      >
        {currentNodeForm}
      </Drawer>
      <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onConnect={onConnect}
        // @ts-ignore
        nodeTypes={arrToMap(
          nodeDef.map(def => def.key),
          key => find(nodeDef, def => isEqual(key, def.key))!.component)
        }
      >
        <Controls/>
        <MiniMap/>
        <Background variant={BackgroundVariant.Cross} gap={20} size={3}/>
      </ReactFlow>
    </FlowableDiv>
  )
}

export default FlowEditor