From ea4a93efb172a82754a342084aa95393a0f11759 Mon Sep 17 00:00:00 2001
From: Daulet Amirkhanov <damirkhanov01@gmail.com>
Date: Wed, 15 Oct 2025 16:57:53 +0100
Subject: [PATCH] Implement count_nodes and count_edges methods for Neo4j

---
 .../databases/graph/neo4j_driver/adapter.py   | 16 +++++++++++
 cognee/tests/test_neo4j.py                    | 27 +++++++++++++++----
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py
index 520295ed2..a61ab6f0b 100644
--- a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py
+++ b/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,
diff --git a/cognee/tests/test_neo4j.py b/cognee/tests/test_neo4j.py
index c74b4ab65..11f6156bd 100644
--- a/cognee/tests/test_neo4j.py
+++ b/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__":