Skip to content

Commit 7177987

Browse files
dhrubo-osdhrubo-os
authored
Merge branch 'main' into agent_bug_fix
2 parents 64b8aaa + 563120b commit 7177987

File tree

126 files changed

+8586
-1112
lines changed

Some content is hidden

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

126 files changed

+8586
-1112
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ public class CommonValue {
4646
public static final String ML_STOP_WORDS_INDEX = ".plugins-ml-stop-words";
4747
public static final String TASK_POLLING_JOB_INDEX = ".ml_commons_task_polling_job";
4848
public static final String MCP_SESSION_MANAGEMENT_INDEX = ".plugins-ml-mcp-session-management";
49+
public static final String MCP_TOOLS_INDEX = ".plugins-ml-mcp-tools";
4950
public static final Set<String> stopWordsIndices = ImmutableSet.of(".plugins-ml-stop-words");
5051
public static final String TOOL_PARAMETERS_PREFIX = "tools.parameters.";
5152

@@ -60,6 +61,7 @@ public class CommonValue {
6061
public static final String ML_MEMORY_META_INDEX_MAPPING_PATH = "index-mappings/ml_memory_meta.json";
6162
public static final String ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH = "index-mappings/ml_memory_message.json";
6263
public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
64+
public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
6365

6466
// Calculate Versions independently of OpenSearch core version
6567
public static final Version VERSION_2_11_0 = Version.fromString("2.11.0");
@@ -98,6 +100,8 @@ public class CommonValue {
98100
public static final String MCP_TOOLS_FIELD = "tools";
99101
public static final String MCP_CONNECTORS_FIELD = "mcp_connectors";
100102
public static final String MCP_CONNECTOR_ID_FIELD = "mcp_connector_id";
103+
public static final String MCP_DEFAULT_SSE_ENDPOINT = "/sse";
104+
public static final String SSE_ENDPOINT_FILED = "sse_endpoint";
101105

102106
// TOOL Constants
103107
public static final String TOOL_INPUT_SCHEMA_FIELD = "input_schema";

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
package org.opensearch.ml.common;
77

8+
import static org.opensearch.ml.common.CommonValue.MCP_SESSION_MANAGEMENT_INDEX;
9+
import static org.opensearch.ml.common.CommonValue.MCP_TOOLS_INDEX;
810
import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX;
911
import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX_MAPPING_PATH;
1012
import static org.opensearch.ml.common.CommonValue.ML_CONFIG_INDEX;
@@ -14,6 +16,7 @@
1416
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
1517
import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX_MAPPING_PATH;
1618
import static org.opensearch.ml.common.CommonValue.ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH;
19+
import static org.opensearch.ml.common.CommonValue.ML_MCP_TOOLS_INDEX_MAPPING_PATH;
1720
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
1821
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH;
1922
import static org.opensearch.ml.common.CommonValue.ML_MEMORY_META_INDEX;
@@ -40,7 +43,8 @@ public enum MLIndex {
4043
AGENT(ML_AGENT_INDEX, false, ML_AGENT_INDEX_MAPPING_PATH),
4144
MEMORY_META(ML_MEMORY_META_INDEX, false, ML_MEMORY_META_INDEX_MAPPING_PATH),
4245
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);
46+
MCP_SESSION_MANAGEMENT(MCP_SESSION_MANAGEMENT_INDEX, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH),
47+
MCP_TOOLS(MCP_TOOLS_INDEX, false, ML_MCP_TOOLS_INDEX_MAPPING_PATH);
4448

4549
private final String indexName;
4650
// whether we use an alias for the index

‎common/src/main/java/org/opensearch/ml/common/agent/MLAgent.java

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,8 @@ public class MLAgent implements ToXContentObject, Writeable {
5151
public static final String APP_TYPE_FIELD = "app_type";
5252
public static final String IS_HIDDEN_FIELD = "is_hidden";
5353

54+
public static final int AGENT_NAME_MAX_LENGTH = 128;
55+
5456
private static final Version MINIMAL_SUPPORTED_VERSION_FOR_HIDDEN_AGENT = CommonValue.VERSION_2_13_0;
5557

5658
private String name;
@@ -102,6 +104,11 @@ private void validate() {
102104
if (name == null) {
103105
throw new IllegalArgumentException("Agent name can't be null");
104106
}
107+
if (name.isBlank() || name.length() > AGENT_NAME_MAX_LENGTH) {
108+
throw new IllegalArgumentException(
109+
String.format("Agent name cannot be empty or exceed max length of %d characters", MLAgent.AGENT_NAME_MAX_LENGTH)
110+
);
111+
}
105112
validateMLAgentType(type);
106113
if (type.equalsIgnoreCase(MLAgentType.CONVERSATIONAL.toString()) && llm == null) {
107114
throw new IllegalArgumentException("We need model information for the conversational agent type");

‎common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@ public abstract class AbstractConnector implements Connector {
3737
public static final String DESCRIPTION_FIELD = "description";
3838
public static final String PROTOCOL_FIELD = "protocol";
3939
public static final String ACTIONS_FIELD = "actions";
40-
public static final String CREDENTIAL_FIELD = "credential";
4140
public static final String PARAMETERS_FIELD = "parameters";
4241
public static final String CREATED_TIME_FIELD = "created_time";
4342
public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time";

‎common/src/main/java/org/opensearch/ml/common/connector/McpConnector.java

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616
import static org.opensearch.ml.common.CommonValue.LAST_UPDATED_TIME_FIELD;
1717
import static org.opensearch.ml.common.CommonValue.NAME_FIELD;
1818
import static org.opensearch.ml.common.CommonValue.OWNER_FIELD;
19+
import static org.opensearch.ml.common.CommonValue.PARAMETERS_FIELD;
1920
import static org.opensearch.ml.common.CommonValue.PROTOCOL_FIELD;
2021
import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
2122
import static org.opensearch.ml.common.CommonValue.URL_FIELD;
23+
import static org.opensearch.ml.common.CommonValue.VERSION_3_1_0;
2224
import static org.opensearch.ml.common.CommonValue.VERSION_FIELD;
2325
import static org.opensearch.ml.common.connector.ConnectorProtocols.MCP_SSE;
2426
import static org.opensearch.ml.common.connector.ConnectorProtocols.validateProtocol;
@@ -35,6 +37,7 @@
3537
import java.util.regex.Pattern;
3638

3739
import org.apache.commons.text.StringSubstitutor;
40+
import org.opensearch.Version;
3841
import org.opensearch.common.io.stream.BytesStreamOutput;
3942
import org.opensearch.commons.authuser.User;
4043
import org.opensearch.core.common.io.stream.StreamInput;
@@ -66,6 +69,7 @@ public class McpConnector implements Connector {
6669

6770
protected Map<String, String> credential;
6871
protected Map<String, String> decryptedHeaders;
72+
protected Map<String, String> parameters;
6973
@Setter
7074
protected Map<String, String> decryptedCredential;
7175
@Setter
@@ -101,7 +105,8 @@ public McpConnector(
101105
ConnectorClientConfig connectorClientConfig,
102106
String tenantId,
103107
String url,
104-
Map<String, String> headers
108+
Map<String, String> headers,
109+
Map<String, String> parameters
105110
) {
106111
validateProtocol(protocol);
107112
this.name = name;
@@ -116,6 +121,7 @@ public McpConnector(
116121
this.tenantId = tenantId;
117122
this.url = url;
118123
this.headers = headers;
124+
this.parameters = parameters;
119125
}
120126

121127
public McpConnector(String protocol, XContentParser parser) throws IOException {
@@ -175,6 +181,10 @@ public McpConnector(String protocol, XContentParser parser) throws IOException {
175181
headers = new HashMap<>();
176182
headers.putAll(parser.mapStrings());
177183
break;
184+
case PARAMETERS_FIELD:
185+
parameters = new HashMap<>();
186+
parameters.putAll(parser.mapStrings());
187+
break;
178188
default:
179189
parser.skipChildren();
180190
break;
@@ -229,6 +239,7 @@ public McpConnector(StreamInput input) throws IOException {
229239
}
230240

231241
private void parseFromStream(StreamInput input) throws IOException {
242+
Version streamInputVersion = input.getVersion();
232243
this.name = input.readOptionalString();
233244
this.version = input.readOptionalString();
234245
this.description = input.readOptionalString();
@@ -252,7 +263,11 @@ private void parseFromStream(StreamInput input) throws IOException {
252263
if (input.readBoolean()) {
253264
this.headers = input.readMap(s -> s.readString(), s -> s.readString());
254265
}
255-
266+
if (streamInputVersion.onOrAfter(VERSION_3_1_0)) {
267+
if (input.readBoolean()) {
268+
this.parameters = input.readMap(s -> s.readString(), s -> s.readString());
269+
}
270+
}
256271
}
257272

258273
@Override
@@ -264,6 +279,7 @@ public void removeCredential() {
264279

265280
@Override
266281
public void writeTo(StreamOutput out) throws IOException {
282+
Version streamOutputVersion = out.getVersion();
267283
out.writeString(protocol);
268284
out.writeOptionalString(name);
269285
out.writeOptionalString(version);
@@ -305,7 +321,14 @@ public void writeTo(StreamOutput out) throws IOException {
305321
} else {
306322
out.writeBoolean(false);
307323
}
308-
324+
if (streamOutputVersion.onOrAfter(VERSION_3_1_0)) {
325+
if (parameters != null) {
326+
out.writeBoolean(true);
327+
out.writeMap(parameters, StreamOutput::writeString, StreamOutput::writeString);
328+
} else {
329+
out.writeBoolean(false);
330+
}
331+
}
309332
}
310333

311334
@Override
@@ -341,6 +364,9 @@ public void update(MLCreateConnectorInput updateContent, BiFunction<String, Stri
341364
if (updateContent.getHeaders() != null) {
342365
this.headers = updateContent.getHeaders();
343366
}
367+
if (updateContent.getParameters() != null) {
368+
this.parameters = updateContent.getParameters();
369+
}
344370
}
345371

346372
@Override
@@ -393,6 +419,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
393419
if (headers != null) {
394420
builder.field(HEADERS_FIELD, headers);
395421
}
422+
if (parameters != null) {
423+
builder.field(PARAMETERS_FIELD, parameters);
424+
}
396425
builder.endObject();
397426
return builder;
398427
}
@@ -415,7 +444,7 @@ public void validateConnectorURL(List<String> urlRegexes) {
415444

416445
@Override
417446
public Map<String, String> getParameters() {
418-
throw new UnsupportedOperationException("Not implemented.");
447+
return parameters;
419448
}
420449

421450
@Override

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED;
1212
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONTROLLER_ENABLED;
1313
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_LOCAL_MODEL_ENABLED;
14+
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MCP_SERVER_ENABLED;
1415
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MULTI_TENANCY_ENABLED;
1516
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED;
1617
import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_OFFLINE_BATCH_INGESTION_ENABLED;
@@ -40,6 +41,8 @@ public class MLFeatureEnabledSetting {
4041
// This is to identify if this node is in multi-tenancy or not.
4142
private volatile Boolean isMultiTenancyEnabled;
4243

44+
private volatile Boolean isMcpServerEnabled;
45+
4346
private final List<SettingsChangeListener> listeners = new ArrayList<>();
4447

4548
public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings) {
@@ -51,6 +54,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
5154
isBatchIngestionEnabled = ML_COMMONS_OFFLINE_BATCH_INGESTION_ENABLED.get(settings);
5255
isBatchInferenceEnabled = ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED.get(settings);
5356
isMultiTenancyEnabled = ML_COMMONS_MULTI_TENANCY_ENABLED.get(settings);
57+
isMcpServerEnabled = ML_COMMONS_MCP_SERVER_ENABLED.get(settings);
5458

5559
clusterService
5660
.getClusterSettings()
@@ -69,6 +73,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
6973
clusterService
7074
.getClusterSettings()
7175
.addSettingsUpdateConsumer(ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED, it -> isBatchInferenceEnabled = it);
76+
clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_MCP_SERVER_ENABLED, it -> isMcpServerEnabled = it);
7277
}
7378

7479
/**
@@ -131,6 +136,14 @@ public boolean isMultiTenancyEnabled() {
131136
return isMultiTenancyEnabled;
132137
}
133138

139+
/**
140+
* Whether the mcp server feature is enabled. If disabled, MCP server APIs in ml-commons will be blocked.
141+
* @return where the MCP server feature is enabled.
142+
*/
143+
public boolean isMcpServerEnabled() {
144+
return isMcpServerEnabled;
145+
}
146+
134147
public void addListener(SettingsChangeListener listener) {
135148
listeners.add(listener);
136149
}

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.transport.agent;
7+
8+
import org.opensearch.action.ActionType;
9+
import org.opensearch.action.update.UpdateResponse;
10+
11+
public class MLAgentUpdateAction extends ActionType<UpdateResponse> {
12+
public static final MLAgentUpdateAction INSTANCE = new MLAgentUpdateAction();
13+
public static final String NAME = "cluster:admin/opensearch/ml/agents/update";
14+
15+
private MLAgentUpdateAction() {
16+
super(NAME, UpdateResponse::new);
17+
}
18+
}

0 commit comments

Comments
 (0)