From 1361127aedd3bda485aae16169325a79eed1938d Mon Sep 17 00:00:00 2001
From: Songkan Tang <songkant@amazon.com>
Date: Mon, 5 Aug 2024 12:08:51 +0800
Subject: [PATCH] Fix parsing model tensors when the input is a list of docs

Signed-off-by: Songkan Tang <songkant@amazon.com>
---
 src/main/java/org/opensearch/agent/tools/RCATool.java | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/main/java/org/opensearch/agent/tools/RCATool.java b/src/main/java/org/opensearch/agent/tools/RCATool.java
index bd5374bc..57068790 100644
--- a/src/main/java/org/opensearch/agent/tools/RCATool.java
+++ b/src/main/java/org/opensearch/agent/tools/RCATool.java
@@ -282,7 +282,9 @@ private List<RealVector> getEmbeddedVector(List<String> docs) {
             MLInput.builder().algorithm(FunctionName.TEXT_EMBEDDING).inputDataset(inputDataSet).build());
         ActionFuture<MLTaskResponse> mlTaskRspFuture =  client.execute(MLPredictionTaskAction.INSTANCE, request);
         ModelTensorOutput modelTensorOutput = (ModelTensorOutput) mlTaskRspFuture.actionGet().getOutput();
-        List<ModelTensor> mlModelOutputs = modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors();
+        List<ModelTensor> mlModelOutputs = modelTensorOutput.getMlModelOutputs().stream()
+            .map(modelTensors -> modelTensors.getMlModelTensors().get(0))
+            .collect(Collectors.toList());
         return mlModelOutputs.stream()
             .map(tensor -> {
                 Number[] data = tensor.getData();