brain/cli/chat.py

"""The Persona Lab REPL — master prompt 4.6.

    python -m brain.cli.chat --mode gf --mood playful
    python -m brain.cli.chat --offline          # no key needed; exercises every rule but the model

P1 needs only a Gemini key: no database, with memories stubbed from a YAML file.

Commands:
    /normal /gf           switch mode
    /mood <id>            set the mood (GF only)
    /prompt               show the assembled system prompt for the last turn
    /memories             show the stubbed memories in scope right now
    /photo /sticker       mark the next message as carrying an attachment
    /proactive            make her message first
    /reset                clear the history
    /quit
"""

import argparse
import asyncio
import os
from datetime import datetime
from pathlib import Path

import yaml
from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel

from .. import config
from ..llm.base import Message
from ..llm.factory import build_provider
from ..llm.fake import FakeLlm
from ..persona import prova
from ..persona.builder import PromptInput
from ..persona.types import Attachment, Mode
from ..turn import run_turn

MEMORIES_FILE = Path(__file__).with_name("memories.example.yaml")

THEME = {Mode.NORMAL: "cyan", Mode.GF: "magenta"}


def load_memories(path: Path) -> tuple[list[str], list[str]]:
    if not path.exists():
        return [], []
    raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
    return list(raw.get("shared", [])), list(raw.get("private", []))


class Repl:
    def __init__(
        self,
        provider,
        settings: config.Settings,
        mode: Mode,
        mood: str,
        memories_file: Path,
    ):
        self.provider = provider
        self.settings = settings
        self.mode = mode
        self.mood = mood
        self.console = Console()
        self.history: list[Message] = []
        self.shared, self.private = load_memories(memories_file)
        self.last_prompt = ""
        self.attachment = Attachment.NONE
        self.proactive = False
        self.tz = config.zone(settings.timezone)

    @property
    def colour(self) -> str:
        return THEME[self.mode]

    def banner(self) -> None:
        tagline = (
            "Executive Assistant"
            if self.mode is Mode.NORMAL
            else "AI Girlfriend cum Professional Assistant"
        )
        mood = f" · mood: {self.mood}" if self.mode is Mode.GF else ""
        self.console.print(
            Panel(
                f"[bold]P.R.O.V.A.[/bold]  {tagline}\n"
                f"mode: [bold]{self.mode.value}[/bold]{mood} · model: {self.settings.chat_model}",
                style=self.colour,
            )
        )

    def handle_command(self, line: str) -> bool:
        """True when the line was a command and no turn should run."""
        command, _, argument = line[1:].partition(" ")
        command = command.strip().lower()
        argument = argument.strip()

        if command in {"quit", "q", "exit"}:
            raise SystemExit(0)
        if command == "normal":
            self.mode = Mode.NORMAL
            self.banner()
        elif command == "gf":
            self.mode = Mode.GF
            self.banner()
        elif command == "mood":
            if argument and argument not in {m.id for m in prova.MOODS}:
                known = ", ".join(m.id for m in prova.MOODS)
                self.console.print(f"[red]Unknown mood.[/red] One of: {known}")
            else:
                self.mood = argument or self.mood
                self.console.print(f"[dim]mood: {self.mood}[/dim]")
        elif command == "prompt":
            self.console.print(Panel(self.last_prompt or "(no turn yet)", title="system prompt"))
        elif command == "memories":
            visible = self.shared + (self.private if self.mode is Mode.GF else [])
            hidden = 0 if self.mode is Mode.GF else len(self.private)
            body = "\n".join(f"- {fact}" for fact in visible) or "(none)"
            note = ""
            if hidden:
                note = f"\n\n[dim]{hidden} private fact(s) withheld in Normal mode.[/dim]"
            self.console.print(Panel(body + note, title="memories in scope"))
        elif command == "photo":
            self.attachment = Attachment.PHOTO
            self.console.print("[dim]next message carries a photo[/dim]")
        elif command == "sticker":
            self.attachment = Attachment.STICKER
            self.console.print("[dim]next message carries a sticker[/dim]")
        elif command == "proactive":
            self.proactive = True
            self.console.print("[dim]next turn: she messages first — just press enter[/dim]")
        elif command == "reset":
            self.history.clear()
            self.console.print("[dim]history cleared[/dim]")
        else:
            self.console.print(f"[red]Unknown command /{command}[/red]. See --help.")
        return True

    def build(self, message: str) -> PromptInput:
        now = datetime.now(self.tz).replace(tzinfo=None)
        return PromptInput(
            ai_name=self.settings.ai_name,
            user_name=self.settings.user_name,
            mode=self.mode,
            now=now,
            mood_id=self.mood,
            memories=tuple(self.shared),
            private_memories=tuple(self.private),
            attachment=self.attachment,
            proactive=self.proactive,
            last_message_at=now if self.history else None,
        )

    async def turn(self, message: str) -> None:
        result = await run_turn(self.provider, self.build(message), message, self.history)
        self.last_prompt = result.prompt
        self.attachment = Attachment.NONE
        self.proactive = False

        if message:
            self.history.append(Message("user", message))
        self.history.append(Message("model", result.text))
        self.history = self.history[-config.SHORT_TERM_CONTEXT_SIZE :]

        self.console.print()
        self.console.print(
            Panel(Markdown(result.display), title=self.settings.ai_name, style=self.colour)
        )

        notes = []
        if result.retries:
            notes.append(f"retries: {result.retries}")
        if result.fell_back:
            notes.append("fell back (leaked twice)")
        if result.stripped:
            notes.append("stripped by the Normal filter")
        if result.violations:
            found = ", ".join(f"{v.kind}={v.text!r}" for v in result.violations)
            notes.append(f"violations: {found}")
        if result.ambiguous:
            notes.append("ambiguity block sent")
        if result.technical:
            notes.append("technical turn (4096 tokens)")
        if result.error:
            notes.append(f"error: {result.error}")
        if result.model:
            notes.append(f"{result.model} · key {result.key_slot + 1} · {result.ms}ms")
        if notes:
            self.console.print(f"[dim]{' · '.join(notes)}[/dim]")

    async def run(self) -> None:
        self.banner()
        self.console.print(
            "[dim]/normal /gf /mood <id> /prompt /memories /photo /proactive /quit[/dim]"
        )
        while True:
            try:
                line = self.console.input(f"[{self.colour}]you ›[/{self.colour}] ").strip()
            except (EOFError, KeyboardInterrupt):
                return

            if line.startswith("/"):
                self.handle_command(line)
                continue
            if not line and not self.proactive:
                continue
            await self.turn(line)


def make_provider(offline: bool, settings: config.Settings):
    if offline:
        return FakeLlm(
            [
                "Aaj ki korle sara din? Ami sara bikel ei project tay atke chilam.",
                "Accha, bolo to — ki hoyeche?",
            ]
        )
    try:
        return build_provider(dict(os.environ), settings, allow_fake=False)
    except RuntimeError:
        raise SystemExit(
            "No GEMINI_KEY_1 in the environment.\n"
            "Copy backend/.env.example to backend/.env and add a key from aistudio.google.com,\n"
            "or run with --offline to exercise every rule except the model itself."
        ) from None


async def main_async(args: argparse.Namespace) -> None:
    settings = config.load_settings()
    provider = make_provider(args.offline, settings)
    repl = Repl(provider, settings, Mode(args.mode), args.mood, Path(args.memories))
    try:
        await repl.run()
    finally:
        if hasattr(provider, "aclose"):
            await provider.aclose()


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--mode", choices=["normal", "gf"], default="normal")
    parser.add_argument("--mood", default=prova.DEFAULT_MOOD_ID)
    parser.add_argument("--memories", default=str(MEMORIES_FILE))
    parser.add_argument("--offline", action="store_true", help="fake provider, no key needed")
    try:
        asyncio.run(main_async(parser.parse_args()))
    except SystemExit:
        raise
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()