Memba/python/memba/cli.py

"""
memba CLI — typer-based entry point.

Sub-commands
------------
    memba chat   --model <gguf> --session <name>  [--gpu-layers N]   → REPL
    memba run    --model <gguf> --prompt  <text>  [--save-state <f>]  → one-shot
    memba list   [--state-dir <dir>]                                   → list sessions
    memba rm     <session>      [--state-dir <dir>]                    → delete session
    memba info   <session>      [--state-dir <dir>]                    → show metadata
"""

from __future__ import annotations

import struct
import sys
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console
from rich.table import Table
from rich.prompt import Prompt

from . import core
from .session import Session

app     = Console()
cli     = typer.Typer(
    name="memba",
    help="Persistent memory layer for SSM-based LLMs (Falcon-Mamba, Zamba).",
    add_completion=False,
)
console = Console()

_DEFAULT_STATE_DIR = "~/.memba/states"

# ── Helpers ────────────────────────────────────────────────────────

def _state_dir_path(state_dir: str) -> Path:
    return Path(state_dir).expanduser()


def _list_sessions(state_dir: Path) -> list[Path]:
    if not state_dir.exists():
        return []
    return sorted(state_dir.glob("*.memb"))


def _parse_memb_header(path: Path) -> dict:
    """Return header fields from a .memb file without loading the full blob."""
    HEADER_FMT  = "<4sI64sIIQ"
    HEADER_SIZE = struct.calcsize(HEADER_FMT)
    raw = path.read_bytes()
    if len(raw) < HEADER_SIZE:
        return {}
    magic, version, model_id, n_ctx, _, data_size = struct.unpack_from(HEADER_FMT, raw)
    return {
        "magic":      magic,
        "version":    version,
        "model_id":   model_id.rstrip(b"\x00").decode("ascii", errors="replace"),
        "n_ctx":      n_ctx,
        "data_size":  data_size,
        "file_size":  path.stat().st_size,
    }


# ── chat — interactive REPL ────────────────────────────────────────

@cli.command()
def chat(
    model:       str = typer.Option(...,  "--model",      "-m", help="Path to GGUF model"),
    session:     str = typer.Option("default", "--session", "-s", help="Session name"),
    gpu_layers:  int = typer.Option(0,    "--gpu-layers",  "-g", help="GPU layers (0=CPU, -1=all)"),
    n_ctx:       int = typer.Option(4096, "--n-ctx",             help="Context size"),
    state_dir:   str = typer.Option(_DEFAULT_STATE_DIR,   "--state-dir"),
    verbose:    bool = typer.Option(False, "--verbose",    "-v"),
    max_tokens:  int = typer.Option(512,  "--max-tokens",        help="Max tokens per turn"),
    chat_format: str = typer.Option("chatml", "--chat-format",
                                    help="Prompt template: chatml (instruct) | raw (base)"),
) -> None:
    """Interactive REPL with auto-save on exit (Ctrl-C or /exit)."""
    console.print(f"[bold cyan]memba[/bold cyan] — session [green]{session!r}[/green]")
    console.print(f"Model : [dim]{model}[/dim]")
    console.print(f"GPU   : {gpu_layers} layers   Format: [magenta]{chat_format}[/magenta]")
    console.print("Type [bold]/save[/bold] to checkpoint, [bold]/exit[/bold] or Ctrl-C to quit.\n")

    with console.status("Loading model…"):
        sess = Session(
            model_path=model,
            session_id=session,
            state_dir=state_dir,
            n_gpu_layers=gpu_layers,
            n_ctx=n_ctx,
            verbose=verbose,
            chat_format=chat_format,
        )
    console.print(f"[dim]State size: {sess.state_size:,} bytes[/dim]\n")

    try:
        while True:
            try:
                user_input = Prompt.ask("[bold]You[/bold]")
            except (EOFError, KeyboardInterrupt):
                break

            if not user_input.strip():
                continue
            if user_input.strip() == "/exit":
                break
            if user_input.strip() == "/save":
                path = sess.save()
                console.print(f"[dim]Saved → {path}[/dim]")
                continue

            with console.status("Thinking…"):
                reply = sess.chat(user_input, max_tokens=max_tokens)
            console.print(f"[bold green]Assistant[/bold green]: {reply}\n")

    except KeyboardInterrupt:
        pass

    console.print("\n[dim]Saving state…[/dim]")
    path = sess.save()
    console.print(f"[bold]Session saved:[/bold] {path}")


# ── run — one-shot with optional save/load ─────────────────────────

@cli.command()
def run(
    model:       str           = typer.Option(..., "--model",      "-m", help="Path to GGUF model"),
    prompt:      str           = typer.Option(..., "--prompt",     "-p", help="Prompt text"),
    save_state:  Optional[str] = typer.Option(None, "--save-state",      help="Write state to this path"),
    load_state:  Optional[str] = typer.Option(None, "--load-state",      help="Load state from this path"),
    gpu_layers:  int           = typer.Option(0,   "--gpu-layers", "-g"),
    n_ctx:       int           = typer.Option(4096, "--n-ctx"),
    max_tokens:  int           = typer.Option(512,  "--max-tokens"),
    verbose:     bool          = typer.Option(False, "--verbose",  "-v"),
) -> None:
    """One-shot generation with optional state save/load."""
    from llama_cpp import Llama

    with console.status("Loading model…"):
        llama = Llama(
            model_path=model,
            n_ctx=n_ctx,
            n_gpu_layers=gpu_layers,
            verbose=verbose,
        )

    if load_state:
        with console.status(f"Loading state from {load_state}…"):
            core.load_state(llama, model, load_state)
        console.print(f"[dim]State loaded: {load_state}[/dim]", file=sys.stderr)

    result = llama(prompt, max_tokens=max_tokens, echo=False)
    print(result["choices"][0]["text"])

    if save_state:
        core.save_state(llama, model, save_state)
        console.print(f"[dim]State saved: {save_state}[/dim]", file=sys.stderr)


# ── list ───────────────────────────────────────────────────────────

@cli.command(name="list")
def list_sessions(
    state_dir: str = typer.Option(_DEFAULT_STATE_DIR, "--state-dir"),
) -> None:
    """List all saved sessions."""
    sdir   = _state_dir_path(state_dir)
    files  = _list_sessions(sdir)

    if not files:
        console.print(f"[dim]No sessions in {sdir}[/dim]")
        return

    table = Table(title=f"Sessions in {sdir}", show_lines=False)
    table.add_column("Session",   style="bold cyan")
    table.add_column("State (B)", justify="right")
    table.add_column("File (B)",  justify="right")
    table.add_column("model_id",  style="dim", no_wrap=True)

    for f in files:
        h = _parse_memb_header(f)
        table.add_row(
            f.stem,
            f"{h.get('data_size', '?'):,}" if isinstance(h.get('data_size'), int) else "?",
            f"{h.get('file_size', '?'):,}"  if isinstance(h.get('file_size'), int) else "?",
            h.get("model_id", "?")[:16] + "…",
        )

    console.print(table)


# ── rm ─────────────────────────────────────────────────────────────

@cli.command()
def rm(
    session:   str = typer.Argument(..., help="Session name to remove"),
    state_dir: str = typer.Option(_DEFAULT_STATE_DIR, "--state-dir"),
    yes:      bool = typer.Option(False, "--yes", "-y", help="Skip confirmation"),
) -> None:
    """Delete a saved session."""
    sdir = _state_dir_path(state_dir)
    path = sdir / f"{session}.memb"

    if not path.exists():
        console.print(f"[red]Session not found:[/red] {path}")
        raise typer.Exit(1)

    if not yes:
        confirmed = typer.confirm(f"Delete {path}?")
        if not confirmed:
            raise typer.Abort()

    path.unlink()
    console.print(f"[dim]Deleted:[/dim] {path}")


# ── info ───────────────────────────────────────────────────────────

@cli.command()
def info(
    session:   str = typer.Argument(..., help="Session name"),
    state_dir: str = typer.Option(_DEFAULT_STATE_DIR, "--state-dir"),
) -> None:
    """Show metadata stored in a session's state file."""
    sdir = _state_dir_path(state_dir)
    path = sdir / f"{session}.memb"

    if not path.exists():
        console.print(f"[red]Session not found:[/red] {path}")
        raise typer.Exit(1)

    h = _parse_memb_header(path)
    if not h:
        console.print("[red]Cannot parse header (truncated file?)[/red]")
        raise typer.Exit(1)

    console.print(f"[bold]Session:[/bold] {session}")
    console.print(f"  File    : {path}")
    console.print(f"  Magic   : {h['magic']!r}")
    console.print(f"  Version : {h['version']}")
    console.print(f"  model_id: {h['model_id']}")
    console.print(f"  n_ctx   : {h['n_ctx']}")
    console.print(f"  State   : {h['data_size']:,} bytes")
    console.print(f"  File    : {h['file_size']:,} bytes")


def main() -> None:
    cli()


if __name__ == "__main__":
    main()