Handle ModelTensorOutput type while parsing response

Signed-off-by: Arjun kumar Giri <arjung@amazon.com>
arjunkumargiri · Dec 8, 2023 · 5cc51ef · 5cc51ef
1 parent fcbfc7e
commit 5cc51ef
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java
@@ -176,6 +176,8 @@ private String parseResponse(Object output) {
 
         } else if (output instanceof ModelTensor) {
             return toJson(((ModelTensor) output).getDataAsMap());
+        } else if (output instanceof ModelTensorOutput)  {
+            return toJson(output.toString());
         } else {
             if (output instanceof String) {
                 return (String) output;