Merge pull request #3350 from RedisInsight/be/feature/RI-5711-context-cache

add cache for context

Author: ArtemHoruzhenko

redisinsight/api/src/modules/ai/query/ai-query.module.ts

@@ -6,12 +6,17 @@ import { AiQueryAuthProvider } from 'src/modules/ai/query/providers/auth/ai-quer
 import { LocalAiQueryAuthProvider } from 'src/modules/ai/query/providers/auth/local.ai-query-auth.provider';
 import { AiQueryMessageRepository } from 'src/modules/ai/query/repositories/ai-query.message.repository';
 import { LocalAiQueryMessageRepository } from 'src/modules/ai/query/repositories/local.ai-query.message.repository';
+import { AiQueryContextRepository } from 'src/modules/ai/query/repositories/ai-query.context.repository';
+import {
+  InMemoryAiQueryContextRepository,
+} from 'src/modules/ai/query/repositories/in-memory.ai-query.context.repository';
 
 @Module({})
 export class AiQueryModule {
   static register(
     aiQueryAuthProvider: Type<AiQueryAuthProvider> = LocalAiQueryAuthProvider,
     aiQueryMessageRepository: Type<AiQueryMessageRepository> = LocalAiQueryMessageRepository,
+    aiQueryContextRepository: Type<AiQueryContextRepository> = InMemoryAiQueryContextRepository,
   ) {
     return {
       module: AiQueryModule,
@@ -27,6 +32,10 @@ export class AiQueryModule {
           provide: AiQueryMessageRepository,
           useClass: aiQueryMessageRepository,
         },
+        {
+          provide: AiQueryContextRepository,
+          useClass: aiQueryContextRepository,
+        },
       ],
     };
   }

redisinsight/api/src/modules/ai/query/ai-query.service.ts

@@ -19,6 +19,7 @@ import { AiQueryMessageRepository } from 'src/modules/ai/query/repositories/ai-q
 import { AiQueryAuthProvider } from 'src/modules/ai/query/providers/auth/ai-query-auth.provider';
 import { classToClass } from 'src/utils';
 import { plainToClass } from 'class-transformer';
+import { AiQueryContextRepository } from 'src/modules/ai/query/repositories/ai-query.context.repository';
 
 const COMMANDS_WHITELIST = {
   'ft.search': true,
@@ -34,6 +35,7 @@ export class AiQueryService {
     private readonly databaseClientFactory: DatabaseClientFactory,
     private readonly aiQueryMessageRepository: AiQueryMessageRepository,
     private readonly aiQueryAuthProvider: AiQueryAuthProvider,
+    private readonly aiQueryContextRepository: AiQueryContextRepository,
   ) {}
 
   static prepareHistoryIntermediateSteps(message: AiQueryMessage): [AiQueryMessageRole, string][] {
@@ -93,7 +95,16 @@ export class AiQueryService {
         context: ClientContext.AI,
       });
 
-      const context = await getFullDbContext(client);
+      let context = await this.aiQueryContextRepository.getFullDbContext(sessionMetadata, databaseId, auth.accountId);
+
+      if (!context) {
+        context = await this.aiQueryContextRepository.setFullDbContext(
+          sessionMetadata,
+          databaseId,
+          auth.accountId,
+          await getFullDbContext(client),
+        );
+      }
 
       const question = classToClass(AiQueryMessage, {
         type: AiQueryMessageType.HumanMessage,
@@ -119,11 +130,27 @@ export class AiQueryService {
 
       socket.on(AiQueryWsEvents.GET_INDEX, async (index, cb) => {
         try {
-          const indexContext = await getIndexContext(client, index);
-          cb(indexContext);
+          const indexContext = await this.aiQueryContextRepository.getIndexContext(
+            sessionMetadata,
+            databaseId,
+            auth.accountId,
+            index,
+          );
+
+          if (!context) {
+            return cb(await this.aiQueryContextRepository.setIndexContext(
+              sessionMetadata,
+              databaseId,
+              auth.accountId,
+              index,
+              await getIndexContext(client, index),
+            ));
+          }
+
+          return cb(indexContext);
         } catch (e) {
           this.logger.warn('Unable to create index content', e);
-          cb(e.message);
+          return cb(e.message);
         }
       });
 
@@ -140,16 +167,6 @@ export class AiQueryService {
         }
       });
 
-      socket.on(AiQueryWsEvents.GET_INDEX, async (index, cb) => {
-        try {
-          const indexContext = await getIndexContext(client, index);
-          return cb(indexContext);
-        } catch (e) {
-          this.logger.warn('Unable to create index content', e);
-          return cb(e.message);
-        }
-      });
-
       socket.on(AiQueryWsEvents.TOOL_CALL, async (data) => {
         answer.steps.push(plainToClass(AiQueryIntermediateStep, {
           type: AiQueryIntermediateStepType.TOOL_CALL,
@@ -187,6 +204,9 @@ export class AiQueryService {
   async clearHistory(sessionMetadata: SessionMetadata, databaseId: string): Promise<void> {
     try {
       const auth = await this.aiQueryAuthProvider.getAuthData(sessionMetadata);
+
+      await this.aiQueryContextRepository.reset(sessionMetadata, databaseId, auth.accountId);
+
       return this.aiQueryMessageRepository.clearHistory(sessionMetadata, databaseId, auth.accountId);
     } catch (e) {
       throw wrapAiQueryError(e, 'Unable to clear history');

redisinsight/api/src/modules/ai/query/repositories/ai-query.context.repository.ts

@@ -0,0 +1,71 @@
+import { SessionMetadata } from 'src/common/models';
+
+export abstract class AiQueryContextRepository {
+  /**
+   * Should return saved db context if exists in particular chat
+   * @param sessionMetadata
+   * @param databaseId
+   * @param accountId
+   */
+  abstract getFullDbContext(
+    sessionMetadata: SessionMetadata,
+    databaseId: string,
+    accountId: string,
+  ): Promise<object>;
+
+  /**
+   * Should save db context for particular chat
+   * @param sessionMetadata
+   * @param databaseId
+   * @param accountId
+   * @param context
+   */
+  abstract setFullDbContext(
+    sessionMetadata: SessionMetadata,
+    databaseId: string,
+    accountId: string,
+    context: object,
+  ): Promise<object>;
+
+  /**
+   * Should return saved index context if exists in particular chat
+   * @param sessionMetadata
+   * @param databaseId
+   * @param accountId
+   * @param index
+   */
+  abstract getIndexContext(
+    sessionMetadata: SessionMetadata,
+    databaseId: string,
+    accountId: string,
+    index: string,
+  ): Promise<object>;
+
+  /**
+   * Should save index context for particular chat
+   * @param sessionMetadata
+   * @param databaseId
+   * @param accountId
+   * @param index
+   * @param context
+   */
+  abstract setIndexContext(
+    sessionMetadata: SessionMetadata,
+    databaseId: string,
+    accountId: string,
+    index: string,
+    context: object,
+  ): Promise<object>;
+
+  /**
+   * Reset all index and db contexts for particular chat
+   * @param sessionMetadata
+   * @param databaseId
+   * @param accountId
+   */
+  abstract reset(
+    sessionMetadata: SessionMetadata,
+    databaseId: string,
+    accountId: string,
+  ): Promise<void>;
+}

redisinsight/api/src/modules/ai/query/repositories/in-memory.ai-query.context.repository.ts

@@ -0,0 +1,86 @@
+import { Injectable } from '@nestjs/common';
+import { get, set, unset } from 'lodash';
+import { SessionMetadata } from 'src/common/models';
+import { AiQueryContextRepository } from 'src/modules/ai/query/repositories/ai-query.context.repository';
+
+@Injectable()
+export class InMemoryAiQueryContextRepository extends AiQueryContextRepository {
+  private chats: Record<string, { index: Record<string, object>, db: object }> = {};
+
+  static getChatId(databaseId: string, accountId: string) {
+    return `${databaseId}_${accountId}`;
+  }
+
+  /**
+   * @inheritdoc
+   */
+  async getFullDbContext(
+    _sessionMetadata: SessionMetadata,
+    databaseId: string,
+    accountId: string,
+  ): Promise<object> {
+    const chatId = InMemoryAiQueryContextRepository.getChatId(databaseId, accountId);
+
+    return get(this.chats, [chatId, 'db'], null);
+  }
+
+  /**
+   * @inheritdoc
+   */
+  async setFullDbContext(
+    _sessionMetadata: SessionMetadata,
+    databaseId: string,
+    accountId: string,
+    context: object,
+  ): Promise<object> {
+    const chatId = InMemoryAiQueryContextRepository.getChatId(databaseId, accountId);
+
+    set(this.chats, [chatId, 'db'], context);
+
+    return context;
+  }
+
+  /**
+   * @inheritdoc
+   */
+  async getIndexContext(
+    _sessionMetadata: SessionMetadata,
+    databaseId: string,
+    accountId: string,
+    index: string,
+  ): Promise<object> {
+    const chatId = InMemoryAiQueryContextRepository.getChatId(databaseId, accountId);
+
+    return get(this.chats, [chatId, 'index', index], null);
+  }
+
+  /**
+   * @inheritdoc
+   */
+  async setIndexContext(
+    _sessionMetadata: SessionMetadata,
+    databaseId: string,
+    accountId: string,
+    index: string,
+    context: object,
+  ): Promise<object> {
+    const chatId = InMemoryAiQueryContextRepository.getChatId(databaseId, accountId);
+
+    set(this.chats, [chatId, 'index', index], context);
+
+    return context;
+  }
+
+  /**
+   * @inheritdoc
+   */
+  async reset(
+    _sessionMetadata:SessionMetadata,
+    databaseId:string,
+    accountId:string,
+  ): Promise<void> {
+    const chatId = InMemoryAiQueryContextRepository.getChatId(databaseId, accountId);
+
+    unset(this.chats, [chatId]);
+  }
+}