grader_agent.py

# grader_agent.py
# Grades Basil's responses using a two-grader architecture:
#   english_grader (0-3): English presence + domain relevance
#   task_grader (0-7):    Task compliance + target word detection
#   Final score = max(english_score, task_score)
#
# The programmatic Englishness floor (score_override.py) and target-word
# safety net are applied downstream in auto_session.py.

import os
import json
import re
from openai import OpenAI

from config import (
    GRADER_MODEL,
    PROMPT_GRADER,
    PROMPT_ENGLISH_GRADER,
    PROMPT_TASK_GRADER,
    SCORE_MIN,
    SCORE_MAX,
)

from llm_client import create_smart_client
client = create_smart_client()


def _load_template(path: str) -> str:
    with open(path, "r") as f:
        return f.read()


def _sanitize_for_json(text: str) -> str:
    """Sanitize text to be safe for JSON embedding."""
    if not text:
        return ""
    sanitized = re.sub(r'[\x00-\x1f\x7f-\x9f]', ' ', text)
    sanitized = re.sub(r' +', ' ', sanitized)
    return sanitized.strip()


def _try_repair_json(raw_output: str) -> dict:
    """Attempt basic JSON repairs for common LLM mistakes."""
    text = raw_output.strip()

    if text.startswith("```"):
        lines = text.split("\n")
        text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])

    text = re.sub(r',\s*([\}\]])', r'\1', text)

    json_match = re.search(r'\{[^{}]*\}', text, re.DOTALL)
    if json_match:
        text = json_match.group(0)

    return json.loads(text)


class _Defaults(dict):
    """Dict subclass that returns '{key}' for missing keys in format_map."""
    def __missing__(self, key):
        return f"{{{key}}}"


# ─── English Grader (0-3) ────────────────────────────────────────────────────

def grade_english(
    basil_output: str,
    subject: str = "",
    lesson: str = "",
) -> dict:
    """
    Score 0-3: English word presence and domain relevance.

    Returns dict with: score, english_words, domain_words, reason
    """
    template = _load_template(PROMPT_ENGLISH_GRADER)
    safe_output = _sanitize_for_json(basil_output) or "(Empty response)"

    prompt = template.format_map(_Defaults(
        subject=subject or "(Not specified)",
        lesson=lesson or "(Not specified)",
        basil_output=safe_output,
    ))

    try:
        response = client.chat.completions.create(
            model=GRADER_MODEL,
            messages=[
                {"role": "system", "content": "Output valid JSON only."},
                {"role": "user", "content": prompt},
            ],
            temperature=0.3,
            max_tokens=200,
        )
        raw = response.choices[0].message.content.strip()
        data = _try_repair_json(raw)
        data["score"] = max(0, min(3, int(data.get("score", 0))))
        return data
    except Exception as e:
        print(f"[EnglishGrader] Error: {e}")
        return {"score": 0, "english_words": [], "domain_words": [], "reason": f"Error: {e}"}


# ─── Task Grader (0-7) ───────────────────────────────────────────────────────

def grade_task(
    task_text: str,
    basil_output: str,
    context: str = "",
    subject: str = "",
    lesson: str = "",
    age_band: int = 0,
    age_band_description: str = "",
) -> dict:
    """
    Score 0-7: Task compliance and target word detection.

    Returns dict with: score, target_word, target_found, target_evidence, reason
    """
    template = _load_template(PROMPT_TASK_GRADER)
    safe_output = _sanitize_for_json(basil_output) or "(Empty response)"

    prompt = template.format_map(_Defaults(
        subject=subject or "(Not specified)",
        lesson=lesson or "(Not specified)",
        task_text=task_text,
        context=context or "(No additional context)",
        basil_output=safe_output,
        age_band=age_band,
        age_band_description=age_band_description or f"age band {age_band}",
    ))

    try:
        response = client.chat.completions.create(
            model=GRADER_MODEL,
            messages=[
                {"role": "system", "content": "Output valid JSON only."},
                {"role": "user", "content": prompt},
            ],
            temperature=0.3,
            max_tokens=250,
        )
        raw = response.choices[0].message.content.strip()
        data = _try_repair_json(raw)
        data["score"] = max(SCORE_MIN, min(SCORE_MAX, int(data.get("score", 0))))
        return data
    except Exception as e:
        print(f"[TaskGrader] Error: {e}")
        return {"score": 0, "target_word": "", "target_found": False,
                "target_evidence": None, "reason": f"Error: {e}"}


# ─── Combined Grade (public API) ─────────────────────────────────────────────

def grade_response(
    task_text: str,
    rubric: dict,
    grader_instructions: str,
    basil_output: str,
    context: str = "",
    age_band: int = 0,
    age_band_description: str = "",
    subject: str = "",
    lesson: str = "",
) -> dict:
    """
    Grade Basil's response using two specialized graders.

    Calls english_grader (0-3) and task_grader (0-7) in sequence,
    returns final_score = max(english_score, task_score).

    Args:
        task_text: The task/prompt given to Basil
        rubric: Dict mapping scores to criteria (kept for API compat, not used by new graders)
        grader_instructions: Additional guidance (kept for API compat)
        basil_output: Basil's raw output
        context: Recent conversation for context
        age_band: Basil's current developmental stage (0-7)
        age_band_description: Human-readable description of age band
        subject: Lesson subject (e.g., "Botany")
        lesson: Lesson title (e.g., "The Hierarchy of Plant Classification")

    Returns:
        dict with keys: score, justification, evidence, notes,
        plus english_grade and task_grade sub-dicts for audit.
    """
    eng = grade_english(basil_output, subject, lesson)
    tsk = grade_task(task_text, basil_output, context, subject, lesson,
                     age_band, age_band_description)

    eng_score = eng.get("score", 0)
    tsk_score = tsk.get("score", 0)

    # English grader gates the task grader: if the english_grader doesn't
    # see domain-relevant words, the task_grader can't claim high task
    # compliance.  This prevents parroting (e.g. "Nice try") from scoring
    # 4-5 when the english_grader correctly identifies no domain content.
    if eng_score <= 1:
        tsk_capped = min(tsk_score, 2)
    elif eng_score == 2:
        tsk_capped = min(tsk_score, 3)
    else:
        tsk_capped = tsk_score

    final_score = max(eng_score, tsk_capped)

    justification = tsk.get("reason", "")
    if eng_score > tsk_capped:
        justification = eng.get("reason", justification)

    notes = f"eng={eng_score} task={tsk_score}"
    if tsk_capped != tsk_score:
        notes += f" task_capped={tsk_capped}"

    evidence = []
    if tsk.get("target_evidence"):
        evidence.append(_sanitize_for_json(str(tsk["target_evidence"]))[:40])
    for dw in eng.get("domain_words", []):
        if isinstance(dw, str):
            evidence.append(_sanitize_for_json(dw)[:40])

    return {
        "score": final_score,
        "justification": justification,
        "evidence": evidence[:5],
        "notes": notes,
        "english_grade": eng,
        "task_grade": tsk,
    }


def _error_grade(error_msg: str) -> dict:
    """Return a grade when grading fails entirely."""
    return {
        "score": 0,
        "justification": f"Grading error: {error_msg}",
        "evidence": [],
        "notes": "Grader failed; score=0 assigned (no auto-upgrade).",
        "grader_error": True,
        "compliance": False,
    }


if __name__ == "__main__":
    print("Testing Two-Grader Architecture...")

    result = grade_response(
        task_text="Say one word that is an animal.",
        rubric={},
        grader_instructions="",
        basil_output="cat the dog",
        context="Tutor was asking about animals.",
        subject="Animals",
        lesson="Common Pets",
    )
    print(json.dumps(result, indent=2))