import React, { useState, useMemo } from 'react'; import { Box, Flex, Button, useTheme, Image } from '@chakra-ui/react'; import { useToast } from '@/hooks/useToast'; import { useConfirm } from '@/hooks/useConfirm'; import { useMutation } from '@tanstack/react-query'; import { getErrText } from '@/utils/tools'; import MyIcon from '@/components/Icon'; import DeleteIcon, { hoverDeleteStyles } from '@/components/Icon/delete'; import { TrainingModeEnum } from '@/constants/plugin'; import FileSelect, { type FileItemType } from './FileSelect'; import { useRouter } from 'next/router'; import { useDatasetStore } from '@/store/dataset'; import { putMarkFilesUsed } from '@/api/core/dataset/file'; import { chunksUpload } from '@/utils/web/core/dataset'; const fileExtension = '.csv'; const CsvImport = ({ kbId }: { kbId: string }) => { const { kbDetail } = useDatasetStore(); const maxToken = kbDetail.vectorModel?.maxToken || 2000; const theme = useTheme(); const router = useRouter(); const { toast } = useToast(); const [files, setFiles] = useState([]); const [successChunks, setSuccessChunks] = useState(0); const totalChunk = useMemo( () => files.reduce((sum, file) => sum + file.chunks.length, 0), [files] ); const emptyFiles = useMemo(() => files.length === 0, [files]); const { openConfirm, ConfirmModal } = useConfirm({ content: `该任务无法终止,需要一定时间生成索引,请确认导入。如果余额不足,未完成的任务会被暂停,充值后可继续进行。` }); const { mutate: onclickUpload, isLoading: uploading } = useMutation({ mutationFn: async () => { // mark the file is used await putMarkFilesUsed({ fileIds: files.map((file) => file.id) }); const chunks = files .map((file) => file.chunks) .flat() .filter((item) => item?.q); const filterChunks = chunks.filter((item) => item.q.length < maxToken * 1.5); if (filterChunks.length !== chunks.length) { toast({ title: `${chunks.length - filterChunks.length}条数据超出长度,已被过滤`, status: 'info' }); } // upload data const { insertLen } = await chunksUpload({ kbId, chunks, mode: TrainingModeEnum.index, onUploading: (insertLen) => { setSuccessChunks(insertLen); } }); toast({ title: `去重后共导入 ${insertLen} 条数据,请耐心等待训练.`, status: 'success' }); router.replace({ query: { kbId, currentTab: 'dataset' } }); }, onError(err) { toast({ title: getErrText(err, '导入文件失败'), status: 'error' }); } }); const filenameStyles = { className: 'textEllipsis', maxW: '400px' }; return ( setFiles((state) => files.concat(state))} showUrlFetch={false} showCreateFile={false} py={emptyFiles ? '100px' : 5} isCsv /> {!emptyFiles && ( <> {files.map((item) => ( {''} {item.filename} { e.stopPropagation(); setFiles((state) => state.filter((file) => file.id !== item.id)); }} /> ))} )} {!emptyFiles && ( 分段预览({totalChunk}组) {totalChunk > 100 && ( 仅展示部分 )} {files.map((file) => file.chunks.slice(0, 100).map((item, i) => ( # {i + 1} {item.source && ({item.source})} { setFiles((state) => state.map((stateFile) => stateFile.id === file.id ? { ...file, chunks: [...file.chunks.slice(0, i), ...file.chunks.slice(i + 1)] } : stateFile ) ); }} /> {`${item.q}\n${item.a}`} )) )} )} ); }; export default CsvImport;