cli.py

"""CLI entrypoint for RL-NPO.

Usage:
    python cli.py --text "Your text here" --target memory --generations 10
    python cli.py --list-targets
    python cli.py --build-targets
"""

import argparse
import sys
import os
import logging

from dotenv import load_dotenv

load_dotenv()


def main():
    parser = argparse.ArgumentParser(
        description="RL-NPO: Reinforcement Learning via Neural Prompt Optimization",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run optimization on a text string
  python cli.py --text "Black holes create gravitational singularities..." --target memory --generations 8

  # Read text from a file
  python cli.py --text-file input.txt --target emotion --generations 10 --mutations 5

  # List available cognitive targets
  python cli.py --list-targets

  # Build target vectors (run once before first optimization)
  python cli.py --build-targets

  # Use a custom target vector
  python cli.py --text "..." --target custom --custom-target targets/custom.npy

  # Build a custom target from your own text file
  python cli.py --build-custom-target my_reference.txt --custom-target-name my_style
        """,
    )

    # Input
    input_group = parser.add_mutually_exclusive_group()
    input_group.add_argument(
        "--text", type=str, default=None,
        help="Input text to optimize.",
    )
    input_group.add_argument(
        "--text-file", type=str, default=None,
        help="Path to a text file containing the input text.",
    )

    # Target
    parser.add_argument(
        "--target", type=str, default="memory",
        choices=["memory", "emotion", "attention", "language", "narrative", "custom"],
        help="Cognitive target ROI (default: memory).",
    )
    parser.add_argument(
        "--custom-target", type=str, default=None,
        help="Path to a custom .npy target vector (required if --target custom).",
    )

    # Optimization parameters
    parser.add_argument(
        "--generations", type=int, default=10,
        help="Number of evolutionary generations (default: 10).",
    )
    parser.add_argument(
        "--mutations", type=int, default=5,
        help="Candidate rewrites per generation (default: 5).",
    )
    parser.add_argument(
        "--run-id", type=str, default="",
        help="Unique run identifier (auto-generated if empty).",
    )
    parser.add_argument(
        "--output-dir", type=str, default="runs",
        help="Output directory for telemetry and results (default: runs).",
    )
    parser.add_argument(
        "--model", type=str, default="anthropic/claude-3.5-haiku",
        help="OpenRouter model for mutation generation (default: anthropic/claude-3.5-haiku).",
    )
    parser.add_argument(
        "--convergence-threshold", type=float, default=0.005,
        help="Minimum score delta to count as improvement (default: 0.005).",
    )
    parser.add_argument(
        "--convergence-patience", type=int, default=3,
        help="Generations without improvement before stopping (default: 3).",
    )

    # Utility flags
    parser.add_argument(
        "--list-targets", action="store_true",
        help="List available cognitive targets with descriptions.",
    )
    parser.add_argument(
        "--build-targets", action="store_true",
        help="Generate target vectors for all predefined ROIs (run once).",
    )
    parser.add_argument(
        "--build-masks", action="store_true",
        help="Generate and save ROI masks.",
    )
    parser.add_argument(
        "--build-custom-target", type=str, default=None, metavar="FILE",
        help="Build a custom target vector from a text (.txt) file. "
             "Output saved to targets/<name>.npy.",
    )
    parser.add_argument(
        "--custom-target-name", type=str, default="custom",
        help="Name for the custom target (default: 'custom'). Output: targets/<name>.npy.",
    )
    parser.add_argument(
        "--verbose", action="store_true",
        help="Enable verbose logging.",
    )

    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.INFO)
    else:
        logging.basicConfig(level=logging.WARNING)

    # ─── Handle utility flags ───────────────────────────────────────

    if args.list_targets:
        from roi_masks import ROI_LABEL_SETS, ROI_DESCRIPTIONS
        from rich.console import Console
        from rich.table import Table
        from rich import box

        console = Console()
        table = Table(
            title="RL-NPO Cognitive Targets",
            box=box.ROUNDED,
            border_style="cyan",
        )
        table.add_column("Target", style="bold yellow")
        table.add_column("Brain Regions", style="dim")
        table.add_column("Description", style="green")

        for target, labels in ROI_LABEL_SETS.items():
            table.add_row(
                target,
                ", ".join(labels),
                ROI_DESCRIPTIONS[target],
            )

        console.print(table)
        return

    if args.build_masks:
        from roi_masks import save_masks
        print("Building ROI masks...")
        save_masks()
        print("Done.")
        return

    if args.build_targets:
        from target_builder import build_all_targets
        build_all_targets(openrouter_model=args.model)
        return

    if args.build_custom_target:
        from target_builder import generate_custom_target
        input_file = args.build_custom_target
        if not os.path.exists(input_file):
            print(f"Error: File not found: {input_file}", file=sys.stderr)
            sys.exit(1)
        name = args.custom_target_name
        output_path = os.path.join("targets", f"{name}.npy")
        print(f"Building custom target '{name}' from: {input_file}")
        print(f"Output: {output_path}")
        generate_custom_target(
            input_path=input_file,
            output_path=output_path,
            name=name,
        )
        print(f"\nDone. Use with: python cli.py --text \"...\" --target custom --custom-target {output_path}")
        return

    # ─── Validate inputs for optimization ───────────────────────────

    input_text = None
    if args.text:
        input_text = args.text
    elif args.text_file:
        if not os.path.exists(args.text_file):
            print(f"Error: File not found: {args.text_file}", file=sys.stderr)
            sys.exit(1)
        with open(args.text_file, "r", encoding="utf-8") as f:
            input_text = f.read().strip()

    if input_text is None:
        parser.print_help()
        print("\nError: --text or --text-file is required for optimization.", file=sys.stderr)
        sys.exit(1)

    if len(input_text) < 50:
        print(
            "Warning: Input text is very short. TRIBE v2 works best with "
            "200+ words. Results may be noisy.",
            file=sys.stderr,
        )

    if args.target == "custom" and not args.custom_target:
        print(
            "Error: --custom-target PATH is required when --target=custom.",
            file=sys.stderr,
        )
        sys.exit(1)

    # Check that target vector exists
    if args.target != "custom":
        target_path = os.path.join("targets", f"{args.target}.npy")
        if not os.path.exists(target_path):
            print(
                f"Error: Target vector not found: {target_path}\n"
                f"Run `python cli.py --build-targets` first.",
                file=sys.stderr,
            )
            sys.exit(1)

    # ─── Run optimization ───────────────────────────────────────────

    from agent import NPOConfig, run_optimization

    config = NPOConfig(
        input_text=input_text,
        target=args.target,
        custom_target_path=args.custom_target or "",
        generations=args.generations,
        mutations=args.mutations,
        run_id=args.run_id,
        output_dir=args.output_dir,
        openrouter_model=args.model,
        convergence_threshold=args.convergence_threshold,
        convergence_patience=args.convergence_patience,
    )

    result = run_optimization(config)

    # Exit with appropriate code
    if result["score_delta"] > 0:
        sys.exit(0)
    else:
        sys.exit(0)  # Still success, just no improvement


if __name__ == "__main__":
    main()