story_genre_generator.py

# story_genre_generator.py
# Dynamically generates a list of story genres/traditions for storytime sessions.
#
# Mirrors the pattern in subject_generator.py:
#   1. LLM generates 50 genres at high temperature
#   2. Python picks a random genre from the list
#   3. The list is cached to disk and refreshed when exhausted or stale
#
# Key functions:
#   generate_story_genres(basil_assessment) -> list[str]
#   pick_genre(basil_assessment) -> str

import os
import json
import random
from datetime import datetime
from openai import OpenAI

from config import (
    TASK_AGENT_MODEL,
    PROMPT_STORY_GENRE_GENERATOR,
    STORY_GENRES_FILE,
    STORYTIME_DIR,
)
from file_lock_utils import get_lock

from llm_client import create_smart_client
client = create_smart_client()

# Age band to human-ish age descriptor (same as subject_generator)
AGE_BAND_TO_DESC = {
    0: "newborn (0-6 months old)",
    1: "infant (6-12 months old)",
    2: "toddler (12-18 months old)",
    3: "toddler (18-24 months old)",
    4: "preschool age (2-3 years old)",
    5: "kindergarten age (4-5 years old)",
    6: "early elementary age (6-8 years old)",
    7: "late elementary / middle school age (9-12 years old)",
}

# Fallback genres if LLM generation fails — Western-tradition focused
FALLBACK_GENRES = [
    "Aesop's Fables",
    "Brothers Grimm Fairy Tales",
    "Hans Christian Andersen Stories",
    "Greek Myths",
    "Norse Myths",
    "Old Testament Stories",
    "New Testament Parables",
    "Lives of Catholic Saints",
    "Arabian Nights Stories",
    "Russian Folktales",
    "Irish Fairy Tales",
    "Beatrix Potter Stories",
    "Rudyard Kipling's Just So Stories",
    "Public Domain Children's Picture Books",
    "Celtic Myths",
    "Roman Myths",
    "Andrew Lang's Fairy Books",
    "Charles Perrault Tales",
    "King Arthur and the Knights of the Round Table",
    "Robin Hood Legends",
    "Appalachian Jack Tales",
    "Brer Rabbit Stories",
    "L. Frank Baum's Oz Stories",
    "Lewis Carroll's Alice Stories",
    "Scandinavian Fairy Tales",
    "Scottish Fairy Tales",
    "English Fairy Tales (Joseph Jacobs)",
    "Stories from the Desert Fathers",
    "La Fontaine's Fables",
    "Tall Tales of the American Frontier",
]


def _load_prompt_template() -> str:
    """Load the story genre generator prompt template."""
    if os.path.exists(PROMPT_STORY_GENRE_GENERATOR):
        with open(PROMPT_STORY_GENRE_GENERATOR, "r") as f:
            return f.read()
    raise FileNotFoundError(
        f"Story genre generator prompt not found: {PROMPT_STORY_GENRE_GENERATOR}"
    )


def _load_cached_genres() -> dict:
    """Load cached genres from disk. Returns full data dict or empty dict."""
    if not os.path.exists(STORY_GENRES_FILE):
        return {}
    try:
        with open(STORY_GENRES_FILE, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, IOError):
        return {}


def _save_genres(genres: list, age_band: int):
    """Save genres list to the cache file."""
    data = {
        "generated_at": datetime.now().isoformat(),
        "age_band": age_band,
        "genres": genres,
    }
    os.makedirs(os.path.dirname(STORY_GENRES_FILE), exist_ok=True)
    with open(STORY_GENRES_FILE, "w") as f:
        json.dump(data, f, indent=2)


def generate_story_genres(
    basil_assessment: dict,
    recent_genres: list = None,
    n: int = 50,
) -> list:
    """
    Generate a diverse list of story genres/traditions via LLM.

    Args:
        basil_assessment: Dict with age_band, capabilities, etc.
        recent_genres: List of recently used genre names to avoid.
        n: Number of genres to request.

    Returns:
        List of genre name strings (deduplicated, filtered).
    """
    age_band = basil_assessment.get("age_band", 0)
    age_desc = AGE_BAND_TO_DESC.get(age_band, AGE_BAND_TO_DESC[0])

    if recent_genres:
        recent_text = ", ".join(recent_genres[-30:])  # Last 30 for prompt
    else:
        recent_text = "(none)"

    template = _load_prompt_template()
    prompt = template.format(
        age_desc=age_desc,
        recent_genres=recent_text,
    )

    system_msg = f"Generate exactly {n} story genres/traditions. Output valid JSON only."

    try:
        response = client.chat.completions.create(
            model=TASK_AGENT_MODEL,
            messages=[
                {"role": "system", "content": system_msg},
                {"role": "user", "content": prompt},
            ],
            temperature=1.0,  # High temp for maximum variety
            max_tokens=2000,
        )

        raw_output = response.choices[0].message.content.strip()

        # Parse JSON, handling potential markdown code blocks
        if raw_output.startswith("```"):
            lines = raw_output.split("\n")
            raw_output = "\n".join(lines[1:-1])

        data = json.loads(raw_output)
        genres_list = data.get("genres", [])

        # Deduplicate (case-insensitive) and filter empties
        seen = set()
        recent_lower = set(g.lower().strip() for g in (recent_genres or []))
        filtered = []

        for genre in genres_list:
            if isinstance(genre, str):
                genre = genre.strip()
                if not genre:
                    continue
                key = genre.lower()
                if key in seen or key in recent_lower:
                    continue
                seen.add(key)
                filtered.append(genre)

        print(f"[StoryGenre] Generated {len(filtered)} genres for age_band={age_band}")

        # Save to cache
        _save_genres(filtered, age_band)

        return filtered

    except json.JSONDecodeError as e:
        print(f"[StoryGenre] JSON parse error: {e}")
        return FALLBACK_GENRES[:]
    except Exception as e:
        print(f"[StoryGenre] Error: {e}")
        return FALLBACK_GENRES[:]


def pick_genre(basil_assessment: dict, used_genres_this_session: list = None) -> str:
    """
    Pick a random story genre for the current storytime session.

    Loads the cached genre list, regenerating it if empty or missing.
    Avoids genres already used in this session (if provided).

    Args:
        basil_assessment: Dict with age_band, capabilities, etc.
        used_genres_this_session: Genres already tried this session (for retries).

    Returns:
        A genre string, e.g. "Aesop's Fables".
    """
    used_this_session = set(
        g.lower().strip() for g in (used_genres_this_session or [])
    )

    # Try loading cached genres first
    with get_lock(STORY_GENRES_FILE):
        cached = _load_cached_genres()
        genres = cached.get("genres", [])

        # Filter out genres already tried this session
        available = [g for g in genres if g.lower().strip() not in used_this_session]

        if len(available) < 5:
            # Cache exhausted or too few remaining — regenerate
            print(f"[StoryGenre] Only {len(available)} genres available, regenerating...")
            genres = generate_story_genres(
                basil_assessment,
                recent_genres=list(used_this_session),
            )
            available = [g for g in genres if g.lower().strip() not in used_this_session]

    if not available:
        # Ultimate fallback
        available = [
            g for g in FALLBACK_GENRES
            if g.lower().strip() not in used_this_session
        ]
        if not available:
            available = FALLBACK_GENRES[:]

    chosen = random.choice(available)
    return chosen


if __name__ == "__main__":
    # Test the genre generator
    print("Testing Story Genre Generator...")
    print()

    test_assessment = {
        "age_band": 0,
        "capabilities": ["Basic pattern recognition"],
    }

    genres = generate_story_genres(test_assessment)
    print(f"\nGenerated {len(genres)} genres:")
    for i, genre in enumerate(genres, 1):
        print(f"  {i:>3}. {genre}")

    print(f"\n--- Random picks ---")
    for _ in range(5):
        genre = pick_genre(test_assessment)
        print(f"  Picked: {genre}")