Implement count_nodes and count_edges methods for Neo4j

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

@@ -87,6 +87,22 @@ class Neo4jAdapter(GraphDBInterface):
         async with self.driver.session(database=self.graph_database_name) as session:
             yield session
 
+    async def count_edges(self) -> int:
+        query = """
+        MATCH ()-[r]->()
+        RETURN COUNT(r) as total_edges;
+        """
+        query_result = await self.query(query)
+        return query_result[0]["total_edges"]
+
+    async def count_nodes(self) -> int:
+        query = """
+        MATCH (n)
+        RETURN COUNT(n) as total_nodes;
+        """
+        query_result = await self.query(query)
+        return query_result[0]["total_nodes"]
+
     @deadlock_retry()
     async def query(
         self,

cognee/tests/test_neo4j.py

@@ -35,6 +35,15 @@ async def main():
     explanation_file_path_nlp = os.path.join(
         pathlib.Path(__file__).parent, "test_data/Natural_language_processing.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, "Graph has to be empty"
+
     await cognee.add([explanation_file_path_nlp], dataset_name)
 
     explanation_file_path_quantum = os.path.join(
@@ -43,8 +52,18 @@ async def main():
 
     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, "Graph has to 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, "Graph shouldn't be empty"
+
     from cognee.infrastructure.databases.vector import get_vector_engine
 
     vector_engine = get_vector_engine()
@@ -117,11 +136,9 @@ 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()
-    nodes, edges = await graph_engine.get_graph_data()
-    assert len(nodes) == 0 and len(edges) == 0, "Neo4j graph database is not empty"
+    edges_count = await graph_engine.count_edges()
+    nodes_count = await graph_engine.count_nodes()
+    assert nodes_count == 0 and edges_count == 0, "Neo4j graph database is not empty"
 
 
 if __name__ == "__main__":