pretensor 0.1.0__py3-none-any.whl

pretensor/benchmark/results.py

@@ -0,0 +1,298 @@
+"""Benchmark results aggregator — schema, I/O, and regression-gate comparator.
+
+The JSON schema this module reads/writes is the canonical contract that
+every benchmark level (L1, L2, L3) emits and that the release-gate script
+consumes. See ``docs/specs/benchmark/spec.md`` §JSON output schema.
+"""
+
+from __future__ import annotations
+
+import csv
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Literal
+
+from pretensor.errors import PretensorError
+
+__all__ = [
+    "BenchmarkResult",
+    "ComparisonError",
+    "ComparisonReport",
+    "Direction",
+    "Metric",
+    "MetricDiff",
+    "compare",
+    "read_json",
+    "write_csv",
+    "write_json",
+]
+
+Direction = Literal["higher_is_better", "lower_is_better"]
+
+
+def _fmt(value: float | None, *, signed: bool = False) -> str:
+    if value is None:
+        return "null"
+    return f"{value:+.6g}" if signed else f"{value:.6g}"
+
+
+@dataclass(frozen=True, slots=True)
+class Metric:
+    """One metric value plus the direction that defines a regression.
+
+    ``value`` is ``None`` when the metric is unavailable on this run
+    (e.g. ``semantic_search`` Recall@K on a build without ``[embeddings]``).
+    ``direction`` encodes regression semantics — a drop on a
+    ``higher_is_better`` metric is a regression; a rise on a
+    ``lower_is_better`` metric is a regression.
+    """
+
+    value: float | None
+    direction: Direction
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"value": self.value, "direction": self.direction}
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Metric:
+        return cls(value=data["value"], direction=data["direction"])
+
+
+@dataclass(frozen=True, slots=True)
+class BenchmarkResult:
+    """One benchmark run's full output — matches the spec's JSON schema.
+
+    ``extra`` is a free-form per-level envelope for fields that don't
+    belong in ``metrics`` or ``per_item`` but must be captured for
+    reproducibility — e.g. L3 records the LLM model, temperature, seed,
+    and prompt hash here. L1 / L2 leave it empty. The spec's JSON-schema
+    section ("downstream consumers must tolerate additional keys")
+    sanctions adding new top-level keys this way.
+    """
+
+    level: str
+    dataset: str
+    pretensor_version: str
+    embeddings_enabled: bool
+    ran_at: str
+    fixture_sha: str
+    metrics: dict[str, Metric]
+    per_item: list[dict[str, Any]]
+    notes: list[str] = field(default_factory=list)
+    extra: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        out: dict[str, Any] = {
+            "level": self.level,
+            "dataset": self.dataset,
+            "pretensor_version": self.pretensor_version,
+            "embeddings_enabled": self.embeddings_enabled,
+            "ran_at": self.ran_at,
+            "fixture_sha": self.fixture_sha,
+            "metrics": {name: m.to_dict() for name, m in self.metrics.items()},
+            "per_item": list(self.per_item),
+            "notes": list(self.notes),
+        }
+        if self.extra:
+            out["extra"] = dict(self.extra)
+        return out
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> BenchmarkResult:
+        return cls(
+            level=data["level"],
+            dataset=data["dataset"],
+            pretensor_version=data["pretensor_version"],
+            embeddings_enabled=data["embeddings_enabled"],
+            ran_at=data["ran_at"],
+            fixture_sha=data["fixture_sha"],
+            metrics={name: Metric.from_dict(m) for name, m in data["metrics"].items()},
+            per_item=list(data["per_item"]),
+            notes=list(data.get("notes", [])),
+            extra=dict(data.get("extra", {})),
+        )
+
+
+def write_json(result: BenchmarkResult, path: Path) -> None:
+    """Serialize ``result`` to ``path`` as deterministic, pretty-printed JSON.
+
+    Output is byte-stable for identical inputs: keys are sorted, indent is 2
+    spaces, and the file ends with a single trailing newline.
+    """
+    text = json.dumps(result.to_dict(), sort_keys=True, indent=2) + "\n"
+    path.write_text(text, encoding="utf-8")
+
+
+def read_json(path: Path) -> BenchmarkResult:
+    """Parse a JSON file into a ``BenchmarkResult``."""
+    data = json.loads(path.read_text(encoding="utf-8"))
+    return BenchmarkResult.from_dict(data)
+
+
+class ComparisonError(PretensorError, ValueError):
+    """Raised when two ``BenchmarkResult`` inputs cannot be compared."""
+
+
+@dataclass(frozen=True, slots=True)
+class MetricDiff:
+    """One metric's before/after delta with its regression direction."""
+
+    name: str
+    baseline: float | None
+    current: float | None
+    direction: Direction
+    delta: float | None
+
+
+@dataclass(frozen=True, slots=True)
+class ComparisonReport:
+    """Result of comparing two ``BenchmarkResult`` instances."""
+
+    has_regression: bool = False
+    regressions: list[MetricDiff] = field(default_factory=list)
+    improvements: list[MetricDiff] = field(default_factory=list)
+    added: list[str] = field(default_factory=list)
+    removed: list[str] = field(default_factory=list)
+
+    def format_diff(self) -> str:
+        """Render a human-readable diff block suitable for stderr output."""
+        lines: list[str] = []
+        if self.has_regression:
+            lines.append("REGRESSION DETECTED")
+        else:
+            lines.append("no regression")
+        if self.regressions:
+            lines.append("")
+            lines.append("Regressions:")
+            for d in self.regressions:
+                arrow = "↓" if d.direction == "higher_is_better" else "↑"
+                lines.append(
+                    f"  {d.name} ({d.direction}): "
+                    f"{_fmt(d.baseline)} → {_fmt(d.current)} "
+                    f"(Δ{_fmt(d.delta, signed=True)}) {arrow}"
+                )
+        if self.removed:
+            lines.append("")
+            lines.append("Metrics removed (present in baseline, missing in current):")
+            for name in self.removed:
+                lines.append(f"  {name}")
+        if self.added:
+            lines.append("")
+            lines.append("New metrics in current (not in baseline):")
+            for name in self.added:
+                lines.append(f"  {name}")
+        if self.improvements:
+            lines.append("")
+            lines.append("Improvements:")
+            for d in self.improvements:
+                lines.append(
+                    f"  {d.name}: {_fmt(d.baseline)} → {_fmt(d.current)} "
+                    f"(Δ{_fmt(d.delta, signed=True)})"
+                )
+        return "\n".join(lines) + "\n"
+
+
+def compare(
+    baseline: BenchmarkResult,
+    current: BenchmarkResult,
+    *,
+    tolerance: float = 1e-6,
+) -> ComparisonReport:
+    """Compare two benchmark results and classify metric changes.
+
+    Raises ``ComparisonError`` if the two inputs disagree on ``dataset`` or
+    ``level`` — comparing across datasets or levels is meaningless. A metric
+    present in ``baseline`` but missing in ``current`` is a regression
+    (``removed``). A metric new in ``current`` is recorded under ``added``
+    but never counts as a regression. Metrics whose value is ``None`` on
+    either side are skipped (no diff emitted).
+    """
+    if baseline.dataset != current.dataset:
+        raise ComparisonError(
+            f"cannot compare across datasets: baseline={baseline.dataset!r}, "
+            f"current={current.dataset!r}"
+        )
+    if baseline.level != current.level:
+        raise ComparisonError(
+            f"cannot compare across levels: baseline={baseline.level!r}, "
+            f"current={current.level!r}"
+        )
+
+    regressions: list[MetricDiff] = []
+    improvements: list[MetricDiff] = []
+    removed: list[str] = []
+    added: list[str] = sorted(set(current.metrics) - set(baseline.metrics))
+
+    for name in sorted(baseline.metrics):
+        base_metric = baseline.metrics[name]
+        cur_metric = current.metrics.get(name)
+        if cur_metric is None:
+            removed.append(name)
+            continue
+        if base_metric.value is None or cur_metric.value is None:
+            continue
+        delta = cur_metric.value - base_metric.value
+        diff = MetricDiff(
+            name=name,
+            baseline=base_metric.value,
+            current=cur_metric.value,
+            direction=base_metric.direction,
+            delta=delta,
+        )
+        if base_metric.direction == "higher_is_better":
+            if delta < -tolerance:
+                regressions.append(diff)
+            elif delta > tolerance:
+                improvements.append(diff)
+        else:  # lower_is_better
+            if delta > tolerance:
+                regressions.append(diff)
+            elif delta < -tolerance:
+                improvements.append(diff)
+
+    return ComparisonReport(
+        has_regression=bool(regressions or removed),
+        regressions=regressions,
+        improvements=improvements,
+        added=added,
+        removed=removed,
+    )
+
+
+_CSV_META_COLUMNS = (
+    "level",
+    "dataset",
+    "pretensor_version",
+    "embeddings_enabled",
+    "ran_at",
+)
+
+
+def write_csv(results: list[BenchmarkResult], path: Path) -> None:
+    """Write one row per ``BenchmarkResult`` as a flat CSV for spreadsheet review.
+
+    Metric columns are the sorted union of metric names across ``results``.
+    A row missing a metric (or whose value is ``None``) writes an empty cell.
+    Metric ``direction`` is JSON-only metadata and is omitted from CSV;
+    reviewers consult the JSON for direction.
+    """
+    metric_names = sorted({name for r in results for name in r.metrics})
+    columns = list(_CSV_META_COLUMNS) + metric_names
+    with path.open("w", encoding="utf-8", newline="") as fh:
+        writer = csv.DictWriter(fh, fieldnames=columns)
+        writer.writeheader()
+        for r in results:
+            row: dict[str, Any] = {
+                "level": r.level,
+                "dataset": r.dataset,
+                "pretensor_version": r.pretensor_version,
+                "embeddings_enabled": r.embeddings_enabled,
+                "ran_at": r.ran_at,
+            }
+            for name in metric_names:
+                metric = r.metrics.get(name)
+                row[name] = (
+                    "" if metric is None or metric.value is None else metric.value
+                )
+            writer.writerow(row)

pretensor/benchmark/runner.py

@@ -0,0 +1,109 @@
+"""Dispatch stubs for the three benchmark levels.
+
+These functions are intentionally minimal: the initial scaffolding lands the
+CLI contract and a stable Python entry point, while the metric
+implementations ship in follow-up changes.
+"""
+
+from __future__ import annotations
+
+import enum
+from pathlib import Path
+
+
+class Dataset(str, enum.Enum):
+    """Closed set of benchmark fixture keys.
+
+    Matches ``docs/specs/benchmark/spec.md`` §CLI contract. Unknown values
+    are rejected by Typer with exit code 2 before reaching ``run_*``.
+    """
+
+    PAGILA = "pagila"
+    TPCH = "tpch"
+    ANALYTICS_DWH = "analytics_dwh"
+    ADVERSARIAL = "adversarial"
+    SAAS_MULTITENANT = "saas_multitenant"
+    ADVENTUREWORKS = "adventureworks"
+    MESSY_WAREHOUSE = "messy_warehouse"
+
+
+class RunnerKind(str, enum.Enum):
+    """L3 runner selector — baseline (raw schema) vs pretensor (MCP)."""
+
+    BASELINE = "baseline"
+    PRETENSOR = "pretensor"
+
+
+def run_l1(
+    dataset: Dataset,
+    out: Path | None,
+    graph_dir: Path,
+    *,
+    embeddings: bool,
+) -> None:
+    """Run L1 (graph-quality) metrics for ``dataset``.
+
+    Computes inferred-join precision / recall against declared FKs,
+    cluster-stability Jaccard across two reindex runs, and role-
+    classification F1 against ``<dataset>_roles.yaml`` when present.
+    ``embeddings=True`` enables the Layer-A embeddings path when the
+    optional extra is installed; missing extra logs an install hint
+    and falls back to the heuristic-only run.
+    """
+    # Imported lazily so the heavy intelligence-pipeline dependency
+    # graph (Kuzu, igraph, etc.) does not load when other run_*
+    # callers import this module.
+    from pretensor.benchmark.l1 import run_l1 as _run_l1
+
+    _run_l1(dataset, out, graph_dir, embeddings=embeddings)
+
+
+def run_l2(
+    dataset: Dataset,
+    out: Path | None,
+    graph_dir: Path,
+    *,
+    embeddings: bool,
+) -> None:
+    """Run L2 (MCP-tool-quality) metrics for ``dataset``.
+
+    Computes ``query`` Recall@K, ``traverse`` path correctness, and
+    ``compile_metric`` SQL correctness against the fixture's gold data.
+    When ``embeddings=True`` and the ``semantic_search`` MCP tool is
+    available, also emits ``semantic_search`` Recall@K (Layer-B gate).
+    """
+    # Imported lazily so the heavy MCP / Kuzu / sqlglot dependency
+    # graph does not load when other run_* callers import this module.
+    from pretensor.benchmark.l2 import run_l2 as _run_l2
+
+    _run_l2(dataset, out, graph_dir, embeddings=embeddings)
+
+
+def run_l3(
+    dataset: Dataset,
+    out: Path | None,
+    graph_dir: Path,
+    *,
+    runner: RunnerKind,
+    model: str,
+    seed: int | None,
+) -> None:
+    """Run L3 (agent-task-success) for ``dataset`` via the chosen runner.
+
+    ``runner=RunnerKind.BASELINE`` runs the agent against the raw schema.
+    ``runner=RunnerKind.PRETENSOR`` runs the agent through the MCP server.
+    The baseline ignores ``graph_dir``; the pretensor runner spawns
+    ``pretensor serve --graph-dir <graph_dir>``.
+    """
+    if runner is RunnerKind.BASELINE:
+        # Imported lazily so the LLM-client / httpx dependency graph does
+        # not load when other run_* callers import this module.
+        from pretensor.benchmark.l3.runner import run_l3_baseline
+
+        run_l3_baseline(dataset, out, model=model, seed=seed)
+        return
+    # Lazy import: the pretensor runner pulls in the MCP stdio client and
+    # asyncio plumbing; other run_* callers should not pay that cost.
+    from pretensor.benchmark.l3.pretensor_runner import run_l3_pretensor
+
+    run_l3_pretensor(dataset, out, graph_dir, model=model, seed=seed)

pretensor/cli/__init__.py

@@ -0,0 +1 @@
+"""Command-line interface for Pretensor Graph."""

pretensor/cli/commands/_source_runner.py

@@ -0,0 +1,147 @@
+"""Shared CLI scaffolding for commands that operate over named sources."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+import typer
+from rich.console import Console
+
+from pretensor.introspection.models.dsn import validate_source_env_vars
+
+if TYPE_CHECKING:
+    from pretensor.cli.config_file import PretensorCliConfig
+
+
+def resolve_embeddings_mode(
+    embeddings: bool | None,
+    console: Console,
+    *,
+    auto_enable_message: str | None = None,
+) -> bool:
+    """Verify or auto-detect the embeddings mode.
+
+    Parameters
+    ----------
+    embeddings:
+        Raw CLI value: ``True`` = explicit opt-in, ``False`` = explicit
+        opt-out, ``None`` = auto-detect.
+    console:
+        Rich console for user-facing output.
+    auto_enable_message:
+        If given, printed when auto-detection enables embeddings.
+        Pass ``None`` (default) to suppress the message.
+
+    Returns
+    -------
+    bool
+        Resolved embeddings flag.
+    """
+    if embeddings:
+        from pretensor.intelligence.embeddings import verify_embeddings_extra_installed
+
+        try:
+            verify_embeddings_extra_installed()
+        except ImportError as e:
+            console.print(str(e), style="red", markup=False)
+            raise typer.Exit(1) from e
+        return True
+    if embeddings is None:
+        from pretensor.intelligence.embeddings import resolve_embeddings_auto
+
+        resolved = resolve_embeddings_auto(None)
+        if resolved and auto_enable_message:
+            console.print(auto_enable_message)
+        return resolved
+    return False
+
+
+def check_source_exists(
+    source: str,
+    cli_config: PretensorCliConfig,
+    console: Console,
+) -> None:
+    """Exit 1 with an informative message if *source* is not in ``cli_config.sources``."""
+    if source not in cli_config.sources:
+        available = ", ".join(cli_config.sources) or "(none)"
+        console.print(f"[red]Unknown source {source!r}.[/red] Available: {available}")
+        raise typer.Exit(1)
+
+
+def print_multi_source_summary(
+    console: Console,
+    results: list[tuple[str, bool, str]],
+) -> None:
+    """Print the tick/cross summary table and exit 1 if any source failed."""
+    console.print(f"\n[bold]{'─' * 40}[/bold]")
+    console.print("[bold]Summary:[/bold]")
+    failed = 0
+    for src_name, ok, msg in results:
+        status = "[green]✓[/green]" if ok else "[red]✗[/red]"
+        console.print(f"  {status} {src_name}: {msg}")
+        if not ok:
+            failed += 1
+    if failed:
+        raise typer.Exit(1)
+
+
+def run_for_all_sources(
+    *,
+    cli_config: PretensorCliConfig,
+    console: Console,
+    logger: logging.Logger,
+    source_banner: str,
+    action_verb: str,
+    run_one: Callable[[str], None],
+) -> None:
+    """Iterate all configured sources, skip ones with missing env vars, exit 1 on any failure.
+
+    Parameters
+    ----------
+    cli_config:
+        Loaded CLI config (must have a non-empty ``sources`` mapping).
+    console:
+        Rich console for user-facing output.
+    logger:
+        Module-level logger for structured warnings/errors.
+    source_banner:
+        Label shown in the per-source header, e.g. ``"Source:"`` or
+        ``"Reindexing source:"``.
+    action_verb:
+        Verb used in error messages, e.g. ``"indexing"`` or ``"reindexing"``.
+    run_one:
+        Callable ``(src_name) -> None`` that performs the per-source
+        action; should raise ``typer.Exit`` on failure.
+    """
+    if not cli_config.sources:
+        console.print(
+            "[red]No sources defined in config.[/red] "
+            "Add a `sources:` section to .pretensor/config.yaml."
+        )
+        raise typer.Exit(1)
+    results: list[tuple[str, bool, str]] = []
+    for src_name, src_cfg in cli_config.sources.items():
+        console.print(
+            f"\n[bold]{'─' * 40}[/bold]\n"
+            f"[bold blue]{source_banner}[/bold blue] {src_name}\n"
+        )
+        missing_vars = validate_source_env_vars(src_cfg)
+        if missing_vars:
+            msg = f"missing env vars: {', '.join(missing_vars)}"
+            console.print(f"[yellow]Skipping {src_name}:[/yellow] {msg}")
+            logger.warning("Skipping source %s: %s", src_name, msg)
+            results.append((src_name, False, f"skipped ({msg})"))
+            continue
+        try:
+            run_one(src_name)
+            results.append((src_name, True, "ok"))
+        except typer.Exit:
+            logger.warning("Source %s exited with failure", src_name)
+            results.append((src_name, False, "failed"))
+        except Exception as e:
+            logger.exception("Error %s source %s", action_verb, src_name)
+            console.print(f"[red]Error {action_verb} {src_name}:[/red] {e}")
+            results.append((src_name, False, str(e)))
+    print_multi_source_summary(console, results)