eval.py

"""Baseline-vs-trained evaluation runner.

Runs N episodes per task with the BASELINE Vendor (untrained Llama 3.2 1B)
and the TRAINED Vendor (LoRA-merged checkpoint), both negotiating against
the real LLMClient (eval_mode=True). Produces:

  - eval_results.json  — per-task metrics
  - eval_comparison.png — side-by-side bar chart for the README

Usage:
    export HF_TOKEN=hf_...
    python eval.py \\
        --base-model unsloth/Llama-3.2-1B-Instruct \\
        --trained-adapter ./vendor_trained_final \\
        --space-url https://ashutosh111-negotiation-arena.hf.space \\
        --episodes-per-task 20

The transformer-loading code path is gated behind ``--enable-llm-vendor``
so this script can also be exercised without GPU / HF model access by
using simple deterministic policies (useful in CI).
"""
from __future__ import annotations

import argparse
import json
import os
import random
import re
from typing import Any, Dict, List, Optional, Tuple

import httpx

TASKS = ["simple_saas", "gdpr_dpa", "enterprise_partnership"]


# ---------------------------------------------------------------------------
# HTTP arena client
# ---------------------------------------------------------------------------
class ArenaClient:
    def __init__(self, base_url: str, timeout: float = 120.0):
        self.base_url = base_url.rstrip("/")
        self._client = httpx.Client(timeout=timeout)

    def reset(self, task_id: str, eval_mode: bool, seed: Optional[int] = None) -> Dict[str, Any]:
        payload: Dict[str, Any] = {"task_id": task_id, "eval_mode": eval_mode}
        if seed is not None:
            payload["seed"] = int(seed)
        r = self._client.post(f"{self.base_url}/reset", json=payload)
        r.raise_for_status()
        return _flatten(r.json())

    def step(self, action: Dict[str, Any]) -> Dict[str, Any]:
        action = {k: v for k, v in action.items() if v is not None}
        r = self._client.post(f"{self.base_url}/step", json={"action": action})
        r.raise_for_status()
        return _flatten(r.json())


def _flatten(payload: Dict[str, Any]) -> Dict[str, Any]:
    if "observation" in payload and isinstance(payload["observation"], dict):
        obs = dict(payload["observation"])
        obs["reward"] = payload.get("reward")
        obs["done"] = payload.get("done", obs.get("done", False))
        return obs
    return payload


# ---------------------------------------------------------------------------
# Vendor policies
# ---------------------------------------------------------------------------
class HeuristicVendor:
    """Deterministic policy used as a CI-friendly stand-in for a real LLM.

    Always proposes the midpoint between vendor and client walkaways on the
    first turn, then accepts whatever the client counters with.
    """

    def __init__(self, seed: int = 0):
        self.rng = random.Random(seed)

    def act(self, obs: Dict[str, Any]) -> Dict[str, Any]:
        history = obs.get("offer_history", [])
        # If the client has already made an offer back to us, accept it
        for entry in reversed(history):
            if entry.get("agent_role") == "client" and entry.get("proposed_terms"):
                return {"action_type": "accept_offer", "agent_role": "vendor"}

        issues = obs.get("open_issues", []) or []
        proposed: Dict[str, Any] = {}
        for issue in issues:
            current = issue.get("current_value")
            proposed[issue["name"]] = current
        return {
            "action_type": "make_offer",
            "agent_role": "vendor",
            "proposed_terms": proposed,
            "reasoning": "heuristic opening offer",
        }


class LLMVendor:
    """Llama-3.2-1B-Instruct backed Vendor (baseline OR trained adapter).

    Loads lazily so we don't require torch/transformers when a user only
    wants to run the heuristic eval path.
    """

    def __init__(
        self,
        base_model: str,
        adapter_path: Optional[str] = None,
        max_new_tokens: int = 512,
        temperature: float = 0.2,
    ):
        self.base_model = base_model
        self.adapter_path = adapter_path
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self._tokenizer = None
        self._model = None

    def _load(self) -> None:
        if self._model is not None:
            return
        from transformers import AutoModelForCausalLM, AutoTokenizer  # local import

        self._tokenizer = AutoTokenizer.from_pretrained(self.base_model)
        self._model = AutoModelForCausalLM.from_pretrained(
            self.base_model, device_map="auto", torch_dtype="auto"
        )
        if self.adapter_path:
            from peft import PeftModel

            self._model = PeftModel.from_pretrained(self._model, self.adapter_path)
        self._model.eval()

    def act(self, obs: Dict[str, Any]) -> Dict[str, Any]:
        self._load()
        prompt = _build_vendor_prompt(obs, int(obs.get("turn_number", 0)) + 1)
        import torch  # local import

        inputs = self._tokenizer(prompt, return_tensors="pt").to(self._model.device)
        with torch.no_grad():
            out = self._model.generate(
                **inputs,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
                do_sample=self.temperature > 0,
                pad_token_id=self._tokenizer.eos_token_id,
            )
        text = self._tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        return _parse_action(text) or {
            "action_type": "walk_away",
            "agent_role": "vendor",
            "reasoning": "llm parse failure",
        }


# ---------------------------------------------------------------------------
# Prompt + parsing helpers (mirrors training notebook)
# ---------------------------------------------------------------------------
def _format_brief(brief: Optional[Dict[str, Any]]) -> str:
    if not brief:
        return "(no brief)"
    parts = [f"Role: {brief.get('role')}"]
    if brief.get("strategy_hints"):
        parts.append(f"Strategy: {brief['strategy_hints']}")
    parts.append("Walkaways: " + json.dumps(brief.get("walkaway_thresholds", {})))
    parts.append("Ideals: " + json.dumps(brief.get("ideal_outcomes", {})))
    parts.append("Priorities: " + json.dumps(brief.get("private_priorities", {})))
    parts.append("Dealbreakers: " + json.dumps(brief.get("dealbreakers", [])))
    return "\n".join(parts)


def _format_issues(issues: List[Dict[str, Any]]) -> str:
    return "\n".join(
        f"- {i.get('name')}: current={i.get('current_value')} ({i.get('type')})"
        for i in issues or []
    ) or "(no open issues)"


def _format_history(history: List[Dict[str, Any]]) -> str:
    if not history:
        return "(no history)"
    return "\n".join(
        f"[{h.get('agent_role')}] {h.get('action_type')} terms={json.dumps(h.get('proposed_terms') or {})}"
        for h in history[-10:]
    )


def _build_vendor_prompt(obs: Dict[str, Any], turn: int) -> str:
    return (
        "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
        "You are an expert negotiation lawyer for the VENDOR.\n\n"
        f"YOUR PRIVATE BRIEF:\n{_format_brief(obs.get('vendor_private_brief'))}\n\n"
        f"CONTRACT CONTEXT:\n{obs.get('contract_draft', '')}\n\n"
        f"OPEN ISSUES:\n{_format_issues(obs.get('open_issues', []))}\n\n"
        f"NEGOTIATION HISTORY:\n{_format_history(obs.get('offer_history', []))}\n\n"
        f"CURRENT TURN: {turn} of {obs.get('max_turns', 30)}\n\n"
        "Respond with ONE JSON action only.<|eot_id|>"
        "<|start_header_id|>user<|end_header_id|>\n\nWhat is your next move?"
        "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
    )


def _parse_action(text: str) -> Optional[Dict[str, Any]]:
    if not text:
        return None
    fence = re.search(r"```(?:json)?\s*([\s\S]*?)```", text.strip())
    if fence:
        text = fence.group(1)
    start, end = text.find("{"), text.rfind("}")
    if start == -1 or end == -1:
        return None
    try:
        parsed = json.loads(text[start : end + 1])
    except json.JSONDecodeError:
        return None
    if isinstance(parsed, dict) and "action_type" in parsed:
        parsed.setdefault("agent_role", "vendor")
        return parsed
    return None


# ---------------------------------------------------------------------------
# Episode driver
# ---------------------------------------------------------------------------
def run_episode(env: ArenaClient, vendor, task_id: str, seed: int, eval_mode: bool) -> Dict[str, Any]:
    obs = env.reset(task_id, eval_mode=eval_mode, seed=seed)
    rewards: List[float] = []
    turns = 0
    max_turns = int(obs.get("max_turns", 30))
    eid = obs.get("episode_id", "")
    while not obs.get("done") and turns < max_turns:
        turns += 1
        action = vendor.act(obs)
        action.setdefault("episode_id", eid)
        action.setdefault("turn_number", turns)
        action.setdefault("agent_role", "vendor")
        obs = env.step(action)
        rewards.append(float(obs.get("reward") or 0.0))
    return {
        "task_id": task_id,
        "turns": turns,
        "deal_struck": bool(obs.get("final_deal_struck", False)),
        "rewards": rewards,
        "final_reward": float(sum(rewards) / len(rewards)) if rewards else 0.0,
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--space-url", default=os.environ.get("SPACE_URL", "http://localhost:8000"))
    p.add_argument("--episodes-per-task", type=int, default=20)
    p.add_argument("--enable-llm-vendor", action="store_true",
                   help="Use real Llama 3.2 1B (requires GPU + transformers/peft)")
    p.add_argument("--base-model", default="unsloth/Llama-3.2-1B-Instruct")
    p.add_argument("--trained-adapter", default=None)
    p.add_argument("--eval-mode", action="store_true",
                   help="Use the LLMClient (70B) as the opposing client; requires HF_TOKEN")
    p.add_argument("--out-dir", default=".")
    args = p.parse_args()

    env = ArenaClient(args.space_url)
    eval_mode = args.eval_mode

    if args.enable_llm_vendor:
        baseline = LLMVendor(args.base_model, adapter_path=None)
        trained = LLMVendor(args.base_model, adapter_path=args.trained_adapter) if args.trained_adapter else baseline
    else:
        baseline = HeuristicVendor(seed=1)
        trained = HeuristicVendor(seed=2)

    results: Dict[str, Dict[str, Dict[str, float]]] = {}
    for task_id in TASKS:
        results[task_id] = {}
        for label, vendor in (("baseline", baseline), ("trained", trained)):
            ep_metrics = []
            for i in range(args.episodes_per_task):
                ep = run_episode(env, vendor, task_id, seed=10_000 + i, eval_mode=eval_mode)
                ep_metrics.append(ep)
            mean_reward = sum(e["final_reward"] for e in ep_metrics) / len(ep_metrics)
            deal_rate = sum(1 for e in ep_metrics if e["deal_struck"]) / len(ep_metrics)
            mean_turns = sum(e["turns"] for e in ep_metrics) / len(ep_metrics)
            results[task_id][label] = {
                "mean_reward": round(mean_reward, 4),
                "deal_rate": round(deal_rate, 4),
                "mean_turns": round(mean_turns, 2),
            }
            print(f"{task_id} / {label}: reward={mean_reward:.3f} deal={deal_rate:.2f} turns={mean_turns:.1f}")

    os.makedirs(args.out_dir, exist_ok=True)
    out_json = os.path.join(args.out_dir, "eval_results.json")
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2)
    print(f"saved {out_json}")

    _plot_comparison(results, os.path.join(args.out_dir, "eval_comparison.png"))
    return 0


def _plot_comparison(results: Dict[str, Dict[str, Dict[str, float]]], out_path: str) -> None:
    try:
        import matplotlib.pyplot as plt
        import numpy as np
    except ImportError:
        print("matplotlib not installed; skipping plot")
        return
    tasks = list(results.keys())
    baseline = [results[t]["baseline"]["mean_reward"] for t in tasks]
    trained = [results[t]["trained"]["mean_reward"] for t in tasks]
    x = np.arange(len(tasks))
    w = 0.35
    fig, ax = plt.subplots(figsize=(9, 5))
    ax.bar(x - w / 2, baseline, w, label="Baseline", color="lightcoral")
    ax.bar(x + w / 2, trained, w, label="Trained", color="seagreen")
    ax.set_xticks(x)
    ax.set_xticklabels(tasks, rotation=15)
    ax.set_ylabel("Mean Reward")
    ax.set_ylim(0, 1)
    ax.set_title("Negotiation Arena: Baseline vs Trained Vendor")
    ax.legend()
    ax.grid(True, axis="y", alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_path, dpi=120, bbox_inches="tight")
    print(f"saved {out_path}")


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