added metadata index put and get support

Author: epinzur

libs/knowledge-store/ragstack_knowledge_store/graph_store.py

@@ -12,10 +12,11 @@
     Sequence,
     Set,
     Tuple,
+    Union,
     cast,
 )
 
-from cassandra.cluster import ConsistencyLevel, Session
+from cassandra.cluster import ConsistencyLevel, Session, ResponseFuture
 from cassio.config import check_resolve_keyspace, check_resolve_session
 
 from ._mmr_helper import MmrHelper
@@ -48,6 +49,20 @@ class SetupMode(Enum):
     ASYNC = 2
     OFF = 3
 
+class MetadataIndexingMode(Enum):
+    DEFAULT_TO_UNSEARCHABLE = 1
+    DEFAULT_TO_SEARCHABLE = 2
+
+MetadataIndexingPolicy = Tuple[MetadataIndexingMode, Set[str]]
+
+def _is_metadata_field_indexed(field_name: str, policy: MetadataIndexingPolicy) -> bool:
+    p_mode, p_fields = policy
+    if p_mode == MetadataIndexingMode.DEFAULT_TO_UNSEARCHABLE:
+        return field_name in p_fields
+    elif p_mode == MetadataIndexingMode.DEFAULT_TO_SEARCHABLE:
+        return field_name not in p_fields
+    else:
+        raise ValueError(f"Unexpected metadata indexing mode {p_mode}")
 
 def _serialize_metadata(md: Dict[str, Any]) -> str:
     if isinstance(md.get("links"), Set):
@@ -90,12 +105,14 @@ def _deserialize_links(json_blob: Optional[str]) -> Set[Link]:
 
 
 def _row_to_node(row: Any) -> Node:
-    metadata = _deserialize_metadata(row.metadata_blob)
+    metadata_s = row.get("metadata_s", {})
+    attributes_blob = row.get("attributes_blob")
+    attributes_dict = _deserialize_metadata(attributes_blob) if attributes_blob is not None else {}
     links = _deserialize_links(row.links_blob)
     return Node(
         id=row.content_id,
         text=row.text_content,
-        metadata=metadata,
+        metadata={**attributes_dict, **metadata_s},
         links=links,
     )
 
@@ -130,6 +147,7 @@ def __init__(
         session: Optional[Session] = None,
         keyspace: Optional[str] = None,
         setup_mode: SetupMode = SetupMode.SYNC,
+        metadata_indexing: Union[Tuple[str, Iterable[str]], str] = "all",
     ):
         session = check_resolve_session(session)
         keyspace = check_resolve_keyspace(keyspace)
@@ -145,6 +163,10 @@ def __init__(
         self._session = session
         self._keyspace = keyspace
 
+        self._metadata_indexing_policy = self._normalize_metadata_indexing_policy(
+            metadata_indexing
+        )
+
         if setup_mode == SetupMode.SYNC:
             self._apply_schema()
         elif setup_mode != SetupMode.OFF:
@@ -158,22 +180,22 @@ def __init__(
             f"""
             INSERT INTO {keyspace}.{node_table} (
                 content_id, kind, text_content, text_embedding, link_to_tags,
-                link_from_tags, metadata_blob, links_blob
-            ) VALUES (?, '{Kind.passage}', ?, ?, ?, ?, ?, ?)
+                link_from_tags, attributes_blob, metadata_s, links_blob
+            ) VALUES (?, '{Kind.passage}', ?, ?, ?, ?, ?, ?, ?)
             """  # noqa: S608
         )
 
         self._query_by_id = session.prepare(
             f"""
-            SELECT content_id, kind, text_content, metadata_blob, links_blob
+            SELECT content_id, kind, text_content, attributes_blob, links_blob
             FROM {keyspace}.{node_table}
             WHERE content_id = ?
             """  # noqa: S608
         )
 
         self._query_by_embedding = session.prepare(
             f"""
-            SELECT content_id, kind, text_content, metadata_blob, links_blob
+            SELECT content_id, kind, text_content, attributes_blob, links_blob
             FROM {keyspace}.{node_table}
             ORDER BY text_embedding ANN OF ?
             LIMIT ?
@@ -255,7 +277,8 @@ def _apply_schema(self) -> None:
 
                 link_to_tags SET<TUPLE<TEXT, TEXT>>,
                 link_from_tags SET<TUPLE<TEXT, TEXT>>,
-                metadata_blob TEXT,
+                attributes_blob TEXT,
+                metadata_s MAP<TEXT,TEXT>,
                 links_blob TEXT,
 
                 PRIMARY KEY (content_id)
@@ -275,6 +298,12 @@ def _apply_schema(self) -> None:
             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))
+            USING 'StorageAttachedIndex';
+        """)
+
     def _concurrent_queries(self) -> ConcurrentQueries:
         return ConcurrentQueries(self._session)
 
@@ -313,7 +342,20 @@ def add_nodes(
                     if tag.direction in {"out", "bidir"}:
                         link_to_tags.add((tag.kind, tag.tag))
 
-                metadata_blob = _serialize_metadata(metadata)
+                attributes_dict = {
+                    k: self._coerce_string(v)
+                    for k, v in metadata.items()
+                    if not _is_metadata_field_indexed(k, self._metadata_indexing_policy)
+                }
+                attributes_blob = _serialize_metadata(attributes_dict)
+
+                metadata_indexed_dict = {
+                    k: v
+                    for k, v in metadata.items()
+                    if _is_metadata_field_indexed(k, self._metadata_indexing_policy)
+                }
+                metadata_s = {k: self._coerce_string(v) for k, v in metadata_indexed_dict.items()}
+
                 links_blob = _serialize_links(links)
                 cq.execute(
                     self._insert_passage,
@@ -323,7 +365,8 @@ def add_nodes(
                         text_embedding,
                         link_to_tags,
                         link_from_tags,
-                        metadata_blob,
+                        attributes_blob,
+                        metadata_s,
                         links_blob,
                     ),
                 )
@@ -670,3 +713,85 @@ def add_targets(rows: Iterable[Any]) -> None:
         # TODO: Consider a combined limit based on the similarity and/or
         # predicated MMR score?
         return targets.values()
+
+    @staticmethod
+    def _normalize_metadata_indexing_policy(
+        metadata_indexing: Union[Tuple[str, Iterable[str]], str]
+    ) -> MetadataIndexingPolicy:
+        mode: MetadataIndexingMode
+        fields: Set[str]
+        # metadata indexing policy normalization:
+        if isinstance(metadata_indexing, str):
+            if metadata_indexing.lower() == "all":
+                mode, fields = (MetadataIndexingMode.DEFAULT_TO_SEARCHABLE, set())
+            elif metadata_indexing.lower() == "none":
+                mode, fields = (MetadataIndexingMode.DEFAULT_TO_UNSEARCHABLE, set())
+            else:
+                raise ValueError(
+                    f"Unsupported metadata_indexing value '{metadata_indexing}'"
+                )
+        else:
+            assert len(metadata_indexing) == 2
+            # it's a 2-tuple (mode, fields) still to normalize
+            _mode, _field_spec = metadata_indexing
+            fields = {_field_spec} if isinstance(_field_spec, str) else set(_field_spec)
+            if _mode.lower() in {
+                "default_to_unsearchable",
+                "allowlist",
+                "allow",
+                "allow_list",
+            }:
+                mode = MetadataIndexingMode.DEFAULT_TO_UNSEARCHABLE
+            elif _mode.lower() in {
+                "default_to_searchable",
+                "denylist",
+                "deny",
+                "deny_list",
+            }:
+                mode = MetadataIndexingMode.DEFAULT_TO_SEARCHABLE
+            else:
+                raise ValueError(
+                    f"Unsupported metadata indexing mode specification '{_mode}'"
+                )
+        return (mode, fields)
+
+    def _split_metadata_fields(self, md_dict: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Split the *indexed* part of the metadata in separate parts,
+        one per Cassandra column.
+
+        Currently: everything gets cast to a string and goes to a single table
+        column. This means:
+            - strings are fine
+            - floats and integers v: they are cast to str(v)
+            - booleans: 'true'/'false' (JSON style)
+            - None => 'null' (JSON style)
+            - anything else v => str(v), no questions asked
+
+        Caveat: one gets strings back when reading metadata
+        """
+
+        # TODO: more care about types here
+        stringy_part = {k: self._coerce_string(v) for k, v in md_dict.items()}
+        return {
+            "metadata_s": stringy_part,
+        }
+
+    @staticmethod
+    def _coerce_string(value: Any) -> str:
+        if isinstance(value, str):
+            return value
+        elif isinstance(value, bool):
+            # bool MUST come before int in this chain of ifs!
+            return json.dumps(value)
+        elif isinstance(value, int):
+            # we don't want to store '1' and '1.0' differently
+            # for the sake of metadata-filtered retrieval:
+            return json.dumps(float(value))
+        elif isinstance(value, float):
+            return json.dumps(value)
+        elif value is None:
+            return json.dumps(value)
+        else:
+            # when all else fails ...
+            return str(value)