weaviategh-303: support maxVectorDistance for hybrid search

bevzzz · Oct 9, 2024 · fcec60b · fcec60b
1 parent 9a99db6
commit fcec60b
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git a/src/main/java/io/weaviate/client/v1/graphql/query/argument/HybridArgument.java b/src/main/java/io/weaviate/client/v1/graphql/query/argument/HybridArgument.java
@@ -3,16 +3,16 @@
 import java.util.LinkedHashSet;
 import java.util.Set;
 
+import org.apache.commons.lang3.ArrayUtils;
+import org.apache.commons.lang3.StringUtils;
+
 import io.weaviate.client.v1.graphql.query.util.Serializer;
-import java.util.stream.Collectors;
 import lombok.AccessLevel;
 import lombok.Builder;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.ToString;
 import lombok.experimental.FieldDefaults;
-import org.apache.commons.lang3.ArrayUtils;
-import org.apache.commons.lang3.StringUtils;
 
 @Getter
 @Builder
@@ -22,6 +22,7 @@
 public class HybridArgument implements Argument {
   String query;
   Float alpha;
+  Float maxVectorDistance;
   Float[] vector;
   String fusionType;
   String[] properties;
@@ -40,14 +41,17 @@ public String build() {
     if (alpha != null) {
       arg.add(String.format("alpha:%s", alpha));
     }
+    if (maxVectorDistance != null) {
+      arg.add(String.format("maxVectorDistance:%s", maxVectorDistance));
+    }
     if (ArrayUtils.isNotEmpty(properties)) {
       arg.add(String.format("properties:%s", Serializer.arrayWithQuotes(properties)));
     }
     if (StringUtils.isNotBlank(fusionType)) {
       arg.add(String.format("fusionType:%s", fusionType));
     }
     if (ArrayUtils.isNotEmpty(targetVectors)) {
-      arg.add(String.format("targetVectors:%s",  Serializer.arrayWithQuotes(targetVectors)));
+      arg.add(String.format("targetVectors:%s", Serializer.arrayWithQuotes(targetVectors)));
     }
     if (searches != null && (searches.nearVector != null || searches.nearText != null)) {
       Set<String> searchesArgs = new LinkedHashSet<>();

diff --git a/src/test/java/io/weaviate/client/v1/graphql/query/argument/HybridArgumentTest.java b/src/test/java/io/weaviate/client/v1/graphql/query/argument/HybridArgumentTest.java
@@ -103,7 +103,18 @@ static Stream<Arguments> testCases() {
                                                                         .build())
                                                                 // @formatter:on
                                                                 .build(),
-                                                "hybrid:{query:\"I'm a simple string\" targets:{combinationMethod:minimum targetVectors:[\"t1\",\"t2\"] weights:{t1:0.8 t2:0.2}}}"));
+                                                "hybrid:{query:\"I'm a simple string\" targets:{combinationMethod:minimum targetVectors:[\"t1\",\"t2\"] weights:{t1:0.8 t2:0.2}}}"),
+                                Arguments.of(
+                                                "max vector distance",
+                                                HybridArgument.builder().query("I'm a simple string")
+                                                // @formatter:off
+                                                .searches(HybridArgument.Searches.builder().nearText(
+                                                        NearTextArgument.builder().concepts(new String[]{"concept"}).build()).build()
+                                                )
+                                                .maxVectorDistance(.5f)
+                                                // @formatter:on
+                                                                .build(),
+                                                "hybrid:{query:\"I'm a simple string\" maxVectorDistance:0.5 searches:{nearText:{concepts:[\"concept\"]}}}"));
         }
 
         @ParameterizedTest