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,
             temperature: float = 0.1, top_k: int = 1) -> str:
        """
        Feed *prompt* to the model (wrapped in the active chat format) and
        return the generated text.

        This uses raw tokenize+eval+sample rather than ``Llama.__call__`` or
        ``create_chat_completion`` because those high-level helpers re-tokenise
        the entire prompt and then reset the KV-cache when the prefix does not
        match — which destroys any state loaded from a .memb file.  Going
        through ``llama.eval()`` directly *appends* new tokens to the live
        state, which is exactly what memba needs.

        For SSM/hybrid 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)

        # Tokenise the wrapped turn.  Add BOS only when the context is fresh
        # (no prior tokens in either the live conversation or a loaded state).
        is_fresh = self._llama.n_tokens == 0
        new_tokens = self._llama.tokenize(
            wrapped.encode("utf-8"),
            add_bos=is_fresh,
            special=True,
        )

        # Append to the live state.  eval() does NOT reset the KV-cache.
        self._llama.eval(new_tokens)

        # Build the set of single-token stops + add EOS.
        stop_token_ids = {self._llama.token_eos()}
        for s in stops:
            for t in self._llama.tokenize(s.encode("utf-8"), add_bos=False, special=True):
                stop_token_ids.add(t)

        out_tokens: list[int] = []
        for _ in range(max_tokens):
            tok = self._llama.sample(top_k=top_k, temp=temperature)
            out_tokens.append(tok)
            # Always eval the sampled token so it lives in the state too —
            # that way the next chat() turn sees the assistant reply as
            # part of the conversation.
            self._llama.eval([tok])
            if tok in stop_token_ids:
                break
            # Multi-token stop-string check (some stops span several BPE pieces)
            if stops:
                snippet = self._llama.detokenize(out_tokens).decode("utf-8", errors="ignore")
                if any(s in snippet for s in stops):
                    break

        text = self._llama.detokenize(out_tokens).decode("utf-8", errors="ignore")
        for s in stops:
            if s in text:
                text = text.split(s)[0]
        return 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})"
        )