FastGPT/packages/service/common/api/frequencyLimit.ts

/* 基于 Team 的限流 */
import { getGlobalRedisConnection } from '../../common/redis';
import { jsonRes } from '../../common/response';
import type { NextApiResponse } from 'next';

export enum LimitTypeEnum {
  chat = 'chat'
}
const limitMap = {
  [LimitTypeEnum.chat]: {
    seconds: 60,
    limit: Number(process.env.CHAT_MAX_QPM || 5000)
  }
};

type FrequencyLimitOption = {
  teamId: string;
  type: LimitTypeEnum;
  res: NextApiResponse;
};

export const teamFrequencyLimit = async ({ teamId, type, res }: FrequencyLimitOption) => {
  const { seconds, limit } = limitMap[type];
  const redis = getGlobalRedisConnection();
  const key = `frequency:${type}:${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;
};