fix: test chunk by @tanstack/react-query #1306 (#1719)

### What problem does this PR solve?

fix: test chunk by @tanstack/react-query #1306

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

web/src/hooks/file-manager-hooks.ts

@@ -7,7 +7,10 @@ import { PaginationProps, UploadFile, message } from 'antd';
 import React, { useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useSearchParams } from 'umi';
-import { useGetNextPagination, useHandleSearchChange } from './logic-hooks';
+import {
+  useGetPaginationWithRouter,
+  useHandleSearchChange,
+} from './logic-hooks';
 import { useSetPaginationParams } from './route-hook';
 
 export const useGetFolderId = () => {
@@ -27,7 +30,7 @@ export interface IListResult {
 
 export const useFetchFileList = (): ResponseType<any> & IListResult => {
   const { searchString, handleInputChange } = useHandleSearchChange();
-  const { pagination, setPagination } = useGetNextPagination();
+  const { pagination, setPagination } = useGetPaginationWithRouter();
   const id = useGetFolderId();
 
   const { data, isFetching: loading } = useQuery({

web/src/hooks/knowledge-hooks.ts

@@ -1,11 +1,19 @@
 import { useShowDeleteConfirm } from '@/hooks/common-hooks';
-import { IKnowledge } from '@/interfaces/database/knowledge';
+import { ResponsePostType } from '@/interfaces/database/base';
+import { IKnowledge, ITestingResult } from '@/interfaces/database/knowledge';
 import i18n from '@/locales/config';
 import kbService from '@/services/knowledge-service';
-import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
+import {
+  useIsMutating,
+  useMutation,
+  useMutationState,
+  useQuery,
+  useQueryClient,
+} from '@tanstack/react-query';
 import { message } from 'antd';
 import { useCallback, useEffect } from 'react';
 import { useDispatch, useSearchParams, useSelector } from 'umi';
+import { useSetPaginationParams } from './route-hook';
 
 export const useKnowledgeBaseId = (): string => {
   const [searchParams] = useSearchParams();
@@ -217,41 +225,65 @@ export const useUpdateKnowledge = () => {
 
 //#region Retrieval testing
 
-export const useTestChunkRetrieval = () => {
-  const dispatch = useDispatch();
+export const useTestChunkRetrieval = (): ResponsePostType<ITestingResult> & {
+  testChunk: (...params: any[]) => void;
+} => {
   const knowledgeBaseId = useKnowledgeBaseId();
+  const { page, size: pageSize } = useSetPaginationParams();
 
-  const testChunk = useCallback(
-    (values: any) => {
-      dispatch({
-        type: 'testingModel/testDocumentChunk',
-        payload: {
-          ...values,
-          kb_id: knowledgeBaseId,
-        },
-      });
-    },
-    [dispatch, knowledgeBaseId],
-  );
-
-  return testChunk;
-};
-
-export const useTestNextChunkRetrieval = () => {
   const {
     data,
     isPending: loading,
     mutateAsync,
   } = useMutation({
-    mutationKey: ['testChunk'],
-    mutationFn: async (canvasIds: string[]) => {
-      const { data } = await kbService.retrieval_test({ canvasIds });
+    mutationKey: ['testChunk'], // This method is invalid
+    mutationFn: async (values: any) => {
+      const { data } = await kbService.retrieval_test({
+        ...values,
+        kb_id: knowledgeBaseId,
+        page,
+        size: pageSize,
+      });
       if (data.retcode === 0) {
+        const res = data.data;
+        return {
+          chunks: res.chunks,
+          documents: res.doc_aggs,
+          total: res.total,
+        };
       }
-      return data?.data ?? [];
+      return (
+        data?.data ?? {
+          chunks: [],
+          documents: [],
+          total: 0,
+        }
+      );
     },
   });
 
-  return { data, loading, testChunk: mutateAsync };
+  return {
+    data: data ?? { chunks: [], documents: [], total: 0 },
+    loading,
+    testChunk: mutateAsync,
+  };
+};
+
+export const useChunkIsTesting = () => {
+  return useIsMutating({ mutationKey: ['testChunk'] }) > 0;
+};
+
+export const useSelectTestingResult = (): ITestingResult => {
+  const data = useMutationState({
+    filters: { mutationKey: ['testChunk'] },
+    select: (mutation) => {
+      return mutation.state.data;
+    },
+  });
+  return (data.at(-1) ?? {
+    chunks: [],
+    documents: [],
+    total: 0,
+  }) as ITestingResult;
 };
 //#endregion

web/src/hooks/logic-hooks.ts

@@ -95,7 +95,7 @@ export const useChangeLanguage = () => {
   return changeLanguage;
 };
 
-export const useGetNextPagination = () => {
+export const useGetPaginationWithRouter = () => {
   const { t } = useTranslate('common');
   const {
     setPaginationParams,
@@ -136,29 +136,32 @@ export const useGetNextPagination = () => {
   };
 };
 
-export const useGetPagination = (
-  total: number,
-  page: number,
-  pageSize: number,
-  onPageChange: PaginationProps['onChange'],
-) => {
+export const useGetPagination = () => {
+  const [pagination, setPagination] = useState({ page: 1, pageSize: 10 });
   const { t } = useTranslate('common');
 
-  const pagination: PaginationProps = useMemo(() => {
+  const onPageChange: PaginationProps['onChange'] = useCallback(
+    (pageNumber: number, pageSize: number) => {
+      setPagination({ page: pageNumber, pageSize });
+    },
+    [],
+  );
+
+  const currentPagination: PaginationProps = useMemo(() => {
     return {
       showQuickJumper: true,
-      total,
+      total: 0,
       showSizeChanger: true,
-      current: page,
-      pageSize: pageSize,
+      current: pagination.page,
+      pageSize: pagination.pageSize,
       pageSizeOptions: [1, 2, 10, 20, 50, 100],
       onChange: onPageChange,
       showTotal: (total) => `${t('total')} ${total}`,
     };
-  }, [t, onPageChange, page, pageSize, total]);
+  }, [t, onPageChange, pagination]);
 
   return {
-    pagination,
+    pagination: currentPagination,
   };
 };
 

web/src/interfaces/database/base.ts

@@ -9,3 +9,9 @@ export interface ResponseGetType<T = any> {
   data: T;
   loading?: boolean;
 }
+
+export interface ResponsePostType<T = any> {
+  data: T;
+  loading?: boolean;
+  [key: string]: unknown;
+}

web/src/interfaces/database/knowledge.ts

@@ -107,6 +107,6 @@ export interface ITestingDocument {
 
 export interface ITestingResult {
   chunks: ITestingChunk[];
-  doc_aggs: Record<string, number>;
+  documents: ITestingDocument[];
   total: number;
 }

web/src/pages/add-knowledge/components/knowledge-testing/index.tsx

@@ -1,7 +1,5 @@
 import { useTestChunkRetrieval } from '@/hooks/knowledge-hooks';
 import { Flex, Form } from 'antd';
-import { useEffect } from 'react';
-import { useDispatch } from 'umi';
 import TestingControl from './testing-control';
 import TestingResult from './testing-result';
 
@@ -9,24 +7,17 @@ import styles from './index.less';
 
 const KnowledgeTesting = () => {
   const [form] = Form.useForm();
-  const testChunk = useTestChunkRetrieval();
+  const { testChunk } = useTestChunkRetrieval();
 
-  const dispatch = useDispatch();
-
-  const handleTesting = async () => {
+  const handleTesting = async (documentIds: string[] = []) => {
     const values = await form.validateFields();
     testChunk({
       ...values,
+      doc_ids: Array.isArray(documentIds) ? documentIds : [],
       vector_similarity_weight: 1 - values.vector_similarity_weight,
     });
   };
 
-  useEffect(() => {
-    return () => {
-      dispatch({ type: 'testingModel/reset' });
-    };
-  }, [dispatch]);
-
   return (
     <Flex className={styles.testingWrapper} gap={16}>
       <TestingControl

web/src/pages/add-knowledge/components/knowledge-testing/model.ts

@@ -1,72 +0,0 @@
-import { BaseState } from '@/interfaces/common';
-import {
-  ITestingChunk,
-  ITestingDocument,
-} from '@/interfaces/database/knowledge';
-import kbService from '@/services/knowledge-service';
-import { DvaModel } from 'umi';
-
-export interface TestingModelState extends Pick<BaseState, 'pagination'> {
-  chunks: ITestingChunk[];
-  documents: ITestingDocument[];
-  total: number;
-  selectedDocumentIds: string[] | undefined;
-}
-
-const initialState = {
-  chunks: [],
-  documents: [],
-  total: 0,
-  pagination: {
-    current: 1,
-    pageSize: 10,
-  },
-  selectedDocumentIds: undefined,
-};
-
-const model: DvaModel<TestingModelState> = {
-  namespace: 'testingModel',
-  state: initialState,
-  reducers: {
-    setChunksAndDocuments(state, { payload }) {
-      return {
-        ...state,
-        ...payload,
-      };
-    },
-    setPagination(state, { payload }) {
-      return { ...state, pagination: { ...state.pagination, ...payload } };
-    },
-    setSelectedDocumentIds(state, { payload }) {
-      return { ...state, selectedDocumentIds: payload };
-    },
-    reset() {
-      return initialState;
-    },
-  },
-  effects: {
-    *testDocumentChunk({ payload = {} }, { call, put, select }) {
-      const { pagination, selectedDocumentIds }: TestingModelState =
-        yield select((state: any) => state.testingModel);
-
-      const { data } = yield call(kbService.retrieval_test, {
-        ...payload,
-        doc_ids: selectedDocumentIds,
-        page: pagination.current,
-        size: pagination.pageSize,
-      });
-      const { retcode, data: res } = data;
-      if (retcode === 0) {
-        yield put({
-          type: 'setChunksAndDocuments',
-          payload: {
-            chunks: res.chunks,
-            documents: res.doc_aggs,
-            total: res.total,
-          },
-        });
-      }
-    },
-  },
-};
-export default model;

web/src/pages/add-knowledge/components/knowledge-testing/testing-control/index.tsx

@@ -1,10 +1,10 @@
 import Rerank from '@/components/rerank';
 import SimilaritySlider from '@/components/similarity-slider';
 import { useTranslate } from '@/hooks/common-hooks';
-import { useOneNamespaceEffectsLoading } from '@/hooks/store-hooks';
 import { Button, Card, Divider, Flex, Form, Input } from 'antd';
 import { FormInstance } from 'antd/lib';
 
+import { useChunkIsTesting } from '@/hooks/knowledge-hooks';
 import styles from './index.less';
 
 type FieldType = {
@@ -20,9 +20,7 @@ interface IProps {
 
 const TestingControl = ({ form, handleTesting }: IProps) => {
   const question = Form.useWatch('question', { form, preserve: true });
-  const loading = useOneNamespaceEffectsLoading('testingModel', [
-    'testDocumentChunk',
-  ]);
+  const loading = useChunkIsTesting();
   const { t } = useTranslate('knowledgeDetails');
 
   const buttonDisabled =

web/src/pages/add-knowledge/components/knowledge-testing/testing-result/index.tsx

@@ -12,10 +12,11 @@ import {
   Space,
 } from 'antd';
 import camelCase from 'lodash/camelCase';
-import { useDispatch, useSelector } from 'umi';
-import { TestingModelState } from '../model';
 import SelectFiles from './select-files';
 
+import { useSelectTestingResult } from '@/hooks/knowledge-hooks';
+import { useGetPaginationWithRouter } from '@/hooks/logic-hooks';
+import { useCallback, useState } from 'react';
 import styles from './index.less';
 
 const similarityList: Array<{ field: keyof ITestingChunk; label: string }> = [
@@ -41,29 +42,28 @@ const ChunkTitle = ({ item }: { item: ITestingChunk }) => {
 };
 
 interface IProps {
-  handleTesting: () => Promise<any>;
+  handleTesting: (documentIds?: string[]) => Promise<any>;
 }
 
 const TestingResult = ({ handleTesting }: IProps) => {
-  const {
-    documents,
-    chunks,
-    total,
-    pagination,
-    selectedDocumentIds,
-  }: TestingModelState = useSelector((state: any) => state.testingModel);
-  const dispatch = useDispatch();
+  const [selectedDocumentIds, setSelectedDocumentIds] = useState<string[]>([]);
+  const { documents, chunks, total } = useSelectTestingResult();
   const { t } = useTranslate('knowledgeDetails');
+  const { pagination, setPagination } = useGetPaginationWithRouter();
 
   const onChange: PaginationProps['onChange'] = (pageNumber, pageSize) => {
-    console.log('Page: ', pageNumber, pageSize);
-    dispatch({
-      type: 'testingModel/setPagination',
-      payload: { current: pageNumber, pageSize },
-    });
-    handleTesting();
+    pagination.onChange?.(pageNumber, pageSize);
+    handleTesting(selectedDocumentIds);
   };
 
+  const onTesting = useCallback(
+    (ids: string[]) => {
+      setPagination({ page: 1 });
+      handleTesting(ids);
+    },
+    [setPagination, handleTesting],
+  );
+
   return (
     <section className={styles.testingResultWrapper}>
       <Collapse
@@ -94,7 +94,10 @@ const TestingResult = ({ handleTesting }: IProps) => {
             ),
             children: (
               <div>
-                <SelectFiles handleTesting={handleTesting}></SelectFiles>
+                <SelectFiles
+                  setSelectedDocumentIds={setSelectedDocumentIds}
+                  handleTesting={onTesting}
+                ></SelectFiles>
               </div>
             ),
           },
@@ -128,12 +131,9 @@ const TestingResult = ({ handleTesting }: IProps) => {
         ))}
       </Flex>
       <Pagination
+        {...pagination}
         size={'small'}
-        showQuickJumper
-        current={pagination.current}
-        pageSize={pagination.pageSize}
         total={total}
-        showSizeChanger
         onChange={onChange}
       />
     </section>

web/src/pages/add-knowledge/components/knowledge-testing/testing-result/select-files.tsx

@@ -1,22 +1,19 @@
 import NewDocumentLink from '@/components/new-document-link';
 import { useTranslate } from '@/hooks/common-hooks';
+import { useSelectTestingResult } from '@/hooks/knowledge-hooks';
 import { ITestingDocument } from '@/interfaces/database/knowledge';
 import { EyeOutlined } from '@ant-design/icons';
 import { Button, Table, TableProps, Tooltip } from 'antd';
-import { useDispatch, useSelector } from 'umi';
 
 interface IProps {
-  handleTesting: () => Promise<any>;
+  handleTesting: (ids: string[]) => void;
+  setSelectedDocumentIds: (ids: string[]) => void;
 }
 
-const SelectFiles = ({ handleTesting }: IProps) => {
-  const documents: ITestingDocument[] = useSelector(
-    (state: any) => state.testingModel.documents,
-  );
+const SelectFiles = ({ setSelectedDocumentIds, handleTesting }: IProps) => {
+  const { documents } = useSelectTestingResult();
   const { t } = useTranslate('fileManager');
 
-  const dispatch = useDispatch();
-
   const columns: TableProps<ITestingDocument>['columns'] = [
     {
       title: 'Name',
@@ -53,11 +50,8 @@ const SelectFiles = ({ handleTesting }: IProps) => {
 
   const rowSelection = {
     onChange: (selectedRowKeys: React.Key[]) => {
-      dispatch({
-        type: 'testingModel/setSelectedDocumentIds',
-        payload: selectedRowKeys,
-      });
-      handleTesting();
+      handleTesting(selectedRowKeys as string[]);
+      setSelectedDocumentIds(selectedRowKeys as string[]);
     },
     getCheckboxProps: (record: ITestingDocument) => ({
       disabled: record.doc_name === 'Disabled User', // Column configuration not to be checked

web/typings.d.ts

@@ -1,6 +1,5 @@
 import { ChunkModelState } from '@/pages/add-knowledge/components/knowledge-chunk/model';
 import { KFModelState } from '@/pages/add-knowledge/components/knowledge-file/model';
-import { TestingModelState } from '@/pages/add-knowledge/components/knowledge-testing/model';
 import { ChatModelState } from '@/pages/chat/model';
 
 declare module 'lodash';
@@ -14,7 +13,6 @@ export interface RootState {
   chatModel: ChatModelState;
   kFModel: KFModelState;
   chunkModel: ChunkModelState;
-  testingModel: TestingModelState;
 }
 
 declare global {