feat: use the agent id from request body

Signed-off-by: SuZhou-Joe <[email protected]>
opensearch-project · Nov 20, 2023 · 7a335b6 · 7a335b6
1 parent 91dce23
commit 7a335b6
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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,
         },