Skip to content

Commit 8ae83bb

Browse files
ylwu-amznb4sjooZhangxunmt
authored
Refactor Agentic Memory (opensearch-project#4218)
* Refactor agentic memory Signed-off-by: Yaliang Wu <[email protected]> * fix namespace issue Signed-off-by: Yaliang Wu <[email protected]> * add more details to tutorial doc Signed-off-by: Yaliang Wu <[email protected]> * add trace data Signed-off-by: Yaliang Wu <[email protected]> * remove long term memory index mapping Signed-off-by: Yaliang Wu <[email protected]> * add index setting Signed-off-by: Yaliang Wu <[email protected]> * add flag to disable session and history Signed-off-by: Yaliang Wu <[email protected]> * Working memory related APIs (opensearch-project#4188) * Aligning agentic memory user experience with Bedrock AgentCore Memory * Rename type to strategy Signed-off-by: Sicheng Song <[email protected]> * Working memory related APIs Signed-off-by: Sicheng Song <[email protected]> --------- Signed-off-by: Sicheng Song <[email protected]> * fix unit test Signed-off-by: Yaliang Wu <[email protected]> * fix bugs Signed-off-by: Yaliang Wu <[email protected]> * fix unit test Signed-off-by: Yaliang Wu <[email protected]> * fix unit test Signed-off-by: Yaliang Wu <[email protected]> * add backendrole in container for access control and update memory container API (opensearch-project#4206) Signed-off-by: Xun Zhang <[email protected]> * tune code Signed-off-by: Yaliang Wu <[email protected]> * fix get working memory api Signed-off-by: Yaliang Wu <[email protected]> * fix update memory container Signed-off-by: Yaliang Wu <[email protected]> * fix system memory index; add owner id field Signed-off-by: Yaliang Wu <[email protected]> * add search memory containers api Signed-off-by: Yaliang Wu <[email protected]> * enhance get/update/delete memory api to support all memory data Signed-off-by: Yaliang Wu <[email protected]> * enhance search memory api to support all memory data Signed-off-by: Yaliang Wu <[email protected]> * Further refactor on working memory related class (opensearch-project#4208) Signed-off-by: Sicheng Song <[email protected]> * fix failed ut in common module Signed-off-by: Yaliang Wu <[email protected]> * fix failed ut in plugin module Signed-off-by: Yaliang Wu <[email protected]> * set use_system_index default as true Signed-off-by: Yaliang Wu <[email protected]> * fix null exception for adding conversational memory when infer as false Signed-off-by: Yaliang Wu <[email protected]> * add more fields to memory history Signed-off-by: Yaliang Wu <[email protected]> * add error to memory history if failed to extract long term memory Signed-off-by: Yaliang Wu <[email protected]> * fix unit test for plugin module Signed-off-by: Yaliang Wu <[email protected]> * fix session index not created bug for working memory Signed-off-by: Yaliang Wu <[email protected]> * tune memory history before/after Signed-off-by: Yaliang Wu <[email protected]> * fix failed ut; tune memory history Signed-off-by: Yaliang Wu <[email protected]> * allow custom prompt in the memory strategy (opensearch-project#4216) Signed-off-by: Xun Zhang <[email protected]> * tune messages to support multiple LLM Signed-off-by: Yaliang Wu <[email protected]> * support multiple LLM Signed-off-by: Yaliang Wu <[email protected]> * fix ut Signed-off-by: Yaliang Wu <[email protected]> * support USER_PREFERENCE and SUMMARY strategies (opensearch-project#4217) * fix customer index bug when no prefix * support USER_PREFERENCE and SUMMARY strategies * fix bug when generating multiple strategies * run spotlessApply Signed-off-by: Yaliang Wu <[email protected]> * revert MLAddMemoriesInput to not fail ut Signed-off-by: Yaliang Wu <[email protected]> * fix failed ut Signed-off-by: Yaliang Wu <[email protected]> * put memory to before field for delete history event Signed-off-by: Yaliang Wu <[email protected]> * Add UT coverages (opensearch-project#4223) * Add more UT coverages (opensearch-project#4224) * use LLM to summarize session Signed-off-by: Yaliang Wu <[email protected]> * fix failed UT Signed-off-by: Yaliang Wu <[email protected]> * Add a generic delete by query API (opensearch-project#4225) * remove get/delete working memory API Signed-off-by: Yaliang Wu <[email protected]> * use correct thread pool Signed-off-by: Yaliang Wu <[email protected]> * remove duplicate error constant Signed-off-by: Yaliang Wu <[email protected]> * fix failed ut Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]> Signed-off-by: Sicheng Song <[email protected]> Signed-off-by: Xun Zhang <[email protected]> Co-authored-by: Sicheng Song <[email protected]> Co-authored-by: Xun Zhang <[email protected]>
1 parent f1e6977 commit 8ae83bb

File tree

121 files changed

+12292
-6777
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

121 files changed

+12292
-6777
lines changed

‎common/src/main/java/org/opensearch/ml/common/CommonValue.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,9 @@ public class CommonValue {
4646
public static final String ML_AGENT_INDEX = ".plugins-ml-agent";
4747
public static final String ML_MEMORY_META_INDEX = ".plugins-ml-memory-meta";
4848
public static final String ML_MEMORY_MESSAGE_INDEX = ".plugins-ml-memory-message";
49-
public static final String ML_MEMORY_CONTAINER_INDEX = ".plugins-ml-memory-container";
49+
public static final String ML_MEMORY_CONTAINER_INDEX = ".plugins-ml-am-memory-container";
50+
public static final String ML_AGENTIC_MEMORY_SYSTEM_INDEX_PREFIX = ".plugins-ml-am";
51+
public static final String ML_AGENTIC_MEMORY_INDEX_PATTERN = ML_AGENTIC_MEMORY_SYSTEM_INDEX_PREFIX + "*";
5052
public static final String ML_STOP_WORDS_INDEX = ".plugins-ml-stop-words";
5153
// index used in 2.19 to track MlTaskBatchUpdate task
5254
public static final String TASK_POLLING_JOB_INDEX = ".ml_commons_task_polling_job";
@@ -68,6 +70,9 @@ public class CommonValue {
6870
public static final String ML_MEMORY_META_INDEX_MAPPING_PATH = "index-mappings/ml_memory_meta.json";
6971
public static final String ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH = "index-mappings/ml_memory_message.json";
7072
public static final String ML_MEMORY_CONTAINER_INDEX_MAPPING_PATH = "index-mappings/ml_memory_container.json";
73+
public static final String ML_MEMORY_SESSION_INDEX_MAPPING_PATH = "index-mappings/ml_memory_sessions.json";
74+
public static final String ML_WORKING_MEMORY_INDEX_MAPPING_PATH = "index-mappings/ml_memory_working.json";
75+
public static final String ML_LONG_MEMORY_HISTORY_INDEX_MAPPING_PATH = "index-mappings/ml_memory_long_term_history.json";
7176
public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
7277
public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
7378
public static final String ML_JOBS_INDEX_MAPPING_PATH = "index-mappings/ml_jobs.json";

‎common/src/main/java/org/opensearch/ml/common/memorycontainer/MLMemory.java‎

Lines changed: 70 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -6,20 +6,19 @@
66
package org.opensearch.ml.common.memorycontainer;
77

88
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
9-
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.AGENT_ID_FIELD;
109
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.CREATED_TIME_FIELD;
1110
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.LAST_UPDATED_TIME_FIELD;
1211
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.MEMORY_EMBEDDING_FIELD;
1312
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.MEMORY_FIELD;
1413
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.MEMORY_TYPE_FIELD;
15-
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.ROLE_FIELD;
16-
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.SESSION_ID_FIELD;
14+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.NAMESPACE_FIELD;
15+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.NAMESPACE_SIZE_FIELD;
16+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.OWNER_ID_FIELD;
17+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.STRATEGY_ID_FIELD;
1718
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.TAGS_FIELD;
18-
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.USER_ID_FIELD;
1919

2020
import java.io.IOException;
2121
import java.time.Instant;
22-
import java.util.HashMap;
2322
import java.util.Map;
2423

2524
import org.opensearch.core.common.io.stream.StreamInput;
@@ -28,6 +27,7 @@
2827
import org.opensearch.core.xcontent.ToXContentObject;
2928
import org.opensearch.core.xcontent.XContentBuilder;
3029
import org.opensearch.core.xcontent.XContentParser;
30+
import org.opensearch.ml.common.utils.StringUtils;
3131

3232
import lombok.Builder;
3333
import lombok.Getter;
@@ -42,14 +42,10 @@
4242
public class MLMemory implements ToXContentObject, Writeable {
4343

4444
// Core fields
45-
private String sessionId;
4645
private String memory;
4746
private MemoryType memoryType;
4847

49-
// Optional fields
50-
private String userId;
51-
private String agentId;
52-
private String role;
48+
private Map<String, String> namespace;
5349
private Map<String, String> tags;
5450

5551
// System fields
@@ -58,55 +54,58 @@ public class MLMemory implements ToXContentObject, Writeable {
5854

5955
// Vector/embedding field (optional, for semantic storage)
6056
private Object memoryEmbedding;
57+
private String ownerId;
58+
private String strategyId;
6159

6260
@Builder
6361
public MLMemory(
64-
String sessionId,
6562
String memory,
6663
MemoryType memoryType,
67-
String userId,
68-
String agentId,
69-
String role,
64+
Map<String, String> namespace,
7065
Map<String, String> tags,
7166
Instant createdTime,
7267
Instant lastUpdatedTime,
73-
Object memoryEmbedding
68+
Object memoryEmbedding,
69+
String ownerId,
70+
String strategyId
7471
) {
75-
this.sessionId = sessionId;
7672
this.memory = memory;
7773
this.memoryType = memoryType;
78-
this.userId = userId;
79-
this.agentId = agentId;
80-
this.role = role;
74+
this.namespace = namespace;
8175
this.tags = tags;
8276
this.createdTime = createdTime;
8377
this.lastUpdatedTime = lastUpdatedTime;
8478
this.memoryEmbedding = memoryEmbedding;
79+
this.ownerId = ownerId;
80+
this.strategyId = strategyId;
8581
}
8682

8783
public MLMemory(StreamInput in) throws IOException {
88-
this.sessionId = in.readString();
8984
this.memory = in.readString();
9085
this.memoryType = in.readEnum(MemoryType.class);
91-
this.userId = in.readOptionalString();
92-
this.agentId = in.readOptionalString();
93-
this.role = in.readOptionalString();
86+
if (in.readBoolean()) {
87+
this.namespace = in.readMap(StreamInput::readString, StreamInput::readString);
88+
}
9489
if (in.readBoolean()) {
9590
this.tags = in.readMap(StreamInput::readString, StreamInput::readString);
9691
}
9792
this.createdTime = in.readInstant();
9893
this.lastUpdatedTime = in.readInstant();
94+
this.ownerId = in.readOptionalString();
95+
this.strategyId = in.readOptionalString();
9996
// Note: memoryEmbedding is not serialized in StreamInput/Output as it's typically handled separately
10097
}
10198

10299
@Override
103100
public void writeTo(StreamOutput out) throws IOException {
104-
out.writeString(sessionId);
105101
out.writeString(memory);
106102
out.writeEnum(memoryType);
107-
out.writeOptionalString(userId);
108-
out.writeOptionalString(agentId);
109-
out.writeOptionalString(role);
103+
if (namespace != null && !namespace.isEmpty()) {
104+
out.writeBoolean(true);
105+
out.writeMap(namespace, StreamOutput::writeString, StreamOutput::writeString);
106+
} else {
107+
out.writeBoolean(false);
108+
}
110109
if (tags != null && !tags.isEmpty()) {
111110
out.writeBoolean(true);
112111
out.writeMap(tags, StreamOutput::writeString, StreamOutput::writeString);
@@ -115,86 +114,70 @@ public void writeTo(StreamOutput out) throws IOException {
115114
}
116115
out.writeInstant(createdTime);
117116
out.writeInstant(lastUpdatedTime);
117+
out.writeOptionalString(ownerId);
118+
out.writeOptionalString(strategyId);
118119
// Note: memoryEmbedding is not serialized in StreamInput/Output as it's typically handled separately
119120
}
120121

121122
@Override
122123
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
123124
builder.startObject();
124-
builder.field(SESSION_ID_FIELD, sessionId);
125125
builder.field(MEMORY_FIELD, memory);
126126
builder.field(MEMORY_TYPE_FIELD, memoryType.getValue());
127127

128-
if (userId != null) {
129-
builder.field(USER_ID_FIELD, userId);
130-
}
131-
if (agentId != null) {
132-
builder.field(AGENT_ID_FIELD, agentId);
133-
}
134-
if (role != null) {
135-
builder.field(ROLE_FIELD, role);
136-
}
137128
if (tags != null && !tags.isEmpty()) {
138129
builder.field(TAGS_FIELD, tags);
139130
}
131+
if (namespace != null && !namespace.isEmpty()) {
132+
builder.field(NAMESPACE_FIELD, namespace);
133+
builder.field(NAMESPACE_SIZE_FIELD, namespace.size());
134+
}
140135

141136
builder.field(CREATED_TIME_FIELD, createdTime.toEpochMilli());
142137
builder.field(LAST_UPDATED_TIME_FIELD, lastUpdatedTime.toEpochMilli());
143138

144139
if (memoryEmbedding != null) {
145140
builder.field(MEMORY_EMBEDDING_FIELD, memoryEmbedding);
146141
}
142+
if (ownerId != null) {
143+
builder.field(OWNER_ID_FIELD, ownerId);
144+
}
145+
if (strategyId != null) {
146+
builder.field(STRATEGY_ID_FIELD, strategyId);
147+
}
147148

148149
builder.endObject();
149150
return builder;
150151
}
151152

152153
public static MLMemory parse(XContentParser parser) throws IOException {
153-
String sessionId = null;
154154
String memory = null;
155155
MemoryType memoryType = null;
156-
String userId = null;
157-
String agentId = null;
158-
String role = null;
156+
Map<String, String> namespace = null;
159157
Map<String, String> tags = null;
160158
Instant createdTime = null;
161159
Instant lastUpdatedTime = null;
162160
Object memoryEmbedding = null;
161+
String ownerId = null;
162+
String strategyId = null;
163163

164164
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
165165
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
166166
String fieldName = parser.currentName();
167167
parser.nextToken();
168168

169169
switch (fieldName) {
170-
case SESSION_ID_FIELD:
171-
sessionId = parser.text();
172-
break;
173170
case MEMORY_FIELD:
174171
memory = parser.text();
175172
break;
176173
case MEMORY_TYPE_FIELD:
177174
memoryType = MemoryType.fromString(parser.text());
178175
break;
179-
case USER_ID_FIELD:
180-
userId = parser.text();
181-
break;
182-
case AGENT_ID_FIELD:
183-
agentId = parser.text();
184-
break;
185-
case ROLE_FIELD:
186-
role = parser.text();
187-
break;
188176
case TAGS_FIELD:
189-
Map<String, Object> tagsMap = parser.map();
190-
if (tagsMap != null) {
191-
tags = new HashMap<>();
192-
for (Map.Entry<String, Object> entry : tagsMap.entrySet()) {
193-
if (entry.getValue() != null) {
194-
tags.put(entry.getKey(), entry.getValue().toString());
195-
}
196-
}
197-
}
177+
tags = StringUtils.getParameterMap(parser.map());
178+
break;
179+
case NAMESPACE_FIELD:
180+
namespace = StringUtils.getParameterMap(parser.map());
198181
break;
199182
case CREATED_TIME_FIELD:
200183
createdTime = Instant.ofEpochMilli(parser.longValue());
@@ -204,7 +187,19 @@ public static MLMemory parse(XContentParser parser) throws IOException {
204187
break;
205188
case MEMORY_EMBEDDING_FIELD:
206189
// Parse embedding as generic object (could be array or sparse map)
207-
memoryEmbedding = parser.map();
190+
if (parser.currentToken() == XContentParser.Token.START_ARRAY) {
191+
memoryEmbedding = parser.list(); // Simple list parsing like ModelTensor
192+
} else if (parser.currentToken() == XContentParser.Token.START_OBJECT) {
193+
memoryEmbedding = parser.map(); // For sparse embeddings
194+
} else {
195+
parser.skipChildren();
196+
}
197+
break;
198+
case OWNER_ID_FIELD:
199+
ownerId = parser.text();
200+
break;
201+
case STRATEGY_ID_FIELD:
202+
strategyId = parser.text();
208203
break;
209204
default:
210205
parser.skipChildren();
@@ -214,16 +209,15 @@ public static MLMemory parse(XContentParser parser) throws IOException {
214209

215210
return MLMemory
216211
.builder()
217-
.sessionId(sessionId)
218212
.memory(memory)
219213
.memoryType(memoryType)
220-
.userId(userId)
221-
.agentId(agentId)
222-
.role(role)
214+
.namespace(namespace)
223215
.tags(tags)
224216
.createdTime(createdTime)
225217
.lastUpdatedTime(lastUpdatedTime)
226218
.memoryEmbedding(memoryEmbedding)
219+
.ownerId(ownerId)
220+
.strategyId(strategyId)
227221
.build();
228222
}
229223

@@ -233,8 +227,6 @@ public static MLMemory parse(XContentParser parser) throws IOException {
233227
public Map<String, Object> toIndexMap() {
234228
Map<String, Object> map = Map
235229
.of(
236-
SESSION_ID_FIELD,
237-
sessionId,
238230
MEMORY_FIELD,
239231
memory,
240232
MEMORY_TYPE_FIELD,
@@ -248,22 +240,22 @@ public Map<String, Object> toIndexMap() {
248240
// Use mutable map for optional fields
249241
Map<String, Object> result = new java.util.HashMap<>(map);
250242

251-
if (userId != null) {
252-
result.put(USER_ID_FIELD, userId);
253-
}
254-
if (agentId != null) {
255-
result.put(AGENT_ID_FIELD, agentId);
256-
}
257-
if (role != null) {
258-
result.put(ROLE_FIELD, role);
243+
if (namespace != null && !namespace.isEmpty()) {
244+
result.put(NAMESPACE_FIELD, namespace);
245+
result.put(NAMESPACE_SIZE_FIELD, namespace.size());
259246
}
260247
if (tags != null && !tags.isEmpty()) {
261248
result.put(TAGS_FIELD, tags);
262249
}
263250
if (memoryEmbedding != null) {
264251
result.put(MEMORY_EMBEDDING_FIELD, memoryEmbedding);
265252
}
266-
253+
if (ownerId != null) {
254+
result.put(OWNER_ID_FIELD, ownerId);
255+
}
256+
if (strategyId != null) {
257+
result.put(STRATEGY_ID_FIELD, strategyId);
258+
}
267259
return result;
268260
}
269261
}

0 commit comments

Comments
 (0)