From cb2a8142964c9c8b416a259911de88c296231932 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Wed, 13 Mar 2024 15:57:47 +0800
Subject: [PATCH] feat: assign output

---
 .../nodes/_base/components/variable/utils.ts        | 13 +++++++++++++
 .../components/variable/var-reference-picker.tsx    |  4 ++++
 .../workflow/nodes/variable-assigner/types.ts       |  4 ++--
 3 files changed, 19 insertions(+), 2 deletions(-)

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 150e6d6d4f..ec3d314472 100644
--- a/web/app/components/workflow/nodes/_base/components/variable/utils.ts
+++ b/web/app/components/workflow/nodes/_base/components/variable/utils.ts
@@ -2,6 +2,7 @@ import type { CodeNodeType } from '../../../code/types'
 import { BlockEnum, InputVarType, VarType } from '@/app/components/workflow/types'
 import type { StartNodeType } from '@/app/components/workflow/nodes/start/types'
 import type { NodeOutPutVar } from '@/app/components/workflow/types'
+import type { VariableAssignerNodeType } from '@/app/components/workflow/nodes/variable-assigner/types'
 import {
   CHAT_QUESTION_CLASSIFIER_OUTPUT_STRUCT,
   COMPLETION_QUESTION_CLASSIFIER_OUTPUT_STRUCT,
@@ -77,6 +78,18 @@ const formatItem = (item: any, isChatMode: boolean): NodeOutPutVar => {
       res.vars = HTTP_REQUEST_OUTPUT_STRUCT
       break
     }
+
+    case BlockEnum.VariableAssigner: {
+      const {
+        output_type,
+      } = data as VariableAssignerNodeType
+      res.vars = [
+        {
+          variable: 'output',
+          type: output_type,
+        },
+      ]
+    }
   }
 
   return res
diff --git a/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx b/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx
index be4cefa8c5..fb240dae9d 100644
--- a/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx
+++ b/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx
@@ -2,6 +2,7 @@
 import type { FC } from 'react'
 import React, { useState } from 'react'
 import cn from 'classnames'
+import { isArray } from 'lodash-es'
 import VarReferencePopup from './var-reference-popup'
 import { toNodeOutputVars } from './utils'
 import type { ValueSelector } from '@/app/components/workflow/types'
@@ -29,6 +30,9 @@ type Props = {
 }
 
 export const getNodeInfoById = (nodes: any, id: string) => {
+  if (!isArray(nodes))
+    return
+
   return nodes.find((node: any) => node.id === id)
 }
 
diff --git a/web/app/components/workflow/nodes/variable-assigner/types.ts b/web/app/components/workflow/nodes/variable-assigner/types.ts
index db0c8b8e31..8a923f0290 100644
--- a/web/app/components/workflow/nodes/variable-assigner/types.ts
+++ b/web/app/components/workflow/nodes/variable-assigner/types.ts
@@ -1,6 +1,6 @@
-import type { CommonNodeType, ValueSelector } from '@/app/components/workflow/types'
+import type { CommonNodeType, ValueSelector, VarType } from '@/app/components/workflow/types'
 
 export type VariableAssignerNodeType = CommonNodeType & {
-  output_type: string
+  output_type: VarType
   variables: ValueSelector[]
 }