feat: add mechannism to register messageParser

Signed-off-by: SuZhou-Joe <[email protected]>

server/plugin.ts

@@ -18,11 +18,12 @@ import { PPLPlugin } from './adaptors/ppl_plugin';
 import './fetch-polyfill';
 import { setupRoutes } from './routes/index';
 import { chatSavedObject } from './saved_objects/chat_saved_object';
-import { AssistantPluginSetup, AssistantPluginStart } from './types';
+import { AssistantPluginSetup, AssistantPluginStart, MessageParser } from './types';
 import { chatConfigSavedObject } from './saved_objects/chat_config_saved_object';
 
 export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPluginStart> {
   private readonly logger: Logger;
+  private messageParsers: MessageParser[] = [];
 
   constructor(initializerContext: PluginInitializerContext) {
     this.logger = initializerContext.logger.get();
@@ -46,7 +47,9 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
     });
 
     // Register server side APIs
-    setupRoutes(router);
+    setupRoutes(router, {
+      messageParsers: this.messageParsers,
+    });
 
     core.savedObjects.registerType(chatSavedObject);
     core.savedObjects.registerType(chatConfigSavedObject);
@@ -57,7 +60,24 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
       },
     }));
 
-    return {};
+    return {
+      registerMessageParser: (messageParser: MessageParser) => {
+        const findItem = this.messageParsers.find((item) => item.id === messageParser.id);
+        if (findItem) {
+          throw new Error(`There is already a messageParser whose id is ${messageParser.id}`);
+        }
+
+        this.messageParsers.push(messageParser);
+      },
+      removeMessageParser: (parserId: MessageParser['id']) => {
+        const findIndex = this.messageParsers.findIndex((item) => item.id === parserId);
+        if (findIndex < 0) {
+          this.logger.error(`There is not a messageParser whose id is ${parserId}`);
+        }
+
+        this.messageParsers.splice(findIndex, 1);
+      },
+    };
   }
 
   public start(core: CoreStart) {

server/routes/chat_routes.ts

@@ -16,6 +16,7 @@ import { OllyChatService } from '../services/chat/olly_chat_service';
 import { SavedObjectsStorageService } from '../services/storage/saved_objects_storage_service';
 import { IMessage, IInput } from '../../common/types/chat_saved_object_attributes';
 import { AgentFrameworkStorageService } from '../services/storage/agent_framework_storage_service';
+import { RoutesOptions } from '../types';
 
 const llmRequestRoute = {
   path: ASSISTANT_API.SEND_MESSAGE,
@@ -104,7 +105,7 @@ const updateSessionRoute = {
   },
 };
 
-export function registerChatRoutes(router: IRouter) {
+export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions) {
   const createStorageService = (context: RequestHandlerContext) =>
     new AgentFrameworkStorageService(context.core.opensearch.client.asCurrentUser);
   const createChatService = () => new OllyChatService();

server/routes/index.ts

@@ -3,11 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { RoutesOptions } from '../types';
 import { IRouter } from '../../../../src/core/server';
 import { registerChatRoutes } from './chat_routes';
 import { registerLangchainRoutes } from './langchain_routes';
 
-export function setupRoutes(router: IRouter) {
-  registerChatRoutes(router);
+export function setupRoutes(router: IRouter, options: RoutesOptions) {
+  registerChatRoutes(router, options);
   registerLangchainRoutes(router);
 }

server/services/storage/saved_objects_storage_service.ts

@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { MessageParser } from '../../types';
 import { SavedObjectsClientContract } from '../../../../../src/core/server';
 import {
   CHAT_SAVED_OBJECT,
@@ -15,7 +16,10 @@ import { GetSessionsSchema } from '../../routes/chat_routes';
 import { StorageService } from './storage_service';
 
 export class SavedObjectsStorageService implements StorageService {
-  constructor(private readonly client: SavedObjectsClientContract) {}
+  constructor(
+    private readonly client: SavedObjectsClientContract,
+    private readonly messageParsers: MessageParser[]
+  ) {}
 
   private convertUpdatedTimeField(updatedAt: string | undefined) {
     return updatedAt ? new Date(updatedAt).getTime() : undefined;

server/types.ts

@@ -3,13 +3,46 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { IMessage } from '../common/types/chat_saved_object_attributes';
 import { ILegacyClusterClient, Logger } from '../../../src/core/server';
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
 export interface AssistantPluginSetup {}
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
 export interface AssistantPluginStart {}
 
+export interface IMessageParserHelper {
+  addMessage: (message: IMessage) => boolean;
+}
+
+export interface Interaction {
+  input: string;
+  response: string;
+}
+
+export interface MessageParser {
+  /**
+   * The id of the parser, should be unique among the parsers.
+   */
+  id: string;
+  /**
+   * Order declare the order message parser will be execute.
+   * parser with order 2 will be execute before parser with order 1.
+   */
+  order?: number;
+  /**
+   * parserProvider is the callback that will be triggered in each message
+   */
+  parserProvider: (
+    interaction: Interaction,
+    messageParserHelper: IMessageParserHelper
+  ) => Promise<unknown>;
+}
+
+export interface RoutesOptions {
+  messageParsers: MessageParser[];
+}
+
 declare module '../../../src/core/server' {
   interface RequestHandlerContext {
     assistant_plugin: {

server/utils/message_parser_helper.test.ts

@@ -0,0 +1,24 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { MessageParserHelper } from './message_parser_helper';
+
+describe('MessageParserHelper', () => {
+  it('return with correct message', async () => {
+    const messageParserHelper = new MessageParserHelper();
+    messageParserHelper.addMessage({
+      type: 'output',
+      contentType: 'markdown',
+      content: 'output',
+    });
+    expect(messageParserHelper.messages).toEqual([
+      {
+        type: 'output',
+        contentType: 'markdown',
+        content: 'output',
+      },
+    ]);
+  });
+});

server/utils/message_parser_helper.ts

@@ -0,0 +1,15 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { IMessage } from '../../common/types/chat_saved_object_attributes';
+import { IMessageParserHelper } from '../types';
+
+export class MessageParserHelper implements IMessageParserHelper {
+  public messages: IMessage[] = [];
+  addMessage(message: IMessage) {
+    this.messages.push(message);
+    return true;
+  }
+}

server/utils/message_parser_runner.test.ts

@@ -0,0 +1,37 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { MessageParserRunner } from './message_parser_runner';
+
+describe('MessageParserRunner', () => {
+  it('run with correct result', async () => {
+    const messageParserRunner = new MessageParserRunner([
+      {
+        id: 'test',
+        parserProvider(interaction, messageParserHelper) {
+          messageParserHelper.addMessage({
+            type: 'output',
+            contentType: 'markdown',
+            content: interaction.response,
+          });
+          return Promise.resolve('');
+        },
+      },
+    ]);
+
+    expect(
+      await messageParserRunner.run({
+        response: 'output',
+        input: 'input',
+      })
+    ).toEqual([
+      {
+        type: 'output',
+        contentType: 'markdown',
+        content: 'output',
+      },
+    ]);
+  });
+});

server/utils/message_parser_runner.ts

@@ -0,0 +1,19 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { IMessage } from '../../common/types/chat_saved_object_attributes';
+import { Interaction, MessageParser } from '../types';
+import { MessageParserHelper } from './message_parser_helper';
+
+export class MessageParserRunner {
+  constructor(private readonly messageParsers: MessageParser[]) {}
+  async run(interaction: Interaction): Promise<IMessage[]> {
+    const messageParserHelper = new MessageParserHelper();
+    for (const messageParser of this.messageParsers) {
+      await messageParser.parserProvider(interaction, messageParserHelper);
+    }
+    return messageParserHelper.messages;
+  }
+}