aptdata 0.0.2__py3-none-any.whl

aptdata/mcp/server.py

@@ -0,0 +1,198 @@
+"""FastMCP server exposing aptdata tools and resources.
+
+The server allows AI agents (Claude Desktop, Copilot, Devin, …) to discover
+and execute aptdata pipelines via the Model Context Protocol.
+"""
+
+from __future__ import annotations
+
+import time
+from threading import Lock
+from typing import Any
+
+from mcp.server.fastmcp import FastMCP
+
+from aptdata.plugins import registry
+from aptdata.plugins.local_fs import (
+    CSVReader,
+    CSVWriter,
+    JSONReader,
+    JSONWriter,
+    ParquetReader,
+    ParquetWriter,
+)
+from aptdata.plugins.manager import plugin_manager
+from aptdata.plugins.postgres import PostgresReader, PostgresWriter
+from aptdata.plugins.rest import APIReader
+from aptdata.plugins.vector import QdrantWriter
+from aptdata.telemetry.instrumentation import mask_telemetry_value
+
+mcp = FastMCP("aptdata")
+_MCP_REQUEST_COUNT = 0
+_MCP_REQUEST_LOCK = Lock()
+
+
+def _mark_request() -> None:
+    global _MCP_REQUEST_COUNT
+    with _MCP_REQUEST_LOCK:
+        _MCP_REQUEST_COUNT += 1
+
+
+def get_mcp_status() -> dict[str, Any]:
+    """Return MCP activity status for TUI and diagnostics."""
+    with _MCP_REQUEST_LOCK:
+        request_count = _MCP_REQUEST_COUNT
+    return {"active": True, "request_count": request_count}
+
+
+def _register_builtin_plugins() -> None:
+    plugin_manager.register_reader("csv_reader", CSVReader)
+    plugin_manager.register_reader("json_reader", JSONReader)
+    plugin_manager.register_reader("parquet_reader", ParquetReader)
+    plugin_manager.register_reader("api_reader", APIReader)
+    plugin_manager.register_reader("postgres_reader", PostgresReader)
+    plugin_manager.register_writer("csv_writer", CSVWriter)
+    plugin_manager.register_writer("json_writer", JSONWriter)
+    plugin_manager.register_writer("parquet_writer", ParquetWriter)
+    plugin_manager.register_writer("postgres_writer", PostgresWriter)
+    plugin_manager.register_writer("qdrant_writer", QdrantWriter)
+
+
+_register_builtin_plugins()
+
+
+@mcp.tool()
+def run_flow(flow_id: str) -> dict[str, Any]:
+    """Execute a registered flow/system and return its status.
+
+    Parameters
+    ----------
+    flow_id:
+        The identifier of a system previously registered in the plugin
+        registry (e.g. ``"pipeline_x"``).
+
+    Returns
+    -------
+    dict
+        A status dict with keys ``status``, ``flow_id``, and
+        ``elapsed_seconds`` on success, or ``status`` and ``error`` on
+        failure.
+    """
+    _mark_request()
+    started_at = time.time()
+    try:
+        system_cls = registry.get(flow_id)
+        if system_cls is None:
+            return {
+                "status": "error",
+                "flow_id": flow_id,
+                "error": f"Flow '{flow_id}' not found in registry.",
+            }
+        instance = system_cls(system_id=flow_id)
+        instance.run()
+        elapsed = round(time.time() - started_at, 3)
+        return {
+            "status": "completed",
+            "flow_id": flow_id,
+            "elapsed_seconds": elapsed,
+        }
+    except Exception as exc:  # noqa: BLE001
+        elapsed = round(time.time() - started_at, 3)
+        return {
+            "status": "error",
+            "flow_id": flow_id,
+            "error": str(exc),
+            "elapsed_seconds": elapsed,
+        }
+
+
+@mcp.tool()
+def list_registered_systems() -> dict[str, Any]:
+    """Return the names of all systems available in the plugin registry.
+
+    Returns
+    -------
+    dict
+        A dict with ``systems`` (list of names) and ``count``.
+    """
+    _mark_request()
+    systems = registry.list_systems()
+    return {"systems": systems, "count": len(systems)}
+
+
+@mcp.tool()
+def list_available_plugins() -> dict[str, Any]:
+    """Return all installed plugins grouped by readers and writers."""
+    _mark_request()
+    plugins = plugin_manager.list_plugins()
+    return {
+        "plugins": plugins,
+        "count": len(plugins["readers"]) + len(plugins["writers"]),
+    }
+
+
+@mcp.tool()
+def get_plugin_schema(plugin_name: str) -> dict[str, Any]:
+    """Return constructor argument schema for a specific plugin."""
+    _mark_request()
+    try:
+        return plugin_manager.get_plugin_schema(plugin_name)
+    except KeyError as exc:
+        return {"status": "error", "error": str(exc), "plugin_name": plugin_name}
+
+
+@mcp.tool()
+def preview_dataset(plugin: str, **reader_config: Any) -> dict[str, Any]:
+    """Execute a reader plugin and return the first five rows."""
+    _mark_request()
+    try:
+        rows = plugin_manager.preview_dataset(plugin, **reader_config)
+        return {
+            "status": "ok",
+            "plugin": plugin,
+            "rows": mask_telemetry_value(rows),
+            "format": "json",
+        }
+    except KeyError as exc:
+        return {
+            "status": "error",
+            "plugin": plugin,
+            "error": str(exc),
+            "error_type": "KeyError",
+        }
+    except ValueError as exc:
+        return {
+            "status": "error",
+            "plugin": plugin,
+            "error": str(exc),
+            "error_type": "ValueError",
+        }
+    except Exception as exc:  # noqa: BLE001
+        return {
+            "status": "error",
+            "plugin": plugin,
+            "error": str(exc),
+            "error_type": type(exc).__name__,
+        }
+
+
+@mcp.resource("schema://datasets/{dataset_name}")
+def get_dataset_schema(dataset_name: str) -> str:
+    """Return metadata for a dataset registered under *dataset_name*.
+
+    This is a placeholder resource – concrete implementations should query
+    a dataset catalogue or registry.  For now it returns a JSON string
+    describing the dataset name so that agents can discover schema
+    information.
+    """
+    import json
+
+    return json.dumps(
+        {
+            "dataset": dataset_name,
+            "fields": [],
+            "description": (
+                f"Schema metadata for '{dataset_name}' (no catalogue loaded)."
+            ),
+        }
+    )

aptdata/plugins/__init__.py

@@ -0,0 +1,77 @@
+"""Plugin registry and plugin manager for aptdata.
+
+Third-party adapters (Spark, REST APIs, databases, …) register concrete
+implementations of :class:`~aptdata.core.system.ISystem` here
+so that the CLI can discover and instantiate them by name.
+
+The module also re-exports the :data:`plugin_manager` singleton from
+:mod:`aptdata.plugins.manager` and the abstract base classes from
+:mod:`aptdata.plugins.base` for convenience.
+
+Usage
+-----
+Register a system::
+
+    from aptdata.plugins import registry
+    from my_package import MySystem
+
+    registry.register("my_system", MySystem)
+
+Look up a system by name::
+
+    system_cls = registry.get("my_system")
+    if system_cls is not None:
+        system_cls(system_id="my_system").run()
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from aptdata.plugins.base import BaseReader, BaseTransformer, BaseWriter
+from aptdata.plugins.manager import PluginDependencyError, PluginManager, plugin_manager
+
+if TYPE_CHECKING:
+    from aptdata.core.system import ISystem
+
+
+class _SystemRegistry:
+    """Simple name → system-class mapping."""
+
+    def __init__(self) -> None:
+        self._store: dict[str, type[ISystem]] = {}
+
+    def register(self, name: str, system_cls: type[ISystem]) -> None:
+        """Register *system_cls* under *name*.
+
+        Parameters
+        ----------
+        name:
+            Unique identifier used on the CLI (e.g. ``"pipeline_x"``).
+        system_cls:
+            A concrete subclass of :class:`~aptdata.core.system.ISystem`.
+        """
+        self._store[name] = system_cls
+
+    def get(self, name: str) -> type[ISystem] | None:
+        """Return the system class registered under *name*, or ``None``."""
+        return self._store.get(name)
+
+    def list_systems(self) -> list[str]:
+        """Return a sorted list of all registered system names."""
+        return sorted(self._store)
+
+
+#: Global singleton registry – import this in adapter modules.
+registry = _SystemRegistry()
+
+__all__ = [
+    "registry",
+    "_SystemRegistry",
+    "BaseReader",
+    "BaseWriter",
+    "BaseTransformer",
+    "PluginManager",
+    "PluginDependencyError",
+    "plugin_manager",
+]

aptdata/plugins/ai/__init__.py

@@ -0,0 +1,6 @@
+"""AI transformation plugins for chunking and embeddings."""
+
+from aptdata.plugins.ai.chunking import TextChunker
+from aptdata.plugins.ai.embeddings import EmbeddingTransformer
+
+__all__ = ["TextChunker", "EmbeddingTransformer"]

aptdata/plugins/ai/chunking.py

@@ -0,0 +1,66 @@
+"""Text chunking plugin for RAG ingestion pipelines."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from opentelemetry import trace
+
+from aptdata.plugins.dataset import InMemoryDataset
+from aptdata.telemetry.instrumentation import record_processed_chunks
+
+
+class TextChunker:
+    """Split long text documents into chunked rows preserving lineage fields."""
+
+    def __init__(
+        self,
+        *,
+        column: str,
+        max_tokens: int = 512,
+        output_column: str | None = None,
+    ) -> None:
+        self.column = column
+        self.max_tokens = max_tokens
+        self.output_column = output_column or f"{column}_chunk"
+
+    def transform(self, dataset: InMemoryDataset) -> InMemoryDataset:
+        """Chunk each row's text and return a new dataset."""
+        rows = dataset.read()
+        chunked_rows: list[dict[str, Any]] = []
+        with trace.get_tracer("aptdata.plugins.ai").start_as_current_span(
+            "TextChunker.transform"
+        ) as span:
+            for row in rows:
+                text = str(row.get(self.column, ""))
+                paragraphs = [
+                    part.strip() for part in text.split("\n\n") if part.strip()
+                ]
+                doc_id = row.get("document_id") or row.get("id")
+                trace_id = row.get("trace_id")
+                chunk_index = 0
+                for paragraph in paragraphs or [""]:
+                    words = paragraph.split()
+                    start = 0
+                    while start < len(words) or (not words and start == 0):
+                        chunk_words = words[start : start + self.max_tokens]
+                        chunk_text = " ".join(chunk_words)
+                        enriched = dict(row)
+                        enriched[self.output_column] = chunk_text
+                        enriched["chunk_index"] = chunk_index
+                        if doc_id is not None:
+                            enriched["document_id"] = doc_id
+                        if trace_id is not None:
+                            enriched["trace_id"] = trace_id
+                        chunked_rows.append(enriched)
+                        chunk_index += 1
+                        if not words:
+                            break
+                        start += self.max_tokens
+            record_processed_chunks(len(chunked_rows))
+            span.set_attribute("aptdata.chunks.generated", len(chunked_rows))
+        out = InMemoryDataset(
+            uri=f"{dataset.uri}#chunked", schema_metadata=dict(dataset.schema_metadata)
+        )
+        out.write(chunked_rows)
+        return out

aptdata/plugins/ai/embeddings.py

@@ -0,0 +1,56 @@
+"""Embedding transformer plugin with token usage telemetry."""
+
+from __future__ import annotations
+
+import hashlib
+from typing import Any
+
+from opentelemetry import trace
+
+from aptdata.plugins.dataset import InMemoryDataset
+from aptdata.telemetry.instrumentation import record_llm_tokens_used
+
+
+class EmbeddingTransformer:
+    """Generate deterministic embeddings for text rows."""
+
+    def __init__(
+        self,
+        *,
+        column: str,
+        model: str = "text-embedding-3-small",
+    ) -> None:
+        self.column = column
+        self.model = model
+
+    def transform(self, dataset: InMemoryDataset) -> InMemoryDataset:
+        """Add embedding vectors and token usage metadata to each row."""
+        rows = dataset.read()
+        total_tokens = 0
+        transformed: list[dict[str, Any]] = []
+        with trace.get_tracer("aptdata.plugins.ai").start_as_current_span(
+            "EmbeddingTransformer.transform"
+        ) as span:
+            for row in rows:
+                text = str(row.get(self.column, ""))
+                tokens = len(text.split())
+                total_tokens += tokens
+                enriched = dict(row)
+                enriched["embedding_model"] = self.model
+                enriched["embedding_tokens"] = tokens
+                enriched[f"{self.column}_embedding"] = self._embed(text)
+                transformed.append(enriched)
+            span.set_attribute("llm.tokens.used", total_tokens)
+            span.set_attribute("llm.model", self.model)
+            span.set_attribute("llm.token_estimation_method", "whitespace")
+        record_llm_tokens_used(total_tokens)
+        out = InMemoryDataset(
+            uri=f"{dataset.uri}#embedded", schema_metadata=dict(dataset.schema_metadata)
+        )
+        out.write(transformed)
+        return out
+
+    @staticmethod
+    def _embed(text: str, *, dimensions: int = 8) -> list[float]:
+        digest = hashlib.sha256(text.encode("utf-8")).digest()
+        return [int(digest[i]) / 255.0 for i in range(dimensions)]

aptdata/plugins/base.py

@@ -0,0 +1,57 @@
+"""Abstract base interfaces for plugin readers, writers, and transformers.
+
+Every concrete reader / writer / transformer must subclass :class:`BaseReader`,
+:class:`BaseWriter`, or :class:`BaseTransformer` and implement the corresponding
+abstract method.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from aptdata.core.dataset import BaseDataset
+
+
+class BaseReader(ABC):
+    """Interface for reading data from an external source.
+
+    Subclasses **must** implement :meth:`read` and return a
+    :class:`~aptdata.core.dataset.BaseDataset` (or compatible subclass).
+    """
+
+    @abstractmethod
+    def read(self, **kwargs: Any) -> BaseDataset:
+        """Read data from the source and return a :class:`BaseDataset`."""
+
+
+class BaseWriter(ABC):
+    """Interface for writing a dataset to an external target.
+
+    Subclasses **must** implement :meth:`write`.
+    """
+
+    @abstractmethod
+    def write(self, dataset: BaseDataset, **kwargs: Any) -> None:
+        """Persist *dataset* to the target."""
+
+
+class BaseTransformer(ABC):
+    """Interface for transforming data using an engine-specific implementation.
+
+    Subclasses **must** implement :attr:`name` and :meth:`transform`.
+    Transformer instances are compatible with :meth:`Workflow.add_step` —
+    pass ``transformer.transform`` as the step callable.
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Human-readable name identifying this transformer."""
+
+    @abstractmethod
+    def transform(self, data: Any) -> Any:
+        """Apply the transformation to *data* and return the result."""
+
+
+__all__ = ["BaseReader", "BaseWriter", "BaseTransformer"]

aptdata/plugins/dataset.py

@@ -0,0 +1,62 @@
+"""In-memory dataset for plugin data exchange.
+
+Provides :class:`InMemoryDataset`, a concrete :class:`BaseDataset`
+subclass that holds tabular data as a list of dictionaries (records).
+Plugin readers produce ``InMemoryDataset`` instances and writers
+consume them.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic.dataclasses import dataclass as pydantic_dataclass
+
+from aptdata.core.dataset import BaseDataset
+
+
+@pydantic_dataclass
+class InMemoryDataset(BaseDataset):
+    """Concrete dataset that stores records in memory.
+
+    Parameters
+    ----------
+    uri:
+        Logical URI describing the data origin (informational).
+    schema_metadata:
+        Optional schema metadata mapping.
+    """
+
+    def __post_init__(self) -> None:
+        self._records: list[dict[str, Any]] = []
+
+    # -- IDataset interface -------------------------------------------------
+
+    def read(self) -> list[dict[str, Any]]:
+        """Return the in-memory records."""
+        return list(self._records)
+
+    def write(self, data: Any) -> None:
+        """Replace the in-memory records with *data*.
+
+        Parameters
+        ----------
+        data:
+            A list of dictionaries (records).
+        """
+        if not isinstance(data, list):
+            raise TypeError("InMemoryDataset expects a list of dicts.")
+        self._records = data
+
+    # -- convenience --------------------------------------------------------
+
+    @property
+    def records(self) -> list[dict[str, Any]]:
+        """Return the stored records (read-only view)."""
+        return list(self._records)
+
+    def __len__(self) -> int:
+        return len(self._records)
+
+
+__all__ = ["InMemoryDataset"]

aptdata/plugins/governance/__init__.py

@@ -0,0 +1,32 @@
+"""Governance plugin package.
+
+Provides business rules registry, dataset catalog, data classification
+policies, and lineage store.
+"""
+
+from __future__ import annotations
+
+from aptdata.plugins.governance.catalog import DatasetCatalog, DatasetCatalogEntry
+from aptdata.plugins.governance.classification import (
+    ColumnClassification,
+    DataClassificationPolicy,
+)
+from aptdata.plugins.governance.lineage_store import LineageStore
+from aptdata.plugins.governance.rules import (
+    BusinessRule,
+    RuleAuditEntry,
+    RuleRegistry,
+    RuleStatus,
+)
+
+__all__ = [
+    "DatasetCatalog",
+    "DatasetCatalogEntry",
+    "ColumnClassification",
+    "DataClassificationPolicy",
+    "LineageStore",
+    "BusinessRule",
+    "RuleAuditEntry",
+    "RuleRegistry",
+    "RuleStatus",
+]

aptdata/plugins/governance/catalog.py

@@ -0,0 +1,115 @@
+"""Dataset catalog for governance and discovery."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+from aptdata.plugins.quality.contract import ColumnClassification, SchemaContract
+
+
+@dataclass
+class DatasetCatalogEntry:
+    """Catalog record for a single dataset.
+
+    Parameters
+    ----------
+    uri:
+        Unique logical URI for the dataset (used as the catalog key).
+    name:
+        Human-readable dataset name.
+    description:
+        Description of the dataset contents.
+    owner:
+        Team or person responsible for this dataset.
+    schema_contract:
+        Optional :class:`~aptdata.plugins.quality.contract.SchemaContract`
+        governing this dataset.
+    tags:
+        Free-form classification tags.
+    classification:
+        Overall data sensitivity classification.
+    created_at:
+        UTC ISO-8601 timestamp when the entry was first registered.
+    updated_at:
+        UTC ISO-8601 timestamp of the most recent update.
+    metadata:
+        Arbitrary extra metadata.
+    """
+
+    uri: str
+    name: str = ""
+    description: str = ""
+    owner: str = ""
+    schema_contract: SchemaContract | None = None
+    tags: list[str] = field(default_factory=list)
+    classification: ColumnClassification = ColumnClassification.INTERNAL
+    created_at: str = field(
+        default_factory=lambda: datetime.now(timezone.utc).isoformat()
+    )
+    updated_at: str = field(
+        default_factory=lambda: datetime.now(timezone.utc).isoformat()
+    )
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class DatasetCatalog:
+    """In-memory catalog of :class:`DatasetCatalogEntry` objects.
+
+    Examples
+    --------
+    ::
+
+        catalog = DatasetCatalog()
+        catalog.register(
+            DatasetCatalogEntry(uri="s3://bucket/data.parquet", name="Sales")
+        )
+        entry = catalog.get("s3://bucket/data.parquet")
+        results = catalog.search(owner="data-team", tag="finance")
+    """
+
+    def __init__(self) -> None:
+        self._entries: dict[str, DatasetCatalogEntry] = {}
+
+    def register(self, entry: DatasetCatalogEntry) -> None:
+        """Register or replace a catalog entry under its
+        :attr:`~DatasetCatalogEntry.uri`."""
+        self._entries[entry.uri] = entry
+
+    def get(self, uri: str) -> DatasetCatalogEntry | None:
+        """Return the entry for *uri*, or ``None`` if not found."""
+        return self._entries.get(uri)
+
+    def search(
+        self,
+        owner: str | None = None,
+        tag: str | None = None,
+        classification: ColumnClassification | None = None,
+    ) -> list[DatasetCatalogEntry]:
+        """Search catalog entries with optional filters.
+
+        Parameters
+        ----------
+        owner:
+            If provided, only entries owned by this owner are returned.
+        tag:
+            If provided, only entries with this tag are returned.
+        classification:
+            If provided, only entries with this classification are returned.
+        """
+        results = list(self._entries.values())
+        if owner is not None:
+            results = [e for e in results if e.owner == owner]
+        if tag is not None:
+            results = [e for e in results if tag in e.tags]
+        if classification is not None:
+            results = [e for e in results if e.classification == classification]
+        return results
+
+    def list_entries(self) -> list[DatasetCatalogEntry]:
+        """Return all registered catalog entries."""
+        return list(self._entries.values())
+
+
+__all__ = ["DatasetCatalogEntry", "DatasetCatalog"]