test_generation_gguf.py

"""
Test the game generation pipeline with llama.cpp and NVIDIA Nemotron 3 Nano 4B GGUF.

Run this script to:
1. Test model availability (llama-cpp-python)
2. Demonstrate GGUF model loading from HuggingFace
3. Generate games with the Nemotron model
4. Fall back to mock generation if model unavailable
5. Validate all outputs against schema
"""

import json
from pathlib import Path
from app.services.retrieval import load_games_dataset, normalize_game_record, retrieve_examples
from app.services.generator import generate_game, build_generation_prompt, NEMOTRON_MODEL_ID, NEMOTRON_GGUF_FILE
from app.services.schema_validator import validate_game_schema


def check_llama_cpp_availability():
    """Check if llama-cpp-python is installed."""
    print("\n" + "=" * 80)
    print("CHECKING ENVIRONMENT")
    print("=" * 80)
    
    try:
        import llama_cpp
        print(f"✓ llama-cpp-python is installed")
        print(f"  Version: {llama_cpp.__version__ if hasattr(llama_cpp, '__version__') else 'unknown'}")
        return True
    except ImportError:
        print("✗ llama-cpp-python not found")
        print("  Install with: pip install llama-cpp-python")
        print("  Or for GPU support: pip install llama-cpp-python[cuda]")
        return False


def main():
    print("\n" + "=" * 80)
    print("PHASE 2, TASK 6: GAME GENERATION WITH NEMOTRON 3 NANO 4B GGUF")
    print("=" * 80)
    
    # Check environment
    llama_cpp_available = check_llama_cpp_availability()
    
    print(f"\nModel Configuration:")
    print(f"  Repository: {NEMOTRON_MODEL_ID}")
    print(f"  File: {NEMOTRON_GGUF_FILE}")
    print(f"  Runtime: llama.cpp (GGUF quantized)")
    print(f"  Benefits:")
    print(f"    • GGUF quantization: 4-bit, memory efficient")
    print(f"    • llama.cpp: Fast CPU/GPU inference")
    print(f"    • Hackathon bonus: Extra credit for llama.cpp runtime")
    
    # Load dataset
    print("\n1. Loading dataset...")
    raw_records = load_games_dataset("app/data/games_dataset.json")
    normalized_records = [normalize_game_record(r) for r in raw_records]
    print(f"✓ Loaded {len(normalized_records)} records")
    
    # Test cases
    test_configs = [
        {
            "name": "Scavenger Hunt - Adults - Medium",
            "config": {
                "game_type": "scavenger_hunt",
                "city": "Paris",
                "area": "Le Marais",
                "location_type": "mixed",
                "duration_minutes": 60,
                "num_players": 4,
                "difficulty": "medium",
                "age_group": "adults"
            }
        },
        {
            "name": "Hide & Seek - Kids - Easy",
            "config": {
                "game_type": "hide_and_seek",
                "city": "Paris",
                "area": "Parc des Buttes-Chaumont",
                "location_type": "park",
                "duration_minutes": 45,
                "num_players": 5,
                "difficulty": "easy",
                "age_group": "kids"
            }
        }
    ]
    
    # Test generation
    results = []
    for test in test_configs:
        print("\n" + "=" * 80)
        print(f"TEST: {test['name']}")
        print("=" * 80)
        
        config = test['config']
        
        # Retrieve similar games
        print("\n2. Retrieving similar games...")
        retrieved = retrieve_examples(config, normalized_records, k=3)
        print(f"✓ Retrieved {len(retrieved)} examples")
        
        # Build prompt
        print("\n3. Building generation prompt...")
        prompt = build_generation_prompt(config, retrieved)
        print(f"✓ Prompt ready ({len(prompt)} chars)")
        
        # Generate game
        print("\n4. Generating game...")
        print(f"   (Using: NVIDIA Nemotron 3 Nano 4B GGUF via llama.cpp)")
        
        try:
            game = generate_game(config, retrieved)
            print(f"✓ Game generated: {game['game_id']}")
            
            # Validate
            print("\n5. Validating against schema...")
            is_valid, errors = validate_game_schema(game)
            
            if is_valid:
                print("✓ Game VALID against schema")
            else:
                print(f"✗ Validation errors: {len(errors)}")
            
            # Display summary
            print("\n6. Game Summary:")
            print(f"  Title: {game['title']}")
            print(f"  Area: {game['setup']['area']}")
            print(f"  Duration: {game['setup']['duration_minutes']} min | Players: {game['setup']['num_players']}")
            print(f"  Tasks: {len(game['tasks'])} | Rules: {len(game['rules'])}")
            print(f"  Tone: {game['story_seed']['tone']}")
            
            results.append({
                'name': test['name'],
                'valid': is_valid,
                'game_id': game['game_id']
            })
            
        except Exception as e:
            print(f"✗ Generation failed: {e}")
            results.append({
                'name': test['name'],
                'valid': False,
                'game_id': None
            })
    
    # Summary
    print("\n" + "=" * 80)
    print("TEST SUMMARY")
    print("=" * 80)
    
    passed = sum(1 for r in results if r['valid'])
    print(f"\nResults: {passed}/{len(results)} tests passed")
    
    for result in results:
        status = "✓ PASS" if result['valid'] else "✗ FAIL"
        print(f"{status}: {result['name']}")
        if result['game_id']:
            print(f"       {result['game_id']}")
    
    print("\n" + "=" * 80)
    print("NOTES")
    print("=" * 80)
    if not llama_cpp_available:
        print("⚠ llama-cpp-python not installed - using mock generation")
        print("  For actual model-based generation, install:")
        print("    pip install llama-cpp-python[cuda]  # For GPU")
        print("    or")
        print("    pip install llama-cpp-python        # For CPU")
    else:
        print("✓ llama-cpp-python available")
        print("✓ NVIDIA Nemotron 3 Nano 4B GGUF ready for download from HuggingFace")
        print("✓ Hackathon extra credit: llama.cpp runtime ✓")
    
    print("\n" + "=" * 80)


if __name__ == "__main__":
    main()