Merge pull request #15 from Kerwin1202/main

feat: support multiple key random usage(Close Chanzhaoyu#155, Close Chanzhaoyu#138)

README.en.md

@@ -21,6 +21,8 @@
 
 [✓] Users manager
 
+[✓] Random Key
+
 </br>
 
 ## Screenshots
@@ -32,6 +34,7 @@
 ![cover3](./docs/basesettings.jpg)
 ![cover3](./docs/prompt_en.jpg)
 ![cover3](./docs/user-manager.jpg)
+![cover3](./docs/key-manager-en.jpg)
 
 - [ChatGPT Web](#chatgpt-web)
 	- [Introduction](#introduction)

README.md

@@ -20,6 +20,8 @@
 [✓] 每个会话设置独有 Prompt
 
 [✓] 用户管理
+
+[✓] 多 Key 随机
 </br>
 
 ## 截图
@@ -31,6 +33,7 @@
 ![cover3](./docs/basesettings.jpg)
 ![cover3](./docs/prompt.jpg)
 ![cover3](./docs/user-manager.jpg)
+![cover3](./docs/key-manager.jpg)
 
 - [ChatGPT Web](#chatgpt-web)
 	- [介绍](#介绍)
@@ -410,7 +413,18 @@ A: 一种可能原因是经过 Nginx 反向代理，开启了 buffer，则 Nginx
 </a>
 
 ## 赞助
-如果你觉得这个项目对你有帮助，请给我点个Star。
+如果你觉得这个项目对你有帮助，请给我点个Star。并且情况允许的话，可以给我一点点支持，总之非常感谢支持～
+
+<div style="display: flex; gap: 20px;">
+	<div style="text-align: center">
+		<img style="max-width: 100%" src="./docs/wechat.png" alt="微信" />
+		<p>WeChat Pay</p>
+	</div>
+	<div style="text-align: center">
+		<img style="max-width: 100%" src="./docs/alipay.png" alt="支付宝" />
+		<p>Alipay</p>
+	</div>
+</div>
 
 ## License
 MIT © [Kerwin1202](./license)

docker-compose/docker-compose.yml

@@ -11,16 +11,6 @@ services:
       - database
     environment:
       TZ: Asia/Shanghai
-      # 二选一
-      OPENAI_API_KEY: 
-      # 二选一
-      OPENAI_ACCESS_TOKEN:
-      # API接口地址，可选，设置 OPENAI_API_KEY 时可用
-      OPENAI_API_BASE_URL:
-      # ChatGPTAPI 或者 ChatGPTUnofficialProxyAPI
-      OPENAI_API_MODEL: 
-      # 反向代理，可选
-      API_REVERSE_PROXY:
       # 访问jwt加密参数，可选 不为空则允许登录 同时需要设置 MONGODB_URL
       AUTH_SECRET_KEY:
       # 每小时最大请求次数，可选，默认无限
@@ -35,8 +25,6 @@ services:
       SOCKS_PROXY_USERNAME:
       # Socks代理密码，可选，和 SOCKS_PROXY_HOST & SOCKS_PROXY_PORT 一起时生效
       SOCKS_PROXY_PASSWORD:
-      # HTTPS_PROXY 代理，可选
-      HTTPS_PROXY: http://xxxx:7890
       # 网站名称
       SITE_TITLE: ChatGpt Web
       # mongodb 的连接字符串

package.json

@@ -1,6 +1,6 @@
 {
   "name": "chatgpt-web",
-  "version": "2.12.4",
+  "version": "2.13.0",
   "private": false,
   "description": "ChatGPT Web",
   "author": "ChenZhaoYu <chenzhaoyu1994@gmail.com>",

service/src/chatgpt/index.ts

@@ -5,14 +5,14 @@ import { ChatGPTAPI, ChatGPTUnofficialProxyAPI } from 'chatgpt'
 import { SocksProxyAgent } from 'socks-proxy-agent'
 import httpsProxyAgent from 'https-proxy-agent'
 import fetch from 'node-fetch'
-import type { AuditConfig, CHATMODEL } from 'src/storage/model'
+import type { AuditConfig, CHATMODEL, KeyConfig, UserInfo } from 'src/storage/model'
 import jwt_decode from 'jwt-decode'
 import dayjs from 'dayjs'
 import type { TextAuditService } from '../utils/textAudit'
 import { textAuditServices } from '../utils/textAudit'
-import { getCacheConfig, getOriginConfig } from '../storage/config'
+import { getCacheApiKeys, getCacheConfig, getOriginConfig } from '../storage/config'
 import { sendResponse } from '../utils'
-import { isNotEmptyString } from '../utils/is'
+import { hasAnyRole, isNotEmptyString } from '../utils/is'
 import type { ChatContext, ChatGPTUnofficialProxyAPIOptions, JWT, ModelConfig } from '../types'
 import { getChatByMessageId } from '../storage/mongo'
 import type { RequestOptions } from './types'
@@ -32,20 +32,17 @@ const ErrorCodeMessage: Record<string, string> = {
 
 let auditService: TextAuditService
 
-export async function initApi(chatModel: CHATMODEL) {
+export async function initApi(key: KeyConfig, chatModel: CHATMODEL) {
   // More Info: https://github.com/transitive-bullshit/chatgpt-api
 
   const config = await getCacheConfig()
-  if (!config.apiKey && !config.accessToken)
-    throw new Error('Missing OPENAI_API_KEY or OPENAI_ACCESS_TOKEN environment variable')
-
   const model = chatModel as string
 
-  if (config.apiModel === 'ChatGPTAPI') {
+  if (key.keyModel === 'ChatGPTAPI') {
     const OPENAI_API_BASE_URL = config.apiBaseUrl
 
     const options: ChatGPTAPIOptions = {
-      apiKey: config.apiKey,
+      apiKey: key.key,
       completionParams: { model },
       debug: !config.apiDisableDebug,
       messageStore: undefined,
@@ -73,7 +70,7 @@ export async function initApi(chatModel: CHATMODEL) {
   }
   else {
     const options: ChatGPTUnofficialProxyAPIOptions = {
-      accessToken: config.accessToken,
+      accessToken: key.key,
       apiReverseProxyUrl: isNotEmptyString(config.reverseProxy) ? config.reverseProxy : 'https://ai.fakeopen.com/api/conversation',
       model,
       debug: !config.apiDisableDebug,
@@ -86,27 +83,30 @@ export async function initApi(chatModel: CHATMODEL) {
 }
 
 async function chatReplyProcess(options: RequestOptions) {
-  const config = await getCacheConfig()
   const model = options.chatModel
+  const key = options.key
+  if (key == null || key === undefined)
+    throw new Error('没有可用的配置。请再试一次 | No available configuration. Please try again.')
+
   const { message, lastContext, process, systemMessage, temperature, top_p } = options
 
   try {
     const timeoutMs = (await getCacheConfig()).timeoutMs
     let options: SendMessageOptions = { timeoutMs }
 
-    if (config.apiModel === 'ChatGPTAPI') {
+    if (key.keyModel === 'ChatGPTAPI') {
       if (isNotEmptyString(systemMessage))
         options.systemMessage = systemMessage
       options.completionParams = { model, temperature, top_p }
     }
 
     if (lastContext != null) {
-      if (config.apiModel === 'ChatGPTAPI')
+      if (key.keyModel === 'ChatGPTAPI')
         options.parentMessageId = lastContext.parentMessageId
       else
         options = { ...lastContext }
     }
-    const api = await initApi(model)
+    const api = await initApi(key, model)
     const response = await api.sendMessage(message, {
       ...options,
       onProgress: (partialResponse) => {
@@ -123,6 +123,9 @@ async function chatReplyProcess(options: RequestOptions) {
       return sendResponse({ type: 'Fail', message: ErrorCodeMessage[code] })
     return sendResponse({ type: 'Fail', message: error.message ?? 'Please check the back-end console' })
   }
+  finally {
+    releaseApiKey(key)
+  }
 }
 
 export function initAuditService(audit: AuditConfig) {
@@ -304,6 +307,39 @@ async function getMessageById(id: string): Promise<ChatMessage | undefined> {
   else { return undefined }
 }
 
+const _lockedKeys: string[] = []
+async function randomKeyConfig(keys: KeyConfig[]): Promise < KeyConfig | null > {
+  if (keys.length <= 0)
+    return null
+  let unsedKeys = keys.filter(d => !_lockedKeys.includes(d.key))
+  const start = Date.now()
+  while (unsedKeys.length <= 0) {
+    if (Date.now() - start > 3000)
+      break
+    await new Promise(resolve => setTimeout(resolve, 1000))
+    unsedKeys = keys.filter(d => !_lockedKeys.includes(d.key))
+  }
+  if (unsedKeys.length <= 0)
+    return null
+  const thisKey = unsedKeys[Math.floor(Math.random() * unsedKeys.length)]
+  _lockedKeys.push(thisKey.key)
+  return thisKey
+}
+
+async function getRandomApiKey(user: UserInfo): Promise<KeyConfig | undefined> {
+  const keys = (await getCacheApiKeys()).filter(d => hasAnyRole(d.userRoles, user.roles))
+  return randomKeyConfig(keys)
+}
+
+async function releaseApiKey(key: KeyConfig) {
+  if (key == null || key === undefined)
+    return
+
+  const index = _lockedKeys.indexOf(key.key)
+  if (index >= 0)
+    _lockedKeys.splice(index, 1)
+}
+
 export type { ChatContext, ChatMessage }
 
-export { chatReplyProcess, chatConfig, containsSensitiveWords }
+export { chatReplyProcess, chatConfig, containsSensitiveWords, getRandomApiKey }

service/src/chatgpt/types.ts

@@ -1,5 +1,5 @@
 import type { ChatMessage } from 'chatgpt'
-import type { CHATMODEL } from 'src/storage/model'
+import type { CHATMODEL, KeyConfig } from 'src/storage/model'
 
 export interface RequestOptions {
   message: string
@@ -9,6 +9,7 @@ export interface RequestOptions {
   temperature?: number
   top_p?: number
   chatModel: CHATMODEL
+  key: KeyConfig
 }
 
 export interface BalanceResponse {

service/src/index.ts

@@ -4,11 +4,11 @@ import * as dotenv from 'dotenv'
 import { ObjectId } from 'mongodb'
 import type { RequestProps } from './types'
 import type { ChatContext, ChatMessage } from './chatgpt'
-import { chatConfig, chatReplyProcess, containsSensitiveWords, initAuditService } from './chatgpt'
+import { chatConfig, chatReplyProcess, containsSensitiveWords, getRandomApiKey, initAuditService } from './chatgpt'
 import { auth } from './middleware/auth'
-import { clearConfigCache, getCacheConfig, getOriginConfig } from './storage/config'
-import type { AuditConfig, CHATMODEL, ChatInfo, ChatOptions, Config, MailConfig, SiteConfig, UsageResponse, UserInfo } from './storage/model'
-import { Status } from './storage/model'
+import { clearConfigCache, getApiKeys, getCacheConfig, getOriginConfig } from './storage/config'
+import type { AuditConfig, CHATMODEL, ChatInfo, ChatOptions, Config, KeyConfig, MailConfig, SiteConfig, UsageResponse, UserInfo } from './storage/model'
+import { Status, UserRole } from './storage/model'
 import {
   clearChat,
   createChatRoom,
@@ -28,6 +28,7 @@ import {
   insertChat,
   insertChatUsage,
   renameChatRoom,
+  updateApiKeyStatus,
   updateChat,
   updateConfig,
   updateRoomPrompt,
@@ -36,6 +37,7 @@ import {
   updateUserInfo,
   updateUserPassword,
   updateUserStatus,
+  upsertKey,
   verifyUser,
 } from './storage/mongo'
 import { limiter } from './middleware/limiter'
@@ -390,7 +392,7 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
     const userId = req.headers.userId.toString()
     const user = await getUserById(userId)
     if (config.auditConfig.enabled || config.auditConfig.customizeEnabled) {
-      if (user.email.toLowerCase() !== process.env.ROOT_USER && await containsSensitiveWords(config.auditConfig, prompt)) {
+      if (!user.roles.includes(UserRole.Admin) && await containsSensitiveWords(config.auditConfig, prompt)) {
         res.send({ status: 'Fail', message: '含有敏感词 | Contains sensitive words', data: null })
         return
       }
@@ -427,12 +429,13 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
       temperature,
       top_p,
       chatModel: user.config.chatModel,
+      key: await getRandomApiKey(user),
     })
     // return the whole response including usage
     res.write(`\n${JSON.stringify(result.data)}`)
   }
   catch (error) {
-    res.write(JSON.stringify(error))
+    res.write(JSON.stringify({ message: error?.message }))
   }
   finally {
     res.end()
@@ -516,9 +519,10 @@ router.post('/user-register', async (req, res) => {
       return
     }
     const newPassword = md5(password)
-    await createUser(username, newPassword)
+    const isRoot = username.toLowerCase() === process.env.ROOT_USER
+    await createUser(username, newPassword, isRoot)
 
-    if (username.toLowerCase() === process.env.ROOT_USER) {
+    if (isRoot) {
       res.send({ status: 'Success', message: '注册成功 | Register success', data: null })
     }
     else {
@@ -536,7 +540,7 @@ router.post('/config', rootAuth, async (req, res) => {
     const userId = req.headers.userId.toString()
 
     const user = await getUserById(userId)
-    if (user == null || user.status !== Status.Normal || user.email.toLowerCase() !== process.env.ROOT_USER)
+    if (user == null || user.status !== Status.Normal || !user.roles.includes(UserRole.Admin))
       throw new Error('无权限 | No permission.')
 
     const response = await chatConfig()
@@ -584,7 +588,7 @@ router.post('/user-login', async (req, res) => {
       avatar: user.avatar,
       description: user.description,
       userId: user._id,
-      root: username.toLowerCase() === process.env.ROOT_USER,
+      root: !user.roles.includes(UserRole.Admin),
       config: user.config,
     }, config.siteConfig.loginSalt.trim())
     res.send({ status: 'Success', message: '登录成功 | Login successfully', data: { token } })
@@ -678,7 +682,10 @@ router.get('/users', rootAuth, async (req, res) => {
 router.post('/user-status', rootAuth, async (req, res) => {
   try {
     const { userId, status } = req.body as { userId: string; status: Status }
+    const user = await getUserById(userId)
     await updateUserStatus(userId, status)
+    if ((user.status === Status.PreVerify || user.status === Status.AdminVerify) && status === Status.Normal)
+      await sendNoticeMail(user.email)
     res.send({ status: 'Success', message: '更新成功 | Update successfully' })
   }
   catch (error) {
@@ -839,6 +846,40 @@ router.post('/audit-test', rootAuth, async (req, res) => {
   }
 })
 
+router.get('/setting-keys', rootAuth, async (req, res) => {
+  try {
+    const result = await getApiKeys()
+    res.send({ status: 'Success', message: null, data: result })
+  }
+  catch (error) {
+    res.send({ status: 'Fail', message: error.message, data: null })
+  }
+})
+
+router.post('/setting-key-status', rootAuth, async (req, res) => {
+  try {
+    const { id, status } = req.body as { id: string; status: Status }
+    await updateApiKeyStatus(id, status)
+    res.send({ status: 'Success', message: '更新成功 | Update successfully' })
+  }
+  catch (error) {
+    res.send({ status: 'Fail', message: error.message, data: null })
+  }
+})
+
+router.post('/setting-key-upsert', rootAuth, async (req, res) => {
+  try {
+    const keyConfig = req.body as KeyConfig
+    if (keyConfig._id !== undefined)
+      keyConfig._id = new ObjectId(keyConfig._id)
+    await upsertKey(keyConfig)
+    res.send({ status: 'Success', message: '成功 | Successfully' })
+  }
+  catch (error) {
+    res.send({ status: 'Fail', message: error.message, data: null })
+  }
+})
+
 router.post('/statistics/by-day', auth, async (req, res) => {
   try {
     const userId = req.headers.userId

service/src/middleware/rootAuth.ts

@@ -1,6 +1,6 @@
 import jwt from 'jsonwebtoken'
 import * as dotenv from 'dotenv'
-import { Status } from '../storage/model'
+import { Status, UserRole } from '../storage/model'
 import { getUserById } from '../storage/mongo'
 import { getCacheConfig } from '../storage/config'
 
@@ -14,7 +14,7 @@ const rootAuth = async (req, res, next) => {
       const info = jwt.verify(token, config.siteConfig.loginSalt.trim())
       req.headers.userId = info.userId
       const user = await getUserById(info.userId)
-      if (user == null || user.status !== Status.Normal || user.email.toLowerCase() !== process.env.ROOT_USER)
+      if (user == null || user.status !== Status.Normal || !user.roles.includes(UserRole.Admin))
         res.send({ status: 'Fail', message: '无权限 | No permission.', data: null })
       else
         next()