Update genesis/graphdb.py

genesis/graphdb.py

@@ -1,81 +1,106 @@
 # genesis/graphdb.py
-from __future__ import annotations
-
 import os
-from typing import Dict, List
-
 from neo4j import GraphDatabase
+from typing import List, Dict, Optional
 
 NEO4J_URI = os.getenv("NEO4J_URI")
 NEO4J_USER = os.getenv("NEO4J_USER")
 NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")
 
+class GraphDB:
+    def __init__(self):
+        if not all([NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD]):
+            raise ValueError("Neo4j credentials are missing from environment variables.")
+        self.driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
+
+    def close(self):
+        self.driver.close()
+
+    # ---------- Basic Helpers ----------
+    def run_query(self, query: str, params: Optional[Dict] = None) -> List[Dict]:
+        with self.driver.session() as session:
+            result = session.run(query, params or {})
+            return [record.data() for record in result]
+
+    def create_node(self, label: str, properties: Dict):
+        props = {k: v for k, v in properties.items() if v is not None}
+        query = f"MERGE (n:{label} {{name: $name}}) SET n += $props RETURN n"
+        return self.run_query(query, {"name": props.get("name"), "props": props})
+
+    def create_relationship(self, start_label: str, start_name: str,
+                             rel_type: str, end_label: str, end_name: str):
+        query = f"""
+        MATCH (a:{start_label} {{name: $start_name}})
+        MATCH (b:{end_label} {{name: $end_name}})
+        MERGE (a)-[r:{rel_type}]->(b)
+        RETURN r
+        """
+        return self.run_query(query, {
+            "start_name": start_name,
+            "end_name": end_name
+        })
+
+    # ---------- Domain-Specific Functions ----------
+    def add_molecule(self, name: str, description: str = None):
+        return self.create_node("Molecule", {"name": name, "description": description})
+
+    def add_pathway(self, name: str, description: str = None):
+        return self.create_node("Pathway", {"name": name, "description": description})
+
+    def add_company(self, name: str, country: str = None, funding: float = None):
+        return self.create_node("Company", {"name": name, "country": country, "funding": funding})
+
+    def add_trial(self, trial_id: str, title: str = None, status: str = None):
+        return self.create_node("Trial", {"name": trial_id, "title": title, "status": status})
+
+    def add_regulation(self, region: str, title: str, summary: str = None):
+        return self.create_node("Regulation", {"name": title, "region": region, "summary": summary})
+
+    def add_biosecurity_alert(self, title: str, severity: str, details: str = None):
+        return self.create_node("BiosecurityAlert", {"name": title, "severity": severity, "details": details})
+
+    # ---------- Relationships ----------
+    def link_molecule_pathway(self, molecule: str, pathway: str):
+        return self.create_relationship("Molecule", molecule, "INVOLVED_IN", "Pathway", pathway)
+
+    def link_company_molecule(self, company: str, molecule: str):
+        return self.create_relationship("Company", company, "DEVELOPS", "Molecule", molecule)
+
+    def link_trial_molecule(self, trial_id: str, molecule: str):
+        return self.create_relationship("Trial", trial_id, "TESTS", "Molecule", molecule)
+
+    def link_company_funder(self, company: str, funder: str):
+        return self.create_relationship("Company", company, "FUNDED_BY", "Funder", funder)
+
+    def link_regulation_pathway(self, regulation: str, pathway: str):
+        return self.create_relationship("Regulation", regulation, "REGULATES", "Pathway", pathway)
+
+    def link_biosecurity_molecule(self, alert: str, molecule: str):
+        return self.create_relationship("BiosecurityAlert", alert, "ASSOCIATED_WITH", "Molecule", molecule)
+
+    # ---------- Queries ----------
+    def get_global_funding_network(self) -> List[Dict]:
+        query = """
+        MATCH (c:Company)-[r:FUNDED_BY]->(f:Funder)
+        RETURN c.name AS company, f.name AS funder, c.country AS country, c.funding AS amount
+        ORDER BY amount DESC
+        """
+        return self.run_query(query)
+
+    def get_pathway_network(self, pathway_name: str) -> List[Dict]:
+        query = """
+        MATCH (m:Molecule)-[:INVOLVED_IN]->(p:Pathway {name: $pathway})
+        OPTIONAL MATCH (m)<-[:DEVELOPS]-(c:Company)
+        RETURN m.name AS molecule, c.name AS company
+        """
+        return self.run_query(query, {"pathway": pathway_name})
+
+    def get_all_entities(self) -> List[Dict]:
+        query = """
+        MATCH (n) RETURN DISTINCT labels(n) AS labels, n.name AS name
+        """
+        return self.run_query(query)
+
 
-def _require_config():
-    if not (NEO4J_URI and NEO4J_USER and NEO4J_PASSWORD):
-        raise RuntimeError(
-            "Neo4j not configured. Set NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD in Space Secrets."
-        )
-
-
-async def write_topic_and_papers(topic: str, citations: List[Dict[str, str]]) -> Dict[str, int]:
-    """
-    Create/merge a Topic node and Paper nodes, and connect them with MENTIONS edges.
-    Schema:
-      (t:Topic {id})-[:MENTIONS]->(p:Paper {url})  with p.title set/updated.
-    Returns: {"nodes": <int>, "rels": <int>}
-    """
-    _require_config()
-    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
-
-    nodes = 0
-    rels = 0
-    try:
-        with driver.session() as sess:
-            # Topic
-            sess.run(
-                "MERGE (t:Topic {id:$id}) SET t.title=$title",
-                id=(topic or "Topic")[:200],
-                title=(topic or "Topic")[:500],
-            )
-            nodes += 1
-
-            # Papers + edges
-            for c in citations or []:
-                title = (c.get("title") or "citation")[:500]
-                url = (c.get("url") or "")[:1000]
-                if not url and not title:
-                    continue
-
-                sess.run(
-                    """
-                    MERGE (p:Paper {url:$url})
-                    ON CREATE SET p.title=$title
-                    ON MATCH SET  p.title = coalesce(p.title, $title)
-                    WITH p
-                    MATCH (t:Topic {id:$topic})
-                    MERGE (t)-[:MENTIONS]->(p)
-                    """,
-                    url=url,
-                    title=title,
-                    topic=(topic or "Topic")[:200],
-                )
-                nodes += 1
-                rels += 1
-
-        return {"nodes": nodes, "rels": rels}
-    finally:
-        driver.close()
-
-
-# Optional helper for debugging connectivity in a notebook/Space shell
-def verify_connection() -> str:
-    _require_config()
-    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
-    try:
-        with driver.session() as sess:
-            res = sess.run("RETURN 1 AS ok").single()
-            assert res and res["ok"] == 1
-        return "Neo4j connection OK."
-    finally:
-        driver.close()
+# Singleton instance for app-wide use
+graphdb = GraphDB()