fix: mark modal cannot select folder (#327)

client/src/components/ChatBox/SelectDataset.tsx

@@ -1,58 +1,43 @@
-import React, { useRef, useState } from 'react';
+import React, { useState } from 'react';
-import {
+import { ModalBody, useTheme, ModalFooter, Button, Box, Card, Flex, Grid } from '@chakra-ui/react';
-  ModalBody,
-  useTheme,
-  ModalFooter,
-  Button,
-  ModalHeader,
-  Box,
-  Card,
-  Flex
-} from '@chakra-ui/react';
-import MyModal from '../MyModal';
 import { useTranslation } from 'next-i18next';
-import { useQuery } from '@tanstack/react-query';
-import { useDatasetStore } from '@/store/dataset';
 import { useToast } from '@/hooks/useToast';
 import Avatar from '../Avatar';
 import MyIcon from '@/components/Icon';
-import { useGlobalStore } from '@/store/global';
+import { KbTypeEnum } from '@/constants/kb';
+import DatasetSelectModal, { useDatasetSelect } from '@/components/core/dataset/SelectModal';
 
 const SelectDataset = ({
+  isOpen,
   onSuccess,
   onClose
 }: {
+  isOpen: boolean;
   onSuccess: (kbId: string) => void;
   onClose: () => void;
 }) => {
   const { t } = useTranslation();
   const theme = useTheme();
-  const { isPc } = useGlobalStore();
   const { toast } = useToast();
-  const { myKbList, loadKbList } = useDatasetStore();
   const [selectedId, setSelectedId] = useState<string>();
-
+  const { paths, parentId, setParentId, datasets } = useDatasetSelect();
-  useQuery(['loadKbList'], () => loadKbList());
 
   return (
-    <MyModal isOpen={true} onClose={onClose} w={'100%'} maxW={['90vw', '900px']} isCentered={!isPc}>
+    <DatasetSelectModal
-      <Flex flexDirection={'column'} h={['90vh', 'auto']}>
+      isOpen={isOpen}
-        <ModalHeader>
+      paths={paths}
-          <Box>{t('chat.Select Mark Kb')}</Box>
+      onClose={onClose}
-          <Box fontSize={'sm'} color={'myGray.500'} fontWeight={'normal'}>
+      parentId={parentId}
-            {t('chat.Select Mark Kb Desc')}
+      setParentId={setParentId}
-          </Box>
+      tips={t('chat.Select Mark Kb Desc')}
-        </ModalHeader>
+    >
-        <ModalBody
+      <ModalBody flex={['1 0 0', '0 0 auto']} maxH={'80vh'} overflowY={'auto'}>
-          flex={['1 0 0', '0 0 auto']}
+        <Grid
-          maxH={'80vh'}
-          overflowY={'auto'}
-          display={'grid'}
           gridTemplateColumns={['repeat(1,1fr)', 'repeat(2,1fr)', 'repeat(3,1fr)']}
           gridGap={3}
           userSelect={'none'}
         >
-          {myKbList.map((item) =>
+          {datasets.map((item) =>
             (() => {
               const selected = selectedId === item._id;
               return (
@@ -72,7 +57,11 @@ const SelectDataset = ({
                       }
                     : {})}
                   onClick={() => {
+                    if (item.type === KbTypeEnum.folder) {
+                      setParentId(item._id);
+                    } else {
                       setSelectedId(item._id);
+                    }
                   }}
                 >
                   <Flex alignItems={'center'} h={'38px'}>
@@ -89,6 +78,15 @@ const SelectDataset = ({
               );
             })()
           )}
+        </Grid>
+        {datasets.length === 0 && (
+          <Flex mt={5} flexDirection={'column'} alignItems={'center'}>
+            <MyIcon name="empty" w={'48px'} h={'48px'} color={'transparent'} />
+            <Box mt={2} color={'myGray.500'}>
+              这个目录已经没东西可选了~
+            </Box>
+          </Flex>
+        )}
       </ModalBody>
       <ModalFooter>
         <Button variant={'base'} mr={2} onClick={onClose}>
@@ -109,8 +107,7 @@ const SelectDataset = ({
           {t('Confirm')}
         </Button>
       </ModalFooter>
-      </Flex>
+    </DatasetSelectModal>
-    </MyModal>
   );
 };
 

client/src/components/ChatBox/index.tsx

@@ -985,13 +985,12 @@ const ChatBox = (
         />
       )}
       {/* select one dataset to insert markData */}
-      {adminMarkData && !adminMarkData.kbId && (
       <SelectDataset
+        isOpen={!!adminMarkData && !adminMarkData.kbId}
         onClose={() => setAdminMarkData(undefined)}
         // @ts-ignore
         onSuccess={(kbId) => setAdminMarkData((state) => ({ ...state, kbId }))}
       />
-      )}
       {/* edit markData modal */}
       {adminMarkData && adminMarkData.kbId && (
         <InputDataModal

client/src/components/Icon/index.tsx

@@ -101,7 +101,7 @@ const MyIcon = ({ name, w = 'auto', h = 'auto', ...props }: { name: IconName } &
       .catch((error) => console.log(error));
   }, [name]);
 
-  return name ? (
+  return !!name && !!iconPaths[name] ? (
     <Icon
       {...IconComponent}
       w={w}

client/src/components/core/dataset/SelectModal.tsx

@@ -0,0 +1,122 @@
+import { getKbList, getKbPaths } from '@/api/plugins/kb';
+import MyModal from '@/components/MyModal';
+import { useQuery } from '@tanstack/react-query';
+import React, { Dispatch, useMemo, useState } from 'react';
+import { useTranslation } from 'react-i18next';
+import { useGlobalStore } from '@/store/global';
+import { Box, Flex, ModalHeader } from '@chakra-ui/react';
+import MyIcon from '@/components/Icon';
+
+type PathItemType = {
+  parentId: string;
+  parentName: string;
+};
+
+const DatasetSelectModal = ({
+  isOpen,
+  parentId,
+  setParentId,
+  paths,
+  onClose,
+  tips,
+  children
+}: {
+  isOpen: boolean;
+  parentId?: string;
+  setParentId: Dispatch<string>;
+  paths: PathItemType[];
+  onClose: () => void;
+  tips?: string | null;
+  children: React.ReactNode;
+}) => {
+  const { t } = useTranslation();
+  const { isPc } = useGlobalStore();
+
+  return (
+    <MyModal
+      isOpen={isOpen}
+      onClose={onClose}
+      w={'100%'}
+      maxW={['90vw', '900px']}
+      isCentered={!isPc}
+    >
+      <Flex flexDirection={'column'} h={['90vh', 'auto']}>
+        <ModalHeader>
+          {!!parentId ? (
+            <Flex
+              flex={1}
+              userSelect={'none'}
+              fontSize={['sm', 'lg']}
+              fontWeight={'normal'}
+              color={'myGray.900'}
+            >
+              {paths.map((item, i) => (
+                <Flex key={item.parentId} mr={2} alignItems={'center'}>
+                  <Box
+                    fontSize={'lg'}
+                    borderRadius={'md'}
+                    {...(i === paths.length - 1
+                      ? {
+                          cursor: 'default'
+                        }
+                      : {
+                          cursor: 'pointer',
+                          _hover: {
+                            color: 'myBlue.600'
+                          },
+                          onClick: () => {
+                            setParentId(item.parentId);
+                          }
+                        })}
+                  >
+                    {item.parentName}
+                  </Box>
+                  {i !== paths.length - 1 && (
+                    <MyIcon name={'rightArrowLight'} color={'myGray.500'} w={['18px', '24px']} />
+                  )}
+                </Flex>
+              ))}
+            </Flex>
+          ) : (
+            <Box>{t('chat.Select Mark Kb')}</Box>
+          )}
+          {!!tips && (
+            <Box fontSize={'sm'} color={'myGray.500'} fontWeight={'normal'}>
+              {tips}
+            </Box>
+          )}
+        </ModalHeader>
+        {children}
+      </Flex>
+    </MyModal>
+  );
+};
+
+export const useDatasetSelect = () => {
+  const { t } = useTranslation();
+  const [parentId, setParentId] = useState<string>();
+
+  const { data } = useQuery(['loadDatasetData', parentId], () =>
+    Promise.all([getKbList({ parentId }), getKbPaths(parentId)])
+  );
+
+  const paths = useMemo(
+    () => [
+      {
+        parentId: '',
+        parentName: t('kb.My Dataset')
+      },
+      ...(data?.[1] || [])
+    ],
+    [data, t]
+  );
+
+  return {
+    parentId,
+    setParentId,
+    datasets: data?.[0] || [],
+    paths
+  };
+};
+
+export default DatasetSelectModal;

client/src/pages/app/detail/components/BasicEdit/index.tsx

@@ -71,7 +71,7 @@ const Settings = ({ appId }: { appId: string }) => {
   const { t } = useTranslation();
   const { toast } = useToast();
   const { appDetail, updateAppDetail } = useUserStore();
-  const { loadAllDatasets, datasets } = useDatasetStore();
+  const { loadAllDatasets, allDatasets } = useDatasetStore();
   const { isPc } = useGlobalStore();
 
   const [editVariable, setEditVariable] = useState<VariableItemType>();
@@ -131,8 +131,8 @@ const Settings = ({ appId }: { appId: string }) => {
     );
   }, [getValues, refresh]);
   const selectedKbList = useMemo(
-    () => datasets.filter((item) => kbList.find((kb) => kb.kbId === item._id)),
+    () => allDatasets.filter((item) => kbList.find((kb) => kb.kbId === item._id)),
-    [datasets, kbList]
+    [allDatasets, kbList]
   );
 
   /* 点击删除 */
@@ -564,8 +564,8 @@ const Settings = ({ appId }: { appId: string }) => {
           defaultData={getValues('chatModel')}
         />
       )}
-      {isOpenKbSelect && (
       <KBSelectModal
+        isOpen={isOpenKbSelect}
         activeKbs={selectedKbList.map((item) => ({
           kbId: item._id,
           vectorModel: item.vectorModel
@@ -573,7 +573,6 @@ const Settings = ({ appId }: { appId: string }) => {
         onClose={onCloseKbSelect}
         onChange={replaceKbList}
       />
-      )}
       {isOpenKbParams && (
         <KbParamsModal
           searchEmptyText={getValues('kb.searchEmptyText')}

client/src/pages/app/detail/components/KBSelectModal.tsx

@@ -5,19 +5,16 @@ import {
   Box,
   Button,
   ModalBody,
-  ModalHeader,
   ModalFooter,
   useTheme,
   Textarea,
   Grid,
   Divider
 } from '@chakra-ui/react';
-import { getKbPaths } from '@/api/plugins/kb';
 import Avatar from '@/components/Avatar';
 import { useForm } from 'react-hook-form';
 import { QuestionOutlineIcon } from '@chakra-ui/icons';
 import type { SelectedKbType } from '@/types/plugin';
-import { useGlobalStore } from '@/store/global';
 import { useToast } from '@/hooks/useToast';
 import MySlider from '@/components/Slider';
 import MyTooltip from '@/components/MyTooltip';
@@ -28,6 +25,7 @@ import { useTranslation } from 'react-i18next';
 import { useQuery } from '@tanstack/react-query';
 import { useDatasetStore } from '@/store/dataset';
 import { feConfigs } from '@/store/static';
+import DatasetSelectModal, { useDatasetSelect } from '@/components/core/dataset/SelectModal';
 
 export type KbParamsType = {
   searchSimilarity: number;
@@ -36,10 +34,12 @@ export type KbParamsType = {
 };
 
 export const KBSelectModal = ({
+  isOpen,
   activeKbs = [],
   onChange,
   onClose
 }: {
+  isOpen: boolean;
   activeKbs: SelectedKbType;
   onChange: (e: SelectedKbType) => void;
   onClose: () => void;
@@ -47,99 +47,30 @@ export const KBSelectModal = ({
   const { t } = useTranslation();
   const theme = useTheme();
   const [selectedKbList, setSelectedKbList] = useState<SelectedKbType>(activeKbs);
-  const { isPc } = useGlobalStore();
   const { toast } = useToast();
-  const [parentId, setParentId] = useState<string>();
+  const { paths, parentId, setParentId, datasets } = useDatasetSelect();
-  const { myKbList, loadKbList, datasets, loadAllDatasets } = useDatasetStore();
+  const { allDatasets, loadAllDatasets } = useDatasetStore();
 
-  const { data } = useQuery(['loadKbList', parentId], () => {
-    return Promise.all([loadKbList(parentId), getKbPaths(parentId)]);
-  });
   useQuery(['loadAllDatasets'], loadAllDatasets);
-  const paths = useMemo(
+
-    () => [
-      {
-        parentId: '',
-        parentName: t('kb.My Dataset')
-      },
-      ...(data?.[1] || [])
-    ],
-    [data, t]
-  );
   const filterKbList = useMemo(() => {
     return {
-      selected: datasets.filter((item) => selectedKbList.find((kb) => kb.kbId === item._id)),
+      selected: allDatasets.filter((item) => selectedKbList.find((kb) => kb.kbId === item._id)),
-      unSelected: myKbList.filter((item) => !selectedKbList.find((kb) => kb.kbId === item._id))
+      unSelected: datasets.filter((item) => !selectedKbList.find((kb) => kb.kbId === item._id))
     };
-  }, [myKbList, datasets, selectedKbList]);
+  }, [datasets, allDatasets, selectedKbList]);
 
   return (
-    <MyModal
+    <DatasetSelectModal
-      isOpen={true}
+      isOpen={isOpen}
-      isCentered={!isPc}
+      paths={paths}
-      maxW={['90vw', '800px']}
+      parentId={parentId}
-      w={'800px'}
+      setParentId={setParentId}
+      tips={'仅能选择同一个索引模型的知识库'}
       onClose={onClose}
     >
-      <Flex flexDirection={'column'} h={['90vh', 'auto']}>
+      <ModalBody flex={['1 0 0', '0 0 auto']} maxH={'80vh'} overflowY={'auto'} userSelect={'none'}>
-        <ModalHeader>
+        <Grid gridTemplateColumns={['repeat(1,1fr)', 'repeat(2,1fr)', 'repeat(3,1fr)']} gridGap={3}>
-          {!!parentId ? (
-            <Flex
-              flex={1}
-              userSelect={'none'}
-              fontSize={['sm', 'lg']}
-              fontWeight={'normal'}
-              color={'myGray.900'}
-            >
-              {paths.map((item, i) => (
-                <Flex key={item.parentId} mr={2} alignItems={'center'}>
-                  <Box
-                    fontSize={'lg'}
-                    borderRadius={'md'}
-                    {...(i === paths.length - 1
-                      ? {
-                          cursor: 'default'
-                        }
-                      : {
-                          cursor: 'pointer',
-                          _hover: {
-                            color: 'myBlue.600'
-                          },
-                          onClick: () => {
-                            setParentId(item.parentId);
-                          }
-                        })}
-                  >
-                    {item.parentName}
-                  </Box>
-                  {i !== paths.length - 1 && (
-                    <MyIcon name={'rightArrowLight'} color={'myGray.500'} w={['18px', '24px']} />
-                  )}
-                </Flex>
-              ))}
-            </Flex>
-          ) : (
-            <Box>关联的知识库({selectedKbList.length})</Box>
-          )}
-          {isPc && (
-            <Box fontSize={'sm'} color={'myGray.500'} fontWeight={'normal'}>
-              仅能选择同一个索引模型的知识库
-            </Box>
-          )}
-        </ModalHeader>
-
-        <ModalBody
-          flex={['1 0 0', '0 0 auto']}
-          maxH={'80vh'}
-          overflowY={'auto'}
-          display={'grid'}
-          userSelect={'none'}
-        >
-          <Grid
-            h={'auto'}
-            gridTemplateColumns={['repeat(1,1fr)', 'repeat(2,1fr)', 'repeat(3,1fr)']}
-            gridGap={3}
-          >
           {filterKbList.selected.map((item) =>
             (() => {
               return (
@@ -173,10 +104,7 @@ export const KBSelectModal = ({
 
         {filterKbList.selected.length > 0 && <Divider my={3} />}
 
-          <Grid
+        <Grid gridTemplateColumns={['repeat(1,1fr)', 'repeat(2,1fr)', 'repeat(3,1fr)']} gridGap={3}>
-            gridTemplateColumns={['repeat(1,1fr)', 'repeat(2,1fr)', 'repeat(3,1fr)']}
-            gridGap={3}
-          >
           {filterKbList.unSelected.map((item) =>
             (() => {
               return (
@@ -268,8 +196,7 @@ export const KBSelectModal = ({
           完成
         </Button>
       </ModalFooter>
-      </Flex>
+    </DatasetSelectModal>
-    </MyModal>
   );
 };
 

client/src/store/dataset.ts

@@ -8,7 +8,7 @@ import { defaultKbDetail } from '@/constants/kb';
 import { KbUpdateParams } from '@/api/request/kb';
 
 type State = {
-  datasets: KbListItemType[];
+  allDatasets: KbListItemType[];
   loadAllDatasets: () => Promise<KbListItemType[]>;
   myKbList: KbListItemType[];
   loadKbList: (parentId?: string) => Promise<any>;
@@ -27,11 +27,11 @@ export const useDatasetStore = create<State>()(
   devtools(
     persist(
       immer((set, get) => ({
-        datasets: [],
+        allDatasets: [],
         async loadAllDatasets() {
           const res = await getAllDataset();
           set((state) => {
-            state.datasets = res;
+            state.allDatasets = res;
           });
           return res;
         },