refactor: RAG Refactor (#985)

Co-authored-by: Aralhi <xiaoping0501@gmail.com>
Co-authored-by: csunny <cfqsunny@163.com>

web/components/knowledge/doc-panel.tsx

@@ -56,7 +56,7 @@ export default function DocPanel(props: IProps) {
       }),
     );
     setDocuments(data?.data);
-    setTotal(data?.total);
+    setTotal(data?.total || 0);
     setIsLoading(false);
   }
 
@@ -175,7 +175,7 @@ export default function DocPanel(props: IProps) {
                 >
                   <p className="mt-2 font-semibold ">{t('Size')}:</p>
                   <p>{document.chunk_size} chunks</p>
-                  <p className="mt-2 font-semibold ">{t('Last_Synch')}:</p>
+                  <p className="mt-2 font-semibold ">{t('Last_Sync')}:</p>
                   <p>{moment(document.last_sync).format('YYYY-MM-DD HH:MM:SS')}</p>
                   <p className="mt-2 mb-2">{renderResultTag(document.status, document.result)}</p>
                 </Card>

web/components/knowledge/doc-type-form.tsx

@@ -12,19 +12,19 @@ export default function DocTypeForm(props: IProps) {
   const { handleStepChange } = props;
   const docTypeList = [
     {
-      type: 'text',
+      type: 'TEXT',
       title: t('Text'),
       subTitle: t('Fill your raw text'),
       iconType: 'TEXT',
     },
     {
-      type: 'webPage',
+      type: 'URL',
       title: t('URL'),
       subTitle: t('Fetch_the_content_of_a_URL'),
       iconType: 'WEBPAGE',
     },
     {
-      type: 'file',
+      type: 'DOCUMENT',
       title: t('Document'),
       subTitle: t('Upload_a_document'),
       iconType: 'DOCUMENT',

web/components/knowledge/doc-upload-form.tsx

@@ -1,10 +1,12 @@
-import { Button, Form, Input, Switch, Upload, message, Spin } from 'antd';
+import { Button, Form, Input, Upload, Spin, message } from 'antd';
 import React, { useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { InboxOutlined } from '@ant-design/icons';
-import { apiInterceptors, addDocument, uploadDocument, syncDocument } from '@/client/api';
+import { apiInterceptors, addDocument, uploadDocument } from '@/client/api';
 import { RcFile, UploadChangeParam } from 'antd/es/upload';
-import { StepChangeParams } from '@/types/knowledge';
+import { File, StepChangeParams } from '@/types/knowledge';
+import { UploadRequestOption as RcCustomRequestOptions } from 'rc-upload/lib/interface';
+import classNames from 'classnames';
 
 type FileParams = {
   file: RcFile;
@@ -12,13 +14,13 @@ type FileParams = {
 };
 
 type IProps = {
+  className: string;
   handleStepChange: (params: StepChangeParams) => void;
   spaceName: string;
   docType: string;
 };
 
 type FieldType = {
-  synchChecked: boolean;
   docName: string;
   textSource: string;
   originFileObj: FileParams;
@@ -30,18 +32,19 @@ const { Dragger } = Upload;
 const { TextArea } = Input;
 
 export default function DocUploadForm(props: IProps) {
-  const { handleStepChange, spaceName, docType } = props;
+  const { className, handleStepChange, spaceName, docType } = props;
   const { t } = useTranslation();
   const [form] = Form.useForm();
   const [spinning, setSpinning] = useState<boolean>(false);
+  const [files, setFiles] = useState<Array<File>>([]);
 
-  const handleFinish = async (data: FieldType) => {
-    const { synchChecked, docName, textSource, originFileObj, text, webPageUrl } = data;
-    let res;
+  const upload = async (data: FieldType) => {
+    const { docName, textSource, text, webPageUrl } = data;
+    let docId;
     setSpinning(true);
     switch (docType) {
-      case 'webPage':
-        res = await apiInterceptors(
+      case 'URL':
+        [, docId] = await apiInterceptors(
           addDocument(spaceName as string, {
             doc_name: docName,
             content: webPageUrl,
@@ -49,16 +52,8 @@ export default function DocUploadForm(props: IProps) {
           }),
         );
         break;
-      case 'file':
-        const formData = new FormData();
-        formData.append('doc_name', docName || originFileObj.file.name);
-        formData.append('doc_file', originFileObj.file);
-        formData.append('doc_type', 'DOCUMENT');
-
-        res = await apiInterceptors(uploadDocument(spaceName as string, formData));
-        break;
-      default:
-        res = await apiInterceptors(
+      case 'TEXT':
+        [, docId] = await apiInterceptors(
           addDocument(spaceName as string, {
             doc_name: docName,
             source: textSource,
@@ -68,38 +63,60 @@ export default function DocUploadForm(props: IProps) {
         );
         break;
     }
-    synchChecked && handleSync?.(spaceName as string, res?.[1] as number);
     setSpinning(false);
-    handleStepChange({ label: 'finish' });
-  };
-
-  const handleSync = async (knowledgeName: string, id: number) => {
-    await apiInterceptors(syncDocument(knowledgeName, { doc_ids: [id] }));
+    if (docType === 'DOCUMENT' && files.length < 1) {
+      return message.error('Upload failed, please re-upload.');
+    } else if (docType !== 'DOCUMENT' && !docId) {
+      return message.error('Upload failed, please re-upload.');
+    }
+    handleStepChange({
+      label: 'forward',
+      files:
+        docType === 'DOCUMENT'
+          ? files
+          : [
+              {
+                name: docName,
+                doc_id: docId || -1,
+              },
+            ],
+    });
   };
 
   const handleFileChange = ({ file, fileList }: UploadChangeParam) => {
-    if (!form.getFieldsValue().docName) {
-      form.setFieldValue('docName', file.name);
-    }
     if (fileList.length === 0) {
       form.setFieldValue('originFileObj', null);
     }
   };
 
-  const beforeUpload = () => {
-    const curFile = form.getFieldsValue().originFileObj;
-    if (!curFile) {
-      return false;
+  const uploadFile = async (options: RcCustomRequestOptions) => {
+    const { onSuccess, onError, file } = options;
+    const formData = new FormData();
+    const filename = file?.name;
+    formData.append('doc_name', filename);
+    formData.append('doc_file', file);
+    formData.append('doc_type', 'DOCUMENT');
+    const [, docId] = await apiInterceptors(uploadDocument(spaceName, formData));
+    if (Number.isInteger(docId)) {
+      onSuccess && onSuccess(docId || 0);
+      setFiles((files) => {
+        files.push({
+          name: filename,
+          doc_id: docId || -1,
+        });
+        return files;
+      });
+    } else {
+      onError && onError({ name: '', message: '' });
     }
-    message.warning(t('Limit_Upload_File_Count_Tips'));
-    return Upload.LIST_IGNORE;
   };
 
-  const renderChooseType = () => {};
-
   const renderText = () => {
     return (
       <>
+        <Form.Item<FieldType> label={`${t('Name')}:`} name="docName" rules={[{ required: true, message: t('Please_input_the_name') }]}>
+          <Input className="mb-5 h-12" placeholder={t('Please_input_the_name')} />
+        </Form.Item>
         <Form.Item<FieldType>
           label={`${t('Text_Source')}:`}
           name="textSource"
@@ -107,7 +124,6 @@ export default function DocUploadForm(props: IProps) {
         >
           <Input className="mb-5  h-12" placeholder={t('Please_input_the_text_source')} />
         </Form.Item>
-
         <Form.Item<FieldType> label={`${t('Text')}:`} name="text" rules={[{ required: true, message: t('Please_input_the_description') }]}>
           <TextArea rows={4} />
         </Form.Item>
@@ -118,7 +134,14 @@ export default function DocUploadForm(props: IProps) {
   const renderWebPage = () => {
     return (
       <>
-        <Form.Item<FieldType> label={`${t('Web_Page_URL')}:`} name="webPageUrl" rules={[{ required: true, message: t('Please_input_the_owner') }]}>
+        <Form.Item<FieldType> label={`${t('Name')}:`} name="docName" rules={[{ required: true, message: t('Please_input_the_name') }]}>
+          <Input className="mb-5 h-12" placeholder={t('Please_input_the_name')} />
+        </Form.Item>
+        <Form.Item<FieldType>
+          label={`${t('Web_Page_URL')}:`}
+          name="webPageUrl"
+          rules={[{ required: true, message: t('Please_input_the_Web_Page_URL') }]}
+        >
           <Input className="mb-5  h-12" placeholder={t('Please_input_the_Web_Page_URL')} />
         </Form.Item>
       </>
@@ -128,8 +151,14 @@ export default function DocUploadForm(props: IProps) {
   const renderDocument = () => {
     return (
       <>
-        <Form.Item<FieldType> name="originFileObj" rules={[{ required: true, message: t('Please_input_the_owner') }]}>
-          <Dragger onChange={handleFileChange} beforeUpload={beforeUpload} multiple={false} accept=".pdf,.ppt,.pptx,.xls,.xlsx,.doc,.docx,.txt,.md">
+        <Form.Item<FieldType> name="originFileObj" rules={[{ required: true, message: t('Please_select_file') }]}>
+          <Dragger
+            multiple
+            onChange={handleFileChange}
+            maxCount={10}
+            accept=".pdf,.ppt,.pptx,.xls,.xlsx,.doc,.docx,.txt,.md"
+            customRequest={uploadFile}
+          >
             <p className="ant-upload-drag-icon">
               <InboxOutlined />
             </p>
@@ -145,9 +174,9 @@ export default function DocUploadForm(props: IProps) {
 
   const renderFormContainer = () => {
     switch (docType) {
-      case 'webPage':
+      case 'URL':
         return renderWebPage();
-      case 'file':
+      case 'DOCUMENT':
         return renderDocument();
       default:
         return renderText();
@@ -159,20 +188,14 @@ export default function DocUploadForm(props: IProps) {
       <Form
         form={form}
         size="large"
-        className="mt-4"
+        className={classNames('mt-4', className)}
         layout="vertical"
         name="basic"
         initialValues={{ remember: true }}
         autoComplete="off"
-        onFinish={handleFinish}
+        onFinish={upload}
       >
-        <Form.Item<FieldType> label={`${t('Name')}:`} name="docName" rules={[{ required: true, message: t('Please_input_the_name') }]}>
-          <Input className="mb-5 h-12" placeholder={t('Please_input_the_name')} />
-        </Form.Item>
         {renderFormContainer()}
-        <Form.Item<FieldType> label={`${t('Synch')}:`} name="synchChecked" initialValue={true}>
-          <Switch className="bg-slate-400" defaultChecked />
-        </Form.Item>
         <Form.Item>
           <Button
             onClick={() => {
@@ -180,8 +203,8 @@ export default function DocUploadForm(props: IProps) {
             }}
             className="mr-4"
           >{`${t('Back')}`}</Button>
-          <Button type="primary" htmlType="submit">
-            {t('Finish')}
+          <Button type="primary" loading={spinning} htmlType="submit">
+            {t('Next')}
           </Button>
         </Form.Item>
       </Form>

web/components/knowledge/segmentation.tsx

@@ -0,0 +1,185 @@
+import { apiInterceptors, getChunkStrategies, getDocumentList, syncBatchDocument } from '@/client/api';
+import { File, IChunkStrategyResponse, ISyncBatchParameter, StepChangeParams } from '@/types/knowledge';
+import { Alert, Button, Collapse, Form, Spin, message } from 'antd';
+import Icon from '@ant-design/icons';
+import { useEffect, useState } from 'react';
+import { useTranslation } from 'react-i18next';
+import StrategyForm from './strategy-form';
+import { DoneIcon, PendingIcon, SyncIcon, FileError } from '@/components/icons';
+
+type IProps = {
+  spaceName: string;
+  docType: string;
+  handleStepChange: (params: StepChangeParams) => void;
+  uploadFiles: Array<File>;
+};
+
+type FieldType = {
+  fileStrategies: Array<ISyncBatchParameter>;
+};
+
+let intervalId: string | number | NodeJS.Timeout | undefined;
+
+export default function Segmentation(props: IProps) {
+  const { spaceName, docType, uploadFiles, handleStepChange } = props;
+  const { t } = useTranslation();
+  const [form] = Form.useForm();
+  const [files, setFiles] = useState(uploadFiles);
+  const [loading, setLoading] = useState<boolean>();
+  const [strategies, setStrategies] = useState<Array<IChunkStrategyResponse>>([]);
+  const [syncStatus, setSyncStatus] = useState<string>('');
+
+  async function getStrategies() {
+    setLoading(true);
+    const [, allStrategies] = await apiInterceptors(getChunkStrategies());
+    setLoading(false);
+    setStrategies((allStrategies || [])?.filter((i) => i.type.indexOf(docType) > -1));
+  }
+
+  useEffect(() => {
+    getStrategies();
+    return () => {
+      intervalId && clearInterval(intervalId);
+    };
+  }, []);
+
+  const handleFinish = async (data: FieldType) => {
+    if (checkParameter(data)) {
+      setLoading(true);
+      const [, result] = await apiInterceptors(syncBatchDocument(spaceName, data.fileStrategies));
+      setLoading(false);
+      if (result?.tasks && result?.tasks?.length > 0) {
+        message.success(`Segemation task start successfully. task id: ${result?.tasks.join(',')}`);
+        setSyncStatus('RUNNING');
+        const docIds = data.fileStrategies.map((i) => i.doc_id);
+        intervalId = setInterval(async () => {
+          const status = await updateSyncStatus(docIds);
+          if (status === 'FINISHED') {
+            clearInterval(intervalId);
+            setSyncStatus('FINISHED');
+            message.success('Congratulation, All files sync successfully.');
+            handleStepChange({
+              label: 'finish',
+            });
+          }
+        }, 3000);
+      }
+    }
+  };
+
+  function checkParameter(data: FieldType) {
+    let checked = true;
+    if (syncStatus === 'RUNNING') {
+      checked = false;
+      message.warning('The task is still running, do not submit it again.');
+    }
+    const { fileStrategies } = data;
+    fileStrategies.map((item) => {
+      if (!item?.chunk_parameters?.chunk_strategy) {
+        message.error(`Please select chunk strategy for ${item.name}.`);
+        checked = false;
+      }
+    });
+    return checked;
+  }
+
+  async function updateSyncStatus(docIds: Array<number>) {
+    const [, docs] = await apiInterceptors(
+      getDocumentList(spaceName, {
+        doc_ids: docIds,
+      }),
+    );
+    if (docs?.data && docs?.data.length > 0) {
+      const copy = [...files!];
+      // set file status one by one
+      docs?.data.map((doc) => {
+        const file = copy?.filter((file) => file.doc_id === doc.id)?.[0];
+        if (file) {
+          file.status = doc.status;
+        }
+      });
+      setFiles(copy);
+      // all doc sync finished
+      if (docs?.data.every((item) => item.status === 'FINISHED' || item.status === 'FAILED')) {
+        return 'FINISHED';
+      }
+    }
+  }
+
+  function renderStrategy() {
+    if (!strategies || !strategies.length) {
+      return <Alert message={`Cannot find one strategy for ${docType} type knowledge.`} type="warning" />;
+    }
+    return (
+      <Form.List name="fileStrategies">
+        {(fields) => {
+          switch (docType) {
+            case 'TEXT':
+            case 'URL':
+              return fields?.map((field) => (
+                <StrategyForm strategies={strategies} docType={docType} fileName={files![field.name].name} field={field} />
+              ));
+            case 'DOCUMENT':
+              return (
+                <Collapse defaultActiveKey={0} size={files.length > 5 ? 'small' : 'middle'}>
+                  {fields?.map((field) => (
+                    // field [{name: 0, key: 0, isListField: true, fieldKey: 0}, {name: 1, key: 1, isListField: true, fieldKey: 1}]
+                    <Collapse.Panel header={`${field.name + 1}. ${files![field.name].name}`} key={field.key} extra={renderSyncStatus(field.name)}>
+                      <StrategyForm strategies={strategies} docType={docType} fileName={files![field.name].name} field={field} />
+                    </Collapse.Panel>
+                  ))}
+                </Collapse>
+              );
+          }
+        }}
+      </Form.List>
+    );
+  }
+
+  function renderSyncStatus(index: number) {
+    const status = files![index].status;
+    switch (status) {
+      case 'FINISHED':
+        return <Icon component={DoneIcon} />;
+      case 'RUNNING':
+        return <Icon className="animate-spin animate-infinite" component={SyncIcon} />;
+      case 'FAILED':
+        return <Icon component={FileError} />;
+      default:
+        return <Icon component={PendingIcon} />;
+    }
+  }
+
+  return (
+    <Spin spinning={loading}>
+      <Form
+        labelCol={{ span: 6 }}
+        wrapperCol={{ span: 18 }}
+        labelAlign="right"
+        form={form}
+        size="large"
+        className="mt-4"
+        layout="horizontal"
+        name="basic"
+        autoComplete="off"
+        initialValues={{
+          fileStrategies: files,
+        }}
+        onFinish={handleFinish}
+      >
+        {renderStrategy()}
+        <Form.Item className="mt-4">
+          <Button
+            onClick={() => {
+              handleStepChange({ label: 'back' });
+            }}
+            className="mr-4"
+          >{`${t('Back')}`}</Button>
+          <Button type="primary" htmlType="submit" loading={loading || syncStatus === 'RUNNING'}>
+            {t('Process')}
+          </Button>
+        </Form.Item>
+      </Form>
+    </Spin>
+  );
+}

web/components/knowledge/strategy-form.tsx

@@ -0,0 +1,79 @@
+import { IChunkStrategyResponse } from '@/types/knowledge';
+import { Alert, Form, FormListFieldData, Input, InputNumber, Radio, RadioChangeEvent } from 'antd';
+import { useState } from 'react';
+import { useTranslation } from 'react-i18next';
+const { TextArea } = Input;
+
+type IProps = {
+  strategies: Array<IChunkStrategyResponse>;
+  docType: string;
+  field: FormListFieldData;
+  fileName: string;
+};
+
+/**
+ * render strategies by doc type and file suffix
+ */
+export default function StrategyForm({ strategies, docType, fileName, field }: IProps) {
+  let filleSuffix = '';
+  if (docType === 'DOCUMENT') {
+    // filter strategy by file suffix
+    const arr = fileName.split('.');
+    filleSuffix = arr[arr.length - 1];
+  }
+  const ableStrategies = filleSuffix ? strategies.filter((i) => i.suffix.indexOf(filleSuffix) > -1) : strategies;
+  const [selectedStrategy, setSelectedStrategy] = useState<string>();
+  const { t } = useTranslation();
+  const DEFAULT_STRATEGY = {
+    strategy: t('Automatic'),
+    name: t('Automatic'),
+    desc: t('Automatic_desc'),
+  };
+
+  function radioChange(e: RadioChangeEvent) {
+    setSelectedStrategy(e.target.value);
+  }
+
+  function renderStrategyParamForm() {
+    if (!selectedStrategy) {
+      return null;
+    }
+    if (selectedStrategy === DEFAULT_STRATEGY.name) {
+      return <p className="my-4">{DEFAULT_STRATEGY.desc}</p>;
+    }
+    const parameters = ableStrategies?.filter((i) => i.strategy === selectedStrategy)[0].parameters;
+    if (!parameters || !parameters.length) {
+      return <Alert className="my-2" type="warning" message={t('No_parameter')} />;
+    }
+    return (
+      <div className="mt-2">
+        {parameters?.map((param) => (
+          <Form.Item
+            key={`param_${param.param_name}`}
+            label={`${param.param_name}: ${param.param_type}`}
+            name={[field!.name, 'chunk_parameters', param.param_name]}
+            rules={[{ required: true, message: t('Please_input_the_name') }]}
+            initialValue={param.default_value}
+          >
+            {param.param_type === 'int' ? <InputNumber className="w-full" min={1} /> : <TextArea className="w-full" rows={2} maxLength={6} />}
+          </Form.Item>
+        ))}
+      </div>
+    );
+  }
+  return (
+    <>
+      <Form.Item name={[field!.name, 'chunk_parameters', 'chunk_strategy']} initialValue={DEFAULT_STRATEGY.strategy}>
+        <Radio.Group style={{ marginTop: 16 }} onChange={radioChange}>
+          <Radio value={DEFAULT_STRATEGY.strategy}>{DEFAULT_STRATEGY.name}</Radio>
+          {ableStrategies.map((strategy) => (
+            <Radio key={`strategy_radio_${strategy.strategy}`} value={strategy.strategy}>
+              {strategy.name}
+            </Radio>
+          ))}
+        </Radio.Group>
+      </Form.Item>
+      {renderStrategyParamForm()}
+    </>
+  );
+}