edda-framework 0.14.1__py3-none-any.whl → 0.15.1__py3-none-any.whl

edda/integrations/graph/__init__.py

@@ -0,0 +1,58 @@
+"""
+Durable Graph Integration for Edda.
+
+This module provides integration between pydantic-graph and Edda's durable
+execution framework, making pydantic-graph execution crash-recoverable and
+supporting durable wait operations.
+
+Example:
+    from dataclasses import dataclass
+    from pydantic_graph import BaseNode, Graph, End
+    from edda import workflow, WorkflowContext
+    from edda.integrations.graph import DurableGraph, DurableGraphContext
+
+    @dataclass
+    class MyState:
+        counter: int = 0
+
+    @dataclass
+    class IncrementNode(BaseNode[MyState, None, int]):
+        async def run(self, ctx: DurableGraphContext) -> "CheckNode":
+            ctx.state.counter += 1
+            return CheckNode()
+
+    @dataclass
+    class CheckNode(BaseNode[MyState, None, int]):
+        async def run(self, ctx: DurableGraphContext) -> IncrementNode | End[int]:
+            if ctx.state.counter >= 5:
+                return End(ctx.state.counter)
+            return IncrementNode()
+
+    graph = Graph(nodes=[IncrementNode, CheckNode])
+    durable = DurableGraph(graph)
+
+    @workflow
+    async def counter_workflow(ctx: WorkflowContext) -> int:
+        return await durable.run(
+            ctx,
+            start_node=IncrementNode(),
+            state=MyState(),
+        )
+
+Installation:
+    pip install 'edda-framework[graph]'
+"""
+
+from .context import DurableGraphContext
+from .exceptions import GraphExecutionError
+from .graph import DurableGraph
+from .nodes import ReceivedEvent, Sleep, WaitForEvent
+
+__all__ = [
+    "DurableGraph",
+    "DurableGraphContext",
+    "GraphExecutionError",
+    "ReceivedEvent",
+    "Sleep",
+    "WaitForEvent",
+]

edda/integrations/graph/context.py

@@ -0,0 +1,81 @@
+"""DurableGraphContext - bridges pydantic-graph and Edda contexts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from edda.context import WorkflowContext
+
+    from .nodes import ReceivedEvent
+
+StateT = TypeVar("StateT")
+DepsT = TypeVar("DepsT")
+
+
+@dataclass
+class DurableGraphContext(Generic[StateT, DepsT]):
+    """
+    Context that bridges pydantic-graph and Edda.
+
+    Provides access to:
+    - pydantic-graph's state and deps via properties
+    - last_event: The most recent event received via WaitForEvent
+
+    This context is passed to node's run() method when executing
+    via DurableGraph.
+
+    For durable wait operations (wait_event, sleep), use the WaitForEvent
+    and Sleep marker nodes instead of calling methods directly:
+
+        from edda.integrations.graph import WaitForEvent, Sleep
+
+        @dataclass
+        class MyNode(BaseNode[MyState, None, str]):
+            async def run(self, ctx: DurableGraphContext) -> WaitForEvent[NextNode]:
+                # Return a marker to wait for an event
+                return WaitForEvent(
+                    event_type="payment.completed",
+                    next_node=NextNode(),
+                    timeout_seconds=3600,
+                )
+
+        @dataclass
+        class NextNode(BaseNode[MyState, None, str]):
+            async def run(self, ctx: DurableGraphContext) -> End[str]:
+                # Access the received event
+                event = ctx.last_event
+                return End(event.data.get("status", "unknown"))
+
+    Attributes:
+        state: The graph state object (mutable, shared across nodes)
+        deps: The dependencies object (immutable)
+        last_event: The most recent event received via WaitForEvent (or None)
+        workflow_ctx: The Edda WorkflowContext
+    """
+
+    _state: StateT
+    _deps: DepsT
+    workflow_ctx: WorkflowContext
+    last_event: ReceivedEvent | None = field(default=None)
+
+    @property
+    def state(self) -> StateT:
+        """Get the graph state object."""
+        return self._state
+
+    @property
+    def deps(self) -> DepsT:
+        """Get the dependencies object."""
+        return self._deps
+
+    @property
+    def instance_id(self) -> str:
+        """Get the workflow instance ID."""
+        return self.workflow_ctx.instance_id
+
+    @property
+    def is_replaying(self) -> bool:
+        """Check if the workflow is currently replaying."""
+        return self.workflow_ctx.is_replaying

edda/integrations/graph/exceptions.py

@@ -0,0 +1,9 @@
+"""Exceptions for durable graph integration."""
+
+
+class GraphExecutionError(Exception):
+    """Raised when a graph node execution fails."""
+
+    def __init__(self, message: str, node_name: str | None = None) -> None:
+        self.node_name = node_name
+        super().__init__(message)

edda/integrations/graph/graph.py

@@ -0,0 +1,385 @@
+"""DurableGraph - makes pydantic-graph execution durable via Edda."""
+
+from __future__ import annotations
+
+import dataclasses
+import importlib
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+from edda.activity import activity
+from edda.pydantic_utils import to_json_dict
+
+from .context import DurableGraphContext
+from .exceptions import GraphExecutionError
+from .nodes import ReceivedEvent, Sleep, WaitForEvent
+
+if TYPE_CHECKING:
+    from edda.context import WorkflowContext
+
+StateT = TypeVar("StateT")
+DepsT = TypeVar("DepsT")
+RunEndT = TypeVar("RunEndT")
+
+
+def _import_pydantic_graph() -> Any:
+    """Import pydantic_graph with helpful error message."""
+    try:
+        import pydantic_graph
+
+        return pydantic_graph
+    except ImportError as e:
+        msg = (
+            "pydantic-graph is not installed. Install with:\n"
+            "  pip install pydantic-graph\n"
+            "or\n"
+            "  pip install 'edda-framework[graph]'"
+        )
+        raise ImportError(msg) from e
+
+
+def _get_class_path(cls: type) -> str:
+    """Get fully qualified class path for serialization."""
+    return f"{cls.__module__}:{cls.__qualname__}"
+
+
+def _import_class(path: str) -> type:
+    """Import a class from its fully qualified path."""
+    module_path, class_name = path.rsplit(":", 1)
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)  # type: ignore[no-any-return]
+
+
+def _serialize_node(node: Any) -> dict[str, Any]:
+    """Serialize a node to a dict."""
+    if dataclasses.is_dataclass(node) and not isinstance(node, type):
+        return {
+            "_class_path": _get_class_path(node.__class__),
+            "_data": dataclasses.asdict(node),
+        }
+    return {
+        "_class_path": _get_class_path(node.__class__),
+        "_data": {},
+    }
+
+
+def _deserialize_node(data: dict[str, Any]) -> Any:
+    """Deserialize a node from a dict."""
+    cls = _import_class(data["_class_path"])
+    return cls(**data.get("_data", {}))
+
+
+def _serialize_state(state: Any) -> dict[str, Any]:
+    """Serialize state to a dict."""
+    if state is None:
+        return {"_none": True}
+    if dataclasses.is_dataclass(state) and not isinstance(state, type):
+        return {
+            "_class_path": _get_class_path(state.__class__),
+            "_data": dataclasses.asdict(state),
+        }
+    if hasattr(state, "model_dump"):
+        return {
+            "_class_path": _get_class_path(state.__class__),
+            "_data": state.model_dump(),
+        }
+    return {"_raw": str(state)}
+
+
+def _serialize_deps(deps: Any) -> dict[str, Any] | None:
+    """Serialize deps to a dict."""
+    if deps is None:
+        return None
+    if dataclasses.is_dataclass(deps) and not isinstance(deps, type):
+        return {
+            "_class_path": _get_class_path(deps.__class__),
+            "_data": dataclasses.asdict(deps),
+        }
+    if hasattr(deps, "model_dump"):
+        return {
+            "_class_path": _get_class_path(deps.__class__),
+            "_data": deps.model_dump(),
+        }
+    # For simple types (int, str, etc.), return as-is wrapped
+    return {"_value": deps}
+
+
+def _deserialize_deps(data: dict[str, Any] | None) -> Any:
+    """Deserialize deps from a dict."""
+    if data is None:
+        return None
+    if "_value" in data:
+        return data["_value"]
+    cls = _import_class(data["_class_path"])
+    if dataclasses.is_dataclass(cls):
+        return cls(**data["_data"])
+    if hasattr(cls, "model_validate"):
+        return cls.model_validate(data["_data"])
+    return cls(**data["_data"])
+
+
+def _deserialize_state(data: dict[str, Any]) -> Any:
+    """Deserialize state from a dict."""
+    if data.get("_none"):
+        return None
+    if "_raw" in data:
+        raise ValueError(f"Cannot deserialize state from raw: {data['_raw']}")
+    cls = _import_class(data["_class_path"])
+    if dataclasses.is_dataclass(cls):
+        return cls(**data["_data"])
+    if hasattr(cls, "model_validate"):
+        return cls.model_validate(data["_data"])
+    return cls(**data["_data"])
+
+
+def _restore_state(source: Any, target: Any) -> None:
+    """Copy state from source to target object."""
+    if dataclasses.is_dataclass(source) and not isinstance(source, type):
+        for field in dataclasses.fields(source):
+            setattr(target, field.name, getattr(source, field.name))
+    elif hasattr(source, "__dict__"):
+        for key, value in source.__dict__.items():
+            if not key.startswith("_"):
+                setattr(target, key, value)
+
+
+@activity
+async def _run_graph_node(
+    ctx: WorkflowContext,
+    node_data: dict[str, Any],
+    state_data: dict[str, Any],
+    deps_data: dict[str, Any] | None,
+    last_event_data: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """
+    Execute a single graph node as a durable activity.
+
+    This activity is the core of DurableGraph - it runs one node and returns
+    the serialized result (next node, End, WaitForEvent, or Sleep) along
+    with the updated state.
+    """
+    pg = _import_pydantic_graph()
+
+    # Deserialize node, state, and deps
+    node = _deserialize_node(node_data)
+    state = _deserialize_state(state_data)
+    deps = _deserialize_deps(deps_data)
+
+    # Reconstruct last_event if provided
+    last_event: ReceivedEvent | None = None
+    if last_event_data:
+        last_event = ReceivedEvent(
+            event_type=last_event_data.get("event_type", ""),
+            data=last_event_data.get("data", {}),
+            metadata=last_event_data.get("metadata", {}),
+        )
+
+    # Create durable context
+    durable_ctx = DurableGraphContext(
+        _state=state,
+        _deps=deps,
+        workflow_ctx=ctx,
+        last_event=last_event,
+    )
+
+    try:
+        # Execute the node
+        result = await node.run(durable_ctx)
+
+        # Serialize result based on type
+        if isinstance(result, pg.End):
+            return {
+                "_type": "End",
+                "_data": to_json_dict(result.data),
+                "_state": _serialize_state(state),
+            }
+        elif isinstance(result, WaitForEvent):
+            return {
+                "_type": "WaitForEvent",
+                "_event_type": result.event_type,
+                "_timeout_seconds": result.timeout_seconds,
+                "_next_node": _serialize_node(result.next_node),
+                "_state": _serialize_state(state),
+            }
+        elif isinstance(result, Sleep):
+            return {
+                "_type": "Sleep",
+                "_seconds": result.seconds,
+                "_next_node": _serialize_node(result.next_node),
+                "_state": _serialize_state(state),
+            }
+        else:
+            # Regular node transition
+            return {
+                "_type": "Node",
+                "_node": _serialize_node(result),
+                "_state": _serialize_state(state),
+            }
+
+    except Exception as e:
+        raise GraphExecutionError(
+            f"Node {node.__class__.__name__} failed: {e}",
+            node.__class__.__name__,
+        ) from e
+
+
+class DurableGraph(Generic[StateT, DepsT, RunEndT]):
+    """
+    Wrapper that makes pydantic-graph execution durable.
+
+    DurableGraph wraps a pydantic-graph Graph and executes it with Edda's
+    durability guarantees:
+
+    - Each node execution is recorded as an Edda Activity
+    - On replay, completed nodes return cached results (no re-execution)
+    - Crash recovery: workflows resume from the last completed node
+    - WaitForEvent/Sleep markers enable durable wait operations
+
+    Example:
+        from dataclasses import dataclass
+        from pydantic_graph import BaseNode, Graph, End
+        from edda import workflow, WorkflowContext
+        from edda.integrations.graph import (
+            DurableGraph,
+            DurableGraphContext,
+            WaitForEvent,
+        )
+
+        @dataclass
+        class OrderState:
+            order_id: str | None = None
+
+        @dataclass
+        class ProcessOrder(BaseNode[OrderState, None, str]):
+            order_id: str
+
+            async def run(self, ctx: DurableGraphContext) -> WaitForEvent[WaitPayment]:
+                ctx.state.order_id = self.order_id
+                return WaitForEvent(
+                    event_type="payment.completed",
+                    next_node=WaitPayment(),
+                )
+
+        @dataclass
+        class WaitPayment(BaseNode[OrderState, None, str]):
+            async def run(self, ctx: DurableGraphContext) -> End[str]:
+                # Access the received event
+                event = ctx.last_event
+                if event and event.data.get("status") == "success":
+                    return End("completed")
+                return End("failed")
+
+        graph = Graph(nodes=[ProcessOrder, WaitPayment])
+        durable = DurableGraph(graph)
+
+        @workflow
+        async def order_workflow(ctx: WorkflowContext, order_id: str) -> str:
+            return await durable.run(
+                ctx,
+                start_node=ProcessOrder(order_id=order_id),
+                state=OrderState(),
+            )
+    """
+
+    def __init__(self, graph: Any) -> None:
+        """
+        Initialize DurableGraph wrapper.
+
+        Args:
+            graph: A pydantic-graph Graph instance
+
+        Raises:
+            TypeError: If graph is not a pydantic-graph Graph instance
+        """
+        pg = _import_pydantic_graph()
+        if not isinstance(graph, pg.Graph):
+            raise TypeError(f"Expected pydantic_graph.Graph, got {type(graph).__name__}")
+        self._graph = graph
+
+    @property
+    def graph(self) -> Any:
+        """Get the underlying pydantic-graph Graph instance."""
+        return self._graph
+
+    async def run(
+        self,
+        ctx: WorkflowContext,
+        start_node: Any,
+        *,
+        state: StateT,
+        deps: DepsT = None,  # type: ignore[assignment]
+    ) -> RunEndT:
+        """
+        Execute the graph durably with Edda crash recovery.
+
+        Args:
+            ctx: Edda WorkflowContext
+            start_node: The initial node to start execution from
+            state: Initial graph state (will be mutated during execution)
+            deps: Optional dependencies accessible via ctx.deps
+
+        Returns:
+            The final result (End.data value)
+
+        Raises:
+            GraphExecutionError: If graph execution fails
+        """
+        from edda.channels import sleep as edda_sleep
+        from edda.channels import wait_event as edda_wait_event
+
+        current_node = start_node
+        last_event_data: dict[str, Any] | None = None
+
+        # Execute nodes until End is reached
+        while True:
+            # Serialize inputs
+            node_data = _serialize_node(current_node)
+            state_data = _serialize_state(state)
+            deps_data = _serialize_deps(deps)
+
+            # Run node as activity (handles replay/caching automatically)
+            # The @activity decorator transforms the function signature
+            result = await _run_graph_node(  # type: ignore[misc,call-arg]
+                ctx,  # type: ignore[arg-type]
+                node_data,
+                state_data,
+                deps_data,
+                last_event_data,
+            )
+
+            # Restore state from result
+            restored_state = _deserialize_state(result["_state"])
+            _restore_state(restored_state, state)
+
+            # Handle result based on type
+            if result["_type"] == "End":
+                return result["_data"]  # type: ignore[no-any-return]
+
+            elif result["_type"] == "WaitForEvent":
+                # Wait for event at workflow level (outside activity)
+                event = await edda_wait_event(
+                    ctx,
+                    result["_event_type"],
+                    timeout_seconds=result.get("_timeout_seconds"),
+                )
+                # Store event data for next node
+                # Note: edda.channels.ReceivedEvent uses 'type' not 'event_type'
+                last_event_data = {
+                    "event_type": getattr(event, "type", result["_event_type"]),
+                    "data": event.data if isinstance(event.data, dict) else {},
+                    "metadata": getattr(event, "extensions", {}),
+                }
+                # Move to next node
+                current_node = _deserialize_node(result["_next_node"])
+
+            elif result["_type"] == "Sleep":
+                # Sleep at workflow level (outside activity)
+                await edda_sleep(ctx, result["_seconds"])
+                # Clear last_event since this wasn't an event wait
+                last_event_data = None
+                # Move to next node
+                current_node = _deserialize_node(result["_next_node"])
+
+            else:
+                # Regular node transition
+                last_event_data = None  # Clear last_event for regular transitions
+                current_node = _deserialize_node(result["_node"])

edda/integrations/graph/nodes.py

@@ -0,0 +1,144 @@
+"""Marker nodes for durable graph operations.
+
+These are special marker classes that tell DurableGraph to perform
+workflow-level operations (wait_event, sleep) outside of activities.
+
+This design keeps activities pure (atomic, retryable units of work)
+while allowing graphs to wait for external events or sleep.
+
+These classes inherit from pydantic-graph's BaseNode so they can be:
+1. Included in return type annotations without type: ignore
+2. Registered in Graph for proper type validation
+3. Detected by graph visualization tools
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+NextT = TypeVar("NextT")
+
+# Try to import BaseNode for inheritance
+# If pydantic-graph is not installed, use a fallback
+try:
+    from pydantic_graph import BaseNode
+
+    _HAS_PYDANTIC_GRAPH = True
+except ImportError:
+    _HAS_PYDANTIC_GRAPH = False
+
+    # Fallback base class when pydantic-graph is not installed
+    class BaseNode:  # type: ignore[no-redef]
+        pass
+
+
+if TYPE_CHECKING:
+    from pydantic_graph import BaseNode as _BaseNode
+
+    # For type checking, we need the actual BaseNode
+    _MarkerBase = _BaseNode[Any, Any, Any]
+else:
+    _MarkerBase = BaseNode
+
+
+@dataclass
+class WaitForEvent(_MarkerBase, Generic[NextT]):  # type: ignore[misc,valid-type]
+    """
+    Marker node that tells DurableGraph to wait for an external event.
+
+    When a node returns this marker, DurableGraph will:
+    1. Complete the current node's activity
+    2. Call wait_event() at the workflow level (outside activities)
+    3. Store the received event data in ctx.last_event
+    4. Continue execution with next_node
+
+    IMPORTANT: Register this class in your Graph for type checking:
+        graph = Graph(nodes=[MyNode1, MyNode2, WaitForEvent])
+
+    Example:
+        @dataclass
+        class WaitForPaymentNode(BaseNode[OrderState, None, str]):
+            async def run(
+                self, ctx: DurableGraphContext
+            ) -> WaitForEvent[ProcessPaymentNode] | End[str]:
+                ctx.state.waiting_for = "payment"
+                return WaitForEvent(
+                    event_type=f"payment.{ctx.state.order_id}",
+                    next_node=ProcessPaymentNode(),
+                    timeout_seconds=3600,
+                )
+
+        @dataclass
+        class ProcessPaymentNode(BaseNode[OrderState, None, str]):
+            async def run(self, ctx: DurableGraphContext) -> End[str]:
+                event = ctx.last_event
+                if event.data.get("status") == "success":
+                    return End("payment_received")
+                return End("payment_failed")
+
+        # Register WaitForEvent in the Graph
+        graph = Graph(nodes=[WaitForPaymentNode, ProcessPaymentNode, WaitForEvent])
+    """
+
+    event_type: str
+    next_node: NextT
+    timeout_seconds: int | None = None
+
+    async def run(self, _ctx: Any) -> Any:
+        """Never called - DurableGraph intercepts this marker."""
+        raise RuntimeError(
+            "WaitForEvent marker should not be executed directly. "
+            "Use DurableGraph.run() instead of Graph.run()."
+        )
+
+
+@dataclass
+class Sleep(_MarkerBase, Generic[NextT]):  # type: ignore[misc,valid-type]
+    """
+    Marker node that tells DurableGraph to sleep before continuing.
+
+    When a node returns this marker, DurableGraph will:
+    1. Complete the current node's activity
+    2. Call sleep() at the workflow level (outside activities)
+    3. Continue execution with next_node
+
+    IMPORTANT: Register this class in your Graph for type checking:
+        graph = Graph(nodes=[MyNode1, MyNode2, Sleep])
+
+    Example:
+        @dataclass
+        class RateLimitNode(BaseNode[ApiState, None, str]):
+            async def run(
+                self, ctx: DurableGraphContext
+            ) -> Sleep[RetryApiNode] | End[str]:
+                if rate_limited:
+                    return Sleep(seconds=60, next_node=RetryApiNode())
+                return End("success")
+
+        # Register Sleep in the Graph
+        graph = Graph(nodes=[RateLimitNode, RetryApiNode, Sleep])
+    """
+
+    seconds: int
+    next_node: NextT
+
+    async def run(self, _ctx: Any) -> Any:
+        """Never called - DurableGraph intercepts this marker."""
+        raise RuntimeError(
+            "Sleep marker should not be executed directly. "
+            "Use DurableGraph.run() instead of Graph.run()."
+        )
+
+
+@dataclass
+class ReceivedEvent:
+    """
+    Event data received from wait_event.
+
+    This is stored in DurableGraphContext.last_event after WaitForEvent completes.
+    """
+
+    event_type: str
+    data: dict[str, Any] = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)