src/chemgraph/eval/runner.py

"""Benchmark runner for ChemGraph multi-model evaluation.

Iterates over ``(model, workflow, query)`` combinations, collects
tool-call outputs, and scores them against ground truth using an
LLM-as-judge approach.
"""

import datetime
import inspect
import json
import os
import traceback
from typing import Any, Dict, List

from chemgraph.agent.llm_agent import ChemGraph
from chemgraph.eval.config import BenchmarkConfig
from chemgraph.eval.datasets import GroundTruthItem, load_dataset
from chemgraph.eval.llm_judge import (
    aggregate_judge_results,
    judge_single_query,
    load_judge_model,
)
from chemgraph.eval.reporter import (
    print_summary_table,
    write_json_report,
    write_markdown_report,
    write_model_detail,
)
from chemgraph.eval.structured_output_judge import (
    aggregate_structured_results,
    judge_structured_output,
)
from chemgraph.utils.get_workflow_from_llm import get_workflow_from_state
from chemgraph.utils.logging_config import setup_logger

logger = setup_logger(__name__)


def _safe_path_component(value: str) -> str:
    """Return a filesystem-safe path component for eval artifacts."""
    text = str(value)
    safe = "".join(ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in text)
    return safe.strip("_") or "unknown"


class ModelBenchmarkRunner:
    """Run evaluation benchmarks across multiple LLM models and workflows.

    Uses an LLM judge to compare the agent's final answer against the
    ground-truth result (binary: correct/wrong).

    Parameters
    ----------
    config : BenchmarkConfig
        Evaluation configuration specifying models, workflows, dataset,
        and output settings.

    Examples
    --------
    >>> from chemgraph.eval import ModelBenchmarkRunner, BenchmarkConfig
    >>> config = BenchmarkConfig(
    ...     models=["gpt-4o-mini", "gemini-2.5-flash"],
    ...     judge_model="gpt-4o",
    ... )
    >>> runner = ModelBenchmarkRunner(config)
    >>> results = asyncio.run(runner.run_all())
    >>> runner.report()
    """

    def __init__(self, config: BenchmarkConfig):
        """Initialize the benchmark runner.

        Parameters
        ----------
        config : BenchmarkConfig
            Validated benchmark configuration.
        """
        self.config = config
        full_dataset: List[GroundTruthItem] = load_dataset(config.dataset)
        # Apply max_queries limit if configured (0 = no limit).
        if config.max_queries > 0:
            self.dataset = full_dataset[: config.max_queries]
            logger.info(
                f"Limiting evaluation to {config.max_queries} of "
                f"{len(full_dataset)} queries"
            )
        else:
            self.dataset = full_dataset
        self.results: Dict[str, Dict[str, dict]] = {}
        self._run_metadata: dict = {}

        # Load judge model only when LLM judge is requested.
        self._judge_llm = None
        if config.judge_type in ("llm", "both"):
            logger.info(f"Loading judge model: {config.judge_model}")
            judge_base_url = config.get_base_url(config.judge_model)
            judge_argo_user = config.get_argo_user()
            self._judge_llm = load_judge_model(
                config.judge_model,
                base_url=judge_base_url,
                argo_user=judge_argo_user,
            )

        if config.judge_type in ("structured", "both"):
            n_with_so = sum(
                1
                for item in self.dataset
                if item.expected_structured_output is not None
            )
            logger.info(
                f"Structured output judge enabled: {n_with_so}/{len(self.dataset)} "
                f"queries have expected structured output"
            )

    # ------------------------------------------------------------------
    # Checkpointing
    # ------------------------------------------------------------------

    def _checkpoint_dir(self) -> str:
        """Return the checkpoint directory path, creating it if needed."""
        d = os.path.join(self.config.output_dir, "checkpoints")
        os.makedirs(d, exist_ok=True)
        return d

    def _checkpoint_path(self, model_name: str, workflow_type: str) -> str:
        """Return the JSONL checkpoint file path for a model/workflow pair.

        Parameters
        ----------
        model_name : str
            Model identifier being evaluated.
        workflow_type : str
            Workflow type being evaluated.

        Returns
        -------
        str
            Checkpoint JSONL file path.
        """
        safe_name = model_name.replace("/", "_").replace(":", "_")
        return os.path.join(
            self._checkpoint_dir(),
            f"{safe_name}_{workflow_type}.jsonl",
        )

    def _save_query_checkpoint(
        self,
        model_name: str,
        workflow_type: str,
        query_id: str,
        query_idx: int,
        query_result: dict,
    ) -> None:
        """Append a single query result to the checkpoint file.

        Each line in the JSONL file is a self-contained JSON object with
        the query ID, index, and full result (raw output + judge scores).
        Append-only writes make this crash-safe: at worst the last line
        may be truncated (one query lost, not all).

        Parameters
        ----------
        model_name : str
            Model identifier being evaluated.
        workflow_type : str
            Workflow type being evaluated.
        query_id : str
            Ground-truth query identifier.
        query_idx : int
            Query index used as the LangGraph thread ID.
        query_result : dict
            Result payload to checkpoint.
        """
        record = {
            "query_id": query_id,
            "query_idx": query_idx,
            **query_result,
        }
        path = self._checkpoint_path(model_name, workflow_type)
        with open(path, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, default=str) + "\n")

    def _load_checkpoint(self, model_name: str, workflow_type: str) -> Dict[str, dict]:
        """Load completed query results from a checkpoint file.

        Parameters
        ----------
        model_name : str
            Model identifier being evaluated.
        workflow_type : str
            Workflow type being evaluated.

        Returns
        -------
        dict
            ``{query_id: {"raw": ..., "judge": ..., "structured_judge": ...}}``
            for each successfully checkpointed query.  Corrupt lines
            (e.g. from a mid-write crash) are silently skipped.
        """
        path = self._checkpoint_path(model_name, workflow_type)
        completed: Dict[str, dict] = {}
        if not os.path.exists(path):
            return completed

        with open(path, "r", encoding="utf-8") as f:
            for line_no, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue
                try:
                    record = json.loads(line)
                    qid = record.get("query_id")
                    if qid is not None:
                        completed[str(qid)] = {
                            "raw": record.get("raw"),
                            "judge": record.get("judge"),
                            "structured_judge": record.get("structured_judge"),
                        }
                except json.JSONDecodeError:
                    logger.warning(
                        f"Skipping corrupt checkpoint line {line_no} in "
                        f"{path} (possible mid-write crash)"
                    )
        if completed:
            logger.info(
                f"Loaded {len(completed)} checkpointed queries for "
                f"{model_name}/{workflow_type}"
            )
        return completed

    def _clear_checkpoint(self, model_name: str, workflow_type: str) -> None:
        """Remove the checkpoint file for a (model, workflow) pair.

        Called when *not* resuming, so that stale checkpoint data from a
        previous run does not leak into the current run.

        Parameters
        ----------
        model_name : str
            Model identifier being evaluated.
        workflow_type : str
            Workflow type being evaluated.
        """
        path = self._checkpoint_path(model_name, workflow_type)
        if os.path.exists(path):
            os.remove(path)
            logger.debug(f"Cleared stale checkpoint: {path}")

    # ------------------------------------------------------------------
    # Core execution
    # ------------------------------------------------------------------

    async def _run_single_model_workflow(
        self,
        model_name: str,
        workflow_type: str,
    ) -> dict:
        """Run all queries for one (model, workflow) pair.

        Parameters
        ----------
        model_name : str
            Model identifier to evaluate.
        workflow_type : str
            Workflow type to evaluate.

        Returns
        -------
        dict
            Contains ``"judge_aggregate"``, ``"judge_details"``, and
            ``"raw_tool_calls"``.
        """
        logger.info(
            f"Starting evaluation: model={model_name}, workflow={workflow_type}"
        )

        # Isolate log directory per model+workflow so parallel runs don't clash.
        run_log_dir = self._workflow_log_dir(model_name, workflow_type)
        os.makedirs(run_log_dir, exist_ok=True)

        try:
            # Resolve per-model base_url and argo_user from config.toml.
            base_url = self.config.get_base_url(model_name)
            argo_user = self.config.get_argo_user()

            # Build desired kwargs and filter to only those accepted by
            # the installed ChemGraph version, so the runner works even
            # against older releases that lack newer parameters.
            desired_kwargs = {
                "model_name": model_name,
                "workflow_type": workflow_type,
                "structured_output": self.config.structured_output,
                "return_option": "state",
                "recursion_limit": self.config.recursion_limit,
                "enable_memory": False,
                "base_url": base_url,
                "argo_user": argo_user,
                "log_dir": run_log_dir,
            }
            sig = inspect.signature(ChemGraph.__init__)
            valid_params = set(sig.parameters.keys()) - {"self"}
            filtered_kwargs = {
                k: v for k, v in desired_kwargs.items() if k in valid_params
            }

            cg = ChemGraph(**filtered_kwargs)
        except Exception as e:
            logger.error(f"Failed to initialise ChemGraph for {model_name}: {e}")
            return self._make_error_result(
                f"Initialisation failed: {e}",
                len(self.dataset),
            )

        raw_tool_calls: List[dict] = []
        per_query_judge_results: List[dict] = []
        per_query_structured_results: List[dict] = []

        # Load checkpoint for resume, or clear stale data for a fresh run.
        checkpoint: Dict[str, dict] = {}
        if self.config.resume:
            checkpoint = self._load_checkpoint(model_name, workflow_type)
        else:
            self._clear_checkpoint(model_name, workflow_type)

        n_skipped = 0
        for idx, item in enumerate(self.dataset):
            # Resume: reuse checkpointed result if available.
            if item.id in checkpoint:
                query_result = checkpoint[item.id]
                n_skipped += 1
                logger.debug(
                    f"Skipping query {idx} ({item.id}): loaded from checkpoint"
                )
            else:
                query_result = await self._run_single_query(
                    cg, item, idx, model_name, workflow_type
                )
                # Checkpoint immediately after each query completes.
                self._save_query_checkpoint(
                    model_name, workflow_type, item.id, idx, query_result
                )

            raw_tool_calls.append(query_result["raw"])
            if query_result.get("judge") is not None:
                per_query_judge_results.append(query_result["judge"])
            if query_result.get("structured_judge") is not None:
                per_query_structured_results.append(query_result["structured_judge"])

        if n_skipped:
            logger.info(
                f"Resumed {model_name}/{workflow_type}: "
                f"{n_skipped} queries from checkpoint, "
                f"{len(self.dataset) - n_skipped} newly evaluated"
            )

        result: Dict[str, Any] = {
            "raw_tool_calls": raw_tool_calls,
        }

        # LLM judge results.
        if self.config.judge_type in ("llm", "both"):
            judge_agg = aggregate_judge_results(per_query_judge_results)
            result["judge_aggregate"] = judge_agg
            result["judge_details"] = per_query_judge_results

        # Structured output judge results.
        if self.config.judge_type in ("structured", "both"):
            struct_agg = aggregate_structured_results(per_query_structured_results)
            result["structured_judge_aggregate"] = struct_agg
            result["structured_judge_details"] = per_query_structured_results

        # Log summary.
        parts = [f"Completed eval {model_name}/{workflow_type}:"]
        if "judge_aggregate" in result:
            jagg = result["judge_aggregate"]
            parts.append(
                f"llm_judge={jagg['accuracy']:.1%} "
                f"({jagg['n_correct']}/{jagg['n_queries']})"
            )
        if "structured_judge_aggregate" in result:
            sagg = result["structured_judge_aggregate"]
            parts.append(
                f"struct_judge={sagg['accuracy']:.1%} "
                f"({sagg['n_correct']}/{sagg['n_queries']})"
            )
        logger.info(" ".join(parts))

        return result

    async def _run_single_query(
        self,
        cg: ChemGraph,
        item: GroundTruthItem,
        idx: int,
        model_name: str,
        workflow_type: str,
    ) -> dict:
        """Execute and evaluate a single query.

        Returns ``{"raw": ..., "judge": ..., "structured_judge": ...}``.

        Parameters
        ----------
        cg : ChemGraph
            Initialized ChemGraph agent.
        item : GroundTruthItem
            Ground-truth query item.
        idx : int
            Query index used as the LangGraph thread ID.
        model_name : str
            Model identifier being evaluated.
        workflow_type : str
            Workflow type being evaluated.

        Returns
        -------
        dict
            Query result containing raw output and judge results.
        """
        try:
            config = {"configurable": {"thread_id": str(idx)}}
            query_log_dir = self._query_log_dir(
                model_name=model_name,
                workflow_type=workflow_type,
                query_idx=idx,
                query_id=item.id,
            )
            os.makedirs(query_log_dir, exist_ok=True)

            old_log_dir = getattr(cg, "log_dir", None)
            old_env_log_dir = os.environ.get("CHEMGRAPH_LOG_DIR")
            cg.log_dir = query_log_dir
            os.environ["CHEMGRAPH_LOG_DIR"] = query_log_dir
            try:
                state = await cg.run(item.query, config)
            finally:
                cg.log_dir = old_log_dir
                if old_env_log_dir is None:
                    os.environ.pop("CHEMGRAPH_LOG_DIR", None)
                else:
                    os.environ["CHEMGRAPH_LOG_DIR"] = old_env_log_dir
            llm_workflow = get_workflow_from_state(state)
            model_tool_calls = llm_workflow.get("tool_calls", [])
            model_result = llm_workflow.get("result", "")
        except Exception as e:
            logger.warning(f"Query {idx} failed for {model_name}/{workflow_type}: {e}")
            logger.debug(traceback.format_exc())
            model_tool_calls = []
            model_result = f"ERROR: {e}"
            llm_workflow = {"tool_calls": [], "result": model_result}

        result: Dict[str, Any] = {"raw": llm_workflow}

        # --- LLM judge ---
        if self.config.judge_type in ("llm", "both") and self._judge_llm is not None:
            judge_result = await judge_single_query(
                judge_llm=self._judge_llm,
                query=item.query,
                expected_result=item.expected_result,
                model_result=model_result,
                expected_tool_calls=item.expected_tool_calls,
                model_tool_calls=model_tool_calls,
            )
            judge_result["query_id"] = item.id
            judge_result["query"] = item.query
            judge_result["category"] = item.category
            result["judge"] = judge_result

        # --- Structured output judge ---
        if self.config.judge_type in ("structured", "both"):
            if item.expected_structured_output is not None:
                struct_result = judge_structured_output(
                    expected=item.expected_structured_output,
                    actual=model_result,
                )
                struct_result["query_id"] = item.id
                struct_result["query"] = item.query
                struct_result["category"] = item.category
                result["structured_judge"] = struct_result
            else:
                logger.debug(
                    f"Query {idx}: no expected_structured_output, "
                    f"skipping structured judge"
                )

        return result

    def _workflow_log_dir(self, model_name: str, workflow_type: str) -> str:
        """Return the base log directory for one model/workflow pair."""
        return os.path.join(
            self.config.output_dir,
            "logs",
            _safe_path_component(model_name),
            _safe_path_component(workflow_type),
        )

    def _query_log_dir(
        self,
        model_name: str,
        workflow_type: str,
        query_idx: int,
        query_id: str,
    ) -> str:
        """Return an isolated artifact directory for one benchmark query."""
        return os.path.join(
            self._workflow_log_dir(model_name, workflow_type),
            f"query_{query_idx:03d}_{_safe_path_component(query_id)}",
        )

    async def run_all(self) -> Dict[str, Dict[str, dict]]:
        """Execute the full benchmark: all models x all workflows.

        Models are run **sequentially** to avoid API rate-limit issues
        and to keep log directories clean.  Within a model, queries run
        sequentially as well (the ``ChemGraph.run`` method already uses
        async streaming internally).

        Returns
        -------
        dict
            ``{model_name: {workflow_type: {"judge_aggregate": ..., ...}}}``
        """
        timestamp = datetime.datetime.now().isoformat()
        self._run_metadata = {
            "timestamp": timestamp,
            "dataset": self.config.dataset,
            "n_queries": len(self.dataset),
            "models": self.config.models,
            "workflow_types": self.config.workflow_types,
            "judge_model": self.config.judge_model,
            "judge_type": self.config.judge_type,
            "structured_output": self.config.structured_output,
            "resume": self.config.resume,
            "tags": self.config.tags,
        }

        self.results = {}

        for model_name in self.config.models:
            self.results[model_name] = {}
            for workflow_type in self.config.workflow_types:
                result = await self._run_single_model_workflow(
                    model_name, workflow_type
                )
                self.results[model_name][workflow_type] = result

                # Write per-model detail file immediately so partial
                # results survive if a later model fails.
                write_model_detail(
                    model_name=model_name,
                    workflow_type=workflow_type,
                    raw_tool_calls=result["raw_tool_calls"],
                    per_query_results=[],
                    output_dir=self.config.output_dir,
                    judge_results=result.get("judge_details"),
                    structured_judge_results=result.get("structured_judge_details"),
                )

                # Write incremental ("running") aggregate report so a
                # usable summary exists even if a later model crashes.
                self._write_running_report()

        return self.results

    # ------------------------------------------------------------------
    # Reporting
    # ------------------------------------------------------------------

    def _write_running_report(self) -> None:
        """Write/overwrite an incremental aggregate report.

        Called after each ``(model, workflow)`` pair completes inside
        ``run_all()``.  The "running" files contain whatever results have
        been collected so far, providing a usable summary even if the
        process crashes before ``report()`` is called.

        The running files are cleaned up by ``report()`` once the final
        timestamped reports are successfully written.
        """
        if not self.results or not self._run_metadata:
            return

        json_path = os.path.join(self.config.output_dir, "benchmark_running.json")
        md_path = os.path.join(self.config.output_dir, "benchmark_running.md")
        try:
            write_json_report(
                results=self.results,
                metadata=self._run_metadata,
                output_path=json_path,
            )
            write_markdown_report(
                results=self.results,
                metadata=self._run_metadata,
                output_path=md_path,
            )
        except Exception as e:
            logger.warning(f"Failed to write running report: {e}")

    def _cleanup_running_report(self) -> None:
        """Remove the incremental running report files.

        Called after ``report()`` has successfully written the final
        timestamped reports.
        """
        for suffix in ("json", "md"):
            path = os.path.join(self.config.output_dir, f"benchmark_running.{suffix}")
            if os.path.exists(path):
                try:
                    os.remove(path)
                    logger.debug(f"Cleaned up running report: {path}")
                except OSError as e:
                    logger.warning(f"Could not remove {path}: {e}")

    def report(self, format: str = "all") -> None:
        """Generate and write evaluation reports.

        Parameters
        ----------
        format : str
            ``"json"``, ``"markdown"``, ``"console"``, or ``"all"``
            (default).
        """
        if not self.results:
            logger.warning("No results to report. Run run_all() first.")
            return

        ts = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

        if format in ("json", "all"):
            write_json_report(
                results=self.results,
                metadata=self._run_metadata,
                output_path=os.path.join(
                    self.config.output_dir, f"benchmark_{ts}.json"
                ),
            )

        if format in ("markdown", "all"):
            write_markdown_report(
                results=self.results,
                metadata=self._run_metadata,
                output_path=os.path.join(self.config.output_dir, f"benchmark_{ts}.md"),
            )

        if format in ("console", "all"):
            print_summary_table(self.results)

        # Clean up incremental running report files now that the final
        # timestamped reports have been written successfully.
        self._cleanup_running_report()

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _make_error_result(error_msg: str, n_queries: int) -> dict:
        """Build an error placeholder result for a failed model init.

        Parameters
        ----------
        error_msg : str
            Error message to store in the aggregate result.
        n_queries : int
            Number of benchmark queries that were skipped.

        Returns
        -------
        dict
            Placeholder aggregate result.
        """
        return {
            "judge_aggregate": {
                "n_queries": n_queries,
                "n_correct": 0,
                "accuracy": 0.0,
                "n_parse_errors": 0,
                "error": error_msg,
            },
            "judge_details": [],
            "raw_tool_calls": [],
        }