Add Hugging Face Chat Completion support to Inference Plugin

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

@@ -78,6 +78,7 @@
 import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceServiceSettings;
+import org.elasticsearch.xpack.inference.services.huggingface.completion.HuggingFaceChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserServiceSettings;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.embeddings.IbmWatsonxEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.rerank.IbmWatsonxRerankServiceSettings;
@@ -353,6 +354,13 @@ private static void addHuggingFaceNamedWriteables(List<NamedWriteableRegistry.En
         namedWriteables.add(
             new NamedWriteableRegistry.Entry(ServiceSettings.class, HuggingFaceServiceSettings.NAME, HuggingFaceServiceSettings::new)
         );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                HuggingFaceChatCompletionServiceSettings.NAME,
+                HuggingFaceChatCompletionServiceSettings::new
+            )
+        );
     }
 
     private static void addGoogleAiStudioNamedWritables(List<NamedWriteableRegistry.Entry> namedWriteables) {

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

@@ -127,6 +127,7 @@
 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;
@@ -361,6 +362,7 @@ 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

@@ -0,0 +1,66 @@
+/*
+ * 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.huggingface;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.external.http.retry.RequestSender;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
+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.huggingface.completion.HuggingFaceChatCompletionModel;
+import org.elasticsearch.xpack.inference.services.huggingface.request.completion.HuggingFaceUnifiedChatCompletionRequest;
+
+import java.util.Objects;
+import java.util.function.Supplier;
+
+public class HuggingFaceCompletionRequestManager extends HuggingFaceRequestManager {
+    private static final Logger logger = LogManager.getLogger(HuggingFaceCompletionRequestManager.class);
+
+    public static HuggingFaceCompletionRequestManager of(
+        HuggingFaceChatCompletionModel model,
+        ResponseHandler responseHandler,
+        ThreadPool threadPool
+    ) {
+        return new HuggingFaceCompletionRequestManager(
+            Objects.requireNonNull(model),
+            Objects.requireNonNull(responseHandler),
+            Objects.requireNonNull(threadPool)
+        );
+    }
+
+    private final HuggingFaceChatCompletionModel model;
+    private final ResponseHandler responseHandler;
+
+    private HuggingFaceCompletionRequestManager(
+        HuggingFaceChatCompletionModel model,
+        ResponseHandler responseHandler,
+        ThreadPool threadPool
+    ) {
+        super(model, threadPool);
+        this.model = model;
+        this.responseHandler = responseHandler;
+    }
+
+    @Override
+    public void execute(
+        InferenceInputs inferenceInputs,
+        RequestSender requestSender,
+        Supplier<Boolean> hasRequestCompletedFunction,
+        ActionListener<InferenceServiceResults> listener
+    ) {
+        var chatCompletionInput = inferenceInputs.castTo(UnifiedChatInput.class);
+        HuggingFaceUnifiedChatCompletionRequest request = new HuggingFaceUnifiedChatCompletionRequest(chatCompletionInput, model);
+
+        execute(new ExecutableInferenceRequest(requestSender, logger, request, responseHandler, hasRequestCompletedFunction, listener));
+    }
+}

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

@@ -0,0 +1,75 @@
+/*
+ * 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.huggingface;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.common.Truncator;
+import org.elasticsearch.xpack.inference.external.http.retry.RequestSender;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
+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 java.util.List;
+import java.util.Objects;
+import java.util.function.Supplier;
+
+import static org.elasticsearch.xpack.inference.common.Truncator.truncate;
+
+public class HuggingFaceEmbeddingsRequestManager extends HuggingFaceRequestManager {
+    private static final Logger logger = LogManager.getLogger(HuggingFaceEmbeddingsRequestManager.class);
+
+    public static HuggingFaceEmbeddingsRequestManager of(
+        HuggingFaceModel model,
+        ResponseHandler responseHandler,
+        Truncator truncator,
+        ThreadPool threadPool
+    ) {
+        return new HuggingFaceEmbeddingsRequestManager(
+            Objects.requireNonNull(model),
+            Objects.requireNonNull(responseHandler),
+            Objects.requireNonNull(truncator),
+            Objects.requireNonNull(threadPool)
+        );
+    }
+
+    private final HuggingFaceModel model;
+    private final ResponseHandler responseHandler;
+    private final Truncator truncator;
+
+    private HuggingFaceEmbeddingsRequestManager(
+        HuggingFaceModel model,
+        ResponseHandler responseHandler,
+        Truncator truncator,
+        ThreadPool threadPool
+    ) {
+        super(model, threadPool);
+        this.model = model;
+        this.responseHandler = responseHandler;
+        this.truncator = truncator;
+    }
+
+    @Override
+    public void execute(
+        InferenceInputs inferenceInputs,
+        RequestSender requestSender,
+        Supplier<Boolean> hasRequestCompletedFunction,
+        ActionListener<InferenceServiceResults> listener
+    ) {
+        List<String> docsInput = EmbeddingsInput.of(inferenceInputs).getStringInputs();
+        var truncatedInput = truncate(docsInput, model.getTokenLimit());
+        var request = new HuggingFaceInferenceRequest(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/HuggingFaceRequestManager.java

@@ -7,66 +7,12 @@
 
 package org.elasticsearch.xpack.inference.services.huggingface;
 
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
-import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.threadpool.ThreadPool;
-import org.elasticsearch.xpack.inference.common.Truncator;
-import org.elasticsearch.xpack.inference.external.http.retry.RequestSender;
-import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
 import org.elasticsearch.xpack.inference.external.http.sender.BaseRequestManager;
-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.HuggingFaceInferenceRequest;
 
-import java.util.List;
-import java.util.Objects;
-import java.util.function.Supplier;
-
-import static org.elasticsearch.xpack.inference.common.Truncator.truncate;
-
-public class HuggingFaceRequestManager extends BaseRequestManager {
-    private static final Logger logger = LogManager.getLogger(HuggingFaceRequestManager.class);
-
-    public static HuggingFaceRequestManager of(
-        HuggingFaceModel model,
-        ResponseHandler responseHandler,
-        Truncator truncator,
-        ThreadPool threadPool
-    ) {
-        return new HuggingFaceRequestManager(
-            Objects.requireNonNull(model),
-            Objects.requireNonNull(responseHandler),
-            Objects.requireNonNull(truncator),
-            Objects.requireNonNull(threadPool)
-        );
-    }
-
-    private final HuggingFaceModel model;
-    private final ResponseHandler responseHandler;
-    private final Truncator truncator;
-
-    private HuggingFaceRequestManager(HuggingFaceModel model, ResponseHandler responseHandler, Truncator truncator, ThreadPool threadPool) {
+public abstract class HuggingFaceRequestManager extends BaseRequestManager {
+    protected HuggingFaceRequestManager(HuggingFaceModel model, ThreadPool threadPool) {
         super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitServiceSettings().rateLimitSettings());
-        this.model = model;
-        this.responseHandler = responseHandler;
-        this.truncator = truncator;
-    }
-
-    @Override
-    public void execute(
-        InferenceInputs inferenceInputs,
-        RequestSender requestSender,
-        Supplier<Boolean> hasRequestCompletedFunction,
-        ActionListener<InferenceServiceResults> listener
-    ) {
-        List<String> docsInput = EmbeddingsInput.of(inferenceInputs).getStringInputs();
-        var truncatedInput = truncate(docsInput, model.getTokenLimit());
-        var request = new HuggingFaceInferenceRequest(truncator, truncatedInput, model);
-
-        execute(new ExecutableInferenceRequest(requestSender, logger, request, responseHandler, hasRequestCompletedFunction, listener));
     }
 
     record RateLimitGrouping(int accountHash) {

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

@@ -33,6 +33,7 @@
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.ServiceUtils;
 import org.elasticsearch.xpack.inference.services.huggingface.action.HuggingFaceActionCreator;
+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.settings.DefaultSecretSettings;
@@ -51,7 +52,11 @@ 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> supportedTaskTypes = EnumSet.of(TaskType.TEXT_EMBEDDING, TaskType.SPARSE_EMBEDDING);
+    private static final EnumSet<TaskType> SUPPORTED_TASK_TYPES = EnumSet.of(
+        TaskType.TEXT_EMBEDDING,
+        TaskType.SPARSE_EMBEDDING,
+        TaskType.COMPLETION
+    );
 
     public HuggingFaceService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) {
         super(factory, serviceComponents);
@@ -78,6 +83,14 @@ protected HuggingFaceModel createModel(
                 context
             );
             case SPARSE_EMBEDDING -> new HuggingFaceElserModel(inferenceEntityId, taskType, NAME, serviceSettings, secretSettings, context);
+            case CHAT_COMPLETION, COMPLETION -> new HuggingFaceChatCompletionModel(
+                inferenceEntityId,
+                taskType,
+                NAME,
+                serviceSettings,
+                secretSettings,
+                context
+            );
             default -> throw new ElasticsearchStatusException(failureMessage, RestStatus.BAD_REQUEST);
         };
     }
@@ -149,7 +162,7 @@ public InferenceServiceConfiguration getConfiguration() {
 
     @Override
     public EnumSet<TaskType> supportedTaskTypes() {
-        return supportedTaskTypes;
+        return SUPPORTED_TASK_TYPES;
     }
 
     @Override
@@ -173,7 +186,7 @@ public static InferenceServiceConfiguration get() {
 
                 configurationMap.put(
                     URL,
-                    new SettingsConfiguration.Builder(supportedTaskTypes).setDefaultValue("https://api.openai.com/v1/embeddings")
+                    new SettingsConfiguration.Builder(SUPPORTED_TASK_TYPES).setDefaultValue("https://api.openai.com/v1/embeddings")
                         .setDescription("The URL endpoint to use for the requests.")
                         .setLabel("URL")
                         .setRequired(true)
@@ -183,12 +196,12 @@ public static InferenceServiceConfiguration get() {
                         .build()
                 );
 
-                configurationMap.putAll(DefaultSecretSettings.toSettingsConfiguration(supportedTaskTypes));
-                configurationMap.putAll(RateLimitSettings.toSettingsConfiguration(supportedTaskTypes));
+                configurationMap.putAll(DefaultSecretSettings.toSettingsConfiguration(SUPPORTED_TASK_TYPES));
+                configurationMap.putAll(RateLimitSettings.toSettingsConfiguration(SUPPORTED_TASK_TYPES));
 
                 return new InferenceServiceConfiguration.Builder().setService(NAME)
                     .setName(SERVICE_NAME)
-                    .setTaskTypes(supportedTaskTypes)
+                    .setTaskTypes(SUPPORTED_TASK_TYPES)
                     .setConfigurations(configurationMap)
                     .build();
             }

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

@@ -11,10 +11,13 @@
 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.HuggingFaceRequestManager;
+import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceCompletionRequestManager;
+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;
 
@@ -26,6 +29,9 @@
  * Provides a way to construct an {@link ExecutableAction} using the visitor pattern based on the hugging face model type.
  */
 public class HuggingFaceActionCreator implements HuggingFaceActionVisitor {
+
+    private static final String FAILED_TO_SEND_REQUEST_ERROR_MESSAGE =
+        "Failed to send Hugging Face %s request from inference entity id [%s]";
     private final Sender sender;
     private final ServiceComponents serviceComponents;
 
@@ -40,34 +46,38 @@ public ExecutableAction create(HuggingFaceEmbeddingsModel model) {
             "hugging face text embeddings",
             HuggingFaceEmbeddingsResponseEntity::fromResponse
         );
-        var requestCreator = HuggingFaceRequestManager.of(
+        var requestCreator = HuggingFaceEmbeddingsRequestManager.of(
             model,
             responseHandler,
             serviceComponents.truncator(),
             serviceComponents.threadPool()
         );
-        var errorMessage = format(
-            "Failed to send Hugging Face %s request from inference entity id [%s]",
-            "text embeddings",
-            model.getInferenceEntityId()
-        );
+        var errorMessage = format(FAILED_TO_SEND_REQUEST_ERROR_MESSAGE, "text embeddings", model.getInferenceEntityId());
         return new SenderExecutableAction(sender, requestCreator, errorMessage);
     }
 
     @Override
     public ExecutableAction create(HuggingFaceElserModel model) {
         var responseHandler = new HuggingFaceResponseHandler("hugging face elser", HuggingFaceElserResponseEntity::fromResponse);
-        var requestCreator = HuggingFaceRequestManager.of(
+        var requestCreator = HuggingFaceEmbeddingsRequestManager.of(
             model,
             responseHandler,
             serviceComponents.truncator(),
             serviceComponents.threadPool()
         );
-        var errorMessage = format(
-            "Failed to send Hugging Face %s request from inference entity id [%s]",
-            "ELSER",
-            model.getInferenceEntityId()
+        var errorMessage = format(FAILED_TO_SEND_REQUEST_ERROR_MESSAGE, "ELSER", model.getInferenceEntityId());
+        return new SenderExecutableAction(sender, requestCreator, errorMessage);
+    }
+
+    @Override
+    public ExecutableAction create(HuggingFaceChatCompletionModel model) {
+        var responseHandler = new HuggingFaceResponseHandler(
+            "hugging face chat completion",
+            HuggingFaceChatCompletionResponseEntity::fromResponse
         );
+
+        var requestCreator = HuggingFaceCompletionRequestManager.of(model, responseHandler, serviceComponents.threadPool());
+        var errorMessage = format(FAILED_TO_SEND_REQUEST_ERROR_MESSAGE, "CHAT COMPLETION", model.getInferenceEntityId());
         return new SenderExecutableAction(sender, requestCreator, errorMessage);
     }
 }