pretensor 0.1.0__py3-none-any.whl

pretensor/intelligence/semantic.py

@@ -0,0 +1,240 @@
+"""Embedding-based ``RelationshipScorer``.
+
+Walks every pair of tables that carry a non-null ``SchemaTable.embedding``
+in the store, computes cosine similarity, and emits ``INFERRED_JOIN``
+candidates with ``source='embedding'`` and ``status='suggested'`` for
+pairs whose cosine ≥ ``threshold`` and whose implied join columns pass
+the heuristic's type-family compatibility check.
+
+Determinism contract:
+
+* Opt-in via ``EmbeddingsConfig.join_threshold`` — ``None`` skips the
+  scorer entirely; the scorer registry behaves byte-identically to the
+  pre-embedding pipeline output.
+* Additive only: candidates pass through the same combiner / shadow-alias
+  filter / explicit-FK dedup as heuristic candidates.  An embedding-only
+  proposal that conflicts with an explicit FK is dropped at the
+  ``discovery.discover`` filter step (line 82); the scorer also skips
+  ``JoinKey``s already in the explicit-FK set up front.
+* Type-family veto: every emitted candidate has ``types_compatible``
+  green for its ``(source_column, target_column)`` pair, reusing the
+  heuristic's check.
+
+Index-time staging: the scorer reads embeddings from the store.  The
+index/reindex paths call ``compute_table_embeddings`` right after schema
+rows are written and *before* relationship discovery, so the scorer sees
+this run's vectors even on the very first index.  Callers that drive
+``run_intelligence_layer`` directly without precomputing still get the
+in-pipeline ``embedding_index`` step, in which case discovery (which runs
+earlier) only sees vectors from the previous cycle.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from pretensor.connectors.models import Column, SchemaSnapshot, Table
+from pretensor.core.ids import table_node_id
+from pretensor.core.store import KuzuStore
+from pretensor.graph_models.relationship import RelationshipCandidate
+from pretensor.intelligence.embeddings import (
+    cosine_similarity,
+    embeddings_disabled_via_env,
+)
+from pretensor.intelligence.heuristic import types_compatible
+from pretensor.intelligence.scoring import JoinKey, RelationshipScorer, ScorerRegistry
+
+__all__ = ["EmbeddingRelationshipScorer", "extend_with_embedding_scorer"]
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingRelationshipScorer(RelationshipScorer):
+    """Cosine-similarity-driven candidate generator for ``INFERRED_JOIN``."""
+
+    def __init__(
+        self,
+        *,
+        store: KuzuStore,
+        database_key: str,
+        threshold: float,
+    ) -> None:
+        if not 0.0 <= threshold <= 1.0:
+            msg = f"threshold must be in [0.0, 1.0], got {threshold}"
+            raise ValueError(msg)
+        self._store = store
+        self._database_key = database_key
+        self._threshold = threshold
+
+    def name(self) -> str:
+        return "embedding"
+
+    def score(
+        self,
+        snapshot: SchemaSnapshot,
+        explicit_fk_keys: set[JoinKey],
+    ) -> list[RelationshipCandidate]:
+        embeddings = self._load_embeddings()
+        if not embeddings:
+            logger.debug(
+                "EmbeddingRelationshipScorer: no embedded tables for %s; skipping",
+                self._database_key,
+            )
+            return []
+
+        table_by_nid = self._index_snapshot_tables(snapshot)
+        if not table_by_nid:
+            return []
+
+        candidates: list[RelationshipCandidate] = []
+        embedded_ids = sorted(nid for nid in embeddings if nid in table_by_nid)
+        for i in range(len(embedded_ids)):
+            a_id = embedded_ids[i]
+            for j in range(i + 1, len(embedded_ids)):
+                b_id = embedded_ids[j]
+                cos = cosine_similarity(embeddings[a_id], embeddings[b_id])
+                if cos < self._threshold:
+                    continue
+                a_tbl = table_by_nid[a_id]
+                b_tbl = table_by_nid[b_id]
+
+                # Both directions: a → b (a is source, b is target) and b → a.
+                # The combiner downstream picks the highest-confidence winner
+                # per directed JoinKey, so emitting both is safe.
+                for src_tbl, dst_tbl, src_id, dst_id in (
+                    (a_tbl, b_tbl, a_id, b_id),
+                    (b_tbl, a_tbl, b_id, a_id),
+                ):
+                    pair = _pick_join_columns(src_tbl, dst_tbl)
+                    if pair is None:
+                        continue
+                    src_col, dst_col = pair
+                    join_key: JoinKey = (src_id, dst_id, src_col, dst_col)
+                    if join_key in explicit_fk_keys:
+                        continue
+                    candidate_id = (
+                        f"embedding::{src_id}::{src_col}::{dst_id}::{dst_col}"
+                    )
+                    candidates.append(
+                        RelationshipCandidate(
+                            candidate_id=candidate_id,
+                            source_node_id=src_id,
+                            target_node_id=dst_id,
+                            source_column=src_col,
+                            target_column=dst_col,
+                            source="embedding",
+                            confidence=float(cos),
+                            status="suggested",
+                            reasoning=(
+                                f"cosine={cos:.3f} between {src_tbl.schema_name}."
+                                f"{src_tbl.name} and {dst_tbl.schema_name}.{dst_tbl.name}"
+                            ),
+                        )
+                    )
+        logger.debug(
+            "EmbeddingRelationshipScorer: emitted %d candidates for %s "
+            "(threshold=%.3f, embedded_tables=%d)",
+            len(candidates),
+            self._database_key,
+            self._threshold,
+            len(embedded_ids),
+        )
+        return candidates
+
+    def _load_embeddings(self) -> dict[str, list[float]]:
+        """One-shot fetch of every embedded ``SchemaTable`` in the database."""
+        out: dict[str, list[float]] = {}
+        for row in self._store.iter_table_embeddings(database=self._database_key):
+            if row.embedding is None:
+                continue
+            out[row.node_id] = [float(x) for x in row.embedding]
+        return out
+
+    def _index_snapshot_tables(self, snapshot: SchemaSnapshot) -> dict[str, Table]:
+        """Map ``SchemaTable.node_id`` → ``Table`` for column lookups."""
+        out: dict[str, Table] = {}
+        for tbl in snapshot.tables:
+            nid = table_node_id(snapshot.connection_name, tbl.schema_name, tbl.name)
+            out[nid] = tbl
+        return out
+
+
+def _pick_join_columns(src: Table, dst: Table) -> tuple[str, str] | None:
+    """Pick a likely (src_col, dst_col) join pair for an embedding candidate.
+
+    Strategy:
+    1.  Prefer the canonical FK shape: ``src.<dst_table>_id`` → ``dst.id``.
+    2.  Fall back to any type-family-compatible ``(FK-like, PK-like)`` pair
+        where FK-like means a column ending in ``_id`` and PK-like means
+        ``id`` or ``<dst_table>_id``.
+    3.  If no compatible pair exists, return ``None`` (the candidate is
+        dropped).
+
+    Type compatibility is enforced via the heuristic scorer's
+    ``types_compatible`` so high-cosine pairs whose join columns disagree
+    on type family (e.g. ``customer_id INT`` ↔ ``customer_id VARCHAR``)
+    never reach the graph.
+    """
+    src_fk_cols: list[Column] = [
+        c for c in src.columns if c.name.lower().endswith("_id")
+    ]
+    dst_lower = dst.name.lower()
+    dst_pk_cols: list[Column] = [
+        c
+        for c in dst.columns
+        if c.name.lower() == "id" or c.name.lower() == f"{dst_lower}_id"
+    ]
+    if not src_fk_cols or not dst_pk_cols:
+        return None
+
+    # Preferred shape: src.<dst_table>_id → dst.id.
+    pref_src = next(
+        (c for c in src_fk_cols if c.name.lower() == f"{dst_lower}_id"), None
+    )
+    pref_dst = next((c for c in dst_pk_cols if c.name.lower() == "id"), None)
+    if (
+        pref_src is not None
+        and pref_dst is not None
+        and types_compatible(pref_src, pref_dst)
+    ):
+        return pref_src.name, pref_dst.name
+
+    # Fallback: any type-compat pair.
+    for sc in src_fk_cols:
+        for dc in dst_pk_cols:
+            if types_compatible(sc, dc):
+                return sc.name, dc.name
+    return None
+
+
+def extend_with_embedding_scorer(
+    base: ScorerRegistry,
+    *,
+    store: KuzuStore,
+    database_key: str,
+    threshold: float | None,
+) -> ScorerRegistry:
+    """Return a new ``ScorerRegistry`` containing ``base``'s scorers in order
+    plus an ``EmbeddingRelationshipScorer`` when ``threshold is not None``.
+
+    When ``threshold is None``, returns ``base`` unchanged (null-path
+    invariant).  ``PRETENSOR_EMBEDDINGS_DISABLED=1`` does the same even
+    when a threshold is set, so the kill switch forces the null path here
+    too.  Otherwise builds a fresh registry, copies every base scorer in
+    registration order, and appends the embedding scorer last so the
+    combiner sees heuristic / statistical candidates before the embedding
+    pass.
+    """
+    if threshold is None or embeddings_disabled_via_env():
+        return base
+    out = ScorerRegistry()
+    for scorer in base:
+        out.register(scorer)
+    out.register(
+        EmbeddingRelationshipScorer(
+            store=store,
+            database_key=database_key,
+            threshold=threshold,
+        )
+    )
+    return out

pretensor/intelligence/shadow_alias.py

@@ -0,0 +1,101 @@
+"""Shadow-alias detection for single-source views.
+
+A *shadow alias* is a view (``table_type = 'view'``) whose sole data source
+is one base table — signaled by having exactly one incoming ``LINEAGE`` edge
+(or at most ``lineage_in_max_for_alias`` edges, configurable via
+:class:`~pretensor.config.GraphConfig`).
+
+These views duplicate information already present on the base table and cause
+noise in clustering, INFERRED_JOIN generation, impact analysis, and traversal.
+The helpers here let each consumer detect and suppress them.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pretensor.config import GraphConfig
+    from pretensor.core.store import KuzuStore
+
+__all__ = [
+    "get_shadow_alias_node_ids",
+    "is_shadow_alias",
+]
+
+logger = logging.getLogger(__name__)
+
+
+def get_shadow_alias_node_ids(
+    store: KuzuStore,
+    database_key: str,
+    *,
+    config: GraphConfig | None = None,
+) -> frozenset[str]:
+    """Return ``node_id`` values for all shadow aliases in *database_key*.
+
+    A table is a shadow alias when:
+
+    1. ``table_type = 'view'``
+    2. ``count(incoming LINEAGE edges) <= lineage_in_max_for_alias`` (default 1)
+    3. It has at least one incoming LINEAGE edge (pure orphan views are not aliases)
+
+    The result is a :class:`frozenset` suitable for fast ``in`` checks by
+    downstream consumers.
+    """
+    if config is not None and not config.collapse_shadow_aliases:
+        return frozenset()
+
+    max_lineage_in = 1 if config is None else config.lineage_in_max_for_alias
+
+    rows = store.query_all_rows(
+        """
+        MATCH (src:SchemaTable)-[:LINEAGE]->(v:SchemaTable {database: $db})
+        WHERE v.table_type = 'view'
+        WITH v.node_id AS vid, count(DISTINCT src.node_id) AS lin
+        WHERE lin >= 1 AND lin <= $max_lin
+        RETURN vid
+        """,
+        {"db": database_key, "max_lin": max_lineage_in},
+    )
+    ids = frozenset(str(r[0]) for r in rows)
+    if ids:
+        logger.debug(
+            "Detected %d shadow alias(es) in %s (lineage_in_max=%d)",
+            len(ids),
+            database_key,
+            max_lineage_in,
+        )
+    return ids
+
+
+def is_shadow_alias(
+    store: KuzuStore,
+    node_id: str,
+    *,
+    config: GraphConfig | None = None,
+) -> bool:
+    """Check whether a single table node is a shadow alias.
+
+    Prefer :func:`get_shadow_alias_node_ids` when checking many tables at once
+    (one query vs. N).
+    """
+    if config is not None and not config.collapse_shadow_aliases:
+        return False
+
+    max_lineage_in = 1 if config is None else config.lineage_in_max_for_alias
+
+    rows = store.query_all_rows(
+        """
+        MATCH (t:SchemaTable {node_id: $nid})
+        WHERE t.table_type = 'view'
+        WITH t
+        MATCH (src:SchemaTable)-[:LINEAGE]->(t)
+        WITH count(DISTINCT src.node_id) AS lin
+        WHERE lin >= 1 AND lin <= $max_lin
+        RETURN lin
+        """,
+        {"nid": node_id, "max_lin": max_lineage_in},
+    )
+    return len(rows) > 0

pretensor/intelligence/statistical.py

@@ -0,0 +1,50 @@
+"""Adjust relationship confidence using sample overlap scores (0.0–1.0)."""
+
+from __future__ import annotations
+
+from pretensor.graph_models.relationship import (
+    RelationshipCandidate,
+    RelationshipSource,
+)
+
+__all__ = ["apply_statistical_adjustment"]
+
+# Weight overlap more than prior confidence when data supports the join.
+_OVERLAP_WEIGHT = 0.65
+_PRIOR_WEIGHT = 0.35
+
+
+def apply_statistical_adjustment(
+    candidate: RelationshipCandidate,
+    overlap_score: float | None,
+) -> RelationshipCandidate:
+    """Blend ``overlap_score`` into ``confidence`` and mark ``source`` as statistical.
+
+    Args:
+        candidate: Incoming hypothesis (heuristic or LLM).
+        overlap_score: Jaccard-like overlap in ``[0, 1]``, or ``None`` to skip.
+
+    Returns:
+        Updated candidate; unchanged if ``overlap_score`` is ``None``.
+    """
+    if overlap_score is None:
+        return candidate
+    o = max(0.0, min(1.0, float(overlap_score)))
+    blended = _PRIOR_WEIGHT * candidate.confidence + _OVERLAP_WEIGHT * o
+    blended = max(0.0, min(1.0, blended))
+    prior = candidate.source
+    note = f"statistical overlap={o:.3f} (prior source={prior})"
+    reasoning = f"{note}. {candidate.reasoning}".strip()
+    return candidate.model_copy(
+        update={
+            "confidence": blended,
+            "source": _statistical_source(candidate.source),
+            "reasoning": reasoning,
+        }
+    )
+
+
+def _statistical_source(prior: RelationshipSource) -> RelationshipSource:
+    """Preserve lineage in reasoning; edge ``source`` is ``statistical`` per contract."""
+    _ = prior
+    return "statistical"

pretensor/intelligence/steps.py

@@ -0,0 +1,191 @@
+"""PipelineStep Protocol and PipelineRunner for the intelligence pipeline.
+
+The runner resolves step execution order from declared dependencies and executes
+each step in sequence.  Plugins can register additional steps (e.g. ``llm_refine``,
+``feedback_score``, ``semantic_propose``) by appending them before calling
+:meth:`PipelineRunner.run`.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Protocol, runtime_checkable
+
+__all__ = [
+    "CyclicDependencyError",
+    "PipelineContext",
+    "PipelineRunner",
+    "PipelineStep",
+]
+# The underscore-prefixed ``_CTX_*`` constants below are intentionally
+# omitted from ``__all__`` — they are private-by-convention names shared
+# across step modules in this package. Importing them by full path is
+# fine for in-package callers; they should not be re-exported as part of
+# the public API surface.
+
+from pretensor.errors import PretensorError
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Canonical ``PipelineContext`` keys — single source of truth.
+# Both :mod:`pretensor.intelligence.pipeline` (the orchestration layer) and
+# every step module that needs to read or write artifacts on the context
+# import from here. Keeping them in one place prevents the silent no-op
+# class of bug where a string diverges between writer and reader.
+# ---------------------------------------------------------------------------
+_CTX_STORE = "store"
+_CTX_DATABASE_KEY = "database_key"
+_CTX_CONFIG = "config"
+_CTX_GRAPH = "graph"
+_CTX_ROLE_BY_TABLE = "role_by_table"
+_CTX_CLUSTERS = "clusters"
+_CTX_PATTERNS = "patterns"
+_CTX_EMBEDDINGS_CONFIG = "embeddings_config"
+_CTX_EMBEDDINGS_PRECOMPUTED = "embeddings_precomputed"
+
+
+class CyclicDependencyError(PretensorError, ValueError):
+    """Raised when the pipeline step dependency graph contains a cycle."""
+
+
+class PipelineContext:
+    """Shared mutable state threaded through every pipeline step.
+
+    Steps read and write named artifacts via :meth:`get` and :meth:`set` so
+    that later steps can consume outputs produced by earlier ones without
+    coupling them directly.
+    """
+
+    def __init__(self, **initial: Any) -> None:
+        self._data: dict[str, Any] = dict(initial)
+
+    def set(self, key: str, value: Any) -> None:  # noqa: A003
+        """Store *value* under *key*."""
+        self._data[key] = value
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Return the value stored under *key*, or *default*."""
+        return self._data.get(key, default)
+
+    def __contains__(self, key: object) -> bool:
+        return key in self._data
+
+
+@runtime_checkable
+class PipelineStep(Protocol):
+    """Structural interface every intelligence pipeline step must satisfy.
+
+    Implementing classes do **not** need to inherit from this class — any
+    object with the right attributes and an ``execute`` coroutine qualifies.
+
+    Attributes:
+        name: Stable, unique identifier for the step (e.g. ``"classify"``).
+        dependencies: Names of steps that must complete before this one runs.
+    """
+
+    name: str
+    dependencies: list[str]
+
+    async def execute(self, ctx: PipelineContext) -> None:
+        """Run this step, reading inputs from and writing outputs to *ctx*.
+
+        Args:
+            ctx: Shared pipeline context populated by upstream steps.
+        """
+        ...
+
+
+class PipelineRunner:
+    """Topological executor for a set of :class:`PipelineStep` objects.
+
+    Steps are registered in any order; the runner resolves execution order
+    from ``step.dependencies`` before the first :meth:`run` call.
+
+    Args:
+        steps: Initial list of steps.  Additional steps can be appended via
+            :meth:`register` before calling :meth:`run`.
+    """
+
+    def __init__(self, steps: list[PipelineStep] | None = None) -> None:
+        self._steps: list[PipelineStep] = list(steps or [])
+
+    def register(self, step: PipelineStep) -> None:
+        """Append *step* to the pipeline.
+
+        Args:
+            step: A :class:`PipelineStep` implementation to add.
+
+        Raises:
+            ValueError: If a step with the same name is already registered.
+        """
+        existing = {s.name for s in self._steps}
+        if step.name in existing:
+            raise ValueError(
+                f"A step named {step.name!r} is already registered in this pipeline."
+            )
+        self._steps.append(step)
+
+    def _ordered(self) -> list[PipelineStep]:
+        """Return steps in topological (dependency-resolved) order.
+
+        Raises:
+            ValueError: If a step lists an unknown dependency name.
+            CyclicDependencyError: If the dependency graph contains a cycle.
+        """
+        by_name = {s.name: s for s in self._steps}
+        for step in self._steps:
+            for dep in step.dependencies:
+                if dep not in by_name:
+                    raise ValueError(
+                        f"Step {step.name!r} declares unknown dependency {dep!r}."
+                    )
+
+        # Kahn's algorithm for topological sort
+        in_degree: dict[str, int] = {name: 0 for name in by_name}
+        dependents: dict[str, list[str]] = {name: [] for name in by_name}
+        for step in self._steps:
+            for dep in step.dependencies:
+                dependents[dep].append(step.name)
+                in_degree[step.name] += 1
+
+        queue = [name for name, deg in in_degree.items() if deg == 0]
+        ordered: list[PipelineStep] = []
+
+        while queue:
+            name = queue.pop(0)
+            ordered.append(by_name[name])
+            for dependent in dependents[name]:
+                in_degree[dependent] -= 1
+                if in_degree[dependent] == 0:
+                    queue.append(dependent)
+
+        if len(ordered) != len(self._steps):
+            cycle_nodes = [n for n, d in in_degree.items() if d > 0]
+            raise CyclicDependencyError(
+                f"Pipeline contains a dependency cycle involving: {cycle_nodes}"
+            )
+
+        return ordered
+
+    async def run(self, ctx: PipelineContext) -> None:
+        """Execute all registered steps in dependency order.
+
+        Args:
+            ctx: Shared pipeline context passed to every step.
+
+        Raises:
+            ValueError: If any step has an unknown dependency.
+            CyclicDependencyError: If the dependency graph contains a cycle.
+        """
+        ordered = self._ordered()
+        logger.info(
+            "PipelineRunner: executing %d step(s): %s",
+            len(ordered),
+            [s.name for s in ordered],
+        )
+        for step in ordered:
+            logger.debug("PipelineRunner: starting step %r", step.name)
+            await step.execute(ctx)
+            logger.debug("PipelineRunner: finished step %r", step.name)