From aff5ab933bfa543086dc383fa513113407f6f010 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Fri, 15 Mar 2024 19:22:07 +0800
Subject: [PATCH] feat: knowledge node var init value and limit

---
 .../workflow/hooks/use-nodes-data.ts          |  2 +-
 .../nodes/knowledge-retrieval/panel.tsx       |  3 ++
 .../nodes/knowledge-retrieval/use-config.ts   | 28 +++++++++++++++++--
 web/app/components/workflow/types.ts          |  1 +
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/web/app/components/workflow/hooks/use-nodes-data.ts b/web/app/components/workflow/hooks/use-nodes-data.ts
index 9eaa40ce3c..417cb3399c 100644
--- a/web/app/components/workflow/hooks/use-nodes-data.ts
+++ b/web/app/components/workflow/hooks/use-nodes-data.ts
@@ -16,7 +16,7 @@ export const useNodesInitialData = () => {
   return useMemo(() => produce(NODES_INITIAL_DATA, (draft) => {
     Object.keys(draft).forEach((key) => {
       draft[key as BlockEnum].title = t(`workflow.blocks.${key}`)
-
+      draft[key as BlockEnum]._isReady = true
       if (nodesDefaultConfigs[key as BlockEnum]) {
         draft[key as BlockEnum] = {
           ...draft[key as BlockEnum],
diff --git a/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx b/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx
index 46bd2a0980..d8acf4fa87 100644
--- a/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx
@@ -26,6 +26,7 @@ const Panel: FC<NodePanelProps<KnowledgeRetrievalNodeType>> = ({
   const {
     inputs,
     handleQueryVarChange,
+    filterVar,
     handleRetrievalModeChange,
     handleMultipleRetrievalConfigChange,
     selectedDatasets,
@@ -52,6 +53,8 @@ const Panel: FC<NodePanelProps<KnowledgeRetrievalNodeType>> = ({
             isShowNodeName
             value={inputs.query_variable_selector}
             onChange={handleQueryVarChange}
+            filterVar={filterVar}
+            width={370}
           />
         </Field>
 
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 7d6e622a4b..e37e1ce9ca 100644
--- a/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts
@@ -1,6 +1,8 @@
 import { useCallback, useEffect, useState } from 'react'
 import produce from 'immer'
-import type { ValueSelector } from '../../types'
+import type { ValueSelector, Var } from '../../types'
+import { BlockEnum, VarType } from '../../types'
+import { useIsChatMode, useWorkflow } from '../../hooks'
 import type { KnowledgeRetrievalNodeType, MultipleRetrievalConfig } from './types'
 import type { RETRIEVE_TYPE } from '@/types/app'
 import type { DataSet } from '@/models/datasets'
@@ -9,10 +11,14 @@ import useNodeCrud from '@/app/components/workflow/nodes/_base/hooks/use-node-cr
 import useOneStepRun from '@/app/components/workflow/nodes/_base/hooks/use-one-step-run'
 
 const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
+  const isChatMode = useIsChatMode()
+  const { getBeforeNodesInSameBranch } = useWorkflow()
+  const startNode = getBeforeNodesInSameBranch(id).find(node => node.data.type === BlockEnum.Start)
+  const startNodeId = startNode?.id
   const { inputs, setInputs } = useNodeCrud<KnowledgeRetrievalNodeType>(id, payload)
-  const handleQueryVarChange = useCallback((newVar: ValueSelector) => {
+  const handleQueryVarChange = useCallback((newVar: ValueSelector | string) => {
     const newInputs = produce(inputs, (draft) => {
-      draft.query_variable_selector = newVar
+      draft.query_variable_selector = newVar as ValueSelector
     })
     setInputs(newInputs)
   }, [inputs, setInputs])
@@ -45,8 +51,19 @@ const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
       })
       setInputs(newInputs)
     })()
+  // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [])
 
+  useEffect(() => {
+    if (inputs._isReady) {
+      if (isChatMode && inputs.query_variable_selector.length === 0 && startNodeId) {
+        handleQueryVarChange(
+          [startNodeId, 'sys.query'],
+        )
+      }
+    }
+  }, [inputs._isReady])
+
   const handleOnDatasetsChange = useCallback((newDatasets: DataSet[]) => {
     const newInputs = produce(inputs, (draft) => {
       draft.dataset_ids = newDatasets.map(d => d.id)
@@ -55,6 +72,10 @@ const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
     setSelectedDatasets(newDatasets)
   }, [inputs, setInputs])
 
+  const filterVar = useCallback((varPayload: Var) => {
+    return varPayload.type === VarType.string
+  }, [])
+
   // single run
   const {
     isShowSingleRun,
@@ -84,6 +105,7 @@ const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
   return {
     inputs,
     handleQueryVarChange,
+    filterVar,
     handleRetrievalModeChange,
     handleMultipleRetrievalConfigChange,
     selectedDatasets,
diff --git a/web/app/components/workflow/types.ts b/web/app/components/workflow/types.ts
index 91d72a1486..bb8efb3bdc 100644
--- a/web/app/components/workflow/types.ts
+++ b/web/app/components/workflow/types.ts
@@ -33,6 +33,7 @@ export type CommonNodeType<T = {}> = {
   _isSingleRun?: boolean
   _runningStatus?: NodeRunningStatus
   _singleRunningStatus?: NodeRunningStatus
+  _isReady?: boolean
   selected?: boolean
   title: string
   desc: string