fix: 优化随机key的逻辑 (Close #204)

Author: Kerwin

service/src/chatgpt/index.ts

@@ -31,6 +31,7 @@ const ErrorCodeMessage: Record<string, string> = {
 }
 
 let auditService: TextAuditService
+const _lockedKeys: { key: string; lockedTime: number }[] = []
 
 export async function initApi(key: KeyConfig, chatModel: CHATMODEL) {
   // More Info: https://github.com/transitive-bullshit/chatgpt-api
@@ -84,8 +85,8 @@ export async function initApi(key: KeyConfig, chatModel: CHATMODEL) {
 const processThreads: { userId: string; abort: AbortController; messageId: string }[] = []
 async function chatReplyProcess(options: RequestOptions) {
   const model = options.chatModel
-  const key = options.key
-  const userId = options.userId
+  const key = await getRandomApiKey(options.user, options.user.config.chatModel)
+  const userId = options.user._id.toString()
   const messageId = options.messageId
   if (key == null || key === undefined)
     throw new Error('没有可用的配置。请再试一次 | No available configuration. Please try again.')
@@ -124,13 +125,20 @@ async function chatReplyProcess(options: RequestOptions) {
   }
   catch (error: any) {
     const code = error.statusCode
-    global.console.log(error)
+    if (code === 429 && (error.message.includes('Too Many Requests') || error.message.includes('Rate limit'))) {
+      // access token  Only one message at a time
+      if (options.tryCount++ < 3) {
+        _lockedKeys.push({ key: key.key, lockedTime: Date.now() })
+        await new Promise(resolve => setTimeout(resolve, 2000))
+        return await chatReplyProcess(options)
+      }
+    }
+    global.console.error(error)
     if (Reflect.has(ErrorCodeMessage, code))
       return sendResponse({ type: 'Fail', message: ErrorCodeMessage[code] })
     return sendResponse({ type: 'Fail', message: error.message ?? 'Please check the back-end console' })
   }
   finally {
-    releaseApiKey(key)
     const index = processThreads.findIndex(d => d.userId === userId)
     if (index > -1)
       processThreads.splice(index, 1)
@@ -326,29 +334,23 @@ async function getMessageById(id: string): Promise<ChatMessage | undefined> {
   else { return undefined }
 }
 
-const _lockedKeys: { key: string; count: number }[] = []
-const _oneTimeCount = 3 // api
 async function randomKeyConfig(keys: KeyConfig[]): Promise<KeyConfig | null> {
   if (keys.length <= 0)
     return null
-  let unsedKeys = keys.filter(d => _lockedKeys.filter(l => d.key === l.key).length <= 0
-    || _lockedKeys.filter(l => d.key === l.key)[0].count < _oneTimeCount)
+  // cleanup old locked keys
+  _lockedKeys.filter(d => d.lockedTime <= Date.now() - 1000 * 20).forEach(d => _lockedKeys.splice(_lockedKeys.indexOf(d), 1))
+
+  let unsedKeys = keys.filter(d => _lockedKeys.filter(l => d.key === l.key).length <= 0)
   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.filter(l => d.key === l.key).length <= 0
-      || _lockedKeys.filter(l => d.key === l.key)[0].count < _oneTimeCount)
+    unsedKeys = keys.filter(d => _lockedKeys.filter(l => d.key === l.key).length <= 0)
   }
   if (unsedKeys.length <= 0)
     return null
   const thisKey = unsedKeys[Math.floor(Math.random() * unsedKeys.length)]
-  const thisLockedKey = _lockedKeys.filter(d => d.key === thisKey.key)
-  if (thisLockedKey.length <= 0)
-    _lockedKeys.push({ key: thisKey.key, count: 1 })
-  else
-    thisLockedKey[0].count++
   return thisKey
 }
 
@@ -357,23 +359,6 @@ async function getRandomApiKey(user: UserInfo, chatModel: CHATMODEL): Promise<Ke
   return randomKeyConfig(keys.filter(d => d.chatModels.includes(chatModel)))
 }
 
-async function releaseApiKey(key: KeyConfig) {
-  if (key == null || key === undefined)
-    return
-
-  const lockedKeys = _lockedKeys.filter(d => d.key === key.key)
-  if (lockedKeys.length > 0) {
-    if (lockedKeys[0].count <= 1) {
-      const index = _lockedKeys.findIndex(item => item.key === key.key)
-      if (index !== -1)
-        _lockedKeys.splice(index, 1)
-    }
-    else {
-      lockedKeys[0].count--
-    }
-  }
-}
-
 export type { ChatContext, ChatMessage }
 
-export { chatReplyProcess, chatConfig, containsSensitiveWords, getRandomApiKey }
+export { chatReplyProcess, chatConfig, containsSensitiveWords }

service/src/chatgpt/types.ts

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

service/src/index.ts

@@ -4,7 +4,7 @@ import * as dotenv from 'dotenv'
 import { ObjectId } from 'mongodb'
 import type { RequestProps } from './types'
 import type { ChatContext, ChatMessage } from './chatgpt'
-import { abortChatProcess, chatConfig, chatReplyProcess, containsSensitiveWords, getRandomApiKey, initAuditService } from './chatgpt'
+import { abortChatProcess, chatConfig, chatReplyProcess, containsSensitiveWords, initAuditService } from './chatgpt'
 import { auth, getUserId } from './middleware/auth'
 import { clearApiKeyCache, clearConfigCache, getApiKeys, getCacheApiKeys, getCacheConfig, getOriginConfig } from './storage/config'
 import type { AuditConfig, CHATMODEL, ChatInfo, ChatOptions, Config, KeyConfig, MailConfig, SiteConfig, UsageResponse, UserInfo } from './storage/model'
@@ -430,9 +430,9 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
       temperature,
       top_p,
       chatModel: user.config.chatModel,
-      key: await getRandomApiKey(user, user.config.chatModel),
-      userId,
+      user,
       messageId: message._id.toString(),
+      tryCount: 0,
     })
     // return the whole response including usage
     res.write(`\n${JSON.stringify(result.data)}`)
@@ -480,7 +480,7 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
       }
     }
     catch (error) {
-      global.console.log(error)
+      global.console.error(error)
     }
   }
 })