fixes query string parsing. adds more tests

Author: justin-cechmanek

redisvl/query/aggregate.py

@@ -112,9 +112,6 @@ def __init__(
         self.set_stopwords(stopwords)
 
         query_string = self._build_query_string()
-        ####
-        print('query_string: ', query_string)
-        ####
         super().__init__(query_string)
 
         self.scorer(text_scorer)  # type: ignore[attr-defined]
@@ -125,7 +122,6 @@ def __init__(
         self.apply(hybrid_score=f"{1-alpha}*@text_score + {alpha}*@vector_similarity")
         self.sort_by(Desc("@hybrid_score"), max=num_results)
         self.dialect(dialect)  # type: ignore[attr-defined]
-
         if return_fields:
             self.load(*return_fields)
 
@@ -216,9 +212,9 @@ def _build_query_string(self) -> str:
         # base KNN query
         knn_query = f"KNN {self._num_results} @{self._vector_field} ${self.VECTOR_PARAM} AS {self.DISTANCE_ID}"
 
-        text = f"(~@{self._text_field}:({self.tokenize_and_escape_query(self._text)}))"
+        text = f"(~@{self._text_field}:({self.tokenize_and_escape_query(self._text)})"
 
         if filter_expression and filter_expression != "*":
-            text += f"({filter_expression})"
+            text += f" AND {filter_expression}"
 
-        return f"{text}=>[{knn_query}]"
+        return f"{text})=>[{knn_query}]"

redisvl/query/query.py

@@ -687,7 +687,7 @@ class RangeQuery(VectorRangeQuery):
     pass
 
 
-class TextQuery(FilterQuery):
+class TextQuery(BaseQuery):
     def __init__(
         self,
         text: str,
@@ -747,22 +747,26 @@ def __init__(
         self.set_stopwords(stopwords)
         self.set_filter(filter_expression)
 
-        query_string = self._build_query_string()
+        if params:
+            self._params = params
 
-        super().__init__(
-            query_string,
-            return_fields=return_fields,
-            num_results=num_results,
-            dialect=dialect,
-            sort_by=sort_by,
-            in_order=in_order,
-            params=params,
-        )
+        self._num_results = num_results
 
-        # Handle query modifiers
-        self.scorer(self._text_scorer)
+        # initialize the base query with the full query string and filter expression
+        query_string = self._build_query_string()
+        super().__init__(query_string)
+
+        # Handle query settings
+        if return_fields:
+            self.return_fields(*return_fields)
         self.paging(0, self._num_results).dialect(dialect)
 
+        if sort_by:
+            self.sort_by(sort_by)
+
+        if in_order:
+            self.in_order()
+
         if return_score:
             self.with_scores()
 
@@ -812,7 +816,7 @@ def _build_query_string(self) -> str:
         else:
             filter_expression = ""
 
-        text = f"(@{self._text_field}:({self.tokenize_and_escape_query(self._text)}))"
+        text = f"@{self._text_field}:({self.tokenize_and_escape_query(self._text)})"
         if filter_expression and filter_expression != "*":
-            text += f"({filter_expression})"
+            text += f" AND {filter_expression}"
         return text

tests/integration/test_aggregation.py

@@ -2,18 +2,10 @@
 from redis.commands.search.aggregation import AggregateResult
 from redis.commands.search.result import Result
 
-from redisvl.redis.connection import compare_versions
-
 from redisvl.index import SearchIndex
 from redisvl.query import HybridAggregationQuery
-from redisvl.query.filter import (
-    FilterExpression,
-    Geo,
-    GeoRadius,
-    Num,
-    Tag,
-    Text,
-)
+from redisvl.query.filter import FilterExpression, Geo, GeoRadius, Num, Tag, Text
+from redisvl.redis.connection import compare_versions
 from redisvl.redis.utils import array_to_buffer
 
 
@@ -131,17 +123,24 @@ def test_empty_query_string():
     # test if text is empty
     with pytest.raises(ValueError):
         hybrid_query = HybridAggregationQuery(
-            text=text, text_field_name=text_field, vector=vector, vector_field_name=vector_field
+            text=text,
+            text_field_name=text_field,
+            vector=vector,
+            vector_field_name=vector_field,
         )
 
     # test if text becomes empty after stopwords are removed
     text = "with a for but and"  # will all be removed as default stopwords
     with pytest.raises(ValueError):
         hybrid_query = HybridAggregationQuery(
-            text=text, text_field_name=text_field, vector=vector, vector_field_name=vector_field
+            text=text,
+            text_field_name=text_field,
+            vector=vector,
+            vector_field_name=vector_field,
         )
 
-def test_aggregation_query_filter(index):
+
+def test_aggregation_query_with_filter(index):
     redis_version = index.client.info()["redis_version"]
     if not compare_versions(redis_version, "7.2.0"):
         pytest.skip("Not using a late enough version of Redis")
@@ -163,7 +162,7 @@ def test_aggregation_query_filter(index):
     )
 
     results = index.aggregate_query(hybrid_query)
-    assert len(results) == 3
+    assert len(results) == 2
     for result in results:
         assert result["credit_score"] == "high"
         assert int(result["age"]) > 30
@@ -179,7 +178,7 @@ def test_aggregation_query_with_geo_filter(index):
     vector = [0.1, 0.1, 0.5]
     vector_field = "user_embedding"
     return_fields = ["user", "credit_score", "age", "job", "location", "description"]
-    filter_expression = Geo("location") == GeoRadius(37.7749, -122.4194, 1000)
+    filter_expression = Geo("location") == GeoRadius(-122.4194, 37.7749, 1000, "m")
 
     hybrid_query = HybridAggregationQuery(
         text=text,
@@ -196,6 +195,36 @@ def test_aggregation_query_with_geo_filter(index):
         assert result["location"] is not None
 
 
+@pytest.mark.parametrize("alpha", [0.1, 0.5, 0.9])
+def test_aggregate_query_alpha(index, alpha):
+    redis_version = index.client.info()["redis_version"]
+    if not compare_versions(redis_version, "7.2.0"):
+        pytest.skip("Not using a late enough version of Redis")
+
+    text = "a medical professional with expertise in lung cancer"
+    text_field = "description"
+    vector = [0.1, 0.1, 0.5]
+    vector_field = "user_embedding"
+
+    hybrid_query = HybridAggregationQuery(
+        text=text,
+        text_field_name=text_field,
+        vector=vector,
+        vector_field_name=vector_field,
+        alpha=alpha,
+    )
+
+    results = index.aggregate_query(hybrid_query)
+    assert len(results) == 7
+    for result in results:
+        score = alpha * float(result["vector_similarity"]) + (1 - alpha) * float(
+            result["text_score"]
+        )
+        assert (
+            float(result["hybrid_score"]) - score <= 0.0001
+        )  # allow for small floating point error
+
+
 def test_aggregate_query_stopwords(index):
     redis_version = index.client.info()["redis_version"]
     if not compare_versions(redis_version, "7.2.0"):
@@ -205,24 +234,34 @@ def test_aggregate_query_stopwords(index):
     text_field = "description"
     vector = [0.1, 0.1, 0.5]
     vector_field = "user_embedding"
+    alpha = 0.5
 
     hybrid_query = HybridAggregationQuery(
         text=text,
         text_field_name=text_field,
         vector=vector,
         vector_field_name=vector_field,
-        alpha=0.5,
+        alpha=alpha,
         stopwords=["medical", "expertise"],
     )
 
+    query_string = hybrid_query._build_query_string()
+
+    assert "medical" not in query_string
+    assert "expertize" not in query_string
+
     results = index.aggregate_query(hybrid_query)
     assert len(results) == 7
-    for r in results:
-        assert r["text_score"] == 0
-        assert r["hybrid_score"] == 0.5 * r["vector_similarity"]
+    for result in results:
+        score = alpha * float(result["vector_similarity"]) + (1 - alpha) * float(
+            result["text_score"]
+        )
+        assert (
+            float(result["hybrid_score"]) - score <= 0.0001
+        )  # allow for small floating point error
 
 
-def test_aggregate_query_text_filter(index):
+def test_aggregate_query_with_text_filter(index):
     redis_version = index.client.info()["redis_version"]
     if not compare_versions(redis_version, "7.2.0"):
         pytest.skip("Not using a late enough version of Redis")
@@ -231,19 +270,39 @@ def test_aggregate_query_text_filter(index):
     text_field = "description"
     vector = [0.1, 0.1, 0.5]
     vector_field = "user_embedding"
-    filter_expression = (Text("description") == ("medical")) | (Text("job") % ("doct*"))
+    filter_expression = Text(text_field) == ("medical")
 
+    # make sure we can still apply filters to the same text field we are querying
     hybrid_query = HybridAggregationQuery(
         text=text,
         text_field_name=text_field,
         vector=vector,
         vector_field_name=vector_field,
         alpha=0.5,
-        filter_expression=filter_expression
-        )
+        filter_expression=filter_expression,
+        return_fields=["job", "description"],
+    )
 
     results = index.aggregate_query(hybrid_query)
-    assert len(results) == 7
+    assert len(results) == 2
+    for result in results:
+        assert "medical" in result[text_field].lower()
+
+    filter_expression = (Text(text_field) == ("medical")) & (
+        (Text(text_field) != ("research"))
+    )
+    hybrid_query = HybridAggregationQuery(
+        text=text,
+        text_field_name=text_field,
+        vector=vector,
+        vector_field_name=vector_field,
+        alpha=0.5,
+        filter_expression=filter_expression,
+        return_fields=["description"],
+    )
+
+    results = index.aggregate_query(hybrid_query)
+    assert len(results) == 2
     for result in results:
-        assert result["text_score"] == 0
-        assert result["hybrid_score"] == 0.5 * result["vector_similarity"]
+        assert "medical" in result[text_field].lower()
+        assert "research" not in result[text_field].lower()

tests/integration/test_query.py

@@ -4,7 +4,13 @@
 from redis.commands.search.result import Result
 
 from redisvl.index import SearchIndex
-from redisvl.query import CountQuery, FilterQuery, VectorQuery, VectorRangeQuery
+from redisvl.query import (
+    CountQuery,
+    FilterQuery,
+    TextQuery,
+    VectorQuery,
+    VectorRangeQuery,
+)
 from redisvl.query.filter import (
     FilterExpression,
     Geo,
@@ -147,6 +153,7 @@ def index(sample_data, redis_url):
                 "storage_type": "hash",
             },
             "fields": [
+                {"name": "description", "type": "text"},
                 {"name": "credit_score", "type": "tag"},
                 {"name": "job", "type": "text"},
                 {"name": "age", "type": "numeric"},
@@ -790,3 +797,78 @@ def test_range_query_normalize_bad_input(index):
             return_fields=["user", "credit_score", "age", "job", "location"],
             distance_threshold=1.2,
         )
+
+
+def test_text_query(index):
+    text = "a medical professional with expertise in lung cancer"
+    text_field = "description"
+    return_fields = ["user", "credit_score", "age", "job", "location", "description"]
+
+    text_query = TextQuery(
+        text=text,
+        text_field_name=text_field,
+        return_fields=return_fields,
+    )
+    results = index.query(text_query)
+
+    assert len(results) == 4
+
+    # make sure at least one word from the query is in the description
+    for result in results:
+        assert any(word in result[text_field] for word in text.split())
+
+
+# test that text queryies work with filter expressions
+def test_text_query_with_filter(index):
+    text = "a medical professional with expertise in lung cancer"
+    text_field = "description"
+    return_fields = ["user", "credit_score", "age", "job", "location", "description"]
+    filter_expression = (Tag("credit_score") == ("high")) & (Num("age") > 30)
+
+    text_query = TextQuery(
+        text=text,
+        text_field_name=text_field,
+        filter_expression=filter_expression,
+        return_fields=return_fields,
+    )
+    results = index.query(text_query)
+    assert len(results) == 2
+    for result in results:
+        assert any(word in result[text_field] for word in text.split())
+        assert result["credit_score"] == "high"
+        assert int(result["age"]) > 30
+
+
+# test that text queryies workt with text filter expressions on the same text field
+def test_text_query_with_text_filter(index):
+    text = "a medical professional with expertise in lung cancer"
+    text_field = "description"
+    return_fields = ["age", "job", "description"]
+    filter_expression = Text(text_field) == ("medical")
+
+    text_query = TextQuery(
+        text=text,
+        text_field_name=text_field,
+        filter_expression=filter_expression,
+        return_fields=return_fields,
+    )
+    results = index.query(text_query)
+    assert len(results) == 2
+    for result in results:
+        assert any(word in result[text_field] for word in text.split())
+        assert "medical" in result[text_field]
+
+    filter_expression = Text(text_field) != ("research")
+
+    text_query = TextQuery(
+        text=text,
+        text_field_name=text_field,
+        filter_expression=filter_expression,
+        return_fields=return_fields,
+    )
+
+    results = index.query(text_query)
+    assert len(results) == 3
+    for result in results:
+        assert any(word in result[text_field] for word in text.split())
+        assert "research" not in result[text_field]