Skip to content

Commit 46481bd

Browse files
pyek-botpyek-bot
authored
Plan, Execute and Reflect Agent Type (#3716)
1 parent 82da998 commit 46481bd

File tree

28 files changed

+1289
-77
lines changed

28 files changed

+1289
-77
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,4 +69,5 @@ public class CommonValue {
6969
public static final Version VERSION_2_17_0 = Version.fromString("2.17.0");
7070
public static final Version VERSION_2_18_0 = Version.fromString("2.18.0");
7171
public static final Version VERSION_2_19_0 = Version.fromString("2.19.0");
72+
public static final Version VERSION_3_0_0 = Version.fromString("3.0.0");
7273
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,8 @@
1010
public enum MLAgentType {
1111
FLOW,
1212
CONVERSATIONAL,
13-
CONVERSATIONAL_FLOW;
13+
CONVERSATIONAL_FLOW,
14+
PLAN_EXECUTE_AND_REFLECT;
1415

1516
public static MLAgentType from(String value) {
1617
if (value == null) {

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

Lines changed: 24 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_0_0;
1112
import static org.opensearch.ml.common.utils.StringUtils.getParameterMap;
1213

1314
import java.io.IOException;
@@ -35,13 +36,15 @@ public class MLToolSpec implements ToXContentObject {
3536
public static final String TOOL_NAME_FIELD = "name";
3637
public static final String DESCRIPTION_FIELD = "description";
3738
public static final String PARAMETERS_FIELD = "parameters";
39+
public static final String ATTRIBUTES_FIELD = "attributes";
3840
public static final String INCLUDE_OUTPUT_IN_AGENT_RESPONSE = "include_output_in_agent_response";
3941
public static final String CONFIG_FIELD = "config";
4042

4143
private String type;
4244
private String name;
4345
private String description;
4446
private Map<String, String> parameters;
47+
private Map<String, String> attributes;
4548
private boolean includeOutputInAgentResponse;
4649
private Map<String, String> configMap;
4750
@Setter
@@ -53,6 +56,7 @@ public MLToolSpec(
5356
String name,
5457
String description,
5558
Map<String, String> parameters,
59+
Map<String, String> attributes,
5660
boolean includeOutputInAgentResponse,
5761
Map<String, String> configMap,
5862
String tenantId
@@ -64,6 +68,7 @@ public MLToolSpec(
6468
this.name = name;
6569
this.description = description;
6670
this.parameters = parameters;
71+
this.attributes = attributes;
6772
this.includeOutputInAgentResponse = includeOutputInAgentResponse;
6873
this.configMap = configMap;
6974
this.tenantId = tenantId;
@@ -82,6 +87,9 @@ public MLToolSpec(StreamInput input) throws IOException {
8287
configMap = input.readMap(StreamInput::readString, StreamInput::readOptionalString);
8388
}
8489
this.tenantId = streamInputVersion.onOrAfter(VERSION_2_19_0) ? input.readOptionalString() : null;
90+
if (input.available() > 0 && input.readBoolean()) {
91+
parameters = input.readMap(StreamInput::readString, StreamInput::readOptionalString);
92+
}
8593
}
8694

8795
public void writeTo(StreamOutput out) throws IOException {
@@ -107,6 +115,14 @@ public void writeTo(StreamOutput out) throws IOException {
107115
if (streamOutputVersion.onOrAfter(VERSION_2_19_0)) {
108116
out.writeOptionalString(tenantId);
109117
}
118+
if (streamOutputVersion.onOrAfter(VERSION_3_0_0)) {
119+
if (attributes != null && !attributes.isEmpty()) {
120+
out.writeBoolean(true);
121+
out.writeMap(attributes, StreamOutput::writeString, StreamOutput::writeOptionalString);
122+
} else {
123+
out.writeBoolean(false);
124+
}
125+
}
110126
}
111127

112128
@Override
@@ -121,6 +137,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
121137
if (description != null) {
122138
builder.field(DESCRIPTION_FIELD, description);
123139
}
140+
if (attributes != null && !attributes.isEmpty()) {
141+
builder.field(ATTRIBUTES_FIELD, attributes);
142+
}
124143
if (parameters != null && !parameters.isEmpty()) {
125144
builder.field(PARAMETERS_FIELD, parameters);
126145
}
@@ -139,6 +158,7 @@ public static MLToolSpec parse(XContentParser parser) throws IOException {
139158
String type = null;
140159
String name = null;
141160
String description = null;
161+
Map<String, String> attributes = null;
142162
Map<String, String> parameters = null;
143163
boolean includeOutputInAgentResponse = false;
144164
Map<String, String> configMap = null;
@@ -159,6 +179,9 @@ public static MLToolSpec parse(XContentParser parser) throws IOException {
159179
case DESCRIPTION_FIELD:
160180
description = parser.text();
161181
break;
182+
case ATTRIBUTES_FIELD:
183+
attributes = getParameterMap(parser.map());
184+
break;
162185
case PARAMETERS_FIELD:
163186
parameters = getParameterMap(parser.map());
164187
break;
@@ -181,6 +204,7 @@ public static MLToolSpec parse(XContentParser parser) throws IOException {
181204
.type(type)
182205
.name(name)
183206
.description(description)
207+
.attributes(attributes)
184208
.parameters(parameters)
185209
.includeOutputInAgentResponse(includeOutputInAgentResponse)
186210
.configMap(configMap)

‎common/src/test/java/org/opensearch/ml/common/agent/MLAgentTest.java

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,16 @@ public class MLAgentTest {
3636
@Rule
3737
public ExpectedException exceptionRule = ExpectedException.none();
3838

39-
MLToolSpec mlToolSpec = new MLToolSpec("test", "test", "test", Collections.emptyMap(), false, Collections.emptyMap(), null);
39+
MLToolSpec mlToolSpec = new MLToolSpec(
40+
"test",
41+
"test",
42+
"test",
43+
Collections.emptyMap(),
44+
Collections.emptyMap(),
45+
false,
46+
Collections.emptyMap(),
47+
null
48+
);
4049

4150
@Test
4251
public void constructor_NullName() {
@@ -249,7 +258,19 @@ public void toXContent() throws IOException {
249258
"CONVERSATIONAL",
250259
"test",
251260
new LLMSpec("test_model", Map.of("test_key", "test_value")),
252-
List.of(new MLToolSpec("test", "test", "test", Map.of("test", "test"), false, Collections.emptyMap(), null)),
261+
List
262+
.of(
263+
new MLToolSpec(
264+
"test",
265+
"test",
266+
"test",
267+
Map.of("test", "test"),
268+
Collections.emptyMap(),
269+
false,
270+
Collections.emptyMap(),
271+
null
272+
)
273+
),
253274
Map.of("test", "test"),
254275
new MLMemorySpec("test", "123", 0),
255276
Instant.EPOCH,

‎common/src/test/java/org/opensearch/ml/common/agent/MLToolSpecTest.java

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,16 @@
2020

2121
public class MLToolSpecTest {
2222

23-
MLToolSpec spec = new MLToolSpec("test", "test", "test", Map.of("test", "test"), false, Map.of("test", "test"), null);
23+
MLToolSpec spec = new MLToolSpec(
24+
"test",
25+
"test",
26+
"test",
27+
Map.of("test", "test"),
28+
Collections.emptyMap(),
29+
false,
30+
Map.of("test", "test"),
31+
null
32+
);
2433

2534
@Test
2635
public void writeTo() throws IOException {
@@ -53,7 +62,16 @@ public void writeToEmptyConfigMap() throws IOException {
5362

5463
@Test
5564
public void writeToNullConfigMap() throws IOException {
56-
MLToolSpec spec = new MLToolSpec("test_type", "test_name", "test_desc", Map.of("test_key", "test_value"), false, null, null);
65+
MLToolSpec spec = new MLToolSpec(
66+
"test_type",
67+
"test_name",
68+
"test_desc",
69+
Map.of("test_key", "test_value"),
70+
Collections.emptyMap(),
71+
false,
72+
null,
73+
null
74+
);
5775
BytesStreamOutput output = new BytesStreamOutput();
5876
spec.writeTo(output);
5977
MLToolSpec spec1 = new MLToolSpec(output.bytes().streamInput());
@@ -170,6 +188,7 @@ public void fromStreamEmptyConfigMap() throws IOException {
170188
"test_name",
171189
"test_desc",
172190
Map.of("test_key", "test_value"),
191+
Collections.emptyMap(),
173192
false,
174193
Collections.emptyMap(),
175194
null
@@ -188,7 +207,16 @@ public void fromStreamEmptyConfigMap() throws IOException {
188207

189208
@Test
190209
public void fromStreamNullConfigMap() throws IOException {
191-
MLToolSpec spec = new MLToolSpec("test_type", "test_name", "test_desc", Map.of("test_key", "test_value"), false, null, null);
210+
MLToolSpec spec = new MLToolSpec(
211+
"test_type",
212+
"test_name",
213+
"test_desc",
214+
Map.of("test_key", "test_value"),
215+
Collections.emptyMap(),
216+
false,
217+
null,
218+
null
219+
);
192220
BytesStreamOutput output = new BytesStreamOutput();
193221
spec.writeTo(output);
194222
MLToolSpec spec1 = MLToolSpec.fromStream(output.bytes().streamInput());

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,19 @@ public void writeTo() throws IOException {
7676
MLAgentType.CONVERSATIONAL.name(),
7777
"test",
7878
new LLMSpec("test_model", Map.of("test_key", "test_value")),
79-
List.of(new MLToolSpec("test", "test", "test", Collections.emptyMap(), false, Collections.emptyMap(), null)),
79+
List
80+
.of(
81+
new MLToolSpec(
82+
"test",
83+
"test",
84+
"test",
85+
Collections.emptyMap(),
86+
Collections.emptyMap(),
87+
false,
88+
Collections.emptyMap(),
89+
null
90+
)
91+
),
8092
Map.of("test", "test"),
8193
new MLMemorySpec("test", "123", 0),
8294
Instant.EPOCH,

0 commit comments

Comments
 (0)