perf: code

This commit is contained in:
archer 2025-12-10 15:48:41 +08:00
parent c3f6aaa5fd
commit d0bac766b0
No known key found for this signature in database
GPG Key ID: 4446499B846D4A9E
3 changed files with 114 additions and 97 deletions

View File

@ -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<typeof FrequencyLimitOptionSchema>;
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<number> => {
// 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}`;

View File

@ -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);

View File

@ -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<number | null> => {
// 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<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);
}
};
export const getCachedTeamQPMLimit = async (teamId: string): Promise<number | null> => {
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<void> => {
const cacheKey = `${CacheKeyEnum.team_qpm_limit}:${teamId}`;
await setRedisCache(cacheKey, limit.toString(), CacheKeyEnumTime.team_qpm_limit);
};
export const clearTeamQPMLimitCache = async (teamId: string): Promise<void> => {
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);
};