feat: add interaction into message_bubble.tsx

Signed-off-by: SuZhou-Joe <[email protected]>
opensearch-project · Nov 22, 2023 · cdfae7a · cdfae7a
1 parent dae5a21
commit cdfae7a
Show file tree

Hide file tree

Showing 7 changed files with 67 additions and 11 deletions.
diff --git a/public/hooks/use_chat_actions.tsx b/public/hooks/use_chat_actions.tsx
@@ -4,7 +4,11 @@
  */
 
 import { ASSISTANT_API } from '../../common/constants/llm';
-import { IMessage, ISuggestedAction } from '../../common/types/chat_saved_object_attributes';
+import {
+  IMessage,
+  ISuggestedAction,
+  Interaction,
+} from '../../common/types/chat_saved_object_attributes';
 import { useChatContext } from '../contexts/chat_context';
 import { useCore } from '../contexts/core_context';
 import { AssistantActions } from '../types';
@@ -14,6 +18,7 @@ interface SendResponse {
   sessionId: string;
   title: string;
   messages: IMessage[];
+  interactions: Interaction[];
 }
 
 interface SetParagraphResponse {
@@ -47,7 +52,13 @@ export const useChatActions = (): AssistantActions => {
       if (!chatContext.title) {
         chatContext.setTitle(response.title);
       }
-      chatStateDispatch({ type: 'receive', payload: response.messages });
+      chatStateDispatch({
+        type: 'receive',
+        payload: {
+          messages: response.messages,
+          interactions: response.interactions,
+        },
+      });
     } catch (error) {
       if (abortController.signal.aborted) return;
       chatStateDispatch({ type: 'error', payload: error });
@@ -70,7 +81,13 @@ export const useChatActions = (): AssistantActions => {
     }
     const session = await core.services.sessionLoad.load(sessionId);
     if (session) {
-      chatStateDispatch({ type: 'receive', payload: session.messages });
+      chatStateDispatch({
+        type: 'receive',
+        payload: {
+          messages: session.messages,
+          interactions: session.interactions,
+        },
+      });
     }
   };
 
@@ -147,7 +164,13 @@ export const useChatActions = (): AssistantActions => {
         if (abortController.signal.aborted) {
           return;
         }
-        chatStateDispatch({ type: 'receive', payload: response.messages });
+        chatStateDispatch({
+          type: 'receive',
+          payload: {
+            messages: response.messages,
+            interactions: response.interactions,
+          },
+        });
       } catch (error) {
         if (abortController.signal.aborted) {
           return;

diff --git a/public/hooks/use_chat_state.tsx b/public/hooks/use_chat_state.tsx
@@ -5,10 +5,11 @@
 
 import { produce } from 'immer';
 import React, { useContext, useMemo, useReducer } from 'react';
-import { IMessage } from '../../common/types/chat_saved_object_attributes';
+import { IMessage, Interaction } from '../../common/types/chat_saved_object_attributes';
 
 interface ChatState {
   messages: IMessage[];
+  interactions: Interaction[];
   llmResponding: boolean;
   llmError?: Error;
 }
@@ -18,7 +19,13 @@ type ChatStateAction =
   | { type: 'abort' }
   | { type: 'reset' }
   | { type: 'send'; payload: IMessage }
-  | { type: 'receive'; payload: ChatState['messages'] }
+  | {
+      type: 'receive';
+      payload: {
+        messages: ChatState['messages'];
+        interactions: ChatState['interactions'];
+      };
+    }
   | {
       type: 'error';
       payload: NonNullable<ChatState['llmError']> | { body: NonNullable<ChatState['llmError']> };
@@ -31,6 +38,7 @@ interface IChatStateContext {
 const ChatStateContext = React.createContext<IChatStateContext | null>(null);
 
 const initialState: ChatState = {
+  interactions: [],
   messages: [],
   llmResponding: false,
 };
@@ -48,7 +56,8 @@ const chatStateReducer: React.Reducer<ChatState, ChatStateAction> = (state, acti
         break;
 
       case 'receive':
-        draft.messages = action.payload;
+        draft.messages = action.payload.messages;
+        draft.interactions = action.payload.interactions;
         draft.llmResponding = false;
         draft.llmError = undefined;
         break;

diff --git a/public/tabs/chat/chat_page.tsx b/public/tabs/chat/chat_page.tsx
@@ -31,7 +31,13 @@ export const ChatPage: React.FC<ChatPageProps> = (props) => {
     }
     const session = await core.services.sessionLoad.load(chatContext.sessionId);
     if (session) {
-      chatStateDispatch({ type: 'receive', payload: session.messages });
+      chatStateDispatch({
+        type: 'receive',
+        payload: {
+          messages: session.messages,
+          interactions: session.interactions,
+        },
+      });
     }
   }, [chatContext.sessionId, chatStateDispatch]);
 

diff --git a/public/tabs/chat/chat_page_content.tsx b/public/tabs/chat/chat_page_content.tsx
@@ -14,7 +14,11 @@ import {
   EuiText,
 } from '@elastic/eui';
 import React, { useLayoutEffect, useRef } from 'react';
-import { IMessage, ISuggestedAction } from '../../../common/types/chat_saved_object_attributes';
+import {
+  IMessage,
+  ISuggestedAction,
+  Interaction,
+} from '../../../common/types/chat_saved_object_attributes';
 import { TermsAndConditions } from '../../components/terms_and_conditions';
 import { useChatContext } from '../../contexts/chat_context';
 import { useChatState } from '../../hooks/use_chat_state';
@@ -120,6 +124,13 @@ export const ChatPageContent: React.FC<ChatPageContentProps> = React.memo((props
         // Only show suggestion on llm outputs after last user input
         const showSuggestions = i > lastInputIndex;
 
+        let interaction: Interaction | undefined;
+        if (message.type === 'output' && message.traceId) {
+          interaction = chatState.interactions.find(
+            (item) => item.interaction_id === message.traceId
+          );
+        }
+
         return (
           <React.Fragment key={i}>
             <ToolsUsed message={message} />
@@ -129,6 +140,7 @@ export const ChatPageContent: React.FC<ChatPageContentProps> = React.memo((props
               showRegenerate={isLatestOutput}
               shouldActionBarVisibleOnHover={!isLatestOutput}
               onRegenerate={chatActions.regenerate}
+              interaction={interaction}
             >
               <MessageContent message={message} />
               {/* <MessageFooter message={message} previousInput={findPreviousInput(array, i)} />*/}

diff --git a/public/tabs/chat/messages/message_bubble.tsx b/public/tabs/chat/messages/message_bubble.tsx
@@ -19,7 +19,11 @@ import React, { useCallback } from 'react';
 import { IconType } from '@elastic/eui/src/components/icon/icon';
 import cx from 'classnames';
 import chatIcon from '../../../assets/chat.svg';
-import { IMessage, IOutput } from '../../../../common/types/chat_saved_object_attributes';
+import {
+  IMessage,
+  IOutput,
+  Interaction,
+} from '../../../../common/types/chat_saved_object_attributes';
 import { useFeedback } from '../../../hooks/use_feed_back';
 
 type MessageBubbleProps = {
@@ -30,6 +34,7 @@ type MessageBubbleProps = {
 } & (
   | {
       message: IMessage;
+      interaction?: Interaction;
     }
   | {
       loading: boolean;

diff --git a/server/routes/chat_routes.ts b/server/routes/chat_routes.ts
@@ -137,6 +137,7 @@ export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions)
             messages: finalMessage.messages,
             sessionId: outputs.memoryId,
             title: finalMessage.title,
+            interactions: finalMessage.interactions,
           },
         });
       } catch (error) {

diff --git a/server/services/storage/agent_framework_storage_service.ts b/server/services/storage/agent_framework_storage_service.ts
@@ -54,7 +54,7 @@ export class AgentFrameworkStorageService implements StorageService {
       createdTimeMs: Date.now(),
       updatedTimeMs: Date.now(),
       messages: finalMessages,
-      interactions: session.body.interactions,
+      interactions: finalInteractions,
     };
   }