Update genesis/visualization.py

genesis/visualization.py

@@ -1,8 +1,21 @@
 # genesis/visualization.py
+"""
+Visualization utilities for GENESIS-AI
+Generates interactive pathway and funding network graphs from Neo4j data.
+"""
+
 import os
-from neo4j import GraphDatabase
+import logging
+from typing import List, Tuple, Optional
+
 import plotly.graph_objects as go
 
+# Optional Neo4j import
+try:
+    from neo4j import GraphDatabase
+except ImportError:
+    GraphDatabase = None
+
 # =========================
 # CONFIGURATION
 # =========================
@@ -17,23 +30,32 @@ COLOR_EDGE = "#AAAAAA"
 # =========================
 # NEO4J CONNECTION
 # =========================
-driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
-
-def run_neo4j_query(query, params=None):
+driver = None
+if GraphDatabase and NEO4J_URI and NEO4J_USER and NEO4J_PASSWORD:
+    try:
+        driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
+        logging.info("[Neo4j] Connected for visualization.")
+    except Exception as e:
+        logging.error(f"[Neo4j] Connection failed: {e}")
+        driver = None
+else:
+    logging.info("[Neo4j] No URI/user/password set — skipping graph_tools connection.")
+
+def run_neo4j_query(query: str, params: dict = None) -> list:
     """Run a Cypher query and return results."""
+    if not driver:
+        logging.warning("[Neo4j] No active connection — returning empty result.")
+        return []
     with driver.session() as session:
         return list(session.run(query, params or {}))
 
 # =========================
 # GRAPH UTILS
 # =========================
-def create_plotly_graph(nodes, edges, title):
+def create_plotly_graph(nodes: List[Tuple[str, str, str]], edges: List[Tuple[str, str]], title: str):
     """Creates an interactive Plotly network graph."""
-    node_x, node_y, node_text, node_color = [], [], [], []
-    edge_x, edge_y = [], []
-
-    # Layout algorithm (spring layout for better spacing)
     import networkx as nx
+
     G = nx.Graph()
     for node_id, label, color in nodes:
         G.add_node(node_id, label=label, color=color)
@@ -42,12 +64,14 @@ def create_plotly_graph(nodes, edges, title):
 
     pos = nx.spring_layout(G, seed=42, k=0.5)
 
+    edge_x, edge_y = [], []
     for src, dst in edges:
         x0, y0 = pos[src]
         x1, y1 = pos[dst]
         edge_x.extend([x0, x1, None])
         edge_y.extend([y0, y1, None])
 
+    node_x, node_y, node_text, node_color = [], [], [], []
     for node_id, label, color in nodes:
         x, y = pos[node_id]
         node_x.append(x)
@@ -90,9 +114,10 @@ def create_plotly_graph(nodes, edges, title):
 # =========================
 # PATHWAY GRAPH
 # =========================
-def generate_pathway_graph(pathway_name):
+def generate_pathway_graph(pathway_name: str) -> Optional[go.Figure]:
     """
     Generate an interactive graph for a given metabolic pathway.
+    Only takes pathway_name — matches pipeline.py signature.
     """
     query = """
     MATCH (p:Pathway {name: $name})-[r:INVOLVES]->(m:Molecule)
@@ -105,9 +130,13 @@ def generate_pathway_graph(pathway_name):
 
     nodes = [(pathway_name, pathway_name, COLOR_PATHWAY_NODE)]
     edges = []
+    seen_nodes = {pathway_name}
+
     for record in results:
         mol_name = record["molecule"]
-        nodes.append((mol_name, mol_name, "#00BFFF"))  # Blue for molecules
+        if mol_name not in seen_nodes:
+            nodes.append((mol_name, mol_name, "#00BFFF"))  # Blue for molecules
+            seen_nodes.add(mol_name)
         edges.append((pathway_name, mol_name))
 
     return create_plotly_graph(nodes, edges, f"Metabolic Pathway: {pathway_name}")
@@ -115,7 +144,7 @@ def generate_pathway_graph(pathway_name):
 # =========================
 # FUNDING NETWORK
 # =========================
-def generate_funding_network(industry_keyword):
+def generate_funding_network(industry_keyword: str) -> Optional[go.Figure]:
     """
     Generate an interactive funding network graph for companies in a given biotech domain.
     """
@@ -152,4 +181,5 @@ def generate_funding_network(industry_keyword):
 # CLEANUP
 # =========================
 def close_driver():
-    driver.close()
+    if driver:
+        driver.close()