aptdata 0.0.2__py3-none-any.whl

aptdata/core/system.py

@@ -0,0 +1,317 @@
+"""System, Component, and Flow — the universal architecture for aptdata.
+
+This module provides the three foundational abstractions:
+
+* :class:`IComponent` / :class:`BaseComponent` — a reusable, metadata-rich
+  unit of work that replaces the legacy ``Step`` abstraction.
+* :class:`IFlow` / :class:`BaseFlow` — a directed execution graph that
+  replaces the DAG management in the legacy ``Pipeline``.
+* :class:`ISystem` / :class:`BaseSystem` — the top-level orchestrator that
+  owns one or more :class:`IFlow` instances.
+
+Design goals
+------------
+* **SOLID** — each class has a single, well-defined responsibility.
+* **Pydantic** — all concrete base classes use ``pydantic.dataclasses`` for
+  runtime field validation.
+* **Metadata-driven** — :class:`ComponentMeta` carries rich information about
+  a component's role, tags, and branching behaviour so that the framework can
+  make decisions without inspecting component internals.
+* **Dependency Injection** — dependencies (datasets, services) are passed in
+  explicitly at call time; nothing is resolved from global state.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from enum import Enum
+from functools import wraps
+from typing import Any
+
+from pydantic.dataclasses import dataclass as pydantic_dataclass
+
+from aptdata.core.dataset import IDataset
+from aptdata.telemetry.instrumentation import get_tracer, mask_telemetry_value
+
+# ---------------------------------------------------------------------------
+# Component metadata
+# ---------------------------------------------------------------------------
+
+
+class ComponentKind(str, Enum):
+    """Supported processing paradigms for a :class:`BaseComponent`."""
+
+    TRANSFORM = "transform"
+    FILTER = "filter"
+    AGGREGATE = "aggregate"
+    EXTRACT = "extract"
+    LOAD = "load"
+    CUSTOM = "custom"
+
+
+@dataclass
+class ComponentMeta:
+    """Rich metadata describing a component's role and branching behaviour.
+
+    Attributes
+    ----------
+    kind:
+        The processing paradigm this component belongs to.
+    tags:
+        Arbitrary string labels for filtering, grouping, or discovery.
+    branch_on:
+        When non-empty, names the output field or condition key on which the
+        flow should branch after this component executes.
+    description:
+        Human-readable summary of what this component does.
+    extra:
+        Open-ended mapping for framework extensions or user-defined metadata.
+    """
+
+    kind: ComponentKind = ComponentKind.CUSTOM
+    tags: list[str] = field(default_factory=list)
+    branch_on: str = ""
+    description: str = ""
+    extra: dict[str, Any] = field(default_factory=dict)
+
+
+# ---------------------------------------------------------------------------
+# Component (replaces Step)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class IComponent(ABC):
+    """Interface for a reusable unit of work.
+
+    A component receives a list of :class:`~aptdata.core.dataset.IDataset`
+    inputs, validates them, executes its logic, and returns a list of
+    :class:`~aptdata.core.dataset.IDataset` outputs.  Unlike the legacy
+    ``IStep``, it may produce *multiple* output datasets to support branching
+    flows.
+    """
+
+    @property
+    @abstractmethod
+    def meta(self) -> ComponentMeta:
+        """Metadata describing this component."""
+
+    @abstractmethod
+    def validate_inputs(self, inputs: list[IDataset]) -> bool:
+        """Return ``True`` when *inputs* are valid for this component."""
+
+    @abstractmethod
+    def execute(self, inputs: list[IDataset]) -> list[IDataset]:
+        """Execute the component logic and return its output datasets."""
+
+
+@pydantic_dataclass
+class BaseComponent(IComponent):
+    """Base component with Pydantic-validated identity and metadata.
+
+    Concrete component implementations must inherit from this class and
+    implement the :meth:`validate_inputs` and :meth:`execute` abstract
+    methods inherited from :class:`IComponent`.
+
+    Parameters
+    ----------
+    component_id:
+        A unique identifier for this component within a flow.
+    metadata:
+        A :class:`ComponentMeta` instance describing the component's role.
+    """
+
+    component_id: str
+    metadata: ComponentMeta = field(default_factory=ComponentMeta)
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        """Wrap subclass execute implementations with telemetry spans."""
+        super().__init_subclass__(**kwargs)
+        execute_fn = cls.__dict__.get("execute")
+        if execute_fn is None or getattr(execute_fn, "_aptdata_instrumented", False):
+            return
+
+        @wraps(execute_fn)
+        def _instrumented_execute(
+            self: BaseComponent, inputs: list[IDataset]
+        ) -> list[IDataset]:
+            span_name = self.component_id or cls.__name__
+            kind = self.meta.kind
+            kind_value = (
+                kind.value if isinstance(kind, ComponentKind) else str(kind or "")
+            )
+            tags = sorted(self.meta.tags) if self.meta.tags else []
+            with get_tracer().start_as_current_span(span_name) as span:
+                span.set_attribute("aptdata.component_id", self.component_id)
+                span.set_attribute("aptdata.kind", kind_value)
+                span.set_attribute("aptdata.tags", tags)
+                span.set_attribute(
+                    "aptdata.branch_on",
+                    mask_telemetry_value(self.meta.branch_on, key="branch_on"),
+                )
+                span.set_attribute(
+                    "aptdata.description",
+                    mask_telemetry_value(self.meta.description, key="description"),
+                )
+                return execute_fn(self, inputs)
+
+        _instrumented_execute.__isabstractmethod__ = getattr(
+            execute_fn, "__isabstractmethod__", False
+        )
+        _instrumented_execute._aptdata_instrumented = True  # type: ignore[attr-defined]
+        cls.execute = _instrumented_execute  # type: ignore[method-assign]
+
+    @property
+    def meta(self) -> ComponentMeta:
+        return self.metadata
+
+
+# ---------------------------------------------------------------------------
+# Flow graph primitives
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class FlowEdge:
+    """A directed edge in a :class:`BaseFlow` execution graph.
+
+    Optionally carries a *condition* callable; when present the edge is only
+    traversed when ``condition(outputs)`` evaluates to ``True``, enabling
+    conditional / branching flows.
+
+    Parameters
+    ----------
+    source_id:
+        The :attr:`~BaseComponent.component_id` of the upstream component.
+    target_id:
+        The :attr:`~BaseComponent.component_id` of the downstream component.
+    condition:
+        Optional predicate evaluated against the source component's outputs.
+    """
+
+    source_id: str
+    target_id: str
+    condition: Callable[[list[IDataset]], bool] | None = None
+
+
+@dataclass
+class FlowNode:
+    """A node wrapping a :class:`IComponent` inside a :class:`IFlow` graph.
+
+    Parameters
+    ----------
+    component:
+        The component held by this node.
+    flow:
+        Back-reference to the owning flow (set by the flow on insertion).
+    """
+
+    component: IComponent
+    flow: IFlow | None = field(default=None, repr=False)
+
+
+# ---------------------------------------------------------------------------
+# Flow (replaces the DAG management in Pipeline)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class IFlow(ABC):
+    """Interface for a directed execution graph of :class:`IComponent` nodes.
+
+    A flow owns a set of components and the directed edges that connect them.
+    It is responsible for validating the graph structure (:meth:`compile`)
+    and driving execution (:meth:`run`).
+    """
+
+    @abstractmethod
+    def add_component(self, component: IComponent) -> None:
+        """Add *component* as a node in this flow."""
+
+    @abstractmethod
+    def connect(
+        self,
+        source_id: str,
+        target_id: str,
+        condition: Callable[[list[IDataset]], bool] | None = None,
+    ) -> None:
+        """Create a directed edge from *source_id* to *target_id*.
+
+        Parameters
+        ----------
+        source_id:
+            The :attr:`~BaseComponent.component_id` of the upstream component.
+        target_id:
+            The :attr:`~BaseComponent.component_id` of the downstream component.
+        condition:
+            Optional predicate that gates traversal of the edge.
+        """
+
+    @abstractmethod
+    def compile(self) -> None:
+        """Validate the graph structure before execution.
+
+        Implementations should raise :exc:`ValueError` when the graph is
+        invalid (e.g. unknown node references, cycles in a DAG-only flow).
+        """
+
+    @abstractmethod
+    def run(self, initial_inputs: list[IDataset]) -> list[IDataset]:
+        """Execute the flow starting with *initial_inputs*.
+
+        Returns the outputs produced by the terminal component(s).
+        """
+
+
+@pydantic_dataclass
+class BaseFlow(IFlow):
+    """Base flow with Pydantic-validated identity and a managed graph.
+
+    Concrete flow implementations must inherit from this class and implement
+    the :meth:`add_component`, :meth:`connect`, :meth:`compile` and
+    :meth:`run` abstract methods inherited from :class:`IFlow`.
+
+    Parameters
+    ----------
+    flow_id:
+        A unique identifier for this flow within a system.
+    """
+
+    flow_id: str
+
+
+# ---------------------------------------------------------------------------
+# System (top-level orchestrator)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ISystem(ABC):
+    """Interface for a system that orchestrates one or more :class:`IFlow` instances."""
+
+    @abstractmethod
+    def register_flow(self, flow: IFlow) -> None:
+        """Register *flow* in this system."""
+
+    @abstractmethod
+    def run(self) -> None:
+        """Execute all registered flows."""
+
+
+@pydantic_dataclass
+class BaseSystem(ISystem):
+    """Base system with Pydantic-validated identity.
+
+    Concrete system implementations must inherit from this class and implement
+    the :meth:`register_flow` and :meth:`run` abstract methods inherited from
+    :class:`ISystem`.
+
+    Parameters
+    ----------
+    system_id:
+        A unique identifier for this system.
+    """
+
+    system_id: str

aptdata/core/workflow.py

@@ -0,0 +1,372 @@
+"""Workflow abstractions with context-aware execution hooks."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from time import sleep, time_ns
+from typing import Any
+from uuid import uuid4
+
+from opentelemetry import trace
+from pydantic.dataclasses import dataclass as pydantic_dataclass
+
+from aptdata.core.context import ExecutionContext
+from aptdata.core.dataset import IDataset
+from aptdata.core.state import StateBackend
+from aptdata.core.system import IComponent
+from aptdata.plugins.dataset import InMemoryDataset
+from aptdata.telemetry.instrumentation import (
+    record_processed_documents,
+    reset_ingestion_metrics,
+    set_ingestion_total_documents,
+)
+
+
+@dataclass
+class WorkflowEdge:
+    """A directed edge in a workflow graph."""
+
+    source_id: str
+    target_id: str
+    condition: Callable[[list[IDataset]], bool] | None = None
+
+
+@dataclass
+class WorkflowNode:
+    """A node that wraps a component inside a workflow."""
+
+    component: IComponent
+    workflow: IWorkflow | None = field(default=None, repr=False)
+
+
+@dataclass
+class IWorkflow(ABC):
+    """Interface for workflow execution."""
+
+    @abstractmethod
+    def add_component(self, component: IComponent) -> None:
+        """Add a component to the workflow."""
+
+    @abstractmethod
+    def connect(
+        self,
+        source_id: str,
+        target_id: str,
+        condition: Callable[[list[IDataset]], bool] | None = None,
+    ) -> None:
+        """Connect components with an optional condition."""
+
+    @abstractmethod
+    def compile(self) -> None:
+        """Validate and prepare workflow execution structures."""
+
+    @abstractmethod
+    def before_run(self, initial_inputs: list[IDataset]) -> None:
+        """Lifecycle hook called before execution."""
+
+    @abstractmethod
+    def after_run(self, outputs: list[IDataset]) -> None:
+        """Lifecycle hook called after execution."""
+
+    @abstractmethod
+    def run(self, initial_inputs: list[IDataset]) -> list[IDataset]:
+        """Execute the workflow."""
+
+
+@pydantic_dataclass
+class BaseWorkflow(IWorkflow):
+    """Default workflow implementation with adjacency compilation and hooks."""
+
+    workflow_id: str
+    context: ExecutionContext = field(default_factory=ExecutionContext)
+
+    def __post_init__(self) -> None:
+        self._nodes: dict[str, WorkflowNode] = {}
+        self._edges: list[WorkflowEdge] = []
+        self._adjacency: dict[str, list[WorkflowEdge]] = {}
+        self._execution_order: list[str] = []
+        self._compiled = False
+
+    def add_component(self, component: IComponent) -> None:
+        self._nodes[component.component_id] = WorkflowNode(
+            component=component, workflow=self
+        )
+        self._compiled = False
+
+    def connect(
+        self,
+        source_id: str,
+        target_id: str,
+        condition: Callable[[list[IDataset]], bool] | None = None,
+    ) -> None:
+        self._edges.append(
+            WorkflowEdge(source_id=source_id, target_id=target_id, condition=condition)
+        )
+        self._compiled = False
+
+    def compile(self) -> None:
+        if not self._nodes:
+            raise ValueError("Workflow has no components.")
+
+        indegree = {component_id: 0 for component_id in self._nodes}
+        adjacency: dict[str, list[WorkflowEdge]] = {
+            component_id: [] for component_id in self._nodes
+        }
+
+        for edge in self._edges:
+            if edge.source_id not in self._nodes:
+                raise ValueError(f"Unknown source_id: {edge.source_id!r}")
+            if edge.target_id not in self._nodes:
+                raise ValueError(f"Unknown target_id: {edge.target_id!r}")
+            adjacency[edge.source_id].append(edge)
+            indegree[edge.target_id] += 1
+
+        queue = deque(
+            component_id
+            for component_id, in_degree in indegree.items()
+            if in_degree == 0
+        )
+        execution_order: list[str] = []
+        while queue:
+            current = queue.popleft()
+            execution_order.append(current)
+            for edge in adjacency[current]:
+                indegree[edge.target_id] -= 1
+                if indegree[edge.target_id] == 0:
+                    queue.append(edge.target_id)
+
+        if len(execution_order) != len(self._nodes):
+            raise ValueError("Workflow graph has a cycle.")
+
+        self._adjacency = adjacency
+        self._execution_order = execution_order
+        self._compiled = True
+
+    def before_run(self, initial_inputs: list[IDataset]) -> None:
+        """Lifecycle hook called before execution."""
+        self.context.set("workflow.last_input_count", len(initial_inputs))
+
+    def after_run(self, outputs: list[IDataset]) -> None:
+        """Lifecycle hook called after execution."""
+        self.context.set("workflow.last_output_count", len(outputs))
+
+    def run(self, initial_inputs: list[IDataset]) -> list[IDataset]:
+        if not self._compiled:
+            raise RuntimeError("Workflow not compiled.")
+
+        self.before_run(initial_inputs)
+        pending_inputs: dict[str, list[IDataset]] = {}
+        for component_id in self._execution_order:
+            pending_inputs[component_id] = []
+        if self._execution_order:
+            pending_inputs[self._execution_order[0]] = list(initial_inputs)
+
+        # Fallback when no component executes or no branch is traversed.
+        terminal_outputs: list[IDataset] = list(initial_inputs)
+        for component_id in self._execution_order:
+            component = self._nodes[component_id].component
+            inputs = pending_inputs.get(component_id, [])
+            if not component.validate_inputs(inputs):
+                continue
+
+            outputs = component.execute(inputs)
+            outgoing = self._adjacency.get(component_id, [])
+            if not outgoing:
+                terminal_outputs = outputs
+                continue
+
+            traversed = False
+            for edge in outgoing:
+                if edge.condition is None or edge.condition(outputs):
+                    pending_inputs[edge.target_id].extend(outputs)
+                    traversed = True
+            if not traversed:
+                terminal_outputs = outputs
+
+        self.after_run(terminal_outputs)
+        return terminal_outputs
+
+
+@dataclass
+class _WorkflowStep:
+    fn: Callable[..., Any]
+    retries: int = 0
+    backoff: float = 0.0
+
+
+class Workflow:
+    """Function-oriented workflow with retries, checkpoints and resume."""
+
+    def __init__(
+        self,
+        name: str,
+        *,
+        enable_checkpointing: bool = False,
+        state_backend: StateBackend | None = None,
+    ) -> None:
+        self.name = name
+        self.enable_checkpointing = enable_checkpointing
+        self.state_backend = state_backend or StateBackend()
+        self._steps: list[_WorkflowStep] = []
+
+    def add_step(
+        self,
+        step: Callable[..., Any],
+        *,
+        retries: int = 0,
+        backoff: float = 0.0,
+    ) -> None:
+        """Add a callable pipeline step with optional retry/backoff policy."""
+        self._steps.append(
+            _WorkflowStep(fn=step, retries=max(0, retries), backoff=max(0.0, backoff))
+        )
+
+    def execute(self, data: Any | None = None) -> Any:
+        """Execute workflow from the first step."""
+        run_id = f"{self.name}_{time_ns()}_{uuid4().hex[:8]}"
+        return self._run(run_id=run_id, start_index=0, data=data)
+
+    def resume(self, run_id: str, data: Any | None = None) -> Any:
+        """Resume execution from the last checkpoint for *run_id*."""
+        state = self.state_backend.load(run_id)
+        restored_data = self._deserialize_payload(state.get("data"))
+        return self._run(
+            run_id=run_id,
+            start_index=int(state.get("next_step_index", 0)),
+            data=restored_data if restored_data is not None else data,
+        )
+
+    def _run(self, *, run_id: str, start_index: int, data: Any | None) -> Any:
+        if not self._steps:
+            return data
+        reset_ingestion_metrics()
+        trace_id = None
+        with trace.get_tracer("aptdata.workflow").start_as_current_span(
+            f"{self.name}.run"
+        ) as span:
+            trace_id = f"{span.get_span_context().trace_id:032x}"
+            payload = self._attach_lineage(data, trace_id=trace_id)
+            set_ingestion_total_documents(self._count_records(payload))
+            for step_index in range(start_index, len(self._steps)):
+                step = self._steps[step_index]
+                payload = self._run_step(
+                    step=step,
+                    payload=payload,
+                    step_index=step_index,
+                    run_id=run_id,
+                    trace_id=trace_id,
+                )
+                record_processed_documents(self._count_records(payload))
+                if self.enable_checkpointing:
+                    self.state_backend.save(
+                        run_id,
+                        {
+                            "run_id": run_id,
+                            "next_step_index": step_index + 1,
+                            "data": self._serialize_payload(payload),
+                        },
+                    )
+            return payload
+
+    def _run_step(
+        self,
+        *,
+        step: _WorkflowStep,
+        payload: Any | None,
+        step_index: int,
+        run_id: str,
+        trace_id: str,
+    ) -> Any:
+        last_error: Exception | None = None
+        for attempt in range(step.retries + 1):
+            with trace.get_tracer("aptdata.workflow").start_as_current_span(
+                f"{self.name}.step.{step_index}"
+            ) as span:
+                span.set_attribute("aptdata.step.index", step_index)
+                span.set_attribute(
+                    "aptdata.step.name", getattr(step.fn, "__name__", "step")
+                )
+                span.set_attribute("aptdata.retry.attempt", attempt + 1)
+                span.set_attribute("aptdata.retry.max_attempts", step.retries + 1)
+                span.set_attribute("aptdata.trace_id", trace_id)
+                try:
+                    result = step.fn() if payload is None else step.fn(payload)
+                    return self._attach_lineage(result, trace_id=trace_id)
+                except Exception as exc:  # noqa: BLE001
+                    last_error = exc
+                    span.record_exception(exc)
+                    if attempt < step.retries:
+                        backoff_seconds = (
+                            step.backoff * (2**attempt) if step.backoff else 0.0
+                        )
+                        sleep(min(backoff_seconds, 30.0))
+        if self.enable_checkpointing:
+            self.state_backend.save(
+                run_id,
+                {
+                    "run_id": run_id,
+                    "next_step_index": step_index,
+                    "data": self._serialize_payload(payload),
+                    "error": str(last_error) if last_error else "unknown error",
+                },
+            )
+        if last_error is not None:
+            raise last_error
+        raise RuntimeError("Workflow step failed with unknown error.")
+
+    @staticmethod
+    def _count_records(payload: Any | None) -> int:
+        if isinstance(payload, InMemoryDataset):
+            return len(payload.read())
+        if isinstance(payload, list):
+            return len(payload)
+        return 0
+
+    @staticmethod
+    def _attach_lineage(payload: Any | None, *, trace_id: str) -> Any | None:
+        if payload is None:
+            return None
+        if isinstance(payload, InMemoryDataset):
+            records = payload.read()
+            for index, record in enumerate(records):
+                if not isinstance(record, dict):
+                    continue
+                record.setdefault("trace_id", trace_id)
+                record.setdefault("document_id", record.get("id") or f"doc-{index}")
+            payload.write(records)
+            return payload
+        if isinstance(payload, list):
+            for index, record in enumerate(payload):
+                if not isinstance(record, dict):
+                    continue
+                record.setdefault("trace_id", trace_id)
+                record.setdefault("document_id", record.get("id") or f"doc-{index}")
+            return payload
+        return payload
+
+    @staticmethod
+    def _serialize_payload(payload: Any | None) -> Any:
+        if isinstance(payload, InMemoryDataset):
+            return {
+                "__type__": "InMemoryDataset",
+                "uri": payload.uri,
+                "schema_metadata": payload.schema_metadata,
+                "records": payload.read(),
+            }
+        return payload
+
+    @staticmethod
+    def _deserialize_payload(payload: Any | None) -> Any | None:
+        if not isinstance(payload, dict):
+            return payload
+        if payload.get("__type__") != "InMemoryDataset":
+            return payload
+        dataset = InMemoryDataset(
+            uri=payload.get("uri", "memory://checkpoint"),
+            schema_metadata=payload.get("schema_metadata", {}),
+        )
+        dataset.write(payload.get("records", []))
+        return dataset

aptdata/mcp/__init__.py

@@ -0,0 +1,5 @@
+"""MCP (Model Context Protocol) server integration for aptdata."""
+
+from aptdata.mcp.server import mcp
+
+__all__ = ["mcp"]