eval_runner.py

"""Example WebSocket client for the LangSmith evals service.

Connects to the /ws/evaluate endpoint, sends one EvaluateRequest payload, and
prints the streamed events as they arrive: per-example progress while the run
is in flight, then the aggregate summary (mean score per evaluator) at the end.

This is a thin reference client, NOT the CI runner. For CI/threshold gating use
run_evals.py (in-process, exits 0/1/2). For a minimal token-free smoke test see
ws_smoke_test.py.

Protocol (see app/server.py + app/runner.py):
    envelope: {"type": ..., "data": {...}}
    order:    started -> dataset_synced -> run_planned
              -> example_completed* (or example_failed) -> completed
              error is terminal and the socket closes after it.

Usage:
    pip install websockets python-dotenv

    # Evaluate an existing LangSmith dataset:
    python eval_runner.py --url ws://<host>/ws/evaluate \
        --dataset my-dataset --evaluators correctness,answer_relevance \
        --prefix my-run --metadata git_sha=abc123 release=v1.2.3

    # Sync inline examples first, then evaluate (requires --allow-overwrite
    # to replace examples in an existing dataset):
    python eval_runner.py --url ws://<host>/ws/evaluate \
        --dataset my-dataset --examples examples.json \
        --evaluators correctness --allow-overwrite
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
from pathlib import Path

import websockets
from dotenv import load_dotenv

load_dotenv()

log = logging.getLogger(__name__)
logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))

# One reconnect attempt budget for transient handshake failures (the server
# only accepts one payload per connection, so this is just connect resilience).
MAX_ATTEMPTS = 3
RETRY_BASE_DELAY = 1.0


def _fmt_scores(results: list[dict]) -> str:
    """Compact per-evaluator score string for one example."""
    return ", ".join(
        f"{r.get('key', '?')}={r.get('score')}" for r in results
    ) or "(no scores)"


def _build_payload(
    dataset_name: str,
    evaluators: list[str],
    examples: list[dict] | None,
    experiment_prefix: str | None,
    metadata: dict,
    allow_overwrite: bool,
) -> dict:
    payload: dict = {
        "dataset_name": dataset_name,
        "evaluators": evaluators,
        "metadata": metadata,
        "allow_overwrite": allow_overwrite,
    }
    if experiment_prefix:
        payload["experiment_prefix"] = experiment_prefix
    if examples is not None:
        payload["examples"] = examples
    return payload


async def _consume(ws, payload: dict) -> str:
    """Send the payload, print streamed events, return the terminal type."""
    await ws.send(json.dumps(payload))
    log.info(
        "sent EvaluateRequest (dataset=%s, evaluators=%s, inline_examples=%s)",
        payload["dataset_name"], payload["evaluators"],
        len(payload.get("examples") or []),
    )

    terminal = "disconnected"
    async for raw in ws:
        msg = json.loads(raw)
        kind = msg.get("type")
        data = msg.get("data", {})

        if kind == "started":
            log.info(
                "[started] dataset=%s target=%s target_provider=%s "
                "judge_provider=%s judge_model=%s",
                data.get("dataset_name"), data.get("target_type"),
                data.get("target_provider"), data.get("judge_provider"),
                data.get("judge_model"),
            )
        elif kind == "dataset_synced":
            log.info("[synced] action=%s count=%s",
                     data.get("action"), data.get("example_count"))
        elif kind == "run_planned":
            log.info("[planned] total=%s evaluators=%s",
                     data.get("total_examples"), data.get("evaluators"))
        elif kind == "example_completed":
            # --- progress message: one per example as the run advances ---
            prog = data.get("progress", {})
            log.info("[progress %s/%s] %s",
                     prog.get("completed"), prog.get("total"),
                     _fmt_scores(data.get("results", [])))
            # Surface per-evaluator errors (score=None) inline.
            for r in data.get("results", []):
                if r.get("score") is None:
                    log.warning("    %s: %s", r.get("key"),
                                (r.get("comment") or "")[:200])
        elif kind == "example_failed":
            log.error("[example_failed idx=%s] %s",
                      data.get("index"), data.get("error"))
        elif kind == "completed":
            # --- final summary: mean score per evaluator ---
            log.info("[completed] examples=%s failures=%s duration=%ss",
                     data.get("total_examples"), data.get("failures"),
                     data.get("duration_s"))
            for name, agg in (data.get("summary") or {}).items():
                log.info("    %s: mean=%s n=%s",
                         name, agg.get("mean"), agg.get("count"))
            terminal = "completed"
            break
        elif kind == "error":
            log.error("[error stage=%s] %s", data.get("stage"),
                      data.get("message") or data.get("errors"))
            terminal = "error"
            break
        else:
            log.warning("[unknown type=%s] %s", kind, data)

    return terminal


async def run(
    url: str,
    dataset_name: str,
    evaluators: list[str],
    examples: list[dict] | None,
    experiment_prefix: str | None,
    metadata: dict,
    allow_overwrite: bool,
) -> int:
    payload = _build_payload(
        dataset_name, evaluators, examples,
        experiment_prefix, metadata, allow_overwrite,
    )

    attempt = 0
    while attempt < MAX_ATTEMPTS:
        attempt += 1
        try:
            log.info("connecting to %s (attempt %s/%s)",
                     url, attempt, MAX_ATTEMPTS)
            async with websockets.connect(url, ping_interval=20,
                                          ping_timeout=20) as ws:
                terminal = await _consume(ws, payload)
            return 0 if terminal == "completed" else 1
        except (OSError, websockets.exceptions.WebSocketException) as exc:
            if attempt >= MAX_ATTEMPTS:
                log.error("giving up after %s attempts: %s", attempt, exc)
                return 1
            delay = RETRY_BASE_DELAY * attempt
            log.warning("connect failed (%s); retrying in %.1fs", exc, delay)
            await asyncio.sleep(delay)

    return 1


def _parse_metadata(pairs: list[str] | None) -> dict:
    """Turn ['k=v', 'a=b'] into {'k': 'v', 'a': 'b'}."""
    out: dict[str, str] = {}
    for item in pairs or []:
        if "=" not in item:
            raise ValueError(f"metadata must be KEY=VALUE, got {item!r}")
        key, value = item.split("=", 1)
        out[key.strip()] = value.strip()
    return out


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", default="ws://localhost:8000/ws/evaluate",
                        help="WebSocket URL of the /ws/evaluate endpoint.")
    parser.add_argument("--dataset", required=True,
                        help="LangSmith dataset name.")
    parser.add_argument("--examples", default=None,
                        help="Path to a JSON array of inline examples to sync "
                             "before evaluating (omit to use the dataset as-is).")
    parser.add_argument("--evaluators", default="correctness",
                        help="Comma-separated evaluator names.")
    parser.add_argument("--prefix", default=None,
                        help="Experiment prefix for this run.")
    parser.add_argument("--metadata", nargs="*", metavar="KEY=VALUE",
                        help="Optional run metadata as key=value pairs.")
    parser.add_argument("--allow-overwrite", action="store_true",
                        help="Allow replacing examples in an existing dataset.")
    args = parser.parse_args()

    examples = None
    if args.examples:
        examples = json.loads(Path(args.examples).read_text(encoding="utf-8"))

    evaluators = [e.strip() for e in args.evaluators.split(",") if e.strip()]
    metadata = _parse_metadata(args.metadata)

    return asyncio.run(run(
        url=args.url,
        dataset_name=args.dataset,
        evaluators=evaluators,
        examples=examples,
        experiment_prefix=args.prefix,
        metadata=metadata,
        allow_overwrite=args.allow_overwrite,
    ))


if __name__ == "__main__":
    raise SystemExit(main())