From 54c788a169eda3ff6d8ca5747caea9400063f76a Mon Sep 17 00:00:00 2001 From: Jackie Han Date: Mon, 22 Jan 2024 10:37:26 -0800 Subject: [PATCH] Adjust ListTool response format (#1887) * Adjust ListTool response format Signed-off-by: Jackie Han * cleanup Signed-off-by: Jackie Han --------- Signed-off-by: Jackie Han --- .../ml/common/transport/tools/MLToolsListResponse.java | 2 ++ .../transport/tools/MLToolsListResponseTests.java | 2 +- .../opensearch/ml/tools/GetToolTransportAction.java | 10 ++++++++-- 3 files changed, 11 insertions(+), 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 index 840981174e..6f2f3cad00 100644 --- 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 index aa41c53ffb..c9ab879731 100644 --- 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 index 39a400007d..07933dcf55 100644 --- 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 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);