Update BaseDocumentStore to not return Nones in result (#19513)

Author: strawgate

llama-index-core/llama_index/core/storage/docstore/types.py

@@ -54,14 +54,10 @@ async def async_add_documents(
     ) -> None: ...
 
     @abstractmethod
-    def get_document(
-        self, doc_id: str, raise_error: bool = True
-    ) -> Optional[BaseNode]: ...
+    def get_document(self, doc_id: str, raise_error: bool = True) -> Optional[BaseNode]: ...
 
     @abstractmethod
-    async def aget_document(
-        self, doc_id: str, raise_error: bool = True
-    ) -> Optional[BaseNode]: ...
+    async def aget_document(self, doc_id: str, raise_error: bool = True) -> Optional[BaseNode]: ...
 
     @abstractmethod
     def delete_document(self, doc_id: str, raise_error: bool = True) -> None:
@@ -130,9 +126,7 @@ async def adelete_ref_doc(self, ref_doc_id: str, raise_error: bool = True) -> No
         """Delete a ref_doc and all it's associated nodes."""
 
     # ===== Nodes =====
-    def get_nodes(
-        self, node_ids: List[str], raise_error: bool = True
-    ) -> List[BaseNode]:
+    def get_nodes(self, node_ids: List[str], raise_error: bool = True) -> List[BaseNode]:
         """
         Get nodes from docstore.
 
@@ -141,15 +135,20 @@ def get_nodes(
             raise_error (bool): raise error if node_id not found
 
         """
-        # if/else needed for type checking
-        if raise_error:
-            return [node for node_id in node_ids if (node := self.get_node(node_id, raise_error=True))]
-        else:
-            return [self.get_node(node_id) for node_id in node_ids]
-
-    async def aget_nodes(
-        self, node_ids: List[str], raise_error: bool = True
-    ) -> List[BaseNode]:
+        nodes: list[BaseNode] = []
+
+        for node_id in node_ids:
+            # if needed for type checking
+            if not raise_error:
+                if node := self.get_node(node_id=node_id, raise_error=False):
+                    nodes.append(node)
+                continue
+
+            nodes.append(self.get_node(node_id=node_id, raise_error=True))
+
+        return nodes
+
+    async def aget_nodes(self, node_ids: List[str], raise_error: bool = True) -> List[BaseNode]:
         """
         Get nodes from docstore.
 
@@ -158,18 +157,18 @@ async def aget_nodes(
             raise_error (bool): raise error if node_id not found
 
         """
-        # if/else needed for type checking
-        if raise_error:
-            return [
-                node
-                for node_id in node_ids
-                if (node := await self.aget_node(node_id, raise_error=True))
-            ]
-        else:
-            return [
-                await self.aget_node(node_id)
-                for node_id in node_ids
-            ]
+        nodes: list[BaseNode] = []
+
+        for node_id in node_ids:
+            # if needed for type checking
+            if not raise_error:
+                if node := await self.aget_node(node_id=node_id, raise_error=False):
+                    nodes.append(node)
+                continue
+
+            nodes.append(await self.aget_node(node_id=node_id, raise_error=True))
+
+        return nodes
 
     @overload
     def get_node(self, node_id: str, raise_error: Literal[True] = True) -> BaseNode: ...
@@ -239,9 +238,7 @@ def get_node_dict(self, node_id_dict: Dict[int, str]) -> Dict[int, BaseNode]:
             node_id_dict (Dict[int, str]): mapping of index to node ids
 
         """
-        return {
-            index: self.get_node(node_id) for index, node_id in node_id_dict.items()
-        }
+        return {index: self.get_node(node_id) for index, node_id in node_id_dict.items()}
 
     async def aget_node_dict(self, node_id_dict: Dict[int, str]) -> Dict[int, BaseNode]:
         """
@@ -251,7 +248,4 @@ async def aget_node_dict(self, node_id_dict: Dict[int, str]) -> Dict[int, BaseNo
             node_id_dict (Dict[int, str]): mapping of index to node ids
 
         """
-        return {
-            index: await self.aget_node(node_id)
-            for index, node_id in node_id_dict.items()
-        }
+        return {index: await self.aget_node(node_id) for index, node_id in node_id_dict.items()}