Apply feedback from code review.

Author: afoucret

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java

@@ -56,7 +56,6 @@
 import org.elasticsearch.xpack.inference.services.elastic.action.ElasticInferenceServiceActionCreator;
 import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationHandler;
 import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationRequestHandler;
-import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceCompletionModel;
 import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.elastic.densetextembeddings.ElasticInferenceServiceDenseTextEmbeddingsModel;
@@ -191,7 +190,7 @@ private static Map<String, DefaultModelConfig> initDefaultEndpoints(
         return Map.of(
             DEFAULT_CHAT_COMPLETION_MODEL_ID_V1,
             new DefaultModelConfig(
-                new ElasticInferenceServiceChatCompletionModel(
+                new ElasticInferenceServiceCompletionModel(
                     DEFAULT_CHAT_COMPLETION_ENDPOINT_ID_V1,
                     TaskType.CHAT_COMPLETION,
                     NAME,
@@ -306,7 +305,7 @@ protected void doUnifiedCompletionInfer(
         TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {
-        if (model instanceof ElasticInferenceServiceChatCompletionModel == false) {
+        if (model instanceof ElasticInferenceServiceCompletionModel == false || model.getTaskType() != TaskType.CHAT_COMPLETION) {
             listener.onFailure(createInvalidModelException(model));
             return;
         }
@@ -316,8 +315,8 @@ protected void doUnifiedCompletionInfer(
         // generating a different "traceparent" as every task and every REST request creates a new span).
         var currentTraceInfo = getCurrentTraceInfo();
 
-        var completionModel = (ElasticInferenceServiceChatCompletionModel) model;
-        var overriddenModel = ElasticInferenceServiceChatCompletionModel.of(completionModel, inputs.getRequest());
+        var completionModel = (ElasticInferenceServiceCompletionModel) model;
+        var overriddenModel = ElasticInferenceServiceCompletionModel.of(completionModel, inputs.getRequest());
         var errorMessage = constructFailedToSendRequestMessage(
             String.format(Locale.ROOT, "%s completions", ELASTIC_INFERENCE_SERVICE_IDENTIFIER)
         );
@@ -509,17 +508,7 @@ private static ElasticInferenceServiceModel createModel(
                 context,
                 chunkingSettings
             );
-            case CHAT_COMPLETION -> new ElasticInferenceServiceChatCompletionModel(
-                inferenceEntityId,
-                taskType,
-                NAME,
-                serviceSettings,
-                taskSettings,
-                secretSettings,
-                elasticInferenceServiceComponents,
-                context
-            );
-            case COMPLETION -> new ElasticInferenceServiceCompletionModel(
+            case CHAT_COMPLETION, COMPLETION -> new ElasticInferenceServiceCompletionModel(
                 inferenceEntityId,
                 taskType,
                 NAME,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceUnifiedCompletionRequestManager.java

@@ -17,7 +17,7 @@
 import org.elasticsearch.xpack.inference.external.http.sender.ExecutableInferenceRequest;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
-import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceChatCompletionModel;
+import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceCompletionModel;
 import org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceUnifiedChatCompletionRequest;
 import org.elasticsearch.xpack.inference.services.openai.response.OpenAiChatCompletionResponseEntity;
 import org.elasticsearch.xpack.inference.telemetry.TraceContext;
@@ -32,7 +32,7 @@ public class ElasticInferenceServiceUnifiedCompletionRequestManager extends Elas
     private static final ResponseHandler HANDLER = createCompletionHandler();
 
     public static ElasticInferenceServiceUnifiedCompletionRequestManager of(
-        ElasticInferenceServiceChatCompletionModel model,
+        ElasticInferenceServiceCompletionModel model,
         ThreadPool threadPool,
         TraceContext traceContext
     ) {
@@ -43,11 +43,11 @@ public static ElasticInferenceServiceUnifiedCompletionRequestManager of(
         );
     }
 
-    private final ElasticInferenceServiceChatCompletionModel model;
+    private final ElasticInferenceServiceCompletionModel model;
     private final TraceContext traceContext;
 
     private ElasticInferenceServiceUnifiedCompletionRequestManager(
-        ElasticInferenceServiceChatCompletionModel model,
+        ElasticInferenceServiceCompletionModel model,
         ThreadPool threadPool,
         TraceContext traceContext
     ) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/action/ElasticInferenceServiceActionCreator.java

@@ -16,16 +16,17 @@
 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.external.http.sender.UnifiedChatInput;
 import org.elasticsearch.xpack.inference.external.response.elastic.ElasticInferenceServiceDenseTextEmbeddingsResponseEntity;
 import org.elasticsearch.xpack.inference.external.response.elastic.ElasticInferenceServiceRerankResponseEntity;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceResponseHandler;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSparseEmbeddingsRequestManager;
 import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceCompletionModel;
 import org.elasticsearch.xpack.inference.services.elastic.densetextembeddings.ElasticInferenceServiceDenseTextEmbeddingsModel;
-import org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceCompletionRequest;
 import org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceDenseTextEmbeddingsRequest;
 import org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceRerankRequest;
+import org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceUnifiedChatCompletionRequest;
 import org.elasticsearch.xpack.inference.services.elastic.rerank.ElasticInferenceServiceRerankModel;
 import org.elasticsearch.xpack.inference.services.elastic.sparseembeddings.ElasticInferenceServiceSparseEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.openai.response.OpenAiChatCompletionResponseEntity;
@@ -40,6 +41,8 @@
 
 public class ElasticInferenceServiceActionCreator implements ElasticInferenceServiceActionVisitor {
 
+    public static final String USER_ROLE = "user";
+
     static final ResponseHandler DENSE_TEXT_EMBEDDINGS_HANDLER = new ElasticInferenceServiceResponseHandler(
         "elastic dense text embedding",
         ElasticInferenceServiceDenseTextEmbeddingsResponseEntity::fromResponse
@@ -127,8 +130,8 @@ public ExecutableAction create(ElasticInferenceServiceCompletionModel model, Map
             threadPool,
             model,
             COMPLETION_HANDLER,
-            (chatCompletionInput) -> new ElasticInferenceServiceCompletionRequest(
-                chatCompletionInput.getInputs(),
+            (chatCompletionInput) -> new ElasticInferenceServiceUnifiedChatCompletionRequest(
+                new UnifiedChatInput(chatCompletionInput.getInputs(), USER_ROLE, false),
                 model,
                 traceContext,
                 extractRequestMetadataFromThreadContext(threadPool.getThreadContext())

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/completion/ElasticInferenceServiceChatCompletionModel.java

@@ -16,6 +16,7 @@
 import org.elasticsearch.inference.SecretSettings;
 import org.elasticsearch.inference.TaskSettings;
 import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.inference.UnifiedCompletionRequest;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
@@ -26,13 +27,22 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.Map;
+import java.util.Objects;
 
-/**
- * Adapter model for COMPLETION task type that converts simple text inputs into chat messages
- * and uses the chat completion endpoint.
- */
 public class ElasticInferenceServiceCompletionModel extends ElasticInferenceServiceExecutableActionModel {
 
+    public static ElasticInferenceServiceCompletionModel of(
+        ElasticInferenceServiceCompletionModel model,
+        UnifiedCompletionRequest request
+    ) {
+        var originalModelServiceSettings = model.getServiceSettings();
+        var overriddenServiceSettings = new ElasticInferenceServiceCompletionServiceSettings(
+            Objects.requireNonNullElse(request.model(), originalModelServiceSettings.modelId())
+        );
+
+        return new ElasticInferenceServiceCompletionModel(model, overriddenServiceSettings);
+    }
+
     private final URI uri;
 
     public ElasticInferenceServiceCompletionModel(
@@ -56,6 +66,14 @@ public ElasticInferenceServiceCompletionModel(
         );
     }
 
+    public ElasticInferenceServiceCompletionModel(
+        ElasticInferenceServiceCompletionModel model,
+        ElasticInferenceServiceCompletionServiceSettings serviceSettings
+    ) {
+        super(model, serviceSettings);
+        this.uri = createUri();
+    }
+
     public ElasticInferenceServiceCompletionModel(
         String inferenceEntityId,
         TaskType taskType,
@@ -74,14 +92,6 @@ public ElasticInferenceServiceCompletionModel(
         this.uri = createUri();
     }
 
-    public ElasticInferenceServiceCompletionModel(
-        ElasticInferenceServiceCompletionModel model,
-        ElasticInferenceServiceCompletionServiceSettings serviceSettings
-    ) {
-        super(model, serviceSettings);
-        this.uri = createUri();
-    }
-
     @Override
     public ElasticInferenceServiceCompletionServiceSettings getServiceSettings() {
         return (ElasticInferenceServiceCompletionServiceSettings) super.getServiceSettings();
@@ -93,7 +103,7 @@ public URI uri() {
 
     private URI createUri() throws ElasticsearchStatusException {
         try {
-            // Use the same chat endpoint as CHAT_COMPLETION
+            // TODO, consider transforming the base URL into a URI for better error handling.
             return new URI(elasticInferenceServiceComponents().elasticInferenceServiceUrl() + "/api/v1/chat");
         } catch (URISyntaxException e) {
             throw new ElasticsearchStatusException(