feat: Use Tree to display the knowledge base list on the search page i…

…nfiniflow#2247

web/src/interfaces/database/knowledge.ts

@@ -21,6 +21,7 @@ export interface IKnowledge {
   update_date: string;
   update_time: number;
   vector_similarity_weight: number;
+  embd_id: string;
 }
 
 export interface Parserconfig {

web/src/pages/search/index.less

@@ -27,6 +27,7 @@
 
 .searchSide {
   position: relative;
+
   :global(.ant-layout-sider-children) {
     height: auto;
   }
@@ -42,9 +43,12 @@
     height: 100%;
   }
   .list {
+    padding-top: 10px;
     width: 100%;
-    height: calc(100vh - 152px);
+    // height: 100%;
+    height: calc(100vh - 76px);
     overflow: auto;
+    background-color: transparent;
     &::-webkit-scrollbar-track {
       background: transparent;
     }
@@ -53,7 +57,10 @@
     width: 100%;
   }
   .knowledgeName {
-    width: 130px;
+    width: 116px;
+  }
+  .embeddingId {
+    width: 170px;
   }
 }
 

web/src/pages/search/index.tsx

@@ -4,7 +4,10 @@ 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 {
+  useNextFetchKnowledgeList,
+  useSelectTestingResult,
+} from '@/hooks/knowledge-hooks';
 import { useGetPaginationWithRouter } from '@/hooks/logic-hooks';
 import { IReference } from '@/interfaces/database/chat';
 import {
@@ -35,6 +38,11 @@ const SearchPage = () => {
   const { t } = useTranslation();
   const [checkedList, setCheckedList] = useState<string[]>([]);
   const { chunks, total } = useSelectTestingResult();
+  const { list: knowledgeList } = useNextFetchKnowledgeList();
+  const checkedWithoutEmbeddingIdList = useMemo(() => {
+    return checkedList.filter((x) => knowledgeList.some((y) => y.id === x));
+  }, [checkedList, knowledgeList]);
+
   const {
     sendQuestion,
     handleClickRelatedQuestion,
@@ -50,7 +58,7 @@ const SearchPage = () => {
     loading,
     isFirstRender,
     selectedDocumentIds,
-  } = useSendQuestion(checkedList);
+  } = useSendQuestion(checkedWithoutEmbeddingIdList);
   const { visible, hideModal, documentId, selectedChunk, clickDocumentButton } =
     useClickDrawer();
   const imgUrl = useFetchBackgroundImage();
@@ -79,7 +87,7 @@ const SearchPage = () => {
       onSearch={sendQuestion}
       size="large"
       loading={sendingLoading}
-      disabled={checkedList.length === 0}
+      disabled={checkedWithoutEmbeddingIdList.length === 0}
       className={isFirstRender ? styles.globalInput : styles.partialInput}
     />
   );
@@ -92,7 +100,7 @@ const SearchPage = () => {
       >
         <SearchSidebar
           isFirstRender={isFirstRender}
-          checkedList={checkedList}
+          checkedList={checkedWithoutEmbeddingIdList}
           setCheckedList={setCheckedList}
         ></SearchSidebar>
         <Layout className={isFirstRender ? styles.mainLayout : ''}>

web/src/pages/search/sidebar.tsx

@@ -1,16 +1,15 @@
 import { useNextFetchKnowledgeList } from '@/hooks/knowledge-hooks';
 import { UserOutlined } from '@ant-design/icons';
-import type { CheckboxProps } from 'antd';
-import { Avatar, Checkbox, Layout, List, Space, Typography } from 'antd';
-import { CheckboxChangeEvent } from 'antd/es/checkbox';
-import { CheckboxValueType } from 'antd/es/checkbox/Group';
+import type { TreeDataNode, TreeProps } from 'antd';
+import { Avatar, Layout, Space, Spin, Tree, Typography } from 'antd';
 import classNames from 'classnames';
 import {
   Dispatch,
   SetStateAction,
   useCallback,
   useEffect,
   useMemo,
+  useState,
 } from 'react';
 
 import styles from './index.less';
@@ -29,30 +28,109 @@ const SearchSidebar = ({
   setCheckedList,
 }: IProps) => {
   const { list, loading } = useNextFetchKnowledgeList();
-  const ids = useMemo(() => list.map((x) => x.id), [list]);
 
-  const checkAll = list.length === checkedList.length;
+  const groupedList = useMemo(() => {
+    return list.reduce((pre: TreeDataNode[], cur) => {
+      const parentItem = pre.find((x) => x.key === cur.embd_id);
+      const childItem: TreeDataNode = {
+        title: cur.name,
+        key: cur.id,
+        isLeaf: true,
+      };
+      if (parentItem) {
+        parentItem.children?.push(childItem);
+      } else {
+        pre.push({
+          title: cur.embd_id,
+          key: cur.embd_id,
+          isLeaf: false,
+          children: [childItem],
+        });
+      }
 
-  const indeterminate =
-    checkedList.length > 0 && checkedList.length < list.length;
+      return pre;
+    }, []);
+  }, [list]);
 
-  const onChange = useCallback(
-    (list: CheckboxValueType[]) => {
-      setCheckedList(list as string[]);
-    },
-    [setCheckedList],
-  );
+  const [expandedKeys, setExpandedKeys] = useState<React.Key[]>([]);
+  const [selectedKeys, setSelectedKeys] = useState<React.Key[]>([]);
+  const [autoExpandParent, setAutoExpandParent] = useState<boolean>(true);
+
+  const onExpand: TreeProps['onExpand'] = (expandedKeysValue) => {
+    // if not set autoExpandParent to false, if children expanded, parent can not collapse.
+    // or, you can remove all expanded children keys.
+    setExpandedKeys(expandedKeysValue);
+    setAutoExpandParent(false);
+  };
+
+  const onCheck: TreeProps['onCheck'] = (checkedKeysValue, info) => {
+    console.log('onCheck', checkedKeysValue, info);
+    const currentCheckedKeysValue = checkedKeysValue as string[];
+
+    let nextSelectedKeysValue: string[] = [];
+    const { isLeaf, checked, key, children } = info.node;
+    if (isLeaf) {
+      const item = list.find((x) => x.id === key);
+      if (!checked) {
+        const embeddingIds = currentCheckedKeysValue
+          .filter((x) => list.some((y) => y.id === x))
+          .map((x) => list.find((y) => y.id === x)?.embd_id);
+
+        if (embeddingIds.some((x) => x !== item?.embd_id)) {
+          nextSelectedKeysValue = [key as string];
+        } else {
+          nextSelectedKeysValue = currentCheckedKeysValue;
+        }
+      } else {
+        nextSelectedKeysValue = currentCheckedKeysValue;
+      }
+    } else {
+      if (!checked) {
+        nextSelectedKeysValue = [
+          key as string,
+          ...(children?.map((x) => x.key as string) ?? []),
+        ];
+      } else {
+        nextSelectedKeysValue = [];
+      }
+    }
 
-  const onCheckAllChange: CheckboxProps['onChange'] = useCallback(
-    (e: CheckboxChangeEvent) => {
-      setCheckedList(e.target.checked ? ids : []);
+    setCheckedList(nextSelectedKeysValue);
+  };
+
+  const onSelect: TreeProps['onSelect'] = (selectedKeysValue, info) => {
+    console.log('onSelect', info);
+
+    setSelectedKeys(selectedKeysValue);
+  };
+
+  const renderTitle = useCallback(
+    (node: TreeDataNode) => {
+      const item = list.find((x) => x.id === node.key);
+      return (
+        <Space>
+          {node.isLeaf && (
+            <Avatar size={24} icon={<UserOutlined />} src={item?.avatar} />
+          )}
+          <Typography.Text
+            ellipsis={{ tooltip: node.title as string }}
+            className={node.isLeaf ? styles.knowledgeName : styles.embeddingId}
+          >
+            {node.title as string}
+          </Typography.Text>
+        </Space>
+      );
     },
-    [ids, setCheckedList],
+    [list],
   );
 
   useEffect(() => {
-    setCheckedList(ids);
-  }, [ids, setCheckedList]);
+    const firstGroup = groupedList[0]?.children?.map((x) => x.key as string);
+    if (firstGroup) {
+      setCheckedList(firstGroup);
+    }
+    setExpandedKeys(groupedList.map((x) => x.key));
+  }, [groupedList, setExpandedKeys, setCheckedList]);
 
   return (
     <Sider
@@ -62,41 +140,21 @@ const SearchSidebar = ({
       theme={'light'}
       width={240}
     >
-      <Checkbox
-        className={styles.modelForm}
-        indeterminate={indeterminate}
-        onChange={onCheckAllChange}
-        checked={checkAll}
-      >
-        All
-      </Checkbox>
-      <Checkbox.Group
-        className={styles.checkGroup}
-        onChange={onChange}
-        value={checkedList}
-      >
-        <List
-          bordered
-          dataSource={list}
+      <Spin spinning={loading}>
+        <Tree
           className={styles.list}
-          loading={loading}
-          renderItem={(item) => (
-            <List.Item>
-              <Checkbox value={item.id} className={styles.checkbox}>
-                <Space>
-                  <Avatar size={30} icon={<UserOutlined />} src={item.avatar} />
-                  <Typography.Text
-                    ellipsis={{ tooltip: item.name }}
-                    className={styles.knowledgeName}
-                  >
-                    {item.name}
-                  </Typography.Text>
-                </Space>
-              </Checkbox>
-            </List.Item>
-          )}
+          checkable
+          onExpand={onExpand}
+          expandedKeys={expandedKeys}
+          autoExpandParent={autoExpandParent}
+          onCheck={onCheck}
+          checkedKeys={checkedList}
+          onSelect={onSelect}
+          selectedKeys={selectedKeys}
+          treeData={groupedList}
+          titleRender={renderTitle}
         />
-      </Checkbox.Group>
+      </Spin>
     </Sider>
   );
 };