pretensor 0.1.0__py3-none-any.whl

pretensor/intelligence/role_exemplars.py

@@ -0,0 +1,263 @@
+"""Curated exemplar tables per ``TableRole`` for embedding-based role voting.
+
+The classifier in :mod:`pretensor.entities.classifier` uses several signals
+(name patterns, column shape, structural degree).  This module adds one
+more: a nearest-centroid vote against a small per-role exemplar set.
+
+For each role, we list a few real-ish ``(qualified_name, columns)`` pairs
+that are *prototypical* of that role.  At runtime, an
+:class:`EmbeddingClient` embeds the exemplar text via
+:func:`format_entity_text` and we mean-pool the vectors per role to get
+one centroid per role.  A new table's role vote is then
+``cosine(table_text_embedding, centroid)``, normalized so the best
+matching role gets a score in roughly ``[0.0, 1.0]``.
+
+The exemplar set ships in source so the wheel is self-contained — no
+fixture files at runtime.  Embeddings are computed lazily (first call)
+and cached on the embedding client via ``_centroid_cache`` so repeated
+classification runs reuse the same centroids without re-embedding.
+
+Determinism contract:
+
+* ``EmbeddingsConfig.role_weight=0.0`` (default) skips this signal entirely.
+* When > 0, the vote is *added* to existing heuristic scores; it never
+  replaces them.  ``role_weight`` should stay well below the typical
+  heuristic score magnitudes (~0.5–2.5) so the heuristic remains primary.
+* If the embedding client is unavailable (extra not installed, or
+  ``embed`` raises), the function returns an all-zeros dict and the
+  classifier falls back to heuristic-only scoring.
+"""
+
+from __future__ import annotations
+
+import logging
+import weakref
+
+from pretensor.intelligence.embeddings import (
+    EmbeddingClient,
+    cosine_similarity,
+    format_entity_text,
+)
+
+__all__ = [
+    "ROLE_EXEMPLARS",
+    "compute_role_centroids",
+    "embedding_role_vote",
+]
+
+logger = logging.getLogger(__name__)
+
+
+# Per-role exemplars: (qualified_name, column_names).  A handful per role is
+# sufficient — mean-pooling smooths out idiosyncrasies of any single example.
+# Names are illustrative; the embeddings encode the schema "shape" via
+# format_entity_text, not the brand-name.
+ROLE_EXEMPLARS: dict[str, list[tuple[str, list[str]]]] = {
+    "fact": [
+        ("public.orders", ["id", "customer_id", "order_date", "total_amount"]),
+        ("public.transactions", ["id", "account_id", "amount", "ts"]),
+        ("sales.invoice_lines", ["id", "invoice_id", "product_id", "qty", "price"]),
+        ("public.events", ["id", "user_id", "event_type", "occurred_at"]),
+        ("public.payments", ["id", "order_id", "amount", "paid_at", "method"]),
+    ],
+    "dimension": [
+        ("public.customers", ["id", "first_name", "last_name", "email"]),
+        ("public.products", ["id", "name", "category", "price"]),
+        ("public.users", ["id", "username", "email", "created_at"]),
+        ("public.locations", ["id", "country", "region", "city"]),
+        ("public.suppliers", ["id", "name", "address", "phone"]),
+    ],
+    "bridge": [
+        ("public.user_roles", ["user_id", "role_id"]),
+        ("public.product_categories", ["product_id", "category_id"]),
+        ("public.order_tags", ["order_id", "tag_id"]),
+        ("public.film_actor", ["film_id", "actor_id"]),
+    ],
+    "junction": [
+        ("public.user_groups", ["user_id", "group_id"]),
+        ("public.role_permissions", ["role_id", "permission_id"]),
+    ],
+    "staging": [
+        ("staging.raw_orders", ["id", "raw_payload", "ingested_at"]),
+        ("staging.tmp_users", ["id", "data", "loaded_at"]),
+        ("public.import_buffer", ["id", "source", "raw_json", "imported_at"]),
+    ],
+    "audit": [
+        ("public.audit_log", ["id", "actor", "action", "target", "ts"]),
+        ("public.access_log", ["id", "user_id", "endpoint", "ip", "ts"]),
+        ("public.history", ["id", "entity_id", "old_value", "new_value", "ts"]),
+    ],
+    "snapshot_scd": [
+        ("public.customer_history", ["id", "customer_id", "valid_from", "valid_to"]),
+        ("public.product_snapshot", ["id", "product_id", "snapshot_date", "price"]),
+        ("public.account_scd2", ["id", "account_id", "is_current", "effective_at"]),
+    ],
+    "aggregate": [
+        ("analytics.daily_revenue", ["day", "total", "order_count"]),
+        ("analytics.monthly_active_users", ["month", "mau", "growth_rate"]),
+        ("public.summary_stats", ["metric", "value", "computed_at"]),
+    ],
+    "system": [
+        ("public.schema_migrations", ["version", "applied_at"]),
+        ("public.ar_internal_metadata", ["key", "value", "updated_at"]),
+    ],
+    "entity_candidate": [
+        ("public.thing", ["id", "name"]),
+        ("public.unknown_table", ["id", "data"]),
+    ],
+    # 'unknown' has no exemplars by design — it's the fallback when nothing
+    # else fires; voting against it would defeat the purpose.
+}
+
+
+# Cache: maps client instance → {role → centroid}.  Keyed on the actual
+# object via WeakKeyDictionary so a freed client's entry vanishes
+# automatically — guards against CPython recycling a deleted client's
+# ``id()`` for a fresh, unrelated allocation, which would let a stale
+# centroid leak into the new client's lookups in long-running processes.
+_CENTROID_CACHE: "weakref.WeakKeyDictionary[EmbeddingClient, dict[str, list[float]]]" = weakref.WeakKeyDictionary()
+# Bounded fallback for client types that aren't weakref-able (e.g. mocks
+# built from raw ``object()``).
+#
+# CAVEAT: this fallback path keys on ``id(client)`` directly and therefore
+# does NOT have the recycling protection of the WeakKeyDictionary above.
+# If a non-weakref-able client A is freed and its id is recycled by a
+# fresh, unrelated client B before A's fallback entry is evicted, B's
+# lookup would return A's stale centroid. The hazard is bounded in
+# practice because (a) only test mocks land here — production
+# ``LocalEmbeddingClient`` is weakref-able, (b) the FIFO ceiling
+# below caps the window in which a stale entry can survive, and (c) tests
+# clear the cache between runs via ``_clear_centroid_cache_for_tests``.
+# Treat the fallback as best-effort for tests, not a load-bearing
+# correctness guarantee.
+_CENTROID_CACHE_FALLBACK: dict[int, dict[str, list[float]]] = {}
+# Maximum entries retained in the non-weakref-able fallback cache. The
+# typical caller is one ``LocalEmbeddingClient`` per process; the only
+# things landing here are short-lived test mocks. A small ceiling
+# (insertion-order eviction, FIFO) is plenty to keep the dict bounded
+# without losing the cache's value on the realistic call patterns.
+_CENTROID_CACHE_FALLBACK_MAX = 32
+
+
+def _cache_get(client: EmbeddingClient) -> dict[str, list[float]] | None:
+    try:
+        return _CENTROID_CACHE.get(client)
+    except TypeError:
+        return _CENTROID_CACHE_FALLBACK.get(id(client))
+
+
+def _cache_set(client: EmbeddingClient, value: dict[str, list[float]]) -> None:
+    try:
+        _CENTROID_CACHE[client] = value
+    except TypeError:
+        # FIFO evict to keep the dict bounded.
+        while len(_CENTROID_CACHE_FALLBACK) >= _CENTROID_CACHE_FALLBACK_MAX:
+            oldest_key = next(iter(_CENTROID_CACHE_FALLBACK))
+            del _CENTROID_CACHE_FALLBACK[oldest_key]
+        _CENTROID_CACHE_FALLBACK[id(client)] = value
+
+
+def compute_role_centroids(
+    client: EmbeddingClient,
+) -> dict[str, list[float]]:
+    """Return one centroid vector per role with at least one exemplar.
+
+    Embeds every exemplar via :func:`format_entity_text` and mean-pools
+    per role.  Cached on the client (weak-ref keyed) so the second call
+    is free without leaking memory if the client is later garbage-collected.
+    Returns an empty dict (no centroids) when the client raises any
+    exception while embedding — caller falls back to heuristic-only.
+    """
+    cached = _cache_get(client)
+    if cached is not None:
+        return cached
+
+    centroids: dict[str, list[float]] = {}
+    try:
+        # Flatten the exemplar set into one batch for one embed() call.
+        flat_texts: list[str] = []
+        flat_owners: list[str] = []
+        for role, exemplars in ROLE_EXEMPLARS.items():
+            for qname, cols in exemplars:
+                flat_texts.append(format_entity_text(qname, cols))
+                flat_owners.append(role)
+        if not flat_texts:
+            _cache_set(client, centroids)
+            return centroids
+
+        vectors = client.embed(flat_texts)
+        if not vectors or len(vectors) != len(flat_texts):
+            # NullEmbeddingClient returns []; treat as no centroids.
+            _cache_set(client, centroids)
+            return centroids
+
+        per_role: dict[str, list[list[float]]] = {}
+        for role, vec in zip(flat_owners, vectors, strict=True):
+            per_role.setdefault(role, []).append(vec)
+
+        for role, vecs in per_role.items():
+            if not vecs:
+                continue
+            dim = len(vecs[0])
+            mean = [0.0] * dim
+            count = 0
+            for v in vecs:
+                if len(v) != dim:
+                    continue
+                for i in range(dim):
+                    mean[i] += float(v[i])
+                count += 1
+            if count == 0:
+                continue
+            centroids[role] = [x / count for x in mean]
+    except ImportError as exc:
+        logger.warning(
+            "role_exemplars: [embeddings] extra not installed (%s); "
+            "skipping role centroids",
+            exc,
+        )
+        centroids = {}
+    except Exception as exc:  # noqa: BLE001 — additive signal, never raise
+        logger.warning(
+            "role_exemplars: failed to compute centroids (%s); "
+            "falling back to heuristic-only role vote",
+            exc,
+        )
+        centroids = {}
+
+    _cache_set(client, centroids)
+    return centroids
+
+
+def embedding_role_vote(
+    table_embedding: list[float] | None,
+    centroids: dict[str, list[float]],
+) -> dict[str, float]:
+    """Cosine vote for the table's embedding against per-role centroids.
+
+    ``table_embedding`` is typically pre-computed at index time via
+    ``EmbeddingIndexStep`` and read back from the store; pass ``None``
+    when the table doesn't carry one and this function returns an
+    all-zeros vote.
+
+    Returns ``{role: cosine}`` for every role with a centroid; roles
+    without a centroid are not present in the output.  Cosine values are
+    in ``[-1, 1]``; the classifier scales by ``role_weight`` before
+    blending.
+    """
+    if table_embedding is None or not centroids:
+        return {role: 0.0 for role in centroids}
+
+    out: dict[str, float] = {}
+    for role, centroid in centroids.items():
+        try:
+            out[role] = cosine_similarity(table_embedding, centroid)
+        except ValueError:
+            out[role] = 0.0
+    return out
+
+
+def _clear_centroid_cache_for_tests() -> None:
+    """Test helper: drop the module-level centroid cache."""
+    _CENTROID_CACHE.clear()
+    _CENTROID_CACHE_FALLBACK.clear()

pretensor/intelligence/schema_classification.py

@@ -0,0 +1,360 @@
+"""Persist table roles and cluster-level schema patterns after clustering."""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections import deque
+from typing import Literal, Protocol, runtime_checkable
+
+from pretensor.core.store import KuzuStore
+from pretensor.entities.classifier import (
+    TableClassification,
+    TableClassifier,
+    TableClassifierInput,
+)
+from pretensor.intelligence.clustering import Cluster
+from pretensor.intelligence.embeddings import EmbeddingClient
+from pretensor.intelligence.role_exemplars import (
+    compute_role_centroids,
+    embedding_role_vote,
+)
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "LlmTableClassificationClient",
+    "classify_database_tables",
+    "classify_database_tables_async",
+    "compute_cluster_schema_patterns",
+    "load_fk_reference_pairs",
+]
+
+SchemaPattern = Literal["star", "snowflake", "constellation", "erd", "unknown"]
+
+_CLASSIFIER = TableClassifier()
+
+# Precedence for ``schema_pattern`` when multiple signals apply:
+# 1. constellation — multiple facts share at least one dimension (direct FK or via bridges)
+# 2. snowflake — star-like hub plus at least one dimension-to-dimension FK in the cluster
+# 3. star — one (or primary) fact with 3+ distinct dimensions via outgoing FKs
+# 4. erd — connected tables with FKs but not a clear star/snowflake/constellation
+# 5. role-only fallback — used when the cluster has no FK edges (introspection gap)
+
+
+@runtime_checkable
+class LlmTableClassificationClient(Protocol):
+    """Async JSON batch classifier for low-confidence tables."""
+
+    async def classify_tables_json(self, user_prompt: str) -> str:
+        """Return JSON array aligned with batch order: {role, confidence, signals?}."""
+        ...
+
+
+def load_fk_reference_pairs(
+    store: KuzuStore, database_key: str
+) -> list[tuple[str, str]]:
+    """Return all ``(src_table_id, dst_table_id)`` for ``FK_REFERENCES`` in one database."""
+    rows = store.query_all_rows(
+        """
+        MATCH (a:SchemaTable {database: $db})-[r:FK_REFERENCES]->(b:SchemaTable {database: $db})
+        RETURN a.node_id, b.node_id
+        """,
+        {"db": database_key},
+    )
+    return [(str(a), str(b)) for a, b in rows]
+
+
+def _fk_degree_rows(store: KuzuStore, database_key: str) -> list[tuple[str, int, int]]:
+    return store.query_all_rows(
+        """
+        MATCH (t:SchemaTable {database: $db})
+        OPTIONAL MATCH (t)-[:FK_REFERENCES]->(o:SchemaTable)
+        OPTIONAL MATCH (i:SchemaTable)-[:FK_REFERENCES]->(t)
+        RETURN t.node_id, count(DISTINCT o.node_id), count(DISTINCT i.node_id)
+        """,
+        {"db": database_key},
+    )
+
+
+def _table_context_rows(store: KuzuStore, database_key: str) -> list[tuple]:
+    return store.query_all_rows(
+        """
+        MATCH (t:SchemaTable {database: $db})
+        RETURN t.node_id, t.table_name, t.schema_name, t.row_count,
+               t.seq_scan_count, t.idx_scan_count, t.insert_count, t.update_count
+        """,
+        {"db": database_key},
+    )
+
+
+def _columns_by_table(store: KuzuStore, database_key: str) -> dict[str, list[str]]:
+    rows = store.query_all_rows(
+        """
+        MATCH (t:SchemaTable {database: $db})-[:HAS_COLUMN]->(c:SchemaColumn)
+        RETURN t.node_id, c.column_name, c.ordinal_position
+        ORDER BY t.node_id, c.ordinal_position, c.column_name
+        """,
+        {"db": database_key},
+    )
+    out: dict[str, list[str]] = {}
+    for tid, cname, _ord in rows:
+        tid_s = str(tid)
+        out.setdefault(tid_s, []).append(str(cname))
+    return out
+
+
+def _initial_classification(
+    inp: TableClassifierInput,
+) -> TableClassification:
+    return _CLASSIFIER.classify(inp)
+
+
+def _fk_pairs_for_cluster(
+    table_ids: frozenset[str], all_pairs: list[tuple[str, str]]
+) -> list[tuple[str, str]]:
+    return [p for p in all_pairs if p[0] in table_ids and p[1] in table_ids]
+
+
+def _dimensions_reachable_from(
+    start: str,
+    table_ids: frozenset[str],
+    adj_out: dict[str, list[str]],
+    dims: set[str],
+) -> set[str]:
+    """Directed reachability from ``start`` within ``table_ids``; return dimension nodes hit."""
+    found: set[str] = set()
+    q: deque[str] = deque([start])
+    seen = {start}
+    while q:
+        u = q.popleft()
+        if u in dims:
+            found.add(u)
+        for v in adj_out.get(u, ()):
+            if v not in table_ids or v in seen:
+                continue
+            seen.add(v)
+            q.append(v)
+    return found
+
+
+def _facts_share_dimension_via_paths(
+    facts: set[str],
+    table_ids: frozenset[str],
+    fk_pairs: list[tuple[str, str]],
+    dims: set[str],
+) -> bool:
+    """True if some pair of facts reaches a common dimension (paths may cross bridges)."""
+    if len(facts) < 2:
+        return False
+    adj_out: dict[str, list[str]] = {}
+    for a, b in fk_pairs:
+        adj_out.setdefault(a, []).append(b)
+    dim_sets: dict[str, set[str]] = {}
+    for f in facts:
+        dim_sets[f] = _dimensions_reachable_from(f, table_ids, adj_out, dims)
+    fact_list = list(facts)
+    for i in range(len(fact_list)):
+        for j in range(i + 1, len(fact_list)):
+            if dim_sets[fact_list[i]] & dim_sets[fact_list[j]]:
+                return True
+    return False
+
+
+def _shared_dimension_direct(
+    facts: set[str], dims: set[str], fk_pairs: list[tuple[str, str]]
+) -> dict[str, set[str]]:
+    """fact -> set of dimensions referenced by a direct FK from the fact."""
+    out: dict[str, set[str]] = {f: set() for f in facts}
+    for a, b in fk_pairs:
+        if a in facts and b in dims:
+            out.setdefault(a, set()).add(b)
+    return out
+
+
+def _schema_pattern_role_only(
+    tset: frozenset[str],
+    roles: dict[str, str],
+    bridges: set[str],
+) -> SchemaPattern:
+    """Heuristic when the cluster has no FK edges (or none inside the cluster)."""
+    facts = {tid for tid in tset if roles.get(tid) == "fact"}
+    dims = {tid for tid in tset if roles.get(tid) == "dimension"}
+    if len(facts) >= 2 and len(dims) >= 2:
+        return "constellation"
+    if facts and len(dims) >= 3:
+        return "star"
+    if bridges and (facts or dims):
+        return "erd"
+    if len(tset) >= 2:
+        return "erd"
+    return "unknown"
+
+
+def compute_cluster_schema_patterns(
+    clusters: list[Cluster],
+    role_by_table: dict[str, TableClassification],
+    fk_pairs: list[tuple[str, str]],
+) -> dict[int, SchemaPattern]:
+    """Map cluster index to ``schema_pattern`` (aligned with ``cluster::{idx}`` ids).
+
+    Args:
+        clusters: Leiden clusters.
+        role_by_table: Latest classification per table node id.
+        fk_pairs: All ``FK_REFERENCES`` (src, dst) for the logical database.
+    """
+    patterns: dict[int, SchemaPattern] = {}
+    for idx, cluster in enumerate(clusters):
+        tset = frozenset(cluster.table_ids)
+        if not tset:
+            patterns[idx] = "unknown"
+            continue
+        roles = {
+            tid: role_by_table.get(tid, TableClassification("unknown", 0.0, ())).role
+            for tid in tset
+        }
+        facts = {tid for tid, r in roles.items() if r == "fact"}
+        dims = {tid for tid, r in roles.items() if r == "dimension"}
+        bridges = {tid for tid, r in roles.items() if r == "bridge"}
+        cluster_pairs = _fk_pairs_for_cluster(tset, fk_pairs)
+
+        if not cluster_pairs:
+            patterns[idx] = _schema_pattern_role_only(tset, roles, bridges)
+            continue
+
+        dim_dim = sum(1 for a, b in cluster_pairs if a in dims and b in dims)
+        fact_to_dim = _shared_dimension_direct(facts, dims, cluster_pairs)
+        facts_per_dim: dict[str, set[str]] = {}
+        for f, dset in fact_to_dim.items():
+            for d in dset:
+                facts_per_dim.setdefault(d, set()).add(f)
+
+        shared_direct = any(len(fs) >= 2 for fs in facts_per_dim.values())
+        multi_hop_share = _facts_share_dimension_via_paths(
+            facts, tset, cluster_pairs, dims
+        )
+        constellation = len(facts) >= 2 and (shared_direct or multi_hop_share)
+
+        max_dim_from_one_fact = max((len(s) for s in fact_to_dim.values()), default=0)
+        star_like = bool(facts and dims and max_dim_from_one_fact >= 3)
+
+        pattern: SchemaPattern = "unknown"
+        if constellation:
+            pattern = "constellation"
+        elif star_like:
+            pattern = "snowflake" if dim_dim >= 1 else "star"
+        elif len(tset) >= 2 and cluster_pairs:
+            pattern = "erd"
+        elif bridges and cluster_pairs:
+            pattern = "erd"
+
+        patterns[idx] = pattern
+    return patterns
+
+
+def classify_database_tables(
+    store: KuzuStore,
+    database_key: str,
+    *,
+    role_weight: float = 0.0,
+    embedding_client: EmbeddingClient | None = None,
+) -> dict[str, TableClassification]:
+    """Write classifier fields and return heuristic results.
+
+    When ``role_weight > 0`` and an ``embedding_client`` is provided, this
+    function computes per-role centroids (cached on the client identity)
+    and bulk-fetches every table's stored embedding.  Each table's role
+    classification then receives an additive ``role_weight * cosine``
+    signal per role; heuristic scores remain primary.
+    Defaults keep behavior byte-identical to the pre-embedding pipeline.
+    """
+    store.ensure_schema()
+    fk_rows = _fk_degree_rows(store, database_key)
+    fk_out = {str(r[0]): int(r[1] or 0) for r in fk_rows}
+    fk_in = {str(r[0]): int(r[2] or 0) for r in fk_rows}
+    col_map = _columns_by_table(store, database_key)
+    ctx_rows = _table_context_rows(store, database_key)
+    results: dict[str, TableClassification] = {}
+
+    # Pre-pass: compute role centroids + bulk-fetch every table's embedding.
+    # Both are no-ops when role_weight is 0, when no embedding_client is
+    # supplied, or when the [embeddings] extra is missing (the client raises
+    # ImportError → role_exemplars logs once and returns {}).
+    centroids: dict[str, list[float]] = {}
+    embeddings_by_id: dict[str, list[float]] = {}
+    if role_weight > 0.0 and embedding_client is not None:
+        centroids = compute_role_centroids(embedding_client)
+        if centroids:
+            for emb_row in store.iter_table_embeddings(database=database_key):
+                if emb_row.embedding is None:
+                    continue
+                embeddings_by_id[emb_row.node_id] = [
+                    float(x) for x in emb_row.embedding
+                ]
+
+    for row in ctx_rows:
+        tid = str(row[0])
+        tname = str(row[1])
+        sname = str(row[2])
+        rcount = int(row[3]) if row[3] is not None else None
+        seq_sc = int(row[4]) if row[4] is not None else None
+        idx_sc = int(row[5]) if row[5] is not None else None
+        ins = int(row[6]) if row[6] is not None else None
+        upd = int(row[7]) if row[7] is not None else None
+        cols = col_map.get(tid, [])
+        inp = TableClassifierInput(
+            name=tname,
+            schema_name=sname,
+            columns=cols,
+            row_count=rcount,
+            fk_out_degree=fk_out.get(tid, 0),
+            fk_in_degree=fk_in.get(tid, 0),
+            seq_scan_count=seq_sc,
+            idx_scan_count=idx_sc,
+            insert_count=ins,
+            update_count=upd,
+        )
+
+        vote: dict[str, float] | None = None
+        emb = embeddings_by_id.get(tid)
+        if centroids and emb is not None:
+            vote = embedding_role_vote(emb, centroids)
+
+        cls = _CLASSIFIER.classify(
+            inp,
+            embedding_vote=vote,
+            role_weight=role_weight if vote is not None else 0.0,
+        )
+        results[tid] = cls
+        store.set_table_classification(
+            tid,
+            role=cls.role,
+            role_confidence=cls.confidence,
+            classification_signals_json=json.dumps(
+                list(cls.signals), ensure_ascii=False
+            ),
+        )
+
+    return results
+
+
+async def classify_database_tables_async(
+    store: KuzuStore,
+    database_key: str,
+    *,
+    llm_client: LlmTableClassificationClient | None = None,
+    role_weight: float = 0.0,
+    embedding_client: EmbeddingClient | None = None,
+) -> dict[str, TableClassification]:
+    """Run heuristic classification with optional embedding role vote.
+
+    ``llm_client`` is accepted but ignored (legacy parameter).  See the
+    sync :func:`classify_database_tables` for ``role_weight`` /
+    ``embedding_client`` semantics.
+    """
+    return classify_database_tables(
+        store,
+        database_key,
+        role_weight=role_weight,
+        embedding_client=embedding_client,
+    )

pretensor/intelligence/scoring.py

@@ -0,0 +1,76 @@
+"""Scoring contracts and registry for relationship discovery candidates."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Iterator
+from typing import TypeAlias
+
+from pretensor.connectors.models import SchemaSnapshot
+from pretensor.graph_models.relationship import RelationshipCandidate
+
+__all__ = ["JoinKey", "RelationshipScorer", "ScorerRegistry"]
+
+JoinKey: TypeAlias = tuple[str, str, str, str]
+
+
+class RelationshipScorer(ABC):
+    """Contract for candidate generators used by relationship discovery."""
+
+    @abstractmethod
+    def name(self) -> str:
+        """Stable scorer name for observability and deduping."""
+
+    @abstractmethod
+    def score(
+        self,
+        snapshot: SchemaSnapshot,
+        explicit_fk_keys: set[JoinKey],
+    ) -> list[RelationshipCandidate]:
+        """Generate inferred relationship candidates for a snapshot."""
+
+
+class ScorerRegistry:
+    """Ordered registry for relationship scorers.
+
+    The first scorer can emit broad candidates while later scorers can refine or
+    add alternatives. Discovery remains deterministic because the insertion order
+    is preserved.
+    """
+
+    def __init__(self, scorers: Iterable[RelationshipScorer] | None = None) -> None:
+        self._scorers_by_name: dict[str, RelationshipScorer] = {}
+        self._order: list[str] = []
+        for scorer in scorers or ():
+            self.register(scorer)
+
+    def register(self, scorer: RelationshipScorer) -> None:
+        """Register a scorer instance by unique name."""
+        name = scorer.name()
+        if name in self._scorers_by_name:
+            msg = f"Relationship scorer already registered: {name}"
+            raise ValueError(msg)
+        self._scorers_by_name[name] = scorer
+        self._order.append(name)
+
+    def score_all(
+        self,
+        snapshot: SchemaSnapshot,
+        explicit_fk_keys: set[JoinKey],
+    ) -> list[RelationshipCandidate]:
+        """Run all scorers in registration order and collect candidates."""
+        out: list[RelationshipCandidate] = []
+        for name in self._order:
+            scorer = self._scorers_by_name[name]
+            out.extend(scorer.score(snapshot, explicit_fk_keys))
+        return out
+
+    def __iter__(self) -> Iterator[RelationshipScorer]:
+        """Iterate scorers in registration order (read-only).
+
+        Used by ``extend_with_embedding_scorer`` (and any future helper
+        that builds a derived registry) to copy scorers without poking at
+        the private ``_scorers_by_name`` / ``_order`` state.
+        """
+        for name in self._order:
+            yield self._scorers_by_name[name]