feat: model related kb

This commit is contained in:
archer
2023-05-17 22:09:38 +08:00
parent a79429fdcd
commit 5bf95bd846
16 changed files with 178 additions and 117 deletions

View File

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

View File

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

View File

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

View File

@@ -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,

View File

@@ -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,

View File

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

View File

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

View File

@@ -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 />
</>

View File

@@ -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,
security: data.security
});
toast({
title: '更新成功',
status: 'success'
share: data.share
});
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>

View File

@@ -31,10 +31,10 @@ const ModelSchema = new Schema({
default: () => new Date()
},
chat: {
useKb: {
// use knowledge base to search
type: Boolean,
default: false
relatedKbs: {
type: [Schema.Types.ObjectId],
ref: 'kb',
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: {}
}
});

View File

@@ -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}`
],

View File

@@ -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 || '',

View File

@@ -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);

View File

@@ -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) {

View File

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

View File

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