import { useTranslation } from 'react-i18next' import { AddDocumentsStep } from './types' import type { DataSourceOption, Datasource } from '@/app/components/rag-pipeline/components/panel/test-run/types' import { useCallback, useMemo, useRef, 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 { DatasourceType } from '@/models/pipeline' import type { CrawlResultItem, DocumentItem, FileItem } from '@/models/datasets' import produce from 'immer' import type { NotionPage } from '@/models/common' export const useAddDocumentsSteps = () => { const { t } = useTranslation() const [currentStep, setCurrentStep] = useState(1) const handleNextStep = useCallback(() => { setCurrentStep(preStep => preStep + 1) }, []) const handleBackStep = useCallback(() => { setCurrentStep(preStep => preStep - 1) }, []) const steps = [ { label: t('datasetPipeline.addDocuments.steps.chooseDatasource'), value: AddDocumentsStep.dataSource, }, { label: t('datasetPipeline.addDocuments.steps.processDocuments'), value: AddDocumentsStep.processDocuments, }, { label: t('datasetPipeline.addDocuments.steps.processingDocuments'), value: AddDocumentsStep.processingDocuments, }, ] return { steps, currentStep, handleNextStep, handleBackStep, } } export const useDatasourceOptions = (pipelineNodes: Node[]) => { const datasourceNodes = pipelineNodes.filter(node => node.data.type === BlockEnum.DataSource) const datasources: Datasource[] = useMemo(() => { return datasourceNodes.map((node) => { return { nodeId: node.id, type: node.data.provider_type as DatasourceType, variables: node.data.variables || [], description: node.data.desc || '', docTitle: '', // todo: Add docTitle and docLink if needed, or remove these properties if not used docLink: '', fileExtensions: node.data.fileExtensions || [], } }) }, [datasourceNodes]) const options = useMemo(() => { const options: DataSourceOption[] = [] datasourceNodes.forEach((node) => { const label = node.data.title options.push({ label, value: node.id, data: node.data, }) }) return options }, [datasourceNodes]) return { datasources, options } } export const useLocalFile = () => { const [fileList, setFileList] = useState([]) const [currentFile, setCurrentFile] = useState() const previewFile = useRef(fileList[0]?.file as DocumentItem) const allFileLoaded = useMemo(() => (fileList.length > 0 && fileList.every(file => file.file.id)), [fileList]) const updateFile = (fileItem: FileItem, progress: number, list: FileItem[]) => { const newList = produce(list, (draft) => { const targetIndex = draft.findIndex(file => file.fileID === fileItem.fileID) draft[targetIndex] = { ...draft[targetIndex], progress, } }) setFileList(newList) } const updateFileList = useCallback((preparedFiles: FileItem[]) => { setFileList(preparedFiles) }, []) const updateCurrentFile = useCallback((file: File) => { setCurrentFile(file) }, []) const hideFilePreview = useCallback(() => { setCurrentFile(undefined) }, []) return { fileList, previewFile, allFileLoaded, updateFile, updateFileList, currentFile, updateCurrentFile, hideFilePreview, } } export const useNotionsPages = () => { const [notionPages, setNotionPages] = useState([]) const [currentNotionPage, setCurrentNotionPage] = useState() const previewNotionPage = useRef(notionPages[0]) const updateNotionPages = (value: NotionPage[]) => { setNotionPages(value) } const updateCurrentPage = useCallback((page: NotionPage) => { setCurrentNotionPage(page) }, []) const hideNotionPagePreview = useCallback(() => { setCurrentNotionPage(undefined) }, []) return { notionPages, previewNotionPage, updateNotionPages, currentNotionPage, updateCurrentPage, hideNotionPagePreview, } } export const useWebsiteCrawl = () => { const [websitePages, setWebsitePages] = useState([]) const [websiteCrawlJobId, setWebsiteCrawlJobId] = useState('') const [currentWebsite, setCurrentWebsite] = useState() const previewWebsitePage = useRef(websitePages[0]) const updateCurrentWebsite = useCallback((website: CrawlResultItem) => { setCurrentWebsite(website) }, []) const hideWebsitePreview = useCallback(() => { setCurrentWebsite(undefined) }, []) return { websitePages, websiteCrawlJobId, previewWebsitePage, setWebsitePages, setWebsiteCrawlJobId, currentWebsite, updateCurrentWebsite, hideWebsitePreview, } }