From 2f7c7905904503c5b12663571b2c516403e91cab Mon Sep 17 00:00:00 2001 From: Duda Nogueira Date: Mon, 24 Aug 2026 20:43:48 -0300 Subject: [PATCH 1/3] fix(query): allow rerank on every search operator rerank(Rerank) was declared on BaseVectorSearchBuilder, so only the near* searches could be reranked. BM25, Hybrid and FetchObjects could not, even though rerank is a top-level field of the search request and the server applies it the same way for all of them. Move the option to BaseQueryOptions, which every operator builder extends, and marshal it from BaseQueryOptions.appendTo. The near* records no longer carry their own rerank component; QueryOperator.rerank() now reads it off the common options, so it stays readable on every operator. Closes #603 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WmY5dAGWCccWDoqkKNC2JU --- .../collections/query/BaseQueryOptions.java | 21 ++++++ .../query/BaseVectorSearchBuilder.java | 10 --- .../v1/api/collections/query/NearAudio.java | 2 - .../v1/api/collections/query/NearDepth.java | 2 - .../v1/api/collections/query/NearImage.java | 2 - .../v1/api/collections/query/NearImu.java | 2 - .../v1/api/collections/query/NearObject.java | 2 - .../v1/api/collections/query/NearText.java | 2 - .../v1/api/collections/query/NearThermal.java | 2 - .../v1/api/collections/query/NearVector.java | 2 - .../v1/api/collections/query/NearVideo.java | 2 - .../api/collections/query/QueryOperator.java | 9 ++- .../api/collections/query/QueryRequest.java | 3 - .../v1/api/collections/query/RerankTest.java | 64 +++++++++++++++++++ 14 files changed, 93 insertions(+), 32 deletions(-) create mode 100644 src/test/java/io/weaviate/client6/v1/api/collections/query/RerankTest.java 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 + * 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/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..dd799883f --- /dev/null +++ b/src/test/java/io/weaviate/client6/v1/api/collections/query/RerankTest.java @@ -0,0 +1,64 @@ +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.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(); + } + + private static WeaviateProtoSearchGet.SearchRequest marshal(QueryOperator operator) { + return QueryRequest.marshal( + new QueryRequest(operator, null), + CollectionDescriptor.ofMap("Things"), + CollectionHandleDefaults.of(ObjectBuilder.identity())); + } +} From 1fa0d25cd86420f3736d7cf6fe846abd2d349a97 Mon Sep 17 00:00:00 2001 From: Duda Nogueira Date: Mon, 24 Aug 2026 20:45:18 -0300 Subject: [PATCH 2/3] fix(query): unmarshal the rerank score from the search reply The server returns a rerank score per object (MetadataResult.rerank_score) and per group (GroupByResult.rerank), but neither was read: a reranked search arrived correctly ordered with the number that produced the order missing, and no way to get it short of re-reading the reply. Add rerankScore to QueryMetadata and QueryResponseGroup and populate both. The score is only set when the reply says it is present -- 0.0 is what the dummy reranker returns for a match, so the value alone cannot distinguish "not reranked" from "reranked with score 0". Closes #604 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WmY5dAGWCccWDoqkKNC2JU --- .../io/weaviate/integration/SearchITest.java | 51 +++++++++++++++-- .../api/collections/query/QueryMetadata.java | 19 ++++++- .../api/collections/query/QueryResponse.java | 5 ++ .../collections/query/QueryResponseGroup.java | 8 +++ .../query/QueryResponseGrouped.java | 1 + .../v1/api/collections/query/RerankTest.java | 56 +++++++++++++++++++ 6 files changed, 134 insertions(+), 6 deletions(-) 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/query/QueryMetadata.java b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryMetadata.java index dd8a54bfb..4a3e007d0 100644 --- a/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryMetadata.java +++ b/src/main/java/io/weaviate/client6/v1/api/collections/query/QueryMetadata.java @@ -9,14 +9,23 @@ public record QueryMetadata( /** BM25 ranking score. */ Float score, /** Components of the BM25 ranking score. */ - String explainScore) { + String explainScore, + /** + * Score assigned by the reranker module. + * + *

+ * 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/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/query/RerankTest.java b/src/test/java/io/weaviate/client6/v1/api/collections/query/RerankTest.java index dd799883f..0aaa170fd 100644 --- 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 @@ -9,6 +9,7 @@ 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; @@ -55,6 +56,61 @@ public void test_rerankWithoutQuery() { 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), From 1814b7a0c2db9e80f0e3676a2219532a0a9303d9 Mon Sep 17 00:00:00 2001 From: Duda Nogueira Date: Tue, 25 Aug 2026 15:12:51 -0300 Subject: [PATCH 3/3] fix(generate): unmarshal the group rerank score in generative search The generative search reuses the query operators, so it accepts rerank the same way -- but GenerativeResponseGrouped dropped the group-level score that QueryResponseGrouped reads, leaving that one path unable to tell how the groups were ranked. Per-object scores were already covered: every unmarshal path funnels through QueryResponse.unmarshalResultObject. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WmY5dAGWCccWDoqkKNC2JU --- .../generate/GenerativeResponseGroup.java | 10 +++ .../generate/GenerativeResponseGrouped.java | 1 + .../generate/GenerativeRerankTest.java | 77 +++++++++++++++++++ 3 files changed, 88 insertions(+) create mode 100644 src/test/java/io/weaviate/client6/v1/api/collections/generate/GenerativeRerankTest.java 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/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())); + } +}