From 66fa68fa18fd19e7b2ed62ae8e27d33b3df5e36a Mon Sep 17 00:00:00 2001
From: jyong <718720800@qq.com>
Date: Wed, 11 Jun 2025 16:36:36 +0800
Subject: [PATCH] r2

---
 api/core/app/apps/common/workflow_response_converter.py    | 7 +++++++
 .../pipeline_template/customized/customized_retrieval.py   | 4 ++++
 .../pipeline_template/database/database_retrieval.py       | 6 ++++--
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/api/core/app/apps/common/workflow_response_converter.py b/api/core/app/apps/common/workflow_response_converter.py
index 6f524a5872..99bb597721 100644
--- a/api/core/app/apps/common/workflow_response_converter.py
+++ b/api/core/app/apps/common/workflow_response_converter.py
@@ -43,10 +43,12 @@ from core.app.entities.task_entities import (
     WorkflowStartStreamResponse,
 )
 from core.file import FILE_MODEL_IDENTITY, File
+from core.plugin.impl.datasource import PluginDatasourceManager
 from core.tools.tool_manager import ToolManager
 from core.workflow.entities.workflow_execution import WorkflowExecution
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecution, WorkflowNodeExecutionStatus
 from core.workflow.nodes import NodeType
+from core.workflow.nodes.datasource.entities import DatasourceNodeData
 from core.workflow.nodes.tool.entities import ToolNodeData
 from models import (
     Account,
@@ -181,6 +183,11 @@ class WorkflowResponseConverter:
                 provider_type=node_data.provider_type,
                 provider_id=node_data.provider_id,
             )
+        elif event.node_type == NodeType.DATASOURCE:
+            node_data = cast(DatasourceNodeData, event.node_data)
+            manager = PluginDatasourceManager()
+            provider_entity = manager.fetch_datasource_provider(self._application_generate_entity.app_config.tenant_id, f"{node_data.plugin_id}/{node_data.provider_name}")
+            response.data.extras["icon"] = provider_entity.declaration.identity.icon
 
         return response
 
diff --git a/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py b/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py
index bb4db81a70..ca94f7f47a 100644
--- a/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py
+++ b/api/services/rag_pipeline/pipeline_template/customized/customized_retrieval.py
@@ -66,6 +66,9 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
         )
         if not pipeline_template:
             return None
+        
+        dsl_data = yaml.safe_load(pipeline_template.yaml_content)
+        graph_data = dsl_data.get("workflow", {}).get("graph", {})
 
         return {
             "id": pipeline_template.id,
@@ -74,5 +77,6 @@ class CustomizedPipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
             "description": pipeline_template.description,
             "chunk_structure": pipeline_template.chunk_structure,
             "export_data": pipeline_template.yaml_content,
+            "graph": graph_data,
             "created_by": pipeline_template.created_user_name,
         }
diff --git a/api/services/rag_pipeline/pipeline_template/database/database_retrieval.py b/api/services/rag_pipeline/pipeline_template/database/database_retrieval.py
index c296111f46..b69a857c3a 100644
--- a/api/services/rag_pipeline/pipeline_template/database/database_retrieval.py
+++ b/api/services/rag_pipeline/pipeline_template/database/database_retrieval.py
@@ -66,13 +66,15 @@ class DatabasePipelineTemplateRetrieval(PipelineTemplateRetrievalBase):
 
         if not pipeline_template:
             return None
-
+        dsl_data = yaml.safe_load(pipeline_template.yaml_content)
+        graph_data = dsl_data.get("workflow", {}).get("graph", {})
         return {
             "id": pipeline_template.id,
             "name": pipeline_template.name,
             "icon_info": pipeline_template.icon,
+            "description": pipeline_template.description,
             "chunk_structure": pipeline_template.chunk_structure,
             "export_data": pipeline_template.yaml_content,
+            "graph": graph_data,
             "created_by": pipeline_template.created_user_name,
-            "description": pipeline_template.description,
         }