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 526a1159c8..05cbabcdbf 100644 --- a/web/app/components/datasets/documents/create-from-pipeline/hooks.ts +++ b/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 { CrawlStep } from '@/models/datasets' -import { useDataSourceStoreWithSelector } from './data-source/store' +import { useDataSourceStore, 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) - }, [setCurrentFile]) + const { setCurrentLocalFile } = dataSourceStore.getState() + 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([]) - const [currentWebsite, setCurrentWebsite] = useState() - const [crawlResult, setCrawlResult] = useState() - const [step, setStep] = useState(CrawlStep.init) - const [previewIndex, setPreviewIndex] = useState(-1) - - const previewWebsitePage = useRef(websitePages[0]) - - const updateCurrentWebsite = useCallback((website: CrawlResultItem, index: number) => { - setCurrentWebsite(website) - setPreviewIndex(index) - }, []) + const websitePages = useDataSourceStoreWithSelector(state => state.websitePages) + const currentWebsite = useDataSourceStoreWithSelector(state => state.currentWebsite) + const previewWebsitePageRef = useDataSourceStoreWithSelector(state => state.previewWebsitePageRef) + const dataSourceStore = useDataSourceStore() const hideWebsitePreview = useCallback(() => { + const { setCurrentWebsite, setPreviewIndex } = dataSourceStore.getState() setCurrentWebsite(undefined) setPreviewIndex(-1) - }, []) - - const updataCheckedCrawlResultChange = useCallback((checkedCrawlResult: CrawlResultItem[]) => { - setWebsitePages(checkedCrawlResult) - previewWebsitePage.current = checkedCrawlResult[0] - }, []) + }, [dataSourceStore]) return { websitePages, - crawlResult, - setCrawlResult, - step, - setStep, - previewWebsitePage, - updataCheckedCrawlResultChange, + previewWebsitePageRef, currentWebsite, - updateCurrentWebsite, - previewIndex, hideWebsitePreview, } } export const useOnlineDrive = () => { - return {} + const fileList = useDataSourceStoreWithSelector(state => state.fileList) + const selectedFileList = useDataSourceStoreWithSelector(state => state.selectedFileList) + + return { + fileList, + selectedFileList, + } } diff --git a/web/app/components/datasets/documents/create-from-pipeline/index.tsx b/web/app/components/datasets/documents/create-from-pipeline/index.tsx index 69d36b5365..9d1f37a486 100644 --- a/web/app/components/datasets/documents/create-from-pipeline/index.tsx +++ b/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(undefined) const [batchId, setBatchId] = useState('') const [documents, setDocuments] = useState([]) + const dataSourceStore = useDataSourceStore() const isPreview = useRef(false) const formRef = useRef(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) => { 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 (