tests/evaluation/eval_runner.py

"""
Evaluation runner for GovBridge golden dataset.

Runs each test case through the REAL graph (live LLM calls) and scores
the output using an LLM judge. Requires API keys in .env.

Usage:
    uv run python tests/evaluation/eval_runner.py
    uv run python tests/evaluation/eval_runner.py --case london-complaint-001
    uv run python tests/evaluation/eval_runner.py --category complaint
    uv run python tests/evaluation/eval_runner.py --region "Tamil Nadu"
"""
import argparse
import json
import os
import uuid
from pathlib import Path

from dotenv import load_dotenv

load_dotenv()

from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.types import Command

from src.workflow.graph import build_graph
from src.integrations.llm_factory import get_llm

DATASET_PATH = Path(__file__).parent / "golden_dataset.json"


def _config():
    return {"configurable": {"thread_id": str(uuid.uuid4())}}


def _initial_state(user_message: str):
    return {
        "messages": [{"role": "user", "content": user_message}],
        "location": "",
        "org": "default",
        "collected_details": {},
        "email_cc": [],
        "conversation_complete": False,
        "draft_approved": False,
        "email_draft": None,
    }


def _stream_all(graph, state_or_command, config):
    events = list(graph.stream(state_or_command, config, stream_mode="values"))
    return events[-1] if events else None


def _get_interrupt_value(graph, config) -> str | None:
    snapshot = graph.get_state(config)
    if snapshot.tasks and snapshot.tasks[0].interrupts:
        return snapshot.tasks[0].interrupts[0].value
    return None


def run_case(graph, case: dict) -> dict:
    """Run a single golden dataset case through the real graph."""
    config = _config()
    print(f"\n{'='*60}")
    print(f"Running: {case['id']} ({case['region']})")
    print(f"Message: {case['user_message'][:80]}...")

    # Stream initial message
    _stream_all(graph, _initial_state(case["user_message"]), config)

    # Work through follow-up answers
    for i, answer in enumerate(case["followup_answers"]):
        interrupt_val = _get_interrupt_value(graph, config)
        if interrupt_val is None:
            break
        print(f"  [turn {i+1}] agent asked: {str(interrupt_val)[:100]}...")
        print(f"  [turn {i+1}] user answers: {answer[:80]}...")
        _stream_all(graph, Command(resume=answer), config)

    # When we reach email_writer, approve the draft
    for _ in range(3):  # allow up to 3 approval cycles
        snapshot = graph.get_state(config)
        if not snapshot.next:
            break
        if "email_writer" not in snapshot.next:
            # still in conversation agent — send a "write it" nudge
            _stream_all(graph, Command(resume="Yes, please write the email now"), config)
            continue
        interrupt_val = _get_interrupt_value(graph, config)
        if interrupt_val:
            print(f"  [email draft] approving...")
            _stream_all(graph, Command(resume="approve"), config)

    final_state = graph.get_state(config).values
    email_body = final_state.get("email_body", "")
    email_to = final_state.get("email_to", "")
    email_subject = final_state.get("email_subject", "")
    location = final_state.get("location", "")
    category = final_state.get("category", "")

    print(f"  Location detected: {location}")
    print(f"  Category: {category}")
    print(f"  To: {email_to}")
    print(f"  Subject: {email_subject}")

    return {
        "case_id": case["id"],
        "email_body": email_body,
        "email_to": email_to,
        "email_subject": email_subject,
        "location_detected": location,
        "category_detected": category,
    }


def judge_output(case: dict, result: dict) -> dict:
    """Use an LLM to score the output against expected properties."""
    llm = get_llm("large")

    expected = case["expected"]
    prompt = f"""You are evaluating the output of a civic email writing assistant.

CASE: {case['id']} — {case['region']}
ORIGINAL REQUEST: {case['user_message']}

EMAIL PRODUCED:
To: {result['email_to']}
Subject: {result['email_subject']}
Body:
{result['email_body']}

EVALUATION CRITERIA:
1. Correct addressee level — should be "{expected['addressee_level']}" (local/city/state/national/federal). The email should be addressed to an appropriate {expected['addressee_level']}-level official.
2. Addressee hints — the To/CC should reference one of: {expected['addressee_hints']}
3. Mentions required terms — the email should mention: {expected['email_must_mention']}
4. Does NOT mention wrong locations — must not contain: {expected['email_must_not_mention']}
5. Formal and professional tone
6. Specific and actionable — clearly states the issue or proposal with details

Score each criterion 0 or 1. Return ONLY a JSON object with this exact format:
{{
  "correct_addressee_level": 0 or 1,
  "addressee_hint_matched": 0 or 1,
  "required_terms_present": 0 or 1,
  "no_wrong_locations": 0 or 1,
  "formal_tone": 0 or 1,
  "specific_and_actionable": 0 or 1,
  "total": <sum of above>,
  "notes": "<one sentence on the main weakness if total < 6>"
}}"""

    response = llm.invoke([SystemMessage(content="You are a strict evaluator. Return only valid JSON."),
                           HumanMessage(content=prompt)])
    try:
        scores = json.loads(response.content)
    except Exception:
        # Try to extract JSON from the response
        import re
        match = re.search(r'\{.*\}', response.content, re.DOTALL)
        scores = json.loads(match.group()) if match else {"total": -1, "notes": "parse error"}

    return scores


def main():
    parser = argparse.ArgumentParser(description="Run GovBridge golden dataset evaluation")
    parser.add_argument("--case", help="Run a specific case by ID")
    parser.add_argument("--category", choices=["complaint", "idea"], help="Filter by category")
    parser.add_argument("--region", help="Filter by region (partial match)")
    parser.add_argument("--no-judge", action="store_true", help="Skip LLM judge scoring")
    args = parser.parse_args()

    with open(DATASET_PATH) as f:
        dataset = json.load(f)

    # Apply filters
    if args.case:
        dataset = [c for c in dataset if c["id"] == args.case]
    if args.category:
        dataset = [c for c in dataset if c["category"] == args.category]
    if args.region:
        dataset = [c for c in dataset if args.region.lower() in c["region"].lower()]

    if not dataset:
        print("No matching cases found.")
        return

    print(f"Running {len(dataset)} case(s)...")
    graph = build_graph()

    results = []
    for case in dataset:
        result = run_case(graph, case)

        if not args.no_judge and result["email_body"]:
            print(f"  Judging output...")
            scores = judge_output(case, result)
            result["scores"] = scores
            total = scores.get("total", 0)
            print(f"  Score: {total}/6  {scores.get('notes', '')}")
        else:
            result["scores"] = {}

        results.append(result)

    # Summary
    print(f"\n{'='*60}")
    print("SUMMARY")
    print(f"{'='*60}")
    scored = [r for r in results if r["scores"].get("total") is not None and r["scores"]["total"] >= 0]
    if scored:
        avg = sum(r["scores"]["total"] for r in scored) / len(scored)
        print(f"Average score: {avg:.1f}/6 across {len(scored)} case(s)")
        print()
        for r in results:
            score = r["scores"].get("total", "n/a")
            note = r["scores"].get("notes", "")
            flag = "✓" if isinstance(score, int) and score >= 5 else "✗"
            print(f"  {flag} {r['case_id']:35s} {score}/6  {note}")
    else:
        for r in results:
            status = "completed" if r["email_body"] else "incomplete"
            print(f"  {r['case_id']:35s} {status}")


if __name__ == "__main__":
    main()