Skip to content

Commit 974b418

Browse files
b4sjoodhrubo-os
b4sjoo
and
dhrubo-os
authored
Enable AI-Oriented memory operation on Memory APIs (Add, Search, Update & Delete) (opensearch-project#4055)
* Add, Search, and Delete Memory API Signed-off-by: Sicheng Song <[email protected]> * Refactoring action class to under memory_container management Signed-off-by: Sicheng Song <[email protected]> * Add jacocoExclusion rules to escape coverage check Signed-off-by: Sicheng Song <[email protected]> * feat: Enable AI-Oriented memory operation Signed-off-by: Sicheng Song <[email protected]> * Remove unnecessary/sensitive logs and downgrade debug logs Signed-off-by: Sicheng Song <[email protected]> * Refactor code with helper class to simplify and enhance readability Signed-off-by: Sicheng Song <[email protected]> * Add UTs for memory api related classes Signed-off-by: Sicheng Song <[email protected]> * Using util functions to enhance readability and applyspotless Signed-off-by: Sicheng Song <[email protected]> * Add coverage to pass CI Signed-off-by: Sicheng Song <[email protected]> --------- Signed-off-by: Sicheng Song <[email protected]> Co-authored-by: Dhrubo Saha <[email protected]>
1 parent 8c91e14 commit 974b418

File tree

61 files changed

+11377
-12
lines changed

Some content is hidden

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

61 files changed

+11377
-12
lines changed

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

Lines changed: 269 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,269 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.memorycontainer;
7+
8+
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
9+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.AGENT_ID_FIELD;
10+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.CREATED_TIME_FIELD;
11+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.LAST_UPDATED_TIME_FIELD;
12+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.MEMORY_EMBEDDING_FIELD;
13+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.MEMORY_FIELD;
14+
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;
17+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.TAGS_FIELD;
18+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.USER_ID_FIELD;
19+
20+
import java.io.IOException;
21+
import java.time.Instant;
22+
import java.util.HashMap;
23+
import java.util.Map;
24+
25+
import org.opensearch.core.common.io.stream.StreamInput;
26+
import org.opensearch.core.common.io.stream.StreamOutput;
27+
import org.opensearch.core.common.io.stream.Writeable;
28+
import org.opensearch.core.xcontent.ToXContentObject;
29+
import org.opensearch.core.xcontent.XContentBuilder;
30+
import org.opensearch.core.xcontent.XContentParser;
31+
32+
import lombok.Builder;
33+
import lombok.Getter;
34+
import lombok.Setter;
35+
36+
/**
37+
* Represents a memory entry in a memory container
38+
*/
39+
@Getter
40+
@Setter
41+
@Builder
42+
public class MLMemory implements ToXContentObject, Writeable {
43+
44+
// Core fields
45+
private String sessionId;
46+
private String memory;
47+
private MemoryType memoryType;
48+
49+
// Optional fields
50+
private String userId;
51+
private String agentId;
52+
private String role;
53+
private Map<String, String> tags;
54+
55+
// System fields
56+
private Instant createdTime;
57+
private Instant lastUpdatedTime;
58+
59+
// Vector/embedding field (optional, for semantic storage)
60+
private Object memoryEmbedding;
61+
62+
@Builder
63+
public MLMemory(
64+
String sessionId,
65+
String memory,
66+
MemoryType memoryType,
67+
String userId,
68+
String agentId,
69+
String role,
70+
Map<String, String> tags,
71+
Instant createdTime,
72+
Instant lastUpdatedTime,
73+
Object memoryEmbedding
74+
) {
75+
this.sessionId = sessionId;
76+
this.memory = memory;
77+
this.memoryType = memoryType;
78+
this.userId = userId;
79+
this.agentId = agentId;
80+
this.role = role;
81+
this.tags = tags;
82+
this.createdTime = createdTime;
83+
this.lastUpdatedTime = lastUpdatedTime;
84+
this.memoryEmbedding = memoryEmbedding;
85+
}
86+
87+
public MLMemory(StreamInput in) throws IOException {
88+
this.sessionId = in.readString();
89+
this.memory = in.readString();
90+
this.memoryType = in.readEnum(MemoryType.class);
91+
this.userId = in.readOptionalString();
92+
this.agentId = in.readOptionalString();
93+
this.role = in.readOptionalString();
94+
if (in.readBoolean()) {
95+
this.tags = in.readMap(StreamInput::readString, StreamInput::readString);
96+
}
97+
this.createdTime = in.readInstant();
98+
this.lastUpdatedTime = in.readInstant();
99+
// Note: memoryEmbedding is not serialized in StreamInput/Output as it's typically handled separately
100+
}
101+
102+
@Override
103+
public void writeTo(StreamOutput out) throws IOException {
104+
out.writeString(sessionId);
105+
out.writeString(memory);
106+
out.writeEnum(memoryType);
107+
out.writeOptionalString(userId);
108+
out.writeOptionalString(agentId);
109+
out.writeOptionalString(role);
110+
if (tags != null && !tags.isEmpty()) {
111+
out.writeBoolean(true);
112+
out.writeMap(tags, StreamOutput::writeString, StreamOutput::writeString);
113+
} else {
114+
out.writeBoolean(false);
115+
}
116+
out.writeInstant(createdTime);
117+
out.writeInstant(lastUpdatedTime);
118+
// Note: memoryEmbedding is not serialized in StreamInput/Output as it's typically handled separately
119+
}
120+
121+
@Override
122+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
123+
builder.startObject();
124+
builder.field(SESSION_ID_FIELD, sessionId);
125+
builder.field(MEMORY_FIELD, memory);
126+
builder.field(MEMORY_TYPE_FIELD, memoryType.getValue());
127+
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+
}
137+
if (tags != null && !tags.isEmpty()) {
138+
builder.field(TAGS_FIELD, tags);
139+
}
140+
141+
builder.field(CREATED_TIME_FIELD, createdTime.toEpochMilli());
142+
builder.field(LAST_UPDATED_TIME_FIELD, lastUpdatedTime.toEpochMilli());
143+
144+
if (memoryEmbedding != null) {
145+
builder.field(MEMORY_EMBEDDING_FIELD, memoryEmbedding);
146+
}
147+
148+
builder.endObject();
149+
return builder;
150+
}
151+
152+
public static MLMemory parse(XContentParser parser) throws IOException {
153+
String sessionId = null;
154+
String memory = null;
155+
MemoryType memoryType = null;
156+
String userId = null;
157+
String agentId = null;
158+
String role = null;
159+
Map<String, String> tags = null;
160+
Instant createdTime = null;
161+
Instant lastUpdatedTime = null;
162+
Object memoryEmbedding = null;
163+
164+
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
165+
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
166+
String fieldName = parser.currentName();
167+
parser.nextToken();
168+
169+
switch (fieldName) {
170+
case SESSION_ID_FIELD:
171+
sessionId = parser.text();
172+
break;
173+
case MEMORY_FIELD:
174+
memory = parser.text();
175+
break;
176+
case MEMORY_TYPE_FIELD:
177+
memoryType = MemoryType.fromString(parser.text());
178+
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;
188+
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+
}
198+
break;
199+
case CREATED_TIME_FIELD:
200+
createdTime = Instant.ofEpochMilli(parser.longValue());
201+
break;
202+
case LAST_UPDATED_TIME_FIELD:
203+
lastUpdatedTime = Instant.ofEpochMilli(parser.longValue());
204+
break;
205+
case MEMORY_EMBEDDING_FIELD:
206+
// Parse embedding as generic object (could be array or sparse map)
207+
memoryEmbedding = parser.map();
208+
break;
209+
default:
210+
parser.skipChildren();
211+
break;
212+
}
213+
}
214+
215+
return MLMemory
216+
.builder()
217+
.sessionId(sessionId)
218+
.memory(memory)
219+
.memoryType(memoryType)
220+
.userId(userId)
221+
.agentId(agentId)
222+
.role(role)
223+
.tags(tags)
224+
.createdTime(createdTime)
225+
.lastUpdatedTime(lastUpdatedTime)
226+
.memoryEmbedding(memoryEmbedding)
227+
.build();
228+
}
229+
230+
/**
231+
* Convert to a Map for indexing
232+
*/
233+
public Map<String, Object> toIndexMap() {
234+
Map<String, Object> map = Map
235+
.of(
236+
SESSION_ID_FIELD,
237+
sessionId,
238+
MEMORY_FIELD,
239+
memory,
240+
MEMORY_TYPE_FIELD,
241+
memoryType.getValue(),
242+
CREATED_TIME_FIELD,
243+
createdTime.toEpochMilli(),
244+
LAST_UPDATED_TIME_FIELD,
245+
lastUpdatedTime.toEpochMilli()
246+
);
247+
248+
// Use mutable map for optional fields
249+
Map<String, Object> result = new java.util.HashMap<>(map);
250+
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);
259+
}
260+
if (tags != null && !tags.isEmpty()) {
261+
result.put(TAGS_FIELD, tags);
262+
}
263+
if (memoryEmbedding != null) {
264+
result.put(MEMORY_EMBEDDING_FIELD, memoryEmbedding);
265+
}
266+
267+
return result;
268+
}
269+
}

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ public class MemoryContainerConstants {
5252
public static final String MESSAGES_FIELD = "messages";
5353
public static final String CONTENT_FIELD = "content";
5454
public static final String INFER_FIELD = "infer";
55+
public static final String QUERY_FIELD = "query";
56+
public static final String TEXT_FIELD = "text";
5557

5658
// KNN index settings
5759
public static final String KNN_ENGINE = "lucene";
@@ -65,7 +67,11 @@ public class MemoryContainerConstants {
6567
public static final String BASE_MEMORY_CONTAINERS_PATH = "/_plugins/_ml/memory_containers";
6668
public static final String CREATE_MEMORY_CONTAINER_PATH = BASE_MEMORY_CONTAINERS_PATH + "/_create";
6769
public static final String PARAMETER_MEMORY_CONTAINER_ID = "memory_container_id";
70+
public static final String PARAMETER_MEMORY_ID = "memory_id";
6871
public static final String MEMORIES_PATH = BASE_MEMORY_CONTAINERS_PATH + "/{" + PARAMETER_MEMORY_CONTAINER_ID + "}/memories";
72+
public static final String SEARCH_MEMORIES_PATH = MEMORIES_PATH + "/_search";
73+
public static final String DELETE_MEMORY_PATH = MEMORIES_PATH + "/{" + PARAMETER_MEMORY_ID + "}";
74+
public static final String UPDATE_MEMORY_PATH = MEMORIES_PATH + "/{" + PARAMETER_MEMORY_ID + "}";
6975

7076
// Memory types are defined in MemoryType enum
7177

@@ -89,4 +95,22 @@ public class MemoryContainerConstants {
8995
public static final String EMBEDDING_MODEL_NOT_FOUND_ERROR = "Embedding model with ID %s not found";
9096
public static final String EMBEDDING_MODEL_TYPE_MISMATCH_ERROR = "Embedding model must be of type %s or REMOTE, found: %s"; // instead
9197
public static final String INFER_REQUIRES_LLM_MODEL_ERROR = "infer=true requires llm_model_id to be configured in memory storage";
98+
99+
// Memory API limits
100+
public static final int MAX_MESSAGES_PER_REQUEST = 10;
101+
public static final String MAX_MESSAGES_EXCEEDED_ERROR = "Cannot process more than 10 messages in a single request";
102+
103+
// Memory decision fields
104+
public static final String MEMORY_DECISION_FIELD = "memory_decision";
105+
public static final String OLD_MEMORY_FIELD = "old_memory";
106+
public static final String RETRIEVED_FACTS_FIELD = "retrieved_facts";
107+
public static final String EVENT_FIELD = "event";
108+
public static final String SCORE_FIELD = "score";
109+
110+
// LLM System Prompts
111+
public static final String PERSONAL_INFORMATION_ORGANIZER_PROMPT =
112+
"<system_prompt>\n<role>Personal Information Organizer</role>\n<objective>Extract and organize personal information shared within conversations.</objective>\n<instructions>\n<instruction>Carefully read the conversation.</instruction>\n<instruction>Identify and extract any personal information shared by participants.</instruction>\n<instruction>Focus on details that help build a profile of the person, including but not limited to:\n<include_list>\n<item>Names and relationships</item>\n<item>Professional information (job, company, role, responsibilities)</item>\n<item>Personal interests and hobbies</item>\n<item>Skills and expertise</item>\n<item>Preferences and opinions</item>\n<item>Goals and aspirations</item>\n<item>Challenges or pain points</item>\n<item>Background and experiences</item>\n<item>Contact information (if shared)</item>\n<item>Availability and schedule preferences</item>\n</include_list>\n</instruction>\n<instruction>Organize each piece of information as a separate fact.</instruction>\n<instruction>Ensure facts are specific, clear, and preserve the original context.</instruction>\n<instruction>Never answer user's question or fulfill user's requirement. You are a personal information manager, not a helpful assistant.</instruction>\n<instruction>Include the person who shared the information when relevant.</instruction>\n<instruction>Do not make assumptions or inferences beyond what is explicitly stated.</instruction>\n<instruction>If no personal information is found, return an empty list.</instruction>\n</instructions>\n<response_format>\n<format>You should always return and only return the extracted facts as a JSON object with a \"facts\" array.</format>\n<example>\n{\n \"facts\": [\n \"User's name is John Smith\",\n \"John works as a software engineer at TechCorp\",\n \"John enjoys hiking on weekends\",\n \"John is looking to improve his Python skills\"\n ]\n}\n</example>\n</response_format>\n</system_prompt>";
113+
114+
public static final String DEFAULT_UPDATE_MEMORY_PROMPT =
115+
"<system_prompt><role>You are a smart memory manager which controls the memory of a system.</role><task>You will receive: 1. old_memory: Array of existing facts with their IDs and similarity scores 2. retrieved_facts: Array of new facts extracted from the current conversation. Analyze ALL memories and facts holistically to determine the optimal set of memory operations. Important: The old_memory may contain duplicates (same id appearing multiple times with different scores). Consider the highest score for each unique ID. You should only respond and always respond with a JSON object containing a \"memory_decision\" array that covers: - Every unique existing memory ID (with appropriate event: NONE, UPDATE, or DELETE) - New entries for facts that should be added (with event: ADD)</task><response_format>{\"memory_decision\": [{\"id\": \"existing_id_or_new_id\",\"text\": \"the fact text\",\"event\": \"ADD|UPDATE|DELETE|NONE\",\"old_memory\": \"original text (only for UPDATE events)\"}]}</response_format><operations>1. **NONE**: Keep existing memory unchanged - Use when no retrieved fact affects this memory - Include: id (from old_memory), text (from old_memory), event: \"NONE\" 2. **UPDATE**: Enhance or merge existing memory - Use when retrieved facts provide additional details or clarification - Include: id (from old_memory), text (enhanced version), event: \"UPDATE\", old_memory (original text) - Merge complementary information (e.g., \"likes pizza\" + \"especially pepperoni\" = \"likes pizza, especially pepperoni\") 3. **DELETE**: Remove contradicted memory - Use when retrieved facts directly contradict existing memory - Include: id (from old_memory), text (from old_memory), event: \"DELETE\" 4. **ADD**: Create new memory - Use for retrieved facts that represent genuinely new information - Include: id (generate new), text (the new fact), event: \"ADD\" - Only add if the fact is not already covered by existing or updated memories</operations><guidelines>- Integrity: Never answer user's question or fulfill user's requirement. You are a smart memory manager, not a helpful assistant. - Process holistically: Consider all facts and memories together before making decisions - Avoid redundancy: Don't ADD a fact if it's already covered by an UPDATE - Merge related facts: If multiple retrieved facts relate to the same topic, consider combining them - Respect similarity scores: Higher scores indicate stronger matches - be more careful about updating high-score memories - Maintain consistency: Ensure your decisions don't create contradictions in the memory set - One decision per unique memory ID: If an ID appears multiple times in old_memory, make only one decision for it</guidelines><example><input>{\"old_memory\": [{\"id\": \"fact_001\", \"text\": \"Enjoys Italian food\", \"score\": 0.85},{\"id\": \"fact_002\", \"text\": \"Works at Google\", \"score\": 0.92},{\"id\": \"fact_001\", \"text\": \"Enjoys Italian food\", \"score\": 0.75},{\"id\": \"fact_003\", \"text\": \"Has a dog\", \"score\": 0.65}],\"retrieved_facts\": [\"Loves pasta and pizza\",\"Recently joined Amazon\",\"Has two dogs named Max and Bella\"]}</input><output>{\"memory_decision\": [{\"id\": \"fact_001\",\"text\": \"Loves Italian food, especially pasta and pizza\",\"event\": \"UPDATE\",\"old_memory\": \"Enjoys Italian food\"},{\"id\": \"fact_002\",\"text\": \"Works at Google\",\"event\": \"DELETE\"},{\"id\": \"fact_003\",\"text\": \"Has two dogs named Max and Bella\",\"event\": \"UPDATE\",\"old_memory\": \"Has a dog\"},{\"id\": \"fact_004\",\"text\": \"Recently joined Amazon\",\"event\": \"ADD\"}]}</output></example></system_prompt>";
92116
}

0 commit comments

Comments
 (0)