Skip to content

Commit 211c517

Browse files
pyek-botpyek-bot
authored
[Bug fixes] Fix ListIndexTool and SearchIndexTool (#3720)
1 parent e4cbf79 commit 211c517

File tree

8 files changed

+206
-121
lines changed

8 files changed

+206
-121
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,8 @@ public MLToolSpec(StreamInput input) throws IOException {
8787
configMap = input.readMap(StreamInput::readString, StreamInput::readOptionalString);
8888
}
8989
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);
90+
if (input.getVersion().onOrAfter(VERSION_3_0_0) && input.available() > 0 && input.readBoolean()) {
91+
attributes = input.readMap(StreamInput::readString, StreamInput::readOptionalString);
9292
}
9393
}
9494

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -693,6 +693,7 @@ private static void runTool(
693693
tools.get(action).run(parameters, toolListener); // run tool
694694
}
695695
} catch (Exception e) {
696+
log.error("Failed to run tool {}", action, e);
696697
nextStepListener
697698
.onResponse(String.format(Locale.ROOT, "Failed to run the tool %s with the error message %s.", action, e.getMessage()));
698699
}

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

Lines changed: 74 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,13 @@
88
import static org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest.DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT;
99
import static org.opensearch.ml.common.utils.StringUtils.gson;
1010

11-
import java.util.Collections;
11+
import java.util.ArrayList;
1212
import java.util.HashMap;
1313
import java.util.List;
1414
import java.util.Map;
1515
import java.util.Map.Entry;
1616

17+
import org.apache.commons.lang3.StringUtils;
1718
import org.apache.logging.log4j.util.Strings;
1819
import org.opensearch.action.admin.indices.get.GetIndexRequest;
1920
import org.opensearch.action.admin.indices.get.GetIndexResponse;
@@ -34,6 +35,8 @@
3435
@ToolAnnotation(IndexMappingTool.TYPE)
3536
public class IndexMappingTool implements Tool {
3637
public static final String TYPE = "IndexMappingTool";
38+
public static final String INPUT_SCHEMA_FIELD = "input_schema";
39+
public static final String STRICT_FIELD = "strict";
3740
private static final String DEFAULT_DESCRIPTION = String
3841
.join(
3942
" ",
@@ -44,6 +47,12 @@ public class IndexMappingTool implements Tool {
4447
"The mappings are in JSON format under the key 'properties' which includes the field name as a key and a JSON object with field type under the key 'type'.",
4548
"The settings are in flattened map with 'index' as the top element and key-value pairs for each setting."
4649
);
50+
public static final String DEFAULT_INPUT_SCHEMA = "{\"type\":\"object\",\""
51+
+ "properties\":{\"index\":{\"type\":\"array\",\"description\":\"OpenSearch index name list, separated by comma. "
52+
+ "for example: [\\\"index1\\\", \\\"index2\\\"]\","
53+
+ "\"items\":{\"type\":\"string\"}}},"
54+
+ "\"required\":[\"index\"],"
55+
+ "\"additionalProperties\":false}";
4756

4857
@Setter
4958
@Getter
@@ -67,12 +76,8 @@ public IndexMappingTool(Client client) {
6776
this.client = client;
6877

6978
this.attributes = new HashMap<>();
70-
attributes
71-
.put(
72-
"input_schema",
73-
"{\"type\":\"object\",\"properties\":{\"index\":{\"type\":\"string\",\"description\":\"OpenSearch index name\"}},\"required\":[\"index\"],\"additionalProperties\":false}"
74-
);
75-
attributes.put("strict", true);
79+
attributes.put(INPUT_SCHEMA_FIELD, DEFAULT_INPUT_SCHEMA);
80+
attributes.put(STRICT_FIELD, true);
7681

7782
outputParser = new Parser<>() {
7883
@Override
@@ -86,75 +91,80 @@ public Object parse(Object o) {
8691

8792
@Override
8893
public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
89-
@SuppressWarnings("unchecked")
90-
List<String> indexList = parameters.containsKey("index")
91-
? gson.fromJson(parameters.get("index"), List.class)
92-
: Collections.emptyList();
93-
if (indexList.isEmpty()) {
94-
@SuppressWarnings("unchecked")
95-
T empty = (T) ("There were no results searching the index parameter [" + parameters.get("index") + "].");
96-
listener.onResponse(empty);
97-
return;
98-
}
94+
try {
95+
List<String> indexList = new ArrayList<>();
96+
if (StringUtils.isNotBlank(parameters.get("index"))) {
97+
indexList = gson.fromJson(parameters.get("index"), List.class);
98+
}
9999

100-
final String[] indices = indexList.toArray(Strings.EMPTY_ARRAY);
100+
if (indexList.isEmpty()) {
101+
@SuppressWarnings("unchecked")
102+
T empty = (T) ("There were no results searching the index parameter [" + parameters.get("index") + "].");
103+
listener.onResponse(empty);
104+
return;
105+
}
101106

102-
final IndicesOptions indicesOptions = IndicesOptions.strictExpand();
103-
final boolean local = Boolean.parseBoolean(parameters.get("local"));
104-
final TimeValue clusterManagerNodeTimeout = DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT;
107+
final String[] indices = indexList.toArray(Strings.EMPTY_ARRAY);
105108

106-
ActionListener<GetIndexResponse> internalListener = new ActionListener<GetIndexResponse>() {
109+
final IndicesOptions indicesOptions = IndicesOptions.strictExpand();
110+
final boolean local = Boolean.parseBoolean(parameters.getOrDefault("local", "false"));
111+
final TimeValue clusterManagerNodeTimeout = DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT;
107112

108-
@Override
109-
public void onResponse(GetIndexResponse getIndexResponse) {
110-
try {
111-
// Handle empty response
112-
if (getIndexResponse.indices().length == 0) {
113-
@SuppressWarnings("unchecked")
114-
T empty = (T) ("There were no results searching the index parameter [" + parameters.get("index") + "].");
115-
listener.onResponse(empty);
116-
return;
117-
}
118-
StringBuilder sb = new StringBuilder();
119-
for (String index : getIndexResponse.indices()) {
120-
sb.append("index: ").append(index).append("\n\n");
121-
122-
MappingMetadata mapping = getIndexResponse.mappings().get(index);
123-
if (mapping != null) {
124-
sb.append("mappings:\n");
125-
for (Entry<String, Object> entry : mapping.sourceAsMap().entrySet()) {
126-
sb.append(entry.getKey()).append("=").append(entry.getValue()).append('\n');
127-
}
128-
sb.append("\n\n");
113+
ActionListener<GetIndexResponse> internalListener = new ActionListener<GetIndexResponse>() {
114+
115+
@Override
116+
public void onResponse(GetIndexResponse getIndexResponse) {
117+
try {
118+
// Handle empty response
119+
if (getIndexResponse.indices().length == 0) {
120+
@SuppressWarnings("unchecked")
121+
T empty = (T) ("There were no results searching the index parameter [" + parameters.get("index") + "].");
122+
listener.onResponse(empty);
123+
return;
129124
}
125+
StringBuilder sb = new StringBuilder();
126+
for (String index : getIndexResponse.indices()) {
127+
sb.append("index: ").append(index).append("\n\n");
128+
129+
MappingMetadata mapping = getIndexResponse.mappings().get(index);
130+
if (mapping != null) {
131+
sb.append("mappings:\n");
132+
for (Entry<String, Object> entry : mapping.sourceAsMap().entrySet()) {
133+
sb.append(entry.getKey()).append("=").append(entry.getValue()).append('\n');
134+
}
135+
sb.append("\n\n");
136+
}
130137

131-
Settings settings = getIndexResponse.settings().get(index);
132-
if (settings != null) {
133-
sb.append("settings:\n").append(settings.toDelimitedString('\n')).append("\n\n");
138+
Settings settings = getIndexResponse.settings().get(index);
139+
if (settings != null) {
140+
sb.append("settings:\n").append(settings.toDelimitedString('\n')).append("\n\n");
141+
}
134142
}
135-
}
136143

137-
@SuppressWarnings("unchecked")
138-
T response = (T) sb.toString();
139-
listener.onResponse(response);
140-
} catch (Exception e) {
141-
onFailure(e);
144+
@SuppressWarnings("unchecked")
145+
T response = (T) sb.toString();
146+
listener.onResponse(response);
147+
} catch (Exception e) {
148+
onFailure(e);
149+
}
142150
}
143-
}
144151

145-
@Override
146-
public void onFailure(final Exception e) {
147-
listener.onFailure(e);
148-
}
152+
@Override
153+
public void onFailure(final Exception e) {
154+
listener.onFailure(e);
155+
}
149156

150-
};
151-
final GetIndexRequest getIndexRequest = new GetIndexRequest()
152-
.indices(indices)
153-
.indicesOptions(indicesOptions)
154-
.local(local)
155-
.clusterManagerNodeTimeout(clusterManagerNodeTimeout);
157+
};
158+
final GetIndexRequest getIndexRequest = new GetIndexRequest()
159+
.indices(indices)
160+
.indicesOptions(indicesOptions)
161+
.local(local)
162+
.clusterManagerNodeTimeout(clusterManagerNodeTimeout);
156163

157-
client.admin().indices().getIndex(getIndexRequest, internalListener);
164+
client.admin().indices().getIndex(getIndexRequest, internalListener);
165+
} catch (Exception e) {
166+
listener.onFailure(e);
167+
}
158168
}
159169

160170
@Override

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

Lines changed: 64 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest.DEFAULT_CLUSTER_MANAGER_NODE_TIMEOUT;
99
import static org.opensearch.ml.common.utils.StringUtils.gson;
1010

11+
import java.util.ArrayList;
1112
import java.util.Collection;
1213
import java.util.Collections;
1314
import java.util.HashMap;
@@ -23,6 +24,7 @@
2324
import java.util.stream.Collectors;
2425
import java.util.stream.StreamSupport;
2526

27+
import org.apache.commons.lang3.StringUtils;
2628
import org.apache.commons.lang3.math.NumberUtils;
2729
import org.apache.logging.log4j.util.Strings;
2830
import org.opensearch.action.admin.cluster.health.ClusterHealthRequest;
@@ -65,17 +67,24 @@
6567
@ToolAnnotation(ListIndexTool.TYPE)
6668
public class ListIndexTool implements Tool {
6769
public static final String TYPE = "ListIndexTool";
70+
public static final String INPUT_SCHEMA_FIELD = "input_schema";
71+
public static final String STRICT_FIELD = "strict";
6872
// This needs to be changed once it's changed in opensearch core in RestIndicesListAction.
6973
private static final int MAX_SUPPORTED_LIST_INDICES_PAGE_SIZE = 5000;
7074
public static final int DEFAULT_PAGE_SIZE = 100;
71-
private static final String DEFAULT_DESCRIPTION = String
75+
public static final String DEFAULT_DESCRIPTION = String
7276
.join(
7377
" ",
7478
"This tool gets index information from the OpenSearch cluster.",
75-
"It takes 2 optional arguments named `index` which is a comma-delimited list of one or more indices to get information from (default is an empty list meaning all indices),",
79+
"It takes 2 optional arguments named `indices` which is a comma-delimited list of one or more indices to get information from (default is an empty list meaning all indices),",
7680
"and `local` which means whether to return information from the local node only instead of the cluster manager node (default is false).",
7781
"The tool returns the indices information, including `health`, `status`, `index`, `uuid`, `pri`, `rep`, `docs.count`, `docs.deleted`, `store.size`, `pri.store. size `, `pri.store.size`, `pri.store`."
7882
);
83+
public static final String DEFAULT_INPUT_SCHEMA = "{\"type\":\"object\","
84+
+ "\"properties\":{\"indices\":{\"type\":\"array\",\"items\": {\"type\": \"string\"},"
85+
+ "\"description\":\"OpenSearch index name list, separated by comma. "
86+
+ "for example: [\\\"index1\\\", \\\"index2\\\"], use empty array [] to list all indices in the cluster\"}},"
87+
+ "\"additionalProperties\":false}";
7988

8089
@Setter
8190
@Getter
@@ -111,62 +120,67 @@ public Object parse(Object o) {
111120
};
112121

113122
this.attributes = new HashMap<>();
114-
attributes
115-
.put(
116-
"input_schema",
117-
"{\"type\":\"object\",\"properties\":{\"indices\":{\"type\":\"string\",\"description\":\"OpenSearch index name list, separated by comma. for example: index1, index2\"}},\"additionalProperties\":false}"
118-
);
119-
attributes.put("strict", false);
123+
attributes.put(INPUT_SCHEMA_FIELD, DEFAULT_INPUT_SCHEMA);
124+
attributes.put(STRICT_FIELD, false);
120125
}
121126

122127
@Override
123128
public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
124129
// TODO: This logic exactly matches the OpenSearch _list/indices REST action. If code at
125130
// o.o.rest/action/list/RestIndicesListAction.java changes those changes need to be reflected here
126-
@SuppressWarnings("unchecked")
127-
List<String> indexList = parameters.containsKey("indices")
128-
? gson.fromJson(parameters.get("indices"), List.class)
129-
: Collections.emptyList();
130-
final String[] indices = indexList.toArray(Strings.EMPTY_ARRAY);
131-
132-
final IndicesOptions indicesOptions = IndicesOptions.strictExpand();
133-
final boolean local = parameters.containsKey("local") && Boolean.parseBoolean(parameters.get("local"));
134-
final boolean includeUnloadedSegments = Boolean.parseBoolean(parameters.get("include_unloaded_segments"));
135-
final int pageSize = parameters.containsKey("page_size")
136-
? NumberUtils.toInt(parameters.get("page_size"), DEFAULT_PAGE_SIZE)
137-
: DEFAULT_PAGE_SIZE;
138-
final PageParams pageParams = new PageParams(null, PageParams.PARAM_ASC_SORT_VALUE, pageSize);
139-
140-
final ActionListener<Table> internalListener = ActionListener.notifyOnce(ActionListener.wrap(table -> {
141-
// Handle empty table
142-
if (table == null || table.getRows().isEmpty()) {
143-
@SuppressWarnings("unchecked")
144-
T empty = (T) ("There were no results searching the indices parameter [" + parameters.get("indices") + "].");
145-
listener.onResponse(empty);
146-
return;
131+
try {
132+
List<String> indexList = new ArrayList<>();
133+
if (StringUtils.isNotBlank(parameters.get("indices"))) {
134+
indexList = parameters.containsKey("indices")
135+
? gson.fromJson(parameters.get("indices"), List.class)
136+
: Collections.emptyList();
147137
}
148-
StringBuilder sb = new StringBuilder(
149-
// Currently using c.value which is short header matching _cat/indices
150-
// May prefer to use c.attr.get("desc") for full description
151-
table.getHeaders().stream().map(c -> c.value.toString()).collect(Collectors.joining(",", "", "\n"))
138+
final String[] indices = indexList.toArray(Strings.EMPTY_ARRAY);
139+
140+
final IndicesOptions indicesOptions = IndicesOptions.strictExpand();
141+
final boolean local = parameters.containsKey("local") && Boolean.parseBoolean(parameters.get("local"));
142+
final boolean includeUnloadedSegments = Boolean.parseBoolean(parameters.get("include_unloaded_segments"));
143+
final int pageSize = parameters.containsKey("page_size")
144+
? NumberUtils.toInt(parameters.get("page_size"), DEFAULT_PAGE_SIZE)
145+
: DEFAULT_PAGE_SIZE;
146+
final PageParams pageParams = new PageParams(null, PageParams.PARAM_ASC_SORT_VALUE, pageSize);
147+
148+
final ActionListener<Table> internalListener = ActionListener.notifyOnce(ActionListener.wrap(table -> {
149+
// Handle empty table
150+
if (table == null || table.getRows().isEmpty()) {
151+
@SuppressWarnings("unchecked")
152+
T empty = (T) ("There were no results searching the indices parameter [" + parameters.get("indices") + "].");
153+
listener.onResponse(empty);
154+
return;
155+
}
156+
StringBuilder sb = new StringBuilder(
157+
// Currently using c.value which is short header matching _cat/indices
158+
// May prefer to use c.attr.get("desc") for full description
159+
table.getHeaders().stream().map(c -> c.value.toString()).collect(Collectors.joining(",", "", "\n"))
160+
);
161+
for (List<Cell> row : table.getRows()) {
162+
sb
163+
.append(
164+
row.stream().map(c -> c.value == null ? null : c.value.toString()).collect(Collectors.joining(",", "", "\n"))
165+
);
166+
}
167+
@SuppressWarnings("unchecked")
168+
T response = (T) sb.toString();
169+
listener.onResponse(response);
170+
}, listener::onFailure));
171+
172+
fetchClusterInfoAndPages(
173+
indices,
174+
local,
175+
includeUnloadedSegments,
176+
pageParams,
177+
indicesOptions,
178+
new ConcurrentLinkedQueue<>(),
179+
internalListener
152180
);
153-
for (List<Cell> row : table.getRows()) {
154-
sb.append(row.stream().map(c -> c.value == null ? null : c.value.toString()).collect(Collectors.joining(",", "", "\n")));
155-
}
156-
@SuppressWarnings("unchecked")
157-
T response = (T) sb.toString();
158-
listener.onResponse(response);
159-
}, listener::onFailure));
160-
161-
fetchClusterInfoAndPages(
162-
indices,
163-
local,
164-
includeUnloadedSegments,
165-
pageParams,
166-
indicesOptions,
167-
new ConcurrentLinkedQueue<>(),
168-
internalListener
169-
);
181+
} catch (Exception e) {
182+
listener.onFailure(e);
183+
}
170184
}
171185

172186
private void fetchClusterInfoAndPages(

0 commit comments

Comments
 (0)