Memba/python/memba/session.py

"""High-level Session API — wraps llama-cpp-python Llama + memba state I/O."""

from __future__ import annotations

from pathlib import Path
from typing import Optional

from . import core


# ── Chat formats ───────────────────────────────────────────────────
# Each entry maps a format name to (prompt_template, stop_strings).
# The template uses {prompt} as the placeholder for the user message.
CHAT_FORMATS: dict[str, tuple[str, list[str]]] = {
    "chatml": (
        "<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n",
        ["<|im_end|>", "<|im_start|>"],
    ),
    "raw": (
        "{prompt}",
        [],
    ),
}


class Session:
    """
    A persistent SSM chat session backed by a memba state file.

    The session auto-loads an existing state on construction (if one exists
    for *session_id*) and accumulates context across calls to chat().
    Call save() to persist the current state.

    Parameters
    ----------
    model_path:   Path to the GGUF model file.
    session_id:   Logical name for this session; determines the state filename.
    state_dir:    Directory where .memb files are stored (created if absent).
    n_gpu_layers: GPU layers to offload (0 = CPU-only, -1 = all layers).
    n_ctx:        Context window size in tokens.
    verbose:      Forward llama.cpp log output to stderr.
    """

    def __init__(
        self,
        model_path: str,
        session_id: str = "default",
        state_dir: str = "~/.memba/states",
        n_gpu_layers: int = 0,
        n_ctx: int = 4096,
        verbose: bool = False,
        chat_format: str = "chatml",
    ) -> None:
        try:
            from llama_cpp import Llama
        except ImportError as e:
            raise ImportError(
                "llama-cpp-python is required: pip install llama-cpp-python"
            ) from e

        if chat_format not in CHAT_FORMATS:
            raise ValueError(
                f"Unknown chat_format {chat_format!r}; choose from {list(CHAT_FORMATS)}"
            )

        self._model_path  = str(Path(model_path).expanduser().resolve())
        self._session_id  = session_id
        self._state_dir   = Path(state_dir).expanduser()
        self._state_dir.mkdir(parents=True, exist_ok=True)
        self._chat_format = chat_format

        self._llama = Llama(
            model_path=self._model_path,
            n_ctx=n_ctx,
            n_gpu_layers=n_gpu_layers,
            verbose=verbose,
        )

        # Auto-load existing state if present
        sp = self._state_path()
        if sp.exists():
            core.load_state(self._llama, self._model_path, str(sp))

    # ── Public methods ─────────────────────────────────────────────

    def chat(self, prompt: str, max_tokens: int = 512) -> str:
        """
        Feed *prompt* to the model (wrapped in the active chat format) and
        return the generated text.

        For SSM models the hidden state accumulates in llama_context across
        calls — there is no explicit message history list, the recurrent
        state IS the memory.  Call save() at any checkpoint you want to
        resume from later.
        """
        template, stops = CHAT_FORMATS[self._chat_format]
        wrapped = template.format(prompt=prompt)
        result = self._llama(
            wrapped,
            max_tokens=max_tokens,
            echo=False,
            stop=stops,
        )
        return result["choices"][0]["text"].strip()

    def save(self, session_id: Optional[str] = None) -> Path:
        """Persist the current state.  Returns the path written."""
        path = self._state_path(session_id)
        core.save_state(self._llama, self._model_path, str(path))
        return path

    def load(self, session_id: Optional[str] = None) -> None:
        """Restore state from a (possibly different) session."""
        path = self._state_path(session_id)
        if not path.exists():
            raise FileNotFoundError(f"No state file found: {path}")
        core.load_state(self._llama, self._model_path, str(path))

    @property
    def state_size(self) -> int:
        """Current serialised byte size of the SSM hidden state."""
        return core.get_state_size(self._llama)

    @property
    def session_id(self) -> str:
        return self._session_id

    @property
    def model_path(self) -> str:
        return self._model_path

    # ── Private ────────────────────────────────────────────────────

    def _state_path(self, session_id: Optional[str] = None) -> Path:
        sid = session_id if session_id is not None else self._session_id
        return self._state_dir / f"{sid}.memb"

    def __repr__(self) -> str:
        return (
            f"Session(model={Path(self._model_path).name!r}, "
            f"session_id={self._session_id!r}, "
            f"state_dir={str(self._state_dir)!r})"
        )