pretensor 0.1.0__py3-none-any.whl

pretensor/introspection/models/semantic.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from enum import StrEnum
+from io import StringIO
+from typing import Literal
+
+from pydantic import Field
+from ruamel.yaml import YAML
+
+from .base import PretensorModel
+
+
+class AttributeRole(StrEnum):
+    IDENTIFIER = "identifier"
+    DIMENSION = "dimension"
+    MEASURE = "measure"
+    TIME_DIMENSION = "time_dimension"
+
+
+class MetricType(StrEnum):
+    COUNT = "count"
+    COUNT_DISTINCT = "count_distinct"
+    SUM = "sum"
+    AVERAGE = "average"
+    DERIVED = "derived"
+
+
+class FormatConfig(PretensorModel):
+    """Per-attribute or per-metric format override."""
+
+    format_type: Literal["number", "currency", "percentage", "date"] | None = None
+    decimals: int | None = None
+    symbol: str | None = None
+    position: Literal["prefix", "suffix"] | None = None
+    suffix: str | None = None
+    thousands_separator: str | None = None
+    format_string: str | None = None
+
+
+class Attribute(PretensorModel):
+    name: str
+    description: str
+    role: AttributeRole
+    source_column: str
+    format_config: FormatConfig | None = None
+
+
+class Metric(PretensorModel):
+    name: str
+    description: str
+    type: MetricType
+    field: str | None = None
+    expression: str | None = None
+    time_dimension: str | None = None
+    granularity: str | None = None
+    format_config: FormatConfig | None = None
+    filters: list[str] = Field(default_factory=list)
+
+
+class Entity(PretensorModel):
+    name: str
+    description: str
+    source_table: str
+    attributes: list[Attribute] = Field(default_factory=list)
+    metrics: list[Metric] = Field(default_factory=list)
+
+
+class Domain(PretensorModel):
+    name: str
+    description: str
+    entities: list[Entity] = Field(default_factory=list)
+
+
+class NumberFormatDefaults(PretensorModel):
+    decimals: int = 0
+    thousands_separator: str = ","
+
+
+class CurrencyFormatDefaults(PretensorModel):
+    symbol: str = "$"
+    position: Literal["prefix", "suffix"] = "prefix"
+    decimals: int = 2
+
+
+class PercentageFormatDefaults(PretensorModel):
+    decimals: int = 1
+    suffix: str = "%"
+
+
+class DateFormatDefaults(PretensorModel):
+    format_string: str = "YYYY-MM-DD"
+
+
+class FormatDefaults(PretensorModel):
+    number: NumberFormatDefaults = Field(default_factory=NumberFormatDefaults)
+    currency: CurrencyFormatDefaults = Field(default_factory=CurrencyFormatDefaults)
+    percentage: PercentageFormatDefaults = Field(
+        default_factory=PercentageFormatDefaults
+    )
+    date: DateFormatDefaults = Field(default_factory=DateFormatDefaults)
+    chart_preferences: dict[str, str] = Field(default_factory=dict)
+
+
+class SemanticLayer(PretensorModel):
+    connection_name: str
+    domains: list[Domain] = Field(default_factory=list)
+    format_defaults: FormatDefaults = Field(default_factory=FormatDefaults)
+
+    def to_yaml(self) -> str:
+        data = self.model_dump(mode="json")
+        yaml = YAML()
+        yaml.default_flow_style = False
+        buf = StringIO()
+        yaml.dump(data, buf)
+        return buf.getvalue()
+
+    @classmethod
+    def from_yaml(cls, yaml_str: str) -> SemanticLayer:
+        yaml = YAML()
+        data = yaml.load(yaml_str)
+        return cls.model_validate(data)

pretensor/introspection/models/validation.py

@@ -0,0 +1,116 @@
+"""Validation report models — shared across all validation layers."""
+
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import Literal
+
+from pydantic import Field
+
+from .base import PretensorModel
+
+__all__ = [
+    "Severity",
+    "RepairState",
+    "ValidationIssue",
+    "ComponentScore",
+    "LayerResult",
+    "AnalyticalQualityResult",
+    "ValidationReport",
+    "LayerStatus",
+]
+
+LayerStatus = Literal["pass", "fail", "repaired", "not_evaluated"]
+
+
+class Severity(StrEnum):
+    ERROR = "error"
+    WARNING = "warning"
+    INFO = "info"
+
+
+class RepairState(StrEnum):
+    PASS = "pass"
+    FAIL = "fail"
+    REPAIRED = "repaired"
+    NOT_EVALUATED = "not_evaluated"
+
+
+class ValidationIssue(PretensorModel):
+    model_config = PretensorModel.model_config.copy()
+    model_config["frozen"] = False
+
+    severity: Severity
+    layer: str
+    path: str
+    message: str
+
+
+class LayerResult(PretensorModel):
+    model_config = PretensorModel.model_config.copy()
+    model_config["frozen"] = False
+
+    layer_name: str
+    status: LayerStatus
+    repair_state: RepairState = RepairState.PASS
+    issues: list[ValidationIssue] = Field(default_factory=list)
+
+    @property
+    def error_count(self) -> int:
+        return sum(1 for i in self.issues if i.severity == Severity.ERROR)
+
+    @property
+    def warning_count(self) -> int:
+        return sum(1 for i in self.issues if i.severity == Severity.WARNING)
+
+
+class ComponentScore(PretensorModel):
+    """Per-chart quality score produced by the L4 analyst review."""
+
+    model_config = PretensorModel.model_config.copy()
+    model_config["frozen"] = False
+
+    score: int
+    answers_main_question: bool
+    priority_rank: int
+    current_position: int
+    note: str
+    subtracting: str | None = None
+    suggested_action: str | None = None
+
+
+class AnalyticalQualityResult(LayerResult):
+    """Extended LayerResult for L4, carrying per-chart ComponentScore breakdown."""
+
+    model_config = LayerResult.model_config.copy()
+
+    overall_score: int = 0
+    overall_summary: str = ""
+    question_answered: bool = False
+    components: dict[str, ComponentScore] = Field(default_factory=dict)
+    repair_tier: int = 1
+
+
+class ValidationReport(PretensorModel):
+    model_config = PretensorModel.model_config.copy()
+    model_config["frozen"] = False
+
+    dashboard_name: str
+    layers: list[LayerResult] = Field(default_factory=list)
+    repair_actions_taken: list[str] = Field(default_factory=list)
+
+    @property
+    def passed(self) -> bool:
+        return all(r.status == "pass" for r in self.layers)
+
+    @property
+    def all_issues(self) -> list[ValidationIssue]:
+        return [issue for layer in self.layers for issue in layer.issues]
+
+    @property
+    def error_count(self) -> int:
+        return sum(r.error_count for r in self.layers)
+
+    @property
+    def warning_count(self) -> int:
+        return sum(r.warning_count for r in self.layers)

pretensor/introspection/snapshot.py

@@ -0,0 +1,46 @@
+"""Snapshot diff — re-exported from ``pretensor.connectors.snapshot``."""
+
+from pathlib import Path
+
+from pretensor.connectors.models import SchemaSnapshot
+from pretensor.connectors.snapshot import (
+    ChangeTarget,
+    ChangeType,
+    SchemaChange,
+    diff_snapshots,
+)
+
+
+def save_snapshot(snapshot: SchemaSnapshot, base_path: Path) -> Path:
+    """Serialize a snapshot to YAML and write it to disk.
+
+    Writes to ``{base_path}/{connection_name}.yaml``.
+
+    Returns:
+        The path to the written file.
+    """
+    base_path.mkdir(parents=True, exist_ok=True)
+    file_path = base_path / f"{snapshot.connection_name}.yaml"
+    file_path.write_text(snapshot.to_yaml())
+    return file_path
+
+
+def load_snapshot(path: Path) -> SchemaSnapshot:
+    """Load a schema snapshot from a YAML file.
+
+    Raises:
+        FileNotFoundError: If the file does not exist.
+    """
+    if not path.exists():
+        raise FileNotFoundError(f"Snapshot file not found: {path}")
+    return SchemaSnapshot.from_yaml(path.read_text())
+
+
+__all__ = [
+    "ChangeTarget",
+    "ChangeType",
+    "SchemaChange",
+    "diff_snapshots",
+    "load_snapshot",
+    "save_snapshot",
+]

pretensor/mcp/__init__.py

@@ -0,0 +1,16 @@
+"""MCP server and service helpers for Pretensor Graph."""
+
+from pretensor.mcp.server import (
+    create_server,
+    print_mcp_config,
+    run_server,
+)
+from pretensor.mcp.tool_registry import McpTool, McpToolRegistry
+
+__all__ = [
+    "create_server",
+    "print_mcp_config",
+    "run_server",
+    "McpTool",
+    "McpToolRegistry",
+]

pretensor/mcp/config_json.py

@@ -0,0 +1,24 @@
+"""MCP client config snippet (Cursor / Claude Desktop)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+
+def mcp_config_json(graph_dir: Path) -> str:
+    """Claude Desktop / Cursor style ``mcpServers`` block for stdio transport."""
+    cmd = "pretensor"
+    graph_dir_resolved = str(graph_dir.resolve())
+    block = {
+        "mcpServers": {
+            "pretensor": {
+                "command": cmd,
+                "args": ["serve", "--state-dir", graph_dir_resolved],
+            }
+        }
+    }
+    return json.dumps(block, indent=2)
+
+
+__all__ = ["mcp_config_json"]

pretensor/mcp/payload_types.py

@@ -0,0 +1,274 @@
+"""Typed payloads and small helpers for MCP tool responses."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any, TypedDict
+
+from pretensor.config import GraphConfig
+from pretensor.mcp.service_context import get_effective_graph_config
+
+DB_TYPE_POSTGRES = "postgresql"
+
+DIALECT_TO_DB_TYPE: dict[str, str] = {
+    "postgres": "postgresql",
+    "mysql": "mysql",
+    "snowflake": "snowflake",
+    "bigquery": "bigquery",
+}
+
+
+class DatabaseListItem(TypedDict, total=False):
+    """One row for ``list_databases`` and the databases resource."""
+
+    name: str
+    db_type: str
+    table_count: int
+    column_count: int
+    row_count: int
+    last_indexed: str
+    is_stale: bool
+    staleness_days: int
+    stale_warning: str
+    graph_path: str
+    database: str
+    schemas: list[str]
+    # Trinary capability state: "not_attempted" | "empty" | "present".
+    # "not_attempted" means the indexer never ran this pass, so absence of signal
+    # is not authoritative. "empty" means the pass ran but found nothing.
+    has_dbt_manifest: str
+    has_llm_enrichment: str
+    has_external_consumers: str
+    # Unified capability list derived from the trinary flags plus always-on
+    # features. Additive convenience over the individual ``has_*`` fields;
+    # scan this to decide whether a graph supports a feature at a glance.
+    capabilities: list[str]
+
+
+class QueryHit(TypedDict, total=False):
+    """One search hit for the ``query`` tool."""
+
+    node_type: str
+    name: str
+    database_name: str
+    connection_name: str
+    description: str
+    snippet: str
+    score: float
+
+
+class SemanticHit(TypedDict, total=False):
+    """One search hit for the ``semantic_search`` tool (cosine over ``SchemaTable.embedding``)."""
+
+    node_type: str
+    name: str
+    database_name: str
+    connection_name: str
+    description: str
+    snippet: str
+    score: float
+    cluster_id: str | None
+
+
+class ColumnInfo(TypedDict, total=False):
+    """Column metadata for ``context`` (from ``SchemaColumn`` nodes when present).
+
+    ``column_name`` is the canonical field.
+    """
+
+    column_name: str
+    data_type: str
+    nullable: bool
+    is_primary_key: bool
+    is_foreign_key: bool
+    description: str
+    is_indexed: bool
+    check_constraints: list[str]
+    most_common_values: list[str]
+    histogram_bounds: list[str]
+    stats_correlation: float
+
+
+class RelationshipInfo(TypedDict, total=False):
+    """FK or inferred join from a table to another table.
+
+    For composite FKs, ``source_columns`` / ``target_columns`` carry the
+    ordered column lists and ``constraint_name`` identifies the constraint.
+    Single-column FKs still set ``source_column`` / ``target_column`` for
+    backward compatibility.
+    """
+
+    target_table: str
+    rel_type: str
+    source_column: str
+    target_column: str
+    source_columns: list[str]
+    target_columns: list[str]
+    constraint_name: str | None
+    source: str | None
+    confidence: float | None
+    reasoning: str | None
+
+
+class LineageRef(TypedDict, total=False):
+    """One table-level LINEAGE hop for ``context``."""
+
+    table: str
+    lineage_type: str
+    confidence: float
+    source: str
+
+
+class ClusterInfo(TypedDict, total=False):
+    """Domain cluster assignment for ``context``."""
+
+    cluster_id: str
+    label: str
+    description: str
+    cohesion_score: float
+    schema_pattern: str
+    stale: bool
+    stale_warning: str
+
+
+class SimilarTable(TypedDict, total=False):
+    """One cross-cluster nearest neighbor in ``context.similar_tables``."""
+
+    table_id: str
+    qualified_name: str
+    score: float
+    cluster_id: str | None
+
+
+class ContextPayload(TypedDict, total=False):
+    """Structured ``context`` tool response."""
+
+    connection_name: str
+    database: str
+    schema_name: str
+    table_name: str
+    table_type: str | None
+    qualified_name: str
+    description: str
+    tags: list[str]
+    has_external_consumers: bool
+    test_count: int
+    staleness_status: str
+    staleness_as_of: str
+    row_count: int | None
+    entity_type: str | None
+    entity_name: str | None
+    entity_description: str | None
+    role: str | None
+    role_confidence: float | None
+    classification_signals: list[str]
+    classification_summary: str
+    cluster: ClusterInfo | None
+    columns: list[ColumnInfo]
+    relationships: list[RelationshipInfo]
+    lineage_in: list[LineageRef]
+    lineage_out: list[LineageRef]
+    lineage_markdown: str
+    aliases: list[str]
+    shadow_of: str | None
+    deprecation_signal: str | None
+    staleness_warning: str | None
+    detail: str
+    usage_stats: dict[str, int]
+    partition: dict[str, object]
+    grants: list[dict[str, str]]
+    access_patterns: dict[str, object]
+    similar_tables: list[SimilarTable]
+    similar_reason: str
+
+
+class TraverseStepPayload(TypedDict, total=False):
+    """One hop in ``traverse`` output."""
+
+    from_table: str
+    to_table: str
+    from_column: str
+    to_column: str
+    edge_type: str
+    from_db: str
+    to_db: str
+    via: str
+    join_columns: str | None
+    constraint_name: str | None
+
+
+class TraversePathPayload(TypedDict, total=False):
+    """One ranked path in ``traverse`` output."""
+
+    confidence: float
+    ambiguous: bool
+    semantic_label: str
+    steps: list[TraverseStepPayload]
+    sql_hint: str
+    stale: bool
+    stale_warning: str
+
+
+class ImpactItemPayload(TypedDict, total=False):
+    """Single dependent in ``impact`` output."""
+
+    type: str
+    name: str
+    via: str
+    confidence: float
+    hop: int
+    # Additive: external code consumers of this table (empty when none). Old
+    # callers that ignore this key are unaffected.
+    consumers: list[dict[str, Any]]
+
+
+def utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def iso_format(dt: datetime) -> str:
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+    return dt.astimezone(timezone.utc).isoformat()
+
+
+def stale_threshold_days(config: GraphConfig | None = None) -> int:
+    return get_effective_graph_config(config).stale_index_warning_days
+
+
+def staleness_days(indexed_at: datetime, now: datetime | None = None) -> int:
+    reference = now or utc_now()
+    if indexed_at.tzinfo is None:
+        indexed_at = indexed_at.replace(tzinfo=timezone.utc)
+    delta = reference - indexed_at.astimezone(timezone.utc)
+    return max(0, int(delta.total_seconds() // 86400))
+
+
+def snippet(text: str, max_len: int = 200) -> str:
+    t = text.strip()
+    if len(t) <= max_len:
+        return t
+    return t[: max_len - 3] + "..."
+
+
+__all__ = [
+    "ClusterInfo",
+    "ColumnInfo",
+    "ContextPayload",
+    "DatabaseListItem",
+    "ImpactItemPayload",
+    "LineageRef",
+    "QueryHit",
+    "RelationshipInfo",
+    "SemanticHit",
+    "SimilarTable",
+    "TraversePathPayload",
+    "TraverseStepPayload",
+    "DB_TYPE_POSTGRES",
+    "DIALECT_TO_DB_TYPE",
+    "iso_format",
+    "snippet",
+    "staleness_days",
+    "stale_threshold_days",
+    "utc_now",
+]

pretensor/mcp/resources/__init__.py

@@ -0,0 +1,17 @@
+"""MCP resource payloads (markdown)."""
+
+from __future__ import annotations
+
+from pretensor.mcp.resources.markdown import (
+    clusters_resource_markdown,
+    databases_resource_markdown,
+    db_overview_resource_markdown,
+    metrics_resource_markdown,
+)
+
+__all__ = [
+    "clusters_resource_markdown",
+    "databases_resource_markdown",
+    "db_overview_resource_markdown",
+    "metrics_resource_markdown",
+]