From db764d124786546aa9321c938b0c8c15caefa125 Mon Sep 17 00:00:00 2001 From: Owais Kazi Date: Wed, 22 Oct 2025 18:47:07 -0700 Subject: [PATCH 1/3] Fix agent type update Signed-off-by: Owais Kazi --- .../transport/agent/MLAgentUpdateInput.java | 17 +++++++++ .../agent/MLAgentUpdateInputTest.java | 37 +++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java index 5dc597d37c..58985845fc 100644 --- a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java +++ b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java @@ -51,12 +51,14 @@ public class MLAgentUpdateInput implements ToXContentObject, Writeable { public static final String MEMORY_TYPE_FIELD = "type"; public static final String MEMORY_SESSION_ID_FIELD = "session_id"; public static final String MEMORY_WINDOW_SIZE_FIELD = "window_size"; + public static final String TYPE_FIELD = "type"; public static final String APP_TYPE_FIELD = "app_type"; public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time"; @Getter private String agentId; private String name; + private String type; private String description; private String llmModelId; private Map llmParameters; @@ -73,6 +75,7 @@ public class MLAgentUpdateInput implements ToXContentObject, Writeable { public MLAgentUpdateInput( String agentId, String name, + String type, String description, String llmModelId, Map llmParameters, @@ -87,6 +90,7 @@ public MLAgentUpdateInput( ) { this.agentId = agentId; this.name = name; + this.type = type; this.description = description; this.llmModelId = llmModelId; this.llmParameters = llmParameters; @@ -105,6 +109,7 @@ public MLAgentUpdateInput(StreamInput in) throws IOException { Version streamInputVersion = in.getVersion(); agentId = in.readString(); name = in.readOptionalString(); + type = in.readOptionalString(); description = in.readOptionalString(); llmModelId = in.readOptionalString(); if (in.readBoolean()) { @@ -135,6 +140,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws if (name != null) { builder.field(AGENT_NAME_FIELD, name); } + if (type != null) { + builder.field(TYPE_FIELD, type); + } if (description != null) { builder.field(DESCRIPTION_FIELD, description); } @@ -185,6 +193,7 @@ public void writeTo(StreamOutput out) throws IOException { Version streamOutputVersion = out.getVersion(); out.writeString(agentId); out.writeOptionalString(name); + out.writeOptionalString(type); out.writeOptionalString(description); out.writeOptionalString(llmModelId); if (llmParameters != null && !llmParameters.isEmpty()) { @@ -221,6 +230,7 @@ public void writeTo(StreamOutput out) throws IOException { public static MLAgentUpdateInput parse(XContentParser parser) throws IOException { String agentId = null; String name = null; + String type = null; String description = null; String llmModelId = null; Map llmParameters = null; @@ -244,6 +254,9 @@ public static MLAgentUpdateInput parse(XContentParser parser) throws IOException case AGENT_NAME_FIELD: name = parser.text(); break; + case TYPE_FIELD: + type = parser.text(); + break; case DESCRIPTION_FIELD: description = parser.text(); break; @@ -314,6 +327,7 @@ public static MLAgentUpdateInput parse(XContentParser parser) throws IOException return new MLAgentUpdateInput( agentId, name, + type, description, llmModelId, llmParameters, @@ -329,6 +343,9 @@ public static MLAgentUpdateInput parse(XContentParser parser) throws IOException } public MLAgent toMLAgent(MLAgent originalAgent) { + if (type != null && !type.equals(originalAgent.getType())) { + throw new IllegalArgumentException("Agent type cannot be updated"); + } LLMSpec finalLlm; if (llmModelId == null && (llmParameters == null || llmParameters.isEmpty())) { finalLlm = originalAgent.getLlm(); diff --git a/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInputTest.java b/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInputTest.java index d7a143fa57..efcc6602a1 100644 --- a/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInputTest.java +++ b/common/src/test/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInputTest.java @@ -392,6 +392,7 @@ public void testParseWithAllFields() throws Exception { { "agent_id": "test-agent-id", "name": "test-agent", + "type": "flow", "description": "test description", "llm": { "model_id": "test-model-id", @@ -423,6 +424,7 @@ public void testParseWithAllFields() throws Exception { """; testParseFromJsonString(inputStr, parsedInput -> { assertEquals("test-agent", parsedInput.getName()); + assertEquals("flow", parsedInput.getType()); assertEquals("test description", parsedInput.getDescription()); assertEquals("test-model-id", parsedInput.getLlmModelId()); assertEquals(1, parsedInput.getTools().size()); @@ -959,6 +961,41 @@ public void testCombinedLLMAndMemoryPartialUpdates() { assertEquals(Integer.valueOf(10), updatedAgent.getMemory().getWindowSize()); // Updated } + @Test + public void testAgentTypeValidation() { + MLAgent originalAgent = MLAgent.builder().type(MLAgentType.FLOW.name()).name("Test Agent").build(); + + // Same type should be allowed + MLAgentUpdateInput sameTypeInput = MLAgentUpdateInput + .builder() + .agentId("test-agent-id") + .type(MLAgentType.FLOW.name()) + .name("Updated Name") + .build(); + + MLAgent updatedAgent = sameTypeInput.toMLAgent(originalAgent); + assertEquals(MLAgentType.FLOW.name(), updatedAgent.getType()); + assertEquals("Updated Name", updatedAgent.getName()); + + // Different type should throw error + MLAgentUpdateInput differentTypeInput = MLAgentUpdateInput + .builder() + .agentId("test-agent-id") + .type(MLAgentType.CONVERSATIONAL.name()) + .name("Updated Name") + .build(); + + IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> { differentTypeInput.toMLAgent(originalAgent); }); + assertEquals("Agent type cannot be updated", e.getMessage()); + + // No type provided should work (original type) + MLAgentUpdateInput noTypeInput = MLAgentUpdateInput.builder().agentId("test-agent-id").name("Updated Name").build(); + + MLAgent originalAgentType = noTypeInput.toMLAgent(originalAgent); + assertEquals(MLAgentType.FLOW.name(), originalAgentType.getType()); + assertEquals("Updated Name", originalAgentType.getName()); + } + @Test public void testStreamInputOutputWithVersion() throws IOException { MLAgentUpdateInput input = MLAgentUpdateInput From e3fe9379c48c243b1d61bef9ff71a2fc29bddb78 Mon Sep 17 00:00:00 2001 From: Owais Kazi Date: Thu, 23 Oct 2025 10:27:47 -0700 Subject: [PATCH 2/3] Added bwc checks Signed-off-by: Owais Kazi --- .../ml/common/transport/agent/MLAgentUpdateInput.java | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java index 58985845fc..9eab72ff57 100644 --- a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java +++ b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java @@ -8,6 +8,7 @@ import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken; import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD; import static org.opensearch.ml.common.CommonValue.VERSION_2_19_0; +import static org.opensearch.ml.common.CommonValue.VERSION_3_3_0; import static org.opensearch.ml.common.utils.StringUtils.getParameterMap; import java.io.IOException; @@ -109,7 +110,7 @@ public MLAgentUpdateInput(StreamInput in) throws IOException { Version streamInputVersion = in.getVersion(); agentId = in.readString(); name = in.readOptionalString(); - type = in.readOptionalString(); + type = streamInputVersion.onOrAfter(VERSION_3_3_0) ? in.readOptionalString() : null; description = in.readOptionalString(); llmModelId = in.readOptionalString(); if (in.readBoolean()) { @@ -193,7 +194,9 @@ public void writeTo(StreamOutput out) throws IOException { Version streamOutputVersion = out.getVersion(); out.writeString(agentId); out.writeOptionalString(name); - out.writeOptionalString(type); + if (streamOutputVersion.onOrAfter(VERSION_3_3_0)) { + out.writeOptionalString(type); + } out.writeOptionalString(description); out.writeOptionalString(llmModelId); if (llmParameters != null && !llmParameters.isEmpty()) { From cfc354a0f146eff6896d8b1ca44c067fd93d61aa Mon Sep 17 00:00:00 2001 From: Owais Kazi Date: Thu, 20 Nov 2025 10:37:35 -0800 Subject: [PATCH 3/3] Updated version Signed-off-by: Owais Kazi --- .../src/main/java/org/opensearch/ml/common/CommonValue.java | 1 + .../ml/common/transport/agent/MLAgentUpdateInput.java | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/common/src/main/java/org/opensearch/ml/common/CommonValue.java b/common/src/main/java/org/opensearch/ml/common/CommonValue.java index 351171ede6..026f7461e7 100644 --- a/common/src/main/java/org/opensearch/ml/common/CommonValue.java +++ b/common/src/main/java/org/opensearch/ml/common/CommonValue.java @@ -97,6 +97,7 @@ public class CommonValue { public static final Version VERSION_3_1_0 = Version.fromString("3.1.0"); public static final Version VERSION_3_2_0 = Version.fromString("3.2.0"); public static final Version VERSION_3_3_0 = Version.fromString("3.3.0"); + public static final Version VERSION_3_4_0 = Version.fromString("3.4.0"); // Connector Constants public static final String NAME_FIELD = "name"; diff --git a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java index 9eab72ff57..d9b7c5299d 100644 --- a/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java +++ b/common/src/main/java/org/opensearch/ml/common/transport/agent/MLAgentUpdateInput.java @@ -8,7 +8,7 @@ import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken; import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD; import static org.opensearch.ml.common.CommonValue.VERSION_2_19_0; -import static org.opensearch.ml.common.CommonValue.VERSION_3_3_0; +import static org.opensearch.ml.common.CommonValue.VERSION_3_4_0; import static org.opensearch.ml.common.utils.StringUtils.getParameterMap; import java.io.IOException; @@ -110,7 +110,7 @@ public MLAgentUpdateInput(StreamInput in) throws IOException { Version streamInputVersion = in.getVersion(); agentId = in.readString(); name = in.readOptionalString(); - type = streamInputVersion.onOrAfter(VERSION_3_3_0) ? in.readOptionalString() : null; + type = streamInputVersion.onOrAfter(VERSION_3_4_0) ? in.readOptionalString() : null; description = in.readOptionalString(); llmModelId = in.readOptionalString(); if (in.readBoolean()) { @@ -194,7 +194,7 @@ public void writeTo(StreamOutput out) throws IOException { Version streamOutputVersion = out.getVersion(); out.writeString(agentId); out.writeOptionalString(name); - if (streamOutputVersion.onOrAfter(VERSION_3_3_0)) { + if (streamOutputVersion.onOrAfter(VERSION_3_4_0)) { out.writeOptionalString(type); } out.writeOptionalString(description);