From 754a1d1197c95dcd1572cf7f335dc92e1fff5688 Mon Sep 17 00:00:00 2001
From: twwu <twwu@dify.ai>
Date: Tue, 27 May 2025 14:17:55 +0800
Subject: [PATCH] refactor: add DatasourceIcon component and update related
 hooks and options

---
 .../data-source-options/datasource-icon.tsx   | 40 +++++++++++++++++++
 .../data-source-options/hooks.tsx             | 31 ++++++++++++++
 .../data-source-options/index.tsx             | 15 +------
 .../data-source-options/option-card.tsx       | 15 +++----
 .../documents/create-from-pipeline/hooks.ts   |  3 +-
 5 files changed, 81 insertions(+), 23 deletions(-)
 create mode 100644 web/app/components/datasets/documents/create-from-pipeline/data-source-options/datasource-icon.tsx
 create mode 100644 web/app/components/datasets/documents/create-from-pipeline/data-source-options/hooks.tsx

diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source-options/datasource-icon.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source-options/datasource-icon.tsx
new file mode 100644
index 0000000000..a2fc0caca7
--- /dev/null
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source-options/datasource-icon.tsx
@@ -0,0 +1,40 @@
+import type { FC } from 'react'
+import { memo } from 'react'
+import cn from '@/utils/classnames'
+
+type DatasourceIconProps = {
+  size?: string
+  className?: string
+  iconUrl: string
+}
+
+const ICON_CONTAINER_CLASSNAME_SIZE_MAP: Record<string, string> = {
+  xs: 'w-4 h-4 rounded-[5px] shadow-xs',
+  sm: 'w-5 h-5 rounded-md shadow-xs',
+  md: 'w-6 h-6 rounded-lg shadow-md',
+}
+
+const DatasourceIcon: FC<DatasourceIconProps> = ({
+  size = 'sm',
+  className,
+  iconUrl,
+}) => {
+  return (
+    <div className={
+      cn(
+        'flex items-center justify-center shadow-none',
+        ICON_CONTAINER_CLASSNAME_SIZE_MAP[size],
+        className,
+      )}
+    >
+      <div
+        className='h-full w-full shrink-0 rounded-md bg-cover bg-center'
+        style={{
+          backgroundImage: `url(${iconUrl})`,
+        }}
+      />
+    </div>
+  )
+}
+
+export default memo(DatasourceIcon)
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
new file mode 100644
index 0000000000..4ceec18733
--- /dev/null
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source-options/hooks.tsx
@@ -0,0 +1,31 @@
+import { useCallback, useMemo, useState } 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 handleUpdateDataSourceList = useCallback((dataSourceList: DataSourceItem[]) => {
+    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.name === data.provider_id)?.icon
+  }, [data, dataSourceList])
+
+  return datasourceIcon
+}
diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source-options/index.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source-options/index.tsx
index 1421cefad8..b56e393fca 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/data-source-options/index.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source-options/index.tsx
@@ -1,11 +1,6 @@
 import { useCallback, useEffect } from 'react'
 import { useDatasourceOptions } from '../hooks'
 import OptionCard from './option-card'
-import { File, Watercrawl } from '@/app/components/base/icons/src/public/knowledge'
-import { Notion } from '@/app/components/base/icons/src/public/common'
-import { Jina } from '@/app/components/base/icons/src/public/llm'
-import { DataSourceType } from '@/models/datasets'
-import { DataSourceProvider } from '@/models/common'
 import type { Datasource } from '@/app/components/rag-pipeline/components/panel/test-run/types'
 import type { Node } from '@/app/components/workflow/types'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
@@ -16,14 +11,6 @@ type DataSourceOptionsProps = {
   onSelect: (option: Datasource) => void
 }
 
-const DATA_SOURCE_ICONS = {
-  [DataSourceType.FILE]: File as React.FC<React.SVGProps<SVGSVGElement>>,
-  [DataSourceType.NOTION]: Notion as React.FC<React.SVGProps<SVGSVGElement>>,
-  [DataSourceProvider.fireCrawl]: '🔥',
-  [DataSourceProvider.jinaReader]: Jina as React.FC<React.SVGProps<SVGSVGElement>>,
-  [DataSourceProvider.waterCrawl]: Watercrawl as React.FC<React.SVGProps<SVGSVGElement>>,
-}
-
 const DataSourceOptions = ({
   pipelineNodes,
   datasourceNodeId,
@@ -51,7 +38,7 @@ const DataSourceOptions = ({
           key={option.value}
           label={option.label}
           selected={datasourceNodeId === option.value}
-          Icon={DATA_SOURCE_ICONS[option.type as keyof typeof DATA_SOURCE_ICONS]}
+          nodeData={option.data}
           onClick={handelSelect.bind(null, option.value)}
         />
       ))}
diff --git a/web/app/components/datasets/documents/create-from-pipeline/data-source-options/option-card.tsx b/web/app/components/datasets/documents/create-from-pipeline/data-source-options/option-card.tsx
index 4e54853282..04c8e08e9c 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/data-source-options/option-card.tsx
+++ b/web/app/components/datasets/documents/create-from-pipeline/data-source-options/option-card.tsx
@@ -1,19 +1,24 @@
 import React from 'react'
 import cn from '@/utils/classnames'
+import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
+import DatasourceIcon from './datasource-icon'
+import { useDatasourceIcon } from './hooks'
 
 type OptionCardProps = {
   label: string
-  Icon: React.FC<React.SVGProps<SVGSVGElement>> | string
   selected: boolean
+  nodeData: DataSourceNodeType
   onClick?: () => void
 }
 
 const OptionCard = ({
   label,
-  Icon,
   selected,
+  nodeData,
   onClick,
 }: OptionCardProps) => {
+  const iconUrl = useDatasourceIcon(nodeData) as string
+
   return (
     <div
       className={cn(
@@ -25,11 +30,7 @@ const OptionCard = ({
       onClick={onClick}
     >
       <div className='flex size-8 items-center justify-center rounded-lg border-[0.5px] border-components-panel-border bg-background-default-dodge p-1.5'>
-        {
-          typeof Icon === 'string'
-            ? <div className='text-[18px] leading-[18px]'>{Icon}</div>
-            : <Icon className='size-5' />
-        }
+        <DatasourceIcon iconUrl={iconUrl} />
       </div>
       <div className={cn('system-sm-medium text-text-secondary', selected && 'text-primary')}>
         {label}
diff --git a/web/app/components/datasets/documents/create-from-pipeline/hooks.ts b/web/app/components/datasets/documents/create-from-pipeline/hooks.ts
index 67134e9058..de046868b1 100644
--- a/web/app/components/datasets/documents/create-from-pipeline/hooks.ts
+++ b/web/app/components/datasets/documents/create-from-pipeline/hooks.ts
@@ -43,12 +43,11 @@ export const useDatasourceOptions = (pipelineNodes: Node<DataSourceNodeType>[])
   const options = useMemo(() => {
     const options: DataSourceOption[] = []
     datasourceNodes.forEach((node) => {
-      const type = node.data.provider_type as DatasourceType
       const label = node.data.title
       options.push({
         label,
         value: node.id,
-        type,
+        data: node.data,
       })
     })
     return options