working multivector query class and tests

Author: justin-cechmanek

redisvl/query/__init__.py

@@ -1,4 +1,4 @@
-from redisvl.query.aggregate import AggregationQuery, HybridQuery
+from redisvl.query.aggregate import AggregationQuery, HybridQuery, MultiVectorQuery
 from redisvl.query.query import (
     BaseQuery,
     BaseVectorQuery,
@@ -21,4 +21,5 @@
     "TextQuery",
     "AggregationQuery",
     "HybridQuery",
+    "MultiVectorQuery",
 ]

redisvl/query/aggregate.py

@@ -231,76 +231,35 @@ def __str__(self) -> str:
 
 class MultiVectorQuery(AggregationQuery):
     """
-        MultiVectorQuery allows for search over multiple vector fields in a document simulateously.
-        The final score will be a weighted combination of the individual vector similarity scores
-        following the formula:
+    MultiVectorQuery allows for search over multiple vector fields in a document simulateously.
+    The final score will be a weighted combination of the individual vector similarity scores
+    following the formula:
 
-        score = (w_1 * score_1 + w_2 * score_2 + w_3 * score_3 + ... ) / (w_1 + w_2 + w_3 + ...)
-
-        Vectors may be of different size and datatype.
-
-        .. code-block:: python
-
-            from redisvl.query import MultiVectorQuery
-            from redisvl.index import SearchIndex
-
-            index = SearchIndex.from_yaml("path/to/index.yaml")
-
-            query = MultiVectorQuery(
-                vectors=[[0.1, 0.2, 0.3], [0.5, 0.5], [0.1, 0.1, 0.1, 0.1]],
-                vector_field_names=["text_vector", "image_vector", "feature_vector"]
-                filter_expression=None,
-                weights=[0.7],
-                dtypes=["float32", "float32", "float32"],
-                num_results=10,
-                return_fields=["field1", "field2"],
-                dialect=2,
-            )
-
-            results = index.query(query)
-
-
-
-        FT.AGGREGATE multi_vector_test 
-        "@user_embedding:[VECTOR_RANGE 2.0 $vector_0]=>{$YIELD_DISTANCE_AS: distance_0}
-        | @image_embedding:[VECTOR_RANGE 2.0 $vector_1]=>{$YIELD_DISTANCE_AS: distance_1}" 
-        PARAMS 4
-        vector_0 "\xcd\xcc\xcc=\xcd\xcc\xcc=\x00\x00\x00?" 
-        vector_1 "\x9a\x99\x99\x99\x99\x99\xb9?\x9a\x99\x99\x99\x99\x99\xb9?\x9a\x99\x99\x99\x99\x99\xb9?\x9a\x99\x99\x99\x99\x99\xb9?\x9a\x99\x99\x99\x99\x99\xb9?" 
-        APPLY "(2 - @distance_0)/2" AS score_0
-        APPLY "(2 - @distance_1)/2" AS score_1 
-        DIALECT 2
-        APPLY "(@score_0 + @score_1)" AS combined_score
-        SORTBY 2 @combined_score 
-        ASC 
-        MAX 10 
-        LOAD 2 score_0 score_1
+    score = (w_1 * score_1 + w_2 * score_2 + w_3 * score_3 + ... )
 
+    Vectors may be of different size and datatype, but must be indexed using the 'cosine' distance_metric.
 
+    .. code-block:: python
 
+        from redisvl.query import MultiVectorQuery
+        from redisvl.index import SearchIndex
 
+        index = SearchIndex.from_yaml("path/to/index.yaml")
 
-    FT.AGGREGATE 'idx:characters'
-     "@embedding1:[VECTOR_RANGE .7 $vector1]=>{$YIELD_DISTANCE_AS: vector_distance1}
-     | @embedding2:[VECTOR_RANGE 1.0 $vector2]=>{$YIELD_DISTANCE_AS: vector_distance2}
-     | @embedding3:[VECTOR_RANGE 1.7 $vector3]=>{$YIELD_DISTANCE_AS: vector_distance3}
-     | @name:(James)"
-     ### ADDSCORES
-     ### SCORER BM25STD.NORM
-     ### LOAD 2 created_at @embedding
-     APPLY '(2 - @vector_distance1)/2' as v1
-     APPLY '(2 - @vector_distance2)/2' as v2
-     APPLY '(2 - @vector_distance3)/2' as v3
-     APPLY '(@__score * 0.3 + (@v1 * 0.3) + (@v2 * 1.2) + (@v3 * 0.1))' AS final_score
-     PARAMS 6 vector1 "\xe4\xd6..." vector2 "\x89\xa0..." vector3 "\x3c\x19..."
-     SORTBY 2 @final_score DESC
-     DIALECT 2
-     LIMIT 0 100
+        query = MultiVectorQuery(
+            vectors=[[0.1, 0.2, 0.3], [0.5, 0.5], [0.1, 0.1, 0.1, 0.1]],
+            vector_field_names=["text_vector", "image_vector", "feature_vector"]
+            filter_expression=None,
+            weights=[0.7, 0.2, 0.5],
+            dtypes=["float32", "bfloat16", "float64"],
+            num_results=10,
+            return_fields=["field1", "field2"],
+            dialect=2,
+        )
 
+        results = index.query(query)
     """
 
-    DISTANCE_ID: str = "vector_distance"
-
     def __init__(
         self,
         vectors: Union[bytes, List[bytes], List[float], List[List[float]]],
@@ -340,58 +299,69 @@ def __init__(
         self._dtypes = dtypes
         self._num_results = num_results
 
-        if len(vectors) == 0 or len(vector_field_names) == 0 or len(weights) == 0:
+        if any([len(x) == 0 for x in [vectors, vector_field_names, weights, dtypes]]):
             raise ValueError(
                 f"""The number of vectors and vector field names must be equal.
-                             If weights are specified their number must match the number of vectors and vector field names also.
-                            Length of vectors list: {len(vectors) = }
-                            Length of vector_field_names list: {len(vector_field_names) = }
-                            Length of weights list: {len(weights) = }
-                            """
+                    If weights or dtypes are specified their number must match the number of vectors and vector field names also.
+                    Length of vectors list: {len(vectors) = }
+                    Length of vector_field_names list: {len(vector_field_names) = }
+                    Length of weights list: {len(weights) = }
+                    length of dtypes list: {len(dtypes) = }
+                    """
             )
 
         if isinstance(vectors, bytes) or isinstance(vectors[0], float):
             self._vectors = [vectors]
         else:
-            self._vectors = vectors
+            self._vectors = vectors  # type: ignore
+
         if isinstance(vector_field_names, str):
             self._vector_field_names = [vector_field_names]
         else:
             self._vector_field_names = vector_field_names
+
         if len(weights) == 1:
             self._weights = weights * len(vectors)
         else:
             self._weights = weights
+
         if len(dtypes) == 1:
             self._dtypes = dtypes * len(vectors)
         else:
             self._dtypes = dtypes
 
-        if (len(self._vectors) != len(self._vector_field_names)) or (
-            len(self._vectors) != len(self._weights)
+        num_vectors = len(self._vectors)
+        if any(
+            [
+                len(x) != num_vectors  # type: ignore
+                for x in [self._vector_field_names, self._weights, self._dtypes]
+            ]
         ):
             raise ValueError(
                 f"""The number of vectors and vector field names must be equal.
-                             If weights are specified their number must match the number of vectors and vector field names also.
-                            Length of vectors list: {len(self._vectors) = }
-                            Length of vector_field_names list: {len(self._vector_field_names) = }
-                            Length of weights list: {len(self._weights) = }
-                            """
+                    If weights or dtypes are specified their number must match the number of vectors and vector field names also.
+                    Length of vectors list: {len(self._vectors) = }
+                    Length of vector_field_names list: {len(self._vector_field_names) = }
+                    Length of weights list: {len(self._weights) = }
+                    Length of dtypes list: {len(self._dtypes) = }
+                    """
             )
 
         query_string = self._build_query_string()
         super().__init__(query_string)
 
+        # calculate the respective vector similarities
+        for i in range(len(vectors)):
+            self.apply(**{f"score_{i}": f"(2 - @distance_{i})/2"})
+
         # construct the scoring string based on the vector similarity scores and weights
         combined_scores = []
         for i, w in enumerate(self._weights):
             combined_scores.append(f"@score_{i} * {w}")
         combined_score_string = " + ".join(combined_scores)
-        combined_score_string = f"'({combined_score_string})'"
 
         self.apply(combined_score=combined_score_string)
 
-        # self.add_scores()
         self.sort_by(Desc("@combined_score"), max=num_results)  # type: ignore
         self.dialect(dialect)
         if return_fields:
@@ -405,43 +375,34 @@ def params(self) -> Dict[str, Any]:
             Dict[str, Any]: The parameters for the aggregation.
         """
         params = {}
-        for i, (vector, vector_field, dtype) in enumerate(zip(
-            self._vectors, self._vector_field_names, self._dtypes
-        )):
+        for i, (vector, dtype) in enumerate(zip(self._vectors, self._dtypes)):
             if isinstance(vector, list):
-                vector = array_to_buffer(vector, dtype=dtype)
+                vector = array_to_buffer(vector, dtype=dtype)  # type: ignore
             params[f"vector_{i}"] = vector
         return params
 
     def _build_query_string(self) -> str:
         """Build the full query string for text search with optional filtering."""
 
-        filter_expression = self._filter_expression
-        if isinstance(self._filter_expression, FilterExpression):
-            filter_expression = str(self._filter_expression)
-
         # base KNN query
-        knn_queries = []
         range_queries = []
-        for i, (vector, field) in enumerate(zip(self._vectors, self._vector_field_names)):
-            knn_queries.append(f"[KNN {self._num_results} @{field} $vector_{i} AS distance_{i}]")
-            range_queries.append(f"@{field}:[VECTOR_RANGE 2.0 $vector_{i}]=>{{$YIELD_DISTANCE_AS: distance_{i}}}")
+        for i, (vector, field) in enumerate(
+            zip(self._vectors, self._vector_field_names)
+        ):
+            range_queries.append(
+                f"@{field}:[VECTOR_RANGE 2.0 $vector_{i}]=>{{$YIELD_DISTANCE_AS: distance_{i}}}"
+            )
 
-        knn_query = " | ".join(knn_queries) ## knn_queries format doesn't work
-        knn_query = " | ".join(range_queries)
+        range_query = " | ".join(range_queries)
 
-        # calculate the respective vector similarities
-        apply_string = ""
-        for i, (vector, field_name, weight) in enumerate(
-            zip(self._vectors, self._vector_field_names, self._weights)
-        ):
-            apply_string += f'APPLY "(2 - @distance_{i})/2" AS score_{i} '
+        filter_expression = self._filter_expression
+        if isinstance(self._filter_expression, FilterExpression):
+            filter_expression = str(self._filter_expression)
 
-        return (
-            f"{knn_query} {filter_expression} {apply_string}"
-            if filter_expression
-            else f"{knn_query} {apply_string}"
-        )
+        if filter_expression:
+            return f"({range_query}) AND ({filter_expression})"
+        else:
+            return f"{range_query}"
 
     def __str__(self) -> str:
         """Return the string representation of the query."""