Update index calls to directly use index response without parsing

Signed-off-by: Daniel Widdis <[email protected]>

Author: dbwiddis

plugin/src/main/java/org/opensearch/ml/action/agents/TransportRegisterAgentAction.java

@@ -150,7 +150,7 @@ private void registerAgentToIndex(MLAgent mlAgent, String tenantId, ActionListen
                                 listener.onFailure(cause);
                             } else {
                                 try {
-                                    IndexResponse indexResponse = IndexResponse.fromXContent(r.parser());
+                                    IndexResponse indexResponse = r.indexResponse();
                                     log.info("Agent creation result: {}, Agent id: {}", indexResponse.getResult(), indexResponse.getId());
                                     MLRegisterAgentResponse response = new MLRegisterAgentResponse(r.id());
                                     listener.onResponse(response);

plugin/src/main/java/org/opensearch/ml/action/connector/TransportCreateConnectorAction.java

@@ -9,7 +9,6 @@
 import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX;
 import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_TRUSTED_CONNECTOR_ENDPOINTS_REGEX;
 
-import java.io.IOException;
 import java.time.Instant;
 import java.util.HashSet;
 import java.util.List;
@@ -179,15 +178,15 @@ private void indexConnector(Connector connector, ActionListener<MLCreateConnecto
                             listener.onFailure(cause);
                         } else {
                             try {
-                                IndexResponse indexResponse = IndexResponse.fromXContent(r.parser());
+                                IndexResponse indexResponse = r.indexResponse();
                                 log
                                     .info(
                                         "Connector creation result: {}, connector id: {}",
                                         indexResponse.getResult(),
                                         indexResponse.getId()
                                     );
                                 listener.onResponse(new MLCreateConnectorResponse(indexResponse.getId()));
-                            } catch (IOException e) {
+                            } catch (Exception e) {
                                 listener.onFailure(e);
                             }
                         }

plugin/src/main/java/org/opensearch/ml/model/MLModelGroupManager.java

@@ -145,7 +145,7 @@ public void createModelGroup(MLRegisterModelGroupInput input, ActionListener<Str
                                         wrappedListener.onFailure(cause);
                                     } else {
                                         try {
-                                            IndexResponse indexResponse = IndexResponse.fromXContent(r.parser());
+                                            IndexResponse indexResponse = r.indexResponse();
                                             log
                                                 .info(
                                                     "Model group creation result: {}, model group id: {}",

plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java

@@ -661,19 +661,7 @@ private void indexRemoteModel(
                 });
 
                 ThreadedActionListener<IndexResponse> putListener = threadedActionListener(REGISTER_THREAD_POOL, indexListener);
-                sdkClient.putDataObjectAsync(putModelMetaRequest).whenComplete((r, throwable) -> {
-                    if (throwable == null) {
-                        try {
-                            IndexResponse ir = IndexResponse.fromXContent(r.parser());
-                            putListener.onResponse(ir);
-                        } catch (Exception e) {
-                            putListener.onFailure(e);
-                        }
-                    } else {
-                        Exception e = SdkClientUtils.unwrapAndConvertToException(throwable);
-                        putListener.onFailure(e);
-                    }
-                });
+                sdkClient.putDataObjectAsync(putModelMetaRequest).whenComplete(SdkClientUtils.wrapPutCompletion(putListener));
             }, error -> {
                 // failed to initialize the model index
                 log.error("Failed to init model index", error);
@@ -764,19 +752,7 @@ void indexRemoteModel(MLRegisterModelInput registerModelInput, MLTask mlTask, St
                 });
 
                 ThreadedActionListener<IndexResponse> putListener = threadedActionListener(REGISTER_THREAD_POOL, indexListener);
-                sdkClient.putDataObjectAsync(putModelMetaRequest).whenComplete((r, throwable) -> {
-                    if (throwable == null) {
-                        try {
-                            IndexResponse ir = IndexResponse.fromXContent(r.parser());
-                            putListener.onResponse(ir);
-                        } catch (Exception e) {
-                            putListener.onFailure(e);
-                        }
-                    } else {
-                        Exception e = SdkClientUtils.unwrapAndConvertToException(throwable);
-                        putListener.onFailure(e);
-                    }
-                });
+                sdkClient.putDataObjectAsync(putModelMetaRequest).whenComplete(SdkClientUtils.wrapPutCompletion(putListener));
 
             }, e -> {
                 log.error("Failed to init model index", e);

plugin/src/main/java/org/opensearch/ml/task/MLTaskManager.java

@@ -304,7 +304,7 @@ public void createMLTask(MLTask mlTask, ActionListener<IndexResponse> listener)
                             listener.onFailure(cause);
                         } else {
                             try {
-                                IndexResponse indexResponse = IndexResponse.fromXContent(r.parser());
+                                IndexResponse indexResponse = r.indexResponse();
                                 log.info("Task creation result: {}, Task id: {}", indexResponse.getResult(), indexResponse.getId());
                                 listener.onResponse(indexResponse);
                             } catch (Exception e) {