feat(builder): Visualise data beads on connections (#7791)

* Visualise data beads on edges

* Add `useBezierPath` hook

* Fix edge color on load

* Updates

* Merge branch 'master' into kpczerwinski/open-1580-visualise-data-coming-down-connections

* Add `visualizeBeads` state in `FlowEditor`

* Add `FlowContext`

Allow disabling beads animation

* fix(builder): linting

---------

Co-authored-by: Nicholas Tindle <nicholas.tindle@agpt.co>

rnd/autogpt_builder/src/components/CustomEdge.tsx

@@ -1,79 +1,172 @@
-import React, { FC, memo, useMemo, useState } from "react";
+import React, { FC, memo, useContext, useEffect, useState } from "react";
 import {
   BaseEdge,
   EdgeLabelRenderer,
   EdgeProps,
-  getBezierPath,
   useReactFlow,
   XYPosition,
 } from "reactflow";
 import "./customedge.css";
 import { X } from "lucide-react";
+import { useBezierPath } from "@/hooks/useBezierPath";
+import { FlowContext } from "./Flow";
 
 export type CustomEdgeData = {
   edgeColor: string;
   sourcePos?: XYPosition;
+  beadUp?: number;
+  beadDown?: number;
+  beadData?: any[];
+};
+
+type Bead = {
+  t: number;
+  targetT: number;
+  startTime: number;
 };
 
 const CustomEdgeFC: FC<EdgeProps<CustomEdgeData>> = ({
   id,
   data,
   selected,
-  source,
-  sourcePosition,
   sourceX,
   sourceY,
-  target,
-  targetPosition,
   targetX,
   targetY,
   markerEnd,
 }) => {
   const [isHovered, setIsHovered] = useState(false);
+  const [beads, setBeads] = useState<{
+    beads: Bead[];
+    created: number;
+    destroyed: number;
+  }>({ beads: [], created: 0, destroyed: 0 });
+  const { svgPath, length, getPointForT, getTForDistance } = useBezierPath(
+    sourceX - 5,
+    sourceY,
+    targetX + 3,
+    targetY,
+  );
   const { deleteElements } = useReactFlow<any, CustomEdgeData>();
+  const { visualizeBeads } = useContext(FlowContext) ?? {
+    visualizeBeads: "no",
+  };
 
   const onEdgeRemoveClick = () => {
     deleteElements({ edges: [{ id }] });
   };
 
-  const [path, labelX, labelY] = getBezierPath({
-    sourceX: sourceX - 5,
-    sourceY,
-    sourcePosition,
-    targetX: targetX + 4,
-    targetY,
-    targetPosition,
-  });
+  const animationDuration = 500; // Duration in milliseconds for bead to travel the curve
+  const beadDiameter = 10;
+  const deltaTime = 16;
 
-  // Calculate y difference between source and source node, to adjust self-loop edge
-  const yDifference = useMemo(
-    () => sourceY - (data?.sourcePos?.y || 0),
-    [data?.sourcePos?.y],
-  );
+  function setTargetPositions(beads: Bead[]) {
+    const distanceBetween = Math.min(
+      (length - beadDiameter) / (beads.length + 1),
+      beadDiameter,
+    );
 
-  // Define special edge path for self-loop
-  const edgePath =
-    source === target
-      ? `M ${sourceX - 5} ${sourceY} C ${sourceX + 128} ${sourceY - yDifference - 128} ${targetX - 128} ${sourceY - yDifference - 128} ${targetX + 3}, ${targetY}`
-      : path;
+    return beads.map((bead, index) => {
+      const targetPosition = distanceBetween * index + beadDiameter * 1.3;
+      const t = getTForDistance(-targetPosition);
 
-  console.table({
-    id,
-    sourceX,
-    sourceY,
-    targetX,
-    targetY,
-    sourcePosition,
-    targetPosition,
-    path,
-    labelX,
-    labelY,
-  });
+      return {
+        ...bead,
+        t: visualizeBeads === "animate" ? bead.t : t,
+        targetT: t,
+      } as Bead;
+    });
+  }
+
+  useEffect(() => {
+    if (data?.beadUp === 0 && data?.beadDown === 0) {
+      setBeads({ beads: [], created: 0, destroyed: 0 });
+      return;
+    }
+
+    const beadUp = data?.beadUp!;
+
+    setBeads(({ beads, created, destroyed }) => {
+      const newBeads = [];
+      for (let i = 0; i < beadUp - created; i++) {
+        newBeads.push({ t: 0, targetT: 0, startTime: Date.now() });
+      }
+
+      const b = setTargetPositions([...beads, ...newBeads]);
+      return { beads: b, created: beadUp, destroyed };
+    });
+
+    if (visualizeBeads !== "animate") {
+      setBeads(({ beads, created, destroyed }) => {
+        let destroyedCount = 0;
+
+        const newBeads = beads
+          .map((bead) => ({ ...bead }))
+          .filter((bead, index) => {
+            const beadDown = data?.beadDown!;
+
+            const removeCount = beadDown - destroyed;
+            if (bead.t >= bead.targetT && index < removeCount) {
+              destroyedCount++;
+              return false;
+            }
+            return true;
+          });
+
+        return {
+          beads: setTargetPositions(newBeads),
+          created,
+          destroyed: destroyed + destroyedCount,
+        };
+      });
+      return;
+    }
+
+    const interval = setInterval(() => {
+      setBeads(({ beads, created, destroyed }) => {
+        let destroyedCount = 0;
+
+        const newBeads = beads
+          .map((bead) => {
+            const progressIncrement = deltaTime / animationDuration;
+            const t = Math.min(
+              bead.t + bead.targetT * progressIncrement,
+              bead.targetT,
+            );
+
+            return {
+              ...bead,
+              t,
+            };
+          })
+          .filter((bead, index) => {
+            const beadDown = data?.beadDown!;
+
+            const removeCount = beadDown - destroyed;
+            if (bead.t >= bead.targetT && index < removeCount) {
+              destroyedCount++;
+              return false;
+            }
+            return true;
+          });
+
+        return {
+          beads: setTargetPositions(newBeads),
+          created,
+          destroyed: destroyed + destroyedCount,
+        };
+      });
+    }, deltaTime);
+
+    return () => clearInterval(interval);
+  }, [data]);
+
+  const middle = getPointForT(0.5);
 
   return (
     <>
       <BaseEdge
-        path={edgePath}
+        path={svgPath}
         markerEnd={markerEnd}
         style={{
           strokeWidth: isHovered ? 3 : 2,
@@ -83,7 +176,7 @@ const CustomEdgeFC: FC<EdgeProps<CustomEdgeData>> = ({
         }}
       />
       <path
-        d={edgePath}
+        d={svgPath}
         fill="none"
         strokeOpacity={0}
         strokeWidth={20}
@@ -95,7 +188,7 @@ const CustomEdgeFC: FC<EdgeProps<CustomEdgeData>> = ({
         <div
           style={{
             position: "absolute",
-            transform: `translate(-50%, -50%) translate(${labelX}px,${labelY}px)`,
+            transform: `translate(-50%, -50%) translate(${middle.x}px,${middle.y}px)`,
             pointerEvents: "all",
           }}
           className="edge-label-renderer"
@@ -110,6 +203,18 @@ const CustomEdgeFC: FC<EdgeProps<CustomEdgeData>> = ({
           </button>
         </div>
       </EdgeLabelRenderer>
+      {beads.beads.map((bead, index) => {
+        const pos = getPointForT(bead.t);
+        return (
+          <circle
+            key={index}
+            cx={pos.x}
+            cy={pos.y}
+            r={beadDiameter / 2} // Bead radius
+            fill={data?.edgeColor ?? "#555555"}
+          />
+        );
+      })}
     </>
   );
 };

rnd/autogpt_builder/src/components/Flow.tsx

@@ -6,6 +6,7 @@ import React, {
   useMemo,
   useRef,
   MouseEvent,
+  createContext,
 } from "react";
 import { shallow } from "zustand/vanilla/shallow";
 import ReactFlow, {
@@ -57,6 +58,12 @@ const MINIMUM_MOVE_BEFORE_LOG = 50;
 
 const ajv = new Ajv({ strict: false, allErrors: true });
 
+type FlowContextType = {
+  visualizeBeads: "no" | "static" | "animate";
+};
+
+export const FlowContext = createContext<FlowContextType | null>(null);
+
 const FlowEditor: React.FC<{
   flowID?: string;
   template?: boolean;
@@ -90,6 +97,9 @@ const FlowEditor: React.FC<{
   const [copiedNodes, setCopiedNodes] = useState<Node<CustomNodeData>[]>([]);
   const [copiedEdges, setCopiedEdges] = useState<Edge<CustomEdgeData>[]>([]);
   const [isAnyModalOpen, setIsAnyModalOpen] = useState(false); // Track if any modal is open
+  const [visualizeBeads, setVisualizeBeads] = useState<
+    "no" | "static" | "animate"
+  >("animate");
 
   const apiUrl = process.env.NEXT_PUBLIC_AGPT_SERVER_URL!;
   const api = useMemo(() => new AutoGPTServerAPI(apiUrl), [apiUrl]);
@@ -219,9 +229,10 @@ const FlowEditor: React.FC<{
   };
 
   // Function to clear status, output, and close the output info dropdown of all nodes
+  // and reset data beads on edges
   const clearNodesStatusAndOutput = useCallback(() => {
-    setNodes((nds) =>
-      nds.map((node) => ({
+    setNodes((nds) => {
+      const newNodes = nds.map((node) => ({
         ...node,
         data: {
           ...node.data,
@@ -229,8 +240,10 @@ const FlowEditor: React.FC<{
           output_data: undefined,
           isOutputOpen: false, // Close the output info dropdown
         },
-      })),
-    );
+      }));
+
+      return newNodes;
+    });
   }, [setNodes]);
 
   const onNodesChange = useCallback(
@@ -446,8 +459,8 @@ const FlowEditor: React.FC<{
     setAgentName(graph.name);
     setAgentDescription(graph.description);
 
-    setNodes(
-      graph.nodes.map((node) => {
+    setNodes(() => {
+      const newNodes = graph.nodes.map((node) => {
         const block = availableNodes.find(
           (block) => block.id === node.block_id,
         )!;
@@ -502,30 +515,38 @@ const FlowEditor: React.FC<{
           },
         };
         return newNode;
-      }),
-    );
-
-    setEdges(
-      graph.links.map((link) => ({
-        id: formatEdgeID(link),
-        type: "custom",
-        data: {
-          edgeColor: getTypeColor(
-            getOutputType(link.source_id, link.source_name),
-          ),
-          sourcePos: getNode(link.source_id)?.position,
-        },
-        markerEnd: {
-          type: MarkerType.ArrowClosed,
-          strokeWidth: 2,
-          color: getTypeColor(getOutputType(link.source_id, link.source_name)),
-        },
-        source: link.source_id,
-        target: link.sink_id,
-        sourceHandle: link.source_name || undefined,
-        targetHandle: link.sink_name || undefined,
-      })),
-    );
+      });
+      setEdges(
+        graph.links.map(
+          (link) =>
+            ({
+              id: formatEdgeID(link),
+              type: "custom",
+              data: {
+                edgeColor: getTypeColor(
+                  getOutputType(link.source_id, link.source_name!),
+                ),
+                sourcePos: getNode(link.source_id)?.position,
+                beadUp: 0,
+                beadDown: 0,
+                beadData: [],
+              },
+              markerEnd: {
+                type: MarkerType.ArrowClosed,
+                strokeWidth: 2,
+                color: getTypeColor(
+                  getOutputType(link.source_id, link.source_name!),
+                ),
+              },
+              source: link.source_id,
+              target: link.sink_id,
+              sourceHandle: link.source_name || undefined,
+              targetHandle: link.sink_name || undefined,
+            }) as Edge<CustomEdgeData>,
+        ),
+      );
+      return newNodes;
+    });
   }
 
   const prepareNodeInputData = (node: Node<CustomNodeData>) => {
@@ -593,6 +614,22 @@ const FlowEditor: React.FC<{
         },
       })),
     );
+    // Reset bead count
+    setEdges((edges) => {
+      return edges.map(
+        (edge) =>
+          ({
+            ...edge,
+            data: {
+              ...edge.data,
+              beadUp: 0,
+              beadDown: 0,
+              beadData: [],
+            },
+          }) as Edge<CustomEdgeData>,
+      );
+    });
+
     await new Promise((resolve) => setTimeout(resolve, 100));
 
     const nodes = getNodes();
@@ -763,28 +800,96 @@ const FlowEditor: React.FC<{
     }
   };
 
+  function getFrontendId(nodeId: string, nodes: Node<CustomNodeData>[]) {
+    const node = nodes.find((node) => node.data.backend_id === nodeId);
+    return node?.id;
+  }
+
+  function updateEdges(
+    executionData: NodeExecutionResult[],
+    nodes: Node<CustomNodeData>[],
+  ) {
+    setEdges((edges) => {
+      const newEdges = JSON.parse(
+        JSON.stringify(edges),
+      ) as Edge<CustomEdgeData>[];
+
+      executionData.forEach((exec) => {
+        if (exec.status === "COMPLETED") {
+          // Produce output beads
+          for (let key in exec.output_data) {
+            const outputEdges = newEdges.filter(
+              (edge) =>
+                edge.source === getFrontendId(exec.node_id, nodes) &&
+                edge.sourceHandle === key,
+            );
+            outputEdges.forEach((edge) => {
+              edge.data!.beadUp = (edge.data!.beadUp ?? 0) + 1;
+              //todo kcze this assumes output at key is always array with one element
+              edge.data!.beadData = [
+                exec.output_data[key][0],
+                ...edge.data!.beadData!,
+              ];
+            });
+          }
+        } else if (exec.status === "RUNNING") {
+          // Consume input beads
+          for (let key in exec.input_data) {
+            const inputEdges = newEdges.filter(
+              (edge) =>
+                edge.target === getFrontendId(exec.node_id, nodes) &&
+                edge.targetHandle === key,
+            );
+
+            inputEdges.forEach((edge) => {
+              if (
+                edge.data!.beadData![edge.data!.beadData!.length - 1] !==
+                exec.input_data[key]
+              ) {
+                return;
+              }
+              edge.data!.beadDown = (edge.data!.beadDown ?? 0) + 1;
+              edge.data!.beadData! = edge.data!.beadData!.slice(0, -1);
+            });
+          }
+        }
+      });
+
+      return newEdges;
+    });
+  }
+
   const updateNodesWithExecutionData = (
     executionData: NodeExecutionResult[],
   ) => {
-    setNodes((nds) =>
-      nds.map((node) => {
+    console.log("Updating nodes with execution data:", executionData);
+    setNodes((nodes) => {
+      if (visualizeBeads !== "no") {
+        updateEdges(executionData, nodes);
+      }
+
+      const updatedNodes = nodes.map((node) => {
         const nodeExecution = executionData.find(
           (exec) => exec.node_id === node.data.backend_id,
         );
-        if (nodeExecution) {
-          return {
-            ...node,
-            data: {
-              ...node.data,
-              status: nodeExecution.status,
-              output_data: nodeExecution.output_data,
-              isOutputOpen: true,
-            },
-          };
+
+        if (!nodeExecution || node.data.status === nodeExecution.status) {
+          return node;
         }
-        return node;
-      }),
-    );
+
+        return {
+          ...node,
+          data: {
+            ...node.data,
+            status: nodeExecution.status,
+            output_data: nodeExecution.output_data,
+            isOutputOpen: true,
+          },
+        };
+      });
+
+      return updatedNodes;
+    });
   };
 
   const handleKeyDown = useCallback(
@@ -897,34 +1002,36 @@ const FlowEditor: React.FC<{
   ];
 
   return (
-    <div className={className}>
-      <ReactFlow
-        nodeTypes={nodeTypes}
-        edgeTypes={edgeTypes}
-        connectionLineComponent={ConnectionLine}
-        onConnect={onConnect}
-        onNodesChange={onNodesChange}
-        onNodesDelete={onNodesDelete}
-        onEdgesChange={onEdgesChange}
-        onNodeDragStop={onNodeDragEnd}
-        onNodeDragStart={onNodeDragStart}
-        deleteKeyCode={["Backspace", "Delete"]}
-        minZoom={0.2}
-        maxZoom={2}
-      >
-        <Controls />
-        <Background />
-        <ControlPanel className="absolute z-10" controls={editorControls}>
-          <BlocksControl blocks={availableNodes} addBlock={addNode} />
-          <SaveControl
-            agentMeta={savedAgent}
-            onSave={saveAgent}
-            onDescriptionChange={setAgentDescription}
-            onNameChange={setAgentName}
-          />
-        </ControlPanel>
-      </ReactFlow>
-    </div>
+    <FlowContext.Provider value={{ visualizeBeads }}>
+      <div className={className}>
+        <ReactFlow
+          nodeTypes={nodeTypes}
+          edgeTypes={edgeTypes}
+          connectionLineComponent={ConnectionLine}
+          onConnect={onConnect}
+          onNodesChange={onNodesChange}
+          onNodesDelete={onNodesDelete}
+          onEdgesChange={onEdgesChange}
+          onNodeDragStop={onNodeDragEnd}
+          onNodeDragStart={onNodeDragStart}
+          deleteKeyCode={["Backspace", "Delete"]}
+          minZoom={0.2}
+          maxZoom={2}
+        >
+          <Controls />
+          <Background />
+          <ControlPanel className="absolute z-10" controls={editorControls}>
+            <BlocksControl blocks={availableNodes} addBlock={addNode} />
+            <SaveControl
+              agentMeta={savedAgent}
+              onSave={saveAgent}
+              onDescriptionChange={setAgentDescription}
+              onNameChange={setAgentName}
+            />
+          </ControlPanel>
+        </ReactFlow>
+      </div>
+    </FlowContext.Provider>
   );
 };
 

rnd/autogpt_builder/src/hooks/useBezierPath.ts

@@ -0,0 +1,157 @@
+import { useCallback, useMemo } from "react";
+
+type XYPosition = {
+  x: number;
+  y: number;
+};
+
+export type BezierPath = {
+  sourcePosition: XYPosition;
+  control1: XYPosition;
+  control2: XYPosition;
+  targetPosition: XYPosition;
+};
+
+export function useBezierPath(
+  sourceX: number,
+  sourceY: number,
+  targetX: number,
+  targetY: number,
+) {
+  const path: BezierPath = useMemo(() => {
+    const xDifference = Math.abs(sourceX - targetX);
+    const yDifference = Math.abs(sourceY - targetY);
+    const xControlDistance =
+      sourceX < targetX ? 64 : Math.max(xDifference / 2, 64);
+    const yControlDistance = yDifference < 128 && sourceX > targetX ? -64 : 0;
+
+    return {
+      sourcePosition: { x: sourceX, y: sourceY },
+      control1: {
+        x: sourceX + xControlDistance,
+        y: sourceY + yControlDistance,
+      },
+      control2: {
+        x: targetX - xControlDistance,
+        y: targetY + yControlDistance,
+      },
+      targetPosition: { x: targetX, y: targetY },
+    };
+  }, [sourceX, sourceY, targetX, targetY]);
+
+  const svgPath = useMemo(
+    () =>
+      `M ${path.sourcePosition.x} ${path.sourcePosition.y} ` +
+      `C ${path.control1.x} ${path.control1.y} ${path.control2.x} ${path.control2.y} ` +
+      `${path.targetPosition.x}, ${path.targetPosition.y}`,
+    [path],
+  );
+
+  const getPointForT = useCallback(
+    (t: number) => {
+      // Bezier formula: (1-t)^3 * p0 + 3*(1-t)^2*t*p1 + 3*(1-t)*t^2*p2 + t^3*p3
+      const x =
+        Math.pow(1 - t, 3) * path.sourcePosition.x +
+        3 * Math.pow(1 - t, 2) * t * path.control1.x +
+        3 * (1 - t) * Math.pow(t, 2) * path.control2.x +
+        Math.pow(t, 3) * path.targetPosition.x;
+
+      const y =
+        Math.pow(1 - t, 3) * path.sourcePosition.y +
+        3 * Math.pow(1 - t, 2) * t * path.control1.y +
+        3 * (1 - t) * Math.pow(t, 2) * path.control2.y +
+        Math.pow(t, 3) * path.targetPosition.y;
+
+      return { x, y };
+    },
+    [path],
+  );
+
+  const getArcLength = useCallback(
+    (t: number, samples: number = 100) => {
+      let length = 0;
+      let prevPoint = getPointForT(0);
+
+      for (let i = 1; i <= samples; i++) {
+        const currT = (i / samples) * t;
+        const currPoint = getPointForT(currT);
+        length += Math.sqrt(
+          Math.pow(currPoint.x - prevPoint.x, 2) +
+            Math.pow(currPoint.y - prevPoint.y, 2),
+        );
+        prevPoint = currPoint;
+      }
+
+      return length;
+    },
+    [path],
+  );
+
+  const length = useMemo(() => {
+    return getArcLength(1);
+  }, [path]);
+
+  const getBezierDerivative = useCallback(
+    (t: number) => {
+      const mt = 1 - t;
+      const x =
+        3 *
+        (mt * mt * (path.control1.x - path.sourcePosition.x) +
+          2 * mt * t * (path.control2.x - path.control1.x) +
+          t * t * (path.targetPosition.x - path.control2.x));
+      const y =
+        3 *
+        (mt * mt * (path.control1.y - path.sourcePosition.y) +
+          2 * mt * t * (path.control2.y - path.control1.y) +
+          t * t * (path.targetPosition.y - path.control2.y));
+      return { x, y };
+    },
+    [path],
+  );
+
+  const getTForDistance = useCallback(
+    (distance: number, epsilon: number = 0.0001) => {
+      if (distance < 0) {
+        distance = length + distance; // If distance is negative, calculate from the end of the curve
+      }
+
+      let t = distance / getArcLength(1);
+      let prevT = 0;
+
+      while (Math.abs(t - prevT) > epsilon) {
+        prevT = t;
+        const length = getArcLength(t);
+        const derivative = Math.sqrt(
+          Math.pow(getBezierDerivative(t).x, 2) +
+            Math.pow(getBezierDerivative(t).y, 2),
+        );
+        t -= (length - distance) / derivative;
+        t = Math.max(0, Math.min(1, t)); // Clamp t between 0 and 1
+      }
+
+      return t;
+    },
+    [path],
+  );
+
+  const getPointAtDistance = useCallback(
+    (distance: number) => {
+      if (distance < 0) {
+        distance = length + distance; // If distance is negative, calculate from the end of the curve
+      }
+
+      const t = getTForDistance(distance);
+      return getPointForT(t);
+    },
+    [path],
+  );
+
+  return {
+    path,
+    svgPath,
+    length,
+    getPointForT,
+    getTForDistance,
+    getPointAtDistance,
+  };
+}