From 6c4026ccef1b71084b9bcc1cffaa158d7cc25aea Mon Sep 17 00:00:00 2001 From: archer <545436317@qq.com> Date: Mon, 3 Apr 2023 10:20:17 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E6=96=87=E4=BB=B6=E7=BB=93=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pages/api/chat/chatGpt.ts | 3 +- src/pages/api/model/del.ts | 3 +- src/pages/api/model/train/putTrainStatus.ts | 4 +- src/pages/api/model/train/train.ts | 3 +- src/pages/api/timer/updateTraining.ts | 3 +- src/service/events/generateAbstract.ts | 3 +- src/service/events/generateQA.ts | 3 +- src/service/utils/chat.ts | 3 +- src/service/utils/openai.ts | 89 +++++++++++++++++++++ src/service/utils/tools.ts | 88 -------------------- 10 files changed, 103 insertions(+), 99 deletions(-) create mode 100644 src/service/utils/openai.ts diff --git a/src/pages/api/chat/chatGpt.ts b/src/pages/api/chat/chatGpt.ts index f363318fb..118263715 100644 --- a/src/pages/api/chat/chatGpt.ts +++ b/src/pages/api/chat/chatGpt.ts @@ -2,7 +2,7 @@ import type { NextApiRequest, NextApiResponse } from 'next'; import { createParser, ParsedEvent, ReconnectInterval } from 'eventsource-parser'; import { connectToDatabase } from '@/service/mongo'; import { getOpenAIApi, authChat } from '@/service/utils/chat'; -import { httpsAgent } from '@/service/utils/tools'; +import { httpsAgent, openaiChatFilter } from '@/service/utils/tools'; import { ChatCompletionRequestMessage, ChatCompletionRequestMessageRoleEnum } from 'openai'; import { ChatItemType } from '@/types/chat'; import { jsonRes } from '@/service/response'; @@ -10,7 +10,6 @@ import type { ModelSchema } from '@/types/mongoSchema'; import { PassThrough } from 'stream'; import { modelList } from '@/constants/model'; import { pushChatBill } from '@/service/events/pushBill'; -import { openaiChatFilter } from '@/service/utils/tools'; /* 发送提示词 */ export default async function handler(req: NextApiRequest, res: NextApiResponse) { diff --git a/src/pages/api/model/del.ts b/src/pages/api/model/del.ts index aeb367f44..e97d5557a 100644 --- a/src/pages/api/model/del.ts +++ b/src/pages/api/model/del.ts @@ -1,7 +1,8 @@ import type { NextApiRequest, NextApiResponse } from 'next'; import { jsonRes } from '@/service/response'; import { Chat, Model, Training, connectToDatabase } from '@/service/mongo'; -import { authToken, getUserApiOpenai } from '@/service/utils/tools'; +import { authToken } from '@/service/utils/tools'; +import { getUserApiOpenai } from '@/service/utils/openai'; import { TrainingStatusEnum } from '@/constants/model'; import { TrainingItemType } from '@/types/training'; import { httpsAgent } from '@/service/utils/tools'; diff --git a/src/pages/api/model/train/putTrainStatus.ts b/src/pages/api/model/train/putTrainStatus.ts index 922bab885..413055c90 100644 --- a/src/pages/api/model/train/putTrainStatus.ts +++ b/src/pages/api/model/train/putTrainStatus.ts @@ -1,8 +1,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 { authToken, getUserApiOpenai } from '@/service/utils/tools'; +import { authToken } from '@/service/utils/tools'; +import { getUserApiOpenai } from '@/service/utils/openai'; import type { ModelSchema } from '@/types/mongoSchema'; import { TrainingItemType } from '@/types/training'; import { ModelStatusEnum, TrainingStatusEnum } from '@/constants/model'; diff --git a/src/pages/api/model/train/train.ts b/src/pages/api/model/train/train.ts index 8e2e76cab..814f31a75 100644 --- a/src/pages/api/model/train/train.ts +++ b/src/pages/api/model/train/train.ts @@ -3,7 +3,8 @@ import type { NextApiRequest, NextApiResponse } from 'next'; import { jsonRes } from '@/service/response'; import { connectToDatabase, Model, Training } from '@/service/mongo'; import formidable from 'formidable'; -import { authToken, getUserApiOpenai } from '@/service/utils/tools'; +import { authToken } from '@/service/utils/tools'; +import { getUserApiOpenai } from '@/service/utils/openai'; import { join } from 'path'; import fs from 'fs'; import type { ModelSchema } from '@/types/mongoSchema'; diff --git a/src/pages/api/timer/updateTraining.ts b/src/pages/api/timer/updateTraining.ts index dc418f704..a6140554a 100644 --- a/src/pages/api/timer/updateTraining.ts +++ b/src/pages/api/timer/updateTraining.ts @@ -4,8 +4,7 @@ import { jsonRes } from '@/service/response'; 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 { getUserApiOpenai } from '@/service/utils/tools'; +import { getUserApiOpenai } from '@/service/utils/openai'; import { OpenAiTuneStatusEnum } from '@/service/constants/training'; import { sendTrainSucceed } from '@/service/utils/sendEmail'; import { httpsAgent } from '@/service/utils/tools'; diff --git a/src/service/events/generateAbstract.ts b/src/service/events/generateAbstract.ts index 4039a2527..a6b50c6b8 100644 --- a/src/service/events/generateAbstract.ts +++ b/src/service/events/generateAbstract.ts @@ -1,6 +1,7 @@ import { DataItem } from '@/service/mongo'; import { getOpenAIApi } from '@/service/utils/chat'; -import { httpsAgent, getOpenApiKey } from '@/service/utils/tools'; +import { httpsAgent } from '@/service/utils/tools'; +import { getOpenApiKey } from '../utils/openai'; import type { ChatCompletionRequestMessage } from 'openai'; import { DataItemSchema } from '@/types/mongoSchema'; import { ChatModelNameEnum } from '@/constants/model'; diff --git a/src/service/events/generateQA.ts b/src/service/events/generateQA.ts index d9e0940aa..3ec568de8 100644 --- a/src/service/events/generateQA.ts +++ b/src/service/events/generateQA.ts @@ -1,6 +1,7 @@ import { SplitData } from '@/service/mongo'; import { getOpenAIApi } from '@/service/utils/chat'; -import { httpsAgent, getOpenApiKey } from '@/service/utils/tools'; +import { httpsAgent } from '@/service/utils/tools'; +import { getOpenApiKey } from '../utils/openai'; import type { ChatCompletionRequestMessage } from 'openai'; import { ChatModelNameEnum } from '@/constants/model'; import { pushSplitDataBill } from '@/service/events/pushBill'; diff --git a/src/service/utils/chat.ts b/src/service/utils/chat.ts index 5f5e11b85..237404f89 100644 --- a/src/service/utils/chat.ts +++ b/src/service/utils/chat.ts @@ -1,7 +1,8 @@ import { Configuration, OpenAIApi } from 'openai'; import { Chat } from '../mongo'; import type { ChatPopulate } from '@/types/mongoSchema'; -import { authToken, getOpenApiKey } from './tools'; +import { authToken } from './tools'; +import { getOpenApiKey } from './openai'; export const getOpenAIApi = (apiKey: string) => { const configuration = new Configuration({ diff --git a/src/service/utils/openai.ts b/src/service/utils/openai.ts new file mode 100644 index 000000000..f373bc965 --- /dev/null +++ b/src/service/utils/openai.ts @@ -0,0 +1,89 @@ +import axios from 'axios'; +import { getOpenAIApi } from '@/service/utils/chat'; +import { httpsAgent } from './tools'; +import { User } from '../models/user'; +import { formatPrice } from '@/utils/user'; + +/* 判断 apikey 是否还有余额 */ +export const checkKeyGrant = async (apiKey: string) => { + const grant = await axios.get('https://api.openai.com/dashboard/billing/credit_grants', { + headers: { + Authorization: `Bearer ${apiKey}` + }, + httpsAgent + }); + 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, 无法请求'); + } + + // 余额校验 + const hasGrant = await checkKeyGrant(userApiKey); + if (!hasGrant) { + return Promise.reject({ + code: 501, + message: 'API 余额不足' + }); + } + + return { + user, + openai: getOpenAIApi(userApiKey), + apiKey: userApiKey + }; +}; + +/* 获取 open api key，如果用户没有自己的key，就用平台的，用平台记得加账单 */ +export const getOpenApiKey = async (userId: string, checkGrant = false) => { + const user = await User.findById(userId); + if (!user) { + return Promise.reject('找不到用户'); + } + + const userApiKey = user?.accounts?.find((item: any) => item.type === 'openai')?.value; + + // 有自己的key + if (userApiKey) { + // api 余额校验 + if (checkGrant) { + const hasGrant = await checkKeyGrant(userApiKey); + if (!hasGrant) { + return Promise.reject({ + code: 501, + message: 'API 余额不足' + }); + } + } + + return { + user, + userApiKey, + systemKey: '' + }; + } + + // 平台账号余额校验 + if (formatPrice(user.balance) <= 0) { + return Promise.reject({ + code: 501, + message: '账号余额不足' + }); + } + + return { + user, + userApiKey: '', + systemKey: process.env.OPENAIKEY as string + }; +}; diff --git a/src/service/utils/tools.ts b/src/service/utils/tools.ts index a208b5cc5..7cd27d909 100644 --- a/src/service/utils/tools.ts +++ b/src/service/utils/tools.ts @@ -1,12 +1,8 @@ import crypto from 'crypto'; import jwt from 'jsonwebtoken'; -import { User } from '../models/user'; import tunnel from 'tunnel'; -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'; /* 密码加密 */ export const hashPassword = (psw: string) => { @@ -56,90 +52,6 @@ export const httpsAgent = }) : undefined; -/* 判断 apikey 是否还有余额 */ -export const checkKeyGrant = async (apiKey: string) => { - const grant = await axios.get('https://api.openai.com/dashboard/billing/credit_grants', { - headers: { - Authorization: `Bearer ${apiKey}` - }, - httpsAgent - }); - 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, 无法请求'); - } - - // 余额校验 - const hasGrant = await checkKeyGrant(userApiKey); - if (!hasGrant) { - return Promise.reject({ - code: 501, - message: 'API 余额不足' - }); - } - - return { - user, - openai: getOpenAIApi(userApiKey), - apiKey: userApiKey - }; -}; - -/* 获取 open api key，如果用户没有自己的key，就用平台的，用平台记得加账单 */ -export const getOpenApiKey = async (userId: string, checkGrant = false) => { - const user = await User.findById(userId); - if (!user) { - return Promise.reject('找不到用户'); - } - - const userApiKey = user?.accounts?.find((item: any) => item.type === 'openai')?.value; - - // 有自己的key - if (userApiKey) { - // api 余额校验 - if (checkGrant) { - const hasGrant = await checkKeyGrant(userApiKey); - if (!hasGrant) { - return Promise.reject({ - code: 501, - message: 'API 余额不足' - }); - } - } - - return { - user, - userApiKey, - systemKey: '' - }; - } - - // 平台账号余额校验 - if (formatPrice(user.balance) <= 0) { - return Promise.reject({ - code: 501, - message: '账号余额不足' - }); - } - - return { - user, - userApiKey: '', - systemKey: process.env.OPENAIKEY as string - }; -}; - /* tokens 截断 */ export const openaiChatFilter = (prompts: ChatItemType[], maxTokens: number) => { let res: ChatItemType[] = [];