feat: Eliminate the targets table (#588)

* feat: Eliminate the targets table

This seems to be unnecessary. An SAI index supporting contains seems to
be sufficient. Eliminating it removes the need to create the
denormalized copy of the data, speeding up insertion and reducing
storage needs.

Author: bjchambers

libs/knowledge-store/notebooks/astra_support.ipynb

@@ -258,7 +258,6 @@
     "graph_store = CassandraGraphStore(\n",
     "    embeddings,\n",
     "    node_table=f\"{SITE_PREFIX}_nodes\",\n",
-    "    targets_table=f\"{SITE_PREFIX}_targets\",\n",
     ")"
    ]
   },

libs/knowledge-store/ragstack_knowledge_store/graph_store.py

@@ -127,7 +127,6 @@ def __init__(
         embedding: EmbeddingModel,
         *,
         node_table: str = "graph_nodes",
-        targets_table: str = "graph_targets",
         session: Optional[Session] = None,
         keyspace: Optional[str] = None,
         setup_mode: SetupMode = SetupMode.SYNC,
@@ -141,12 +140,8 @@ def __init__(
         if not _CQL_IDENTIFIER_PATTERN.fullmatch(node_table):
             raise ValueError(f"Invalid node table name: {node_table}")
 
-        if not _CQL_IDENTIFIER_PATTERN.fullmatch(targets_table):
-            raise ValueError(f"Invalid node table name: {targets_table}")
-
         self._embedding = embedding
         self._node_table = node_table
-        self._targets_table = targets_table
         self._session = session
         self._keyspace = keyspace
 
@@ -163,16 +158,8 @@ def __init__(
             f"""
             INSERT INTO {keyspace}.{node_table} (
                 content_id, kind, text_content, text_embedding, link_to_tags,
-                metadata_blob, links_blob
-            ) VALUES (?, '{Kind.passage}', ?, ?, ?, ?, ?)
-            """  # noqa: S608
-        )
-
-        self._insert_tag = session.prepare(
-            f"""
-            INSERT INTO {keyspace}.{targets_table} (
-                target_content_id, kind, tag, target_text_embedding, target_link_to_tags
-            ) VALUES (?, ?, ?, ?, ?)
+                link_from_tags, metadata_blob, links_blob
+            ) VALUES (?, '{Kind.passage}', ?, ?, ?, ?, ?, ?)
             """  # noqa: S608
         )
 
@@ -236,56 +223,44 @@ def __init__(
 
         self._query_targets_embeddings_by_kind_and_tag_and_embedding = session.prepare(
             f"""
-            SELECT target_content_id, target_text_embedding, tag, target_link_to_tags
-            FROM {keyspace}.{targets_table}
-            WHERE kind = ? AND tag = ?
-            ORDER BY target_text_embedding ANN of ?
+            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 ?
-            """  # noqa: S608
+            """
         )
 
         self._query_targets_by_kind_and_value = session.prepare(
             f"""
-            SELECT target_content_id, kind, tag
-            FROM {keyspace}.{targets_table}
-            WHERE kind = ? AND tag = ?
-            """  # noqa: S608
+            SELECT
+                content_id AS target_content_id
+            FROM {keyspace}.{node_table}
+            WHERE link_from_tags CONTAINS (?, ?)
+            """
         )
 
     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} (
+        self._session.execute(f"""
+            CREATE TABLE IF NOT EXISTS {self._keyspace}.{self._node_table} (
                 content_id TEXT,
                 kind TEXT,
                 text_content TEXT,
                 text_embedding VECTOR<FLOAT, {embedding_dim}>,
 
                 link_to_tags SET<TUPLE<TEXT, TEXT>>,
+                link_from_tags SET<TUPLE<TEXT, TEXT>>,
                 metadata_blob TEXT,
                 links_blob TEXT,
 
                 PRIMARY KEY (content_id)
             )
-            """
-        )
-
-        self._session.execute(
-            f"""CREATE TABLE IF NOT EXISTS {self._keyspace}.{self._targets_table} (
-                target_content_id TEXT,
-                kind TEXT,
-                tag TEXT,
-
-                -- text_embedding of target node. allows MMR to be applied without
-                -- fetching nodes.
-                target_text_embedding VECTOR<FLOAT, {embedding_dim}>,
-                target_link_to_tags SET<TUPLE<TEXT, TEXT>>,
-
-                PRIMARY KEY ((kind, tag), target_content_id)
-            )
-            """
-        )
+        """)
 
         # Index on text_embedding (for similarity search)
         self._session.execute(f"""
@@ -294,12 +269,11 @@ def _apply_schema(self) -> None:
             USING 'StorageAttachedIndex';
         """)
 
-        # Index on target_text_embedding (for similarity search)
         self._session.execute(f"""
-            CREATE CUSTOM INDEX IF NOT EXISTS {self._targets_table}_target_text_embedding_index
-            ON {self._keyspace}.{self._targets_table}(target_text_embedding)
+            CREATE CUSTOM INDEX IF NOT EXISTS {self._node_table}_link_from_tags
+            ON {self._keyspace}.{self._node_table}(link_from_tags)
             USING 'StorageAttachedIndex';
-        """)  # noqa: E501
+        """)
 
     def _concurrent_queries(self) -> ConcurrentQueries:
         return ConcurrentQueries(self._session)
@@ -348,17 +322,12 @@ def add_nodes(
                         text,
                         text_embedding,
                         link_to_tags,
+                        link_from_tags,
                         metadata_blob,
                         links_blob,
                     ),
                 )
 
-                for kind, value in link_from_tags:
-                    cq.execute(
-                        self._insert_tag,
-                        parameters=(node_id, kind, value, text_embedding, link_to_tags),
-                    )
-
         return node_ids
 
     def _nodes_with_ids(

libs/knowledge-store/tests/integration_tests/test_graph_store.py

@@ -51,13 +51,11 @@ def _make_graph_store() -> GraphStore:
         name = secrets.token_hex(8)
 
         node_table = f"nodes_{name}"
-        targets_table = f"targets_{name}"
         return GraphStore(
             embedding,
             session=session,
             keyspace=KEYSPACE,
             node_table=node_table,
-            targets_table=targets_table,
         )
 
     yield _make_graph_store

libs/langchain/ragstack_langchain/graph_store/cassandra.py

@@ -39,7 +39,6 @@ def __init__(
         embedding: Embeddings,
         *,
         node_table: str = "graph_nodes",
-        targets_table: str = "graph_targets",
         session: Optional[Session] = None,
         keyspace: Optional[str] = None,
         setup_mode: SetupMode = SetupMode.SYNC,
@@ -63,7 +62,6 @@ def __init__(
         self.store = graph_store.GraphStore(
             embedding=_EmbeddingModelAdapter(embedding),
             node_table=node_table,
-            targets_table=targets_table,
             session=session,
             keyspace=keyspace,
             setup_mode=_setup_mode,

libs/langchain/tests/integration_tests/test_graph_store.py

@@ -20,7 +20,6 @@ def __init__(self, session: Session, keyspace: str, embedding: Embeddings) -> No
         self.keyspace = keyspace
         self.uid = secrets.token_hex(8)
         self.node_table = f"nodes_{self.uid}"
-        self.targets_table = f"targets_{self.uid}"
         self.embedding = embedding
         self._store = None
 
@@ -39,17 +38,13 @@ def store(
                 session=self.session,
                 keyspace=self.keyspace,
                 node_table=self.node_table,
-                targets_table=self.targets_table,
                 ids=ids,
             )
 
         return self._store
 
     def drop(self):
         self.session.execute(f"DROP TABLE IF EXISTS {self.keyspace}.{self.node_table};")
-        self.session.execute(
-            f"DROP TABLE IF EXISTS {self.keyspace}.{self.targets_table};"
-        )
 
 
 @pytest.fixture(scope="session")