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 getEmbeddedVector(List docs) { MLInput.builder().algorithm(FunctionName.TEXT_EMBEDDING).inputDataset(inputDataSet).build()); ActionFuture mlTaskRspFuture = client.execute(MLPredictionTaskAction.INSTANCE, request); ModelTensorOutput modelTensorOutput = (ModelTensorOutput) mlTaskRspFuture.actionGet().getOutput(); - List mlModelOutputs = modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors(); + List mlModelOutputs = modelTensorOutput.getMlModelOutputs().stream() + .map(modelTensors -> modelTensors.getMlModelTensors().get(0)) + .collect(Collectors.toList()); return mlModelOutputs.stream() .map(tensor -> { Number[] data = tensor.getData();