feat: Integrate useOnlineDrive hook and enhance datasource handling in CreateFormPipeline

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

@@ -1,12 +1,10 @@
 import { useTranslation } from 'react-i18next'
 import { AddDocumentsStep } from './types'
 import type { DataSourceOption } from '@/app/components/rag-pipeline/components/panel/test-run/types'
-import { useCallback, useMemo, useRef, useState } from 'react'
+import { useCallback, useMemo, useState } from 'react'
 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 { useDataSourceStore, useDataSourceStoreWithSelector } from './data-source/store'
-import { CrawlStep } from '@/models/datasets'
-import { useDataSourceStoreWithSelector } from './data-source/store'
 
 export const useAddDocumentsSteps = () => {
   const { t } = useTranslation()
@@ -66,13 +64,14 @@ export const useLocalFile = () => {
   const fileList = useDataSourceStoreWithSelector(state => state.localFileList)
   const previewFileRef = useDataSourceStoreWithSelector(state => state.previewLocalFileRef)
   const currentLocalFile = useDataSourceStoreWithSelector(state => state.currentLocalFile)
-  const setCurrentFile = useDataSourceStoreWithSelector(state => state.setCurrentLocalFile)
+  const dataSourceStore = useDataSourceStore()
 
   const allFileLoaded = useMemo(() => (fileList.length > 0 && fileList.every(file => file.file.id)), [fileList])
 
   const hidePreviewLocalFile = useCallback(() => {
-    setCurrentFile(undefined)
+    const { setCurrentLocalFile } = dataSourceStore.getState()
-  }, [setCurrentFile])
+    setCurrentLocalFile(undefined)
+  }, [dataSourceStore])
 
   return {
     fileList,
@@ -86,12 +85,13 @@ export const useLocalFile = () => {
 export const useOnlineDocuments = () => {
   const onlineDocuments = useDataSourceStoreWithSelector(state => state.onlineDocuments)
   const previewOnlineDocumentRef = useDataSourceStoreWithSelector(state => state.previewOnlineDocumentRef)
-  const setCurrentDocument = useDataSourceStoreWithSelector(state => state.setCurrentDocument)
   const currentDocument = useDataSourceStoreWithSelector(state => state.currentDocument)
+  const dataSourceStore = useDataSourceStore()
 
   const hidePreviewOnlineDocument = useCallback(() => {
+    const { setCurrentDocument } = dataSourceStore.getState()
     setCurrentDocument(undefined)
-  }, [setCurrentDocument])
+  }, [dataSourceStore])
 
   return {
     onlineDocuments,
@@ -102,44 +102,31 @@ export const useOnlineDocuments = () => {
 }
 
 export const useWebsiteCrawl = () => {
-  const [websitePages, setWebsitePages] = useState<CrawlResultItem[]>([])
+  const websitePages = useDataSourceStoreWithSelector(state => state.websitePages)
-  const [currentWebsite, setCurrentWebsite] = useState<CrawlResultItem | undefined>()
+  const currentWebsite = useDataSourceStoreWithSelector(state => state.currentWebsite)
-  const [crawlResult, setCrawlResult] = useState<CrawlResult | undefined>()
+  const previewWebsitePageRef = useDataSourceStoreWithSelector(state => state.previewWebsitePageRef)
-  const [step, setStep] = useState<CrawlStep>(CrawlStep.init)
+  const dataSourceStore = useDataSourceStore()
-  const [previewIndex, setPreviewIndex] = useState<number>(-1)
-
-  const previewWebsitePage = useRef<CrawlResultItem>(websitePages[0])
-
-  const updateCurrentWebsite = useCallback((website: CrawlResultItem, index: number) => {
-    setCurrentWebsite(website)
-    setPreviewIndex(index)
-  }, [])
 
   const hideWebsitePreview = useCallback(() => {
+    const { setCurrentWebsite, setPreviewIndex } = dataSourceStore.getState()
     setCurrentWebsite(undefined)
     setPreviewIndex(-1)
-  }, [])
+  }, [dataSourceStore])
-
-  const updataCheckedCrawlResultChange = useCallback((checkedCrawlResult: CrawlResultItem[]) => {
-    setWebsitePages(checkedCrawlResult)
-    previewWebsitePage.current = checkedCrawlResult[0]
-  }, [])
 
   return {
     websitePages,
-    crawlResult,
+    previewWebsitePageRef,
-    setCrawlResult,
-    step,
-    setStep,
-    previewWebsitePage,
-    updataCheckedCrawlResultChange,
     currentWebsite,
-    updateCurrentWebsite,
-    previewIndex,
     hideWebsitePreview,
   }
 }
 
 export const useOnlineDrive = () => {
-  return {}
+  const fileList = useDataSourceStoreWithSelector(state => state.fileList)
+  const selectedFileList = useDataSourceStoreWithSelector(state => state.selectedFileList)
+
+  return {
+    fileList,
+    selectedFileList,
+  }
 }

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

@@ -27,8 +27,9 @@ import Processing from './processing'
 import type { InitialDocumentDetail, PublishedPipelineRunPreviewResponse, PublishedPipelineRunResponse } from '@/models/pipeline'
 import { DatasourceType } from '@/models/pipeline'
 import { TransferMethod } from '@/types/app'
-import { useAddDocumentsSteps, useLocalFile, useOnlineDocuments, useWebsiteCrawl } from './hooks'
+import { useAddDocumentsSteps, useLocalFile, useOnlineDocuments, useOnlineDrive, useWebsiteCrawl } from './hooks'
 import DataSourceProvider from './data-source/store/provider'
+import { useDataSourceStore } from './data-source/store'
 
 const CreateFormPipeline = () => {
   const { t } = useTranslation()
@@ -39,6 +40,7 @@ const CreateFormPipeline = () => {
   const [estimateData, setEstimateData] = useState<FileIndexingEstimateResponse | undefined>(undefined)
   const [batchId, setBatchId] = useState('')
   const [documents, setDocuments] = useState<InitialDocumentDetail[]>([])
+  const dataSourceStore = useDataSourceStore()
 
   const isPreview = useRef(false)
   const formRef = useRef<any>(null)
@@ -66,27 +68,44 @@ const CreateFormPipeline = () => {
   } = useOnlineDocuments()
   const {
     websitePages,
-    previewWebsitePage,
+    previewWebsitePageRef,
     currentWebsite,
     hideWebsitePreview,
   } = useWebsiteCrawl()
-  // const { } = useOnlineDrive()
+  const {
+    fileList: onlineDriveFileList,
+    selectedFileList,
+  } = useOnlineDrive()
 
-  const isVectorSpaceFull = plan.usage.vectorSpace >= plan.total.vectorSpace
-  const isShowVectorSpaceFull = allFileLoaded && isVectorSpaceFull && enableBilling
-  const notSupportBatchUpload = enableBilling && plan.type === 'sandbox'
   const datasourceType = datasource?.nodeData.provider_type
+  const isVectorSpaceFull = plan.usage.vectorSpace >= plan.total.vectorSpace
+  const isShowVectorSpaceFull = useMemo(() => {
+    if (!datasource)
+      return false
+    if (datasourceType === DatasourceType.localFile)
+      return allFileLoaded && isVectorSpaceFull && enableBilling
+    if (datasourceType === DatasourceType.onlineDocument)
+      return onlineDocuments.length > 0 && isVectorSpaceFull && enableBilling
+    if (datasourceType === DatasourceType.websiteCrawl)
+      return websitePages.length > 0 && isVectorSpaceFull && enableBilling
+    if (datasourceType === DatasourceType.onlineDrive)
+      return onlineDriveFileList.length > 0 && isVectorSpaceFull && enableBilling
+    return false
+  }, [allFileLoaded, datasource, datasourceType, enableBilling, isVectorSpaceFull, onlineDocuments.length, onlineDriveFileList.length, websitePages.length])
+  const notSupportBatchUpload = enableBilling && plan.type === 'sandbox'
 
   const nextBtnDisabled = useMemo(() => {
     if (!datasource) return true
     if (datasourceType === DatasourceType.localFile)
-      return isShowVectorSpaceFull || !fileList.length || fileList.some(file => !file.file.id)
+      return isShowVectorSpaceFull || !fileList.length || !allFileLoaded
     if (datasourceType === DatasourceType.onlineDocument)
       return isShowVectorSpaceFull || !onlineDocuments.length
     if (datasourceType === DatasourceType.websiteCrawl)
       return isShowVectorSpaceFull || !websitePages.length
+    if (datasourceType === DatasourceType.onlineDrive)
+      return isShowVectorSpaceFull || !onlineDriveFileList.length
     return false
-  }, [datasource, datasourceType, isShowVectorSpaceFull, fileList, onlineDocuments.length, websitePages.length])
+  }, [datasource, datasourceType, isShowVectorSpaceFull, fileList.length, allFileLoaded, onlineDocuments.length, websitePages.length, onlineDriveFileList.length])
 
   const { mutateAsync: runPublishedPipeline, isIdle, isPending } = useRunPublishedPipeline()
 
@@ -117,7 +136,7 @@ const CreateFormPipeline = () => {
       datasourceInfoList.push(documentInfo)
     }
     if (datasourceType === DatasourceType.websiteCrawl)
-      datasourceInfoList.push(previewWebsitePage.current)
+      datasourceInfoList.push(previewWebsitePageRef.current!)
     await runPublishedPipeline({
       pipeline_id: pipelineId!,
       inputs: data,
@@ -130,7 +149,7 @@ const CreateFormPipeline = () => {
         setEstimateData((res as PublishedPipelineRunPreviewResponse).data.outputs)
       },
     })
-  }, [datasource, datasourceType, pipelineId, previewFileRef, previewOnlineDocumentRef, previewWebsitePage, runPublishedPipeline])
+  }, [datasource, datasourceType, previewWebsitePageRef, runPublishedPipeline, pipelineId, previewFileRef, previewOnlineDocumentRef])
 
   const handleProcess = useCallback(async (data: Record<string, any>) => {
     if (!datasource)
@@ -167,6 +186,17 @@ const CreateFormPipeline = () => {
         datasourceInfoList.push(websitePage)
       })
     }
+    if (datasourceType === DatasourceType.onlineDrive) {
+      if (datasourceType === DatasourceType.onlineDrive) {
+        const { bucket } = dataSourceStore.getState()
+        selectedFileList.forEach((key) => {
+          datasourceInfoList.push({
+            bucket,
+            key,
+          })
+        })
+      }
+    }
     await runPublishedPipeline({
       pipeline_id: pipelineId!,
       inputs: data,
@@ -181,7 +211,7 @@ const CreateFormPipeline = () => {
         handleNextStep()
       },
     })
-  }, [datasource, datasourceType, fileList, handleNextStep, onlineDocuments, pipelineId, runPublishedPipeline, websitePages])
+  }, [dataSourceStore, datasource, datasourceType, fileList, handleNextStep, onlineDocuments, pipelineId, runPublishedPipeline, selectedFileList, websitePages])
 
   const onClickProcess = useCallback(() => {
     isPreview.current = false
@@ -208,9 +238,9 @@ const CreateFormPipeline = () => {
   }, [onClickPreview, previewOnlineDocumentRef])
 
   const handlePreviewWebsiteChange = useCallback((website: CrawlResultItem) => {
-    previewWebsitePage.current = website
+    previewWebsitePageRef.current = website
     onClickPreview()
-  }, [onClickPreview, previewWebsitePage])
+  }, [onClickPreview, previewWebsitePageRef])
 
   if (isFetchingPipelineInfo) {
     return (