Adding BIT support

Author: fzowl

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/voyageai/VoyageAIEmbeddingsResponseEntity.java

@@ -16,6 +16,8 @@
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.core.inference.results.InferenceByteEmbedding;
+import org.elasticsearch.xpack.core.inference.results.InferenceTextEmbeddingBitResults;
 import org.elasticsearch.xpack.core.inference.results.InferenceTextEmbeddingByteResults;
 import org.elasticsearch.xpack.core.inference.results.InferenceTextEmbeddingFloatResults;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
@@ -43,7 +45,9 @@ public class VoyageAIEmbeddingsResponseEntity {
         toLowerCase(VoyageAIEmbeddingType.FLOAT),
         VoyageAIEmbeddingsResponseEntity::parseFloatEmbeddingsArray,
         toLowerCase(VoyageAIEmbeddingType.INT8),
-        VoyageAIEmbeddingsResponseEntity::parseByteEmbeddingsArray
+        VoyageAIEmbeddingsResponseEntity::parseByteEmbeddingsArray,
+        toLowerCase(VoyageAIEmbeddingType.BINARY),
+        VoyageAIEmbeddingsResponseEntity::parseBitEmbeddingsArray
     );
 
     private static final String VALID_EMBEDDING_TYPES_STRING = supportedEmbeddingTypes();
@@ -119,7 +123,7 @@ public static InferenceServiceResults fromResponse(Request request, HttpResult r
 
                 return new InferenceTextEmbeddingFloatResults(embeddingList);
             } else if(embeddingType == VoyageAIEmbeddingType.INT8) {
-                List<InferenceTextEmbeddingByteResults.InferenceByteEmbedding> embeddingList = parseList(
+                List<InferenceByteEmbedding> embeddingList = parseList(
                     jsonParser,
                     VoyageAIEmbeddingsResponseEntity::parseEmbeddingObjectByte
                 );
@@ -144,7 +148,7 @@ private static InferenceTextEmbeddingFloatResults.InferenceFloatEmbedding parseE
         return InferenceTextEmbeddingFloatResults.InferenceFloatEmbedding.of(embeddingValuesList);
     }
 
-    private static InferenceTextEmbeddingByteResults.InferenceByteEmbedding parseEmbeddingObjectByte(XContentParser parser)
+    private static InferenceByteEmbedding parseEmbeddingObjectByte(XContentParser parser)
         throws IOException {
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
 
@@ -154,7 +158,13 @@ private static InferenceTextEmbeddingByteResults.InferenceByteEmbedding parseEmb
         // parse and discard the rest of the object
         consumeUntilObjectEnd(parser);
 
-        return InferenceTextEmbeddingByteResults.InferenceByteEmbedding.of(embeddingValuesList);
+        return InferenceByteEmbedding.of(embeddingValuesList);
+    }
+
+    private static InferenceServiceResults parseBitEmbeddingsArray(XContentParser parser) throws IOException {
+        var embeddingList = parseList(parser, VoyageAIEmbeddingsResponseEntity::parseByteArrayEntry);
+
+        return new InferenceTextEmbeddingBitResults(embeddingList);
     }
 
     private static InferenceServiceResults parseByteEmbeddingsArray(XContentParser parser) throws IOException {
@@ -163,11 +173,11 @@ private static InferenceServiceResults parseByteEmbeddingsArray(XContentParser p
         return new InferenceTextEmbeddingByteResults(embeddingList);
     }
 
-    private static InferenceTextEmbeddingByteResults.InferenceByteEmbedding parseByteArrayEntry(XContentParser parser) throws IOException {
+    private static InferenceByteEmbedding parseByteArrayEntry(XContentParser parser) throws IOException {
         ensureExpectedToken(XContentParser.Token.START_ARRAY, parser.currentToken(), parser);
         List<Byte> embeddingValuesList = parseList(parser, VoyageAIEmbeddingsResponseEntity::parseEmbeddingInt8Entry);
 
-        return InferenceTextEmbeddingByteResults.InferenceByteEmbedding.of(embeddingValuesList);
+        return InferenceByteEmbedding.of(embeddingValuesList);
     }
 
     private static Byte parseEmbeddingInt8Entry(XContentParser parser) throws IOException {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/voyageai/embeddings/VoyageAIEmbeddingType.java

@@ -7,8 +7,6 @@
 
 package org.elasticsearch.xpack.inference.services.voyageai.embeddings;
 
-import org.elasticsearch.TransportVersion;
-import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
 
@@ -36,18 +34,29 @@ public enum VoyageAIEmbeddingType {
     /**
      * This is a synonym for INT8
      */
-    BYTE(DenseVectorFieldMapper.ElementType.BYTE, RequestConstants.INT8);
+    BYTE(DenseVectorFieldMapper.ElementType.BYTE, RequestConstants.INT8),
+    /**
+     * Use this when you want to get back binary embeddings. Valid only for v3 models.
+     */
+    BIT(DenseVectorFieldMapper.ElementType.BIT, RequestConstants.BIT),
+    /**
+     * This is a synonym for BIT
+     */
+    BINARY(DenseVectorFieldMapper.ElementType.BIT, RequestConstants.BIT);
 
     private static final class RequestConstants {
         private static final String FLOAT = "float";
         private static final String INT8 = "int8";
+        private static final String BIT = "binary";
     }
 
     private static final Map<DenseVectorFieldMapper.ElementType, VoyageAIEmbeddingType> ELEMENT_TYPE_TO_VOYAGE_EMBEDDING = Map.of(
         DenseVectorFieldMapper.ElementType.FLOAT,
         FLOAT,
         DenseVectorFieldMapper.ElementType.BYTE,
-        BYTE
+        BYTE,
+        DenseVectorFieldMapper.ElementType.BIT,
+        BIT
     );
     static final EnumSet<DenseVectorFieldMapper.ElementType> SUPPORTED_ELEMENT_TYPES = EnumSet.copyOf(
         ELEMENT_TYPE_TO_VOYAGE_EMBEDDING.keySet()