fix: access token 模式多 key (Close #203)

Author: Kerwin

service/src/chatgpt/index.ts

@@ -14,7 +14,7 @@ import { getCacheApiKeys, getCacheConfig, getOriginConfig } from '../storage/con
 import { sendResponse } from '../utils'
 import { hasAnyRole, isNotEmptyString } from '../utils/is'
 import type { ChatContext, ChatGPTUnofficialProxyAPIOptions, JWT, ModelConfig } from '../types'
-import { getChatByMessageId } from '../storage/mongo'
+import { getChatByMessageId, updateRoomAccountId } from '../storage/mongo'
 import type { RequestOptions } from './types'
 
 const { HttpsProxyAgent } = httpsProxyAgent
@@ -84,13 +84,22 @@ 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 = await getRandomApiKey(options.user, options.user.config.chatModel)
+  const model = options.user.config.chatModel
+  const key = await getRandomApiKey(options.user, options.user.config.chatModel, options.room.accountId)
   const userId = options.user._id.toString()
   const messageId = options.messageId
   if (key == null || key === undefined)
     throw new Error('没有可用的配置。请再试一次 | No available configuration. Please try again.')
 
+  if (key.keyModel === 'ChatGPTUnofficialProxyAPI') {
+    if (!options.room.accountId)
+      updateRoomAccountId(userId, options.room.roomId, getAccountId(key.key))
+
+    if (options.lastContext && ((options.lastContext.conversationId && !options.lastContext.parentMessageId)
+      || (!options.lastContext.conversationId && options.lastContext.parentMessageId)))
+      throw new Error('无法在一个房间同时使用 AccessToken 以及 Api，请联系管理员，或新开聊天室进行对话 | Unable to use AccessToken and Api at the same time in the same room, please contact the administrator or open a new chat room for conversation')
+  }
+
   const { message, lastContext, process, systemMessage, temperature, top_p } = options
 
   try {
@@ -270,10 +279,10 @@ function formatDate(date) {
 
 async function chatConfig() {
   const config = await getOriginConfig() as ModelConfig
-  if (config.apiModel === 'ChatGPTAPI')
-    config.balance = await fetchBalance()
-  else
-    config.accessTokenExpiredTime = await fetchAccessTokenExpiredTime()
+  // if (config.apiModel === 'ChatGPTAPI')
+  //   config.balance = await fetchBalance()
+  // else
+  //   config.accessTokenExpiredTime = await fetchAccessTokenExpiredTime()
   return sendResponse<ModelConfig>({
     type: 'Success',
     data: config,
@@ -354,9 +363,23 @@ async function randomKeyConfig(keys: KeyConfig[]): Promise<KeyConfig | null> {
   return thisKey
 }
 
-async function getRandomApiKey(user: UserInfo, chatModel: CHATMODEL): Promise<KeyConfig | undefined> {
-  const keys = (await getCacheApiKeys()).filter(d => hasAnyRole(d.userRoles, user.roles))
-  return randomKeyConfig(keys.filter(d => d.chatModels.includes(chatModel)))
+async function getRandomApiKey(user: UserInfo, chatModel: CHATMODEL, accountId?: string): Promise<KeyConfig | undefined> {
+  let keys = (await getCacheApiKeys()).filter(d => hasAnyRole(d.userRoles, user.roles))
+    .filter(d => d.chatModels.includes(chatModel))
+  if (accountId)
+    keys = keys.filter(d => d.keyModel === 'ChatGPTUnofficialProxyAPI' && getAccountId(d.key) === accountId)
+
+  return randomKeyConfig(keys)
+}
+
+function getAccountId(accessToken: string): string {
+  try {
+    const jwt = jwt_decode(accessToken) as JWT
+    return jwt['https://api.openai.com/auth'].user_id
+  }
+  catch (error) {
+    return ''
+  }
 }
 
 export type { ChatContext, ChatMessage }

service/src/chatgpt/types.ts

@@ -1,5 +1,5 @@
 import type { ChatMessage } from 'chatgpt'
-import type { CHATMODEL, UserInfo } from 'src/storage/model'
+import type { ChatRoom, UserInfo } from 'src/storage/model'
 
 export interface RequestOptions {
   message: string
@@ -8,10 +8,10 @@ export interface RequestOptions {
   systemMessage?: string
   temperature?: number
   top_p?: number
-  chatModel: CHATMODEL
   user: UserInfo
   messageId: string
   tryCount: number
+  room: ChatRoom
 }
 
 export interface BalanceResponse {

service/src/index.ts

@@ -429,10 +429,10 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
       systemMessage,
       temperature,
       top_p,
-      chatModel: user.config.chatModel,
       user,
       messageId: message._id.toString(),
       tryCount: 0,
+      room,
     })
     // return the whole response including usage
     res.write(`\n${JSON.stringify(result.data)}`)

service/src/storage/model.ts

@@ -96,12 +96,15 @@ export class ChatRoom {
   prompt: string
   usingContext: boolean
   status: Status = Status.Normal
+  // only access token used
+  accountId?: string
   constructor(userId: string, title: string, roomId: number) {
     this.userId = userId
     this.title = title
     this.prompt = undefined
     this.roomId = roomId
     this.usingContext = true
+    this.accountId = null
   }
 }
 

service/src/storage/mongo.ts

@@ -108,6 +108,17 @@ export async function updateRoomUsingContext(userId: string, roomId: number, usi
   return result.modifiedCount > 0
 }
 
+export async function updateRoomAccountId(userId: string, roomId: number, accountId: string) {
+  const query = { userId, roomId }
+  const update = {
+    $set: {
+      accountId,
+    },
+  }
+  const result = await roomCol.updateOne(query, update)
+  return result.modifiedCount > 0
+}
+
 export async function getChatRooms(userId: string) {
   const cursor = await roomCol.find({ userId, status: { $ne: Status.Deleted } })
   const rooms = []