feat: Retrieval chunks by page #2247 (#2373)

### What problem does this PR solve? feat: Retrieval chunks by page #2247 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
infiniflow · Sep 11, 2024 · 8052cbc · 8052cbc
1 parent b0e0e1f
commit 8052cbc
Show file tree

Hide file tree

Showing 5 changed files with 99 additions and 18 deletions.
diff --git a/web/src/components/retrieval-documents/index.tsx b/web/src/components/retrieval-documents/index.tsx
@@ -3,19 +3,22 @@ import { Collapse, Flex, Space } from 'antd';
 import SelectFiles from './select-files';
 
 import { useSelectTestingResult } from '@/hooks/knowledge-hooks';
-import { useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import styles from './index.less';
 
 interface IProps {
-  selectedDocumentIdsLength?: number;
   onTesting(documentIds: string[]): void;
+  setSelectedDocumentIds(documentIds: string[]): void;
+  selectedDocumentIds: string[];
 }
 
-const RetrievalDocuments = ({ onTesting }: IProps) => {
+const RetrievalDocuments = ({
+  onTesting,
+  selectedDocumentIds,
+  setSelectedDocumentIds,
+}: IProps) => {
   const { t } = useTranslation();
   const { documents } = useSelectTestingResult();
-  const [selectedDocumentIds, setSelectedDocumentIds] = useState<string[]>([]);
 
   return (
     <Collapse

diff --git a/web/src/layouts/components/header/index.tsx b/web/src/layouts/components/header/index.tsx
@@ -4,7 +4,7 @@ import { ReactComponent as KnowledgeBaseIcon } from '@/assets/svg/knowledge-base
 import { useTranslate } from '@/hooks/common-hooks';
 import { useFetchAppConf } from '@/hooks/logic-hooks';
 import { useNavigateWithFromState } from '@/hooks/route-hook';
-import { MessageOutlined } from '@ant-design/icons';
+import { MessageOutlined, SearchOutlined } from '@ant-design/icons';
 import { Flex, Layout, Radio, Space, theme } from 'antd';
 import { useCallback, useMemo } from 'react';
 import { useLocation } from 'umi';
@@ -27,7 +27,7 @@ const RagHeader = () => {
     () => [
       { path: '/knowledge', name: t('knowledgeBase'), icon: KnowledgeBaseIcon },
       { path: '/chat', name: t('chat'), icon: MessageOutlined },
-      // { path: '/search', name: t('search'), icon: SearchOutlined },
+      { path: '/search', name: t('search'), icon: SearchOutlined },
       { path: '/flow', name: t('flow'), icon: GraphIcon },
       { path: '/file', name: t('fileManager'), icon: FileIcon },
     ],

diff --git a/web/src/pages/search/hooks.ts b/web/src/pages/search/hooks.ts
@@ -1,6 +1,9 @@
 import { useFetchMindMap, useFetchRelatedQuestions } from '@/hooks/chat-hooks';
 import { useTestChunkRetrieval } from '@/hooks/knowledge-hooks';
-import { useSendMessageWithSse } from '@/hooks/logic-hooks';
+import {
+  useGetPaginationWithRouter,
+  useSendMessageWithSse,
+} from '@/hooks/logic-hooks';
 import { IAnswer } from '@/interfaces/database/chat';
 import api from '@/utils/api';
 import { get, isEmpty, trim } from 'lodash';
@@ -20,6 +23,9 @@ export const useSendQuestion = (kbIds: string[]) => {
   } = useFetchMindMap();
   const [searchStr, setSearchStr] = useState<string>('');
   const [isFirstRender, setIsFirstRender] = useState(true);
+  const [selectedDocumentIds, setSelectedDocumentIds] = useState<string[]>([]);
+
+  const { pagination } = useGetPaginationWithRouter();
 
   const sendQuestion = useCallback(
     (question: string) => {
@@ -55,18 +61,20 @@ export const useSendQuestion = (kbIds: string[]) => {
   );
 
   const handleTestChunk = useCallback(
-    (documentIds: string[]) => {
+    (documentIds: string[], page: number = 1, size: number = 10) => {
       const q = trim(searchStr);
       if (sendingLoading || isEmpty(q)) return;
 
       testChunk({
         kb_id: kbIds,
         highlight: true,
         question: q,
-        doc_ids: Array.isArray(documentIds) ? documentIds : [],
+        doc_ids: documentIds ?? selectedDocumentIds,
+        page,
+        size,
       });
     },
-    [sendingLoading, searchStr, kbIds, testChunk],
+    [sendingLoading, searchStr, kbIds, testChunk, selectedDocumentIds],
   );
 
   useEffect(() => {
@@ -89,6 +97,7 @@ export const useSendQuestion = (kbIds: string[]) => {
     handleSearchStrChange,
     handleClickRelatedQuestion,
     handleTestChunk,
+    setSelectedDocumentIds,
     loading,
     sendingLoading,
     answer: currentAnswer,
@@ -97,6 +106,7 @@ export const useSendQuestion = (kbIds: string[]) => {
     mindMapLoading,
     searchStr,
     isFirstRender,
+    selectedDocumentIds,
   };
 };
 
@@ -124,3 +134,45 @@ export const useFetchBackgroundImage = () => {
 
   return `https://cn.bing.com${imgUrl}`;
 };
+
+export const useTestRetrieval = (
+  kbIds: string[],
+  searchStr: string,
+  sendingLoading: boolean,
+) => {
+  const { testChunk, loading } = useTestChunkRetrieval();
+  const { pagination } = useGetPaginationWithRouter();
+
+  const [selectedDocumentIds, setSelectedDocumentIds] = useState<string[]>([]);
+
+  const handleTestChunk = useCallback(() => {
+    const q = trim(searchStr);
+    if (sendingLoading || isEmpty(q)) return;
+
+    testChunk({
+      kb_id: kbIds,
+      highlight: true,
+      question: q,
+      doc_ids: Array.isArray(selectedDocumentIds) ? selectedDocumentIds : [],
+      page: pagination.current,
+      size: pagination.pageSize,
+    });
+  }, [
+    sendingLoading,
+    searchStr,
+    kbIds,
+    testChunk,
+    selectedDocumentIds,
+    pagination,
+  ]);
+
+  useEffect(() => {
+    handleTestChunk();
+  }, [handleTestChunk]);
+
+  return {
+    loading,
+    selectedDocumentIds,
+    setSelectedDocumentIds,
+  };
+};
diff --git a/web/src/pages/search/index.less b/web/src/pages/search/index.less
@@ -62,6 +62,10 @@
     // background-color: aqua;
     overflow: auto;
     padding: 20px 10px 10px;
+    .chunks {
+      // overflow: auto;
+      // height: 60vh;
+    }
   }
 
   .graph {
@@ -100,6 +104,9 @@
 
 .globalInput {
   width: 600px;
+  position: sticky;
+  top: 0;
+  z-index: 1;
   .input();
 }
 .partialInput {

diff --git a/web/src/pages/search/index.tsx b/web/src/pages/search/index.tsx
@@ -1,7 +1,11 @@
 import HightLightMarkdown from '@/components/highlight-markdown';
 import { ImageWithPopover } from '@/components/image';
 import IndentedTree from '@/components/indented-tree/indented-tree';
+import PdfDrawer from '@/components/pdf-drawer';
+import { useClickDrawer } from '@/components/pdf-drawer/hooks';
+import RetrievalDocuments from '@/components/retrieval-documents';
 import { useSelectTestingResult } from '@/hooks/knowledge-hooks';
+import { useGetPaginationWithRouter } from '@/hooks/logic-hooks';
 import { IReference } from '@/interfaces/database/chat';
 import {
   Card,
@@ -10,19 +14,18 @@ import {
   Input,
   Layout,
   List,
+  Pagination,
+  PaginationProps,
   Skeleton,
   Space,
   Tag,
 } from 'antd';
 import { useState } from 'react';
+import { useTranslation } from 'react-i18next';
 import MarkdownContent from '../chat/markdown-content';
 import { useFetchBackgroundImage, useSendQuestion } from './hooks';
 import SearchSidebar from './sidebar';
 
-import PdfDrawer from '@/components/pdf-drawer';
-import { useClickDrawer } from '@/components/pdf-drawer/hooks';
-import RetrievalDocuments from '@/components/retrieval-documents';
-import { useTranslation } from 'react-i18next';
 import styles from './index.less';
 
 const { Content } = Layout;
@@ -31,13 +34,14 @@ const { Search } = Input;
 const SearchPage = () => {
   const { t } = useTranslation();
   const [checkedList, setCheckedList] = useState<string[]>([]);
-  const list = useSelectTestingResult();
+  const { chunks, total } = useSelectTestingResult();
   // const appConf = useFetchAppConf();
   const {
     sendQuestion,
     handleClickRelatedQuestion,
     handleSearchStrChange,
     handleTestChunk,
+    setSelectedDocumentIds,
     answer,
     sendingLoading,
     relatedQuestions,
@@ -46,10 +50,17 @@ const SearchPage = () => {
     searchStr,
     loading,
     isFirstRender,
+    selectedDocumentIds,
   } = useSendQuestion(checkedList);
   const { visible, hideModal, documentId, selectedChunk, clickDocumentButton } =
     useClickDrawer();
   const imgUrl = useFetchBackgroundImage();
+  const { pagination } = useGetPaginationWithRouter();
+
+  const onChange: PaginationProps['onChange'] = (pageNumber, pageSize) => {
+    pagination.onChange?.(pageNumber, pageSize);
+    handleTestChunk(selectedDocumentIds, pageNumber, pageSize);
+  };
 
   const InputSearch = (
     <Search
@@ -106,14 +117,16 @@ const SearchPage = () => {
                   )}
                   <Divider></Divider>
                   <RetrievalDocuments
-                    selectedDocumentIdsLength={0}
+                    selectedDocumentIds={selectedDocumentIds}
+                    setSelectedDocumentIds={setSelectedDocumentIds}
                     onTesting={handleTestChunk}
                   ></RetrievalDocuments>
                   <Divider></Divider>
-                  {list.chunks.length > 0 && (
+                  {chunks.length > 0 && (
                     <List
-                      dataSource={list.chunks}
+                      dataSource={chunks}
                       loading={loading}
+                      className={styles.chunks}
                       renderItem={(item) => (
                         <List.Item>
                           <Card className={styles.card}>
@@ -145,6 +158,12 @@ const SearchPage = () => {
                       </Flex>
                     </Card>
                   )}
+                  <Divider></Divider>
+                  <Pagination
+                    {...pagination}
+                    total={total}
+                    onChange={onChange}
+                  />
                 </section>
                 <section className={styles.graph}>
                   {mindMapLoading ? (