generate_prompts.py

#!/usr/bin/env python3
"""
Generate prompt.md files for MCode benchmarks.
Reads config.toml to determine which benchmarks to process.
Supports both 'server' and 'cli' benchmark types.

Templates are stored in templates/<version>/ directory as Jinja2 files:
- prompt_cli.md.j2    - Template for CLI benchmark prompts
- prompt_server.md.j2 - Template for server benchmark prompts

Template versions:
- v1: Original minimal template (used by initial experiments)
- v2: Improved template with drop-in replacement emphasis
"""

DEFAULT_TEMPLATE_VERSION = "v1"  # Default to v1 for backward compatibility
import yaml
import tomllib
import argparse
import sys
from pathlib import Path
from typing import Optional
from jinja2 import Environment, FileSystemLoader, TemplateNotFound


def load_config(config_path: Path) -> dict:
    """Load configuration from TOML file."""
    if not config_path.exists():
        return {}
    with open(config_path, 'rb') as f:
        return tomllib.load(f)


def load_benchmark_config(benchmark_dir: Path) -> Optional[dict]:
    """Load benchmark.yml from a benchmark directory."""
    benchmark_file = benchmark_dir / "benchmark.yml"
    metadata_file = benchmark_dir / "metadata.yml"

    if not benchmark_file.exists():
        return None

    with open(benchmark_file) as f:
        config = yaml.safe_load(f)

    # Skip benchmarks missing required fields
    if 'type' not in config.get('benchmark', {}):
        return None

    # Load ID from metadata.yml
    benchmark_id = "000"
    if metadata_file.exists():
        with open(metadata_file) as f:
            metadata = yaml.safe_load(f)
            benchmark_id = metadata.get("id", "000")

    config['_id'] = benchmark_id
    config['_dir'] = benchmark_dir
    return config


def discover_benchmarks(dataset_dir: Path) -> dict:
    """Discover all benchmarks in the dataset directory."""
    benchmarks = {}
    for item in dataset_dir.iterdir():
        if item.is_dir() and (item / "benchmark.yml").exists():
            config = load_benchmark_config(item)
            if config:
                benchmarks[config['_id']] = config
    return benchmarks


# Cache for Jinja2 environments (one per version)
_jinja_envs = {}

def get_jinja_env(version: str = DEFAULT_TEMPLATE_VERSION) -> Environment:
    """Get or create the Jinja2 environment for a specific template version."""
    global _jinja_envs
    if version not in _jinja_envs:
        templates_dir = Path(__file__).parent / "templates" / version
        if not templates_dir.exists():
            raise ValueError(f"Template version '{version}' not found at {templates_dir}")
        _jinja_envs[version] = Environment(
            loader=FileSystemLoader(templates_dir),
            trim_blocks=True,
            lstrip_blocks=True,
        )
    return _jinja_envs[version]


def generate_prompt(config: dict, template_version: str = DEFAULT_TEMPLATE_VERSION) -> str:
    """Generate prompt.md content from benchmark configuration using Jinja2 templates."""
    benchmark_type = config['benchmark']['type']
    source_lang = config['source']['language'].capitalize()
    source_desc = config['source']['description']
    target_lang = config['destination']['language'].capitalize()
    target_desc = config['destination']['description']

    # Build instructions (handle both install_cmd and build_cmd)
    install_cmd = config['destination'].get('install_cmd')
    build_cmd = config['destination'].get('build_cmd')

    build_steps = []
    if install_cmd:
        build_steps.append(f"Install dependencies: `{install_cmd}`")
    if build_cmd:
        build_steps.append(f"Build: `{build_cmd}`")

    if not build_steps:
        build_instructions = "No build steps required."
    else:
        build_instructions = "\n".join(f"- {step}" for step in build_steps)

    # Load and render template based on benchmark type
    env = get_jinja_env(template_version)
    template_name = f"prompt_{benchmark_type}.md.j2"

    try:
        template = env.get_template(template_name)
    except TemplateNotFound:
        raise ValueError(f"Unknown benchmark type: {benchmark_type} (template {template_name} not found)")

    # Prepare template variables
    template_vars = {
        'source_lang': source_lang,
        'source_desc': source_desc,
        'target_lang': target_lang,
        'target_desc': target_desc,
        'build_instructions': build_instructions,
        'run_cmd': config['destination'].get('run_cmd', ''),
    }

    # Add server-specific variables
    if benchmark_type == 'server':
        template_vars['port_env_var'] = config['destination'].get('port_env_var', 'SERVER_PORT')

    return template.render(**template_vars)


def main():
    parser = argparse.ArgumentParser(description='Generate prompt.md files for MCode benchmarks')
    parser.add_argument('--config', type=str, default='config.toml',
                        help='Path to config file (default: config.toml)')
    parser.add_argument('--all', action='store_true',
                        help='Generate prompts for all benchmarks, ignoring config.toml filter')
    parser.add_argument('--benchmark', type=str,
                        help='Generate prompt for a specific benchmark ID only')
    parser.add_argument('--template-version', type=str, default=DEFAULT_TEMPLATE_VERSION,
                        help=f'Template version to use (default: {DEFAULT_TEMPLATE_VERSION})')
    args = parser.parse_args()

    # Setup paths
    root_dir = Path(__file__).parent
    dataset_dir = root_dir / "dataset"
    config_path = root_dir / args.config

    # Load config
    config = load_config(config_path)
    benchmark_ids = config.get('benchmarks', {}).get('ids', [])

    # Discover benchmarks
    benchmarks = discover_benchmarks(dataset_dir)

    if not benchmarks:
        print("No benchmarks found in dataset/")
        sys.exit(1)

    # Filter benchmarks
    if args.benchmark:
        # Single benchmark by ID
        if args.benchmark not in benchmarks:
            print(f"Benchmark ID '{args.benchmark}' not found")
            print(f"Available IDs: {list(benchmarks.keys())}")
            sys.exit(1)
        benchmarks = {args.benchmark: benchmarks[args.benchmark]}
    elif not args.all and benchmark_ids:
        # Filter by config.toml
        benchmarks = {k: v for k, v in benchmarks.items() if k in benchmark_ids}
        if not benchmarks:
            print(f"No benchmarks found matching IDs in config: {benchmark_ids}")
            print(f"Available IDs: {list(discover_benchmarks(dataset_dir).keys())}")
            sys.exit(1)

    # Sort by ID
    sorted_benchmarks = sorted(benchmarks.items(), key=lambda x: x[0])

    print(f"Generating prompts for {len(sorted_benchmarks)} benchmark(s)...")
    print(f"Using template version: {args.template_version}")
    print("=" * 60)

    for benchmark_id, config in sorted_benchmarks:
        benchmark_name = config['benchmark']['name']
        benchmark_type = config['benchmark']['type']
        benchmark_dir = config['_dir']

        print(f"[{benchmark_id}] {benchmark_name} ({benchmark_type})")

        try:
            prompt = generate_prompt(config, args.template_version)
            prompt_path = benchmark_dir / "workspace" / "prompt.md"

            # Ensure workspace directory exists
            prompt_path.parent.mkdir(parents=True, exist_ok=True)

            with open(prompt_path, 'w') as f:
                f.write(prompt)

            print(f"  -> Generated {prompt_path.relative_to(root_dir)}")

        except KeyError as e:
            print(f"  -> ERROR: Missing required field {e}")
        except Exception as e:
            print(f"  -> ERROR: {e}")

    print("=" * 60)
    print("Done.")


if __name__ == "__main__":
    main()