memory.py

"""pit active context management — a delayed-batch memory system.

Two complementary mechanisms:

1. prune_stale_tool_results()  — MECHANICAL, no model. The acute win: old, large
   tool results (a grep dump from 8 steps ago) get collapsed to a one-line stub,
   preserving tool_call_id pairing so the API stays valid. Runs cheaply every step.

2. MemoryStore + batch_process()  — a MemGPT variant. DELAYED BATCH: the
   conversation just accumulates in the live buffer; only when it crosses a
   CHARACTER threshold do we pay the cost ONCE — a single model pass that extracts
   typed items (facts/tasks/context/decisions), classifies each Mem0-style
   (ADD/UPDATE/DELETE/NOOP) against what we already know, files them into
   datetime-stamped SUBJECT THREADS, and merges a PROGRESSIVE HIERARCHICAL summary
   (batch -> session -> rolling, so summaries don't degrade from re-summarizing).
   Then the live buffer is reset to: system + rendered memory + a few recent turns.
   Retrieval is by subject thread ("the squat-form thread"), not vector search.

The pitch vs MemGPT: MemGPT pays a constant per-turn tax editing memory mid-
conversation. This pays it lazily, after the fact — a secretary who takes raw
notes during the meeting and organizes them afterward.
"""
import hashlib
import json
import os
import re
import time
from datetime import datetime
from pathlib import Path

import httpx

HERE = Path(__file__).parent
# buffer trigger (by design: char count, not tokens). Env-overridable so
# experiments can isolate variables (set huge to disable batch, low to force it).
BATCH_CHARS = int(os.environ.get("PIT_BATCH_CHARS", "24000"))
PRUNE_KEEP_RECENT = 3         # keep this many newest tool results verbatim
PRUNE_MIN_CHARS = 1_200       # only stub tool results bigger than this
KEEP_RECENT_TURNS = 6         # raw messages preserved after a batch compaction
BUCKETS = ("facts", "tasks", "context", "decisions")


# --- 1. mechanical stale-tool-result pruning ---------------------------------
def prune_stale_tool_results(messages, keep_recent=PRUNE_KEEP_RECENT,
                             min_chars=PRUNE_MIN_CHARS):
    """Collapse old/large tool results to stubs in place. Preserves each tool
    message's role + tool_call_id (so the assistant tool_call it answers stays
    paired and the API stays valid). Returns chars saved."""
    # map tool_call_id -> tool name, from assistant tool_calls
    id_to_name = {}
    for m in messages:
        for tc in m.get("tool_calls") or []:
            id_to_name[tc.get("id")] = tc.get("function", {}).get("name", "tool")
    tool_idxs = [i for i, m in enumerate(messages) if m.get("role") == "tool"]
    protect = set(tool_idxs[-keep_recent:]) if keep_recent else set()
    saved = 0
    for i in tool_idxs:
        if i in protect:
            continue
        m = messages[i]
        c = m.get("content") or ""
        if len(c) <= min_chars or c.startswith("[pruned"):
            continue
        name = id_to_name.get(m.get("tool_call_id"), "tool")
        head = c.splitlines()[0][:80] if c.splitlines() else ""
        stub = (f"[pruned stale {name} result — {len(c)} chars elided. "
                f"First line: {head!r}. Re-run the call if you need it again.]")
        saved += len(c) - len(stub)
        m["content"] = stub
    return saved


# --- 2a. the store -----------------------------------------------------------
def _slug(s):
    return re.sub(r"[^a-z0-9]+", "-", s.lower()).strip("-")[:60] or "misc"


def _now():
    return datetime.now().strftime("%Y-%m-%d %H:%M")


class MemoryStore:
    """Disk-backed typed buckets + subject threads + hierarchical summary,
    keyed per project (cwd). Plain JSON — inspectable, no vector index."""

    def __init__(self, project_dir):
        key = hashlib.md5(str(Path(project_dir).resolve()).encode()).hexdigest()[:10]
        self.root = HERE / "memory_store" / f"{_slug(Path(project_dir).name)}-{key}"
        self.root.mkdir(parents=True, exist_ok=True)
        self.path = self.root / "store.json"
        self.data = self._load()

    def _load(self):
        if self.path.exists():
            try:
                return json.loads(self.path.read_text(encoding="utf-8"))
            except Exception:                    # noqa: BLE001
                pass
        return {"buckets": {b: [] for b in BUCKETS},   # each: {id, ts, text}
                "threads": {},                          # subject -> [{ts, text}]
                "summary": {"batch": "", "session": "", "rolling": ""},
                "batch_count": 0}

    def save(self):
        self.path.write_text(json.dumps(self.data, ensure_ascii=False, indent=2),
                             encoding="utf-8")
        self._render_markdown()

    # -- Mem0-style ops on the typed buckets --
    def apply_ops(self, ops):
        """ops: [{action, bucket, id?, text, subject?}]. Returns a count summary."""
        tally = {"ADD": 0, "UPDATE": 0, "DELETE": 0, "NOOP": 0}
        for op in ops:
            action = (op.get("action") or "NOOP").upper()
            bucket = op.get("bucket")
            if bucket not in BUCKETS or action not in tally:
                tally["NOOP"] += 1
                continue
            items = self.data["buckets"][bucket]
            if action == "ADD":
                items.append({"id": op.get("id") or _mkid(op.get("text", "")),
                              "ts": _now(), "text": op.get("text", "").strip(),
                              "subject": op.get("subject", "")})
            elif action == "UPDATE":
                for it in items:
                    if it["id"] == op.get("id"):
                        it["text"] = op.get("text", it["text"]).strip()
                        it["ts"] = _now()
                        break
                else:
                    items.append({"id": op.get("id") or _mkid(op.get("text", "")),
                                  "ts": _now(), "text": op.get("text", "").strip(),
                                  "subject": op.get("subject", "")})
            elif action == "DELETE":
                self.data["buckets"][bucket] = [
                    it for it in items if it["id"] != op.get("id")]
            tally[action] += 1
        return tally

    def add_thread_entries(self, subjects):
        """subjects: {subject_name: [entry_text, ...]}. Datetime-stamped append."""
        for subj, entries in (subjects or {}).items():
            key = _slug(subj)
            thread = self.data["threads"].setdefault(key, [])
            for e in entries:
                if e and e.strip():
                    thread.append({"ts": _now(), "text": e.strip()})

    def merge_summary(self, batch_summary):
        """Progressive + hierarchical. Newest batch summary becomes L1; every 4
        batches L1 rolls into the session summary; the rolling summary always
        reflects session-level. Merge (not re-summarize-from-scratch) so detail
        doesn't erode."""
        s = self.data["summary"]
        s["batch"] = batch_summary.strip()
        self.data["batch_count"] += 1
        # append batch into session, keep session bounded
        s["session"] = (s["session"] + "\n- " + batch_summary.strip()).strip()
        if self.data["batch_count"] % 4 == 0:
            s["rolling"] = (s["rolling"] + "\n\n== rollup " + _now() + " ==\n"
                            + s["session"]).strip()[-6000:]

    def open_items(self):
        return {b: self.data["buckets"][b] for b in BUCKETS}

    def relevant_threads(self, text, k=3):
        """Keyword-overlap retrieval (NOT vector search) — pull the subject
        threads whose name/content best overlaps the query text."""
        words = set(re.findall(r"[a-z0-9]{4,}", (text or "").lower()))
        scored = []
        for subj, entries in self.data["threads"].items():
            blob = (subj + " " + " ".join(e["text"] for e in entries[-6:])).lower()
            hits = sum(1 for w in words if w in blob)
            if hits:
                scored.append((hits, subj, entries))
        scored.sort(reverse=True, key=lambda x: x[0])
        return scored[:k]

    def render_context(self, focus_text=""):
        """The compacted memory injected back into the live buffer."""
        s = self.data["summary"]
        parts = ["# Working memory (from prior conversation)"]
        if s["rolling"]:
            parts.append("## Rolling summary\n" + s["rolling"])
        if s["session"]:
            parts.append("## This session\n" + s["session"][-3000:])
        for b in BUCKETS:
            items = self.data["buckets"][b]
            if items:
                lines = "\n".join(f"- {it['text']}  ({it['ts']})"
                                  for it in items[-12:])
                parts.append(f"## {b.capitalize()}\n{lines}")
        threads = self.relevant_threads(focus_text)
        if threads:
            parts.append("## Relevant subject threads")
            for _, subj, entries in threads:
                recent = "\n".join(f"  - [{e['ts']}] {e['text']}"
                                   for e in entries[-5:])
                parts.append(f"### {subj}\n{recent}")
        return "\n\n".join(parts)

    def _render_markdown(self):
        """Human-readable mirror for inspection."""
        try:
            (self.root / "MEMORY.md").write_text(
                self.render_context(), encoding="utf-8")
        except Exception:                        # noqa: BLE001
            pass


def _mkid(text):
    return hashlib.md5(text.encode("utf-8")).hexdigest()[:8]


# --- 2b. the batch pass ------------------------------------------------------
def _complete(model_cfg, prompt, timeout=180):
    """One non-streaming completion, thinking disabled, returns text."""
    url = model_cfg["base_url"].rstrip("/") + "/chat/completions"
    headers = {"Content-Type": "application/json"}
    import os
    kenv = model_cfg.get("api_key_env")
    if kenv and os.environ.get(kenv):
        headers["Authorization"] = f"Bearer {os.environ[kenv]}"
    body = {"model": model_cfg["model"], "stream": False,
            "messages": [{"role": "user", "content": prompt}]}
    body.update(model_cfg.get("extra_body", {}))
    body.setdefault("chat_template_kwargs", {})["enable_thinking"] = False
    r = httpx.post(url, headers=headers, json=body, timeout=timeout)
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"] or ""


def _extract_json(text):
    """Pull the first JSON object out of a model reply (strip fences/prose)."""
    text = re.sub(r"```(?:json)?", "", text)
    start = text.find("{")
    if start < 0:
        return None
    depth = 0
    for i in range(start, len(text)):
        depth += (text[i] == "{") - (text[i] == "}")
        if depth == 0:
            try:
                return json.loads(text[start:i + 1])
            except json.JSONDecodeError:
                return None
    return None


def _conversation_text(messages):
    """Flatten the buffer (skip the system msg) into readable transcript."""
    out = []
    for m in messages:
        role = m.get("role")
        if role == "system":
            continue
        c = m.get("content") or ""
        if role == "assistant" and m.get("tool_calls"):
            names = ", ".join(tc["function"]["name"] for tc in m["tool_calls"])
            c = (c + f"  [called: {names}]").strip()
        if role == "tool":
            c = "TOOL RESULT: " + c[:800]
        if c:
            out.append(f"{role.upper()}: {c}")
    return "\n\n".join(out)


def batch_process(messages, model_cfg, store, echo=False):
    """The delayed batch pass. Returns a NEW compacted message list; does not
    mutate `messages`. messages[0] must be the system prompt."""
    system = messages[0]
    # mechanical prune first so the LLM sees less noise too
    prune_stale_tool_results(messages, keep_recent=0, min_chars=PRUNE_MIN_CHARS)
    convo = _conversation_text(messages)
    existing = json.dumps(
        {"buckets": {b: [{"id": it["id"], "text": it["text"]}
                          for it in store.data["buckets"][b]] for b in BUCKETS},
         "subjects": list(store.data["threads"].keys())},
        ensure_ascii=False)
    prompt = (HERE / "prompts" / "batch.md").read_text(encoding="utf-8")
    prompt = prompt.replace("{{EXISTING}}", existing[:6000]).replace(
        "{{CONVERSATION}}", convo[:18000])

    result = None
    try:
        raw = _complete(model_cfg, prompt)
        result = _extract_json(raw)
    except Exception as e:                       # noqa: BLE001
        if echo:
            print(f"[memory] batch model call failed: {e}")

    if result:
        tally = store.apply_ops(result.get("ops", []))
        store.add_thread_entries(result.get("subjects", {}))
        store.merge_summary(result.get("summary", ""))
        store.save()
        if echo:
            print(f"[memory] batch #{store.data['batch_count']}: "
                  f"ADD {tally['ADD']} UPDATE {tally['UPDATE']} "
                  f"DELETE {tally['DELETE']} NOOP {tally['NOOP']}; "
                  f"{len(store.data['threads'])} threads")
    elif echo:
        print("[memory] batch produced no parseable result — keeping raw buffer")
        return messages  # safe fallback: don't lose context

    # rebuild the live buffer: system + rendered memory + recent raw turns
    focus = _conversation_text(messages[-4:])
    mem_msg = {"role": "user",
               "content": store.render_context(focus)
               + "\n\n(Continue. Above is your organized memory of earlier context.)"}
    tail = [m for m in messages[1:][-KEEP_RECENT_TURNS:]
            if m.get("role") in ("user", "assistant")]  # drop orphan tool msgs
    return [system, mem_msg, *tail]


def buffer_chars(messages):
    n = 0
    for m in messages:
        n += len(m.get("content") or "")
        for tc in m.get("tool_calls") or []:
            n += len(tc.get("function", {}).get("arguments", ""))
    return n