Fix/upload document limit (#1033)

web/app/components/base/notion-page-selector/base.tsx

@@ -1,4 +1,6 @@
 import { useCallback, useEffect, useMemo, useState } from 'react'
+import { useTranslation } from 'react-i18next'
+import { useContext } from 'use-context-selector'
 import useSWR from 'swr'
 import cn from 'classnames'
 import s from './base.module.css'
@@ -9,6 +11,7 @@ import { preImportNotionPages } from '@/service/datasets'
 import AccountSetting from '@/app/components/header/account-setting'
 import { NotionConnector } from '@/app/components/datasets/create/step-one'
 import type { DataSourceNotionPage, DataSourceNotionPageMap, DataSourceNotionWorkspace } from '@/models/common'
+import { ToastContext } from '@/app/components/base/toast'
 
 export type NotionPageSelectorValue = DataSourceNotionPage & { workspace_id: string }
 
@@ -19,6 +22,8 @@ type NotionPageSelectorProps = {
   previewPageId?: string
   onPreview?: (selectedPage: NotionPageSelectorValue) => void
   datasetId?: string
+  countLimit: number
+  countUsed: number
 }
 
 const NotionPageSelector = ({
@@ -28,7 +33,11 @@ const NotionPageSelector = ({
   previewPageId,
   onPreview,
   datasetId = '',
+  countLimit,
+  countUsed,
 }: NotionPageSelectorProps) => {
+  const { t } = useTranslation()
+  const { notify } = useContext(ToastContext)
   const { data, mutate } = useSWR({ url: '/notion/pre-import/pages', datasetId }, preImportNotionPages)
   const [prevData, setPrevData] = useState(data)
   const [searchValue, setSearchValue] = useState('')
@@ -71,9 +80,13 @@ const NotionPageSelector = ({
   const handleSelectWorkspace = useCallback((workspaceId: string) => {
     setCurrentWorkspaceId(workspaceId)
   }, [])
-  const handleSelecPages = (selectedPagesId: Set<string>) => {
-    setSelectedPagesId(new Set(Array.from(selectedPagesId)))
-    const selectedPages = Array.from(selectedPagesId).map(pageId => getPagesMapAndSelectedPagesId[0][pageId])
+  const handleSelecPages = (newSelectedPagesId: Set<string>) => {
+    const selectedPages = Array.from(newSelectedPagesId).map(pageId => getPagesMapAndSelectedPagesId[0][pageId])
+    if (selectedPages.length > countLimit - countUsed) {
+      notify({ type: 'error', message: t('datasetCreation.stepOne.overCountLimit', { countLimit }) })
+      return false
+    }
+    setSelectedPagesId(new Set(Array.from(newSelectedPagesId)))
     onSelect(selectedPages)
   }
   const handlePreviewPage = (previewPageId: string) => {

web/app/components/base/notion-page-selector/page-selector/index.tsx

@@ -228,29 +228,30 @@ const PageSelector = ({
     setDataList(newDataList)
   }
 
+  const copyValue = new Set([...value])
   const handleCheck = (index: number) => {
     const current = currentDataList[index]
     const pageId = current.page_id
     const currentWithChildrenAndDescendants = listMapWithChildrenAndDescendants[pageId]
 
-    if (value.has(pageId)) {
+    if (copyValue.has(pageId)) {
       if (!searchValue) {
         for (const item of currentWithChildrenAndDescendants.descendants)
-          value.delete(item)
+          copyValue.delete(item)
       }
 
-      value.delete(pageId)
+      copyValue.delete(pageId)
     }
     else {
       if (!searchValue) {
         for (const item of currentWithChildrenAndDescendants.descendants)
-          value.add(item)
+          copyValue.add(item)
       }
 
-      value.add(pageId)
+      copyValue.add(pageId)
     }
 
-    onSelect(new Set([...value]))
+    onSelect(new Set([...copyValue]))
   }
 
   const handlePreview = (index: number) => {

web/app/components/datasets/create/file-uploader/index.tsx

@@ -18,6 +18,8 @@ type IFileUploaderProps = {
   onFileUpdate: (fileItem: FileItem, progress: number, list: FileItem[]) => void
   onFileListUpdate?: (files: any) => void
   onPreview: (file: File) => void
+  countLimit: number
+  countUsed: number
 }
 
 const ACCEPTS = [
@@ -39,6 +41,8 @@ const FileUploader = ({
   onFileUpdate,
   onFileListUpdate,
   onPreview,
+  countLimit,
+  countUsed,
 }: IFileUploaderProps) => {
   const { t } = useTranslation()
   const { notify } = useContext(ToastContext)
@@ -145,6 +149,10 @@ const FileUploader = ({
   const initialUpload = useCallback((files: File[]) => {
     if (!files.length)
       return false
+    if (files.length > countLimit - countUsed) {
+      notify({ type: 'error', message: t('datasetCreation.stepOne.overCountLimit', { countLimit }) })
+      return false
+    }
     const preparedFiles = files.map((file, index) => ({
       fileID: `file${index}-${Date.now()}`,
       file,

web/app/components/datasets/create/step-one/index.tsx

@@ -1,5 +1,6 @@
 'use client'
 import React, { useMemo, useState } from 'react'
+import useSWR from 'swr'
 import { useTranslation } from 'react-i18next'
 import cn from 'classnames'
 import FilePreview from '../file-preview'
@@ -13,6 +14,7 @@ import { DataSourceType } from '@/models/datasets'
 import Button from '@/app/components/base/button'
 import { NotionPageSelector } from '@/app/components/base/notion-page-selector'
 import { useDatasetDetailContext } from '@/context/dataset-detail'
+import { fetchDocumentsLimit } from '@/service/common'
 
 type IStepOneProps = {
   datasetId?: string
@@ -61,6 +63,7 @@ const StepOne = ({
   notionPages = [],
   updateNotionPages,
 }: IStepOneProps) => {
+  const { data: limitsData } = useSWR('/datasets/limit', fetchDocumentsLimit)
   const { dataset } = useDatasetDetailContext()
   const [showModal, setShowModal] = useState(false)
   const [currentFile, setCurrentFile] = useState<File | undefined>()
@@ -151,7 +154,7 @@ const StepOne = ({
               </div>
             )
           }
-          {dataSourceType === DataSourceType.FILE && (
+          {dataSourceType === DataSourceType.FILE && limitsData && (
             <>
               <FileUploader
                 fileList={files}
@@ -160,6 +163,8 @@ const StepOne = ({
                 onFileListUpdate={updateFileList}
                 onFileUpdate={updateFile}
                 onPreview={updateCurrentFile}
+                countLimit={limitsData.documents_limit}
+                countUsed={limitsData.documents_count}
               />
               <Button disabled={nextDisabled} className={s.submitButton} type='primary' onClick={onStepChange}>{t('datasetCreation.stepOne.button')}</Button>
             </>
@@ -167,10 +172,16 @@ const StepOne = ({
           {dataSourceType === DataSourceType.NOTION && (
             <>
               {!hasConnection && <NotionConnector onSetting={onSetting} />}
-              {hasConnection && (
+              {hasConnection && limitsData && (
                 <>
                   <div className='mb-8 w-[640px]'>
-                    <NotionPageSelector value={notionPages.map(page => page.page_id)} onSelect={updateNotionPages} onPreview={updateCurrentPage} />
+                    <NotionPageSelector
+                      value={notionPages.map(page => page.page_id)}
+                      onSelect={updateNotionPages}
+                      onPreview={updateCurrentPage}
+                      countLimit={limitsData.documents_limit}
+                      countUsed={limitsData.documents_count}
+                    />
                   </div>
                   <Button disabled={!notionPages.length} className={s.submitButton} type='primary' onClick={onStepChange}>{t('datasetCreation.stepOne.button')}</Button>
                 </>

web/i18n/lang/dataset-creation.en.ts

@@ -49,6 +49,7 @@ const translation = {
       confirmButton: 'Create',
       failed: 'Creation failed',
     },
+    overCountLimit: 'All your documents have overed limit {{countLimit}}.',
   },
   stepTwo: {
     segmentation: 'Segmentation settings',

web/i18n/lang/dataset-creation.zh.ts

@@ -49,6 +49,7 @@ const translation = {
       confirmButton: '创建',
       failed: '创建失败',
     },
+    overCountLimit: '您的文件总数已超出限制 {{countLimit}}。',
   },
   stepTwo: {
     segmentation: '分段设置',

web/models/common.ts

@@ -173,3 +173,8 @@ export type FileUploadConfigResponse = {
   file_size_limit: number
   batch_count_limit: number
 }
+
+export type DocumentsLimitResponse = {
+  documents_count: number
+  documents_limit: number
+}

web/service/common.ts

@@ -2,6 +2,7 @@ import type { Fetcher } from 'swr'
 import { del, get, patch, post, put } from './base'
 import type {
   AccountIntegrate, CommonResponse, DataSourceNotion,
+  DocumentsLimitResponse,
   FileUploadConfigResponse,
   ICurrentWorkspace,
   IWorkspace, LangGeniusVersionResponse, Member,
@@ -179,3 +180,7 @@ export const submitFreeQuota: Fetcher<{ type: string; redirect_url?: string; res
 export const fetchFileUploadConfig: Fetcher<FileUploadConfigResponse, { url: string }> = ({ url }) => {
   return get(url) as Promise<FileUploadConfigResponse>
 }
+
+export const fetchDocumentsLimit: Fetcher<DocumentsLimitResponse, string> = (url) => {
+  return get(url) as Promise<DocumentsLimitResponse>
+}