feat: test buildNodesAndEdgesFromDSLComponents (#940)

### What problem does this PR solve?
 feat: test buildNodesAndEdgesFromDSLComponents #918

### Type of change

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

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

@@ -8,7 +8,6 @@ import ReactFlow, {
   OnConnect,
   OnEdgesChange,
   OnNodesChange,
-  Position,
   addEdge,
   applyEdgeChanges,
   applyNodeChanges,
@@ -19,47 +18,24 @@ 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 { TextUpdaterNode } from './node';
 
 const nodeTypes = { textUpdater: TextUpdaterNode };
 
-const initialNodes = [
-  {
-    sourcePosition: Position.Left,
-    targetPosition: Position.Right,
-    id: 'node-1',
-    type: 'textUpdater',
-    position: { x: 400, y: 100 },
-    data: { label: 123 },
-  },
-  {
-    sourcePosition: Position.Right,
-    targetPosition: Position.Left,
-    id: '1',
-    data: { label: 'Hello' },
-    position: { x: 0, y: 50 },
-    type: 'input',
-  },
-  {
-    sourcePosition: Position.Right,
-    targetPosition: Position.Left,
-    id: '2',
-    data: { label: 'World' },
-    position: { x: 200, y: 50 },
-  },
-];
-
-const initialEdges = [
-  { id: '1-2', source: '1', target: '2', label: 'to the', type: 'step' },
-];
-
 interface IProps {
   sideWidth: number;
 }
 
 function FlowCanvas({ sideWidth }: IProps) {
-  const [nodes, setNodes] = useState<Node[]>(initialNodes);
-  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+  const { nodes: layoutedNodes, edges: layoutedEdges } = getLayoutedElements(
+    initialNodes,
+    initialEdges,
+    'LR',
+  );
+  const [nodes, setNodes] = useState<Node[]>(layoutedNodes);
+  const [edges, setEdges] = useState<Edge[]>(layoutedEdges);
   const { ref, menu, onNodeContextMenu, onPaneClick } =
     useHandleNodeContextMenu(sideWidth);
   const { drawerVisible, hideDrawer, showDrawer } = useShowDrawer();

web/src/pages/flow/mock.tsx

@@ -3,6 +3,7 @@ import {
   RocketOutlined,
   SendOutlined,
 } from '@ant-design/icons';
+import { Position } from 'reactflow';
 
 export const componentList = [
   { name: 'Begin', icon: <SendOutlined />, description: '' },
@@ -10,6 +11,39 @@ export const componentList = [
   { name: 'Generate', icon: <MergeCellsOutlined />, description: '' },
 ];
 
+export const initialNodes = [
+  {
+    sourcePosition: Position.Left,
+    targetPosition: Position.Right,
+    id: 'node-1',
+    type: 'textUpdater',
+    position: { x: 0, y: 0 },
+    // position: { x: 400, y: 100 },
+    data: { label: 123 },
+  },
+  {
+    sourcePosition: Position.Right,
+    targetPosition: Position.Left,
+    id: '1',
+    data: { label: 'Hello' },
+    position: { x: 0, y: 0 },
+    // position: { x: 0, y: 50 },
+    type: 'input',
+  },
+  {
+    sourcePosition: Position.Right,
+    targetPosition: Position.Left,
+    id: '2',
+    data: { label: 'World' },
+    position: { x: 0, y: 0 },
+    // position: { x: 200, y: 50 },
+  },
+];
+
+export const initialEdges = [
+  { id: '1-2', source: '1', target: '2', label: 'to the', type: 'step' },
+];
+
 export const dsl = {
   components: {
     begin: {
@@ -17,8 +51,8 @@ export const dsl = {
         component_name: 'Begin',
         params: {},
       },
-      downstream: ['Answer:China'],
-      upstream: [],
+      downstream: ['Answer:China'], // other edge target is downstream, edge source is current node id
+      upstream: [], // edge source is upstream, edge target is current node id
     },
     'Answer:China': {
       obj: {

web/src/pages/flow/utils.test.ts

@@ -0,0 +1,30 @@
+import { dsl } from './mock';
+import { buildNodesAndEdgesFromDSLComponents } from './utils';
+
+test('buildNodesAndEdgesFromDSLComponents', () => {
+  const { edges, nodes } = buildNodesAndEdgesFromDSLComponents(dsl.components);
+
+  expect(nodes.length).toEqual(4);
+  expect(edges.length).toEqual(4);
+
+  expect(edges).toEqual(
+    expect.arrayContaining([
+      expect.objectContaining({
+        source: 'begin',
+        target: 'Answer:China',
+      }),
+      expect.objectContaining({
+        source: 'Answer:China',
+        target: 'Retrieval:China',
+      }),
+      expect.objectContaining({
+        source: 'Retrieval:China',
+        target: 'Generate:China',
+      }),
+      expect.objectContaining({
+        source: 'Generate:China',
+        target: 'Answer:China',
+      }),
+    ]),
+  );
+});

web/src/pages/flow/utils.ts

@@ -1,10 +1,32 @@
 import { DSLComponents } from '@/interfaces/database/flow';
+import dagre from 'dagre';
 import { Edge, Node, Position } from 'reactflow';
 import { v4 as uuidv4 } from 'uuid';
 
-export const buildNodesFromDSLComponents = (data: DSLComponents) => {
+const buildEdges = (
+  operatorIds: string[],
+  currentId: string,
+  allEdges: Edge[],
+  isUpstream = false,
+) => {
+  operatorIds.forEach((cur) => {
+    const source = isUpstream ? cur : currentId;
+    const target = isUpstream ? currentId : cur;
+    if (!allEdges.some((e) => e.source === source && e.target === target)) {
+      allEdges.push({
+        id: uuidv4(),
+        label: '',
+        type: 'step',
+        source: source,
+        target: target,
+      });
+    }
+  });
+};
+
+export const buildNodesAndEdgesFromDSLComponents = (data: DSLComponents) => {
   const nodes: Node[] = [];
-  const edges: Edge[] = [];
+  let edges: Edge[] = [];
 
   Object.entries(data).forEach(([key, value]) => {
     const downstream = [...value.downstream];
@@ -23,22 +45,51 @@ export const buildNodesFromDSLComponents = (data: DSLComponents) => {
       targetPosition: Position.Right,
     });
 
-    // intermediate node
-    // The first and last nodes do not need to be considered
-    if (upstream.length > 0 && downstream.length > 0) {
-      for (let i = 0; i < upstream.length; i++) {
-        const up = upstream[i];
-        for (let j = 0; j < downstream.length; j++) {
-          const down = downstream[j];
-          edges.push({
-            id: uuidv4(),
-            label: '',
-            type: 'step',
-            source: up,
-            target: down,
-          });
-        }
-      }
-    }
+    buildEdges(upstream, key, edges, true);
+    buildEdges(downstream, key, edges, false);
   });
+
+  return { nodes, edges };
+};
+
+const dagreGraph = new dagre.graphlib.Graph();
+dagreGraph.setDefaultEdgeLabel(() => ({}));
+
+const nodeWidth = 172;
+const nodeHeight = 36;
+
+export const getLayoutedElements = (
+  nodes: Node[],
+  edges: Edge[],
+  direction = 'TB',
+) => {
+  const isHorizontal = direction === 'LR';
+  dagreGraph.setGraph({ rankdir: direction });
+
+  nodes.forEach((node) => {
+    dagreGraph.setNode(node.id, { width: nodeWidth, height: nodeHeight });
+  });
+
+  edges.forEach((edge) => {
+    dagreGraph.setEdge(edge.source, edge.target);
+  });
+
+  dagre.layout(dagreGraph);
+
+  nodes.forEach((node) => {
+    const nodeWithPosition = dagreGraph.node(node.id);
+    node.targetPosition = isHorizontal ? Position.Left : Position.Top;
+    node.sourcePosition = isHorizontal ? Position.Right : Position.Bottom;
+
+    // We are shifting the dagre node position (anchor=center center) to the top left
+    // so it matches the React Flow node anchor point (top left).
+    node.position = {
+      x: nodeWithPosition.x - nodeWidth / 2,
+      y: nodeWithPosition.y - nodeHeight / 2,
+    };
+
+    return node;
+  });
+
+  return { nodes, edges };
 };