Use more generic types in API

Author: cbornet

libs/knowledge-store/ragstack_knowledge_store/knowledge_store.py

@@ -325,7 +325,7 @@ def _concurrent_queries(self) -> ConcurrentQueries:
     def add_nodes(
         self,
         nodes: Iterable[Node] = None,
-    ):
+    ) -> Iterable[str]:
         texts = []
         metadatas = []
         for node in nodes:
@@ -483,7 +483,7 @@ def add_edges_for_targets(
     def _query_by_ids(
         self,
         ids: Iterable[str],
-    ) -> Iterable[TextNode]:
+    ) -> List[TextNode]:
         results = []
         with self._concurrent_queries() as cq:
 

libs/knowledge-store/ragstack_knowledge_store/langchain/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(