reactome_chatbot/export_csvs.py at main · reactome/reactome_chatbot · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
#!/usr/bin/env python3
import csv
import os
from pathlib import Path

from neo4j import Driver, GraphDatabase

NEO4J_URI = os.environ.get("NEO4J_URI", "bolt://localhost:7687")
NEO4J_USER = os.environ.get("NEO4J_USERNAME", "neo4j")
NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD")
if not NEO4J_PASSWORD:
    # Without this the driver constructs fine and fails later at connect time
    # with an error that does not mention the missing variable.
    raise SystemExit(
        "NEO4J_PASSWORD is not set. Export it before running this script; it is "
        "no longer hardcoded here."
    )

OUTPUT_DIR = Path("./embeddings/openai/bge-m3/plantreactome/Release68/csv_files")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

QUERIES = {
    "reactions": """
    MATCH (pathway:Pathway)-[:hasEvent]->(reaction:ReactionLikeEvent)
    OPTIONAL MATCH (reaction)-[:input]->(input:PhysicalEntity)
    OPTIONAL MATCH (reaction)-[:output]->(output:PhysicalEntity)
    OPTIONAL MATCH (reaction)-[:catalystActivity]->(cat:CatalystActivity)-[:physicalEntity]->(catalyst:PhysicalEntity)
    RETURN reaction.stId AS st_id, reaction.displayName AS display_name,
      pathway.stId AS pathway_id, pathway.displayName AS pathway_name,
      pathway.speciesName AS species,
      COLLECT(DISTINCT input.stId) AS input_id,
      COLLECT(DISTINCT input.displayName) AS input_name,
      COLLECT(DISTINCT output.stId) AS output_id,
      COLLECT(DISTINCT output.displayName) AS output_name,
      COLLECT(DISTINCT catalyst.stId) AS catalyst_id,
      COLLECT(DISTINCT catalyst.displayName) AS catalyst_name,
      "https://plantreactome.gramene.org/content/detail/" + reaction.stId AS url
    """,
    "summations": """
    MATCH (e)-[:summation]->(s:Summation)
    WHERE (e:Pathway OR e:ReactionLikeEvent)
    RETURN e.stId AS st_id, e.displayName AS display_name, labels(e) AS labels,
      e.speciesName AS species,
      CASE WHEN size(s.text) > 10000 THEN LEFT(s.text, 10000) + '...' ELSE s.text END AS summation,
      "https://plantreactome.gramene.org/content/detail/" + e.stId AS url
    """,
    "complexes": """
    MATCH (complex:Complex)-[:hasComponent]->(component)
    RETURN complex.speciesName AS species, complex.stId AS st_id,
      complex.name AS display_name, component.stId AS component_id, component.name AS component_name,
      "https://plantreactome.gramene.org/content/detail/" + complex.stId AS url
    """,
    "ewas": """
    MATCH (db:ReferenceDatabase)<-[:referenceDatabase]-(gene:ReferenceEntity)<-[:referenceEntity]-(prot:PhysicalEntity)
    RETURN DISTINCT
      prot.stId AS st_id,
      prot.displayName AS display_name,
      gene.geneName AS canonical_gene_name,
      '' AS synonyms_gene_name,
      gene.url AS uniprot_link,
      "https://plantreactome.gramene.org/content/detail/" + prot.stId AS url
    """,
}


def clean_value(v: object) -> str:
    if v is None:
        return ""
    if isinstance(v, list):
        return "|".join(str(x) for x in v)
    s = str(v).strip()
    if s.startswith("[") and s.endswith("]"):
        inner = s[1:-1]
        items = [i.strip().strip('"').strip("'") for i in inner.split(",") if i.strip()]
        return "|".join(items)
    return s


def run_query(driver: Driver, query: str) -> list[dict[str, str]]:
    with driver.session() as session:
        result = session.run(query)
        records = [r.data() for r in result]
    cleaned = []
    for row in records:
        new_row = {}
        for k, v in row.items():
            if k is None:
                continue
            new_row[k.strip()] = clean_value(v)
        cleaned.append(new_row)
    return cleaned


def main() -> None:
    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
    for name, query in QUERIES.items():
        print(f"Exporting {name}...")
        rows = run_query(driver, query)
        if not rows:
            print(f"  WARNING: No rows returned for {name}")
            continue
        fieldnames = list(rows[0].keys())
        outfile = OUTPUT_DIR / f"{name}.csv"
        with open(outfile, "w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(rows)
        print(f"  Saved {len(rows)} rows to {outfile}")
    driver.close()
    print("Done.")


if __name__ == "__main__":
    main()