feat: 增加账单

42c26bd1 · archer · dc467c26 · 42c26bd1 · 42c26bd1 · 42c26bd1
Commit 42c26bd1 authored Mar 21, 2023 by archer
15 changed files
--- a/src/constants/model.ts
+++ b/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[]> = {
+export const ModelList: ModelConstantsData[] = [
-  openai: [
+  {
-    {
+    serviceCompany: 'openai',
-      name: 'chatGPT',
+    name: 'chatGPT',
-      model: ChatModelNameEnum.GPT35,
+    model: ChatModelNameEnum.GPT35,
-      trainName: 'turbo',
+    trainName: 'turbo',
-      maxToken: 4000,
+    maxToken: 4000,
-      maxTemperature: 2
+    maxTemperature: 2,
-    },
+    price: 2
-    {
+  },
-      name: 'GPT3',
+  {
-      model: ChatModelNameEnum.GPT3,
+    serviceCompany: 'openai',
-      trainName: 'davinci',
+    name: 'GPT3',
-      maxToken: 4000,
+    model: ChatModelNameEnum.GPT3,
-      maxTemperature: 2
+    trainName: 'davinci',
-    }
+    maxToken: 4000,
-  ]
+    maxTemperature: 2,
-};
+    price: 20
+  }
+];
 export enum TrainingStatusEnum {
  pending = 'pending',

--- a/src/pages/api/chat/chatGpt.ts
+++ b/src/pages/api/chat/chatGpt.ts
 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(
+    const modelConstantsData = ModelList.find((item) => item.model === model.service.modelName);
-      (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',
+    console.log('api response time:', `time: ${(Date.now() - startTime) / 1000}s`);
-      `time: ${(Date.now() - startTime) / 1000}s`,
-      `promptLen: ${formatPrompts.length}`,
-      `contentLen: ${formatPrompts.reduce((sum, item) => sum + item.content.length, 0)}`
-    );
    // 创建响应流
    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, {

--- a/src/pages/api/chat/gpt3.ts
+++ b/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(
+    const modelConstantsData = ModelList.find((item) => item.model === model.service.modelName);
-      (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

--- a/src/pages/api/model/create.ts
+++ b/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 {
+    const { name, serviceModelName } = req.body as {
-      name,
-      serviceModelName,
-      serviceModelCompany = 'openai'
-    } = 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(
+    const modelItem = ModelList.find((item) => item.model === serviceModelName);
-      (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

--- a/src/pages/chat/components/SlideBar.tsx
+++ b/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}>

--- a/src/pages/model/components/CreateModel.tsx
+++ b/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>

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

--- a/src/pages/number/setting.tsx
+++ b/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>
    </>
  );
 };

--- a/src/service/events/bill.ts
+++ b/src/service/events/bill.ts
+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);
+  }
+};
--- a/src/service/models/bill.ts
+++ b/src/service/models/bill.ts
+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);
--- a/src/service/mongo.ts
+++ b/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';
--- a/src/service/utils/chat.ts
+++ b/src/service/utils/chat.ts
 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) {
+  if (!userApiKey && formatPrice(user.balance) <= -1) {
-    return Promise.reject('缺少ApiKey, 无法请求');
+    return Promise.reject('该账号余额不足');
  }
  return {
    userApiKey,
-    chat
+    systemKey: process.env.OPENAIKEY as string,
+    chat,
+    userId: user._id
  };
 };
--- a/src/store/chat.ts
+++ b/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);
          });
        },

--- a/src/store/user.ts
+++ b/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);
      },

--- a/src/utils/user.ts
+++ b/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;
+};