refactor: Replace useDataSourceStore with useDataSourceStoreWithSelector for improved state selection across components

web/app/components/datasets/documents/create-from-pipeline/data-source/local-file/index.tsx

@@ -15,7 +15,7 @@ import { IS_CE_EDITION } from '@/config'
 import { Theme } from '@/types/app'
 import useTheme from '@/hooks/use-theme'
 import { useFileUploadConfig } from '@/service/use-common'
-import { useDataSourceStore } from '../store'
+import { useDataSourceStoreWithSelector } from '../store'
 import produce from 'immer'
 
 const FILES_NUMBER_LIMIT = 20
@@ -32,10 +32,10 @@ const LocalFile = ({
   const { t } = useTranslation()
   const { notify } = useContext(ToastContext)
   const { locale } = useContext(I18n)
-  const fileList = useDataSourceStore(state => state.localFileList)
+  const fileList = useDataSourceStoreWithSelector(state => state.localFileList)
-  const setFileList = useDataSourceStore(state => state.setLocalFileList)
+  const setFileList = useDataSourceStoreWithSelector(state => state.setLocalFileList)
-  const setCurrentFile = useDataSourceStore(state => state.setCurrentLocalFile)
+  const setCurrentFile = useDataSourceStoreWithSelector(state => state.setCurrentLocalFile)
-  const previewFileRef = useDataSourceStore(state => state.previewLocalFileRef)
+  const previewFileRef = useDataSourceStoreWithSelector(state => state.previewLocalFileRef)
   const [dragging, setDragging] = useState(false)
 
   const dropRef = useRef<HTMLDivElement>(null)

web/app/components/datasets/documents/create-from-pipeline/data-source/online-documents/index.tsx

@@ -10,7 +10,7 @@ import { ssePost } from '@/service/base'
 import Toast from '@/app/components/base/toast'
 import type { DataSourceNodeCompletedResponse } from '@/types/pipeline'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
-import { useDataSourceStore } from '../store'
+import { useDataSourceStoreWithSelector } from '../store'
 
 type OnlineDocumentsProps = {
   isInPipeline?: boolean
@@ -24,16 +24,16 @@ const OnlineDocuments = ({
   nodeData,
 }: OnlineDocumentsProps) => {
   const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
-  const documentsData = useDataSourceStore(state => state.documentsData)
+  const documentsData = useDataSourceStoreWithSelector(state => state.documentsData)
-  const setDocumentsData = useDataSourceStore(state => state.setDocumentsData)
+  const setDocumentsData = useDataSourceStoreWithSelector(state => state.setDocumentsData)
-  const searchValue = useDataSourceStore(state => state.searchValue)
+  const searchValue = useDataSourceStoreWithSelector(state => state.searchValue)
-  const setSearchValue = useDataSourceStore(state => state.setSearchValue)
+  const setSearchValue = useDataSourceStoreWithSelector(state => state.setSearchValue)
-  const currentWorkspaceId = useDataSourceStore(state => state.currentWorkspaceId)
+  const currentWorkspaceId = useDataSourceStoreWithSelector(state => state.currentWorkspaceId)
-  const setCurrentWorkspaceId = useDataSourceStore(state => state.setCurrentWorkspaceId)
+  const setCurrentWorkspaceId = useDataSourceStoreWithSelector(state => state.setCurrentWorkspaceId)
-  const setOnlineDocuments = useDataSourceStore(state => state.setOnlineDocuments)
+  const setOnlineDocuments = useDataSourceStoreWithSelector(state => state.setOnlineDocuments)
-  const setCurrentDocument = useDataSourceStore(state => state.setCurrentDocument)
+  const setCurrentDocument = useDataSourceStoreWithSelector(state => state.setCurrentDocument)
-  const selectedPagesId = useDataSourceStore(state => state.selectedPagesId)
+  const selectedPagesId = useDataSourceStoreWithSelector(state => state.selectedPagesId)
-  const setSelectedPagesId = useDataSourceStore(state => state.setSelectedPagesId)
+  const setSelectedPagesId = useDataSourceStoreWithSelector(state => state.setSelectedPagesId)
 
   const PagesMapAndSelectedPagesId: DataSourceNotionPageMap = useMemo(() => {
     const pagesMap = (documentsData || []).reduce((prev: DataSourceNotionPageMap, next: DataSourceNotionWorkspace) => {

web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/index.tsx

@@ -9,6 +9,7 @@ type FileListProps = {
   selectedFileList: string[]
   prefix: string[]
   keywords: string
+  isInPipeline: boolean
   resetKeywords: () => void
   updateKeywords: (keywords: string) => void
   searchResultsLength: number
@@ -26,6 +27,7 @@ const FileList = ({
   searchResultsLength,
   handleSelectFile,
   handleOpenFolder,
+  isInPipeline,
 }: FileListProps) => {
   const [inputValue, setInputValue] = useState(keywords)
 
@@ -64,6 +66,7 @@ const FileList = ({
         handleResetKeywords={handleResetKeywords}
         handleOpenFolder={handleOpenFolder}
         handleSelectFile={handleSelectFile}
+        isInPipeline={isInPipeline}
       />
     </div>
   )

web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/file-list/list/index.tsx

@@ -7,6 +7,7 @@ type FileListProps = {
   fileList: OnlineDriveFile[]
   selectedFileList: string[]
   keywords: string
+  isInPipeline: boolean
   handleResetKeywords: () => void
   handleSelectFile: (file: OnlineDriveFile) => void
   handleOpenFolder: (file: OnlineDriveFile) => void
@@ -19,6 +20,7 @@ const List = ({
   handleResetKeywords,
   handleSelectFile,
   handleOpenFolder,
+  isInPipeline,
 }: FileListProps) => {
   const isEmptyFolder = fileList.length === 0 && keywords.length === 0
   const isSearchResultEmpty = fileList.length === 0 && keywords.length > 0
@@ -47,6 +49,7 @@ const List = ({
                   isSelected={isSelected}
                   onSelect={handleSelectFile}
                   onOpen={handleOpenFolder}
+                  isMultipleChoice={!isInPipeline}
                 />
               )
             })

web/app/components/datasets/documents/create-from-pipeline/data-source/online-drive/index.tsx

@@ -8,7 +8,7 @@ import { useDatasetDetailContextWithSelector } from '@/context/dataset-detail'
 import { ssePost } from '@/service/base'
 import type { DataSourceNodeCompletedResponse, DataSourceNodeErrorResponse } from '@/types/pipeline'
 import Toast from '@/app/components/base/toast'
-import { useDataSourceStore } from '../store'
+import { useDataSourceStoreWithSelector } from '../store'
 import { convertOnlineDriveDataToFileList } from './utils'
 import produce from 'immer'
 
@@ -24,18 +24,18 @@ const OnlineDrive = ({
   isInPipeline = false,
 }: OnlineDriveProps) => {
   const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
-  const prefix = useDataSourceStore(state => state.prefix)
+  const prefix = useDataSourceStoreWithSelector(state => state.prefix)
-  const setPrefix = useDataSourceStore(state => state.setPrefix)
+  const setPrefix = useDataSourceStoreWithSelector(state => state.setPrefix)
-  const keywords = useDataSourceStore(state => state.keywords)
+  const keywords = useDataSourceStoreWithSelector(state => state.keywords)
-  const setKeywords = useDataSourceStore(state => state.setKeywords)
+  const setKeywords = useDataSourceStoreWithSelector(state => state.setKeywords)
-  const bucket = useDataSourceStore(state => state.bucket)
+  const bucket = useDataSourceStoreWithSelector(state => state.bucket)
-  const setBucket = useDataSourceStore(state => state.setBucket)
+  const setBucket = useDataSourceStoreWithSelector(state => state.setBucket)
-  const startAfter = useDataSourceStore(state => state.startAfter)
+  const startAfter = useDataSourceStoreWithSelector(state => state.startAfter)
-  const setStartAfter = useDataSourceStore(state => state.setStartAfter)
+  const setStartAfter = useDataSourceStoreWithSelector(state => state.setStartAfter)
-  const selectedFileList = useDataSourceStore(state => state.selectedFileList)
+  const selectedFileList = useDataSourceStoreWithSelector(state => state.selectedFileList)
-  const setSelectedFileList = useDataSourceStore(state => state.setSelectedFileList)
+  const setSelectedFileList = useDataSourceStoreWithSelector(state => state.setSelectedFileList)
-  const fileList = useDataSourceStore(state => state.fileList)
+  const fileList = useDataSourceStoreWithSelector(state => state.fileList)
-  const setFileList = useDataSourceStore(state => state.setFileList)
+  const setFileList = useDataSourceStoreWithSelector(state => state.setFileList)
 
   const datasourceNodeRunURL = !isInPipeline
     ? `/rag/pipelines/${pipelineId}/workflows/published/datasource/nodes/${nodeId}/run`
@@ -128,6 +128,7 @@ const OnlineDrive = ({
         searchResultsLength={fileList.length}
         handleSelectFile={handleSelectFile}
         handleOpenFolder={handleOpenFolder}
+        isInPipeline={isInPipeline}
       />
     </div>
   )

web/app/components/datasets/documents/create-from-pipeline/data-source/store/index.ts

@@ -25,10 +25,18 @@ export const createDataSourceStore = () => {
   }))
 }
 
-export const useDataSourceStore = <T>(selector: (state: DataSourceShape) => T): T => {
+export const useDataSourceStoreWithSelector = <T>(selector: (state: DataSourceShape) => T): T => {
   const store = useContext(DataSourceContext)
   if (!store)
     throw new Error('Missing DataSourceContext.Provider in the tree')
 
   return useStore(store, selector)
 }
+
+export const useDataSourceStore = () => {
+  const store = useContext(DataSourceContext)
+  if (!store)
+    throw new Error('Missing DataSourceContext.Provider in the tree')
+
+  return store
+}

web/app/components/datasets/documents/create-from-pipeline/data-source/website-crawl/index.tsx

@@ -20,7 +20,7 @@ import type {
   DataSourceNodeProcessingResponse,
 } from '@/types/pipeline'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
-import { useDataSourceStore } from '../store'
+import { useDataSourceStoreWithSelector } from '../store'
 
 const I18N_PREFIX = 'datasetCreation.stepOne.website'
 
@@ -41,16 +41,16 @@ const WebsiteCrawl = ({
   const [crawledNum, setCrawledNum] = useState(0)
   const [crawlErrorMessage, setCrawlErrorMessage] = useState('')
   const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
-  const crawlResult = useDataSourceStore(state => state.crawlResult)
+  const crawlResult = useDataSourceStoreWithSelector(state => state.crawlResult)
-  const setCrawlResult = useDataSourceStore(state => state.setCrawlResult)
+  const setCrawlResult = useDataSourceStoreWithSelector(state => state.setCrawlResult)
-  const step = useDataSourceStore(state => state.step)
+  const step = useDataSourceStoreWithSelector(state => state.step)
-  const setStep = useDataSourceStore(state => state.setStep)
+  const setStep = useDataSourceStoreWithSelector(state => state.setStep)
-  const checkedCrawlResult = useDataSourceStore(state => state.websitePages)
+  const checkedCrawlResult = useDataSourceStoreWithSelector(state => state.websitePages)
-  const setWebsitePages = useDataSourceStore(state => state.setWebsitePages)
+  const setWebsitePages = useDataSourceStoreWithSelector(state => state.setWebsitePages)
-  const previewWebsitePageRef = useDataSourceStore(state => state.previewWebsitePageRef)
+  const previewWebsitePageRef = useDataSourceStoreWithSelector(state => state.previewWebsitePageRef)
-  const previewIndex = useDataSourceStore(state => state.previewIndex)
+  const previewIndex = useDataSourceStoreWithSelector(state => state.previewIndex)
-  const setCurrentWebsite = useDataSourceStore(state => state.setCurrentWebsite)
+  const setCurrentWebsite = useDataSourceStoreWithSelector(state => state.setCurrentWebsite)
-  const setPreviewIndex = useDataSourceStore(state => state.setPreviewIndex)
+  const setPreviewIndex = useDataSourceStoreWithSelector(state => state.setPreviewIndex)
 
   const usePreProcessingParams = useRef(!isInPipeline ? usePublishedPipelinePreProcessingParams : useDraftPipelinePreProcessingParams)
   const { data: paramsConfig, isFetching: isFetchingParams } = usePreProcessingParams.current({

web/app/components/datasets/documents/create-from-pipeline/hooks.ts

@@ -6,7 +6,7 @@ import { BlockEnum, type Node } from '@/app/components/workflow/types'
 import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
 import type { CrawlResult, CrawlResultItem } from '@/models/datasets'
 import { CrawlStep } from '@/models/datasets'
-import { useDataSourceStore } from './data-source/store'
+import { useDataSourceStoreWithSelector } from './data-source/store'
 
 export const useAddDocumentsSteps = () => {
   const { t } = useTranslation()
@@ -63,10 +63,10 @@ export const useDatasourceOptions = (pipelineNodes: Node<DataSourceNodeType>[])
 }
 
 export const useLocalFile = () => {
-  const fileList = useDataSourceStore(state => state.localFileList)
+  const fileList = useDataSourceStoreWithSelector(state => state.localFileList)
-  const previewFileRef = useDataSourceStore(state => state.previewLocalFileRef)
+  const previewFileRef = useDataSourceStoreWithSelector(state => state.previewLocalFileRef)
-  const currentLocalFile = useDataSourceStore(state => state.currentLocalFile)
+  const currentLocalFile = useDataSourceStoreWithSelector(state => state.currentLocalFile)
-  const setCurrentFile = useDataSourceStore(state => state.setCurrentLocalFile)
+  const setCurrentFile = useDataSourceStoreWithSelector(state => state.setCurrentLocalFile)
 
   const allFileLoaded = useMemo(() => (fileList.length > 0 && fileList.every(file => file.file.id)), [fileList])
 
@@ -84,10 +84,10 @@ export const useLocalFile = () => {
 }
 
 export const useOnlineDocuments = () => {
-  const onlineDocuments = useDataSourceStore(state => state.onlineDocuments)
+  const onlineDocuments = useDataSourceStoreWithSelector(state => state.onlineDocuments)
-  const previewOnlineDocumentRef = useDataSourceStore(state => state.previewOnlineDocumentRef)
+  const previewOnlineDocumentRef = useDataSourceStoreWithSelector(state => state.previewOnlineDocumentRef)
-  const setCurrentDocument = useDataSourceStore(state => state.setCurrentDocument)
+  const setCurrentDocument = useDataSourceStoreWithSelector(state => state.setCurrentDocument)
-  const currentDocument = useDataSourceStore(state => state.currentDocument)
+  const currentDocument = useDataSourceStoreWithSelector(state => state.currentDocument)
 
   const hidePreviewOnlineDocument = useCallback(() => {
     setCurrentDocument(undefined)

web/app/components/rag-pipeline/components/panel/test-run/index.tsx

@@ -16,14 +16,15 @@ import CloseButton from './close-button'
 import Header from './header'
 import FooterTips from './footer-tips'
 import DataSourceProvider from '@/app/components/datasets/documents/create-from-pipeline/data-source/store/provider'
-import { useDataSourceStore } from '@/app/components/datasets/documents/create-from-pipeline/data-source/store'
+import { useDataSourceStore, useDataSourceStoreWithSelector } from '@/app/components/datasets/documents/create-from-pipeline/data-source/store'
 
 const TestRunPanel = () => {
   const setShowDebugAndPreviewPanel = useWorkflowStoreWithSelector(state => state.setShowDebugAndPreviewPanel)
-  const fileList = useDataSourceStore(state => state.localFileList)
+  const fileList = useDataSourceStoreWithSelector(state => state.localFileList)
-  const onlineDocuments = useDataSourceStore(state => state.onlineDocuments)
+  const onlineDocuments = useDataSourceStoreWithSelector(state => state.onlineDocuments)
-  const websitePages = useDataSourceStore(state => state.websitePages)
+  const websitePages = useDataSourceStoreWithSelector(state => state.websitePages)
-  const selectedFileList = useDataSourceStore(state => state.selectedFileList)
+  const selectedFileList = useDataSourceStoreWithSelector(state => state.selectedFileList)
+  const { bucket } = useDataSourceStore().getState()
   const [datasource, setDatasource] = useState<Datasource>()
 
   const {
@@ -82,13 +83,19 @@ const TestRunPanel = () => {
     }
     if (datasourceType === DatasourceType.websiteCrawl)
       datasourceInfoList.push(websitePages[0])
+    if (datasourceType === DatasourceType.onlineDrive) {
+      datasourceInfoList.push({
+        bucket,
+        file: selectedFileList[0],
+      })
+    }
     handleRun({
       inputs: data,
       start_node_id: datasource.nodeId,
       datasource_type: datasourceType,
       datasource_info_list: datasourceInfoList,
     })
-  }, [datasource, datasourceType, fileList, handleRun, onlineDocuments, websitePages])
+  }, [bucket, datasource, datasourceType, fileList, handleRun, onlineDocuments, selectedFileList, websitePages])
 
   return (
     <div