feat(llms): finish stream call

sd0xdev · Aug 27, 2023 · 7fb8de9 · 7fb8de9
1 parent d0196fe
commit 7fb8de9
Show file tree

Hide file tree

Showing 13 changed files with 187 additions and 80 deletions.
diff --git a/apps/yggdrasil-core-engine/src/app/controllers/llm-ai/llm-ai.controller.spec.ts b/apps/yggdrasil-core-engine/src/app/controllers/llm-ai/llm-ai.controller.spec.ts
@@ -1,8 +1,7 @@
 import { Test, TestingModule } from '@nestjs/testing';
 import { LLMAIController } from './llm-ai.controller';
-import { AsgardLoggerSupplement } from '@asgard-hub/nest-winston';
 import { AuthGuard } from '../../auth/auth.guard';
-import { LangChainService } from '../../llm-ai/lang-chain/lang-chain.service';
+import { LLMAIService } from '../../llm-ai/service/llmai/llmai.service';
 
 describe('LlmAiController', () => {
   let controller: LLMAIController;
@@ -11,12 +10,10 @@ describe('LlmAiController', () => {
     const module: TestingModule = await Test.createTestingModule({
       providers: [
         {
-          provide: AsgardLoggerSupplement.LOGGER_HELPER_SERVICE,
-          useValue: {},
-        },
-        {
-          provide: LangChainService,
-          useValue: {},
+          provide: LLMAIService,
+          useValue: {
+            chat: jest.fn(),
+          },
         },
       ],
       controllers: [LLMAIController],

diff --git a/apps/yggdrasil-core-engine/src/app/controllers/llm-ai/llm-ai.controller.ts b/apps/yggdrasil-core-engine/src/app/controllers/llm-ai/llm-ai.controller.ts
@@ -1,10 +1,12 @@
 import { Controller, Inject, UseGuards } from '@nestjs/common';
-import { LangChainService } from '../../llm-ai/lang-chain/lang-chain.service';
 import { AuthGuard } from '../../auth/auth.guard';
 import { FunctionSelection } from '@asgard-hub/utils';
-import { TranslateToEnglishPrompt } from '../../llm-ai/prompt/translate-to-english.prompt';
-import { GrpcMethod } from '@nestjs/microservices';
-import { AsgardLoggerSupplement, AsgardLogger } from '@asgard-hub/nest-winston';
+import { GrpcMethod, GrpcStreamMethod } from '@nestjs/microservices';
+import {
+  GrpcResponse,
+  LLMAIService,
+} from '../../llm-ai/service/llmai/llmai.service';
+import { Observable, Subject, switchMap } from 'rxjs';
 
 export interface ChatOptions {
   userInput: string;
@@ -14,43 +16,26 @@ export interface ChatOptions {
 @UseGuards(AuthGuard)
 @Controller('llm-ai')
 export class LLMAIController {
-  @Inject(AsgardLoggerSupplement.LOGGER_HELPER_SERVICE)
-  private readonly asgardLogger: AsgardLogger;
   @Inject()
-  private readonly langChainService: LangChainService;
+  private readonly llmAIService: LLMAIService;
 
   @GrpcMethod('LLMAIService')
-  async chat(data: ChatOptions) {
-    const type = this.promptSelect(data.key);
-    let response = '';
-    if (!type) {
-      response = await this.langChainService.getGeneralChatResponse(
-        data.userInput
-      );
-    } else {
-      response = (
-        await this.langChainService.getFunctionChainResponse(
-          data.userInput,
-          type
-        )
-      ).message;
-    }
+  chat(data: ChatOptions) {
+    return this.llmAIService.chat(data);
+  }
 
-    this.asgardLogger.debug(response);
+  @GrpcStreamMethod('LLMAIService', 'ChatStream')
+  chatStream(data: Observable<ChatOptions>) {
+    const subject = new Subject<GrpcResponse>();
 
-    return {
-      response,
-    };
-  }
+    data
+      .pipe(
+        switchMap((data) => {
+          return this.llmAIService.chat(data, subject);
+        })
+      )
+      .subscribe();
 
-  promptSelect(key: FunctionSelection) {
-    switch (key) {
-      case FunctionSelection.generalChat:
-        return undefined;
-      case FunctionSelection.anyTranslateToEnglish:
-        return TranslateToEnglishPrompt;
-      default:
-        return undefined;
-    }
+    return subject.asObservable();
   }
 }
diff --git a/apps/yggdrasil-core-engine/src/app/llm-ai/llm-ai.module.ts b/apps/yggdrasil-core-engine/src/app/llm-ai/llm-ai.module.ts
@@ -8,10 +8,11 @@ import { IAppConfig, ConfigPath } from '../config/app.config';
 import { isDev, isStaging, isProd } from '../constants/common.constant';
 import { EventEmitterModule } from '@nestjs/event-emitter';
 import { LLMAIController } from '../controllers/llm-ai/llm-ai.controller';
-import { LangChainService } from './lang-chain/lang-chain.service';
+import { LangChainService } from './service/langchain/langchain.service';
 import { NestLangChainModule } from '@sd0x/nest-langchain';
 import { OpenAIProvider } from '../provider/open-ai/open-ai';
 import { ProviderModule } from '../provider/provider.module';
+import { LLMAIService } from './service/llmai/llmai.service';
 
 @Module({
   imports: [
@@ -43,8 +44,8 @@ import { ProviderModule } from '../provider/provider.module';
     MongoModule,
     DataSourceAdapterModule,
   ],
-  providers: [LangChainService],
-  exports: [LangChainService],
+  providers: [LangChainService, LLMAIService],
+  exports: [LLMAIService],
   controllers: [LLMAIController],
 })
 export class LLMAIModule {}
diff --git a/...-ai/lang-chain/lang-chain.service.spec.ts → ...rvice/langchain/langchain.service.spec.ts b/...-ai/lang-chain/lang-chain.service.spec.ts → ...rvice/langchain/langchain.service.spec.ts
@@ -1,5 +1,5 @@
 import { Test, TestingModule } from '@nestjs/testing';
-import { LangChainService } from './lang-chain.service';
+import { LangChainService } from './langchain.service';
 import { NestLangchainOptionsSupplement } from '@sd0x/nest-langchain';
 import { AsgardLoggerSupplement } from '@asgard-hub/nest-winston';
 

diff --git a/...p/llm-ai/lang-chain/lang-chain.service.ts → ...ai/service/langchain/langchain.service.ts b/...p/llm-ai/lang-chain/lang-chain.service.ts → ...ai/service/langchain/langchain.service.ts
@@ -20,23 +20,24 @@ export class LangChainService {
     userInput: string,
     langChainPromptType: { new (): BaseLangChainPrompt<T> }
   ) {
-    const result = await this.nestLangchainService.getFunctionChainResponse<T>({
+    return this.nestLangchainService.getFunctionChainResponse<T>({
       input: {
         userInput,
       },
       langChainPromptType,
     });
-
-    this.asgardLogger.debug(result);
-
-    return result;
   }
 
   async getGeneralChatResponse(userInput: string) {
     const message = await this.nestLangchainService.getGeneralChatResponse(
       userInput
     );
 
-    return message.content;
+    return {
+      metaOutput: {
+        ...message,
+      },
+      response: message.content,
+    };
   }
 }
diff --git a/apps/yggdrasil-core-engine/src/app/llm-ai/service/llmai/llmai.service.spec.ts b/apps/yggdrasil-core-engine/src/app/llm-ai/service/llmai/llmai.service.spec.ts
@@ -0,0 +1,30 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { LLMAIService } from './llmai.service';
+import { AsgardLoggerSupplement } from '@asgard-hub/nest-winston';
+import { LangChainService } from '../langchain/langchain.service';
+
+describe('LlmaiService', () => {
+  let service: LLMAIService;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [
+        {
+          provide: AsgardLoggerSupplement.LOGGER_HELPER_SERVICE,
+          useValue: {},
+        },
+        {
+          provide: LangChainService,
+          useValue: {},
+        },
+        LLMAIService,
+      ],
+    }).compile();
+
+    service = module.get<LLMAIService>(LLMAIService);
+  });
+
+  it('should be defined', () => {
+    expect(service).toBeDefined();
+  });
+});
diff --git a/apps/yggdrasil-core-engine/src/app/llm-ai/service/llmai/llmai.service.ts b/apps/yggdrasil-core-engine/src/app/llm-ai/service/llmai/llmai.service.ts
@@ -0,0 +1,72 @@
+import { Inject, Injectable } from '@nestjs/common';
+import { FunctionSelection } from '@asgard-hub/utils';
+import { AsgardLogger, AsgardLoggerSupplement } from '@asgard-hub/nest-winston';
+import { LangChainService } from '../langchain/langchain.service';
+import { ChatOptions } from '../../../controllers/llm-ai/llm-ai.controller';
+import { TranslateToEnglishPrompt } from '../../prompt/translate-to-english.prompt';
+import { Subject, of, switchMap } from 'rxjs';
+import { randomUUID } from 'crypto';
+
+export interface GrpcResponse<T = unknown> {
+  id: string;
+  timestamp: number;
+  uuid: string;
+  event: string;
+  data: string | T;
+}
+
+@Injectable()
+export class LLMAIService {
+  @Inject(AsgardLoggerSupplement.LOGGER_HELPER_SERVICE)
+  private readonly asgardLogger: AsgardLogger;
+  @Inject()
+  private readonly langChainService: LangChainService;
+  chat(data: ChatOptions, subject?: Subject<GrpcResponse>) {
+    return of(this.promptSelect(data.key)).pipe(
+      switchMap((prompt) => {
+        subject?.next({
+          id: 'prompt',
+          timestamp: Date.now(),
+          uuid: randomUUID(),
+          event: 'notification',
+          data: prompt?.name ?? FunctionSelection.generalChat,
+        });
+        return of(prompt);
+      }),
+      switchMap(async (type) => {
+        if (!type) {
+          return await this.langChainService.getGeneralChatResponse(
+            data.userInput
+          );
+        } else {
+          return await this.langChainService.getFunctionChainResponse(
+            data.userInput,
+            type
+          );
+        }
+      }),
+      switchMap((response) => {
+        subject?.next({
+          id: 'response',
+          timestamp: Date.now(),
+          uuid: randomUUID(),
+          event: 'message',
+          data: JSON.stringify(response),
+        });
+        this.asgardLogger.debug(response);
+        return of(subject.complete());
+      })
+    );
+  }
+
+  private promptSelect(key: FunctionSelection) {
+    switch (key) {
+      case FunctionSelection.generalChat:
+        return undefined;
+      case FunctionSelection.anyTranslateToEnglish:
+        return TranslateToEnglishPrompt;
+      default:
+        return undefined;
+    }
+  }
+}
diff --git a/apps/yggdrasil-core-engine/src/assets/llmai/proto/llmai.proto b/apps/yggdrasil-core-engine/src/assets/llmai/proto/llmai.proto
@@ -4,6 +4,7 @@ package llmai;
 
 service LLMAIService {
   rpc Chat (ChatOptions) returns (ChatCompletionResponse);
+  rpc ChatStream (stream ChatOptions) returns (stream ChatCompletionResponse);
 }
 
 message ChatOptions {
@@ -12,6 +13,9 @@ message ChatOptions {
 }
 
 message ChatCompletionResponse {
-    string response = 1;
-}
-
+  string id = 1;
+  int32 timestamp = 2;
+  string uuid = 3;
+  string event = 4;
+  string data = 5;
+}
diff --git a/apps/yggdrasil-discord-client/src/app/config/core.engine.config.ts b/apps/yggdrasil-discord-client/src/app/config/core.engine.config.ts
@@ -40,7 +40,7 @@ export const coreEngineLLMAIConfig = registerAs(
 );
 
 function getProtoPath(fileName: string) {
-  return isDev
-    ? resolve(__dirname, 'assets', fileName)
+  return process.env['IS_DEBUG']
+    ? resolve(__dirname, '..', '..', 'assets', fileName)
     : resolve(__dirname, '.', 'assets', fileName);
 }
diff --git a/apps/yggdrasil-discord-client/src/app/discord-bot/interface/chatgpt.service.interface.ts b/apps/yggdrasil-discord-client/src/app/discord-bot/interface/chatgpt.service.interface.ts
@@ -118,12 +118,21 @@ export interface ChatOptions {
 }
 
 export interface ChatCompletionResponse {
-  response: string;
+  id: string;
+  timestamp: number;
+  event: string;
+  uuid: string;
+  data: string;
 }
 
 export interface LLMAIService {
   chat(
     options: ChatOptions,
     metadata?: Metadata
   ): Observable<ChatCompletionResponse>;
+
+  chatStream(
+    options: Observable<ChatOptions>,
+    metadata?: Metadata
+  ): Observable<ChatCompletionResponse>;
 }
diff --git a/apps/yggdrasil-discord-client/src/app/discord-bot/services/message/message.service.ts b/apps/yggdrasil-discord-client/src/app/discord-bot/services/message/message.service.ts
@@ -8,7 +8,7 @@ import {
   ThreadChannel,
 } from 'discord.js';
 import { List, Stack } from 'immutable';
-import { lastValueFrom } from 'rxjs';
+import { ReplaySubject, lastValueFrom, switchMap } from 'rxjs';
 import {
   DISCORD_BOT_MODULE_OPTIONS,
   DiSCORD_SPLIT_MESSAGE_TARGET,
@@ -23,6 +23,7 @@ import { splitString } from '../../utils/split-string';
 import { delay } from '../../utils/delay';
 import {
   ChatGPTService,
+  ChatOptions,
   LLMAIService,
 } from '../../interface/chatgpt.service.interface';
 import { DiscordClientService } from '../discord-client/discord-client.service';
@@ -297,21 +298,24 @@ export class MessageService implements OnModuleInit {
     try {
       // show typing
       await (message.channel as TextChannel).sendTyping();
-
-      const result = await lastValueFrom(
-        this.llmAIService.chat(
-          {
-            userInput: message.content,
-          },
-          this.metadata
+      const request$ = new ReplaySubject<ChatOptions>();
+      request$.next({
+        userInput: message.content,
+      });
+      request$.complete();
+      this.llmAIService
+        .chatStream(request$, this.metadata)
+        .pipe(
+          switchMap(async (result) => {
+            await this.sendMessageReply(result?.data, message);
+
+            this.asgardLogger.log(
+              `successfully send message: ${result?.data}`
+            );
+            this.asgardLogger.log(`successfully send message`);
+          })
         )
-      );
-
-      const response = result.response;
-
-      await this.sendMessageReply(response, message);
-      this.asgardLogger.log(`successfully send message: ${response}`);
-      this.asgardLogger.log(`successfully send message`);
+        .subscribe();
     } catch (e) {
       this.asgardLogger.error(e);
       await this.sendMessageReply(