loom-core 0.1.0__py3-none-any.whl

src/core/state.py

@@ -0,0 +1,96 @@
+from contextlib import asynccontextmanager
+from typing import Any, Awaitable, Callable, Generic
+
+from ..common.errors import StopReplay
+from ..schemas.workflow import InputT, StateT
+
+
+class StateProxy(Generic[InputT, StateT]):
+    """
+    Proxy class for managing state interactions.
+    Provides methods to get and set state values in the database.
+    """
+
+    _data: StateT
+    _ctx: Any
+    _batch = None
+
+    def __init__(self, ctx: Any, data: StateT) -> None:
+        self._data = data
+        self._ctx = ctx
+
+    def __getattr__(self, name: str) -> Any:
+        return self._data.get(name)
+
+    def get(self, name: str, default: Any = None) -> Any:
+        return self._data.get(name, default)
+
+    def snapshot(self) -> StateT:
+        return self._data
+
+    async def set(self, name: str, value: Any) -> None:
+        event = self._ctx._peek()
+        if event and event["type"] == "STATE_SET" and event["payload"]["key"] == name:
+            self._ctx._consume()
+            return
+
+        event = ("STATE_SET", {"key": name, "value": value})
+
+        if self._batch is not None:
+            self._batch.append(event)
+        else:
+            await self._ctx._append_event(*event)
+            raise StopReplay
+
+    async def update(self, **updaters: Callable[..., Awaitable[Any]]) -> None:
+        """
+        Example:
+            await ctx.state.update(
+                count=lambda c: (c or 0) + 1,
+                name=lambda _: "Satadeep",
+            )
+        """
+        event = self._ctx._peek()
+
+        if event and event["type"] == "STATE_UPDATE":
+            payload = event["payload"]
+            if set(payload["values"].keys()) == set(updaters.keys()):
+                self._ctx._consume()
+                return
+        new_values = {}
+
+        for key, fn in updaters.items():
+            old = self._data.get(key)
+            new_values[key] = await fn(old)
+
+        event = ("STATE_UPDATE", {"values": new_values})
+
+        if self._batch is not None:
+            self._batch.append(event)
+        else:
+            await self._ctx._append_event(*event)
+            raise StopReplay
+
+    @asynccontextmanager
+    async def batch(self):
+        """
+        Context manager to batch multiple state updates into a single event.
+        Example:
+            async with ctx.state.batch():
+                await ctx.state.set("a", 1)
+                await ctx.state.set("b", 2)
+                await ctx.state.update(
+                    count=lambda c: (c or 0) + 1,
+                )
+        """
+        if self._batch is not None:
+            raise RuntimeError("Nested batches are not supported.")
+        self._batch = [] # type: ignore
+
+        try:
+            yield
+        finally:
+            for type, payload in self._batch:
+                await self._ctx._append_event(type, payload)
+            self._batch = None  # type: ignore
+            raise StopReplay

src/core/worker.py

@@ -0,0 +1,147 @@
+"""Distributed workflow worker with graceful shutdown and concurrent task execution."""
+
+import asyncio
+import signal
+import sys
+from datetime import datetime, timezone
+
+from .runner import run_once
+
+
+class WorkflowWorker:
+    """Distributed worker for processing workflow tasks with concurrency control.
+
+    Features:
+    - Concurrent task processing with configurable worker count
+    - Graceful shutdown on SIGINT/SIGTERM
+    - Automatic retry on transient failures
+    - Configurable polling interval
+    - Health monitoring and statistics
+    """
+
+    def __init__(
+        self,
+        workers: int = 4,
+        poll_interval: float = 0.5,
+        shutdown_timeout: float = 30.0,
+    ):
+        """Initialize the workflow worker.
+
+        Args:
+            workers: Number of concurrent task processors (default: 4)
+            poll_interval: Seconds between task queue polls (default: 0.5)
+            shutdown_timeout: Max seconds to wait for graceful shutdown (default: 30)
+        """
+        self.workers = workers
+        self.poll_interval = poll_interval
+        self.shutdown_timeout = shutdown_timeout
+        self._shutdown_event = asyncio.Event()
+        self._tasks: set[asyncio.Task] = set()
+        self._stats: dict[str, int | datetime | None] = {
+            "tasks_completed": 0,
+            "tasks_failed": 0,
+            "started_at": None,
+        }
+
+    async def start(self) -> None:
+        """Start the worker and process tasks until shutdown signal."""
+        self._stats["started_at"] = datetime.now(timezone.utc)
+
+        # Register signal handlers for graceful shutdown
+        self._register_signal_handlers()
+
+        print(f"Workflow worker started with {self.workers} concurrent workers")
+        print(f"Polling interval: {self.poll_interval}s")
+
+        try:
+            # Start worker tasks
+            for i in range(self.workers):
+                task = asyncio.create_task(
+                    self._worker_loop(worker_id=i), name=f"worker-{i}"
+                )
+                self._tasks.add(task)
+
+            # Wait for shutdown signal
+            await self._shutdown_event.wait()
+
+        finally:
+            await self._graceful_shutdown()
+
+    async def _worker_loop(self, worker_id: int) -> None:
+        """Main processing loop for a single worker.
+
+        Args:
+            worker_id: Unique identifier for this worker instance
+        """
+        while not self._shutdown_event.is_set():
+            try:
+                # Try to claim and execute a task
+                task_executed = await run_once()
+
+                if task_executed:
+                    self._stats["tasks_completed"] = int(self._stats["tasks_completed"]) + 1  # type: ignore
+                else:
+                    # No tasks available, wait before polling again
+                    await asyncio.sleep(self.poll_interval)
+
+            except asyncio.CancelledError:
+                # Graceful shutdown requested
+                break
+            except Exception as e:
+                self._stats["tasks_failed"] = int(self._stats["tasks_failed"]) + 1  # type: ignore
+                print(f"Worker {worker_id} error: {e}")
+                # Brief pause before retrying
+                await asyncio.sleep(1.0)
+
+    async def _graceful_shutdown(self) -> None:
+        """Gracefully shut down all worker tasks."""
+        print("\nShutting down gracefully...")
+
+        # Cancel all worker tasks
+        for task in self._tasks:
+            task.cancel()
+
+        # Wait for tasks to complete with timeout
+        try:
+            await asyncio.wait_for(
+                asyncio.gather(*self._tasks, return_exceptions=True),
+                timeout=self.shutdown_timeout,
+            )
+        except asyncio.TimeoutError:
+            print(f"Shutdown timeout reached ({self.shutdown_timeout}s)")
+
+        self._print_stats()
+        print("Worker shutdown complete")
+
+    def _register_signal_handlers(self) -> None:
+        """Register handlers for SIGINT and SIGTERM."""
+
+        def signal_handler(sig, frame):
+            print(f"\nReceived signal {signal.Signals(sig).name}")
+            self._shutdown_event.set()
+
+        # Handle Ctrl+C and termination signals
+        signal.signal(signal.SIGINT, signal_handler)
+        if sys.platform != "win32":
+            signal.signal(signal.SIGTERM, signal_handler)
+
+    def _print_stats(self) -> None:
+        """Print worker statistics."""
+        started_at = self._stats["started_at"]
+        if started_at and isinstance(started_at, datetime):
+            uptime = datetime.now(timezone.utc) - started_at
+            print("\nWorker Statistics:")
+            print(f"   Uptime: {uptime}")
+            print(f"   Tasks completed: {self._stats['tasks_completed']}")
+            print(f"   Tasks failed: {self._stats['tasks_failed']}")
+
+
+async def start_worker(workers: int = 4, poll_interval: float = 0.5) -> None:
+    """Start a workflow worker process.
+
+    Args:
+        workers: Number of concurrent task processors
+        poll_interval: Seconds between task queue polls
+    """
+    worker = WorkflowWorker(workers=workers, poll_interval=poll_interval)
+    await worker.start()

src/core/workflow.py

@@ -0,0 +1,168 @@
+import inspect
+from typing import Generic, List, TypeVar
+
+from ..schemas.workflow import InputT, StateT, Step
+from .compiled import CompiledWorkflow
+
+# For better type inference in classmethods
+Self = TypeVar("Self", bound="Workflow")
+
+
+class Workflow(Generic[InputT, StateT]):
+    """
+    Abstract base class for defining typed workflows.
+
+    This class provides the foundation for creating deterministic, durable workflows
+    with strong typing support. Workflows are parameterized by:
+    - InputT: The immutable input type for the workflow
+    - StateT: The mutable state type that evolves during execution
+
+    Example:
+        @dataclass
+        class MyInput:
+            user_id: str
+
+        @dataclass
+        class MyState:
+            processed: bool = False
+            result: str = ""
+
+        @loom.workflow
+        class MyWorkflow(Workflow[MyInput, MyState]):
+            @loom.step
+            async def process(self, ctx: WorkflowContext[MyState]):
+                # Workflow logic here
+                pass
+    """
+
+    @classmethod
+    def compile(cls) -> CompiledWorkflow[InputT, StateT]:
+        """
+        Compile the workflow definition directly from the class.
+
+        This is a convenience method that allows calling SomeWorkflow.compile()
+        instead of SomeWorkflow().compile().
+
+        Returns:
+            CompiledWorkflow: A compiled, immutable workflow definition ready for execution
+
+        Raises:
+            ValueError: If the workflow has no steps defined or is malformed
+        """
+        # Create instance and delegate to instance method
+        instance = cls()
+        return instance._compile_instance()
+
+    def _compile_instance(self) -> CompiledWorkflow[InputT, StateT]:
+        """
+        Internal instance compilation method.
+
+        This method introspects the class to extract:
+        - Workflow metadata (name, description, version)
+        - Step definitions and their order
+        - Validation of workflow structure
+
+        Returns:
+            CompiledWorkflow: A compiled, immutable workflow definition ready for execution
+
+        Raises:
+            ValueError: If the workflow has no steps defined or is malformed
+        """
+        # Extract workflow metadata with sensible defaults
+        name = self._get_workflow_name()
+        description = self._get_workflow_description()
+        version = self._get_workflow_version()
+        module = self._get_workflow_module()
+
+        # Discover and validate workflow steps
+        steps = self._discover_workflow_steps()
+
+        # Validate workflow structure
+        self._validate_workflow(steps)
+
+        return CompiledWorkflow[InputT, StateT](
+            name=name,
+            description=description,
+            version=version,
+            module=module,
+            steps=steps,
+        )
+
+    def _get_workflow_name(self) -> str:
+        """Get the workflow name from metadata or class name."""
+        return getattr(self, "_workflow_name", self.__class__.__name__)
+
+    def _get_workflow_description(self) -> str:
+        """Get the workflow description from metadata or docstring."""
+        explicit_desc = getattr(self, "_workflow_description", "")
+        if explicit_desc:
+            return explicit_desc
+
+        # Fallback to class docstring first line
+        docstring = self.__class__.__doc__
+        if docstring:
+            return docstring.strip().split("\n")[0]
+
+        return ""
+
+    def _get_workflow_version(self) -> str:
+        """Get the workflow version from metadata."""
+        return getattr(self, "_workflow_version", "1.0.0")
+
+    def _get_workflow_module(self) -> str:
+        """Get the workflow module path."""
+        return getattr(self, "_workflow_module", self.__class__.__module__)
+
+    def _discover_workflow_steps(self) -> List[Step]:
+        """
+        Discover all workflow steps by introspecting decorated methods.
+
+        Returns:
+            List[Step]: Ordered list of step definitions
+        """
+        steps: List[Step] = []
+        # Get all callable attributes that are decorated as steps
+        for attr_name in self.__class__.__dict__:
+            if attr_name.startswith("_"):
+                continue
+
+            attr = getattr(self, attr_name)
+            if callable(attr) and hasattr(attr, "_step_name"):
+                step_info: Step = {
+                    "name": getattr(attr, "_step_name"),
+                    "description": getattr(attr, "_step_description", ""),
+                    "fn": attr.__name__,
+                }
+                steps.append(step_info)
+
+        return steps
+
+    def _validate_workflow(self, steps: List[Step]) -> None:
+        """
+        Validate the workflow structure and step signatures.
+        Args:
+            steps (List[Step]): The list of discovered
+        Raises:
+            ValueError: If the workflow is malformed
+        """
+        if not steps:
+            raise ValueError(
+                f"Workflow '{self.__class__.__name__}' must have at least one step"
+            )
+        seen = set()
+
+        for step in steps:
+            name = step["name"]
+            if name in seen:
+                raise ValueError(f"Duplicate step name: {name}")
+            seen.add(name)
+
+            fn = getattr(self, step["fn"])
+            sig = inspect.signature(fn)
+            params = list(sig.parameters.values())
+
+            # bound method → first param is self
+            if len(params) != 1:
+                raise ValueError(
+                    f"Step '{name}' must have signature (self, ctx), " f"got {sig}"
+                )