[ML] Adding support for Cohere inference service

docs/changelog/104559.yaml

@@ -0,0 +1,5 @@
+pr: 104559
+summary: Adding support for Cohere inference service
+area: Machine Learning
+type: enhancement
+issues: []

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

@@ -190,6 +190,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_MULTI_CLUSTERS_ENRICH = def(8_576_00_0);
     public static final TransportVersion NESTED_KNN_MORE_INNER_HITS = def(8_577_00_0);
     public static final TransportVersion REQUIRE_DATA_STREAM_ADDED = def(8_578_00_0);
+    public static final TransportVersion ML_INFERENCE_COHERE_EMBEDDINGS_ADDED = def(8_579_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/inference/InputType.java

@@ -8,17 +8,12 @@
 
 package org.elasticsearch.inference;
 
-import org.elasticsearch.common.io.stream.StreamInput;
-import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.common.io.stream.Writeable;
-
-import java.io.IOException;
 import java.util.Locale;
 
 /**
  * Defines the type of request, whether the request is to ingest a document or search for a document.
  */
-public enum InputType implements Writeable {
+public enum InputType {
     INGEST,
     SEARCH;
 
@@ -29,12 +24,7 @@ public String toString() {
         return name().toLowerCase(Locale.ROOT);
     }
 
-    public static InputType fromStream(StreamInput in) throws IOException {
-        return in.readEnum(InputType.class);
-    }
-
-    @Override
-    public void writeTo(StreamOutput out) throws IOException {
-        out.writeEnum(this);
+    public static InputType fromString(String name) {
+        return valueOf(name.trim().toUpperCase(Locale.ROOT));
     }
 }

server/src/main/java/org/elasticsearch/inference/Model.java

@@ -23,6 +23,20 @@ public Model(ModelConfigurations configurations, ModelSecrets secrets) {
         this.secrets = Objects.requireNonNull(secrets);
     }
 
+    public Model(Model model, TaskSettings taskSettings) {
+        Objects.requireNonNull(model);
+
+        configurations = ModelConfigurations.of(model, taskSettings);
+        secrets = model.getSecrets();
+    }
+
+    public Model(Model model, ServiceSettings serviceSettings) {
+        Objects.requireNonNull(model);
+
+        configurations = ModelConfigurations.of(model, serviceSettings);
+        secrets = model.getSecrets();
+    }
+
     public Model(ModelConfigurations configurations) {
         this(configurations, new ModelSecrets());
     }

server/src/main/java/org/elasticsearch/inference/ModelConfigurations.java

@@ -27,6 +27,32 @@ public class ModelConfigurations implements ToXContentObject, VersionedNamedWrit
     public static final String TASK_SETTINGS = "task_settings";
     private static final String NAME = "inference_model";
 
+    public static ModelConfigurations of(Model model, TaskSettings taskSettings) {
+        Objects.requireNonNull(model);
+        Objects.requireNonNull(taskSettings);
+
+        return new ModelConfigurations(
+            model.getConfigurations().getModelId(),
+            model.getConfigurations().getTaskType(),
+            model.getConfigurations().getService(),
+            model.getServiceSettings(),
+            taskSettings
+        );
+    }
+
+    public static ModelConfigurations of(Model model, ServiceSettings serviceSettings) {
+        Objects.requireNonNull(model);
+        Objects.requireNonNull(serviceSettings);
+
+        return new ModelConfigurations(
+            model.getConfigurations().getModelId(),
+            model.getConfigurations().getTaskType(),
+            model.getConfigurations().getService(),
+            serviceSettings,
+            model.getTaskSettings()
+        );
+    }
+
     private final String modelId;
     private final TaskType taskType;
     private final String service;

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/InferenceAction.java

@@ -88,7 +88,7 @@ public Request(StreamInput in) throws IOException {
             }
             this.taskSettings = in.readGenericMap();
             if (in.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_REQUEST_INPUT_TYPE_ADDED)) {
-                this.inputType = InputType.fromStream(in);
+                this.inputType = in.readEnum(InputType.class);
             } else {
                 this.inputType = InputType.INGEST;
             }
@@ -141,7 +141,7 @@ public void writeTo(StreamOutput out) throws IOException {
             }
             out.writeGenericMap(taskSettings);
             if (out.getTransportVersion().onOrAfter(TransportVersions.ML_INFERENCE_REQUEST_INPUT_TYPE_ADDED)) {
-                inputType.writeTo(out);
+                out.writeEnum(inputType);
             }
         }
 

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/EmbeddingInt.java

@@ -0,0 +1,12 @@
+/*
+ * 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.core.inference.results;
+
+public interface EmbeddingInt {
+    int getSize();
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbedding.java

@@ -0,0 +1,18 @@
+/*
+ * 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.core.inference.results;
+
+public interface TextEmbedding {
+
+    /**
+     * Returns the first text embedding entry in the result list's array size.
+     * @return the size of the text embedding
+     * @throws IllegalStateException if the list of embeddings is empty
+     */
+    int getFirstEmbeddingSize() throws IllegalStateException;
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingByteResults.java

@@ -0,0 +1,146 @@
+/*
+ * 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.core.inference.results;
+
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.inference.InferenceResults;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/**
+ * Writes a text embedding result in the follow json format
+ * {
+ *     "text_embedding": [
+ *         {
+ *             "embedding": [
+ *                 23
+ *             ]
+ *         },
+ *         {
+ *             "embedding": [
+ *                 -23
+ *             ]
+ *         }
+ *     ]
+ * }
+ */
+public record TextEmbeddingByteResults(List<Embedding> embeddings) implements InferenceServiceResults, TextEmbedding {
+    public static final String NAME = "text_embedding_service_byte_results";
+    public static final String TEXT_EMBEDDING = TaskType.TEXT_EMBEDDING.toString();
+
+    public TextEmbeddingByteResults(StreamInput in) throws IOException {
+        this(in.readCollectionAsList(Embedding::new));
+    }
+
+    @Override
+    public int getFirstEmbeddingSize() {
+        return TextEmbeddingUtils.getFirstEmbeddingSize(new ArrayList<>(embeddings));
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startArray(TEXT_EMBEDDING);
+        for (Embedding embedding : embeddings) {
+            embedding.toXContent(builder, params);
+        }
+        builder.endArray();
+        return builder;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeCollection(embeddings);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    @Override
+    public List<? extends InferenceResults> transformToCoordinationFormat() {
+        return embeddings.stream()
+            .map(embedding -> embedding.values.stream().mapToDouble(value -> value).toArray())
+            .map(values -> new org.elasticsearch.xpack.core.ml.inference.results.TextEmbeddingResults(TEXT_EMBEDDING, values, false))
+            .toList();
+    }
+
+    @Override
+    @SuppressWarnings("deprecation")
+    public List<? extends InferenceResults> transformToLegacyFormat() {
+        var legacyEmbedding = new LegacyTextEmbeddingResults(
+            embeddings.stream().map(embedding -> new LegacyTextEmbeddingResults.Embedding(embedding.toFloats())).toList()
+        );
+
+        return List.of(legacyEmbedding);
+    }
+
+    public Map<String, Object> asMap() {
+        Map<String, Object> map = new LinkedHashMap<>();
+        map.put(TEXT_EMBEDDING, embeddings.stream().map(Embedding::asMap).collect(Collectors.toList()));
+
+        return map;
+    }
+
+    public record Embedding(List<Byte> values) implements Writeable, ToXContentObject, EmbeddingInt {
+        public static final String EMBEDDING = "embedding";
+
+        public Embedding(StreamInput in) throws IOException {
+            this(in.readCollectionAsImmutableList(StreamInput::readByte));
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeCollection(values, StreamOutput::writeByte);
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            builder.startObject();
+
+            builder.startArray(EMBEDDING);
+            for (Byte value : values) {
+                builder.value(value);
+            }
+            builder.endArray();
+
+            builder.endObject();
+            return builder;
+        }
+
+        @Override
+        public String toString() {
+            return Strings.toString(this);
+        }
+
+        public Map<String, Object> asMap() {
+            return Map.of(EMBEDDING, values);
+        }
+
+        public List<Float> toFloats() {
+            return values.stream().map(Byte::floatValue).toList();
+        }
+
+        @Override
+        public int getSize() {
+            return values().size();
+        }
+    }
+}

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingResults.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -40,7 +41,7 @@
  *     ]
  * }
  */
-public record TextEmbeddingResults(List<Embedding> embeddings) implements InferenceServiceResults {
+public record TextEmbeddingResults(List<Embedding> embeddings) implements InferenceServiceResults, TextEmbedding {
     public static final String NAME = "text_embedding_service_results";
     public static final String TEXT_EMBEDDING = TaskType.TEXT_EMBEDDING.toString();
 
@@ -58,6 +59,11 @@ public TextEmbeddingResults(StreamInput in) throws IOException {
         );
     }
 
+    @Override
+    public int getFirstEmbeddingSize() {
+        return TextEmbeddingUtils.getFirstEmbeddingSize(new ArrayList<>(embeddings));
+    }
+
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startArray(TEXT_EMBEDDING);
@@ -103,13 +109,18 @@ public Map<String, Object> asMap() {
         return map;
     }
 
-    public record Embedding(List<Float> values) implements Writeable, ToXContentObject {
+    public record Embedding(List<Float> values) implements Writeable, ToXContentObject, EmbeddingInt {
         public static final String EMBEDDING = "embedding";
 
         public Embedding(StreamInput in) throws IOException {
             this(in.readCollectionAsImmutableList(StreamInput::readFloat));
         }
 
+        @Override
+        public int getSize() {
+            return values.size();
+        }
+
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeCollection(values, StreamOutput::writeFloat);

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingUtils.java

@@ -0,0 +1,30 @@
+/*
+ * 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.core.inference.results;
+
+import java.util.List;
+
+public class TextEmbeddingUtils {
+
+    /**
+     * Returns the first text embedding entry's array size.
+     * @param embeddings the list of embeddings
+     * @return the size of the text embedding
+     * @throws IllegalStateException if the list of embeddings is empty
+     */
+    public static int getFirstEmbeddingSize(List<EmbeddingInt> embeddings) throws IllegalStateException {
+        if (embeddings.isEmpty()) {
+            throw new IllegalStateException("Embeddings list is empty");
+        }
+
+        return embeddings.get(0).getSize();
+    }
+
+    private TextEmbeddingUtils() {}
+
+}

x-pack/plugin/inference/src/main/java/module-info.java

@@ -22,10 +22,5 @@
     exports org.elasticsearch.xpack.inference.registry;
     exports org.elasticsearch.xpack.inference.rest;
     exports org.elasticsearch.xpack.inference.services;
-    exports org.elasticsearch.xpack.inference.external.http.sender;
-    exports org.elasticsearch.xpack.inference.external.http;
-    exports org.elasticsearch.xpack.inference.services.elser;
-    exports org.elasticsearch.xpack.inference.services.huggingface.elser;
-    exports org.elasticsearch.xpack.inference.services.openai;
     exports org.elasticsearch.xpack.inference;
 }