Skip to content

Commit 3fa3a1e

Browse files
b4sjoodhrubo-os
b4sjoo
and
dhrubo-os
authored
Implement create and add memory container API (opensearch-project#4050)
* feat: implement create and add memory container API Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * remove memory api related for pr review Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Split validation methods and clean up unused fields Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Readability enhancement on memory container api Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * add jacoco exlude rule Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Sicheng Song <sicheng.song@outlook.com> Co-authored-by: Dhrubo Saha <dhrubo@amazon.com>
1 parent 068658c commit 3fa3a1e

19 files changed

+1858
-1
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
@@ -43,6 +43,7 @@ public class CommonValue {
4343
public static final String ML_AGENT_INDEX = ".plugins-ml-agent";
4444
public static final String ML_MEMORY_META_INDEX = ".plugins-ml-memory-meta";
4545
public static final String ML_MEMORY_MESSAGE_INDEX = ".plugins-ml-memory-message";
46+
public static final String ML_MEMORY_CONTAINER_INDEX = ".plugins-ml-memory-container";
4647
public static final String ML_STOP_WORDS_INDEX = ".plugins-ml-stop-words";
4748
// index used in 2.19 to track MlTaskBatchUpdate task
4849
public static final String TASK_POLLING_JOB_INDEX = ".ml_commons_task_polling_job";
@@ -63,6 +64,7 @@ public class CommonValue {
6364
public static final String ML_AGENT_INDEX_MAPPING_PATH = "index-mappings/ml_agent.json";
6465
public static final String ML_MEMORY_META_INDEX_MAPPING_PATH = "index-mappings/ml_memory_meta.json";
6566
public static final String ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH = "index-mappings/ml_memory_message.json";
67+
public static final String ML_MEMORY_CONTAINER_INDEX_MAPPING_PATH = "index-mappings/ml_memory_container.json";
6668
public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
6769
public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
6870
public static final String ML_JOBS_INDEX_MAPPING_PATH = "index-mappings/ml_jobs.json";

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX_MAPPING_PATH;
2020
import static org.opensearch.ml.common.CommonValue.ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH;
2121
import static org.opensearch.ml.common.CommonValue.ML_MCP_TOOLS_INDEX_MAPPING_PATH;
22+
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_CONTAINER_INDEX;
23+
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_CONTAINER_INDEX_MAPPING_PATH;
2224
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
2325
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH;
2426
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_META_INDEX;
@@ -45,6 +47,7 @@ public enum MLIndex {
4547
AGENT(ML_AGENT_INDEX, false, ML_AGENT_INDEX_MAPPING_PATH),
4648
MEMORY_META(ML_MEMORY_META_INDEX, false, ML_MEMORY_META_INDEX_MAPPING_PATH),
4749
MEMORY_MESSAGE(ML_MEMORY_MESSAGE_INDEX, false, ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH),
50+
MEMORY_CONTAINER(ML_MEMORY_CONTAINER_INDEX, false, ML_MEMORY_CONTAINER_INDEX_MAPPING_PATH),
4851
MCP_SESSION_MANAGEMENT(MCP_SESSION_MANAGEMENT_INDEX, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH),
4952
MCP_TOOLS(MCP_TOOLS_INDEX, false, ML_MCP_TOOLS_INDEX_MAPPING_PATH),
5053
JOBS(ML_JOBS_INDEX, false, ML_JOBS_INDEX_MAPPING_PATH);

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

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
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.CommonValue.TENANT_ID_FIELD;
10+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.CREATED_TIME_FIELD;
11+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.DESCRIPTION_FIELD;
12+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.LAST_UPDATED_TIME_FIELD;
13+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.MEMORY_STORAGE_CONFIG_FIELD;
14+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.NAME_FIELD;
15+
import static org.opensearch.ml.common.memorycontainer.MemoryContainerConstants.OWNER_FIELD;
16+
17+
import java.io.IOException;
18+
import java.time.Instant;
19+
20+
import org.opensearch.commons.authuser.User;
21+
import org.opensearch.core.common.io.stream.StreamInput;
22+
import org.opensearch.core.common.io.stream.StreamOutput;
23+
import org.opensearch.core.common.io.stream.Writeable;
24+
import org.opensearch.core.xcontent.ToXContent;
25+
import org.opensearch.core.xcontent.ToXContentObject;
26+
import org.opensearch.core.xcontent.XContentBuilder;
27+
import org.opensearch.core.xcontent.XContentParser;
28+
29+
import lombok.Builder;
30+
import lombok.EqualsAndHashCode;
31+
import lombok.Getter;
32+
import lombok.Setter;
33+
34+
/**
35+
* ML Memory Container data model that stores metadata about memory-related objects
36+
*/
37+
@Getter
38+
@Setter
39+
@Builder
40+
@EqualsAndHashCode
41+
public class MLMemoryContainer implements ToXContentObject, Writeable {
42+
43+
private String name;
44+
private String description;
45+
private User owner;
46+
private String tenantId;
47+
private Instant createdTime;
48+
private Instant lastUpdatedTime;
49+
private MemoryStorageConfig memoryStorageConfig;
50+
51+
public MLMemoryContainer(
52+
String name,
53+
String description,
54+
User owner,
55+
String tenantId,
56+
Instant createdTime,
57+
Instant lastUpdatedTime,
58+
MemoryStorageConfig memoryStorageConfig
59+
) {
60+
this.name = name;
61+
this.description = description;
62+
this.owner = owner;
63+
this.tenantId = tenantId;
64+
this.createdTime = createdTime;
65+
this.lastUpdatedTime = lastUpdatedTime;
66+
this.memoryStorageConfig = memoryStorageConfig;
67+
}
68+
69+
public MLMemoryContainer(StreamInput input) throws IOException {
70+
this.name = input.readOptionalString();
71+
this.description = input.readOptionalString();
72+
if (input.readBoolean()) {
73+
this.owner = new User(input);
74+
}
75+
this.tenantId = input.readOptionalString();
76+
this.createdTime = input.readOptionalInstant();
77+
this.lastUpdatedTime = input.readOptionalInstant();
78+
if (input.readBoolean()) {
79+
this.memoryStorageConfig = new MemoryStorageConfig(input);
80+
}
81+
}
82+
83+
@Override
84+
public void writeTo(StreamOutput out) throws IOException {
85+
out.writeOptionalString(name);
86+
out.writeOptionalString(description);
87+
if (owner != null) {
88+
out.writeBoolean(true);
89+
owner.writeTo(out);
90+
} else {
91+
out.writeBoolean(false);
92+
}
93+
out.writeOptionalString(tenantId);
94+
out.writeOptionalInstant(createdTime);
95+
out.writeOptionalInstant(lastUpdatedTime);
96+
if (memoryStorageConfig != null) {
97+
out.writeBoolean(true);
98+
memoryStorageConfig.writeTo(out);
99+
} else {
100+
out.writeBoolean(false);
101+
}
102+
}
103+
104+
@Override
105+
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
106+
builder.startObject();
107+
if (name != null) {
108+
builder.field(NAME_FIELD, name);
109+
}
110+
if (description != null) {
111+
builder.field(DESCRIPTION_FIELD, description);
112+
}
113+
if (owner != null) {
114+
builder.field(OWNER_FIELD, owner);
115+
}
116+
if (tenantId != null) {
117+
builder.field(TENANT_ID_FIELD, tenantId);
118+
}
119+
if (createdTime != null) {
120+
builder.field(CREATED_TIME_FIELD, createdTime.toEpochMilli());
121+
}
122+
if (lastUpdatedTime != null) {
123+
builder.field(LAST_UPDATED_TIME_FIELD, lastUpdatedTime.toEpochMilli());
124+
}
125+
if (memoryStorageConfig != null) {
126+
builder.field(MEMORY_STORAGE_CONFIG_FIELD, memoryStorageConfig);
127+
}
128+
builder.endObject();
129+
return builder;
130+
}
131+
132+
public static MLMemoryContainer parse(XContentParser parser) throws IOException {
133+
String name = null;
134+
String description = null;
135+
User owner = null;
136+
String tenantId = null;
137+
Instant createdTime = null;
138+
Instant lastUpdatedTime = null;
139+
MemoryStorageConfig memoryStorageConfig = null;
140+
141+
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
142+
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
143+
String fieldName = parser.currentName();
144+
parser.nextToken();
145+
146+
switch (fieldName) {
147+
case NAME_FIELD:
148+
name = parser.text();
149+
break;
150+
case DESCRIPTION_FIELD:
151+
description = parser.text();
152+
break;
153+
case OWNER_FIELD:
154+
owner = User.parse(parser);
155+
break;
156+
case TENANT_ID_FIELD:
157+
tenantId = parser.text();
158+
break;
159+
case CREATED_TIME_FIELD:
160+
createdTime = Instant.ofEpochMilli(parser.longValue());
161+
break;
162+
case LAST_UPDATED_TIME_FIELD:
163+
lastUpdatedTime = Instant.ofEpochMilli(parser.longValue());
164+
break;
165+
case MEMORY_STORAGE_CONFIG_FIELD:
166+
memoryStorageConfig = MemoryStorageConfig.parse(parser);
167+
break;
168+
default:
169+
parser.skipChildren();
170+
break;
171+
}
172+
}
173+
174+
return MLMemoryContainer
175+
.builder()
176+
.name(name)
177+
.description(description)
178+
.owner(owner)
179+
.tenantId(tenantId)
180+
.createdTime(createdTime)
181+
.lastUpdatedTime(lastUpdatedTime)
182+
.memoryStorageConfig(memoryStorageConfig)
183+
.build();
184+
}
185+
}

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

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.memorycontainer;
7+
8+
/**
9+
* Constants for Memory Container feature
10+
*/
11+
public class MemoryContainerConstants {
12+
13+
// Field names for MemoryContainer
14+
public static final String MEMORY_CONTAINER_ID_FIELD = "memory_container_id";
15+
public static final String NAME_FIELD = "name";
16+
public static final String DESCRIPTION_FIELD = "description";
17+
public static final String OWNER_FIELD = "owner";
18+
public static final String CREATED_TIME_FIELD = "created_time";
19+
public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time";
20+
public static final String MEMORY_STORAGE_CONFIG_FIELD = "memory_storage_config";
21+
22+
// Field names for MemoryStorageConfig
23+
public static final String MEMORY_INDEX_NAME_FIELD = "memory_index_name";
24+
public static final String SEMANTIC_STORAGE_ENABLED_FIELD = "semantic_storage_enabled";
25+
public static final String EMBEDDING_MODEL_TYPE_FIELD = "embedding_model_type";
26+
public static final String EMBEDDING_MODEL_ID_FIELD = "embedding_model_id";
27+
public static final String LLM_MODEL_ID_FIELD = "llm_model_id";
28+
public static final String DIMENSION_FIELD = "dimension";
29+
public static final String MAX_INFER_SIZE_FIELD = "max_infer_size";
30+
31+
// Default values
32+
public static final int MAX_INFER_SIZE_DEFAULT_VALUE = 5;
33+
34+
// Memory index type prefixes
35+
public static final String STATIC_MEMORY_INDEX_PREFIX = "ml-static-memory-";
36+
public static final String KNN_MEMORY_INDEX_PREFIX = "ml-knn-memory-";
37+
public static final String SPARSE_MEMORY_INDEX_PREFIX = "ml-sparse-memory-";
38+
39+
// Memory data index field names
40+
public static final String USER_ID_FIELD = "user_id";
41+
public static final String AGENT_ID_FIELD = "agent_id";
42+
public static final String SESSION_ID_FIELD = "session_id";
43+
public static final String MEMORY_FIELD = "memory";
44+
public static final String MEMORY_EMBEDDING_FIELD = "memory_embedding";
45+
public static final String TAGS_FIELD = "tags";
46+
public static final String MEMORY_ID_FIELD = "memory_id";
47+
public static final String MEMORY_TYPE_FIELD = "memory_type";
48+
public static final String ROLE_FIELD = "role";
49+
50+
// Request body field names (different from storage field names)
51+
public static final String MESSAGE_FIELD = "message";
52+
public static final String MESSAGES_FIELD = "messages";
53+
public static final String CONTENT_FIELD = "content";
54+
public static final String INFER_FIELD = "infer";
55+
56+
// KNN index settings
57+
public static final String KNN_ENGINE = "lucene";
58+
public static final String KNN_SPACE_TYPE = "cosinesimil";
59+
public static final String KNN_METHOD_NAME = "hnsw";
60+
public static final int KNN_EF_SEARCH = 100;
61+
public static final int KNN_EF_CONSTRUCTION = 100;
62+
public static final int KNN_M = 16;
63+
64+
// REST API paths
65+
public static final String BASE_MEMORY_CONTAINERS_PATH = "/_plugins/_ml/memory_containers";
66+
public static final String CREATE_MEMORY_CONTAINER_PATH = BASE_MEMORY_CONTAINERS_PATH + "/_create";
67+
public static final String PARAMETER_MEMORY_CONTAINER_ID = "memory_container_id";
68+
public static final String MEMORIES_PATH = BASE_MEMORY_CONTAINERS_PATH + "/{" + PARAMETER_MEMORY_CONTAINER_ID + "}/memories";
69+
70+
// Memory types are defined in MemoryType enum
71+
72+
// Response fields
73+
public static final String STATUS_FIELD = "status";
74+
75+
// Error messages
76+
public static final String SEMANTIC_STORAGE_EMBEDDING_MODEL_TYPE_REQUIRED_ERROR =
77+
"Embedding model type is required when embedding model ID is provided";
78+
public static final String SEMANTIC_STORAGE_EMBEDDING_MODEL_ID_REQUIRED_ERROR =
79+
"Embedding model ID is required when embedding model type is provided";
80+
public static final String TEXT_EMBEDDING_DIMENSION_REQUIRED_ERROR = "Dimension is required for TEXT_EMBEDDING";
81+
public static final String SPARSE_ENCODING_DIMENSION_NOT_ALLOWED_ERROR = "Dimension is not allowed for SPARSE_ENCODING";
82+
public static final String INVALID_EMBEDDING_MODEL_TYPE_ERROR = "Embedding model type must be either TEXT_EMBEDDING or SPARSE_ENCODING";
83+
public static final String MAX_INFER_SIZE_LIMIT_ERROR = "Maximum infer size cannot exceed 10";
84+
public static final String FIELD_NOT_ALLOWED_SEMANTIC_DISABLED_ERROR = "Field %s is not allowed when semantic storage is disabled";
85+
86+
// Model validation error messages
87+
public static final String LLM_MODEL_NOT_FOUND_ERROR = "LLM model with ID %s not found";
88+
public static final String LLM_MODEL_NOT_REMOTE_ERROR = "LLM model must be a REMOTE model, found: %s";
89+
public static final String EMBEDDING_MODEL_NOT_FOUND_ERROR = "Embedding model with ID %s not found";
90+
public static final String EMBEDDING_MODEL_TYPE_MISMATCH_ERROR = "Embedding model must be of type %s or REMOTE, found: %s"; // instead
91+
public static final String INFER_REQUIRES_LLM_MODEL_ERROR = "infer=true requires llm_model_id to be configured in memory storage";
92+
}

0 commit comments

Comments
 (0)