restored full metadata_blob

epinzur · epinzur · commit 5b0bca4df88b · 2024-07-23T10:19:49.000-04:00
diff --git a/libs/knowledge-store/ragstack_knowledge_store/graph_store.py b/libs/knowledge-store/ragstack_knowledge_store/graph_store.py
@@ -29,9 +29,7 @@
 
 CONTENT_ID = "content_id"
 
-CONTENT_COLUMNS = (
-    "content_id, kind, text_content, attributes_blob, metadata_s, links_blob"
-)
+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};"
@@ -119,18 +117,12 @@ def _deserialize_links(json_blob: Optional[str]) -> Set[Link]:
 
 
 def _row_to_node(row: Any) -> Node:
-    metadata_s = row.metadata_s
-    if metadata_s is None:
-        metadata_s = {}
-    attributes_blob = row.attributes_blob
-    attributes_dict = (
-        _deserialize_metadata(attributes_blob) if attributes_blob is not None else {}
-    )
+    metadata = _deserialize_metadata(row.metadata_blob)
     links = _deserialize_links(row.links_blob)
     return Node(
         id=row.content_id,
         text=row.text_content,
-        metadata={**attributes_dict, **metadata_s},
+        metadata=metadata,
         links=links,
     )
 
@@ -198,7 +190,7 @@ def __init__(
             f"""
             INSERT INTO {keyspace}.{node_table} (
                 content_id, kind, text_content, text_embedding, link_to_tags,
-                link_from_tags, attributes_blob, metadata_s, links_blob
+                link_from_tags, links_blob, metadata_blob, metadata_s
             ) VALUES (?, '{Kind.passage}', ?, ?, ?, ?, ?, ?, ?)
             """  # noqa: S608
         )
@@ -265,9 +257,9 @@ def _apply_schema(self) -> None:
 
                 link_to_tags SET<TUPLE<TEXT, TEXT>>,
                 link_from_tags SET<TUPLE<TEXT, TEXT>>,
-                attributes_blob TEXT,
-                metadata_s MAP<TEXT,TEXT>,
                 links_blob TEXT,
+                metadata_blob TEXT,
+                metadata_s MAP<TEXT,TEXT>,
 
                 PRIMARY KEY (content_id)
             )
@@ -287,36 +279,14 @@ def _apply_schema(self) -> None:
         """)
 
         self._session.execute(f"""
-            CREATE CUSTOM INDEX IF NOT EXISTS {self._node_table}_metadata_index
+            CREATE CUSTOM INDEX IF NOT EXISTS {self._node_table}_metadata_s_index
             ON {self.table_name()}(ENTRIES(metadata_s))
             USING 'StorageAttachedIndex';
         """)
 
     def _concurrent_queries(self) -> ConcurrentQueries:
         return ConcurrentQueries(self._session)
 
-    def _parse_metadata(
-        self, metadata: Dict[str, Any], is_query: bool
-    ) -> Tuple[str, Dict[str, str]]:
-        attributes_dict = {
-            k: self._coerce_string(v)
-            for k, v in metadata.items()
-            if not _is_metadata_field_indexed(k, self._metadata_indexing_policy)
-        }
-        if is_query and len(attributes_dict) > 0:
-            raise ValueError("Non-indexed metadata fields cannot be used in queries.")
-        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()
-        }
-        return (attributes_blob, metadata_s)
-
     # TODO: Async (aadd_nodes)
     def add_nodes(
         self,
@@ -352,10 +322,13 @@ def add_nodes(
                     if tag.direction in {"out", "bidir"}:
                         link_to_tags.add((tag.kind, tag.tag))
 
-                attributes_blob, metadata_s = self._parse_metadata(
-                    metadata=metadata, is_query=False
-                )
+                metadata_s = {
+                    k: self._coerce_string(v)
+                    for k, v in metadata.items()
+                    if _is_metadata_field_indexed(k, self._metadata_indexing_policy)
+                }
 
+                metadata_blob = _serialize_metadata(metadata)
                 links_blob = _serialize_links(links)
                 cq.execute(
                     self._insert_passage,
@@ -365,9 +338,9 @@ def add_nodes(
                         text_embedding,
                         link_to_tags,
                         link_from_tags,
-                        attributes_blob,
-                        metadata_s,
                         links_blob,
+                        metadata_blob,
+                        metadata_s,
                     ),
                 )
 
@@ -413,7 +386,7 @@ def mmr_traversal_search(
         adjacent_k: int = 10,
         lambda_mult: float = 0.5,
         score_threshold: float = float("-inf"),
-        metadata: Dict[str, Any] = {},
+        metadata_filter: Dict[str, Any] = {},
     ) -> Iterable[Node]:
         """Retrieve documents from this graph store using MMR-traversal.
 
@@ -439,7 +412,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.
+            metadata_filter: Optional metadata to filter the results.
         """
         query_embedding = self._embedding.embed_query(query)
         helper = MmrHelper(
@@ -458,7 +431,7 @@ def fetch_initial_candidates() -> None:
             query, params = self._get_search_cql(
                 limit=fetch_k,
                 columns="content_id, text_embedding, link_to_tags",
-                metadata=metadata,
+                metadata=metadata_filter,
                 embedding=query_embedding,
             )
 
@@ -539,7 +512,7 @@ def traversal_search(
         *,
         k: int = 4,
         depth: int = 1,
-        metadata: Dict[str, Any] = {},
+        metadata_filter: Dict[str, Any] = {},
     ) -> Iterable[Node]:
         """Retrieve documents from this knowledge store.
 
@@ -552,7 +525,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.
+            metadata_filter: Optional metadata to filter the results.
 
         Returns:
             Collection of retrieved documents.
@@ -639,8 +612,9 @@ 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",
                 limit=k,
-                metadata=metadata,
+                metadata=metadata_filter,
                 embedding=query_embedding,
             )
 
@@ -656,11 +630,11 @@ def similarity_search(
         self,
         embedding: List[float],
         k: int = 4,
-        metadata: Dict[str, Any] = {},
+        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(
-            embedding=embedding, limit=k, metadata=metadata
+            embedding=embedding, limit=k, metadata=metadata_filter
         )
 
         for row in self._session.execute(query, params):
@@ -813,17 +787,20 @@ def _coerce_string(value: Any) -> str:
     def _extract_where_clause_blocks(
         self, metadata: Dict[str, Any]
     ) -> Tuple[str, List[Any]]:
-        _, metadata_s = self._parse_metadata(metadata=metadata, is_query=True)
-
-        if len(metadata_s) == 0:
-            return "", []
-
         wc_blocks: List[str] = []
         vals_list: List[Any] = []
 
-        for k, v in sorted(metadata_s.items()):
-            wc_blocks.append(f"metadata_s['{k}'] = ?")
-            vals_list.append(v)
+        for key, value in sorted(metadata.items()):
+            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 "", []
 
         where_clause = "WHERE " + " AND ".join(wc_blocks)
         return where_clause, vals_list
diff --git a/libs/knowledge-store/tests/integration_tests/test_graph_store.py b/libs/knowledge-store/tests/integration_tests/test_graph_store.py
@@ -1,7 +1,7 @@
 # ruff: noqa: PT011, RUF015
 import math
 import secrets
-from typing import Callable, Iterable, Iterator, List, Optional
+from typing import Callable, Iterable, Iterator, List
 
 import numpy as np
 import pytest
@@ -272,8 +272,8 @@ def test_metadata(graph_store_factory: Callable[[], GraphStore]) -> None:
         Link(direction="in", kind="hyperlink", tag="http://a"),
         Link(direction="bidir", kind="other", tag="foo"),
     }
-def test_graph_store_metadata(graph_store_factory: Callable[[str], GraphStore]) -> None:
-    gs = graph_store_factory()
+
+
 def test_graph_store_metadata(
     graph_store_factory: Callable[[MetadataIndexingType], GraphStore],
 ) -> None:
@@ -288,16 +288,14 @@ def test_graph_store_metadata(
     assert gotten2 == Node(text="", id="row2", metadata={})
 
     md3 = {"a": 1, "b": "Bee", "c": True}
-    md3_string = {"a": "1.0", "b": "Bee", "c": "true"}
     gs.add_nodes([Node(text="", id="row3", metadata=md3)])
     gotten3 = gs.get_node(content_id="row3")
-    assert gotten3 == Node(text="", id="row3", metadata=md3_string)
+    assert gotten3 == Node(text="", id="row3", metadata=md3)
 
     md4 = {"c1": True, "c2": True, "c3": True}
-    md4_string = {"c1": "true", "c2": "true", "c3": "true"}
     gs.add_nodes([Node(text="", id="row4", metadata=md4)])
     gotten4 = gs.get_node(content_id="row4")
-    assert gotten4 == Node(text="", id="row4", metadata=md4_string)
+    assert gotten4 == Node(text="", id="row4", metadata=md4)
 
     # metadata searches:
     md_gotten3a = list(gs.metadata_search(metadata={"a": 1}))[0]
@@ -319,8 +317,8 @@ def test_graph_store_metadata(
         key=lambda res: int(float(res.metadata["index"])),
     )
     expected = [
-        Node(text="", id="twin_a", metadata={"twin": "true", "index": "0.0"}),
-        Node(text="", id="twin_b", metadata={"twin": "true", "index": "1.0"}),
+        Node(text="", id="twin_a", metadata={"twin": True, "index": 0}),
+        Node(text="", id="twin_b", metadata={"twin": True, "index": 1}),
     ]
     assert md_twins_gotten == expected
     assert list(gs.metadata_search(metadata={"fake": True})) == []
@@ -330,20 +328,19 @@ def test_graph_store_metadata_routing(
     graph_store_factory: Callable[[MetadataIndexingType], GraphStore],
 ) -> None:
     test_md = {"mds": "string", "mdn": 255, "mdb": True}
-    test_md_string = {"mds": "string", "mdn": "255.0", "mdb": "true"}
 
     gs_all = graph_store_factory("all")
     gs_all.add_nodes([Node(id="row1", text="bb1", metadata=test_md)])
     gotten_all = list(gs_all.metadata_search(metadata={"mds": "string"}))[0]
-    assert gotten_all.metadata == test_md_string
+    assert gotten_all.metadata == test_md
     gs_none = graph_store_factory("none")
     gs_none.add_nodes([Node(id="row1", text="bb1", metadata=test_md)])
     with pytest.raises(ValueError):
         # querying on non-indexed metadata fields:
         list(gs_none.metadata_search(metadata={"mds": "string"}))
     gotten_none = gs_none.get_node(content_id="row1")
     assert gotten_none is not None
-    assert gotten_none.metadata == test_md_string
+    assert gotten_none.metadata == test_md
     test_md_allowdeny = {
         "mdas": "MDAS",
         "mdds": "MDDS",
@@ -352,23 +349,15 @@ def test_graph_store_metadata_routing(
         "mdab": True,
         "mddb": True,
     }
-    test_md_allowdeny_string = {
-        "mdas": "MDAS",
-        "mdds": "MDDS",
-        "mdan": "255.0",
-        "mddn": "127.0",
-        "mdab": "true",
-        "mddb": "true",
-    }
     gs_allow = graph_store_factory(("allow", {"mdas", "mdan", "mdab"}))
     gs_allow.add_nodes([Node(id="row1", text="bb1", metadata=test_md_allowdeny)])
     with pytest.raises(ValueError):
         list(gs_allow.metadata_search(metadata={"mdds": "MDDS"}))
     gotten_allow = list(gs_allow.metadata_search(metadata={"mdas": "MDAS"}))[0]
-    assert gotten_allow.metadata == test_md_allowdeny_string
+    assert gotten_allow.metadata == test_md_allowdeny
     gs_deny = graph_store_factory(("deny", {"mdds", "mddn", "mddb"}))
     gs_deny.add_nodes([Node(id="row1", text="bb1", metadata=test_md_allowdeny)])
     with pytest.raises(ValueError):
         list(gs_deny.metadata_search(metadata={"mdds": "MDDS"}))
     gotten_deny = list(gs_deny.metadata_search(metadata={"mdas": "MDAS"}))[0]
-    assert gotten_deny.metadata == test_md_allowdeny_string
+    assert gotten_deny.metadata == test_md_allowdeny
