Avoid NPE

Signed-off-by: Heng Qian <[email protected]>
opensearch-project · Jul 17, 2024 · 39d1bcf · 39d1bcf
1 parent c28de3b
commit 39d1bcf
Showing 1 changed file with 13 additions and 7 deletions.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/MLModelTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/MLModelTool.java
@@ -8,6 +8,7 @@
 import java.util.List;
 import java.util.Map;
 
+import java.util.NoSuchElementException;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.client.Client;
 import org.opensearch.core.action.ActionListener;
@@ -67,13 +68,18 @@ public MLModelTool(Client client, String modelId, String responseField) {
         this.responseField = responseField;
 
         outputParser = o -> {
-            List<ModelTensors> mlModelOutputs = (List<ModelTensors>) o;
-            Map<String, ?> dataAsMap = mlModelOutputs.get(0).getMlModelTensors().get(0).getDataAsMap();
-            // Return the response field if it exists, otherwise return the whole response a json string.
-            if (dataAsMap.containsKey(responseField)) {
-                return dataAsMap.get(responseField);
-            } else {
-                return StringUtils.toJson(dataAsMap);
+            try {
+                List<ModelTensors> mlModelOutputs = (List<ModelTensors>) o;
+                Map<String, ?> dataAsMap = mlModelOutputs.getFirst().getMlModelTensors().getFirst()
+                    .getDataAsMap();
+                // Return the response field if it exists, otherwise return the whole response as json string.
+                if (dataAsMap.containsKey(responseField)) {
+                    return dataAsMap.get(responseField);
+                } else {
+                    return StringUtils.toJson(dataAsMap);
+                }
+            } catch (Exception e) {
+                throw new IllegalStateException("LLM returns wrong or empty tensors", e);
             }
         };
     }