token刷新请求并发合并处理

This commit is contained in:
Vinlic 2024-03-13 18:41:07 +08:00
parent b4e8b489ec
commit 28bd87e7f0
2 changed files with 54 additions and 29 deletions

View File

@ -34,6 +34,8 @@ const FAKE_HEADERS = {
const FILE_MAX_SIZE = 100 * 1024 * 1024; const FILE_MAX_SIZE = 100 * 1024 * 1024;
// access_token映射 // access_token映射
const accessTokenMap = new Map(); const accessTokenMap = new Map();
// access_token请求队列映射
const accessTokenRequestQueueMap: Record<string, Function[]> = {};
/** /**
* access_token * access_token
@ -43,24 +45,48 @@ const accessTokenMap = new Map();
* @param refreshToken access_token的refresh_token * @param refreshToken access_token的refresh_token
*/ */
async function requestToken(refreshToken: string) { async function requestToken(refreshToken: string) {
const result = await axios.get('https://kimi.moonshot.cn/api/auth/token/refresh', { if (accessTokenRequestQueueMap[refreshToken])
headers: { return new Promise(resolve => accessTokenRequestQueueMap[refreshToken].push(resolve));
Authorization: `Bearer ${refreshToken}`, accessTokenRequestQueueMap[refreshToken] = [];
Referer: 'https://kimi.moonshot.cn', logger.info(`Refresh token: ${refreshToken}`);
...FAKE_HEADERS const result = await (async () => {
}, const result = await axios.get('https://kimi.moonshot.cn/api/auth/token/refresh', {
timeout: 15000, headers: {
validateStatus: () => true Authorization: `Bearer ${refreshToken}`,
}); Referer: 'https://kimi.moonshot.cn',
const { ...FAKE_HEADERS
access_token, },
refresh_token timeout: 15000,
} = checkResult(result, refreshToken); validateStatus: () => true
return { });
accessToken: access_token, const {
refreshToken: refresh_token, access_token,
refreshTime: util.unixTimestamp() + ACCESS_TOKEN_EXPIRES refresh_token
} } = checkResult(result, refreshToken);
return {
accessToken: access_token,
refreshToken: refresh_token,
refreshTime: util.unixTimestamp() + ACCESS_TOKEN_EXPIRES
}
})()
.then(result => {
if(accessTokenRequestQueueMap[refreshToken]) {
accessTokenRequestQueueMap[refreshToken].forEach(resolve => resolve(result));
delete accessTokenRequestQueueMap[refreshToken];
}
logger.success(`Refresh successful`);
return result;
})
.catch(err => {
if(accessTokenRequestQueueMap[refreshToken]) {
accessTokenRequestQueueMap[refreshToken].forEach(resolve => resolve(err));
delete accessTokenRequestQueueMap[refreshToken];
}
return err;
});
if(_.isError(result))
throw result;
return result;
} }
/** /**
@ -224,15 +250,15 @@ async function createCompletionStream(messages: any[], refreshToken: string, use
*/ */
function extractRefFileUrls(messages: any[]) { function extractRefFileUrls(messages: any[]) {
return messages.reduce((urls, message) => { return messages.reduce((urls, message) => {
if(_.isArray(message.content)) { if (_.isArray(message.content)) {
message.content.forEach(v => { message.content.forEach(v => {
if(!_.isObject(v) || !['file', 'image_url'].includes(v['type'])) if (!_.isObject(v) || !['file', 'image_url'].includes(v['type']))
return; return;
// kimi-free-api支持格式 // kimi-free-api支持格式
if(v['type'] == 'file' && _.isObject(v['file_url']) && _.isString(v['file_url']['url'])) if (v['type'] == 'file' && _.isObject(v['file_url']) && _.isString(v['file_url']['url']))
urls.push(v['file_url']['url']); urls.push(v['file_url']['url']);
// 兼容gpt-4-vision-preview API格式 // 兼容gpt-4-vision-preview API格式
else if(v['type'] == 'image_url' && _.isObject(v['image_url']) && _.isString(v['image_url']['url'])) else if (v['type'] == 'image_url' && _.isObject(v['image_url']) && _.isString(v['image_url']['url']))
urls.push(v['image_url']['url']); urls.push(v['image_url']['url']);
}); });
} }
@ -254,7 +280,7 @@ function messagesPrepare(messages: any[]) {
const content = messages.reduce((content, message) => { const content = messages.reduce((content, message) => {
if (_.isArray(message.content)) { if (_.isArray(message.content)) {
return message.content.reduce((_content, v) => { return message.content.reduce((_content, v) => {
if(!_.isObject(v) || v['type'] != 'text') if (!_.isObject(v) || v['type'] != 'text')
return _content; return _content;
return _content + (v['text'] || ''); return _content + (v['text'] || '');
}, content); }, content);
@ -307,18 +333,18 @@ async function preSignUrl(filename: string, refreshToken: string) {
* @param fileUrl URL * @param fileUrl URL
*/ */
async function checkFileUrl(fileUrl: string) { async function checkFileUrl(fileUrl: string) {
if(util.isBASE64Data(fileUrl)) if (util.isBASE64Data(fileUrl))
return; return;
const result = await axios.head(fileUrl, { const result = await axios.head(fileUrl, {
timeout: 15000, timeout: 15000,
validateStatus: () => true validateStatus: () => true
}); });
if(result.status >= 400) if (result.status >= 400)
throw new APIException(EX.API_FILE_URL_INVALID, `File ${fileUrl} is not valid: [${result.status}] ${result.statusText}`); throw new APIException(EX.API_FILE_URL_INVALID, `File ${fileUrl} is not valid: [${result.status}] ${result.statusText}`);
// 检查文件大小 // 检查文件大小
if (result.headers && result.headers['content-length']) { if (result.headers && result.headers['content-length']) {
const fileSize = parseInt(result.headers['content-length'], 10); const fileSize = parseInt(result.headers['content-length'], 10);
if(fileSize > FILE_MAX_SIZE) if (fileSize > FILE_MAX_SIZE)
throw new APIException(EX.API_FILE_EXECEEDS_SIZE, `File ${fileUrl} is not valid`); throw new APIException(EX.API_FILE_EXECEEDS_SIZE, `File ${fileUrl} is not valid`);
} }
} }
@ -335,7 +361,7 @@ async function uploadFile(fileUrl: string, refreshToken: string) {
let filename, fileData, mimeType; let filename, fileData, mimeType;
// 如果是BASE64数据则直接转换为Buffer // 如果是BASE64数据则直接转换为Buffer
if(util.isBASE64Data(fileUrl)) { if (util.isBASE64Data(fileUrl)) {
mimeType = util.extractBASE64DataFormat(fileUrl); mimeType = util.extractBASE64DataFormat(fileUrl);
const ext = mime.getExtension(mimeType); const ext = mime.getExtension(mimeType);
filename = `${util.uuid()}.${ext}`; filename = `${util.uuid()}.${ext}`;
@ -358,7 +384,7 @@ async function uploadFile(fileUrl: string, refreshToken: string) {
url: uploadUrl, url: uploadUrl,
object_name: objectName object_name: objectName
} = await preSignUrl(filename, refreshToken); } = await preSignUrl(filename, refreshToken);
// 获取文件的MIME类型 // 获取文件的MIME类型
mimeType = mimeType || mime.getType(filename); mimeType = mimeType || mime.getType(filename);
// 上传文件到目标OSS // 上传文件到目标OSS

View File

@ -17,7 +17,6 @@ export default {
.validate('headers.authorization', _.isString) .validate('headers.authorization', _.isString)
const token = request.headers.authorization; const token = request.headers.authorization;
const refreshToken = token.replace('Bearer ', ''); const refreshToken = token.replace('Bearer ', '');
logger.info(`Refresh token: ${refreshToken}`);
const messages = request.body.messages; const messages = request.body.messages;
if (request.body.stream) { if (request.body.stream) {
const stream = await chat.createCompletionStream(request.body.messages, refreshToken, request.body.use_search); const stream = await chat.createCompletionStream(request.body.messages, refreshToken, request.body.use_search);