Add AI21 support to Inference Plugin (#131238)

* Add Ai21ActionVisitor interface

* Implement AI21 chat completion service and related components

* Make MODEL_FIELD and OBJECT_FIELD optional in OpenAiUnifiedStreamingProcessor

* Add changelog

* Add unit tests for Ai21 chat completion service and settings

* Add unit tests for Ai21ChatCompletionResponseHandler

* Make API_COMPLETIONS_PATH public and add unit tests for Ai21ChatCompletionRequestEntity and Ai21ChatCompletionRequest

* Refactor error handling and add unit tests for Ai21 service

* Refactor Ai21 service constructors and error handling for improved clarity and functionality

* Add Ai21ApiConstants class for API endpoint definitions

* Refactor Ai21ChatCompletionModel to use URIBuilder for endpoint URL construction

* Make fromString method public for improved accessibility

* Use Strings utility for null and blank check in ErrorResponse.fromString method

* Fix typos

---------

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

Author: Jan-Kazlouski-elastic

docs/changelog/131238.yaml

@@ -0,0 +1,5 @@
+pr: 113757
+summary: Added AI21 Completion and Chat Completion support to the Inference Plugin
+area: Machine Learning
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -356,6 +356,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion PIPELINE_TRACKING_INFO = def(9_131_0_00);
     public static final TransportVersion COMPONENT_TEMPLATE_TRACKING_INFO = def(9_132_0_00);
     public static final TransportVersion TO_CHILD_BLOCK_JOIN_QUERY = def(9_133_0_00);
+    public static final TransportVersion ML_INFERENCE_AI21_COMPLETION_ADDED = def(9_134_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -31,6 +31,7 @@
 import org.elasticsearch.xpack.inference.chunking.SentenceBoundaryChunkingSettings;
 import org.elasticsearch.xpack.inference.chunking.WordBoundaryChunkingSettings;
 import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
+import org.elasticsearch.xpack.inference.services.ai21.completion.Ai21ChatCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.AlibabaCloudSearchServiceSettings;
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.completion.AlibabaCloudSearchCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.completion.AlibabaCloudSearchCompletionTaskSettings;
@@ -178,6 +179,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         addVoyageAINamedWriteables(namedWriteables);
         addCustomNamedWriteables(namedWriteables);
         addLlamaNamedWriteables(namedWriteables);
+        addAi21NamedWriteables(namedWriteables);
 
         addUnifiedNamedWriteables(namedWriteables);
 
@@ -298,6 +300,17 @@ private static void addLlamaNamedWriteables(List<NamedWriteableRegistry.Entry> n
         // no task settings for Llama
     }
 
+    private static void addAi21NamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                Ai21ChatCompletionServiceSettings.NAME,
+                Ai21ChatCompletionServiceSettings::new
+            )
+        );
+        // no task settings for AI21
+    }
+
     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

@@ -113,6 +113,7 @@
 import org.elasticsearch.xpack.inference.rest.RestStreamInferenceAction;
 import org.elasticsearch.xpack.inference.rest.RestUpdateInferenceModelAction;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
+import org.elasticsearch.xpack.inference.services.ai21.Ai21Service;
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.AlibabaCloudSearchService;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.AmazonBedrockService;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.client.AmazonBedrockRequestSender;
@@ -413,6 +414,7 @@ public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {
             context -> new VoyageAIService(httpFactory.get(), serviceComponents.get(), context),
             context -> new DeepSeekService(httpFactory.get(), serviceComponents.get(), context),
             context -> new LlamaService(httpFactory.get(), serviceComponents.get(), context),
+            context -> new Ai21Service(httpFactory.get(), serviceComponents.get(), context),
             ElasticsearchInternalService::new,
             context -> new CustomService(httpFactory.get(), serviceComponents.get(), context)
         );

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/ErrorResponse.java

@@ -7,6 +7,10 @@
 
 package org.elasticsearch.xpack.inference.external.http.retry;
 
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.xpack.inference.external.http.HttpResult;
+
+import java.nio.charset.StandardCharsets;
 import java.util.Objects;
 
 public class ErrorResponse {
@@ -46,4 +50,39 @@ public boolean equals(Object o) {
     public int hashCode() {
         return Objects.hash(errorMessage, errorStructureFound);
     }
+
+    /**
+     * Creates an ErrorResponse from the given HttpResult.
+     * Attempts to read the body as a UTF-8 string and constructs an ErrorResponse.
+     * If reading fails, returns a generic UNDEFINED_ERROR.
+     *
+     * @param response the HttpResult containing the error response
+     * @return an ErrorResponse instance
+     */
+    public static ErrorResponse fromResponse(HttpResult response) {
+        try {
+            String errorMessage = new String(response.body(), StandardCharsets.UTF_8);
+            return new ErrorResponse(errorMessage);
+        } catch (Exception e) {
+            // swallow the error
+        }
+
+        return ErrorResponse.UNDEFINED_ERROR;
+    }
+
+    /**
+     * Parses a string response into an ErrorResponse.
+     * If the string is not blank, creates a new ErrorResponse with the string as the error message.
+     * If the string is blank, returns UNDEFINED_ERROR.
+     *
+     * @param response the error response as a string
+     * @return an ErrorResponse instance
+     */
+    public static ErrorResponse fromString(String response) {
+        if (Strings.isNullOrBlank(response) == false) {
+            return new ErrorResponse(response);
+        } else {
+            return ErrorResponse.UNDEFINED_ERROR;
+        }
+    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ai21/Ai21Model.java

@@ -0,0 +1,64 @@
+/*
+ * 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.ai21;
+
+import org.elasticsearch.inference.ModelConfigurations;
+import org.elasticsearch.inference.ModelSecrets;
+import org.elasticsearch.inference.ServiceSettings;
+import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
+import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
+import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
+
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Objects;
+
+/**
+ * Represents a AI21 model that can be used for inference tasks.
+ * This class extends RateLimitGroupingModel to handle rate limiting based on model and API key.
+ */
+public abstract class Ai21Model extends RateLimitGroupingModel {
+    protected URI uri;
+    protected RateLimitSettings rateLimitSettings;
+
+    protected Ai21Model(ModelConfigurations configurations, ModelSecrets secrets) {
+        super(configurations, secrets);
+    }
+
+    protected Ai21Model(RateLimitGroupingModel model, ServiceSettings serviceSettings) {
+        super(model, serviceSettings);
+    }
+
+    public URI uri() {
+        return this.uri;
+    }
+
+    @Override
+    public RateLimitSettings rateLimitSettings() {
+        return this.rateLimitSettings;
+    }
+
+    @Override
+    public int rateLimitGroupingHash() {
+        return Objects.hash(getServiceSettings().modelId(), getSecretSettings().apiKey());
+    }
+
+    // Needed for testing only
+    public void setURI(String newUri) {
+        try {
+            this.uri = new URI(newUri);
+        } catch (URISyntaxException e) {
+            // swallow any error
+        }
+    }
+
+    @Override
+    public DefaultSecretSettings getSecretSettings() {
+        return (DefaultSecretSettings) super.getSecretSettings();
+    }
+}