From f3ec1801025eb5cc1c2dc899a8aa3eca02ae4165 Mon Sep 17 00:00:00 2001 From: Daulet Amirkhanov Date: Wed, 15 Oct 2025 16:39:25 +0100 Subject: [PATCH] Implement count_edges and count_methods for Kuzu --- .../databases/graph/kuzu/adapter.py | 16 ++++++++++ cognee/tests/test_kuzu.py | 29 ++++++++++++++++--- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/cognee/infrastructure/databases/graph/kuzu/adapter.py b/cognee/infrastructure/databases/graph/kuzu/adapter.py index 7b772097f..a31726c9a 100644 --- a/cognee/infrastructure/databases/graph/kuzu/adapter.py +++ b/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. diff --git a/cognee/tests/test_kuzu.py b/cognee/tests/test_kuzu.py index 8749e42d0..e39edd06a 100644 --- a/cognee/tests/test_kuzu.py +++ b/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() - nodes, edges = await graph_engine.get_graph_data() - assert len(nodes) == 0 and len(edges) == 0, "Kuzu graph database is not empty" + 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" finally: # Ensure cleanup even if tests fail