Use more generic types in API

Author: cbornet

libs/knowledge-store/ragstack_knowledge_store/base.py

@@ -102,7 +102,7 @@ def add_nodes(
         self,
         nodes: Iterable[Node],
         **kwargs: Any,
-    ) -> List[str]:
+    ) -> Iterable[str]:
         """Add nodes to the knowledge store
 
         Args:
@@ -113,13 +113,19 @@ async def aadd_nodes(
         self,
         nodes: Iterable[Node],
         **kwargs: Any,
-    ) -> List[str]:
+    ) -> AsyncIterable[str]:
         """Add nodes to the knowledge store
 
         Args:
             nodes: the nodes to add.
         """
-        return await run_in_executor(None, self.add_nodes, nodes, **kwargs)
+        iterator = iter(await run_in_executor(None, self.add_nodes, nodes, **kwargs))
+        done = object()
+        while True:
+            doc = await run_in_executor(None, next, iterator, done)
+            if doc is done:
+                break
+            yield doc
 
     def add_texts(
         self,
@@ -130,7 +136,7 @@ def add_texts(
         **kwargs: Any,
     ) -> List[str]:
         nodes = _texts_to_nodes(texts, metadatas, ids)
-        return self.add_nodes(nodes, **kwargs)
+        return list(self.add_nodes(nodes, **kwargs))
 
     async def aadd_texts(
         self,
@@ -141,7 +147,7 @@ async def aadd_texts(
         **kwargs: Any,
     ) -> List[str]:
         nodes = _texts_to_nodes(texts, metadatas, ids)
-        return await self.aadd_nodes(nodes, **kwargs)
+        return [_id async for _id in self.aadd_nodes(nodes, **kwargs)]
 
     def add_documents(
         self,
@@ -151,7 +157,7 @@ def add_documents(
         **kwargs: Any,
     ) -> List[str]:
         nodes = _documents_to_nodes(documents, ids)
-        return self.add_nodes(nodes, **kwargs)
+        return list(self.add_nodes(nodes, **kwargs))
 
     async def aadd_documents(
         self,
@@ -161,7 +167,7 @@ async def aadd_documents(
         **kwargs: Any,
     ) -> List[str]:
         nodes = _documents_to_nodes(documents, ids)
-        return await self.aadd_nodes(nodes, **kwargs)
+        return [_id async for _id in self.aadd_nodes(nodes, **kwargs)]
 
     @abstractmethod
     def traversal_search(
@@ -209,9 +215,16 @@ async def atraversal_search(
         Returns:
             Retrieved documents.
         """
-        for doc in await run_in_executor(
-            None, self.traversal_search, query, k=k, depth=depth, **kwargs
-        ):
+        iterator = iter(
+            await run_in_executor(
+                None, self.traversal_search, query, k=k, depth=depth, **kwargs
+            )
+        )
+        done = object()
+        while True:
+            doc = await run_in_executor(None, next, iterator, done)
+            if doc is done:
+                break
             yield doc
 
     @abstractmethod
@@ -284,17 +297,24 @@ async def ammr_traversal_search(
             score_threshold: Only documents with a score greater than or equal
                 this threshold will be chosen. Defaults to negative infinity.
         """
-        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,
-        ):
+        iterator = iter(
+            await run_in_executor(
+                None,
+                self.mmr_traversal_search,
+                query,
+                k=k,
+                fetch_k=fetch_k,
+                depth=depth,
+                lambda_mult=lambda_mult,
+                score_threshold=score_threshold,
+                **kwargs,
+            )
+        )
+        done = object()
+        while True:
+            doc = await run_in_executor(None, next, iterator, done)
+            if doc is done:
+                break
             yield doc
 
     def similarity_search(

libs/knowledge-store/ragstack_knowledge_store/cassandra.py

@@ -312,7 +312,7 @@ def add_nodes(
         self,
         nodes: Iterable[Node] = None,
         **kwargs: Any,
-    ):
+    ) -> List[str]:
         texts = []
         metadatas = []
         for node in nodes:
@@ -512,7 +512,7 @@ def similarity_search_by_vector(
     def _query_by_ids(
         self,
         ids: Iterable[str],
-    ) -> Iterable[Document]:
+    ) -> List[Document]:
         results = []
         with self._concurrent_queries() as cq:
 
@@ -545,7 +545,8 @@ def mmr_traversal_search(
         fetch_k: int = 100,
         lambda_mult: float = 0.5,
         score_threshold: float = float("-inf"),
-    ) -> Iterable[Document]:
+        **kwargs: Any,
+    ) -> List[Document]:
         """Retrieve documents from this knowledge store using MMR-traversal.
 
         This strategy first retrieves the top `fetch_k` results by similarity to
@@ -645,8 +646,8 @@ def mmr_traversal_search(
         return self._query_by_ids(selected_ids)
 
     def traversal_search(
-        self, query: str, *, k: int = 4, depth: int = 1
-    ) -> Iterable[Document]:
+        self, query: str, *, k: int = 4, depth: int = 1, **kwargs: Any
+    ) -> List[Document]:
         """Retrieve documents from this knowledge store.
 
         First, `k` nodes are retrieved using a vector search for the `query` string.