diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 5672bf0df2..a33ec7ceaf 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -26,6 +26,7 @@ Each stage is a single function in its own module. They communicate through plai | `cache.py` | `check_semantic_cache / save_semantic_cache` | files → (cached, uncached) split | | `security.py` | validation helpers | URL / path / label → validated or raises | | `validate.py` | `validate_extraction(data)` | extraction dict → raises on schema errors | +| `storage.py` | `init_db / ingest_extraction / ingest_communities` | extraction dict → NeuG `graph.db` (optional, requires `neug`) | | `serve.py` | `start_server(graph_path)` | graph file path → MCP stdio server | | `watch.py` | `watch(root, flag_path)` | directory → writes flag file on change | | `benchmark.py` | `run_benchmark(graph_path)` | graph file → corpus vs subgraph token comparison | diff --git a/README.md b/README.md index 465a5bb054..739c3f9b75 100644 --- a/README.md +++ b/README.md @@ -164,6 +164,7 @@ Install only what you need: | `video` | Video/audio transcription (faster-whisper + yt-dlp) | `uv tool install "graphifyy[video]"` | | `mcp` | MCP stdio server | `uv tool install "graphifyy[mcp]"` | | `neo4j` | Neo4j push support | `uv tool install "graphifyy[neo4j]"` | +| `neug` | [NeuG](https://github.com/alibaba/neug) embedded graph database — Cypher queries on your graph | `uv tool install "graphifyy[neug]"` | | `svg` | SVG graph export | `uv tool install "graphifyy[svg]"` | | `leiden` | Leiden community detection (Python < 3.13 only) | `uv tool install "graphifyy[leiden]"` | | `ollama` | Ollama local inference | `uv tool install "graphifyy[ollama]"` | @@ -503,6 +504,9 @@ graphify install # overwrites the skill file /graphify ./raw --graphml # export for Gephi / yEd /graphify ./raw --neo4j # generate cypher.txt for Neo4j /graphify ./raw --neo4j-push bolt://localhost:7687 + +graphify cypher "MATCH (n) RETURN n LIMIT 10" # query graph.db with Cypher (requires neug) +graphify cypher "MATCH (n:code)-[e]->(m) RETURN n.id, e, m.id LIMIT 10" --db path/to/graph.db # default: graphify-out/graph.db /graphify ./raw --watch # auto-sync as files change /graphify ./raw --mcp # start MCP stdio server diff --git a/graphify/__main__.py b/graphify/__main__.py index 759d913b91..81784a09ad 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -2133,6 +2133,8 @@ def main() -> None: print(" --backend= backend to use for community naming (default: auto-detect)") print(" label (re)name communities with the configured LLM backend, regenerate report") print(" --backend= backend to use (default: auto-detect from API keys)") + print(" cypher \"MATCH ...\" execute a Cypher query against graph.db (requires neug)") + print(" --db path to graph.db (default graphify-out/graph.db)") print(" query \"\" BFS traversal of graph.json for a question") print(" --dfs use depth-first instead of breadth-first") print(" --context C explicit edge-context filter (repeatable)") @@ -2641,6 +2643,31 @@ def main() -> None: else: print("Usage: graphify hook [install|uninstall|status]", file=sys.stderr) sys.exit(1) + elif cmd == "cypher": + if len(sys.argv) < 3: + print('Usage: graphify cypher "MATCH ..." [--db path]', file=sys.stderr) + sys.exit(1) + query_str = sys.argv[2] + db_path = str(Path(_GRAPHIFY_OUT) / "graph.db") + args = sys.argv[3:] + for i, a in enumerate(args): + if a == "--db" and i + 1 < len(args): + db_path = args[i + 1] + try: + from graphify.storage import init_db, execute_cypher, close_db + except ImportError: + print("error: neug is not installed. Run: pip install neug", file=sys.stderr) + sys.exit(1) + if not Path(db_path).exists(): + print(f"error: database not found: {db_path}", file=sys.stderr) + sys.exit(1) + db, conn = init_db(db_path) + try: + results = execute_cypher(conn, query_str) + for row in results: + print("\t".join(str(v) for v in row)) + finally: + close_db(db, conn) elif cmd == "query": if len(sys.argv) < 3: print("Usage: graphify query \"\" [--dfs] [--context C] [--budget N] [--graph path]", file=sys.stderr) @@ -4302,6 +4329,21 @@ def _progress(idx: int, total: int, _result: dict) -> None: graph_json_path.write_text( json.dumps(merged, indent=2), encoding="utf-8" ) + try: + from graphify.storage import init_db as _init_db, ensure_schema as _ensure_schema, ingest_extraction as _ingest, close_db as _close_db + _db_path = str(graphify_out / "graph.db") + _is_inc = Path(_db_path).exists() + _db, _conn = _init_db(_db_path) + _known = _ensure_schema(_conn, create_tables=not _is_inc) + _ingest(_conn, merged, incremental=_is_inc, + prune_sources=deleted_files or None, root=target, + known_tables=_known) + _close_db(_db, _conn) + print("[graphify extract] graph.db written (powered by NeuG)") + except ImportError: + pass + except Exception as _exc: + print(f"[graphify extract] warning: NeuG write failed: {_exc}", file=sys.stderr) cost = _estimate_cost( backend, merged["input_tokens"], merged["output_tokens"] ) @@ -4379,6 +4421,22 @@ def _progress(idx: int, total: int, _result: dict) -> None: from graphify.export import backup_if_protected as _backup _backup(graphify_out) _to_json(G, communities, str(graph_json_path), force=True) + try: + from graphify.storage import init_db as _init_db, ensure_schema as _ensure_schema, ingest_extraction as _ingest, ingest_communities as _ingest_comm, close_db as _close_db + _db_path = str(graphify_out / "graph.db") + _is_inc = Path(_db_path).exists() + _db, _conn = _init_db(_db_path) + _known = _ensure_schema(_conn, create_tables=not _is_inc) + _ntypes = _ingest(_conn, merged, incremental=_is_inc, + prune_sources=deleted_files or None, root=target, + known_tables=_known) + _ingest_comm(_conn, communities, node_types=_ntypes) + _close_db(_db, _conn) + print("[graphify extract] graph.db written (powered by NeuG)") + except ImportError: + pass + except Exception as _exc: + print(f"[graphify extract] warning: NeuG write failed: {_exc}", file=sys.stderr) if merged.get("output_tokens", 0) > 0: (graphify_out / ".graphify_semantic_marker").write_text( json.dumps({"output_tokens": merged["output_tokens"]}), encoding="utf-8" diff --git a/graphify/serve.py b/graphify/serve.py index 6e05b6a98e..d437893342 100644 --- a/graphify/serve.py +++ b/graphify/serve.py @@ -524,6 +524,20 @@ def _build_server(graph_path: str): G = _load_graph(graph_path) communities = _communities_from_graph(G) + _neug_conn = None + _neug_db = None + _neug_execute = None + try: + from graphify.storage import init_db as _neug_init, execute_cypher as _neug_exec, close_db as _neug_close + _neug_db_path = str(Path(graph_path).parent / "graph.db") + if Path(_neug_db_path).exists(): + _neug_db, _neug_conn = _neug_init(_neug_db_path) + _neug_execute = _neug_exec + except ImportError: + pass + except Exception: + pass + # Hot-reload state: mtime+size key lets us detect graph.json changes without # polling. Initialised from the file stat at startup so the first tool call # never triggers a redundant reload. @@ -683,6 +697,20 @@ async def list_tools() -> list[types.Tool]: }, }, ), + types.Tool( + name="cypher_query", + description=( + "Execute a Cypher query against the NeuG graph database. " + "Returns tabular results. Requires neug to be installed and graph.db to exist." + ), + inputSchema={ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Cypher query string"}, + }, + "required": ["query"], + }, + ), ] def _tool_query_graph(arguments: dict) -> str: @@ -933,6 +961,22 @@ def _tool_triage_prs(arguments: dict) -> str: ) return "\n\n".join(lines) + def _tool_cypher_query(arguments: dict) -> str: + if _neug_conn is None: + return "NeuG not available (not installed or graph.db not found)." + query = arguments["query"] + from graphify.storage import execute_cypher as _exec_cypher + try: + results = _exec_cypher(_neug_conn, query) + except RuntimeError as exc: + return f"Cypher error: {exc}" + if not results: + return "No results." + lines = [] + for row in results: + lines.append("\t".join(str(v) for v in row)) + return "\n".join(lines) + _handlers = { "query_graph": _tool_query_graph, "get_node": _tool_get_node, @@ -944,6 +988,7 @@ def _tool_triage_prs(arguments: dict) -> str: "list_prs": _tool_list_prs, "get_pr_impact": _tool_get_pr_impact, "triage_prs": _tool_triage_prs, + "cypher_query": _tool_cypher_query, } def _load_community_labels() -> dict[int, str]: diff --git a/graphify/storage.py b/graphify/storage.py new file mode 100644 index 0000000000..07aec39afc --- /dev/null +++ b/graphify/storage.py @@ -0,0 +1,540 @@ +"""NeuG graph database adapter for graphify. + +Provides an optional parallel storage engine alongside NetworkX. +NeuG is lazily imported — when not installed, callers should catch +ImportError at the call site and skip silently. + +All property values interpolated into Cypher statements use NeuG's native +parameterised queries ($param syntax) to prevent injection. Table/label +names (which come from a fixed internal set, not user input) are still +interpolated as identifiers. +""" +from __future__ import annotations + +import csv +import os +import re +import tempfile +from pathlib import Path + +from .build import _FILE_TYPE_SYNONYMS, _normalize_id, _norm_source_file +from .validate import VALID_FILE_TYPES + +# --------------------------------------------------------------------------- +# Node tables (one per file_type) +# --------------------------------------------------------------------------- + +_NODE_TABLES = { + "code": """CREATE NODE TABLE IF NOT EXISTS code ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, source_location STRING, community INT64)""", + "document": """CREATE NODE TABLE IF NOT EXISTS document ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", + "paper": """CREATE NODE TABLE IF NOT EXISTS paper ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", + "image": """CREATE NODE TABLE IF NOT EXISTS image ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", + "concept": """CREATE NODE TABLE IF NOT EXISTS concept ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", + "rationale": """CREATE NODE TABLE IF NOT EXISTS rationale ( + id STRING PRIMARY KEY, label STRING, + source_file STRING, community INT64)""", +} + +_NODE_COLUMNS = { + "code": ["id", "label", "source_file", "source_location", "community"], + "document": ["id", "label", "source_file", "community"], + "paper": ["id", "label", "source_file", "community"], + "image": ["id", "label", "source_file", "community"], + "concept": ["id", "label", "source_file", "community"], + "rationale": ["id", "label", "source_file", "community"], +} + +_EDGE_COLUMNS = ["from_id", "to_id", "relation", "confidence", + "confidence_score", "source_file", "weight"] + +# --------------------------------------------------------------------------- +# Edge tables — split by (src_type, tgt_type, relation). +# --------------------------------------------------------------------------- + +_EDGE_DDL_TEMPLATE = """CREATE REL TABLE IF NOT EXISTS {tbl}( + FROM {src} TO {tgt}, + relation STRING, confidence STRING, + confidence_score DOUBLE, source_file STRING, weight DOUBLE)""" + +# Known relation types per (src, tgt) pair — pre-built at init time. +_KNOWN_RELATIONS: dict[tuple[str, str], list[str]] = { + ("code", "code"): [ + "calls", "contains", "method", "uses", "inherits", "defines", + "references", "imports", "imports_from", "listened_by", "case_of", + "references_constant", "bound_to", "uses_static_prop", "uses_config", + ], + ("rationale", "code"): ["rationale_for"], +} + + +def _sanitize_rel_name(relation: str) -> str: + """Normalize a relation string into a safe table-name suffix.""" + r = relation.lower().strip() + r = re.sub(r"[^a-z0-9_]", "_", r) + r = re.sub(r"_+", "_", r).strip("_") + return r or "rel" + + +def _edge_table_name(src_type: str, tgt_type: str, relation: str) -> str: + return f"edge_{src_type}_{tgt_type}_{_sanitize_rel_name(relation)}" + + +# --------------------------------------------------------------------------- +# CSV helpers for bulk COPY FROM +# --------------------------------------------------------------------------- + +def _sanitize_csv_value(v: object) -> str: + if isinstance(v, str): + return v.replace("\n", "\\n").replace("\r", "") + return str(v) + + +def _write_csv(path: str, rows: list[dict], columns: list[str]) -> int: + if not rows: + return 0 + with open(path, "w", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=columns, extrasaction="ignore", + quoting=csv.QUOTE_ALL) + w.writeheader() + for row in rows: + w.writerow({k: _sanitize_csv_value(row.get(k, "")) for k in columns}) + return len(rows) + + +def _copy_node_csv(conn: object, csv_path: str, table: str) -> None: + conn.execute( + f'COPY {table} FROM "{csv_path}" (header=true, delim=",", escaping=false)' + ) + + +def _copy_rel_csv(conn: object, csv_path: str, tbl: str, + src_table: str, tgt_table: str) -> None: + conn.execute( + f'COPY {tbl} FROM "{csv_path}" ' + f'(from="{src_table}", to="{tgt_table}", ' + f'header=true, delim=",", escaping=false)' + ) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def init_db(db_path: str) -> tuple: + """Open (or create) a NeuG database and connect. + + Returns (db, conn). Raises ImportError if neug is not installed. + """ + import neug + db = neug.Database(db_path) + conn = db.connect() + return db, conn + + +def ensure_schema(conn: object, *, create_tables: bool = True) -> set[str]: + """Populate known table registry; optionally execute DDL. + + create_tables=True (first build): run CREATE TABLE statements. + create_tables=False (incremental): only build the registry set + so _ensure_rel_table() knows what exists. + + Returns the set of known rel table names (per-connection registry). + """ + created: set[str] = set() + + if create_tables: + for ddl in _NODE_TABLES.values(): + conn.execute(ddl) + + for (src, tgt), rels in _KNOWN_RELATIONS.items(): + for rel in rels: + tbl = _edge_table_name(src, tgt, rel) + if create_tables: + conn.execute(_EDGE_DDL_TEMPLATE.format(tbl=tbl, src=src, tgt=tgt)) + created.add(tbl) + + return created + + +def _ensure_rel_table( + conn: object, src_type: str, tgt_type: str, relation: str, + known: set[str], +) -> str: + """Resolve edge table name, creating on-the-fly if needed. Returns table name.""" + tbl = _edge_table_name(src_type, tgt_type, relation) + if tbl in known: + return tbl + conn.execute(_EDGE_DDL_TEMPLATE.format(tbl=tbl, src=src_type, tgt=tgt_type)) + known.add(tbl) + return tbl + + +def _fix_file_type(ft: str | None) -> str: + """Canonicalize file_type, matching build.py:138-146 logic.""" + if not ft or ft not in VALID_FILE_TYPES: + return _FILE_TYPE_SYNONYMS.get(ft, "concept") if ft else "concept" + return ft + + +def _bulk_ingest( + conn: object, + extraction: dict, + *, + root: str | None = None, + known_tables: set[str] | None = None, +) -> dict[str, str]: + """Full build via COPY FROM — much faster than per-row Cypher CREATE.""" + _known = known_tables if known_tables is not None else set() + nodes = extraction.get("nodes") or [] + edges = extraction.get("edges") or [] + + # --- collect node rows grouped by file_type --- + node_types: dict[str, str] = {} + node_buckets: dict[str, list[dict]] = {ft: [] for ft in _NODE_TABLES} + written_ids: set[str] = set() + + for node in nodes: + nid = _normalize_id(node.get("id", "")) + if not nid or nid in written_ids: + continue + written_ids.add(nid) + ft = _fix_file_type(node.get("file_type")) + node_types[nid] = ft + row: dict = { + "id": nid, + "label": node.get("label", ""), + "source_file": _norm_source_file(node.get("source_file"), root) or "", + "community": 0, + } + if ft == "code": + row["source_location"] = node.get("source_location") or "" + node_buckets.setdefault(ft, []).append(row) + + # --- collect edge rows grouped by rel table --- + edge_buckets: dict[str, list[dict]] = {} + edge_table_types: dict[str, tuple[str, str]] = {} + + for edge in edges: + src_id = _normalize_id(edge.get("source") or edge.get("from", "")) + tgt_id = _normalize_id(edge.get("target") or edge.get("to", "")) + if not src_id or not tgt_id: + continue + src_ft = node_types.get(src_id) + tgt_ft = node_types.get(tgt_id) + if not src_ft or not tgt_ft: + continue + + rel_raw = edge.get("relation", "") + tbl = _ensure_rel_table(conn, src_ft, tgt_ft, rel_raw, _known) + edge_table_types[tbl] = (src_ft, tgt_ft) + edge_buckets.setdefault(tbl, []).append({ + "from_id": src_id, + "to_id": tgt_id, + "relation": rel_raw, + "confidence": edge.get("confidence", ""), + "confidence_score": float(edge.get("confidence_score", 0.0)), + "source_file": _norm_source_file(edge.get("source_file"), root) or "", + "weight": float(edge.get("weight", 1.0)), + }) + + # --- write CSV + COPY FROM in a temp dir --- + tmp_dir = tempfile.mkdtemp(prefix="graphify_bulk_") + try: + for ft, rows in node_buckets.items(): + if not rows: + continue + csv_path = os.path.join(tmp_dir, f"node_{ft}.csv") + _write_csv(csv_path, rows, _NODE_COLUMNS[ft]) + _copy_node_csv(conn, csv_path, ft) + + for tbl, rows in edge_buckets.items(): + if not rows: + continue + csv_path = os.path.join(tmp_dir, f"edge_{tbl}.csv") + _write_csv(csv_path, rows, _EDGE_COLUMNS) + src_ft, tgt_ft = edge_table_types[tbl] + _copy_rel_csv(conn, csv_path, tbl, src_ft, tgt_ft) + finally: + import shutil + shutil.rmtree(tmp_dir, ignore_errors=True) + + return node_types + + +def _incremental_ingest( + conn: object, + extraction: dict, + *, + prune_sources: list[str] | None = None, + root: str | None = None, + known_tables: set[str] | None = None, +) -> dict[str, str]: + """Incremental update via DELETE affected source_files + COPY FROM. + + Much faster than per-row MERGE: deletes nodes whose source_file appears + in the incoming extraction (or in prune_sources), then bulk-inserts the + new data via COPY FROM. Incoming cross-file edges (from unchanged files + into affected nodes) are saved before deletion and restored afterwards. + """ + _known = known_tables if known_tables is not None else set() + nodes = extraction.get("nodes") or [] + edges = extraction.get("edges") or [] + + # --- collect affected source_files from the incoming data --- + affected_sfs: set[str] = set() + if prune_sources: + for sf in prune_sources: + sf_norm = _norm_source_file(sf, root) or sf + affected_sfs.add(sf_norm) + + node_types: dict[str, str] = {} + node_buckets: dict[str, list[dict]] = {ft: [] for ft in _NODE_TABLES} + written_ids: set[str] = set() + + for node in nodes: + nid = _normalize_id(node.get("id", "")) + if not nid or nid in written_ids: + continue + written_ids.add(nid) + ft = _fix_file_type(node.get("file_type")) + node_types[nid] = ft + sf = _norm_source_file(node.get("source_file"), root) or "" + if sf: + affected_sfs.add(sf) + row: dict = { + "id": nid, + "label": node.get("label", ""), + "source_file": sf, + "community": 0, + } + if ft == "code": + row["source_location"] = node.get("source_location") or "" + node_buckets.setdefault(ft, []).append(row) + + # --- resolve types for non-delta edge endpoints (before DELETE) --- + unknown_ids: set[str] = set() + for edge in edges: + for key in ("source", "from", "target", "to"): + eid = _normalize_id(edge.get(key, "")) + if eid and eid not in node_types: + unknown_ids.add(eid) + for nid in unknown_ids: + for tbl in _NODE_TABLES: + try: + rows = list(conn.execute( + f"MATCH (n:{tbl} {{id: $nid}}) RETURN 1", + parameters={"nid": nid}, + )) + if rows: + node_types[nid] = tbl + break + except RuntimeError: + pass + + # --- save incoming cross-file edges before DELETE --- + # Collect IDs of nodes that will be deleted. + affected_node_ids: set[str] = set() + for sf in affected_sfs: + for tbl in _NODE_TABLES: + try: + for row in conn.execute( + f"MATCH (n:{tbl}) WHERE n.source_file = $sf RETURN n.id", + parameters={"sf": sf}, + ): + affected_node_ids.add(row[0]) + except RuntimeError: + pass + + # For each known edge table, find edges where the target is in an affected + # source_file but the source is NOT (incoming from unchanged files). + saved_edge_buckets: dict[str, list[dict]] = {} + saved_edge_types: dict[str, tuple[str, str]] = {} + + for tbl in list(_known): + parts = tbl.split("_", 3) + if len(parts) < 4 or parts[0] != "edge": + continue + src_type, tgt_type = parts[1], parts[2] + + for sf in affected_sfs: + try: + rows = list(conn.execute( + f"MATCH (a:{src_type})-[e:{tbl}]->(b:{tgt_type}) " + f"WHERE b.source_file = $sf " + f"RETURN a.id, b.id, e.relation, e.confidence, " + f"e.confidence_score, e.source_file, e.weight", + parameters={"sf": sf}, + )) + except RuntimeError: + continue + + for row in rows: + if row[0] in affected_node_ids: + continue + saved_edge_types[tbl] = (src_type, tgt_type) + saved_edge_buckets.setdefault(tbl, []).append({ + "from_id": row[0], "to_id": row[1], + "relation": row[2] or "", + "confidence": row[3] or "", + "confidence_score": float(row[4] or 0.0), + "source_file": row[5] or "", + "weight": float(row[6] or 1.0), + }) + + # --- DELETE nodes from affected source_files --- + for sf in affected_sfs: + for tbl in _NODE_TABLES: + conn.execute( + f"MATCH (n:{tbl}) WHERE n.source_file = $sf DETACH DELETE n", + parameters={"sf": sf}, + ) + + # --- collect delta edge rows --- + edge_buckets: dict[str, list[dict]] = {} + edge_table_types: dict[str, tuple[str, str]] = {} + + for edge in edges: + src_id = _normalize_id(edge.get("source") or edge.get("from", "")) + tgt_id = _normalize_id(edge.get("target") or edge.get("to", "")) + if not src_id or not tgt_id: + continue + src_ft = node_types.get(src_id) + tgt_ft = node_types.get(tgt_id) + if not src_ft or not tgt_ft: + continue + + rel_raw = edge.get("relation", "") + tbl = _ensure_rel_table(conn, src_ft, tgt_ft, rel_raw, _known) + edge_table_types[tbl] = (src_ft, tgt_ft) + edge_buckets.setdefault(tbl, []).append({ + "from_id": src_id, + "to_id": tgt_id, + "relation": rel_raw, + "confidence": edge.get("confidence", ""), + "confidence_score": float(edge.get("confidence_score", 0.0)), + "source_file": _norm_source_file(edge.get("source_file"), root) or "", + "weight": float(edge.get("weight", 1.0)), + }) + + # --- merge saved incoming edges back --- + for tbl, rows in saved_edge_buckets.items(): + edge_buckets.setdefault(tbl, []).extend(rows) + if tbl not in edge_table_types: + edge_table_types[tbl] = saved_edge_types[tbl] + + # --- COPY FROM bulk insert --- + tmp_dir = tempfile.mkdtemp(prefix="graphify_inc_") + try: + for ft, rows in node_buckets.items(): + if not rows: + continue + csv_path = os.path.join(tmp_dir, f"node_{ft}.csv") + _write_csv(csv_path, rows, _NODE_COLUMNS[ft]) + _copy_node_csv(conn, csv_path, ft) + + for tbl, rows in edge_buckets.items(): + if not rows: + continue + csv_path = os.path.join(tmp_dir, f"edge_{tbl}.csv") + _write_csv(csv_path, rows, _EDGE_COLUMNS) + src_ft, tgt_ft = edge_table_types[tbl] + _copy_rel_csv(conn, csv_path, tbl, src_ft, tgt_ft) + finally: + import shutil + shutil.rmtree(tmp_dir, ignore_errors=True) + + return node_types + + +def ingest_extraction( + conn: object, + extraction: dict, + *, + incremental: bool = False, + prune_sources: list[str] | None = None, + root: str | Path | None = None, + known_tables: set[str] | None = None, +) -> dict[str, str]: + """Write an extraction dict into NeuG. + + incremental=False: first build — uses COPY FROM bulk loading. + incremental=True: update — uses MERGE (upsert) per row. + + Returns node_types dict (id -> file_type) for use by ingest_communities. + """ + _root = str(Path(root).resolve()) if root else None + + if incremental: + return _incremental_ingest( + conn, extraction, + prune_sources=prune_sources, root=_root, + known_tables=known_tables, + ) + else: + return _bulk_ingest( + conn, extraction, + root=_root, known_tables=known_tables, + ) + + +def ingest_communities( + conn: object, + communities: dict[int, list[str]], + community_labels: dict[int, str] | None = None, + node_types: dict[str, str] | None = None, +) -> None: + """Write community assignments into NeuG node properties. + + If node_types is provided (id -> file_type mapping from ingest_extraction), + each node is looked up in its specific table directly. Otherwise falls + back to probing all 6 tables (slower). + + Note: NeuG does not support parameterised SET for non-string values, + so community ID is interpolated as an integer literal. The id value + uses a parameterised query. + """ + for cid, node_ids in communities.items(): + cid_int = int(cid) + for nid in node_ids: + nid_norm = _normalize_id(nid) + if not nid_norm: + continue + if node_types and nid_norm in node_types: + tbl = node_types[nid_norm] + conn.execute( + f"MATCH (n:{tbl}) WHERE n.id = $nid " + f"SET n.community = {cid_int}", + parameters={"nid": nid_norm}, + ) + else: + for tbl in _NODE_TABLES: + conn.execute( + f"MATCH (n:{tbl}) WHERE n.id = $nid " + f"SET n.community = {cid_int}", + parameters={"nid": nid_norm}, + ) + + +def execute_cypher(conn: object, query: str) -> list[list]: + """Execute a Cypher query and return results as list of lists.""" + try: + return list(conn.execute(query)) + except RuntimeError as exc: + raise RuntimeError(f"Cypher query failed: {exc}") from exc + + +def close_db(db: object, conn: object) -> None: + """Close the NeuG connection and database.""" + conn.close() + db.close() diff --git a/pyproject.toml b/pyproject.toml index 3653bae52c..ecacc1faf9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,12 +66,13 @@ gemini = ["openai", "tiktoken"] openai = ["openai", "tiktoken"] chinese = ["jieba"] sql = ["tree-sitter-sql"] +neug = ["neug>=0.1.2,<0.2"] # tree-sitter-dm (BYOND DreamMaker) ships only a Windows wheel, so on Linux/Mac it # must compile from source (needs a C toolchain + python3-dev). Keeping it optional # avoids breaking the default `uv tool install graphifyy` for everyone (#1104). dm = ["tree-sitter-dm"] terraform = ["tree-sitter-hcl"] -all = ["mcp", "neo4j", "pypdf", "markdownify", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper; python_version >= '3.11'", "yt-dlp", "matplotlib", "numpy>=2.0; python_version >= '3.13'", "openai", "tiktoken", "boto3", "anthropic", "tree-sitter-sql", "jieba", "tree-sitter-dm", "tree-sitter-hcl"] +all = ["mcp", "neo4j", "neug>=0.1.2,<0.2", "pypdf", "markdownify", "watchdog", "graspologic; python_version < '3.13'", "python-docx", "openpyxl", "faster-whisper; python_version >= '3.11'", "yt-dlp", "matplotlib", "numpy>=2.0; python_version >= '3.13'", "openai", "tiktoken", "boto3", "anthropic", "tree-sitter-sql", "jieba", "tree-sitter-dm", "tree-sitter-hcl"] [project.scripts] graphify = "graphify.__main__:main" diff --git a/tests/test_cypher_cli.py b/tests/test_cypher_cli.py new file mode 100644 index 0000000000..8e18a31250 --- /dev/null +++ b/tests/test_cypher_cli.py @@ -0,0 +1,59 @@ +"""Tests for the `graphify cypher` CLI command.""" +import json +import subprocess +import sys +import tempfile +from pathlib import Path + +import pytest + +try: + import neug + _has_neug = True +except ImportError: + _has_neug = False + +pytestmark = pytest.mark.skipif(not _has_neug, reason="neug not installed") + +FIXTURES = Path(__file__).parent / "fixtures" +EXTRACTION_JSON = FIXTURES / "extraction.json" + + +def _build_db(tmp_path) -> str: + from graphify.storage import init_db, ensure_schema, ingest_extraction, close_db + db_path = str(tmp_path / "graph.db") + ext = json.loads(EXTRACTION_JSON.read_text()) + db, conn = init_db(db_path) + known = ensure_schema(conn) + ingest_extraction(conn, ext, incremental=False, known_tables=known) + close_db(db, conn) + return db_path + + +def test_cypher_command_basic(tmp_path): + db_path = _build_db(tmp_path) + result = subprocess.run( + [sys.executable, "-m", "graphify", "cypher", + "MATCH (n:code) RETURN count(n)", "--db", db_path], + capture_output=True, text=True, timeout=30, + ) + assert result.returncode == 0 + assert "3" in result.stdout + + +def test_cypher_command_db_not_found(tmp_path): + result = subprocess.run( + [sys.executable, "-m", "graphify", "cypher", + "MATCH (n) RETURN n", "--db", str(tmp_path / "nonexistent.db")], + capture_output=True, text=True, timeout=30, + ) + assert result.returncode != 0 + assert "not found" in result.stderr.lower() or "error" in result.stderr.lower() + + +def test_cypher_command_no_query(): + result = subprocess.run( + [sys.executable, "-m", "graphify", "cypher"], + capture_output=True, text=True, timeout=30, + ) + assert result.returncode != 0 diff --git a/tests/test_storage.py b/tests/test_storage.py new file mode 100644 index 0000000000..ad6bcfd02d --- /dev/null +++ b/tests/test_storage.py @@ -0,0 +1,180 @@ +"""Tests for graphify.storage — NeuG adapter layer.""" +import json +import shutil +import tempfile +from pathlib import Path + +import pytest + +try: + import neug + _has_neug = True +except ImportError: + _has_neug = False + +pytestmark = pytest.mark.skipif(not _has_neug, reason="neug not installed") + +FIXTURES = Path(__file__).parent / "fixtures" +EXTRACTION_JSON = FIXTURES / "extraction.json" + + +def _load_extraction() -> dict: + return json.loads(EXTRACTION_JSON.read_text()) + + +@pytest.fixture() +def tmp_db(tmp_path): + db_path = str(tmp_path / "test.db") + yield db_path + + +def _init(db_path): + from graphify.storage import init_db, ensure_schema + db, conn = init_db(db_path) + ensure_schema(conn) + return db, conn + + +def _close(db, conn): + from graphify.storage import close_db + close_db(db, conn) + + +def _query(conn, cypher): + from graphify.storage import execute_cypher + return execute_cypher(conn, cypher) + + +# --- init_db --- + +def test_init_db_creates_tables(tmp_db): + db, conn = _init(tmp_db) + for tbl in ("code", "document", "paper", "image", "concept", "rationale"): + rows = _query(conn, f"MATCH (n:{tbl}) RETURN count(n)") + assert rows == [[0]] + _close(db, conn) + + +# --- ingest_extraction: CREATE mode --- + +def test_ingest_extraction_create_mode(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + rows = _query(conn, "MATCH (n:code) RETURN n.id ORDER BY n.id") + ids = sorted([r[0] for r in rows]) + assert "n_attention" in ids + assert "n_transformer" in ids + assert "n_layernorm" in ids + edge_rows = _query(conn, "MATCH (a:code)-[e:edge_code_code_contains]->(b:code) RETURN count(e)") + assert edge_rows[0][0] == 2 + _close(db, conn) + + +# --- ingest_extraction: MERGE mode --- + +def test_ingest_extraction_merge_mode(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + ext["nodes"][0]["label"] = "TransformerV2" + ingest_extraction(conn, ext, incremental=True) + rows = _query(conn, "MATCH (n:code) WHERE n.id = 'n_transformer' RETURN n.label") + assert rows[0][0] == "TransformerV2" + count = _query(conn, "MATCH (n:code) RETURN count(n)") + assert count[0][0] == 3 + _close(db, conn) + + +# --- file_type routing --- + +def test_ingest_extraction_file_type_routing(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + doc_rows = _query(conn, "MATCH (n:document) RETURN n.id") + assert len(doc_rows) == 1 + assert doc_rows[0][0] == "n_concept_attn" + _close(db, conn) + + +# --- prune_sources --- + +def test_ingest_extraction_prune(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + before = _query(conn, "MATCH (n:code) RETURN count(n)")[0][0] + assert before == 3 + ingest_extraction(conn, ext, incremental=True, prune_sources=["model.py"]) + after_prune = _query(conn, "MATCH (n:code) RETURN count(n)")[0][0] + assert after_prune == 3 + _close(db, conn) + + +# --- fallback rel table --- + +def test_fallback_rel_table(tmp_db): + from graphify.storage import _ensure_rel_table, ensure_schema + db, conn = _init(tmp_db) + known = ensure_schema(conn) + tbl = _ensure_rel_table(conn, "paper", "document", "cites", known) + assert tbl == "edge_paper_document_cites" + assert tbl in known + _close(db, conn) + + +# --- communities --- + +def test_ingest_communities(tmp_db): + from graphify.storage import ingest_extraction, ingest_communities + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + communities = {0: ["n_transformer", "n_attention"], 1: ["n_layernorm"]} + ingest_communities(conn, communities) + rows = _query(conn, "MATCH (n:code) WHERE n.id = 'n_transformer' RETURN n.community") + assert rows[0][0] == 0 + rows = _query(conn, "MATCH (n:code) WHERE n.id = 'n_layernorm' RETURN n.community") + assert rows[0][0] == 1 + _close(db, conn) + + +# --- execute_cypher --- + +def test_execute_cypher(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + rows = _query(conn, "MATCH (n:code) RETURN n.label ORDER BY n.id") + labels = [r[0] for r in rows] + assert "MultiHeadAttention" in labels + assert "Transformer" in labels + _close(db, conn) + + +def test_execute_cypher_bad_query(tmp_db): + db, conn = _init(tmp_db) + with pytest.raises(RuntimeError): + _query(conn, "THIS IS NOT VALID CYPHER") + _close(db, conn) + + +# --- roundtrip consistency --- + +def test_roundtrip_node_count(tmp_db): + from graphify.storage import ingest_extraction + db, conn = _init(tmp_db) + ext = _load_extraction() + ingest_extraction(conn, ext, incremental=False) + total = 0 + for tbl in ("code", "document", "paper", "image", "concept", "rationale"): + rows = _query(conn, f"MATCH (n:{tbl}) RETURN count(n)") + total += rows[0][0] + assert total == len(ext["nodes"]) + _close(db, conn)