evaluate.py

#!/usr/bin/env python3
"""
Benchmark evaluator: scores collected LLM responses against the rubric using an Ollama judge model.

One Ollama call per rubric criterion per question (single-criterion prompts for reliability).
Results saved incrementally — safe to stop and restart.

Usage:
  python evaluate.py --model <response_model> --trial <N> --judge <judge_model> [--questions ...] [--host <url>]

--questions selection rules (same as run_benchmark.py):
  omitted             → all available questions
  --questions 5       → only EX5
  --questions 3 7     → EX3 through EX7 inclusive
  --questions 1 3 5 7 → exactly EX1, EX3, EX5, EX7
"""

import argparse
import json
import os
import re
import sys
import urllib.error
import urllib.request
from datetime import datetime

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "EvaluatorTools"))
from conciseness_analyze import count_stats


# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

BENCHMARK_JSONL = os.path.join(SCRIPT_DIR, "Dataset", "benchmark.jsonl")
RESPONSES_DIR = os.path.join(SCRIPT_DIR, "Responses")
EVALUATIONS_DIR = os.path.join(SCRIPT_DIR, "Evaluations")

JUDGE_SYSTEM_PROMPT = (
    "You are an expert evaluator scoring AI tutoring responses. "
    "You will be given a student question, evaluator context, a rubric criterion with its "
    "scoring guide, and the response to score. "
    'Reply ONLY with valid JSON in this exact format: {"thinking": "<reason through whether the response meets the criterion before committing to a score>", "score": <0, 1, or 2>, "reasoning": "<one sentence summary of your decision>"}'
)


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

def parse_args():
    parser = argparse.ArgumentParser(description="Score benchmark responses using an Ollama judge model.")
    parser.add_argument("--model", required=True, help="Response model name (as stored in Responses/)")
    parser.add_argument("--trial", required=True, type=int, help="Trial number to evaluate (e.g. 1)")
    parser.add_argument("--judge", required=True, help="Ollama judge model name (e.g. phi4:latest)")
    parser.add_argument(
        "--questions",
        nargs="+",
        type=int,
        metavar="N",
        help="Question IDs: single int, two ints (range), or 3+ ints (explicit list)",
    )
    parser.add_argument("--host", default="http://localhost:11434", help="Ollama base URL")
    return parser.parse_args()


def resolve_question_ids(questions_arg):
    """Convert --questions arg to a list of 'EX{N}' strings, or None for all."""
    if questions_arg is None:
        return None
    if len(questions_arg) == 1:
        return [f"EX{questions_arg[0]}"]
    if len(questions_arg) == 2:
        lo, hi = questions_arg
        return [f"EX{n}" for n in range(lo, hi + 1)]
    return [f"EX{n}" for n in questions_arg]


# ---------------------------------------------------------------------------
# Slug helper
# ---------------------------------------------------------------------------

def model_slug(model_name):
    return re.sub(r"[:/\\]", "_", model_name)


# ---------------------------------------------------------------------------
# Response file parsing
# ---------------------------------------------------------------------------

def extract_response_text(filepath):
    """Extract the [RESPONSE]...[/RESPONSE] block from a response file."""
    content = open(filepath, encoding="utf-8").read()
    match = re.search(r"\[RESPONSE\](.*?)\[/RESPONSE\]", content, re.DOTALL)
    if match:
        return match.group(1).strip()
    # Fallback: return full content minus the Time: line and any THINKING block
    content = re.sub(r"^Time:.*\n?", "", content, flags=re.MULTILINE)
    content = re.sub(r"\[THINKING\].*?\[/THINKING\]", "", content, flags=re.DOTALL)
    return content.strip()


# ---------------------------------------------------------------------------
# Judge prompt construction
# ---------------------------------------------------------------------------

def build_judge_message(entry, criterion_obj, response_text, extra_context=""):
    student_question = entry["student_question"]
    evaluator_notes = entry.get("evaluator_notes", "")
    criterion = criterion_obj["criterion"]
    broad_desc = criterion_obj["broad_description"]
    entry_specific = criterion_obj.get("entry_specific")
    guide = criterion_obj["scoring_guide"]

    entry_specific_section = ""
    if entry_specific:
        entry_specific_section = f"\nAdditional context for this specific question:\n{entry_specific}\n"

    extra_section = f"\n{extra_context}\n" if extra_context else ""

    return (
        f"## Student Question\n{student_question}\n\n"
        f"## Evaluator Notes (background context — not shown to student)\n{evaluator_notes}\n\n"
        f"## Criterion: {criterion}\n{broad_desc}{entry_specific_section}{extra_section}\n"
        f"Scoring guide:\n"
        f"  2 points: {guide['2']}\n"
        f"  1 point:  {guide['1']}\n"
        f"  0 points: {guide['0']}\n\n"
        f"## Response to Score\n{response_text}\n\n"
        f'Score the above response on the "{criterion}" criterion. Reply with JSON only: {{"thinking": "...", "score": <0/1/2>, "reasoning": "..."}}'
    )


# ---------------------------------------------------------------------------
# Ollama API call
# ---------------------------------------------------------------------------

def call_ollama(host, model, system_prompt, user_message):
    url = f"{host.rstrip('/')}/api/chat"
    body = {
        "model": model,
        "stream": False,
        "options": {"temperature": 0},
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
    }
    data = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})

    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            raw = resp.read().decode("utf-8")
    except urllib.error.URLError as e:
        raise ConnectionError(f"Could not reach Ollama at {host}: {e.reason}") from e
    except urllib.error.HTTPError as e:
        body_text = e.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {e.code}: {body_text}") from e

    return json.loads(raw)


# ---------------------------------------------------------------------------
# Judge output parsing
# ---------------------------------------------------------------------------

def parse_judge_output(raw_content):
    """
    Parse thinking, score, and reasoning from judge output.
    Returns (score_int_or_None, reasoning_str, thinking_str).
    """
    # Attempt 1: direct parse
    try:
        obj = json.loads(raw_content.strip())
        score = obj.get("score")
        reasoning = obj.get("reasoning", "")
        thinking = obj.get("thinking", "")
        if score in (0, 1, 2):
            return score, reasoning, thinking
    except (json.JSONDecodeError, AttributeError):
        pass

    # Attempt 2: find JSON object inside the output (thinking field may contain nested quotes)
    match = re.search(r'\{.*?"score".*?\}', raw_content, re.DOTALL)
    if match:
        try:
            obj = json.loads(match.group())
            score = obj.get("score")
            reasoning = obj.get("reasoning", "")
            thinking = obj.get("thinking", "")
            if score in (0, 1, 2):
                return score, reasoning, thinking
        except json.JSONDecodeError:
            pass

    # Attempt 3: extract fields individually via regex
    score_match = re.search(r'"score"\s*:\s*([012])', raw_content)
    reasoning_match = re.search(r'"reasoning"\s*:\s*"([^"]*)"', raw_content)
    thinking_match = re.search(r'"thinking"\s*:\s*"([^"]*)"', raw_content)
    if score_match:
        score = int(score_match.group(1))
        reasoning = reasoning_match.group(1) if reasoning_match else ""
        thinking = thinking_match.group(1) if thinking_match else ""
        return score, reasoning, thinking

    # Failed
    truncated = raw_content[:200].replace("\n", " ")
    return None, f"PARSE_ERROR: {truncated}", ""


# ---------------------------------------------------------------------------
# Summary computation
# ---------------------------------------------------------------------------

def compute_summary(response_model, trial, judge_model, results):
    """
    results: list of per-question result dicts (from EX*.json content).
    Only counts non-null scores in percentages.
    """
    def pct(score, max_val):
        if max_val == 0:
            return None
        return round(score / max_val * 100, 1)

    overall_score = 0
    overall_max = 0

    by_criterion = {}
    by_question_type = {}
    by_bug_category = {}
    by_question = {}

    for r in results:
        qid = r["id"]
        qtype = r["question_type"]
        bcat = r["bug_category"]

        q_score = 0
        q_max = 0

        for c in r["criteria"]:
            crit = c["criterion"]
            score = c["score"]
            max_pts = c["max_points"]

            if crit not in by_criterion:
                by_criterion[crit] = {"total_score": 0, "max": 0}

            if score is not None:
                q_score += score
                q_max += max_pts
                overall_score += score
                overall_max += max_pts
                by_criterion[crit]["total_score"] += score
                by_criterion[crit]["max"] += max_pts

        # by_question_type
        if qtype not in by_question_type:
            by_question_type[qtype] = {"total_score": 0, "max_possible": 0, "count": 0}
        by_question_type[qtype]["total_score"] += q_score
        by_question_type[qtype]["max_possible"] += q_max
        by_question_type[qtype]["count"] += 1

        # by_bug_category
        if bcat not in by_bug_category:
            by_bug_category[bcat] = {"total_score": 0, "max_possible": 0}
        by_bug_category[bcat]["total_score"] += q_score
        by_bug_category[bcat]["max_possible"] += q_max

        by_question[qid] = {
            "score": q_score,
            "max": q_max,
            "percent": pct(q_score, q_max),
            "question_type": qtype,
            "bug_category": bcat,
        }

    # Add percentages
    for v in by_criterion.values():
        v["percent"] = pct(v["total_score"], v["max"])
    for v in by_question_type.values():
        v["percent"] = pct(v["total_score"], v["max_possible"])
    for v in by_bug_category.values():
        v["percent"] = pct(v["total_score"], v["max_possible"])

    return {
        "response_model": response_model,
        "trial": trial,
        "judge_model": judge_model,
        "evaluated_at": datetime.now().isoformat(timespec="seconds"),
        "questions_evaluated": len(results),
        "questions_skipped": 0,  # updated by caller
        "overall": {
            "total_score": overall_score,
            "max_possible": overall_max,
            "percent": pct(overall_score, overall_max),
        },
        "by_criterion": by_criterion,
        "by_question_type": by_question_type,
        "by_bug_category": by_bug_category,
        "by_question": by_question,
    }


# ---------------------------------------------------------------------------
# Console summary printer
# ---------------------------------------------------------------------------

def print_summary(summary):
    print("\n" + "=" * 60)
    print(f"  EVALUATION SUMMARY")
    print(f"  Model:   {summary['response_model']}  Trial {summary['trial']}")
    print(f"  Judge:   {summary['judge_model']}")
    print(f"  Date:    {summary['evaluated_at']}")
    print("=" * 60)
    o = summary["overall"]
    print(f"  Overall: {o['total_score']}/{o['max_possible']}  ({o['percent']}%)")
    print(f"  Evaluated: {summary['questions_evaluated']} question(s)  Skipped: {summary['questions_skipped']}")

    print("\n  --- By Criterion ---")
    for crit, v in summary["by_criterion"].items():
        bar = _bar(v["percent"])
        print(f"  {crit:<28} {v['total_score']:>3}/{v['max']:<3}  {v['percent']:>5.1f}%  {bar}")

    print("\n  --- By Question Type ---")
    for qtype, v in summary["by_question_type"].items():
        print(f"  {qtype:<24} {v['total_score']:>4}/{v['max_possible']:<4}  {v['percent']:>5.1f}%  (n={v['count']})")

    print("\n  --- By Question ---")
    for qid, v in summary["by_question"].items():
        flag = "✓" if (v["percent"] or 0) >= 75 else "✗"
        print(f"  {flag} {qid:<5}  {v['score']:>2}/{v['max']:<2}  {(v['percent'] or 0):>5.1f}%  [{v['bug_category']}]")

    print("=" * 60)


def _bar(pct, width=10):
    if pct is None:
        return ""
    filled = round((pct / 100) * width)
    return "[" + "█" * filled + "░" * (width - filled) + "]"


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    args = parse_args()

    selected_ids = resolve_question_ids(args.questions)

    # Load benchmark entries
    entries = []
    with open(BENCHMARK_JSONL, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                entries.append(json.loads(line))

    if selected_ids is not None:
        id_set = set(selected_ids)
        entries = [e for e in entries if e["id"] in id_set]
        missing = id_set - {e["id"] for e in entries}
        if missing:
            print(f"WARNING: IDs not found in benchmark: {sorted(missing)}", file=sys.stderr)

    if not entries:
        print("No matching entries found. Exiting.", file=sys.stderr)
        sys.exit(1)

    # Determine directories
    resp_slug = model_slug(args.model)
    judge_slug = model_slug(args.judge)
    trial_responses_dir = os.path.join(RESPONSES_DIR, resp_slug, f"Trial{args.trial}")
    eval_dir = os.path.join(EVALUATIONS_DIR, resp_slug, f"Trial{args.trial}", judge_slug)
    os.makedirs(eval_dir, exist_ok=True)

    print(f"Evaluating: {args.model} Trial{args.trial} | Judge: {args.judge} | Questions: {len(entries)}")
    print(f"Output: {eval_dir}")

    # Check Ollama reachable
    try:
        req = urllib.request.Request(f"{args.host.rstrip('/')}/api/tags")
        with urllib.request.urlopen(req):
            pass
    except Exception as e:
        print(f"\nERROR: Cannot reach Ollama at {args.host}\n  {e}", file=sys.stderr)
        sys.exit(1)

    results = []
    skipped = 0

    for entry in entries:
        eid = entry["id"]
        response_file = os.path.join(trial_responses_dir, f"{eid}.txt")

        if not os.path.exists(response_file):
            print(f"  ✗ {eid} — SKIPPED: response file not found at {response_file}")
            skipped += 1
            continue

        response_text = extract_response_text(response_file)

        criteria_results = []
        entry_total = 0
        entry_max = 0

        for criterion_obj in entry["rubric"]:
            crit_name = criterion_obj["criterion"]
            max_pts = criterion_obj["max_points"]

            extra_context = ""
            if crit_name == "conciseness":
                stats = count_stats(response_text)
                extra_context = f"Character count: {stats['characters']}"
            user_message = build_judge_message(entry, criterion_obj, response_text, extra_context)

            try:
                print(f"    [{datetime.now().strftime('%H:%M:%S')}] Calling judge for {eid}/{crit_name}...", flush=True)
                result = call_ollama(args.host, args.judge, JUDGE_SYSTEM_PROMPT, user_message)
                print(f"    [{datetime.now().strftime('%H:%M:%S')}] Got response for {eid}/{crit_name}", flush=True)
            except ConnectionError as e:
                criteria_results.append({
                    "criterion": crit_name,
                    "score": None,
                    "max_points": max_pts,
                    "thinking": "",
                    "reasoning": f"CONNECTION_ERROR: {e}",
                })
                print(f"    [{datetime.now().strftime('%H:%M:%S')}] {crit_name}: TIMEOUT/CONNECTION ERROR — skipping", flush=True)
                continue    
            except RuntimeError as e:
                criteria_results.append({
                    "criterion": crit_name,
                    "score": None,
                    "max_points": max_pts,
                    "thinking": "",
                    "reasoning": f"API_ERROR: {e}",
                })
                print(f"    {crit_name}: ERROR")
                continue

            raw_content = result.get("message", {}).get("content", "")
            score, reasoning, thinking = parse_judge_output(raw_content)

            criteria_results.append({
                "criterion": crit_name,
                "score": score,
                "max_points": max_pts,
                "thinking": thinking,
                "reasoning": reasoning,
            })

            score_display = str(score) if score is not None else "?"
            print(f"    {crit_name:<28} {score_display}/2  — {reasoning[:60]}")

            if score is not None:
                entry_total += score
                entry_max += max_pts

        percent = round(entry_total / entry_max * 100, 1) if entry_max > 0 else None

        question_result = {
            "id": eid,
            "question_type": entry["question_type"],
            "bug_category": entry["bug_category"],
            "total_score": entry_total,
            "max_total_points": entry_max,
            "percent": percent,
            "criteria": criteria_results,
        }

        # Save immediately
        out_path = os.path.join(eval_dir, f"{eid}.json")
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump(question_result, f, indent=2)

        pct_str = f"{percent:.0f}%" if percent is not None else "n/a"
        print(f"  ✓ {eid} — {entry_total}/{entry_max} ({pct_str})")
        results.append(question_result)

    # Compute and write summary
    summary = compute_summary(args.model, args.trial, args.judge, results)
    summary["questions_skipped"] = skipped

    summary_path = os.path.join(eval_dir, "summary.json")
    with open(summary_path, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)

    print_summary(summary)
    print(f"\nSummary saved to {summary_path}")


if __name__ == "__main__":
    main()