feat: model related kb

lumingxu2020 · May 17, 2023 · 5bf95bd · 5bf95bd
1 parent a79429f
commit 5bf95bd
Show file tree

Hide file tree

Showing 16 changed files with 178 additions and 117 deletions.
diff --git a/src/api/model.ts b/src/api/model.ts
@@ -2,7 +2,6 @@ import { GET, POST, DELETE, PUT } from './request';
 import type { ModelSchema } from '@/types/mongoSchema';
 import type { ModelUpdateParams, ShareModelItem } from '@/types/model';
 import { RequestPaging } from '../types/index';
-import { Obj2Query } from '@/utils/tools';
 import type { ModelListResponse } from './response/model';
 
 /**

diff --git a/src/constants/model.ts b/src/constants/model.ts
@@ -1,6 +1,6 @@
 import { getSystemModelList } from '@/api/system';
-import type { ModelSchema } from '@/types/mongoSchema';
 import type { ShareChatEditType } from '@/types/model';
+import type { ModelSchema } from '@/types/mongoSchema';
 
 export const embeddingModel = 'text-embedding-ada-002';
 export type EmbeddingModelType = 'text-embedding-ada-002';
@@ -142,7 +142,7 @@ export const defaultModel: ModelSchema = {
   status: ModelStatusEnum.pending,
   updateTime: Date.now(),
   chat: {
-    useKb: false,
+    relatedKbs: [],
     searchMode: ModelVectorSearchModeEnum.hightSimilarity,
     systemPrompt: '',
     temperature: 0,
@@ -153,13 +153,6 @@ export const defaultModel: ModelSchema = {
     isShareDetail: false,
     intro: '',
     collection: 0
-  },
-  security: {
-    domain: ['*'],
-    contextMaxLen: 1,
-    contentMaxLen: 1,
-    expiredTime: 9999,
-    maxLoadAmount: 1
   }
 };
 

diff --git a/src/constants/theme.ts b/src/constants/theme.ts
@@ -1,6 +1,6 @@
 import { extendTheme, defineStyleConfig, ComponentStyleConfig } from '@chakra-ui/react';
 // @ts-ignore
-import { modalAnatomy, switchAnatomy, selectAnatomy } from '@chakra-ui/anatomy';
+import { modalAnatomy, switchAnatomy, selectAnatomy, checkboxAnatomy } from '@chakra-ui/anatomy';
 // @ts-ignore
 import { createMultiStyleConfigHelpers } from '@chakra-ui/styled-system';
 
@@ -11,6 +11,8 @@ const { definePartsStyle: switchPart, defineMultiStyleConfig: switchMultiStyle }
   createMultiStyleConfigHelpers(switchAnatomy.keys);
 const { definePartsStyle: selectPart, defineMultiStyleConfig: selectMultiStyle } =
   createMultiStyleConfigHelpers(selectAnatomy.keys);
+const { definePartsStyle: checkboxPart, defineMultiStyleConfig: checkboxMultiStyle } =
+  createMultiStyleConfigHelpers(checkboxAnatomy.keys);
 
 // modal 弹窗
 const ModalTheme = defineMultiStyleConfig({

diff --git a/src/pages/api/chat/chat.ts b/src/pages/api/chat/chat.ts
@@ -54,7 +54,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     const prompts = [...content, prompt];
 
     // 使用了知识库搜索
-    if (model.chat.useKb) {
+    if (model.chat.relatedKbs.length > 0) {
       const { code, searchPrompts } = await searchKb({
         userOpenAiKey,
         prompts,

diff --git a/src/pages/api/chat/shareChat/chat.ts b/src/pages/api/chat/shareChat/chat.ts
@@ -50,7 +50,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     const modelConstantsData = ChatModelMap[model.chat.chatModel];
 
     // 使用了知识库搜索
-    if (model.chat.useKb) {
+    if (model.chat.relatedKbs.length > 0) {
       const { code, searchPrompts } = await searchKb({
         userOpenAiKey,
         prompts,

diff --git a/src/pages/api/model/update.ts b/src/pages/api/model/update.ts
@@ -9,10 +9,10 @@ import { authModel } from '@/service/utils/auth';
 /* 获取我的模型 */
 export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
   try {
-    const { name, avatar, chat, share, security } = req.body as ModelUpdateParams;
+    const { name, avatar, chat, share } = req.body as ModelUpdateParams;
     const { modelId } = req.query as { modelId: string };
 
-    if (!name || !chat || !security || !modelId) {
+    if (!name || !chat || !modelId) {
       throw new Error('参数错误');
     }
 
@@ -38,8 +38,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
         chat,
         'share.isShare': share.isShare,
         'share.isShareDetail': share.isShareDetail,
-        'share.intro': share.intro,
-        security
+        'share.intro': share.intro
       }
     );
 

diff --git a/src/pages/api/openapi/chat/chat.ts b/src/pages/api/openapi/chat/chat.ts
@@ -70,7 +70,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     const modelConstantsData = ChatModelMap[model.chat.chatModel];
 
     // 使用了知识库搜索
-    if (model.chat.useKb) {
+    if (model.chat.relatedKbs.length > 0) {
       const similarity = ModelVectorSearchModeMap[model.chat.searchMode]?.similarity || 0.22;
 
       const { code, searchPrompts } = await searchKb({

diff --git a/src/pages/model/components/detail/components/ModelEditForm.tsx b/src/pages/model/components/detail/components/ModelEditForm.tsx
@@ -32,11 +32,9 @@ import {
   Th,
   Td,
   TableContainer,
-  IconButton
+  Checkbox
 } from '@chakra-ui/react';
-import { DeleteIcon } from '@chakra-ui/icons';
 import { QuestionOutlineIcon } from '@chakra-ui/icons';
-import type { ModelSchema } from '@/types/mongoSchema';
 import { useForm, UseFormReturn } from 'react-hook-form';
 import { ChatModelMap, ModelVectorSearchModeMap, getChatModelList } from '@/constants/model';
 import { formatPrice } from '@/utils/user';
@@ -49,9 +47,12 @@ import { getShareChatList, createShareChat, delShareChatById } from '@/api/chat'
 import { useRouter } from 'next/router';
 import { defaultShareChat } from '@/constants/model';
 import type { ShareChatEditType } from '@/types/model';
+import type { ModelSchema } from '@/types/mongoSchema';
 import { formatTimeToChatTime, useCopyData } from '@/utils/tools';
 import MyIcon from '@/components/Icon';
 import { useGlobalStore } from '@/store/global';
+import { useUserStore } from '@/store/user';
+import type { KbItemType } from '@/types/plugin';
 
 const ModelEditForm = ({
   formHooks,
@@ -62,10 +63,11 @@ const ModelEditForm = ({
   isOwner: boolean;
   handleDelModel: () => void;
 }) => {
-  const { toast } = useToast();
   const { modelId } = useRouter().query as { modelId: string };
-  const { setLoading } = useGlobalStore();
   const [refresh, setRefresh] = useState(false);
+  const { toast } = useToast();
+  const { setLoading } = useGlobalStore();
+  const { loadKbList } = useUserStore();
 
   const { openConfirm, ConfirmChild } = useConfirm({
     content: '确认删除该AI助手?'
@@ -86,6 +88,11 @@ const ModelEditForm = ({
     onOpen: onOpenCreateShareChat,
     onClose: onCloseCreateShareChat
   } = useDisclosure();
+  const {
+    isOpen: isOpenKbSelect,
+    onOpen: onOpenKbSelect,
+    onClose: onCloseKbSelect
+  } = useDisclosure();
   const { File, onOpen: onOpenSelectFile } = useSelectFile({
     fileType: '.jpg,.png',
     multiple: false
@@ -153,11 +160,41 @@ ${e.password ? `密码为: ${e.password}` : ''}`;
     ]
   );
 
+  // format share used token
   const formatTokens = (tokens: number) => {
     if (tokens < 10000) return tokens;
     return `${(tokens / 10000).toFixed(2)}万`;
   };
 
+  // init kb select list
+  const { data: kbList = [] } = useQuery(['loadKbList'], () => loadKbList());
+  const RenderSelectedKbList = useCallback(() => {
+    const kbs = getValues('chat.relatedKbs').map((id) => kbList.find((kb) => kb._id === id));
+
+    return (
+      <>
+        {kbs.map((item) =>
+          item ? (
+            <Card key={item._id} p={3} mt={3}>
+              <Flex alignItems={'center'}>
+                <Image
+                  src={item.avatar}
+                  fallbackSrc="/icon/logo.png"
+                  w={'20px'}
+                  h={'20px'}
+                  alt=""
+                ></Image>
+                <Box ml={3} fontWeight={'bold'}>
+                  {item.name}
+                </Box>
+              </Flex>
+            </Card>
+          ) : null
+        )}
+      </>
+    );
+  }, [getValues, kbList]);
+
   return (
     <>
       {/* basic info */}
@@ -292,18 +329,7 @@ ${e.password ? `密码为: ${e.password}` : ''}`;
             </Slider>
           </Flex>
         </FormControl>
-        <Flex mt={4} alignItems={'center'}>
-          <Box mr={4}>知识库搜索</Box>
-          <Switch
-            isDisabled={!isOwner}
-            isChecked={getValues('chat.useKb')}
-            onChange={() => {
-              setValue('chat.useKb', !getValues('chat.useKb'));
-              setRefresh(!refresh);
-            }}
-          />
-        </Flex>
-        {getValues('chat.useKb') && (
+        {getValues('chat.relatedKbs').length > 0 && (
           <Flex mt={4} alignItems={'center'}>
             <Box mr={4} whiteSpace={'nowrap'}>
               搜索模式&emsp;
@@ -339,7 +365,9 @@ ${e.password ? `密码为: ${e.password}` : ''}`;
             <Box fontWeight={'bold'}>分享设置</Box>
             <Box>
               <Flex mt={5} alignItems={'center'}>
-                <Box mr={1}>模型分享:</Box>
+                <Box mr={1} fontSize={['sm', 'md']}>
+                  模型分享:
+                </Box>
                 <Tooltip label="开启模型分享后，你的模型将会出现在共享市场，可供 FastGpt 所有用户使用。用户使用时不会消耗你的 tokens，而是消耗使用者的 tokens。">
                   <QuestionOutlineIcon mr={3} />
                 </Tooltip>
@@ -350,7 +378,8 @@ ${e.password ? `密码为: ${e.password}` : ''}`;
                     setRefresh(!refresh);
                   }}
                 />
-                <Box ml={12} mr={1}>
+
+                <Box ml={12} mr={1} fontSize={['sm', 'md']}>
                   分享模型细节:
                 </Box>
                 <Tooltip label="开启分享详情后，其他用户可以查看该模型的特有数据：温度、提示词和数据集。">
@@ -376,8 +405,22 @@ ${e.password ? `密码为: ${e.password}` : ''}`;
               </Box>
             </Box>
           </Card>
-          {/* shareChat */}
           <Card p={4}>
+            <Flex justifyContent={'space-between'}>
+              <Box fontWeight={'bold'}>关联的知识库</Box>
+              <Button
+                size={'sm'}
+                variant={'outline'}
+                colorScheme={'myBlue'}
+                onClick={onOpenKbSelect}
+              >
+                选择
+              </Button>
+            </Flex>
+            <RenderSelectedKbList />
+          </Card>
+          {/* shareChat */}
+          <Card p={4} gridColumnStart={1} gridColumnEnd={[2, 3]}>
             <Flex justifyContent={'space-between'}>
               <Box fontWeight={'bold'}>
                 免登录聊天窗口
@@ -410,7 +453,7 @@ ${e.password ? `密码为: ${e.password}` : ''}`;
                     <Th>最大上下文</Th>
                     <Th>tokens消耗</Th>
                     <Th>最后使用时间</Th>
-                    <Th></Th>
+                    <Th>操作</Th>
                   </Tr>
                 </Thead>
                 <Tbody>
@@ -539,6 +582,52 @@ ${e.password ? `密码为: ${e.password}` : ''}`;
           </ModalFooter>
         </ModalContent>
       </Modal>
+      {/* select kb modal */}
+      <Modal isOpen={isOpenKbSelect} onClose={onCloseKbSelect}>
+        <ModalOverlay />
+        <ModalContent>
+          <ModalHeader>选择关联的知识库</ModalHeader>
+          <ModalCloseButton />
+          <ModalBody>
+            {kbList.map((item) => (
+              <Card key={item._id} p={3} mb={3}>
+                <Checkbox
+                  isChecked={getValues('chat.relatedKbs').includes(item._id)}
+                  onChange={(e) => {
+                    const ids = getValues('chat.relatedKbs');
+                    // toggle to true
+                    if (e.target.checked) {
+                      setValue('chat.relatedKbs', ids.concat(item._id));
+                    } else {
+                      const i = ids.findIndex((id) => id === item._id);
+                      ids.splice(i, 1);
+                      setValue('chat.relatedKbs', ids);
+                    }
+                    setRefresh(!refresh);
+                  }}
+                >
+                  <Flex alignItems={'center'}>
+                    <Image
+                      src={item.avatar}
+                      fallbackSrc="/icon/logo.png"
+                      w={'20px'}
+                      h={'20px'}
+                      alt=""
+                    ></Image>
+                    <Box ml={3} fontWeight={'bold'}>
+                      {item.name}
+                    </Box>
+                  </Flex>
+                </Checkbox>
+              </Card>
+            ))}
+          </ModalBody>
+
+          <ModalFooter>
+            <Button onClick={onCloseKbSelect}>完成,记得点保存修改</Button>
+          </ModalFooter>
+        </ModalContent>
+      </Modal>
       <File onSelect={onSelectFile} />
       <ConfirmChild />
     </>