Update model to return correct model for CHAT_COMPLETION task type (#120326)

* Update model to return correct model for CHAT_COMPLETION task type

* Update docs/changelog/120326.yaml

* Delete docs/changelog/120326.yaml

* Fixing chat completion functionality

* Fixing tests

* naming

* Fixing tests

* Adding more tests

---------

Co-authored-by: Jonathan Buttner <[email protected]>
Co-authored-by: Jonathan Buttner <[email protected]>

Author: jaybcee

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);
     }
 }