refactor memory interface; add agentic conversation memory (#4434)

Signed-off-by: Yaliang Wu <ylwu@amazon.com>

Author: ylwu-amzn

common/src/main/java/org/opensearch/ml/common/MLMemoryType.java

@@ -0,0 +1,25 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common;
+
+import java.util.Locale;
+
+public enum MLMemoryType {
+    CONVERSATION_INDEX,
+    AGENTIC_MEMORY,
+    REMOTE_AGENTIC_MEMORY;
+
+    public static MLMemoryType from(String value) {
+        if (value != null) {
+            try {
+                return MLMemoryType.valueOf(value.toUpperCase(Locale.ROOT));
+            } catch (Exception e) {
+                throw new IllegalArgumentException("Wrong Memory type");
+            }
+        }
+        return null;
+    }
+}

common/src/main/java/org/opensearch/ml/common/agent/MLMemorySpec.java

@@ -26,32 +26,37 @@ public class MLMemorySpec implements ToXContentObject {
     public static final String MEMORY_TYPE_FIELD = "type";
     public static final String WINDOW_SIZE_FIELD = "window_size";
     public static final String SESSION_ID_FIELD = "session_id";
+    public static final String MEMORY_CONTAINER_ID_FIELD = "memory_container_id";
 
     private String type;
     @Setter
     private String sessionId;
     private Integer windowSize;
+    private String memoryContainerId;
 
     @Builder(toBuilder = true)
-    public MLMemorySpec(String type, String sessionId, Integer windowSize) {
+    public MLMemorySpec(String type, String sessionId, Integer windowSize, String memoryContainerId) {
         if (type == null) {
             throw new IllegalArgumentException("agent name is null");
         }
         this.type = type;
         this.sessionId = sessionId;
         this.windowSize = windowSize;
+        this.memoryContainerId = memoryContainerId;
     }
 
     public MLMemorySpec(StreamInput input) throws IOException {
         type = input.readString();
         sessionId = input.readOptionalString();
         windowSize = input.readOptionalInt();
+        memoryContainerId = input.readOptionalString();
     }
 
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(type);
         out.writeOptionalString(sessionId);
         out.writeOptionalInt(windowSize);
+        out.writeOptionalString(memoryContainerId);
     }
 
     @Override
@@ -64,6 +69,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (sessionId != null) {
             builder.field(SESSION_ID_FIELD, sessionId);
         }
+        if (memoryContainerId != null) {
+            builder.field(MEMORY_CONTAINER_ID_FIELD, memoryContainerId);
+        }
         builder.endObject();
         return builder;
     }
@@ -72,6 +80,7 @@ public static MLMemorySpec parse(XContentParser parser) throws IOException {
         String type = null;
         String sessionId = null;
         Integer windowSize = null;
+        String memoryContainerId = null;
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -88,12 +97,15 @@ public static MLMemorySpec parse(XContentParser parser) throws IOException {
                 case WINDOW_SIZE_FIELD:
                     windowSize = parser.intValue();
                     break;
+                case MEMORY_CONTAINER_ID_FIELD:
+                    memoryContainerId = parser.text();
+                    break;
                 default:
                     parser.skipChildren();
                     break;
             }
         }
-        return MLMemorySpec.builder().type(type).sessionId(sessionId).windowSize(windowSize).build();
+        return MLMemorySpec.builder().type(type).sessionId(sessionId).windowSize(windowSize).memoryContainerId(memoryContainerId).build();
     }
 
     public static MLMemorySpec fromStream(StreamInput in) throws IOException {

common/src/main/java/org/opensearch/ml/common/conversation/Interaction.java

@@ -28,6 +28,7 @@
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.ml.common.CommonValue;
+import org.opensearch.ml.common.memory.Message;
 import org.opensearch.search.SearchHit;
 
 import lombok.AllArgsConstructor;
@@ -39,7 +40,7 @@
  */
 @Builder
 @AllArgsConstructor
-public class Interaction implements Writeable, ToXContentObject {
+public class Interaction implements Writeable, ToXContentObject, Message {
 
     @Getter
     private String id;
@@ -275,4 +276,13 @@ public String toString() {
             + "}";
     }
 
+    @Override
+    public String getType() {
+        return "";
+    }
+
+    @Override
+    public String getContent() {
+        return "";
+    }
 }

common/src/main/java/org/opensearch/ml/common/memory/Memory.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.memory;
+
+import java.util.List;
+import java.util.Map;
+
+import org.opensearch.core.action.ActionListener;
+
+/**
+ * A general memory interface.
+ * @param <T> Message type
+ * @param <R> Save response type
+ * @param <S> Update response type
+ */
+public interface Memory<T extends Message, R, S> {
+
+    /**
+     * Get memory type.
+     * @return memory type
+     */
+    String getType();
+
+    /**
+     * Get memory ID.
+     * @return memory ID
+     */
+    String getId();
+
+    default void save(Message message, String parentId, Integer traceNum, String action) {}
+
+    default void save(Message message, String parentId, Integer traceNum, String action, ActionListener<R> listener) {}
+
+    default void update(String messageId, Map<String, Object> updateContent, ActionListener<S> updateListener) {}
+
+    default void getMessages(int size, ActionListener<List<T>> listener) {}
+
+    /**
+     * Clear all memory.
+     */
+    void clear();
+
+    void deleteInteractionAndTrace(String regenerateInteractionId, ActionListener<Boolean> wrap);
+
+    interface Factory<M extends Memory> {
+        /**
+         * Create an instance of this Memory.
+         *
+         * @param params Parameters for the memory
+         * @param listener Action listener for the memory creation action
+         */
+        void create(Map<String, Object> params, ActionListener<M> listener);
+    }
+}

spi/src/main/java/org/opensearch/ml/common/spi/memory/Message.java renamed to common/src/main/java/org/opensearch/ml/common/memory/Message.java

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.ml.common.spi.memory;
+package org.opensearch.ml.common.memory;
 
 /**
  * General message interface.
@@ -12,13 +12,13 @@ public interface Message {
 
     /**
      * Get message type.
-     * @return
+     * @return message type
      */
     String getType();
 
     /**
      * Get message content.
-     * @return
+     * @return message content
      */
     String getContent();
 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/agents/AgentContextUtil.java

@@ -17,7 +17,7 @@
 import org.opensearch.ml.common.hooks.EnhancedPostToolEvent;
 import org.opensearch.ml.common.hooks.HookRegistry;
 import org.opensearch.ml.common.hooks.PreLLMEvent;
-import org.opensearch.ml.common.spi.memory.Memory;
+import org.opensearch.ml.common.memory.Memory;
 import org.opensearch.ml.common.utils.StringUtils;
 import org.opensearch.ml.engine.memory.ConversationIndexMemory;
 

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java

@@ -10,6 +10,7 @@
 import static org.opensearch.ml.common.CommonValue.MCP_CONNECTORS_FIELD;
 import static org.opensearch.ml.common.CommonValue.MCP_CONNECTOR_ID_FIELD;
 import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX;
+import static org.opensearch.ml.common.agent.MLMemorySpec.MEMORY_CONTAINER_ID_FIELD;
 import static org.opensearch.ml.common.utils.StringUtils.getParameterMap;
 import static org.opensearch.ml.common.utils.StringUtils.gson;
 import static org.opensearch.ml.common.utils.StringUtils.isJson;
@@ -29,6 +30,7 @@
 import static org.opensearch.ml.engine.algorithms.agent.MLChatAgentRunner.TOOL_DESCRIPTIONS;
 import static org.opensearch.ml.engine.algorithms.agent.MLChatAgentRunner.TOOL_NAMES;
 import static org.opensearch.ml.engine.algorithms.agent.MLPlanExecuteAndReflectAgentRunner.RESPONSE_FIELD;
+import static org.opensearch.ml.engine.memory.ConversationIndexMemory.APP_TYPE;
 import static org.opensearch.ml.engine.memory.ConversationIndexMemory.LAST_N_INTERACTIONS;
 
 import java.io.IOException;
@@ -83,6 +85,7 @@
 import org.opensearch.ml.engine.algorithms.remote.McpStreamableHttpConnectorExecutor;
 import org.opensearch.ml.engine.encryptor.Encryptor;
 import org.opensearch.ml.engine.function_calling.FunctionCalling;
+import org.opensearch.ml.engine.memory.ConversationIndexMemory;
 import org.opensearch.ml.engine.tools.McpSseTool;
 import org.opensearch.ml.engine.tools.McpStreamableHttpTool;
 import org.opensearch.remote.metadata.client.GetDataObjectRequest;
@@ -1014,4 +1017,22 @@ public static Tool createTool(Map<String, Tool.Factory> toolFactories, Map<Strin
 
         return tool;
     }
+
+    public static Map<String, Object> createMemoryParams(
+        String question,
+        String memoryId,
+        String appType,
+        MLAgent mlAgent,
+        String memoryContainerId
+    ) {
+        Map<String, Object> memoryParams = new HashMap<>();
+        memoryParams.put(ConversationIndexMemory.MEMORY_NAME, question);
+        memoryParams.put(ConversationIndexMemory.MEMORY_ID, memoryId);
+        memoryParams.put(APP_TYPE, appType);
+        if (mlAgent.getMemory().getMemoryContainerId() != null) {
+            memoryParams.put(MEMORY_CONTAINER_ID_FIELD, mlAgent.getMemory().getMemoryContainerId());
+        }
+        memoryParams.putIfAbsent(MEMORY_CONTAINER_ID_FIELD, memoryContainerId);
+        return memoryParams;
+    }
 }