edda-framework 0.14.0__py3-none-any.whl → 0.15.0__py3-none-any.whl

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)

edda/integrations/llamaindex/__init__.py

@@ -0,0 +1,51 @@
+"""
+LlamaIndex Workflow Integration for Edda.
+
+This module provides integration between LlamaIndex Workflow and Edda's durable
+execution framework, making workflow execution crash-recoverable and supporting
+durable wait operations.
+
+Example:
+    from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent
+    from edda import workflow, WorkflowContext
+    from edda.integrations.llamaindex import DurableWorkflowRunner, DurableSleepEvent
+
+    # Define events
+    class ProcessedEvent(Event):
+        data: str
+
+    # Define workflow
+    class MyWorkflow(Workflow):
+        @step
+        async def process(self, ctx: Context, ev: StartEvent) -> ProcessedEvent:
+            return ProcessedEvent(data=f"processed: {ev.input}")
+
+        @step
+        async def finalize(self, ctx: Context, ev: ProcessedEvent) -> StopEvent:
+            return StopEvent(result={"status": "done", "data": ev.data})
+
+    # Create durable runner
+    runner = DurableWorkflowRunner(MyWorkflow)
+
+    # Use in Edda workflow
+    @workflow
+    async def my_workflow(ctx: WorkflowContext, input_data: str) -> dict:
+        result = await runner.run(ctx, input=input_data)
+        return result
+
+Installation:
+    pip install 'edda-framework[llamaindex]'
+"""
+
+from .events import DurableSleepEvent, DurableWaitEvent, ResumeEvent
+from .exceptions import WorkflowExecutionError, WorkflowReplayError
+from .workflow import DurableWorkflowRunner
+
+__all__ = [
+    "DurableWorkflowRunner",
+    "DurableSleepEvent",
+    "DurableWaitEvent",
+    "ResumeEvent",
+    "WorkflowExecutionError",
+    "WorkflowReplayError",
+]

edda/integrations/llamaindex/events.py

@@ -0,0 +1,160 @@
+"""Durable events for LlamaIndex Workflow integration.
+
+These events signal to the DurableWorkflow that a durable operation
+(sleep or wait for external event) should be performed.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+# Lazy import to avoid requiring llama-index at import time
+if TYPE_CHECKING:
+    pass
+
+
+def _import_event_class() -> type[Any]:
+    """Import Event class with helpful error message."""
+    try:
+        from llama_index.core.workflow import Event  # type: ignore[import-not-found]
+
+        return Event  # type: ignore[no-any-return]
+    except ImportError as e:
+        msg = (
+            "llama-index-core is not installed. Install with:\n"
+            "  pip install llama-index-core\n"
+            "or\n"
+            "  pip install 'edda-framework[llamaindex]'"
+        )
+        raise ImportError(msg) from e
+
+
+class DurableSleepEvent:
+    """
+    Event that signals a durable sleep operation.
+
+    When a step returns this event, the DurableWorkflow will:
+    1. Record the step completion
+    2. Call Edda's sleep() function (durable timer)
+    3. Resume with the specified resume_data after the sleep completes
+
+    Example:
+        @step
+        async def rate_limited_step(self, ctx: Context, ev: SomeEvent) -> DurableSleepEvent:
+            # Hit rate limit, need to wait
+            return DurableSleepEvent(
+                seconds=60,
+                resume_data={"retry_count": ev.retry_count + 1},
+            )
+    """
+
+    def __init__(
+        self,
+        seconds: float,
+        resume_data: dict[str, Any] | None = None,
+    ) -> None:
+        """
+        Initialize a durable sleep event.
+
+        Args:
+            seconds: Number of seconds to sleep
+            resume_data: Data to include when resuming after sleep
+        """
+        self.seconds = seconds
+        self.resume_data = resume_data or {}
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to dictionary."""
+        return {
+            "_type": "DurableSleepEvent",
+            "seconds": self.seconds,
+            "resume_data": self.resume_data,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> DurableSleepEvent:
+        """Deserialize from dictionary."""
+        return cls(
+            seconds=data["seconds"],
+            resume_data=data.get("resume_data", {}),
+        )
+
+
+class DurableWaitEvent:
+    """
+    Event that signals waiting for an external event.
+
+    When a step returns this event, the DurableWorkflow will:
+    1. Record the step completion
+    2. Call Edda's wait_event() function (durable event subscription)
+    3. Resume with the received event data after the event arrives
+
+    Example:
+        @step
+        async def wait_for_approval(self, ctx: Context, ev: OrderEvent) -> DurableWaitEvent:
+            return DurableWaitEvent(
+                event_type=f"approval.{ev.order_id}",
+                timeout_seconds=3600,  # 1 hour timeout
+            )
+    """
+
+    def __init__(
+        self,
+        event_type: str,
+        timeout_seconds: float | None = None,
+    ) -> None:
+        """
+        Initialize a durable wait event.
+
+        Args:
+            event_type: The event type to wait for (e.g., "payment.completed")
+            timeout_seconds: Optional timeout in seconds
+        """
+        self.event_type = event_type
+        self.timeout_seconds = timeout_seconds
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to dictionary."""
+        return {
+            "_type": "DurableWaitEvent",
+            "event_type": self.event_type,
+            "timeout_seconds": self.timeout_seconds,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> DurableWaitEvent:
+        """Deserialize from dictionary."""
+        return cls(
+            event_type=data["event_type"],
+            timeout_seconds=data.get("timeout_seconds"),
+        )
+
+
+class ResumeEvent:
+    """
+    Event used to resume workflow after a durable operation.
+
+    This is an internal event type used by DurableWorkflow to resume
+    execution after a DurableSleepEvent or DurableWaitEvent completes.
+    """
+
+    def __init__(self, data: dict[str, Any] | None = None) -> None:
+        """
+        Initialize a resume event.
+
+        Args:
+            data: Data from the completed operation (sleep resume_data or received event)
+        """
+        self.data = data or {}
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to dictionary."""
+        return {
+            "_type": "ResumeEvent",
+            "data": self.data,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ResumeEvent:
+        """Deserialize from dictionary."""
+        return cls(data=data.get("data", {}))

edda/integrations/llamaindex/exceptions.py

@@ -0,0 +1,15 @@
+"""Exceptions for LlamaIndex Workflow integration."""
+
+
+class WorkflowExecutionError(Exception):
+    """Error during workflow execution."""
+
+    def __init__(self, message: str, step_name: str | None = None) -> None:
+        self.step_name = step_name
+        super().__init__(message)
+
+
+class WorkflowReplayError(Exception):
+    """Error during workflow replay."""
+
+    pass

edda/integrations/llamaindex/workflow.py

@@ -0,0 +1,306 @@
+"""DurableWorkflow - makes LlamaIndex Workflow execution durable via Edda.
+
+This module provides integration between LlamaIndex Workflow and Edda's durable
+execution framework, making workflow execution crash-recoverable and supporting
+durable wait operations.
+"""
+
+from __future__ import annotations
+
+import importlib
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from edda.activity import activity
+from edda.pydantic_utils import to_json_dict
+
+from .events import DurableSleepEvent, DurableWaitEvent, ResumeEvent
+from .exceptions import WorkflowExecutionError
+
+if TYPE_CHECKING:
+    from edda.context import WorkflowContext
+
+T = TypeVar("T")
+
+
+def _import_llamaindex_workflow() -> Any:
+    """Import llama_index.core.workflow with helpful error message."""
+    try:
+        from llama_index.core import workflow  # type: ignore[import-not-found]
+
+        return workflow
+    except ImportError as e:
+        msg = (
+            "llama-index-core is not installed. Install with:\n"
+            "  pip install llama-index-core\n"
+            "or\n"
+            "  pip install 'edda-framework[llamaindex]'"
+        )
+        raise ImportError(msg) from e
+
+
+def _serialize_event(event: Any) -> dict[str, Any]:
+    """Serialize a LlamaIndex Event to a dictionary."""
+    if isinstance(event, DurableSleepEvent):
+        return event.to_dict()
+    if isinstance(event, DurableWaitEvent):
+        return event.to_dict()
+    if isinstance(event, ResumeEvent):
+        return event.to_dict()
+
+    # For LlamaIndex events, use model_dump if available (Pydantic)
+    if hasattr(event, "model_dump"):
+        data = event.model_dump()
+    elif hasattr(event, "__dict__"):
+        data = {k: v for k, v in event.__dict__.items() if not k.startswith("_")}
+    else:
+        data = {}
+
+    return {
+        "_type": f"{event.__class__.__module__}:{event.__class__.__qualname__}",
+        "_data": to_json_dict(data),
+    }
+
+
+def _deserialize_event(data: dict[str, Any]) -> Any:
+    """Deserialize a dictionary to a LlamaIndex Event."""
+    event_type = data.get("_type", "")
+
+    # Handle our special events
+    if event_type == "DurableSleepEvent":
+        return DurableSleepEvent.from_dict(data)
+    if event_type == "DurableWaitEvent":
+        return DurableWaitEvent.from_dict(data)
+    if event_type == "ResumeEvent":
+        return ResumeEvent.from_dict(data)
+
+    # Handle LlamaIndex events
+    if ":" in event_type:
+        module_path, class_name = event_type.rsplit(":", 1)
+        module = importlib.import_module(module_path)
+        event_class = getattr(module, class_name)
+        event_data = data.get("_data", {})
+
+        # Use model_validate for Pydantic models
+        if hasattr(event_class, "model_validate"):
+            return event_class.model_validate(event_data)
+        return event_class(**event_data)
+
+    raise ValueError(f"Unknown event type: {event_type}")
+
+
+@dataclass
+class StepResult:
+    """Result of a step execution."""
+
+    event: Any
+    step_name: str
+
+
+@activity
+async def _run_step(
+    ctx: WorkflowContext,  # noqa: ARG001 - Used by @activity decorator
+    workflow_class_path: str,
+    step_name: str,
+    event_data: dict[str, Any],
+    context_data: dict[str, Any],
+) -> dict[str, Any]:
+    """
+    Execute a single workflow step as a durable activity.
+
+    This activity is the core of DurableWorkflow - it runs one step and returns
+    the serialized result event.
+    """
+    # Import the workflow class
+    module_path, class_name = workflow_class_path.rsplit(":", 1)
+    module = importlib.import_module(module_path)
+    workflow_class = getattr(module, class_name)
+
+    # Create workflow instance
+    workflow_instance = workflow_class()
+
+    # Deserialize the input event
+    input_event = _deserialize_event(event_data)
+
+    # Get the step method
+    step_method = getattr(workflow_instance, step_name, None)
+    if step_method is None:
+        raise WorkflowExecutionError(f"Step '{step_name}' not found", step_name)
+
+    # Create a minimal context for the step
+    # Note: LlamaIndex Context is created per-run, we create a simple mock
+
+    # Execute the step
+    try:
+        # The step method signature is: async def step(self, ctx, event) -> Event
+        # We need to provide a context - use a simple object with store
+        @dataclass
+        class SimpleContext:
+            store: dict[str, Any]
+
+        simple_ctx = SimpleContext(store=context_data.get("store", {}))
+        result_event = await step_method(simple_ctx, input_event)
+    except Exception as e:
+        raise WorkflowExecutionError(f"Step '{step_name}' failed: {e}", step_name) from e
+
+    # Serialize the result
+    return {
+        "event": _serialize_event(result_event),
+        "context_store": simple_ctx.store,
+    }
+
+
+class DurableWorkflowRunner:
+    """
+    Runner that executes a LlamaIndex Workflow with Edda durability.
+
+    This class wraps a LlamaIndex Workflow and executes it step-by-step,
+    recording each step as an Edda Activity for crash recovery.
+
+    Example:
+        from llama_index.core.workflow import Workflow, step, StartEvent, StopEvent
+
+        class MyWorkflow(Workflow):
+            @step
+            async def process(self, ctx: Context, ev: StartEvent) -> StopEvent:
+                return StopEvent(result="done")
+
+        runner = DurableWorkflowRunner(MyWorkflow)
+
+        @workflow
+        async def my_edda_workflow(ctx: WorkflowContext) -> str:
+            result = await runner.run(ctx, input_data="hello")
+            return result
+    """
+
+    def __init__(self, workflow_class: type) -> None:
+        """
+        Initialize DurableWorkflowRunner.
+
+        Args:
+            workflow_class: A LlamaIndex Workflow class (not instance)
+        """
+        self._workflow_class = workflow_class
+        self._class_path = f"{workflow_class.__module__}:{workflow_class.__qualname__}"
+
+        # Validate it's a Workflow subclass
+        llamaindex_workflow = _import_llamaindex_workflow()
+        if not issubclass(workflow_class, llamaindex_workflow.Workflow):
+            raise TypeError(f"Expected a Workflow subclass, got {type(workflow_class).__name__}")
+
+    async def run(
+        self,
+        ctx: WorkflowContext,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Execute the workflow durably with Edda crash recovery.
+
+        Args:
+            ctx: Edda WorkflowContext
+            **kwargs: Arguments passed to StartEvent
+
+        Returns:
+            The result from StopEvent
+        """
+        from edda.channels import sleep as edda_sleep
+        from edda.channels import wait_event as edda_wait_event
+
+        llamaindex_workflow = _import_llamaindex_workflow()
+
+        # Create a workflow instance to analyze its steps
+        workflow_instance = self._workflow_class()
+
+        # Build step registry from the workflow
+        step_registry = self._build_step_registry(workflow_instance)
+
+        # Start with StartEvent
+        start_event_class = llamaindex_workflow.StartEvent
+        current_event = start_event_class(**kwargs)
+        context_store: dict[str, Any] = {}
+
+        # Main execution loop
+        while True:
+            # Find the step that handles this event type
+            event_type = type(current_event)
+            step_name = self._find_step_for_event(step_registry, event_type)
+
+            if step_name is None:
+                # No step found - check if it's a stop event
+                if isinstance(current_event, llamaindex_workflow.StopEvent):
+                    return current_event.result
+                raise WorkflowExecutionError(f"No step found for event type: {event_type.__name__}")
+
+            # Execute the step as an activity
+            result = await _run_step(  # type: ignore[misc,call-arg]
+                ctx,  # type: ignore[arg-type]
+                self._class_path,
+                step_name,
+                _serialize_event(current_event),
+                {"store": context_store},
+            )
+
+            # Update context store
+            context_store = result.get("context_store", {})
+
+            # Deserialize the result event
+            result_event = _deserialize_event(result["event"])
+
+            # Handle special durable events
+            if isinstance(result_event, DurableSleepEvent):
+                # Durable sleep
+                await edda_sleep(ctx, int(result_event.seconds))
+                # Resume with ResumeEvent containing the sleep's resume_data
+                current_event = ResumeEvent(data=result_event.resume_data)
+
+            elif isinstance(result_event, DurableWaitEvent):
+                # Durable wait for external event
+                received = await edda_wait_event(
+                    ctx,
+                    result_event.event_type,
+                    timeout_seconds=(
+                        int(result_event.timeout_seconds) if result_event.timeout_seconds else None
+                    ),
+                )
+                # Resume with ResumeEvent containing the received data
+                current_event = ResumeEvent(data=received.data if hasattr(received, "data") else {})
+
+            elif isinstance(result_event, llamaindex_workflow.StopEvent):
+                # Workflow completed
+                return result_event.result
+
+            else:
+                # Normal event transition
+                current_event = result_event
+
+    def _build_step_registry(self, workflow_instance: Any) -> dict[str, list[type]]:
+        """Build a registry mapping step names to their input event types."""
+        registry: dict[str, list[type]] = {}
+
+        # Look for methods decorated with @step in the class
+        workflow_class = type(workflow_instance)
+        for name in dir(workflow_class):
+            if name.startswith("_"):
+                continue
+
+            # Get the raw function from the class (not bound method)
+            method = getattr(workflow_class, name, None)
+            if method is None or not callable(method):
+                continue
+
+            # Check if it's a step (has _step_config attribute from @step decorator)
+            if hasattr(method, "_step_config"):
+                step_config = method._step_config
+                # Get accepted event types directly from step config
+                if hasattr(step_config, "accepted_events"):
+                    registry[name] = list(step_config.accepted_events)
+
+        return registry
+
+    def _find_step_for_event(self, registry: dict[str, list[type]], event_type: type) -> str | None:
+        """Find the step that handles the given event type."""
+        for step_name, accepted_types in registry.items():
+            for accepted_type in accepted_types:
+                if issubclass(event_type, accepted_type):
+                    return step_name
+        return None