Skip to content

Commit 103de42

Browse files
zane-neozane-neo
committed
Address critical comments in #3781
Signed-off-by: zane-neo <[email protected]>
1 parent c0cf540 commit 103de42

File tree

13 files changed

+143
-133
lines changed

13 files changed

+143
-133
lines changed

‎common/src/main/java/org/opensearch/ml/common/transport/mcpserver/requests/register/McpTool.java

Lines changed: 35 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -30,42 +30,43 @@
3030
@Log4j2
3131
@Data
3232
public class McpTool implements ToXContentObject, Writeable {
33-
private static final String NAME_FIELD = "name";
33+
private static final String TYPE_FIELD = "type";
3434
private static final String DESCRIPTION_FIELD = "description";
35-
private static final String PARAMS_FIELD = "params";
36-
private static final String SCHEMA_FIELD = "schema";
37-
private final String name;
35+
private static final String PARAMS_FIELD = "parameters";
36+
private static final String ATTRIBUTES_FIELD = "attributes";
37+
public static final String SCHEMA_FIELD = "input_schema";
38+
private final String type;
3839
private final String description;
39-
private Map<String, Object> params;
40-
private Map<String, Object> schema;
41-
private static final String nameNotShownExceptionMessage = "name field required";
40+
private Map<String, Object> parameters;
41+
private Map<String, Object> attributes;
42+
private static final String TYPE_NOT_SHOWN_EXCEPTION_MESSAGE = "type field required";
4243

4344
public McpTool(StreamInput streamInput) throws IOException {
44-
name = streamInput.readString();
45-
if (name == null) {
46-
throw new IllegalArgumentException(nameNotShownExceptionMessage);
45+
type = streamInput.readString();
46+
if (type == null) {
47+
throw new IllegalArgumentException(TYPE_NOT_SHOWN_EXCEPTION_MESSAGE);
4748
}
4849
description = streamInput.readOptionalString();
4950
if (streamInput.readBoolean()) {
50-
params = streamInput.readMap(StreamInput::readString, StreamInput::readGenericValue);
51+
parameters = streamInput.readMap(StreamInput::readString, StreamInput::readGenericValue);
5152
}
5253
if (streamInput.readBoolean()) {
53-
schema = streamInput.readMap(StreamInput::readString, StreamInput::readGenericValue);
54+
attributes = streamInput.readMap(StreamInput::readString, StreamInput::readGenericValue);
5455
}
5556
}
5657

57-
public McpTool(String name, String description, Map<String, Object> params, Map<String, Object> schema) {
58-
if (name == null) {
59-
throw new IllegalArgumentException(nameNotShownExceptionMessage);
58+
public McpTool(String type, String description, Map<String, Object> parameters, Map<String, Object> attributes) {
59+
if (type == null) {
60+
throw new IllegalArgumentException(TYPE_NOT_SHOWN_EXCEPTION_MESSAGE);
6061
}
61-
this.name = name;
62+
this.type = type;
6263
this.description = description;
63-
this.params = params;
64-
this.schema = schema;
64+
this.parameters = parameters;
65+
this.attributes = attributes;
6566
}
6667

6768
public static McpTool parse(XContentParser parser) throws IOException {
68-
String name = null;
69+
String type = null;
6970
String description = null;
7071
Map<String, Object> params = null;
7172
Map<String, Object> schema = null;
@@ -75,8 +76,8 @@ public static McpTool parse(XContentParser parser) throws IOException {
7576
parser.nextToken();
7677

7778
switch (fieldName) {
78-
case NAME_FIELD:
79-
name = parser.text();
79+
case TYPE_FIELD:
80+
type = parser.text();
8081
break;
8182
case DESCRIPTION_FIELD:
8283
description = parser.text();
@@ -92,26 +93,26 @@ public static McpTool parse(XContentParser parser) throws IOException {
9293
break;
9394
}
9495
}
95-
if (name == null) {
96-
throw new IllegalArgumentException(nameNotShownExceptionMessage);
96+
if (type == null) {
97+
throw new IllegalArgumentException(TYPE_NOT_SHOWN_EXCEPTION_MESSAGE);
9798
}
98-
return new McpTool(name, description, params, schema);
99+
return new McpTool(type, description, params, schema);
99100
}
100101

101102
@Override
102103
public void writeTo(StreamOutput streamOutput) throws IOException {
103-
streamOutput.writeString(name);
104+
streamOutput.writeString(type);
104105
streamOutput.writeOptionalString(description);
105-
if (params != null) {
106+
if (parameters != null) {
106107
streamOutput.writeBoolean(true);
107-
streamOutput.writeMap(params, StreamOutput::writeString, StreamOutput::writeGenericValue);
108+
streamOutput.writeMap(parameters, StreamOutput::writeString, StreamOutput::writeGenericValue);
108109
} else {
109110
streamOutput.writeBoolean(false);
110111
}
111112

112-
if (schema != null) {
113+
if (attributes != null) {
113114
streamOutput.writeBoolean(true);
114-
streamOutput.writeMap(schema, StreamOutput::writeString, StreamOutput::writeGenericValue);
115+
streamOutput.writeMap(attributes, StreamOutput::writeString, StreamOutput::writeGenericValue);
115116
} else {
116117
streamOutput.writeBoolean(false);
117118
}
@@ -120,15 +121,15 @@ public void writeTo(StreamOutput streamOutput) throws IOException {
120121
@Override
121122
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params xcontentParams) throws IOException {
122123
builder.startObject();
123-
builder.field(NAME_FIELD, name);
124+
builder.field(TYPE_FIELD, type);
124125
if (description != null) {
125126
builder.field(DESCRIPTION_FIELD, description);
126127
}
127-
if (params != null && !params.isEmpty()) {
128-
builder.field(PARAMS_FIELD, params);
128+
if (parameters != null && !parameters.isEmpty()) {
129+
builder.field(PARAMS_FIELD, parameters);
129130
}
130-
if (schema != null && !schema.isEmpty()) {
131-
builder.field(SCHEMA_FIELD, schema);
131+
if (attributes != null && !attributes.isEmpty()) {
132+
builder.field(SCHEMA_FIELD, attributes);
132133
}
133134
builder.endObject();
134135
return builder;

‎common/src/test/java/org/opensearch/ml/common/transport/mcpserver/requests/register/MLMcpToolsRegisterNodeRequestTest.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ public void testStreamSerialization() throws IOException {
5959

6060
assertEquals(1, deserializedRequest.getMcpTools().getTools().size());
6161
assertEquals(timestamp, deserializedRequest.getMcpTools().getCreatedTime());
62-
assertEquals("test_tool", deserializedRequest.getMcpTools().getTools().get(0).getName());
62+
assertEquals("test_tool", deserializedRequest.getMcpTools().getTools().get(0).getType());
6363
}
6464

6565
@Test
@@ -82,7 +82,7 @@ public void writeTo(StreamOutput out) throws IOException {
8282
MLMcpToolsRegisterNodeRequest result = MLMcpToolsRegisterNodeRequest.fromActionRequest(transportRequest);
8383

8484
assertNotNull("Converted request should not be null", result);
85-
assertEquals("test_tool", result.getMcpTools().getTools().get(0).getName());
85+
assertEquals("test_tool", result.getMcpTools().getTools().get(0).getType());
8686
assertEquals(timestamp, result.getMcpTools().getCreatedTime());
8787
}
8888

‎common/src/test/java/org/opensearch/ml/common/transport/mcpserver/requests/register/MLMcpToolsRegisterNodesRequestTest.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ public void testConstructorWithNodeIds() {
4343

4444
assertArrayEquals(nodeIds, request.nodesIds());
4545
assertEquals(1, request.getMcpTools().getTools().size());
46-
assertEquals("metric_analyzer", request.getMcpTools().getTools().get(0).getName());
46+
assertEquals("metric_analyzer", request.getMcpTools().getTools().get(0).getType());
4747
}
4848

4949
@Test
@@ -58,7 +58,7 @@ public void testStreamSerialization() throws IOException {
5858

5959
assertArrayEquals(nodeIds, deserialized.nodesIds());
6060
assertEquals(sampleTools.getCreatedTime(), deserialized.getMcpTools().getCreatedTime());
61-
assertEquals("metric_analyzer", deserialized.getMcpTools().getTools().get(0).getName());
61+
assertEquals("metric_analyzer", deserialized.getMcpTools().getTools().get(0).getType());
6262
}
6363

6464
@Test
@@ -89,7 +89,7 @@ public void writeTo(StreamOutput out) throws IOException {
8989

9090
MLMcpToolsRegisterNodesRequest converted = MLMcpToolsRegisterNodesRequest.fromActionRequest(wrappedRequest);
9191

92-
assertEquals("metric_analyzer", converted.getMcpTools().getTools().get(0).getName());
92+
assertEquals("metric_analyzer", converted.getMcpTools().getTools().get(0).getType());
9393
assertArrayEquals(nodeIds, converted.nodesIds());
9494
}
9595

‎common/src/test/java/org/opensearch/ml/common/transport/mcpserver/requests/register/McpToolTest.java

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -51,16 +51,16 @@ public void setUp() {
5151

5252
@Test
5353
public void testConstructor_Success() {
54-
assertEquals(toolName, mcptool.getName());
54+
assertEquals(toolName, mcptool.getType());
5555
assertEquals(description, mcptool.getDescription());
56-
assertEquals(params, mcptool.getParams());
57-
assertEquals(schema, mcptool.getSchema());
56+
assertEquals(params, mcptool.getParameters());
57+
assertEquals(schema, mcptool.getAttributes());
5858
}
5959

6060
@Test
6161
public void testParse_AllFields() throws Exception {
62-
String jsonStr = "{\"name\":\"stock_tool\",\"description\":\"Stock data tool\","
63-
+ "\"params\":{\"exchange\":\"NYSE\"},\"schema\":{\"properties\":{\"symbol\":{\"type\":\"string\"}}}}";
62+
String jsonStr = "{\"type\":\"stock_tool\",\"description\":\"Stock data tool\","
63+
+ "\"parameters\":{\"exchange\":\"NYSE\"},\"input_schema\":{\"properties\":{\"symbol\":{\"type\":\"string\"}}}}";
6464

6565
XContentParser parser = XContentType.JSON
6666
.xContent()
@@ -72,10 +72,10 @@ public void testParse_AllFields() throws Exception {
7272
parser.nextToken();
7373

7474
McpTool parsed = McpTool.parse(parser);
75-
assertEquals("stock_tool", parsed.getName());
75+
assertEquals("stock_tool", parsed.getType());
7676
assertEquals("Stock data tool", parsed.getDescription());
77-
assertEquals(Collections.singletonMap("exchange", "NYSE"), parsed.getParams());
78-
assertTrue(parsed.getSchema().containsKey("properties"));
77+
assertEquals(Collections.singletonMap("exchange", "NYSE"), parsed.getParameters());
78+
assertTrue(parsed.getAttributes().containsKey("properties"));
7979
}
8080

8181
@Test
@@ -91,7 +91,7 @@ public void testParse_MissingNameField() throws Exception {
9191
parser.nextToken();
9292

9393
exceptionRule.expect(IllegalArgumentException.class);
94-
exceptionRule.expectMessage("name field required");
94+
exceptionRule.expectMessage("type field required");
9595
McpTool.parse(parser);
9696
}
9797

@@ -101,10 +101,10 @@ public void testToXContent_AllFields() throws Exception {
101101
mcptool.toXContent(builder, ToXContent.EMPTY_PARAMS);
102102
String jsonStr = builder.toString();
103103

104-
assertTrue(jsonStr.contains("\"name\":\"weather_tool\""));
104+
assertTrue(jsonStr.contains("\"type\":\"weather_tool\""));
105105
assertTrue(jsonStr.contains("\"description\":\"Fetch weather data\""));
106-
assertTrue(jsonStr.contains("\"params\":{\"unit\":\"celsius\"}"));
107-
assertTrue(jsonStr.contains("\"schema\":{\"type\":\"object\"}"));
106+
assertTrue(jsonStr.contains("\"parameters\":{\"unit\":\"celsius\"}"));
107+
assertTrue(jsonStr.contains("\"input_schema\":{\"type\":\"object\"}"));
108108
}
109109

110110
@Test
@@ -114,7 +114,7 @@ public void testToXContent_MinimalFields() throws Exception {
114114
minimalTool.toXContent(builder, ToXContent.EMPTY_PARAMS);
115115
String jsonStr = builder.toString();
116116

117-
assertTrue(jsonStr.contains("\"name\":\"minimal_tool\""));
117+
assertTrue(jsonStr.contains("\"type\":\"minimal_tool\""));
118118
assertFalse(jsonStr.contains("description"));
119119
assertFalse(jsonStr.contains("params"));
120120
assertFalse(jsonStr.contains("schema"));
@@ -128,10 +128,10 @@ public void testStreamInputOutput_Success() throws IOException {
128128
StreamInput input = output.bytes().streamInput();
129129
McpTool parsed = new McpTool(input);
130130

131-
assertEquals(toolName, parsed.getName());
131+
assertEquals(toolName, parsed.getType());
132132
assertEquals(description, parsed.getDescription());
133-
assertEquals(params, parsed.getParams());
134-
assertEquals(schema, parsed.getSchema());
133+
assertEquals(params, parsed.getParameters());
134+
assertEquals(schema, parsed.getAttributes());
135135
}
136136

137137
@Test
@@ -143,10 +143,10 @@ public void testStreamInputOutput_WithNullFields() throws IOException {
143143
StreamInput input = output.bytes().streamInput();
144144
McpTool parsed = new McpTool(input);
145145

146-
assertEquals("null_tool", parsed.getName());
146+
assertEquals("null_tool", parsed.getType());
147147
assertNull(parsed.getDescription());
148-
assertNull(parsed.getParams());
149-
assertNull(parsed.getSchema());
148+
assertNull(parsed.getParameters());
149+
assertNull(parsed.getAttributes());
150150
}
151151

152152
@Test
@@ -166,8 +166,8 @@ public void testComplexParameters() throws Exception {
166166
complexTool.writeTo(output);
167167
McpTool parsed = new McpTool(output.bytes().streamInput());
168168

169-
assertEquals(complexParams, parsed.getParams());
170-
assertEquals(complexSchema, parsed.getSchema());
169+
assertEquals(complexParams, parsed.getParameters());
170+
assertEquals(complexSchema, parsed.getAttributes());
171171

172172
// XContent测试
173173
XContentBuilder builder = MediaTypeRegistry.contentBuilder(XContentType.JSON);

‎common/src/test/java/org/opensearch/ml/common/transport/mcpserver/requests/register/McpToolsTest.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ public void testConstructor_Success() {
6868
public void testParse_FullData() throws Exception {
6969
String jsonStr = "{"
7070
+ "\"tools\":[{"
71-
+ "\"name\":\"weather_tool\","
71+
+ "\"type\":\"weather_tool\","
7272
+ "\"description\":\"Fetch weather data\","
7373
+ "\"params\":{\"unit\":\"celsius\"},"
7474
+ "\"schema\":{\"type\":\"object\"}"
@@ -130,7 +130,7 @@ public void testEmptyTools() throws IOException {
130130

131131
@Test
132132
public void testPartialData() throws Exception {
133-
String jsonStr = "{" + "\"tools\":[{" + "\"name\":\"minimal_tool\"" + "}]," + "\"create_time\":1745836800000" + "}";
133+
String jsonStr = "{" + "\"tools\":[{" + "\"type\":\"minimal_tool\"" + "}]," + "\"create_time\":1745836800000" + "}";
134134

135135
XContentParser parser = XContentType.JSON
136136
.xContent()

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/SearchIndexTool.java

Lines changed: 6 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -56,22 +56,12 @@ public class SearchIndexTool implements Tool {
5656
private static final String DEFAULT_DESCRIPTION =
5757
"Use this tool to search an index by providing two parameters: 'index' for the index name, and 'query' for the OpenSearch DSL formatted query. Only use this tool when both index name and DSL query is available.";
5858

59-
public static final String DEFAULT_INPUT_SCHEMA =
60-
"""
61-
{
62-
"input": {
63-
"index": {
64-
"type": "string",
65-
"description": "OpenSearch index name. for example: index1"
66-
},
67-
"query": {
68-
"type": "object",
69-
"description": "OpenSearch search index query. You need to get index mapping to write correct search query. It must be a valid OpenSearch query. Valid value:\\n{\\"query\\":{\\"match\\":{\\"population_description\\":\\"seattle 2023 population\\"}},\\"size\\":2,\\"_source\\":\\"population_description\\"}\\nInvalid value: \\n{\\"match\\":{\\"population_description\\":\\"seattle 2023 population\\"}}\\nThe value is invalid because the match not wrapped by \\"query\\".",
70-
"additionalProperties": false
71-
}
72-
}
73-
}
74-
""";
59+
public static final String DEFAULT_INPUT_SCHEMA = "{\"type\":\"object\","
60+
+ "\"properties\":{\"index\":{\"type\":\"string\",\"description\":\"OpenSearch index name. for example: index1\"},"
61+
+ "\"query\":{\"type\":\"object\",\"description\":\"OpenSearch search index query. You need to get index mapping to write correct search query. It must be a valid OpenSearch query."
62+
+ " Valid value:\\n{\\\"query\\\":{\\\"match\\\":{\\\"population_description\\\":\\\"seattle 2023 population\\\"}},\\\"size\\\":2,\\\"_source\\\":\\\"population_description\\\"}\\n"
63+
+ "Invalid value: \\n{\\\"match\\\":{\\\"population_description\\\":\\\"seattle 2023 population\\\"}}\\nThe value is invalid because the match not wrapped by \\\"query\\\".\","
64+
+ "\"additionalProperties\":false}},\"required\":[\"index\",\"query\"],\"additionalProperties\":false}";
7565

7666
private static final Gson GSON = new GsonBuilder().serializeSpecialFloatingPointValues().create();
7767

‎ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/SearchIndexToolTests.java

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -79,24 +79,17 @@ public void testGetType() {
7979
@Test
8080
@SneakyThrows
8181
public void testDefaultAttributes() {
82-
String expected =
83-
"""
84-
{
85-
"input": {
86-
"index": {
87-
"type": "string",
88-
"description": "OpenSearch index name. for example: index1"
89-
},
90-
"query": {
91-
"type": "object",
92-
"description": "OpenSearch search index query. You need to get index mapping to write correct search query. It must be a valid OpenSearch query. Valid value:\\n{\\"query\\":{\\"match\\":{\\"population_description\\":\\"seattle 2023 population\\"}},\\"size\\":2,\\"_source\\":\\"population_description\\"}\\nInvalid value: \\n{\\"match\\":{\\"population_description\\":\\"seattle 2023 population\\"}}\\nThe value is invalid because the match not wrapped by \\"query\\".",
93-
"additionalProperties": false
94-
}
95-
}
96-
}
97-
""";
9882
Map<String, Object> attributes = mockedSearchIndexTool.getAttributes();
99-
assertEquals(expected, attributes.get(INPUT_SCHEMA_FIELD));
83+
assertEquals(
84+
"{\"type\":\"object\",\"properties\":"
85+
+ "{\"index\":{\"type\":\"string\",\"description\":\"OpenSearch index name. for example: index1\"},"
86+
+ "\"query\":{\"type\":\"object\",\"description\":\"OpenSearch search index query. "
87+
+ "You need to get index mapping to write correct search query. It must be a valid OpenSearch query. "
88+
+ "Valid value:\\n{\\\"query\\\":{\\\"match\\\":{\\\"population_description\\\":\\\"seattle 2023 population\\\"}},\\\"size\\\":2,\\\"_source\\\":\\\"population_description\\\"}"
89+
+ "\\nInvalid value: \\n{\\\"match\\\":{\\\"population_description\\\":\\\"seattle 2023 population\\\"}}\\nThe value is invalid because the match not wrapped by \\\"query\\\".\","
90+
+ "\"additionalProperties\":false}},\"required\":[\"index\",\"query\"],\"additionalProperties\":false}",
91+
attributes.get(INPUT_SCHEMA_FIELD)
92+
);
10093
assertEquals(false, attributes.get(STRICT_FIELD));
10194
}
10295

‎plugin/src/main/java/org/opensearch/ml/action/mcpserver/McpAsyncServerHolder.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,6 @@ public class McpAsyncServerHolder {
3131
public static final McpAsyncServer asyncServer = McpServer
3232
.async(mcpServerTransportProvider)
3333
.capabilities(McpSchema.ServerCapabilities.builder().tools(true).logging().build())
34-
.serverInfo("OpenSearch-MCP-Server", "1.0.0")
34+
.serverInfo("OpenSearch-MCP-Server", "0.1.0")
3535
.build();
3636
}

‎plugin/src/main/java/org/opensearch/ml/action/mcpserver/TransportMcpToolsRegisterOnNodesAction.java

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

66
package org.opensearch.ml.action.mcpserver;
77

8+
import static org.opensearch.ml.common.transport.mcpserver.requests.register.McpTool.SCHEMA_FIELD;
9+
810
import java.io.IOException;
911
import java.util.List;
1012
import java.util.Map;
@@ -110,11 +112,12 @@ protected MLMcpRegisterNodeResponse nodeOperation(MLMcpToolsRegisterNodeRequest
110112
private MLMcpRegisterNodeResponse registerToolsOnNode(McpTools mcpTools) {
111113
Flux.fromStream(mcpTools.getTools().stream()).flatMap(tool -> {
112114
// check if user request contains tools that not in our system.
113-
String toolName = tool.getName();
115+
String toolName = tool.getType();
114116
Tool.Factory factory = toolFactoryWrapper.getToolsFactories().get(toolName);
115-
Tool actualTool = factory.create(tool.getParams());
117+
Tool actualTool = factory.create(tool.getParameters());
116118
Map<String, Object> mSchema = Optional
117-
.ofNullable(tool.getSchema())
119+
.ofNullable(tool.getAttributes())
120+
.map(x -> (Map<String, Object>) x.get(SCHEMA_FIELD))
118121
.orElse(
119122
Optional
120123
.ofNullable(actualTool.getAttributes())

0 commit comments

Comments
 (0)