Add support for full text queries and hybrid search queries (#303)

Co-authored-by: Robert Shelton <[email protected]>
Co-authored-by: Andrew Brookins <[email protected]>
Co-authored-by: Tyler Hutcherson <[email protected]>
Co-authored-by: Robert Shelton <[email protected]>

Author: 4 people

poetry.lock

@@ -33,8 +33,8 @@ tenacity = ">=8.2.2"
 tabulate = "^0.9.0"
 ml-dtypes = "^0.4.0"
 python-ulid = "^3.0.0"
+nltk = { version = "^3.8.1", optional = true }
 jsonpath-ng = "^1.5.0"
-
 openai = { version = "^1.13.0", optional = true }
 sentence-transformers = { version = "^3.4.0", optional = true }
 scipy = [
@@ -58,6 +58,7 @@ mistralai = ["mistralai"]
 voyageai = ["voyageai"]
 ranx = ["ranx"]
 bedrock = ["boto3"]
+nltk = ["nltk"]
 
 [tool.poetry.group.dev.dependencies]
 black = "^25.1.0"

pyproject.toml

@@ -18,6 +18,7 @@
     Union,
 )
 
+from redisvl.redis.utils import convert_bytes, make_dict
 from redisvl.utils.utils import deprecated_argument, deprecated_function, sync_wrapper
 
 if TYPE_CHECKING:
@@ -39,7 +40,14 @@
     SchemaValidationError,
 )
 from redisvl.index.storage import BaseStorage, HashStorage, JsonStorage
-from redisvl.query import BaseQuery, BaseVectorQuery, CountQuery, FilterQuery
+from redisvl.query import (
+    AggregationQuery,
+    BaseQuery,
+    BaseVectorQuery,
+    CountQuery,
+    FilterQuery,
+    HybridQuery,
+)
 from redisvl.query.filter import FilterExpression
 from redisvl.redis.connection import (
     RedisConnectionFactory,
@@ -138,6 +146,34 @@ def _process(doc: "Document") -> Dict[str, Any]:
     return [_process(doc) for doc in results.docs]
 
 
+def process_aggregate_results(
+    results: "AggregateResult", query: AggregationQuery, storage_type: StorageType
+) -> List[Dict[str, Any]]:
+    """Convert an aggregate reslt object into a list of document dictionaries.
+
+    This function processes results from Redis, handling different storage
+    types and query types. For JSON storage with empty return fields, it
+    unpacks the JSON object while retaining the document ID. The 'payload'
+    field is also removed from all resulting documents for consistency.
+
+    Args:
+        results (AggregarteResult): The aggregart results from Redis.
+        query (AggregationQuery): The aggregation query object used for the aggregation.
+        storage_type (StorageType): The storage type of the search
+            index (json or hash).
+
+    Returns:
+        List[Dict[str, Any]]: A list of processed document dictionaries.
+    """
+
+    def _process(row):
+        result = make_dict(convert_bytes(row))
+        result.pop("__score", None)
+        return result
+
+    return [_process(r) for r in results.rows]
+
+
 class BaseSearchIndex:
     """Base search engine class"""
 
@@ -650,6 +686,17 @@ def fetch(self, id: str) -> Optional[Dict[str, Any]]:
             return convert_bytes(obj[0])
         return None
 
+    def _aggregate(self, aggregation_query: AggregationQuery) -> List[Dict[str, Any]]:
+        """Execute an aggretation query and processes the results."""
+        results = self.aggregate(
+            aggregation_query, query_params=aggregation_query.params  # type: ignore[attr-defined]
+        )
+        return process_aggregate_results(
+            results,
+            query=aggregation_query,
+            storage_type=self.schema.index.storage_type,
+        )
+
     def aggregate(self, *args, **kwargs) -> "AggregateResult":
         """Perform an aggregation operation against the index.
 
@@ -772,14 +819,14 @@ def _query(self, query: BaseQuery) -> List[Dict[str, Any]]:
         results = self.search(query.query, query_params=query.params)
         return process_results(results, query=query, schema=self.schema)
 
-    def query(self, query: BaseQuery) -> List[Dict[str, Any]]:
+    def query(self, query: Union[BaseQuery, AggregationQuery]) -> List[Dict[str, Any]]:
         """Execute a query on the index.
 
-        This method takes a BaseQuery object directly, runs the search, and
+        This method takes a BaseQuery or AggregationQuery object directly, and
         handles post-processing of the search.
 
         Args:
-            query (BaseQuery): The query to run.
+            query (Union[BaseQuery, AggregateQuery]): The query to run.
 
         Returns:
             List[Result]: A list of search results.
@@ -797,7 +844,10 @@ def query(self, query: BaseQuery) -> List[Dict[str, Any]]:
             results = index.query(query)
 
         """
-        return self._query(query)
+        if isinstance(query, AggregationQuery):
+            return self._aggregate(query)
+        else:
+            return self._query(query)
 
     def paginate(self, query: BaseQuery, page_size: int = 30) -> Generator:
         """Execute a given query against the index and return results in
@@ -1303,6 +1353,19 @@ async def fetch(self, id: str) -> Optional[Dict[str, Any]]:
             return convert_bytes(obj[0])
         return None
 
+    async def _aggregate(
+        self, aggregation_query: AggregationQuery
+    ) -> List[Dict[str, Any]]:
+        """Execute an aggretation query and processes the results."""
+        results = await self.aggregate(
+            aggregation_query, query_params=aggregation_query.params  # type: ignore[attr-defined]
+        )
+        return process_aggregate_results(
+            results,
+            query=aggregation_query,
+            storage_type=self.schema.index.storage_type,
+        )
+
     async def aggregate(self, *args, **kwargs) -> "AggregateResult":
         """Perform an aggregation operation against the index.
 
@@ -1426,14 +1489,16 @@ async def _query(self, query: BaseQuery) -> List[Dict[str, Any]]:
         results = await self.search(query.query, query_params=query.params)
         return process_results(results, query=query, schema=self.schema)
 
-    async def query(self, query: BaseQuery) -> List[Dict[str, Any]]:
+    async def query(
+        self, query: Union[BaseQuery, AggregationQuery]
+    ) -> List[Dict[str, Any]]:
         """Asynchronously execute a query on the index.
 
-        This method takes a BaseQuery object directly, runs the search, and
-        handles post-processing of the search.
+        This method takes a BaseQuery or AggregationQuery object directly, runs
+        the search, and handles post-processing of the search.
 
         Args:
-            query (BaseQuery): The query to run.
+            query (Union[BaseQuery, AggregateQuery]): The query to run.
 
         Returns:
             List[Result]: A list of search results.
@@ -1450,7 +1515,10 @@ async def query(self, query: BaseQuery) -> List[Dict[str, Any]]:
 
             results = await index.query(query)
         """
-        return await self._query(query)
+        if isinstance(query, AggregationQuery):
+            return await self._aggregate(query)
+        else:
+            return await self._query(query)
 
     async def paginate(self, query: BaseQuery, page_size: int = 30) -> AsyncGenerator:
         """Execute a given query against the index and return results in

redisvl/index/index.py

@@ -1,9 +1,11 @@
+from redisvl.query.aggregate import AggregationQuery, HybridQuery
 from redisvl.query.query import (
     BaseQuery,
     BaseVectorQuery,
     CountQuery,
     FilterQuery,
     RangeQuery,
+    TextQuery,
     VectorQuery,
     VectorRangeQuery,
 )
@@ -16,4 +18,7 @@
     "RangeQuery",
     "VectorRangeQuery",
     "CountQuery",
+    "TextQuery",
+    "AggregationQuery",
+    "HybridQuery",
 ]