[8.x] [ML] Refactor OpenAI request managers (#124144)

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

@@ -10,9 +10,18 @@
 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.sender.OpenAiCompletionRequestManager;
-import org.elasticsearch.xpack.inference.external.http.sender.OpenAiEmbeddingsRequestManager;
+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.GenericRequestManager;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
+import org.elasticsearch.xpack.inference.external.http.sender.TruncatingRequestManager;
+import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
+import org.elasticsearch.xpack.inference.external.openai.OpenAiChatCompletionResponseHandler;
+import org.elasticsearch.xpack.inference.external.openai.OpenAiResponseHandler;
+import org.elasticsearch.xpack.inference.external.request.openai.OpenAiEmbeddingsRequest;
+import org.elasticsearch.xpack.inference.external.request.openai.OpenAiUnifiedChatCompletionRequest;
+import org.elasticsearch.xpack.inference.external.response.openai.OpenAiChatCompletionResponseEntity;
+import org.elasticsearch.xpack.inference.external.response.openai.OpenAiEmbeddingsResponseEntity;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsModel;
@@ -27,6 +36,18 @@
  */
 public class OpenAiActionCreator implements OpenAiActionVisitor {
     public static final String COMPLETION_ERROR_PREFIX = "OpenAI chat completions";
+    public static final String USER_ROLE = "user";
+
+    static final ResponseHandler COMPLETION_HANDLER = new OpenAiChatCompletionResponseHandler(
+        "openai completion",
+        OpenAiChatCompletionResponseEntity::fromResponse
+    );
+    public static final ResponseHandler EMBEDDINGS_HANDLER = new OpenAiResponseHandler(
+        "openai text embedding",
+        OpenAiEmbeddingsResponseEntity::fromResponse,
+        false
+    );
+
     private final Sender sender;
     private final ServiceComponents serviceComponents;
 
@@ -38,20 +59,30 @@ public OpenAiActionCreator(Sender sender, ServiceComponents serviceComponents) {
     @Override
     public ExecutableAction create(OpenAiEmbeddingsModel model, Map<String, Object> taskSettings) {
         var overriddenModel = OpenAiEmbeddingsModel.of(model, taskSettings);
-        var requestCreator = OpenAiEmbeddingsRequestManager.of(
+        var manager = new TruncatingRequestManager(
+            serviceComponents.threadPool(),
             overriddenModel,
-            serviceComponents.truncator(),
-            serviceComponents.threadPool()
+            EMBEDDINGS_HANDLER,
+            (truncationResult) -> new OpenAiEmbeddingsRequest(serviceComponents.truncator(), truncationResult, overriddenModel),
+            overriddenModel.getServiceSettings().maxInputTokens()
         );
+
         var errorMessage = constructFailedToSendRequestMessage("OpenAI embeddings");
-        return new SenderExecutableAction(sender, requestCreator, errorMessage);
+        return new SenderExecutableAction(sender, manager, errorMessage);
     }
 
     @Override
     public ExecutableAction create(OpenAiChatCompletionModel model, Map<String, Object> taskSettings) {
         var overriddenModel = OpenAiChatCompletionModel.of(model, taskSettings);
-        var requestCreator = OpenAiCompletionRequestManager.of(overriddenModel, serviceComponents.threadPool());
+        var manager = new GenericRequestManager<>(
+            serviceComponents.threadPool(),
+            overriddenModel,
+            COMPLETION_HANDLER,
+            (inputs) -> new OpenAiUnifiedChatCompletionRequest(new UnifiedChatInput(inputs, USER_ROLE), overriddenModel),
+            ChatCompletionInput.class
+        );
+
         var errorMessage = constructFailedToSendRequestMessage(COMPLETION_ERROR_PREFIX);
-        return new SingleInputSenderExecutableAction(sender, requestCreator, errorMessage, COMPLETION_ERROR_PREFIX);
+        return new SingleInputSenderExecutableAction(sender, manager, errorMessage, COMPLETION_ERROR_PREFIX);
     }
 }

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

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.inference.external.http.sender;
 
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
 import java.util.Objects;
@@ -17,14 +18,31 @@
 abstract class BaseRequestManager implements RequestManager {
     private final ThreadPool threadPool;
     private final String inferenceEntityId;
-    private final Object rateLimitGroup;
+    // It's possible that two inference endpoints have the same information defining the group but have different
+    // rate limits then they should be in different groups otherwise whoever initially created the group will set
+    // the rate and the other inference endpoint's rate will be ignored
+    private final EndpointGrouping endpointGrouping;
     private final RateLimitSettings rateLimitSettings;
 
     BaseRequestManager(ThreadPool threadPool, String inferenceEntityId, Object rateLimitGroup, RateLimitSettings rateLimitSettings) {
         this.threadPool = Objects.requireNonNull(threadPool);
         this.inferenceEntityId = Objects.requireNonNull(inferenceEntityId);
-        this.rateLimitGroup = Objects.requireNonNull(rateLimitGroup);
-        this.rateLimitSettings = Objects.requireNonNull(rateLimitSettings);
+
+        Objects.requireNonNull(rateLimitSettings);
+        this.endpointGrouping = new EndpointGrouping(Objects.requireNonNull(rateLimitGroup).hashCode(), rateLimitSettings);
+        this.rateLimitSettings = rateLimitSettings;
+    }
+
+    BaseRequestManager(ThreadPool threadPool, RateLimitGroupingModel rateLimitGroupingModel) {
+        this.threadPool = Objects.requireNonNull(threadPool);
+        Objects.requireNonNull(rateLimitGroupingModel);
+
+        this.inferenceEntityId = rateLimitGroupingModel.inferenceEntityId();
+        this.endpointGrouping = new EndpointGrouping(
+            rateLimitGroupingModel.rateLimitGroupingHash(),
+            rateLimitGroupingModel.rateLimitSettings()
+        );
+        this.rateLimitSettings = rateLimitGroupingModel.rateLimitSettings();
     }
 
     protected void execute(Runnable runnable) {
@@ -38,16 +56,13 @@ public String inferenceEntityId() {
 
     @Override
     public Object rateLimitGrouping() {
-        // It's possible that two inference endpoints have the same information defining the group but have different
-        // rate limits then they should be in different groups otherwise whoever initially created the group will set
-        // the rate and the other inference endpoint's rate will be ignored
-        return new EndpointGrouping(rateLimitGroup, rateLimitSettings);
+        return endpointGrouping;
     }
 
     @Override
     public RateLimitSettings rateLimitSettings() {
         return rateLimitSettings;
     }
 
-    private record EndpointGrouping(Object group, RateLimitSettings settings) {}
+    private record EndpointGrouping(int group, RateLimitSettings settings) {}
 }

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

@@ -0,0 +1,59 @@
+/*
+ * 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.http.sender;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.external.http.retry.RequestSender;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
+import org.elasticsearch.xpack.inference.external.request.Request;
+import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
+
+import java.util.Objects;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+/**
+ * This is a temporary class to use while we refactor all the request managers. After all the request managers extend
+ * this class we'll move this functionality directly into the {@link BaseRequestManager}.
+ */
+public class GenericRequestManager<T extends InferenceInputs> extends BaseRequestManager {
+    private static final Logger logger = LogManager.getLogger(GenericRequestManager.class);
+
+    protected final ResponseHandler responseHandler;
+    protected final Function<T, Request> requestCreator;
+    protected final Class<T> inputType;
+
+    public GenericRequestManager(
+        ThreadPool threadPool,
+        RateLimitGroupingModel rateLimitGroupingModel,
+        ResponseHandler responseHandler,
+        Function<T, Request> requestCreator,
+        Class<T> inputType
+    ) {
+        super(threadPool, rateLimitGroupingModel);
+        this.responseHandler = Objects.requireNonNull(responseHandler);
+        this.requestCreator = Objects.requireNonNull(requestCreator);
+        this.inputType = Objects.requireNonNull(inputType);
+    }
+
+    @Override
+    public void execute(
+        InferenceInputs inferenceInputs,
+        RequestSender requestSender,
+        Supplier<Boolean> hasRequestCompletedFunction,
+        ActionListener<InferenceServiceResults> listener
+    ) {
+        var request = requestCreator.apply(inferenceInputs.castTo(inputType));
+
+        execute(new ExecutableInferenceRequest(requestSender, logger, request, responseHandler, hasRequestCompletedFunction, listener));
+    }
+}

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

@@ -22,7 +22,7 @@ public static IllegalArgumentException createUnsupportedTypeException(InferenceI
         );
     }
 
-    public <T> T castTo(Class<T> clazz) {
+    public <T extends InferenceInputs> T castTo(Class<T> clazz) {
         if (clazz.isInstance(this) == false) {
             throw createUnsupportedTypeException(this, clazz);
         }