add processor chain and add support for model and tool

Signed-off-by: Yaliang Wu <[email protected]>

Author: ylwu-amzn

common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java

@@ -111,6 +111,8 @@ public <T> void parseResponse(T response, List<ModelTensor> modelTensors, boolea
         if (response instanceof String && isJson((String) response)) {
             Map<String, Object> data = StringUtils.fromJson((String) response, ML_MAP_RESPONSE_KEY);
             modelTensors.add(ModelTensor.builder().name("response").dataAsMap(data).build());
+        } else if (response instanceof Map) {
+            modelTensors.add(ModelTensor.builder().name("response").dataAsMap((Map<String, ?>) response).build());
         } else {
             Map<String, Object> map = new HashMap<>();
             map.put("response", response);

common/src/main/java/org/opensearch/ml/common/output/model/ModelTensor.java

@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Map;
 
+import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
@@ -289,4 +290,14 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeBoolean(false);
         }
     }
+
+    @Override
+    public String toString() {
+        try {
+            return this.toXContent(JsonXContent.contentBuilder(), null).toString();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
 }

common/src/main/java/org/opensearch/ml/common/output/model/ModelTensorOutput.java

@@ -11,6 +11,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -102,4 +103,13 @@ public static ModelTensorOutput parse(XContentParser parser) throws IOException
 
         return ModelTensorOutput.builder().mlModelOutputs(mlModelOutputs).build();
     }
+
+    @Override
+    public String toString() {
+        try {
+            return this.toXContent(JsonXContent.contentBuilder(), null).toString();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }

common/src/main/java/org/opensearch/ml/common/output/model/ModelTensors.java

@@ -13,6 +13,7 @@
 import java.util.List;
 
 import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -171,4 +172,13 @@ public static ModelTensors parse(XContentParser parser) throws IOException {
         modelTensors.setStatusCode(statusCode);
         return modelTensors;
     }
+
+    @Override
+    public String toString() {
+        try {
+            return this.toXContent(JsonXContent.contentBuilder(), null).toString();
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
 }

common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java

@@ -33,11 +33,15 @@
 import org.json.JSONObject;
 import org.opensearch.OpenSearchParseException;
 import org.opensearch.action.ActionRequestValidationException;
+import org.opensearch.ml.common.output.model.ModelTensor;
+import org.opensearch.ml.common.output.model.ModelTensorOutput;
+import org.opensearch.ml.common.output.model.ModelTensors;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
 import com.google.gson.JsonElement;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonParser;
@@ -74,7 +78,12 @@ public class StringUtils {
     public static final Gson gson;
 
     static {
-        gson = new Gson();
+        gson = new GsonBuilder()
+            .disableHtmlEscaping()
+            .registerTypeAdapter(ModelTensor.class, new ToStringTypeAdapter<>(ModelTensor.class))
+            .registerTypeAdapter(ModelTensorOutput.class, new ToStringTypeAdapter<>(ModelTensorOutput.class))
+            .registerTypeAdapter(ModelTensors.class, new ToStringTypeAdapter<>(ModelTensors.class))
+            .create();
     }
     public static final String TO_STRING_FUNCTION_NAME = ".toString()";
 

common/src/main/java/org/opensearch/ml/common/utils/ToStringTypeAdapter.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.utils;
+
+import java.io.IOException;
+
+import com.google.gson.TypeAdapter;
+import com.google.gson.stream.JsonReader;
+import com.google.gson.stream.JsonWriter;
+
+public class ToStringTypeAdapter<T> extends TypeAdapter<T> {
+
+    private final Class<T> clazz;
+
+    public ToStringTypeAdapter(Class<T> clazz) {
+        this.clazz = clazz;
+    }
+
+    @Override
+    public void write(JsonWriter out, T value) throws IOException {
+        if (value == null) {
+            out.nullValue();
+            return;
+        }
+        String json = value.toString();
+        out.jsonValue(json);
+    }
+
+    @Override
+    public T read(JsonReader in) throws IOException {
+        throw new UnsupportedOperationException("Deserialization not supported");
+    }
+}

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/ToolUtils.java renamed to common/src/main/java/org/opensearch/ml/common/utils/ToolUtils.java

@@ -3,12 +3,11 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.ml.engine.tools;
+package org.opensearch.ml.common.utils;
 
 import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
 import static org.opensearch.ml.common.utils.StringUtils.gson;
 
-import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -19,8 +18,6 @@
 import org.opensearch.ml.common.output.model.ModelTensor;
 import org.opensearch.ml.common.output.model.ModelTensorOutput;
 import org.opensearch.ml.common.output.model.ModelTensors;
-import org.opensearch.ml.common.spi.tools.Tool;
-import org.opensearch.ml.common.utils.StringUtils;
 
 import com.google.gson.reflect.TypeToken;
 import com.jayway.jsonpath.JsonPath;
@@ -38,6 +35,19 @@ public class ToolUtils {
     public static final String TOOL_OUTPUT_FILTERS_FIELD = "output_filter";
     public static final String TOOL_REQUIRED_PARAMS = "required_parameters";
 
+    /**
+     * Extracts required parameters based on tool attributes specification.
+     * <p>
+     * The method performs the following:
+     * <ul>
+     *     <li>If required parameters are specified in attributes, only those parameters are extracted</li>
+     *     <li>If no required parameters are specified, all parameters are returned</li>
+     * </ul>
+     *
+     * @param parameters The input parameters map to extract from
+     * @param attributes The attributes map containing required parameter specifications
+     * @return Map containing only the required parameters
+     */
     public static Map<String, String> extractRequiredParameters(Map<String, String> parameters, Map<String, ?> attributes) {
         Map<String, String> extractedParameters = new HashMap<>();
         if (parameters == null) {
@@ -56,6 +66,26 @@ public static Map<String, String> extractRequiredParameters(Map<String, String>
         return extractedParameters;
     }
 
+    /**
+     * Extracts and processes input parameters, including handling "input" parameter.
+     * <p>
+     * The method performs the following steps:
+     * <ol>
+     *     <li>Extracts required parameters based on tool attributes specification</li>
+     *     <li>If an "input" parameter exists:
+     *         <ul>
+     *             <li>Substitutes any parameter placeholders</li>
+     *             <li>Parses it as a JSON map</li>
+     *             <li>Merges the parsed values with other parameters</li>
+     *         </ul>
+     *     </li>
+     * </ol>
+     *
+     * @param parameters The raw input parameters
+     * @param attributes The tool attributes containing parameter specifications
+     * @return Map of processed input parameters
+     * @throws IllegalArgumentException if input JSON parsing fails
+     */
     public static Map<String, String> extractInputParameters(Map<String, String> parameters, Map<String, ?> attributes) {
         Map<String, String> extractedParameters = ToolUtils.extractRequiredParameters(parameters, attributes);
         if (extractedParameters.containsKey("input")) {
@@ -73,6 +103,22 @@ public static Map<String, String> extractInputParameters(Map<String, String> par
         return extractedParameters;
     }
 
+    /**
+     * Builds the final parameter map for tool execution.
+     * <p>
+     * The method performs the following steps:
+     * <ol>
+     *     <li>Combines tool specification parameters with input parameters</li>
+     *     <li>Processes tool-specific parameter prefixes</li>
+     *     <li>Applies configuration overrides from tool specification</li>
+     *     <li>Adds tenant identification</li>
+     * </ol>
+     *
+     * @param parameters The input parameters to process
+     * @param toolSpec The tool specification containing default parameters and configuration
+     * @param tenantId The identifier for the tenant
+     * @return Map of processed parameters ready for tool execution
+     */
     public static Map<String, String> buildToolParameters(Map<String, String> parameters, MLToolSpec toolSpec, String tenantId) {
         Map<String, String> executeParams = new HashMap<>();
         if (toolSpec.getParameters() != null) {
@@ -102,30 +148,14 @@ public static Map<String, String> buildToolParameters(Map<String, String> parame
         return executeParams;
     }
 
-    public static Tool createTool(Map<String, Tool.Factory> toolFactories, Map<String, String> executeParams, MLToolSpec toolSpec) {
-        if (!toolFactories.containsKey(toolSpec.getType())) {
-            throw new IllegalArgumentException("Tool not found: " + toolSpec.getType());
-        }
-        Map<String, Object> toolParams = new HashMap<>();
-        toolParams.putAll(executeParams);
-        Map<String, Object> runtimeResources = toolSpec.getRuntimeResources();
-        if (runtimeResources != null) {
-            toolParams.putAll(runtimeResources);
-        }
-        Tool tool = toolFactories.get(toolSpec.getType()).create(toolParams);
-        String toolName = getToolName(toolSpec);
-        tool.setName(toolName);
-
-        if (toolSpec.getDescription() != null) {
-            tool.setDescription(toolSpec.getDescription());
-        }
-        if (executeParams.containsKey(toolName + ".description")) {
-            tool.setDescription(executeParams.get(toolName + ".description"));
-        }
-
-        return tool;
-    }
-
+    /**
+     * Filters tool output based on specified output filters in tool parameters.
+     * Uses JSONPath expressions to extract specific portions of the response.
+     *
+     * @param toolParams The tool parameters containing output filter specifications
+     * @param response The raw tool response to filter
+     * @return Filtered output if successful, original response if filtering fails
+     */
     public static Object filterToolOutput(Map<String, String> toolParams, Object response) {
         if (toolParams != null && toolParams.containsKey(TOOL_OUTPUT_FILTERS_FIELD)) {
             try {
@@ -142,6 +172,20 @@ public static Object filterToolOutput(Map<String, String> toolParams, Object res
         return response;
     }
 
+    /**
+     * Parses different types of tool responses into a JSON string representation.
+     * <p>
+     * Handles the following special cases:
+     * <ul>
+     *     <li>ModelTensors - converts to XContent JSON representation</li>
+     *     <li>ModelTensor - converts to XContent JSON representation</li>
+     *     <li>ModelTensorOutput - converts to XContent JSON representation</li>
+     *     <li>Other types - converts to generic JSON string</li>
+     * </ul>
+     *
+     * @param output The tool output object to parse
+     * @return JSON string representation of the output
+     */
     public static String parseResponse(Object output) {
         try {
             if (output instanceof List && !((List) output).isEmpty() && ((List) output).get(0) instanceof ModelTensors) {
@@ -159,16 +203,49 @@ public static String parseResponse(Object output) {
         }
     }
 
-    public static List<String> getToolNames(Map<String, Tool> tools) {
-        final List<String> inputTools = new ArrayList<>();
-        for (Map.Entry<String, Tool> entry : tools.entrySet()) {
-            String toolName = entry.getValue().getName();
-            inputTools.add(toolName);
-        }
-        return inputTools;
-    }
-
+    /**
+     * Gets the tool name from a tool specification.
+     * Returns the specified name if available, otherwise returns the tool type.
+     *
+     * @param toolSpec The tool specification
+     * @return The name of the tool
+     */
     public static String getToolName(MLToolSpec toolSpec) {
         return toolSpec.getName() != null ? toolSpec.getName() : toolSpec.getType();
     }
+
+    /**
+     * Converts various types of tool output into a standardized ModelTensor format.
+     * The conversion logic depends on the type of input:
+     * <ul>
+     *     <li>For Map inputs: directly uses the map as data</li>
+     *     <li>For List inputs: wraps the list in a map with "output" as the key</li>
+     *     <li>For other types: converts to JSON string and attempts to parse as map,
+     *         if parsing fails, wraps the original output in a map with "output" as the key</li>
+     * </ul>
+     *
+     * @param output    The output object to be converted. Can be a Map, List, or any other object
+     * @param outputKey The key/name to be assigned to the resulting ModelTensor
+     * @return A ModelTensor containing the formatted output data
+     */
+    public static ModelTensor convertOutputToModelTensor(Object output, String outputKey) {
+        ModelTensor modelTensor;
+        if (output instanceof Map) {
+            modelTensor = ModelTensor.builder().name(outputKey).dataAsMap((Map) output).build();
+        } else if (output instanceof List) {
+            Map<String, Object> resultMap = new HashMap<>();
+            resultMap.put("output", output);
+            modelTensor = ModelTensor.builder().name(outputKey).dataAsMap(resultMap).build();
+        } else {
+            String outputJson = StringUtils.toJson(output);
+            Map<String, Object> resultMap;
+            if (StringUtils.isJson(outputJson)) {
+                resultMap = StringUtils.fromJson(outputJson, "output");
+            } else {
+                resultMap = Map.of("output", output);
+            }
+            modelTensor = ModelTensor.builder().name(outputKey).dataAsMap(resultMap).build();
+        }
+        return modelTensor;
+    }
 }