cli.py

"""
Command-line interface for LLM Playground.

For users who prefer terminal-based interaction.
"""

import argparse
import sys
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent))

import config
from models import get_model
from logger import get_logger
from experiments.zero_shot import run_zero_shot_experiment, EXAMPLE_TASKS
from experiments.few_shot import run_few_shot_experiment, EXAMPLE_SCENARIOS
from experiments.sampling_params import run_temperature_experiment
from experiments.prompt_sensitivity import run_prompt_sensitivity_experiment


def main():
    """CLI entry point."""
    parser = argparse.ArgumentParser(
        description="LLM Playground - Experiment with Large Language Models",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    
    # Model configuration
    parser.add_argument(
        "--provider",
        choices=["ollama", "openai"],
        default="ollama",
        help="Model provider (default: ollama)",
    )
    
    parser.add_argument(
        "--model",
        default=config.DEFAULT_MODEL,
        help=f"Model name (default: {config.DEFAULT_MODEL})",
    )
    
    # Sampling parameters
    parser.add_argument(
        "--temperature",
        type=float,
        default=config.DEFAULT_TEMPERATURE,
        help=f"Sampling temperature (default: {config.DEFAULT_TEMPERATURE})",
    )
    
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=config.DEFAULT_MAX_TOKENS,
        help=f"Maximum tokens to generate (default: {config.DEFAULT_MAX_TOKENS})",
    )
    
    parser.add_argument(
        "--top-p",
        type=float,
        default=config.DEFAULT_TOP_P,
        help=f"Nucleus sampling parameter (default: {config.DEFAULT_TOP_P})",
    )
    
    # Subcommands
    subparsers = parser.add_subparsers(dest="command", help="Command to run")
    
    # Generate command
    generate_parser = subparsers.add_parser("generate", help="Generate text from prompt")
    generate_parser.add_argument("prompt", help="Input prompt")
    
    # Zero-shot command
    zero_shot_parser = subparsers.add_parser("zero-shot", help="Run zero-shot experiment")
    zero_shot_parser.add_argument(
        "--task",
        choices=list(EXAMPLE_TASKS.keys()),
        help="Example task to run",
    )
    zero_shot_parser.add_argument(
        "--prompt",
        help="Custom prompt",
    )
    
    # Temperature command
    temp_parser = subparsers.add_parser("temperature", help="Test temperature effects")
    temp_parser.add_argument("prompt", help="Input prompt")
    temp_parser.add_argument(
        "--temps",
        nargs="+",
        type=float,
        default=[0.1, 0.7, 1.5],
        help="Temperatures to test (default: 0.1 0.7 1.5)",
    )
    temp_parser.add_argument(
        "--samples",
        type=int,
        default=1,
        help="Samples per temperature (default: 1)",
    )
    
    # List models command
    list_parser = subparsers.add_parser("list-models", help="List available models")
    
    args = parser.parse_args()
    
    # Initialize logger
    logger = get_logger()
    
    # Handle commands
    if args.command == "generate":
        generate_command(args, logger)
    
    elif args.command == "zero-shot":
        zero_shot_command(args, logger)
    
    elif args.command == "temperature":
        temperature_command(args, logger)
    
    elif args.command == "list-models":
        list_models_command(args)
    
    else:
        parser.print_help()


def generate_command(args, logger):
    """Generate text from a prompt."""
    print(f"\n🚀 Generating with {args.provider}:{args.model}")
    print(f"   Temperature: {args.temperature}, Max Tokens: {args.max_tokens}")
    print(f"\n📝 Prompt:\n{args.prompt}\n")
    
    try:
        # Create model
        model = get_model(args.provider, args.model)
        
        # Generate
        response = model.generate(
            prompt=args.prompt,
            temperature=args.temperature,
            max_tokens=args.max_tokens,
            top_p=args.top_p,
        )
        
        # Log
        logger.log_interaction(
            prompt=args.prompt,
            response=response,
            parameters={
                "temperature": args.temperature,
                "max_tokens": args.max_tokens,
                "top_p": args.top_p,
            },
            experiment_type="cli_generate",
        )
        
        # Display
        print(f"💬 Response:\n{response.text}\n")
        print(f"📊 Metrics:")
        print(f"   Tokens: {response.total_tokens} "
              f"(prompt: {response.prompt_tokens}, completion: {response.completion_tokens})")
        print(f"   Latency: {response.latency_ms:.0f}ms")
        print(f"   Speed: {response.latency_ms/response.completion_tokens:.1f}ms/token")
        if response.cost_usd > 0:
            print(f"   Cost: ${response.cost_usd:.4f}")
        
        print(f"\n✅ Logged to: {logger.get_log_file_path()}")
        
    except Exception as e:
        print(f"\n❌ Error: {e}")
        sys.exit(1)


def zero_shot_command(args, logger):
    """Run zero-shot experiment."""
    # Get prompt
    if args.prompt:
        prompt = args.prompt
    elif args.task:
        prompt = EXAMPLE_TASKS[args.task]
    else:
        print("❌ Error: Provide either --task or --prompt")
        sys.exit(1)
    
    print(f"\n🎯 Zero-Shot Experiment with {args.provider}:{args.model}")
    print(f"\n📝 Task:\n{prompt}\n")
    
    try:
        model = get_model(args.provider, args.model)
        
        response = run_zero_shot_experiment(
            model=model,
            task=prompt,
            logger=logger,
            temperature=args.temperature,
            max_tokens=args.max_tokens,
        )
        
        print(f"💬 Response:\n{response.text}\n")
        print(f"📊 Tokens: {response.total_tokens} | Latency: {response.latency_ms:.0f}ms")
        print(f"\n✅ Logged to: {logger.get_log_file_path()}")
        
    except Exception as e:
        print(f"\n❌ Error: {e}")
        sys.exit(1)


def temperature_command(args, logger):
    """Test temperature effects."""
    print(f"\n🌡️ Temperature Experiment with {args.provider}:{args.model}")
    print(f"   Temperatures: {args.temps}")
    print(f"   Samples: {args.samples}")
    print(f"\n📝 Prompt:\n{args.prompt}\n")
    
    try:
        model = get_model(args.provider, args.model)
        
        results = run_temperature_experiment(
            model=model,
            prompt=args.prompt,
            temperatures=args.temps,
            logger=logger,
            max_tokens=args.max_tokens,
            num_samples=args.samples,
        )
        
        # Display results
        for temp in sorted(results.keys()):
            print(f"\n{'='*80}")
            print(f"Temperature: {temp}")
            print(f"{'='*80}")
            
            for i, response in enumerate(results[temp], 1):
                print(f"\nSample {i}:")
                print(response.text)
                print(f"[{response.completion_tokens} tokens, {response.latency_ms:.0f}ms]")
        
        print(f"\n✅ Logged to: {logger.get_log_file_path()}")
        
    except Exception as e:
        print(f"\n❌ Error: {e}")
        sys.exit(1)


def list_models_command(args):
    """List available models."""
    print(f"\n📋 Available Models for {args.provider}:\n")
    
    try:
        if args.provider == "ollama":
            from models.ollama_model import OllamaModel
            model = OllamaModel()
            models = model.list_available_models()
            
            for i, model_name in enumerate(models, 1):
                print(f"   {i}. {model_name}")
        
        else:
            print("   OpenAI models:")
            print("   1. gpt-3.5-turbo")
            print("   2. gpt-4")
            print("   3. gpt-4-turbo")
        
        print()
        
    except Exception as e:
        print(f"\n❌ Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()