memory.py

"""Chapter 8 reference: small filesystem-backed memory index."""

from __future__ import annotations

from pathlib import Path


MEMORY_DIR = Path.cwd() / "memory"
MEMORY_INDEX_NAME = "MEMORY.md"


def parse_frontmatter(text: str) -> tuple[dict[str, str], str]:
    """Parse a tiny YAML-like frontmatter block from a memory file."""
    if not text.startswith("---\n"):
        return {}, text

    end = text.find("\n---", 4)
    if end == -1:
        return {}, text

    raw = text[4:end]
    body = text[end + 4 :].lstrip("\n")
    metadata: dict[str, str] = {}
    for line in raw.splitlines():
        if ":" not in line:
            continue
        key, value = line.split(":", 1)
        metadata[key.strip()] = value.strip().strip("\"'")

    return metadata, body


def refresh_memory_index() -> str:
    """Rebuild MEMORY.md from individual memory files."""
    MEMORY_DIR.mkdir(parents=True, exist_ok=True)
    lines = ["# Memory Index", ""]

    for path in sorted(MEMORY_DIR.glob("*.md")):
        if path.name == MEMORY_INDEX_NAME:
            continue

        try:
            metadata, _ = parse_frontmatter(path.read_text(encoding="utf-8"))
        except OSError:
            continue

        memory_type = metadata.get("type", "note")
        summary = metadata.get("summary", path.stem)
        updated = metadata.get("updated")
        suffix = f" - updated {updated}" if updated else ""
        lines.append(f"- {memory_type}: {summary} (`{path.name}`){suffix}")

    index = "\n".join(lines).rstrip() + "\n"
    (MEMORY_DIR / MEMORY_INDEX_NAME).write_text(index, encoding="utf-8")
    return index


def load_memory_index() -> str:
    """Load the current memory index, creating it when needed."""
    index_path = MEMORY_DIR / MEMORY_INDEX_NAME
    if not index_path.exists():
        return refresh_memory_index()
    return index_path.read_text(encoding="utf-8")


def build_memory_prompt_section() -> str:
    """Return the memory rules and current index for the system prompt."""
    index = load_memory_index()
    return (
        "# Memory System\n"
        "Use memory files only when they are relevant to the current task.\n"
        "To save durable preferences or project facts, write Markdown files "
        f"under `{MEMORY_DIR}` with frontmatter fields `type`, `summary`, "
        "and `updated`.\n\n"
        f"{index}"
    )