From 284bb988210fe11dc2888ba5b1e4b41626f47f8a Mon Sep 17 00:00:00 2001
From: zane-neo <zaniu@amazon.com>
Date: Wed, 8 May 2024 22:33:04 +0800
Subject: [PATCH] Add validation to predictAction in connector

Signed-off-by: zane-neo <zaniu@amazon.com>
---
 .../ml/engine/algorithms/remote/RemoteConnectorExecutor.java   | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java
index 3d6281bfdb..e6fa6331bc 100644
--- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java
+++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java
@@ -103,6 +103,9 @@ private Tuple<Integer, Integer> calculateChunkSize(TextDocsInputDataSet textDocs
             }
         } else {
             Optional<ConnectorAction> predictAction = getConnector().findPredictAction();
+            if (predictAction.isEmpty()) {
+                throw new IllegalArgumentException("no predict action found");
+            }
             String preProcessFunction = predictAction.get().getPreProcessFunction();
             if (preProcessFunction != null && !MLPreProcessFunction.contains(preProcessFunction)) {
                 // user defined preprocess script, this case, the chunk size is always equals to text docs length.