From ae6a558662bf147fdd13986bfbb777fe83209a87 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Thu, 14 Mar 2024 20:29:47 +0800
Subject: [PATCH] feat: add prev and next nodes

---
 web/app/components/workflow/constants.ts            |  3 +++
 web/app/components/workflow/nodes/answer/default.ts |  6 ++++--
 web/app/components/workflow/nodes/code/default.ts   | 13 ++++++++-----
 web/app/components/workflow/nodes/end/default.ts    |  6 ++++--
 web/app/components/workflow/nodes/http/default.ts   | 11 +++++++----
 .../components/workflow/nodes/if-else/default.ts    | 11 +++++++----
 .../workflow/nodes/knowledge-retrieval/default.ts   | 12 ++++++++----
 web/app/components/workflow/nodes/llm/default.ts    | 11 +++++++----
 .../workflow/nodes/question-classifier/default.ts   | 11 +++++++----
 web/app/components/workflow/nodes/start/default.ts  | 11 +++++++----
 .../workflow/nodes/template-transform/default.ts    | 11 +++++++----
 web/app/components/workflow/nodes/tool/default.ts   | 11 +++++++----
 .../workflow/nodes/variable-assigner/default.ts     | 11 +++++++----
 web/app/components/workflow/types.ts                |  4 ++--
 14 files changed, 85 insertions(+), 47 deletions(-)

diff --git a/web/app/components/workflow/constants.ts b/web/app/components/workflow/constants.ts
index 50cf03f50b..323f800c26 100644
--- a/web/app/components/workflow/constants.ts
+++ b/web/app/components/workflow/constants.ts
@@ -64,6 +64,9 @@ export const NODES_EXTRA_DATA = {
   },
 }
 
+export const ALL_CHAT_AVAILABLE_BLOCKS = Object.keys(NODES_EXTRA_DATA).filter(key => key !== BlockEnum.End) as BlockEnum[]
+export const ALL_COMPLETION_AVAILABLE_BLOCKS = Object.keys(NODES_EXTRA_DATA).filter(key => key !== BlockEnum.Answer) as BlockEnum[]
+
 export const NODES_INITIAL_DATA = {
   [BlockEnum.Start]: {
     type: BlockEnum.Start,
diff --git a/web/app/components/workflow/nodes/answer/default.ts b/web/app/components/workflow/nodes/answer/default.ts
index d36cf78bbf..12725f1b14 100644
--- a/web/app/components/workflow/nodes/answer/default.ts
+++ b/web/app/components/workflow/nodes/answer/default.ts
@@ -1,13 +1,15 @@
 import type { NodeDefault } from '../../types'
 import type { AnswerNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<AnswerNodeType> = {
   defaultValue: {
     variables: [],
     answer: '',
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
   getAvailableNextNodes() {
     return []
diff --git a/web/app/components/workflow/nodes/code/default.ts b/web/app/components/workflow/nodes/code/default.ts
index 7a255188ba..a1b77e3ed1 100644
--- a/web/app/components/workflow/nodes/code/default.ts
+++ b/web/app/components/workflow/nodes/code/default.ts
@@ -1,18 +1,21 @@
 import type { NodeDefault } from '../../types'
 import { CodeLanguage, type CodeNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<CodeNodeType> = {
   defaultValue: {
     code: '',
     code_language: CodeLanguage.python3,
     variables: [],
-    outputs: [],
+    outputs: {},
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/nodes/end/default.ts b/web/app/components/workflow/nodes/end/default.ts
index 5f3db6960f..8193e146a5 100644
--- a/web/app/components/workflow/nodes/end/default.ts
+++ b/web/app/components/workflow/nodes/end/default.ts
@@ -1,12 +1,14 @@
 import type { NodeDefault } from '../../types'
 import { type EndNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<EndNodeType> = {
   defaultValue: {
     outputs: [],
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
   getAvailableNextNodes() {
     return []
diff --git a/web/app/components/workflow/nodes/http/default.ts b/web/app/components/workflow/nodes/http/default.ts
index aade82383a..090ff6b035 100644
--- a/web/app/components/workflow/nodes/http/default.ts
+++ b/web/app/components/workflow/nodes/http/default.ts
@@ -1,5 +1,6 @@
 import type { NodeDefault } from '../../types'
 import { AuthorizationType, BodyType, type HttpNodeType, Method } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<HttpNodeType> = {
   defaultValue: {
@@ -17,11 +18,13 @@ const nodeDefault: NodeDefault<HttpNodeType> = {
       data: '',
     },
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/nodes/if-else/default.ts b/web/app/components/workflow/nodes/if-else/default.ts
index 4368922426..763977ffa1 100644
--- a/web/app/components/workflow/nodes/if-else/default.ts
+++ b/web/app/components/workflow/nodes/if-else/default.ts
@@ -1,5 +1,6 @@
 import type { NodeDefault } from '../../types'
 import { type IfElseNodeType, LogicalOperator } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<IfElseNodeType> = {
   defaultValue: {
@@ -16,11 +17,13 @@ const nodeDefault: NodeDefault<IfElseNodeType> = {
     logical_operator: LogicalOperator.and,
     conditions: [],
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/nodes/knowledge-retrieval/default.ts b/web/app/components/workflow/nodes/knowledge-retrieval/default.ts
index ba396e18b4..1074975972 100644
--- a/web/app/components/workflow/nodes/knowledge-retrieval/default.ts
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/default.ts
@@ -1,5 +1,7 @@
 import type { NodeDefault } from '../../types'
 import type { KnowledgeRetrievalNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
+
 import { RETRIEVE_TYPE } from '@/types/app'
 
 const nodeDefault: NodeDefault<KnowledgeRetrievalNodeType> = {
@@ -8,11 +10,13 @@ const nodeDefault: NodeDefault<KnowledgeRetrievalNodeType> = {
     dataset_ids: [],
     retrieval_mode: RETRIEVE_TYPE.oneWay,
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/nodes/llm/default.ts b/web/app/components/workflow/nodes/llm/default.ts
index 758176658d..4c4a48cea6 100644
--- a/web/app/components/workflow/nodes/llm/default.ts
+++ b/web/app/components/workflow/nodes/llm/default.ts
@@ -1,5 +1,6 @@
 import { type NodeDefault, PromptRole } from '../../types'
 import type { LLMNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<LLMNodeType> = {
   defaultValue: {
@@ -24,11 +25,13 @@ const nodeDefault: NodeDefault<LLMNodeType> = {
       enabled: false,
     },
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/nodes/question-classifier/default.ts b/web/app/components/workflow/nodes/question-classifier/default.ts
index 599ef2c594..754030eab8 100644
--- a/web/app/components/workflow/nodes/question-classifier/default.ts
+++ b/web/app/components/workflow/nodes/question-classifier/default.ts
@@ -1,5 +1,6 @@
 import type { NodeDefault } from '../../types'
 import type { QuestionClassifierNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<QuestionClassifierNodeType> = {
   defaultValue: {
@@ -14,11 +15,13 @@ const nodeDefault: NodeDefault<QuestionClassifierNodeType> = {
     },
     classes: [],
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/nodes/start/default.ts b/web/app/components/workflow/nodes/start/default.ts
index 3f196b5067..875406e680 100644
--- a/web/app/components/workflow/nodes/start/default.ts
+++ b/web/app/components/workflow/nodes/start/default.ts
@@ -1,15 +1,18 @@
 import type { NodeDefault } from '../../types'
 import type { StartNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<StartNodeType> = {
   defaultValue: {
     variables: [],
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/nodes/template-transform/default.ts b/web/app/components/workflow/nodes/template-transform/default.ts
index a6e4a8848c..e2dd3f357b 100644
--- a/web/app/components/workflow/nodes/template-transform/default.ts
+++ b/web/app/components/workflow/nodes/template-transform/default.ts
@@ -1,15 +1,18 @@
 import type { NodeDefault } from '../../types'
 import type { TemplateTransformNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<TemplateTransformNodeType> = {
   defaultValue: {
     variables: [],
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/nodes/tool/default.ts b/web/app/components/workflow/nodes/tool/default.ts
index e1c81de97d..67fb1f579f 100644
--- a/web/app/components/workflow/nodes/tool/default.ts
+++ b/web/app/components/workflow/nodes/tool/default.ts
@@ -1,16 +1,19 @@
 import type { NodeDefault } from '../../types'
 import type { ToolNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<ToolNodeType> = {
   defaultValue: {
     tool_parameters: [],
     tool_configurations: {},
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/nodes/variable-assigner/default.ts b/web/app/components/workflow/nodes/variable-assigner/default.ts
index ce7785fc3f..fbbde4f3f2 100644
--- a/web/app/components/workflow/nodes/variable-assigner/default.ts
+++ b/web/app/components/workflow/nodes/variable-assigner/default.ts
@@ -1,16 +1,19 @@
 import type { NodeDefault } from '../../types'
 import type { VariableAssignerNodeType } from './types'
+import { ALL_CHAT_AVAILABLE_BLOCKS, ALL_COMPLETION_AVAILABLE_BLOCKS } from '@/app/components/workflow/constants'
 
 const nodeDefault: NodeDefault<VariableAssignerNodeType> = {
   defaultValue: {
     output_type: 'string',
     variables: [],
   },
-  getAvailablePrevNodes() {
-    return []
+  getAvailablePrevNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
-  getAvailableNextNodes() {
-    return []
+  getAvailableNextNodes(isChatMode: boolean) {
+    const nodes = isChatMode ? ALL_CHAT_AVAILABLE_BLOCKS : ALL_COMPLETION_AVAILABLE_BLOCKS
+    return nodes
   },
 }
 
diff --git a/web/app/components/workflow/types.ts b/web/app/components/workflow/types.ts
index 1d112ee3da..a8038171be 100644
--- a/web/app/components/workflow/types.ts
+++ b/web/app/components/workflow/types.ts
@@ -157,8 +157,8 @@ export type Block = {
 
 export type NodeDefault<T> = {
   defaultValue: Partial<T>
-  getAvailablePrevNodes: () => BlockEnum[]
-  getAvailableNextNodes: () => BlockEnum[]
+  getAvailablePrevNodes: (isChatMode: boolean) => BlockEnum[]
+  getAvailableNextNodes: (isChatMode: boolean) => BlockEnum[]
 }
 
 export type OnSelectBlock = (type: BlockEnum, toolDefaultValue?: ToolDefaultValue) => void