diff --git a/cognee/modules/visualization/cognee_network_visualization.py b/cognee/modules/visualization/cognee_network_visualization.py index 3ac596494..dde2fe98d 100644 --- a/cognee/modules/visualization/cognee_network_visualization.py +++ b/cognee/modules/visualization/cognee_network_visualization.py @@ -1,5 +1,6 @@ import os import json +import networkx from cognee.shared.logging_utils import get_logger from cognee.infrastructure.files.storage.LocalFileStorage import LocalFileStorage @@ -7,37 +8,7 @@ from cognee.infrastructure.files.storage.LocalFileStorage import LocalFileStorag logger = get_logger() -def _import_networkx(): - """Dynamically import networkx with helpful error message if not available.""" - try: - import networkx - return networkx - except ImportError: - raise ImportError( - "NetworkX is required for graph visualization but is not installed. " - "Please install it with: pip install 'cognee[visualization]' or pip install networkx" - ) - - async def cognee_network_visualization(graph_data, destination_file_path: str = None): - """ - Generate an interactive HTML visualization of the graph data. - - This function requires NetworkX to be installed. If you don't have NetworkX installed, - you can install it with: pip install 'cognee[visualization]' or pip install networkx - - Args: - graph_data: Tuple of (nodes_data, edges_data) - destination_file_path: Optional path to save the HTML file - - Returns: - str: Path to the generated HTML file - - Raises: - ImportError: If NetworkX is not installed - """ - networkx = _import_networkx() - nodes_data, edges_data = graph_data G = networkx.DiGraph() diff --git a/pyproject.toml b/pyproject.toml index c07790e9e..095349ff3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "pypdf>=4.1.0,<6.0.0", "jinja2>=3.1.3,<4", "matplotlib>=3.8.3,<4", + "networkx>=3.4.2,<4", "lancedb>=0.24.0,<1.0.0", "alembic>=1.13.3,<2", "pre-commit>=4.0.1,<5", @@ -119,7 +120,6 @@ gui = [ "qasync>=0.27.1,<0.28", ] graphiti = ["graphiti-core>=0.7.0,<0.8"] -visualization = ["networkx>=3.4.2,<4"] # Note: New s3fs and boto3 versions don't work well together # Always use comaptible fixed versions of these two dependencies aws = ["s3fs[boto3]==2025.3.2"]