feat: Add RetrievalDocuments to SearchPage #2247 (#2327)

### What problem does this PR solve? feat: Add RetrievalDocuments to SearchPage #2247 feat: Click on the link in the reference to display the pdf drawer #2247 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
infiniflow · Sep 9, 2024 · 42eeb38 · 42eeb38
1 parent 7241c73
commit 42eeb38
Show file tree

Hide file tree

Showing 14 changed files with 390 additions and 143 deletions.
diff --git a/web/src/components/pdf-drawer/hooks.ts b/web/src/components/pdf-drawer/hooks.ts
@@ -0,0 +1,27 @@
+import { useSetModalState } from '@/hooks/common-hooks';
+import { IChunk } from '@/interfaces/database/knowledge';
+import { useCallback, useState } from 'react';
+
+export const useClickDrawer = () => {
+  const { visible, showModal, hideModal } = useSetModalState();
+  const [selectedChunk, setSelectedChunk] = useState<IChunk>({} as IChunk);
+  const [documentId, setDocumentId] = useState<string>('');
+
+  const clickDocumentButton = useCallback(
+    (documentId: string, chunk: IChunk) => {
+      showModal();
+      setSelectedChunk(chunk);
+      setDocumentId(documentId);
+    },
+    [showModal],
+  );
+
+  return {
+    clickDocumentButton,
+    visible,
+    showModal,
+    hideModal,
+    selectedChunk,
+    documentId,
+  };
+};
diff --git a/web/src/components/pdf-drawer/index.tsx b/web/src/components/pdf-drawer/index.tsx
@@ -0,0 +1,33 @@
+import { IModalProps } from '@/interfaces/common';
+import { IChunk } from '@/interfaces/database/knowledge';
+import { Drawer } from 'antd';
+import DocumentPreviewer from '../pdf-previewer';
+
+interface IProps extends IModalProps<any> {
+  documentId: string;
+  chunk: IChunk;
+}
+
+export const PdfDrawer = ({
+  visible = false,
+  hideModal,
+  documentId,
+  chunk,
+}: IProps) => {
+  return (
+    <Drawer
+      title="Document Previewer"
+      onClose={hideModal}
+      open={visible}
+      width={'50vw'}
+    >
+      <DocumentPreviewer
+        documentId={documentId}
+        chunk={chunk}
+        visible={visible}
+      ></DocumentPreviewer>
+    </Drawer>
+  );
+};
+
+export default PdfDrawer;
diff --git a/web/src/components/retrieval-documents/index.less b/web/src/components/retrieval-documents/index.less
@@ -0,0 +1,11 @@
+.selectFilesCollapse {
+  :global(.ant-collapse-header) {
+    padding-left: 22px;
+  }
+  margin-bottom: 32px;
+  overflow-y: auto;
+}
+
+.selectFilesTitle {
+  padding-right: 10px;
+}
diff --git a/web/src/components/retrieval-documents/index.tsx b/web/src/components/retrieval-documents/index.tsx
@@ -0,0 +1,55 @@
+import { ReactComponent as SelectedFilesCollapseIcon } from '@/assets/svg/selected-files-collapse.svg';
+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;
+}
+
+const RetrievalDocuments = ({ onTesting }: IProps) => {
+  const { t } = useTranslation();
+  const { documents } = useSelectTestingResult();
+  const [selectedDocumentIds, setSelectedDocumentIds] = useState<string[]>([]);
+
+  return (
+    <Collapse
+      expandIcon={() => <SelectedFilesCollapseIcon></SelectedFilesCollapseIcon>}
+      className={styles.selectFilesCollapse}
+      items={[
+        {
+          key: '1',
+          label: (
+            <Flex
+              justify={'space-between'}
+              align="center"
+              className={styles.selectFilesTitle}
+            >
+              <Space>
+                <span>
+                  {selectedDocumentIds.length ?? 0}/{documents.length}
+                </span>
+                {t('knowledgeDetails.filesSelected')}
+              </Space>
+            </Flex>
+          ),
+          children: (
+            <div>
+              <SelectFiles
+                setSelectedDocumentIds={setSelectedDocumentIds}
+                handleTesting={onTesting}
+              ></SelectFiles>
+            </div>
+          ),
+        },
+      ]}
+    />
+  );
+};
+
+export default RetrievalDocuments;
diff --git a/web/src/components/retrieval-documents/select-files.tsx b/web/src/components/retrieval-documents/select-files.tsx
@@ -0,0 +1,73 @@
+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';
+
+interface IProps {
+  handleTesting: (ids: string[]) => void;
+  setSelectedDocumentIds: (ids: string[]) => void;
+}
+
+const SelectFiles = ({ setSelectedDocumentIds, handleTesting }: IProps) => {
+  const { documents } = useSelectTestingResult();
+  const { t } = useTranslate('fileManager');
+
+  const columns: TableProps<ITestingDocument>['columns'] = [
+    {
+      title: 'Name',
+      dataIndex: 'doc_name',
+      key: 'doc_name',
+      render: (text) => <p>{text}</p>,
+    },
+
+    {
+      title: 'Hits',
+      dataIndex: 'count',
+      key: 'count',
+      width: 80,
+    },
+    {
+      title: 'View',
+      key: 'view',
+      width: 50,
+      render: (_, { doc_id, doc_name }) => (
+        <NewDocumentLink
+          documentName={doc_name}
+          documentId={doc_id}
+          prefix="document"
+        >
+          <Tooltip title={t('preview')}>
+            <Button type="text">
+              <EyeOutlined size={20} />
+            </Button>
+          </Tooltip>
+        </NewDocumentLink>
+      ),
+    },
+  ];
+
+  const rowSelection = {
+    onChange: (selectedRowKeys: React.Key[]) => {
+      handleTesting(selectedRowKeys as string[]);
+      setSelectedDocumentIds(selectedRowKeys as string[]);
+    },
+    getCheckboxProps: (record: ITestingDocument) => ({
+      disabled: record.doc_name === 'Disabled User', // Column configuration not to be checked
+      name: record.doc_name,
+    }),
+  };
+
+  return (
+    <Table
+      columns={columns}
+      dataSource={documents}
+      showHeader={false}
+      rowSelection={rowSelection}
+      rowKey={'doc_id'}
+    />
+  );
+};
+
+export default SelectFiles;
diff --git a/web/src/locales/en.ts b/web/src/locales/en.ts
@@ -646,7 +646,7 @@ The above is the content you need to summarize.`,
       operation: 'operation',
       run: 'Run',
       save: 'Save',
-      title: 'Title:',
+      title: 'ID:',
       beginDescription: 'This is where the flow begins.',
       answerDescription: `A component that serves as the interface between human and bot, receiving user inputs and displaying the agent's responses.`,
       retrievalDescription: `A component that retrieves information from a specified knowledge base and returns 'Empty response' if no information is found. Ensure the correct knowledge base is selected.`,

diff --git a/web/src/locales/zh-traditional.ts b/web/src/locales/zh-traditional.ts
@@ -602,7 +602,7 @@ export default {
       operation: '操作',
       run: '運行',
       save: '儲存',
-      title: '標題：',
+      title: 'ID：',
 
       beginDescription: '這是流程開始的地方',
       answerDescription: `該組件用作機器人與人類之間的介面。它接收使用者的輸入並顯示機器人的計算結果。`,

diff --git a/web/src/locales/zh.ts b/web/src/locales/zh.ts
@@ -621,7 +621,7 @@ export default {
       operation: '操作',
       run: '运行',
       save: '保存',
-      title: '标题：',
+      title: 'ID：',
       beginDescription: '这是流程开始的地方',
       answerDescription: `该组件用作机器人与人类之间的接口。它接收用户的输入并显示机器人的计算结果。`,
       retrievalDescription: `此组件用于从知识库中检索相关信息。选择知识库。如果没有检索到任何内容，将返回“空响应”。`,

diff --git a/web/src/pages/chat/chat-container/index.tsx b/web/src/pages/chat/chat-container/index.tsx
@@ -1,9 +1,7 @@
 import MessageItem from '@/components/message-item';
-import DocumentPreviewer from '@/components/pdf-previewer';
 import { MessageType } from '@/constants/chat';
-import { Drawer, Flex, Spin } from 'antd';
+import { Flex, Spin } from 'antd';
 import {
-  useClickDrawer,
   useCreateConversationBeforeUploadDocument,
   useGetFileIcon,
   useGetSendButtonDisabled,
@@ -13,6 +11,8 @@ import {
 import { buildMessageItemReference } from '../utils';
 
 import MessageInput from '@/components/message-input';
+import PdfDrawer from '@/components/pdf-drawer';
+import { useClickDrawer } from '@/components/pdf-drawer/hooks';
 import {
   useFetchNextConversation,
   useGetChatSearchParams,
@@ -96,18 +96,12 @@ const ChatContainer = () => {
           }
         ></MessageInput>
       </Flex>
-      <Drawer
-        title="Document Previewer"
-        onClose={hideModal}
-        open={visible}
-        width={'50vw'}
-      >
-        <DocumentPreviewer
-          documentId={documentId}
-          chunk={selectedChunk}
-          visible={visible}
-        ></DocumentPreviewer>
-      </Drawer>
+      <PdfDrawer
+        visible={visible}
+        hideModal={hideModal}
+        documentId={documentId}
+        chunk={selectedChunk}
+      ></PdfDrawer>
     </>
   );
 };

diff --git a/web/src/pages/chat/hooks.ts b/web/src/pages/chat/hooks.ts
@@ -23,7 +23,6 @@ import {
   useSendMessageWithSse,
 } from '@/hooks/logic-hooks';
 import { IConversation, IDialog, Message } from '@/interfaces/database/chat';
-import { IChunk } from '@/interfaces/database/knowledge';
 import { getFileExtension } from '@/utils';
 import { useMutationState } from '@tanstack/react-query';
 import { get } from 'lodash';
@@ -545,30 +544,6 @@ export const useRenameConversation = () => {
   };
 };
 
-export const useClickDrawer = () => {
-  const { visible, showModal, hideModal } = useSetModalState();
-  const [selectedChunk, setSelectedChunk] = useState<IChunk>({} as IChunk);
-  const [documentId, setDocumentId] = useState<string>('');
-
-  const clickDocumentButton = useCallback(
-    (documentId: string, chunk: IChunk) => {
-      showModal();
-      setSelectedChunk(chunk);
-      setDocumentId(documentId);
-    },
-    [showModal],
-  );
-
-  return {
-    clickDocumentButton,
-    visible,
-    showModal,
-    hideModal,
-    selectedChunk,
-    documentId,
-  };
-};
-
 export const useGetSendButtonDisabled = () => {
   const { dialogId, conversationId } = useGetChatSearchParams();
 

diff --git a/web/src/pages/flow/chat/box.tsx b/web/src/pages/flow/chat/box.tsx
@@ -1,13 +1,14 @@
 import MessageItem from '@/components/message-item';
-import DocumentPreviewer from '@/components/pdf-previewer';
 import { MessageType } from '@/constants/chat';
 import { useTranslate } from '@/hooks/common-hooks';
-import { useClickDrawer, useGetFileIcon } from '@/pages/chat/hooks';
+import { useGetFileIcon } from '@/pages/chat/hooks';
 import { buildMessageItemReference } from '@/pages/chat/utils';
-import { Button, Drawer, Flex, Input, Spin } from 'antd';
+import { Button, Flex, Input, Spin } from 'antd';
 
 import { useSendNextMessage } from './hooks';
 
+import PdfDrawer from '@/components/pdf-drawer';
+import { useClickDrawer } from '@/components/pdf-drawer/hooks';
 import { useFetchUserInfo } from '@/hooks/user-setting-hooks';
 import styles from './index.less';
 
@@ -79,19 +80,12 @@ const FlowChatBox = () => {
           onChange={handleInputChange}
         />
       </Flex>
-      <Drawer
-        title="Document Previewer"
-        onClose={hideModal}
-        open={visible}
-        width={'50vw'}
-        mask={false}
-      >
-        <DocumentPreviewer
-          documentId={documentId}
-          chunk={selectedChunk}
-          visible={visible}
-        ></DocumentPreviewer>
-      </Drawer>
+      <PdfDrawer
+        visible={visible}
+        hideModal={hideModal}
+        documentId={documentId}
+        chunk={selectedChunk}
+      ></PdfDrawer>
     </>
   );
 };

diff --git a/web/src/pages/search/hooks.ts b/web/src/pages/search/hooks.ts
@@ -46,10 +46,27 @@ export const useSendQuestion = (kbIds: string[]) => {
 
   const handleClickRelatedQuestion = useCallback(
     (question: string) => () => {
+      if (sendingLoading) return;
+
       setSearchStr(question);
       sendQuestion(question);
     },
-    [sendQuestion],
+    [sendQuestion, sendingLoading],
+  );
+
+  const handleTestChunk = useCallback(
+    (documentIds: string[]) => {
+      const q = trim(searchStr);
+      if (sendingLoading || isEmpty(q)) return;
+
+      testChunk({
+        kb_id: kbIds,
+        highlight: true,
+        question: q,
+        doc_ids: Array.isArray(documentIds) ? documentIds : [],
+      });
+    },
+    [sendingLoading, searchStr, kbIds, testChunk],
   );
 
   useEffect(() => {
@@ -71,6 +88,7 @@ export const useSendQuestion = (kbIds: string[]) => {
     sendQuestion,
     handleSearchStrChange,
     handleClickRelatedQuestion,
+    handleTestChunk,
     loading,
     sendingLoading,
     answer: currentAnswer,