perf: api调用和余额校验

56dab7ab · archer · ed1f93d8 · 56dab7ab · 56dab7ab · 56dab7ab
Commit 56dab7ab authored Mar 31, 2023 by archer
10 changed files
--- a/src/pages/api/model/data/splitData.ts
+++ b/src/pages/api/model/data/splitData.ts
@@ -45,6 +45,10 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
      }
    });
+    if (splitText) {
+      textList.push(splitText);
+    }
    // 批量插入数据
    await SplitData.create({
      userId,
@@ -55,9 +59,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
    generateQA();
-    jsonRes(res, {
+    jsonRes(res);
-      data: { chunks, replaceText }
-    });
  } catch (err) {
    jsonRes(res, {
      code: 500,

--- a/src/pages/api/model/del.ts
+++ b/src/pages/api/model/del.ts
 import type { NextApiRequest, NextApiResponse } from 'next';
 import { jsonRes } from '@/service/response';
 import { Chat, Model, Training, connectToDatabase, ModelData } from '@/service/mongo';
-import { authToken, getUserOpenaiKey } from '@/service/utils/tools';
+import { authToken, getUserApiOpenai } from '@/service/utils/tools';
 import { TrainingStatusEnum } from '@/constants/model';
 import { getOpenAIApi } from '@/service/utils/chat';
 import { TrainingItemType } from '@/types/training';
@@ -67,7 +67,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse<
    // 如果正在训练，需要删除openai上的相关信息
    if (training) {
-      const openai = getOpenAIApi(await getUserOpenaiKey(userId));
+      const { openai } = await getUserApiOpenai(userId);
      // 获取训练记录
      const tuneRecord = await openai.retrieveFineTune(training.tuneId, { httpsAgent });

--- a/src/pages/api/model/train/putTrainStatus.ts
+++ b/src/pages/api/model/train/putTrainStatus.ts
@@ -2,7 +2,7 @@ import type { NextApiRequest, NextApiResponse } from 'next';
 import { jsonRes } from '@/service/response';
 import { connectToDatabase, Model, Training } from '@/service/mongo';
 import { getOpenAIApi } from '@/service/utils/chat';
-import { authToken, getUserOpenaiKey } from '@/service/utils/tools';
+import { authToken, getUserApiOpenai } from '@/service/utils/tools';
 import type { ModelSchema } from '@/types/mongoSchema';
 import { TrainingItemType } from '@/types/training';
 import { ModelStatusEnum, TrainingStatusEnum } from '@/constants/model';
@@ -43,7 +43,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
    }
    // 用户的 openai 实例
-    const openai = getOpenAIApi(await getUserOpenaiKey(userId));
+    const { openai } = await getUserApiOpenai(userId);
    // 获取 openai 的训练情况
    const { data } = await openai.retrieveFineTune(training.tuneId, { httpsAgent });

--- a/src/pages/api/model/train/train.ts
+++ b/src/pages/api/model/train/train.ts
@@ -2,9 +2,8 @@
 import type { NextApiRequest, NextApiResponse } from 'next';
 import { jsonRes } from '@/service/response';
 import { connectToDatabase, Model, Training } from '@/service/mongo';
-import { getOpenAIApi } from '@/service/utils/chat';
 import formidable from 'formidable';
-import { authToken, getUserOpenaiKey } from '@/service/utils/tools';
+import { authToken, getUserApiOpenai } from '@/service/utils/tools';
 import { join } from 'path';
 import fs from 'fs';
 import type { ModelSchema } from '@/types/mongoSchema';
@@ -49,7 +48,8 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
    const trainingType = model.service.trainId; // 目前都默认是 openai text-davinci-03
    // 获取用户的 API Key 实例化后的对象
-    openai = getOpenAIApi(await getUserOpenaiKey(userId));
+    const user = await getUserApiOpenai(userId);
+    openai = user.openai;
    // 接收文件并保存
    const form = formidable({

--- a/src/pages/api/timer/updateTraining.ts
+++ b/src/pages/api/timer/updateTraining.ts
@@ -5,7 +5,7 @@ import { connectToDatabase, Training, Model } from '@/service/mongo';
 import type { TrainingItemType } from '@/types/training';
 import { TrainingStatusEnum, ModelStatusEnum } from '@/constants/model';
 import { getOpenAIApi } from '@/service/utils/chat';
-import { getUserOpenaiKey } from '@/service/utils/tools';
+import { getUserApiOpenai } from '@/service/utils/tools';
 import { OpenAiTuneStatusEnum } from '@/service/constants/training';
 import { sendTrainSucceed } from '@/service/utils/sendEmail';
 import { httpsAgent } from '@/service/utils/tools';
@@ -23,7 +23,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
      status: TrainingStatusEnum.pending
    });
-    const openai = getOpenAIApi(await getUserOpenaiKey('63f9a14228d2a688d8dc9e1b'));
+    const { openai } = await getUserApiOpenai('63f9a14228d2a688d8dc9e1b');
    const response = await Promise.all(
      trainingRecords.map(async (item) => {

--- a/src/service/events/generateAbstract.ts
+++ b/src/service/events/generateAbstract.ts
@@ -43,12 +43,15 @@ export async function generateAbstract(next = false): Promise<any> {
      const key = await getOpenApiKey(dataItem.userId);
      userApiKey = key.userApiKey;
      systemKey = key.systemKey;
-    } catch (error) {
+    } catch (error: any) {
+      if (error?.code === 501) {
        // 余额不够了, 把用户所有记录改成闲置
        await DataItem.updateMany({
          userId: dataItem.userId,
          status: 0
        });
+      }
      throw new Error('获取 openai key 失败');
    }

--- a/src/service/events/generateQA.ts
+++ b/src/service/events/generateQA.ts
@@ -40,12 +40,15 @@ export async function generateQA(next = false): Promise<any> {
      const key = await getOpenApiKey(dataItem.userId);
      userApiKey = key.userApiKey;
      systemKey = key.systemKey;
-    } catch (error) {
+    } catch (error: any) {
+      if (error?.code === 501) {
        // 余额不够了, 清空该记录
        await SplitData.findByIdAndUpdate(dataItem._id, {
          textList: [],
-        errorText: '余额不足，生成数据集任务终止'
+          errorText: error.message
        });
+      }
      throw new Error('获取 openai key 失败');
    }
@@ -121,7 +124,7 @@ export async function generateQA(next = false): Promise<any> {
    setTimeout(() => {
      generateQA(true);
-    }, 10000);
+    }, 5000);
  }
 }

--- a/src/service/events/generateVector.ts
+++ b/src/service/events/generateVector.ts
@@ -85,5 +85,9 @@ export async function generateVector(next = false): Promise<any> {
        generateVector(true);
      }, 60000);
    }
+    setTimeout(() => {
+      generateVector(true);
+    }, 3000);
  }
 }
--- 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';
+import { authToken, getOpenApiKey } from './tools';
-import { authToken } from './tools';
 export const getOpenAIApi = (apiKey: string) => {
  const configuration = new Configuration({
@@ -14,18 +13,11 @@ export const getOpenAIApi = (apiKey: string) => {
 export const authChat = async (chatId: string, authorization?: string) => {
  // 获取 chat 数据
-  const chat = await Chat.findById<ChatPopulate>(chatId)
+  const chat = await Chat.findById<ChatPopulate>(chatId).populate({
-    .populate({
    path: 'modelId',
    options: {
      strictPopulate: false
    }
-    })
-    .populate({
-      path: 'userId',
-      options: {
-        strictPopulate: false
-      }
  });
  if (!chat || !chat.modelId || !chat.userId) {
@@ -43,21 +35,14 @@ export const authChat = async (chatId: string, authorization?: string) => {
  }
  // 获取 user 的 apiKey
-  const user = chat.userId;
+  const { user, userApiKey, systemKey } = await getOpenApiKey(chat.userId as unknown as string);
-  const userApiKey = user.accounts?.find((item: any) => item.type === 'openai')?.value;
-  // 没有 apikey ，校验余额
-  if (!userApiKey && formatPrice(user.balance) <= 0) {
-    return Promise.reject('该账号余额不足');
-  }
  // filter 掉被 deleted 的内容
  chat.content = chat.content.filter((item) => item.deleted !== true);
  return {
    userApiKey,
-    systemKey: process.env.OPENAIKEY as string,
+    systemKey,
    chat,
    userId: user._id
  };

--- a/src/service/utils/tools.ts
+++ b/src/service/utils/tools.ts
@@ -2,10 +2,12 @@ import crypto from 'crypto';
 import jwt from 'jsonwebtoken';
 import { User } from '../models/user';
 import tunnel from 'tunnel';
-import type { UserModelSchema } from '@/types/mongoSchema';
 import { formatPrice } from '@/utils/user';
 import { ChatItemType } from '@/types/chat';
 import { encode } from 'gpt-token-utils';
+import { getOpenAIApi } from '@/service/utils/chat';
+import axios from 'axios';
+import { UserModelSchema } from '@/types/mongoSchema';
 /* 密码加密 */
 export const hashPassword = (psw: string) => {
@@ -44,40 +46,83 @@ export const authToken = (token?: string): Promise<string> => {
  });
 };
-/* 获取用户的 openai APIkey */
+/* 判断 apikey 是否还有余额 */
-export const getUserOpenaiKey = async (userId: string) => {
+export const checkKeyGrant = async (apiKey: string) => {
+  const grant = await axios.get('https://api.openai.com/dashboard/billing/credit_grants', {
+    headers: {
+      Authorization: `Bearer ${apiKey}`
+    }
+  });
+  console.log(grant.data?.total_available);
+  if (grant.data?.total_available <= 0.2) {
+    return false;
+  }
+  return true;
+};
+/* 获取用户 api 的 openai 信息 */
+export const getUserApiOpenai = async (userId: string) => {
  const user = await User.findById(userId);
  const userApiKey = user?.accounts?.find((item: any) => item.type === 'openai')?.value;
  if (!userApiKey) {
    return Promise.reject('缺少ApiKey, 无法请求');
  }
-  return Promise.resolve(userApiKey as string);
+  // 余额校验
+  const hasGrant = await checkKeyGrant(userApiKey);
+  if (!hasGrant) {
+    return Promise.reject({
+      code: 501,
+      message: 'API 余额不足'
+    });
+  }
+  return {
+    user,
+    openai: getOpenAIApi(userApiKey),
+    apiKey: userApiKey
+  };
 };
-/* 获取key，如果没有就用平台的，用平台记得加账单 */
+/* 获取 open api key，如果用户没有自己的key，就用平台的，用平台记得加账单 */
 export const getOpenApiKey = async (userId: string) => {
-  const user = await User.findById<UserModelSchema>(userId);
+  const user = await User.findById(userId);
+  if (!user) {
-  if (!user) return Promise.reject('用户不存在');
+    return Promise.reject('找不到用户');
+  }
-  const userApiKey = user.accounts?.find((item: any) => item.type === 'openai')?.value;
+  const userApiKey = user?.accounts?.find((item: any) => item.type === 'openai')?.value;
-  // 有自己的key， 直接使用
+  // 有自己的key
  if (userApiKey) {
+    // api 余额校验
+    const hasGrant = await checkKeyGrant(userApiKey);
+    if (!hasGrant) {
+      return Promise.reject({
+        code: 501,
+        message: 'API 余额不足'
+      });
+    }
    return {
-      userApiKey: await getUserOpenaiKey(userId),
+      user,
+      userApiKey,
      systemKey: ''
    };
  }
-  // 余额校验
+  // 平台账号余额校验
  if (formatPrice(user.balance) <= 0) {
-    return Promise.reject('该账号余额不足');
+    return Promise.reject({
+      code: 501,
+      message: '账号余额不足'
+    });
  }
  return {
+    user,
    userApiKey: '',
    systemKey: process.env.OPENAIKEY as string
  };