From d0bac766b05b965d6541a7a9664eba2d957dc389 Mon Sep 17 00:00:00 2001 From: archer <545436317@qq.com> Date: Wed, 10 Dec 2025 15:48:41 +0800 Subject: [PATCH] perf: code --- packages/service/common/api/frequencyLimit.ts | 63 ++++---- .../service/support/permission/teamLimit.ts | 4 +- packages/service/support/wallet/sub/utils.ts | 144 ++++++++++-------- 3 files changed, 114 insertions(+), 97 deletions(-) diff --git a/packages/service/common/api/frequencyLimit.ts b/packages/service/common/api/frequencyLimit.ts index 60fdb1528..be552d0db 100644 --- a/packages/service/common/api/frequencyLimit.ts +++ b/packages/service/common/api/frequencyLimit.ts @@ -2,51 +2,44 @@ import { getGlobalRedisConnection } from '../../common/redis'; import { jsonRes } from '../../common/response'; import type { NextApiResponse } from 'next'; -import { - getCachedTeamQPMLimit, - setCachedTeamQPMLimit, - getTeamPlanStatus -} from '../../support/wallet/sub/utils'; -import { SubTypeEnum } from '@fastgpt/global/support/wallet/sub/constants'; +import { teamQPM } from '../../support/wallet/sub/utils'; +import z from 'zod'; export enum LimitTypeEnum { chat = 'chat' } -const limitSecondsMap = { - [LimitTypeEnum.chat]: 60 -}; +const FrequencyLimitOptionSchema = z.union([ + z.object({ + type: z.literal(LimitTypeEnum.chat), + teamId: z.string() + }) +]); +type FrequencyLimitOption = z.infer; -type FrequencyLimitOption = { - teamId: string; - type: LimitTypeEnum; - res: NextApiResponse; -}; +const getLimitData = async (data: FrequencyLimitOption) => { + if (data.type === LimitTypeEnum.chat) { + const qpm = await teamQPM.getTeamQPMLimit(data.teamId); -// Get team's dynamic QPM limit with caching -export const getTeamQPMLimit = async (teamId: string): Promise => { - // 1. Try to get from cache first - const cachedLimit = await getCachedTeamQPMLimit(teamId); - if (cachedLimit !== null) { - return cachedLimit; + if (!qpm) return; + + return { + limit: qpm, + seconds: 60 + }; } - - // 2. Cache miss, compute from database - const teamPlanStatus = await getTeamPlanStatus({ teamId }); - const limit = - teamPlanStatus[SubTypeEnum.standard]?.requestsPerMinute ?? - teamPlanStatus.standardConstants?.requestsPerMinute ?? - 30; - - // 3. Write to cache - await setCachedTeamQPMLimit(teamId, limit); - - return limit; + return; }; -export const teamFrequencyLimit = async ({ teamId, type, res }: FrequencyLimitOption) => { - const limit = await getTeamQPMLimit(teamId); - const seconds = limitSecondsMap[type]; +export const teamFrequencyLimit = async ({ + teamId, + type, + res +}: FrequencyLimitOption & { res: NextApiResponse }) => { + const data = await getLimitData({ type, teamId }); + if (!data) return; + + const { limit, seconds } = data; const redis = getGlobalRedisConnection(); const key = `frequency:${type}:${teamId}`; diff --git a/packages/service/support/permission/teamLimit.ts b/packages/service/support/permission/teamLimit.ts index 6d026c9ba..0292c065c 100644 --- a/packages/service/support/permission/teamLimit.ts +++ b/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); diff --git a/packages/service/support/wallet/sub/utils.ts b/packages/service/support/wallet/sub/utils.ts index 1a27f1018..a89949457 100644 --- a/packages/service/support/wallet/sub/utils.ts +++ b/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,72 +223,96 @@ export const getTeamPlanStatus = async ({ }; }; -export const clearTeamPointsCache = async (teamId: string) => { - 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}`; - await Promise.all([delRedisCache(surplusCacheKey), delRedisCache(totalCacheKey)]); -}; + const [surplusCacheStr, totalCacheStr] = await Promise.all([ + getRedisCache(surplusCacheKey), + getRedisCache(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}`; + 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) - ]); -}; - -export const 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); + const planStatus = await getTeamPlanStatus({ teamId }); return { - totalPoints, - surplusPoints, - usedPoints: totalPoints - surplusPoints + 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 => { + // 1. 尝试从缓存中获取 + const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`; + const cached = await getRedisCache(cacheKey); - const planStatus = await getTeamPlanStatus({ teamId }); - return { - totalPoints: planStatus.totalPoints, - surplusPoints: planStatus.totalPoints - planStatus.usedPoints, - usedPoints: planStatus.usedPoints - }; + if (cached) { + return Number(cached); + } + + // 2. Computed + const teamPlanStatus = await getTeamPlanStatus({ teamId }); + const limit = + teamPlanStatus[SubTypeEnum.standard]?.requestsPerMinute ?? + teamPlanStatus.standardConstants?.requestsPerMinute; + + if (!limit) return null; + + // 3. Set cache + await teamQPM.setCachedTeamQPMLimit(teamId, limit); + + return limit; + }, + setCachedTeamQPMLimit: async (teamId: string, limit: number): Promise => { + const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`; + await setRedisCache(cacheKey, limit.toString(), CacheKeyEnumTime.team_qpm_limit); + }, + clearTeamQPMLimitCache: async (teamId: string): Promise => { + const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`; + await delRedisCache(cacheKey); + } }; -export const getCachedTeamQPMLimit = async (teamId: string): Promise => { - const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`; - const cached = await getRedisCache(cacheKey); - return cached ? Number(cached) : null; -}; -export const setCachedTeamQPMLimit = async (teamId: string, limit: number): Promise => { - const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`; - await setRedisCache(cacheKey, limit.toString(), CacheKeyEnumTime.team_qpm_limit); -}; -export const clearTeamQPMLimitCache = async (teamId: string): Promise => { - const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`; - await delRedisCache(cacheKey); +// controler +export const clearTeamPlanCache = async (teamId: string) => { + await teamPoint.clearTeamPointsCache(teamId); + await teamQPM.clearTeamQPMLimitCache(teamId); };