feat: 增加账单

src/constants/model.ts

@@ -7,31 +7,35 @@ export enum ChatModelNameEnum {
 }
 
 export type ModelConstantsData = {
+  serviceCompany: `${ServiceName}`;
   name: string;
   model: `${ChatModelNameEnum}`;
   trainName: string; // 空字符串代表不能训练
   maxToken: number;
   maxTemperature: number;
+  price: number; // 多少钱 / 1字，单位: 0.00001元
 };
 
-export const ModelList: Record<ServiceName, ModelConstantsData[]> = {
-  openai: [
-    {
-      name: 'chatGPT',
-      model: ChatModelNameEnum.GPT35,
-      trainName: 'turbo',
-      maxToken: 4000,
-      maxTemperature: 2
-    },
-    {
-      name: 'GPT3',
-      model: ChatModelNameEnum.GPT3,
-      trainName: 'davinci',
-      maxToken: 4000,
-      maxTemperature: 2
-    }
-  ]
-};
+export const ModelList: ModelConstantsData[] = [
+  {
+    serviceCompany: 'openai',
+    name: 'chatGPT',
+    model: ChatModelNameEnum.GPT35,
+    trainName: 'turbo',
+    maxToken: 4000,
+    maxTemperature: 2,
+    price: 2
+  },
+  {
+    serviceCompany: 'openai',
+    name: 'GPT3',
+    model: ChatModelNameEnum.GPT3,
+    trainName: 'davinci',
+    maxToken: 4000,
+    maxTemperature: 2,
+    price: 20
+  }
+];
 
 export enum TrainingStatusEnum {
   pending = 'pending',

src/pages/api/chat/chatGpt.ts

@@ -1,6 +1,6 @@
 import type { NextApiRequest, NextApiResponse } from 'next';
 import { createParser, ParsedEvent, ReconnectInterval } from 'eventsource-parser';
-import { connectToDatabase, Chat } from '@/service/mongo';
+import { connectToDatabase } from '@/service/mongo';
 import { getOpenAIApi, authChat } from '@/service/utils/chat';
 import { httpsAgent } from '@/service/utils/tools';
 import { ChatCompletionRequestMessage, ChatCompletionRequestMessageRoleEnum } from 'openai';
@@ -9,6 +9,7 @@ import { jsonRes } from '@/service/response';
 import type { ModelSchema } from '@/types/mongoSchema';
 import { PassThrough } from 'stream';
 import { ModelList } from '@/constants/model';
+import { pushBill } from '@/service/events/bill';
 
 /* 发送提示词 */
 export default async function handler(req: NextApiRequest, res: NextApiResponse) {
@@ -24,7 +25,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
 
     await connectToDatabase();
 
-    const { chat, userApiKey } = await authChat(chatId);
+    const { chat, userApiKey, systemKey, userId } = await authChat(chatId);
 
     const model: ModelSchema = chat.modelId;
 
@@ -58,16 +59,14 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     }
 
     // 计算温度
-    const modelConstantsData = ModelList['openai'].find(
-      (item) => item.model === model.service.modelName
-    );
+    const modelConstantsData = ModelList.find((item) => item.model === model.service.modelName);
     if (!modelConstantsData) {
       throw new Error('模型异常');
     }
     const temperature = modelConstantsData.maxTemperature * (model.temperature / 10);
 
     // 获取 chatAPI
-    const chatAPI = getOpenAIApi(userApiKey);
+    const chatAPI = getOpenAIApi(userApiKey || systemKey);
     let startTime = Date.now();
     // 发出请求
     const chatResponse = await chatAPI.createChatCompletion(
@@ -84,12 +83,8 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
         httpsAgent
       }
     );
-    console.log(
-      'response success',
-      `time: ${(Date.now() - startTime) / 1000}s`,
-      `promptLen: ${formatPrompts.length}`,
-      `contentLen: ${formatPrompts.reduce((sum, item) => sum + item.content.length, 0)}`
-    );
+
+    console.log('api response time:', `time: ${(Date.now() - startTime) / 1000}s`);
 
     // 创建响应流
     res.setHeader('Content-Type', 'text/event-stream;charset-utf-8');
@@ -97,6 +92,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     res.setHeader('X-Accel-Buffering', 'no');
     res.setHeader('Cache-Control', 'no-cache, no-transform');
 
+    let responseContent = '';
     const pass = new PassThrough();
     pass.pipe(res);
 
@@ -108,6 +104,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
         const json = JSON.parse(data);
         const content: string = json?.choices?.[0].delta.content || '';
         if (!content) return;
+        responseContent += content;
         // console.log('content:', content)
         pass.push(content.replace(/\n/g, '<br/>'));
       } catch (error) {
@@ -125,6 +122,17 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
       console.log('pipe error', error);
     }
     pass.push(null);
+
+    const promptsLen = formatPrompts.reduce((sum, item) => sum + item.content.length, 0);
+    console.log(`responseLen: ${responseContent.length}`, `promptLen: ${promptsLen}`);
+    // 只有使用平台的 key 才计费
+    !userApiKey &&
+      pushBill({
+        modelName: model.service.modelName,
+        userId,
+        chatId,
+        textLen: promptsLen + responseContent.length
+      });
   } catch (err: any) {
     res.status(500);
     jsonRes(res, {

src/pages/api/chat/gpt3.ts

@@ -6,6 +6,7 @@ import { getOpenAIApi, authChat } from '@/service/utils/chat';
 import { ChatItemType } from '@/types/chat';
 import { httpsAgent } from '@/service/utils/tools';
 import { ModelList } from '@/constants/model';
+import { pushBill } from '@/service/events/bill';
 
 /* 发送提示词 */
 export default async function handler(req: NextApiRequest, res: NextApiResponse) {
@@ -18,20 +19,18 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
 
     await connectToDatabase();
 
-    const { chat, userApiKey } = await authChat(chatId);
+    const { chat, userApiKey, systemKey, userId } = await authChat(chatId);
 
     const model = chat.modelId;
 
     // 获取 chatAPI
-    const chatAPI = getOpenAIApi(userApiKey);
+    const chatAPI = getOpenAIApi(userApiKey || systemKey);
 
     // prompt处理
-    const formatPrompt = prompt.map((item) => `${item.value}\n\n###\n\n`).join('');
+    const formatPrompts = prompt.map((item) => `${item.value}\n\n###\n\n`).join('');
 
     // 计算温度
-    const modelConstantsData = ModelList['openai'].find(
-      (item) => item.model === model.service.modelName
-    );
+    const modelConstantsData = ModelList.find((item) => item.model === model.service.modelName);
     if (!modelConstantsData) {
       throw new Error('模型异常');
     }
@@ -41,7 +40,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     const response = await chatAPI.createCompletion(
       {
         model: model.service.modelName,
-        prompt: formatPrompt,
+        prompt: formatPrompts,
         temperature: temperature,
         // max_tokens: modelConstantsData.maxToken,
         top_p: 1,
@@ -54,7 +53,18 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
       }
     );
 
-    const responseMessage = response.data.choices[0]?.text;
+    const responseMessage = response.data.choices[0]?.text || '';
+
+    const promptsLen = prompt.reduce((sum, item) => sum + item.value.length, 0);
+    console.log(`responseLen: ${responseMessage.length}`, `promptLen: ${promptsLen}`);
+    // 只有使用平台的 key 才计费
+    !userApiKey &&
+      pushBill({
+        modelName: model.service.modelName,
+        userId,
+        chatId,
+        textLen: promptsLen + responseMessage.length
+      });
 
     jsonRes(res, {
       data: responseMessage

src/pages/api/model/create.ts

@@ -4,19 +4,13 @@ import { jsonRes } from '@/service/response';
 import { connectToDatabase } from '@/service/mongo';
 import { authToken } from '@/service/utils/tools';
 import { ModelStatusEnum, ModelList, ChatModelNameEnum } from '@/constants/model';
-import type { ServiceName } from '@/types/mongoSchema';
 import { Model } from '@/service/models/model';
 
 export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
   try {
-    const {
-      name,
-      serviceModelName,
-      serviceModelCompany = 'openai'
-    } = req.body as {
+    const { name, serviceModelName } = req.body as {
       name: string;
       serviceModelName: `${ChatModelNameEnum}`;
-      serviceModelCompany: ServiceName;
     };
     const { authorization } = req.headers;
 
@@ -24,16 +18,14 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
       throw new Error('无权操作');
     }
 
-    if (!name || !serviceModelName || !serviceModelCompany) {
+    if (!name || !serviceModelName) {
       throw new Error('缺少参数');
     }
 
     // 凭证校验
     const userId = await authToken(authorization);
 
-    const modelItem = ModelList[serviceModelCompany].find(
-      (item) => item.model === serviceModelName
-    );
+    const modelItem = ModelList.find((item) => item.model === serviceModelName);
 
     if (!modelItem) {
       throw new Error('模型不存在');
@@ -64,7 +56,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
       userId,
       status: ModelStatusEnum.running,
       service: {
-        company: serviceModelCompany,
+        company: modelItem.serviceCompany,
         trainId: modelItem.trainName,
         chatModel: modelItem.model,
         modelName: modelItem.model

src/pages/chat/components/SlideBar.tsx

@@ -156,7 +156,7 @@ const SlideBar = ({
 
       {/* 我的模型 & 历史记录 折叠框*/}
       <Box flex={'1 0 0'} px={3} h={0} overflowY={'auto'}>
-        <Accordion defaultIndex={[0]} allowToggle allowMultiple>
+        <Accordion defaultIndex={[0]} allowMultiple>
           {isSuccess && (
             <AccordionItem borderTop={0} borderBottom={0}>
               <AccordionButton borderRadius={'md'} pl={1}>

src/pages/model/components/CreateModel.tsx

@@ -42,7 +42,7 @@ const CreateModel = ({
     formState: { errors }
   } = useForm<CreateFormType>({
     defaultValues: {
-      serviceModelName: ModelList['openai'][0].model
+      serviceModelName: ModelList[0].model
     }
   });
 
@@ -95,7 +95,7 @@ const CreateModel = ({
                   required: '底层模型不能为空'
                 })}
               >
-                {ModelList['openai'].map((item) => (
+                {ModelList.map((item) => (
                   <option key={item.model} value={item.model}>
                     {item.name}
                   </option>

src/pages/model/detail.tsx

@@ -38,9 +38,7 @@ const ModelDetail = ({ modelId }: { modelId: string }) => {
   });
 
   const canTrain = useMemo(() => {
-    const openai = ModelList[model.service.company].find(
-      (item) => item.model === model?.service.modelName
-    );
+    const openai = ModelList.find((item) => item.model === model?.service.modelName);
     return openai && openai.trainName;
   }, [model]);
 

src/pages/number/setting.tsx

@@ -68,17 +68,17 @@ const NumberSetting = () => {
             <Box>{userInfo?.email}</Box>
           </Flex>
         </Box>
-        {/* <Box mt={6}>
+        <Box mt={6}>
           <Flex alignItems={'center'}>
             <Box flex={'0 0 60px'}>余额:</Box>
             <Box>
               <strong>{userInfo?.balance}</strong> 元
             </Box>
-            <Button size={'sm'} w={'80px'} ml={5}>
+            {/* <Button size={'sm'} w={'80px'} ml={5}>
               充值
-            </Button>
+            </Button> */}
           </Flex>
-        </Box> */}
+        </Box>
       </Card>
       <Card mt={6} px={6} py={4}>
         <Flex mb={5} justifyContent={'space-between'}>
@@ -148,6 +148,55 @@ const NumberSetting = () => {
           </Table>
         </TableContainer>
       </Card>
+      <Card mt={6} px={6} py={4}>
+        <Box fontSize={'xl'} fontWeight={'bold'}>
+          使用记录
+        </Box>
+        <TableContainer>
+          <Table>
+            <Thead>
+              <Tr>
+                <Th>账号类型</Th>
+                <Th>值</Th>
+                <Th></Th>
+              </Tr>
+            </Thead>
+            <Tbody>
+              {accounts.map((item, i) => (
+                <Tr key={item.id}>
+                  <Td minW={'200px'}>
+                    <Select
+                      {...register(`accounts.${i}.type`, {
+                        required: '类型不能为空'
+                      })}
+                    >
+                      <option value="openai">openai</option>
+                    </Select>
+                  </Td>
+                  <Td minW={'200px'} whiteSpace="pre-wrap" wordBreak={'break-all'}>
+                    <Input
+                      {...register(`accounts.${i}.value`, {
+                        required: '账号不能为空'
+                      })}
+                    ></Input>
+                  </Td>
+                  <Td>
+                    <IconButton
+                      aria-label="删除账号"
+                      icon={<DeleteIcon />}
+                      colorScheme={'red'}
+                      onClick={() => {
+                        removeAccount(i);
+                        handleSubmit(onclickSave)();
+                      }}
+                    />
+                  </Td>
+                </Tr>
+              ))}
+            </Tbody>
+          </Table>
+        </TableContainer>
+      </Card>
     </>
   );
 };

src/service/events/bill.ts

@@ -0,0 +1,41 @@
+import { connectToDatabase, Bill, User } from '../mongo';
+import { ModelList } from '@/constants/model';
+
+export const pushBill = async ({
+  modelName,
+  userId,
+  chatId,
+  textLen
+}: {
+  modelName: string;
+  userId: string;
+  chatId: string;
+  textLen: number;
+}) => {
+  await connectToDatabase();
+
+  const modelItem = ModelList.find((item) => item.model === modelName);
+
+  if (!modelItem) return;
+
+  const price = modelItem.price * textLen;
+
+  let billId;
+  try {
+    // 插入 Bill 记录
+    const res = await Bill.create({
+      userId,
+      chatId,
+      textLen,
+      price
+    });
+    billId = res._id;
+
+    // 扣费
+    await User.findByIdAndUpdate(userId, {
+      $inc: { balance: -price }
+    });
+  } catch (error) {
+    Bill.findByIdAndDelete(billId);
+  }
+};

src/service/models/bill.ts

@@ -0,0 +1,29 @@
+import { Schema, model, models } from 'mongoose';
+
+const BillSchema = new Schema({
+  userId: {
+    type: Schema.Types.ObjectId,
+    ref: 'user',
+    required: true
+  },
+  chatId: {
+    type: Schema.Types.ObjectId,
+    ref: 'chat',
+    required: true
+  },
+  time: {
+    type: Number,
+    default: () => Date.now()
+  },
+  textLen: {
+    // 提示词+响应的总字数
+    type: Number,
+    required: true
+  },
+  price: {
+    type: Number,
+    required: true
+  }
+});
+
+export const Bill = models['bill'] || model('bill', BillSchema);

src/service/mongo.ts

@@ -30,3 +30,4 @@ export * from './models/chat';
 export * from './models/model';
 export * from './models/user';
 export * from './models/training';
+export * from './models/bill';

src/service/utils/chat.ts

@@ -1,6 +1,7 @@
 import { Configuration, OpenAIApi } from 'openai';
 import { Chat } from '../mongo';
 import type { ChatPopulate } from '@/types/mongoSchema';
+import { formatPrice } from '@/utils/user';
 
 export const getOpenAIApi = (apiKey: string) => {
   const configuration = new Configuration({
@@ -40,12 +41,14 @@ export const authChat = async (chatId: string) => {
 
   const userApiKey = user.accounts?.find((item: any) => item.type === 'openai')?.value;
 
-  if (!userApiKey) {
-    return Promise.reject('缺少ApiKey, 无法请求');
+  if (!userApiKey && formatPrice(user.balance) <= -1) {
+    return Promise.reject('该账号余额不足');
   }
 
   return {
     userApiKey,
-    chat
+    systemKey: process.env.OPENAIKEY as string,
+    chat,
+    userId: user._id
   };
 };

src/store/chat.ts

@@ -19,6 +19,7 @@ export const useChatStore = create<Props>()(
         chatHistory: [],
         pushChatHistory(item: HistoryItem) {
           set((state) => {
+            if (state.chatHistory.find((history) => history.chatId === item.chatId)) return;
             state.chatHistory = [item, ...state.chatHistory].slice(0, 20);
           });
         },

src/store/user.ts

@@ -5,6 +5,7 @@ import type { UserType, UserUpdateParams } from '@/types/user';
 import type { ModelSchema } from '@/types/mongoSchema';
 import { setToken } from '@/utils/user';
 import { getMyModels } from '@/api/model';
+import { formatPrice } from '@/utils/user';
 
 type State = {
   userInfo: UserType | null;
@@ -21,7 +22,10 @@ export const useUserStore = create<State>()(
       userInfo: null,
       setUserInfo(user: UserType, token?: string) {
         set((state) => {
-          state.userInfo = user;
+          state.userInfo = {
+            ...user,
+            balance: formatPrice(user.balance)
+          };
         });
         token && setToken(token);
       },

src/utils/user.ts

@@ -9,3 +9,10 @@ export const getToken = () => {
 export const clearToken = () => {
   localStorage.removeItem(tokenKey);
 };
+
+/**
+ * 把数据库读取到的price，转化成元
+ */
+export const formatPrice = (val: number) => {
+  return val / 100000;
+};