From 7a335b6ee8eac6979330c75e106cacf0ca6de8a0 Mon Sep 17 00:00:00 2001
From: SuZhou-Joe <suzhou@amazon.com>
Date: Mon, 20 Nov 2023 11:15:15 +0800
Subject: [PATCH] feat: use the agent id from request body

Signed-off-by: SuZhou-Joe <suzhou@amazon.com>
---
 public/chat_header_button.tsx             | 4 ++++
 public/contexts/chat_context.tsx          | 1 +
 public/hooks/use_chat_actions.tsx         | 1 +
 server/routes/chat_routes.ts              | 1 +
 server/services/chat/olly_chat_service.ts | 4 ++--
 5 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/public/chat_header_button.tsx b/public/chat_header_button.tsx
index ce21f666..26e76373 100644
--- a/public/chat_header_button.tsx
+++ b/public/chat_header_button.tsx
@@ -38,6 +38,9 @@ export const HeaderChatButton: React.FC<HeaderChatButtonProps> = (props) => {
   const [traceId, setTraceId] = useState<string | undefined>(undefined);
   const [chatSize, setChatSize] = useState<number | 'fullscreen' | 'dock-right'>('dock-right');
   const flyoutFullScreen = chatSize === 'fullscreen';
+  const [rootAgentId, setRootAgentId] = useState<string>(
+    new URL(window.location.href).searchParams.get('agent_id') || ''
+  );
 
   if (!flyoutLoaded && flyoutVisible) flyoutLoaded = true;
 
@@ -73,6 +76,7 @@ export const HeaderChatButton: React.FC<HeaderChatButtonProps> = (props) => {
       setTitle,
       traceId,
       setTraceId,
+      rootAgentId,
     }),
     [
       appId,
diff --git a/public/contexts/chat_context.tsx b/public/contexts/chat_context.tsx
index c0807be3..24d3a0f7 100644
--- a/public/contexts/chat_context.tsx
+++ b/public/contexts/chat_context.tsx
@@ -25,6 +25,7 @@ export interface IChatContext {
   setTitle: React.Dispatch<React.SetStateAction<string | undefined>>;
   traceId?: string;
   setTraceId: React.Dispatch<React.SetStateAction<string | undefined>>;
+  rootAgentId?: string;
 }
 export const ChatContext = React.createContext<IChatContext | null>(null);
 
diff --git a/public/hooks/use_chat_actions.tsx b/public/hooks/use_chat_actions.tsx
index c71836b3..6781cc0b 100644
--- a/public/hooks/use_chat_actions.tsx
+++ b/public/hooks/use_chat_actions.tsx
@@ -36,6 +36,7 @@ export const useChatActions = (): AssistantActions => {
         // do not send abort signal to http client to allow LLM call run in background
         body: JSON.stringify({
           sessionId: chatContext.sessionId,
+          rootAgentId: chatContext.rootAgentId,
           ...(!chatContext.sessionId && { messages: chatState.messages }), // include all previous messages for new chats
           input,
         }),
diff --git a/server/routes/chat_routes.ts b/server/routes/chat_routes.ts
index e857b645..01329e1f 100644
--- a/server/routes/chat_routes.ts
+++ b/server/routes/chat_routes.ts
@@ -23,6 +23,7 @@ const llmRequestRoute = {
     body: schema.object({
       sessionId: schema.maybe(schema.string()),
       messages: schema.maybe(schema.arrayOf(schema.any())),
+      rootAgentId: schema.string(),
       input: schema.object({
         type: schema.literal('input'),
         context: schema.object({
diff --git a/server/services/chat/olly_chat_service.ts b/server/services/chat/olly_chat_service.ts
index d7dd4abe..68a92ed4 100644
--- a/server/services/chat/olly_chat_service.ts
+++ b/server/services/chat/olly_chat_service.ts
@@ -31,7 +31,7 @@ export class OllyChatService implements ChatService {
     messages: IMessage[];
     memoryId: string;
   }> {
-    const { input, sessionId } = payload;
+    const { input, sessionId, rootAgentId } = request.body;
     const opensearchClient = context.core.opensearch.client.asCurrentUser;
 
     if (payload.sessionId) {
@@ -57,7 +57,7 @@ export class OllyChatService implements ChatService {
       }
       const agentFrameworkResponse = (await opensearchClient.transport.request({
         method: 'POST',
-        path: '/_plugins/_ml/agents/-jld3IsBXlmiPBu-5dDC/_execute',
+        path: `/_plugins/_ml/agents/${rootAgentId}/_execute`,
         body: {
           parameters: parametersPayload,
         },