feat: 增加账单

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