inference generics

Author: jan-elastic

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

@@ -41,5 +41,5 @@ private static BytesReference toBytesReference(XContent xContent, byte[] value)
         return BytesReference.bytes(builder);
     }
 
-    public record ByteEmbeddingChunk(byte[] embedding, String matchedText, TextOffset offset) {}
+    public record ByteEmbeddingChunk(byte[] embedding, String matchedText, TextOffset offset) implements EmbeddingResults.EmbeddingChunk {}
 }

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

@@ -41,5 +41,7 @@ private static BytesReference toBytesReference(XContent xContent, float[] value)
         return BytesReference.bytes(b);
     }
 
-    public record FloatEmbeddingChunk(float[] embedding, String matchedText, TextOffset offset) {}
+    public record FloatEmbeddingChunk(float[] embedding, String matchedText, TextOffset offset)
+        implements
+            EmbeddingResults.EmbeddingChunk {}
 }

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

@@ -63,5 +63,7 @@ private static BytesReference toBytesReference(XContent xContent, List<WeightedT
         return BytesReference.bytes(b);
     }
 
-    public record SparseEmbeddingChunk(List<WeightedToken> weightedTokens, String matchedText, TextOffset offset) {}
+    public record SparseEmbeddingChunk(List<WeightedToken> weightedTokens, String matchedText, TextOffset offset)
+        implements
+            EmbeddingResults.EmbeddingChunk {}
 }

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

@@ -0,0 +1,26 @@
+/*
+ * 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.inference.ChunkedInference;
+import org.elasticsearch.inference.InferenceServiceResults;
+
+import java.util.List;
+
+public interface EmbeddingResults<C extends EmbeddingResults.EmbeddingChunk, E extends EmbeddingResults.EmbeddingResult<C>>
+    extends
+        InferenceServiceResults {
+
+    interface EmbeddingChunk {}
+
+    interface EmbeddingResult<C extends EmbeddingResults.EmbeddingChunk> {
+        C toEmbeddingChunk(String text, ChunkedInference.TextOffset offset);
+    }
+
+    List<E> embeddings();
+}

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

@@ -13,14 +13,20 @@
 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.ChunkedInference;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
 
-public record InferenceByteEmbedding(byte[] values) implements Writeable, ToXContentObject, EmbeddingInt {
+public record InferenceByteEmbedding(byte[] values)
+    implements
+        Writeable,
+        ToXContentObject,
+        EmbeddingInt,
+        EmbeddingResults.EmbeddingResult<ChunkedInferenceEmbeddingByte.ByteEmbeddingChunk> {
     public static final String EMBEDDING = "embedding";
 
     public InferenceByteEmbedding(StreamInput in) throws IOException {
@@ -92,4 +98,9 @@ public boolean equals(Object o) {
     public int hashCode() {
         return Arrays.hashCode(values);
     }
+
+    @Override
+    public ChunkedInferenceEmbeddingByte.ByteEmbeddingChunk toEmbeddingChunk(String text, ChunkedInference.TextOffset offset) {
+        return new ChunkedInferenceEmbeddingByte.ByteEmbeddingChunk(values, text, offset);
+    }
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/InferenceTextEmbeddingByteResults.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;
 
@@ -42,7 +41,10 @@
  *     ]
  * }
  */
-public record InferenceTextEmbeddingByteResults(List<InferenceByteEmbedding> embeddings) implements InferenceServiceResults, TextEmbedding {
+public record InferenceTextEmbeddingByteResults(List<InferenceByteEmbedding> embeddings)
+    implements
+        EmbeddingResults<ChunkedInferenceEmbeddingByte.ByteEmbeddingChunk, InferenceByteEmbedding>,
+        TextEmbedding {
     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/InferenceTextEmbeddingFloatResults.java

@@ -15,8 +15,8 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
+import org.elasticsearch.inference.ChunkedInference;
 import org.elasticsearch.inference.InferenceResults;
-import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xcontent.ToXContent;
@@ -53,7 +53,7 @@
  */
 public record InferenceTextEmbeddingFloatResults(List<InferenceFloatEmbedding> embeddings)
     implements
-        InferenceServiceResults,
+        EmbeddingResults<ChunkedInferenceEmbeddingFloat.FloatEmbeddingChunk, InferenceTextEmbeddingFloatResults.InferenceFloatEmbedding>,
         TextEmbedding {
     public static final String NAME = "text_embedding_service_results";
     public static final String TEXT_EMBEDDING = TaskType.TEXT_EMBEDDING.toString();
@@ -151,7 +151,12 @@ public int hashCode() {
         return Objects.hash(embeddings);
     }
 
-    public record InferenceFloatEmbedding(float[] values) implements Writeable, ToXContentObject, EmbeddingInt {
+    public record InferenceFloatEmbedding(float[] values)
+        implements
+            Writeable,
+            ToXContentObject,
+            EmbeddingInt,
+            EmbeddingResults.EmbeddingResult<ChunkedInferenceEmbeddingFloat.FloatEmbeddingChunk> {
         public static final String EMBEDDING = "embedding";
 
         public InferenceFloatEmbedding(StreamInput in) throws IOException {
@@ -220,5 +225,10 @@ public boolean equals(Object o) {
         public int hashCode() {
             return Arrays.hashCode(values);
         }
+
+        @Override
+        public ChunkedInferenceEmbeddingFloat.FloatEmbeddingChunk toEmbeddingChunk(String text, ChunkedInference.TextOffset offset) {
+            return new ChunkedInferenceEmbeddingFloat.FloatEmbeddingChunk(values, text, offset);
+        }
     }
 }

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

@@ -13,8 +13,8 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
+import org.elasticsearch.inference.ChunkedInference;
 import org.elasticsearch.inference.InferenceResults;
-import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xcontent.ToXContent;
@@ -33,7 +33,9 @@
 
 import static org.elasticsearch.xpack.core.ml.inference.trainedmodel.InferenceConfig.DEFAULT_RESULTS_FIELD;
 
-public record SparseEmbeddingResults(List<Embedding> embeddings) implements InferenceServiceResults {
+public record SparseEmbeddingResults(List<Embedding> embeddings)
+    implements
+        EmbeddingResults<ChunkedInferenceEmbeddingSparse.SparseEmbeddingChunk, SparseEmbeddingResults.Embedding> {
 
     public static final String NAME = "sparse_embedding_results";
     public static final String SPARSE_EMBEDDING = TaskType.SPARSE_EMBEDDING.toString();
@@ -114,7 +116,11 @@ public List<? extends InferenceResults> transformToLegacyFormat() {
             .toList();
     }
 
-    public record Embedding(List<WeightedToken> tokens, boolean isTruncated) implements Writeable, ToXContentObject {
+    public record Embedding(List<WeightedToken> tokens, boolean isTruncated)
+        implements
+            Writeable,
+            ToXContentObject,
+            EmbeddingResult<ChunkedInferenceEmbeddingSparse.SparseEmbeddingChunk> {
 
         public static final String EMBEDDING = "embedding";
         public static final String IS_TRUNCATED = "is_truncated";
@@ -163,5 +169,10 @@ public Map<String, Object> asMap() {
         public String toString() {
             return Strings.toString(this);
         }
+
+        @Override
+        public ChunkedInferenceEmbeddingSparse.SparseEmbeddingChunk toEmbeddingChunk(String text, ChunkedInference.TextOffset offset) {
+            return new ChunkedInferenceEmbeddingSparse.SparseEmbeddingChunk(tokens, text, offset);
+        }
     }
 }