feat: refactor memory API url path and fields (#106)

* feat: refactor memory API url path and fields Signed-off-by: SuZhou-Joe <[email protected]> * fix: update snapshot Signed-off-by: SuZhou-Joe <[email protected]> * feat: remove useless type and variables Signed-off-by: SuZhou-Joe <[email protected]> --------- Signed-off-by: SuZhou-Joe <[email protected]>
opensearch-project · Jan 23, 2024 · 87e5b43 · 87e5b43
1 parent 4387510
commit 87e5b43
Show file tree

Hide file tree

Showing 9 changed files with 61 additions and 42 deletions.
diff --git a/common/types/chat_saved_object_attributes.ts b/common/types/chat_saved_object_attributes.ts
@@ -13,9 +13,16 @@ export interface Interaction {
   interaction_id: string;
   create_time: string;
   additional_info?: { feedback?: SendFeedbackBody; [key: string]: unknown };
-  parent_interaction_id?: string;
 }
 
+export type InteractionFromAgentFramework = Omit<
+  Interaction,
+  'interaction_id' | 'conversation_id'
+> & {
+  message_id: string;
+  memory_id: string;
+};
+
 export interface IConversation {
   title: string;
   version?: number;

diff --git a/common/utils/llm_chat/traces.ts b/common/utils/llm_chat/traces.ts
@@ -5,7 +5,6 @@
 
 export interface AgentFrameworkTrace {
   interactionId: string;
-  parentInteractionId: string;
   createTime: string;
   input: string;
   output: string;

diff --git a/public/components/agent_framework_traces.test.tsx b/public/components/agent_framework_traces.test.tsx
@@ -14,7 +14,6 @@ describe('<AgentFrameworkTraces/> spec', () => {
     const traces = [
       {
         interactionId: 'test_interactionId',
-        parentInteractionId: 'test_parent_interactionId',
         input: 'input',
         output: 'output',
         createTime: '',
@@ -23,7 +22,6 @@ describe('<AgentFrameworkTraces/> spec', () => {
       },
       {
         interactionId: 'test_interactionId',
-        parentInteractionId: 'test_parent_interactionId',
         input: 'input',
         output: 'output',
         createTime: '',
@@ -32,7 +30,6 @@ describe('<AgentFrameworkTraces/> spec', () => {
       },
       {
         interactionId: 'test_interactionId',
-        parentInteractionId: 'test_parent_interactionId',
         input: 'input',
         output: '',
         createTime: '',
@@ -41,7 +38,6 @@ describe('<AgentFrameworkTraces/> spec', () => {
       },
       {
         interactionId: 'test_interactionId',
-        parentInteractionId: 'test_parent_interactionId',
         input: '',
         output: 'output',
         createTime: '',

diff --git a/public/hooks/use_fetch_agentframework_traces.test.ts b/public/hooks/use_fetch_agentframework_traces.test.ts
@@ -28,7 +28,6 @@ describe('useFetchAgentFrameworkTraces hook', () => {
     const traces = [
       {
         interactionId: 'test_interactionId',
-        parentInteractionId: 'test_parent_interactionId',
         input: 'input',
         output: 'output',
         createTime: '',

diff --git a/server/routes/chat_routes.test.ts b/server/routes/chat_routes.test.ts
@@ -138,7 +138,6 @@ describe('chat routes', () => {
       const getTraceResultMock = [
         {
           interactionId: 'interaction-1',
-          parentInteractionId: '',
           createTime: '',
           input: 'foo',
           output: 'bar',

diff --git a/server/services/storage/agent_framework_storage_service.test.ts b/server/services/storage/agent_framework_storage_service.test.ts
@@ -6,7 +6,6 @@
 import { AgentFrameworkStorageService } from './agent_framework_storage_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('AgentFrameworkStorageService', () => {
   const coreContext = new CoreRouteHandlerContext(
@@ -24,7 +23,7 @@ describe('AgentFrameworkStorageService', () => {
   });
   it('getConversation', async () => {
     mockedTransport.mockImplementation(async (params) => {
-      if (params.path.includes('/_list?max_results=1000')) {
+      if (params.path.includes('/messages?max_results=1000')) {
         return {
           body: {
             interactions: [
@@ -39,7 +38,7 @@ describe('AgentFrameworkStorageService', () => {
 
       return {
         body: {
-          conversation_id: 'conversation_id',
+          memory_id: 'conversation_id',
           create_time: 0,
           updated_time: 0,
           name: 'foo',
@@ -52,7 +51,9 @@ describe('AgentFrameworkStorageService', () => {
         "createdTimeMs": 0,
         "interactions": Array [
           Object {
+            "conversation_id": "",
             "input": "input",
+            "interaction_id": "",
             "response": "response",
           },
         ],
@@ -66,13 +67,13 @@ describe('AgentFrameworkStorageService', () => {
         Array [
           Object {
             "method": "GET",
-            "path": "/_plugins/_ml/memory/conversation/_mock/_list?max_results=1000",
+            "path": "/_plugins/_ml/memory/_mock/messages?max_results=1000",
           },
         ],
         Array [
           Object {
             "method": "GET",
-            "path": "/_plugins/_ml/memory/conversation/_mock",
+            "path": "/_plugins/_ml/memory/_mock",
           },
         ],
       ]
@@ -163,7 +164,7 @@ describe('AgentFrameworkStorageService', () => {
               "size": 10,
             },
             "method": "GET",
-            "path": "/_plugins/_ml/memory/conversation/_search",
+            "path": "/_plugins/_ml/memory/_search",
           },
         ],
         Array [
@@ -181,7 +182,7 @@ describe('AgentFrameworkStorageService', () => {
               "size": 10,
             },
             "method": "GET",
-            "path": "/_plugins/_ml/memory/conversation/_search",
+            "path": "/_plugins/_ml/memory/_search",
           },
         ],
       ]
@@ -257,13 +258,11 @@ describe('AgentFrameworkStorageService', () => {
       body: {
         traces: [
           {
-            conversation_id: 'conversation_id',
-            interaction_id: 'interaction_id',
+            message_id: 'interaction_id',
             create_time: 'create_time',
             input: 'input',
             response: 'response',
             origin: 'origin',
-            parent_interaction_id: 'parent_interaction_id',
             trace_number: 1,
           },
         ],
@@ -277,7 +276,6 @@ describe('AgentFrameworkStorageService', () => {
           "interactionId": "interaction_id",
           "origin": "origin",
           "output": "response",
-          "parentInteractionId": "parent_interaction_id",
           "traceNumber": 1,
         },
       ]
@@ -353,7 +351,9 @@ describe('AgentFrameworkStorageService', () => {
     expect(agentFrameworkService.getInteraction('_id', 'interaction_id')).resolves
       .toMatchInlineSnapshot(`
       Object {
+        "conversation_id": "",
         "input": "input",
+        "interaction_id": "",
         "response": "response",
       }
     `);

diff --git a/server/services/storage/agent_framework_storage_service.ts b/server/services/storage/agent_framework_storage_service.ts
@@ -11,12 +11,14 @@ import {
   IConversation,
   IConversationFindResponse,
   Interaction,
+  InteractionFromAgentFramework,
 } from '../../../common/types/chat_saved_object_attributes';
 import { GetConversationsSchema } from '../../routes/chat_routes';
 import { StorageService } from './storage_service';
 import { MessageParser } from '../../types';
 import { MessageParserRunner } from '../../utils/message_parser_runner';
 import { ML_COMMONS_BASE_API } from '../../utils/constants';
+import { formatInteractionFromBackend } from '../../utils/format';
 
 export interface ConversationOptResponse {
   success: boolean;
@@ -33,15 +35,15 @@ export class AgentFrameworkStorageService implements StorageService {
     const [interactionsResp, conversation] = await Promise.all([
       this.client.transport.request({
         method: 'GET',
-        path: `${ML_COMMONS_BASE_API}/memory/conversation/${conversationId}/_list?max_results=1000`,
+        path: `${ML_COMMONS_BASE_API}/memory/${conversationId}/messages?max_results=1000`,
       }) as TransportRequestPromise<
         ApiResponse<{
-          interactions: Interaction[];
+          interactions: InteractionFromAgentFramework[];
         }>
       >,
       this.client.transport.request({
         method: 'GET',
-        path: `${ML_COMMONS_BASE_API}/memory/conversation/${conversationId}`,
+        path: `${ML_COMMONS_BASE_API}/memory/${conversationId}`,
       }) as TransportRequestPromise<
         ApiResponse<{
           conversation_id: string;
@@ -51,7 +53,9 @@ export class AgentFrameworkStorageService implements StorageService {
         }>
       >,
     ]);
-    const finalInteractions = interactionsResp.body.interactions;
+    const finalInteractions = interactionsResp.body.interactions.map((item) =>
+      formatInteractionFromBackend(item)
+    );
 
     return {
       title: conversation.body.name,
@@ -99,7 +103,7 @@ export class AgentFrameworkStorageService implements StorageService {
 
     const conversations = await this.client.transport.request({
       method: 'GET',
-      path: `${ML_COMMONS_BASE_API}/memory/conversation/_search`,
+      path: `${ML_COMMONS_BASE_API}/memory/_search`,
       body: requestParams,
     });
 
@@ -143,7 +147,7 @@ export class AgentFrameworkStorageService implements StorageService {
     try {
       const response = await this.client.transport.request({
         method: 'DELETE',
-        path: `${ML_COMMONS_BASE_API}/memory/conversation/${conversationId}/_delete`,
+        path: `${ML_COMMONS_BASE_API}/memory/${conversationId}`,
       });
       if (response.statusCode === 200) {
         return {
@@ -168,7 +172,7 @@ export class AgentFrameworkStorageService implements StorageService {
     try {
       const response = await this.client.transport.request({
         method: 'PUT',
-        path: `${ML_COMMONS_BASE_API}/memory/conversation/${conversationId}/_update`,
+        path: `${ML_COMMONS_BASE_API}/memory/${conversationId}`,
         body: {
           name: title,
         },
@@ -193,23 +197,20 @@ export class AgentFrameworkStorageService implements StorageService {
     try {
       const response = (await this.client.transport.request({
         method: 'GET',
-        path: `${ML_COMMONS_BASE_API}/memory/trace/${interactionId}/_list`,
+        path: `${ML_COMMONS_BASE_API}/memory/message/${interactionId}/traces`,
       })) as ApiResponse<{
         traces: Array<{
-          conversation_id: string;
-          interaction_id: string;
+          message_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,
+        interactionId: item.message_id,
         input: item.input,
         output: item.response,
         createTime: item.create_time,
@@ -228,7 +229,7 @@ export class AgentFrameworkStorageService implements StorageService {
     try {
       const response = await this.client.transport.request({
         method: 'PUT',
-        path: `${ML_COMMONS_BASE_API}/memory/interaction/${interactionId}/_update`,
+        path: `${ML_COMMONS_BASE_API}/memory/message/${interactionId}`,
         body: {
           additional_info: additionalInfo,
         },
@@ -275,8 +276,8 @@ export class AgentFrameworkStorageService implements StorageService {
     }
     const interactionsResp = (await this.client.transport.request({
       method: 'GET',
-      path: `${ML_COMMONS_BASE_API}/memory/conversation/${conversationId}/${interactionId}`,
-    })) as ApiResponse<Interaction>;
-    return interactionsResp.body;
+      path: `${ML_COMMONS_BASE_API}/memory/message/${interactionId}`,
+    })) as ApiResponse<InteractionFromAgentFramework>;
+    return formatInteractionFromBackend(interactionsResp.body);
   }
 }
diff --git a/server/utils/format.ts b/server/utils/format.ts
@@ -0,0 +1,24 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  Interaction,
+  InteractionFromAgentFramework,
+} from '../../common/types/chat_saved_object_attributes';
+
+/**
+ * Agent framework changes interaction_id to message_id
+ * In FE we need to change this back.
+ */
+export const formatInteractionFromBackend = (
+  interaction: InteractionFromAgentFramework
+): Interaction => {
+  const { message_id: messageId, memory_id: memoryId, ...others } = interaction || {};
+  return {
+    ...others,
+    interaction_id: messageId || '',
+    conversation_id: memoryId || '',
+  };
+};
diff --git a/server/utils/message_parser_runner.test.ts b/server/utils/message_parser_runner.test.ts
@@ -32,7 +32,6 @@ describe('MessageParserRunner', () => {
           interaction_id: '',
           create_time: '',
           additional_info: {},
-          parent_interaction_id: '',
         },
         {
           interactions: [
@@ -43,7 +42,6 @@ describe('MessageParserRunner', () => {
               interaction_id: '',
               create_time: '',
               additional_info: {},
-              parent_interaction_id: '',
             },
           ],
         }
@@ -122,7 +120,6 @@ describe('MessageParserRunner', () => {
           interaction_id: '',
           create_time: '',
           additional_info: {},
-          parent_interaction_id: '',
         },
         {
           interactions: [
@@ -133,7 +130,6 @@ describe('MessageParserRunner', () => {
               interaction_id: '',
               create_time: '',
               additional_info: {},
-              parent_interaction_id: '',
             },
           ],
         }
@@ -199,7 +195,6 @@ describe('MessageParserRunner', () => {
           interaction_id: '',
           create_time: '',
           additional_info: {},
-          parent_interaction_id: '',
         },
         {
           interactions: [
@@ -210,7 +205,6 @@ describe('MessageParserRunner', () => {
               interaction_id: '',
               create_time: '',
               additional_info: {},
-              parent_interaction_id: '',
             },
           ],
         }