diff --git a/server/plugin.ts b/server/plugin.ts index 3b3fc35f..47107178 100644 --- a/server/plugin.ts +++ b/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 { private readonly logger: Logger; + private messageParsers: MessageParser[] = []; constructor(initializerContext: PluginInitializerContext) { this.logger = initializerContext.logger.get(); @@ -46,7 +47,9 @@ export class AssistantPlugin implements Plugin { + 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) { diff --git a/server/routes/chat_routes.ts b/server/routes/chat_routes.ts index 01329e1f..18bb4680 100644 --- a/server/routes/chat_routes.ts +++ b/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(); diff --git a/server/routes/index.ts b/server/routes/index.ts index ae33e1c3..0b2d7eee 100644 --- a/server/routes/index.ts +++ b/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); } diff --git a/server/services/storage/saved_objects_storage_service.ts b/server/services/storage/saved_objects_storage_service.ts index 78fcffb4..f85bba48 100644 --- a/server/services/storage/saved_objects_storage_service.ts +++ b/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; diff --git a/server/types.ts b/server/types.ts index bb72cc4b..3c1ffca5 100644 --- a/server/types.ts +++ b/server/types.ts @@ -3,6 +3,7 @@ * 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 @@ -10,6 +11,38 @@ 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; +} + +export interface RoutesOptions { + messageParsers: MessageParser[]; +} + declare module '../../../src/core/server' { interface RequestHandlerContext { assistant_plugin: { diff --git a/server/utils/message_parser_helper.test.ts b/server/utils/message_parser_helper.test.ts new file mode 100644 index 00000000..7799892a --- /dev/null +++ b/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', + }, + ]); + }); +}); diff --git a/server/utils/message_parser_helper.ts b/server/utils/message_parser_helper.ts new file mode 100644 index 00000000..36ff0963 --- /dev/null +++ b/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; + } +} diff --git a/server/utils/message_parser_runner.test.ts b/server/utils/message_parser_runner.test.ts new file mode 100644 index 00000000..03ce3a1b --- /dev/null +++ b/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', + }, + ]); + }); +}); diff --git a/server/utils/message_parser_runner.ts b/server/utils/message_parser_runner.ts new file mode 100644 index 00000000..3545d02d --- /dev/null +++ b/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 { + const messageParserHelper = new MessageParserHelper(); + for (const messageParser of this.messageParsers) { + await messageParser.parserProvider(interaction, messageParserHelper); + } + return messageParserHelper.messages; + } +}