Update app.py

app.py

@@ -252,31 +252,50 @@ h1, h2, h3 { font-family: 'IBM Plex Sans', sans-serif; font-weight: 600; }
 # ─────────────────────────────────────────────
 # NEO4J CONNECTION
 # ─────────────────────────────────────────────
-@st.cache_resource
-def get_neo4j_driver():
-    # Try secrets first, then env vars, then hardcoded fallback
+def _get_neo4j_creds():
     uri      = (st.secrets.get("NEO4J_URI")      or os.getenv("NEO4J_URI")      or "neo4j+s://e3b8c8ec.databases.neo4j.io")
     user     = (st.secrets.get("NEO4J_USERNAME")  or os.getenv("NEO4J_USERNAME")  or "e3b8c8ec")
     password = (st.secrets.get("NEO4J_PASSWORD")  or os.getenv("NEO4J_PASSWORD")  or "")
     database = (st.secrets.get("NEO4J_DATABASE")  or os.getenv("NEO4J_DATABASE")  or "e3b8c8ec")
+    return (
+        uri.strip().strip('"').strip("'"),
+        user.strip().strip('"').strip("'"),
+        password.strip().strip('"').strip("'"),
+        database.strip().strip('"').strip("'"),
+    )
 
-    # Strip whitespace and accidental quotes
-    uri      = uri.strip().strip('"').strip("'")
-    user     = user.strip().strip('"').strip("'")
-    password = password.strip().strip('"').strip("'")
-    database = database.strip().strip('"').strip("'")
-
+@st.cache_resource
+def get_neo4j_driver():
+    uri, user, password, database = _get_neo4j_creds()
     if not password:
         st.sidebar.warning("NEO4J_PASSWORD not set in Secrets.")
-        return None
+        return None, None
     try:
-        driver = GraphDatabase.driver(uri, auth=(user, password))
+        driver = GraphDatabase.driver(
+            uri, auth=(user, password),
+            max_connection_lifetime=200,
+            keep_alive=True,
+        )
         driver.verify_connectivity()
-        driver._default_database = database
-        return driver
+        return driver, database
     except Exception as e:
         st.sidebar.error(f"Neo4j error: {e}")
-        return None
+        return None, None
+
+def get_session(driver, database):
+    """Always get a fresh session, reconnecting if needed."""
+    try:
+        return driver.session(database=database)
+    except Exception:
+        # Driver is defunct — clear cache and rebuild
+        get_neo4j_driver.clear()
+        uri, user, password, database = _get_neo4j_creds()
+        new_driver = GraphDatabase.driver(
+            uri, auth=(user, password),
+            max_connection_lifetime=200,
+            keep_alive=True,
+        )
+        return new_driver.session(database=database)
 
 @st.cache_resource
 def get_groq_client():
@@ -361,7 +380,7 @@ MERGE (a5)-[:TRANSACTION {amount:2.00, type:'card_test'}]->(m3)
 """
 
 def seed_database(driver):
-    with driver.session(database=getattr(driver, "_default_database", "neo4j")) as session:
+    with get_session(driver, "e3b8c8ec") as session:
         for stmt in SEED_CYPHER.strip().split(';'):
             stmt = stmt.strip()
             if stmt:
@@ -408,7 +427,8 @@ Respond ONLY with a valid JSON object, no markdown, no explanation:
 # ─────────────────────────────────────────────
 def run_cypher(driver, cypher: str):
     try:
-        with driver.session(database=getattr(driver, "_default_database", "neo4j")) as session:
+        _, database = get_neo4j_driver()
+        with get_session(driver, database or "e3b8c8ec") as session:
             result = session.run(cypher)
             records = [dict(r) for r in result]
             return records, None
@@ -504,7 +524,8 @@ def build_graph_html(driver, limit=60):
         "Customer": "#42A5F5", "Account": "#69F0AE", "Device": "#FFD54F",
         "IP": "#FF6B6B", "Merchant": "#CE93D8",
     }
-    with driver.session(database=getattr(driver, "_default_database", "neo4j")) as session:
+    _, database = get_neo4j_driver()
+    with get_session(driver, database or "e3b8c8ec") as session:
         nodes_q = "MATCH (n) RETURN n, labels(n) AS lbl LIMIT $limit"
         edges_q = "MATCH (a)-[r]->(b) RETURN id(a) AS src, id(b) AS tgt, type(r) AS rel, r LIMIT $limit"
 
@@ -558,10 +579,10 @@ with st.sidebar:
     st.markdown("---")
     st.markdown("**⚙️ Connections**")
 
-    driver = get_neo4j_driver()
+    driver, _db = get_neo4j_driver()
     groq_client = get_groq_client()
 
-    neo4j_ok = driver is not None
+    neo4j_ok = driver is not None and _db is not None
     groq_ok = groq_client is not None
 
     st.markdown(f"{'🟢' if neo4j_ok else '🔴'} Neo4j Aura {'Connected' if neo4j_ok else 'Not connected'}")