pretensor 0.1.0__py3-none-any.whl

pretensor/cli/plugin.py

@@ -0,0 +1,52 @@
+"""CLI plugin discovery via Python entry_points.
+
+Plugins declare themselves under the ``pretensor.cli_plugins`` entry-point
+group in their ``pyproject.toml``::
+
+    [project.entry-points."pretensor.cli_plugins"]
+    my-plugin = "my_package.cli_plugin:register"
+
+Each entry point must be a callable with the signature::
+
+    def register(app: typer.Typer) -> None: ...
+"""
+
+from __future__ import annotations
+
+import importlib.metadata
+import logging
+
+import typer
+
+__all__ = ["discover_cli_plugins"]
+
+_ENTRY_POINT_GROUP = "pretensor.cli_plugins"
+
+logger = logging.getLogger(__name__)
+
+
+def discover_cli_plugins(app: typer.Typer) -> None:
+    """Scan the ``pretensor.cli_plugins`` entry-point group and register commands.
+
+    Each entry point must resolve to a callable that accepts a single
+    ``typer.Typer`` argument.  Errors during load or registration are logged
+    and skipped so a broken plugin never prevents the CLI from starting.
+
+    Args:
+        app: The root Typer application to register plugin commands into.
+    """
+    eps = importlib.metadata.entry_points(group=_ENTRY_POINT_GROUP)
+    for ep in eps:
+        try:
+            register = ep.load()
+        except Exception:
+            logger.warning(
+                "pretensor: failed to load CLI plugin %r — skipping", ep.name, exc_info=True
+            )
+            continue
+        try:
+            register(app)
+        except Exception:
+            logger.warning(
+                "pretensor: failed to register CLI plugin %r — skipping", ep.name, exc_info=True
+            )

pretensor/config.py

@@ -0,0 +1,226 @@
+"""Graph package configuration (clustering, intelligence defaults)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pretensor.core.store import KuzuStore
+    from pretensor.intelligence.combining import ConfidenceCombiner
+    from pretensor.intelligence.scoring import ScorerRegistry
+    from pretensor.search.base import BaseSearchIndex
+    from pretensor.semantic.base import SemanticLayer
+    from pretensor.semantic.yaml_layer import YamlSemanticLayer
+
+__all__ = [
+    "EmbeddingsConfig",
+    "GraphConfig",
+    "PretensorConfig",
+    "load_semantic_layer",
+]
+
+
+@dataclass(frozen=True, slots=True)
+class EmbeddingsConfig:
+    """Opt-in toggles for the ``[embeddings]`` extra."""
+
+    index_tables: bool = False
+    """When True and ``[embeddings]`` is installed, ``_EmbeddingIndexStep``
+    computes one embedding per :class:`SchemaTable` during intelligence
+    indexing.  Default ``False`` keeps the null path byte-identical to the
+    pre-embedding pipeline output."""
+
+    cluster_blend: float = 0.0
+    """Weight (in [0.0, 1.0]) for blending embedding cosine similarity into
+    existing FK / INFERRED_JOIN edge weights before community detection.
+
+    For each existing structural edge ``(a, b)`` where both endpoints
+    carry an embedding, the edge weight is incremented by
+    ``cluster_blend * cosine(emb_a, emb_b)``.  No new edges are introduced
+    from cosine alone — the blend only nudges already-structurally-adjacent
+    pairs.  Default ``0.0`` keeps cluster membership byte-identical to the
+    pre-embedding pipeline output (Invariant #6).  A small positive value
+    (~0.25) lets domain-related tables with weak lexical overlap drift
+    into the same cluster.
+    """
+
+    join_threshold: float | None = None
+    """Cosine threshold for proposing ``INFERRED_JOIN`` candidates from
+    embedding similarity.  ``None`` (default) disables the
+    embedding scorer entirely → null-path identical.  When set (typical
+    value ~0.85), an ``EmbeddingRelationshipScorer`` is registered after
+    the heuristic scorer; pairs whose cosine ≥ threshold are emitted as
+    ``status="suggested"`` candidates with ``source="embedding"``.  The
+    scorer applies a type-family compatibility gate per column pair
+    (reusing the heuristic's check) so high-cosine pairs with
+    incompatible types are dropped.
+    """
+
+    role_weight: float = 0.0
+    """Weight for the embedding-based role-classification vote.
+    ``0.0`` (default) skips the embedding signal entirely → null-path
+    identical.  When > 0, ``classify_database_tables_async`` embeds each
+    table and adds ``role_weight * cosine(table, role_centroid)`` to
+    every role's heuristic score; the heuristic remains the primary
+    signal.  Centroids are computed once per process from a small curated
+    exemplar set in ``pretensor.intelligence.role_exemplars``.
+    """
+
+    def __post_init__(self) -> None:
+        # Validate the float knobs at construction time so misuse surfaces
+        # close to its source rather than mid-pipeline. All three numeric
+        # toggles are checked here for one consistent contract on the
+        # dataclass; ``EmbeddingRelationshipScorer.__init__`` keeps its own
+        # belt-and-braces check for callers that construct the scorer
+        # directly with an out-of-band threshold.
+        if not 0.0 <= self.cluster_blend <= 1.0:
+            msg = (
+                f"EmbeddingsConfig.cluster_blend must be in [0.0, 1.0]; "
+                f"got {self.cluster_blend!r}"
+            )
+            raise ValueError(msg)
+        if not 0.0 <= self.role_weight <= 1.0:
+            msg = (
+                f"EmbeddingsConfig.role_weight must be in [0.0, 1.0]; "
+                f"got {self.role_weight!r}"
+            )
+            raise ValueError(msg)
+        if self.join_threshold is not None and not (
+            0.0 <= self.join_threshold <= 1.0
+        ):
+            msg = (
+                f"EmbeddingsConfig.join_threshold must be in [0.0, 1.0] "
+                f"or None; got {self.join_threshold!r}"
+            )
+            raise ValueError(msg)
+
+
+@dataclass(frozen=True, slots=True)
+class GraphConfig:
+    """Tunable parameters for graph intelligence (clustering, join paths)."""
+
+    stale_index_warning_days: int = 7
+    """Warn in MCP responses when ``last_indexed`` is older than this many days."""
+
+    clustering_resolution_override: float | None = None
+    """If set, skip automatic resolution heuristic and use this value for Leiden."""
+
+    join_path_max_depth: int = 4
+    """Maximum hop count for precomputed intra-cluster join paths."""
+
+    min_cluster_size_merge: int = 3
+    """Clusters with fewer tables are merged into a neighbor cluster."""
+
+    collapse_shadow_aliases: bool = True
+    """Treat single-source views (``table_type='view'`` with exactly one
+    incoming ``LINEAGE`` edge) as transparent aliases of their base table.
+    When enabled, intelligence consumers (clustering, INFERRED_JOIN, impact,
+    traverse, context) suppress or pass through these aliases."""
+
+    lineage_in_max_for_alias: int = 1
+    """Maximum incoming LINEAGE edge count for a view to be considered a
+    shadow alias.  The default of 1 matches pure 1:1 projections."""
+
+
+def _default_scorer_registry() -> ScorerRegistry:
+    from pretensor.intelligence.heuristic import HeuristicScorer
+    from pretensor.intelligence.scoring import ScorerRegistry
+
+    return ScorerRegistry([HeuristicScorer()])
+
+
+def _default_combiner() -> ConfidenceCombiner:
+    from pretensor.intelligence.combining import MaxScoreCombiner
+
+    return MaxScoreCombiner()
+
+
+def _default_search_index_cls() -> type[BaseSearchIndex]:
+    from pretensor.search.index import KeywordSearchIndex
+
+    return KeywordSearchIndex
+
+
+def _default_semantic_layer() -> SemanticLayer:
+    from pretensor.semantic.base import NullSemanticLayer
+
+    return NullSemanticLayer()
+
+
+@dataclass
+class PretensorConfig:
+    """Central pluggable configuration for the Pretensor graph system.
+
+    Bundles all tunable and swappable components so that downstream packages can ship a
+    ``PretensorConfig`` subclass with its own implementations
+    without patching individual files.
+
+    All fields default to the OSS implementations, so constructing
+    ``PretensorConfig()`` with no arguments is always valid for basic use.
+
+    Attributes:
+        graph: Graph intelligence tuning parameters (clustering, join paths).
+        embeddings: Opt-in toggles for the ``[embeddings]`` extra.
+            Default: ``EmbeddingsConfig()`` (all toggles off).
+        scorer_registry: Ordered registry of relationship scorers used by discovery.
+            Default: ``ScorerRegistry([HeuristicScorer()])``.
+        combiner: Strategy for merging scored relationship candidates.
+            Default: ``MaxScoreCombiner()``.
+        search_index_cls: Class used to build or load the keyword/vector search index.
+            Default: ``KeywordSearchIndex``.
+        semantic_layer: Semantic enrichment layer for entities, metrics, and dimensions.
+            Default: ``NullSemanticLayer()`` (no-op for OSS).
+    """
+
+    graph: GraphConfig = field(default_factory=GraphConfig)
+    embeddings: EmbeddingsConfig = field(default_factory=EmbeddingsConfig)
+    scorer_registry: ScorerRegistry = field(default_factory=_default_scorer_registry)
+    combiner: ConfidenceCombiner = field(default_factory=_default_combiner)
+    search_index_cls: type[BaseSearchIndex] = field(
+        default_factory=_default_search_index_cls
+    )
+    semantic_layer: SemanticLayer = field(default_factory=_default_semantic_layer)
+
+
+def load_semantic_layer(
+    yaml_path: Path,
+    *,
+    store: KuzuStore,
+    database_key: str,
+    dialect: str = "postgres",
+) -> YamlSemanticLayer:
+    """Read a YAML file and return a ready-to-use :class:`YamlSemanticLayer`.
+
+    The YAML must match the
+    :class:`pretensor.introspection.models.semantic.SemanticLayer` schema.
+    The returned layer uses ``store`` for graph lookups (metric compilation,
+    query validation, impact) scoped to ``database_key``.
+
+    Args:
+        yaml_path: Filesystem path to the semantic layer YAML.
+        store: Open Kuzu store for the same ``database_key``.
+        database_key: Logical database key as recorded on ``SchemaTable``.
+        dialect: sqlglot dialect for SQL parsing (default PostgreSQL).
+
+    Returns:
+        A :class:`YamlSemanticLayer` instance.
+
+    Raises:
+        FileNotFoundError: If ``yaml_path`` does not exist.
+        pydantic.ValidationError: If the YAML fails schema validation.
+    """
+    from pretensor.introspection.models.semantic import (
+        SemanticLayer as SemanticLayerModel,
+    )
+    from pretensor.semantic.yaml_layer import YamlSemanticLayer
+
+    text = yaml_path.read_text(encoding="utf-8")
+    model = SemanticLayerModel.from_yaml(text)
+    return YamlSemanticLayer(
+        model,
+        store=store,
+        database_key=database_key,
+        dialect=dialect,
+    )

pretensor/connectors/__init__.py

@@ -0,0 +1,29 @@
+"""Connector-layer schema types, introspection, and snapshot diffing."""
+
+from pretensor.connectors.inspect import inspect
+from pretensor.connectors.models import (
+    Column,
+    ConnectionConfig,
+    ForeignKey,
+    SchemaSnapshot,
+    Table,
+)
+from pretensor.connectors.snapshot import (
+    ChangeTarget,
+    ChangeType,
+    SchemaChange,
+    diff_snapshots,
+)
+
+__all__ = [
+    "ChangeTarget",
+    "ChangeType",
+    "Column",
+    "ConnectionConfig",
+    "ForeignKey",
+    "SchemaChange",
+    "SchemaSnapshot",
+    "Table",
+    "diff_snapshots",
+    "inspect",
+]

pretensor/connectors/base.py

@@ -0,0 +1,165 @@
+"""Abstract connector interface for database introspection.
+
+Every database connector implements BaseConnector. All database-specific
+imports and logic live inside the concrete connector files — nothing
+above the connector layer knows which database produced the data.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from pydantic import Field
+
+from pretensor.connectors.models import ViewDependency
+from pretensor.introspection.models.base import PretensorModel
+from pretensor.introspection.models.config import ConnectionConfig, SchemaFilter
+
+
+class TableInfo(PretensorModel):
+    """Table metadata returned by a connector (no column details)."""
+
+    name: str
+    schema_name: str
+    row_count: int | None = None
+    row_count_source: str | None = Field(
+        default=None,
+        description=(
+            "Provenance of `row_count`: 'stat' (catalog statistic), "
+            "'view_count' (live SELECT COUNT(*) on a view), "
+            "'view_timeout' (view count exceeded the connector timeout — "
+            "row_count = -1)."
+        ),
+    )
+    comment: str | None = None
+    table_type: str | None = None
+
+
+class ColumnInfo(PretensorModel):
+    """Column metadata returned by a connector (no stats)."""
+
+    name: str
+    data_type: str
+    nullable: bool = True
+    is_primary_key: bool = False
+    default_value: str | None = None
+    comment: str | None = None
+    is_indexed: bool = False
+    check_constraints: list[str] = Field(default_factory=list)
+    ordinal_position: int | None = None
+    column_cardinality: int | None = None
+    index_type: str | None = None
+    index_is_unique: bool | None = None
+    parent_column: str | None = None
+    is_array: bool = False
+
+
+class ForeignKeyInfo(PretensorModel):
+    """Foreign key relationship returned by a connector."""
+
+    constraint_name: str | None = None
+    source_schema: str
+    source_table: str
+    source_column: str
+    target_schema: str
+    target_table: str
+    target_column: str
+
+
+class ColumnStats(PretensorModel):
+    """Statistical metadata for a single column."""
+
+    distinct_count: int | None = None
+    min_value: str | None = None
+    max_value: str | None = None
+    null_percentage: float | None = None
+    sample_distinct_values: list[str] | None = None
+
+
+class TableGrant(PretensorModel):
+    """SELECT grant on a table for an effective grantee (role or login)."""
+
+    grantee: str
+    schema_name: str
+    table_name: str
+
+
+class BaseConnector(ABC):
+    """Abstract interface that every database connector must implement.
+
+    Usage as a context manager is preferred::
+
+        with PostgresConnector(config) as conn:
+            tables = conn.get_tables()
+    """
+
+    def __init__(self, config: ConnectionConfig) -> None:
+        self.config = config
+
+    def __enter__(self) -> BaseConnector:
+        self.connect()
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.disconnect()
+
+    @abstractmethod
+    def connect(self) -> None: ...
+
+    @abstractmethod
+    def disconnect(self) -> None: ...
+
+    @abstractmethod
+    def get_tables(
+        self, schema_filter: SchemaFilter | None = None
+    ) -> list[TableInfo]: ...
+
+    @abstractmethod
+    def get_columns(self, table_name: str, schema_name: str) -> list[ColumnInfo]: ...
+
+    @abstractmethod
+    def get_foreign_keys(self) -> list[ForeignKeyInfo]: ...
+
+    @abstractmethod
+    def get_table_row_count(self, table_name: str, schema_name: str) -> int: ...
+
+    @abstractmethod
+    def get_column_stats(
+        self, table_name: str, column_name: str, schema_name: str
+    ) -> ColumnStats: ...
+
+    @abstractmethod
+    def execute_query(self, sql: str) -> list[dict[str, Any]]: ...
+
+    def load_deep_catalog(
+        self, schema_filter: SchemaFilter
+    ) -> tuple[
+        dict[tuple[str, str], dict[str, Any]],
+        dict[tuple[str, str, str], dict[str, Any]],
+    ]:
+        """Optional bulk catalog enrichment (usage, stats, grants). Default: none.
+
+        Callers treat this as best-effort: failures must not block indexing (log and
+        continue without enrichment).
+        """
+        _ = schema_filter
+        return {}, {}
+
+    def load_view_dependencies(self, schema_filter: SchemaFilter) -> list[ViewDependency]:
+        """Table-level lineage from views, triggers, or warehouse objects. Default: none."""
+        _ = schema_filter
+        return []
+
+    def get_table_grants(
+        self, schema_filter: SchemaFilter | None = None
+    ) -> list[TableGrant]:
+        """SELECT table grants, including connector-specific expansion (e.g. role inheritance).
+
+        Default: no grant introspection (empty list).
+
+        Callers treat this as best-effort: failures must not block indexing (log and
+        continue without grants), same contract as :meth:`load_deep_catalog`.
+        """
+        _ = schema_filter
+        return []