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

@@ -35,6 +35,9 @@
 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;
@@ -77,7 +80,6 @@ public class StringUtils {
 
     public static final String SAFE_INPUT_DESCRIPTION = "can only contain letters, numbers, spaces, and basic punctuation (.,!?():@-_'/\")";
 
-    public static final Gson gson = new Gson();
     public static final Gson PLAIN_NUMBER_GSON = new GsonBuilder()
         .serializeNulls()
         .registerTypeAdapter(Float.class, new PlainFloatAdapter())
@@ -86,6 +88,15 @@ public class StringUtils {
         .registerTypeAdapter(double.class, new PlainDoubleAdapter())
         .create();
 
+    public static final Gson gson;
+    static {
+        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()";
 
     public static final ObjectMapper MAPPER = new ObjectMapper();

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");
+    }
+}

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

@@ -215,4 +215,38 @@ 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;
+    }
 }