Feat: Plan cache (#5052)

* add cache for team member,app,dataset (#5007)

* add cache for team member,app,dataset

* update for hook

* update redis

* update

* perf: cache code

* fix: i18n

---------

Co-authored-by: gggaaallleee <91131304+gggaaallleee@users.noreply.github.com>

packages/service/common/redis/cache.ts

@@ -6,7 +6,16 @@ const redisPrefix = 'cache:';
 const getCacheKey = (key: string) => `${redisPrefix}${key}`;
 
 export enum CacheKeyEnum {
-  team_vector_count = 'team_vector_count'
+  team_vector_count = 'team_vector_count',
+  team_point_surplus = 'team_point_surplus',
+  team_point_total = 'team_point_total'
+}
+
+// Seconds
+export enum CacheKeyEnumTime {
+  team_vector_count = 30 * 60,
+  team_point_surplus = 1 * 60,
+  team_point_total = 1 * 60
 }
 
 export const setRedisCache = async (
@@ -17,7 +26,6 @@ export const setRedisCache = async (
   return await retryFn(async () => {
     try {
       const redis = getGlobalRedisConnection();
-
       if (expireSeconds) {
         await redis.set(getCacheKey(key), data, 'EX', expireSeconds);
       } else {
@@ -32,7 +40,19 @@ export const setRedisCache = async (
 
 export const getRedisCache = async (key: string) => {
   const redis = getGlobalRedisConnection();
-  return await retryFn(() => redis.get(getCacheKey(key)));
+  return retryFn(() => redis.get(getCacheKey(key)));
+};
+
+// Add value to cache
+export const incrValueToCache = async (key: string, increment: number) => {
+  if (!increment || increment === 0) return;
+  const redis = getGlobalRedisConnection();
+  try {
+    const exists = await redis.exists(getCacheKey(key));
+    if (!exists) return;
+
+    await retryFn(() => redis.incrbyfloat(getCacheKey(key), increment));
+  } catch (error) {}
 };
 
 export const delRedisCache = async (key: string) => {

packages/service/common/vectorDB/controller.ts

@@ -6,7 +6,14 @@ import { type DelDatasetVectorCtrlProps, type InsertVectorProps } from './contro
 import { type EmbeddingModelItemType } from '@fastgpt/global/core/ai/model.d';
 import { MILVUS_ADDRESS, PG_ADDRESS, OCEANBASE_ADDRESS } from './constants';
 import { MilvusCtrl } from './milvus';
-import { setRedisCache, getRedisCache, delRedisCache, CacheKeyEnum } from '../redis/cache';
+import {
+  setRedisCache,
+  getRedisCache,
+  delRedisCache,
+  incrValueToCache,
+  CacheKeyEnum,
+  CacheKeyEnumTime
+} from '../redis/cache';
 import { throttle } from 'lodash';
 import { retryFn } from '@fastgpt/global/common/system/utils';
 
@@ -23,6 +30,7 @@ const onDelCache = throttle((teamId: string) => delRedisCache(getChcheKey(teamId
   leading: true,
   trailing: true
 });
+const onIncrCache = (teamId: string) => incrValueToCache(getChcheKey(teamId), 1);
 
 const Vector = getVectorObj();
 
@@ -40,7 +48,7 @@ export const getVectorCountByTeamId = async (teamId: string) => {
 
   const count = await Vector.getVectorCountByTeamId(teamId);
 
-  await setRedisCache(key, count, 30 * 60);
+  await setRedisCache(key, count, CacheKeyEnumTime.team_vector_count);
 
   return count;
 };
@@ -67,7 +75,7 @@ export const insertDatasetDataVector = async ({
       vector: vectors[0]
     });
 
-    onDelCache(props.teamId);
+    onIncrCache(props.teamId);
 
     return {
       tokens,

packages/service/support/permission/teamLimit.ts

@@ -1,4 +1,4 @@
-import { getTeamPlanStatus, getTeamStandPlan } from '../../support/wallet/sub/utils';
+import { getTeamPlanStatus, getTeamStandPlan, getTeamPoints } 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';
@@ -10,11 +10,9 @@ import { TeamMemberStatusEnum } from '@fastgpt/global/support/user/team/constant
 import { getVectorCountByTeamId } from '../../common/vectorDB/controller';
 
 export const checkTeamAIPoints = async (teamId: string) => {
-  const { standardConstants, totalPoints, usedPoints } = await getTeamPlanStatus({
-    teamId
-  });
+  if (!global.subPlans?.standard) return;
 
-  if (!standardConstants) return;
+  const { totalPoints, usedPoints } = await getTeamPoints({ teamId });
 
   if (usedPoints >= totalPoints) {
     return Promise.reject(TeamErrEnum.aiPointsNotEnough);

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

@@ -13,6 +13,14 @@ import dayjs from 'dayjs';
 import { type ClientSession } from '../../../common/mongo';
 import { addMonths } from 'date-fns';
 import { readFromSecondary } from '../../../common/mongo/utils';
+import {
+  setRedisCache,
+  getRedisCache,
+  delRedisCache,
+  CacheKeyEnum,
+  CacheKeyEnumTime,
+  incrValueToCache
+} from '../../../common/redis/cache';
 
 export const getStandardPlansConfig = () => {
   return global?.subPlans?.standard;
@@ -168,6 +176,8 @@ export const getTeamPlanStatus = async ({
       ? standardPlans[standardPlan.currentSubLevel]
       : undefined;
 
+  updateTeamPointsCache({ teamId, totalPoints, surplusPoints });
+
   return {
     [SubTypeEnum.standard]: standardPlan,
     standardConstants: standardConstants
@@ -185,3 +195,59 @@ export const getTeamPlanStatus = async ({
     datasetMaxSize: totalDatasetSize
   };
 };
+
+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}`;
+
+  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);
+    return {
+      totalPoints,
+      surplusPoints,
+      usedPoints: totalPoints - surplusPoints
+    };
+  }
+
+  const planStatus = await getTeamPlanStatus({ teamId });
+  return {
+    totalPoints: planStatus.totalPoints,
+    surplusPoints: planStatus.totalPoints - planStatus.usedPoints,
+    usedPoints: planStatus.usedPoints
+  };
+};