feat: add pages to ChunkMethodModal (infiniflow#143)

web/src/components/max-token-number.tsx

@@ -0,0 +1,32 @@
+import { Flex, Form, InputNumber, Slider } from 'antd';
+
+const MaxTokenNumber = () => {
+  return (
+    <Form.Item
+      label="Token number"
+      tooltip="It determine the token number of a chunk approximately."
+    >
+      <Flex gap={20} align="center">
+        <Flex flex={1}>
+          <Form.Item
+            name={['parser_config', 'chunk_token_num']}
+            noStyle
+            initialValue={128}
+            rules={[{ required: true, message: 'Province is required' }]}
+          >
+            <Slider max={2048} style={{ width: '100%' }} />
+          </Form.Item>
+        </Flex>
+        <Form.Item
+          name={['parser_config', 'chunk_token_num']}
+          noStyle
+          rules={[{ required: true, message: 'Street is required' }]}
+        >
+          <InputNumber max={2048} min={0} />
+        </Form.Item>
+      </Flex>
+    </Form.Item>
+  );
+};
+
+export default MaxTokenNumber;

web/src/hooks/documentHooks.ts

@@ -1,4 +1,5 @@
 import { IChunk, IKnowledgeFile } from '@/interfaces/database/knowledge';
+import { IChangeParserConfigRequestBody } from '@/interfaces/request/document';
 import { api_host } from '@/utils/api';
 import { buildChunkHighlights } from '@/utils/documentUtils';
 import { useCallback, useMemo } from 'react';
@@ -117,14 +118,19 @@ export const useSetDocumentParser = () => {
   const { knowledgeId } = useGetKnowledgeSearchParams();
 
   const setDocumentParser = useCallback(
-    (parserId: string, documentId: string) => {
+    (
+      parserId: string,
+      documentId: string,
+      parserConfig: IChangeParserConfigRequestBody,
+    ) => {
       try {
         return dispatch<any>({
           type: 'kFModel/document_change_parser',
           payload: {
             parser_id: parserId,
             doc_id: documentId,
             kb_id: knowledgeId,
+            parser_config: parserConfig,
           },
         });
       } catch (errorInfo) {

web/src/interfaces/database/knowledge.ts

@@ -28,6 +28,12 @@ export interface Parserconfig {
   to_page: number;
 }
 
+export interface IKnowledgeFileParserConfig {
+  chunk_token_num: number;
+  layout_recognize: boolean;
+  pages: number[][];
+  task_page_size: number;
+}
 export interface IKnowledgeFile {
   chunk_num: number;
   create_date: string;
@@ -51,6 +57,7 @@ export interface IKnowledgeFile {
   type: string;
   update_date: string;
   update_time: number;
+  parser_config: IKnowledgeFileParserConfig;
 }
 
 export interface ITenantInfo {

web/src/interfaces/request/document.ts

@@ -0,0 +1,12 @@
+export interface IChangeParserConfigRequestBody {
+  pages: number[][];
+  chunk_token_num: number;
+  layout_recognize: boolean;
+  task_page_size: number;
+}
+
+export interface IChangeParserRequestBody {
+  parser_id: string;
+  doc_id: string;
+  parser_config: IChangeParserConfigRequestBody;
+}

web/src/pages/add-knowledge/components/knowledge-file/chunk-method-modal.tsx

@@ -1,52 +1,98 @@
 import { IModalManagerChildrenProps } from '@/components/modal-manager';
 import {
-  useFetchTenantInfo,
-  useSelectParserList,
-} from '@/hooks/userSettingHook';
-import { Modal, Space, Tag } from 'antd';
-import React, { useEffect, useState } from 'react';
+  Button,
+  Divider,
+  Form,
+  InputNumber,
+  Modal,
+  Space,
+  Switch,
+  Tag,
+} from 'antd';
+import React, { useEffect, useMemo } from 'react';
+
+import MaxTokenNumber from '@/components/max-token-number';
+import { IKnowledgeFileParserConfig } from '@/interfaces/database/knowledge';
+import { IChangeParserConfigRequestBody } from '@/interfaces/request/document';
+import { MinusCircleOutlined, PlusOutlined } from '@ant-design/icons';
+import omit from 'lodash/omit';
+import {} from 'module';
+import { useFetchParserListOnMount } from './hooks';
 
 import styles from './index.less';
 
 const { CheckableTag } = Tag;
 
 interface IProps extends Omit<IModalManagerChildrenProps, 'showModal'> {
   loading: boolean;
-  onOk: (parserId: string) => void;
+  onOk: (
+    parserId: string,
+    parserConfig: IChangeParserConfigRequestBody,
+  ) => void;
   showModal?(): void;
-  parser_id: string;
+  parserId: string;
+  parserConfig: IKnowledgeFileParserConfig;
+  documentType: string;
 }
 
+const hidePagesChunkMethods = ['qa', 'table', 'picture', 'resume', 'one'];
+
 const ChunkMethodModal: React.FC<IProps> = ({
-  parser_id,
+  parserId,
   onOk,
   hideModal,
   visible,
+  documentType,
+  parserConfig,
 }) => {
-  const [selectedTag, setSelectedTag] = useState('');
-  const parserList = useSelectParserList();
-
-  useFetchTenantInfo();
-
-  useEffect(() => {
-    setSelectedTag(parser_id);
-  }, [parser_id]);
+  const { parserList, handleChange, selectedTag } =
+    useFetchParserListOnMount(parserId);
+  const [form] = Form.useForm();
 
   const handleOk = async () => {
-    onOk(selectedTag);
+    const values = await form.validateFields();
+    console.info(values);
+    const parser_config = {
+      ...values.parser_config,
+      pages: values.pages?.map((x: any) => [x.from, x.to]) ?? [],
+    };
+    console.info(parser_config);
+    onOk(selectedTag, parser_config);
   };
 
-  const handleChange = (tag: string, checked: boolean) => {
-    const nextSelectedTag = checked ? tag : selectedTag;
-    setSelectedTag(nextSelectedTag);
+  const showPages = useMemo(() => {
+    return (
+      documentType === 'pdf' &&
+      hidePagesChunkMethods.every((x) => x !== selectedTag)
+    );
+  }, [documentType, selectedTag]);
+
+  const showOne = useMemo(() => {
+    return showPages || selectedTag === 'one';
+  }, [showPages, selectedTag]);
+
+  const afterClose = () => {
+    form.resetFields();
   };
 
+  useEffect(() => {
+    if (visible) {
+      const pages =
+        parserConfig.pages?.map((x) => ({ from: x[0], to: x[1] })) ?? [];
+      form.setFieldsValue({
+        pages: pages.length > 0 ? pages : [{ from: 1, to: 1024 }],
+        parser_config: omit(parserConfig, 'pages'),
+      });
+    }
+  }, [form, parserConfig, visible]);
+
   return (
     <Modal
       title="Chunk Method"
       open={visible}
       onOk={handleOk}
       onCancel={hideModal}
+      afterClose={afterClose}
     >
       <Space size={[0, 8]} wrap>
         <div className={styles.tags}>
@@ -63,6 +109,138 @@ const ChunkMethodModal: React.FC<IProps> = ({
           })}
         </div>
       </Space>
+      <Divider></Divider>
+      {
+        <Form name="dynamic_form_nest_item" autoComplete="off" form={form}>
+          {showPages && (
+            <>
+              <Form.List name="pages">
+                {(fields, { add, remove }) => (
+                  <>
+                    {fields.map(({ key, name, ...restField }) => (
+                      <Space
+                        key={key}
+                        style={{
+                          display: 'flex',
+                        }}
+                        align="baseline"
+                      >
+                        <Form.Item
+                          {...restField}
+                          name={[name, 'from']}
+                          dependencies={name > 0 ? [name - 1, 'to'] : []}
+                          rules={[
+                            {
+                              required: true,
+                              message: 'Missing start page number',
+                            },
+                            ({ getFieldValue }) => ({
+                              validator(_, value) {
+                                if (
+                                  name === 0 ||
+                                  !value ||
+                                  getFieldValue(['pages', name - 1, 'to']) <
+                                    value
+                                ) {
+                                  return Promise.resolve();
+                                }
+                                return Promise.reject(
+                                  new Error(
+                                    'The current value must be greater than the previous to!',
+                                  ),
+                                );
+                              },
+                            }),
+                          ]}
+                        >
+                          <InputNumber
+                            placeholder="from"
+                            min={0}
+                            precision={0}
+                            className={styles.pageInputNumber}
+                          />
+                        </Form.Item>
+                        <Form.Item
+                          {...restField}
+                          name={[name, 'to']}
+                          dependencies={[name, 'from']}
+                          rules={[
+                            {
+                              required: true,
+                              message: 'Missing end page number(excluding)',
+                            },
+                            ({ getFieldValue }) => ({
+                              validator(_, value) {
+                                if (
+                                  !value ||
+                                  getFieldValue(['pages', name, 'from']) < value
+                                ) {
+                                  return Promise.resolve();
+                                }
+                                return Promise.reject(
+                                  new Error(
+                                    'The current value must be greater than to!',
+                                  ),
+                                );
+                              },
+                            }),
+                          ]}
+                        >
+                          <InputNumber
+                            placeholder="to"
+                            min={0}
+                            precision={0}
+                            className={styles.pageInputNumber}
+                          />
+                        </Form.Item>
+                        {name > 0 && (
+                          <MinusCircleOutlined onClick={() => remove(name)} />
+                        )}
+                      </Space>
+                    ))}
+                    <Form.Item>
+                      <Button
+                        type="dashed"
+                        onClick={() => add()}
+                        block
+                        icon={<PlusOutlined />}
+                      >
+                        Add page
+                      </Button>
+                    </Form.Item>
+                  </>
+                )}
+              </Form.List>
+              <Form.Item
+                name={['parser_config', 'task_page_size']}
+                label="Task page size"
+                tooltip={'coming soon'}
+                initialValue={2}
+                rules={[
+                  {
+                    required: true,
+                    message: 'Please input your task page size!',
+                  },
+                ]}
+              >
+                <InputNumber min={1} max={128} />
+              </Form.Item>
+            </>
+          )}
+          {showOne && (
+            <Form.Item
+              name={['parser_config', 'layout_recognize']}
+              label="Layout recognize"
+              initialValue={true}
+              valuePropName="checked"
+              tooltip={'coming soon'}
+            >
+              <Switch />
+            </Form.Item>
+          )}
+          {selectedTag === 'naive' && <MaxTokenNumber></MaxTokenNumber>}
+        </Form>
+      }
     </Modal>
   );
 };

web/src/pages/add-knowledge/components/knowledge-file/hooks.ts

@@ -7,9 +7,13 @@ import {
 } from '@/hooks/documentHooks';
 import { useGetKnowledgeSearchParams } from '@/hooks/routeHook';
 import { useOneNamespaceEffectsLoading } from '@/hooks/storeHooks';
-import { useFetchTenantInfo } from '@/hooks/userSettingHook';
+import {
+  useFetchTenantInfo,
+  useSelectParserList,
+} from '@/hooks/userSettingHook';
 import { Pagination } from '@/interfaces/common';
 import { IKnowledgeFile } from '@/interfaces/database/knowledge';
+import { IChangeParserConfigRequestBody } from '@/interfaces/request/document';
 import { PaginationProps } from 'antd';
 import { useCallback, useEffect, useMemo, useState } from 'react';
 import { useDispatch, useNavigate, useSelector } from 'umi';
@@ -222,8 +226,8 @@ export const useChangeDocumentParser = (documentId: string) => {
   ]);
 
   const onChangeParserOk = useCallback(
-    async (parserId: string) => {
-      const ret = await setDocumentParser(parserId, documentId);
+    async (parserId: string, parserConfig: IChangeParserConfigRequestBody) => {
+      const ret = await setDocumentParser(parserId, documentId, parserConfig);
       if (ret === 0) {
         hideChangeParserModal();
       }
@@ -239,3 +243,21 @@ export const useChangeDocumentParser = (documentId: string) => {
     showChangeParserModal,
   };
 };
+
+export const useFetchParserListOnMount = (parserId: string) => {
+  const [selectedTag, setSelectedTag] = useState('');
+  const parserList = useSelectParserList();
+
+  useFetchTenantInfo();
+
+  useEffect(() => {
+    setSelectedTag(parserId);
+  }, [parserId]);
+
+  const handleChange = (tag: string, checked: boolean) => {
+    const nextSelectedTag = checked ? tag : selectedTag;
+    setSelectedTag(nextSelectedTag);
+  };
+
+  return { parserList, handleChange, selectedTag };
+};