From d598bff4e1c0bf616c6214bb54c954749f272c6a Mon Sep 17 00:00:00 2001
From: SuZhou-Joe <suzhou@amazon.com>
Date: Wed, 22 Nov 2023 15:34:26 +0800
Subject: [PATCH] feat: optimize

Signed-off-by: SuZhou-Joe <suzhou@amazon.com>
---
 public/tabs/chat/chat_page.tsx |  2 +-
 server/plugin.ts               |  4 +---
 server/routes/chat_routes.ts   | 11 +++++++----
 server/routes/index.ts         |  4 ++--
 server/types.ts                |  8 ++++----
 5 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/public/tabs/chat/chat_page.tsx b/public/tabs/chat/chat_page.tsx
index 6ca8893f..f8101d2b 100644
--- a/public/tabs/chat/chat_page.tsx
+++ b/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 {
diff --git a/server/plugin.ts b/server/plugin.ts
index 47107178..7559d9af 100644
--- a/server/plugin.ts
+++ b/server/plugin.ts
@@ -47,9 +47,7 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
     });
 
     // Register server side APIs
-    setupRoutes(router, {
-      messageParsers: this.messageParsers,
-    });
+    setupRoutes(router);
 
     core.savedObjects.registerType(chatSavedObject);
     core.savedObjects.registerType(chatConfigSavedObject);
diff --git a/server/routes/chat_routes.ts b/server/routes/chat_routes.ts
index 18bb4680..de4a4dbf 100644
--- a/server/routes/chat_routes.ts
+++ b/server/routes/chat_routes.ts
@@ -105,7 +105,7 @@ const updateSessionRoute = {
   },
 };
 
-export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions) {
+export function registerChatRoutes(router: IRouter) {
   const createStorageService = (context: RequestHandlerContext) =>
     new AgentFrameworkStorageService(context.core.opensearch.client.asCurrentUser);
   const createChatService = () => new OllyChatService();
@@ -134,7 +134,7 @@ export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions)
           body: {
             messages: finalMessage.messages,
             sessionId: outputs.memoryId,
-            title: finalMessage.title
+            title: finalMessage.title,
           },
         });
       } catch (error) {
@@ -272,13 +272,16 @@ export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions)
         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 },
diff --git a/server/routes/index.ts b/server/routes/index.ts
index 0b2d7eee..52426cf2 100644
--- a/server/routes/index.ts
+++ b/server/routes/index.ts
@@ -8,7 +8,7 @@ import { IRouter } from '../../../../src/core/server';
 import { registerChatRoutes } from './chat_routes';
 import { registerLangchainRoutes } from './langchain_routes';
 
-export function setupRoutes(router: IRouter, options: RoutesOptions) {
-  registerChatRoutes(router, options);
+export function setupRoutes(router: IRouter) {
+  registerChatRoutes(router);
   registerLangchainRoutes(router);
 }
diff --git a/server/types.ts b/server/types.ts
index dbcda88d..625f51ba 100644
--- a/server/types.ts
+++ b/server/types.ts
@@ -11,13 +11,13 @@ 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;
+  conversation_id: string;
+  interaction_id: string;
+  create_time: string;
+  additional_info: Record<string, unknown>;
 }
 
 export interface MessageParser {