From df9e2e478f36cd00b85f8638de0d5a1ba9f74389 Mon Sep 17 00:00:00 2001
From: StyleZhang <jasonapring2015@outlook.com>
Date: Mon, 1 Apr 2024 16:38:13 +0800
Subject: [PATCH] workflow template

---
 .../workflow/hooks/use-workflow-template.ts   | 73 +++++++++++++++++++
 .../components/workflow/hooks/use-workflow.ts | 23 ++----
 web/app/components/workflow/utils.ts          |  4 +-
 3 files changed, 83 insertions(+), 17 deletions(-)
 create mode 100644 web/app/components/workflow/hooks/use-workflow-template.ts

diff --git a/web/app/components/workflow/hooks/use-workflow-template.ts b/web/app/components/workflow/hooks/use-workflow-template.ts
new file mode 100644
index 0000000000..3af3f733f1
--- /dev/null
+++ b/web/app/components/workflow/hooks/use-workflow-template.ts
@@ -0,0 +1,73 @@
+import { generateNewNode } from '../utils'
+import {
+  NODE_WIDTH_X_OFFSET,
+  START_INITIAL_POSITION,
+} from '../constants'
+import { useIsChatMode } from './use-workflow'
+import { useNodesInitialData } from './use-nodes-data'
+
+export const useWorkflowTemplate = () => {
+  const isChatMode = useIsChatMode()
+  const nodesInitialData = useNodesInitialData()
+
+  const startNode = generateNewNode({
+    data: nodesInitialData.start,
+    position: START_INITIAL_POSITION,
+  })
+
+  if (isChatMode) {
+    const llmNode = generateNewNode({
+      id: 'llm',
+      data: {
+        ...nodesInitialData.llm,
+        memory: {
+          window: { enabled: false, size: 10 },
+        },
+        selected: true,
+      },
+      position: {
+        x: START_INITIAL_POSITION.x + NODE_WIDTH_X_OFFSET,
+        y: START_INITIAL_POSITION.y,
+      },
+    } as any)
+
+    const answerNode = generateNewNode({
+      id: 'answer',
+      data: {
+        ...nodesInitialData.answer,
+        answer: `{{#${llmNode.id}.text#}}`,
+      },
+      position: {
+        x: START_INITIAL_POSITION.x + NODE_WIDTH_X_OFFSET * 2,
+        y: START_INITIAL_POSITION.y,
+      },
+    } as any)
+
+    const startToLlmEdge = {
+      id: `${startNode.id}-${llmNode.id}`,
+      source: startNode.id,
+      sourceHandle: 'source',
+      target: llmNode.id,
+      targetHandle: 'target',
+    }
+
+    const llmToAnswerEdge = {
+      id: `${llmNode.id}-${answerNode.id}`,
+      source: llmNode.id,
+      sourceHandle: 'source',
+      target: answerNode.id,
+      targetHandle: 'target',
+    }
+
+    return {
+      nodes: [startNode, llmNode, answerNode],
+      edges: [startToLlmEdge, llmToAnswerEdge],
+    }
+  }
+  else {
+    return {
+      nodes: [startNode],
+      edges: [],
+    }
+  }
+}
diff --git a/web/app/components/workflow/hooks/use-workflow.ts b/web/app/components/workflow/hooks/use-workflow.ts
index 9d8d3d4190..0093f3fd6b 100644
--- a/web/app/components/workflow/hooks/use-workflow.ts
+++ b/web/app/components/workflow/hooks/use-workflow.ts
@@ -19,7 +19,6 @@ import type {
   Viewport,
 } from 'reactflow'
 import {
-  generateNewNode,
   getLayoutByDagre,
   initialEdges,
   initialNodes,
@@ -39,14 +38,11 @@ import {
 } from '../store'
 import {
   AUTO_LAYOUT_OFFSET,
-  START_INITIAL_POSITION,
   SUPPORT_OUTPUT_VARS_NODE,
 } from '../constants'
 import { findUsedVarNodes, getNodeOutputVars, updateNodeVars } from '../nodes/_base/components/variable/utils'
-import {
-  useNodesExtraData,
-  useNodesInitialData,
-} from './use-nodes-data'
+import { useNodesExtraData } from './use-nodes-data'
+import { useWorkflowTemplate } from './use-workflow-template'
 import { useNodesSyncDraft } from './use-nodes-sync-draft'
 import { useStore as useAppStore } from '@/app/components/app/store'
 import {
@@ -393,7 +389,10 @@ export const useFetchToolsData = () => {
 
 export const useWorkflowInit = () => {
   const workflowStore = useWorkflowStore()
-  const nodesInitialData = useNodesInitialData()
+  const {
+    nodes: nodesTemplate,
+    edges: edgesTemplate,
+  } = useWorkflowTemplate()
   const { handleFetchAllTools } = useFetchToolsData()
   const appDetail = useAppStore(state => state.appDetail)!
   const { data, isLoading, error, mutate } = useSWR(`/apps/${appDetail.id}/workflows/draft`, fetchWorkflowDraft)
@@ -435,14 +434,8 @@ export const useWorkflowInit = () => {
           url: `/apps/${appDetail.id}/workflows/draft`,
           params: {
             graph: {
-              nodes: [generateNewNode({
-                data: {
-                  ...nodesInitialData.start,
-                  selected: true,
-                },
-                position: START_INITIAL_POSITION,
-              })],
-              edges: [],
+              nodes: nodesTemplate,
+              edges: edgesTemplate,
             },
             features: {},
           },
diff --git a/web/app/components/workflow/utils.ts b/web/app/components/workflow/utils.ts
index f3367f3eb4..7176f8b7ca 100644
--- a/web/app/components/workflow/utils.ts
+++ b/web/app/components/workflow/utils.ts
@@ -152,9 +152,9 @@ export const getNodesConnectedSourceOrTargetHandleIdsMap = (changes: ConnectedSo
   return nodesConnectedSourceOrTargetHandleIdsMap
 }
 
-export const generateNewNode = ({ data, position }: Pick<Node, 'data' | 'position'>) => {
+export const generateNewNode = ({ data, position, id }: Pick<Node, 'data' | 'position'> & { id?: string }) => {
   return {
-    id: `${Date.now()}`,
+    id: id || `${Date.now()}`,
     type: 'custom',
     data,
     position,