opensearch-project · SuZhou-Joe · Nov 30, 2023 · Nov 29, 2023 · Nov 30, 2023 · Nov 30, 2023
diff --git a/common/constants/llm.ts b/common/constants/llm.ts
@@ -18,6 +18,7 @@ export const ASSISTANT_API = {
   FEEDBACK: `${API_BASE}/feedback`,
   ABORT_AGENT_EXECUTION: `${API_BASE}/abort`,
   REGENERATE: `${API_BASE}/regenerate`,
+  TRACE: `${API_BASE}/trace`,
 } as const;
 
 export const LLM_INDEX = {

diff --git a/common/utils/llm_chat/traces.ts b/common/utils/llm_chat/traces.ts
@@ -7,6 +7,16 @@ import { Run } from 'langchain/callbacks';
 import { AgentRun } from 'langchain/dist/callbacks/handlers/tracer';
 import _ from 'lodash';
 
+export interface AgentFrameworkTrace {
+  interactionId: string;
+  parentInteractionId: string;
+  createTime: string;
+  input: string;
+  output: string;
+  origin: string;
+  traceNumber: number;
+}
+
 export interface LangchainTrace {
   id: Run['id'];
   parentRunId?: Run['parent_run_id'];

diff --git a/public/chat_flyout.tsx b/public/chat_flyout.tsx
@@ -10,7 +10,8 @@ import { useChatContext } from './contexts/chat_context';
 import { ChatPage } from './tabs/chat/chat_page';
 import { ChatWindowHeader } from './tabs/chat_window_header';
 import { ChatHistoryPage } from './tabs/history/chat_history_page';
-import { LangchainTracesFlyoutBody } from './components/langchain_traces_flyout_body';
+import { AgentFrameworkTracesFlyoutBody } from './components/agent_framework_traces_flyout_body';
+import { TAB_ID } from './utils/constants';
 
 let chatHistoryPageLoaded = false;
 
@@ -31,15 +32,15 @@ export const ChatFlyout: React.FC<ChatFlyoutProps> = (props) => {
 
   if (!props.overrideComponent) {
     switch (chatContext.selectedTabId) {
-      case 'chat':
+      case TAB_ID.CHAT:
         chatPageVisible = true;
         break;
 
-      case 'history':
+      case TAB_ID.HISTORY:
         chatHistoryPageVisible = true;
         break;
 
-      case 'trace':
+      case TAB_ID.TRACE:
         chatTraceVisible = true;
         break;
 
@@ -134,7 +135,7 @@ export const ChatFlyout: React.FC<ChatFlyoutProps> = (props) => {
                       className={cs({ 'llm-chat-hidden': !chatHistoryPageVisible })}
                     />
                   )}
-                  {chatTraceVisible && chatContext.traceId && <LangchainTracesFlyoutBody />}
+                  {chatTraceVisible && chatContext.traceId && <AgentFrameworkTracesFlyoutBody />}
                 </Panel>
               </>
             </>

diff --git a/public/chat_header_button.tsx b/public/chat_header_button.tsx
@@ -15,6 +15,7 @@ import { ChatStateProvider } from './hooks/use_chat_state';
 import './index.scss';
 import chatIcon from './assets/chat.svg';
 import { ActionExecutor, AssistantActions, ContentRenderer, UserAccount, TabId } from './types';
+import { TAB_ID } from './utils/constants';
 
 interface HeaderChatButtonProps {
   application: ApplicationStart;
@@ -33,7 +34,7 @@ export const HeaderChatButton: React.FC<HeaderChatButtonProps> = (props) => {
   const [title, setTitle] = useState<string>();
   const [flyoutVisible, setFlyoutVisible] = useState(false);
   const [flyoutComponent, setFlyoutComponent] = useState<React.ReactNode | null>(null);
-  const [selectedTabId, setSelectedTabId] = useState<TabId>('chat');
+  const [selectedTabId, setSelectedTabId] = useState<TabId>(TAB_ID.CHAT);
   const [preSelectedTabId, setPreSelectedTabId] = useState<TabId | undefined>(undefined);
   const [traceId, setTraceId] = useState<string | undefined>(undefined);
   const [chatSize, setChatSize] = useState<number | 'fullscreen' | 'dock-right'>('dock-right');

diff --git a/public/components/agent_framework_traces.tsx b/public/components/agent_framework_traces.tsx
@@ -0,0 +1,92 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  EuiAccordion,
+  EuiCodeBlock,
+  EuiEmptyPrompt,
+  EuiLoadingContent,
+  EuiSpacer,
+  EuiText,
+  EuiMarkdownFormat,
+  EuiHorizontalRule,
+} from '@elastic/eui';
+import React from 'react';
+import { useFetchAgentFrameworkTraces } from '../hooks/use_fetch_agentframework_traces';
+
+interface AgentFrameworkTracesProps {
+  traceId: string;
+}
+
+export const AgentFrameworkTraces: React.FC<AgentFrameworkTracesProps> = (props) => {
+  const { data: traces, loading, error } = useFetchAgentFrameworkTraces(props.traceId);
+
+  if (loading) {
+    return (
+      <>
+        <EuiText>Loading...</EuiText>
+        <EuiLoadingContent lines={10} />
+      </>
+    );
+  }
+  if (error) {
+    return (
+      <EuiEmptyPrompt
+        iconType="alert"
+        iconColor="danger"
+        title={<h2>Error loading details</h2>}
+        body={error.toString()}
+      />
+    );
+  }
+  if (!traces?.length) {
+    return <EuiText>Data not available.</EuiText>;
+  }
+
+  const question = traces[traces.length - 1].input;
+  const result = traces[traces.length - 1].output;
+  const questionAndResult = `# How was this generated
+#### Question
+${question}
+#### Result
+${result}
+`;
+
+  return (
+    <>
+      <EuiMarkdownFormat>{questionAndResult}</EuiMarkdownFormat>
+
+      <EuiSpacer size="l" />
+
+      <EuiText>
+        <h3>Response</h3>
+      </EuiText>
+      {traces
+        // if origin exists, it indicates that the trace was generated by a tool, we only show the non-empty traces of tools
+        .filter((trace) => trace.origin && (trace.input || trace.output))
+        .map((trace, i) => {
+          const stepContent = `Step ${i + 1}`;
+          return (
+            <div key={trace.interactionId}>
+              <EuiSpacer size="s" />
+              <EuiAccordion id={stepContent} buttonContent={stepContent}>
+                {trace.input && (
+                  <EuiCodeBlock fontSize="m" paddingSize="s">
+                    Input: {trace.input}
+                  </EuiCodeBlock>
+                )}
+                {trace.output && (
+                  <EuiCodeBlock fontSize="m" paddingSize="s">
+                    Output: {trace.output}
+                  </EuiCodeBlock>
+                )}
+              </EuiAccordion>
+              <EuiHorizontalRule margin="xs" />
+            </div>
+          );
+        })}
+    </>
+  );
+};
diff --git a/public/components/agent_framework_traces_flyout_body.tsx b/public/components/agent_framework_traces_flyout_body.tsx
@@ -0,0 +1,73 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  EuiButtonEmpty,
+  EuiFlyoutBody,
+  EuiPage,
+  EuiPageBody,
+  EuiPageContentBody,
+  EuiPageHeader,
+  EuiButtonIcon,
+  EuiPageHeaderSection,
+} from '@elastic/eui';
+import React from 'react';
+import { useChatContext } from '../contexts/chat_context';
+import { AgentFrameworkTraces } from './agent_framework_traces';
+import { TAB_ID } from '../utils/constants';
+
+export const AgentFrameworkTracesFlyoutBody: React.FC = () => {
+  const chatContext = useChatContext();
+  const traceId = chatContext.traceId;
+  if (!traceId) {
+    return null;
+  }
+
+  // docked right or fullscreen with history open
+  const showBack = !chatContext.flyoutFullScreen || chatContext.preSelectedTabId === TAB_ID.HISTORY;
+
+  return (
+    <EuiFlyoutBody className="llm-chat-flyout llm-chat-flyout-body">
+      <EuiPage>
+        <EuiPageBody>
+          <EuiPageHeader>
+            <EuiPageHeaderSection>
+              {showBack && (
+                <EuiButtonEmpty
+                  size="xs"
+                  flush="left"
+                  onClick={() => {
+                    chatContext.setSelectedTabId(
+                      chatContext.flyoutFullScreen ? TAB_ID.HISTORY : TAB_ID.CHAT
+                    );
+                  }}
+                  iconType="arrowLeft"
+                >
+                  Back
+                </EuiButtonEmpty>
+              )}
+            </EuiPageHeaderSection>
+            <EuiPageHeaderSection>
+              {!showBack && (
+                <EuiButtonIcon
+                  aria-label="close"
+                  size="xs"
+                  color="text"
+                  iconType="cross"
+                  onClick={() => {
+                    chatContext.setSelectedTabId(TAB_ID.CHAT);
+                  }}
+                />
+              )}
+            </EuiPageHeaderSection>
+          </EuiPageHeader>
+          <EuiPageContentBody>
+            <AgentFrameworkTraces traceId={traceId} />
+          </EuiPageContentBody>
+        </EuiPageBody>
+      </EuiPage>
+    </EuiFlyoutBody>
+  );
+};
diff --git a/public/hooks/use_chat_actions.tsx b/public/hooks/use_chat_actions.tsx
@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { TAB_ID } from '../utils/constants';
 import { ASSISTANT_API } from '../../common/constants/llm';
 import {
   IMessage,
@@ -52,7 +53,7 @@ export const useChatActions = (): AssistantActions => {
         !chatContext.sessionId &&
         response.sessionId &&
         core.services.sessions.options?.page === 1 &&
-        chatContext.selectedTabId === 'history'
+        chatContext.selectedTabId === TAB_ID.HISTORY
       ) {
         core.services.sessions.reload();
       }
@@ -81,7 +82,7 @@ export const useChatActions = (): AssistantActions => {
     chatContext.setTitle(title);
     // Chat page will always visible in fullscreen mode, we don't need to change the tab anymore
     if (!chatContext.flyoutFullScreen) {
-      chatContext.setSelectedTabId('chat');
+      chatContext.setSelectedTabId(TAB_ID.CHAT);
     }
     chatContext.setFlyoutComponent(null);
     if (!sessionId) {
@@ -102,7 +103,7 @@ export const useChatActions = (): AssistantActions => {
 
   const openChatUI = () => {
     chatContext.setFlyoutVisible(true);
-    chatContext.setSelectedTabId('chat');
+    chatContext.setSelectedTabId(TAB_ID.CHAT);
   };
 
   const executeAction = async (suggestedAction: ISuggestedAction, message: IMessage) => {

diff --git a/public/hooks/use_fetch_agentframework_traces.ts b/public/hooks/use_fetch_agentframework_traces.ts
@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { useEffect, useReducer } from 'react';
+import { ASSISTANT_API } from '../../common/constants/llm';
+import { AgentFrameworkTrace } from '../../common/utils/llm_chat/traces';
+import { useCore } from '../contexts/core_context';
+import { GenericReducer, genericReducer } from './fetch_reducer';
+
+export const useFetchAgentFrameworkTraces = (traceId: string) => {
+  const core = useCore();
+  const reducer: GenericReducer<AgentFrameworkTrace[]> = genericReducer;
+  const [state, dispatch] = useReducer(reducer, { loading: false });
+
+  useEffect(() => {
+    const abortController = new AbortController();
+    dispatch({ type: 'request' });
+    if (!traceId) {
+      dispatch({ type: 'success', payload: undefined });
+      return;
+    }
+
+    core.services.http
+      .get<AgentFrameworkTrace[]>(`${ASSISTANT_API.TRACE}/${traceId}`)
+      .then((payload) =>
+        dispatch({
+          type: 'success',
+          payload,
+        })
+      )
+      .catch((error) => dispatch({ type: 'failure', error }));
+
+    return () => abortController.abort();
+  }, [traceId]);
+
+  return { ...state };
+};
diff --git a/public/tabs/chat/messages/message_bubble.tsx b/public/tabs/chat/messages/message_bubble.tsx
@@ -188,6 +188,7 @@ export const MessageBubble: React.FC<MessageBubbleProps> = React.memo((props) =>
                     {feedbackResult !== false ? (
                       <EuiFlexItem grow={false}>
                         <EuiButtonIcon
+                          aria-label="feedback"
                           color={feedbackResult === true ? 'primary' : 'text'}
                           iconType="thumbsUp"
                           onClick={() => feedbackOutput(true, feedbackResult)}
@@ -197,6 +198,7 @@ export const MessageBubble: React.FC<MessageBubbleProps> = React.memo((props) =>
                     {feedbackResult !== true ? (
                       <EuiFlexItem grow={false}>
                         <EuiButtonIcon
+                          aria-label="feedback result"
                           color={feedbackResult === false ? 'primary' : 'text'}
                           iconType="thumbsDown"
                           onClick={() => feedbackOutput(false, feedbackResult)}

diff --git a/public/tabs/chat/messages/message_footer.tsx b/public/tabs/chat/messages/message_footer.tsx
@@ -7,9 +7,9 @@ import { EuiButtonEmpty, EuiFlexGroup, EuiFlexItem, EuiHorizontalRule } from '@e
 import React from 'react';
 import { IMessage } from '../../../../common/types/chat_saved_object_attributes';
 import { FeedbackModal } from '../../../components/feedback_modal';
-import { LangchainTracesFlyoutBody } from '../../../components/langchain_traces_flyout_body';
 import { useChatContext } from '../../../contexts/chat_context';
 import { useCore } from '../../../contexts/core_context';
+import { AgentFrameworkTracesFlyoutBody } from '../../../components/agent_framework_traces_flyout_body';
 
 interface MessageFooterProps {
   message: IMessage;
@@ -31,12 +31,7 @@ export const MessageFooter: React.FC<MessageFooterProps> = React.memo((props) =>
           size="xs"
           flush="left"
           onClick={() => {
-            chatContext.setFlyoutComponent(
-              <LangchainTracesFlyoutBody
-                closeFlyout={() => chatContext.setFlyoutComponent(null)}
-                traceId={traceId}
-              />
-            );
+            chatContext.setFlyoutComponent(<AgentFrameworkTracesFlyoutBody />);
           }}
         >
           How was this generated?

diff --git a/public/tabs/chat_window_header.tsx b/public/tabs/chat_window_header.tsx
@@ -22,6 +22,7 @@ import { useCore } from '../contexts/core_context';
 import { useChatState } from '../hooks/use_chat_state';
 import { useSaveChat } from '../hooks/use_save_chat';
 import chatIcon from '../assets/chat.svg';
+import { TAB_ID } from '../utils/constants';
 interface ChatWindowHeaderProps {
   flyoutFullScreen: boolean;
   toggleFlyoutFullScreen: () => void;
@@ -157,10 +158,10 @@ export const ChatWindowHeader: React.FC<ChatWindowHeaderProps> = React.memo((pro
                   chatContext.setFlyoutComponent(undefined);
                   // Back to chat tab if history page already visible
                   chatContext.setSelectedTabId(
-                    chatContext.selectedTabId === 'history' ? 'chat' : 'history'
+                    chatContext.selectedTabId === TAB_ID.HISTORY ? TAB_ID.CHAT : TAB_ID.HISTORY
                   );
                 }}
-                display={chatContext.selectedTabId === 'history' ? 'fill' : undefined}
+                display={chatContext.selectedTabId === TAB_ID.HISTORY ? 'fill' : undefined}
               />
             </EuiFlexItem>
           </EuiFlexGroup>