From 8e0d8fdb3f4c6ff4844a34e81208cfae8a316fca Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Thu, 21 Mar 2024 16:06:35 +0800
Subject: [PATCH] feat: other nodes support rename and fix knonw set var bug

---
 .../components/workflow/hooks/use-workflow.ts |  2 +-
 .../nodes/_base/components/variable/utils.ts  | 78 +++++++++++++++++++
 .../nodes/knowledge-retrieval/use-config.ts   |  3 +-
 3 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/web/app/components/workflow/hooks/use-workflow.ts b/web/app/components/workflow/hooks/use-workflow.ts
index 0a7f71973c..85bfbf9df8 100644
--- a/web/app/components/workflow/hooks/use-workflow.ts
+++ b/web/app/components/workflow/hooks/use-workflow.ts
@@ -216,7 +216,7 @@ export const useWorkflow = () => {
     const { getNodes, setNodes } = store.getState()
     const afterNodes = getAfterNodesInSameBranch(nodeId)
     const effectNodes = findUsedVarNodes(oldValeSelector, afterNodes)
-    console.log(effectNodes)
+    // console.log(effectNodes)
     if (effectNodes.length > 0) {
       // debugger
       const newNodes = getNodes().map((node) => {
diff --git a/web/app/components/workflow/nodes/_base/components/variable/utils.ts b/web/app/components/workflow/nodes/_base/components/variable/utils.ts
index 73bc790e89..4e45b90d4e 100644
--- a/web/app/components/workflow/nodes/_base/components/variable/utils.ts
+++ b/web/app/components/workflow/nodes/_base/components/variable/utils.ts
@@ -336,6 +336,84 @@ export const updateNodeVars = (oldNode: Node, oldVarSelector: ValueSelector, new
 
         break
       }
+      case BlockEnum.KnowledgeRetrieval: {
+        const payload = data as KnowledgeRetrievalNodeType
+        if (payload.query_variable_selector.join('.') === oldVarSelector.join('.'))
+          payload.query_variable_selector = newVarSelector
+        break
+      }
+      case BlockEnum.IfElse: {
+        const payload = data as IfElseNodeType
+        if (payload.conditions) {
+          payload.conditions = payload.conditions.map((c) => {
+            if (c.variable_selector.join('.') === oldVarSelector.join('.'))
+              c.variable_selector = newVarSelector
+            return c
+          })
+        }
+        break
+      }
+      case BlockEnum.Code: {
+        const payload = data as CodeNodeType
+        if (payload.variables) {
+          payload.variables = payload.variables.map((v) => {
+            if (v.value_selector.join('.') === oldVarSelector.join('.'))
+              v.value_selector = newVarSelector
+            return v
+          })
+        }
+        break
+      }
+      case BlockEnum.TemplateTransform: {
+        const payload = data as TemplateTransformNodeType
+        if (payload.variables) {
+          payload.variables = payload.variables.map((v: any) => {
+            if (v.value_selector.join('.') === oldVarSelector.join('.'))
+              v.value_selector = newVarSelector
+            return v
+          })
+        }
+        break
+      }
+      case BlockEnum.QuestionClassifier: {
+        const payload = data as QuestionClassifierNodeType
+        if (payload.query_variable_selector.join('.') === oldVarSelector.join('.'))
+          payload.query_variable_selector = newVarSelector
+        break
+      }
+      case BlockEnum.HttpRequest: {
+        const payload = data as HttpNodeType
+        if (payload.variables) {
+          payload.variables = payload.variables.map((v) => {
+            if (v.value_selector.join('.') === oldVarSelector.join('.'))
+              v.value_selector = newVarSelector
+            return v
+          })
+        }
+        break
+      }
+      case BlockEnum.Tool: {
+        const payload = data as ToolNodeType
+        if (payload.tool_parameters) {
+          payload.tool_parameters = payload.tool_parameters.map((v) => {
+            if (v.value_selector?.join('.') === oldVarSelector.join('.'))
+              v.value_selector = newVarSelector
+            return v
+          })
+        }
+        break
+      }
+      case BlockEnum.VariableAssigner: {
+        const payload = data as VariableAssignerNodeType
+        if (payload.variables) {
+          payload.variables = payload.variables.map((v) => {
+            if (v.join('.') === oldVarSelector.join('.'))
+              v = newVarSelector
+            return v
+          })
+        }
+        break
+      }
     }
   })
   return newNode
diff --git a/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts b/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts
index efa1274bc2..b9c312e9fa 100644
--- a/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts
@@ -182,6 +182,7 @@ const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
   const [selectedDatasets, setSelectedDatasets] = useState<DataSet[]>([])
   useEffect(() => {
     (async () => {
+      const inputs = inputRef.current
       const datasetIds = inputs.dataset_ids
       if (datasetIds?.length > 0) {
         const { data: dataSetsWithDetail } = await fetchDatasets({ url: '/datasets', params: { page: 1, ids: datasetIds } })
@@ -196,7 +197,7 @@ const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
   }, [])
 
   useEffect(() => {
-    let query_variable_selector: ValueSelector = []
+    let query_variable_selector: ValueSelector = inputs.query_variable_selector
     if (isChatMode && inputs.query_variable_selector.length === 0 && startNodeId)
       query_variable_selector = [startNodeId, 'sys.query']