pretensor 0.1.0__py3-none-any.whl

pretensor/intelligence/steps_embedding.py

@@ -0,0 +1,168 @@
+"""Index-time embedding computation for SchemaTable nodes.
+
+Opt-in via ``PretensorConfig.embeddings.index_tables``.  When the toggle is
+off (default) nothing runs — zero reads, zero writes — so the null path
+stays byte-identical to the pre-embedding pipeline output.
+
+:func:`compute_table_embeddings` is the single computation entry point.
+Index and reindex call it right after schema rows are written and *before*
+relationship discovery, so the embedding relationship scorer, the
+clustering blend, and the classify step's role vote all consume vectors
+from the same run — no one-cycle staging lag.  :class:`EmbeddingIndexStep`
+remains in the pipeline for callers that invoke
+``run_intelligence_layer`` directly; it delegates to the same function and
+skips when the context marks vectors as already computed this run.
+
+Every entry point honors the ``PRETENSOR_EMBEDDINGS_DISABLED`` environment
+variable: when set to a truthy value (``1``, ``true``, ``yes`` —
+case-insensitive), the embedding pass is forced to the null path even when
+the ``[embeddings]`` extra is installed and ``index_tables`` is True.  CI
+exercises this to assert that "extra installed but user opted out"
+produces output byte-identical to "extra absent" — the determinism
+contract behind the release gate.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from pretensor.intelligence.embeddings import (
+    embeddings_disabled_via_env,
+    format_entity_text,
+    get_default_embedding_client,
+)
+
+if TYPE_CHECKING:
+    from pretensor.config import EmbeddingsConfig
+    from pretensor.core.store import KuzuStore
+    from pretensor.intelligence.steps import PipelineContext
+
+# ``embeddings_disabled_via_env`` is canonically defined in
+# ``embeddings.py``; it is imported above for this module's own use but is
+# deliberately NOT re-exported here, so there is one authoritative import
+# path for every caller.
+__all__ = [
+    "EmbeddingIndexStep",
+    "compute_table_embeddings",
+]
+
+logger = logging.getLogger(__name__)
+
+
+# Canonical context keys live in ``steps.py`` (also home to PipelineStep
+# / PipelineRunner / PipelineContext). Importing them keeps writer and
+# reader on one source of truth — drift between sibling string literals
+# would silently turn this step into a no-op.
+from pretensor.intelligence.steps import (  # noqa: E402  (post-imports to keep the section clear)
+    _CTX_DATABASE_KEY,
+    _CTX_EMBEDDINGS_CONFIG,
+    _CTX_EMBEDDINGS_PRECOMPUTED,
+    _CTX_STORE,
+)
+
+
+def compute_table_embeddings(store: KuzuStore, database_key: str) -> None:
+    """Compute one embedding per ``SchemaTable`` and persist it to Kuzu.
+
+    Callers gate on ``EmbeddingsConfig.index_tables`` and
+    :func:`embeddings_disabled_via_env` — this function assumes the user
+    opted in.  When the ``[embeddings]`` extra is absent it logs once at
+    ``WARNING`` and returns without raising.  Per-table write failures are
+    tolerated: the offending table is left with ``embedding=None`` and the
+    rest of the pass continues.
+    """
+    rows = store.query_all_rows(
+        """
+        MATCH (t:SchemaTable {database: $db})
+        OPTIONAL MATCH (t)-[:HAS_COLUMN]->(c:SchemaColumn)
+        RETURN t.node_id, t.schema_name, t.table_name, collect(c.column_name)
+        """,
+        {"db": database_key},
+    )
+    if not rows:
+        return
+
+    client = get_default_embedding_client()
+
+    # Build the per-table input texts and a parallel ``node_id`` list so
+    # we can issue ONE batched ``embed()`` call instead of N. ONNX
+    # inference is dominated by per-call overhead (model warm-up,
+    # tensor allocation), so a single 50-row batch beats 50 single-row
+    # calls by a wide margin on every realistic schema size.
+    nids: list[str] = []
+    texts: list[str] = []
+    for row in rows:
+        nids.append(str(row[0]))
+        qualified_name = f"{row[1]}.{row[2]}"
+        column_names = [str(c) for c in (row[3] or []) if c is not None]
+        texts.append(format_entity_text(qualified_name, column_names))
+
+    try:
+        vectors = client.embed(texts)
+    except ImportError as exc:
+        logger.warning(
+            "embedding_index: skipping, [embeddings] extra not installed: %s",
+            exc,
+        )
+        return
+    except Exception as exc:  # noqa: BLE001 — batch failure is tolerated
+        logger.warning(
+            "embedding_index: batched embed failed (%s); leaving every "
+            "table's embedding=None",
+            exc,
+        )
+        return
+
+    if not vectors:
+        return
+    if len(vectors) != len(texts):
+        logger.warning(
+            "embedding_index: client returned %d vectors for %d texts; "
+            "skipping write to avoid mis-aligned assignment",
+            len(vectors),
+            len(texts),
+        )
+        return
+
+    for nid, vec in zip(nids, vectors, strict=True):
+        try:
+            store.set_table_embedding(nid, vec)
+        except ValueError as exc:
+            logger.warning("embedding_index: invalid vector for %s: %s", nid, exc)
+
+
+class EmbeddingIndexStep:
+    """Pipeline wrapper around :func:`compute_table_embeddings`.
+
+    Kept for callers that drive ``run_intelligence_layer`` directly (and
+    for plugins that anchor on the ``embedding_index`` step name).  The
+    index/reindex CLI paths compute embeddings *before* relationship
+    discovery and mark the context with ``embeddings_precomputed=True``,
+    in which case this step is a no-op.
+
+    Setting ``PRETENSOR_EMBEDDINGS_DISABLED=1`` forces the null path even
+    when ``index_tables`` is True; used by CI to enforce null-path parity.
+    """
+
+    name = "embedding_index"
+    # No dependencies: the computation reads only SchemaTable + HAS_COLUMN
+    # rows, which exist before the intelligence pipeline starts.  Running
+    # first lets the classify step's role vote consume same-run vectors
+    # instead of being one index cycle behind.
+    dependencies: list[str] = []
+
+    async def execute(self, ctx: PipelineContext) -> None:
+        emb_cfg: EmbeddingsConfig | None = ctx.get(_CTX_EMBEDDINGS_CONFIG)
+        if emb_cfg is None or not emb_cfg.index_tables:
+            return
+        if embeddings_disabled_via_env():
+            logger.info("embedding_index: skipping (PRETENSOR_EMBEDDINGS_DISABLED set)")
+            return
+        if ctx.get(_CTX_EMBEDDINGS_PRECOMPUTED, False):
+            logger.debug("embedding_index: vectors already computed this run")
+            return
+
+        store: KuzuStore = ctx.get(_CTX_STORE)
+        database_key: str = ctx.get(_CTX_DATABASE_KEY)
+        compute_table_embeddings(store, database_key)

pretensor/introspection/__init__.py

@@ -0,0 +1,6 @@
+"""Schema introspection: project config, DSN parsing, and compatibility shims.
+
+Schema snapshot DTOs, PostgreSQL introspection, and ``diff_snapshots`` live in
+``pretensor.connectors``. This package re-exports ``inspector.inspect``,
+``snapshot`` helpers, and ``models.schema`` for backward compatibility.
+"""

pretensor/introspection/inspector.py

@@ -0,0 +1,5 @@
+"""Schema inspector — re-exported from ``pretensor.connectors.inspect``."""
+
+from pretensor.connectors.inspect import inspect
+
+__all__ = ["inspect"]

pretensor/introspection/models/base.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class PretensorModel(BaseModel):
+    model_config = ConfigDict(frozen=True)

pretensor/introspection/models/config.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+import os
+import re
+from enum import StrEnum
+from pathlib import Path
+from typing import Any, Literal
+
+from pydantic import ConfigDict, Field, model_validator
+from ruamel.yaml import YAML
+
+from .base import PretensorModel
+
+ENV_VAR_PATTERN = re.compile(r"\$\{([^}]+)\}")
+
+
+def _resolve_env_vars(value: str) -> str:
+    """Replace ${VAR_NAME} references with their environment variable values."""
+
+    def _replace(match: re.Match[str]) -> str:
+        var_name = match.group(1)
+        env_value = os.environ.get(var_name)
+        if env_value is None:
+            raise ValueError(f"Environment variable '{var_name}' is not set")
+        return env_value
+
+    return ENV_VAR_PATTERN.sub(_replace, value)
+
+
+class DatabaseType(StrEnum):
+    POSTGRES = "postgres"
+    BIGQUERY = "bigquery"
+    SNOWFLAKE = "snowflake"
+    DUCKDB = "duckdb"
+    MYSQL = "mysql"
+
+
+class PrivacyMode(StrEnum):
+    SCHEMA_ONLY = "schema_only"
+    SYNTHETIC_SAMPLES = "synthetic_samples"
+    REAL_SAMPLES = "real_samples"
+
+
+class SchemaFilter(PretensorModel):
+    include: list[str] = Field(default_factory=list)
+    exclude: list[str] = Field(default_factory=list)
+
+
+class ConnectionConfig(PretensorModel):
+    # Mutable config (env resolution, optional connector-specific fields).
+    model_config = ConfigDict(frozen=False)
+
+    name: str
+    type: DatabaseType
+    host: str | None = None
+    port: int | str | None = None
+    database: str | None = None
+    user: str | None = None
+    password: str | None = None
+    schema_filter: SchemaFilter = Field(default_factory=SchemaFilter)
+    metadata_extra: dict[str, Any] = Field(default_factory=dict)
+
+    @model_validator(mode="after")
+    def resolve_env_vars(self) -> ConnectionConfig:
+        for field_name in self.__class__.model_fields:
+            raw = getattr(self, field_name)
+            if isinstance(raw, str) and ENV_VAR_PATTERN.search(raw):
+                object.__setattr__(self, field_name, _resolve_env_vars(raw))
+        # Also resolve env vars inside metadata_extra string values.
+        for key, val in self.metadata_extra.items():
+            if isinstance(val, str) and ENV_VAR_PATTERN.search(val):
+                self.metadata_extra[key] = _resolve_env_vars(val)
+        # Coerce port back to int after env var substitution (${PORT} → "5432").
+        raw_port = getattr(self, "port")
+        if isinstance(raw_port, str):
+            object.__setattr__(self, "port", int(raw_port))
+        return self
+
+
+class LLMProvider(StrEnum):
+    """Registered LLM backends (see `LLMConfig.provider`)."""
+
+    ANTHROPIC = "anthropic"
+    OPENROUTER = "openrouter"
+
+
+class LLMStepKey(StrEnum):
+    """Pipeline step identifiers for per-step LLM overrides (`llm_steps`)."""
+
+    SEMANTIC = "semantic"
+    PLANNER = "planner"
+    CODEGEN_QUERIES = "codegen_queries"
+    CODEGEN_PAGES = "codegen_pages"
+    VALIDATION = "validation"
+    REPAIR = "repair"
+
+
+class StepLLMOverride(PretensorModel):
+    """Partial LLM settings merged over the global `llm` block for one step.
+
+    Omitted fields keep the global default. ``http_headers`` replaces the
+    global header map when set (non-``None``).
+    """
+
+    provider: LLMProvider | None = None
+    model: str | None = None
+    temperature: float | None = None
+    max_tokens: int | None = None
+    base_url: str | None = None
+    http_headers: dict[str, str] | None = None
+
+
+class LLMConfig(PretensorModel):
+    provider: LLMProvider = LLMProvider.ANTHROPIC
+    model: str = "claude-sonnet-4-20250514"
+    temperature: float = 0.0
+    max_tokens: int = 4096
+    base_url: str | None = None
+    http_headers: dict[str, str] = Field(default_factory=dict)
+
+    def merged_with(self, override: StepLLMOverride | None) -> LLMConfig:
+        """Return a copy with non-``None`` fields from ``override`` applied."""
+        if override is None:
+            return self.model_copy(deep=True)
+        data = self.model_dump()
+        if override.provider is not None:
+            data["provider"] = override.provider
+        if override.model is not None:
+            data["model"] = override.model
+        if override.temperature is not None:
+            data["temperature"] = override.temperature
+        if override.max_tokens is not None:
+            data["max_tokens"] = override.max_tokens
+        if override.base_url is not None:
+            data["base_url"] = override.base_url
+        if override.http_headers is not None:
+            data["http_headers"] = dict(override.http_headers)
+        return LLMConfig.model_validate(data)
+
+
+class SyntheticOptions(PretensorModel):
+    hash_strings: bool = True
+    randomize_numbers: bool = True
+    preserve_magnitude: bool = True
+
+
+class PrivacyConfig(PretensorModel):
+    mode: PrivacyMode = PrivacyMode.SCHEMA_ONLY
+    synthetic: SyntheticOptions = Field(default_factory=SyntheticOptions)
+
+
+class MCPConfig(PretensorModel):
+    enabled: bool = False
+    url: str = "https://evidence.studio/mcp"
+
+
+class EvidenceConfig(PretensorModel):
+    mcp: MCPConfig = Field(default_factory=MCPConfig)
+
+
+class RepairConfig(PretensorModel):
+    max_code_retries: int = 3
+    max_plan_retries: int = 2
+    max_total_iterations: int = 10
+
+
+class StagingConfig(PretensorModel):
+    annotations_enabled: bool = True
+    staging_env_var: str = "PRETENSOR_STAGING"
+
+
+class ValidationConfig(PretensorModel):
+    threshold: int = 7
+    top_components_check: int = 2
+    repair: RepairConfig = Field(default_factory=RepairConfig)
+    staging: StagingConfig = Field(default_factory=StagingConfig)
+    # Layer 1 — skip LLM self-review when SQL lint, page markdown, and Evidence
+    # build (if run) produced no ERROR issues; still run when any deterministic
+    # step reports ERROR.
+    l1_skip_llm_self_review_on_clean_deterministic: bool = True
+    # Layer 3 — deterministic structure vs plan + optional presentation hints.
+    l3_structural_checks_enabled: bool = True
+    l3_orphan_chart_severity: Literal["error", "warning"] = "error"
+    l3_structural_presentation_warnings: bool = True
+
+
+class ProjectConfig(PretensorModel):
+    project_name: str | None = None
+    connections: list[ConnectionConfig] | dict[str, ConnectionConfig] = Field(
+        default_factory=list
+    )
+    llm: LLMConfig = Field(default_factory=LLMConfig)
+    llm_steps: dict[LLMStepKey, StepLLMOverride] = Field(default_factory=dict)
+    privacy: PrivacyConfig = Field(default_factory=PrivacyConfig)
+    evidence: EvidenceConfig = Field(default_factory=EvidenceConfig)
+    validation: ValidationConfig = Field(default_factory=ValidationConfig)
+
+    @model_validator(mode="before")
+    @classmethod
+    def _preprocess_config(cls, data: dict) -> dict:
+        """Handle legacy and shorthand config formats."""
+        if not isinstance(data, dict):
+            return data
+
+        # Infer project_name if missing
+        if not data.get("project_name"):
+            data["project_name"] = Path.cwd().name
+
+        # Handle connections as dict
+        connections = data.get("connections")
+        if isinstance(connections, dict):
+            connection_list = []
+            for name, config in connections.items():
+                if isinstance(config, dict):
+                    config["name"] = name
+                    connection_list.append(config)
+            data["connections"] = connection_list
+
+        return data
+
+    def resolved_llm_for_step(self, step: LLMStepKey) -> LLMConfig:
+        """Global ``llm`` merged with ``llm_steps[step]`` when present."""
+        override = self.llm_steps.get(step)
+        return self.llm.merged_with(override)
+
+    @classmethod
+    def from_yaml(cls, yaml_str: str) -> ProjectConfig:
+        yaml = YAML()
+        data = yaml.load(yaml_str)
+        return cls.model_validate(data)
+
+    @classmethod
+    def load(cls, path: Path) -> ProjectConfig:
+        """Load project config from a YAML file on disk."""
+        if not path.exists():
+            raise FileNotFoundError(f"Config file not found: {path}")
+        return cls.from_yaml(path.read_text())