Cleanup knowledge store code (#475)

Author: cbornet

libs/knowledge-store/pyproject.toml

@@ -13,7 +13,6 @@ packages = [{ include = "ragstack_knowledge_store" }]
 python = ">=3.10,<3.13"
 langchain-core = "^0.2"
 cassio = "^0.1.7"
-asyncstdlib = "^3.12.3"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "*"

libs/knowledge-store/ragstack_knowledge_store/_utils.py

@@ -19,12 +19,14 @@ def batched(iterable: Iterable[T], n: int) -> Iterator[Iterator[T]]:
         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):

libs/knowledge-store/ragstack_knowledge_store/base.py

@@ -223,7 +223,7 @@ def mmr_traversal_search(
         depth: int = 2,
         fetch_k: int = 100,
         lambda_mult: float = 0.5,
-        score_threshold: float = float('-inf'),
+        score_threshold: float = float("-inf"),
         **kwargs: Any,
     ) -> Iterable[Document]:
         """Retrieve documents from this knowledge store using MMR-traversal.
@@ -258,7 +258,7 @@ async def ammr_traversal_search(
         depth: int = 2,
         fetch_k: int = 100,
         lambda_mult: float = 0.5,
-        score_threshold: float = float('-inf'),
+        score_threshold: float = float("-inf"),
         **kwargs: Any,
     ) -> AsyncIterable[Document]:
         """Retrieve documents from this knowledge store using MMR-traversal.
@@ -297,17 +297,23 @@ async def ammr_traversal_search(
         ):
             yield doc
 
-    def similarity_search(self, query: str, k: int = 4, **kwargs: Any) -> List[Document]:
+    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]:
+    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)
@@ -322,7 +328,9 @@ def search(self, query: str, search_type: str, **kwargs: Any) -> List[Document]:
                 "'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":
@@ -420,7 +428,9 @@ def _get_relevant_documents(
         if self.search_type == "traversal":
             return list(self.vectorstore.traversal_search(query, **self.search_kwargs))
         elif self.search_type == "mmr_traversal":
-            return list(self.vectorstore.mmr_traversal_search(query, **self.search_kwargs))
+            return list(
+                self.vectorstore.mmr_traversal_search(query, **self.search_kwargs)
+            )
         else:
             return super()._get_relevant_documents(query, run_manager=run_manager)
 
@@ -430,12 +440,18 @@ async def _aget_relevant_documents(
         if self.search_type == "traversal":
             return [
                 doc
-                async for doc in self.vectorstore.atraversal_search(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)
+                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
+            )

libs/knowledge-store/ragstack_knowledge_store/cassandra.py

@@ -338,26 +338,35 @@ def add_nodes(
                 id = metadata[CONTENT_ID]
                 ids.append(id)
 
-                link_to_tags = set() # link to these tags
-                link_from_tags = set() # link from these tags
+                link_to_tags = set()  # link to these tags
+                link_from_tags = set()  # link from these tags
 
                 for tag in get_link_tags(metadata):
                     tag_str = f"{tag.kind}:{tag.tag}"
                     if tag.direction == "incoming" or tag.direction == "bidir":
                         # An incom`ing link should be linked *from* nodes with the given tag.
                         link_from_tags.add(tag_str)
-                        tag_to_new_targets.setdefault(tag_str, dict())[id] = (tag.kind, text_embedding)
+                        tag_to_new_targets.setdefault(tag_str, dict())[id] = (
+                            tag.kind,
+                            text_embedding,
+                        )
                     if tag.direction == "outgoing" or tag.direction == "bidir":
                         link_to_tags.add(tag_str)
-                        tag_to_new_sources.setdefault(tag_str, list()).append((tag.kind, id))
+                        tag_to_new_sources.setdefault(tag_str, list()).append(
+                            (tag.kind, id)
+                        )
 
-                cq.execute(self._insert_passage, (id, text, text_embedding, link_to_tags, link_from_tags))
+                cq.execute(
+                    self._insert_passage,
+                    (id, text, text_embedding, link_to_tags, link_from_tags),
+                )
 
         # Step 2: Query information about those tags to determine the edges to add.
         # Add edges as needed.
         id_set = set(ids)
         with self._concurrent_queries() as cq:
             edges = []
+
             def add_edge(source_id, target_id, kind, target_embedding):
                 nonlocal added_edges
                 if source_id == target_id:
@@ -399,27 +408,31 @@ def add_edges_for_targets(
                         # Don't add here (will be handled later).
                         continue
 
-                    for (kind, source_id) in sources:
-                         add_edge(source_id, target.content_id, kind, target.text_embedding)
+                    for kind, source_id in sources:
+                        add_edge(
+                            source_id, target.content_id, kind, target.text_embedding
+                        )
 
             for tag, new_target_embs in tag_to_new_targets.items():
                 # For each new node with a `link_from_tag`, find the source
                 # nodes with that `link_to_tag`` and create the edges.
                 cq.execute(
                     self._query_ids_by_link_to_tag,
-                    parameters=(tag, ),
+                    parameters=(tag,),
                     callback=lambda sources, targets=new_target_embs: add_edges_for_sources(
-                        sources, targets)
+                        sources, targets
+                    ),
                 )
 
             for tag, new_sources in tag_to_new_sources.items():
                 # For each new node with a `link_to_tag`, find the target
                 # nodes with that `link_from_tag` tag and create the edges.
                 cq.execute(
                     self._query_ids_and_embedding_by_link_from_tag,
-                    parameters=(tag, ),
+                    parameters=(tag,),
                     callback=lambda targets, sources=new_sources: add_edges_for_targets(
-                        sources, targets)
+                        sources, targets
+                    ),
                 )
 
         # Step 3: Add edges.
@@ -429,27 +442,29 @@ def add_edges_for_targets(
         # more than |max concurency| edges.
         added_edges = 0
         with self._concurrent_queries() as cq:
-            print("Adding edges")
             # Add edges from query results (should be one new node and one old node)
             for edge in edges:
                 added_edges += 1
                 cq.execute(self._insert_edge, edge)
 
             # Add edges for new nodes
             for tag, new_sources in tag_to_new_sources.items():
-                for (kind, source_id) in new_sources:
+                for kind, source_id in new_sources:
                     new_targets = tag_to_new_targets.get(tag, None)
                     if new_targets is None:
                         continue
 
-                    for (target_id, (target_kind, target_embedding)) in new_targets.items():
+                    for target_id, (
+                        target_kind,
+                        target_embedding,
+                    ) in new_targets.items():
                         # TODO: Improve the structures so this can be a lookup?
                         if target_kind == kind and source_id != target_id:
                             added_edges += 1
-                            cq.execute(self._insert_edge,
-                                       (source_id, target_id, kind, target_embedding))
-
-        print(f"Added {added_edges} edges")
+                            cq.execute(
+                                self._insert_edge,
+                                (source_id, target_id, kind, target_embedding),
+                            )
 
         return ids
 
@@ -530,7 +545,7 @@ def mmr_traversal_search(
         depth: int = 2,
         fetch_k: int = 100,
         lambda_mult: float = 0.5,
-        score_threshold: float = float('-inf'),
+        score_threshold: float = float("-inf"),
     ) -> Iterable[Document]:
         """Retrieve documents from this knowledge store using MMR-traversal.
 
@@ -588,7 +603,7 @@ def mmr_traversal_search(
             selected_embedding = next_selected.embedding
             selected_embeddings.append(selected_embedding)
 
-            best_score = float('-inf')
+            best_score = float("-inf")
             next_id = None
 
             # Update unselected scores.

libs/knowledge-store/ragstack_knowledge_store/edge_extractor.py

@@ -1,13 +1,20 @@
 from __future__ import annotations
 
 import abc
-import dataclasses
 from abc import ABC, abstractmethod
-from typing import Any, AsyncIterator, Dict, Generic, Iterable, Iterator, Literal, Set, Sequence, TypeVar, Union
-
-import asyncstdlib
-from langchain_core.runnables import run_in_executor
-from langchain_core.documents import Document, BaseDocumentTransformer
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    Iterator,
+    Literal,
+    Set,
+    TypeVar,
+    Union,
+)
+
+from langchain_core.documents import Document
 from pydantic import BaseModel
 from ._utils import strict_zip
 
@@ -20,17 +27,22 @@ class LinkTag(BaseModel, abc.ABC):
     def __hash__(self):
         return hash((type(self),) + tuple(self.__dict__.values()))
 
+
 class OutgoingLinkTag(LinkTag):
     direction: Literal["outgoing"] = "outgoing"
 
+
 class IncomingLinkTag(LinkTag):
     direction: Literal["incoming"] = "incoming"
 
+
 class BidirLinkTag(LinkTag):
     direction: Literal["bidir"] = "bidir"
 
+
 LINK_TAGS = "link_tags"
 
+
 def get_link_tags(doc_or_md: Union[Document, Dict[str, Any]]) -> Set[LinkTag]:
     """Get the link-tag set from a document or metadata.
 
@@ -49,7 +61,10 @@ def get_link_tags(doc_or_md: Union[Document, Dict[str, Any]]) -> Set[LinkTag]:
         doc_or_md[LINK_TAGS] = link_tags
     return link_tags
 
+
 InputT = TypeVar("InputT")
+
+
 class EdgeExtractor(ABC, Generic[InputT]):
     @abstractmethod
     def extract_one(self, document: Document, input: InputT):
@@ -60,12 +75,14 @@ def extract_one(self, document: Document, input: InputT):
             inputs: The input content to extract edges from.
         """
 
-    def extract(self, documents: Iterable[Document], inputs: Iterable[InputT]) -> Iterator[Set[LinkTag]]:
+    def extract(
+        self, documents: Iterable[Document], inputs: Iterable[InputT]
+    ) -> Iterator[Set[LinkTag]]:
         """Add edges from each `input` to the corresponding documents.
 
         Args:
             documents: The documents to add the link tags to.
             inputs: The input content to extract edges from.
         """
-        for (document, input) in strict_zip(documents, inputs):
-            self.extract_one(document, input)
+        for document, input in strict_zip(documents, inputs):
+            self.extract_one(document, input)