Update model to return correct model for CHAT_COMPLETION task type

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/UnifiedCompletionAction.java

@@ -92,9 +92,9 @@ public ActionRequestValidationException validate() {
                 return e;
             }
 
-            if (taskType.isAnyOrSame(TaskType.COMPLETION) == false) {
+            if (taskType.isAnyOrSame(TaskType.CHAT_COMPLETION) == false) {
                 var e = new ActionRequestValidationException();
-                e.addValidationError("Field [taskType] must be [completion]");
+                e.addValidationError("Field [taskType] must be [chat_completion]");
                 return e;
             }
 

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/action/UnifiedCompletionActionRequestTests.java

@@ -52,7 +52,7 @@ public void testValidation_ReturnsException_When_TaskType_IsNot_Completion() {
             TimeValue.timeValueSeconds(10)
         );
         var exception = request.validate();
-        assertThat(exception.getMessage(), is("Validation Failed: 1: Field [taskType] must be [completion];"));
+        assertThat(exception.getMessage(), is("Validation Failed: 1: Field [taskType] must be [chat_completion];"));
     }
 
     public void testValidation_ReturnsNull_When_TaskType_IsAny() {

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java

@@ -272,9 +272,9 @@ public void testGetServicesWithChatCompletionTaskType() throws IOException {
         List<Object> services = getServices(TaskType.CHAT_COMPLETION);
         if ((ElasticInferenceServiceFeature.DEPRECATED_ELASTIC_INFERENCE_SERVICE_FEATURE_FLAG.isEnabled()
             || ElasticInferenceServiceFeature.ELASTIC_INFERENCE_SERVICE_FEATURE_FLAG.isEnabled())) {
-            assertThat(services.size(), equalTo(2));
+            assertThat(services.size(), equalTo(3));
         } else {
-            assertThat(services.size(), equalTo(1));
+            assertThat(services.size(), equalTo(2));
         }
 
         String[] providers = new String[services.size()];
@@ -283,7 +283,7 @@ public void testGetServicesWithChatCompletionTaskType() throws IOException {
             providers[i] = (String) serviceConfig.get("service");
         }
 
-        var providerList = new ArrayList<>(List.of("openai"));
+        var providerList = new ArrayList<>(List.of("openai", "streaming_completion_test_service"));
 
         if ((ElasticInferenceServiceFeature.DEPRECATED_ELASTIC_INFERENCE_SERVICE_FEATURE_FLAG.isEnabled()
             || ElasticInferenceServiceFeature.ELASTIC_INFERENCE_SERVICE_FEATURE_FLAG.isEnabled())) {
@@ -519,14 +519,19 @@ public void testSupportedStream() throws Exception {
 
     public void testUnifiedCompletionInference() throws Exception {
         String modelId = "streaming";
-        putModel(modelId, mockCompletionServiceModelConfig(TaskType.COMPLETION));
+        putModel(modelId, mockCompletionServiceModelConfig(TaskType.CHAT_COMPLETION));
         var singleModel = getModel(modelId);
         assertEquals(modelId, singleModel.get("inference_id"));
-        assertEquals(TaskType.COMPLETION.toString(), singleModel.get("task_type"));
+        assertEquals(TaskType.CHAT_COMPLETION.toString(), singleModel.get("task_type"));
 
         var input = IntStream.range(1, 2 + randomInt(8)).mapToObj(i -> randomAlphanumericOfLength(5)).toList();
         try {
-            var events = unifiedCompletionInferOnMockService(modelId, TaskType.COMPLETION, input, VALIDATE_ELASTIC_PRODUCT_HEADER_CONSUMER);
+            var events = unifiedCompletionInferOnMockService(
+                modelId,
+                TaskType.CHAT_COMPLETION,
+                input,
+                VALIDATE_ELASTIC_PRODUCT_HEADER_CONSUMER
+            );
             var expectedResponses = expectedResultsIterator(input);
             assertThat(events.size(), equalTo((input.size() + 1) * 2));
             events.forEach(event -> {

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java

@@ -55,9 +55,9 @@ public List<Factory> getInferenceServiceFactories() {
 
     public static class TestInferenceService extends AbstractTestInferenceService {
         private static final String NAME = "streaming_completion_test_service";
-        private static final Set<TaskType> supportedStreamingTasks = Set.of(TaskType.COMPLETION);
+        private static final Set<TaskType> supportedStreamingTasks = Set.of(TaskType.COMPLETION, TaskType.CHAT_COMPLETION);
 
-        private static final EnumSet<TaskType> supportedTaskTypes = EnumSet.of(TaskType.COMPLETION);
+        private static final EnumSet<TaskType> supportedTaskTypes = EnumSet.of(TaskType.COMPLETION, TaskType.CHAT_COMPLETION);
 
         public TestInferenceService(InferenceServiceExtension.InferenceServiceFactoryContext context) {}
 
@@ -129,7 +129,7 @@ public void unifiedCompletionInfer(
             ActionListener<InferenceServiceResults> listener
         ) {
             switch (model.getConfigurations().getTaskType()) {
-                case COMPLETION -> listener.onResponse(makeUnifiedResults(request));
+                case CHAT_COMPLETION -> listener.onResponse(makeUnifiedResults(request));
                 default -> listener.onFailure(
                     new ElasticsearchStatusException(
                         TaskType.unsupportedTaskTypeErrorMsg(model.getConfigurations().getTaskType(), name()),

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

@@ -52,7 +52,7 @@ public TransportUnifiedCompletionInferenceAction(
 
     @Override
     protected boolean isInvalidTaskTypeForInferenceEndpoint(UnifiedCompletionAction.Request request, UnparsedModel unparsedModel) {
-        return request.getTaskType().isAnyOrSame(TaskType.COMPLETION) == false || unparsedModel.taskType() != TaskType.COMPLETION;
+        return request.getTaskType().isAnyOrSame(TaskType.CHAT_COMPLETION) == false || unparsedModel.taskType() != TaskType.CHAT_COMPLETION;
     }
 
     @Override
@@ -64,7 +64,7 @@ protected ElasticsearchStatusException createInvalidTaskTypeException(
             "Incompatible task_type for unified API, the requested type [{}] must be one of [{}]",
             RestStatus.BAD_REQUEST,
             request.getTaskType(),
-            TaskType.COMPLETION.toString()
+            TaskType.CHAT_COMPLETION.toString()
         );
     }
 

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

@@ -26,7 +26,7 @@ public class OpenAiCompletionRequestManager extends OpenAiRequestManager {
 
     private static final Logger logger = LogManager.getLogger(OpenAiCompletionRequestManager.class);
     private static final ResponseHandler HANDLER = createCompletionHandler();
-    static final String USER_ROLE = "user";
+    public static final String USER_ROLE = "user";
 
     public static OpenAiCompletionRequestManager of(OpenAiChatCompletionModel model, ThreadPool threadPool) {
         return new OpenAiCompletionRequestManager(Objects.requireNonNull(model), Objects.requireNonNull(threadPool));

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

@@ -251,7 +251,7 @@ private static ElasticInferenceServiceModel createModel(
                 eisServiceComponents,
                 context
             );
-            case COMPLETION -> new ElasticInferenceServiceCompletionModel(
+            case CHAT_COMPLETION -> new ElasticInferenceServiceCompletionModel(
                 inferenceEntityId,
                 taskType,
                 NAME,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/validation/ModelValidatorBuilder.java

@@ -23,6 +23,9 @@ public static ModelValidator buildModelValidator(TaskType taskType) {
             case COMPLETION -> {
                 return new ChatCompletionModelValidator(new SimpleServiceIntegrationValidator());
             }
+            case CHAT_COMPLETION -> {
+                return new ChatCompletionModelValidator(new SimpleChatCompletionServiceIntegrationValidator());
+            }
             case SPARSE_EMBEDDING, RERANK, ANY -> {
                 return new SimpleModelValidator(new SimpleServiceIntegrationValidator());
             }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/validation/SimpleChatCompletionServiceIntegrationValidator.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.services.validation;
+
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.inference.InferenceService;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.inference.Model;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
+
+import java.util.List;
+
+import static org.elasticsearch.xpack.inference.external.http.sender.OpenAiCompletionRequestManager.USER_ROLE;
+
+/**
+ * This class uses the unified chat completion method to perform validation.
+ */
+public class SimpleChatCompletionServiceIntegrationValidator implements ServiceIntegrationValidator {
+    private static final List<String> TEST_INPUT = List.of("how big");
+
+    @Override
+    public void validate(InferenceService service, Model model, ActionListener<InferenceServiceResults> listener) {
+        var chatCompletionInput = new UnifiedChatInput(TEST_INPUT, USER_ROLE, false);
+        service.unifiedCompletionInfer(
+            model,
+            chatCompletionInput.getRequest(),
+            InferenceAction.Request.DEFAULT_TIMEOUT,
+            ActionListener.wrap(r -> {
+                if (r != null) {
+                    listener.onResponse(r);
+                } else {
+                    listener.onFailure(
+                        new ElasticsearchStatusException(
+                            "Could not complete inference endpoint creation as validation call to service returned null response.",
+                            RestStatus.BAD_REQUEST
+                        )
+                    );
+                }
+            }, e -> {
+                listener.onFailure(
+                    new ElasticsearchStatusException(
+                        "Could not complete inference endpoint creation as validation call to service threw an exception.",
+                        RestStatus.BAD_REQUEST,
+                        e
+                    )
+                );
+            })
+        );
+    }
+}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/BaseTransportInferenceActionTestCase.java

@@ -57,11 +57,15 @@ public abstract class BaseTransportInferenceActionTestCase<Request extends BaseI
     private BaseTransportInferenceAction<Request> action;
 
     protected static final String serviceId = "serviceId";
-    protected static final TaskType taskType = TaskType.COMPLETION;
+    protected final TaskType taskType;
     protected static final String inferenceId = "inferenceEntityId";
     protected InferenceServiceRegistry serviceRegistry;
     protected InferenceStats inferenceStats;
 
+    public BaseTransportInferenceActionTestCase(TaskType taskType) {
+        this.taskType = taskType;
+    }
+
     @Before
     public void setUp() throws Exception {
         super.setUp();
@@ -377,7 +381,7 @@ protected void mockModelAndServiceRegistry(InferenceService service) {
         when(serviceRegistry.getService(any())).thenReturn(Optional.of(service));
     }
 
-    protected void mockValidLicenseState(){
+    protected void mockValidLicenseState() {
         when(licenseState.isAllowed(InferencePlugin.INFERENCE_API_FEATURE)).thenReturn(true);
     }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/TransportInferenceActionTests.java

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.inference.InferenceServiceRegistry;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.license.MockLicenseState;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
@@ -20,6 +21,10 @@
 
 public class TransportInferenceActionTests extends BaseTransportInferenceActionTestCase<InferenceAction.Request> {
 
+    public TransportInferenceActionTests() {
+        super(TaskType.COMPLETION);
+    }
+
     @Override
     protected BaseTransportInferenceAction<InferenceAction.Request> createAction(
         TransportService transportService,

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/TransportUnifiedCompletionActionTests.java

@@ -33,6 +33,10 @@
 
 public class TransportUnifiedCompletionActionTests extends BaseTransportInferenceActionTestCase<UnifiedCompletionAction.Request> {
 
+    public TransportUnifiedCompletionActionTests() {
+        super(TaskType.CHAT_COMPLETION);
+    }
+
     @Override
     protected BaseTransportInferenceAction<UnifiedCompletionAction.Request> createAction(
         TransportService transportService,
@@ -71,7 +75,7 @@ public void testThrows_IncompatibleTaskTypeException_WhenUsingATextEmbeddingInfe
             assertThat(e, isA(ElasticsearchStatusException.class));
             assertThat(
                 e.getMessage(),
-                is("Incompatible task_type for unified API, the requested type [" + requestTaskType + "] must be one of [completion]")
+                is("Incompatible task_type for unified API, the requested type [" + requestTaskType + "] must be one of [chat_completion]")
             );
             assertThat(((ElasticsearchStatusException) e).status(), is(RestStatus.BAD_REQUEST));
         }));
@@ -96,7 +100,7 @@ public void testThrows_IncompatibleTaskTypeException_WhenUsingRequestIsAny_Model
             assertThat(e, isA(ElasticsearchStatusException.class));
             assertThat(
                 e.getMessage(),
-                is("Incompatible task_type for unified API, the requested type [" + requestTaskType + "] must be one of [completion]")
+                is("Incompatible task_type for unified API, the requested type [" + requestTaskType + "] must be one of [chat_completion]")
             );
             assertThat(((ElasticsearchStatusException) e).status(), is(RestStatus.BAD_REQUEST));
         }));

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiActionCreatorTests.java

@@ -44,7 +44,7 @@
 import static org.elasticsearch.xpack.inference.results.ChatCompletionResultsTests.buildExpectationCompletion;
 import static org.elasticsearch.xpack.inference.results.TextEmbeddingResultsTests.buildExpectationFloat;
 import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings;
-import static org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModelTests.createChatCompletionModel;
+import static org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModelTests.createCompletionModel;
 import static org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionRequestTaskSettingsTests.getChatCompletionRequestTaskSettingsMap;
 import static org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsModelTests.createModel;
 import static org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsRequestTaskSettingsTests.createRequestTaskSettingsMap;
@@ -325,7 +325,7 @@ public void testCreate_OpenAiChatCompletionModel() throws IOException {
 
             webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
 
-            var model = createChatCompletionModel(getUrl(webServer), "org", "secret", "model", "user");
+            var model = createCompletionModel(getUrl(webServer), "org", "secret", "model", "user");
             var actionCreator = new OpenAiActionCreator(sender, createWithEmptySettings(threadPool));
             var overriddenTaskSettings = getChatCompletionRequestTaskSettingsMap("overridden_user");
             var action = actionCreator.create(model, overriddenTaskSettings);
@@ -389,7 +389,7 @@ public void testCreate_OpenAiChatCompletionModel_WithoutUser() throws IOExceptio
 
             webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
 
-            var model = createChatCompletionModel(getUrl(webServer), "org", "secret", "model", null);
+            var model = createCompletionModel(getUrl(webServer), "org", "secret", "model", null);
             var actionCreator = new OpenAiActionCreator(sender, createWithEmptySettings(threadPool));
             var overriddenTaskSettings = getChatCompletionRequestTaskSettingsMap(null);
             var action = actionCreator.create(model, overriddenTaskSettings);
@@ -452,7 +452,7 @@ public void testCreate_OpenAiChatCompletionModel_WithoutOrganization() throws IO
 
             webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
 
-            var model = createChatCompletionModel(getUrl(webServer), null, "secret", "model", null);
+            var model = createCompletionModel(getUrl(webServer), null, "secret", "model", null);
             var actionCreator = new OpenAiActionCreator(sender, createWithEmptySettings(threadPool));
             var overriddenTaskSettings = getChatCompletionRequestTaskSettingsMap("overridden_user");
             var action = actionCreator.create(model, overriddenTaskSettings);
@@ -521,7 +521,7 @@ public void testCreate_OpenAiChatCompletionModel_FailsFromInvalidResponseFormat(
                 """;
             webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
 
-            var model = createChatCompletionModel(getUrl(webServer), null, "secret", "model", null);
+            var model = createCompletionModel(getUrl(webServer), null, "secret", "model", null);
             var actionCreator = new OpenAiActionCreator(sender, createWithEmptySettings(threadPool));
             var overriddenTaskSettings = getChatCompletionRequestTaskSettingsMap("overridden_user");
             var action = actionCreator.create(model, overriddenTaskSettings);

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiChatCompletionActionTests.java

@@ -51,7 +51,7 @@
 import static org.elasticsearch.xpack.inference.external.request.openai.OpenAiUtils.ORGANIZATION_HEADER;
 import static org.elasticsearch.xpack.inference.results.ChatCompletionResultsTests.buildExpectationCompletion;
 import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings;
-import static org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModelTests.createChatCompletionModel;
+import static org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModelTests.createCompletionModel;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
@@ -284,7 +284,7 @@ public void testExecute_ThrowsException_WhenInputIsGreaterThanOne() throws IOExc
     }
 
     private ExecutableAction createAction(String url, String org, String apiKey, String modelName, @Nullable String user, Sender sender) {
-        var model = createChatCompletionModel(url, org, apiKey, modelName, user);
+        var model = createCompletionModel(url, org, apiKey, modelName, user);
         var requestCreator = OpenAiCompletionRequestManager.of(model, threadPool);
         var errorMessage = constructFailedToSendRequestMessage(model.getServiceSettings().uri(), "OpenAI chat completions");
         return new SingleInputSenderExecutableAction(sender, requestCreator, errorMessage, "OpenAI chat completions");

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/elastic/ElasticInferenceServiceUnifiedChatCompletionRequestEntityTests.java

@@ -21,7 +21,7 @@
 import java.util.ArrayList;
 
 import static org.elasticsearch.xpack.inference.Utils.assertJsonEquals;
-import static org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModelTests.createChatCompletionModel;
+import static org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionModelTests.createCompletionModel;
 
 public class ElasticInferenceServiceUnifiedChatCompletionRequestEntityTests extends ESTestCase {
 
@@ -40,7 +40,7 @@ public void testModelUserFieldsSerialization() throws IOException {
         var unifiedRequest = UnifiedCompletionRequest.of(messageList);
 
         UnifiedChatInput unifiedChatInput = new UnifiedChatInput(unifiedRequest, true);
-        OpenAiChatCompletionModel model = createChatCompletionModel("test-url", "organizationId", "api-key", "test-endpoint", null);
+        OpenAiChatCompletionModel model = createCompletionModel("test-url", "organizationId", "api-key", "test-endpoint", null);
 
         OpenAiUnifiedChatCompletionRequestEntity entity = new OpenAiUnifiedChatCompletionRequestEntity(unifiedChatInput, model);