Skip to content

Commit 79626b5

Browse files
opensearch-trigger-bot[bot]owaiskazi19b4sjoo
authored
[Backport 3.3] Fix agent type update (#4481)
* Fix agent type update (#4341) * Fix agent type update Signed-off-by: Owais Kazi <owaiskazi19@gmail.com> * Added bwc checks Signed-off-by: Owais Kazi <owaiskazi19@gmail.com> * Updated version Signed-off-by: Owais Kazi <owaiskazi19@gmail.com> --------- Signed-off-by: Owais Kazi <owaiskazi19@gmail.com> (cherry picked from commit 79cbcfb) * Update common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java Co-authored-by: Owais Kazi <owaiskazi19@gmail.com> Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Update common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java Co-authored-by: Owais Kazi <owaiskazi19@gmail.com> Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Apply suggestions from code review Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Add version import in MLAgentUpdateInput Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Owais Kazi <owaiskazi19@gmail.com> Signed-off-by: Sicheng Song <sicheng.song@outlook.com> Co-authored-by: Owais Kazi <owaiskazi19@gmail.com> Co-authored-by: Sicheng Song <sicheng.song@outlook.com>
1 parent 2b7719c commit 79626b5

File tree

2 files changed

+57
-0
lines changed

2 files changed

+57
-0
lines changed

‎common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
99
import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
1010
import static org.opensearch.ml.common.CommonValue.VERSION_2_19_0;
11+
import static org.opensearch.ml.common.CommonValue.VERSION_3_3_0;
1112
import static org.opensearch.ml.common.utils.StringUtils.getParameterMap;
1213

1314
import java.io.IOException;
@@ -51,12 +52,14 @@ public class MLAgentUpdateInput implements ToXContentObject, Writeable {
5152
public static final String MEMORY_TYPE_FIELD = "type";
5253
public static final String MEMORY_SESSION_ID_FIELD = "session_id";
5354
public static final String MEMORY_WINDOW_SIZE_FIELD = "window_size";
55+
public static final String TYPE_FIELD = "type";
5456
public static final String APP_TYPE_FIELD = "app_type";
5557
public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time";
5658

5759
@Getter
5860
private String agentId;
5961
private String name;
62+
private String type;
6063
private String description;
6164
private String llmModelId;
6265
private Map<String, String> llmParameters;
@@ -73,6 +76,7 @@ public class MLAgentUpdateInput implements ToXContentObject, Writeable {
7376
public MLAgentUpdateInput(
7477
String agentId,
7578
String name,
79+
String type,
7680
String description,
7781
String llmModelId,
7882
Map<String, String> llmParameters,
@@ -87,6 +91,7 @@ public MLAgentUpdateInput(
8791
) {
8892
this.agentId = agentId;
8993
this.name = name;
94+
this.type = type;
9095
this.description = description;
9196
this.llmModelId = llmModelId;
9297
this.llmParameters = llmParameters;
@@ -105,6 +110,7 @@ public MLAgentUpdateInput(StreamInput in) throws IOException {
105110
Version streamInputVersion = in.getVersion();
106111
agentId = in.readString();
107112
name = in.readOptionalString();
113+
type = streamInputVersion.onOrAfter(VERSION_3_3_0) ? in.readOptionalString() : null;
108114
description = in.readOptionalString();
109115
llmModelId = in.readOptionalString();
110116
if (in.readBoolean()) {
@@ -135,6 +141,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
135141
if (name != null) {
136142
builder.field(AGENT_NAME_FIELD, name);
137143
}
144+
if (type != null) {
145+
builder.field(TYPE_FIELD, type);
146+
}
138147
if (description != null) {
139148
builder.field(DESCRIPTION_FIELD, description);
140149
}
@@ -185,6 +194,9 @@ public void writeTo(StreamOutput out) throws IOException {
185194
Version streamOutputVersion = out.getVersion();
186195
out.writeString(agentId);
187196
out.writeOptionalString(name);
197+
if (streamOutputVersion.onOrAfter(VERSION_3_3_0)) {
198+
out.writeOptionalString(type);
199+
}
188200
out.writeOptionalString(description);
189201
out.writeOptionalString(llmModelId);
190202
if (llmParameters != null && !llmParameters.isEmpty()) {
@@ -221,6 +233,7 @@ public void writeTo(StreamOutput out) throws IOException {
221233
public static MLAgentUpdateInput parse(XContentParser parser) throws IOException {
222234
String agentId = null;
223235
String name = null;
236+
String type = null;
224237
String description = null;
225238
String llmModelId = null;
226239
Map<String, String> llmParameters = null;
@@ -244,6 +257,9 @@ public static MLAgentUpdateInput parse(XContentParser parser) throws IOException
244257
case AGENT_NAME_FIELD:
245258
name = parser.text();
246259
break;
260+
case TYPE_FIELD:
261+
type = parser.text();
262+
break;
247263
case DESCRIPTION_FIELD:
248264
description = parser.text();
249265
break;
@@ -314,6 +330,7 @@ public static MLAgentUpdateInput parse(XContentParser parser) throws IOException
314330
return new MLAgentUpdateInput(
315331
agentId,
316332
name,
333+
type,
317334
description,
318335
llmModelId,
319336
llmParameters,
@@ -329,6 +346,9 @@ public static MLAgentUpdateInput parse(XContentParser parser) throws IOException
329346
}
330347

331348
public MLAgent toMLAgent(MLAgent originalAgent) {
349+
if (type != null && !type.equals(originalAgent.getType())) {
350+
throw new IllegalArgumentException("Agent type cannot be updated");
351+
}
332352
LLMSpec finalLlm;
333353
if (llmModelId == null && (llmParameters == null || llmParameters.isEmpty())) {
334354
finalLlm = originalAgent.getLlm();

‎common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInputTest.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -392,6 +392,7 @@ public void testParseWithAllFields() throws Exception {
392392
{
393393
"agent_id": "test-agent-id",
394394
"name": "test-agent",
395+
"type": "flow",
395396
"description": "test description",
396397
"llm": {
397398
"model_id": "test-model-id",
@@ -423,6 +424,7 @@ public void testParseWithAllFields() throws Exception {
423424
""";
424425
testParseFromJsonString(inputStr, parsedInput -> {
425426
assertEquals("test-agent", parsedInput.getName());
427+
assertEquals("flow", parsedInput.getType());
426428
assertEquals("test description", parsedInput.getDescription());
427429
assertEquals("test-model-id", parsedInput.getLlmModelId());
428430
assertEquals(1, parsedInput.getTools().size());
@@ -959,6 +961,41 @@ public void testCombinedLLMAndMemoryPartialUpdates() {
959961
assertEquals(Integer.valueOf(10), updatedAgent.getMemory().getWindowSize()); // Updated
960962
}
961963

964+
@Test
965+
public void testAgentTypeValidation() {
966+
MLAgent originalAgent = MLAgent.builder().type(MLAgentType.FLOW.name()).name("Test Agent").build();
967+
968+
// Same type should be allowed
969+
MLAgentUpdateInput sameTypeInput = MLAgentUpdateInput
970+
.builder()
971+
.agentId("test-agent-id")
972+
.type(MLAgentType.FLOW.name())
973+
.name("Updated Name")
974+
.build();
975+
976+
MLAgent updatedAgent = sameTypeInput.toMLAgent(originalAgent);
977+
assertEquals(MLAgentType.FLOW.name(), updatedAgent.getType());
978+
assertEquals("Updated Name", updatedAgent.getName());
979+
980+
// Different type should throw error
981+
MLAgentUpdateInput differentTypeInput = MLAgentUpdateInput
982+
.builder()
983+
.agentId("test-agent-id")
984+
.type(MLAgentType.CONVERSATIONAL.name())
985+
.name("Updated Name")
986+
.build();
987+
988+
IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> { differentTypeInput.toMLAgent(originalAgent); });
989+
assertEquals("Agent type cannot be updated", e.getMessage());
990+
991+
// No type provided should work (original type)
992+
MLAgentUpdateInput noTypeInput = MLAgentUpdateInput.builder().agentId("test-agent-id").name("Updated Name").build();
993+
994+
MLAgent originalAgentType = noTypeInput.toMLAgent(originalAgent);
995+
assertEquals(MLAgentType.FLOW.name(), originalAgentType.getType());
996+
assertEquals("Updated Name", originalAgentType.getName());
997+
}
998+
962999
@Test
9631000
public void testStreamInputOutputWithVersion() throws IOException {
9641001
MLAgentUpdateInput input = MLAgentUpdateInput

0 commit comments

Comments
 (0)