Skip to content

Commit 15ed82d

Browse files
dhrubo-ospyek-botrithin-pullela-awszane-neoopensearch-trigger-bot[bot]
authored
rebasing main branch to the feature prompt branch. (#3840)
* [BUG] Agent Framework: Handle model response when toolUse is not accompanied by text (#3755) * fix: handle model response when toolUse is not accompanied by text Signed-off-by: Pavan Yekbote <[email protected]> * feat: add test case for parseLLMOutput Signed-off-by: Pavan Yekbote <[email protected]> --------- Signed-off-by: Pavan Yekbote <[email protected]> * [BUG] Allow user to control react agent max_interations value to prevent empty response (#3756) * fix: expose max_iteration for react Signed-off-by: Pavan Yekbote <[email protected]> * fix: defaults for agent execution and differentiate between step and step result Signed-off-by: Pavan Yekbote <[email protected]> * fix: return react agent id in agent response to expose more details Signed-off-by: Pavan Yekbote <[email protected]> * spotless Signed-off-by: Pavan Yekbote <[email protected]> * fix: remove test prompt from react system prompt Signed-off-by: Pavan Yekbote <[email protected]> * refactor: rename parameters exposed to user to executor Signed-off-by: Pavan Yekbote <[email protected]> * fix: give user complete control over planner system prompt Signed-off-by: Pavan Yekbote <[email protected]> --------- Signed-off-by: Pavan Yekbote <[email protected]> * Clean up JSM from MCP (#3773) Signed-off-by: rithin-pullela-aws <[email protected]> * [Bug] ListTools call does not return tool attributes (#3785) * initial commit for MCP server in OpenSearch (#3781) * initial commit for MCP server in OpenSearch Signed-off-by: zane-neo <[email protected]> * Make change to support register or remove tools across cluster Signed-off-by: zane-neo <[email protected]> * format code Signed-off-by: zane-neo <[email protected]> * fix UT failure caused by code change Signed-off-by: zane-neo <[email protected]> * format code Signed-off-by: zane-neo <[email protected]> * format code Signed-off-by: zane-neo <[email protected]> * add license header Signed-off-by: zane-neo <[email protected]> * fix notifications initialized not respond issue Signed-off-by: zane-neo <[email protected]> * fix minor issues and add UTs Signed-off-by: zane-neo <[email protected]> * Add more UTs Signed-off-by: zane-neo <[email protected]> --------- Signed-off-by: zane-neo <[email protected]> * Remove beta1 qualifier (#3794) (#3795) (cherry picked from commit 3f503f1) Signed-off-by: Peter Zhu <[email protected]> Co-authored-by: Peter Zhu <[email protected]> * [AUTO] Increment version to 3.1.0-SNAPSHOT (#3789) * Increment version to 3.1.0-SNAPSHOT Signed-off-by: opensearch-ci-bot <[email protected]> * Update build.gradle Signed-off-by: Peter Zhu <[email protected]> --------- Signed-off-by: opensearch-ci-bot <[email protected]> Signed-off-by: Peter Zhu <[email protected]> Co-authored-by: opensearch-ci-bot <[email protected]> Co-authored-by: Peter Zhu <[email protected]> * add release note for 3.0 (#3792) Signed-off-by: Mingshi Liu <[email protected]> * support MCP session management (#3803) * support MCP session management Signed-off-by: zane-neo <[email protected]> * Addressing comments Signed-off-by: zane-neo <[email protected]> * add feature flag for mcp server and renaming mcp connector feature flag Signed-off-by: zane-neo <[email protected]> * Address critical comments in #3781 Signed-off-by: zane-neo <[email protected]> --------- Signed-off-by: zane-neo <[email protected]> * upgrade http client to version align with core (#3809) * upgrade http client to versoin align with core Signed-off-by: zane-neo <[email protected]> * upgrade httpclient-h2 to correct versiono Signed-off-by: zane-neo <[email protected]> * use placeholder approach Signed-off-by: zane-neo <[email protected]> --------- Signed-off-by: zane-neo <[email protected]> * support customized message endpoint and addressing comments (#3810) * support customized message endpoint and addressing comments Signed-off-by: zane-neo <[email protected]> * fix UT failures Signed-off-by: zane-neo <[email protected]> * add files to jacoco exception Signed-off-by: zane-neo <[email protected]> * fix tool name issue and optimize register tool api Signed-off-by: zane-neo <[email protected]> * fix schema not parsed correctly issue and NPE when parameters is null Signed-off-by: zane-neo <[email protected]> * fix failure UT Signed-off-by: zane-neo <[email protected]> --------- Signed-off-by: zane-neo <[email protected]> * excluding circuit breaker for Agent (#3814) Signed-off-by: Dhrubo Saha <[email protected]> * change release note (#3811) * change release note Signed-off-by: zane-neo <[email protected]> * Update opensearch-ml-common.release-notes-3.0.0.0.md * Update opensearch-ml-common.release-notes-3.0.0.0.md * Update opensearch-ml-common.release-notes-3.0.0.0.md --------- Signed-off-by: zane-neo <[email protected]> Co-authored-by: Peter Zhu <[email protected]> * Downgrade MCP version to 0.9 (#3821) Signed-off-by: rithin-pullela-aws <[email protected]> * remove libs folder (#3824) Signed-off-by: Yaliang Wu <[email protected]> * add more logging to deploy/undeploy flows for better debugging (#3825) * add more logging to deploy/undeploy flows for better debugging Signed-off-by: Bhavana Goud Ramaram <[email protected]> * Fix python client not able to connect to MCP server issue (#3822) Signed-off-by: zane-neo <[email protected]> Co-authored-by: Dhrubo Saha <[email protected]> * exclude trusted connector check for hidden model (#3838) Signed-off-by: Dhrubo Saha <[email protected]> * adding tenantId to the connector executor when this is inline connector (#3837) * adding tenantId to the connector executor when this is inline connector Signed-off-by: Dhrubo Saha <[email protected]> * added more unit tests Signed-off-by: Dhrubo Saha <[email protected]> --------- Signed-off-by: Dhrubo Saha <[email protected]> --------- Signed-off-by: Pavan Yekbote <[email protected]> Signed-off-by: rithin-pullela-aws <[email protected]> Signed-off-by: zane-neo <[email protected]> Signed-off-by: Peter Zhu <[email protected]> Signed-off-by: opensearch-ci-bot <[email protected]> Signed-off-by: Mingshi Liu <[email protected]> Signed-off-by: Dhrubo Saha <[email protected]> Signed-off-by: Yaliang Wu <[email protected]> Signed-off-by: Bhavana Goud Ramaram <[email protected]> Co-authored-by: Pavan Yekbote <[email protected]> Co-authored-by: Rithin Pullela <[email protected]> Co-authored-by: zane-neo <[email protected]> Co-authored-by: opensearch-trigger-bot[bot] <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Co-authored-by: Peter Zhu <[email protected]> Co-authored-by: opensearch-ci-bot <[email protected]> Co-authored-by: Mingshi Liu <[email protected]> Co-authored-by: Yaliang Wu <[email protected]> Co-authored-by: Bhavana Goud Ramaram <[email protected]>
1 parent 9795399 commit 15ed82d

File tree

88 files changed

+5399
-216
lines changed

Some content is hidden

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

88 files changed

+5399
-216
lines changed

‎build.gradle

+2-2
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ buildscript {
1111
ext {
1212
opensearch_group = "org.opensearch"
1313
isSnapshot = "true" == System.getProperty("build.snapshot", "true")
14-
opensearch_version = System.getProperty("opensearch.version", "3.0.0-SNAPSHOT")
15-
buildVersionQualifier = System.getProperty("build.version_qualifier", "beta1")
14+
opensearch_version = System.getProperty("opensearch.version", "3.1.0-SNAPSHOT")
15+
buildVersionQualifier = System.getProperty("build.version_qualifier", "")
1616
asm_version = "9.7"
1717

1818
// 2.0.0-rc1-SNAPSHOT -> 2.0.0.0-rc1-SNAPSHOT

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

+2-7
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import java.util.Set;
99

1010
import org.opensearch.Version;
11-
import org.opensearch.common.settings.Setting;
1211

1312
import com.google.common.collect.ImmutableSet;
1413

@@ -46,6 +45,7 @@ public class CommonValue {
4645
public static final String ML_MEMORY_MESSAGE_INDEX = ".plugins-ml-memory-message";
4746
public static final String ML_STOP_WORDS_INDEX = ".plugins-ml-stop-words";
4847
public static final String TASK_POLLING_JOB_INDEX = ".ml_commons_task_polling_job";
48+
public static final String MCP_SESSION_MANAGEMENT_INDEX = ".plugins-ml-mcp-session-management";
4949
public static final Set<String> stopWordsIndices = ImmutableSet.of(".plugins-ml-stop-words");
5050
public static final String TOOL_PARAMETERS_PREFIX = "tools.parameters.";
5151

@@ -59,6 +59,7 @@ public class CommonValue {
5959
public static final String ML_AGENT_INDEX_MAPPING_PATH = "index-mappings/ml_agent.json";
6060
public static final String ML_MEMORY_META_INDEX_MAPPING_PATH = "index-mappings/ml_memory_meta.json";
6161
public static final String ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH = "index-mappings/ml_memory_message.json";
62+
public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
6263

6364
// Calculate Versions independently of OpenSearch core version
6465
public static final Version VERSION_2_11_0 = Version.fromString("2.11.0");
@@ -93,16 +94,10 @@ public class CommonValue {
9394
public static final String MCP_TOOL_DESCRIPTION_FIELD = "description";
9495
public static final String MCP_TOOL_INPUT_SCHEMA_FIELD = "inputSchema";
9596
public static final String MCP_SYNC_CLIENT = "mcp_sync_client";
96-
public static final String MCP_EXECUTOR_SERVICE = "mcp_executor_service";
9797
public static final String MCP_TOOLS_FIELD = "tools";
9898
public static final String MCP_CONNECTORS_FIELD = "mcp_connectors";
9999
public static final String MCP_CONNECTOR_ID_FIELD = "mcp_connector_id";
100100

101-
public static final Setting<Boolean> ML_COMMONS_MCP_FEATURE_ENABLED = Setting
102-
.boolSetting("plugins.ml_commons.mcp_feature_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
103-
public static final String ML_COMMONS_MCP_FEATURE_DISABLED_MESSAGE =
104-
"The MCP feature is not enabled. To enable, please update the setting " + ML_COMMONS_MCP_FEATURE_ENABLED.getKey();
105-
106101
// TOOL Constants
107102
public static final String TOOL_INPUT_SCHEMA_FIELD = "input_schema";
108103
}

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX_MAPPING_PATH;
1414
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
1515
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX_MAPPING_PATH;
16+
import static org.opensearch.ml.common.CommonValue.ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH;
1617
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
1718
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH;
1819
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_META_INDEX;
@@ -38,7 +39,8 @@ public enum MLIndex {
3839
CONTROLLER(ML_CONTROLLER_INDEX, false, ML_CONTROLLER_INDEX_MAPPING_PATH),
3940
AGENT(ML_AGENT_INDEX, false, ML_AGENT_INDEX_MAPPING_PATH),
4041
MEMORY_META(ML_MEMORY_META_INDEX, false, ML_MEMORY_META_INDEX_MAPPING_PATH),
41-
MEMORY_MESSAGE(ML_MEMORY_MESSAGE_INDEX, false, ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH);
42+
MEMORY_MESSAGE(ML_MEMORY_MESSAGE_INDEX, false, ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH),
43+
MCP_SESSION_MANAGEMENT(ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH);
4244

4345
private final String indexName;
4446
// whether we use an alias for the index

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

+29-2
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@
77
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
88

99
import java.io.IOException;
10+
import java.util.Map;
1011

12+
import org.opensearch.Version;
1113
import org.opensearch.core.common.io.stream.StreamInput;
1214
import org.opensearch.core.common.io.stream.StreamOutput;
1315
import org.opensearch.core.common.io.stream.Writeable;
@@ -25,6 +27,9 @@ public class ToolMetadata implements ToXContentObject, Writeable {
2527
public static final String TOOL_DESCRIPTION_FIELD = "description";
2628
public static final String TOOL_TYPE_FIELD = "type";
2729
public static final String TOOL_VERSION_FIELD = "version";
30+
public static final String TOOL_ATTRIBUTES_FIELD = "attributes";
31+
32+
private static final Version MINIMUM_VERSION_FOR_TOOL_ATTRIBUTES = Version.V_3_0_0;
2833

2934
@Getter
3035
private String name;
@@ -34,27 +39,43 @@ public class ToolMetadata implements ToXContentObject, Writeable {
3439
private String type;
3540
@Getter
3641
private String version;
42+
@Getter
43+
private Map<String, Object> attributes;
3744

3845
@Builder(toBuilder = true)
39-
public ToolMetadata(String name, String description, String type, String version) {
46+
public ToolMetadata(String name, String description, String type, String version, Map<String, Object> attributes) {
4047
this.name = name;
4148
this.description = description;
4249
this.type = type;
4350
this.version = version;
51+
this.attributes = attributes;
4452
}
4553

4654
public ToolMetadata(StreamInput input) throws IOException {
55+
Version byteStreamVersion = input.getVersion();
4756
name = input.readString();
4857
description = input.readString();
4958
type = input.readString();
5059
version = input.readOptionalString();
60+
if (byteStreamVersion.onOrAfter(MINIMUM_VERSION_FOR_TOOL_ATTRIBUTES) && input.readBoolean()) {
61+
attributes = input.readMap(StreamInput::readString, StreamInput::readGenericValue);
62+
}
5163
}
5264

5365
public void writeTo(StreamOutput output) throws IOException {
66+
Version byteStreamVersion = output.getVersion();
5467
output.writeString(name);
5568
output.writeString(description);
5669
output.writeString(type);
5770
output.writeOptionalString(version);
71+
if (byteStreamVersion.onOrAfter(MINIMUM_VERSION_FOR_TOOL_ATTRIBUTES)) {
72+
if (attributes != null) {
73+
output.writeBoolean(true);
74+
output.writeMap(attributes, StreamOutput::writeString, StreamOutput::writeGenericValue);
75+
} else {
76+
output.writeBoolean(false);
77+
}
78+
}
5879
}
5980

6081
@Override
@@ -70,6 +91,9 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
7091
builder.field(TOOL_TYPE_FIELD, type);
7192
}
7293
builder.field(TOOL_VERSION_FIELD, version != null ? version : "undefined");
94+
if (attributes != null) {
95+
builder.field(TOOL_ATTRIBUTES_FIELD, attributes);
96+
}
7397
builder.endObject();
7498
return builder;
7599
}
@@ -79,6 +103,7 @@ public static ToolMetadata parse(XContentParser parser) throws IOException {
79103
String description = null;
80104
String type = null;
81105
String version = null;
106+
Map<String, Object> attributes = null;
82107

83108
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
84109
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -97,12 +122,14 @@ public static ToolMetadata parse(XContentParser parser) throws IOException {
97122
break;
98123
case TOOL_VERSION_FIELD:
99124
version = parser.text();
125+
case TOOL_ATTRIBUTES_FIELD:
126+
attributes = parser.map();
100127
default:
101128
parser.skipChildren();
102129
break;
103130
}
104131
}
105-
return ToolMetadata.builder().name(name).description(description).type(type).version(version).build();
132+
return ToolMetadata.builder().name(name).description(description).type(type).version(version).attributes(attributes).build();
106133
}
107134

108135
public static ToolMetadata fromStream(StreamInput in) throws IOException {

‎common/src/main/java/org/opensearch/ml/common/settings/MLCommonsSettings.java

+9-2
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
import org.opensearch.common.settings.Setting;
1717
import org.opensearch.core.common.unit.ByteSizeUnit;
1818
import org.opensearch.core.common.unit.ByteSizeValue;
19-
import org.opensearch.ml.common.CommonValue;
2019

2120
import com.google.common.collect.ImmutableList;
2221

@@ -217,7 +216,15 @@ private MLCommonsSettings() {}
217216
public static final Setting<Boolean> ML_COMMONS_MEMORY_FEATURE_ENABLED = Setting
218217
.boolSetting("plugins.ml_commons.memory_feature_enabled", true, Setting.Property.NodeScope, Setting.Property.Dynamic);
219218

220-
public static final Setting<Boolean> ML_COMMONS_MCP_FEATURE_ENABLED = CommonValue.ML_COMMONS_MCP_FEATURE_ENABLED;
219+
public static final Setting<Boolean> ML_COMMONS_MCP_CONNECTOR_ENABLED = Setting
220+
.boolSetting("plugins.ml_commons.mcp_connector_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
221+
public static final String ML_COMMONS_MCP_CONNECTOR_DISABLED_MESSAGE =
222+
"The MCP connector is not enabled. To enable, please update the setting " + ML_COMMONS_MCP_CONNECTOR_ENABLED.getKey();
223+
224+
public static final Setting<Boolean> ML_COMMONS_MCP_SERVER_ENABLED = Setting
225+
.boolSetting("plugins.ml_commons.mcp_server_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
226+
public static final String ML_COMMONS_MCP_SERVER_DISABLED_MESSAGE =
227+
"The MCP server is not enabled. To enable, please update the setting " + ML_COMMONS_MCP_SERVER_ENABLED.getKey();
221228

222229
// Feature flag for enabling search processors for Retrieval Augmented Generation using OpenSearch and Remote Inference.
223230
public static final Setting<Boolean> ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED = Setting

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpMessageAction.java

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.action;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.action.support.clustermanager.AcknowledgedResponse;
10+
11+
public class MLMcpMessageAction extends ActionType<AcknowledgedResponse> {
12+
public static MLMcpMessageAction INSTANCE = new MLMcpMessageAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/mcp/message";
14+
15+
private MLMcpMessageAction() {
16+
super(NAME, AcknowledgedResponse::new);
17+
}
18+
19+
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpMessageDispatchAction.java

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.action;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.action.support.clustermanager.AcknowledgedResponse;
10+
11+
public class MLMcpMessageDispatchAction extends ActionType<AcknowledgedResponse> {
12+
public static MLMcpMessageDispatchAction INSTANCE = new MLMcpMessageDispatchAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/mcp/message/dispatch";
14+
15+
private MLMcpMessageDispatchAction() {
16+
super(NAME, AcknowledgedResponse::new);
17+
}
18+
19+
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpToolsRegisterOnNodesAction.java

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.action;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.ml.common.transport.mcpserver.responses.register.MLMcpRegisterNodesResponse;
10+
11+
public class MLMcpToolsRegisterOnNodesAction extends ActionType<MLMcpRegisterNodesResponse> {
12+
public static MLMcpToolsRegisterOnNodesAction INSTANCE = new MLMcpToolsRegisterOnNodesAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/register_on_nodes";
14+
15+
private MLMcpToolsRegisterOnNodesAction() {
16+
super(NAME, MLMcpRegisterNodesResponse::new);
17+
}
18+
19+
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/action/MLMcpToolsRemoveOnNodesAction.java

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.action;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.ml.common.transport.mcpserver.responses.remove.MLMcpRemoveNodesResponse;
10+
11+
public class MLMcpToolsRemoveOnNodesAction extends ActionType<MLMcpRemoveNodesResponse> {
12+
public static MLMcpToolsRemoveOnNodesAction INSTANCE = new MLMcpToolsRemoveOnNodesAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/mcp_tools/remove_on_nodes";
14+
15+
private MLMcpToolsRemoveOnNodesAction() {
16+
super(NAME, MLMcpRemoveNodesResponse::new);
17+
}
18+
19+
}

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/requests/message/MLMcpMessageRequest.java

+82
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.mcpserver.requests.message;
7+
8+
import java.io.ByteArrayInputStream;
9+
import java.io.ByteArrayOutputStream;
10+
import java.io.IOException;
11+
import java.io.UncheckedIOException;
12+
13+
import org.apache.commons.lang3.StringUtils;
14+
import org.opensearch.action.ActionRequest;
15+
import org.opensearch.action.ActionRequestValidationException;
16+
import org.opensearch.core.common.io.stream.InputStreamStreamInput;
17+
import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
18+
import org.opensearch.core.common.io.stream.StreamInput;
19+
import org.opensearch.core.common.io.stream.StreamOutput;
20+
import org.opensearch.transport.TransportRequest;
21+
22+
import lombok.Builder;
23+
import lombok.Getter;
24+
25+
@Getter
26+
public class MLMcpMessageRequest extends ActionRequest {
27+
28+
private final String nodeId;
29+
30+
private final String sessionId;
31+
32+
private final String requestBody;
33+
34+
public MLMcpMessageRequest(StreamInput in) throws IOException {
35+
super(in);
36+
this.nodeId = in.readString();
37+
this.sessionId = in.readString();
38+
this.requestBody = in.readString();
39+
if (StringUtils.isEmpty(nodeId) || StringUtils.isEmpty(sessionId) || StringUtils.isEmpty(requestBody)) {
40+
throw new IllegalStateException("nodeId, sessionId and requestBody must not be null");
41+
}
42+
}
43+
44+
@Builder
45+
public MLMcpMessageRequest(String nodeId, String sessionId, String requestBody) {
46+
super();
47+
if (StringUtils.isEmpty(nodeId) || StringUtils.isEmpty(sessionId) || StringUtils.isEmpty(requestBody)) {
48+
throw new IllegalStateException("nodeId, sessionId and requestBody must not be null");
49+
}
50+
this.nodeId = nodeId;
51+
this.sessionId = sessionId;
52+
this.requestBody = requestBody;
53+
}
54+
55+
@Override
56+
public void writeTo(StreamOutput out) throws IOException {
57+
super.writeTo(out);
58+
out.writeString(nodeId);
59+
out.writeString(sessionId);
60+
out.writeString(requestBody);
61+
}
62+
63+
public static MLMcpMessageRequest fromActionRequest(TransportRequest actionRequest) {
64+
if (actionRequest instanceof MLMcpMessageRequest) {
65+
return (MLMcpMessageRequest) actionRequest;
66+
}
67+
68+
try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
69+
actionRequest.writeTo(osso);
70+
try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
71+
return new MLMcpMessageRequest(input);
72+
}
73+
} catch (IOException e) {
74+
throw new UncheckedIOException("Failed to parse ActionRequest into MLMcpMessageRequest", e);
75+
}
76+
}
77+
78+
@Override
79+
public ActionRequestValidationException validate() {
80+
return null;
81+
}
82+
}

0 commit comments

Comments
 (0)