fix: mark modal cannot select folder (#327)

client/src/components/ChatBox/SelectDataset.tsx

-import React, { useRef, useState } from 'react';
-import {
-  ModalBody,
-  useTheme,
-  ModalFooter,
-  Button,
-  ModalHeader,
-  Box,
-  Card,
-  Flex
-} from '@chakra-ui/react';
-import MyModal from '../MyModal';
+import React, { useState } from 'react';
+import { ModalBody, useTheme, ModalFooter, Button, Box, Card, Flex, Grid } from '@chakra-ui/react';
 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>();
-
-  useQuery(['loadKbList'], () => loadKbList());
+  const { paths, parentId, setParentId, datasets } = useDatasetSelect();
 
   return (
-    <MyModal isOpen={true} onClose={onClose} w={'100%'} maxW={['90vw', '900px']} isCentered={!isPc}>
-      <Flex flexDirection={'column'} h={['90vh', 'auto']}>
-        <ModalHeader>
-          <Box>{t('chat.Select Mark Kb')}</Box>
-          <Box fontSize={'sm'} color={'myGray.500'} fontWeight={'normal'}>
-            {t('chat.Select Mark Kb Desc')}
-          </Box>
-        </ModalHeader>
-        <ModalBody
-          flex={['1 0 0', '0 0 auto']}
-          maxH={'80vh'}
-          overflowY={'auto'}
-          display={'grid'}
+    <DatasetSelectModal
+      isOpen={isOpen}
+      paths={paths}
+      onClose={onClose}
+      parentId={parentId}
+      setParentId={setParentId}
+      tips={t('chat.Select Mark Kb Desc')}
+    >
+      <ModalBody flex={['1 0 0', '0 0 auto']} maxH={'80vh'} overflowY={'auto'}>
+        <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={() => {
-                    setSelectedId(item._id);
+                    if (item.type === KbTypeEnum.folder) {
+                      setParentId(item._id);
+                    } else {
+                      setSelectedId(item._id);
+                    }
                   }}
                 >
                   <Flex alignItems={'center'} h={'38px'}>
@@ -89,28 +78,36 @@ const SelectDataset = ({
               );
             })()
           )}
-        </ModalBody>
-        <ModalFooter>
-          <Button variant={'base'} mr={2} onClick={onClose}>
-            {t('Cancel')}
-          </Button>
-          <Button
-            onClick={() => {
-              if (!selectedId) {
-                return toast({
-                  status: 'warning',
-                  title: t('Select value is empty')
-                });
-              }
+        </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}>
+          {t('Cancel')}
+        </Button>
+        <Button
+          onClick={() => {
+            if (!selectedId) {
+              return toast({
+                status: 'warning',
+                title: t('Select value is empty')
+              });
+            }
 
-              onSuccess(selectedId);
-            }}
-          >
-            {t('Confirm')}
-          </Button>
-        </ModalFooter>
-      </Flex>
-    </MyModal>
+            onSuccess(selectedId);
+          }}
+        >
+          {t('Confirm')}
+        </Button>
+      </ModalFooter>
+    </DatasetSelectModal>
   );
 };
 

client/src/components/ChatBox/index.tsx

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

+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)),
-    [datasets, kbList]
+    () => allDatasets.filter((item) => kbList.find((kb) => kb.kbId === item._id)),
+    [allDatasets, kbList]
   );
 
   /* 点击删除 */
@@ -564,16 +564,15 @@ const Settings = ({ appId }: { appId: string }) => {
           defaultData={getValues('chatModel')}
         />
       )}
-      {isOpenKbSelect && (
-        <KBSelectModal
-          activeKbs={selectedKbList.map((item) => ({
-            kbId: item._id,
-            vectorModel: item.vectorModel
-          }))}
-          onClose={onCloseKbSelect}
-          onChange={replaceKbList}
-        />
-      )}
+      <KBSelectModal
+        isOpen={isOpenKbSelect}
+        activeKbs={selectedKbList.map((item) => ({
+          kbId: item._id,
+          vectorModel: item.vectorModel
+        }))}
+        onClose={onCloseKbSelect}
+        onChange={replaceKbList}
+      />
       {isOpenKbParams && (
         <KbParamsModal
           searchEmptyText={getValues('kb.searchEmptyText')}

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;
         },