feat: preserve manual component data if code unchanged (#5482)

* chore: preserve manual component data if code unchanged * Fixed tool mode switch not changing with code * Removed console.log * Made tool mode not be lost when mutating template * use setNode instead of setNodes on component code * refactor: tool_mode return * Fix tool mode being reverted on toggle * fix: tool_mode outputs return * change: logic for setting `tool_mode` --------- Co-authored-by: Lucas Oliveira <[email protected]> Co-authored-by: Lucas Oliveira <[email protected]> Co-authored-by: Gabriel Luiz Freitas Almeida <[email protected]>
langflow-ai · Jan 20, 2025 · c8eb408 · c8eb408
1 parent c1cec32
commit c8eb408
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 30 deletions.
diff --git a/src/backend/base/langflow/template/utils.py b/src/backend/base/langflow/template/utils.py
@@ -78,5 +78,28 @@ def update_frontend_node_with_template_values(frontend_node, raw_frontend_node):
     old_code = raw_frontend_node["template"]["code"]["value"]
     new_code = frontend_node["template"]["code"]["value"]
     frontend_node["edited"] = raw_frontend_node["edited"] or (old_code != new_code)
+    frontend_node["tool_mode"] = raw_frontend_node.get("tool_mode", False)
+
+    if any(extract_tool_modes(raw_frontend_node["template"])):
+        frontend_node["tool_mode"] = False
+
+    if not frontend_node.get("edited", False):
+        frontend_node["display_name"] = raw_frontend_node["display_name"]
+        frontend_node["description"] = raw_frontend_node["description"]
 
     return frontend_node
+
+
+def extract_tool_modes(data: dict | list) -> list[bool | None]:
+    tool_models = []
+    if isinstance(data, dict):
+        for key, value in data.items():
+            if key == "tool_mode":
+                tool_models.append(value)
+            else:
+                tool_models.extend(extract_tool_modes(value))
+    elif isinstance(data, list):
+        for item in data:
+            tool_models.extend(extract_tool_modes(item))
+
+    return tool_models
diff --git a/src/frontend/src/CustomNodes/helpers/mutate-template.ts b/src/frontend/src/CustomNodes/helpers/mutate-template.ts
@@ -32,7 +32,7 @@ export const mutateTemplate = debounce(
       if (newTemplate) {
         newNode.template = newTemplate.template;
         newNode.outputs = newTemplate.outputs;
-        newNode.tool_mode = toolMode;
+        newNode.tool_mode = toolMode ?? node.tool_mode;
       }
       setNodeClass(newNode);
       callback?.();

diff --git a/...c/components/core/parameterRenderComponent/components/refreshParameterComponent/index.tsx b/...c/components/core/parameterRenderComponent/components/refreshParameterComponent/index.tsx
@@ -30,7 +30,6 @@ export function RefreshParameterComponent({
     parameterId: name,
     nodeId: nodeId,
     node: nodeClass,
-    tool_mode: nodeClass.tool_mode ?? false,
   });
 
   const setErrorData = useAlertStore((state) => state.setErrorData);

diff --git a/src/frontend/src/modals/codeAreaModal/index.tsx b/src/frontend/src/modals/codeAreaModal/index.tsx
@@ -7,7 +7,6 @@ import "ace-builds/src-noconflict/theme-twilight";
 import { usePostValidateCode } from "@/controllers/API/queries/nodes/use-post-validate-code";
 import { usePostValidateComponentCode } from "@/controllers/API/queries/nodes/use-post-validate-component-code";
 import useFlowStore from "@/stores/flowStore";
-import { cloneDeep } from "lodash";
 import { useEffect, useRef, useState } from "react";
 import AceEditor from "react-ace";
 import ReactAce from "react-ace/lib/ace";
@@ -61,9 +60,7 @@ export default function CodeAreaModal({
   } | null>(null);
 
   const { mutate: validateComponentCode } = usePostValidateComponentCode();
-  const currentFlow = useFlowStore((state) => state.currentFlow);
-  const nodes = useFlowStore((state) => state.nodes);
-  const setNodes = useFlowStore((state) => state.setNodes);
+  const setNode = useFlowStore((state) => state.setNode);
 
   useEffect(() => {
     // if nodeClass.template has more fields other than code and dynamic is true
@@ -126,17 +123,6 @@ export default function CodeAreaModal({
           if (data && type) {
             setValue(code);
             setNodeClass(data, type);
-            const currentNode = nodes.find((node) => node.id === componentId);
-            const currentNodeIndex = nodes.findIndex(
-              (node) => node.id === componentId,
-            );
-            const currentNodes = cloneDeep(nodes);
-
-            if (currentNode) {
-              currentNodes[currentNodeIndex].data.node = data;
-            }
-            setNodes(currentNodes);
-
             setError({ detail: { error: undefined, traceback: undefined } });
             setOpen(false);
           }

diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -143,20 +143,26 @@ const NodeToolbarComponent = memo(
       [data, numberOfOutputHandles],
     );
 
-    const [toolMode, setToolMode] = useState(() => {
-      // Check if tool mode is explicitly set on the node
-      const hasToolModeProperty = data.node?.tool_mode;
-      if (hasToolModeProperty) {
-        return hasToolModeProperty;
-      }
-
-      // Otherwise check if node has component_as_tool output
-      const hasComponentAsTool = data.node?.outputs?.some(
-        (output) => output.name === "component_as_tool",
-      );
+    const [toolMode, setToolMode] = useState(
+      () =>
+        data.node?.tool_mode ||
+        data.node?.outputs?.some(
+          (output) => output.name === "component_as_tool",
+        ) ||
+        false,
+    );
 
-      return hasComponentAsTool ?? false;
-    });
+    useEffect(() => {
+      if (data.node?.tool_mode !== undefined) {
+        setToolMode(
+          data.node?.tool_mode ||
+            data.node?.outputs?.some(
+              (output) => output.name === "component_as_tool",
+            ) ||
+            false,
+        );
+      }
+    }, [data.node?.tool_mode, data.node?.outputs]);
 
     const { handleNodeClass: handleNodeClassHook } = useHandleNodeClass(
       data.id,