Merge pull request #5 from opensearch-project/feature/add-builder-reg…

…istration

feat: add mechanism to register messageParser

CHANGELOG.md

@@ -0,0 +1,7 @@
+# CHANGELOG
+
+Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
+
+### 📈 Features/Enhancements
+
+- Add support for registerMessageParser ([#5](https://github.com/opensearch-project/dashboards-assistant/pull/5))

public/tabs/chat/chat_page.tsx

@@ -6,11 +6,11 @@
 import { EuiFlyoutBody, EuiFlyoutFooter, EuiPage, EuiPageBody, EuiSpacer } from '@elastic/eui';
 import React, { useCallback, useState } from 'react';
 import cs from 'classnames';
+import { useObservable } from 'react-use';
 import { useChatContext } from '../../contexts/chat_context';
 import { useChatState } from '../../hooks/use_chat_state';
 import { ChatPageContent } from './chat_page_content';
 import { ChatInputControls } from './controls/chat_input_controls';
-import { useObservable } from 'react-use';
 import { useCore } from '../../contexts/core_context';
 
 interface ChatPageProps {

server/README.md

@@ -0,0 +1,27 @@
+# `registerMessageParser` — Register your customized parser logic into Chatbot.
+
+**Interaction** refers to a question-answer pair in Chatbot application. In most cases, an interaction consists of two messages: an `Input` message and an `Output` message. However, as the Chatbot evolves to become more powerful, it may display new messages such as visualizations, data explorers, or data grids. Therefore, it is crucial to implement a mechanism that allows other plugins to register their customized parser logic based on each interaction body.
+
+![message parser](https://github.com/opensearch-project/dashboards-assistant/assets/13493605/b4ec1ff8-5339-4119-ad20-b2c31057bb0b)
+
+## API
+
+### registerMessageParser
+
+```typescript
+dashboardAssistant.registerMessageParser({
+  id: 'foo_parser',
+  parserProvider: async (interaction) => {
+    if (interaction.input) {
+      return [
+        {
+          type: 'input',
+          contentType: 'text',
+          content: interaction.input,
+        },
+      ];
+    }
+    return [];
+  },
+});
+```

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();
@@ -57,7 +58,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,
@@ -133,7 +134,7 @@ export function registerChatRoutes(router: IRouter) {
           body: {
             messages: finalMessage.messages,
             sessionId: outputs.memoryId,
-            title: finalMessage.title
+            title: finalMessage.title,
           },
         });
       } catch (error) {
@@ -271,13 +272,16 @@ export function registerChatRoutes(router: IRouter) {
         const outputs = await chatService.requestLLM(
           { messages, input, sessionId },
           context,
-          request as any
+          // @ts-ignore
+          request
         );
         const title = input.content.substring(0, 50);
         const saveMessagesResponse = await storageService.saveMessages(
           title,
           sessionId,
-          [...messages, input, ...outputs.messages].filter((message) => message.content !== 'AbortError')
+          [...messages, input, ...outputs.messages].filter(
+            (message) => message.content !== 'AbortError'
+          )
         );
         return response.ok({
           body: { ...saveMessagesResponse, title },

server/routes/index.ts

@@ -3,6 +3,7 @@
  * 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';

server/types.ts

@@ -3,13 +3,45 @@
  * 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 Interaction {
+  input: string;
+  response: string;
+  conversation_id: string;
+  interaction_id: string;
+  create_time: string;
+  additional_info: Record<string, unknown>;
+}
+
+export interface MessageParser {
+  /**
+   * The id of the parser, should be unique among the parsers.
+   */
+  id: string;
+  /**
+   * Order field declares the order message parser will be execute.
+   * parser with order 2 will be executed after parser with order 1.
+   * If not specified, the default order will be 999.
+   * @default 999
+   */
+  order?: number;
+  /**
+   * parserProvider is the callback that will be triggered in each message
+   */
+  parserProvider: (interaction: Interaction) => Promise<IMessage[]>;
+}
+
+export interface RoutesOptions {
+  messageParsers: MessageParser[];
+}
+
 declare module '../../../src/core/server' {
   interface RequestHandlerContext {
     assistant_plugin: {

server/utils/message_parser_runner.test.ts

@@ -0,0 +1,150 @@
+/*
+ * 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) {
+          return Promise.resolve([
+            {
+              type: 'output',
+              contentType: 'markdown',
+              content: interaction.response,
+            },
+          ]);
+        },
+      },
+    ]);
+
+    expect(
+      await messageParserRunner.run({
+        response: 'output',
+        input: 'input',
+      })
+    ).toEqual([
+      {
+        type: 'output',
+        contentType: 'markdown',
+        content: 'output',
+      },
+    ]);
+  });
+
+  it('run with correct result when different order is present', async () => {
+    const messageParserRunner = new MessageParserRunner([
+      {
+        id: 'testA',
+        order: 2,
+        parserProvider() {
+          return Promise.resolve([
+            {
+              type: 'output',
+              contentType: 'markdown',
+              content: 'A',
+            },
+          ]);
+        },
+      },
+      {
+        id: 'testOrder1000',
+        order: 1000,
+        parserProvider() {
+          return Promise.resolve([
+            {
+              type: 'output',
+              contentType: 'markdown',
+              content: 'order1000',
+            },
+          ]);
+        },
+      },
+      {
+        id: 'testNoOrder',
+        parserProvider(interaction) {
+          return Promise.resolve([
+            {
+              type: 'output',
+              contentType: 'markdown',
+              content: 'NoOrder',
+            },
+          ]);
+        },
+      },
+      {
+        id: 'testB',
+        order: 1,
+        parserProvider() {
+          return Promise.resolve([
+            {
+              type: 'output',
+              contentType: 'markdown',
+              content: 'B',
+            },
+          ]);
+        },
+      },
+    ]);
+
+    expect(
+      await messageParserRunner.run({
+        response: 'output',
+        input: 'input',
+      })
+    ).toEqual([
+      {
+        type: 'output',
+        contentType: 'markdown',
+        content: 'B',
+      },
+      {
+        type: 'output',
+        contentType: 'markdown',
+        content: 'A',
+      },
+      {
+        type: 'output',
+        contentType: 'markdown',
+        content: 'NoOrder',
+      },
+      {
+        type: 'output',
+        contentType: 'markdown',
+        content: 'order1000',
+      },
+    ]);
+  });
+
+  it('Do not append messages that are throwed with error or not an array', async () => {
+    const messageParserRunner = new MessageParserRunner([
+      {
+        id: 'test_with_error',
+        parserProvider() {
+          throw new Error('error');
+        },
+      },
+      {
+        id: 'test_with_incorrect_format_of_return',
+        parserProvider() {
+          return Promise.resolve({
+            type: 'output',
+            contentType: 'markdown',
+            content: 'order1000',
+          });
+        },
+      },
+    ]);
+
+    expect(
+      await messageParserRunner.run({
+        response: 'output',
+        input: 'input',
+      })
+    ).toEqual([]);
+  });
+});

server/utils/message_parser_runner.ts

@@ -0,0 +1,36 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { IMessage } from '../../common/types/chat_saved_object_attributes';
+import { Interaction, MessageParser } from '../types';
+
+export class MessageParserRunner {
+  constructor(private readonly messageParsers: MessageParser[]) {}
+  async run(interaction: Interaction): Promise<IMessage[]> {
+    const sortedParsers = [...this.messageParsers];
+    sortedParsers.sort((parserA, parserB) => {
+      const { order: orderA = 999 } = parserA;
+      const { order: orderB = 999 } = parserB;
+      return orderA - orderB;
+    });
+    let results: IMessage[] = [];
+    for (const messageParser of sortedParsers) {
+      let tempResult: IMessage[] = [];
+      try {
+        tempResult = await messageParser.parserProvider(interaction);
+        /**
+         * Make sure the tempResult is an array.
+         */
+        if (!Array.isArray(tempResult)) {
+          tempResult = [];
+        }
+      } catch (e) {
+        tempResult = [];
+      }
+      results = [...results, ...tempResult];
+    }
+    return results;
+  }
+}