[ML] Consolidate ExecutableActions

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/ActionUtils.java

@@ -24,13 +24,13 @@ public static ActionListener<InferenceServiceResults> wrapFailuresInElasticsearc
         String errorMessage,
         ActionListener<InferenceServiceResults> listener
     ) {
-        return ActionListener.wrap(listener::onResponse, e -> {
+        return listener.delegateResponse((l, e) -> {
             var unwrappedException = ExceptionsHelper.unwrapCause(e);
 
             if (unwrappedException instanceof ElasticsearchException esException) {
-                listener.onFailure(esException);
+                l.onFailure(esException);
             } else {
-                listener.onFailure(createInternalServerError(unwrappedException, errorMessage));
+                l.onFailure(createInternalServerError(unwrappedException, errorMessage));
             }
         });
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/amazonbedrock/AmazonBedrockEmbeddingsAction.java → x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/SenderExecutableAction.java

@@ -5,44 +5,38 @@
  * 2.0.
  */
 
-package org.elasticsearch.xpack.inference.external.action.amazonbedrock;
+package org.elasticsearch.xpack.inference.external.action;
 
-import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.inference.InferenceServiceResults;
-import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.RequestManager;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 
 import java.util.Objects;
 
-import static org.elasticsearch.xpack.inference.external.action.ActionUtils.createInternalServerError;
 import static org.elasticsearch.xpack.inference.external.action.ActionUtils.wrapFailuresInElasticsearchException;
 
-public class AmazonBedrockEmbeddingsAction implements ExecutableAction {
+public class SenderExecutableAction implements ExecutableAction {
 
     private final Sender sender;
     private final RequestManager requestManager;
-    private final String errorMessage;
+    private final String failedToSendRequestErrorMessage;
 
-    public AmazonBedrockEmbeddingsAction(Sender sender, RequestManager requestManager, String errorMessage) {
+    public SenderExecutableAction(Sender sender, RequestManager requestManager, String failedToSendRequestErrorMessage) {
         this.sender = Objects.requireNonNull(sender);
         this.requestManager = Objects.requireNonNull(requestManager);
-        this.errorMessage = Objects.requireNonNull(errorMessage);
+        this.failedToSendRequestErrorMessage = Objects.requireNonNull(failedToSendRequestErrorMessage);
     }
 
     @Override
     public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
+        var wrappedListener = wrapFailuresInElasticsearchException(failedToSendRequestErrorMessage, listener);
         try {
-            ActionListener<InferenceServiceResults> wrappedListener = wrapFailuresInElasticsearchException(errorMessage, listener);
-
             sender.send(requestManager, inferenceInputs, timeout, wrappedListener);
-        } catch (ElasticsearchException e) {
-            listener.onFailure(e);
         } catch (Exception e) {
-            listener.onFailure(createInternalServerError(e, errorMessage));
+            wrappedListener.onFailure(e);
         }
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/SingleInputSenderExecutableAction.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.external.action;
+
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.xpack.inference.external.http.sender.DocumentsOnlyInput;
+import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
+import org.elasticsearch.xpack.inference.external.http.sender.RequestManager;
+import org.elasticsearch.xpack.inference.external.http.sender.Sender;
+
+import java.util.Objects;
+
+public class SingleInputSenderExecutableAction extends SenderExecutableAction {
+    private final String requestTypeForInputValidationError;
+
+    public SingleInputSenderExecutableAction(
+        Sender sender,
+        RequestManager requestManager,
+        String failedToSendRequestErrorMessage,
+        String requestTypeForInputValidationError
+    ) {
+        super(sender, requestManager, failedToSendRequestErrorMessage);
+        this.requestTypeForInputValidationError = Objects.requireNonNull(requestTypeForInputValidationError);
+    }
+
+    @Override
+    public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
+        if (inferenceInputs instanceof DocumentsOnlyInput == false) {
+            listener.onFailure(new ElasticsearchStatusException("Invalid inference input type", RestStatus.INTERNAL_SERVER_ERROR));
+            return;
+        }
+
+        var docsOnlyInput = (DocumentsOnlyInput) inferenceInputs;
+        if (docsOnlyInput.getInputs().size() > 1) {
+            listener.onFailure(
+                new ElasticsearchStatusException(requestTypeForInputValidationError + " only accepts 1 input", RestStatus.BAD_REQUEST)
+            );
+            return;
+        }
+
+        super.execute(inferenceInputs, timeout, listener);
+    }
+
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/amazonbedrock/AmazonBedrockActionCreator.java

@@ -10,6 +10,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
+import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
 import org.elasticsearch.xpack.inference.external.http.sender.AmazonBedrockChatCompletionRequestManager;
 import org.elasticsearch.xpack.inference.external.http.sender.AmazonBedrockEmbeddingsRequestManager;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
@@ -43,14 +44,14 @@ public ExecutableAction create(AmazonBedrockEmbeddingsModel embeddingsModel, Map
             timeout
         );
         var errorMessage = constructFailedToSendRequestMessage(null, "Amazon Bedrock embeddings");
-        return new AmazonBedrockEmbeddingsAction(sender, requestManager, errorMessage);
+        return new SenderExecutableAction(sender, requestManager, errorMessage);
     }
 
     @Override
     public ExecutableAction create(AmazonBedrockChatCompletionModel completionModel, Map<String, Object> taskSettings) {
         var overriddenModel = AmazonBedrockChatCompletionModel.of(completionModel, taskSettings);
         var requestManager = new AmazonBedrockChatCompletionRequestManager(overriddenModel, serviceComponents.threadPool(), timeout);
         var errorMessage = constructFailedToSendRequestMessage(null, "Amazon Bedrock completion");
-        return new AmazonBedrockChatCompletionAction(sender, requestManager, errorMessage);
+        return new SenderExecutableAction(sender, requestManager, errorMessage);
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/anthropic/AnthropicActionCreator.java

@@ -8,17 +8,22 @@
 package org.elasticsearch.xpack.inference.external.action.anthropic;
 
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
+import org.elasticsearch.xpack.inference.external.action.SingleInputSenderExecutableAction;
+import org.elasticsearch.xpack.inference.external.http.sender.AnthropicCompletionRequestManager;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.anthropic.completion.AnthropicChatCompletionModel;
 
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.xpack.inference.external.action.ActionUtils.constructFailedToSendRequestMessage;
+
 /**
  * Provides a way to construct an {@link ExecutableAction} using the visitor pattern based on the anthropic model type.
  */
 public class AnthropicActionCreator implements AnthropicActionVisitor {
+    private static final String ERROR_PREFIX = "Anthropic chat completions";
     private final Sender sender;
     private final ServiceComponents serviceComponents;
 
@@ -30,7 +35,8 @@ public AnthropicActionCreator(Sender sender, ServiceComponents serviceComponents
     @Override
     public ExecutableAction create(AnthropicChatCompletionModel model, Map<String, Object> taskSettings) {
         var overriddenModel = AnthropicChatCompletionModel.of(model, taskSettings);
-
-        return new AnthropicChatCompletionAction(sender, overriddenModel, serviceComponents);
+        var requestCreator = AnthropicCompletionRequestManager.of(overriddenModel, serviceComponents.threadPool());
+        var errorMessage = constructFailedToSendRequestMessage(overriddenModel.getUri(), ERROR_PREFIX);
+        return new SingleInputSenderExecutableAction(sender, requestCreator, errorMessage, ERROR_PREFIX);
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/azureaistudio/AzureAiStudioActionCreator.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.inference.external.action.azureaistudio;
 
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
+import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
 import org.elasticsearch.xpack.inference.external.http.sender.AzureAiStudioChatCompletionRequestManager;
 import org.elasticsearch.xpack.inference.external.http.sender.AzureAiStudioEmbeddingsRequestManager;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
@@ -34,7 +35,7 @@ public ExecutableAction create(AzureAiStudioChatCompletionModel completionModel,
         var overriddenModel = AzureAiStudioChatCompletionModel.of(completionModel, taskSettings);
         var requestManager = new AzureAiStudioChatCompletionRequestManager(overriddenModel, serviceComponents.threadPool());
         var errorMessage = constructFailedToSendRequestMessage(completionModel.uri(), "Azure AI Studio completion");
-        return new AzureAiStudioAction(sender, requestManager, errorMessage);
+        return new SenderExecutableAction(sender, requestManager, errorMessage);
     }
 
     @Override
@@ -46,6 +47,6 @@ public ExecutableAction create(AzureAiStudioEmbeddingsModel embeddingsModel, Map
             serviceComponents.threadPool()
         );
         var errorMessage = constructFailedToSendRequestMessage(embeddingsModel.uri(), "Azure AI Studio embeddings");
-        return new AzureAiStudioAction(sender, requestManager, errorMessage);
+        return new SenderExecutableAction(sender, requestManager, errorMessage);
     }
 }