added metadata filter support to graph methods

Author: epinzur

libs/knowledge-store/ragstack_knowledge_store/graph_store.py

@@ -18,7 +18,7 @@
     cast,
 )
 
-from cassandra.cluster import ConsistencyLevel, Session
+from cassandra.cluster import ConsistencyLevel, PreparedStatement, Session
 from cassio.config import check_resolve_keyspace, check_resolve_session
 
 from ._mmr_helper import MmrHelper
@@ -32,7 +32,7 @@
 CONTENT_COLUMNS = "content_id, kind, text_content, links_blob, metadata_blob"
 
 SELECT_CQL_TEMPLATE = (
-    "SELECT {columns} FROM {table_name} {where_clause} {order_clause} {limit_clause};"
+    "SELECT {columns} FROM {table_name}{where_clause}{order_clause}{limit_clause};"
 )
 
 
@@ -172,6 +172,7 @@ def __init__(
         self._node_table = node_table
         self._session = session
         self._keyspace = keyspace
+        self._prepared_query_cache: Dict[str, PreparedStatement] = {}
 
         self._metadata_indexing_policy = self._normalize_metadata_indexing_policy(
             metadata_indexing=metadata_indexing,
@@ -219,28 +220,6 @@ def __init__(
             """  # noqa: S608
         )
 
-        self._query_targets_embeddings_by_kind_and_tag_and_embedding = session.prepare(
-            f"""
-            SELECT
-                content_id AS target_content_id,
-                text_embedding AS target_text_embedding,
-                link_to_tags AS target_link_to_tags
-            FROM {keyspace}.{node_table}
-            WHERE link_from_tags CONTAINS (?, ?)
-            ORDER BY text_embedding ANN of ?
-            LIMIT ?
-            """
-        )
-
-        self._query_targets_by_kind_and_value = session.prepare(
-            f"""
-            SELECT
-                content_id AS target_content_id
-            FROM {keyspace}.{node_table}
-            WHERE link_from_tags CONTAINS (?, ?)
-            """
-        )
-
     def table_name(self) -> str:
         """Returns the fully qualified table name."""
         return f"{self._keyspace}.{self._node_table}"
@@ -427,15 +406,23 @@ def mmr_traversal_search(
 
         # Fetch the initial candidates and add them to the helper and
         # outgoing_tags.
+        initial_candidates_query = self._get_search_cql(
+            has_limit=True,
+            columns="content_id, text_embedding, link_to_tags",
+            metadata_keys=list(metadata_filter.keys()),
+            has_embedding=True,
+        )
+
         def fetch_initial_candidates() -> None:
-            query, params = self._get_search_cql(
+            params = self._get_search_params(
                 limit=fetch_k,
-                columns="content_id, text_embedding, link_to_tags",
                 metadata=metadata_filter,
                 embedding=query_embedding,
             )
 
-            fetched = self._session.execute(query=query, parameters=params)
+            fetched = self._session.execute(
+                query=initial_candidates_query, parameters=params
+            )
             candidates = {}
             for row in fetched:
                 candidates[row.content_id] = row.text_embedding
@@ -474,6 +461,7 @@ def fetch_initial_candidates() -> None:
                     link_to_tags,
                     query_embedding=query_embedding,
                     k_per_tag=adjacent_k,
+                    metadata_filter=metadata_filter,
                 )
 
                 # Record the link_to_tags as visited.
@@ -541,6 +529,19 @@ def traversal_search(
         #
         # ...
 
+        traversal_query = self._get_search_cql(
+            columns="content_id, link_to_tags",
+            has_limit=True,
+            metadata_keys=list(metadata_filter.keys()),
+            has_embedding=True,
+        )
+
+        visit_nodes_query = self._get_search_cql(
+            columns="content_id AS target_content_id",
+            has_link_from_tags=True,
+            metadata_keys=list(metadata_filter.keys()),
+        )
+
         with self._concurrent_queries() as cq:
             # Map from visited ID to depth
             visited_ids: Dict[str, int] = {}
@@ -583,12 +584,12 @@ def visit_nodes(d: int, nodes: Sequence[Any]) -> None:
                     # If there are new tags to visit at the next depth, query for the
                     # node IDs.
                     for kind, value in outgoing_tags:
+                        params = self._get_search_params(
+                            link_from_tags=(kind, value), metadata=metadata_filter
+                        )
                         cq.execute(
-                            self._query_targets_by_kind_and_value,
-                            parameters=(
-                                kind,
-                                value,
-                            ),
+                            query=visit_nodes_query,
+                            parameters=params,
                             callback=lambda rows, d=d: visit_targets(d, rows),
                         )
 
@@ -611,15 +612,14 @@ def visit_targets(d: int, targets: Sequence[Any]) -> None:
                         )
 
             query_embedding = self._embedding.embed_query(query)
-            query, params = self._get_search_cql(
-                columns="content_id, link_to_tags",
+            params = self._get_search_params(
                 limit=k,
                 metadata=metadata_filter,
                 embedding=query_embedding,
             )
 
             cq.execute(
-                query,
+                traversal_query,
                 parameters=params,
                 callback=lambda nodes: visit_nodes(0, nodes),
             )
@@ -633,7 +633,7 @@ def similarity_search(
         metadata_filter: Dict[str, Any] = {},
     ) -> Iterable[Node]:
         """Retrieve nodes similar to the given embedding, optionally filtered by metadata."""  # noqa: E501
-        query, params = self._get_search_cql(
+        query, params = self._get_search_cql_and_params(
             embedding=embedding, limit=k, metadata=metadata_filter
         )
 
@@ -644,7 +644,7 @@ def metadata_search(
         self, metadata: Dict[str, Any] = {}, n: int = 5
     ) -> Iterable[Node]:
         """Retrieve nodes based on their metadata."""
-        query, params = self._get_search_cql(metadata=metadata, limit=n)
+        query, params = self._get_search_cql_and_params(metadata=metadata, limit=n)
 
         for row in self._session.execute(query, params):
             yield _row_to_node(row)
@@ -681,19 +681,35 @@ def _get_adjacent(
         tags: Set[Tuple[str, str]],
         query_embedding: List[float],
         k_per_tag: Optional[int] = None,
+        metadata_filter: Dict[str, Any] = {},
     ) -> Iterable[_Edge]:
         """Return the target nodes with incoming links from any of the given tags.
 
         Args:
             tags: The tags to look for links *from*.
             query_embedding: The query embedding. Used to rank target nodes.
             k_per_tag: The number of target nodes to fetch for each outgoing tag.
+            metadata_filter: Optional metadata to filter the results.
 
         Returns:
             List of adjacent edges.
         """
         targets: Dict[str, _Edge] = {}
 
+        columns = """
+            content_id AS target_content_id,
+            text_embedding AS target_text_embedding,
+            link_to_tags AS target_link_to_tags
+        """
+
+        adjacent_query = self._get_search_cql(
+            has_limit=True,
+            columns=columns,
+            metadata_keys=list(metadata_filter.keys()),
+            has_embedding=True,
+            has_link_from_tags=True,
+        )
+
         def add_targets(rows: Iterable[Any]) -> None:
             # TODO: Figure out how to use the "kind" on the edge.
             # This is tricky, since we currently issue one query for anything
@@ -709,14 +725,16 @@ def add_targets(rows: Iterable[Any]) -> None:
 
         with self._concurrent_queries() as cq:
             for kind, value in tags:
+                params = self._get_search_params(
+                    limit=k_per_tag or 10,
+                    metadata=metadata_filter,
+                    embedding=query_embedding,
+                    link_from_tags=(kind, value),
+                )
+
                 cq.execute(
-                    self._query_targets_embeddings_by_kind_and_tag_and_embedding,
-                    parameters=(
-                        kind,
-                        value,
-                        query_embedding,
-                        k_per_tag or 10,
-                    ),
+                    query=adjacent_query,
+                    parameters=params,
                     callback=add_targets,
                 )
 
@@ -784,55 +802,116 @@ def _coerce_string(value: Any) -> str:
         # when all else fails ...
         return str(value)
 
-    def _extract_where_clause_blocks(
-        self, metadata: Dict[str, Any]
-    ) -> Tuple[str, List[Any]]:
+    def _extract_where_clause_cql(
+        self,
+        metadata_keys: List[str] = [],
+        has_link_from_tags: bool = False,
+    ) -> str:
         wc_blocks: List[str] = []
-        vals_list: List[Any] = []
 
-        for key, value in sorted(metadata.items()):
+        if has_link_from_tags:
+            wc_blocks.append("link_from_tags CONTAINS (?, ?)")
+
+        for key in sorted(metadata_keys):
             if _is_metadata_field_indexed(key, self._metadata_indexing_policy):
                 wc_blocks.append(f"metadata_s['{key}'] = ?")
-                vals_list.append(self._coerce_string(value=value))
             else:
                 raise ValueError(
                     "Non-indexed metadata fields cannot be used in queries."
                 )
 
         if len(wc_blocks) == 0:
-            return "", []
+            return ""
 
-        where_clause = "WHERE " + " AND ".join(wc_blocks)
-        return where_clause, vals_list
+        return " WHERE " + " AND ".join(wc_blocks)
+
+    def _extract_where_clause_params(
+        self,
+        metadata: Dict[str, Any],
+        link_from_tags: Optional[Tuple[str, str]] = None,
+    ) -> List[Any]:
+        params: List[Any] = []
+
+        if link_from_tags is not None:
+            params.append(link_from_tags[0])
+            params.append(link_from_tags[1])
+
+        for key, value in sorted(metadata.items()):
+            if _is_metadata_field_indexed(key, self._metadata_indexing_policy):
+                params.append(self._coerce_string(value=value))
+            else:
+                raise ValueError(
+                    "Non-indexed metadata fields cannot be used in queries."
+                )
+
+        return params
 
     def _get_search_cql(
         self,
-        limit: int,
+        has_limit: bool = False,
         columns: Optional[str] = CONTENT_COLUMNS,
-        metadata: Dict[str, Any] = {},
-        embedding: Optional[List[float]] = None,
-    ) -> Tuple[str, Tuple[Any, ...]]:
-        where_clause, get_cql_vals = self._extract_where_clause_blocks(
-            metadata=metadata
+        metadata_keys: List[str] = [],
+        has_embedding: bool = False,
+        has_link_from_tags: bool = False,
+    ) -> PreparedStatement:
+        where_clause = self._extract_where_clause_cql(
+            metadata_keys=metadata_keys, has_link_from_tags=has_link_from_tags
         )
-        limit_clause = "LIMIT ?"
-        limit_cql_vals = [limit]
+        limit_clause = " LIMIT ?" if has_limit else ""
+        order_clause = " ORDER BY text_embedding ANN OF ?" if has_embedding else ""
 
-        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=columns,
             table_name=self.table_name(),
             where_clause=where_clause,
             order_clause=order_clause,
             limit_clause=limit_clause,
         )
+
+        if select_cql in self._prepared_query_cache:
+            return self._prepared_query_cache[select_cql]
+
         prepared_query = self._session.prepare(select_cql)
         prepared_query.consistency_level = ConsistencyLevel.ONE
+        self._prepared_query_cache[select_cql] = prepared_query
 
-        return prepared_query, select_vals
+        return prepared_query
+
+    def _get_search_params(
+        self,
+        limit: Optional[int] = None,
+        metadata: Dict[str, Any] = {},
+        embedding: Optional[List[float]] = None,
+        link_from_tags: Optional[Tuple[str, str]] = None,
+    ) -> Tuple[PreparedStatement, Tuple[Any, ...]]:
+        where_params = self._extract_where_clause_params(
+            metadata=metadata, link_from_tags=link_from_tags
+        )
+
+        limit_params = [limit] if limit is not None else []
+        order_params = [embedding] if embedding is not None else []
+
+        return tuple(list(where_params) + order_params + limit_params)
+
+    def _get_search_cql_and_params(
+        self,
+        limit: Optional[int] = None,
+        columns: Optional[str] = CONTENT_COLUMNS,
+        metadata: Dict[str, Any] = {},
+        embedding: Optional[List[float]] = None,
+        link_from_tags: Optional[Tuple[str, str]] = None,
+    ) -> Tuple[PreparedStatement, Tuple[Any, ...]]:
+        query = self._get_search_cql(
+            has_limit=limit is not None,
+            columns=columns,
+            metadata_keys=list(metadata.keys()),
+            has_embedding=embedding is not None,
+            has_link_from_tags=link_from_tags is not None,
+        )
+        params = self._get_search_params(
+            limit=limit,
+            metadata=metadata,
+            embedding=embedding,
+            link_from_tags=link_from_tags,
+        )
+        return query, params