Skip to content

Commit ce7f3f5

Browse files
b4sjoob4sjoo
authored
Apply retry logic in remote agentic conversation memory (#4420)
* backport structured_data_blob changes from remote * apply retry logic in remote agentic conversation memory
1 parent 05ea0cb commit ce7f3f5

File tree

9 files changed

+471
-69
lines changed

9 files changed

+471
-69
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ public class CommonValue {
2828
public static final String CREATE_TIME_FIELD = "create_time";
2929
public static final String LAST_UPDATE_TIME_FIELD = "last_update_time";
3030

31+
public static final String AGENTIC_MEMORY_THREAD_POOL = "opensearch_ml_agentic_memory";
32+
3133
public static final String BOX_TYPE_KEY = "box_type";
3234
// hot node
3335
public static String HOT_BOX_TYPE = "hot";

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.OWNER_ID_FIELD;
2222
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PAYLOAD_TYPE_FIELD;
2323
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.SESSION_ID_FIELD;
24+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.STRUCTURED_DATA_BLOB_FIELD;
2425
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.STRUCTURED_DATA_FIELD;
2526
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.TAGS_FIELD;
2627

@@ -58,6 +59,7 @@ public class MLWorkingMemory implements ToXContentObject, Writeable {
5859
private Integer messageId;
5960
private String binaryData;
6061
private Map<String, Object> structuredData;
62+
private Map<String, Object> structuredDataBlob;
6163

6264
// Optional fields
6365
private Map<String, String> namespace;
@@ -80,6 +82,7 @@ public MLWorkingMemory(
8082
Integer messageId,
8183
String binaryData,
8284
Map<String, Object> structuredData,
85+
Map<String, Object> structuredDataBlob,
8386
Map<String, String> namespace,
8487
boolean infer,
8588
Map<String, String> metadata,
@@ -97,6 +100,7 @@ public MLWorkingMemory(
97100
this.messageId = messageId;
98101
this.binaryData = binaryData;
99102
this.structuredData = structuredData;
103+
this.structuredDataBlob = structuredDataBlob;
100104
this.namespace = namespace;
101105
this.namespaceSize = namespace == null ? null : namespace.size();
102106
this.infer = infer; // default infer is false
@@ -123,6 +127,9 @@ public MLWorkingMemory(StreamInput in) throws IOException {
123127
if (in.readBoolean()) {
124128
this.structuredData = in.readMap();
125129
}
130+
if (in.readBoolean()) {
131+
this.structuredDataBlob = in.readMap();
132+
}
126133
if (in.readBoolean()) {
127134
this.namespace = in.readMap(StreamInput::readString, StreamInput::readString);
128135
}
@@ -161,6 +168,12 @@ public void writeTo(StreamOutput out) throws IOException {
161168
} else {
162169
out.writeBoolean(false);
163170
}
171+
if (structuredDataBlob != null) {
172+
out.writeBoolean(true);
173+
out.writeMap(structuredDataBlob);
174+
} else {
175+
out.writeBoolean(false);
176+
}
164177
if (namespace != null && !namespace.isEmpty()) {
165178
out.writeBoolean(true);
166179
out.writeMap(namespace, StreamOutput::writeString, StreamOutput::writeString);
@@ -211,6 +224,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
211224
if (structuredData != null) {
212225
builder.field(STRUCTURED_DATA_FIELD, structuredData);
213226
}
227+
if (structuredDataBlob != null) {
228+
builder.field(STRUCTURED_DATA_BLOB_FIELD, structuredDataBlob);
229+
}
214230
if (namespace != null && !namespace.isEmpty()) {
215231
builder.field(NAMESPACE_FIELD, namespace);
216232
}
@@ -247,6 +263,7 @@ public static MLWorkingMemory parse(XContentParser parser) throws IOException {
247263
Integer messageId = null;
248264
String binaryData = null;
249265
Map<String, Object> structuredData = null;
266+
Map<String, Object> structuredDataBlob = null;
250267
Map<String, String> namespace = null;
251268
boolean infer = false;
252269
Map<String, String> metadata = null;
@@ -284,6 +301,9 @@ public static MLWorkingMemory parse(XContentParser parser) throws IOException {
284301
case STRUCTURED_DATA_FIELD:
285302
structuredData = parser.map();
286303
break;
304+
case STRUCTURED_DATA_BLOB_FIELD:
305+
structuredDataBlob = parser.map();
306+
break;
287307
case NAMESPACE_FIELD:
288308
namespace = StringUtils.getParameterMap(parser.map());
289309
break;
@@ -322,6 +342,7 @@ public static MLWorkingMemory parse(XContentParser parser) throws IOException {
322342
.messageId(messageId)
323343
.binaryData(binaryData)
324344
.structuredData(structuredData)
345+
.structuredDataBlob(structuredDataBlob)
325346
.namespace(namespace)
326347
.infer(infer)
327348
.metadata(metadata)

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ public class MemoryContainerConstants {
6363
public static final String STRATEGY_CONFIG_FIELD = "configuration";
6464
public static final String BINARY_DATA_FIELD = "binary_data";
6565
public static final String STRUCTURED_DATA_FIELD = "structured_data";
66+
public static final String STRUCTURED_DATA_BLOB_FIELD = "structured_data_blob";
6667
public static final String NAMESPACE_SIZE_FIELD = "namespace_size";
6768
public static final String MEMORY_FIELD = "memory";
6869
public static final String MEMORY_EMBEDDING_FIELD = "memory_embedding";

‎common/src/main/java/org/opensearch/ml/common/transport/memorycontainer/memory/MLAddMemoriesInput.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PARAMETERS_FIELD;
2424
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.PAYLOAD_TYPE_FIELD;
2525
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.SESSION_ID_FIELD;
26+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.STRUCTURED_DATA_BLOB_FIELD;
2627
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.STRUCTURED_DATA_FIELD;
2728
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.TAGS_FIELD;
2829

@@ -62,6 +63,7 @@ public class MLAddMemoriesInput implements ToXContentObject, Writeable {
6263
private Integer messageId;
6364
private String binaryData;
6465
private Map<String, Object> structuredData;
66+
private Map<String, Object> structuredDataBlob;
6567

6668
// Optional fields
6769
private Map<String, String> namespace;
@@ -82,6 +84,7 @@ public MLAddMemoriesInput(
8284
Integer messageId,
8385
String binaryData,
8486
Map<String, Object> structuredData,
87+
Map<String, Object> structuredDataBlob,
8588
Map<String, String> namespace,
8689
boolean infer,
8790
Map<String, String> metadata,
@@ -99,6 +102,7 @@ public MLAddMemoriesInput(
99102
this.messageId = messageId;
100103
this.binaryData = binaryData;
101104
this.structuredData = structuredData;
105+
this.structuredDataBlob = structuredDataBlob;
102106
this.namespace = namespace;
103107
this.infer = infer; // default infer is false
104108
this.metadata = metadata;
@@ -140,6 +144,9 @@ public MLAddMemoriesInput(StreamInput in) throws IOException {
140144
if (in.readBoolean()) {
141145
this.structuredData = in.readMap();
142146
}
147+
if (in.readBoolean()) {
148+
this.structuredDataBlob = in.readMap();
149+
}
143150
if (in.readBoolean()) {
144151
this.namespace = in.readMap(StreamInput::readString, StreamInput::readString);
145152
}
@@ -179,6 +186,12 @@ public void writeTo(StreamOutput out) throws IOException {
179186
} else {
180187
out.writeBoolean(false);
181188
}
189+
if (structuredDataBlob != null) {
190+
out.writeBoolean(true);
191+
out.writeMap(structuredDataBlob);
192+
} else {
193+
out.writeBoolean(false);
194+
}
182195
if (namespace != null && !namespace.isEmpty()) {
183196
out.writeBoolean(true);
184197
out.writeMap(namespace, StreamOutput::writeString, StreamOutput::writeString);
@@ -236,6 +249,9 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
236249
if (structuredData != null) {
237250
builder.field(STRUCTURED_DATA_FIELD, structuredData);
238251
}
252+
if (structuredDataBlob != null) {
253+
builder.field(STRUCTURED_DATA_BLOB_FIELD, structuredDataBlob);
254+
}
239255
if (namespace != null && !namespace.isEmpty()) {
240256
builder.field(NAMESPACE_FIELD, namespace);
241257
builder.field(NAMESPACE_SIZE_FIELD, namespace.size());
@@ -274,6 +290,7 @@ public static MLAddMemoriesInput parse(XContentParser parser, String memoryConta
274290
Integer messageId = null;
275291
String binaryData = null;
276292
Map<String, Object> structuredData = null;
293+
Map<String, Object> structuredDataBlob = null;
277294
Map<String, String> namespace = null;
278295
boolean infer = false;
279296
Map<String, String> metadata = null;
@@ -310,6 +327,9 @@ public static MLAddMemoriesInput parse(XContentParser parser, String memoryConta
310327
case STRUCTURED_DATA_FIELD:
311328
structuredData = parser.map();
312329
break;
330+
case STRUCTURED_DATA_BLOB_FIELD:
331+
structuredDataBlob = parser.map();
332+
break;
313333
case NAMESPACE_FIELD:
314334
namespace = StringUtils.getParameterMap(parser.map());
315335
break;
@@ -347,6 +367,7 @@ public static MLAddMemoriesInput parse(XContentParser parser, String memoryConta
347367
.messageId(messageId)
348368
.binaryData(binaryData)
349369
.structuredData(structuredData)
370+
.structuredDataBlob(structuredDataBlob)
350371
.namespace(namespace)
351372
.infer(infer)
352373
.metadata(metadata)

‎common/src/main/resources/index-mappings/ml_memory_working.json‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"_meta": {
3-
"schema_version": 1
3+
"schema_version": 2
44
},
55
"properties": {
66
"owner_id": {
@@ -35,6 +35,10 @@
3535
"structured_data": {
3636
"type": "flat_object"
3737
},
38+
"structured_data_blob": {
39+
"type": "object",
40+
"enabled": false
41+
},
3842
"namespace": {
3943
"type": "flat_object"
4044
},

0 commit comments

Comments
 (0)