From 3e0a10b7ed0b86480d3e9842b74f863c9e7e30e5 Mon Sep 17 00:00:00 2001
From: jyong <718720800@qq.com>
Date: Thu, 5 Jun 2025 11:45:53 +0800
Subject: [PATCH] r2

---
 api/services/rag_pipeline/rag_pipeline.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py
index 80b961851a..b6de92f201 100644
--- a/api/services/rag_pipeline/rag_pipeline.py
+++ b/api/services/rag_pipeline/rag_pipeline.py
@@ -430,11 +430,16 @@ class RagPipelineService:
             raise ValueError("Workflow not initialized")
 
         # run draft workflow node
+        datasource_node_data = None
         start_at = time.perf_counter()
-
-        datasource_node_data = published_workflow.graph_dict.get("nodes", {}).get(node_id, {}).get("data", {})
+        datasource_nodes = published_workflow.graph_dict.get("nodes", [])
+        for datasource_node in datasource_nodes:
+            if datasource_node.get("id") == node_id:
+                datasource_node_data = datasource_node.get("data", {})
+                break
         if not datasource_node_data:
             raise ValueError("Datasource node data not found")
+        
         datasource_parameters = datasource_node_data.get("datasource_parameters", {})
         for key, value in datasource_parameters.items():
             if not user_inputs.get(key):
@@ -665,7 +670,12 @@ class RagPipelineService:
             raise ValueError("Workflow not initialized")
 
         # get second step node
-        datasource_node_data = published_workflow.graph_dict.get("nodes", {}).get(node_id, {}).get("data", {})
+        datasource_node_data = None
+        datasource_nodes = published_workflow.graph_dict.get("nodes", [])
+        for datasource_node in datasource_nodes:
+            if datasource_node.get("id") == node_id:
+                datasource_node_data = datasource_node.get("data", {})
+                break
         if not datasource_node_data:
             raise ValueError("Datasource node data not found")
         datasource_parameters = datasource_node_data.get("datasource_parameters", {})
@@ -692,7 +702,12 @@ class RagPipelineService:
             raise ValueError("Workflow not initialized")
 
         # get second step node
-        datasource_node_data = draft_workflow.graph_dict.get("nodes", {}).get(node_id, {}).get("data", {})
+        datasource_node_data = None
+        datasource_nodes = draft_workflow.graph_dict.get("nodes", [])
+        for datasource_node in datasource_nodes:
+            if datasource_node.get("id") == node_id:
+                datasource_node_data = datasource_node.get("data", {})
+                break
         if not datasource_node_data:
             raise ValueError("Datasource node data not found")
         datasource_parameters = datasource_node_data.get("datasource_parameters", {})