memory.py

"""Memory system — persistent file-based memory for the Wcode agent.

Memories are stored as individual markdown files with YAML frontmatter inside
``.memory/``.  An index file (``MEMORY.md``) lists all memories for quick
lookup.  The system supports relevance-based selection, extraction from
conversations, and periodic consolidation.
"""

import json
import re
import time
from pathlib import Path

import litellm

from config import MY_MODEL, get_logger
from utils import parse_frontmatter

logger = get_logger(__name__)


class MemoryManager:
    """Manages persistent file-based memories for the agent.

    Memories live as markdown files in *memory_dir*, indexed by
    ``MEMORY.md``.  Provides CRUD, relevance-based retrieval, extraction
    from conversations, and consolidation.

    Usage::

        mgr = MemoryManager(Path(".memory"))
        relevant = mgr.load_relevant(messages)
        mgr.extract(messages)
        mgr.consolidate()
    """

    def __init__(self, memory_dir: Path) -> None:
        self.memory_dir = memory_dir
        self.memory_dir.mkdir(exist_ok=True)
        self.index_path = self.memory_dir / "MEMORY.md"

    # ── Index & CRUD ───────────────────────────────────────────────────

    def read_index(self) -> str:
        """Return the contents of the memory index file.

        Returns an empty string if the index does not exist or is empty.
        """
        if not self.index_path.exists():
            return ""
        text = self.index_path.read_text().strip()
        return text if text else ""

    def read_file(self, filename: str) -> str | None:
        """Read a single memory file by name.

        Args:
            filename: Basename of the memory file (e.g. ``"preferences.md"``).

        Returns:
            Full raw text of the file, or ``None`` if not found.
        """
        filepath = self.memory_dir / filename
        if not filepath.exists():
            return None
        return filepath.read_text()

    def list_files(self) -> list[dict]:
        """Return metadata for all memory files (excludes ``MEMORY.md``).

        Each dict has keys: ``filename``, ``name``, ``description``,
        ``type``, ``body``.
        """
        result: list[dict] = []
        for f in sorted(self.memory_dir.glob("*.md")):
            if f.name == "MEMORY.md":
                continue
            raw = f.read_text()
            meta, body = parse_frontmatter(raw)
            result.append({
                "filename": f.name,
                "name": meta.get("name", f.stem),
                "description": meta.get("description", ""),
                "type": meta.get("type", "user"),
                "body": body,
            })
        return result

    def write_file(self, name: str, mem_type: str, description: str, body: str) -> Path:
        """Write a single memory ``.md`` file and rebuild the index.

        Args:
            name: Human-readable name (used for filename slug and frontmatter).
            mem_type: Type of memory (``"user"``, ``"project"``, ``"feedback"``, ``"reference"``).
            description: One-line description for the index.
            body: Full markdown body of the memory.

        Returns:
            Path to the newly created file.
        """
        slug = name.lower().replace(" ", "-").replace("/", "-")
        filename = f"{slug}.md"
        filepath = self.memory_dir / filename
        filepath.write_text(
            f"---\nname: {name}\ndescription: {description}\ntype: {mem_type}\n---\n\n{body}\n"
        )
        self._rebuild_index()
        return filepath

    # ── Relevance selection ─────────────────────────────────────────────

    def select_relevant(self, messages: list, max_items: int = 5) -> list[str]:
        """Use the LLM to select which memory files are relevant to the conversation.

        Args:
            messages: Conversation history (list of message dicts).
            max_items: Maximum number of filenames to return.

        Returns:
            List of filenames, or empty list if nothing matches.
        """
        files = self.list_files()
        if not files:
            return []

        recent_queries: list[str] = []
        cnt = 3
        for m in reversed(messages):
            if m.get("role") == "user":
                recent_queries.append(m.get("content", ""))
                cnt -= 1
                if cnt <= 0:
                    break
        recent = "\n".join(reversed(recent_queries))[:2000]
        if not recent.strip():
            return []

        catalog_lines = [
            f"{i}: {f['name']} — {f['description']}" for i, f in enumerate(files)
        ]
        catalog = "\n".join(catalog_lines)

        prompt = (
            "Given the recent conversation and the memory catalog below, "
            "select the indices of memories that are clearly relevant. "
            "Return ONLY a JSON array of integers, e.g. [0, 3]. "
            "If none are relevant, return [].\n\n"
            f"Recent conversation:\n{recent}\n\n"
            f"Memory catalog:\n{catalog}"
        )

        response = litellm.completion(
            messages=[{"role": "user", "content": prompt}],
            model=MY_MODEL,
            max_tokens=2000,
            num_ctx=4096,
        )
        text = response.choices[0].message.content.strip()
        match = re.search(r"\[.*?\]", text, re.DOTALL)
        if match:
            indices = json.loads(match.group())
            selected: list[str] = []
            for idx in indices:
                if isinstance(idx, int) and 0 <= idx < len(files):
                    selected.append(files[idx]["filename"])
                    if len(selected) >= max_items:
                        break
            return selected
        return []

    def load_relevant(self, messages: list) -> str:
        """Select and load relevant memories, returning an XML snippet.

        Returns a ``<relevant_memories>`` block, or ``""`` if nothing is relevant.
        """
        selected_files = self.select_relevant(messages)
        logger.info(f"Selected memories: {selected_files}")
        if not selected_files:
            return ""

        parts = ["<relevant_memories>"]
        for filename in selected_files:
            content = self.read_file(filename)
            if content:
                parts.append(content)
        parts.append("</relevant_memories>")
        return "\n".join(parts)

    # ── Writing & maintenance ───────────────────────────────────────────

    def consolidate(self, max_memory_size: int = 5) -> None:
        """Use the LLM to merge duplicates and remove outdated memories.

        If the number of memory files exceeds *max_memory_size*, all
        memories are sent to the LLM for consolidation.  Existing files
        are replaced with the cleaned set.
        """
        files = self.list_files()
        if len(files) < max_memory_size:
            return

        catalog = "\n\n".join(
            f"## {f['filename']}\nname: {f['name']}\ndescription: {f['description']}\n{f['body']}"
            for f in files
        )

        prompt = (
            "Consolidate the following memory files. Rules:\n"
            "1. Merge duplicates into one\n"
            "2. Remove outdated/contradicted memories\n"
            "3. Keep the total under 30 memories\n"
            "4. Preserve important user preferences above all\n"
            "Return a JSON array. Each item: {name, type, description, body}.\n\n"
            f"{catalog[:16000]}"
        )

        response = litellm.completion(
            messages=[{"role": "user", "content": prompt}],
            model=MY_MODEL,
            max_tokens=8000,
        )
        text = response.choices[0].message.content
        match = re.search(r"\[.*\]", text, re.DOTALL)
        if not match:
            return
        items = json.loads(match.group())

        for f in self.memory_dir.glob("*.md"):
            if f.name != "MEMORY.md":
                f.unlink()

        for mem in items:
            mem_name = mem.get("name", f"memory_{int(time.time())}")
            mem_type = mem.get("type", "user")
            desc = mem.get("description", "")
            mem_body = mem.get("body", "")
            if desc and mem_body:
                self.write_file(mem_name, mem_type, desc, mem_body)

        logger.info(f"[Memory: consolidated {len(files)} → {len(items)} memories]")

    def extract(self, messages: list) -> None:
        """Extract new preferences, constraints, or project facts from the
        conversation and persist them as memory files.

        Analyses the last 10 user/assistant messages, compares them against
        existing memories, and asks the LLM to extract novel information.
        """
        dialogue_parts: list[str] = []
        msg_cnt = 10
        for msg in reversed(messages):
            role = msg.get("role", "?")
            content = msg.get("content", "").strip()
            if role not in ["assistant", "user"] or content == "":
                continue
            dialogue_parts.append(f"{role}: {content}")
            msg_cnt -= 1
            if msg_cnt == 0:
                break
        if len(dialogue_parts) == 0:
            return
        dialogue = "\n".join(dialogue_parts)

        existing = self.list_files()
        existing_desc = (
            "\n".join(f"- {m['name']}: {m['description']}" for m in existing)
            if existing
            else "(none)"
        )

        prompt = (
            "Extract user preferences, constraints, or project facts from this dialogue.\n"
            "Return a JSON array. Each item: {name, type, description, body}.\n"
            "- name: short kebab-case identifier (e.g. 'user-preference-tabs')\n"
            "- type: one of 'user' (user preference), 'feedback' (guidance), "
            "'project' (project fact), 'reference' (external pointer)\n"
            "- description: one-line summary for index lookup\n"
            "- body: full detail in markdown\n"
            "If nothing new or already covered by existing memories, return [].\n\n"
            f"Existing memories:\n{existing_desc}\n\n"
            f"Dialogue:\n{dialogue[:4000]}"
        )

        response = litellm.completion(
            messages=[{"role": "user", "content": prompt}],
            model=MY_MODEL,
            max_tokens=8000,
        )
        text = response.choices[0].message.content
        match = re.search(r"\[.*\]", text, re.DOTALL)
        if not match:
            return
        items = json.loads(match.group())
        if not items:
            return
        count = 0
        for mem in items:
            mem_name = mem.get("name", f"memory_{int(time.time())}")
            mem_type = mem.get("type", "user")
            desc = mem.get("description", "")
            mem_body = mem.get("body", "")
            if desc and mem_body:
                self.write_file(mem_name, mem_type, desc, mem_body)
                count += 1
        if count:
            logger.info(f"[Memory: extracted {count} new memories]")

    # ── Internal ────────────────────────────────────────────────────────

    def _rebuild_index(self) -> None:
        """Rebuild MEMORY.md from the current set of memory files."""
        lines: list[str] = []
        for f in sorted(self.memory_dir.glob("*.md")):
            if f.name == "MEMORY.md":
                continue
            raw = f.read_text()
            meta, body = parse_frontmatter(raw)
            name = meta.get("name", f.stem)
            desc = meta.get("description", body.split("\n")[0][:80])
            lines.append(f"- [{name}]({f.name}) — {desc}")
        self.index_path.write_text("\n".join(lines) + "\n" if lines else "")


# ── Module-level shims (delegate to a default registry) ──────────────────

_default_mgr: MemoryManager | None = None


def _get_default() -> MemoryManager:
    global _default_mgr
    if _default_mgr is None:
        from config import MEMORY_DIR
        _default_mgr = MemoryManager(MEMORY_DIR)
    return _default_mgr


def read_memory_index() -> str:
    return _get_default().read_index()


def read_memory_file(filename: str) -> str | None:
    return _get_default().read_file(filename)


def list_memory_files() -> list[dict]:
    return _get_default().list_files()


def select_relevant_memories(messages: list, max_items: int = 5) -> list[str]:
    return _get_default().select_relevant(messages, max_items)


def load_memories(messages: list) -> str:
    return _get_default().load_relevant(messages)


def write_memory_file(name: str, mem_type: str, description: str, body: str) -> Path:
    return _get_default().write_file(name, mem_type, description, body)


def consolidate_memories(max_memory_size: int = 5) -> None:
    _get_default().consolidate(max_memory_size)


def extract_memories(messages: list) -> None:
    _get_default().extract(messages)