Add support for streaming chat completion task for HuggingFace

Author: Jan-Kazlouski-elastic

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

@@ -127,7 +127,6 @@
 import org.elasticsearch.xpack.inference.services.googleaistudio.GoogleAiStudioService;
 import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiService;
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceService;
-import org.elasticsearch.xpack.inference.services.huggingface.completion.HuggingFaceChatCompletionService;
 import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserService;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.IbmWatsonxService;
 import org.elasticsearch.xpack.inference.services.jinaai.JinaAIService;
@@ -362,7 +361,6 @@ public void loadExtensions(ExtensionLoader loader) {
     public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {
         return List.of(
             context -> new HuggingFaceElserService(httpFactory.get(), serviceComponents.get()),
-            context -> new HuggingFaceChatCompletionService(httpFactory.get(), serviceComponents.get()),
             context -> new HuggingFaceService(httpFactory.get(), serviceComponents.get()),
             context -> new OpenAiService(httpFactory.get(), serviceComponents.get()),
             context -> new CohereService(httpFactory.get(), serviceComponents.get()),

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceCompletionRequestManager.java renamed to x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceChatCompletionRequestManager.java

@@ -23,15 +23,22 @@
 import java.util.Objects;
 import java.util.function.Supplier;
 
-public class HuggingFaceCompletionRequestManager extends HuggingFaceRequestManager {
-    private static final Logger logger = LogManager.getLogger(HuggingFaceCompletionRequestManager.class);
+/**
+ * Manages the execution of chat completion requests for Hugging Face models.
+ * <p>
+ * This class is responsible for creating and executing requests to Hugging Face's chat completion API.
+ * It extends {@link HuggingFaceRequestManager} to provide specific functionality for chat completion models.
+ * </p>
+ */
+public class HuggingFaceChatCompletionRequestManager extends HuggingFaceRequestManager {
+    private static final Logger logger = LogManager.getLogger(HuggingFaceChatCompletionRequestManager.class);
 
-    public static HuggingFaceCompletionRequestManager of(
+    public static HuggingFaceChatCompletionRequestManager of(
         HuggingFaceChatCompletionModel model,
         ResponseHandler responseHandler,
         ThreadPool threadPool
     ) {
-        return new HuggingFaceCompletionRequestManager(
+        return new HuggingFaceChatCompletionRequestManager(
             Objects.requireNonNull(model),
             Objects.requireNonNull(responseHandler),
             Objects.requireNonNull(threadPool)
@@ -41,7 +48,7 @@ public static HuggingFaceCompletionRequestManager of(
     private final HuggingFaceChatCompletionModel model;
     private final ResponseHandler responseHandler;
 
-    private HuggingFaceCompletionRequestManager(
+    private HuggingFaceChatCompletionRequestManager(
         HuggingFaceChatCompletionModel model,
         ResponseHandler responseHandler,
         ThreadPool threadPool

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceEmbeddingsRequestManager.java

@@ -18,17 +18,30 @@
 import org.elasticsearch.xpack.inference.external.http.sender.EmbeddingsInput;
 import org.elasticsearch.xpack.inference.external.http.sender.ExecutableInferenceRequest;
 import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
-import org.elasticsearch.xpack.inference.services.huggingface.request.embeddings.HuggingFaceInferenceRequest;
+import org.elasticsearch.xpack.inference.services.huggingface.request.embeddings.HuggingFaceEmbeddingsRequest;
 
 import java.util.List;
 import java.util.Objects;
 import java.util.function.Supplier;
 
 import static org.elasticsearch.xpack.inference.common.Truncator.truncate;
 
+/**
+ * This class is responsible for managing requests to the Hugging Face API for generating embeddings.
+ * It handles the execution of requests, including truncation of input data and response handling.
+ */
 public class HuggingFaceEmbeddingsRequestManager extends HuggingFaceRequestManager {
     private static final Logger logger = LogManager.getLogger(HuggingFaceEmbeddingsRequestManager.class);
 
+    /**
+     * Creates a new instance of HuggingFaceEmbeddingsRequestManager.
+     *
+     * @param model           The Hugging Face model to be used for generating embeddings.
+     * @param responseHandler  The response handler for processing the API responses.
+     * @param truncator       The truncator for handling input data truncation.
+     * @param threadPool      The thread pool for executing requests.
+     * @return A new instance of HuggingFaceEmbeddingsRequestManager.
+     */
     public static HuggingFaceEmbeddingsRequestManager of(
         HuggingFaceModel model,
         ResponseHandler responseHandler,
@@ -68,7 +81,7 @@ public void execute(
     ) {
         List<String> docsInput = EmbeddingsInput.of(inferenceInputs).getStringInputs();
         var truncatedInput = truncate(docsInput, model.getTokenLimit());
-        var request = new HuggingFaceInferenceRequest(truncator, truncatedInput, model);
+        var request = new HuggingFaceEmbeddingsRequest(truncator, truncatedInput, model);
 
         execute(new ExecutableInferenceRequest(requestSender, logger, request, responseHandler, hasRequestCompletedFunction, listener));
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceService.java

@@ -43,19 +43,24 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.elasticsearch.xpack.inference.services.ServiceFields.URL;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
-import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwUnsupportedUnifiedCompletionOperation;
 
+/**
+ * This class is responsible for managing the Hugging Face inference service.
+ * It handles the creation of models, chunked inference, and unified completion inference.
+ */
 public class HuggingFaceService extends HuggingFaceBaseService {
     public static final String NAME = "hugging_face";
 
     private static final String SERVICE_NAME = "Hugging Face";
     private static final EnumSet<TaskType> SUPPORTED_TASK_TYPES = EnumSet.of(
         TaskType.TEXT_EMBEDDING,
         TaskType.SPARSE_EMBEDDING,
-        TaskType.COMPLETION
+        TaskType.COMPLETION,
+        TaskType.CHAT_COMPLETION
     );
 
     public HuggingFaceService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) {
@@ -152,7 +157,21 @@ protected void doUnifiedCompletionInfer(
         TimeValue timeout,
         ActionListener<InferenceServiceResults> listener
     ) {
-        throwUnsupportedUnifiedCompletionOperation(NAME);
+        if (model instanceof HuggingFaceChatCompletionModel == false) {
+            listener.onFailure(createInvalidModelException(model));
+            return;
+        }
+        HuggingFaceChatCompletionModel huggingFaceChatCompletionModel = (HuggingFaceChatCompletionModel) model;
+        var actionCreator = new HuggingFaceActionCreator(getSender(), getServiceComponents());
+        var overriddenModel = HuggingFaceChatCompletionModel.of(huggingFaceChatCompletionModel, inputs.getRequest());
+        var action = overriddenModel.accept(actionCreator);
+
+        action.execute(inputs, timeout, listener);
+    }
+
+    @Override
+    public Set<TaskType> supportedStreamingTasks() {
+        return EnumSet.of(TaskType.COMPLETION, TaskType.CHAT_COMPLETION);
     }
 
     @Override
@@ -180,6 +199,9 @@ public static InferenceServiceConfiguration get() {
             return configuration.getOrCompute();
         }
 
+        private Configuration() {
+        }
+
         private static final LazyInitializable<InferenceServiceConfiguration, RuntimeException> configuration = new LazyInitializable<>(
             () -> {
                 var configurationMap = new HashMap<String, SettingsConfiguration>();

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/action/HuggingFaceActionCreator.java

@@ -11,15 +11,16 @@
 import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
-import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceCompletionRequestManager;
+import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceChatCompletionRequestManager;
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceEmbeddingsRequestManager;
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceResponseHandler;
 import org.elasticsearch.xpack.inference.services.huggingface.completion.HuggingFaceChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserModel;
 import org.elasticsearch.xpack.inference.services.huggingface.embeddings.HuggingFaceEmbeddingsModel;
-import org.elasticsearch.xpack.inference.services.huggingface.response.HuggingFaceChatCompletionResponseEntity;
 import org.elasticsearch.xpack.inference.services.huggingface.response.HuggingFaceElserResponseEntity;
 import org.elasticsearch.xpack.inference.services.huggingface.response.HuggingFaceEmbeddingsResponseEntity;
+import org.elasticsearch.xpack.inference.services.openai.OpenAiUnifiedChatCompletionResponseHandler;
+import org.elasticsearch.xpack.inference.services.openai.response.OpenAiChatCompletionResponseEntity;
 
 import java.util.Objects;
 
@@ -71,12 +72,12 @@ public ExecutableAction create(HuggingFaceElserModel model) {
 
     @Override
     public ExecutableAction create(HuggingFaceChatCompletionModel model) {
-        var responseHandler = new HuggingFaceResponseHandler(
+        var responseHandler = new OpenAiUnifiedChatCompletionResponseHandler(
             "hugging face chat completion",
-            HuggingFaceChatCompletionResponseEntity::fromResponse
+            OpenAiChatCompletionResponseEntity::fromResponse
         );
 
-        var requestCreator = HuggingFaceCompletionRequestManager.of(model, responseHandler, serviceComponents.threadPool());
+        var requestCreator = HuggingFaceChatCompletionRequestManager.of(model, responseHandler, serviceComponents.threadPool());
         var errorMessage = format(FAILED_TO_SEND_REQUEST_ERROR_MESSAGE, "CHAT COMPLETION", model.getInferenceEntityId());
         return new SenderExecutableAction(sender, requestCreator, errorMessage);
     }