support chunking for text embedding bits

Author: jan-elastic

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

@@ -13,7 +13,6 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
 import org.elasticsearch.inference.InferenceResults;
-import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xpack.core.ml.inference.results.MlTextEmbeddingResults;
 
@@ -41,11 +40,9 @@
  *     ]
  * }
  */
-// TODO: let this implement EmbeddingResults, so that this also supports chunking.
 public record TextEmbeddingBitResults(List<TextEmbeddingByteResults.Embedding> embeddings)
     implements
-        InferenceServiceResults,
-        TextEmbedding {
+        TextEmbeddingResults<TextEmbeddingByteResults.Chunk, TextEmbeddingByteResults.Embedding> {
     public static final String NAME = "text_embedding_service_bit_results";
     public static final String TEXT_EMBEDDING_BITS = "text_embedding_bits";
 

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

@@ -50,8 +50,7 @@
  */
 public record TextEmbeddingByteResults(List<Embedding> embeddings)
     implements
-        EmbeddingResults<TextEmbeddingByteResults.Chunk, TextEmbeddingByteResults.Embedding>,
-        TextEmbedding {
+        TextEmbeddingResults<TextEmbeddingByteResults.Chunk, TextEmbeddingByteResults.Embedding> {
     public static final String NAME = "text_embedding_service_byte_results";
     public static final String TEXT_EMBEDDING_BYTES = "text_embedding_bytes";
 

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

@@ -55,8 +55,7 @@
  */
 public record TextEmbeddingFloatResults(List<Embedding> embeddings)
     implements
-        EmbeddingResults<TextEmbeddingFloatResults.Chunk, TextEmbeddingFloatResults.Embedding>,
-        TextEmbedding {
+        TextEmbeddingResults<TextEmbeddingFloatResults.Chunk, TextEmbeddingFloatResults.Embedding> {
     public static final String NAME = "text_embedding_service_results";
     public static final String TEXT_EMBEDDING = TaskType.TEXT_EMBEDDING.toString();
 

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

@@ -7,7 +7,9 @@
 
 package org.elasticsearch.xpack.core.inference.results;
 
-public interface TextEmbedding {
+public interface TextEmbeddingResults<C extends EmbeddingResults.Chunk, E extends EmbeddingResults.Embedding<C>>
+    extends
+        EmbeddingResults<C, E> {
 
     /**
      * Returns the first text embedding entry in the result list's array size.

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

@@ -22,8 +22,8 @@
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
-import org.elasticsearch.xpack.core.inference.results.TextEmbedding;
 import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingResults;
 import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.inference.services.settings.ApiKeySecrets;
 
@@ -741,7 +741,7 @@ public static void getEmbeddingSize(Model model, InferenceService service, Actio
             InputType.INGEST,
             InferenceAction.Request.DEFAULT_TIMEOUT,
             listener.delegateFailureAndWrap((delegate, r) -> {
-                if (r instanceof TextEmbedding embeddingResults) {
+                if (r instanceof TextEmbeddingResults<?, ?> embeddingResults) {
                     try {
                         delegate.onResponse(embeddingResults.getFirstEmbeddingSize());
                     } catch (Exception e) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/validation/TextEmbeddingModelValidator.java

@@ -14,8 +14,8 @@
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.rest.RestStatus;
-import org.elasticsearch.xpack.core.inference.results.TextEmbedding;
 import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingResults;
 
 public class TextEmbeddingModelValidator implements ModelValidator {
 
@@ -33,7 +33,7 @@ public void validate(InferenceService service, Model model, ActionListener<Model
     }
 
     private Model postValidate(InferenceService service, Model model, InferenceServiceResults results) {
-        if (results instanceof TextEmbedding embeddingResults) {
+        if (results instanceof TextEmbeddingResults<?, ?> embeddingResults) {
             var serviceSettings = model.getServiceSettings();
             var dimensions = serviceSettings.dimensions();
             int embeddingSize = getEmbeddingSize(embeddingResults);
@@ -67,7 +67,7 @@ private Model postValidate(InferenceService service, Model model, InferenceServi
         }
     }
 
-    private int getEmbeddingSize(TextEmbedding embeddingResults) {
+    private int getEmbeddingSize(TextEmbeddingResults<?, ?> embeddingResults) {
         int embeddingSize;
         try {
             embeddingSize = embeddingResults.getFirstEmbeddingSize();