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,7 @@ 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';
 
 let chatHistoryPageLoaded = false;
 
@@ -134,7 +134,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/components/agent_framework_traces.tsx b/public/components/agent_framework_traces.tsx
@@ -0,0 +1,91 @@
+/*
+ * 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[0].input;
+  const result = traces[0].output;
+  const questionAndResult = `# How was this generated
+#### Question
+${question}
+#### Result
+${result}
+`;
+
+  return (
+    <>
+      <EuiMarkdownFormat>{questionAndResult}</EuiMarkdownFormat>
+
+      <EuiSpacer size="l" />
+
+      <EuiText>
+        <h3>Response</h3>
+      </EuiText>
+      {traces
+        .filter((trace) => trace.origin?.includes('Tool') && (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,70 @@
+/*
+ * 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';
+
+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 === '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 ? 'history' : 'chat');
+                  }}
+                  iconType="arrowLeft"
+                >
+                  Back
+                </EuiButtonEmpty>
+              )}
+            </EuiPageHeaderSection>
+            <EuiPageHeaderSection>
+              {!showBack && (
+                <EuiButtonIcon
+                  aria-label="close"
+                  size="xs"
+                  color="text"
+                  iconType="cross"
+                  onClick={() => {
+                    chatContext.setSelectedTabId('chat');
+                  }}
+                />
+              )}
+            </EuiPageHeaderSection>
+          </EuiPageHeader>
+          <EuiPageContentBody>
+            <AgentFrameworkTraces traceId={traceId} />
+          </EuiPageContentBody>
+        </EuiPageBody>
+      </EuiPage>
+    </EuiFlyoutBody>
+  );
+};
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;
@@ -32,7 +32,7 @@ export const MessageFooter: React.FC<MessageFooterProps> = React.memo((props) =>
           flush="left"
           onClick={() => {
             chatContext.setFlyoutComponent(
-              <LangchainTracesFlyoutBody
+              <AgentFrameworkTracesFlyoutBody
                 closeFlyout={() => chatContext.setFlyoutComponent(null)}
                 traceId={traceId}
               />

diff --git a/public/tabs/history/chat_history_page.tsx b/public/tabs/history/chat_history_page.tsx
@@ -105,7 +105,11 @@ export const ChatHistoryPage: React.FC<ChatHistoryPageProps> = React.memo((props
             {flyoutFullScreen ? (
               <EuiFlexGroup gutterSize="none" justifyContent="flexEnd">
                 <EuiFlexItem grow={false}>
-                  <EuiButtonIcon iconType="cross" onClick={handleBack} />
+                  <EuiButtonIcon
+                    aria-label="full screen back"
+                    iconType="cross"
+                    onClick={handleBack}
+                  />
                 </EuiFlexItem>
               </EuiFlexGroup>
             ) : (

diff --git a/server/routes/chat_routes.ts b/server/routes/chat_routes.ts
@@ -105,6 +105,15 @@ const updateSessionRoute = {
   },
 };
 
+const getTracesRoute = {
+  path: `${ASSISTANT_API.TRACE}/{traceId}`,
+  validate: {
+    params: schema.object({
+      traceId: schema.string(),
+    }),
+  },
+};
+
 export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions) {
   const createStorageService = (context: RequestHandlerContext) =>
     new AgentFrameworkStorageService(
@@ -226,6 +235,25 @@ export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions)
     }
   );
 
+  router.get(
+    getTracesRoute,
+    async (
+      context,
+      request,
+      response
+    ): Promise<IOpenSearchDashboardsResponse<HttpResponsePayload | ResponseError>> => {
+      const storageService = createStorageService(context);
+
+      try {
+        const getResponse = await storageService.getTraces(request.params.traceId);
+        return response.ok({ body: getResponse });
+      } catch (error) {
+        context.assistant_plugin.logger.error(error);
+        return response.custom({ statusCode: error.statusCode || 500, body: error.message });
+      }
+    }
+  );
+
   router.post(
     abortAgentExecutionRoute,
     async (

diff --git a/server/services/chat/olly_chat_service.ts b/server/services/chat/olly_chat_service.ts
@@ -12,7 +12,6 @@ import { LLMModelFactory } from '../../olly/models/llm_model_factory';
 import { PPLTools } from '../../olly/tools/tool_sets/ppl';
 import { PPLGenerationRequestSchema } from '../../routes/langchain_routes';
 import { ChatService } from './chat_service';
-import { LLMRequestSchema } from '../../routes/chat_routes';
 
 const MEMORY_ID_FIELD = 'memory_id';
 

diff --git a/server/services/storage/agent_framework_storage_service.ts b/server/services/storage/agent_framework_storage_service.ts
@@ -4,6 +4,7 @@
  */
 
 import { ApiResponse } from '@opensearch-project/opensearch/.';
+import { AgentFrameworkTrace } from 'common/utils/llm_chat/traces';
 import { OpenSearchClient } from '../../../../../src/core/server';
 import {
   IMessage,
@@ -181,4 +182,38 @@ export class AgentFrameworkStorageService implements StorageService {
       throw new Error('update converstaion failed, reason:' + JSON.stringify(error.meta?.body));
     }
   }
+
+  async getTraces(interactionId: string): Promise<AgentFrameworkTrace[]> {
+    try {
+      const response = (await this.client.transport.request({
+        method: 'GET',
+        path: `/_plugins/_ml/memory/trace/${interactionId}/_list`,
+      })) as ApiResponse<{
+        traces: Array<{
+          conversation_id: string;
+          interaction_id: string;
+          create_time: string;
+          input: string;
+          response: string;
+          origin: string;
+          parent_interaction_id: string;
+          trace_number: number;
+        }>;
+      }>;
+
+      return response.body.traces
+        .map((item) => ({
+          interactionId: item.interaction_id,
+          parentInteractionId: item.parent_interaction_id,
+          input: item.input,
+          output: item.response,
+          createTime: item.create_time,
+          origin: item.origin,
+          traceNumber: item.trace_number,
+        }))
+        .reverse();
+    } catch (error) {
+      throw new Error('get traces failed, reason:' + JSON.stringify(error.meta?.body));
+    }
+  }
 }