diff --git a/cognee/api/v1/search/search.py b/cognee/api/v1/search/search.py index 880a57b99..0caca619a 100644 --- a/cognee/api/v1/search/search.py +++ b/cognee/api/v1/search/search.py @@ -8,7 +8,7 @@ from cognee.modules.search.types import SearchResult, SearchType, CombinedSearch from cognee.modules.users.methods import get_default_user from cognee.modules.search.methods import search as search_function from cognee.modules.data.methods import get_authorized_existing_datasets -from cognee.modules.data.exceptions import DatasetNotFoundError, SearchOnEmptyGraphError +from cognee.modules.data.exceptions import DatasetNotFoundError async def search( @@ -177,12 +177,10 @@ async def search( raise DatasetNotFoundError(message="No datasets found.") graph_engine = await get_graph_engine() - nodes_count = await graph_engine.count_nodes() + is_empty = await graph_engine.is_empty() - if nodes_count == 0: - raise SearchOnEmptyGraphError( - message="Knowledge graph is empty, please ensure data is added and cognified." - ) + if is_empty: + return [] filtered_search_results = await search_function( query_text=query_text, diff --git a/cognee/infrastructure/databases/graph/graph_db_interface.py b/cognee/infrastructure/databases/graph/graph_db_interface.py index d7542eac6..67df1a27c 100644 --- a/cognee/infrastructure/databases/graph/graph_db_interface.py +++ b/cognee/infrastructure/databases/graph/graph_db_interface.py @@ -160,9 +160,9 @@ class GraphDBInterface(ABC): """ @abstractmethod - async def count_nodes(self) -> int: - logger.warning("count_nodes is not implemented") - return 1 # dummy value to not fail search() + async def is_empty(self) -> bool: + logger.warning("is_empty() is not implemented") + return True @abstractmethod async def query(self, query: str, params: dict) -> List[Any]: diff --git a/cognee/infrastructure/databases/graph/kuzu/adapter.py b/cognee/infrastructure/databases/graph/kuzu/adapter.py index 04c163efa..29ff92247 100644 --- a/cognee/infrastructure/databases/graph/kuzu/adapter.py +++ b/cognee/infrastructure/databases/graph/kuzu/adapter.py @@ -185,13 +185,14 @@ class KuzuAdapter(GraphDBInterface): except FileNotFoundError: logger.warning(f"Kuzu S3 storage file not found: {self.db_path}") - async def count_nodes(self) -> int: + async def is_empty(self) -> bool: query = """ MATCH (n) - RETURN COUNT(n); + RETURN true + LIMIT 1; """ query_result = await self.query(query) - return query_result[0][0] + return len(query_result) == 0 async def query(self, query: str, params: Optional[dict] = None) -> List[Tuple]: """ diff --git a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py index ac19069f4..5861b69cb 100644 --- a/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py +++ b/cognee/infrastructure/databases/graph/neo4j_driver/adapter.py @@ -87,13 +87,14 @@ class Neo4jAdapter(GraphDBInterface): async with self.driver.session(database=self.graph_database_name) as session: yield session - async def count_nodes(self) -> int: + async def is_empty(self) -> bool: query = """ + RETURN EXISTS { MATCH (n) - RETURN COUNT(n) as total_nodes; + } AS node_exists; """ query_result = await self.query(query) - return query_result[0]["total_nodes"] + return not query_result[0]["node_exists"] @deadlock_retry() async def query( diff --git a/cognee/modules/data/exceptions/__init__.py b/cognee/modules/data/exceptions/__init__.py index ba943634d..54af81070 100644 --- a/cognee/modules/data/exceptions/__init__.py +++ b/cognee/modules/data/exceptions/__init__.py @@ -9,5 +9,4 @@ from .exceptions import ( UnauthorizedDataAccessError, DatasetNotFoundError, DatasetTypeError, - SearchOnEmptyGraphError, ) diff --git a/cognee/modules/data/exceptions/exceptions.py b/cognee/modules/data/exceptions/exceptions.py index c2921750a..ac3b68e64 100644 --- a/cognee/modules/data/exceptions/exceptions.py +++ b/cognee/modules/data/exceptions/exceptions.py @@ -35,16 +35,6 @@ class DatasetNotFoundError(CogneeValidationError): super().__init__(message, name, status_code) -class SearchOnEmptyGraphError(CogneeValidationError): - def __init__( - self, - message: str = "Knowledge graph is empty, please ensure data is added and cognified.", - name: str = "SearchOnEmptyGraphError", - status_code=status.HTTP_400_BAD_REQUEST, - ): - super().__init__(message, name, status_code) - - class DatasetTypeError(CogneeValidationError): def __init__( self, diff --git a/cognee/tests/test_kuzu.py b/cognee/tests/test_kuzu.py index c07a51104..fe9da6dcb 100644 --- a/cognee/tests/test_kuzu.py +++ b/cognee/tests/test_kuzu.py @@ -51,21 +51,21 @@ async def main(): graph_engine = await get_graph_engine() - nodes_count = await graph_engine.count_nodes() + is_empty = await graph_engine.is_empty() - assert nodes_count == 0, "Kuzu graph database is not empty" + assert is_empty, "Kuzu graph database is not empty" await cognee.add([explanation_file_path_quantum], dataset_name) - nodes_count = await graph_engine.count_nodes() + is_empty = await graph_engine.is_empty() - assert nodes_count == 0, "Kuzu graph database should be empty before cognify" + assert is_empty, "Kuzu graph database should be empty before cognify" await cognee.cognify([dataset_name]) - nodes_count = await graph_engine.count_nodes() + is_empty = await graph_engine.is_empty() - assert nodes_count != 0, "Kuzu graph database should not be empty" + assert not is_empty, "Kuzu graph database should not be empty" from cognee.infrastructure.databases.vector import get_vector_engine @@ -131,9 +131,9 @@ async def main(): await cognee.prune.prune_system(metadata=True) - nodes_count = await graph_engine.count_nodes() + is_empty = await graph_engine.is_empty() - assert nodes_count == 0, "Kuzu graph database is not empty" + assert is_empty, "Kuzu graph database is not empty" finally: # Ensure cleanup even if tests fail diff --git a/cognee/tests/test_neo4j.py b/cognee/tests/test_neo4j.py index 6f1fcf975..925614e67 100644 --- a/cognee/tests/test_neo4j.py +++ b/cognee/tests/test_neo4j.py @@ -39,9 +39,9 @@ async def main(): graph_engine = await get_graph_engine() - nodes_count = await graph_engine.count_nodes() + is_empty = await graph_engine.is_empty() - assert nodes_count == 0, "Graph has to be empty" + assert is_empty, "Graph has to be empty" await cognee.add([explanation_file_path_nlp], dataset_name) @@ -50,15 +50,15 @@ async def main(): ) await cognee.add([explanation_file_path_quantum], dataset_name) - nodes_count = await graph_engine.count_nodes() + is_empty = await graph_engine.is_empty() - assert nodes_count == 0, "Graph has to be empty before cognify" + assert is_empty, "Graph has to be empty before cognify" await cognee.cognify([dataset_name]) - nodes_count = await graph_engine.count_nodes() + is_empty = await graph_engine.is_empty() - assert nodes_count != 0, "Graph shouldn't be empty" + assert not is_empty, "Graph shouldn't be empty" from cognee.infrastructure.databases.vector import get_vector_engine @@ -132,8 +132,8 @@ async def main(): assert not os.path.isdir(data_root_directory), "Local data files are not deleted" await cognee.prune.prune_system(metadata=True) - nodes_count = await graph_engine.count_nodes() - assert nodes_count == 0, "Neo4j graph database is not empty" + is_empty = await graph_engine.is_empty() + assert is_empty, "Neo4j graph database is not empty" if __name__ == "__main__":