From b320ebe2ba1f90f5a2db097459aa1b22a617eccb Mon Sep 17 00:00:00 2001
From: zxhlyh <jasonapring2015@outlook.com>
Date: Tue, 27 May 2025 18:44:04 +0800
Subject: [PATCH] datasource variables

---
 .../nodes/_base/components/variable/utils.ts   |  3 ++-
 .../workflow/nodes/data-source/default.ts      | 18 +++++++++++++++++-
 .../workflow/nodes/data-source/utils.ts        | 10 ++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 web/app/components/workflow/nodes/data-source/utils.ts

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 8c2c7811a4..65f32ba230 100644
--- a/web/app/components/workflow/nodes/_base/components/variable/utils.ts
+++ b/web/app/components/workflow/nodes/_base/components/variable/utils.ts
@@ -34,6 +34,7 @@ import {
   TOOL_OUTPUT_STRUCT,
 } from '@/app/components/workflow/constants'
 import DataSourceNodeDefault from '@/app/components/workflow/nodes/data-source/default'
+import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
 import type { PromptItem } from '@/models/debug'
 import { VAR_REGEX } from '@/config'
 import type { AgentNodeType } from '../../../agent/types'
@@ -459,7 +460,7 @@ const formatItem = (
     }
 
     case BlockEnum.DataSource: {
-      res.vars = DataSourceNodeDefault.getOutputVars?.(data as any) || []
+      res.vars = DataSourceNodeDefault.getOutputVars?.(data as DataSourceNodeType) || []
       break
     }
 
diff --git a/web/app/components/workflow/nodes/data-source/default.ts b/web/app/components/workflow/nodes/data-source/default.ts
index e27ccc155b..bc3f1a543b 100644
--- a/web/app/components/workflow/nodes/data-source/default.ts
+++ b/web/app/components/workflow/nodes/data-source/default.ts
@@ -3,6 +3,7 @@ import type { DataSourceNodeType } from './types'
 import { genNodeMetaData } from '@/app/components/workflow/utils'
 import { BlockEnum } from '@/app/components/workflow/types'
 import { OUTPUT_VARIABLES_MAP } from './constants'
+import { inputVarTypeToVarType } from './utils'
 
 const metaData = genNodeMetaData({
   sort: -1,
@@ -22,8 +23,12 @@ const nodeDefault: NodeDefault<DataSourceNodeType> = {
     }
   },
   getOutputVars(payload) {
-    const { provider_type } = payload
+    const {
+      provider_type,
+      variables,
+    } = payload
     const isLocalFile = provider_type === 'local_file'
+    const hasUserInputFields = !!variables?.length
     return [
       {
         variable: OUTPUT_VARIABLES_MAP.datasource_type.name,
@@ -39,6 +44,17 @@ const nodeDefault: NodeDefault<DataSourceNodeType> = {
           ]
           : []
       ),
+      ...(
+        hasUserInputFields
+          ? variables.map((field) => {
+            return {
+              variable: field.variable,
+              type: inputVarTypeToVarType(field.type),
+              isUserInputField: true,
+            }
+          })
+          : []
+      ),
     ]
   },
 }
diff --git a/web/app/components/workflow/nodes/data-source/utils.ts b/web/app/components/workflow/nodes/data-source/utils.ts
new file mode 100644
index 0000000000..4240d6f464
--- /dev/null
+++ b/web/app/components/workflow/nodes/data-source/utils.ts
@@ -0,0 +1,10 @@
+import { PipelineInputVarType } from '@/models/pipeline'
+import { VarType } from '@/app/components/workflow/types'
+
+export const inputVarTypeToVarType = (type: PipelineInputVarType): VarType => {
+  return ({
+    [PipelineInputVarType.number]: VarType.number,
+    [PipelineInputVarType.singleFile]: VarType.file,
+    [PipelineInputVarType.multiFiles]: VarType.arrayFile,
+  } as any)[type] || VarType.string
+}