Research-Stack/scripts/load_module_graph.py

#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "gremlinpython",
#   "python-dotenv",
# ]
# ///
"""
load_module_graph.py — Parse Lean import statements → Gremlin graph

Walks all .lean files under Semantics/Semantics/, extracts import lines,
creates vertices (modules) and edges (imports), loads into mathblob.

Run with:  uv run scripts/load_module_graph.py
Requires:  .env.gremlin (written by setup_mathblob.py)
"""

import os
import re
import time
from pathlib import Path
from dotenv import load_dotenv
from gremlin_python.driver import client as gremlin_client, serializer

# ── Config ────────────────────────────────────────────────────────────────────

ROOT     = Path(__file__).parent.parent
LEAN_DIR = ROOT / "0-Core-Formalism/lean/Semantics/Semantics"
ENV_FILE = ROOT / ".env.gremlin"

load_dotenv(ENV_FILE)

ENDPOINT = os.environ["GREMLIN_ENDPOINT"]
USERNAME = os.environ["GREMLIN_USERNAME"]
PASSWORD = os.environ["GREMLIN_PASSWORD"]

BATCH_SIZE = 50   # Cosmos DB free tier: small batches to avoid RU exhaustion

# ── Parse imports ─────────────────────────────────────────────────────────────

def parse_lean_file(path: Path) -> dict:
    """Return {module_id, label, imports: [str]} for a .lean file."""
    # Derive module id from filename: Semantics/Semantics/Foo.lean → Semantics.Foo
    rel = path.relative_to(ROOT / "0-Core-Formalism/lean/Semantics")
    module_id = ".".join(rel.with_suffix("").parts)  # e.g. Semantics.Foo

    imports = []
    try:
        text = path.read_text(errors="replace")
        for line in text.splitlines():
            line = line.strip()
            if not line.startswith("import "):
                break  # imports always at top; stop at first non-import non-blank
            if line == "import":
                continue
            target = line.removeprefix("import ").strip()
            if target:
                imports.append(target)
    except Exception as e:
        print(f"  WARN: could not read {path.name}: {e}")

    return {"id": module_id, "label": "module", "imports": imports}

def collect_graph() -> tuple[dict, list]:
    """Walk LEAN_DIR, return (vertices dict, edges list)."""
    files = sorted(LEAN_DIR.rglob("*.lean"))
    print(f"Scanning {len(files)} .lean files...")

    vertices = {}  # id → {id, label, kind}
    edges    = []  # (src_id, dst_id, label)

    for f in files:
        mod = parse_lean_file(f)
        mid = mod["id"]
        vertices[mid] = {"id": mid, "label": "module", "kind": "internal"}

        for imp in mod["imports"]:
            # Ensure target vertex exists (may be external like Mathlib)
            if imp not in vertices:
                kind = "external" if not imp.startswith("Semantics.") else "internal"
                vertices[imp] = {"id": imp, "label": "module", "kind": kind}
            edges.append((mid, imp, "imports"))

    print(f"  {len(vertices)} vertices, {len(edges)} edges")
    return vertices, edges

# ── Gremlin helpers ───────────────────────────────────────────────────────────

def make_client():
    return gremlin_client.Client(
        ENDPOINT, "g",
        username=USERNAME,
        password=PASSWORD,
        message_serializer=serializer.GraphSONSerializersV2d0(),
    )

def submit(c, query: str, bindings: dict = None):
    """Submit a single Gremlin query, return results."""
    try:
        cb = c.submitAsync(query, bindings or {})
        return cb.result().all().result()
    except Exception as e:
        print(f"  ERR: {e!s:.120}")
        return None

def upsert_vertex(c, v: dict):
    """Add vertex if it doesn't exist. Cosmos DB Gremlin requires a 'pk' property
    matching the /pk partition key path; 'id' is the vertex identifier."""
    q = (
        "g.V().has('module','id',vid).fold()"
        ".coalesce(unfold(),"
        "addV('module').property('id',vid).property('pk',vid).property('kind',kind))"
        ".property('kind',kind)"
    )
    submit(c, q, {"vid": v["id"], "kind": v["kind"]})

def upsert_edge(c, src: str, dst: str, label: str = "imports"):
    """Add edge if it doesn't exist."""
    q = (
        "g.V().has('module','id',src).as('s')"
        ".V().has('module','id',dst).as('d')"
        ".coalesce("
        "  select('s').outE(lbl).where(inV().as('d')),"
        "  addE(lbl).from('s').to('d')"
        ")"
    )
    submit(c, q, {"src": src, "dst": dst, "lbl": label})

# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    vertices, edges = collect_graph()

    print(f"\nConnecting to {ENDPOINT}...")
    c = make_client()

    # Smoke test
    count = submit(c, "g.V().count()")
    print(f"Current vertex count: {count}")

    # Load vertices
    print(f"\nLoading {len(vertices)} vertices (batch {BATCH_SIZE})...")
    vlist = list(vertices.values())
    for i in range(0, len(vlist), BATCH_SIZE):
        batch = vlist[i:i + BATCH_SIZE]
        for v in batch:
            upsert_vertex(c, v)
        done = min(i + BATCH_SIZE, len(vlist))
        print(f"  vertices {done}/{len(vlist)}", end="\r")
        time.sleep(0.1)  # gentle on free-tier RUs
    print(f"\n  vertices done.")

    # Load edges
    print(f"\nLoading {len(edges)} edges (batch {BATCH_SIZE})...")
    for i in range(0, len(edges), BATCH_SIZE):
        batch = edges[i:i + BATCH_SIZE]
        for src, dst, lbl in batch:
            upsert_edge(c, src, dst, lbl)
        done = min(i + BATCH_SIZE, len(edges))
        print(f"  edges {done}/{len(edges)}", end="\r")
        time.sleep(0.1)
    print(f"\n  edges done.")

    # Final count
    v_count = submit(c, "g.V().count()")
    e_count = submit(c, "g.E().count()")
    print(f"\nGraph loaded: {v_count} vertices, {e_count} edges")

    # Sample queries
    print("\n── Sample queries ──────────────────────────────────────────────")
    print("Most imported modules:")
    top = submit(c,
        "g.V().hasLabel('module').order().by(__.in('imports').count(), decr)"
        ".limit(10).project('name','in_degree')"
        ".by('id').by(__.in('imports').count())"
    )
    if top:
        for row in top:
            print(f"  {row.get('name','?'):50s}  ← {row.get('in_degree',0)}")

    print("\nDirect dependencies of HydrogenicPhiTorsionBraid:")
    deps = submit(c,
        "g.V().has('module','id','Semantics.HydrogenicPhiTorsionBraid')"
        ".out('imports').values('id')"
    )
    if deps:
        for d in deps:
            print(f"  {d}")

    c.close()
    print("\nDone.")

if __name__ == "__main__":
    main()