Add NVIDIA support to Inference Plugin (#132388)

Creation of new NVIDIA inference provider integration allowing:
text_embedding, completion, chat_completion, rerank

Additional changes:
* Refactor to replace CohereTruncation with Truncation across multiple classes for improved consistency and clarity

---------

Co-authored-by: elasticsearchmachine <[email protected]>

Author: Jan-Kazlouski-elastic

docs/changelog/132388.yaml

@@ -0,0 +1,5 @@
+pr: 132388
+summary: Added NVIDIA support to Inference Plugin
+area: Machine Learning
+type: enhancement
+issues: []

rest-api-spec/src/main/resources/rest-api-spec/api/inference.put_nvidia.json

@@ -0,0 +1,49 @@
+{
+  "inference.put_nvidia": {
+    "documentation": {
+      "url": "https://www.elastic.co/docs/api/doc/elasticsearch/operation/operation-inference-put-nvidia",
+      "description": "Create an Nvidia inference endpoint"
+    },
+    "stability": "stable",
+    "visibility": "public",
+    "headers": {
+      "accept": ["application/json"],
+      "content_type": ["application/json"]
+    },
+    "url": {
+      "paths": [
+        {
+          "path": "/_inference/{task_type}/{nvidia_inference_id}",
+          "methods": ["PUT"],
+          "parts": {
+            "task_type": {
+              "type": "enum",
+              "description": "The task type",
+              "options": [
+                "rerank",
+                "text_embedding",
+                "completion",
+                "chat_completion"
+              ]
+            },
+            "nvidia_inference_id": {
+              "type": "string",
+              "description": "The inference ID"
+            }
+          }
+        }
+      ]
+    },
+    "body": {
+      "description": "The inference endpoint's task and service settings",
+      "required": true
+    },
+    "params": {
+      "timeout": {
+        "type": "time",
+        "description": "Specifies the amount of time to wait for the inference endpoint to be created.",
+        "default": "30s"
+      }
+    }
+  }
+}

server/src/main/resources/transport/definitions/referable/ml_inference_nvidia_added.csv

@@ -0,0 +1 @@
+9233000

server/src/main/resources/transport/upper_bounds/9.3.csv

@@ -1 +1 @@
-downsample_add_multi_field_sources,9232000
+ml_inference_nvidia_added,9233000

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

@@ -70,6 +70,7 @@ public void testGetServicesWithoutTaskType() throws IOException {
                     "jinaai",
                     "llama",
                     "mistral",
+                    "nvidia",
                     "openai",
                     "openshift_ai",
                     "streaming_completion_test_service",
@@ -117,6 +118,7 @@ public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
                     "jinaai",
                     "llama",
                     "mistral",
+                    "nvidia",
                     "openai",
                     "openshift_ai",
                     "text_embedding_test_service",
@@ -143,6 +145,7 @@ public void testGetServicesWithRerankTaskType() throws IOException {
                     "elasticsearch",
                     "googlevertexai",
                     "jinaai",
+                    "nvidia",
                     "openshift_ai",
                     "test_reranking_service",
                     "voyageai",
@@ -177,6 +180,7 @@ public void testGetServicesWithCompletionTaskType() throws IOException {
                     "hugging_face",
                     "amazon_sagemaker",
                     "mistral",
+                    "nvidia",
                     "watsonxai"
                 ).toArray()
             )
@@ -200,6 +204,7 @@ public void testGetServicesWithChatCompletionTaskType() throws IOException {
                     "amazon_sagemaker",
                     "googlevertexai",
                     "mistral",
+                    "nvidia",
                     "watsonxai"
                 ).toArray()
             )

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

@@ -114,6 +114,10 @@
 import org.elasticsearch.xpack.inference.services.llama.embeddings.LlamaEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.mistral.completion.MistralChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsServiceSettings;
+import org.elasticsearch.xpack.inference.services.nvidia.completion.NvidiaChatCompletionServiceSettings;
+import org.elasticsearch.xpack.inference.services.nvidia.embeddings.NvidiaEmbeddingsServiceSettings;
+import org.elasticsearch.xpack.inference.services.nvidia.embeddings.NvidiaEmbeddingsTaskSettings;
+import org.elasticsearch.xpack.inference.services.nvidia.rerank.NvidiaRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.openai.completion.OpenAiChatCompletionTaskSettings;
 import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsServiceSettings;
@@ -181,6 +185,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         addLlamaNamedWriteables(namedWriteables);
         addAi21NamedWriteables(namedWriteables);
         addOpenShiftAiNamedWriteables(namedWriteables);
+        addNvidiaNamedWriteables(namedWriteables);
 
         addUnifiedNamedWriteables(namedWriteables);
 
@@ -329,6 +334,29 @@ private static void addAi21NamedWriteables(List<NamedWriteableRegistry.Entry> na
         // no task settings for AI21
     }
 
+    private static void addNvidiaNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                NvidiaChatCompletionServiceSettings.NAME,
+                NvidiaChatCompletionServiceSettings::new
+            )
+        );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                NvidiaEmbeddingsServiceSettings.NAME,
+                NvidiaEmbeddingsServiceSettings::new
+            )
+        );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(ServiceSettings.class, NvidiaRerankServiceSettings.NAME, NvidiaRerankServiceSettings::new)
+        );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(TaskSettings.class, NvidiaEmbeddingsTaskSettings.NAME, NvidiaEmbeddingsTaskSettings::new)
+        );
+    }
+
     private static void addAzureAiStudioNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
         namedWriteables.add(
             new NamedWriteableRegistry.Entry(

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

@@ -173,6 +173,7 @@
 import org.elasticsearch.xpack.inference.services.jinaai.JinaAIService;
 import org.elasticsearch.xpack.inference.services.llama.LlamaService;
 import org.elasticsearch.xpack.inference.services.mistral.MistralService;
+import org.elasticsearch.xpack.inference.services.nvidia.NvidiaService;
 import org.elasticsearch.xpack.inference.services.openai.OpenAiService;
 import org.elasticsearch.xpack.inference.services.openshiftai.OpenShiftAiService;
 import org.elasticsearch.xpack.inference.services.sagemaker.SageMakerClient;
@@ -590,6 +591,7 @@ public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {
             context -> new LlamaService(httpFactory.get(), serviceComponents.get(), context),
             context -> new Ai21Service(httpFactory.get(), serviceComponents.get(), context),
             context -> new OpenShiftAiService(httpFactory.get(), serviceComponents.get(), context),
+            context -> new NvidiaService(httpFactory.get(), serviceComponents.get(), context),
             ElasticsearchInternalService::new,
             context -> new CustomService(httpFactory.get(), serviceComponents.get(), context)
         );

…ce/services/cohere/CohereTruncation.java …k/inference/common/model/Truncation.javax-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereTruncation.java renamed to x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/model/Truncation.java x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereTruncation.java renamed to x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/common/model/Truncation.java

@@ -5,20 +5,19 @@
  * 2.0.
  */
 
-package org.elasticsearch.xpack.inference.services.cohere;
+package org.elasticsearch.xpack.inference.common.model;
 
 import java.util.EnumSet;
 import java.util.Locale;
 
 /**
- * Defines the type of truncation for a cohere request. The specified value determines how the Cohere API will handle inputs
+ * Defines the type of truncation for an embeddings request. The specified value determines how the provider's API will handle inputs
  * longer than the maximum token length.
- *
  * <p>
- * <a href="https://docs.cohere.com/reference/embed">See api docs for details.</a>
+ * <a href="https://docs.cohere.com/reference/embed">Details can be found in Cohere embeddings API docs.</a>
  * </p>
  */
-public enum CohereTruncation {
+public enum Truncation {
     /**
      * When the input exceeds the maximum input token length an error will be returned.
      */
@@ -32,14 +31,14 @@ public enum CohereTruncation {
      */
     END;
 
-    public static final EnumSet<CohereTruncation> ALL = EnumSet.allOf(CohereTruncation.class);
+    public static final EnumSet<Truncation> ALL = EnumSet.allOf(Truncation.class);
 
     @Override
     public String toString() {
         return name().toLowerCase(Locale.ROOT);
     }
 
-    public static CohereTruncation fromString(String name) {
+    public static Truncation fromString(String name) {
         return valueOf(name.trim().toUpperCase(Locale.ROOT));
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockService.java

@@ -370,7 +370,7 @@ private static void checkProviderForTask(TaskType taskType, AmazonBedrockProvide
     }
 
     private static void checkTaskSettingsForTextEmbeddingModel(AmazonBedrockEmbeddingsModel model) {
-        if (model.provider() != AmazonBedrockProvider.COHERE && model.getTaskSettings().cohereTruncation() != null) {
+        if (model.provider() != AmazonBedrockProvider.COHERE && model.getTaskSettings().truncation() != null) {
             throw new ElasticsearchStatusException(
                 "The [{}] task type for provider [{}] does not allow [truncate] field",
                 RestStatus.BAD_REQUEST,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/embeddings/AmazonBedrockEmbeddingsTaskSettings.java

@@ -15,7 +15,7 @@
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.TaskSettings;
 import org.elasticsearch.xcontent.XContentBuilder;
-import org.elasticsearch.xpack.inference.services.cohere.CohereTruncation;
+import org.elasticsearch.xpack.inference.common.model.Truncation;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -24,8 +24,8 @@
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalEnum;
 import static org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockConstants.TRUNCATE_FIELD;
 
-public record AmazonBedrockEmbeddingsTaskSettings(@Nullable CohereTruncation cohereTruncation) implements TaskSettings {
-    public static final AmazonBedrockEmbeddingsTaskSettings EMPTY = new AmazonBedrockEmbeddingsTaskSettings((CohereTruncation) null);
+public record AmazonBedrockEmbeddingsTaskSettings(@Nullable Truncation truncation) implements TaskSettings {
+    public static final AmazonBedrockEmbeddingsTaskSettings EMPTY = new AmazonBedrockEmbeddingsTaskSettings((Truncation) null);
     public static final String NAME = "amazon_bedrock_embeddings_task_settings";
     private static final TransportVersion AMAZON_BEDROCK_TASK_SETTINGS = TransportVersion.fromName("amazon_bedrock_task_settings");
 
@@ -36,36 +36,36 @@ public static AmazonBedrockEmbeddingsTaskSettings fromMap(Map<String, Object> ma
 
         ValidationException validationException = new ValidationException();
 
-        var cohereTruncation = extractOptionalEnum(
+        var extractedTruncation = extractOptionalEnum(
             map,
             TRUNCATE_FIELD,
             ModelConfigurations.TASK_SETTINGS,
-            CohereTruncation::fromString,
-            CohereTruncation.ALL,
+            Truncation::fromString,
+            Truncation.ALL,
             validationException
         );
 
         if (validationException.validationErrors().isEmpty() == false) {
             throw validationException;
         }
 
-        return new AmazonBedrockEmbeddingsTaskSettings(cohereTruncation);
+        return new AmazonBedrockEmbeddingsTaskSettings(extractedTruncation);
     }
 
     public AmazonBedrockEmbeddingsTaskSettings(StreamInput in) throws IOException {
-        this(in.readOptionalEnum(CohereTruncation.class));
+        this(in.readOptionalEnum(Truncation.class));
     }
 
     @Override
     public boolean isEmpty() {
-        return cohereTruncation() == null;
+        return truncation() == null;
     }
 
     @Override
     public AmazonBedrockEmbeddingsTaskSettings updatedTaskSettings(Map<String, Object> newSettings) {
         var newTaskSettings = fromMap(new HashMap<>(newSettings));
 
-        return new AmazonBedrockEmbeddingsTaskSettings(firstNonNullOrNull(newTaskSettings.cohereTruncation(), cohereTruncation()));
+        return new AmazonBedrockEmbeddingsTaskSettings(firstNonNullOrNull(newTaskSettings.truncation(), truncation()));
     }
 
     private static <T> T firstNonNullOrNull(T first, T second) {
@@ -90,14 +90,14 @@ public boolean supportsVersion(TransportVersion version) {
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeOptionalEnum(cohereTruncation());
+        out.writeOptionalEnum(truncation());
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
-        if (cohereTruncation != null) {
-            builder.field(TRUNCATE_FIELD, cohereTruncation);
+        if (truncation != null) {
+            builder.field(TRUNCATE_FIELD, truncation);
         }
         return builder.endObject();
     }