Memba/experiments/recall_poc.py

"""
recall_poc.py — proof of concept for the 'recall' product idea.

A daemon would normally feed events into Mamba state continuously.  For the
POC we batch-feed the last N days of git activity from a set of repos, save
the resulting state, then in a SEPARATE invocation query the saved state.

If the query in mode 2 produces an answer that genuinely reflects the input,
the concept is viable.  If it produces vague/wrong answers — rethink.

Usage:
    # Step 1: ingest git history into a memba state
    python recall_poc.py build --model <gguf>

    # Step 2 (separate process — proves persistence):
    python recall_poc.py query --model <gguf>
    python recall_poc.py query --model <gguf> --q "Which project saw the most activity?"
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
import time
from datetime import datetime, timedelta
from pathlib import Path

# ── Config ───────────────────────────────────────────────────────
SCAN_ROOTS = [
    Path.home() / "Desktop/Coding",
    Path.home() / "Desktop/Coding/AI",
]
DAYS_BACK    = 30
STATE_DIR    = Path.home() / ".recall/states"
SESSION_ID   = "recall_poc_nemotron"
DEFAULT_Q    = (
    "Summarise what I've been working on over the last month. "
    "Group by project. Mention the main themes per project. "
    "Be concise — one short paragraph per project."
)


# ── Git scraping ─────────────────────────────────────────────────

def find_repos(roots: list[Path]) -> list[Path]:
    repos = []
    seen  = set()
    for root in roots:
        if not root.exists():
            continue
        for entry in sorted(root.iterdir()):
            if not entry.is_dir() or not (entry / ".git").exists():
                continue
            real = entry.resolve()
            if real in seen:
                continue
            seen.add(real)
            repos.append(entry)
    return repos


def git_log(repo: Path, since_days: int) -> str:
    """Return human-readable commit log for last `since_days` days."""
    try:
        out = subprocess.check_output(
            [
                "git", "-C", str(repo), "log",
                f"--since={since_days} days ago",
                "--no-merges",
                "--date=short",
                "--pretty=format:%ad  %s",
            ],
            text=True, stderr=subprocess.DEVNULL,
        )
    except subprocess.CalledProcessError:
        return ""
    return out.strip()


def build_corpus(repos: list[Path], since_days: int) -> tuple[str, int]:
    """Build a single human-readable summary block.  Returns (text, n_commits)."""
    blocks: list[str] = []
    total = 0
    for repo in repos:
        log = git_log(repo, since_days)
        if not log:
            continue
        n = log.count("\n") + 1
        total += n
        blocks.append(f"=== {repo.name} ({n} commits) ===\n{log}")
    return "\n\n".join(blocks), total


# ── memba interaction ───────────────────────────────────────────

def open_session(model_path: str, n_gpu_layers: int) -> "Session":  # noqa: F821
    # Imported lazily so query mode doesn't pay the cost when only listing
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
    from memba import Session
    return Session(
        model_path=model_path,
        session_id=SESSION_ID,
        state_dir=str(STATE_DIR),
        n_gpu_layers=n_gpu_layers,
        n_ctx=8192,           # bigger than default so a month of commits fits comfortably
        chat_format="chatml",
    )


# ── build mode ──────────────────────────────────────────────────

def cmd_build(args: argparse.Namespace) -> int:
    repos = find_repos(SCAN_ROOTS)
    print(f"[build] discovered {len(repos)} git repos under {[str(r) for r in SCAN_ROOTS]}")

    corpus, n_commits = build_corpus(repos, args.days)
    if not corpus:
        print(f"[build] no commits in the last {args.days} days — abort", file=sys.stderr)
        return 1

    char_count = len(corpus)
    print(f"[build] corpus: {n_commits} commits, {char_count:,} characters")
    print(f"[build] first 400 chars:\n---\n{corpus[:400]}\n---")

    # Wipe any previous state for a clean test
    state_file = STATE_DIR / f"{SESSION_ID}.memb"
    if state_file.exists():
        state_file.unlink()
        print(f"[build] removed previous state file")

    print(f"[build] loading model ({Path(args.model).name}) on GPU layers={args.gpu_layers}…")
    t0 = time.time()
    sess = open_session(args.model, args.gpu_layers)
    print(f"[build] model loaded in {time.time()-t0:.1f}s, initial state {sess.state_size:,} B")

    prompt = (
        "I am going to give you my git commit history from the last month, "
        "across several of my personal projects. Please READ it and remember "
        "the overall picture — which projects I worked on, the kinds of "
        "changes I made, and any themes that emerge. Just reply 'noted' when "
        "you have processed it; I'll ask questions in a follow-up.\n\n"
        f"=== git log (last {args.days} days, {n_commits} commits) ===\n\n"
        f"{corpus}"
    )

    print(f"[build] feeding {len(prompt):,} characters into state…")
    t0 = time.time()
    ack = sess.chat(prompt, max_tokens=32)
    elapsed = time.time() - t0
    print(f"[build] processed in {elapsed:.1f}s   (~{len(prompt)/elapsed:,.0f} char/s)")
    print(f"[build] model ack: {ack!r}")
    print(f"[build] state after ingest: {sess.state_size:,} B")

    path = sess.save()
    on_disk = path.stat().st_size
    print(f"[build] state saved → {path}   ({on_disk:,} B on disk)")
    return 0


# ── query mode ──────────────────────────────────────────────────

def cmd_query(args: argparse.Namespace) -> int:
    state_file = STATE_DIR / f"{SESSION_ID}.memb"
    if not state_file.exists():
        print(f"[query] no state file at {state_file} — run `build` first", file=sys.stderr)
        return 1

    on_disk = state_file.stat().st_size
    print(f"[query] state file: {state_file}  ({on_disk:,} B)")
    print(f"[query] loading model ({Path(args.model).name}) on GPU layers={args.gpu_layers}…")
    t0 = time.time()
    sess = open_session(args.model, args.gpu_layers)
    print(f"[query] model+state loaded in {time.time()-t0:.1f}s, state size {sess.state_size:,} B")

    question = args.q or DEFAULT_Q
    print(f"\n[query] question:\n{question}\n")
    print(f"[query] response:\n---")
    t0 = time.time()
    reply = sess.chat(question, max_tokens=args.max_tokens)
    print(reply)
    print(f"---  ({time.time()-t0:.1f}s)")
    return 0


# ── main ────────────────────────────────────────────────────────

def main() -> int:
    p = argparse.ArgumentParser(description="recall POC — git log → mamba state → cross-process query")
    sub = p.add_subparsers(dest="cmd", required=True)

    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--model", required=True, help="Path to GGUF model")
    common.add_argument("--gpu-layers", type=int, default=-1, dest="gpu_layers",
                        help="Layers to offload (-1=all, 0=CPU)")

    b = sub.add_parser("build", parents=[common], help="Ingest git history into state")
    b.add_argument("--days", type=int, default=DAYS_BACK)
    b.set_defaults(func=cmd_build)

    q = sub.add_parser("query", parents=[common], help="Query the saved state")
    q.add_argument("--q", help="Question to ask (default: monthly summary)")
    q.add_argument("--max-tokens", type=int, default=512, dest="max_tokens")
    q.set_defaults(func=cmd_query)

    args = p.parse_args()
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())