pretensor 0.1.0__py3-none-any.whl

pretensor/benchmark/l2/metrics.py

@@ -0,0 +1,146 @@
+"""Pure L2 metric functions.
+
+Each function takes already-collected observations and returns a scalar.
+The runner owns all I/O — these helpers stay testable without spinning
+up Kuzu, the keyword index, or any MCP machinery.
+
+Boundary cases (empty input, missing gold per item) are handled here
+so the runner never divides by zero.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+
+__all__ = [
+    "JoinPair",
+    "RankedHit",
+    "compile_metric_correctness",
+    "query_recall_at_k",
+    "semantic_search_recall_at_k",
+    "top_k_with_ties",
+    "traverse_correctness",
+]
+
+
+JoinPair = tuple[str, str]
+"""``(from_table, to_table)`` — schema-qualified bare names."""
+
+RankedHit = tuple[float, str]
+"""``(score, table_name)`` — one search hit; higher score = better."""
+
+
+_SCORE_ROUNDING_DECIMALS = 4
+"""Bucket precision applied to scores before sorting/cutoff.
+
+Upstream BM25 (SQLite FTS5) has been observed to emit micro-noise — e.g.
+``-1e-06`` versus ``-0.0`` — for hits that are functionally tied in
+ranking. Rounding to 4 decimals collapses that noise into clean zero
+buckets while preserving meaningful score differences (anything ≥ 1e-4).
+"""
+
+
+def top_k_with_ties(
+    ranked: Sequence[RankedHit],
+    k: int,
+) -> list[RankedHit]:
+    """Return the top-K hits plus any extra items tied at the K-th score.
+
+    Sorts defensively by ``(-score, name)`` before slicing so the caller
+    needn't pre-sort. AC #5: a tied top-K includes ALL tied items so a
+    near-cutoff hit is not missed when comparing against gold.
+
+    Scores are rounded to :data:`_SCORE_ROUNDING_DECIMALS` decimals before
+    comparison so that floating-point noise (e.g. SQLite FTS5's ``-1e-06``
+    vs ``-0.0`` for functionally-tied BM25 hits) doesn't make the cutoff
+    flip between runs. The returned tuples carry the rounded score; the
+    metric only consumes the table-name field, so this loss of precision
+    is invisible downstream.
+    """
+    if k <= 0 or not ranked:
+        return []
+    rounded: list[RankedHit] = [
+        (round(score, _SCORE_ROUNDING_DECIMALS), name) for score, name in ranked
+    ]
+    ordered = sorted(rounded, key=lambda h: (-h[0], h[1]))
+    if len(ordered) <= k:
+        return ordered
+    cutoff_score = ordered[k - 1][0]
+    return [h for h in ordered if h[0] >= cutoff_score]
+
+
+def query_recall_at_k(
+    observations: Iterable[tuple[Iterable[str], Sequence[RankedHit]]],
+    k: int = 5,
+) -> float:
+    """Mean Recall@K across observations.
+
+    Each observation is ``(gold_tables, ranked_results)``. ``ranked_results``
+    may arrive in any order — :func:`top_k_with_ties` re-sorts it. Recall
+    per observation is ``|gold ∩ retrieved_top_k| / |gold|``. Observations
+    with empty ``gold_tables`` are skipped (recall undefined). Returns
+    ``0.0`` when no observation has gold data.
+    """
+    total = 0.0
+    n = 0
+    for gold, ranked in observations:
+        gold_set = set(gold)
+        if not gold_set:
+            continue
+        top_k = top_k_with_ties(list(ranked), k)
+        retrieved_set = {name for _, name in top_k}
+        recall = len(gold_set & retrieved_set) / len(gold_set)
+        total += recall
+        n += 1
+    return total / n if n > 0 else 0.0
+
+
+def semantic_search_recall_at_k(
+    observations: Iterable[tuple[Iterable[str], Sequence[RankedHit]]],
+    k: int = 5,
+) -> float:
+    """Mean Recall@K against ``semantic_search`` retrievals.
+
+    Same shape and semantics as :func:`query_recall_at_k` — kept distinct
+    so the two retrieval systems can diverge later (different ranking
+    semantics, different post-processing) without forcing one to fit the
+    other's API.
+    """
+    return query_recall_at_k(observations, k=k)
+
+
+def traverse_correctness(
+    observations: Iterable[tuple[Sequence[JoinPair], Sequence[Sequence[JoinPair]]]],
+) -> float:
+    """Fraction of items whose gold path matches at least one returned path.
+
+    The ``traverse`` tool emits all top-ranked paths when ties exist;
+    this metric counts the item as correct when ANY returned path
+    equals the gold sequence (tuple-wise comparison). Returns ``0.0``
+    when ``observations`` is empty.
+    """
+    obs_list = list(observations)
+    if not obs_list:
+        return 0.0
+
+    correct = 0
+    for gold_path, returned_paths in obs_list:
+        gold_t = tuple(tuple(p) for p in gold_path)
+        for path in returned_paths:
+            if tuple(tuple(p) for p in path) == gold_t:
+                correct += 1
+                break
+    return correct / len(obs_list)
+
+
+def compile_metric_correctness(observations: Iterable[bool]) -> float:
+    """Fraction of metric templates that compiled to valid SQL.
+
+    Each observation is the boolean ``valid`` flag returned by
+    ``compile_metric_payload``. Returns ``0.0`` when ``observations`` is
+    empty.
+    """
+    obs_list = list(observations)
+    if not obs_list:
+        return 0.0
+    return sum(1 for ok in obs_list if ok) / len(obs_list)

pretensor/benchmark/l2/pipeline.py

@@ -0,0 +1,124 @@
+"""Per-call graph construction for the L2 benchmark.
+
+Each invocation of :func:`run_l2` indexes the fixture's schema YAML into
+a fresh, throwaway Kuzu store + ``registry.json`` under a temp dir,
+then hands that directory to the MCP tool payload functions
+(``query_payload``, ``traverse_payload``, ``compile_metric_payload``).
+This guarantees determinism — no leftover state from previous runs.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from pathlib import Path
+
+from pretensor.config import EmbeddingsConfig, PretensorConfig
+from pretensor.connectors.models import SchemaSnapshot
+from pretensor.core.builder import GraphBuilder
+from pretensor.core.registry import GraphRegistry
+from pretensor.core.store import KuzuStore
+
+__all__ = ["build_l2_graph_dir"]
+
+
+# Pinned so the persisted ``last_indexed_at`` doesn't churn the registry
+# bytes between runs. The benchmark only consults graph contents — the
+# timestamp is metadata that has no semantic effect on tool output, but
+# we still want byte-identical state across invocations for any future
+# code that hashes the registry.
+_DETERMINISTIC_INDEXED_AT = datetime(1970, 1, 1, tzinfo=timezone.utc)
+
+
+def build_l2_graph_dir(
+    snapshot: SchemaSnapshot, *, work_dir: Path, embeddings: bool = False
+) -> Path:
+    """Build a Kuzu graph + registry under ``work_dir`` and return the dir.
+
+    When ``embeddings`` is True, tables are indexed with
+    ``EmbeddingsConfig(index_tables=True)`` so the ``semantic_search`` /
+    hybrid-rerank legs measure real cosine retrieval instead of the
+    no-vectors fallback (which would peg semantic recall at 0.0 and make
+    the embeddings CI lane indistinguishable from the plain one).  The
+    embedding model revision is pinned, so vectors are deterministic
+    across runs.
+
+    The returned path is suitable as the ``graph_dir`` argument to every
+    MCP tool payload function. ``work_dir`` is created if it doesn't
+    exist; the caller (typically ``tempfile.TemporaryDirectory``) is
+    responsible for cleanup.
+
+    Notes:
+        ``run_relationship_discovery=False`` matches the determinism
+        policy — L2 metrics measure the tools' behaviour over the
+        declared FK / structural shape, not over heuristic discovery
+        output. (L1 covers the discovery quality leg via its own
+        :func:`pretensor.benchmark.l1.pipeline.discover_inferred_joins_blind`.)
+
+        After build, this helper strips :class:`Cluster` nodes (and the
+        ``IN_CLUSTER`` edges that link them to tables). Reason: the
+        intelligence layer that ``GraphBuilder.build`` runs unconditionally
+        produces cluster labels via Louvain/Leiden community detection,
+        and even with a seeded RNG the labels can differ across
+        environments (different igraph builds, different platform RNG
+        wiring). Those labels feed into the FTS5 keyword index's
+        ``cluster_context`` column and shift BM25 scores across machines.
+        L2 measures ``query`` / ``semantic_search`` / ``traverse`` /
+        ``compile_metric`` quality — none of those tools depend on
+        clustering being present, so removing it from the graph keeps
+        the benchmark cross-environment-deterministic without changing
+        what's being measured.
+    """
+    work_dir.mkdir(parents=True, exist_ok=True)
+    graph_path = work_dir / "graphs" / f"{snapshot.connection_name}.kuzu"
+    graph_path.parent.mkdir(parents=True, exist_ok=True)
+
+    store = KuzuStore(graph_path)
+    try:
+        cfg = PretensorConfig(
+            embeddings=EmbeddingsConfig(index_tables=embeddings),
+        )
+        GraphBuilder().build(
+            snapshot, store, run_relationship_discovery=False, config=cfg
+        )
+        if embeddings and not store.has_any_table_embeddings():
+            # The production embed path tolerates per-run failures (a
+            # degraded index beats an aborted one), but a benchmark lane
+            # that silently measures the no-vectors fallback would report
+            # a bogus regression against the embeddings baseline. Fail
+            # loudly instead — the usual culprit is a failed model
+            # download (e.g. HF Hub rate-limiting in CI).
+            msg = (
+                "L2 was invoked with --embeddings but no table vectors were "
+                "computed; the embedding model is likely unavailable "
+                "(download failure / rate limit). Fix the model fetch or "
+                "re-run without --embeddings."
+            )
+            raise RuntimeError(msg)
+        # See docstring: clusters introduce cross-environment nondeterminism
+        # via FTS5's cluster_context column. L2 doesn't measure clustering
+        # quality (L1 does), so strip them before the runner queries the graph.
+        _strip_clusters(store)
+    finally:
+        store.close()
+
+    reg = GraphRegistry(work_dir / "registry.json").load()
+    reg.upsert(
+        connection_name=snapshot.connection_name,
+        database=snapshot.database,
+        dsn="",
+        graph_path=graph_path,
+        indexed_at=_DETERMINISTIC_INDEXED_AT,
+    )
+    reg.save()
+
+    return work_dir
+
+
+def _strip_clusters(store: KuzuStore) -> None:
+    """Remove all ``Cluster`` nodes (and ``IN_CLUSTER`` edges) from the graph.
+
+    Kuzu doesn't expose ``DETACH DELETE`` semantics on every type, so
+    we wipe the relationship rows first and then the cluster nodes.
+    """
+    store.execute_write("MATCH ()-[r:IN_CLUSTER]->() DELETE r")
+    store.execute_write("MATCH (c:Cluster) DELETE c")