perf: add Redis cache for vector counting

packages/service/common/redis/index.ts

@@ -1,3 +1,4 @@
+import { addLog } from '../system/log';
 import Redis from 'ioredis';
 
 const REDIS_URL = process.env.REDIS_URL ?? 'redis://localhost:6379';
@@ -25,3 +26,37 @@ export function newWorkerRedisConnection() {
   });
   return redis;
 }
+
+export function getGlobalRedisCacheConnection() {
+  if (global.redisCache) return global.redisCache;
+  const redis = new Redis(REDIS_URL, { keyPrefix: 'fastgpt:cache:' });
+  redis.on('connect', () => {
+    console.log('Redis connected');
+  });
+  redis.on('error', (error) => {
+    console.error('Redis connection error', error);
+  });
+  global.redisCache = redis;
+  return redis;
+}
+
+export async function checkAndIncr(redis: Redis, key: string, retry = 0) {
+  // start a transaction
+  await redis.watch(key);
+
+  const exists = await redis.exists(key);
+  if (exists) {
+    const result = await redis.multi().incr(key).exec();
+    const err = result?.[0][0];
+    if (err) {
+      addLog.error('redis opt checkAndIncr', err);
+      // retry
+      if (retry < 10) {
+        await checkAndIncr(redis, key, retry + 1);
+      } else {
+        // try to delete the dirty cache
+        await redis.del(key);
+      }
+    }
+  }
+}

packages/service/common/redis/type.d.ts

@@ -0,0 +1,5 @@
+import Redis from 'ioredis';
+
+declare global {
+  var redisCache: Redis | null;
+}

packages/service/common/vectorStore/controller.ts

@@ -2,10 +2,11 @@
 import { PgVectorCtrl } from './pg/class';
 import { ObVectorCtrl } from './oceanbase/class';
 import { getVectorsByText } from '../../core/ai/embedding';
-import { InsertVectorProps } from './controller.d';
+import { DelDatasetVectorCtrlProps, InsertVectorProps } from './controller.d';
 import { EmbeddingModelItemType } from '@fastgpt/global/core/ai/model.d';
 import { MILVUS_ADDRESS, PG_ADDRESS, OCEANBASE_ADDRESS } from './constants';
 import { MilvusCtrl } from './milvus/class';
+import { getGlobalRedisCacheConnection, checkAndIncr } from '../redis';
 
 const getVectorObj = () => {
   if (PG_ADDRESS) return new PgVectorCtrl();
@@ -17,14 +18,28 @@ const getVectorObj = () => {
 
 const Vector = getVectorObj();
 
+const redis = getGlobalRedisCacheConnection();
+
 export const initVectorStore = Vector.init;
-export const deleteDatasetDataVector = Vector.delete;
 export const recallFromVectorStore = Vector.embRecall;
 export const getVectorDataByTime = Vector.getVectorDataByTime;
-export const getVectorCountByTeamId = Vector.getVectorCountByTeamId;
+
+export const getVectorCountByTeamId = async (teamId: string) => {
+  const countStr = await redis.getex(`countByTeamId:${teamId}`, 'EX', 1800);
+  if (!countStr) {
+    return await initCache(teamId);
+  }
+  return parseInt(countStr);
+};
 export const getVectorCountByDatasetId = Vector.getVectorCountByDatasetId;
 export const getVectorCountByCollectionId = Vector.getVectorCountByCollectionId;
 
+async function initCache(teamId: string) {
+  const count = await Vector.getVectorCountByTeamId(teamId);
+  await redis.setex(`countByTeamId:${teamId}`, 1800, count);
+  return count;
+}
+
 export const insertDatasetDataVector = async ({
   model,
   query,
@@ -43,8 +58,16 @@ export const insertDatasetDataVector = async ({
     vector: vectors[0]
   });
 
+  await checkAndIncr(redis, `countByTeamId:${props.teamId}`);
+
   return {
     tokens,
     insertId
   };
 };
+
+export const deleteDatasetDataVector = async (props: DelDatasetVectorCtrlProps) => {
+  const result = await Vector.delete(props);
+  await initCache(props.teamId);
+  return result;
+};