feat: add rerank() argument to nearText / nearAudio

Author: bevzzz

src/it/java/io/weaviate/integration/DataITest.java

@@ -24,11 +24,11 @@
 import io.weaviate.client6.v1.api.collections.data.BatchReference;
 import io.weaviate.client6.v1.api.collections.data.DeleteManyResponse;
 import io.weaviate.client6.v1.api.collections.data.Reference;
+import io.weaviate.client6.v1.api.collections.query.Filter;
 import io.weaviate.client6.v1.api.collections.query.Metadata;
 import io.weaviate.client6.v1.api.collections.query.Metadata.MetadataField;
 import io.weaviate.client6.v1.api.collections.query.QueryMetadata;
 import io.weaviate.client6.v1.api.collections.query.QueryReference;
-import io.weaviate.client6.v1.api.collections.query.Filter;
 import io.weaviate.containers.Container;
 
 public class DataITest extends ConcurrentTest {

src/it/java/io/weaviate/integration/SearchITest.java

@@ -21,9 +21,11 @@
 import io.weaviate.ConcurrentTest;
 import io.weaviate.client6.v1.api.WeaviateApiException;
 import io.weaviate.client6.v1.api.WeaviateClient;
+import io.weaviate.client6.v1.api.collections.Generative;
 import io.weaviate.client6.v1.api.collections.ObjectMetadata;
 import io.weaviate.client6.v1.api.collections.Property;
 import io.weaviate.client6.v1.api.collections.ReferenceProperty;
+import io.weaviate.client6.v1.api.collections.Reranker;
 import io.weaviate.client6.v1.api.collections.VectorConfig;
 import io.weaviate.client6.v1.api.collections.Vectors;
 import io.weaviate.client6.v1.api.collections.WeaviateMetadata;
@@ -37,8 +39,10 @@
 import io.weaviate.client6.v1.api.collections.query.Metadata;
 import io.weaviate.client6.v1.api.collections.query.QueryMetadata;
 import io.weaviate.client6.v1.api.collections.query.QueryResponseGroup;
+import io.weaviate.client6.v1.api.collections.query.Rerank;
 import io.weaviate.client6.v1.api.collections.query.SortBy;
 import io.weaviate.client6.v1.api.collections.query.Target;
+import io.weaviate.client6.v1.api.collections.rerankers.DummyReranker;
 import io.weaviate.client6.v1.api.collections.vectorindex.Hnsw;
 import io.weaviate.client6.v1.api.collections.vectorindex.MultiVector;
 import io.weaviate.containers.Container;
@@ -52,7 +56,7 @@ public class SearchITest extends ConcurrentTest {
       Weaviate.custom()
           .withModel2VecUrl(Model2Vec.URL)
           .withImageInference(Img2VecNeural.URL, Img2VecNeural.MODULE)
-          .addModules("generative-dummy")
+          .addModules(Generative.Kind.DUMMY.jsonValue(), Reranker.Kind.DUMMY.jsonValue())
           .build(),
       Container.IMG2VEC_NEURAL,
       Container.MODEL2VEC);
@@ -741,7 +745,7 @@ public void teset_filterPropertyLength() throws IOException {
     // Assertions
     Assertions.assertThat(got.objects()).hasSize(2);
   }
-  
+
   /**
    * Ensure the client respects server's configuration for max gRPC size:
    * we create a server with 1-byte message size and try to send a large payload
@@ -768,4 +772,30 @@ public void test_maxGrpcMessageSize() throws Exception {
       }).isInstanceOf(io.grpc.StatusRuntimeException.class);
     }
   }
+
+  @Test
+  public void test_rerankQueries() throws IOException {
+    // Arrange
+    var nsThigns = ns("Things");
+
+    var things = client.collections.create(nsThigns,
+        c -> c
+            .properties(Property.text("title"), Property.integer("price"))
+            .vectorConfig(VectorConfig.text2vecModel2Vec(
+                t2v -> t2v.sourceProperties("title", "price")))
+            .rerankerModules(new DummyReranker()));
+
+    things.data.insertMany(
+        Map.of("title", "Ergonomic chair", "price", 269),
+        Map.of("title", "Height-adjustable desk", "price", 349));
+
+    // Act
+    var got = things.query.nearText(
+        "office supplies",
+        nt -> nt.rerank(Rerank.by("price",
+            rank -> rank.query("cheaper first"))));
+
+    // Assert: ranking not important really, just that the request was valid.
+    Assertions.assertThat(got.objects()).hasSize(2);
+  }
 }

src/main/java/io/weaviate/client6/v1/api/collections/CollectionConfig.java

@@ -276,19 +276,23 @@ public void write(JsonWriter out, CollectionConfig value) throws IOException {
           // Reranker and Generative module configs belong to the "moduleConfig".
           var rerankerModules = jsonObject.remove("rerankerModules").getAsJsonArray();
           var generativeModule = jsonObject.remove("generativeModule");
-          if (!rerankerModules.isEmpty() || !generativeModule.isJsonNull()) {
-            var modules = new JsonObject();
 
+          var modules = new JsonObject();
+          if (!rerankerModules.isEmpty()) {
             // Copy configuration for each reranker module.
             rerankerModules.forEach(reranker -> {
               reranker.getAsJsonObject().entrySet()
                   .stream().forEach(entry -> modules.add(entry.getKey(), entry.getValue()));
             });
+          }
 
+          if (!generativeModule.isJsonNull()) {
             // Copy configuration for each generative module.
             generativeModule.getAsJsonObject().entrySet()
                 .stream().forEach(entry -> modules.add(entry.getKey(), entry.getValue()));
+          }
 
+          if (!modules.isEmpty()) {
             jsonObject.add("moduleConfig", modules);
           }
 

src/main/java/io/weaviate/client6/v1/api/collections/Reranker.java

@@ -14,6 +14,7 @@
 import com.google.gson.stream.JsonWriter;
 
 import io.weaviate.client6.v1.api.collections.rerankers.CohereReranker;
+import io.weaviate.client6.v1.api.collections.rerankers.DummyReranker;
 import io.weaviate.client6.v1.api.collections.rerankers.JinaAiReranker;
 import io.weaviate.client6.v1.api.collections.rerankers.NvidiaReranker;
 import io.weaviate.client6.v1.api.collections.rerankers.TransformersReranker;
@@ -24,6 +25,7 @@
 
 public interface Reranker extends TaggedUnion<Reranker.Kind, Object> {
   public enum Kind implements JsonEnum<Kind> {
+    DUMMY("reranker-dummy"),
     JINAAI("reranker-jinaai"),
     VOYAGEAI("reranker-voyageai"),
     NVIDIA("reranker-nvidia"),
@@ -120,6 +122,11 @@ private final void addAdapter(Gson gson, Reranker.Kind kind, Class<? extends Rer
 
     private final void init(Gson gson) {
       addAdapter(gson, Reranker.Kind.COHERE, CohereReranker.class);
+      addAdapter(gson, Reranker.Kind.JINAAI, JinaAiReranker.class);
+      addAdapter(gson, Reranker.Kind.NVIDIA, NvidiaReranker.class);
+      addAdapter(gson, Reranker.Kind.TRANSFORMERS, TransformersReranker.class);
+      addAdapter(gson, Reranker.Kind.VOYAGEAI, VoyageAiReranker.class);
+      addAdapter(gson, Reranker.Kind.DUMMY, DummyReranker.class);
     }
 
     @SuppressWarnings("unchecked")

src/main/java/io/weaviate/client6/v1/api/collections/query/BaseVectorSearchBuilder.java

@@ -6,6 +6,7 @@ abstract class BaseVectorSearchBuilder<SelfT extends BaseVectorSearchBuilder<Sel
   // Optional query parameters.
   Float distance;
   Float certainty;
+  Rerank rerank;
 
   /**
    * Discard objects whose vectors are further away
@@ -35,4 +36,13 @@ public SelfT certainty(float certainty) {
     this.certainty = certainty;
     return (SelfT) this;
   }
+
+  /**
+   * Control the ranking of the query results.
+   */
+  @SuppressWarnings("unchecked")
+  public SelfT rerank(Rerank rerank) {
+    this.rerank = rerank;
+    return (SelfT) this;
+  }
 }

src/main/java/io/weaviate/client6/v1/api/collections/query/NearAudio.java

@@ -10,7 +10,7 @@
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoBaseSearch;
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoSearchGet;
 
-public record NearAudio(Target searchTarget, Float distance, Float certainty, BaseQueryOptions common)
+public record NearAudio(Target searchTarget, Float distance, Float certainty, Rerank rerank, BaseQueryOptions common)
     implements QueryOperator, AggregateObjectFilter {
 
   public static NearAudio of(String audio) {
@@ -34,6 +34,7 @@ public NearAudio(Builder builder) {
         builder.media,
         builder.distance,
         builder.certainty,
+        builder.rerank,
         builder.baseOptions());
   }
 
@@ -80,6 +81,7 @@ private WeaviateProtoBaseSearch.NearAudioSearch.Builder protoBuilder() {
     } else if (distance != null) {
       nearAudio.setDistance(distance);
     }
+
     return nearAudio;
   }
 }

src/main/java/io/weaviate/client6/v1/api/collections/query/NearText.java

@@ -13,7 +13,12 @@
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoBaseSearch;
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoSearchGet;
 
-public record NearText(Target searchTarget, Float distance, Float certainty, Move moveTo,
+public record NearText(
+    Target searchTarget,
+    Float distance,
+    Float certainty,
+    Rerank rerank,
+    Move moveTo,
     Move moveAway,
     BaseQueryOptions common) implements QueryOperator, AggregateObjectFilter {
 
@@ -38,6 +43,7 @@ public NearText(Builder builder) {
         builder.searchTarget,
         builder.distance,
         builder.certainty,
+        builder.rerank,
         builder.moveTo,
         builder.moveAway,
         builder.baseOptions());

src/main/java/io/weaviate/client6/v1/api/collections/query/QueryOperator.java

@@ -3,6 +3,14 @@
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoSearchGet;
 
 public interface QueryOperator {
+  default BaseQueryOptions common() {
+    return null;
+  }
+
+  default Rerank rerank() {
+    return null;
+  }
+
   /** Append QueryOperator to the request message. */
   void appendTo(WeaviateProtoSearchGet.SearchRequest.Builder req);
 }

src/main/java/io/weaviate/client6/v1/api/collections/query/QueryRequest.java

@@ -28,6 +28,13 @@ public static <PropertiesT> WeaviateProtoSearchGet.SearchRequest marshal(
     message.setUses125Api(true);
     message.setUses123Api(true);
     message.setCollection(collection.collectionName());
+
+    if (request.operator.common() != null) {
+      request.operator.common().appendTo(message);
+    }
+    if (request.operator.rerank() != null) {
+      request.operator.rerank().appendTo(message);
+    }
     request.operator.appendTo(message);
 
     if (defaults.tenant() != null) {

src/main/java/io/weaviate/client6/v1/api/collections/query/Rerank.java

@@ -0,0 +1,50 @@
+package io.weaviate.client6.v1.api.collections.query;
+
+import java.util.function.Function;
+
+import io.weaviate.client6.v1.internal.ObjectBuilder;
+import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoSearchGet;
+
+public record Rerank(String property, String query) {
+
+  public static Rerank by(String property) {
+    return by(property, ObjectBuilder.identity());
+  }
+
+  public static Rerank by(String property, Function<Builder, ObjectBuilder<Rerank>> fn) {
+    return fn.apply(new Builder(property)).build();
+  }
+
+  void appendTo(WeaviateProtoSearchGet.SearchRequest.Builder req) {
+    var rerank = WeaviateProtoSearchGet.Rerank.newBuilder()
+        .setProperty(property);
+
+    if (query != null) {
+      rerank.setQuery(query);
+    }
+    req.setRerank(rerank);
+  }
+
+  public Rerank(Builder builder) {
+    this(builder.property, builder.query);
+  }
+
+  public static class Builder implements ObjectBuilder<Rerank> {
+    private final String property;
+    private String query;
+
+    public Builder(String property) {
+      this.property = property;
+    }
+
+    public Builder query(String query) {
+      this.query = query;
+      return this;
+    }
+
+    @Override
+    public Rerank build() {
+      return new Rerank(this);
+    }
+  }
+}