pretensor 0.1.0__py3-none-any.whl

pretensor/__init__.py

@@ -0,0 +1,50 @@
+"""Pretensor Graph — schema knowledge graph backed by Kuzu.
+
+Heavy imports (e.g. :class:`GraphBuilder`) are lazy so ``import pretensor``
+works in tests that only touch the registry.
+"""
+
+# pyright: reportUnsupportedDunderAll=false
+# Names in __all__ are provided via __getattr__; Pyright does not model that.
+
+from __future__ import annotations
+
+from typing import Any
+
+__all__ = [
+    "GraphBuilder",
+    "GraphEdge",
+    "GraphNode",
+    "GraphRegistry",
+    "KuzuStore",
+    "RelationshipCandidate",
+]
+
+
+def __getattr__(name: str) -> Any:
+    if name == "GraphBuilder":
+        from pretensor.core.builder import GraphBuilder
+
+        return GraphBuilder
+    if name == "GraphRegistry":
+        from pretensor.core.registry import GraphRegistry
+
+        return GraphRegistry
+    if name == "KuzuStore":
+        from pretensor.core.store import KuzuStore
+
+        return KuzuStore
+    if name == "GraphNode":
+        from pretensor.graph_models.node import GraphNode
+
+        return GraphNode
+    if name == "GraphEdge":
+        from pretensor.graph_models.edge import GraphEdge
+
+        return GraphEdge
+    if name == "RelationshipCandidate":
+        from pretensor.graph_models.relationship import RelationshipCandidate
+
+        return RelationshipCandidate
+    msg = f"module {__name__!r} has no attribute {name!r}"
+    raise AttributeError(msg)

pretensor/benchmark/__init__.py

@@ -0,0 +1,54 @@
+"""``pretensor benchmark`` subsystem — 3-level benchmark harness.
+
+The CLI contract and metric definitions are specified in
+``pretensor-ai/pretensor-specs`` at ``docs/specs/benchmark/spec.md``. This
+package contains the CLI dispatch / I/O wiring, the results aggregator
+(``BenchmarkResult``, ``compare``, JSON / CSV I/O), and the ``compare``
+regression-gate. L1 / L2 metric computation and the L3 agent-task runners
+land in follow-up changes.
+"""
+
+from __future__ import annotations
+
+from pretensor.benchmark.cli import register_benchmark_command
+from pretensor.benchmark.fixtures import Fixture, load_dataset
+from pretensor.benchmark.results import (
+    BenchmarkResult,
+    ComparisonError,
+    ComparisonReport,
+    Direction,
+    Metric,
+    MetricDiff,
+    compare,
+    read_json,
+    write_csv,
+    write_json,
+)
+from pretensor.benchmark.runner import (
+    Dataset,
+    RunnerKind,
+    run_l1,
+    run_l2,
+    run_l3,
+)
+
+__all__ = [
+    "BenchmarkResult",
+    "ComparisonError",
+    "ComparisonReport",
+    "Dataset",
+    "Direction",
+    "Fixture",
+    "Metric",
+    "MetricDiff",
+    "RunnerKind",
+    "compare",
+    "load_dataset",
+    "read_json",
+    "register_benchmark_command",
+    "run_l1",
+    "run_l2",
+    "run_l3",
+    "write_csv",
+    "write_json",
+]

pretensor/benchmark/cli.py

@@ -0,0 +1,294 @@
+"""``pretensor benchmark`` command group — L1 / L2 / L3 dispatch.
+
+Contract: see ``docs/specs/benchmark/spec.md`` §CLI contract. stdout carries
+only the JSON document when ``--out`` is absent; all human-facing text goes
+to stderr.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import typer
+from rich.console import Console
+
+if TYPE_CHECKING:
+    # Type-only import — runtime cost is paid lazily inside l3_command.
+    from pretensor.benchmark.l3.mcp_client import McpClientError
+
+from pretensor.benchmark.results import (
+    ComparisonError,
+    compare,
+    read_json,
+)
+from pretensor.benchmark.runner import (
+    Dataset,
+    RunnerKind,
+    run_l1,
+    run_l2,
+    run_l3,
+)
+
+__all__ = ["register_benchmark_command"]
+
+
+# Exit 1 covers both stub run failures and detected regressions — the spec
+# §CLI contract collapses them under a single "release-gating assertion
+# failed" code; downstream scripts inspect stderr to disambiguate.
+_EXIT_FAILED = 1
+_EXIT_BAD_INPUT = 2
+_DEFAULT_GRAPH_DIR = Path(".pretensor")
+_DATASET_HELP = (
+    "Fixture key (pagila, tpch, analytics_dwh, adversarial, "
+    "saas_multitenant, adventureworks)."
+)
+_OUT_HELP = "Write JSON output to this path; stdout when omitted."
+_GRAPH_DIR_HELP = "Existing indexed graph directory. Defaults to .pretensor."
+_BASELINE_HELP = "Path to the baseline benchmark result JSON."
+_CURRENT_HELP = (
+    "Path to the current benchmark result JSON to compare against the baseline."
+)
+
+
+def register_benchmark_command(app: typer.Typer) -> None:
+    """Register ``pretensor benchmark`` onto ``app``.
+
+    The benchmark commands route all human output to a dedicated stderr
+    console (the spec reserves stdout for the JSON document). Unlike most
+    ``register_*`` callers, this function therefore does not accept the
+    shared stdout console from ``cli.main``.
+    """
+    err_console = Console(stderr=True)
+
+    benchmark_app = typer.Typer(
+        name="benchmark",
+        help=(
+            "Run the 3-level benchmark harness (L1 graph quality, "
+            "L2 MCP tool quality, L3 agent task success)."
+        ),
+        no_args_is_help=True,
+    )
+
+    def _handle_not_implemented(exc: NotImplementedError) -> None:
+        err_console.print(f"[yellow]{exc}[/yellow]")
+        raise typer.Exit(_EXIT_FAILED) from exc
+
+    def _handle_input_error(
+        exc: LookupError | FileNotFoundError | McpClientError,
+    ) -> None:
+        """Friendly exit for the L3 runner's user-recoverable errors.
+
+        Three families land here:
+
+        * ``LookupError`` — missing per-dataset DB URL env var or the
+          LLM-provider key (the L3 runner converts the constructor's
+          ``LlmCallError`` to ``LookupError`` so this handler picks it
+          up too).
+        * ``FileNotFoundError`` — missing DDL bundle, missing gold
+          file, or the pretensor runner's missing-indexed-graph guard.
+        * ``McpClientError`` — the pretensor runner failed to spawn or
+          handshake with ``pretensor serve`` (broken binary, version
+          mismatch, transport fault). Distinct from missing-graph-dir
+          so an operator can tell apart "I forgot to index" from "the
+          serve binary is broken" in the output.
+
+        All three are user-recoverable misconfigurations — print the
+        message as-is (the exception text is already actionable) and
+        exit with the same release-gate code as a stub failure so
+        wrappers don't need to special-case L3 vs L1/L2.
+        """
+        err_console.print(f"[red]{exc}[/red]")
+        raise typer.Exit(_EXIT_FAILED) from exc
+
+    @benchmark_app.command("l1")
+    def l1_command(
+        dataset: Dataset = typer.Option(
+            ...,
+            "--dataset",
+            help=_DATASET_HELP,
+            case_sensitive=False,
+        ),
+        out: Path | None = typer.Option(
+            None,
+            "--out",
+            help=_OUT_HELP,
+            file_okay=True,
+            dir_okay=False,
+            writable=True,
+            resolve_path=True,
+        ),
+        graph_dir: Path = typer.Option(
+            _DEFAULT_GRAPH_DIR,
+            "--graph-dir",
+            help=_GRAPH_DIR_HELP,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+        embeddings: bool = typer.Option(
+            False,
+            "--embeddings/--no-embeddings",
+            help="Run the Layer-A embeddings path (requires the "
+            "'embeddings' optional extra).",
+        ),
+    ) -> None:
+        """Run L1 graph-quality metrics against a fixture."""
+        try:
+            run_l1(dataset, out, graph_dir, embeddings=embeddings)
+        except NotImplementedError as exc:
+            _handle_not_implemented(exc)
+
+    @benchmark_app.command("l2")
+    def l2_command(
+        dataset: Dataset = typer.Option(
+            ...,
+            "--dataset",
+            help=_DATASET_HELP,
+            case_sensitive=False,
+        ),
+        out: Path | None = typer.Option(
+            None,
+            "--out",
+            help=_OUT_HELP,
+            file_okay=True,
+            dir_okay=False,
+            writable=True,
+            resolve_path=True,
+        ),
+        graph_dir: Path = typer.Option(
+            _DEFAULT_GRAPH_DIR,
+            "--graph-dir",
+            help=_GRAPH_DIR_HELP,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+        embeddings: bool = typer.Option(
+            False,
+            "--embeddings/--no-embeddings",
+            help="Include the semantic_search Recall@K metric "
+            "(requires the 'embeddings' optional extra).",
+        ),
+    ) -> None:
+        """Run L2 MCP-tool-quality metrics against a fixture."""
+        try:
+            run_l2(dataset, out, graph_dir, embeddings=embeddings)
+        except NotImplementedError as exc:
+            _handle_not_implemented(exc)
+
+    @benchmark_app.command("l3")
+    def l3_command(
+        dataset: Dataset = typer.Option(
+            ...,
+            "--dataset",
+            help=_DATASET_HELP,
+            case_sensitive=False,
+        ),
+        out: Path | None = typer.Option(
+            None,
+            "--out",
+            help=_OUT_HELP,
+            file_okay=True,
+            dir_okay=False,
+            writable=True,
+            resolve_path=True,
+        ),
+        graph_dir: Path = typer.Option(
+            _DEFAULT_GRAPH_DIR,
+            "--graph-dir",
+            help=_GRAPH_DIR_HELP,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+        ),
+        runner: RunnerKind = typer.Option(
+            ...,
+            "--runner",
+            help="Which L3 runner to execute: baseline (agent + raw "
+            "schema) or pretensor (agent + MCP).",
+            case_sensitive=False,
+        ),
+        model: str = typer.Option(
+            ...,
+            "--model",
+            help="LLM model identifier (e.g. claude-haiku-4-5).",
+        ),
+        seed: int | None = typer.Option(
+            None,
+            "--seed",
+            help="Optional integer seed for reproducible agent runs.",
+        ),
+    ) -> None:
+        """Run L3 agent-task-success evaluation against a fixture."""
+        # Lazy import: ``McpClientError`` lives in the L3 MCP-client
+        # module, which transitively pulls in the MCP SDK. Importing it
+        # at CLI module load would defeat the lazy-import strategy used
+        # by the run_l3 dispatcher (and slow `pretensor benchmark
+        # --help` for everyone).
+        from pretensor.benchmark.l3.mcp_client import McpClientError
+
+        try:
+            run_l3(
+                dataset,
+                out,
+                graph_dir,
+                runner=runner,
+                model=model,
+                seed=seed,
+            )
+        except NotImplementedError as exc:
+            _handle_not_implemented(exc)
+        except (LookupError, FileNotFoundError, McpClientError) as exc:
+            _handle_input_error(exc)
+
+    @benchmark_app.command("compare")
+    def compare_command(
+        baseline: Path = typer.Option(
+            ...,
+            "--baseline",
+            help=_BASELINE_HELP,
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            readable=True,
+            resolve_path=True,
+        ),
+        current: Path = typer.Option(
+            ...,
+            "--current",
+            help=_CURRENT_HELP,
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            readable=True,
+            resolve_path=True,
+        ),
+        tolerance: float = typer.Option(
+            1e-6,
+            "--tolerance",
+            help="Numeric tolerance below which metric drift is treated as noise.",
+        ),
+    ) -> None:
+        """Compare two benchmark result JSONs and exit non-zero on regression."""
+        try:
+            baseline_result = read_json(baseline)
+        except (json.JSONDecodeError, KeyError) as exc:
+            err_console.print(f"[red]Failed to read baseline: {exc}[/red]")
+            raise typer.Exit(_EXIT_BAD_INPUT) from exc
+        try:
+            current_result = read_json(current)
+        except (json.JSONDecodeError, KeyError) as exc:
+            err_console.print(f"[red]Failed to read current: {exc}[/red]")
+            raise typer.Exit(_EXIT_BAD_INPUT) from exc
+        try:
+            report = compare(baseline_result, current_result, tolerance=tolerance)
+        except ComparisonError as exc:
+            err_console.print(f"[red]{exc}[/red]")
+            raise typer.Exit(_EXIT_BAD_INPUT) from exc
+        err_console.print(report.format_diff())
+        if report.has_regression:
+            raise typer.Exit(_EXIT_FAILED)
+
+    app.add_typer(benchmark_app)

pretensor/benchmark/fixtures.py

@@ -0,0 +1,84 @@
+"""Benchmark-fixture loader.
+
+Every L1 / L2 / L3 runner resolves datasets through :func:`load_dataset`, so
+the set of supported keys and the physical file layout lives in one place.
+Fixtures are resolved relative to the repo root (this module's grandparent of
+``src``) — callers never need to know where the files live.
+
+The returned :class:`Fixture` names the dataset, points to its schema-snapshot
+YAML (always present), and optionally to a DDL-only ``schema.sql`` and a gold
+NL-to-SQL question set. Datasets without DDL or a question set (``adversarial``,
+``analytics_dwh``, ``saas_multitenant``, ``messy_warehouse``) return ``None`` for
+those fields.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from pretensor.benchmark.runner import Dataset
+
+__all__ = ["Fixture", "load_dataset"]
+
+# Walk up from src/pretensor/benchmark/fixtures.py to the repo root.
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_SCHEMAS_DIR = _REPO_ROOT / "tests" / "fixtures" / "schemas"
+_DATA_DIR = _REPO_ROOT / "scripts" / "data"
+
+
+@dataclass(frozen=True, slots=True)
+class Fixture:
+    """One benchmark dataset's on-disk layout.
+
+    ``name`` is the canonical :class:`Dataset` enum key. ``schema_yaml_path``
+    is always set and points to a file that :meth:`SchemaSnapshot.from_yaml`
+    can parse. The remaining fields are ``None`` for datasets that have no
+    DDL dump, no gold question set, or no metric-templates file checked in.
+    """
+
+    name: Dataset
+    schema_yaml_path: Path
+    ddl_sql_path: Path | None
+    questions_path: Path | None
+    metric_templates_path: Path | None
+
+
+def load_dataset(name: str | Dataset) -> Fixture:
+    """Resolve a benchmark dataset's fixture paths.
+
+    Accepts a :class:`Dataset` enum value or its string form. Unknown keys
+    raise :class:`ValueError`. The schema YAML must exist — if missing,
+    :class:`FileNotFoundError` is raised. DDL and question-set paths are
+    returned as ``None`` when the dataset has no such file on disk.
+    """
+    if isinstance(name, Dataset):
+        key = name
+    else:
+        try:
+            key = Dataset(name)
+        except ValueError as exc:
+            raise ValueError(
+                f"Unknown benchmark dataset: {name!r}. "
+                f"Known: {[d.value for d in Dataset]}"
+            ) from exc
+
+    schema_path = _SCHEMAS_DIR / f"{key.value}.yaml"
+    if not schema_path.exists():
+        raise FileNotFoundError(
+            f"Schema YAML missing for dataset {key.value!r}: {schema_path}"
+        )
+
+    ddl_path = _DATA_DIR / key.value / "schema.sql"
+    questions_path = _DATA_DIR / f"{key.value}_nl2sql_bench.json"
+    metric_templates_path = _DATA_DIR / f"{key.value}_metric_templates.yaml"
+
+    return Fixture(
+        name=key,
+        schema_yaml_path=schema_path,
+        ddl_sql_path=ddl_path if ddl_path.exists() else None,
+        questions_path=questions_path if questions_path.exists() else None,
+        metric_templates_path=(
+            metric_templates_path if metric_templates_path.exists() else None
+        ),
+    )

pretensor/benchmark/l1/__init__.py

@@ -0,0 +1,23 @@
+"""L1 graph-quality benchmark — metric implementations and runner.
+
+See ``docs/specs/benchmark/spec.md`` §L1 metric definitions for the
+contract this module fulfils. Metric values are wrapped in the
+``{value, direction}`` envelope from
+:mod:`pretensor.benchmark.results`.
+"""
+
+from __future__ import annotations
+
+from pretensor.benchmark.l1.metrics import (
+    cluster_stability_jaccard,
+    inferred_join_pr,
+    role_f1,
+)
+from pretensor.benchmark.l1.runner import run_l1
+
+__all__ = [
+    "cluster_stability_jaccard",
+    "inferred_join_pr",
+    "role_f1",
+    "run_l1",
+]

pretensor/benchmark/l1/metrics.py

@@ -0,0 +1,141 @@
+"""Pure L1 metric functions.
+
+Each function takes only data — no graph store, no I/O. The runner
+collects predictions from the intelligence pipeline, then hands them
+to these helpers. Boundary cases (empty ground truth, identical
+inputs, etc.) are handled here so the runner never divides by zero.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+
+__all__ = [
+    "canonicalise_join_key",
+    "cluster_stability_jaccard",
+    "inferred_join_pr",
+    "role_f1",
+]
+
+
+JoinKey = tuple[str, str, str, str]
+"""``(src_table, src_column, dst_table, dst_column)`` — bare names.
+
+Schema is folded into the table name (``schema.table``) so callers can
+key on either form consistently.
+"""
+
+
+def inferred_join_pr(
+    inferred: Iterable[JoinKey],
+    declared_fks: Iterable[JoinKey],
+) -> tuple[float, float]:
+    """Precision and recall of inferred joins against declared FK ground truth.
+
+    Both inputs are sets of ``(src_table, src_column, dst_table, dst_column)``
+    tuples. Keys are canonicalised to unordered pairs before comparison —
+    declared FKs are directional, but the heuristic discovery layer may
+    emit either orientation for the same physical relationship, and an
+    inferred edge ``A.x ↔ B.y`` should match the declared FK ``A.x → B.y``
+    regardless of orientation.
+
+    Returns:
+        ``(precision, recall)``. Both are ``0.0`` when the corresponding
+        denominator is zero (no candidates ⇒ precision 0; no ground truth
+        ⇒ recall 0). Identical sets ⇒ ``(1.0, 1.0)``.
+    """
+    inferred_set = {canonicalise_join_key(k) for k in inferred}
+    declared_set = {canonicalise_join_key(k) for k in declared_fks}
+
+    true_positives = inferred_set & declared_set
+    precision = len(true_positives) / len(inferred_set) if inferred_set else 0.0
+    recall = len(true_positives) / len(declared_set) if declared_set else 0.0
+    return precision, recall
+
+
+def cluster_stability_jaccard(
+    run_a: Iterable[Iterable[str]],
+    run_b: Iterable[Iterable[str]],
+) -> float:
+    """Mean pairwise Jaccard between two cluster partitions of the same input.
+
+    Each cluster is a collection of table node IDs. A score of ``1.0``
+    means the two partitions are identical (modulo cluster ordering); a
+    score of ``0.0`` means no cluster from ``run_a`` matches any cluster
+    in ``run_b``.
+
+    For each cluster in ``run_a``, this picks the best-matching cluster
+    in ``run_b`` (highest Jaccard) and averages those scores. Empty
+    inputs on both sides ⇒ ``1.0`` (vacuously stable). Empty input on
+    one side only ⇒ ``0.0``.
+    """
+    a_sets = [frozenset(c) for c in run_a if c]
+    b_sets = [frozenset(c) for c in run_b if c]
+
+    if not a_sets and not b_sets:
+        return 1.0
+    if not a_sets or not b_sets:
+        return 0.0
+
+    total = 0.0
+    for cluster_a in a_sets:
+        best = 0.0
+        for cluster_b in b_sets:
+            # Both sides are filtered to non-empty frozensets above, so
+            # the union is always non-empty — no zero-divide guard needed.
+            jacc = len(cluster_a & cluster_b) / len(cluster_a | cluster_b)
+            if jacc > best:
+                best = jacc
+        total += best
+    return total / len(a_sets)
+
+
+def role_f1(
+    predicted: Mapping[str, str],
+    gold: Mapping[str, str],
+) -> float:
+    """Macro-averaged F1 across the gold role labels.
+
+    For each role appearing in ``gold``, computes precision and recall
+    against ``predicted`` and averages the per-class F1 scores. Tables
+    in ``gold`` that the predictor did not classify count as misses;
+    tables in ``predicted`` whose gold role is absent are ignored
+    (they're outside the evaluation set).
+
+    Returns ``1.0`` when ``gold`` and ``predicted`` are both empty
+    (vacuously correct). Returns ``0.0`` when ``gold`` is non-empty but
+    ``predicted`` is empty, and vice versa.
+    """
+    if not gold and not predicted:
+        return 1.0
+    if not gold or not predicted:
+        return 0.0
+
+    gold_roles = sorted(set(gold.values()))
+    f1_sum = 0.0
+    for role in gold_roles:
+        gold_for_role = {t for t, r in gold.items() if r == role}
+        pred_for_role = {t for t, r in predicted.items() if r == role}
+        true_positives = gold_for_role & pred_for_role
+        precision = len(true_positives) / len(pred_for_role) if pred_for_role else 0.0
+        recall = len(true_positives) / len(gold_for_role) if gold_for_role else 0.0
+        if precision + recall == 0.0:
+            f1 = 0.0
+        else:
+            f1 = 2 * precision * recall / (precision + recall)
+        f1_sum += f1
+    return f1_sum / len(gold_roles)
+
+
+def canonicalise_join_key(key: JoinKey) -> tuple[tuple[str, str], tuple[str, str]]:
+    """Collapse a directional ``JoinKey`` to an unordered pair.
+
+    Treats ``(A.x → B.y)`` and ``(B.y → A.x)`` as the same edge by
+    sorting the two endpoint ``(table, column)`` tuples. Exposed so
+    other modules (e.g. the runner's per-item collector) can canonicalise
+    keys against the same convention without duplicating the logic.
+    """
+    src_t, src_c, dst_t, dst_c = key
+    a = (src_t, src_c)
+    b = (dst_t, dst_c)
+    return (a, b) if a <= b else (b, a)