subject_generator.py

# subject_generator.py
# Dynamically generates a list of age-appropriate school subjects per session.
#
# Key function:
#   generate_subject_candidates(basil_assessment, recent_subjects, n_min, n_max) -> list[str]

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

from config import TASK_AGENT_MODEL, PROMPT_SUBJECT_GENERATOR, SUBJECTS_FILE
from curriculum_manager import normalize_subject

from llm_client import create_smart_client
client = create_smart_client()

# Age band to human-ish age descriptor
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)",
}


def load_prompt_template() -> str:
    """Load the subject generator prompt template."""
    if os.path.exists(PROMPT_SUBJECT_GENERATOR):
        with open(PROMPT_SUBJECT_GENERATOR, "r") as f:
            return f.read()
    raise FileNotFoundError(f"Subject generator prompt not found: {PROMPT_SUBJECT_GENERATOR}")


def generate_subject_candidates(
    basil_assessment: dict,
    recent_subjects: list,
    n_min: int = 20,
    n_max: int = 50,
) -> list:
    """
    Generate a list of age-appropriate school subjects for this session.
    
    Args:
        basil_assessment: Dict with age_band, capabilities, etc.
        recent_subjects: List of recently used subject names to avoid
        n_min: Minimum number of candidates to return (will retry if needed)
        n_max: Maximum number of candidates to request from LLM
    
    Returns:
        List of subject name strings (deduplicated, filtered, display versions)
    """
    age_band = basil_assessment.get("age_band", 0)
    age_desc = AGE_BAND_TO_DESC.get(age_band, AGE_BAND_TO_DESC[0])
    
    # Format capabilities
    capabilities = basil_assessment.get("capabilities", [])
    if isinstance(capabilities, list):
        capabilities_text = "\n".join(f"- {c}" for c in capabilities)
    else:
        capabilities_text = str(capabilities)
    
    # Normalize recent subjects for comparison
    recent_normalized = set(normalize_subject(s) for s in recent_subjects)
    
    # Format recent subjects for prompt
    if recent_subjects:
        recent_subjects_text = ", ".join(recent_subjects[-25:])  # Last 25 for prompt
    else:
        recent_subjects_text = "(none)"
    
    # First attempt
    candidates = _call_subject_generator(
        age_desc=age_desc,
        capabilities=capabilities_text,
        recent_subjects_text=recent_subjects_text,
        n_max=n_max,
    )
    
    # Filter out subjects that match recent subjects (normalized comparison)
    filtered = []
    seen_normalized = set()
    
    for subj in candidates:
        normalized = normalize_subject(subj)
        
        # Skip if matches recent subject
        if normalized in recent_normalized:
            continue
        
        # Skip if duplicate within this batch
        if normalized in seen_normalized:
            continue
        
        seen_normalized.add(normalized)
        filtered.append(subj)
    
    # Check if we have enough candidates
    if len(filtered) < n_min:
        print(f"[Subject Generator] Only {len(filtered)} candidates after filtering, retrying with broader request...")
        
        # Retry with instruction to be more general
        retry_candidates = _call_subject_generator(
            age_desc=age_desc,
            capabilities=capabilities_text,
            recent_subjects_text=recent_subjects_text,
            n_max=n_max,
            broader=True,
        )
        
        # Add new candidates that aren't duplicates
        for subj in retry_candidates:
            normalized = normalize_subject(subj)
            if normalized not in recent_normalized and normalized not in seen_normalized:
                seen_normalized.add(normalized)
                filtered.append(subj)
    
    # If still not enough, relax blacklist (use only last 10)
    if len(filtered) < n_min and len(recent_subjects) > 10:
        print(f"[Subject Generator] Still only {len(filtered)} candidates, relaxing blacklist...")
        recent_relaxed = set(normalize_subject(s) for s in recent_subjects[-10:])
        
        # Re-filter with relaxed blacklist
        for subj in candidates:
            normalized = normalize_subject(subj)
            if normalized not in recent_relaxed and normalized not in seen_normalized:
                seen_normalized.add(normalized)
                filtered.append(subj)
    
    print(f"[Subject Generator] Generated {len(filtered)} subject candidates for age_band={age_band}")
    
    # Save candidates to JSON for debugging visibility
    try:
        data = {
            "generated_at": datetime.now().isoformat(),
            "age_band": age_band,
            "candidates": filtered,
        }
        os.makedirs(os.path.dirname(SUBJECTS_FILE), exist_ok=True)
        with open(SUBJECTS_FILE, "w") as f:
            json.dump(data, f, indent=2)
    except Exception as e:
        print(f"[Subject Generator] Warning: Failed to save candidates to JSON: {e}")
    
    return filtered


def _call_subject_generator(
    age_desc: str,
    capabilities: str,
    recent_subjects_text: str,
    n_max: int,
    broader: bool = False,
) -> list:
    """
    Call the LLM to generate subject candidates.
    
    Args:
        age_desc: Human-readable age description
        capabilities: Formatted capabilities text
        recent_subjects_text: Comma-separated recent subjects
        n_max: Max subjects to request
        broader: If True, request broader/more general subjects
    
    Returns:
        List of subject name strings (raw from LLM, not filtered)
    """
    template = load_prompt_template()
    
    prompt = template.format(
        age_desc=age_desc,
        capabilities=capabilities,
        recent_subjects=recent_subjects_text,
    )
    
    # Add broader instruction if needed
    if broader:
        prompt += "\n\nIMPORTANT: Be more general and creative. Include broader subject categories and less common but still child-appropriate subjects."
    
    system_msg = f"Generate exactly {n_max} school subjects. 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=0.9,  # Higher temp for variety
            max_tokens=1500,
        )
        
        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)
        
        # Extract subject names
        subjects_list = data.get("subjects", [])
        names = []
        for item in subjects_list:
            if isinstance(item, dict) and "name" in item:
                name = item["name"].strip()
                if name:
                    names.append(name)
            elif isinstance(item, str):
                name = item.strip()
                if name:
                    names.append(name)
        
        return names
        
    except json.JSONDecodeError as e:
        print(f"[Subject Generator] JSON parse error: {e}")
        print(f"[Subject Generator] Raw output: {raw_output[:500]}...")
        return _fallback_subjects()
    except Exception as e:
        print(f"[Subject Generator] Error: {e}")
        return _fallback_subjects()


def _fallback_subjects() -> list:
    """Return fallback subjects if generation fails."""
    return [
        "Language Arts",
        "Reading",
        "Writing",
        "Mathematics",
        "Numbers",
        "Shapes",
        "Science",
        "Nature",
        "Animals",
        "Weather",
        "Art",
        "Music",
        "Health",
        "Safety",
        "Feelings",
        "Manners",
        "Colors",
        "Counting",
        "Patterns",
        "Stories",
        "Bible Stories",
        "Saints",
        "Prayer",
        "Angels",
    ]


if __name__ == "__main__":
    # Test the subject generator
    print("Testing Subject Generator...")
    print()
    
    # Test with different age bands
    for age_band in [0, 2, 4]:
        print(f"=== Age Band {age_band} ===")
        
        test_assessment = {
            "age_band": age_band,
            "capabilities": [
                "Basic pattern recognition",
                "Some word production",
            ],
        }
        
        recent = ["Mathematics", "Science", "Reading"]
        
        candidates = generate_subject_candidates(
            basil_assessment=test_assessment,
            recent_subjects=recent,
            n_min=10,
            n_max=25,
        )
        
        print(f"Generated {len(candidates)} candidates:")
        for i, subj in enumerate(candidates, 1):
            print(f"  {i:>3}. {subj}")
        print()