feat: integrate regenerate API

public/hooks/use_chat_actions.tsx

@@ -160,15 +160,19 @@ export const useChatActions = (): AssistantActions => {
     }
   };
 
-  const regenerate = async () => {
+  const regenerate = async (interactionId: string) => {
     if (chatContext.sessionId) {
       const abortController = new AbortController();
       abortControllerRef = abortController;
       chatStateDispatch({ type: 'regenerate' });
 
       try {
         const response = await core.services.http.put(`${ASSISTANT_API.REGENERATE}`, {
-          body: JSON.stringify({ sessionId: chatContext.sessionId }),
+          body: JSON.stringify({
+            sessionId: chatContext.sessionId,
+            rootAgentId: chatContext.rootAgentId,
+            interactionId,
+          }),
         });
 
         if (abortController.signal.aborted) {

public/tabs/chat/messages/message_bubble.test.tsx

@@ -130,6 +130,13 @@ describe('<MessageBubble />', () => {
           contentType: 'markdown',
           content: 'here are the indices in your cluster: .alert',
         }}
+        interaction={{
+          input: 'foo',
+          response: 'bar',
+          conversation_id: 'foo',
+          interaction_id: 'bar',
+          create_time: new Date().toLocaleString(),
+        }}
       />
     );
     expect(screen.queryAllByTitle('regenerate message')).toHaveLength(1);

public/tabs/chat/messages/message_bubble.tsx

@@ -30,7 +30,7 @@ type MessageBubbleProps = {
   showActionBar: boolean;
   showRegenerate?: boolean;
   shouldActionBarVisibleOnHover?: boolean;
-  onRegenerate?: () => void;
+  onRegenerate?: (interactionId: string) => void;
 } & (
   | {
       message: IMessage;
@@ -192,17 +192,17 @@ export const MessageBubble: React.FC<MessageBubbleProps> = React.memo((props) =>
                     </EuiCopy>
                   </EuiFlexItem>
                 )}
-                {props.showRegenerate && (
+                {props.showRegenerate && props.interaction?.interaction_id ? (
                   <EuiFlexItem grow={false}>
                     <EuiButtonIcon
                       aria-label="regenerate message"
-                      onClick={props.onRegenerate}
+                      onClick={() => props.onRegenerate?.(props.interaction?.interaction_id || '')}
                       title="regenerate message"
                       color="text"
                       iconType="refresh"
                     />
                   </EuiFlexItem>
-                )}
+                ) : null}
                 {showFeedback && (
                   // After feedback, only corresponding thumb icon will be kept and disabled.
                   <>

public/types.ts

@@ -16,7 +16,7 @@ export interface AssistantActions {
   openChatUI: (sessionId?: string) => void;
   executeAction: (suggestedAction: ISuggestedAction, message: IMessage) => void;
   abortAction: (sessionId?: string) => void;
-  regenerate: () => void;
+  regenerate: (interactionId: string) => void;
 }
 
 export interface AppPluginStartDependencies {

server/routes/chat_routes.ts

@@ -13,7 +13,6 @@ import {
 } from '../../../../src/core/server';
 import { ASSISTANT_API } from '../../common/constants/llm';
 import { OllyChatService } from '../services/chat/olly_chat_service';
-import { IMessage, IInput } from '../../common/types/chat_saved_object_attributes';
 import { AgentFrameworkStorageService } from '../services/storage/agent_framework_storage_service';
 import { RoutesOptions } from '../types';
 import { ChatService } from '../services/chat/chat_service';
@@ -64,6 +63,7 @@ const regenerateRoute = {
     body: schema.object({
       sessionId: schema.string(),
       rootAgentId: schema.string(),
+      interactionId: schema.string(),
     }),
   },
 };
@@ -314,42 +314,39 @@ export function registerChatRoutes(router: IRouter, routeOptions: RoutesOptions)
       request,
       response
     ): Promise<IOpenSearchDashboardsResponse<HttpResponsePayload | ResponseError>> => {
-      const { sessionId, rootAgentId } = request.body;
+      const { sessionId: sessionIdInRequestBody, rootAgentId, interactionId } = request.body;
       const storageService = createStorageService(context);
-      let messages: IMessage[] = [];
       const chatService = createChatService();
 
+      let outputs: Awaited<ReturnType<ChatService['regenerate']>> | undefined;
+
+      /**
+       * Get final answer from Agent framework
+       */
       try {
-        const session = await storageService.getSession(sessionId);
-        messages.push(...session.messages);
+        outputs = await chatService.regenerate(
+          { sessionId: sessionIdInRequestBody, rootAgentId, interactionId },
+          context
+        );
       } catch (error) {
-        return response.custom({ statusCode: error.statusCode || 500, body: error.message });
+        context.assistant_plugin.logger.error(error);
       }
 
-      const lastInputIndex = messages.findLastIndex((msg) => msg.type === 'input');
-      // Find last input message
-      const input = messages[lastInputIndex] as IInput;
-      // Take the messages before last input message as memory as regenerate will exclude the last outputs
-      messages = messages.slice(0, lastInputIndex);
-
+      /**
+       * Retrieve latest interactions from memory
+       */
+      const sessionId = sessionIdInRequestBody;
       try {
-        const outputs = await chatService.requestLLM(
-          { messages, input, sessionId, rootAgentId },
-          context
-        );
-        const title = input.content.substring(0, 50);
-        const saveMessagesResponse = await storageService.saveMessages(
-          title,
-          sessionId,
-          [...messages, input, ...outputs.messages].filter(
-            (message) => message.content !== 'AbortError'
-          )
-        );
+        const conversation = await storageService.getSession(sessionId);
+
         return response.ok({
-          body: { ...saveMessagesResponse, title },
+          body: {
+            ...conversation,
+            sessionId,
+          },
         });
       } catch (error) {
-        context.assistant_plugin.logger.warn(error);
+        context.assistant_plugin.logger.error(error);
         return response.custom({ statusCode: error.statusCode || 500, body: error.message });
       }
     }

server/services/chat/chat_service.ts

@@ -10,8 +10,15 @@ import { LLMRequestSchema } from '../../routes/chat_routes';
 export interface ChatService {
   requestLLM(
     payload: { messages: IMessage[]; input: IInput; sessionId?: string },
-    context: RequestHandlerContext,
-    request: OpenSearchDashboardsRequest<unknown, unknown, LLMRequestSchema, 'post'>
+    context: RequestHandlerContext
+  ): Promise<{
+    messages: IMessage[];
+    memoryId: string;
+  }>;
+
+  regenerate(
+    payload: { sessionId: string; interactionId: string; rootAgentId: string },
+    context: RequestHandlerContext
   ): Promise<{
     messages: IMessage[];
     memoryId: string;

server/services/chat/olly_chat_service.test.ts

@@ -0,0 +1,176 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { OllyChatService } from './olly_chat_service';
+import { CoreRouteHandlerContext } from '../../../../../src/core/server/core_route_handler_context';
+import { coreMock, httpServerMock } from '../../../../../src/core/server/mocks';
+import { loggerMock } from '../../../../../src/core/server/logging/logger.mock';
+
+describe('OllyChatService', () => {
+  const ollyChatService = new OllyChatService();
+  const coreContext = new CoreRouteHandlerContext(
+    coreMock.createInternalStart(),
+    httpServerMock.createOpenSearchDashboardsRequest()
+  );
+  const mockedTransport = coreContext.opensearch.client.asCurrentUser.transport
+    .request as jest.Mock;
+  const contextMock = {
+    core: coreContext,
+    assistant_plugin: {
+      logger: loggerMock.create(),
+    },
+  };
+  beforeEach(() => {
+    mockedTransport.mockClear();
+  });
+  it('requestLLM should invoke client call with correct params', async () => {
+    mockedTransport.mockImplementationOnce(() => {
+      return {
+        body: {
+          inference_results: [
+            {
+              output: [
+                {
+                  name: 'memory_id',
+                  result: 'foo',
+                },
+              ],
+            },
+          ],
+        },
+      };
+    });
+    const result = await ollyChatService.requestLLM(
+      {
+        messages: [],
+        input: {
+          type: 'input',
+          contentType: 'text',
+          content: 'content',
+        },
+        sessionId: '',
+        rootAgentId: 'rootAgentId',
+      },
+      contextMock
+    );
+    expect(mockedTransport.mock.calls).toMatchInlineSnapshot(`
+      Array [
+        Array [
+          Object {
+            "body": Object {
+              "parameters": Object {
+                "question": "content",
+                "verbose": true,
+              },
+            },
+            "method": "POST",
+            "path": "/_plugins/_ml/agents/rootAgentId/_execute",
+          },
+          Object {
+            "maxRetries": 0,
+            "requestTimeout": 300000,
+          },
+        ],
+      ]
+    `);
+    expect(result).toMatchInlineSnapshot(`
+      Object {
+        "memoryId": "foo",
+        "messages": Array [],
+      }
+    `);
+  });
+
+  it('requestLLM should throw error when transport.request throws error', async () => {
+    mockedTransport.mockImplementationOnce(() => {
+      throw new Error('error');
+    });
+    expect(
+      ollyChatService.requestLLM(
+        {
+          messages: [],
+          input: {
+            type: 'input',
+            contentType: 'text',
+            content: 'content',
+          },
+          sessionId: '',
+          rootAgentId: 'rootAgentId',
+        },
+        contextMock
+      )
+    ).rejects.toMatchInlineSnapshot(`[Error: error]`);
+  });
+
+  it('regenerate should invoke client call with correct params', async () => {
+    mockedTransport.mockImplementationOnce(() => {
+      return {
+        body: {
+          inference_results: [
+            {
+              output: [
+                {
+                  name: 'memory_id',
+                  result: 'foo',
+                },
+              ],
+            },
+          ],
+        },
+      };
+    });
+    const result = await ollyChatService.regenerate(
+      {
+        sessionId: 'sessionId',
+        rootAgentId: 'rootAgentId',
+        interactionId: 'interactionId',
+      },
+      contextMock
+    );
+    expect(mockedTransport.mock.calls).toMatchInlineSnapshot(`
+      Array [
+        Array [
+          Object {
+            "body": Object {
+              "parameters": Object {
+                "memory_id": "sessionId",
+                "regenerate_interaction_id": "interactionId",
+                "verbose": true,
+              },
+            },
+            "method": "POST",
+            "path": "/_plugins/_ml/agents/rootAgentId/_execute",
+          },
+          Object {
+            "maxRetries": 0,
+            "requestTimeout": 300000,
+          },
+        ],
+      ]
+    `);
+    expect(result).toMatchInlineSnapshot(`
+      Object {
+        "memoryId": "foo",
+        "messages": Array [],
+      }
+    `);
+  });
+
+  it('regenerate should throw error when transport.request throws error', async () => {
+    mockedTransport.mockImplementationOnce(() => {
+      throw new Error('error');
+    });
+    expect(
+      ollyChatService.regenerate(
+        {
+          sessionId: 'sessionId',
+          rootAgentId: 'rootAgentId',
+          interactionId: 'interactionId',
+        },
+        contextMock
+      )
+    ).rejects.toMatchInlineSnapshot(`[Error: error]`);
+  });
+});