pretensor 0.1.0a0__py3-none-any.whl

pretensor/mcp/tools/impact.py

@@ -0,0 +1,165 @@
+"""MCP ``impact`` tool payload."""
+
+from __future__ import annotations
+
+import json
+from collections import deque
+from pathlib import Path
+from typing import Any
+
+from pretensor.intelligence.shadow_alias import get_shadow_alias_node_ids
+from pretensor.visibility.filter import VisibilityFilter
+
+from ..payload_types import ImpactItemPayload
+from ..service_context import get_effective_visibility_filter
+from ..service_registry import (
+    graph_path_for_entry,
+    load_registry,
+    open_store_for_entry,
+    resolve_registry_entry,
+)
+from .context import qualified_from_node_id, resolve_table_node_id
+
+
+def impact_payload(
+    graph_dir: Path,
+    *,
+    table: str,
+    database: str,
+    column: str | None = None,
+    max_depth: int = 3,
+    visibility_filter: VisibilityFilter | None = None,
+) -> dict[str, Any]:
+    """Downstream tables grouped by hop depth (FK + inferred edges)."""
+    reg = load_registry(graph_dir)
+    entry = resolve_registry_entry(reg, database)
+    if entry is None:
+        return {
+            "error": "Unknown database connection or name; pass `database`.",
+        }
+    gp = graph_path_for_entry(entry)
+    if not gp.exists():
+        return {"error": f"Graph file missing: {gp}"}
+
+    db_key = str(entry.database)
+    vf = visibility_filter or get_effective_visibility_filter()
+    cn = str(entry.connection_name)
+    store = open_store_for_entry(entry)
+    try:
+        start_id, err = resolve_table_node_id(
+            store, table, db_key, visibility_filter=vf
+        )
+        if err is not None:
+            try:
+                return json.loads(err)
+            except json.JSONDecodeError:
+                return {"error": err}
+        if start_id is None:
+            return {"error": "Could not resolve starting table; check the table name and database."}
+
+        fk_rows = store.query_all_rows(
+            """
+            MATCH (a:SchemaTable)-[r:FK_REFERENCES]->(b:SchemaTable)
+            WHERE a.database = $db AND b.database = $db
+            RETURN a.node_id, b.node_id, r.source_column, r.target_column, 1.0 AS conf
+            """,
+            {"db": db_key},
+        )
+        inf_rows = store.query_all_rows(
+            """
+            MATCH (a:SchemaTable)-[r:INFERRED_JOIN]->(b:SchemaTable)
+            WHERE a.database = $db AND b.database = $db
+            RETURN a.node_id, b.node_id, r.source_column, r.target_column, r.confidence
+            """,
+            {"db": db_key},
+        )
+        shadow_ids = get_shadow_alias_node_ids(store, db_key)
+
+        rev: dict[str, list[tuple[str, str, str, float]]] = {}
+        for row in fk_rows + inf_rows:
+            child, parent, sc, tc, conf = row
+            if child is None or parent is None:
+                continue
+            sa, sb = str(child), str(parent)
+            if sa in shadow_ids or sb in shadow_ids:
+                continue
+            if vf is not None:
+                cmeta = store.query_all_rows(
+                    """
+                    MATCH (t:SchemaTable {node_id: $id})
+                    RETURN t.schema_name, t.table_name
+                    """,
+                    {"id": sa},
+                )
+                if cmeta:
+                    sn, tn = cmeta[0]
+                    if not vf.is_table_visible(cn, str(sn or ""), str(tn or "")):
+                        continue
+            w = float(conf) if conf is not None else 1.0
+            rev.setdefault(sb, []).append((sa, str(sc or ""), str(tc or ""), w))
+
+        q: deque[tuple[str, int, float, str, str]] = deque()
+        for dependent, scol, _tcol, w in rev.get(start_id, []):
+            if column is not None and scol != column:
+                continue
+            dep_short = qualified_from_node_id(store, dependent).split(".")[-1]
+            via = f"{dep_short}.{scol}"
+            q.append((dependent, 1, w, via, start_id))
+
+        seen: dict[str, tuple[int, float, str]] = {}
+        while q:
+            cur, depth, path_min, via, _parent = q.popleft()
+            if cur in seen:
+                continue
+            seen[cur] = (depth, path_min, via)
+            if depth >= max_depth:
+                continue
+            for dependent, scol, _tcol, w in rev.get(cur, []):
+                if dependent in seen:
+                    continue
+                new_d = depth + 1
+                new_min = min(path_min, w)
+                dep_short = qualified_from_node_id(store, dependent).split(".")[-1]
+                new_via = f"{via} → {dep_short}.{scol}"
+                q.append((dependent, new_d, new_min, new_via, cur))
+
+        direct: list[ImpactItemPayload] = []
+        two_hop: list[ImpactItemPayload] = []
+        three_hop: list[ImpactItemPayload] = []
+        low_confidence: list[ImpactItemPayload] = []
+        for node_id, (depth, conf, via) in seen.items():
+            name = qualified_from_node_id(store, node_id)
+            item: ImpactItemPayload = {
+                "type": "table",
+                "name": name,
+                "via": via,
+                "confidence": conf,
+            }
+            if conf < 0.5:
+                item["hop"] = depth
+                low_confidence.append(item)
+                continue
+            if depth == 1:
+                direct.append(item)
+            elif depth == 2:
+                two_hop.append(item)
+            elif depth == 3:
+                three_hop.append(item)
+
+        total = len(seen)
+        return {
+            "target": {"table": table, "column": column},
+            "database": database,
+            "impact": {
+                "direct": direct,
+                "two_hop": two_hop,
+                "three_hop": three_hop,
+                "low_confidence": low_confidence,
+            },
+            "total_affected": total,
+        }
+    finally:
+        store.close()
+
+
+__all__ = ["impact_payload"]

pretensor/mcp/tools/list.py

@@ -0,0 +1,110 @@
+"""MCP ``list_databases`` tool payload."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from pretensor.config import GraphConfig
+from pretensor.visibility.filter import VisibilityFilter
+
+from ..payload_types import (
+    DIALECT_TO_DB_TYPE,
+    DatabaseListItem,
+    iso_format,
+    stale_threshold_days,
+    staleness_days,
+    utc_now,
+)
+from ..service_context import (
+    get_effective_graph_config,
+    get_effective_visibility_filter,
+)
+from ..service_registry import (
+    counts_for_graph,
+    graph_path_for_entry,
+    load_registry,
+    open_store_for_entry,
+)
+
+
+def list_databases_payload(
+    graph_dir: Path,
+    *,
+    config: GraphConfig | None = None,
+    visibility_filter: VisibilityFilter | None = None,
+) -> dict[str, Any]:
+    """Build JSON-serializable payload for ``list_databases``."""
+    cfg = get_effective_graph_config(config)
+    vf = visibility_filter or get_effective_visibility_filter()
+    reg = load_registry(graph_dir)
+    items: list[DatabaseListItem] = []
+    now = utc_now()
+    threshold = stale_threshold_days(cfg)
+    for entry in reg.list_entries():
+        table_count = 0
+        column_count = 0
+        row_sum = 0
+        schemas: list[str] = []
+        has_dbt_manifest: str = "not_attempted"
+        has_llm_enrichment: str = "not_attempted"
+        has_external_consumers: str = "not_attempted"
+        gp = graph_path_for_entry(entry)
+        if gp.exists():
+            store = open_store_for_entry(entry)
+            try:
+                counts = counts_for_graph(
+                    store,
+                    connection_name=entry.connection_name,
+                    visibility_filter=vf,
+                    entry=entry,
+                )
+                table_count = counts.table_count
+                column_count = counts.column_count
+                row_sum = counts.row_count
+                schemas = counts.schemas
+                has_dbt_manifest = counts.has_dbt_manifest
+                has_llm_enrichment = counts.has_llm_enrichment
+                has_external_consumers = counts.has_external_consumers
+            finally:
+                store.close()
+        days = staleness_days(entry.last_indexed_at, now)
+        stale = days > threshold
+        capabilities = ["schema", "fk_edges", "inferred_joins", "clustering"]
+        if has_dbt_manifest == "present":
+            capabilities.append("dbt_lineage")
+            capabilities.append("dbt_metrics")
+        if has_llm_enrichment == "present":
+            capabilities.append("llm_enrichment")
+        if has_external_consumers == "present":
+            capabilities.append("external_consumers")
+        row: DatabaseListItem = {
+            "name": entry.connection_name,
+            "db_type": DIALECT_TO_DB_TYPE.get(entry.dialect, "postgresql"),
+            "table_count": table_count,
+            "column_count": column_count,
+            "row_count": row_sum,
+            "schemas": schemas,
+            "has_dbt_manifest": has_dbt_manifest,
+            "has_llm_enrichment": has_llm_enrichment,
+            "has_external_consumers": has_external_consumers,
+            "capabilities": capabilities,
+            "last_indexed": iso_format(entry.last_indexed_at),
+            "is_stale": stale,
+            "staleness_days": days,
+            "graph_path": str(graph_path_for_entry(entry)),
+        }
+        # Only emit ``database`` when it differs from ``name`` — otherwise the
+        # duplicate field is noise (e.g. ``{"name": "pagila", "database": "pagila"}``).
+        if entry.database and entry.database != entry.connection_name:
+            row["database"] = entry.database
+        if stale:
+            row["stale_warning"] = (
+                f"Graph index is {days} days old (threshold {threshold} days). "
+                "Consider `pretensor reindex <dsn>` or `pretensor index`."
+            )
+        items.append(row)
+    return {"databases": items}
+
+
+__all__ = ["list_databases_payload"]

pretensor/mcp/tools/schema.py

@@ -0,0 +1,128 @@
+"""MCP ``schema`` tool: introspect node labels and edge types in the Kuzu graph."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+from pretensor.core.schema import CATALOG_EDGE_TYPES, CATALOG_NODE_LABELS
+from pretensor.core.store import KuzuStore
+
+from ..service_registry import (
+    graph_path_for_entry,
+    load_registry,
+    resolve_registry_entry,
+)
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["schema_payload"]
+
+_NODE_LABEL_DESC = {label: desc for label, desc in CATALOG_NODE_LABELS}
+_EDGE_INFO = {name: (src, dst, desc) for name, src, dst, desc in CATALOG_EDGE_TYPES}
+
+
+def _list_catalog_tables(store: KuzuStore) -> list[tuple[str, str]]:
+    """Return ``[(name, type), ...]`` from ``CALL show_tables()``.
+
+    ``type`` is ``'NODE'`` or ``'REL'`` per Kuzu.
+    """
+    raw = store.execute("CALL show_tables() RETURN name, type")
+    if isinstance(raw, list):
+        return []
+    rows = raw.rows_as_dict()
+    out: list[tuple[str, str]] = []
+    while rows.has_next():
+        row: Any = rows.get_next()
+        if isinstance(row, dict):
+            out.append((str(row.get("name", "")), str(row.get("type", "")).upper()))
+    return out
+
+
+def _table_properties(store: KuzuStore, label: str) -> list[dict[str, str]]:
+    """Return ``[{"name": ..., "type": ...}, ...]`` for one node/rel table."""
+    try:
+        raw = store.execute(f"CALL table_info('{label}') RETURN name, type")
+    except Exception:
+        return []
+    if isinstance(raw, list):
+        return []
+    rows = raw.rows_as_dict()
+    out: list[dict[str, str]] = []
+    while rows.has_next():
+        row: Any = rows.get_next()
+        if isinstance(row, dict):
+            out.append(
+                {
+                    "name": str(row.get("name", "")),
+                    "type": str(row.get("type", "")),
+                }
+            )
+    return out
+
+
+def schema_payload(
+    graph_dir: Path,
+    *,
+    database: str,
+    label: str | None = None,
+) -> dict[str, Any]:
+    """Return the node and edge catalog for the indexed graph of ``database``.
+
+    Args:
+        graph_dir: Graph workspace directory (registry + Kuzu files).
+        database: Connection name or logical database name.
+        label: If set, return only the matching node label or edge type.
+
+    Returns:
+        ``{"nodes": [...], "edges": [...]}`` on success, or ``{"error": "..."}``.
+    """
+    db = (database or "").strip()
+    if not db:
+        return {"error": "Missing `database`"}
+
+    reg = load_registry(graph_dir)
+    entry = resolve_registry_entry(reg, db)
+    if entry is None:
+        return {"error": f"Unknown database: {db!r}"}
+
+    graph_path = graph_path_for_entry(entry)
+    if not graph_path.exists():
+        return {"error": f"Graph file not found for database: {db!r}"}
+
+    store = KuzuStore(graph_path)
+    try:
+        store.ensure_schema()
+        catalog = _list_catalog_tables(store)
+        nodes: list[dict[str, Any]] = []
+        edges: list[dict[str, Any]] = []
+        for name, kind in catalog:
+            if label and name != label:
+                continue
+            properties = _table_properties(store, name)
+            if kind == "NODE":
+                nodes.append(
+                    {
+                        "label": name,
+                        "description": _NODE_LABEL_DESC.get(name, ""),
+                        "properties": properties,
+                    }
+                )
+            elif kind == "REL":
+                src, dst, desc = _EDGE_INFO.get(name, ("", "", ""))
+                edges.append(
+                    {
+                        "type": name,
+                        "from": src,
+                        "to": dst,
+                        "description": desc,
+                        "properties": properties,
+                    }
+                )
+    finally:
+        store.close()
+
+    if label and not nodes and not edges:
+        return {"error": f"Unknown label: {label!r}"}
+    return {"nodes": nodes, "edges": edges}

pretensor/mcp/tools/search.py

@@ -0,0 +1,108 @@
+"""MCP ``query`` tool (BM25 search) payload."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any, cast
+
+from pretensor.observability import log_timed_operation
+from pretensor.search.base import BaseSearchIndex
+from pretensor.search.index import KeywordSearchIndex
+from pretensor.visibility.filter import VisibilityFilter
+
+from ..payload_types import QueryHit, snippet
+from ..service_context import (
+    get_effective_search_index_cls,
+    get_effective_visibility_filter,
+)
+from ..service_registry import load_registry
+
+logger = logging.getLogger(__name__)
+
+
+def _hit_visible(hit: QueryHit, vf: VisibilityFilter | None) -> bool:
+    if vf is None or hit.get("node_type") != "SchemaTable":
+        return True
+    name = str(hit.get("name", "")).strip()
+    if "." not in name:
+        return vf.is_table_visible(str(hit.get("connection_name", "")), "", name)
+    sn, _, tn = name.partition(".")
+    return vf.is_table_visible(str(hit.get("connection_name", "")), sn, tn)
+
+
+def _load_search_index(
+    graph_dir: Path,
+    *,
+    search_index_cls: type[BaseSearchIndex],
+) -> BaseSearchIndex:
+    """Load/build a search index using class-level hooks when available."""
+    reg = load_registry(graph_dir)
+    search_index_cls_any = search_index_cls
+    default_path = getattr(search_index_cls_any, "default_path", None)
+    index_path = (
+        Path(cast(Path, default_path(graph_dir)))
+        if callable(default_path)
+        else KeywordSearchIndex.default_path(graph_dir)
+    )
+    load_or_build = getattr(search_index_cls_any, "load_or_build", None)
+    if callable(load_or_build):
+        built = load_or_build(reg, index_path)
+        if not isinstance(built, BaseSearchIndex):
+            raise TypeError(
+                f"{search_index_cls.__name__}.load_or_build must return BaseSearchIndex"
+            )
+        return built
+    index = cast(Any, search_index_cls_any)(index_path)
+    if not isinstance(index, BaseSearchIndex):
+        raise TypeError(
+            f"{search_index_cls.__name__} constructor must return BaseSearchIndex"
+        )
+    index.index_graph(reg)
+    return index
+
+
+def query_payload(
+    graph_dir: Path,
+    *,
+    q: str,
+    db: str | None = None,
+    limit: int = 10,
+    search_index_cls: type[BaseSearchIndex] | None = None,
+    visibility_filter: VisibilityFilter | None = None,
+) -> dict[str, Any]:
+    """BM25 search over indexed metadata."""
+    with log_timed_operation(
+        logger,
+        event="mcp.query_payload",
+        query=q,
+        db=db,
+        limit=limit,
+        graph_dir=str(graph_dir),
+    ):
+        vf = visibility_filter or get_effective_visibility_filter()
+        index_cls = get_effective_search_index_cls(search_index_cls)
+        idx = _load_search_index(graph_dir, search_index_cls=index_cls)
+        raw = idx.search(q, db=db, limit=limit * 4)
+        hits: list[QueryHit] = []
+        for row in raw:
+            if row.node_type != "SchemaTable":
+                continue
+            item: QueryHit = {
+                "node_type": row.node_type,
+                "name": row.name,
+                "database_name": row.database_name,
+                "connection_name": row.connection_name,
+                "description": row.description,
+                "snippet": snippet(row.description),
+                "score": row.score,
+            }
+            if not _hit_visible(item, vf):
+                continue
+            hits.append(item)
+            if len(hits) >= limit:
+                break
+        return {"query": q, "db": db, "results": hits}
+
+
+__all__ = ["query_payload"]