feat: Search for the answers you want based on the selected knowledge base #2247 (#2287)

cike8899 · web-flow · commit 925dd2aa85dd · 2024-09-06T15:42:55.000+08:00
### What problem does this PR solve? feat: Search for the answers you want based on the selected knowledge base #2247 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
diff --git a/web/src/components/message-item/index.tsx b/web/src/components/message-item/index.tsx
@@ -1,6 +1,6 @@
 import { ReactComponent as AssistantIcon } from '@/assets/svg/assistant.svg';
 import { MessageType } from '@/constants/chat';
-import { useSetModalState, useTranslate } from '@/hooks/common-hooks';
+import { useSetModalState } from '@/hooks/common-hooks';
 import { useSelectFileThumbnails } from '@/hooks/knowledge-hooks';
 import { IReference } from '@/interfaces/database/chat';
 import { IChunk } from '@/interfaces/database/knowledge';
@@ -50,7 +50,6 @@ const MessageItem = ({
 }: IProps) => {
   const isAssistant = item.role === MessageType.Assistant;
   const isUser = item.role === MessageType.User;
-  const { t } = useTranslate('chat');
   const fileThumbnails = useSelectFileThumbnails();
   const { data: documentList, setDocumentIds } = useFetchDocumentInfosByIds();
   const { data: documentThumbnails, setDocumentIds: setIds } =
@@ -62,14 +61,6 @@ const MessageItem = ({
     return reference?.doc_aggs ?? [];
   }, [reference?.doc_aggs]);
 
-  const content = useMemo(() => {
-    let text = item.content;
-    if (text === '') {
-      text = t('searching');
-    }
-    return loading ? text?.concat('~~2$$') : text;
-  }, [item.content, loading, t]);
-
   const handleUserDocumentClick = useCallback(
     (id: string) => () => {
       setClickedDocumentId(id);
@@ -154,7 +145,8 @@ const MessageItem = ({
               }
             >
               <MarkdownContent
-                content={content}
+                loading={loading}
+                content={item.content}
                 reference={reference}
                 clickDocumentButton={clickDocumentButton}
               ></MarkdownContent>
diff --git a/web/src/hooks/knowledge-hooks.ts b/web/src/hooks/knowledge-hooks.ts
@@ -206,6 +206,7 @@ export const useTestChunkRetrieval = (): ResponsePostType<ITestingResult> & {
     mutateAsync,
   } = useMutation({
     mutationKey: ['testChunk'], // This method is invalid
+    gcTime: 0,
     mutationFn: async (values: any) => {
       const { data } = await kbService.retrieval_test({
         ...values,
diff --git a/web/src/hooks/logic-hooks.ts b/web/src/hooks/logic-hooks.ts
@@ -297,69 +297,6 @@ export const useSpeechWithSse = (url: string = api.tts) => {
   return { read };
 };
 
-export const useFetchAudioWithSse = (url: string = api.tts) => {
-  // const [answer, setAnswer] = useState<IAnswer>({} as IAnswer);
-  const [done, setDone] = useState(true);
-
-  const read = useCallback(
-    async (
-      body: any,
-    ): Promise<{ response: Response; data: ResponseType } | undefined> => {
-      try {
-        setDone(false);
-        const response = await fetch(url, {
-          method: 'POST',
-          headers: {
-            [Authorization]: getAuthorization(),
-            'Content-Type': 'application/json',
-          },
-          body: JSON.stringify(body),
-        });
-
-        const res = response.clone().json();
-
-        const reader = response?.body?.getReader();
-
-        while (true) {
-          const x = await reader?.read();
-          if (x) {
-            const { done, value } = x;
-            try {
-              // const val = JSON.parse(value || '');
-              const val = value;
-              // const d = val?.data;
-              // if (typeof d !== 'boolean') {
-              // console.info('data:', d);
-              // setAnswer({
-              //   ...d,
-              //   conversationId: body?.conversation_id,
-              // });
-              // }
-            } catch (e) {
-              console.warn(e);
-            }
-            if (done) {
-              console.info('done');
-              break;
-            }
-          }
-        }
-        console.info('done?');
-        setDone(true);
-        // setAnswer({} as IAnswer);
-        return { data: await res, response };
-      } catch (e) {
-        setDone(true);
-        // setAnswer({} as IAnswer);
-        console.warn(e);
-      }
-    },
-    [url],
-  );
-
-  return { read, done, setDone };
-};
-
 //#region chat hooks
 
 export const useScrollToBottom = (messages?: unknown) => {
diff --git a/web/src/layouts/components/header/index.tsx b/web/src/layouts/components/header/index.tsx
@@ -9,7 +9,7 @@ import { useLocation } from 'umi';
 import Toolbar from '../right-toolbar';
 
 import { useFetchAppConf } from '@/hooks/logic-hooks';
-import { MessageOutlined } from '@ant-design/icons';
+import { MessageOutlined, SearchOutlined } from '@ant-design/icons';
 import styles from './index.less';
 
 const { Header } = Layout;
@@ -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/chat/markdown-content/index.tsx b/web/src/pages/chat/markdown-content/index.tsx
@@ -7,13 +7,14 @@ import { getExtension } from '@/utils/document-util';
 import { InfoCircleOutlined } from '@ant-design/icons';
 import { Button, Flex, Popover, Space } from 'antd';
 import DOMPurify from 'dompurify';
-import { useCallback } from 'react';
+import { useCallback, useMemo } from 'react';
 import Markdown from 'react-markdown';
 import reactStringReplace from 'react-string-replace';
 import SyntaxHighlighter from 'react-syntax-highlighter';
 import remarkGfm from 'remark-gfm';
 import { visitParents } from 'unist-util-visit-parents';
 
+import { useTranslation } from 'react-i18next';
 import styles from './index.less';
 
 const reg = /(#{2}\d+\${2})/g;
@@ -25,11 +26,22 @@ const MarkdownContent = ({
   reference,
   clickDocumentButton,
   content,
+  loading,
 }: {
   content: string;
+  loading: boolean;
   reference: IReference;
   clickDocumentButton?: (documentId: string, chunk: IChunk) => void;
 }) => {
+  const { t } = useTranslation();
+  const contentWithCursor = useMemo(() => {
+    let text = content;
+    if (text === '') {
+      text = t('chat.searching');
+    }
+    return loading ? text?.concat('~~2$$') : text;
+  }, [content, loading, t]);
+
   const fileThumbnails = useSelectFileThumbnails();
 
   const handleDocumentButtonClick = useCallback(
@@ -173,7 +185,7 @@ const MarkdownContent = ({
         } as any
       }
     >
-      {content}
+      {contentWithCursor}
     </Markdown>
   );
 };
diff --git a/web/src/pages/search/hooks.ts b/web/src/pages/search/hooks.ts
@@ -1,38 +1,37 @@
-import { MessageType } from '@/constants/chat';
 import { useTestChunkRetrieval } from '@/hooks/knowledge-hooks';
 import { useSendMessageWithSse } from '@/hooks/logic-hooks';
+import { IAnswer } from '@/interfaces/database/chat';
 import api from '@/utils/api';
-import { useCallback, useEffect, useMemo, useState } from 'react';
-import { IMessage } from '../chat/interface';
+import { isEmpty } from 'lodash';
+import { useCallback, useEffect, useState } from 'react';
 
 export const useSendQuestion = (kbIds: string[]) => {
   const { send, answer, done } = useSendMessageWithSse(api.ask);
   const { testChunk, loading } = useTestChunkRetrieval();
   const [sendingLoading, setSendingLoading] = useState(false);
-
-  const message: IMessage = useMemo(() => {
-    return {
-      id: '',
-      content: answer.answer,
-      role: MessageType.Assistant,
-      reference: answer.reference,
-    };
-  }, [answer]);
+  const [currentAnswer, setCurrentAnswer] = useState({} as IAnswer);
 
   const sendQuestion = useCallback(
     (question: string) => {
+      setCurrentAnswer({} as IAnswer);
       setSendingLoading(true);
       send({ kb_ids: kbIds, question });
       testChunk({ kb_id: kbIds, highlight: true, question });
     },
     [send, testChunk, kbIds],
   );
 
+  useEffect(() => {
+    if (!isEmpty(answer)) {
+      setCurrentAnswer(answer);
+    }
+  }, [answer]);
+
   useEffect(() => {
     if (done) {
       setSendingLoading(false);
     }
   }, [done]);
 
-  return { sendQuestion, message, loading, sendingLoading };
+  return { sendQuestion, loading, sendingLoading, answer: currentAnswer };
 };
diff --git a/web/src/pages/search/index.less b/web/src/pages/search/index.less
@@ -1,5 +1,7 @@
 .searchPage {
-  // height: 100%;
+  .card {
+    width: 100%;
+  }
 }
 
 .searchSide {
diff --git a/web/src/pages/search/index.tsx b/web/src/pages/search/index.tsx
@@ -1,10 +1,10 @@
 import HightLightMarkdown from '@/components/highlight-markdown';
 import { ImageWithPopover } from '@/components/image';
-import MessageItem from '@/components/message-item';
 import { useSelectTestingResult } from '@/hooks/knowledge-hooks';
 import { IReference } from '@/interfaces/database/chat';
 import { Card, Flex, Input, Layout, List, Space } from 'antd';
 import { useState } from 'react';
+import MarkdownContent from '../chat/markdown-content';
 import { useSendQuestion } from './hooks';
 import SearchSidebar from './sidebar';
 
@@ -16,8 +16,7 @@ const { Search } = Input;
 const SearchPage = () => {
   const [checkedList, setCheckedList] = useState<string[]>([]);
   const list = useSelectTestingResult();
-  const { sendQuestion, message, sendingLoading } =
-    useSendQuestion(checkedList);
+  const { sendQuestion, answer, sendingLoading } = useSendQuestion(checkedList);
 
   return (
     <Layout className={styles.searchPage}>
@@ -33,19 +32,20 @@ const SearchPage = () => {
                 placeholder="input search text"
                 onSearch={sendQuestion}
                 size="large"
+                loading={sendingLoading}
+                disabled={checkedList.length === 0}
               />
-              <MessageItem
-                item={message}
-                nickname="You"
-                reference={message.reference ?? ({} as IReference)}
+              <MarkdownContent
                 loading={sendingLoading}
-                index={0}
-              ></MessageItem>
+                content={answer.answer}
+                reference={answer.reference ?? ({} as IReference)}
+                clickDocumentButton={() => {}}
+              ></MarkdownContent>
               <List
                 dataSource={list.chunks}
                 renderItem={(item) => (
                   <List.Item>
-                    <Card>
+                    <Card className={styles.card}>
                       <Space>
                         <ImageWithPopover id={item.img_id}></ImageWithPopover>
                         <HightLightMarkdown>

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,7 @@`
`1`	`1`	`.searchPage {`
`2`		`- // height: 100%;`
	`2`	`+ .card {`
	`3`	`+ width: 100%;`
	`4`	`+ }`
`3`	`5`	`}`
`4`	`6`
`5`	`7`	`.searchSide {`