ayushbharadva/patchpilot-api

1#!/usr/bin/env python32"""Zero-cost reseed snapshot — tar save/restore of .patchpilot_memory/ (Pitfall 4).3 4Only `cognify()` incurs meaningful LLM cost (graph extraction); `add()`,5`search()`, and `forget()` are cheap or free. Re-running the full seed CLI6from scratch on every demo rehearsal re-bills `cognify()` for all three7datasets every time. Snapshotting `.patchpilot_memory/` AFTER the first8successful `--seed` lets `seed/seed_cli.py --reset` restore a fully9cognified, demo-ready state for $0 instead of re-cognifying.10 11This module is pure filesystem tar work — it deliberately does NOT import12`cognee` or `backend.cognee_config`, since a tarball save/restore never13needs to resolve Cognee's LLM/embedding provider. It does, however, need to14resolve the same storage-root override convention `backend/cognee_config.py`15uses (`SYSTEM_ROOT_DIRECTORY` in `.env`), so it re-reads that one env var16directly instead of hardcoding the default path.17 18The resulting `*.snapshot.tar` is gitignored (Plan 01) — it may contain19cognified seed content and must never be committed (see 01-04-PLAN.md20threat T-01-02).21 22Usage:23    .venv/bin/python scripts/snapshot_memory.py --save     # tar current .patchpilot_memory/24    .venv/bin/python scripts/snapshot_memory.py --restore  # untar, replacing the live tree25"""26 27import argparse28import os29import shutil30import sqlite331import sys32import tarfile33from pathlib import Path34 35from dotenv import load_dotenv36 37# Resolve the same SYSTEM_ROOT_DIRECTORY override backend/cognee_config.py38# supports, without importing cognee itself (see module docstring).39load_dotenv()40REPO_ROOT = Path(__file__).resolve().parent.parent41MEMORY_ROOT = Path(os.environ.get("SYSTEM_ROOT_DIRECTORY", str(REPO_ROOT / ".patchpilot_memory")))42SNAPSHOT_PATH = REPO_ROOT / "patchpilot_memory.snapshot.tar"43 44 45def snapshot_exists() -> bool:46    return SNAPSHOT_PATH.exists()47 48 49def save() -> None:50    """Tar the entire .patchpilot_memory/ tree into a gitignored *.snapshot.tar."""51    if not MEMORY_ROOT.exists() or not any(MEMORY_ROOT.iterdir()):52        print(f"Nothing to save -- {MEMORY_ROOT} does not exist or is empty. Run --seed first.")53        sys.exit(1)54    with tarfile.open(SNAPSHOT_PATH, "w") as tar:55        tar.add(MEMORY_ROOT, arcname=MEMORY_ROOT.name)56    print(f"Saved snapshot: {SNAPSHOT_PATH}")57 58 59def restore() -> None:60    """Extract the tarball back, replacing the live .patchpilot_memory/ tree.61 62    Raises FileNotFoundError (a normal Exception subclass) rather than63    calling sys.exit() directly -- this function is called both from the64    CLI (main(), below) and from backend/reset.py's POST /reset handler.65    sys.exit() raises SystemExit, a BaseException that is NOT caught by66    reset.py's `except Exception:` guard nor by asyncio's task machinery,67    which would crash the entire running uvicorn worker (CR-01) instead of68    returning a graceful D-24 error response. main() below translates this69    back into a CLI-appropriate sys.exit(1).70    """71    if not snapshot_exists():72        raise FileNotFoundError(73            f"No snapshot found at {SNAPSHOT_PATH}. Run --save first (after a successful --seed)."74        )75    if MEMORY_ROOT.exists():76        shutil.rmtree(MEMORY_ROOT)77    with tarfile.open(SNAPSHOT_PATH, "r") as tar:78        tar.extractall(REPO_ROOT, filter="data")79    _rewrite_absolute_paths()80    print(f"Restored {MEMORY_ROOT} from {SNAPSHOT_PATH}")81 82 83def _rewrite_absolute_paths() -> None:84    """Rewrite machine-specific absolute paths inside the restored sqlite DBs.85 86    Cognee records absolute paths from the machine that CREATED the snapshot87    (dataset_database.vector_database_url, data.raw_data_location, node88    properties JSON — e.g. /Users/<dev>/.../.patchpilot_memory/databases/...).89    Restored onto a different machine those paths are dead: observed live on90    the HF Space (Jul 6), every search failed with LanceDB "Unable to create91    lance dataset at /Users/... Permission denied (os error 13)" because the92    Linux container obediently tried the Mac path. Replace every occurrence93    of the snapshot's original memory root with THIS machine's MEMORY_ROOT,94    in every text-typed column of every restored sqlite database. No-op when95    the snapshot was created on this machine (old root == new root).96    """97    # The marker is the memory root's directory name as recorded INSIDE the98    # snapshot (save() tars with arcname=MEMORY_ROOT.name, which is99    # .patchpilot_memory everywhere this project runs) — NOT the runtime100    # root's name, which an env override may have pointed elsewhere.101    marker = ".patchpilot_memory"102    for db_path in (MEMORY_ROOT / "databases").glob("*"):103        if not db_path.is_file() or db_path.name.endswith(("-shm", "-wal")):104            continue105        try:106            con = sqlite3.connect(db_path)107        except sqlite3.Error:108            continue109        try:110            tables = [111                r[0]112                for r in con.execute(113                    "SELECT name FROM sqlite_master WHERE type='table'"114                ).fetchall()115            ]116            # Discover every distinct stored prefix ending at the marker —117            # values appear both as plain paths (/Users/.../.patchpilot_memory)118            # and as URIs (file:///Users/.../.patchpilot_memory), and each119            # form must be rewritten while keeping its scheme.120            old_roots: set[str] = set()121            text_cols: dict[str, list[str]] = {}122            for table in tables:123                cols = con.execute(f'PRAGMA table_info("{table}")').fetchall()124                text_cols[table] = [125                    c[1]126                    for c in cols127                    if any(t in (c[2] or "").upper() for t in ("CHAR", "TEXT", "JSON", "CLOB"))128                ]129                for col in text_cols[table]:130                    for (value,) in con.execute(131                        f'SELECT "{col}" FROM "{table}" WHERE "{col}" LIKE ?',132                        (f"%{marker}%",),133                    ).fetchall():134                        if isinstance(value, str):135                            start = 0136                            while (idx := value.find(marker, start)) != -1:137                                # Prefix runs from the start of the embedded138                                # path/URI; for JSON blobs fall back to the139                                # nearest quote boundary.140                                boundary = max(141                                    value.rfind('"', 0, idx), value.rfind(" ", 0, idx)142                                )143                                old_roots.add(value[boundary + 1 : idx + len(marker)])144                                start = idx + len(marker)145            old_roots.discard(str(MEMORY_ROOT))146            rewritten = 0147            # Longest first so URI forms never get half-eaten by their plain148            # substring form.149            for old_root in sorted(old_roots, key=len, reverse=True):150                if "://" in old_root:151                    new_root = old_root.split("://", 1)[0] + "://" + str(MEMORY_ROOT)152                else:153                    new_root = str(MEMORY_ROOT)154                if old_root == new_root:155                    continue156                for table, cols in text_cols.items():157                    for col in cols:158                        cur = con.execute(159                            f'UPDATE "{table}" SET "{col}" = REPLACE("{col}", ?, ?) '160                            f'WHERE "{col}" LIKE ?',161                            (old_root, new_root, f"%{old_root}%"),162                        )163                        rewritten += cur.rowcount164                print(f"Path rewrite in {db_path.name}: {old_root} -> {new_root}")165            con.commit()166            if rewritten:167                print(f"Rewrote {rewritten} row(s) in {db_path.name}")168        except sqlite3.Error as exc:169            print(f"WARNING: path rewrite skipped for {db_path.name}: {exc}")170        finally:171            con.close()172 173 174def main() -> int:175    parser = argparse.ArgumentParser(description="Save/restore .patchpilot_memory/ for zero-cost reseeds")176    group = parser.add_mutually_exclusive_group(required=True)177    group.add_argument("--save", action="store_true", help="Tar the current .patchpilot_memory/ tree")178    group.add_argument("--restore", action="store_true", help="Restore .patchpilot_memory/ from the tarball")179    args = parser.parse_args()180 181    try:182        if args.save:183            save()184        else:185            restore()186    except FileNotFoundError as exc:187        print(exc)188        return 1189    return 0190 191 192if __name__ == "__main__":193    sys.exit(main())194