Implement count_edges and count_methods for Kuzu

cognee/infrastructure/databases/graph/kuzu/adapter.py

@@ -185,6 +185,22 @@ class KuzuAdapter(GraphDBInterface):
         except FileNotFoundError:
             logger.warning(f"Kuzu S3 storage file not found: {self.db_path}")
 
+    async def count_edges(self) -> int:
+        query = """
+        MATCH ()-[r]->()
+        RETURN COUNT(r);
+        """
+        query_result = await self.query(query)
+        return query_result[0][0]
+
+    async def count_nodes(self) -> int:
+        query = """
+        MATCH (n)
+        RETURN COUNT(n);
+        """
+        query_result = await self.query(query)
+        return query_result[0][0]
+
     async def query(self, query: str, params: Optional[dict] = None) -> List[Tuple]:
         """
         Execute a Kuzu query asynchronously with automatic reconnection.

cognee/tests/test_kuzu.py

@@ -47,10 +47,31 @@ async def main():
             pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
         )
 
+        from cognee.infrastructure.databases.graph import get_graph_engine
+
+        graph_engine = await get_graph_engine()
+
+        edges_count = await graph_engine.count_edges()
+        nodes_count = await graph_engine.count_nodes()
+
+        assert edges_count == 0 and nodes_count == 0, "Kuzu graph database is not empty"
+
         await cognee.add([explanation_file_path_quantum], dataset_name)
 
+        edges_count = await graph_engine.count_edges()
+        nodes_count = await graph_engine.count_nodes()
+
+        assert edges_count == 0 and nodes_count == 0, (
+            "Kuzu graph database should be empty before cognify"
+        )
+
         await cognee.cognify([dataset_name])
 
+        edges_count = await graph_engine.count_edges()
+        nodes_count = await graph_engine.count_nodes()
+
+        assert edges_count != 0 and nodes_count != 0, "Kuzu graph database should not be empty"
+
         from cognee.infrastructure.databases.vector import get_vector_engine
 
         vector_engine = get_vector_engine()
@@ -114,11 +135,11 @@ async def main():
         assert not os.path.isdir(data_root_directory), "Local data files are not deleted"
 
         await cognee.prune.prune_system(metadata=True)
-        from cognee.infrastructure.databases.graph import get_graph_engine
 
-        graph_engine = await get_graph_engine()
+        edges_count = await graph_engine.count_edges()
-        nodes, edges = await graph_engine.get_graph_data()
+        nodes_count = await graph_engine.count_nodes()
-        assert len(nodes) == 0 and len(edges) == 0, "Kuzu graph database is not empty"
+
+        assert edges_count == 0 and nodes_count == 0, "Kuzu graph database is not empty"
 
     finally:
         # Ensure cleanup even if tests fail