eval/intent/run_eval.py

"""Intent-routing eval runner (E3).

Feeds each golden case in `intent_dataset.json` to the live 6-intent router
(`OrchestratorAgent.classify`), then scores correctness + records latency and
token usage. Prints a per-case detail table and an aggregate summary, and
writes a timestamped JSON report under `results/` (never overwritten — one file
per run, so runs can be diffed over time).

Run before every deploy that touches the router prompt or its few-shots.
Invoke as a module (`-m`) so the repo root is on `sys.path` and `src` imports
resolve — running the file path directly (`python eval/intent/run_eval.py`)
puts only `eval/intent/` on the path and fails:

    uv run python -m eval.intent.run_eval
    uv run python -m eval.intent.run_eval --limit 6     # quick smoke test

Tokens come straight from the model response (LangChain `usage_metadata` via a
callback) — no Langfuse needed. The router is called unmodified: it already
accepts a `callbacks=` list and forwards it into the chain config.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import statistics
import time
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult

from src.agents.orchestration import OrchestratorAgent

_HERE = Path(__file__).resolve().parent
DATASET = _HERE / "intent_dataset.json"
RESULTS_DIR = _HERE / "results"

INTENTS = [
    "chat",
    "help",
    "problem_statement",
    "check",
    "unstructured_flow",
    "structured_flow",
]

# Short labels so the EXPECT->GOT column stays narrow in the detail table.
_ABBR = {
    "chat": "chat",
    "help": "help",
    "problem_statement": "prob_stmt",
    "check": "check",
    "unstructured_flow": "unstruct",
    "structured_flow": "structF",
}


class _UsageCollector(BaseCallbackHandler):
    """Sums token usage across the LLM calls made during one classify().

    Reads `usage_metadata` off each returned message (the canonical LangChain
    field), falling back to `llm_output['token_usage']` for providers that only
    populate the legacy field.
    """

    def __init__(self) -> None:
        self.input_tokens = 0
        self.output_tokens = 0
        self.total_tokens = 0

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        before = self.total_tokens
        for generation_list in response.generations:
            for generation in generation_list:
                message = getattr(generation, "message", None)
                usage = getattr(message, "usage_metadata", None) if message else None
                if usage:
                    self.input_tokens += usage.get("input_tokens", 0)
                    self.output_tokens += usage.get("output_tokens", 0)
                    self.total_tokens += usage.get("total_tokens", 0)
        if self.total_tokens == before and response.llm_output:
            usage = response.llm_output.get("token_usage") or {}
            self.input_tokens += usage.get("prompt_tokens", 0)
            self.output_tokens += usage.get("completion_tokens", 0)
            self.total_tokens += usage.get("total_tokens", 0)

    @property
    def tokens(self) -> dict[str, int]:
        return {
            "input": self.input_tokens,
            "output": self.output_tokens,
            "total": self.total_tokens,
        }


@dataclass
class CaseResult:
    id: str
    lang: str
    message: str
    expected: str
    got: str
    correct: bool
    latency_ms: int
    tokens: dict[str, int]


def load_cases(path: Path) -> list[dict[str, Any]]:
    """Read the `cases` array, skipping the leading `_*` doc keys and `schema`."""
    data = json.loads(path.read_text(encoding="utf-8"))
    return list(data["cases"])


@dataclass
class _LangfuseCtx:
    """Optional Langfuse sink — one session groups all cases of a run."""

    session_id: str
    client: Any


def _new_langfuse_handler(lf_ctx: _LangfuseCtx, case: dict[str, Any]) -> Any:
    """Per-case LangChain callback so each trace carries the case's labels."""
    from langfuse.callback import CallbackHandler

    from src.config.settings import settings

    return CallbackHandler(
        public_key=settings.LANGFUSE_PUBLIC_KEY,
        secret_key=settings.LANGFUSE_SECRET_KEY,
        host=settings.LANGFUSE_HOST,
        session_id=lf_ctx.session_id,
        trace_name=f"intent_eval/{case['id']}",
        metadata={
            "case_id": case["id"],
            "expected": case["expected_intent"],
            "lang": case["lang"],
        },
        tags=["intent-eval", case["expected_intent"], case["lang"]],
    )


def _score_langfuse(lf_ctx: _LangfuseCtx, handler: Any, result: CaseResult) -> None:
    """Attach a 1/0 correctness score to the case's trace. Best-effort."""
    try:
        lf_ctx.client.score(
            trace_id=handler.get_trace_id(),
            name="intent_correct",
            value=1 if result.correct else 0,
            comment=f"{result.expected} -> {result.got}",
        )
    except Exception:  # noqa: BLE001, S110 — scoring must never break the run
        pass


async def run_case(
    agent: OrchestratorAgent,
    case: dict[str, Any],
    lf_ctx: _LangfuseCtx | None = None,
) -> CaseResult:
    """Classify one message; never throws — a failed call is recorded as ERROR."""
    collector = _UsageCollector()
    callbacks: list[Any] = [collector]
    lf_handler = _new_langfuse_handler(lf_ctx, case) if lf_ctx else None
    if lf_handler is not None:
        callbacks.append(lf_handler)

    start = time.perf_counter()
    got: str
    try:
        decision = await agent.classify(case["message"], callbacks=callbacks)
        got = decision.intent
    except Exception as exc:  # noqa: BLE001 — one bad case shouldn't kill the run
        got = f"ERROR:{type(exc).__name__}"
    latency_ms = round((time.perf_counter() - start) * 1000)

    result = CaseResult(
        id=case["id"],
        lang=case["lang"],
        message=case["message"],
        expected=case["expected_intent"],
        got=got,
        correct=got == case["expected_intent"],
        latency_ms=latency_ms,
        tokens=collector.tokens,
    )
    if lf_ctx is not None and lf_handler is not None:
        _score_langfuse(lf_ctx, lf_handler, result)
    return result


def _group_accuracy(results: list[CaseResult], key: str) -> dict[str, dict[str, Any]]:
    out: dict[str, dict[str, Any]] = {}
    keys = INTENTS if key == "expected" else sorted({getattr(r, key) for r in results})
    for k in keys:
        sub = [r for r in results if getattr(r, key) == k]
        if not sub:
            continue
        passed = sum(r.correct for r in sub)
        out[k] = {
            "n": len(sub),
            "passed": passed,
            "accuracy": round(passed / len(sub), 3),
        }
    return out


def summarize(results: list[CaseResult]) -> dict[str, Any]:
    n = len(results)
    passed = sum(r.correct for r in results)
    latencies = [r.latency_ms for r in results]
    tok_in = sum(r.tokens["input"] for r in results)
    tok_out = sum(r.tokens["output"] for r in results)
    tok_total = sum(r.tokens["total"] for r in results)
    return {
        "total": n,
        "passed": passed,
        "accuracy": round(passed / n, 3) if n else 0.0,
        "runtime_avg_ms": round(statistics.mean(latencies)) if latencies else 0,
        "runtime_total_s": round(sum(latencies) / 1000, 1),
        "tokens": {
            "input": tok_in,
            "output": tok_out,
            "total": tok_total,
            "avg_total_per_case": round(tok_total / n) if n else 0,
        },
        "by_intent": _group_accuracy(results, "expected"),
        "by_lang": _group_accuracy(results, "lang"),
    }


def _truncate(text: str, width: int) -> str:
    text = text.replace("\n", " ")
    return text if len(text) <= width else text[: width - 3] + "..."


def format_table(results: list[CaseResult]) -> str:
    header = (
        f"{'ID':<15} {'L':<3} {'QUESTION':<40} "
        f"{'EXPECT->GOT':<20} {'OK':<3} {'MS':>5} {'TOK':>6}"
    )
    rule = "-" * len(header)
    lines = [rule, header, rule]
    for r in results:
        exp_got = f"{_ABBR.get(r.expected, r.expected)}->{_ABBR.get(r.got, r.got)}"
        ok = "ok" if r.correct else "X"
        lines.append(
            f"{r.id:<15} {r.lang:<3} {_truncate(r.message, 40):<40} "
            f"{_truncate(exp_got, 20):<20} {ok:<3} {r.latency_ms:>5} {r.tokens['total']:>6}"
        )
    lines.append(rule)
    return "\n".join(lines)


def format_summary(summary: dict[str, Any], results: list[CaseResult]) -> str:
    lines = ["SUMMARY"]
    lines.append(
        f"  Overall    {summary['passed']}/{summary['total']} correct"
        f"   ({summary['accuracy'] * 100:.1f}%)"
    )
    lines.append(
        f"  Runtime    avg {summary['runtime_avg_ms']} ms"
        f"   |  total {summary['runtime_total_s']} s"
    )
    tok = summary["tokens"]
    lines.append(
        f"  Tokens     avg {tok['avg_total_per_case']}"
        f"   |  total {tok['total']}   (in {tok['input']} / out {tok['output']})"
    )
    lines.append("")
    lines.append("  By intent")
    for intent, m in summary["by_intent"].items():
        lines.append(
            f"    {intent:<18} {m['passed']}/{m['n']}  {m['accuracy'] * 100:.0f}%"
        )
    lines.append("  By language")
    for lang, m in summary["by_lang"].items():
        lines.append(
            f"    {lang:<18} {m['passed']}/{m['n']}  {m['accuracy'] * 100:.0f}%"
        )
    failures = [r for r in results if not r.correct]
    lines.append("")
    lines.append(f"  FAILURES ({len(failures)})")
    for r in failures:
        lines.append(f"    {r.id:<14} [{r.lang}]  {r.expected:<12} -> {r.got}")
    return "\n".join(lines)


def build_report(
    results: list[CaseResult], summary: dict[str, Any], meta: dict[str, Any]
) -> dict[str, Any]:
    run = {**meta, **{k: summary[k] for k in (
        "total", "passed", "accuracy", "runtime_avg_ms", "runtime_total_s", "tokens"
    )}}
    return {
        "run": run,
        "by_intent": summary["by_intent"],
        "by_lang": summary["by_lang"],
        "cases": [asdict(r) for r in results],
    }


def _model_name() -> str:
    try:
        from src.config.settings import settings

        return str(settings.azureai_deployment_name_4o)
    except Exception:  # noqa: BLE001 — meta only; .env may be absent
        return "gpt-4o"


async def main() -> None:
    parser = argparse.ArgumentParser(description="Intent-routing eval (E3)")
    parser.add_argument("--dataset", type=Path, default=DATASET)
    parser.add_argument("--limit", type=int, default=0, help="run first N cases only")
    parser.add_argument("--prompt-version", default="intent_router.md")
    parser.add_argument("--no-table", action="store_true", help="skip the detail table")
    parser.add_argument(
        "--langfuse", action="store_true",
        help="also send each case as a Langfuse trace + correctness score",
    )
    args = parser.parse_args()

    cases = load_cases(args.dataset)
    if args.limit:
        cases = cases[: args.limit]

    started = datetime.now()
    print(f"Intent Routing Eval -- {started:%Y-%m-%d %H:%M:%S}")
    print(f"dataset: {args.dataset.name} ({len(cases)})  model: {_model_name()}  "
          f"prompt: {args.prompt_version}")

    lf_ctx: _LangfuseCtx | None = None
    if args.langfuse:
        try:
            from src.observability.langfuse.langfuse import get_langfuse

            lf_ctx = _LangfuseCtx(
                session_id=f"intent_eval_{started:%Y%m%d_%H%M%S}",
                client=get_langfuse(),  # type: ignore[no-untyped-call]
            )
            print(f"langfuse: enabled (session {lf_ctx.session_id})")
        except Exception as exc:  # noqa: BLE001 — Langfuse is optional
            print(f"langfuse: disabled ({type(exc).__name__}: {exc})")

    agent = OrchestratorAgent()
    results: list[CaseResult] = []
    for case in cases:
        results.append(await run_case(agent, case, lf_ctx))

    if lf_ctx is not None:
        try:
            lf_ctx.client.flush()
        except Exception:  # noqa: BLE001, S110 — flush failure shouldn't fail the run
            pass

    summary = summarize(results)
    if not args.no_table:
        print(format_table(results))
    print(format_summary(summary, results))

    meta = {
        "timestamp": started.isoformat(timespec="seconds"),
        "dataset": args.dataset.name,
        "model": _model_name(),
        "prompt_version": args.prompt_version,
        "langfuse_session": lf_ctx.session_id if lf_ctx else None,
    }
    report = build_report(results, summary, meta)
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    out_path = RESULTS_DIR / f"eval_result_{started:%Y-%m-%d_%H%M%S}.json"
    out_path.write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(f"\n-> saved: {out_path.relative_to(_HERE.parent.parent)}")


if __name__ == "__main__":
    asyncio.run(main())