ellements 0.2.0__py3-none-any.whl

ellements/benchmarking/harness.py

@@ -0,0 +1,342 @@
+"""BenchmarkModel adapter and orchestration for `lm-evaluation-harness`.
+
+Wraps an :class:`~ellements.core.LLMClient` (optionally with a
+:class:`~ellements.execution.Strategy`) as an ``lm_eval.api.model.LM``
+so any prompting strategy can be evaluated against standard benchmarks
+(MMLU, GSM8K, HellaSwag, ARC, HumanEval, …).
+
+Key design points:
+
+- **Single persistent event loop owned by the harness.** All async LLM
+  calls are dispatched onto one background thread running a single
+  ``asyncio`` loop. No per-call ``asyncio.run`` and no thread-pool
+  re-entry. The loop is closed via :meth:`BenchmarkModel.shutdown` or
+  by using the model as a context manager.
+- **Provider prefix is preserved.** The full ``"openai/gpt-4o"`` (or
+  ``"anthropic/..."``, etc.) is passed unchanged to ``LLMClient``;
+  litellm needs the prefix for routing.
+- **Hard-fail on missing log-probabilities.** When a task requires
+  ``loglikelihood`` and the configured model does not return token
+  log-probs, the harness raises :class:`LogprobsUnsupportedError`
+  immediately rather than silently returning ``0.0``.
+- **Real per-sample progress.** ``on_progress`` fires before each task
+  starts (per strategy/task pair) so callers can render a progress bar
+  that actually reflects work happening.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+from collections.abc import Callable, Coroutine, Iterable
+from typing import Any, TypeVar
+
+from ellements.core import LLMClient, LogprobsUnsupportedError
+from ellements.execution import Strategy, StrategyConfigInput
+
+try:
+    from lm_eval.api.model import LM
+except ImportError as exc:  # pragma: no cover
+    raise ImportError(
+        "lm-evaluation-harness is required for benchmarking. "
+        "Install it with: pip install 'ellements[benchmarking]'"
+    ) from exc
+
+from .results import BenchmarkComparison
+
+_logger = logging.getLogger(__name__)
+T = TypeVar("T")
+
+OnBenchmarkProgress = Callable[[str, str, int, int], None]
+"""Callback signature: ``(strategy_name, task_name, current_idx, total)``."""
+
+
+class _LoopRunner:
+    """Owns a persistent asyncio event loop on a dedicated thread."""
+
+    def __init__(self) -> None:
+        self._loop = asyncio.new_event_loop()
+        self._thread = threading.Thread(
+            target=self._loop.run_forever,
+            name="ellements-benchmarking-loop",
+            daemon=True,
+        )
+        self._thread.start()
+
+    def submit(self, coro: Coroutine[Any, Any, T]) -> T:
+        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        return future.result()
+
+    def close(self) -> None:
+        if not self._loop.is_closed():
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            self._thread.join(timeout=5.0)
+            self._loop.close()
+
+
+class BenchmarkModel(LM):  # type: ignore[misc]
+    """Wraps an :class:`LLMClient` (optionally with a :class:`Strategy`)
+    as an ``lm_eval.api.model.LM``.
+
+    Args:
+        client: The LLM client used for generation.
+        strategy: Optional execution strategy. When set, ``generate_until``
+            routes through ``strategy.execute()``. Strategies do **not**
+            affect ``loglikelihood`` — those are an intrinsic model
+            property.
+        strategy_config: Extra config forwarded to ``strategy.execute()``.
+        requires_logprobs: When True, the harness validates at startup
+            that the underlying model returns token log-probabilities,
+            raising :class:`LogprobsUnsupportedError` if not.
+    """
+
+    def __init__(
+        self,
+        client: LLMClient,
+        strategy: Strategy | None = None,
+        strategy_config: StrategyConfigInput = None,
+        *,
+        requires_logprobs: bool = False,
+    ) -> None:
+        super().__init__()
+        self._client = client
+        self._strategy = strategy
+        self._strategy_config: StrategyConfigInput = strategy_config or {}
+        self._loop = _LoopRunner()
+        self._closed = False
+
+        if requires_logprobs:
+            self._verify_logprobs_support()
+
+    def __enter__(self) -> BenchmarkModel:
+        return self
+
+    def __exit__(self, *_exc: object) -> None:
+        self.shutdown()
+
+    def shutdown(self) -> None:
+        """Shut down the persistent event loop. Idempotent."""
+        if not self._closed:
+            self._loop.close()
+            self._closed = True
+
+    # ── generate_until (generative benchmarks) ───────────────────────
+
+    def generate_until(self, requests: list[Any]) -> list[str]:
+        results: list[str] = []
+        for req in requests:
+            context, gen_kwargs = req.args
+            stop_seqs: Iterable[str] = gen_kwargs.get("until", []) or []
+
+            if self._strategy is not None:
+                output = self._loop.submit(self._strategy_call(context))
+            else:
+                output = self._loop.submit(self._direct_call(context))
+
+            for stop in stop_seqs:
+                if stop in output:
+                    output = output[: output.index(stop)]
+
+            results.append(output)
+        return results
+
+    async def _direct_call(self, prompt: str) -> str:
+        return await self._client.complete(
+            [{"role": "user", "content": prompt}]
+        )
+
+    async def _strategy_call(self, prompt: str) -> str:
+        assert self._strategy is not None
+        result = await self._strategy.execute(
+            prompts={"default": prompt},
+            client=self._client,
+            tools=None,
+            config=self._strategy_config,
+        )
+        return result.output
+
+    # ── loglikelihood (multiple-choice benchmarks) ───────────────────
+
+    def loglikelihood(self, requests: list[Any]) -> list[tuple[float, bool]]:
+        """Compute log-likelihoods. Hard-fails when not supported."""
+        results: list[tuple[float, bool]] = []
+        for req in requests:
+            context, continuation = req.args
+            ll, is_greedy = self._loop.submit(
+                self._loglikelihood_call(context, continuation)
+            )
+            results.append((ll, is_greedy))
+        return results
+
+    def loglikelihood_rolling(self, requests: list[Any]) -> list[float]:
+        results: list[float] = []
+        for req in requests:
+            (text,) = req.args
+            ll, _ = self._loop.submit(self._loglikelihood_call("", text))
+            results.append(ll)
+        return results
+
+    async def _loglikelihood_call(
+        self, context: str, continuation: str
+    ) -> tuple[float, bool]:
+        return await self._client.loglikelihood(
+            context=context,
+            continuation=continuation,
+            max_tokens=max(1, len(continuation.split()) * 3),
+        )
+
+    def _verify_logprobs_support(self) -> None:
+        """Make a one-token probe call; raise if logprobs aren't returned."""
+        try:
+            self._loop.submit(self._loglikelihood_call("Hello", " world"))
+        except LogprobsUnsupportedError:
+            raise
+        except Exception as exc:
+            raise LogprobsUnsupportedError(
+                f"Probe loglikelihood call failed for model {self._client.model!r}: {exc}"
+            ) from exc
+
+
+# ── public convenience functions ─────────────────────────────────────
+
+
+def _tasks_require_logprobs(tasks: list[str]) -> bool:
+    """Best-effort introspection of lm-eval tasks for loglikelihood output type."""
+    try:
+        from lm_eval import tasks as lm_tasks
+    except ImportError:  # pragma: no cover
+        return False
+
+    try:
+        task_dict = lm_tasks.get_task_dict(tasks)
+    except Exception:  # pragma: no cover
+        return False
+
+    for value in task_dict.values():
+        output_type = getattr(value, "OUTPUT_TYPE", None) or getattr(
+            getattr(value, "_config", None), "output_type", None
+        )
+        if isinstance(output_type, str) and output_type.startswith(
+            "loglikelihood"
+        ):
+            return True
+        if isinstance(output_type, str) and output_type == "multiple_choice":
+            return True
+    return False
+
+
+def run_benchmark(
+    *,
+    model: str,
+    strategy: Strategy | None = None,
+    strategy_config: StrategyConfigInput = None,
+    tasks: list[str],
+    num_fewshot: int | None = None,
+    limit: int | float | None = None,
+    batch_size: int | None = 1,
+    on_progress: OnBenchmarkProgress | None = None,
+    strategy_name: str = "default",
+    client: LLMClient | None = None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Run lm-eval benchmarks, optionally through an ellements strategy.
+
+    Args:
+        model: Litellm-format model id (e.g. ``"openai/gpt-4o"``). The
+            full provider prefix is preserved.
+        strategy: Optional :class:`Strategy`; routes generative tasks
+            through ``strategy.execute()``.
+        strategy_config: Extra config forwarded to the strategy.
+        tasks: lm-eval task names (required and non-empty).
+        num_fewshot: Number of few-shot examples.
+        limit: Sample limit per task (int = count; float in [0,1] = ratio).
+        batch_size: Batch size for lm-eval (defaults to 1 for API models).
+        on_progress: Fires before each task starts.
+        strategy_name: Logical label used in progress callbacks.
+        client: Optional pre-built :class:`LLMClient`. Useful for sharing
+            observers or providing API keys/headers. When None, one is
+            constructed from *model*.
+
+    Returns:
+        The raw lm-eval results dict.
+    """
+    from lm_eval.evaluator import simple_evaluate
+
+    if not tasks:
+        raise ValueError("tasks must be a non-empty list (e.g. ['gsm8k'])")
+
+    requires_logprobs = _tasks_require_logprobs(tasks)
+
+    owned_client = client is None
+    llm_client = client or LLMClient(model=model)
+
+    with BenchmarkModel(
+        client=llm_client,
+        strategy=strategy,
+        strategy_config=strategy_config,
+        requires_logprobs=requires_logprobs,
+    ) as benchmark_model:
+        for i, task in enumerate(tasks):
+            if on_progress is not None:
+                on_progress(strategy_name, task, i, len(tasks))
+
+        results = simple_evaluate(
+            model=benchmark_model,
+            tasks=tasks,
+            num_fewshot=num_fewshot,
+            limit=limit,
+            batch_size=batch_size,
+            **kwargs,
+        )
+
+    if owned_client:
+        _logger.debug("BenchmarkModel finished; owned client released.")
+    typed_results: dict[str, Any] = results
+    return typed_results
+
+
+def compare_benchmarks(
+    *,
+    model: str,
+    strategies: dict[str, Strategy | None],
+    tasks: list[str],
+    num_fewshot: int | None = None,
+    limit: int | float | None = None,
+    batch_size: int | None = 1,
+    on_progress: OnBenchmarkProgress | None = None,
+    client: LLMClient | None = None,
+    **kwargs: Any,
+) -> BenchmarkComparison:
+    """Run the same benchmarks across multiple strategies and compare results."""
+    if not tasks:
+        raise ValueError("tasks must be a non-empty list (e.g. ['gsm8k'])")
+
+    all_results: dict[str, dict[str, Any]] = {}
+
+    for name, strategy in strategies.items():
+        _logger.info("Running benchmarks for strategy: %s", name)
+        all_results[name] = run_benchmark(
+            model=model,
+            strategy=strategy,
+            tasks=tasks,
+            num_fewshot=num_fewshot,
+            limit=limit,
+            batch_size=batch_size,
+            on_progress=on_progress,
+            strategy_name=name,
+            client=client,
+            **kwargs,
+        )
+
+    return BenchmarkComparison(
+        model=model, tasks=tasks, results=all_results
+    )
+
+
+__all__ = [
+    "BenchmarkModel",
+    "OnBenchmarkProgress",
+    "compare_benchmarks",
+    "run_benchmark",
+]

ellements/benchmarking/results.py

@@ -0,0 +1,173 @@
+"""Benchmark comparison results with formatting and export."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class BenchmarkComparison:
+    """Holds results from :func:`compare_benchmarks` and provides
+    formatting, comparison, and export utilities.
+
+    Parameters
+    ----------
+    model : str
+        The model used across all strategies.
+    tasks : list[str]
+        The benchmark tasks that were run.
+    results : dict[str, dict]
+        Map of strategy name → raw lm-eval results dict.
+    """
+
+    model: str
+    tasks: list[str]
+    results: dict[str, dict[str, Any]]
+
+    # -- Extraction helpers ------------------------------------------------
+
+    def _extract_scores(self) -> dict[str, dict[str, float | None]]:
+        """Extract metric scores per strategy per task.
+
+        Returns ``{strategy: {task_metric: score}}``.
+        """
+        scores: dict[str, dict[str, float | None]] = {}
+        for strategy_name, raw in self.results.items():
+            strategy_scores: dict[str, float | None] = {}
+            task_results = raw.get("results", {})
+            for task_name, metrics in task_results.items():
+                for metric_key, value in metrics.items():
+                    if metric_key.endswith(",none"):
+                        metric_key = metric_key[: -len(",none")]
+                    if isinstance(value, (int, float)):
+                        key = f"{task_name}/{metric_key}"
+                        strategy_scores[key] = round(float(value), 4)
+            scores[strategy_name] = strategy_scores
+        return scores
+
+    # -- Formatting --------------------------------------------------------
+
+    def print_table(self) -> str:
+        """Print a formatted comparison table and return it as a string.
+
+        Rows = metric (task/metric_name), Columns = strategies.
+        """
+        scores = self._extract_scores()
+        strategies = list(scores.keys())
+
+        # Collect all unique metric keys across strategies
+        all_metrics: list[str] = []
+        seen: set[str] = set()
+        for strat_scores in scores.values():
+            for k in strat_scores:
+                if k not in seen:
+                    all_metrics.append(k)
+                    seen.add(k)
+
+        if not all_metrics:
+            msg = "(No benchmark results to display)"
+            print(msg)
+            return msg
+
+        # Column widths
+        metric_col_width = max(len(m) for m in all_metrics)
+        metric_col_width = max(metric_col_width, len("Metric"))
+        strat_col_width = max(max((len(s) for s in strategies), default=8), 10)
+
+        # Header
+        header = f"{'Metric':<{metric_col_width}}"
+        for s in strategies:
+            header += f"  {s:>{strat_col_width}}"
+        separator = "-" * len(header)
+
+        lines = [
+            f"Model: {self.model}",
+            f"Tasks: {', '.join(self.tasks)}",
+            "",
+            header,
+            separator,
+        ]
+
+        # Find best score per metric for highlighting
+        for metric in all_metrics:
+            vals = {s: scores[s].get(metric) for s in strategies}
+            best_val = max(
+                (v for v in vals.values() if v is not None), default=None
+            )
+
+            row = f"{metric:<{metric_col_width}}"
+            for s in strategies:
+                v = vals[s]
+                if v is None:
+                    cell = "—"
+                elif v == best_val and len(strategies) > 1:
+                    cell = f"{v:.4f} ★"
+                else:
+                    cell = f"{v:.4f}"
+                row += f"  {cell:>{strat_col_width}}"
+            lines.append(row)
+
+        lines.append(separator)
+        table = "\n".join(lines)
+        print(table)
+        return table
+
+    def best_strategy(self, metric: str) -> str | None:
+        """Return the strategy name with the highest score for *metric*.
+
+        Parameters
+        ----------
+        metric : str
+            Metric key like ``"gsm8k/acc"`` or ``"mmlu/acc"``.
+        """
+        scores = self._extract_scores()
+        best_name: str | None = None
+        best_val: float | None = None
+        for strat, strat_scores in scores.items():
+            val = strat_scores.get(metric)
+            if val is not None and (best_val is None or val > best_val):
+                best_val = val
+                best_name = strat
+        return best_name
+
+    # -- Export ------------------------------------------------------------
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return a JSON-serialisable dictionary."""
+        return {
+            "model": self.model,
+            "tasks": self.tasks,
+            "scores": self._extract_scores(),
+        }
+
+    def to_json(self, path: str) -> None:
+        """Write comparison results to a JSON file."""
+        with open(path, "w") as f:
+            json.dump(self.to_dict(), f, indent=2)
+
+    def to_csv(self, path: str) -> None:
+        """Write comparison results to a CSV file."""
+        import csv
+
+        scores = self._extract_scores()
+        strategies = list(scores.keys())
+
+        all_metrics: list[str] = []
+        seen: set[str] = set()
+        for strat_scores in scores.values():
+            for k in strat_scores:
+                if k not in seen:
+                    all_metrics.append(k)
+                    seen.add(k)
+
+        with open(path, "w", newline="") as f:
+            writer = csv.writer(f)
+            writer.writerow(["metric"] + strategies)
+            for metric in all_metrics:
+                row = [metric]
+                for s in strategies:
+                    val = scores[s].get(metric)
+                    row.append(f"{val:.4f}" if val is not None else "")
+                writer.writerow(row)

ellements/cli/__init__.py

@@ -0,0 +1,80 @@
+"""CLI, TUI, and terminal interaction helpers for ellements."""
+
+from .adapters import (
+    guideline_provider_from_library,
+    persona_provider_from_library,
+)
+from .agent_tui import (
+    BUILTIN_COMMANDS,
+    AgentRunner,
+    AgentTUI,
+    FileSaveHandler,
+    GuidelineProvider,
+    Mode,
+    PersonaProvider,
+    RunResult,
+    SaveHandler,
+    SlashCommand,
+    TuiConfig,
+    TuiControls,
+    run_agent_tui,
+)
+from .components import (
+    ISSUE_STYLES,
+    ROLE_STYLES,
+    app_header,
+    chat_input_prompt,
+    chat_message,
+    error,
+    issue_line,
+    issues_block,
+    key_value_table,
+    result_json,
+    result_markdown,
+    section_rule,
+    stats_footer,
+    step_log,
+    success,
+    tool_use_badge,
+    warning,
+    welcome_banner,
+)
+from .printer import CliPrinter, EventRenderer
+
+__all__ = [
+    "AgentRunner",
+    "AgentTUI",
+    "BUILTIN_COMMANDS",
+    "CliPrinter",
+    "EventRenderer",
+    "FileSaveHandler",
+    "GuidelineProvider",
+    "ISSUE_STYLES",
+    "Mode",
+    "PersonaProvider",
+    "ROLE_STYLES",
+    "RunResult",
+    "SaveHandler",
+    "SlashCommand",
+    "TuiConfig",
+    "TuiControls",
+    "app_header",
+    "chat_input_prompt",
+    "chat_message",
+    "error",
+    "guideline_provider_from_library",
+    "issue_line",
+    "issues_block",
+    "key_value_table",
+    "persona_provider_from_library",
+    "result_json",
+    "result_markdown",
+    "run_agent_tui",
+    "section_rule",
+    "stats_footer",
+    "step_log",
+    "success",
+    "tool_use_badge",
+    "warning",
+    "welcome_banner",
+]

ellements/cli/adapters.py

@@ -0,0 +1,73 @@
+"""Adapters between ellements core types and TUI provider Protocols.
+
+The :mod:`ellements.cli.agent_tui` module defines ``PersonaProvider`` /
+``GuidelineProvider`` Protocols so the TUI never imports an
+``AgentController`` directly. This module supplies the canonical
+adapters that bridge a :class:`~ellements.core.PersonaLibrary` /
+:class:`~ellements.core.GuidelineLibrary` + controller pair to those
+Protocols.
+
+Apps that don't use ``AgentController`` can write their own adapters
+and pass them to :class:`~ellements.cli.agent_tui.TuiConfig` directly.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .agent_tui import GuidelineProvider, PersonaProvider
+
+
+class _LibraryPersonaProvider:
+    """Bridge a ``PersonaLibrary`` + controller to ``PersonaProvider``."""
+
+    def __init__(self, library: Any, controller: Any) -> None:
+        self._library = library
+        self._controller = controller
+
+    def list_ids(self) -> list[str]:
+        return list(self._library.list_ids())
+
+    def current_id(self) -> str | None:
+        return getattr(self._controller, "current_persona_id", None)
+
+    def select(self, persona_id: str) -> None:
+        self._controller.set_persona_by_id(persona_id)
+
+
+class _LibraryGuidelineProvider:
+    """Bridge a ``GuidelineLibrary`` + controller to ``GuidelineProvider``."""
+
+    def __init__(self, library: Any, controller: Any) -> None:
+        self._library = library
+        self._controller = controller
+
+    def list_ids(self) -> list[str]:
+        return list(self._library.list_ids())
+
+    def current_id(self) -> str | None:
+        return getattr(self._controller, "current_guideline_id", None)
+
+    def select(self, guideline_id: str) -> None:
+        self._controller.set_guideline_by_id(guideline_id)
+
+    def clear(self) -> None:
+        self._controller.clear_guideline()
+
+
+def persona_provider_from_library(library: Any, controller: Any) -> PersonaProvider:
+    """Adapt a :class:`PersonaLibrary` + controller pair to the TUI."""
+    return _LibraryPersonaProvider(library, controller)
+
+
+def guideline_provider_from_library(
+    library: Any, controller: Any
+) -> GuidelineProvider:
+    """Adapt a :class:`GuidelineLibrary` + controller pair to the TUI."""
+    return _LibraryGuidelineProvider(library, controller)
+
+
+__all__ = [
+    "guideline_provider_from_library",
+    "persona_provider_from_library",
+]