Add Llama support to Inference Plugin

docs/changelog/130092.yaml

@@ -0,0 +1,5 @@
+pr: 130092
+summary: "Added Llama provider support to the Inference Plugin"
+area: Machine Learning
+type: enhancement
+issues: []

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

@@ -212,6 +212,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_PROFILE_INCLUDE_PLAN_8_19 = def(8_841_0_62);
     public static final TransportVersion ESQL_SPLIT_ON_BIG_VALUES_8_19 = def(8_841_0_63);
     public static final TransportVersion ESQL_FIXED_INDEX_LIKE_8_19 = def(8_841_0_64);
+    public static final TransportVersion ML_INFERENCE_LLAMA_ADDED_8_19 = def(8_841_0_65);
     public static final TransportVersion V_9_0_0 = def(9_000_0_09);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_1 = def(9_000_0_10);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_2 = def(9_000_0_11);
@@ -334,7 +335,13 @@ static TransportVersion def(int id) {
     public static final TransportVersion PROJECT_STATE_REGISTRY_RECORDS_DELETIONS = def(9_113_0_00);
     public static final TransportVersion ESQL_SERIALIZE_TIMESERIES_FIELD_TYPE = def(9_114_0_00);
     public static final TransportVersion ML_INFERENCE_IBM_WATSONX_COMPLETION_ADDED = def(9_115_0_00);
-    public static final TransportVersion ML_INFERENCE_LLAMA_ADDED = def(9_116_0_00);
+    public static final TransportVersion ESQL_SPLIT_ON_BIG_VALUES = def(9_116_0_00);
+    public static final TransportVersion ESQL_LOCAL_RELATION_WITH_NEW_BLOCKS = def(9_117_0_00);
+    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_EMBEDDING_TYPE = def(9_118_0_00);
+    public static final TransportVersion ESQL_FIXED_INDEX_LIKE = def(9_119_0_00);
+    public static final TransportVersion LOOKUP_JOIN_CCS = def(9_120_0_00);
+    public static final TransportVersion NODE_USAGE_STATS_FOR_THREAD_POOLS_IN_CLUSTER_INFO = def(9_121_0_00);
+    public static final TransportVersion ML_INFERENCE_LLAMA_ADDED = def(9_122_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/LlamaModel.java

@@ -12,7 +12,9 @@
 import org.elasticsearch.inference.ModelSecrets;
 import org.elasticsearch.inference.SecretSettings;
 import org.elasticsearch.inference.ServiceSettings;
+import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
+import org.elasticsearch.xpack.inference.services.llama.action.LlamaActionVisitor;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
@@ -26,7 +28,6 @@
  * This class extends RateLimitGroupingModel and provides common functionality for Llama models.
  */
 public abstract class LlamaModel extends RateLimitGroupingModel {
-    protected String modelId;
     protected URI uri;
     protected RateLimitSettings rateLimitSettings;
 
@@ -49,10 +50,6 @@ protected LlamaModel(RateLimitGroupingModel model, ServiceSettings serviceSettin
         super(model, serviceSettings);
     }
 
-    public String model() {
-        return this.modelId;
-    }
-
     public URI uri() {
         return this.uri;
     }
@@ -64,7 +61,7 @@ public RateLimitSettings rateLimitSettings() {
 
     @Override
     public int rateLimitGroupingHash() {
-        return Objects.hash(modelId, uri, getSecretSettings());
+        return Objects.hash(getServiceSettings().modelId(), uri, getSecretSettings());
     }
 
     // Needed for testing only
@@ -87,4 +84,6 @@ public void setURI(String newUri) {
     protected static SecretSettings retrieveSecretSettings(Map<String, Object> secrets) {
         return (secrets != null && secrets.isEmpty()) ? EmptySecretSettings.INSTANCE : DefaultSecretSettings.fromMap(secrets);
     }
+
+    protected abstract ExecutableAction accept(LlamaActionVisitor creator);
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/LlamaService.java

@@ -106,11 +106,8 @@ protected void doInfer(
         ActionListener<InferenceServiceResults> listener
     ) {
         var actionCreator = new LlamaActionCreator(getSender(), getServiceComponents());
-
-        if (model instanceof LlamaEmbeddingsModel llamaEmbeddingsModel) {
-            llamaEmbeddingsModel.accept(actionCreator).execute(inputs, timeout, listener);
-        } else if (model instanceof LlamaChatCompletionModel llamaChatCompletionModel) {
-            llamaChatCompletionModel.accept(actionCreator).execute(inputs, timeout, listener);
+        if (model instanceof LlamaModel llamaModel) {
+            llamaModel.accept(actionCreator).execute(inputs, timeout, listener);
         } else {
             listener.onFailure(createInvalidModelException(model));
         }
@@ -127,7 +124,6 @@ protected void validateInputType(InputType inputType, Model model, ValidationExc
      * @param inferenceId the unique identifier for the inference entity
      * @param taskType the type of task this model is designed for
      * @param serviceSettings the settings for the inference service
-     * @param taskSettings the settings specific to the task
      * @param chunkingSettings the settings for chunking, if applicable
      * @param secretSettings the secret settings for the model, such as API keys or tokens
      * @param failureMessage the message to use in case of failure
@@ -138,24 +134,14 @@ protected LlamaModel createModel(
         String inferenceId,
         TaskType taskType,
         Map<String, Object> serviceSettings,
-        Map<String, Object> taskSettings,
         ChunkingSettings chunkingSettings,
         Map<String, Object> secretSettings,
         String failureMessage,
         ConfigurationParseContext context
     ) {
         switch (taskType) {
             case TEXT_EMBEDDING:
-                return new LlamaEmbeddingsModel(
-                    inferenceId,
-                    taskType,
-                    NAME,
-                    serviceSettings,
-                    taskSettings,
-                    chunkingSettings,
-                    secretSettings,
-                    context
-                );
+                return new LlamaEmbeddingsModel(inferenceId, taskType, NAME, serviceSettings, chunkingSettings, secretSettings, context);
             case CHAT_COMPLETION, COMPLETION:
                 return new LlamaChatCompletionModel(inferenceId, taskType, NAME, serviceSettings, secretSettings, context);
             default:
@@ -168,7 +154,7 @@ public Model updateModelWithEmbeddingDetails(Model model, int embeddingSize) {
         if (model instanceof LlamaEmbeddingsModel embeddingsModel) {
             var serviceSettings = embeddingsModel.getServiceSettings();
             var similarityFromModel = serviceSettings.similarity();
-            var similarityToUse = similarityFromModel == null ? SimilarityMeasure.COSINE : similarityFromModel;
+            var similarityToUse = similarityFromModel == null ? SimilarityMeasure.DOT_PRODUCT : similarityFromModel;
 
             var updatedServiceSettings = new LlamaEmbeddingsServiceSettings(
                 serviceSettings.modelId(),
@@ -283,7 +269,6 @@ public void parseRequestConfig(
                 modelId,
                 taskType,
                 serviceSettingsMap,
-                taskSettingsMap,
                 chunkingSettings,
                 serviceSettingsMap,
                 TaskType.unsupportedTaskTypeErrorMsg(taskType, NAME),
@@ -308,7 +293,7 @@ public Model parsePersistedConfigWithSecrets(
         Map<String, Object> secrets
     ) {
         Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
-        Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
+        removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
         Map<String, Object> secretSettingsMap = removeFromMapOrDefaultEmpty(secrets, ModelSecrets.SECRET_SETTINGS);
 
         ChunkingSettings chunkingSettings = null;
@@ -320,7 +305,6 @@ public Model parsePersistedConfigWithSecrets(
             modelId,
             taskType,
             serviceSettingsMap,
-            taskSettingsMap,
             chunkingSettings,
             secretSettingsMap,
             parsePersistedConfigErrorMsg(modelId, NAME)
@@ -331,7 +315,6 @@ private LlamaModel createModelFromPersistent(
         String inferenceEntityId,
         TaskType taskType,
         Map<String, Object> serviceSettings,
-        Map<String, Object> taskSettings,
         ChunkingSettings chunkingSettings,
         Map<String, Object> secretSettings,
         String failureMessage
@@ -340,7 +323,6 @@ private LlamaModel createModelFromPersistent(
             inferenceEntityId,
             taskType,
             serviceSettings,
-            taskSettings,
             chunkingSettings,
             secretSettings,
             failureMessage,
@@ -351,7 +333,7 @@ private LlamaModel createModelFromPersistent(
     @Override
     public Model parsePersistedConfig(String modelId, TaskType taskType, Map<String, Object> config) {
         Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
-        Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
+        removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
 
         ChunkingSettings chunkingSettings = null;
         if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
@@ -362,7 +344,6 @@ public Model parsePersistedConfig(String modelId, TaskType taskType, Map<String,
             modelId,
             taskType,
             serviceSettingsMap,
-            taskSettingsMap,
             chunkingSettings,
             null,
             parsePersistedConfigErrorMsg(modelId, NAME)
@@ -374,6 +355,12 @@ public TransportVersion getMinimalSupportedVersion() {
         return TransportVersions.ML_INFERENCE_LLAMA_ADDED;
     }
 
+    @Override
+    public boolean hideFromConfigurationApi() {
+        // The Llama service is very configurable so we're going to hide it from being exposed in the service API.
+        return true;
+    }
+
     /**
      * Configuration class for the Llama inference service.
      * It provides the settings and configurations required for the service.

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/action/LlamaActionCreator.java

@@ -18,13 +18,13 @@
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.external.http.sender.UnifiedChatInput;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
+import org.elasticsearch.xpack.inference.services.huggingface.response.HuggingFaceEmbeddingsResponseEntity;
 import org.elasticsearch.xpack.inference.services.llama.completion.LlamaChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.llama.completion.LlamaCompletionResponseHandler;
 import org.elasticsearch.xpack.inference.services.llama.embeddings.LlamaEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.llama.embeddings.LlamaEmbeddingsResponseHandler;
 import org.elasticsearch.xpack.inference.services.llama.request.completion.LlamaChatCompletionRequest;
 import org.elasticsearch.xpack.inference.services.llama.request.embeddings.LlamaEmbeddingsRequest;
-import org.elasticsearch.xpack.inference.services.llama.response.embeddings.LlamaEmbeddingsResponseEntity;
 import org.elasticsearch.xpack.inference.services.openai.response.OpenAiChatCompletionResponseEntity;
 
 import java.util.Objects;
@@ -44,7 +44,7 @@ public class LlamaActionCreator implements LlamaActionVisitor {
 
     private static final ResponseHandler EMBEDDINGS_HANDLER = new LlamaEmbeddingsResponseHandler(
         "llama text embedding",
-        LlamaEmbeddingsResponseEntity::fromResponse
+        HuggingFaceEmbeddingsResponseEntity::fromResponse
     );
     private static final ResponseHandler COMPLETION_HANDLER = new LlamaCompletionResponseHandler(
         "llama completion",

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/completion/LlamaChatCompletionModel.java

@@ -68,7 +68,7 @@ public LlamaChatCompletionModel(
         SecretSettings secrets
     ) {
         super(
-            new ModelConfigurations(inferenceEntityId, taskType, service, serviceSettings, new EmptyTaskSettings()),
+            new ModelConfigurations(inferenceEntityId, taskType, service, serviceSettings, EmptyTaskSettings.INSTANCE),
             new ModelSecrets(secrets)
         );
         setPropertiesFromServiceSettings(serviceSettings);
@@ -105,7 +105,6 @@ public static LlamaChatCompletionModel of(LlamaChatCompletionModel model, Unifie
     }
 
     private void setPropertiesFromServiceSettings(LlamaChatCompletionServiceSettings serviceSettings) {
-        this.modelId = serviceSettings.modelId();
         this.uri = serviceSettings.uri();
         this.rateLimitSettings = serviceSettings.rateLimitSettings();
     }
@@ -126,6 +125,7 @@ public LlamaChatCompletionServiceSettings getServiceSettings() {
      * @param creator the visitor that creates the executable action
      * @return an ExecutableAction representing this model
      */
+    @Override
     public ExecutableAction accept(LlamaActionVisitor creator) {
         return creator.create(this);
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/completion/LlamaChatCompletionServiceSettings.java

@@ -116,9 +116,16 @@ public String getWriteableName() {
 
     @Override
     public TransportVersion getMinimalSupportedVersion() {
+        assert false : "should never be called when supportsVersion is used";
         return TransportVersions.ML_INFERENCE_LLAMA_ADDED;
     }
 
+    @Override
+    public boolean supportsVersion(TransportVersion version) {
+        return version.onOrAfter(TransportVersions.ML_INFERENCE_LLAMA_ADDED)
+            || version.isPatchFrom(TransportVersions.ML_INFERENCE_LLAMA_ADDED_8_19);
+    }
+
     @Override
     public String modelId() {
         return this.modelId;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/embeddings/LlamaEmbeddingsModel.java

@@ -12,7 +12,6 @@
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
 import org.elasticsearch.inference.SecretSettings;
-import org.elasticsearch.inference.TaskSettings;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
@@ -42,7 +41,6 @@ public LlamaEmbeddingsModel(
         TaskType taskType,
         String service,
         Map<String, Object> serviceSettings,
-        Map<String, Object> taskSettings,
         ChunkingSettings chunkingSettings,
         Map<String, Object> secrets,
         ConfigurationParseContext context
@@ -52,7 +50,6 @@ public LlamaEmbeddingsModel(
             taskType,
             service,
             LlamaEmbeddingsServiceSettings.fromMap(serviceSettings, context),
-            EmptyTaskSettings.INSTANCE,    // no task settings for Llama embeddings
             chunkingSettings,
             retrieveSecretSettings(secrets)
         );
@@ -75,7 +72,6 @@ public LlamaEmbeddingsModel(LlamaEmbeddingsModel model, LlamaEmbeddingsServiceSe
      * @param serviceSettings the service settings to extract properties from
      */
     private void setPropertiesFromServiceSettings(LlamaEmbeddingsServiceSettings serviceSettings) {
-        this.modelId = serviceSettings.modelId();
         this.uri = serviceSettings.uri();
         this.rateLimitSettings = serviceSettings.rateLimitSettings();
     }
@@ -87,7 +83,6 @@ private void setPropertiesFromServiceSettings(LlamaEmbeddingsServiceSettings ser
      * @param taskType the type of task this model is designed for
      * @param service the name of the inference service
      * @param serviceSettings the settings for the inference service, specific to embeddings
-     * @param taskSettings the task settings for the model
      * @param chunkingSettings the chunking settings for processing input data
      * @param secrets the secret settings for the model, such as API keys or tokens
      */
@@ -96,7 +91,6 @@ public LlamaEmbeddingsModel(
         TaskType taskType,
         String service,
         LlamaEmbeddingsServiceSettings serviceSettings,
-        TaskSettings taskSettings,
         ChunkingSettings chunkingSettings,
         SecretSettings secrets
     ) {
@@ -118,6 +112,7 @@ public LlamaEmbeddingsServiceSettings getServiceSettings() {
      * @param creator the visitor that creates the executable action
      * @return an ExecutableAction representing the Llama embeddings model
      */
+    @Override
     public ExecutableAction accept(LlamaActionVisitor creator) {
         return creator.create(this);
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/embeddings/LlamaEmbeddingsServiceSettings.java

@@ -154,9 +154,16 @@ public String getWriteableName() {
 
     @Override
     public TransportVersion getMinimalSupportedVersion() {
+        assert false : "should never be called when supportsVersion is used";
         return TransportVersions.ML_INFERENCE_LLAMA_ADDED;
     }
 
+    @Override
+    public boolean supportsVersion(TransportVersion version) {
+        return version.onOrAfter(TransportVersions.ML_INFERENCE_LLAMA_ADDED)
+            || version.isPatchFrom(TransportVersions.ML_INFERENCE_LLAMA_ADDED_8_19);
+    }
+
     @Override
     public String modelId() {
         return this.modelId;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/request/completion/LlamaChatCompletionRequest.java

@@ -60,8 +60,7 @@ public HttpRequest createHttpRequest() {
         httpPost.setEntity(byteEntity);
 
         httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaTypeWithoutParameters());
-        if (model.getSecretSettings() instanceof DefaultSecretSettings) {
-            var secretSettings = (DefaultSecretSettings) model.getSecretSettings();
+        if (model.getSecretSettings() instanceof DefaultSecretSettings secretSettings) {
             httpPost.setHeader(createAuthBearerHeader(secretSettings.apiKey()));
         }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/llama/request/embeddings/LlamaEmbeddingsRequest.java

@@ -60,13 +60,13 @@ public HttpRequest createHttpRequest() {
         HttpPost httpPost = new HttpPost(this.uri);
 
         ByteArrayEntity byteEntity = new ByteArrayEntity(
-            Strings.toString(new LlamaEmbeddingsRequestEntity(model.model(), truncationResult.input())).getBytes(StandardCharsets.UTF_8)
+            Strings.toString(new LlamaEmbeddingsRequestEntity(model.getServiceSettings().modelId(), truncationResult.input()))
+                .getBytes(StandardCharsets.UTF_8)
         );
         httpPost.setEntity(byteEntity);
 
         httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaTypeWithoutParameters());
-        if (model.getSecretSettings() instanceof DefaultSecretSettings) {
-            var secretSettings = (DefaultSecretSettings) model.getSecretSettings();
+        if (model.getSecretSettings() instanceof DefaultSecretSettings secretSettings) {
             httpPost.setHeader(createAuthBearerHeader(secretSettings.apiKey()));
         }