Refactor EmbeddingRequestChunker (#122818)

* refactor

* inference generics

* more refactor

* unify naming

* remove interface "EmbeddingInt"

* more renaming

* javadoc

* revert accidental changeas

* remove ununsed EmbeddingRequestChunker.EmbeddingType

* polish

* support chunking for text embedding bits

* Polish error messagex

* fix VoyageAI conflicts

Author: jan-elastic

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

@@ -7,12 +7,8 @@
 
 package org.elasticsearch.xpack.core.inference.results;
 
-import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.inference.ChunkedInference;
-import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContent;
-import org.elasticsearch.xcontent.XContentBuilder;
-import org.elasticsearch.xpack.core.ml.search.WeightedToken;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -21,17 +17,17 @@
 
 import static org.elasticsearch.xpack.core.inference.results.TextEmbeddingUtils.validateInputSizeAgainstEmbeddings;
 
-public record ChunkedInferenceEmbeddingSparse(List<SparseEmbeddingChunk> chunks) implements ChunkedInference {
+public record ChunkedInferenceEmbedding(List<? extends EmbeddingResults.Chunk> chunks) implements ChunkedInference {
 
     public static List<ChunkedInference> listOf(List<String> inputs, SparseEmbeddingResults sparseEmbeddingResults) {
         validateInputSizeAgainstEmbeddings(inputs, sparseEmbeddingResults.embeddings().size());
 
         var results = new ArrayList<ChunkedInference>(inputs.size());
         for (int i = 0; i < inputs.size(); i++) {
             results.add(
-                new ChunkedInferenceEmbeddingSparse(
+                new ChunkedInferenceEmbedding(
                     List.of(
-                        new SparseEmbeddingChunk(
+                        new SparseEmbeddingResults.Chunk(
                             sparseEmbeddingResults.embeddings().get(i).tokens(),
                             inputs.get(i),
                             new TextOffset(0, inputs.get(i).length())
@@ -47,21 +43,9 @@ public static List<ChunkedInference> listOf(List<String> inputs, SparseEmbedding
     @Override
     public Iterator<Chunk> chunksAsMatchedTextAndByteReference(XContent xcontent) throws IOException {
         var asChunk = new ArrayList<Chunk>();
-        for (var chunk : chunks) {
-            asChunk.add(new Chunk(chunk.matchedText(), chunk.offset(), toBytesReference(xcontent, chunk.weightedTokens())));
+        for (var chunk : chunks()) {
+            asChunk.add(chunk.toChunk(xcontent));
         }
         return asChunk.iterator();
     }
-
-    private static BytesReference toBytesReference(XContent xContent, List<WeightedToken> tokens) throws IOException {
-        XContentBuilder b = XContentBuilder.builder(xContent);
-        b.startObject();
-        for (var weightedToken : tokens) {
-            weightedToken.toXContent(b, ToXContent.EMPTY_PARAMS);
-        }
-        b.endObject();
-        return BytesReference.bytes(b);
-    }
-
-    public record SparseEmbeddingChunk(List<WeightedToken> weightedTokens, String matchedText, TextOffset offset) {}
 }

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

@@ -0,0 +1,51 @@
+/*
+ * 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 org.elasticsearch.xcontent.XContent;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * The results of a call to the inference service that contains embeddings (sparse or dense).
+ * A call to the inference service may contain multiple input texts, so this results may
+ * contain multiple results.
+ */
+public interface EmbeddingResults<C extends EmbeddingResults.Chunk, E extends EmbeddingResults.Embedding<C>>
+    extends
+        InferenceServiceResults {
+
+    /**
+     * A resulting embedding together with its input text.
+     */
+    interface Chunk {
+        ChunkedInference.Chunk toChunk(XContent xcontent) throws IOException;
+
+        String matchedText();
+
+        ChunkedInference.TextOffset offset();
+    }
+
+    /**
+     * A resulting embedding for one of the input texts to the inference service.
+     */
+    interface Embedding<C extends Chunk> {
+        /**
+         * Combines the resulting embedding with the input into a chunk.
+         */
+        C toChunk(String text, ChunkedInference.TextOffset offset);
+    }
+
+    /**
+     * The resulting list of embeddings for the input texts to the inference service.
+     */
+    List<E> embeddings();
+}