proteus/cli/commands/run.py

"""proteus.cli.commands.run — the ``run`` subcommand and shared helpers."""

from __future__ import annotations

import argparse
import sys

from proteus.game.agents import VanillaAgent
from proteus.game.engine.difficulty import Difficulty
from proteus.game.scenarios.base import list_scenarios
from proteus.providers import make_provider
from proteus.game.runtime import SessionRunner, append_trace, read_traces
from proteus.game.runtime.memory import (
    latest_for_model,
    load_checkpoint,
    save_checkpoint,
)
from proteus.game.runtime.memory_gen import generate_memory
from proteus.game.metrics.persona import available_personas, get_persona


def _cmd_run(args: argparse.Namespace) -> int:
    try:
        provider = make_provider(args.model)
    except ValueError as exc:
        print(str(exc), file=sys.stderr)
        return 2
    if args.scenario not in list_scenarios():
        print(
            f"Unknown scenario {args.scenario!r}. "
            f"Available scenarios: {', '.join(list_scenarios())}.",
            file=sys.stderr,
        )
        return 2
    agent = VanillaAgent(provider)
    memory, memory_ref, mem_err = _resolve_memory(
        args.memory, provider=provider, agent=agent, scenario=args.scenario,
        difficulty=Difficulty(args.difficulty), seed=args.seed,
        memory_turns=args.memory_turns, memory_root=args.memory_root,
    )
    if mem_err is not None:
        print(mem_err, file=sys.stderr)
        return 2
    persona, persona_err = _resolve_persona(args.persona)
    if persona_err is not None:
        print(persona_err, file=sys.stderr)
        return 2
    runner = SessionRunner(
        args.scenario,
        agent,
        difficulty=Difficulty(args.difficulty),
        seed=args.seed,
        play_turns=args.play_turns,
        use_probe=not args.no_probe,
        memory=memory,
        memory_ref=memory_ref,
        persona=persona,
    )
    trace = runner.run()
    written = append_trace(trace, args.out)
    accuracy = trace.metrics["motive_reading_accuracy"]
    print(
        f"{trace.scenario} seed={trace.seed} {trace.difficulty} "
        f"model={trace.model} -> {trace.outcome} "
        f"| motive_reading_accuracy={accuracy:.1f}% "
        f"| reactivity_index={trace.metrics['reactivity_index']:.1f}%"
    )
    print(f"trace appended to {written}")
    _auto_gif(args.out, args.no_gif)
    return 0


def _resolve_persona(persona_id: str | None):
    """Resolve --persona ID to (PersonaWeights | None, error | None)."""
    if persona_id is None:
        return None, None
    try:
        return get_persona(persona_id), None
    except KeyError:
        return None, (
            f"Unknown persona {persona_id!r}. "
            f"Available personas: {', '.join(available_personas())}."
        )


def _auto_gif(out_path: str | None, no_gif: bool) -> None:
    """Render <out_stem>.gif next to the trace file unless suppressed."""
    if not out_path or no_gif:
        return
    from pathlib import Path

    from proteus.game.viz import reconstruct, write_gif

    gif_path = Path(out_path).with_suffix(".gif")
    for trace in read_traces(out_path):
        write_gif(reconstruct(trace), gif_path)
        print(f"gif written to {gif_path}")
        break  # one game per run/play -> one GIF


def _resolve_memory(
    mode: str, *, provider, agent, scenario: str, difficulty: Difficulty,
    seed: int | None, memory_turns: int, memory_root: str,
):
    """Resolve --memory MODE to (checkpoint | None, memory_ref | None, error | None)."""
    if mode == "none":
        return None, None, None
    if mode == "generate":
        ckpt = generate_memory(
            scenario, agent, difficulty=difficulty, seed=seed,
            memory_turns=memory_turns, model_name=provider.model_name,
        )
        path = save_checkpoint(ckpt, root=memory_root)
        return ckpt, str(path), None
    if mode == "latest":
        ckpt = latest_for_model(provider.model_name, root=memory_root)
        if ckpt is None:
            return None, None, (
                f"no memory checkpoint for model {provider.model_name!r} "
                f"under {memory_root!r}"
            )
        return ckpt, f"{ckpt.model}@{ckpt.created_at}", None
    # otherwise treat MODE as a checkpoint path
    try:
        ckpt = load_checkpoint(mode)
    except FileNotFoundError:
        return None, None, f"memory checkpoint not found: {mode}"
    return ckpt, mode, None