added metadata filtering to existing search methods

Author: epinzur

libs/knowledge-store/ragstack_knowledge_store/graph_store.py

@@ -29,7 +29,7 @@
 
 CONTENT_COLUMNS = "content_id, kind, text_content, attributes_blob, metadata_s, links_blob"
 
-SELECT_CQL_TEMPLATE = "SELECT {columns} FROM {table_name} {where_clause} {limit_clause};"
+SELECT_CQL_TEMPLATE = "SELECT {columns} FROM {table_name} {where_clause} {order_clause} {limit_clause};"
 
 @dataclass
 class Node:
@@ -198,28 +198,6 @@ def __init__(
             """  # noqa: S608
         )
 
-        self._query_by_embedding = session.prepare(
-            f"""
-            SELECT {CONTENT_COLUMNS}
-            FROM {keyspace}.{node_table}
-            ORDER BY text_embedding ANN OF ?
-            LIMIT ?
-            """  # noqa: S608
-        )
-        self._query_by_embedding.consistency_level = ConsistencyLevel.ONE
-
-        self._query_ids_and_link_to_tags_by_embedding = session.prepare(
-            f"""
-            SELECT content_id, link_to_tags
-            FROM {keyspace}.{node_table}
-            ORDER BY text_embedding ANN OF ?
-            LIMIT ?
-            """  # noqa: S608
-        )
-        self._query_ids_and_link_to_tags_by_embedding.consistency_level = (
-            ConsistencyLevel.ONE
-        )
-
         self._query_ids_and_link_to_tags_by_id = session.prepare(
             f"""
             SELECT content_id, link_to_tags
@@ -228,18 +206,6 @@ def __init__(
             """  # noqa: S608
         )
 
-        self._query_ids_and_embedding_by_embedding = session.prepare(
-            f"""
-            SELECT content_id, text_embedding, link_to_tags
-            FROM {keyspace}.{node_table}
-            ORDER BY text_embedding ANN OF ?
-            LIMIT ?
-            """  # noqa: S608
-        )
-        self._query_ids_and_embedding_by_embedding.consistency_level = (
-            ConsistencyLevel.ONE
-        )
-
         self._query_source_tags_by_id = session.prepare(
             f"""
             SELECT link_to_tags
@@ -270,11 +236,14 @@ def __init__(
             """
         )
 
+    def table_name(self) -> str:
+        return f"{self._keyspace}.{self._node_table}"
+
     def _apply_schema(self) -> None:
         """Apply the schema to the database."""
         embedding_dim = len(self._embedding.embed_query("Test Query"))
         self._session.execute(f"""
-            CREATE TABLE IF NOT EXISTS {self._keyspace}.{self._node_table} (
+            CREATE TABLE IF NOT EXISTS {self.table_name()} (
                 content_id TEXT,
                 kind TEXT,
                 text_content TEXT,
@@ -293,19 +262,19 @@ def _apply_schema(self) -> None:
         # Index on text_embedding (for similarity search)
         self._session.execute(f"""
             CREATE CUSTOM INDEX IF NOT EXISTS {self._node_table}_text_embedding_index
-            ON {self._keyspace}.{self._node_table}(text_embedding)
+            ON {self.table_name()}(text_embedding)
             USING 'StorageAttachedIndex';
         """)
 
         self._session.execute(f"""
             CREATE CUSTOM INDEX IF NOT EXISTS {self._node_table}_link_from_tags
-            ON {self._keyspace}.{self._node_table}(link_from_tags)
+            ON {self.table_name()}(link_from_tags)
             USING 'StorageAttachedIndex';
         """)
 
         self._session.execute(f"""
             CREATE CUSTOM INDEX IF NOT EXISTS {self._node_table}_metadata_index
-            ON {self._keyspace}.{self._node_table}(ENTRIES(metadata_s))
+            ON {self.table_name()}(ENTRIES(metadata_s))
             USING 'StorageAttachedIndex';
         """)
 
@@ -425,6 +394,7 @@ def mmr_traversal_search(
         adjacent_k: int = 10,
         lambda_mult: float = 0.5,
         score_threshold: float = float("-inf"),
+        metadata: Optional[Dict[str, Any]] = [],
     ) -> Iterable[Node]:
         """Retrieve documents from this graph store using MMR-traversal.
 
@@ -450,6 +420,7 @@ def mmr_traversal_search(
                 diversity and 1 to minimum diversity. Defaults to 0.5.
             score_threshold: Only documents with a score greater than or equal
                 this threshold will be chosen. Defaults to -infinity.
+            metadata: Optional metadata to filter the results.
         """
         query_embedding = self._embedding.embed_query(query)
         helper = MmrHelper(
@@ -465,10 +436,14 @@ def mmr_traversal_search(
         # Fetch the initial candidates and add them to the helper and
         # outgoing_tags.
         def fetch_initial_candidates() -> None:
-            fetched = self._session.execute(
-                self._query_ids_and_embedding_by_embedding,
-                (query_embedding, fetch_k),
+            query, params = self._get_search_cql(
+                limit=fetch_k,
+                columns="content_id, text_embedding, link_to_tags",
+                metadata=metadata,
+                embedding=query_embedding
             )
+
+            fetched = self._session.execute(query=query, parameters=params)
             candidates = {}
             for row in fetched:
                 candidates[row.content_id] = row.text_embedding
@@ -540,7 +515,7 @@ def fetch_initial_candidates() -> None:
         return self._nodes_with_ids(helper.selected_ids)
 
     def traversal_search(
-        self, query: str, *, k: int = 4, depth: int = 1
+        self, query: str, *, k: int = 4, depth: int = 1, metadata: Optional[Dict[str, Any]] = [],
     ) -> Iterable[Node]:
         """Retrieve documents from this knowledge store.
 
@@ -553,6 +528,7 @@ def traversal_search(
             k: The number of Documents to return from the initial vector search.
                 Defaults to 4.
             depth: The maximum depth of edges to traverse. Defaults to 1.
+            metadata: Optional metadata to filter the results.
 
         Returns:
             Collection of retrieved documents.
@@ -638,9 +614,15 @@ def visit_targets(d: int, targets: Sequence[Any]) -> None:
                         )
 
             query_embedding = self._embedding.embed_query(query)
+            query, params = self._get_search_cql(
+                limit=k,
+                metadata=metadata,
+                embedding=query_embedding,
+            )
+
             cq.execute(
-                self._query_ids_and_link_to_tags_by_embedding,
-                parameters=(query_embedding, k),
+                query,
+                parameters=params,
                 callback=lambda nodes: visit_nodes(0, nodes),
             )
 
@@ -650,17 +632,18 @@ def similarity_search(
         self,
         embedding: List[float],
         k: int = 4,
+        metadata: Optional[Dict[str, Any]] = [],
     ) -> Iterable[Node]:
-        """Retrieve nodes similar to the given embedding."""
-        for row in self._session.execute(self._query_by_embedding, (embedding, k)):
+        """Retrieve nodes similar to the given embedding, optionally filtered by metadata"""
+        query, params = self._get_search_cql(embedding=embedding, limit=k, metadata=metadata)
+
+        for row in self._session.execute(query, params):
             yield _row_to_node(row)
 
     def metadata_search(self, metadata: Dict[str, Any] = {}, n: Optional[int] = 5)-> Iterable[Node]:
-        query, params = self._get_metadata_search_cql(metadata=metadata, n=n)
-
-        prepared_query = self._session.prepare(query)
+        query, params = self._get_search_cql(metadata=metadata, limit=n)
 
-        for row in self._session.execute(prepared_query, params):
+        for row in self._session.execute(query, params):
             yield _row_to_node(row)
 
     def get_node(self, id: str) -> Node:
@@ -802,7 +785,7 @@ def _extract_where_clause_blocks(
         self, metadata: Dict[str, Any]
     ) -> Tuple[str, List[Any]]:
 
-        attributes_blob, metadata_s = self._parse_metadata(metadata=metadata, is_query=True)
+        _, metadata_s = self._parse_metadata(metadata=metadata, is_query=True)
 
         if len(metadata_s) == 0:
             return "", []
@@ -818,17 +801,27 @@ def _extract_where_clause_blocks(
         return where_clause, vals_list
 
 
-    def _get_metadata_search_cql(self, n: int, metadata: Dict[str, Any]) -> Tuple[str, Tuple[Any, ...]]:
+    def _get_search_cql(self, limit: int, columns: Optional[str] = CONTENT_COLUMNS, metadata: Optional[Dict[str, Any]] = {}, embedding: Optional[List[float]] = None) -> Tuple[str, Tuple[Any, ...]]:
         where_clause, get_cql_vals = self._extract_where_clause_blocks(metadata=metadata)
         limit_clause = "LIMIT ?"
-        limit_cql_vals = [n]
-        select_vals = tuple(list(get_cql_vals) + limit_cql_vals)
-        #
+        limit_cql_vals = [limit]
+
+        order_clause=""
+        order_cql_vals = []
+        if embedding is not None:
+            order_clause = "ORDER BY text_embedding ANN OF ?"
+            order_cql_vals = [embedding]
+
+        select_vals = tuple(list(get_cql_vals) + order_cql_vals + limit_cql_vals)
         select_cql = SELECT_CQL_TEMPLATE.format(
-            columns=CONTENT_COLUMNS,
-            table_name=f"{self._keyspace}.{self._node_table}",
+            columns=columns,
+            table_name=self.table_name(),
             where_clause=where_clause,
+            order_clause=order_clause,
             limit_clause=limit_clause,
 
         )
-        return select_cql, select_vals
+        prepared_query = self._session.prepare(select_cql)
+        prepared_query.consistency_level = ConsistencyLevel.ONE
+
+        return prepared_query, select_vals