[knowledge-store] feat: Dynamic, MMR-based traversal (#460)

This required denormalizing the `text_embedding` of target nodes into
the edge, making it easier to guide traversal based on distance to a
query.

Also fixed a bug in keyword linking which didn't create links to old
nodes with the given keyword.

Author: bjchambers

libs/knowledge-store/ragstack_knowledge_store/__init__.py

@@ -1,4 +1,4 @@
-from .cassandra import CassandraKnowledgeStore
 from .base import KnowledgeStore
+from .cassandra import CassandraKnowledgeStore
 
 __all__ = ["CassandraKnowledgeStore", "KnowledgeStore"]

libs/knowledge-store/ragstack_knowledge_store/_utils.py

@@ -1,3 +1,5 @@
+import sys
+
 try:
     # Try importing the function from itertools (Python 3.12+)
     from itertools import batched
@@ -16,3 +18,39 @@ def batched(iterable: Iterable[T], n: int) -> Iterator[Iterator[T]]:
         it = iter(iterable)
         while batch := tuple(islice(it, n)):
             yield batch
+
+# TODO: Remove the "polyfill" when we required python is >= 3.10.
+
+if sys.version_info >= (3, 10):
+
+    def strict_zip(*iterables):
+        return zip(*iterables, strict=True)
+else:
+
+    def strict_zip(*iterables):
+        # Custom implementation for Python versions older than 3.10
+        if not iterables:
+            return
+
+        iterators = tuple(iter(iterable) for iterable in iterables)
+        try:
+            while True:
+                items = []
+                for iterator in iterators:
+                    items.append(next(iterator))
+                yield tuple(items)
+        except StopIteration:
+            pass
+
+        if items:
+            i = len(items)
+            plural = " " if i == 1 else "s 1-"
+            msg = f"strict_zip() argument {i+1} is shorter than argument{plural}{i}"
+            raise ValueError(msg)
+
+        sentinel = object()
+        for i, iterator in enumerate(iterators[1:], 1):
+            if next(iterator, sentinel) is not sentinel:
+                plural = " " if i == 1 else "s 1-"
+                msg = f"strict_zip() argument {i+1} is longer than argument{plural}{i}"
+                raise ValueError(msg)

libs/knowledge-store/ragstack_knowledge_store/base.py

@@ -36,7 +36,7 @@ class Node(Serializable):
     id: Optional[str]
     """Unique ID for the node. Shall be generated by the KnowledgeStore if not set"""
     metadata: dict = Field(default_factory=dict)
-    """Metadata for the node. May contain information used to link this node 
+    """Metadata for the node. May contain information used to link this node
     with other nodes."""
 
 
@@ -164,7 +164,7 @@ async def aadd_documents(
         return await self.aadd_nodes(nodes, **kwargs)
 
     @abstractmethod
-    def traversing_retrieve(
+    def traversal_search(
         self,
         query: str,
         *,
@@ -187,7 +187,7 @@ def traversing_retrieve(
             Retrieved documents.
         """
 
-    async def atraversing_retrieve(
+    async def atraversal_search(
         self,
         query: str,
         *,
@@ -210,42 +210,121 @@ async def atraversing_retrieve(
             Retrieved documents.
         """
         for doc in await run_in_executor(
-            None, self.traversing_retrieve, query, k=k, depth=depth, **kwargs
+            None, self.traversal_search, query, k=k, depth=depth, **kwargs
         ):
             yield doc
 
-    def similarity_search(
-        self, query: str, k: int = 4, **kwargs: Any
-    ) -> List[Document]:
-        return list(self.traversing_retrieve(query, k=k, depth=0))
+    @abstractmethod
+    def mmr_traversal_search(
+        self,
+        query: str,
+        *,
+        k: int = 4,
+        depth: int = 2,
+        fetch_k: int = 100,
+        lambda_mult: float = 0.5,
+        score_threshold: float = 0.0,
+        **kwargs: Any,
+    ) -> Iterable[Document]:
+        """Retrieve documents from this knowledge store using MMR-traversal.
 
-    async def asimilarity_search(
-        self, query: str, k: int = 4, **kwargs: Any
-    ) -> List[Document]:
-        return [doc async for doc in self.atraversing_retrieve(query, k=k, depth=0)]
+        This strategy first retrieves the top `fetch_k` results by similarity to
+        the question. It then selects the top `k` results based on
+        maximum-marginal relevance using the given `lambda_mult`.
+
+        At each step, it considers the (remaining) documents from `fetch_k` as
+        well as any documents connected by edges to a selected document
+        retrieved based on similarity (a "root").
+
+        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.
+                Defaults to 10.
+            depth: Maximum depth of a node (number of edges) from a node
+                retrieved via similarity. Defaults to 2.
+            lambda_mult: Number between 0 and 1 that determines the degree
+                of diversity among the results with 0 corresponding to maximum
+                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 0.0 so all scores are
+                taken.
+        """
+
+    async def ammr_traversal_search(
+        self,
+        query: str,
+        *,
+        k: int = 4,
+        depth: int = 2,
+        fetch_k: int = 100,
+        lambda_mult: float = 0.5,
+        score_threshold: float = 0.0,
+        **kwargs: Any,
+    ) -> AsyncIterable[Document]:
+        """Retrieve documents from this knowledge store using MMR-traversal.
+
+        This strategy first retrieves the top `fetch_k` results by similarity to
+        the question. It then selects the top `k` results based on
+        maximum-marginal relevance using the given `lambda_mult`.
+
+        At each step, it considers the (remaining) documents from `fetch_k` as
+        well as any documents connected by edges to a selected document
+        retrieved based on similarity (a "root").
+
+        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.
+                Defaults to 10.
+            depth: Maximum depth of a node (number of edges) from a node
+                retrieved via similarity. Defaults to 2.
+            lambda_mult: Number between 0 and 1 that determines the degree
+                of diversity among the results with 0 corresponding to maximum
+                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 0.0 so all scores are
+                taken.
+        """
+        for doc in await run_in_executor(
+            None,
+            self.traversal_search,
+            query,
+            k=k,
+            fetch_k=fetch_k,
+            depth=depth,
+            lambda_mult=lambda_mult,
+            score_threshold=score_threshold,
+            **kwargs,
+        ):
+            yield doc
+
+    def similarity_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Document]:
+        return list(self.traversal_search(query, k=k, depth=0))
+
+    async def asimilarity_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Document]:
+        return [doc async for doc in self.atraversal_search(query, k=k, depth=0)]
 
     def search(self, query: str, search_type: str, **kwargs: Any) -> List[Document]:
         if search_type == "similarity":
             return self.similarity_search(query, **kwargs)
         elif search_type == "similarity_score_threshold":
-            docs_and_similarities = self.similarity_search_with_relevance_scores(
-                query, **kwargs
-            )
+            docs_and_similarities = self.similarity_search_with_relevance_scores(query, **kwargs)
             return [doc for doc, _ in docs_and_similarities]
         elif search_type == "mmr":
             return self.max_marginal_relevance_search(query, **kwargs)
         elif search_type == "traversal":
-            return list(self.traversing_retrieve(query, **kwargs))
+            return list(self.traversal_search(query, **kwargs))
+        elif search_type == "mmr_traversal":
+            return list(self.mmr_traversal_search(query, **kwargs))
         else:
             raise ValueError(
                 f"search_type of {search_type} not allowed. Expected "
                 "search_type to be 'similarity', 'similarity_score_threshold', "
                 "'mmr' or 'traversal'."
             )
 
-    async def asearch(
-        self, query: str, search_type: str, **kwargs: Any
-    ) -> List[Document]:
+    async def asearch(self, query: str, search_type: str, **kwargs: Any) -> List[Document]:
         if search_type == "similarity":
             return await self.asimilarity_search(query, **kwargs)
         elif search_type == "similarity_score_threshold":
@@ -256,7 +335,7 @@ async def asearch(
         elif search_type == "mmr":
             return await self.amax_marginal_relevance_search(query, **kwargs)
         elif search_type == "traversal":
-            return [doc async for doc in self.atraversing_retrieve(query, **kwargs)]
+            return [doc async for doc in self.atraversal_search(query, **kwargs)]
         else:
             raise ValueError(
                 f"search_type of {search_type} not allowed. Expected "
@@ -334,15 +413,16 @@ class KnowledgeStoreRetriever(VectorStoreRetriever):
         "similarity_score_threshold",
         "mmr",
         "traversal",
+        "mmr_traversal",
     )
 
     def _get_relevant_documents(
         self, query: str, *, run_manager: CallbackManagerForRetrieverRun
     ) -> List[Document]:
         if self.search_type == "traversal":
-            return list(
-                self.vectorstore.traversing_retrieve(query, **self.search_kwargs)
-            )
+            return list(self.vectorstore.traversal_search(query, **self.search_kwargs))
+        elif self.search_type == "mmr_traversal":
+            return list(self.vectorstore.traversal_search(query, **self.search_kwargs))
         else:
             return super()._get_relevant_documents(query, run_manager=run_manager)
 
@@ -352,11 +432,12 @@ async def _aget_relevant_documents(
         if self.search_type == "traversal":
             return [
                 doc
-                async for doc in self.vectorstore.atraversing_retrieve(
-                    query, **self.search_kwargs
-                )
+                async for doc in self.vectorstore.atraversal_search(query, **self.search_kwargs)
+            ]
+        elif self.search_type == "mmr_traversal":
+            return [
+                doc
+                async for doc in self.vectorstore.ammr_traversal_search(query, **self.search_kwargs)
             ]
         else:
-            return await super()._aget_relevant_documents(
-                query, run_manager=run_manager
-            )
+            return await super()._aget_relevant_documents(query, run_manager=run_manager)