Merge pull request #446 from kreneskyp/graph_ui

UI for LangGraph state machines

frontend/chains/ChainGraphEditor.js

@@ -304,6 +304,7 @@ const ChainGraphEditor = ({ graph }) => {
       // create reactflow edge
       const id = uuid4();
       const source = reactFlowInstance.getNode(params.source);
+      const target = reactFlowInstance.getNode(params.target);
       const flowNodeType =
         source.id === "root" ? "root" : source.data.type.type;
       const style = getEdgeStyle(colorMode, flowNodeType);
@@ -314,8 +315,14 @@ const ChainGraphEditor = ({ graph }) => {
       const is_out = params.sourceHandle === "out";
       const is_in = params.targetHandle === "in";
       const from_branch = source.data.type.type === "branch";
+      const with_graph =
+        source.data.type.type === "graph" ||
+        (target.data.type.type === "graph" && params.targetHandle === "loop");
+
       const displayRelation =
-        from_root || is_out || is_in || from_branch ? "LINK" : "PROP";
+        from_root || is_out || is_in || from_branch || with_graph
+          ? "LINK"
+          : "PROP";
       setEdges((els) =>
         addEdge({ ...params, ...style[displayRelation], data: { id } }, els)
       );
@@ -330,7 +337,13 @@ const ChainGraphEditor = ({ graph }) => {
         root_node_ids.push(params.target);
         api.setRoots(chain.id, { node_ids: root_node_ids });
       } else {
-        const relation = from_root || is_out || from_branch ? "LINK" : "PROP";
+        let relation = "PROP";
+        if (with_graph) {
+          relation = "GRAPH";
+        } else if (from_root || is_out || from_branch) {
+          relation = "LINK";
+        }
+
         const data = {
           id,
           source_id: params.source,

frontend/chains/LinkEdge.js

@@ -12,8 +12,6 @@ export const LinkEdge = (props) => {
   // Apply scaling in addition to the translation
   const transform = `translate(${markerOffsetX}, ${markerOffsetY}) scale(1.3)`;
 
-  console.log("props", props);
-
   return (
     <BaseEdge
       {...props}

frontend/chains/editor/styles.js

@@ -23,6 +23,8 @@ import { RootNode } from "chains/flow/RootNode";
 import { LoopIcon } from "icons/LoopIcon";
 import { SplitIcon } from "icons/SplitIcon";
 import { MergeIcon } from "icons/MergeIcon";
+import { GraphIcon } from "icons/GraphIcon";
+import { GraphNode } from "chains/flow/GraphNode";
 
 export const NODE_STYLES = {
   llm: {
@@ -105,6 +107,12 @@ export const NODE_STYLES = {
     },
     component: BranchNode,
   },
+  graph: {
+    icon: {
+      component: GraphIcon,
+    },
+    component: GraphNode,
+  },
   map: {
     icon: {
       component: MergeIcon,

frontend/chains/editor/useColorMode.js

@@ -72,6 +72,7 @@ export const useEditorColorMode = () => {
       memory_backend: isLight ? "purple.500" : "purple.300",
       chain: isLight ? "blue.500" : "blue.300",
       agent: isLight ? "green.500" : "green.600",
+      graph: isLight ? "green.500" : "green.600",
       parser: retrieval,
       retriever: retrieval,
       schema: data,

frontend/chains/flow/BranchNode.js

@@ -8,9 +8,19 @@ import { ConnectorPopover } from "chains/editor/ConnectorPopover";
 const useFlowConnectors = (node) => {
   const edges = useEdges();
   return useMemo(() => {
+    const input = {
+      key: "in",
+      type: "target",
+      required: true,
+      connected: edges?.find((edge) => edge.target === node.id),
+      source_type: ["root", "agent", "chain", "flow"],
+    };
+
     return {
-      input: {
-        key: "in",
+      input: input,
+      in: input,
+      loop: {
+        key: "loop",
         type: "target",
         required: true,
         connected: edges?.find((edge) => edge.target === node.id),
@@ -40,26 +50,32 @@ const useFlowConnectors = (node) => {
   }, [edges, node.id, node?.config?.branches]);
 };
 
-export const InputConnector = ({ type, node }) => {
-  const { input } = useFlowConnectors(node);
-  const intputColor = useConnectorColor(node, input);
+export const InputConnector = ({
+  type,
+  node,
+  required = true,
+  id = "in",
+  label = "Input",
+}) => {
+  const connector = useFlowConnectors(node)[id];
+  const inputColor = useConnectorColor(node, connector);
   return (
     <Box position="relative">
       <Handle
-        id="in"
+        id={id}
         type="target"
         position="left"
         style={{ top: "50%", transform: "translateY(-50%)" }}
       />
-      <Heading fontSize="xs" px={2} color={intputColor}>
+      <Heading fontSize="xs" px={2} color={inputColor}>
         <ConnectorPopover
           type={type}
           node={node}
-          connector={input}
-          label={"Input"}
+          connector={connector}
+          label={label}
           placement={"left"}
         />{" "}
-        <RequiredAsterisk color={intputColor} />
+        {required && <RequiredAsterisk color={inputColor} />}
       </Heading>
     </Box>
   );

frontend/chains/flow/GraphNode.js

@@ -0,0 +1,25 @@
+import React from "react";
+import { BranchConnectors, InputConnector } from "chains/flow/BranchNode";
+import { VStack, Flex } from "@chakra-ui/react";
+import { OutputConnector } from "chains/flow/ChainNode";
+
+export const GraphNode = ({ type, node, config, onFieldChange }) => {
+  return (
+    <VStack spacing={0} alignItems="stretch" fontSize="xs">
+      <Flex mt={1} mb={3} justify={"space-between"}>
+        <InputConnector type={type} node={node} />
+        <OutputConnector type={type} node={node} />
+      </Flex>
+      <Flex mt={1} mb={3} justify={"space-between"}>
+        <InputConnector
+          type={type}
+          node={node}
+          id={"loop"}
+          label={"Loop"}
+          required={false}
+        />
+        <BranchConnectors node={node} />
+      </Flex>
+    </VStack>
+  );
+};

frontend/components/HashListForm.js

@@ -20,10 +20,10 @@ export const HashListForm = ({
 
   const handleInputChange = (e, index) => {
     const { value } = e.target;
-    const updatedItems = [...list];
-    const updatedHashes = [...hash_list];
+    const updatedItems = [...(list || [])];
+    const updatedHashes = [...(hash_list || [])];
 
-    if (index < list.length) {
+    if (index < list?.length || 0) {
       updatedItems[index] = value;
     } else {
       updatedItems.push(value);
@@ -63,8 +63,8 @@ export const HashListForm = ({
             <FontAwesomeIcon
               icon={faTrash}
               onClick={() => handleRemoveClick(index)}
-              cursor={index < list.length ? "pointer" : "default"}
-              color={index < list.length ? "inherit" : "transparent"}
+              cursor={index < list?.length ? "pointer" : "default"}
+              color={index < list?.length ? "inherit" : "transparent"}
             />
           </HStack>
         ))}