mirror of
https://github.com/labring/FastGPT.git
synced 2025-12-25 20:02:47 +00:00
perf: code
This commit is contained in:
parent
c3f6aaa5fd
commit
d0bac766b0
|
|
@ -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}`;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
};
|
||||
|
|
|
|||
Loading…
Reference in New Issue