@@ -5,14 +5,14 @@ import { ChatGPTAPI, ChatGPTUnofficialProxyAPI } from 'chatgpt'
5
5
import { SocksProxyAgent } from 'socks-proxy-agent'
6
6
import httpsProxyAgent from 'https-proxy-agent'
7
7
import fetch from 'node-fetch'
8
- import type { AuditConfig , CHATMODEL } from 'src/storage/model'
8
+ import type { AuditConfig , CHATMODEL , KeyConfig , UserInfo } from 'src/storage/model'
9
9
import jwt_decode from 'jwt-decode'
10
10
import dayjs from 'dayjs'
11
11
import type { TextAuditService } from '../utils/textAudit'
12
12
import { textAuditServices } from '../utils/textAudit'
13
- import { getCacheConfig , getOriginConfig } from '../storage/config'
13
+ import { getCacheApiKeys , getCacheConfig , getOriginConfig } from '../storage/config'
14
14
import { sendResponse } from '../utils'
15
- import { isNotEmptyString } from '../utils/is'
15
+ import { hasAnyRole , isNotEmptyString } from '../utils/is'
16
16
import type { ChatContext , ChatGPTUnofficialProxyAPIOptions , JWT , ModelConfig } from '../types'
17
17
import { getChatByMessageId } from '../storage/mongo'
18
18
import type { RequestOptions } from './types'
@@ -32,20 +32,17 @@ const ErrorCodeMessage: Record<string, string> = {
32
32
33
33
let auditService : TextAuditService
34
34
35
- export async function initApi ( chatModel : CHATMODEL ) {
35
+ export async function initApi ( key : KeyConfig , chatModel : CHATMODEL ) {
36
36
// More Info: https://github.com/transitive-bullshit/chatgpt-api
37
37
38
38
const config = await getCacheConfig ( )
39
- if ( ! config . apiKey && ! config . accessToken )
40
- throw new Error ( 'Missing OPENAI_API_KEY or OPENAI_ACCESS_TOKEN environment variable' )
41
-
42
39
const model = chatModel as string
43
40
44
- if ( config . apiModel === 'ChatGPTAPI' ) {
41
+ if ( key . keyModel === 'ChatGPTAPI' ) {
45
42
const OPENAI_API_BASE_URL = config . apiBaseUrl
46
43
47
44
const options : ChatGPTAPIOptions = {
48
- apiKey : config . apiKey ,
45
+ apiKey : key . key ,
49
46
completionParams : { model } ,
50
47
debug : ! config . apiDisableDebug ,
51
48
messageStore : undefined ,
@@ -73,7 +70,7 @@ export async function initApi(chatModel: CHATMODEL) {
73
70
}
74
71
else {
75
72
const options : ChatGPTUnofficialProxyAPIOptions = {
76
- accessToken : config . accessToken ,
73
+ accessToken : key . key ,
77
74
apiReverseProxyUrl : isNotEmptyString ( config . reverseProxy ) ? config . reverseProxy : 'https://ai.fakeopen.com/api/conversation' ,
78
75
model,
79
76
debug : ! config . apiDisableDebug ,
@@ -86,27 +83,30 @@ export async function initApi(chatModel: CHATMODEL) {
86
83
}
87
84
88
85
async function chatReplyProcess ( options : RequestOptions ) {
89
- const config = await getCacheConfig ( )
90
86
const model = options . chatModel
87
+ const key = options . key
88
+ if ( key == null || key === undefined )
89
+ throw new Error ( '没有可用的配置。请再试一次 | No available configuration. Please try again.' )
90
+
91
91
const { message, lastContext, process, systemMessage, temperature, top_p } = options
92
92
93
93
try {
94
94
const timeoutMs = ( await getCacheConfig ( ) ) . timeoutMs
95
95
let options : SendMessageOptions = { timeoutMs }
96
96
97
- if ( config . apiModel === 'ChatGPTAPI' ) {
97
+ if ( key . keyModel === 'ChatGPTAPI' ) {
98
98
if ( isNotEmptyString ( systemMessage ) )
99
99
options . systemMessage = systemMessage
100
100
options . completionParams = { model, temperature, top_p }
101
101
}
102
102
103
103
if ( lastContext != null ) {
104
- if ( config . apiModel === 'ChatGPTAPI' )
104
+ if ( key . keyModel === 'ChatGPTAPI' )
105
105
options . parentMessageId = lastContext . parentMessageId
106
106
else
107
107
options = { ...lastContext }
108
108
}
109
- const api = await initApi ( model )
109
+ const api = await initApi ( key , model )
110
110
const response = await api . sendMessage ( message , {
111
111
...options ,
112
112
onProgress : ( partialResponse ) => {
@@ -123,6 +123,9 @@ async function chatReplyProcess(options: RequestOptions) {
123
123
return sendResponse ( { type : 'Fail' , message : ErrorCodeMessage [ code ] } )
124
124
return sendResponse ( { type : 'Fail' , message : error . message ?? 'Please check the back-end console' } )
125
125
}
126
+ finally {
127
+ releaseApiKey ( key )
128
+ }
126
129
}
127
130
128
131
export function initAuditService ( audit : AuditConfig ) {
@@ -304,6 +307,39 @@ async function getMessageById(id: string): Promise<ChatMessage | undefined> {
304
307
else { return undefined }
305
308
}
306
309
310
+ const _lockedKeys : string [ ] = [ ]
311
+ async function randomKeyConfig ( keys : KeyConfig [ ] ) : Promise < KeyConfig | null > {
312
+ if ( keys . length <= 0 )
313
+ return null
314
+ let unsedKeys = keys . filter ( d => ! _lockedKeys . includes ( d . key ) )
315
+ const start = Date . now ( )
316
+ while ( unsedKeys . length <= 0 ) {
317
+ if ( Date . now ( ) - start > 3000 )
318
+ break
319
+ await new Promise ( resolve => setTimeout ( resolve , 1000 ) )
320
+ unsedKeys = keys . filter ( d => ! _lockedKeys . includes ( d . key ) )
321
+ }
322
+ if ( unsedKeys . length <= 0 )
323
+ return null
324
+ const thisKey = unsedKeys [ Math . floor ( Math . random ( ) * unsedKeys . length ) ]
325
+ _lockedKeys . push ( thisKey . key )
326
+ return thisKey
327
+ }
328
+
329
+ async function getRandomApiKey ( user : UserInfo ) : Promise < KeyConfig | undefined > {
330
+ const keys = ( await getCacheApiKeys ( ) ) . filter ( d => hasAnyRole ( d . userRoles , user . roles ) )
331
+ return randomKeyConfig ( keys )
332
+ }
333
+
334
+ async function releaseApiKey ( key : KeyConfig ) {
335
+ if ( key == null || key === undefined )
336
+ return
337
+
338
+ const index = _lockedKeys . indexOf ( key . key )
339
+ if ( index >= 0 )
340
+ _lockedKeys . splice ( index , 1 )
341
+ }
342
+
307
343
export type { ChatContext , ChatMessage }
308
344
309
- export { chatReplyProcess , chatConfig , containsSensitiveWords }
345
+ export { chatReplyProcess , chatConfig , containsSensitiveWords , getRandomApiKey }
0 commit comments