From bcf090d7218e354f3ffcb06b7777a043301e4aac Mon Sep 17 00:00:00 2001 From: FinleyGe Date: Sat, 15 Nov 2025 22:06:05 +0800 Subject: [PATCH] feat: team rate limitation --- packages/service/common/api/frequencyLimit.ts | 49 +++++++++++++++++++ .../app/src/pages/api/v1/chat/completions.ts | 14 +++++- 2 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 packages/service/common/api/frequencyLimit.ts diff --git a/packages/service/common/api/frequencyLimit.ts b/packages/service/common/api/frequencyLimit.ts new file mode 100644 index 000000000000..c5016a288039 --- /dev/null +++ b/packages/service/common/api/frequencyLimit.ts @@ -0,0 +1,49 @@ +import { getGlobalRedisConnection } from '../../common/redis'; +import { jsonRes } from '../../common/response'; +import type { NextApiResponse } from 'next'; + +type FrequencyLimitOption = { + teamId: string; + seconds: number; + limit: number; + keyPrefix: string; + res: NextApiResponse; +}; + +export const teamFrequencyLimit = async ({ + teamId, + seconds, + limit, + keyPrefix, + res +}: FrequencyLimitOption) => { + const redis = getGlobalRedisConnection(); + const key = `${keyPrefix}:${teamId}`; + + const result = await redis + .multi() + .incr(key) + .expire(key, seconds, 'NX') // 只在key不存在时设置过期时间 + .exec(); + + if (!result) { + return Promise.reject(new Error('Redis connection error')); + } + + const currentCount = result[0][1] as number; + + if (currentCount > limit) { + const remainingTime = await redis.ttl(key); + jsonRes(res, { + code: 429, + error: `Rate limit exceeded. Maximum ${limit} requests per ${seconds} seconds for this team. Please try again in ${remainingTime} seconds.` + }); + return false; + } + + // 在响应头中添加限流信息 + res.setHeader('X-RateLimit-Limit', limit); + res.setHeader('X-RateLimit-Remaining', Math.max(0, limit - currentCount)); + res.setHeader('X-RateLimit-Reset', Date.now() + seconds * 1000); + return true; +}; diff --git a/projects/app/src/pages/api/v1/chat/completions.ts b/projects/app/src/pages/api/v1/chat/completions.ts index 11b6cbecf37e..f971ab920c2d 100644 --- a/projects/app/src/pages/api/v1/chat/completions.ts +++ b/projects/app/src/pages/api/v1/chat/completions.ts @@ -33,7 +33,6 @@ import { removeEmptyUserInput } from '@fastgpt/global/core/chat/utils'; import { updateApiKeyUsage } from '@fastgpt/service/support/openapi/tools'; -import { getUserChatInfo } from '@fastgpt/service/support/user/team/utils'; import { getRunningUserInfoByTmbId } from '@fastgpt/service/support/user/team/utils'; import { AuthUserTypeEnum } from '@fastgpt/global/support/permission/constant'; import { MongoApp } from '@fastgpt/service/core/app/schema'; @@ -61,6 +60,7 @@ import { getWorkflowToolInputsFromStoreNodes } from '@fastgpt/global/core/app/to import { UserError } from '@fastgpt/global/common/error/utils'; import { getLocale } from '@fastgpt/service/common/middle/i18n'; import { formatTime2YMDHM } from '@fastgpt/global/common/string/time'; +import { teamFrequencyLimit } from '@fastgpt/service/common/api/frequencyLimit'; type FastGptWebChatProps = { chatId?: string; // undefined: get histories from messages, '': new chat, 'xxxxx': get histories from db @@ -185,6 +185,18 @@ async function handler(req: NextApiRequest, res: NextApiResponse) { chatId }); })(); + + if ( + !(await teamFrequencyLimit({ + teamId, + keyPrefix: 'chat:completions', + seconds: 60, + limit: 5000, + res + })) + ) { + return {}; + } retainDatasetCite = retainDatasetCite && !!responseDetail; const isPlugin = app.type === AppTypeEnum.workflowTool;