diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source-options/hooks.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source-options/hooks.tsx
index 7f3a736da8..09a02ce7f4 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/data-source-options/hooks.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source-options/hooks.tsx
@@ -1,31 +1,23 @@
-import { useCallback, useMemo, useState } from 'react'
+import { useMemo } from 'react'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
-import type { DataSourceItem } from '@/app/components/workflow/block-selector/types'
 import { basePath } from '@/utils/var'
 import { useDataSourceList } from '@/service/use-pipeline'
-import { useDatasetDetailContextWithSelector } from '@/context/dataset-detail'
-import type { ToolWithProvider } from '@/app/components/workflow/types'
 import { transformDataSourceToTool } from '@/app/components/workflow/block-selector/utils'
 
 export const useDatasourceIcon = (data: DataSourceNodeType) => {
-  const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
-  const [dataSourceList, setDataSourceList] = useState<ToolWithProvider[]>([])
+  const { data: dataSourceListData, isSuccess } = useDataSourceList(true)
 
-  const handleUpdateDataSourceList = useCallback((dataSourceList: DataSourceItem[]) => {
+  const datasourceIcon = useMemo(() => {
+    if (!isSuccess) return
+    const dataSourceList = [...(dataSourceListData || [])]
     dataSourceList.forEach((item) => {
       const icon = item.declaration.identity.icon
       if (typeof icon == 'string' && !icon.includes(basePath))
         item.declaration.identity.icon = `${basePath}${icon}`
     })
     const formattedDataSourceList = dataSourceList.map(item => transformDataSourceToTool(item))
-    setDataSourceList!(formattedDataSourceList)
-  }, [])
-
-  useDataSourceList(!!pipelineId, handleUpdateDataSourceList)
-
-  const datasourceIcon = useMemo(() => {
-    return dataSourceList?.find(toolWithProvider => toolWithProvider.plugin_id === data.plugin_id)?.icon
-  }, [data, dataSourceList])
+    return formattedDataSourceList?.find(toolWithProvider => toolWithProvider.plugin_id === data.plugin_id)?.icon
+  }, [data.plugin_id, dataSourceListData, isSuccess])
 
   return datasourceIcon
 }
diff --git a/web/app/components/workflow/nodes/knowledge-retrieval/node.tsx b/web/app/components/workflow/nodes/knowledge-retrieval/node.tsx
index c24c3e089c..da84baf201 100644
--- a/web/app/components/workflow/nodes/knowledge-retrieval/node.tsx
+++ b/web/app/components/workflow/nodes/knowledge-retrieval/node.tsx
@@ -1,10 +1,10 @@
 import { type FC, useEffect, useState } from 'react'
 import React from 'react'
 import type { KnowledgeRetrievalNodeType } from './types'
-import { Folder } from '@/app/components/base/icons/src/vender/solid/files'
 import type { NodeProps } from '@/app/components/workflow/types'
 import type { DataSet } from '@/models/datasets'
 import { useDatasetsDetailStore } from '../../datasets-detail-store/store'
+import AppIcon from '@/app/components/base/app-icon'
 
 const Node: FC<NodeProps<KnowledgeRetrievalNodeType>> = ({
   data,
@@ -32,11 +32,16 @@ const Node: FC<NodeProps<KnowledgeRetrievalNodeType>> = ({
   return (
     <div className='mb-1 px-3 py-1'>
       <div className='space-y-0.5'>
-        {selectedDatasets.map(({ id, name }) => (
-          <div key={id} className='flex h-[26px] items-center rounded-md bg-workflow-block-parma-bg  px-1 text-xs font-normal text-gray-700'>
-            <div className='mr-1 shrink-0 rounded-md border-[0.5px] border-[#E0EAFF] bg-[#F5F8FF] p-1'>
-              <Folder className='h-3 w-3 text-[#444CE7]' />
-            </div>
+        {selectedDatasets.map(({ id, name, icon_info }) => (
+          <div key={id} className='flex h-[26px] items-center gap-x-1 rounded-md bg-workflow-block-parma-bg px-1'>
+            <AppIcon
+              size='xs'
+              iconType={icon_info.icon_type}
+              icon={icon_info.icon}
+              background={icon_info.icon_type === 'image' ? undefined : icon_info.icon_background}
+              imageUrl={icon_info.icon_type === 'image' ? icon_info.icon_url : undefined}
+              className='shrink-0'
+            />
             <div className='system-xs-regular w-0 grow truncate text-text-secondary'>
               {name}
             </div>