fix the tests

Author: davidkyle

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/QueryAndDocsInputs.java

@@ -27,10 +27,6 @@ public static QueryAndDocsInputs of(InferenceInputs inferenceInputs) {
     private final Boolean returnDocuments;
     private final Integer topN;
 
-    public QueryAndDocsInputs(String query, List<String> chunks) {
-        this(query, chunks, null, null, false);
-    }
-
     public QueryAndDocsInputs(
         String query,
         List<String> chunks,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereCompletionRequestManager.java

@@ -9,21 +9,22 @@
 
 import org.elasticsearch.common.settings.SecureString;
 import org.elasticsearch.core.Nullable;
-import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
 import org.elasticsearch.inference.ServiceSettings;
 import org.elasticsearch.inference.TaskSettings;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
+import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
 import org.elasticsearch.xpack.inference.services.ServiceUtils;
 import org.elasticsearch.xpack.inference.services.cohere.action.CohereActionVisitor;
 import org.elasticsearch.xpack.inference.services.settings.ApiKeySecrets;
+import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
 import java.net.URI;
 import java.util.Map;
 import java.util.Objects;
 
-public abstract class CohereModel extends Model {
+public abstract class CohereModel extends RateLimitGroupingModel {
     private final SecureString apiKey;
     private final CohereRateLimitServiceSettings rateLimitServiceSettings;
 
@@ -63,5 +64,15 @@ public CohereRateLimitServiceSettings rateLimitServiceSettings() {
 
     public abstract ExecutableAction accept(CohereActionVisitor creator, Map<String, Object> taskSettings);
 
-    public abstract URI uri();
+    public RateLimitSettings rateLimitSettings() {
+        return rateLimitServiceSettings.rateLimitSettings();
+    }
+
+    public int rateLimitGroupingHash() {
+        return apiKey().hashCode();
+    }
+
+    public URI uri() {
+        return rateLimitServiceSettings.uri();
+    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereEmbeddingsRequestManager.java

@@ -9,7 +9,10 @@
 
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
+import java.net.URI;
+
 public interface CohereRateLimitServiceSettings {
     RateLimitSettings rateLimitSettings();
 
+    URI uri();
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereModel.java

@@ -10,14 +10,23 @@
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
 import org.elasticsearch.xpack.inference.external.action.SingleInputSenderExecutableAction;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
+import org.elasticsearch.xpack.inference.external.http.sender.ChatCompletionInput;
+import org.elasticsearch.xpack.inference.external.http.sender.EmbeddingsInput;
+import org.elasticsearch.xpack.inference.external.http.sender.GenericRequestManager;
+import org.elasticsearch.xpack.inference.external.http.sender.QueryAndDocsInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
-import org.elasticsearch.xpack.inference.services.cohere.CohereCompletionRequestManager;
-import org.elasticsearch.xpack.inference.services.cohere.CohereEmbeddingsRequestManager;
-import org.elasticsearch.xpack.inference.services.cohere.CohereRerankRequestManager;
+import org.elasticsearch.xpack.inference.services.cohere.CohereResponseHandler;
 import org.elasticsearch.xpack.inference.services.cohere.completion.CohereCompletionModel;
 import org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingsModel;
+import org.elasticsearch.xpack.inference.services.cohere.request.v1.CohereV1CompletionRequest;
+import org.elasticsearch.xpack.inference.services.cohere.request.v1.CohereV1EmbeddingsRequest;
+import org.elasticsearch.xpack.inference.services.cohere.request.v1.CohereV1RerankRequest;
 import org.elasticsearch.xpack.inference.services.cohere.rerank.CohereRerankModel;
+import org.elasticsearch.xpack.inference.services.cohere.response.CohereCompletionResponseEntity;
+import org.elasticsearch.xpack.inference.services.cohere.response.CohereEmbeddingsResponseEntity;
+import org.elasticsearch.xpack.inference.services.cohere.response.CohereRankedResponseEntity;
 
 import java.util.Map;
 import java.util.Objects;
@@ -28,12 +37,30 @@
  * Provides a way to construct an {@link ExecutableAction} using the visitor pattern based on the cohere model type.
  */
 public class CohereActionCreator implements CohereActionVisitor {
+
+    private static final ResponseHandler EMBEDDINGS_HANDLER = new CohereResponseHandler(
+        "cohere text embedding",
+        CohereEmbeddingsResponseEntity::fromResponse,
+        false
+    );
+
+    private static final ResponseHandler RERANK_HANDLER = new CohereResponseHandler(
+        "cohere rerank",
+        (request, response) -> CohereRankedResponseEntity.fromResponse(response),
+        false
+    );
+
+    private static final ResponseHandler COMPLETION_HANDLER = new CohereResponseHandler(
+        "cohere completion",
+        CohereCompletionResponseEntity::fromResponse,
+        true
+    );
+
     private static final String COMPLETION_ERROR_PREFIX = "Cohere completion";
     private final Sender sender;
     private final ServiceComponents serviceComponents;
 
     public CohereActionCreator(Sender sender, ServiceComponents serviceComponents) {
-        // TODO Batching - accept a class that can handle batching
         this.sender = Objects.requireNonNull(sender);
         this.serviceComponents = Objects.requireNonNull(serviceComponents);
     }
@@ -42,24 +69,53 @@ public CohereActionCreator(Sender sender, ServiceComponents serviceComponents) {
     public ExecutableAction create(CohereEmbeddingsModel model, Map<String, Object> taskSettings) {
         var overriddenModel = CohereEmbeddingsModel.of(model, taskSettings);
         var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage("Cohere embeddings");
-        // TODO - Batching pass the batching class on to the CohereEmbeddingsRequestManager
-        var requestCreator = CohereEmbeddingsRequestManager.of(overriddenModel, serviceComponents.threadPool());
+        var requestCreator = new GenericRequestManager<>(
+            serviceComponents.threadPool(),
+            model,
+            EMBEDDINGS_HANDLER,
+            (inferenceInputs -> new CohereV1EmbeddingsRequest(
+                inferenceInputs.getStringInputs(),
+                inferenceInputs.getInputType(),
+                overriddenModel
+            )),
+            EmbeddingsInput.class
+        );
         return new SenderExecutableAction(sender, requestCreator, failedToSendRequestErrorMessage);
     }
 
     @Override
     public ExecutableAction create(CohereRerankModel model, Map<String, Object> taskSettings) {
         var overriddenModel = CohereRerankModel.of(model, taskSettings);
-        var requestCreator = CohereRerankRequestManager.of(overriddenModel, serviceComponents.threadPool());
+        var requestCreator = new GenericRequestManager<>(
+            serviceComponents.threadPool(),
+            overriddenModel,
+            RERANK_HANDLER,
+            (inferenceInputs -> new CohereV1RerankRequest(
+                inferenceInputs.getQuery(),
+                inferenceInputs.getChunks(),
+                inferenceInputs.getReturnDocuments(),
+                inferenceInputs.getTopN(),
+                overriddenModel
+            )),
+            QueryAndDocsInputs.class
+        );
+
         var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage("Cohere rerank");
         return new SenderExecutableAction(sender, requestCreator, failedToSendRequestErrorMessage);
     }
 
     @Override
     public ExecutableAction create(CohereCompletionModel model, Map<String, Object> taskSettings) {
         // no overridden model as task settings are always empty for cohere completion model
-        var requestManager = CohereCompletionRequestManager.of(model, serviceComponents.threadPool());
+        var requestCreator = new GenericRequestManager<>(
+            serviceComponents.threadPool(),
+            model,
+            COMPLETION_HANDLER,
+            (completionInput) -> new CohereV1CompletionRequest(completionInput.getInputs(), model, completionInput.stream()),
+            ChatCompletionInput.class
+        );
+
         var failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(COMPLETION_ERROR_PREFIX);
-        return new SingleInputSenderExecutableAction(sender, requestManager, failedToSendRequestErrorMessage, COMPLETION_ERROR_PREFIX);
+        return new SingleInputSenderExecutableAction(sender, requestCreator, failedToSendRequestErrorMessage, COMPLETION_ERROR_PREFIX);
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereRateLimitServiceSettings.java

@@ -20,7 +20,6 @@
 import org.elasticsearch.xpack.inference.services.cohere.action.CohereActionVisitor;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 
-import java.net.URI;
 import java.util.Map;
 
 public class CohereCompletionModel extends CohereModel {
@@ -73,9 +72,4 @@ public DefaultSecretSettings getSecretSettings() {
     public ExecutableAction accept(CohereActionVisitor visitor, Map<String, Object> taskSettings) {
         return visitor.create(this, taskSettings);
     }
-
-    @Override
-    public URI uri() {
-        return getServiceSettings().uri();
-    }
 }