From ca6ab4421a649d06d665f5c2c3875f8df4fc18e7 Mon Sep 17 00:00:00 2001 From: gaobinlong Date: Thu, 30 Nov 2023 17:53:44 +0800 Subject: [PATCH] Integrate list traces api of agent framework (#25) * Integrate list traces api of agent framework Signed-off-by: gaobinlong * Fix import issue Signed-off-by: gaobinlong * Fix some issues Signed-off-by: gaobinlong * Add some comments Signed-off-by: gaobinlong --------- Signed-off-by: gaobinlong --- common/constants/llm.ts | 1 + common/utils/llm_chat/traces.ts | 10 ++ public/chat_flyout.tsx | 11 +- public/chat_header_button.tsx | 3 +- public/components/agent_framework_traces.tsx | 92 ++++++++++++++ .../agent_framework_traces_flyout_body.tsx | 73 +++++++++++ public/hooks/use_chat_actions.tsx | 7 +- .../hooks/use_fetch_agentframework_traces.ts | 39 ++++++ public/tabs/chat/messages/message_bubble.tsx | 2 + public/tabs/chat/messages/message_footer.tsx | 9 +- public/tabs/chat_window_header.tsx | 113 +++++++++--------- public/tabs/history/chat_history_page.tsx | 9 +- public/utils/constants.ts | 12 ++ server/routes/chat_routes.ts | 28 +++++ server/services/chat/olly_chat_service.ts | 1 - .../agent_framework_storage_service.ts | 33 +++++ 16 files changed, 369 insertions(+), 74 deletions(-) create mode 100644 public/components/agent_framework_traces.tsx create mode 100644 public/components/agent_framework_traces_flyout_body.tsx create mode 100644 public/hooks/use_fetch_agentframework_traces.ts create mode 100644 public/utils/constants.ts diff --git a/common/constants/llm.ts b/common/constants/llm.ts index 53029273..8a68dba2 100644 --- 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 index 8ac0d314..73ee3088 100644 --- 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 index 04e9c34b..42d6e32c 100644 --- 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 = (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 = (props) => { className={cs({ 'llm-chat-hidden': !chatHistoryPageVisible })} /> )} - {chatTraceVisible && chatContext.traceId && } + {chatTraceVisible && chatContext.traceId && } diff --git a/public/chat_header_button.tsx b/public/chat_header_button.tsx index b2a9739b..9b3d1503 100644 --- a/public/chat_header_button.tsx +++ b/public/chat_header_button.tsx @@ -15,6 +15,7 @@ import { ChatStateProvider } from './hooks'; 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 = (props) => { const [title, setTitle] = useState(); const [flyoutVisible, setFlyoutVisible] = useState(false); const [flyoutComponent, setFlyoutComponent] = useState(null); - const [selectedTabId, setSelectedTabId] = useState('chat'); + const [selectedTabId, setSelectedTabId] = useState(TAB_ID.CHAT); const [preSelectedTabId, setPreSelectedTabId] = useState(undefined); const [traceId, setTraceId] = useState(undefined); const [chatSize, setChatSize] = useState('dock-right'); diff --git a/public/components/agent_framework_traces.tsx b/public/components/agent_framework_traces.tsx new file mode 100644 index 00000000..903b494a --- /dev/null +++ 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 = (props) => { + const { data: traces, loading, error } = useFetchAgentFrameworkTraces(props.traceId); + + if (loading) { + return ( + <> + Loading... + + + ); + } + if (error) { + return ( + Error loading details} + body={error.toString()} + /> + ); + } + if (!traces?.length) { + return Data not available.; + } + + 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 ( + <> + {questionAndResult} + + + + +

Response

+ + {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 ( +
+ + + {trace.input && ( + + Input: {trace.input} + + )} + {trace.output && ( + + Output: {trace.output} + + )} + + +
+ ); + })} + + ); +}; diff --git a/public/components/agent_framework_traces_flyout_body.tsx b/public/components/agent_framework_traces_flyout_body.tsx new file mode 100644 index 00000000..3aefdbfd --- /dev/null +++ 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 ( + + + + + + {showBack && ( + { + chatContext.setSelectedTabId( + chatContext.flyoutFullScreen ? TAB_ID.HISTORY : TAB_ID.CHAT + ); + }} + iconType="arrowLeft" + > + Back + + )} + + + {!showBack && ( + { + chatContext.setSelectedTabId(TAB_ID.CHAT); + }} + /> + )} + + + + + + + + + ); +}; diff --git a/public/hooks/use_chat_actions.tsx b/public/hooks/use_chat_actions.tsx index 2dd26964..e1975b46 100644 --- 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 new file mode 100644 index 00000000..443e0960 --- /dev/null +++ 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 = 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(`${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 index 91c18909..18853603 100644 --- a/public/tabs/chat/messages/message_bubble.tsx +++ b/public/tabs/chat/messages/message_bubble.tsx @@ -203,6 +203,7 @@ export const MessageBubble: React.FC = React.memo((props) => {feedbackResult !== false ? ( feedbackOutput(true, feedbackResult)} @@ -212,6 +213,7 @@ export const MessageBubble: React.FC = React.memo((props) => {feedbackResult !== true ? ( feedbackOutput(false, feedbackResult)} diff --git a/public/tabs/chat/messages/message_footer.tsx b/public/tabs/chat/messages/message_footer.tsx index aa00fa89..8fcdd24b 100644 --- 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 = React.memo((props) => size="xs" flush="left" onClick={() => { - chatContext.setFlyoutComponent( - chatContext.setFlyoutComponent(null)} - traceId={traceId} - /> - ); + chatContext.setFlyoutComponent(); }} > How was this generated? diff --git a/public/tabs/chat_window_header.tsx b/public/tabs/chat_window_header.tsx index 033acdaa..c0b1d5e2 100644 --- a/public/tabs/chat_window_header.tsx +++ b/public/tabs/chat_window_header.tsx @@ -8,6 +8,7 @@ import React from 'react'; import { useChatContext } from '../contexts/chat_context'; import { ChatWindowHeaderTitle } from '../components/chat_window_header_title'; import chatIcon from '../assets/chat.svg'; +import { TAB_ID } from '../utils/constants'; interface ChatWindowHeaderProps { flyoutFullScreen: boolean; toggleFlyoutFullScreen: () => void; @@ -35,60 +36,62 @@ export const ChatWindowHeader: React.FC = React.memo((pro ); return ( - - - - - - - - - - - { - chatContext.setFlyoutComponent(undefined); - // Back to chat tab if history page already visible - chatContext.setSelectedTabId( - chatContext.selectedTabId === 'history' ? 'chat' : 'history' - ); - }} - display={chatContext.selectedTabId === 'history' ? 'fill' : undefined} - /> - - - - - - - - { - chatContext.setFlyoutVisible(false); - }} - /> - - - + <> + + + + + + + + + + + { + chatContext.setFlyoutComponent(undefined); + // Back to chat tab if history page already visible + chatContext.setSelectedTabId( + chatContext.selectedTabId === TAB_ID.HISTORY ? TAB_ID.CHAT : TAB_ID.HISTORY + ); + }} + display={chatContext.selectedTabId === TAB_ID.HISTORY ? 'fill' : undefined} + /> + + + + + + + + { + chatContext.setFlyoutVisible(false); + }} + /> + + + + ); }); diff --git a/public/tabs/history/chat_history_page.tsx b/public/tabs/history/chat_history_page.tsx index 05fbf3b5..99adc4aa 100644 --- a/public/tabs/history/chat_history_page.tsx +++ b/public/tabs/history/chat_history_page.tsx @@ -22,6 +22,7 @@ import { useDebounce, useObservable } from 'react-use'; import cs from 'classnames'; import { useChatActions, useChatState } from '../../hooks'; import { useChatContext, useCore } from '../../contexts'; +import { TAB_ID } from '../../utils/constants'; import { ChatHistorySearchList } from './chat_history_search_list'; interface ChatHistoryPageProps { @@ -70,7 +71,7 @@ export const ChatHistoryPage: React.FC = React.memo((props }, []); const handleBack = useCallback(() => { - setSelectedTabId('chat'); + setSelectedTabId(TAB_ID.CHAT); }, [setSelectedTabId]); const handleHistoryDeleted = useCallback( @@ -113,7 +114,11 @@ export const ChatHistoryPage: React.FC = React.memo((props {flyoutFullScreen ? ( - + ) : ( diff --git a/public/utils/constants.ts b/public/utils/constants.ts new file mode 100644 index 00000000..9c35746e --- /dev/null +++ b/public/utils/constants.ts @@ -0,0 +1,12 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +export enum TAB_ID { + CHAT = 'chat', + COMPOSE = 'compose', + INSIGHTS = 'insights', + HISTORY = 'history', + TRACE = 'trace', +} diff --git a/server/routes/chat_routes.ts b/server/routes/chat_routes.ts index 632babe4..c234a42d 100644 --- 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> => { + 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 index 8dadc991..7e3723d1 100644 --- 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 index d80f5b7e..06e7777a 100644 --- 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,36 @@ export class AgentFrameworkStorageService implements StorageService { throw new Error('update converstaion failed, reason:' + JSON.stringify(error.meta?.body)); } } + + async getTraces(interactionId: string): Promise { + 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, + })); + } catch (error) { + throw new Error('get traces failed, reason:' + JSON.stringify(error.meta?.body)); + } + } }