From 425e162a910f3f9f3cdc51d4f286a3be58bd0dfd Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Thu, 7 Mar 2024 18:36:45 +0800
Subject: [PATCH] feat: knowledge support single run

---
 .../(commonLayout)/workflow/nodes/page.tsx    |  2 +-
 .../nodes/_base/hooks/use-one-step-run.ts     |  9 +++++
 .../nodes/knowledge-retrieval/panel.tsx       | 31 ++++++++++++++++-
 .../nodes/knowledge-retrieval/use-config.ts   | 33 +++++++++++++++++++
 .../components/workflow/nodes/llm/panel.tsx   |  1 -
 .../workflow/nodes/llm/use-config.ts          | 10 ++----
 6 files changed, 75 insertions(+), 11 deletions(-)

diff --git a/web/app/(commonLayout)/workflow/nodes/page.tsx b/web/app/(commonLayout)/workflow/nodes/page.tsx
index bd8ef8f23a..e742f63c68 100644
--- a/web/app/(commonLayout)/workflow/nodes/page.tsx
+++ b/web/app/(commonLayout)/workflow/nodes/page.tsx
@@ -32,7 +32,7 @@ const allMockData = {
   [BlockEnum.End]: EndNodeMock,
 }
 const nodes = [
-  BlockEnum.LLM/* 3 */, BlockEnum.Start/* 1 */, BlockEnum.DirectAnswer/* 2 */, BlockEnum.KnowledgeRetrieval/* 4 */, BlockEnum.QuestionClassifier/* 5 */,
+  BlockEnum.KnowledgeRetrieval/* 4 */, BlockEnum.Start/* 1 */, BlockEnum.DirectAnswer/* 2 */, BlockEnum.LLM/* 3 */, BlockEnum.QuestionClassifier/* 5 */,
   BlockEnum.IfElse/* 6 */, BlockEnum.Code/* 7 */, BlockEnum.TemplateTransform/* 8 */, BlockEnum.HttpRequest/* 9 */, BlockEnum.Tool/* 10 */,
   BlockEnum.VariableAssigner/* 11 */, BlockEnum.End/* 12 */,
 ].map((item, i) => {
diff --git a/web/app/components/workflow/nodes/_base/hooks/use-one-step-run.ts b/web/app/components/workflow/nodes/_base/hooks/use-one-step-run.ts
index 1de22548aa..6c1ee036b4 100644
--- a/web/app/components/workflow/nodes/_base/hooks/use-one-step-run.ts
+++ b/web/app/components/workflow/nodes/_base/hooks/use-one-step-run.ts
@@ -23,6 +23,13 @@ const useOneStepRun = <T>({ id, data, defaultRunInputData }: Params<T>) => {
   }
 
   const [runningStatus, setRunningStatus] = useState('un started')
+  const handleRun = () => {
+    setRunningStatus('running')
+  }
+
+  const handleStop = () => {
+    setRunningStatus('not started')
+  }
 
   // TODO: store to node
   const [runInputData, setRunInputData] = useState<Record<string, any>>(defaultRunInputData || {})
@@ -50,6 +57,8 @@ const useOneStepRun = <T>({ id, data, defaultRunInputData }: Params<T>) => {
     hideSingleRun,
     toVarInputs,
     runningStatus,
+    handleRun,
+    handleStop,
     setRunningStatus,
     runInputData,
     setRunInputData,
diff --git a/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx b/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx
index 46230209eb..3cfe8a6e8e 100644
--- a/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/panel.tsx
@@ -10,7 +10,8 @@ import type { KnowledgeRetrievalNodeType } from './types'
 import Field from '@/app/components/workflow/nodes/_base/components/field'
 import Split from '@/app/components/workflow/nodes/_base/components/split'
 import OutputVars, { VarItem } from '@/app/components/workflow/nodes/_base/components/output-vars'
-import type { NodePanelProps } from '@/app/components/workflow/types'
+import { InputVarType, type NodePanelProps } from '@/app/components/workflow/types'
+import BeforeRunForm from '@/app/components/workflow/nodes/_base/components/before-run-form'
 
 const i18nPrefix = 'workflow.nodes.knowledgeRetrieval'
 
@@ -28,6 +29,13 @@ const Panel: FC<NodePanelProps<KnowledgeRetrievalNodeType>> = ({
     handleMultipleRetrievalConfigChange,
     selectedDatasets,
     handleOnDatasetsChange,
+    isShowSingleRun,
+    hideSingleRun,
+    runningStatus,
+    handleRun,
+    handleStop,
+    query,
+    setQuery,
   } = useConfig(id, data)
 
   return (
@@ -111,6 +119,27 @@ const Panel: FC<NodePanelProps<KnowledgeRetrievalNodeType>> = ({
 
           </>
         </OutputVars>
+        {isShowSingleRun && (
+          <BeforeRunForm
+            nodeName={inputs.title}
+            onHide={hideSingleRun}
+            forms={[
+              {
+                inputs: [{
+                  label: t(`${i18nPrefix}.queryVariable`)!,
+                  variable: 'query',
+                  type: InputVarType.paragraph,
+                  required: true,
+                }],
+                values: { query },
+                onChange: keyValue => setQuery((keyValue as any).query),
+              },
+            ]}
+            runningStatus={runningStatus}
+            onRun={handleRun}
+            onStop={handleStop}
+          />
+        )}
       </div>
     </div>
   )
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 0914c39ff5..025eedb885 100644
--- a/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/use-config.ts
@@ -6,6 +6,7 @@ import type { RETRIEVE_TYPE } from '@/types/app'
 import type { DataSet } from '@/models/datasets'
 import { fetchDatasets } from '@/service/datasets'
 import useNodeCrud from '@/app/components/workflow/nodes/_base/hooks/use-node-crud'
+import useOneStepRun from '@/app/components/workflow/nodes/_base/hooks/use-one-step-run'
 
 const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
   const { inputs, setInputs } = useNodeCrud<KnowledgeRetrievalNodeType>(id, payload)
@@ -54,6 +55,31 @@ const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
     setSelectedDatasets(newDatasets)
   }, [inputs, setInputs])
 
+  // single run
+  const {
+    isShowSingleRun,
+    hideSingleRun,
+    runningStatus,
+    handleRun,
+    handleStop,
+    runInputData,
+    setRunInputData,
+  } = useOneStepRun<KnowledgeRetrievalNodeType>({
+    id,
+    data: inputs,
+    defaultRunInputData: {
+      query: 'What is the capital of France?',
+    },
+  })
+
+  const query = runInputData.query
+  const setQuery = useCallback((newQuery: string) => {
+    setRunInputData({
+      ...runInputData,
+      query: newQuery,
+    })
+  }, [runInputData, setRunInputData])
+
   return {
     inputs,
     handleQueryVarChange,
@@ -61,6 +87,13 @@ const useConfig = (id: string, payload: KnowledgeRetrievalNodeType) => {
     handleMultipleRetrievalConfigChange,
     selectedDatasets,
     handleOnDatasetsChange,
+    isShowSingleRun,
+    hideSingleRun,
+    runningStatus,
+    handleRun,
+    handleStop,
+    query,
+    setQuery,
   }
 }
 
diff --git a/web/app/components/workflow/nodes/llm/panel.tsx b/web/app/components/workflow/nodes/llm/panel.tsx
index 0bdfc71a2a..9e31a6ab07 100644
--- a/web/app/components/workflow/nodes/llm/panel.tsx
+++ b/web/app/components/workflow/nodes/llm/panel.tsx
@@ -216,7 +216,6 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({
           onStop={handleStop}
         />
       )}
-
     </div>
   )
 }
diff --git a/web/app/components/workflow/nodes/llm/use-config.ts b/web/app/components/workflow/nodes/llm/use-config.ts
index 78dc309e89..0b4da65daa 100644
--- a/web/app/components/workflow/nodes/llm/use-config.ts
+++ b/web/app/components/workflow/nodes/llm/use-config.ts
@@ -95,7 +95,8 @@ const useConfig = (id: string, payload: LLMNodeType) => {
     hideSingleRun,
     toVarInputs,
     runningStatus,
-    setRunningStatus,
+    handleRun,
+    handleStop,
     runInputData,
     setRunInputData,
   } = useOneStepRun<LLMNodeType>({
@@ -145,13 +146,6 @@ const useConfig = (id: string, payload: LLMNodeType) => {
   }, [runInputData, setRunInputData])
 
   const varInputs = toVarInputs(inputs.variables)
-  const handleRun = () => {
-    setRunningStatus('running')
-  }
-
-  const handleStop = () => {
-    setRunningStatus('not started')
-  }
 
   return {
     inputs,