more advanced planning features

Author: luke-beamcloud

examples/workflows.py

@@ -2,7 +2,17 @@
 Workflow examples - function-based and class-based styles.
 """
 
-from fastloop import FastLoop, LoopContext, LoopEvent, Workflow, WorkflowBlock
+from typing import Any
+
+from fastloop import (
+    BlockPlan,
+    FastLoop,
+    LoopContext,
+    LoopEvent,
+    ScheduleType,
+    Workflow,
+    WorkflowBlock,
+)
 
 app = FastLoop(name="workflow-demo")
 
@@ -127,6 +137,95 @@ async def execute(
                 print(f"Survey results: {answers}")
 
 
+@app.event("start_email_monitor")
+class StartEmailMonitor(LoopEvent):
+    pass
+
+
+async def email_monitor_plan(
+    _ctx: LoopContext,
+    _blocks: list[WorkflowBlock],
+    current_block: WorkflowBlock,
+    block_output: Any,
+) -> BlockPlan:
+    """Plan function decides next block and scheduling based on block output."""
+    match current_block.type:
+        case "search_emails":
+            if block_output and block_output.get("found_emails"):
+                return BlockPlan(schedule_type=ScheduleType.IMMEDIATE)
+            return BlockPlan(
+                next_block_index=0,
+                schedule_type=ScheduleType.DELAY,
+                delay_seconds=600,
+                reason="No emails found, retry in 10 minutes",
+            )
+        case "notify":
+            return BlockPlan(schedule_type=ScheduleType.STOP)
+        case _:
+            return BlockPlan(schedule_type=ScheduleType.IMMEDIATE)
+
+
+@app.workflow(
+    name="email_monitor",
+    start_event=StartEmailMonitor,
+    plan=email_monitor_plan,
+)
+async def email_monitor_workflow(
+    ctx: LoopContext,
+    _blocks: list[WorkflowBlock],
+    current_block: WorkflowBlock,
+) -> dict[str, Any]:
+    """Workflow that returns output for the plan function to decide next steps."""
+    print(f"[{ctx.block_index + 1}/{ctx.block_count}] {current_block.type}")
+
+    match current_block.type:
+        case "search_emails":
+            emails = []  # simulate search
+            return {"found_emails": len(emails) > 0, "emails": emails}
+        case "summarize":
+            return {"summary": "Email summary"}
+        case "notify":
+            return {"notified": True}
+        case _:
+            return {}
+
+
+@app.event("start_retry_demo")
+class StartRetryDemo(LoopEvent):
+    pass
+
+
+@app.workflow(name="retry_demo", start_event=StartRetryDemo)
+class RetryDemoWorkflow(Workflow):
+    """Class-based workflow with plan method for rate limit handling."""
+
+    async def plan(
+        self,
+        ctx: LoopContext,
+        _blocks: list[WorkflowBlock],
+        _current_block: WorkflowBlock,
+        block_output: Any,
+    ) -> BlockPlan | None:
+        if block_output and block_output.get("rate_limited"):
+            return BlockPlan(
+                next_block_index=ctx.block_index,
+                schedule_type=ScheduleType.DELAY,
+                delay_seconds=block_output.get("retry_after", 60),
+            )
+        if ctx.block_index >= ctx.block_count - 1:
+            return BlockPlan(schedule_type=ScheduleType.STOP)
+        return None
+
+    async def execute(
+        self,
+        ctx: LoopContext,
+        _blocks: list[WorkflowBlock],
+        current_block: WorkflowBlock,
+    ) -> dict[str, Any]:
+        print(f"[{ctx.block_index + 1}/{ctx.block_count}] {current_block.type}")
+        return {"success": True}
+
+
 if __name__ == "__main__":
     app.run(port=8111)
 
@@ -156,3 +255,21 @@ async def execute(
 #    curl http://localhost:8111/onboarding/<workflow_id>
 #
 # 5. If service restarts, workflow resumes from current block automatically
+#
+# 6. Start email monitor workflow with plan function:
+#    curl -X POST http://localhost:8111/email_monitor \
+#      -H "Content-Type: application/json" \
+#      -d '{
+#        "type": "start_email_monitor",
+#        "blocks": [
+#          {"type": "search_emails", "text": "Search for customer emails"},
+#          {"type": "summarize", "text": "Summarize found emails"},
+#          {"type": "notify", "text": "Send notification"}
+#        ]
+#      }'
+#
+#    The plan function will:
+#    - Retry the search block with a 10-minute delay if no emails found
+#    - Proceed immediately when emails are found
+#    - Stop the workflow after notification
+#    - Delays use Redis scheduling, so the workflow survives restarts

fastloop/__init__.py

@@ -2,14 +2,16 @@
 from .context import LoopContext
 from .fastloop import FastLoop
 from .loop import Loop, LoopEvent, Workflow, WorkflowBlock
-from .types import RetryPolicy
+from .types import BlockPlan, RetryPolicy, ScheduleType
 
 __all__ = [
+    "BlockPlan",
     "FastLoop",
     "Loop",
     "LoopContext",
     "LoopEvent",
     "RetryPolicy",
+    "ScheduleType",
     "Workflow",
     "WorkflowBlock",
     "integrations",

fastloop/context.py

@@ -10,6 +10,7 @@
     LoopContextSwitchError,
     LoopPausedError,
     LoopStoppedError,
+    WorkflowGotoError,
     WorkflowNextError,
     WorkflowRepeatError,
 )
@@ -254,3 +255,12 @@ def next(self, payload: dict | None = None) -> None:
     def repeat(self) -> None:
         """Re-execute the current block in a workflow."""
         raise WorkflowRepeatError()
+
+    def goto(
+        self,
+        block_index: int,
+        delay_seconds: float | None = None,
+        reason: str | None = None,
+    ) -> None:
+        """Jump to a specific block in a workflow, optionally after a delay."""
+        raise WorkflowGotoError(block_index, delay_seconds, reason)

fastloop/exceptions.py

@@ -50,6 +50,18 @@ class WorkflowRepeatError(Exception):
     pass
 
 
+class WorkflowGotoError(Exception):
+    def __init__(
+        self,
+        block_index: int,
+        delay_seconds: float | None = None,
+        reason: str | None = None,
+    ):
+        self.block_index = block_index
+        self.delay_seconds = delay_seconds
+        self.reason = reason
+
+
 class WorkflowNotFoundError(Exception):
     pass
 

fastloop/fastloop.py

@@ -460,6 +460,7 @@ def workflow(
         on_stop: Callable[..., Any] | None = None,
         on_block_complete: Callable[..., Any] | None = None,
         on_error: Callable[..., Any] | None = None,
+        plan: Callable[..., Any] | None = None,
         retry: RetryPolicy | None = None,
     ) -> Callable[
         [Callable[..., Any] | type[Workflow]], Callable[..., Any] | type[Workflow]
@@ -478,13 +479,15 @@ def _decorator(
                 workflow_on_stop = workflow_instance.on_stop
                 workflow_on_block_complete = workflow_instance.on_block_complete
                 workflow_on_error = workflow_instance.on_error
+                workflow_plan = getattr(workflow_instance, "plan", None) or plan
             else:
                 workflow_instance = None  # type: ignore
                 func = func_or_class  # type: ignore
                 workflow_on_start = on_start
                 workflow_on_stop = on_stop
                 workflow_on_block_complete = on_block_complete
                 workflow_on_error = on_error
+                workflow_plan = plan
 
             start_event_key = _resolve_event_key(start_event)
 
@@ -497,6 +500,7 @@ def _decorator(
                 "on_stop": workflow_on_stop,
                 "on_block_complete": workflow_on_block_complete,
                 "on_error": workflow_on_error,
+                "plan": workflow_plan,
                 "workflow_instance": workflow_instance,
                 "retry_policy": retry,
             }
@@ -561,6 +565,7 @@ async def _start_handler(request: dict[str, Any]):
                     on_stop=workflow_on_stop,
                     on_block_complete=workflow_on_block_complete,
                     on_error=workflow_on_error,
+                    plan=workflow_plan,
                     retry_policy=retry,
                 )
                 return (
@@ -745,6 +750,7 @@ async def restart_workflow(self, workflow_id: str) -> bool:
                 on_stop=metadata.get("on_stop"),
                 on_block_complete=metadata.get("on_block_complete"),
                 on_error=metadata.get("on_error"),
+                plan=metadata.get("plan"),
                 retry_policy=metadata.get("retry_policy"),
             )
 
@@ -789,12 +795,26 @@ def __init__(
     def stop(self) -> None:
         self._stop_event.set()
 
-    async def _process_wake(self, loop_id: str) -> None:
-        if await self.state_manager.has_claim(loop_id):
-            return
-        logger.info("Loop woke up, restarting", extra={"loop_id": loop_id})
-        if not await self.restart_callback(loop_id):
-            await self.state_manager.update_loop_status(loop_id, LoopStatus.STOPPED)
+    async def _process_wake(self, wake_id: str) -> None:
+        if wake_id.startswith("workflow:"):
+            workflow_id = wake_id[9:]
+            if await self.state_manager.workflow_has_claim(workflow_id):
+                return
+            logger.info(
+                "Workflow woke up, restarting", extra={"workflow_id": workflow_id}
+            )
+            await self.state_manager.clear_workflow_wake_time(workflow_id)
+            if not await self.fastloop_instance.restart_workflow(workflow_id):
+                await self.state_manager.update_workflow_status(
+                    workflow_id, LoopStatus.STOPPED
+                )
+        else:
+            loop_id = wake_id
+            if await self.state_manager.has_claim(loop_id):
+                return
+            logger.info("Loop woke up, restarting", extra={"loop_id": loop_id})
+            if not await self.restart_callback(loop_id):
+                await self.state_manager.update_loop_status(loop_id, LoopStatus.STOPPED)
 
     async def _check_orphaned_loops(self) -> None:
         running_loops = await self.state_manager.get_all_loops(