team qpm limit & plan tracks (#6066)

* team qpm limit & plan tracks

* api entry qpm

* perf: computed days

* Revert "api entry qpm"

This reverts commit 1210c07.

* perf: code

* system qpm limit

* system qpm limit

---------

Co-authored-by: archer <545436317@qq.com>

Author: newfish-cmyk

packages/global/common/middle/tracks/constants.ts

@@ -10,5 +10,7 @@ export enum TrackEnum {
   readSystemAnnouncement = 'readSystemAnnouncement',
   clickOperationalAd = 'clickOperationalAd',
   closeOperationalAd = 'closeOperationalAd',
-  teamChatQPM = 'teamChatQPM'
+  teamChatQPM = 'teamChatQPM',
+  subscriptionDeleted = 'subscriptionDeleted',
+  freeAccountCleanup = 'freeAccountCleanup'
 }

packages/service/common/api/frequencyLimit.ts

@@ -2,25 +2,50 @@
 import { getGlobalRedisConnection } from '../../common/redis';
 import { jsonRes } from '../../common/response';
 import type { NextApiResponse } from 'next';
+import { teamQPM } from '../../support/wallet/sub/utils';
+import z from 'zod';
+import { addLog } from 'common/system/log';
 
 export enum LimitTypeEnum {
   chat = 'chat'
 }
-const limitMap = {
-  [LimitTypeEnum.chat]: {
-    seconds: 60,
-    limit: Number(process.env.CHAT_MAX_QPM || 5000)
+
+const FrequencyLimitOptionSchema = z.union([
+  z.object({
+    type: z.literal(LimitTypeEnum.chat),
+    teamId: z.string()
+  })
+]);
+type FrequencyLimitOption = z.infer<typeof FrequencyLimitOptionSchema>;
+
+const getLimitData = async (data: FrequencyLimitOption) => {
+  if (data.type === LimitTypeEnum.chat) {
+    const qpm = await teamQPM.getTeamQPMLimit(data.teamId);
+
+    if (!qpm) return;
+
+    return {
+      limit: qpm,
+      seconds: 60
+    };
   }
+  return;
 };
 
-type FrequencyLimitOption = {
-  teamId: string;
-  type: LimitTypeEnum;
-  res: NextApiResponse;
-};
+/* 
+  true: 未达到限制
+  false: 达到了限制
+*/
+export const teamFrequencyLimit = async ({
+  teamId,
+  type,
+  res
+}: FrequencyLimitOption & { res: NextApiResponse }) => {
+  const data = await getLimitData({ type, teamId });
+  if (!data) return true;
+
+  const { limit, seconds } = data;
 
-export const teamFrequencyLimit = async ({ teamId, type, res }: FrequencyLimitOption) => {
-  const { seconds, limit } = limitMap[type];
   const redis = getGlobalRedisConnection();
   const key = `frequency:${type}:${teamId}`;
 
@@ -31,13 +56,16 @@ export const teamFrequencyLimit = async ({ teamId, type, res }: FrequencyLimitOp
     .exec();
 
   if (!result) {
-    return Promise.reject(new Error('Redis connection error'));
+    return true;
   }
 
   const currentCount = result[0][1] as number;
 
   if (currentCount > limit) {
     const remainingTime = await redis.ttl(key);
+    addLog.info(
+      `[Completion Limit] Team ${teamId} reached the limit of ${limit} requests per ${seconds} seconds. Remaining time: ${remainingTime} seconds.`
+    );
     jsonRes(res, {
       code: 429,
       error: `Rate limit exceeded. Maximum ${limit} requests per ${seconds} seconds for this team. Please try again in ${remainingTime} seconds.`

packages/service/common/middle/tracks/utils.ts

@@ -8,6 +8,7 @@ import type { DatasetTypeEnum } from '@fastgpt/global/core/dataset/constants';
 import { getAppLatestVersion } from '../../../core/app/version/controller';
 import { type ShortUrlParams } from '@fastgpt/global/support/marketing/type';
 import { getRedisCache, setRedisCache } from '../../redis/cache';
+import { differenceInDays } from 'date-fns';
 
 const createTrack = ({ event, data }: { event: TrackEnum; data: Record<string, any> }) => {
   if (!global.feConfigs?.isPlus) return;
@@ -156,5 +157,35 @@ export const pushTrack = {
         teamId: data.teamId
       }
     });
+  },
+  subscriptionDeleted: (data: {
+    teamId: string;
+    subscriptionType: string;
+    subLevel?: string;
+    totalPoints: number;
+    usedPoints: number;
+    startTime: Date;
+    expiredTime: Date;
+  }) => {
+    return createTrack({
+      event: TrackEnum.subscriptionDeleted,
+      data: {
+        teamId: data.teamId,
+        subscriptionType: data.subscriptionType,
+        subLevel: data.subLevel,
+        totalPoints: data.totalPoints,
+        usedPoints: data.usedPoints,
+        activeDays: differenceInDays(data.expiredTime, data.startTime)
+      }
+    });
+  },
+  freeAccountCleanup: (data: { teamId: string; expiredTime: Date }) => {
+    return createTrack({
+      event: TrackEnum.freeAccountCleanup,
+      data: {
+        teamId: data.teamId,
+        expiredTime: data.expiredTime
+      }
+    });
   }
 };

packages/service/common/redis/cache.ts

@@ -8,14 +8,16 @@ const getCacheKey = (key: string) => `${redisPrefix}${key}`;
 export enum CacheKeyEnum {
   team_vector_count = 'team_vector_count',
   team_point_surplus = 'team_point_surplus',
-  team_point_total = 'team_point_total'
+  team_point_total = 'team_point_total',
+  team_qpm_limit = 'team_qpm_limit'
 }
 
 // Seconds
 export enum CacheKeyEnumTime {
   team_vector_count = 30 * 60,
   team_point_surplus = 1 * 60,
-  team_point_total = 1 * 60
+  team_point_total = 1 * 60,
+  team_qpm_limit = 60 * 60
 }
 
 export const setRedisCache = async (

packages/service/common/system/timerLock/constants.ts

@@ -10,7 +10,9 @@ export enum TimerIdEnum {
   clearExpiredRawTextBuffer = 'clearExpiredRawTextBuffer',
   clearExpiredDatasetImage = 'clearExpiredDatasetImage',
   clearExpiredMinioFiles = 'clearExpiredMinioFiles',
-  recordTeamQPM = 'recordTeamQPM'
+  recordTeamQPM = 'recordTeamQPM',
+  auditLogCleanup = 'auditLogCleanup',
+  chatHistoryCleanup = 'chatHistoryCleanup'
 }
 
 export enum LockNotificationEnum {

packages/service/support/permission/teamLimit.ts

@@ -1,4 +1,4 @@
-import { getTeamPlanStatus, getTeamStandPlan, getTeamPoints } from '../../support/wallet/sub/utils';
+import { getTeamPlanStatus, getTeamStandPlan, teamPoint } from '../../support/wallet/sub/utils';
 import { MongoApp } from '../../core/app/schema';
 import { MongoDataset } from '../../core/dataset/schema';
 import { DatasetTypeEnum } from '@fastgpt/global/core/dataset/constants';
@@ -12,7 +12,7 @@ import { getVectorCountByTeamId } from '../../common/vectorDB/controller';
 export const checkTeamAIPoints = async (teamId: string) => {
   if (!global.subPlans?.standard) return;
 
-  const { totalPoints, usedPoints } = await getTeamPoints({ teamId });
+  const { totalPoints, usedPoints } = await teamPoint.getTeamPoints({ teamId });
 
   if (usedPoints >= totalPoints) {
     return Promise.reject(TeamErrEnum.aiPointsNotEnough);

packages/service/support/wallet/sub/utils.ts

@@ -190,7 +190,7 @@ export const getTeamPlanStatus = async ({
       ? standardPlans[standardPlan.currentSubLevel]
       : undefined;
 
-  updateTeamPointsCache({ teamId, totalPoints, surplusPoints });
+  teamPoint.updateTeamPointsCache({ teamId, totalPoints, surplusPoints });
 
   return {
     [SubTypeEnum.standard]: standardPlan,
@@ -223,58 +223,99 @@ export const getTeamPlanStatus = async ({
   };
 };
 
-export const clearTeamPointsCache = async (teamId: string) => {
-  const surplusCacheKey = `${CacheKeyEnum.team_point_surplus}:${teamId}`;
-  const totalCacheKey = `${CacheKeyEnum.team_point_total}:${teamId}`;
-
-  await Promise.all([delRedisCache(surplusCacheKey), delRedisCache(totalCacheKey)]);
-};
-
-export const incrTeamPointsCache = async ({ teamId, value }: { teamId: string; value: number }) => {
-  const surplusCacheKey = `${CacheKeyEnum.team_point_surplus}:${teamId}`;
-  await incrValueToCache(surplusCacheKey, value);
-};
-export const updateTeamPointsCache = async ({
-  teamId,
-  totalPoints,
-  surplusPoints
-}: {
-  teamId: string;
-  totalPoints: number;
-  surplusPoints: number;
-}) => {
-  const surplusCacheKey = `${CacheKeyEnum.team_point_surplus}:${teamId}`;
-  const totalCacheKey = `${CacheKeyEnum.team_point_total}:${teamId}`;
+export const teamPoint = {
+  getTeamPoints: async ({ teamId }: { teamId: string }) => {
+    const surplusCacheKey = `${CacheKeyEnum.team_point_surplus}:${teamId}`;
+    const totalCacheKey = `${CacheKeyEnum.team_point_total}:${teamId}`;
+
+    const [surplusCacheStr, totalCacheStr] = await Promise.all([
+      getRedisCache(surplusCacheKey),
+      getRedisCache(totalCacheKey)
+    ]);
+
+    if (surplusCacheStr && totalCacheStr) {
+      const totalPoints = Number(totalCacheStr);
+      const surplusPoints = Number(surplusCacheStr);
+      return {
+        totalPoints,
+        surplusPoints,
+        usedPoints: totalPoints - surplusPoints
+      };
+    }
 
-  await Promise.all([
-    setRedisCache(surplusCacheKey, surplusPoints, CacheKeyEnumTime.team_point_surplus),
-    setRedisCache(totalCacheKey, totalPoints, CacheKeyEnumTime.team_point_total)
-  ]);
+    const planStatus = await getTeamPlanStatus({ teamId });
+    return {
+      totalPoints: planStatus.totalPoints,
+      surplusPoints: planStatus.totalPoints - planStatus.usedPoints,
+      usedPoints: planStatus.usedPoints
+    };
+  },
+  incrTeamPointsCache: async ({ teamId, value }: { teamId: string; value: number }) => {
+    const surplusCacheKey = `${CacheKeyEnum.team_point_surplus}:${teamId}`;
+    await incrValueToCache(surplusCacheKey, value);
+  },
+  updateTeamPointsCache: async ({
+    teamId,
+    totalPoints,
+    surplusPoints
+  }: {
+    teamId: string;
+    totalPoints: number;
+    surplusPoints: number;
+  }) => {
+    const surplusCacheKey = `${CacheKeyEnum.team_point_surplus}:${teamId}`;
+    const totalCacheKey = `${CacheKeyEnum.team_point_total}:${teamId}`;
+
+    await Promise.all([
+      setRedisCache(surplusCacheKey, surplusPoints, CacheKeyEnumTime.team_point_surplus),
+      setRedisCache(totalCacheKey, totalPoints, CacheKeyEnumTime.team_point_total)
+    ]);
+  },
+  clearTeamPointsCache: async (teamId: string) => {
+    const surplusCacheKey = `${CacheKeyEnum.team_point_surplus}:${teamId}`;
+    const totalCacheKey = `${CacheKeyEnum.team_point_total}:${teamId}`;
+
+    await Promise.all([delRedisCache(surplusCacheKey), delRedisCache(totalCacheKey)]);
+  }
 };
+export const teamQPM = {
+  getTeamQPMLimit: async (teamId: string): Promise<number | null> => {
+    // 1. 尝试从缓存中获取
+    const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`;
+    const cached = await getRedisCache(cacheKey);
+
+    if (cached) {
+      return Number(cached);
+    }
 
-export const getTeamPoints = async ({ teamId }: { teamId: string }) => {
-  const surplusCacheKey = `${CacheKeyEnum.team_point_surplus}:${teamId}`;
-  const totalCacheKey = `${CacheKeyEnum.team_point_total}:${teamId}`;
+    // 2. Computed
+    const teamPlanStatus = await getTeamPlanStatus({ teamId });
+    const limit =
+      teamPlanStatus[SubTypeEnum.standard]?.requestsPerMinute ??
+      teamPlanStatus.standardConstants?.requestsPerMinute;
 
-  const [surplusCacheStr, totalCacheStr] = await Promise.all([
-    getRedisCache(surplusCacheKey),
-    getRedisCache(totalCacheKey)
-  ]);
+    if (!limit) {
+      if (process.env.CHAT_MAX_QPM) return Number(process.env.CHAT_MAX_QPM);
+      return null;
+    }
 
-  if (surplusCacheStr && totalCacheStr) {
-    const totalPoints = Number(totalCacheStr);
-    const surplusPoints = Number(surplusCacheStr);
-    return {
-      totalPoints,
-      surplusPoints,
-      usedPoints: totalPoints - surplusPoints
-    };
+    // 3. Set cache
+    await teamQPM.setCachedTeamQPMLimit(teamId, limit);
+
+    return limit;
+  },
+  setCachedTeamQPMLimit: async (teamId: string, limit: number): Promise<void> => {
+    const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`;
+    await setRedisCache(cacheKey, limit.toString(), CacheKeyEnumTime.team_qpm_limit);
+  },
+  clearTeamQPMLimitCache: async (teamId: string): Promise<void> => {
+    const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`;
+    await delRedisCache(cacheKey);
   }
+};
 
-  const planStatus = await getTeamPlanStatus({ teamId });
-  return {
-    totalPoints: planStatus.totalPoints,
-    surplusPoints: planStatus.totalPoints - planStatus.usedPoints,
-    usedPoints: planStatus.usedPoints
-  };
+// controler
+export const clearTeamPlanCache = async (teamId: string) => {
+  await teamPoint.clearTeamPointsCache(teamId);
+  await teamQPM.clearTeamQPMLimitCache(teamId);
 };