refactor tool use formatting

Signed-off-by: Jiaping Zeng <[email protected]>

Author: jiapingzeng

common/src/main/java/org/opensearch/ml/common/agui/AGUIInputConverter.java

@@ -24,6 +24,7 @@
 import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_TOOLS;
 import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_PARAM_TOOL_CALL_RESULTS;
 import static org.opensearch.ml.common.agui.AGUIConstants.AGUI_ROLE_USER;
+import static org.opensearch.ml.common.utils.StringUtils.getStringField;
 
 import java.util.HashMap;
 import java.util.List;
@@ -124,11 +125,6 @@ public static AgentMLInput convertFromAGUIInput(String aguiInputJson, String age
         }
     }
 
-    private static String getStringField(JsonObject obj, String fieldName) {
-        JsonElement element = obj.get(fieldName);
-        return element != null && !element.isJsonNull() ? element.getAsString() : null;
-    }
-
     private static void extractUserQuestion(JsonElement messages, Map<String, String> parameters) {
         if (messages == null || !messages.isJsonArray()) {
             throw new IllegalArgumentException("Invalid AG-UI messages");

common/src/main/java/org/opensearch/ml/common/output/AGUIOutput.java

@@ -9,8 +9,6 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UncheckedIOException;
-import java.util.ArrayList;
-import java.util.List;
 
 import org.opensearch.core.action.ActionResponse;
 import org.opensearch.core.common.io.stream.InputStreamStreamInput;
@@ -21,12 +19,7 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.MLCommonsClassLoader;
-import org.opensearch.ml.common.output.AGUIOutput;
 import org.opensearch.ml.common.output.Output;
-import org.opensearch.ml.common.output.model.ModelTensorOutput;
-import org.opensearch.ml.common.output.model.ModelTensors;
-
-import com.google.gson.Gson;
 
 import lombok.Builder;
 import lombok.Getter;
@@ -82,34 +75,6 @@ public static MLExecuteTaskResponse fromActionResponse(ActionResponse actionResp
 
     @Override
     public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException {
-        if (functionName == FunctionName.AGENT && output instanceof ModelTensorOutput) {
-            ModelTensorOutput modelOutput = (ModelTensorOutput) output;
-            if (isAGUIOutput(modelOutput)) {
-                return extractAGUIOutput(modelOutput).toXContent(builder, params);
-            }
-        }
         return output.toXContent(builder, params);
     }
-
-    private boolean isAGUIOutput(ModelTensorOutput modelOutput) {
-        if (modelOutput.getMlModelOutputs() != null && modelOutput.getMlModelOutputs().size() == 1) {
-            ModelTensors modelTensors = modelOutput.getMlModelOutputs().get(0);
-            if (modelTensors.getMlModelTensors() != null && modelTensors.getMlModelTensors().size() == 1) {
-                return "ag_ui_events".equals(modelTensors.getMlModelTensors().get(0).getName());
-            }
-        }
-        return false;
-    }
-
-    private AGUIOutput extractAGUIOutput(ModelTensorOutput modelOutput) {
-        try {
-            ModelTensors modelTensors = modelOutput.getMlModelOutputs().get(0);
-            String eventsJson = modelTensors.getMlModelTensors().get(0).getResult();
-            Gson gson = new Gson();
-            List<Object> events = gson.fromJson(eventsJson, List.class);
-            return AGUIOutput.builder().events(events).build();
-        } catch (Exception e) {
-            return AGUIOutput.builder().events(new ArrayList<>()).build();
-        }
-    }
 }

common/src/main/java/org/opensearch/ml/common/transport/execute/MLExecuteTaskResponse.java

@@ -726,4 +726,9 @@ public Float read(JsonReader in) throws IOException {
             return f;
         }
     }
+
+    public static String getStringField(JsonObject obj, String fieldName) {
+        JsonElement element = obj.get(fieldName);
+        return element != null && !element.isJsonNull() ? element.getAsString() : null;
+    }
 }