feat: add FlowHeader and delete edge (#959)

### What problem does this PR solve?
feat: add FlowHeader and delete edge #918

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

web/package-lock.json

@@ -16,6 +16,7 @@
         "classnames": "^2.5.1",
         "dagre": "^0.8.5",
         "dayjs": "^1.11.10",
+        "elkjs": "^0.9.3",
         "eventsource-parser": "^1.1.2",
         "i18next": "^23.7.16",
         "i18next-browser-languagedetector": "^8.0.0",
@@ -11110,6 +11111,11 @@
       "resolved": "https://registry.npmmirror.com/electron-to-chromium/-/electron-to-chromium-1.4.767.tgz",
       "integrity": "sha512-nzzHfmQqBss7CE3apQHkHjXW77+8w3ubGCIoEijKCJebPufREaFETgGXWTkh32t259F3Kcq+R8MZdFdOJROgYw=="
     },
+    "node_modules/elkjs": {
+      "version": "0.9.3",
+      "resolved": "https://registry.npmmirror.com/elkjs/-/elkjs-0.9.3.tgz",
+      "integrity": "sha512-f/ZeWvW/BCXbhGEf1Ujp29EASo/lk1FDnETgNKwJrsVvGZhUWCZyg3xLJjAsxfOmt8KjswHmI5EwCQcPMpOYhQ=="
+    },
     "node_modules/elliptic": {
       "version": "6.5.5",
       "resolved": "https://registry.npmmirror.com/elliptic/-/elliptic-6.5.5.tgz",

web/package.json

@@ -21,6 +21,7 @@
     "classnames": "^2.5.1",
     "dagre": "^0.8.5",
     "dayjs": "^1.11.10",
+    "elkjs": "^0.9.3",
     "eventsource-parser": "^1.1.2",
     "i18next": "^23.7.16",
     "i18next-browser-languagedetector": "^8.0.0",

web/src/locales/zh-traditional.ts

@@ -24,7 +24,7 @@ export default {
       copied: '複製成功',
       comingSoon: '即將推出',
       download: '下載',
-      close: '关闭',
+      close: '關閉',
       preview: '預覽',
     },
     login: {

web/src/pages/flow/canvas/context-menu/index.tsx

@@ -84,9 +84,6 @@ export const useHandleNodeContextMenu = (sideWidth: number) => {
       //     event.clientY >= pane.height - 200 ? pane.height - event.clientY : 0,
       // });
 
-      console.info('clientX:', event.clientX);
-      console.info('clientY:', event.clientY);
-
       setMenu({
         id: node.id,
         top: event.clientY - 72,

web/src/pages/flow/canvas/index.tsx

@@ -17,9 +17,13 @@ import 'reactflow/dist/style.css';
 import { NodeContextMenu, useHandleNodeContextMenu } from './context-menu';
 
 import FlowDrawer from '../flow-drawer';
-import { useHandleDrop, useShowDrawer } from '../hooks';
-import { initialEdges, initialNodes } from '../mock';
-import { getLayoutedElements } from '../utils';
+import {
+  useHandleDrop,
+  useHandleKeyUp,
+  useHandleSelectionChange,
+  useShowDrawer,
+} from '../hooks';
+import { dsl } from '../mock';
 import { TextUpdaterNode } from './node';
 
 const nodeTypes = { textUpdater: TextUpdaterNode };
@@ -29,13 +33,11 @@ interface IProps {
 }
 
 function FlowCanvas({ sideWidth }: IProps) {
-  const { nodes: layoutedNodes, edges: layoutedEdges } = getLayoutedElements(
-    initialNodes,
-    initialEdges,
-    'LR',
-  );
-  const [nodes, setNodes] = useState<Node[]>(layoutedNodes);
-  const [edges, setEdges] = useState<Edge[]>(layoutedEdges);
+  const [nodes, setNodes] = useState<Node[]>(dsl.graph.nodes);
+  const [edges, setEdges] = useState<Edge[]>(dsl.graph.edges);
+
+  const { selectedEdges, selectedNodes } = useHandleSelectionChange();
+
   const { ref, menu, onNodeContextMenu, onPaneClick } =
     useHandleNodeContextMenu(sideWidth);
   const { drawerVisible, hideDrawer, showDrawer } = useShowDrawer();
@@ -60,6 +62,8 @@ function FlowCanvas({ sideWidth }: IProps) {
 
   const { onDrop, onDragOver, setReactFlowInstance } = useHandleDrop(setNodes);
 
+  const { handleKeyUp } = useHandleKeyUp(selectedEdges, selectedNodes);
+
   useEffect(() => {
     console.info('nodes:', nodes);
     console.info('edges:', edges);
@@ -82,6 +86,7 @@ function FlowCanvas({ sideWidth }: IProps) {
         onDragOver={onDragOver}
         onNodeClick={onNodeClick}
         onInit={setReactFlowInstance}
+        onKeyUp={handleKeyUp}
       >
         <Background />
         <Controls />

web/src/pages/flow/elk-hooks.ts

@@ -0,0 +1,35 @@
+import { useCallback, useLayoutEffect } from 'react';
+import { getLayoutedElements } from './elk-utils';
+
+export const elkOptions = {
+  'elk.algorithm': 'layered',
+  'elk.layered.spacing.nodeNodeBetweenLayers': '100',
+  'elk.spacing.nodeNode': '80',
+};
+
+export const useLayoutGraph = (
+  initialNodes,
+  initialEdges,
+  setNodes,
+  setEdges,
+) => {
+  const onLayout = useCallback(({ direction, useInitialNodes = false }) => {
+    const opts = { 'elk.direction': direction, ...elkOptions };
+    const ns = initialNodes;
+    const es = initialEdges;
+
+    getLayoutedElements(ns, es, opts).then(
+      ({ nodes: layoutedNodes, edges: layoutedEdges }) => {
+        setNodes(layoutedNodes);
+        setEdges(layoutedEdges);
+
+        // window.requestAnimationFrame(() => fitView());
+      },
+    );
+  }, []);
+
+  // Calculate the initial layout on mount.
+  useLayoutEffect(() => {
+    onLayout({ direction: 'RIGHT', useInitialNodes: true });
+  }, [onLayout]);
+};

web/src/pages/flow/elk-utils.ts

@@ -0,0 +1,42 @@
+import ELK from 'elkjs/lib/elk.bundled.js';
+import { Edge, Node } from 'reactflow';
+
+const elk = new ELK();
+
+export const getLayoutedElements = (
+  nodes: Node[],
+  edges: Edge[],
+  options = {},
+) => {
+  const isHorizontal = options?.['elk.direction'] === 'RIGHT';
+  const graph = {
+    id: 'root',
+    layoutOptions: options,
+    children: nodes.map((node) => ({
+      ...node,
+      // Adjust the target and source handle positions based on the layout
+      // direction.
+      targetPosition: isHorizontal ? 'left' : 'top',
+      sourcePosition: isHorizontal ? 'right' : 'bottom',
+
+      // Hardcode a width and height for elk to use when layouting.
+      width: 150,
+      height: 50,
+    })),
+    edges: edges,
+  };
+
+  return elk
+    .layout(graph)
+    .then((layoutedGraph) => ({
+      nodes: layoutedGraph.children.map((node) => ({
+        ...node,
+        // React Flow expects a position property on the node instead of `x`
+        // and `y` fields.
+        position: { x: node.x, y: node.y },
+      })),
+
+      edges: layoutedGraph.edges,
+    }))
+    .catch(console.error);
+};

web/src/pages/flow/header/index.less

@@ -0,0 +1,3 @@
+.flowHeader {
+  padding: 20px;
+}

web/src/pages/flow/header/index.tsx

@@ -0,0 +1,26 @@
+import { Button, Flex } from 'antd';
+
+import { useSaveGraph } from '../hooks';
+import styles from './index.less';
+
+const FlowHeader = () => {
+  const { saveGraph } = useSaveGraph();
+
+  return (
+    <Flex
+      align="center"
+      justify="end"
+      gap={'large'}
+      className={styles.flowHeader}
+    >
+      <Button>
+        <b>Debug</b>
+      </Button>
+      <Button type="primary" onClick={saveGraph}>
+        <b>Save</b>
+      </Button>
+    </Flex>
+  );
+};
+
+export default FlowHeader;

web/src/pages/flow/hooks.ts

@@ -1,6 +1,18 @@
 import { useSetModalState } from '@/hooks/commonHooks';
-import React, { Dispatch, SetStateAction, useCallback, useState } from 'react';
-import { Node, Position, ReactFlowInstance } from 'reactflow';
+import React, {
+  Dispatch,
+  KeyboardEventHandler,
+  SetStateAction,
+  useCallback,
+  useState,
+} from 'react';
+import {
+  Node,
+  Position,
+  ReactFlowInstance,
+  useOnSelectionChange,
+  useReactFlow,
+} from 'reactflow';
 import { v4 as uuidv4 } from 'uuid';
 
 export const useHandleDrag = () => {
@@ -75,3 +87,52 @@ export const useShowDrawer = () => {
     showDrawer,
   };
 };
+
+export const useHandleSelectionChange = () => {
+  const [selectedNodes, setSelectedNodes] = useState<string[]>([]);
+  const [selectedEdges, setSelectedEdges] = useState<string[]>([]);
+
+  useOnSelectionChange({
+    onChange: ({ nodes, edges }) => {
+      setSelectedNodes(nodes.map((node) => node.id));
+      setSelectedEdges(edges.map((edge) => edge.id));
+    },
+  });
+
+  return { selectedEdges, selectedNodes };
+};
+
+export const useDeleteEdge = (selectedEdges: string[]) => {
+  const { setEdges } = useReactFlow();
+
+  const deleteEdge = useCallback(() => {
+    setEdges((edges) =>
+      edges.filter((edge) => selectedEdges.every((x) => x !== edge.id)),
+    );
+  }, [setEdges, selectedEdges]);
+
+  return deleteEdge;
+};
+
+export const useHandleKeyUp = (
+  selectedEdges: string[],
+  selectedNodes: string[],
+) => {
+  const deleteEdge = useDeleteEdge(selectedEdges);
+  const handleKeyUp: KeyboardEventHandler = useCallback(
+    (e) => {
+      if (e.code === 'Delete') {
+        deleteEdge();
+      }
+    },
+    [deleteEdge],
+  );
+
+  return { handleKeyUp };
+};
+
+export const useSaveGraph = () => {
+  const saveGraph = useCallback(() => {}, []);
+
+  return { saveGraph };
+};

web/src/pages/flow/index.tsx

@@ -3,6 +3,7 @@ import { useState } from 'react';
 import { ReactFlowProvider } from 'reactflow';
 import FlowCanvas from './canvas';
 import Sider from './flow-sider';
+import FlowHeader from './header';
 
 const { Content } = Layout;
 
@@ -14,6 +15,7 @@ function RagFlow() {
       <ReactFlowProvider>
         <Sider setCollapsed={setCollapsed} collapsed={collapsed}></Sider>
         <Layout>
+          <FlowHeader></FlowHeader>
           <Content style={{ margin: '0 16px' }}>
             <FlowCanvas sideWidth={collapsed ? 0 : 200}></FlowCanvas>
           </Content>

web/src/pages/flow/mock.tsx

@@ -45,6 +45,100 @@ export const initialEdges = [
 ];
 
 export const dsl = {
+  graph: {
+    nodes: [
+      {
+        id: 'begin',
+        type: 'textUpdater',
+        position: {
+          x: 50,
+          y: 200,
+        },
+        data: {
+          label: 'Begin',
+        },
+        sourcePosition: 'left',
+        targetPosition: 'right',
+      },
+      {
+        id: 'Answer:China',
+        type: 'textUpdater',
+        position: {
+          x: 150,
+          y: 200,
+        },
+        data: {
+          label: 'Answer',
+        },
+        sourcePosition: 'left',
+        targetPosition: 'right',
+      },
+      {
+        id: 'Retrieval:China',
+        type: 'textUpdater',
+        position: {
+          x: 250,
+          y: 200,
+        },
+        data: {
+          label: 'Retrieval',
+        },
+        sourcePosition: 'left',
+        targetPosition: 'right',
+      },
+      {
+        id: 'Generate:China',
+        type: 'textUpdater',
+        position: {
+          x: 100,
+          y: 100,
+        },
+        data: {
+          label: 'Generate',
+        },
+        sourcePosition: 'left',
+        targetPosition: 'right',
+      },
+    ],
+    edges: [
+      {
+        id: '7facb53d-65c9-43b3-ac55-339c445d3891',
+        label: '',
+        source: 'begin',
+        target: 'Answer:China',
+        markerEnd: {
+          type: 'arrow',
+        },
+      },
+      {
+        id: '7ac83631-502d-410f-a6e7-bec6866a5e99',
+        label: '',
+        source: 'Generate:China',
+        target: 'Answer:China',
+        markerEnd: {
+          type: 'arrow',
+        },
+      },
+      {
+        id: '0aaab297-5779-43ed-9281-2c4d3741566f',
+        label: '',
+        source: 'Answer:China',
+        target: 'Retrieval:China',
+        markerEnd: {
+          type: 'arrow',
+        },
+      },
+      {
+        id: '3477f9f3-0a7d-400e-af96-a11ea7673183',
+        label: '',
+        source: 'Retrieval:China',
+        target: 'Generate:China',
+        markerEnd: {
+          type: 'arrow',
+        },
+      },
+    ],
+  },
   components: {
     begin: {
       obj: {

web/src/pages/flow/utils.ts

@@ -1,6 +1,6 @@
 import { DSLComponents } from '@/interfaces/database/flow';
 import dagre from 'dagre';
-import { Edge, Node, Position } from 'reactflow';
+import { Edge, MarkerType, Node, Position } from 'reactflow';
 import { v4 as uuidv4 } from 'uuid';
 
 const buildEdges = (
@@ -16,9 +16,12 @@ const buildEdges = (
       allEdges.push({
         id: uuidv4(),
         label: '',
-        type: 'step',
+        // type: 'step',
         source: source,
         target: target,
+        markerEnd: {
+          type: MarkerType.Arrow,
+        },
       });
     }
   });