feat: team rate limitation

Author: FinleyGe

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;
+};

packages/service/common/api/teamFrequencyLimit.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;