task_agent.py

# task_agent.py
# Generates tasks and rubrics for evaluating Basil's responses.

import os
import json
import re
from openai import OpenAI

from config import (
    TASK_AGENT_MODEL,
    PROMPT_TASK_AGENT,
    PROMPT_TASK_GENERATOR,
    PROMPT_TASK_SELECTOR,
    BASIL_MAX_TOKENS,
    AGE_BAND_DESCRIPTIONS,
    get_task_guidance_for_age_band,
    get_age_band_natural_description,
)
from task_contract import normalize_token, normalize_targets

from llm_client import create_smart_client
client = create_smart_client()



def load_prompt_template() -> str:
    """Load the task agent prompt template."""
    with open(PROMPT_TASK_AGENT, "r") as f:
        return f.read()


def generate_task(
    recent_log_snippet: str,
    content_snippet: str = "",
    subject_of_the_day: str = None,
    lesson_of_the_day: str = None,
    age_band: int = 0,
    basil_max_tokens: int = None,
    allowed_keywords: list = None,
    previous_askers: list = None,
    recent_targets: list = None,
) -> dict:
    """
    Generate a task and rubric for Basil.
    
    Args:
        recent_log_snippet: Recent conversation lines (last ~20-60 lines)
        content_snippet: Optional content (fable, story) to reference
        subject_of_the_day: Current session subject
        lesson_of_the_day: Current session lesson
        age_band: Basil's current developmental stage (0-7)
        basil_max_tokens: Maximum tokens for Basil's response
        allowed_keywords: List of on-topic keywords for this lesson
        previous_askers: List of recent askers for variety
        recent_targets: List of recent target words used (for variety)
    
    Returns:
        dict with keys: task_text, task_category, stage_hint, asker, rubric, 
                       grader_instructions, max_tokens
    """
    template = load_prompt_template()
    
    if basil_max_tokens is None:
        basil_max_tokens = BASIL_MAX_TOKENS
    
    if allowed_keywords is None:
        allowed_keywords = []
    
    if previous_askers is None:
        previous_askers = []
    
    if recent_targets is None:
        recent_targets = []
    
    age_band_desc = AGE_BAND_DESCRIPTIONS.get(age_band, "unknown")
    
    # Format allowed keywords for prompt
    keywords_str = ", ".join(allowed_keywords) if allowed_keywords else "(derive from lesson)"
    
    # Format previous askers for variety guidance
    askers_str = ", ".join(previous_askers[-3:]) if previous_askers else "(none yet)"
    
    # Format recent targets for variety guidance
    targets_str = ", ".join(recent_targets[-4:]) if recent_targets else "(none yet)"
    
    prompt = template.format(
        tutor_summary="(No summary available)",  # Placeholder for removed story files
        recent_log_snippet=recent_log_snippet or "(No conversation yet)",
        content_snippet=content_snippet or "(No content snippet provided)",
        subject_of_the_day=subject_of_the_day or "(No subject)",
        lesson_of_the_day=lesson_of_the_day or "(No lesson)",
        age_band=age_band,
        age_band_description=age_band_desc,
        basil_max_tokens=basil_max_tokens,
        allowed_keywords=keywords_str,
        previous_askers=askers_str,
        recent_targets=targets_str,
    )
    
    # Debug log
    print(f"[Task Agent] Generating task for age_band={age_band}, subject={subject_of_the_day}, lesson={lesson_of_the_day}")
    
    try:
        response = client.chat.completions.create(
            model=TASK_AGENT_MODEL,
            messages=[
                {"role": "system", "content": "You are a task generation agent. Output valid JSON only."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.7,
            max_tokens=500,
        )
        
        raw_output = response.choices[0].message.content.strip()
        
        # Parse JSON, handling potential markdown code blocks
        if raw_output.startswith("```"):
            # Remove markdown code block
            lines = raw_output.split("\n")
            raw_output = "\n".join(lines[1:-1])
        
        task_data = json.loads(raw_output)
        
        # Validate required fields
        required_fields = ["task_text", "task_category", "stage_hint", "rubric", "grader_instructions"]
        for field in required_fields:
            if field not in task_data:
                raise ValueError(f"Missing required field: {field}")
        
        # Set default max_tokens if not provided (use session's basil_max_tokens)
        if "max_tokens" not in task_data:
            task_data["max_tokens"] = basil_max_tokens
        
        # Set default asker if not provided (backward compatibility)
        if "asker" not in task_data:
            task_data["asker"] = "tutor"
        else:
            # Normalize asker to lowercase
            task_data["asker"] = task_data["asker"].lower().strip()
            if task_data["asker"] not in ("tutor", "sophie"):
                task_data["asker"] = "tutor"
        
        # Normalize targets using central normalize_targets function
        raw_targets = task_data.get("targets", [])
        if not isinstance(raw_targets, list):
            raw_targets = [raw_targets] if raw_targets else []
        
        task_data["targets"] = normalize_targets(raw_targets)
        
        # If targets still empty after normalization, use fallback extraction
        if not task_data["targets"]:
            task_data["targets"] = _extract_targets_fallback(task_data["task_text"])
            print(f"[Task Agent] WARNING: LLM didn't provide targets, used fallback: {task_data['targets']}")
        
        # Normalize choices using central normalize_targets function
        raw_choices = task_data.get("choices", [])
        if not isinstance(raw_choices, list):
            raw_choices = [raw_choices] if raw_choices else []
        
        task_data["choices"] = normalize_targets(raw_choices)
        
        # Post-normalization defensive checks
        task_data = _enforce_target_invariants(task_data)
        
        return task_data
        
    except json.JSONDecodeError as e:
        print(f"[Task Agent] JSON parse error: {e}")
        print(f"[Task Agent] Raw output: {raw_output}")
        # Return a safe default task
        return _default_task()
    except Exception as e:
        print(f"[Task Agent] Error: {e}")
        return _default_task()


def generate_task_candidates(
    episode_transcript: str,
    subject_of_the_day: str = None,
    lesson_of_the_day: str = None,
    age_band: int = 0,
    basil_max_tokens: int = None,
    previous_askers: list = None,
    recent_targets: list = None,
) -> list:
    """
    Generate 5-8 candidate tasks for Basil.
    
    Args:
        episode_transcript: Current episode transcript only (includes teaching paragraph and Sophie reaction for this episode)
        subject_of_the_day: Current session subject
        lesson_of_the_day: Current session lesson
        age_band: Basil's current developmental stage (0-7)
        basil_max_tokens: Maximum tokens for Basil's response
        previous_askers: List of recent askers for variety
        recent_targets: List of recent target words used (for variety)
    
    Returns:
        List of candidate task dicts, each with: task_text, task_category, stage_hint,
        asker, targets, choices
    """
    try:
        with open(PROMPT_TASK_GENERATOR, "r") as f:
            template = f.read()
    except FileNotFoundError:
        print(f"[Task Generator] Prompt file not found: {PROMPT_TASK_GENERATOR}")
        return []
    
    if basil_max_tokens is None:
        basil_max_tokens = BASIL_MAX_TOKENS
    
    if previous_askers is None:
        previous_askers = []
    
    if recent_targets is None:
        recent_targets = []
    
    age_band_desc = AGE_BAND_DESCRIPTIONS.get(age_band, "unknown")
    basil_age_description = get_age_band_natural_description(age_band)
    task_guidance = get_task_guidance_for_age_band(age_band)
    
    # Format previous askers for variety guidance
    askers_str = ", ".join(previous_askers[-3:]) if previous_askers else "(none yet)"
    
    # Format recent targets for variety guidance
    targets_str = ", ".join(recent_targets[-4:]) if recent_targets else "(none yet)"
    
    prompt = template.format(
        episode_transcript=episode_transcript or "(No conversation yet)",
        subject_of_the_day=subject_of_the_day or "(No subject)",
        lesson_of_the_day=lesson_of_the_day or "(No lesson)",
        age_band=age_band,
        age_band_description=age_band_desc,
        basil_age_description=basil_age_description,
        task_guidance=task_guidance,
        basil_max_tokens=basil_max_tokens,
        previous_askers=askers_str,
        recent_targets=targets_str,
    )
    
    if episode_transcript:  # Only log if we have context
        print(f"[Task Generator] Generating candidates for age_band={age_band}, subject={subject_of_the_day}, lesson={lesson_of_the_day}")
    
    try:
        response = client.chat.completions.create(
            model=TASK_AGENT_MODEL,
            messages=[
                {"role": "system", "content": "You are a task generation agent. Output valid JSON only."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.9,  # Higher temperature for diversity
            max_tokens=1500,  # More tokens for multiple candidates
        )
        
        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])
        
        result = json.loads(raw_output)
        
        candidates = result.get("candidates", [])
        if not isinstance(candidates, list):
            print(f"[Task Generator] Invalid candidates format: expected list, got {type(candidates)}")
            return []
        
        # Validate and normalize each candidate
        validated_candidates = []
        for i, candidate in enumerate(candidates):
            if not isinstance(candidate, dict):
                print(f"[Task Generator] Skipping invalid candidate {i}: not a dict")
                continue
            
            # Normalize targets
            raw_targets = candidate.get("targets", [])
            if not isinstance(raw_targets, list):
                raw_targets = [raw_targets] if raw_targets else []
            candidate["targets"] = normalize_targets(raw_targets)
            
            # Normalize choices
            raw_choices = candidate.get("choices", [])
            if not isinstance(raw_choices, list):
                raw_choices = [raw_choices] if raw_choices else []
            candidate["choices"] = normalize_targets(raw_choices)
            
            # Ensure required fields
            if "task_text" not in candidate or "task_category" not in candidate:
                print(f"[Task Generator] Skipping candidate {i}: missing required fields")
                continue
            
            # Set defaults
            if "stage_hint" not in candidate:
                candidate["stage_hint"] = age_band
            if "asker" not in candidate:
                candidate["asker"] = "tutor"
            else:
                candidate["asker"] = candidate["asker"].lower().strip()
                if candidate["asker"] not in ("tutor", "sophie"):
                    candidate["asker"] = "tutor"
            
            validated_candidates.append(candidate)
        
        print(f"[Task Generator] Generated {len(validated_candidates)} valid candidates")
        return validated_candidates
        
    except json.JSONDecodeError as e:
        print(f"[Task Generator] JSON parse error: {e}")
        print(f"[Task Generator] Raw output: {raw_output[:200]}...")
        return []
    except Exception as e:
        print(f"[Task Generator] Error: {e}")
        return []


def select_best_task(
    candidates: list,
    teaching_snippet: str = "",
    subject_of_the_day: str = None,
    lesson_of_the_day: str = None,
    age_band: int = 0,
    recent_targets: list = None,
    basil_max_tokens: int = None,
) -> dict:
    """
    Select the best task from a list of candidates.
    
    Args:
        candidates: List of candidate task dicts (from generate_task_candidates)
        teaching_snippet: The Tutor's teaching paragraph for this episode
        subject_of_the_day: Current session subject
        lesson_of_the_day: Current session lesson
        age_band: Basil's current developmental stage (0-7)
        recent_targets: List of recent target words used (for variety)
        basil_max_tokens: Maximum tokens for Basil's response
    
    Returns:
        Complete task spec dict with rubric and grader_instructions
    """
    if not candidates:
        print("[Task Selector] No candidates provided, using default task")
        return _default_task()
    
    try:
        with open(PROMPT_TASK_SELECTOR, "r") as f:
            template = f.read()
    except FileNotFoundError:
        print(f"[Task Selector] Prompt file not found: {PROMPT_TASK_SELECTOR}")
        # Fallback: just return first candidate with basic rubric
        if candidates:
            return _complete_task_spec(candidates[0], basil_max_tokens)
        else:
            return _default_task()
    
    if basil_max_tokens is None:
        basil_max_tokens = BASIL_MAX_TOKENS
    
    if recent_targets is None:
        recent_targets = []
    
    age_band_desc = AGE_BAND_DESCRIPTIONS.get(age_band, "unknown")
    task_guidance = get_task_guidance_for_age_band(age_band)
    
    # Format candidates as JSON for prompt
    candidates_json = json.dumps(candidates, indent=2)
    
    # Format recent targets
    targets_str = ", ".join(recent_targets[-4:]) if recent_targets else "(none yet)"
    
    prompt = template.format(
        subject_of_the_day=subject_of_the_day or "(No subject)",
        lesson_of_the_day=lesson_of_the_day or "(No lesson)",
        age_band=age_band,
        age_band_description=age_band_desc,
        task_guidance=task_guidance,
        teaching_snippet=teaching_snippet or "(No teaching paragraph yet)",
        recent_targets=targets_str,
        candidates_json=candidates_json,
        basil_max_tokens=basil_max_tokens,
    )
    
    print(f"[Task Selector] Selecting best task from {len(candidates)} candidates")
    
    try:
        response = client.chat.completions.create(
            model=TASK_AGENT_MODEL,
            messages=[
                {"role": "system", "content": "You are a task selection agent. Output valid JSON only."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.3,  # Lower temperature for consistent selection
            max_tokens=800,
        )
        
        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])
        
        task_data = json.loads(raw_output)
        
        # Validate required fields
        required_fields = ["task_text", "task_category", "stage_hint", "rubric", "grader_instructions"]
        for field in required_fields:
            if field not in task_data:
                raise ValueError(f"Missing required field: {field}")
        
        # Set default max_tokens if not provided
        if "max_tokens" not in task_data:
            task_data["max_tokens"] = basil_max_tokens
        
        # Normalize asker
        if "asker" not in task_data:
            task_data["asker"] = "tutor"
        else:
            task_data["asker"] = task_data["asker"].lower().strip()
            if task_data["asker"] not in ("tutor", "sophie"):
                task_data["asker"] = "tutor"
        
        # Normalize targets
        raw_targets = task_data.get("targets", [])
        if not isinstance(raw_targets, list):
            raw_targets = [raw_targets] if raw_targets else []
        task_data["targets"] = normalize_targets(raw_targets)
        
        # Normalize choices
        raw_choices = task_data.get("choices", [])
        if not isinstance(raw_choices, list):
            raw_choices = [raw_choices] if raw_choices else []
        task_data["choices"] = normalize_targets(raw_choices)
        
        # Post-normalization defensive checks
        task_data = _enforce_target_invariants(task_data)
        
        selection_reason = task_data.get("selection_reason", "")
        if selection_reason:
            print(f"[Task Selector] Selected: '{task_data['task_text'][:50]}...' - {selection_reason}")
        
        return task_data
        
    except json.JSONDecodeError as e:
        print(f"[Task Selector] JSON parse error: {e}")
        print(f"[Task Selector] Raw output: {raw_output[:200]}...")
        # Fallback: return first candidate with basic rubric
        return _complete_task_spec(candidates[0], basil_max_tokens)
    except Exception as e:
        print(f"[Task Selector] Error: {e}")
        # Fallback: return first candidate with basic rubric
        return _complete_task_spec(candidates[0], basil_max_tokens)


def _complete_task_spec(candidate: dict, basil_max_tokens: int) -> dict:
    """
    Complete a candidate task spec with rubric and grader_instructions.
    Used as fallback when selector fails.
    """
    task_text = candidate.get("task_text", "")
    targets = candidate.get("targets", [])
    
    # Generate basic rubric (0-7 scale)
    rubric = {
        "0": "No output or completely unintelligible (random tokens, no English)",
        "1": "Recognizable English, but nothing related to the lesson topic or question",
        "2": "Multiple different English words, but none related to the lesson topic",
        "3": "English word(s) from the lesson/teaching paragraph, but does not address the specific question",
        "4": "On-topic vocabulary directly related to the question asked",
        "5": "Attempts the right response format or uses a concept-word from the question, but target answer absent or wrong",
        "6": "Contains the target/correct answer, but with extra noise or filler words",
        "7": "Clean, correct answer that directly addresses the task"
    }
    
    # Generate grader instructions
    if targets:
        if len(targets) == 1:
            grader_instructions = f"Award 7 if '{targets[0]}' appears cleanly. Award 6 if '{targets[0]}' appears with noise. Award 3-4 if output uses on-topic words from the lesson."
        else:
            targets_str = " or ".join([f"'{t}'" for t in targets])
            grader_instructions = f"Award 7 if {targets_str} appears cleanly. Award 6 if target appears with noise. Award 3-4 if output uses on-topic words from the lesson."
    else:
        grader_instructions = "Award 7 for clean correct answer. Award 5-6 for partial answers. Award 3-4 for on-topic vocabulary. Award 1-2 for any English words."
    
    return {
        "task_text": task_text,
        "task_category": candidate.get("task_category", "vocab"),
        "stage_hint": candidate.get("stage_hint", 0),
        "asker": candidate.get("asker", "tutor"),
        "targets": targets,
        "choices": candidate.get("choices", []),
        "rubric": rubric,
        "grader_instructions": grader_instructions,
        "max_tokens": basil_max_tokens or BASIL_MAX_TOKENS,
    }


def _enforce_target_invariants(task_data: dict) -> dict:
    """
    Enforce invariants on targets and choices after normalization.
    
    Rules:
    1. If targets is empty, derive from choices or use ["hello"] as fallback
    2. For MCQ tasks (choices non-empty), ensure all choices are in targets
    3. For yes/no tasks, ensure targets contains ["yes", "no"] or one of them
    4. Ensure targets is never empty
    """
    targets = task_data.get("targets", [])
    choices = task_data.get("choices", [])
    task_text_lower = (task_data.get("task_text", "") or "").lower()
    
    # Detect yes/no task patterns
    is_yes_no = "yes or no" in task_text_lower or \
                "yes or no?" in task_text_lower or \
                (set(choices) == {"yes", "no"})
    
    # Rule 1: If targets empty but choices exist, use choices as targets
    if not targets and choices:
        targets = choices[:]
        print(f"[Task Agent] Derived targets from choices: {targets}")
    
    # Rule 2: For yes/no tasks, enforce correct targets
    if is_yes_no:
        if "yes" not in targets and "no" not in targets:
            targets = ["yes", "no"]
            print(f"[Task Agent] Fixed yes/no task targets: {targets}")
        elif "yes" in targets and "no" not in targets:
            pass  # "Say yes" type task - single target is OK
        elif "no" in targets and "yes" not in targets:
            pass  # "Say no" type task - single target is OK
        # Else: both are present, which is correct
    
    # Rule 3: For MCQ tasks, ensure choices are subset of targets
    if choices:
        for choice in choices:
            if choice not in targets:
                targets.append(choice)
                print(f"[Task Agent] Added choice '{choice}' to targets")
    
    # Rule 4: Ensure targets is never empty
    if not targets:
        targets = ["yes"]
        print(f"[Task Agent] WARNING: Targets was empty, using fallback ['yes']")
    
    task_data["targets"] = targets
    task_data["choices"] = choices
    return task_data


def _extract_targets_fallback(task_text: str) -> list:
    """
    Fallback target extraction when LLM doesn't provide targets.
    
    Minimal implementation - just handles the most common patterns.
    This should rarely be needed since TaskAgent now emits targets explicitly.
    """
    targets = []
    task_lower = task_text.lower()
    
    # Pattern: "say yes or no" takes priority
    if "say yes or no" in task_lower or "yes or no?" in task_lower:
        return ["yes", "no"]
    
    # Pattern: Quoted words
    quoted = re.findall(r"['\"](\w+)['\"]", task_text)
    if quoted:
        return [w.lower() for w in quoted]
    
    # Pattern: "X or Y?" at end
    ab_match = re.search(r"(\w+)\s+or\s+(\w+)\??\s*$", task_text, re.IGNORECASE)
    if ab_match:
        a, b = ab_match.group(1).lower(), ab_match.group(2).lower()
        skip = {"say", "is", "it", "a", "an", "the", "can", "you"}
        if a not in skip:
            targets.append(a)
        if b not in skip:
            targets.append(b)
        if targets:
            return targets
    
    # Pattern: "Say X" (single word)
    say_match = re.search(r"\bsay\s+(\w+)", task_text, re.IGNORECASE)
    if say_match:
        word = say_match.group(1).lower()
        if word not in {"one", "a", "an", "the", "yes", "no", "or"}:
            return [word]
    
    return targets


def _default_task() -> dict:
    """Return a safe default task if generation fails."""
    import random
    return {
        "task_text": "Say yes.",
        "task_category": "control",
        "stage_hint": 0,
        "asker": random.choice(["tutor", "sophie"]),  # Vary asker even in fallback
        "targets": ["yes"],  # Explicit targets
        "choices": [],
        "rubric": {
            "0": "No output or completely unintelligible (random tokens, no English)",
            "1": "Recognizable English, but nothing related to the question",
            "2": "Multiple different English words, but none related to the question",
            "3": "English word(s) from the lesson domain, but does not answer the question",
            "4": "On-topic vocabulary directly related to the question",
            "5": "Attempts a response format, but 'yes' not present",
            "6": "'yes' appears but with extra noise or filler words",
            "7": "Said 'yes' clearly"
        },
        "grader_instructions": "Award 7 if 'yes' appears cleanly. Award 6 if 'yes' appears with noise. Award 3-4 if output shows topical awareness or attempts a response format.",
        "max_tokens": BASIL_MAX_TOKENS
    }


if __name__ == "__main__":
    # Test the task agent
    print("Testing Task Agent...")
    
    test_snippet = """
Tutor: Hello Basil! Can you say hello back?
Basil: agkjh the hello wor
Tutor: Good try! I heard 'hello' in there. Can you say just one word?
Basil: word the is a
"""
    
    task = generate_task(test_snippet)
    print(json.dumps(task, indent=2))