382 lines
14 KiB
TypeScript
Raw Normal View History

'use client'
import { useCallback, useMemo, useRef, useState } from 'react'
2025-05-22 14:49:40 +08:00
import DataSourceOptions from './data-source-options'
import type { CrawlResultItem, DocumentItem, CustomFile as File, FileIndexingEstimateResponse, FileItem } from '@/models/datasets'
2025-05-22 14:49:40 +08:00
import LocalFile from '@/app/components/rag-pipeline/components/panel/test-run/data-source/local-file'
import produce from 'immer'
import { useProviderContextSelector } from '@/context/provider-context'
import type { NotionPage } from '@/models/common'
import Notion from '@/app/components/rag-pipeline/components/panel/test-run/data-source/notion'
2025-05-22 14:49:40 +08:00
import VectorSpaceFull from '@/app/components/billing/vector-space-full'
import WebsiteCrawl from '@/app/components/rag-pipeline/components/panel/test-run/data-source/website-crawl'
2025-05-22 14:49:40 +08:00
import Actions from './data-source/actions'
import { useTranslation } from 'react-i18next'
import type { Datasource } from '@/app/components/rag-pipeline/components/panel/test-run/types'
import LeftHeader from './left-header'
import { usePublishedPipelineInfo, useRunPublishedPipeline } from '@/service/use-pipeline'
2025-05-22 14:49:40 +08:00
import { useDatasetDetailContextWithSelector } from '@/context/dataset-detail'
import Loading from '@/app/components/base/loading'
import type { Node } from '@/app/components/workflow/types'
import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
import FilePreview from './preview/file-preview'
import NotionPagePreview from './preview/notion-page-preview'
import WebsitePreview from './preview/web-preview'
import ProcessDocuments from './process-documents'
import ChunkPreview from './preview/chunk-preview'
import Processing from './processing'
import { DatasourceType } from '@/models/pipeline'
const TestRunPanel = () => {
const { t } = useTranslation()
const [currentStep, setCurrentStep] = useState(1)
const [datasource, setDatasource] = useState<Datasource>()
const [fileList, setFiles] = useState<FileItem[]>([])
const [notionPages, setNotionPages] = useState<NotionPage[]>([])
const [websitePages, setWebsitePages] = useState<CrawlResultItem[]>([])
const [websiteCrawlJobId, setWebsiteCrawlJobId] = useState('')
2025-05-22 14:49:40 +08:00
const [currentFile, setCurrentFile] = useState<File | undefined>()
const [currentNotionPage, setCurrentNotionPage] = useState<NotionPage | undefined>()
const [currentWebsite, setCurrentWebsite] = useState<CrawlResultItem | undefined>()
const [estimateData, setEstimateData] = useState<FileIndexingEstimateResponse | undefined>(undefined)
const plan = useProviderContextSelector(state => state.plan)
const enableBilling = useProviderContextSelector(state => state.enableBilling)
const datasetId = useDatasetDetailContextWithSelector(s => s.dataset?.id)
2025-05-22 14:49:40 +08:00
const pipelineId = useDatasetDetailContextWithSelector(s => s.dataset?.pipeline_id)
const indexingType = useDatasetDetailContextWithSelector(s => s.dataset?.indexing_technique)
const retrievalMethod = useDatasetDetailContextWithSelector(s => s.dataset?.retrieval_model_dict.search_method)
const isPreview = useRef(false)
const formRef = useRef<any>(null)
const previewFile = useRef<DocumentItem>(fileList[0].file as DocumentItem)
const previewNotionPage = useRef<NotionPage>(notionPages[0])
const previewWebsitePage = useRef<CrawlResultItem>(websitePages[0])
2025-05-22 14:49:40 +08:00
const { data: pipelineInfo, isFetching: isFetchingPipelineInfo } = usePublishedPipelineInfo(pipelineId || '')
const allFileLoaded = (fileList.length > 0 && fileList.every(file => file.file.id))
const isVectorSpaceFull = plan.usage.vectorSpace >= plan.total.vectorSpace
const isShowVectorSpaceFull = allFileLoaded && isVectorSpaceFull && enableBilling
const notSupportBatchUpload = enableBilling && plan.type === 'sandbox'
const nextDisabled = useMemo(() => {
if (!fileList.length)
return true
if (fileList.some(file => !file.file.id))
return true
return isShowVectorSpaceFull
}, [fileList, isShowVectorSpaceFull])
const nextBtnDisabled = useMemo(() => {
if (!datasource) return true
if (datasource.type === DatasourceType.localFile)
return nextDisabled
if (datasource.type === DatasourceType.onlineDocument)
return isShowVectorSpaceFull || !notionPages.length
if (datasource.type === DatasourceType.websiteCrawl)
return isShowVectorSpaceFull || !websitePages.length
return false
}, [datasource, nextDisabled, isShowVectorSpaceFull, notionPages.length, websitePages.length])
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,
}
})
setFiles(newList)
}
2025-05-22 14:49:40 +08:00
const updateFileList = useCallback((preparedFiles: FileItem[]) => {
setFiles(preparedFiles)
2025-05-22 14:49:40 +08:00
}, [])
2025-05-22 14:49:40 +08:00
const updateNotionPages = useCallback((value: NotionPage[]) => {
setNotionPages(value)
2025-05-22 14:49:40 +08:00
}, [])
const updateCurrentFile = useCallback((file: File) => {
setCurrentFile(file)
}, [])
const hideFilePreview = useCallback(() => {
setCurrentFile(undefined)
}, [])
const updateCurrentPage = useCallback((page: NotionPage) => {
setCurrentNotionPage(page)
}, [])
const hideNotionPagePreview = useCallback(() => {
setCurrentNotionPage(undefined)
}, [])
const updateCurrentWebsite = useCallback((website: CrawlResultItem) => {
setCurrentWebsite(website)
}, [])
const hideWebsitePreview = useCallback(() => {
setCurrentWebsite(undefined)
}, [])
const handleNextStep = useCallback(() => {
setCurrentStep(preStep => preStep + 1)
}, [])
const handleBackStep = useCallback(() => {
setCurrentStep(preStep => preStep - 1)
}, [])
const { mutateAsync: runPublishedPipeline, isIdle, isPending } = useRunPublishedPipeline()
const handlePreviewChunks = useCallback(async (data: Record<string, any>) => {
if (!datasource)
return
const datasourceInfoList: Record<string, any>[] = []
if (datasource.type === DatasourceType.localFile) {
const { id, name, type, size, extension, mime_type } = previewFile.current as File
const documentInfo = {
upload_file_id: id,
name,
type,
size,
extension,
mime_type,
}
datasourceInfoList.push(documentInfo)
}
if (datasource.type === DatasourceType.onlineDocument) {
const { workspace_id, ...rest } = previewNotionPage.current
const documentInfo = {
workspace_id,
page: rest,
}
datasourceInfoList.push(documentInfo)
}
if (datasource.type === DatasourceType.websiteCrawl) {
const documentInfo = {
job_id: websiteCrawlJobId,
result: previewWebsitePage.current,
}
datasourceInfoList.push(documentInfo)
}
await runPublishedPipeline({
pipeline_id: pipelineId!,
inputs: data,
start_node_id: datasource.nodeId,
datasource_type: datasource.type,
datasource_info_list: datasourceInfoList,
is_preview: true,
}, {
onSuccess: (res) => {
setEstimateData(res.data.outputs as FileIndexingEstimateResponse)
},
})
}, [datasource, pipelineId, runPublishedPipeline, websiteCrawlJobId])
const handleProcess = useCallback(async (data: Record<string, any>) => {
if (!datasource)
return
const datasourceInfoList: Record<string, any>[] = []
if (datasource.type === DatasourceType.localFile) {
fileList.forEach((file) => {
const { id, name, type, size, extension, mime_type } = file.file
const documentInfo = {
upload_file_id: id,
name,
type,
size,
extension,
mime_type,
}
datasourceInfoList.push(documentInfo)
})
}
if (datasource.type === DatasourceType.onlineDocument) {
notionPages.forEach((page) => {
const { workspace_id, ...rest } = page
const documentInfo = {
workspace_id,
page: rest,
}
datasourceInfoList.push(documentInfo)
})
}
if (datasource.type === DatasourceType.websiteCrawl) {
const documentInfo = {
job_id: websiteCrawlJobId,
result: websitePages,
}
datasourceInfoList.push(documentInfo)
}
await runPublishedPipeline({
pipeline_id: pipelineId!,
inputs: data,
start_node_id: datasource.nodeId,
datasource_type: datasource.type,
datasource_info_list: datasourceInfoList,
is_preview: false,
}, {
onSuccess: () => {
handleNextStep()
},
})
}, [datasource, fileList, handleNextStep, notionPages, pipelineId, runPublishedPipeline, websiteCrawlJobId, websitePages])
const onClickProcess = useCallback(() => {
isPreview.current = false
formRef.current?.submit()
}, [])
const onClickPreview = useCallback(() => {
isPreview.current = true
formRef.current?.submit()
}, [])
const handleSubmit = useCallback((data: Record<string, any>) => {
isPreview.current ? handlePreviewChunks(data) : handleProcess(data)
}, [handlePreviewChunks, handleProcess])
const handlePreviewFileChange = useCallback((file: DocumentItem) => {
previewFile.current = file
onClickPreview()
}, [onClickPreview])
const handlePreviewNotionPageChange = useCallback((page: NotionPage) => {
previewNotionPage.current = page
onClickPreview()
}, [onClickPreview])
const handlePreviewWebsiteChange = useCallback((website: CrawlResultItem) => {
previewWebsitePage.current = website
onClickPreview()
}, [onClickPreview])
2025-05-22 14:49:40 +08:00
if (isFetchingPipelineInfo) {
return (
<Loading type='app' />
)
}
return (
<div
className='relative flex h-[calc(100vh-56px)] overflow-x-auto rounded-t-2xl border-t border-effects-highlight bg-background-default-subtle'
>
<div className='flex h-full min-w-[760px] flex-1 flex-col px-14'>
<LeftHeader
title={t('datasetPipeline.addDocuments.title')}
currentStep={currentStep}
/>
<div className='grow overflow-y-auto'>
{
currentStep === 1 && (
2025-05-22 14:49:40 +08:00
<div className='flex flex-col gap-y-5 pt-4'>
<DataSourceOptions
datasourceNodeId={datasource?.nodeId || ''}
onSelect={setDatasource}
2025-05-22 14:49:40 +08:00
pipelineNodes={(pipelineInfo?.graph.nodes || []) as Node<DataSourceNodeType>[]}
/>
{datasource?.type === DatasourceType.localFile && (
2025-05-22 14:49:40 +08:00
<LocalFile
files={fileList}
allowedExtensions={datasource?.fileExtensions || []}
2025-05-22 14:49:40 +08:00
updateFile={updateFile}
updateFileList={updateFileList}
onPreview={updateCurrentFile}
notSupportBatchUpload={notSupportBatchUpload}
/>
)}
{datasource?.type === DatasourceType.onlineDocument && (
2025-05-22 14:49:40 +08:00
<Notion
nodeId={datasource?.nodeId || ''}
notionPages={notionPages}
updateNotionPages={updateNotionPages}
canPreview
onPreview={updateCurrentPage}
/>
)}
{datasource?.type === DatasourceType.websiteCrawl && (
<WebsiteCrawl
2025-05-22 14:49:40 +08:00
nodeId={datasource?.nodeId || ''}
variables={datasource?.variables}
headerInfo={{
title: datasource.description,
docTitle: datasource.docTitle || '',
docLink: datasource.docLink || '',
}}
2025-05-22 14:49:40 +08:00
checkedCrawlResult={websitePages}
onCheckedCrawlResultChange={setWebsitePages}
onJobIdChange={setWebsiteCrawlJobId}
onPreview={updateCurrentWebsite}
/>
)}
{isShowVectorSpaceFull && (
<VectorSpaceFull />
)}
<Actions disabled={nextBtnDisabled} handleNextStep={handleNextStep} />
2025-05-22 14:49:40 +08:00
</div>
)
}
{
currentStep === 2 && (
<ProcessDocuments
ref={formRef}
dataSourceNodeId={datasource?.nodeId || ''}
onProcess={onClickProcess}
onPreview={onClickPreview}
onSubmit={handleSubmit}
onBack={handleBackStep}
/>
)
}
{
currentStep === 3 && (
<Processing
datasetId={datasetId!}
batchId={''}
documents={[]}
indexingType={indexingType!}
retrievalMethod={retrievalMethod!}
/>
)
}
</div>
</div>
{/* Preview */}
{
currentStep === 1 && (
<div className='flex h-full w-[752px] shrink-0 pl-2 pt-2'>
{currentFile && <FilePreview file={currentFile} hidePreview={hideFilePreview} />}
{currentNotionPage && <NotionPagePreview currentPage={currentNotionPage} hidePreview={hideNotionPagePreview} />}
{currentWebsite && <WebsitePreview payload={currentWebsite} hidePreview={hideWebsitePreview} />}
</div>
)
}
{
currentStep === 2 && (
<div className='flex h-full w-[752px] shrink-0 pl-2 pt-2'>
{estimateData && (
<ChunkPreview
datasource={datasource!}
files={fileList.map(file => file.file)}
notionPages={notionPages}
websitePages={websitePages}
isIdle={isIdle}
isPending={isPending}
estimateData={estimateData}
onPreview={onClickPreview}
handlePreviewFileChange={handlePreviewFileChange}
handlePreviewNotionPageChange={handlePreviewNotionPageChange}
handlePreviewWebsitePageChange={handlePreviewWebsiteChange}
/>
)}
</div>
)
}
</div>
)
}
export default TestRunPanel