feat: Limit items retrieved during MMR traversal (#514)

* feat: Limit items retrieved during MMR traversal

* remove unused method / unneeded test

* remove unused query

Author: bjchambers

libs/knowledge-store/ragstack_knowledge_store/graph_store.py

@@ -260,11 +260,13 @@ def __init__(
             """
         )
 
-        self._query_targets_embeddings_by_kind_and_tag = session.prepare(
+        self._query_targets_embeddings_by_kind_and_tag_and_embedding = session.prepare(
             f"""
             SELECT target_content_id, target_text_embedding, tag
             FROM {keyspace}.{targets_table}
             WHERE kind = ? AND tag = ?
+            ORDER BY target_text_embedding ANN of ?
+            LIMIT ?
             """
         )
 
@@ -317,6 +319,14 @@ def _apply_schema(self):
             """
         )
 
+        # 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)
+            USING 'StorageAttachedIndex';
+            """
+        )
+
     def _concurrent_queries(self) -> ConcurrentQueries:
         return ConcurrentQueries(self._session)
 
@@ -393,20 +403,14 @@ def add_nodes(rows):
 
         return [results[id] for id in ids]
 
-    def _linked_ids(
-        self,
-        source_id: str,
-    ) -> Iterable[str]:
-        adjacent = self._get_adjacent([source_id])
-        return [edge.target_content_id for edge in adjacent]
-
     def mmr_traversal_search(
         self,
         query: str,
         *,
         k: int = 4,
         depth: int = 2,
         fetch_k: int = 100,
+        adjacent_k: int = 10,
         lambda_mult: float = 0.5,
         score_threshold: float = float("-inf"),
     ) -> Iterable[Node]:
@@ -423,7 +427,9 @@ def mmr_traversal_search(
         Args:
             query: The query string to search for.
             k: Number of Documents to return. Defaults to 4.
-            fetch_k: Number of Documents to fetch via similarity.
+            fetch_k: Number of initial Documents to fetch via similarity.
+                Defaults to 100.
+            adjacent_k: Number of adjacent Documents to fetch.
                 Defaults to 10.
             depth: Maximum depth of a node (number of edges) from a node
                 retrieved via similarity. Defaults to 2.
@@ -446,9 +452,9 @@ def mmr_traversal_search(
             (query_embedding, fetch_k),
         )
 
-        query_embedding = emb_to_ndarray(query_embedding)
+        query_embedding_ndarray = emb_to_ndarray(query_embedding)
         unselected = {
-            row.content_id: _Candidate(row.text_embedding, lambda_mult, query_embedding)
+            row.content_id: _Candidate(row.text_embedding, lambda_mult, query_embedding_ndarray)
             for row in fetched
         }
         best_score, next_id = max(
@@ -479,7 +485,9 @@ def mmr_traversal_search(
             # Add unselected edges if reached nodes are within `depth`:
             next_depth = next_selected.distance + 1
             if next_depth < depth:
-                adjacents = self._get_adjacent([selected_id])
+                adjacents = self._get_adjacent([selected_id],
+                                               query_embedding=query_embedding,
+                                               k_per_tag=adjacent_k)
                 for adjacent in adjacents:
                     target_id = adjacent.target_content_id
                     if target_id in selected_set:
@@ -494,7 +502,7 @@ def mmr_traversal_search(
                         continue
 
                     candidate = _Candidate(
-                        adjacent.target_text_embedding, lambda_mult, query_embedding
+                        adjacent.target_text_embedding, lambda_mult, query_embedding_ndarray
                     )
                     for selected_embedding in selected_embeddings:
                         candidate.update_for_selection(lambda_mult, selected_embedding)
@@ -621,8 +629,19 @@ def similarity_search(
     def _get_adjacent(
         self,
         source_ids: Iterable[str],
+        query_embedding: List[float],
+        k_per_tag: Optional[int] = None,
     ) -> Iterable[_Edge]:
-        """Return the target nodes adjacent to any of the source nodes."""
+        """Return the target nodes adjacent to any of the source nodes.
+
+        Args:
+            source_ids: The source IDs to start from when retrieving adjacent nodes.
+            query_embedding: The query embedding. Used to rank target nodes.
+            k_per_tag: The number of target nodes to fetch for each outgoing tag.
+
+        Returns:
+            List of adjacent edges.
+        """
 
         link_to_tags = set()
         targets = dict()
@@ -632,9 +651,10 @@ def add_sources(rows):
                 for new_tag in row.link_to_tags or []:
                     if new_tag not in link_to_tags:
                         link_to_tags.add(new_tag)
+
                         cq.execute(
-                            self._query_targets_embeddings_by_kind_and_tag,
-                            new_tag,
+                            self._query_targets_embeddings_by_kind_and_tag_and_embedding,
+                            parameters = (new_tag[0], new_tag[1], query_embedding, k_per_tag or 10),
                             callback=add_targets,
                         )
                         link_to_tags.add(new_tag)
@@ -653,6 +673,7 @@ def add_targets(rows):
                     self._query_source_tags_by_id, (source_id,), callback=add_sources
                 )
 
+        # TODO: Consider a combined limit based on the similarity and/or predicated MMR score?
         return [
             _Edge(target_content_id=content_id, target_text_embedding=embedding)
             for (content_id, embedding) in targets.items()

libs/langchain/tests/integration_tests/test_graph_store.py

@@ -112,54 +112,6 @@ def _result_ids(docs: Iterable[Document]) -> List[str]:
     return list(map(lambda d: d.metadata[METADATA_CONTENT_ID_KEY], docs))
 
 
-def test_link_directed(cassandra: GraphStoreFactory) -> None:
-    a = Document(
-        page_content="A",
-        metadata={
-            METADATA_CONTENT_ID_KEY: "a",
-            METADATA_LINKS_KEY: {
-                Link.incoming(kind="hyperlink", tag="http://a"),
-            },
-        },
-    )
-    b = Document(
-        page_content="B",
-        metadata={
-            METADATA_CONTENT_ID_KEY: "b",
-            METADATA_LINKS_KEY: {
-                Link.incoming(kind="hyperlink", tag="http://b"),
-                Link.outgoing(kind="hyperlink", tag="http://a"),
-            },
-        },
-    )
-    c = Document(
-        page_content="C",
-        metadata={
-            METADATA_CONTENT_ID_KEY: "c",
-            METADATA_LINKS_KEY: {
-                Link.outgoing(kind="hyperlink", tag="http://a"),
-            },
-        },
-    )
-    d = Document(
-        page_content="D",
-        metadata={
-            METADATA_CONTENT_ID_KEY: "d",
-            METADATA_LINKS_KEY: {
-                Link.outgoing(kind="hyperlink", tag="http://a"),
-                Link.outgoing(kind="hyperlink", tag="http://b"),
-            },
-        },
-    )
-
-    store = cassandra.store([a, b, c, d])
-
-    assert list(store.store._linked_ids("a")) == []
-    assert list(store.store._linked_ids("b")) == ["a"]
-    assert list(store.store._linked_ids("c")) == ["a"]
-    assert sorted(store.store._linked_ids("d")) == ["a", "b"]
-
-
 @pytest.mark.parametrize("gs_factory", ["cassandra", "astra_db"])
 def test_mmr_traversal(request, gs_factory: str):
     """