Mortdecai/training/prompts/loader.py

"""
Prompt template loader — shared between self-play and the chat app.

Loads prompt templates from per-category JSONL files with placeholder
expansion ({player}, {target}, {region}, etc.) and filtering by mode
(sudo/god/god_system) and call_type (model/gateway).

Usage:
    from training.prompts.loader import load_prompts, load_manifest, expand_template

    # Self-play: load all sudo prompts
    prompts = load_prompts(mode_filter="sudo")

    # Chat app: get manifest for template selector UI
    manifest = load_manifest()

    # Expand a template with random values
    expanded = expand_template("sudo protect {region} for {target}")
"""

import json
import random
from pathlib import Path

PROMPTS_DIR = Path(__file__).resolve().parent

TEMPLATE_VARS = {
    "player": ["slingshooter08", "Ace13245", "TheBigBoss", "xXDragonSlayerXx"],
    "target": ["Ace13245", "TheBigBoss", "xXDragonSlayerXx", "slingshooter08"],
    "region": ["my-base", "spawn-zone", "pvp-arena", "vip-lounge", "farm-area"],
    "warp": ["arena", "spawn", "shop", "nether", "farm", "end"],
    "group": ["vip", "builder", "moderator", "default"],
    "world": ["world", "world_nether", "world_the_end"],
}


def expand_template(prompt: str, overrides: dict = None) -> str:
    """Replace {placeholder} tokens with random values.

    Args:
        prompt: Template string with {placeholder} tokens.
        overrides: Optional dict of placeholder -> value to use instead of random.
                   e.g. {"player": "slingshooter08"} forces that player name.
    """
    overrides = overrides or {}
    for key, values in TEMPLATE_VARS.items():
        token = "{" + key + "}"
        while token in prompt:
            value = overrides.get(key, random.choice(values))
            prompt = prompt.replace(token, value, 1)
    return prompt


def load_manifest(prompts_dir: Path = PROMPTS_DIR) -> dict:
    """Load the prompt manifest with full metadata.

    Returns dict of category -> {file, mode, call_type, count}.
    Used by the chat app for template selection UI.
    """
    manifest_path = prompts_dir / "manifest.json"
    if not manifest_path.exists():
        return {}
    with open(manifest_path) as f:
        return json.load(f)


def load_prompts(prompts_dir: Path = PROMPTS_DIR,
                 mode_filter: str = None,
                 call_type_filter: str = None,
                 categories: list = None) -> dict[str, list[str]]:
    """Load prompt templates from per-category JSONL files.

    Args:
        prompts_dir: Directory containing manifest.json and prompt JSONL files.
        mode_filter: Only load categories matching this mode (sudo/god/god_system).
        call_type_filter: Only load categories matching call type (model/gateway).
        categories: Only load these specific category names.

    Returns:
        Dict mapping category name -> list of prompt template strings.
    """
    manifest = load_manifest(prompts_dir)
    if not manifest:
        return {}

    prompts = {}
    for category, meta in manifest.items():
        if categories and category not in categories:
            continue
        if mode_filter and meta.get("mode") not in (mode_filter, "mixed"):
            continue
        if call_type_filter and meta.get("call_type") != call_type_filter:
            continue

        filepath = prompts_dir / meta["file"]
        if not filepath.exists():
            continue
        cat_prompts = []
        with open(filepath) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                entry = json.loads(line)
                cat_prompts.append(entry["prompt"])
        if cat_prompts:
            prompts[category] = cat_prompts

    return prompts


def load_prompt_entries(category: str,
                        prompts_dir: Path = PROMPTS_DIR) -> list[dict]:
    """Load full prompt entries (with all metadata) for a single category.

    Returns list of dicts with prompt, category, mode, etc.
    Used by the chat app to populate the template editor.
    """
    manifest = load_manifest(prompts_dir)
    meta = manifest.get(category)
    if not meta:
        return []

    filepath = prompts_dir / meta["file"]
    if not filepath.exists():
        return []

    entries = []
    with open(filepath) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            entries.append(json.loads(line))
    return entries