diff --git a/src/it/java/io/weaviate/integration/SearchITest.java b/src/it/java/io/weaviate/integration/SearchITest.java index 3c8d7f359..9fc570e1a 100644 --- a/src/it/java/io/weaviate/integration/SearchITest.java +++ b/src/it/java/io/weaviate/integration/SearchITest.java @@ -801,13 +801,56 @@ public void test_rerankQueries() throws IOException { 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")))); + var rerank = Rerank.by("price", rank -> rank.query("cheaper first")); + var got = things.query.nearText("office supplies", nt -> nt.rerank(rerank)); // Assert: ranking not important really, just that the request was valid. Assertions.assertThat(got.objects()).hasSize(2); + + // Assert: rerank is not exclusive to vector search -- BM25, hybrid and + // fetchObjects accept it too, and the server reranks for all of them. + Assertions.assertThat(things.query.bm25("chair", bm25 -> bm25.rerank(rerank)).objects()) + .as("bm25").isNotEmpty().allSatisfy(SearchITest::assertReranked); + Assertions.assertThat(things.query.hybrid("chair", hybrid -> hybrid.rerank(rerank)).objects()) + .as("hybrid").isNotEmpty().allSatisfy(SearchITest::assertReranked); + Assertions.assertThat(things.query.fetchObjects(fetch -> fetch.rerank(rerank)).objects()) + .as("fetchObjects").hasSize(2).allSatisfy(SearchITest::assertReranked); + } + + private static void assertReranked(WeaviateObject> object) { + Assertions.assertThat(object.queryMetadata().rerankScore()) + .as("rerank score of %s", object.uuid()).isNotNull(); + } + + @Test + public void test_rerankScoreIsReturned() throws IOException { + // Arrange + var nsThings = ns("Things"); + + var things = client.collections.create(nsThings, + 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.fetchObjects( + fetch -> fetch.rerank(Rerank.by("title", rank -> rank.query("chair")))); + + // Assert: the score which produced the ordering is readable. + Assertions.assertThat(got.objects()).hasSize(2) + .allSatisfy(obj -> Assertions.assertThat(obj.queryMetadata().rerankScore()) + .as("rerank score of %s", obj.uuid()).isNotNull()); + + // Assert: a query without rerank leaves the score unset. + Assertions.assertThat(things.query.fetchObjects().objects()).hasSize(2) + .allSatisfy(obj -> Assertions.assertThat(obj.queryMetadata().rerankScore()) + .as("not reranked").isNull()); } @Test diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/generate/GenerativeResponseGroup.java b/src/main/java/io/weaviate/client6/v1/api/collections/generate/GenerativeResponseGroup.java index 8ab444d11..dfe9beb83 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/generate/GenerativeResponseGroup.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/generate/GenerativeResponseGroup.java @@ -2,7 +2,9 @@ import java.util.List; +import io.weaviate.client6.v1.api.collections.query.BaseQueryOptions; import io.weaviate.client6.v1.api.collections.query.QueryObjectGrouped; +import io.weaviate.client6.v1.api.collections.query.Rerank; public record GenerativeResponseGroup( /** Group name. */ @@ -19,6 +21,14 @@ public record GenerativeResponseGroup( Float maxDistance, /** The size of the group. */ long numberOfObjects, + /** + * Score assigned to this group by the reranker module. + * + *

+ * Only present if the query requested reranking, see + * {@link BaseQueryOptions.Builder#rerank(Rerank)}. + */ + Double rerankScore, /** Objects retrieved in the query. */ List> objects, /** Output of the summary task for this group. */ diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/generate/GenerativeResponseGrouped.java b/src/main/java/io/weaviate/client6/v1/api/collections/generate/GenerativeResponseGrouped.java index 95710580c..7ba0984a0 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/generate/GenerativeResponseGrouped.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/generate/GenerativeResponseGrouped.java @@ -58,6 +58,7 @@ static GenerativeResponseGrouped unmarshal( group.getMinDistance(), group.getMaxDistance(), group.getNumberOfObjects(), + group.hasRerank() ? group.getRerank().getScore() : null, objects, generative); }) diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/query/BaseQueryOptions.java b/src/main/java/io/weaviate/client6/v1/api/collections/query/BaseQueryOptions.java index ebe2cfbdb..e7e44af29 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/query/BaseQueryOptions.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/query/BaseQueryOptions.java @@ -21,6 +21,7 @@ public record BaseQueryOptions( ConsistencyLevel consistencyLevel, Filter filters, Boost boost, + Rerank rerank, GenerativeSearch generativeSearch, List returnProperties, List returnReferences, @@ -40,6 +41,7 @@ private BaseQueryOptions(Builder, T> builder.consistencyLevel, builder.filter, builder.boost, + builder.rerank, builder.generativeSearch, builder.returnProperties, builder.returnReferences, @@ -57,6 +59,7 @@ public static abstract class Builder, T extends private ConsistencyLevel consistencyLevel; private Filter filter; private Boost boost; + private Rerank rerank; private GenerativeSearch generativeSearch; private List returnProperties = new ArrayList<>(); private List returnReferences = new ArrayList<>(); @@ -152,6 +155,20 @@ public final SelfT boost(Boost boost) { return (SelfT) this; } + /** + * Control the ranking of the query results. + * + *

+ * Reranking is applied by the server on top of the result set produced by the + * search operator, so it works with every operator: {@link NearText} and the + * other {@code near*} searches, {@link Bm25}, {@link Hybrid} and + * {@link FetchObjects}. + */ + public final SelfT rerank(Rerank rerank) { + this.rerank = rerank; + return (SelfT) this; + } + /** Select properties to include in the query result. */ public final SelfT returnProperties(String... properties) { return returnProperties(Arrays.asList(properties)); @@ -243,6 +260,10 @@ final void appendTo(WeaviateProtoSearchGet.SearchRequest.Builder req) { req.setBoost(boost.toProto()); } + if (rerank != null) { + rerank.appendTo(req); + } + if (generativeSearch != null) { var generative = WeaviateProtoGenerative.GenerativeSearch.newBuilder(); generativeSearch.appendTo(generative); diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/query/BaseVectorSearchBuilder.java b/src/main/java/io/weaviate/client6/v1/api/collections/query/BaseVectorSearchBuilder.java index 1d78b4698..b54152ca2 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/query/BaseVectorSearchBuilder.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/query/BaseVectorSearchBuilder.java @@ -6,7 +6,6 @@ abstract class BaseVectorSearchBuilder + * Only present if the query requested reranking, see + * {@link BaseQueryOptions.Builder#rerank(Rerank)}. + */ + Double rerankScore) { private QueryMetadata(Builder builder) { this( builder.distance, builder.certainty, builder.score, - builder.explainScore); + builder.explainScore, + builder.rerankScore); } static class Builder implements ObjectBuilder { @@ -24,6 +33,7 @@ static class Builder implements ObjectBuilder { private Float certainty; private Float score; private String explainScore; + private Double rerankScore; final Builder distance(Float distance) { this.distance = distance; @@ -45,6 +55,11 @@ final Builder explainScore(String explainScore) { return this; } + final Builder rerankScore(Double rerankScore) { + this.rerankScore = rerankScore; + return this; + } + @Override public final QueryMetadata build() { return new QueryMetadata(this); diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryOperator.java b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryOperator.java index 886fd5b19..51f0c6ab2 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryOperator.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryOperator.java @@ -7,8 +7,15 @@ default BaseQueryOptions common() { return null; } + /** + * Reranking requested for this query, if any. + * + *

+ * Reranking is a common query option, so operators that carry + * {@link BaseQueryOptions} read it from there. + */ default Rerank rerank() { - return null; + return common() != null ? common().rerank() : null; } /** Append QueryOperator to the request message. */ diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryRequest.java b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryRequest.java index 9b031ddc2..40f92b62e 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryRequest.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryRequest.java @@ -33,9 +33,6 @@ public static WeaviateProtoSearchGet.SearchRequest marshal( if (request.operator.common() != null) { request.operator.common().appendTo(message); } - if (request.operator.rerank() != null) { - request.operator.rerank().appendTo(message); - } request.operator.appendTo(message); defaults.tenant().ifPresent(message::setTenant); diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponse.java b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponse.java index 25313d2f6..d982067d6 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponse.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponse.java @@ -78,6 +78,11 @@ public static WeaviateObject unmarshalResultObject( if (metadataResult.getExplainScorePresent()) { metadata.explainScore(metadataResult.getExplainScore()); } + // 0.0 is a legitimate score, so the presence flag is the only way to tell + // "not reranked" from "reranked with score 0". + if (metadataResult.getRerankScorePresent()) { + metadata.rerankScore(metadataResult.getRerankScore()); + } return new WeaviateObject<>( object.uuid(), collection.collectionName(), diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponseGroup.java b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponseGroup.java index 5ad750051..f473eb981 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponseGroup.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponseGroup.java @@ -17,6 +17,14 @@ public record QueryResponseGroup( Float maxDistance, /** The size of the group. */ long numberOfObjects, + /** + * Score assigned to this group by the reranker module. + * + *

+ * Only present if the query requested reranking, see + * {@link BaseQueryOptions.Builder#rerank(Rerank)}. + */ + Double rerankScore, /** Objects retrieved in the query. */ List> objects) { } diff --git a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponseGrouped.java b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponseGrouped.java index bd23738c4..ed210c742 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponseGrouped.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponseGrouped.java @@ -39,6 +39,7 @@ static QueryResponseGrouped unmarshal( group.getMinDistance(), group.getMaxDistance(), group.getNumberOfObjects(), + group.hasRerank() ? group.getRerank().getScore() : null, objects); }) // Collectors.toMap() throws an NPE if either key or value in the map are null. diff --git a/src/test/java/io/weaviate/client6/v1/api/collections/generate/GenerativeRerankTest.java b/src/test/java/io/weaviate/client6/v1/api/collections/generate/GenerativeRerankTest.java new file mode 100644 index 000000000..1bacbc056 --- /dev/null +++ b/src/test/java/io/weaviate/client6/v1/api/collections/generate/GenerativeRerankTest.java @@ -0,0 +1,77 @@ +package io.weaviate.client6.v1.api.collections.generate; + +import java.util.Map; + +import org.assertj.core.api.Assertions; +import org.junit.Test; + +import io.weaviate.client6.v1.api.collections.CollectionHandleDefaults; +import io.weaviate.client6.v1.api.collections.query.QueryMetadata; +import io.weaviate.client6.v1.api.collections.query.QueryObjectGrouped; +import io.weaviate.client6.v1.internal.ObjectBuilder; +import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoSearchGet; +import io.weaviate.client6.v1.internal.orm.CollectionDescriptor; + +/** + * The generative search reuses the query operators, so it is reranked the same + * way and has to read the scores back the same way. + */ +public class GenerativeRerankTest { + + @Test + public void test_groupRerankScoreIsUnmarshalled() { + var reply = WeaviateProtoSearchGet.SearchReply.newBuilder() + .addGroupByResults(group("fish").setRerank( + WeaviateProtoSearchGet.RerankReply.newBuilder().setScore(0.42))) + .build(); + + var response = unmarshal(reply); + + Assertions.assertThat(response.groups()).extractingByKey("fish") + .returns(0.42, GenerativeResponseGroup::rerankScore); + } + + @Test + public void test_groupRerankScoreIsAbsentWithoutRerank() { + var reply = WeaviateProtoSearchGet.SearchReply.newBuilder() + .addGroupByResults(group("fish")) + .build(); + + var response = unmarshal(reply); + + Assertions.assertThat(response.groups()).extractingByKey("fish") + .returns(null, GenerativeResponseGroup::rerankScore); + } + + /** Objects in the group carry their own score, unmarshalled by the query package. */ + @Test + public void test_objectRerankScoreIsUnmarshalled() { + var reply = WeaviateProtoSearchGet.SearchReply.newBuilder() + .addGroupByResults(group("fish")) + .build(); + + var response = unmarshal(reply); + + Assertions.assertThat(response.objects()).first() + .extracting(QueryObjectGrouped::metadata) + .returns(0.7811474, QueryMetadata::rerankScore); + } + + private static WeaviateProtoSearchGet.GroupByResult.Builder group(String name) { + return WeaviateProtoSearchGet.GroupByResult.newBuilder() + .setName(name) + .setNumberOfObjects(1) + .addObjects(WeaviateProtoSearchGet.SearchResult.newBuilder() + .setMetadata(WeaviateProtoSearchGet.MetadataResult.newBuilder() + .setId("d3b07384-d113-4ec4-92e5-1e0f0ab84d43") + .setRerankScore(0.7811474) + .setRerankScorePresent(true))); + } + + private static GenerativeResponseGrouped> unmarshal( + WeaviateProtoSearchGet.SearchReply reply) { + return GenerativeResponseGrouped.unmarshal(reply, + CollectionDescriptor.ofMap("Things"), + CollectionHandleDefaults.of(ObjectBuilder.identity())); + } +} diff --git a/src/test/java/io/weaviate/client6/v1/api/collections/query/RerankTest.java b/src/test/java/io/weaviate/client6/v1/api/collections/query/RerankTest.java new file mode 100644 index 000000000..0aaa170fd --- /dev/null +++ b/src/test/java/io/weaviate/client6/v1/api/collections/query/RerankTest.java @@ -0,0 +1,120 @@ +package io.weaviate.client6.v1.api.collections.query; + +import org.assertj.core.api.Assertions; +import org.junit.Test; +import org.junit.runner.RunWith; + +import com.jparams.junit4.JParamsTestRunner; +import com.jparams.junit4.data.DataMethod; +import com.jparams.junit4.description.Name; + +import io.weaviate.client6.v1.api.collections.CollectionHandleDefaults; +import io.weaviate.client6.v1.api.collections.WeaviateObject; +import io.weaviate.client6.v1.internal.ObjectBuilder; +import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoSearchGet; +import io.weaviate.client6.v1.internal.orm.CollectionDescriptor; + +@RunWith(JParamsTestRunner.class) +public class RerankTest { + private static final Rerank RERANK = Rerank.by("title", rank -> rank.query("fish")); + + /** Every search operator can be reranked, not only the vector searches. */ + public static Object[][] operators() { + return new Object[][] { + { "bm25", Bm25.of("animal", q -> q.rerank(RERANK)) }, + { "hybrid", Hybrid.of("animal", q -> q.rerank(RERANK)) }, + { "fetchObjects", FetchObjects.of(q -> q.rerank(RERANK)) }, + { "nearText", NearText.of("animal", q -> q.rerank(RERANK)) }, + { "nearVector", NearVector.of(new float[] { 1, 2 }, q -> q.rerank(RERANK)) }, + { "nearObject", NearObject.of("d3b07384-d113-4ec4-92e5-1e0f0ab84d43", q -> q.rerank(RERANK)) }, + }; + } + + @Name("{0}") + @DataMethod(source = RerankTest.class, method = "operators") + @Test + public void test_rerankIsMarshalled(String __, QueryOperator operator) { + var request = marshal(operator); + + Assertions.assertThat(request.hasRerank()).as("has rerank").isTrue(); + Assertions.assertThat(request.getRerank().getProperty()).isEqualTo("title"); + Assertions.assertThat(request.getRerank().getQuery()).isEqualTo("fish"); + Assertions.assertThat(operator.rerank()).as("readable back off the operator").isEqualTo(RERANK); + } + + @Test + public void test_noRerankByDefault() { + Assertions.assertThat(marshal(Bm25.of("animal")).hasRerank()).isFalse(); + Assertions.assertThat(Bm25.of("animal").rerank()).isNull(); + } + + @Test + public void test_rerankWithoutQuery() { + var request = marshal(Bm25.of("animal", q -> q.rerank(Rerank.by("title")))); + + Assertions.assertThat(request.getRerank().getProperty()).isEqualTo("title"); + Assertions.assertThat(request.getRerank().hasQuery()).as("query is optional").isFalse(); + } + + @Test + public void test_rerankScoreIsUnmarshalled() { + var reply = WeaviateProtoSearchGet.SearchReply.newBuilder() + .addResults(searchResult(0.7811474)) + .build(); + + var response = QueryResponse.unmarshal(reply, CollectionDescriptor.ofMap("Things")); + + Assertions.assertThat(response.objects()).first() + .extracting(WeaviateObject::queryMetadata) + .returns(0.7811474, QueryMetadata::rerankScore); + } + + /** 0.0 is a legitimate score, so absence has to come from the presence flag. */ + @Test + public void test_rerankScoreZeroIsNotAbsent() { + var reply = WeaviateProtoSearchGet.SearchReply.newBuilder() + .addResults(searchResult(0.0)) + .addResults(WeaviateProtoSearchGet.SearchResult.newBuilder() + .setMetadata(WeaviateProtoSearchGet.MetadataResult.newBuilder() + .setId("2a5f8b3c-0f1e-4d6a-9c8b-7e2d1a0f3b4c"))) + .build(); + + var objects = QueryResponse.unmarshal(reply, CollectionDescriptor.ofMap("Things")).objects(); + + Assertions.assertThat(objects.get(0).queryMetadata().rerankScore()).as("reranked with 0").isEqualTo(0.0); + Assertions.assertThat(objects.get(1).queryMetadata().rerankScore()).as("not reranked").isNull(); + } + + @Test + public void test_groupRerankScoreIsUnmarshalled() { + var reply = WeaviateProtoSearchGet.SearchReply.newBuilder() + .addGroupByResults(WeaviateProtoSearchGet.GroupByResult.newBuilder() + .setName("fish") + .setNumberOfObjects(1) + .addObjects(searchResult(0.7811474)) + .setRerank(WeaviateProtoSearchGet.RerankReply.newBuilder().setScore(0.42))) + .build(); + + var response = QueryResponseGrouped.unmarshal(reply, + CollectionDescriptor.ofMap("Things"), + CollectionHandleDefaults.of(ObjectBuilder.identity())); + + Assertions.assertThat(response.groups()).extractingByKey("fish") + .returns(0.42, QueryResponseGroup::rerankScore); + } + + private static WeaviateProtoSearchGet.SearchResult.Builder searchResult(double rerankScore) { + return WeaviateProtoSearchGet.SearchResult.newBuilder() + .setMetadata(WeaviateProtoSearchGet.MetadataResult.newBuilder() + .setId("d3b07384-d113-4ec4-92e5-1e0f0ab84d43") + .setRerankScore(rerankScore) + .setRerankScorePresent(true)); + } + + private static WeaviateProtoSearchGet.SearchRequest marshal(QueryOperator operator) { + return QueryRequest.marshal( + new QueryRequest(operator, null), + CollectionDescriptor.ofMap("Things"), + CollectionHandleDefaults.of(ObjectBuilder.identity())); + } +}