feat: model related kb

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';
 
 /**

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

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({

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,

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,

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,
+        'share.intro': share.intro
-        security
       }
     );
 

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({

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'}>
+        {getValues('chat.relatedKbs').length > 0 && (
-          <Box mr={4}>知识库搜索</Box>
-          <Switch
-            isDisabled={!isOwner}
-            isChecked={getValues('chat.useKb')}
-            onChange={() => {
-              setValue('chat.useKb', !getValues('chat.useKb'));
-              setRefresh(!refresh);
-            }}
-          />
-        </Flex>
-        {getValues('chat.useKb') && (
           <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 />
     </>

src/pages/model/components/detail/index.tsx

@@ -2,10 +2,9 @@ import React, { useCallback, useState, useMemo, useEffect } from 'react';
 import { useRouter } from 'next/router';
 import { delModelById, putModelById } from '@/api/model';
 import type { ModelSchema } from '@/types/mongoSchema';
-import { Card, Box, Flex, Button, Tag, Grid } from '@chakra-ui/react';
+import { Card, Box, Flex, Button, Grid } from '@chakra-ui/react';
 import { useToast } from '@/hooks/useToast';
 import { useForm } from 'react-hook-form';
-import { formatModelStatus } from '@/constants/model';
 import { useQuery } from '@tanstack/react-query';
 import { useUserStore } from '@/store/user';
 import { useLoading } from '@/hooks/useLoading';
@@ -18,7 +17,7 @@ const ModelDetail = ({ modelId, isPc }: { modelId: string; isPc: boolean }) => {
   const { Loading, setIsLoading } = useLoading();
   const [btnLoading, setBtnLoading] = useState(false);
 
-  const formHooks = useForm<ModelSchema>({
+  const formHooks = useForm({
     defaultValues: modelDetail
   });
 
@@ -84,13 +83,9 @@ const ModelDetail = ({ modelId, isPc }: { modelId: string; isPc: boolean }) => {
           name: data.name,
           avatar: data.avatar || '/icon/logo.png',
           chat: data.chat,
-          share: data.share,
+          share: data.share
-          security: data.security
-        });
-        toast({
-          title: '更新成功',
-          status: 'success'
         });
+
         refreshModel.updateModelDetail(data);
       } catch (err: any) {
         toast({
@@ -120,18 +115,16 @@ const ModelDetail = ({ modelId, isPc }: { modelId: string; isPc: boolean }) => {
     });
   }, [formHooks.formState.errors, toast]);
 
+  const saveUpdateModel = useCallback(
+    () => formHooks.handleSubmit(saveSubmitSuccess, saveSubmitError)(),
+    [formHooks, saveSubmitError, saveSubmitSuccess]
+  );
+
   useEffect(() => {
-    router.prefetch('/chat');
-
-    window.onbeforeunload = (e) => {
-      e.preventDefault();
-      e.returnValue = '内容已修改，确认离开页面吗？';
-    };
-
     return () => {
-      window.onbeforeunload = null;
+      saveUpdateModel();
     };
-  }, [router]);
+  }, []);
 
   return canRead ? (
     <Box h={'100%'} p={5} overflow={'overlay'} position={'relative'}>
@@ -142,13 +135,6 @@ const ModelDetail = ({ modelId, isPc }: { modelId: string; isPc: boolean }) => {
             <Box fontSize={'xl'} fontWeight={'bold'}>
               {modelDetail.name}
             </Box>
-            <Tag
-              ml={2}
-              variant="solid"
-              colorScheme={formatModelStatus[modelDetail.status].colorTheme}
-            >
-              {formatModelStatus[modelDetail.status].text}
-            </Tag>
             <Box flex={1} />
             <Button variant={'outline'} onClick={handlePreviewChat}>
               开始对话
@@ -157,7 +143,18 @@ const ModelDetail = ({ modelId, isPc }: { modelId: string; isPc: boolean }) => {
               <Button
                 isLoading={btnLoading}
                 ml={4}
-                onClick={formHooks.handleSubmit(saveSubmitSuccess, saveSubmitError)}
+                onClick={async () => {
+                  try {
+                    await saveUpdateModel();
+                    toast({
+                      title: '更新成功',
+                      status: 'success'
+                    });
+                  } catch (error) {
+                    console.log(error);
+                    error;
+                  }
+                }}
               >
                 保存修改
               </Button>
@@ -169,9 +166,6 @@ const ModelDetail = ({ modelId, isPc }: { modelId: string; isPc: boolean }) => {
               <Box as={'h3'} fontSize={'xl'} fontWeight={'bold'} flex={1}>
                 {modelDetail.name}
               </Box>
-              <Tag ml={2} colorScheme={formatModelStatus[modelDetail.status].colorTheme}>
-                {formatModelStatus[modelDetail.status].text}
-              </Tag>
             </Flex>
             <Box mt={4} textAlign={'right'}>
               <Button variant={'outline'} size={'sm'} onClick={handlePreviewChat}>
@@ -182,7 +176,18 @@ const ModelDetail = ({ modelId, isPc }: { modelId: string; isPc: boolean }) => {
                   ml={4}
                   size={'sm'}
                   isLoading={btnLoading}
-                  onClick={formHooks.handleSubmit(saveSubmitSuccess, saveSubmitError)}
+                  onClick={async () => {
+                    try {
+                      await saveUpdateModel();
+                      toast({
+                        title: '更新成功',
+                        status: 'success'
+                      });
+                    } catch (error) {
+                      console.log(error);
+                      error;
+                    }
+                  }}
                 >
                   保存修改
                 </Button>

src/service/models/model.ts

@@ -31,10 +31,10 @@ const ModelSchema = new Schema({
     default: () => new Date()
   },
   chat: {
-    useKb: {
+    relatedKbs: {
-      // use knowledge base to search
+      type: [Schema.Types.ObjectId],
-      type: Boolean,
+      ref: 'kb',
-      default: false
+      default: []
     },
     searchMode: {
       // knowledge base search mode
@@ -79,33 +79,6 @@ const ModelSchema = new Schema({
       type: Number,
       default: 0
     }
-  },
-  security: {
-    type: {
-      domain: {
-        type: [String],
-        default: ['*']
-      },
-      contextMaxLen: {
-        type: Number,
-        default: 20
-      },
-      contentMaxLen: {
-        type: Number,
-        default: 4000
-      },
-      expiredTime: {
-        type: Number,
-        default: 1,
-        set: (val: number) => val * (60 * 60 * 1000)
-      },
-      maxLoadAmount: {
-        // 负数代表不限制
-        type: Number,
-        default: -1
-      }
-    },
-    default: {}
   }
 });
 

src/service/plugins/searchKb.ts

@@ -48,7 +48,7 @@ export const searchKb = async ({
           where: [
             ['status', ModelDataStatusEnum.ready],
             'AND',
-            ['model_id', model._id],
+            `kb_id IN (${model.chat.relatedKbs.map((item) => `'${item}'`).join(',')})`,
             'AND',
             `vector <=> '[${promptVector}]' < ${similarity}`
           ],

src/service/utils/auth.ts

@@ -34,6 +34,13 @@ export const authToken = (req: NextApiRequest): Promise<string> => {
   });
 };
 
+export const getOpenAiKey = () => {
+  // 纯字符串类型
+  const keys = process.env.OPENAIKEY?.split(',') || [];
+  const i = Math.floor(Math.random() * keys.length);
+  return keys[i] || (process.env.OPENAIKEY as string);
+};
+
 /* 获取 api 请求的 key */
 export const getApiKey = async ({
   model,
@@ -52,7 +59,7 @@ export const getApiKey = async ({
   const keyMap = {
     [OpenAiChatEnum.GPT35]: {
       userOpenAiKey: user.openaiKey || '',
-      systemAuthKey: process.env.OPENAIKEY as string
+      systemAuthKey: getOpenAiKey() as string
     },
     [OpenAiChatEnum.GPT4]: {
       userOpenAiKey: user.openaiKey || '',

src/service/utils/chat/openai.ts

@@ -7,6 +7,7 @@ import { adaptChatItem_openAI } from '@/utils/chat/openai';
 import { modelToolMap } from '@/utils/chat';
 import { ChatCompletionType, ChatContextFilter, StreamResponseType } from './index';
 import { ChatRoleEnum } from '@/constants/chat';
+import { getOpenAiKey } from '../auth';
 
 export const getOpenAIApi = (apiKey: string) => {
   const configuration = new Configuration({
@@ -27,7 +28,7 @@ export const openaiCreateEmbedding = async ({
   userId: string;
   textArr: string[];
 }) => {
-  const systemAuthKey = process.env.OPENAIKEY as string;
+  const systemAuthKey = getOpenAiKey();
 
   // 获取 chatAPI
   const chatAPI = getOpenAIApi(userOpenAiKey || systemAuthKey);

src/store/user.ts

@@ -2,7 +2,6 @@ import { create } from 'zustand';
 import { devtools, persist } from 'zustand/middleware';
 import { immer } from 'zustand/middleware/immer';
 import type { UserType, UserUpdateParams } from '@/types/user';
-import type { ModelSchema } from '@/types/mongoSchema';
 import { getMyModels, getModelById } from '@/api/model';
 import { formatPrice } from '@/utils/user';
 import { getTokenLogin } from '@/api/user';
@@ -11,6 +10,7 @@ import { ModelListItemType } from '@/types/model';
 import { KbItemType } from '@/types/plugin';
 import { getKbList } from '@/api/plugins/kb';
 import { defaultKbDetail } from '@/constants/kb';
+import type { ModelSchema } from '@/types/mongoSchema';
 
 type State = {
   userInfo: UserType | null;
@@ -34,7 +34,7 @@ type State = {
   lastKbId: string;
   setLastKbId: (id: string) => void;
   myKbList: KbItemType[];
-  loadKbList: (init?: boolean) => Promise<null>;
+  loadKbList: (init?: boolean) => Promise<KbItemType[]>;
   KbDetail: KbItemType;
   getKbDetail: (id: string) => KbItemType;
 };
@@ -123,12 +123,12 @@ export const useUserStore = create<State>()(
         },
         myKbList: [],
         async loadKbList(init = false) {
-          if (get().myKbList.length > 0 && !init) return null;
+          if (get().myKbList.length > 0 && !init) return get().myKbList;
           const res = await getKbList();
           set((state) => {
             state.myKbList = res;
           });
-          return null;
+          return res;
         },
         KbDetail: defaultKbDetail,
         getKbDetail(id: string) {

src/types/model.d.ts

@@ -1,5 +1,6 @@
 import { ModelStatusEnum } from '@/constants/model';
-import type { ModelSchema } from './mongoSchema';
+import type { ModelSchema, kbSchema } from './mongoSchema';
+import { ChatModelType, ModelVectorSearchModeEnum } from '@/constants/model';
 
 export type ModelListItemType = {
   _id: string;
@@ -13,7 +14,6 @@ export interface ModelUpdateParams {
   avatar: string;
   chat: ModelSchema['chat'];
   share: ModelSchema['share'];
-  security: ModelSchema['security'];
 }
 
 export interface ShareModelItem {

src/types/mongoSchema.d.ts

@@ -38,7 +38,7 @@ export interface ModelSchema {
   status: `${ModelStatusEnum}`;
   updateTime: number;
   chat: {
-    useKb: boolean;
+    relatedKbs: string[];
     searchMode: `${ModelVectorSearchModeEnum}`;
     systemPrompt: string;
     temperature: number;
@@ -50,13 +50,6 @@ export interface ModelSchema {
     intro: string;
     collection: number;
   };
-  security: {
-    domain: string[];
-    contextMaxLen: number;
-    contentMaxLen: number;
-    expiredTime: number;
-    maxLoadAmount: number;
-  };
 }
 
 export interface ModelPopulate extends ModelSchema {