change argument parsing

Signed-off-by: xinyual <xinyual@amazon.com>

Author: xinyual

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java

@@ -326,7 +326,7 @@ private void runReAct(
                     }
                     String thought = String.valueOf(dataAsMap.get("thought"));
                     String action = String.valueOf(dataAsMap.get("action"));
-                    String actionInput = gson.toJson(dataAsMap.get("action_input"));
+                    String actionInput = parseInputFromLLMReturn(dataAsMap);
                     String finalAnswer = (String) dataAsMap.get("final_answer");
                     if (!dataAsMap.containsKey("thought")) {
                         String response = (String) dataAsMap.get("response");
@@ -337,7 +337,7 @@ private void runReAct(
                             Map map = gson.fromJson(jsonBlock, Map.class);
                             thought = String.valueOf(map.get("thought"));
                             action = String.valueOf(map.get("action"));
-                            actionInput = gson.toJson(map.get("action_input"));
+                            actionInput = parseInputFromLLMReturn(map);
                             finalAnswer = (String) map.get("final_answer");
                         } else {
                             finalAnswer = response;
@@ -675,4 +675,16 @@ private String outputToOutputString(Object output) throws PrivilegedActionExcept
         return outputString;
     }
 
+    private String parseInputFromLLMReturn(Map<String, ?> retMap){
+        Object actionInput = retMap.get("action_input");
+        if (actionInput instanceof Map)
+        {
+            return gson.toJson(actionInput);
+        }
+        else {
+            return String.valueOf(actionInput);
+        }
+
+    }
+
 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/AgentTool.java

@@ -53,7 +53,7 @@ public AgentTool(Client client, String agentId) {
 
     @Override
     public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
-        parameters = extractFromChatParameters(parameters);
+        extractFromChatParameters(parameters);
         AgentMLInput agentMLInput = AgentMLInput
             .AgentMLInputBuilder()
             .agentId(agentId)
@@ -139,15 +139,14 @@ public String getDefaultVersion() {
         }
     }
 
-    private Map<String, String> extractFromChatParameters(Map<String, String> parameters) {
+    private void extractFromChatParameters(Map<String, String> parameters) {
         if (parameters.containsKey("input")) {
             try {
                 Map<String, String> chatParameters = gson.fromJson(parameters.get("input"), Map.class);
                 parameters.putAll(chatParameters);
             } finally {
-                return parameters;
+                return ;
             }
         }
-        return parameters;
     }
 }