Adjust ListTool response format (opensearch-project#1887)

* Adjust ListTool response format Signed-off-by: Jackie Han <[email protected]> * cleanup Signed-off-by: Jackie Han <[email protected]> --------- Signed-off-by: Jackie Han <[email protected]>
rbhavna · Jan 22, 2024 · 54c788a · 54c788a
1 parent 0232f1a
commit 54c788a
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/transport/tools/MLToolsListResponse.java b/common/src/main/java/org/opensearch/ml/common/transport/tools/MLToolsListResponse.java
@@ -45,6 +45,7 @@ public void writeTo(StreamOutput out) throws IOException {
 
     @Override
     public XContentBuilder toXContent(XContentBuilder xContentBuilder, ToXContent.Params params) throws IOException {
+        xContentBuilder.startArray();
         for (ToolMetadata toolMetadata : toolMetadataList) {
             xContentBuilder.startObject();
             xContentBuilder.field(ToolMetadata.TOOL_NAME_FIELD, toolMetadata.getName());
@@ -53,6 +54,7 @@ public XContentBuilder toXContent(XContentBuilder xContentBuilder, ToXContent.Pa
             xContentBuilder.field(ToolMetadata.TOOL_VERSION_FIELD, toolMetadata.getVersion() != null ? toolMetadata.getVersion() : "undefined");
             xContentBuilder.endObject();
         }
+        xContentBuilder.endArray();
         return xContentBuilder;
     }
 

diff --git a/common/src/test/java/org/opensearch/ml/common/transport/tools/MLToolsListResponseTests.java b/common/src/test/java/org/opensearch/ml/common/transport/tools/MLToolsListResponseTests.java
@@ -67,7 +67,7 @@ public void toXContentTest() throws IOException {
         mlToolsListResponse.toXContent(builder, ToXContent.EMPTY_PARAMS);
         assertNotNull(builder);
         String jsonStr = builder.toString();
-        assertEquals("{\"name\":\"SearchWikipediaTool\",\"description\":\"Useful when you need to use this tool to search general knowledge on wikipedia.\",\"type\":\"SearchWikipediaTool\",\"version\":\"undefined\"} {\"name\":\"MathTool\",\"description\":\"Use this tool to calculate any math problem.\",\"type\":\"MathTool\",\"version\":\"test\"}", jsonStr);
+        assertEquals("[{\"name\":\"SearchWikipediaTool\",\"description\":\"Useful when you need to use this tool to search general knowledge on wikipedia.\",\"type\":\"SearchWikipediaTool\",\"version\":\"undefined\"},{\"name\":\"MathTool\",\"description\":\"Use this tool to calculate any math problem.\",\"type\":\"MathTool\",\"version\":\"test\"}]", jsonStr);
     }
 
     @Test

diff --git a/plugin/src/main/java/org/opensearch/ml/tools/GetToolTransportAction.java b/plugin/src/main/java/org/opensearch/ml/tools/GetToolTransportAction.java
@@ -6,13 +6,14 @@
 package org.opensearch.ml.tools;
 
 import java.util.List;
-import java.util.NoSuchElementException;
 
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.ml.common.ToolMetadata;
 import org.opensearch.ml.common.transport.tools.*;
 import org.opensearch.tasks.Task;
@@ -45,7 +46,12 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLTool
                 .stream()
                 .filter(tool -> tool.getName().equals(toolName))
                 .findFirst()
-                .orElseThrow(NoSuchElementException::new);
+                .orElseThrow(
+                    () -> new OpenSearchStatusException(
+                        "Failed to find tool information with the provided tool name: " + toolName,
+                        RestStatus.NOT_FOUND
+                    )
+                );
             listener.onResponse(MLToolGetResponse.builder().toolMetadata(theTool).build());
         } catch (Exception e) {
             log.error("Failed to get tool", e);