Fixing chat completion functionality

elastic · Jan 16, 2025 · 08e0f50 · 08e0f50
1 parent f264c67
commit 08e0f50
Show file tree

Hide file tree

Showing 10 changed files with 76 additions and 12 deletions.
diff --git a/.../src/main/java/org/elasticsearch/xpack/core/inference/action/UnifiedCompletionAction.java b/.../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;
             }
 

diff --git a/...va/org/elasticsearch/xpack/core/inference/action/UnifiedCompletionActionRequestTests.java b/...va/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() {

diff --git a/...ervice-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java b/...ervice-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java
@@ -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 -> {

diff --git a/.../java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java b/.../java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java
@@ -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()),

diff --git a/...a/org/elasticsearch/xpack/inference/action/TransportUnifiedCompletionInferenceAction.java b/...a/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()
         );
     }
 

diff --git a/...rg/elasticsearch/xpack/inference/external/http/sender/OpenAiCompletionRequestManager.java b/...rg/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));

diff --git a/...ain/java/org/elasticsearch/xpack/inference/services/validation/ModelValidatorBuilder.java b/...ain/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());
             }

diff --git a/.../xpack/inference/services/validation/SimpleChatCompletionServiceIntegrationValidator.java b/.../xpack/inference/services/validation/SimpleChatCompletionServiceIntegrationValidator.java
@@ -0,0 +1,54 @@
+
+/*
+ * 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.
+ */
+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 a = new UnifiedChatInput(TEST_INPUT, USER_ROLE, false);
+        service.unifiedCompletionInfer(model, a.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
+                )
+            );
+        }));
+    }
+}
diff --git a/.../java/org/elasticsearch/xpack/inference/action/TransportUnifiedCompletionActionTests.java b/.../java/org/elasticsearch/xpack/inference/action/TransportUnifiedCompletionActionTests.java
@@ -71,7 +71,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 +96,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));
         }));

diff --git a/...ava/org/elasticsearch/xpack/inference/services/validation/ModelValidatorBuilderTests.java b/...ava/org/elasticsearch/xpack/inference/services/validation/ModelValidatorBuilderTests.java
@@ -35,6 +35,8 @@ private Map<TaskType, Class<? extends ModelValidator>> taskTypeToModelValidatorC
             SimpleModelValidator.class,
             TaskType.COMPLETION,
             ChatCompletionModelValidator.class,
+            TaskType.CHAT_COMPLETION,
+            ChatCompletionModelValidator.class,
             TaskType.ANY,
             SimpleModelValidator.class
         );