Python: [BREAKING] Python: Make executor ID required, improvements around handling rehydrating checkpoints (#832)

* Make executor ID required, improvements around handling rehydrating checkpoints.

* Duplicate executor validation added

* fix remaining issues

---------

Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
This commit is contained in:
Evan Mattson
2025-09-20 03:57:09 +09:00
committed by GitHub
Unverified
parent 7cd45e313b
commit aba094b5cf
33 changed files with 1967 additions and 275 deletions
@@ -43,6 +43,7 @@ Once comfortable with these, explore the rest of the samples below.
| Sample | File | Concepts |
|---|---|---|
| Checkpoint & Resume | [checkpoint/checkpoint_with_resume.py](./checkpoint/checkpoint_with_resume.py) | Create checkpoints, inspect them, and resume execution |
| Checkpoint & HITL Resume | [checkpoint/checkpoint_with_human_in_the_loop.py](./checkpoint/checkpoint_with_human_in_the_loop.py) | Combine checkpointing with human approvals and resume pending HITL requests |
### composition
| Sample | File | Concepts |
@@ -50,7 +50,7 @@ Prerequisites
# - Compute a result
# - Forward that result to downstream node(s) using ctx.send_message(result)
class UpperCase(Executor):
def __init__(self, id: str | None = None):
def __init__(self, id: str):
super().__init__(id=id)
@handler
@@ -1,108 +1,34 @@
# Copyright (c) Microsoft. All rights reserved.
# import asyncio
# from agent_framework.foundry import FoundryChatClient
# from agent_framework import AgentRunUpdateEvent, WorkflowBuilder, WorkflowCompletedEvent
# from azure.identity.aio import AzureCliCredential
# """
# Sample: Agents in a workflow with streaming
# A Writer agent generates content, then a Reviewer agent critiques it.
# The workflow uses streaming so you can observe incremental AgentRunUpdateEvent chunks as each agent produces tokens.
# Purpose:
# Show how to wire chat agents directly into a WorkflowBuilder pipeline where agents are auto wrapped as executors.
# Demonstrate:
# - Automatic streaming of agent deltas via AgentRunUpdateEvent.
# - A simple console aggregator that groups updates by executor id and prints them as they arrive.
# - A final WorkflowCompletedEvent that contains the reviewer outcome after both agents finish.
# Prerequisites:
# - Foundry Agent Service configured, along with the required environment variables.
# - Authentication via azure-identity. Use AzureCliCredential and run az login before executing the sample.
# - Basic familiarity with WorkflowBuilder, edges, events, and streaming runs.
# """
# async def main():
# """Build and run a simple two node agent workflow: Writer then Reviewer."""
# # Create the Foundry chat client.
# async with (
# AzureCliCredential() as credential,
# FoundryChatClient(async_credential=credential).create_agent(
# name="Writer",
# instructions=(
# "You are an excellent content writer.You create new content and edit contents based on the feedback."
# ),
# ) as writer_agent,
# FoundryChatClient(async_credential=credential).create_agent(
# name="Reviewer",
# instructions=(
# "You are an excellent content reviewer."
# "Provide actionable feedback to the writer about the provided content."
# "Provide the feedback in the most concise manner possible."
# ),
# ) as reviewer_agent,
# ):
# # Build the workflow using the fluent builder.
# # Set the start node and connect an edge from writer to reviewer.
# workflow = WorkflowBuilder().set_start_executor(writer_agent).add_edge(writer_agent, reviewer_agent).build()
# # Stream events from the workflow. We aggregate partial token updates per executor for readable output.
# completed_event: WorkflowCompletedEvent | None = None
# last_executor_id = None
# async for event in workflow.run_stream(
# "Create a slogan for a new electric SUV that is affordable and fun to drive."
# ):
# if isinstance(event, AgentRunUpdateEvent):
# # AgentRunUpdateEvent contains incremental text deltas from the underlying agent.
# # Print a prefix when the executor changes, then append updates on the same line.
# eid = event.executor_id
# if eid != last_executor_id:
# if last_executor_id is not None:
# print()
# print(f"{eid}:", end=" ", flush=True)
# last_executor_id = eid
# print(event.data, end="", flush=True)
# elif isinstance(event, WorkflowCompletedEvent):
# # Terminal event with the final reviewer output.
# completed_event = event
# # Print the final consolidated reviewer result.
# if completed_event:
# print("\n===== Final Output =====")
# print(completed_event.data)
# """
# Sample Output:
# writer_agent: Charge Up Your Journey. Fun, Affordable, Electric.
# reviewer_agent: Clear message, but consider highlighting SUV specific benefits
# (space, versatility) for stronger impact. Try more vivid language to evoke
# excitement. Example: "Big on Space. Big on Fun. Electric for Everyone."
# ===== Final Output =====
# Clear message, but consider highlighting SUV specific benefits (space, versatility)
# for stronger impact. Try more vivid language to evoke excitement. Example:
# "Big on Space. Big on Fun. Electric for Everyone."
# """
# if __name__ == "__main__":
# asyncio.run(main())
import asyncio
from collections.abc import Awaitable, Callable
from contextlib import AsyncExitStack
from typing import Any
from collections.abc import Awaitable, Callable
from agent_framework.foundry import FoundryChatClient
from agent_framework import AgentRunUpdateEvent, WorkflowBuilder, WorkflowCompletedEvent
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
"""
Sample: Agents in a workflow with streaming
A Writer agent generates content, then a Reviewer agent critiques it.
The workflow uses streaming so you can observe incremental AgentRunUpdateEvent chunks as each agent produces tokens.
Purpose:
Show how to wire chat agents directly into a WorkflowBuilder pipeline where agents are auto wrapped as executors.
Demonstrate:
- Automatic streaming of agent deltas via AgentRunUpdateEvent.
- A simple console aggregator that groups updates by executor id and prints them as they arrive.
- A final WorkflowCompletedEvent that contains the reviewer outcome after both agents finish.
Prerequisites:
- Foundry Agent Service configured, along with the required environment variables.
- Authentication via azure-identity. Use AzureCliCredential and run az login before executing the sample.
- Basic familiarity with WorkflowBuilder, edges, events, and streaming runs.
"""
async def create_foundry_agent() -> tuple[Callable[..., Awaitable[Any]], Callable[[], Awaitable[None]]]:
"""Helper method to create a Foundry agent factory and a close function.
@@ -1,27 +1,31 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import sys
from dataclasses import dataclass
from pathlib import Path
from agent_framework import (
# Ensure local getting_started package can be imported when running as a script.
_SAMPLES_ROOT = Path(__file__).resolve().parents[3]
if str(_SAMPLES_ROOT) not in sys.path:
sys.path.insert(0, str(_SAMPLES_ROOT))
from agent_framework import ( # noqa: E402
ChatMessage,
Executor,
FunctionCallContent,
FunctionResultContent,
Role,
)
from agent_framework.openai import OpenAIChatClient
from agent_framework import (
Executor,
RequestInfoExecutor,
RequestInfoMessage,
RequestResponse,
Role,
WorkflowAgent,
WorkflowBuilder,
WorkflowContext,
handler,
)
from samples.getting_started.workflow.agents.workflow_as_agent_reflection_pattern import (
from agent_framework.openai import OpenAIChatClient # noqa: E402
from getting_started.workflow.agents.workflow_as_agent_reflection_pattern import ( # noqa: E402
ReviewRequest,
ReviewResponse,
Worker,
@@ -56,8 +60,9 @@ class HumanReviewRequest(RequestInfoMessage):
class ReviewerWithHumanInTheLoop(Executor):
"""Executor that always escalates reviews to a human manager."""
def __init__(self, worker_id: str, request_info_id: str) -> None:
super().__init__()
def __init__(self, worker_id: str, request_info_id: str, reviewer_id: str | None = None) -> None:
unique_id = reviewer_id or f"{worker_id}-reviewer"
super().__init__(id=unique_id)
self._worker_id = worker_id
self._request_info_id = request_info_id
@@ -96,8 +101,8 @@ async def main() -> None:
# Create executors for the workflow.
print("Creating chat client and executors...")
mini_chat_client = OpenAIChatClient(ai_model_id="gpt-4.1-nano")
worker = Worker(chat_client=mini_chat_client)
request_info_executor = RequestInfoExecutor()
worker = Worker(id="sub-worker", chat_client=mini_chat_client)
request_info_executor = RequestInfoExecutor(id="request_info")
reviewer = ReviewerWithHumanInTheLoop(worker_id=worker.id, request_info_id=request_info_executor.id)
print("Building workflow with Worker ↔ Reviewer cycle...")
@@ -4,9 +4,19 @@ import asyncio
from dataclasses import dataclass
from uuid import uuid4
from agent_framework import AgentRunResponseUpdate, ChatClientProtocol, ChatMessage, Contents, Role
from agent_framework import (
AgentRunResponseUpdate,
AgentRunUpdateEvent,
ChatClientProtocol,
ChatMessage,
Contents,
Executor,
Role,
WorkflowBuilder,
WorkflowContext,
handler,
)
from agent_framework.openai import OpenAIChatClient
from agent_framework import AgentRunUpdateEvent, Executor, WorkflowBuilder, WorkflowContext, handler
from pydantic import BaseModel
"""
@@ -54,8 +64,8 @@ class ReviewResponse:
class Reviewer(Executor):
"""Executor that reviews agent responses and provides structured feedback."""
def __init__(self, chat_client: ChatClientProtocol) -> None:
super().__init__()
def __init__(self, id: str, chat_client: ChatClientProtocol) -> None:
super().__init__(id=id)
self._chat_client = chat_client
@handler
@@ -106,8 +116,8 @@ class Reviewer(Executor):
class Worker(Executor):
"""Executor that generates responses and incorporates feedback when necessary."""
def __init__(self, chat_client: ChatClientProtocol) -> None:
super().__init__()
def __init__(self, id: str, chat_client: ChatClientProtocol) -> None:
super().__init__(id=id)
self._chat_client = chat_client
self._pending_requests: dict[str, tuple[ReviewRequest, list[ChatMessage]]] = {}
@@ -189,8 +199,8 @@ async def main() -> None:
print("Creating chat client and executors...")
mini_chat_client = OpenAIChatClient(ai_model_id="gpt-4.1-nano")
chat_client = OpenAIChatClient(ai_model_id="gpt-4.1")
reviewer = Reviewer(chat_client=chat_client)
worker = Worker(chat_client=mini_chat_client)
reviewer = Reviewer(id="reviewer", chat_client=chat_client)
worker = Worker(id="worker", chat_client=mini_chat_client)
print("Building workflow with Worker ↔ Reviewer cycle...")
agent = (
@@ -0,0 +1,483 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from collections.abc import AsyncIterable
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
from agent_framework import (
AgentExecutor,
AgentExecutorRequest,
AgentExecutorResponse,
ChatMessage,
Executor,
FileCheckpointStorage,
RequestInfoEvent,
RequestInfoExecutor,
RequestInfoMessage,
RequestResponse,
Role,
WorkflowBuilder,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
from agent_framework.azure import AzureChatClient
from azure.identity import AzureCliCredential
# NOTE: the Azure client imports above are real dependencies. When running this
# sample outside of Azure-enabled environments you may wish to swap in the
# `agent_framework.builtin` chat client or mock the writer executor. We keep the
# concrete import here so readers can see an end-to-end configuration.
if TYPE_CHECKING:
from agent_framework import Workflow
from agent_framework._workflow._checkpoint import WorkflowCheckpoint
"""
Sample: Checkpoint + human-in-the-loop quickstart.
This getting-started sample keeps the moving pieces to a minimum:
1. A brief is turned into a consistent prompt for an AI copywriter.
2. The copywriter (an `AgentExecutor`) drafts release notes.
3. A reviewer gateway routes every draft through `RequestInfoExecutor` so a human
can approve or request tweaks.
4. The workflow records checkpoints between each superstep so you can stop the
program, restart later, and optionally pre-supply human answers on resume.
Key concepts demonstrated
-------------------------
- Minimal executor pipeline with checkpoint persistence.
- Human-in-the-loop pause/resume by pairing `RequestInfoExecutor` with
checkpoint restoration.
- Supplying responses at restore time (`run_stream_from_checkpoint(..., responses=...)`).
Typical pause/resume flow
-------------------------
1. Run the workflow until a human approval request is emitted.
2. If the human is offline, exit the program. A checkpoint with
``status=awaiting human response`` now exists.
3. Later, restart the script, select that checkpoint, and provide the stored
human decision when prompted to pre-supply responses.
Doing so applies the answer immediately on resume, so the system does **not**
re-emit the same `RequestInfoEvent`.
"""
# Directory used for the sample's temporary checkpoint files. We isolate the
# demo artefacts so that repeated runs do not collide with other samples and so
# the clean-up step at the end of the script can simply delete the directory.
TEMP_DIR = Path(__file__).with_suffix("").parent / "tmp" / "checkpoints_hitl"
TEMP_DIR.mkdir(parents=True, exist_ok=True)
class BriefPreparer(Executor):
"""Normalises the user brief and sends a single AgentExecutorRequest."""
# The first executor in the workflow. By keeping it tiny we make it easier
# to reason about the state that will later be captured in the checkpoint.
# It is responsible for tidying the human-provided brief and kicking off the
# agent run with a deterministic prompt structure.
def __init__(self, id: str, agent_id: str) -> None:
super().__init__(id=id)
self._agent_id = agent_id
@handler
async def prepare(self, brief: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
# Collapse errant whitespace so the prompt is stable between runs.
normalized = " ".join(brief.split()).strip()
if not normalized.endswith("."):
normalized += "."
# Persist the cleaned brief in shared state so downstream executors and
# future checkpoints can recover the original intent.
await ctx.set_shared_state("brief", normalized)
prompt = (
"You are drafting product release notes. Summarise the brief below in two sentences. "
"Keep it positive and end with a call to action.\n\n"
f"BRIEF: {normalized}"
)
# Hand the prompt to the writer agent. We always route through the
# workflow context so the runtime can capture messages for checkpointing.
await ctx.send_message(
AgentExecutorRequest(messages=[ChatMessage(Role.USER, text=prompt)], should_respond=True),
target_id=self._agent_id,
)
@dataclass
class HumanApprovalRequest(RequestInfoMessage):
"""Message sent to the human reviewer via RequestInfoExecutor."""
# These fields are intentionally simple because they are serialised into
# checkpoints. Keeping them primitive types guarantees the new
# `pending_requests_from_checkpoint` helper can reconstruct them on resume.
prompt: str = ""
draft: str = ""
iteration: int = 0
class ReviewGateway(Executor):
"""Routes agent drafts to humans and optionally back for revisions."""
def __init__(self, id: str, reviewer_id: str, writer_id: str, finalize_id: str) -> None:
super().__init__(id=id)
self._reviewer_id = reviewer_id
self._writer_id = writer_id
self._finalize_id = finalize_id
@handler
async def on_agent_response(
self,
response: AgentExecutorResponse,
ctx: WorkflowContext[HumanApprovalRequest],
) -> None:
# Capture the agent output so we can surface it to the reviewer and
# persist iterations. The `RequestInfoExecutor` relies on this state to
# rehydrate when checkpoints are restored.
draft = response.agent_run_response.text or ""
iteration = int((await ctx.get_state() or {}).get("iteration", 0)) + 1
await ctx.set_state({"iteration": iteration, "last_draft": draft})
# Emit a human approval request. Because this flows through
# RequestInfoExecutor it will pause the workflow until an answer is
# supplied either interactively or via pre-supplied responses.
await ctx.send_message(
HumanApprovalRequest(
prompt="Review the draft. Reply 'approve' or provide edit instructions.",
draft=draft,
iteration=iteration,
),
target_id=self._reviewer_id,
)
@handler
async def on_human_feedback(
self,
feedback: RequestResponse[HumanApprovalRequest, str],
ctx: WorkflowContext[AgentExecutorRequest | str],
) -> None:
# The RequestResponse wrapper gives us both the human data and the
# original request message, even when resuming from checkpoints.
reply = (feedback.data or "").strip()
state = await ctx.get_state() or {}
draft = state.get("last_draft") or (feedback.original_request.draft if feedback.original_request else "")
if reply.lower() == "approve":
# When the human signs off we can short-circuit the workflow and
# send the approved draft to the final executor.
await ctx.send_message(draft, target_id=self._finalize_id)
return
# Any other response loops us back to the writer with fresh guidance.
guidance = reply or "Tighten the copy and emphasise customer benefit."
iteration = int(state.get("iteration", 1)) + 1
await ctx.set_state({"iteration": iteration, "last_draft": draft})
prompt = (
"Revise the launch note. Respond with the new copy only.\n\n"
f"Previous draft:\n{draft}\n\n"
f"Human guidance: {guidance}"
)
await ctx.send_message(
AgentExecutorRequest(messages=[ChatMessage(Role.USER, text=prompt)], should_respond=True),
target_id=self._writer_id,
)
class FinaliseExecutor(Executor):
"""Publishes the approved text."""
@handler
async def publish(self, text: str, ctx: WorkflowContext[Any]) -> None:
# Store the output so diagnostics or a UI could fetch the final copy.
await ctx.set_state({"published_text": text})
# Emit a workflow completion event so the runner stops cleanly.
await ctx.add_event(WorkflowCompletedEvent(text))
def create_workflow(*, checkpoint_storage: FileCheckpointStorage | None = None) -> "Workflow":
"""Assemble the workflow graph used by both the initial run and resume."""
# The Azure client is created once so our agent executor can issue calls to
# the hosted model. The agent id is stable across runs which keeps
# checkpoints deterministic.
chat_client = AzureChatClient(credential=AzureCliCredential())
writer = AgentExecutor(
chat_client.create_agent(
instructions="Write concise, warm release notes that sound human and helpful.",
),
id="writer",
)
# RequestInfoExecutor is the lynchpin for human-in-the-loop: every draft is
# routed through it so checkpoints can pause while waiting for responses.
review = RequestInfoExecutor(id="request_info")
finalise = FinaliseExecutor(id="finalise")
gateway = ReviewGateway(
id="review_gateway",
reviewer_id=review.id,
writer_id=writer.id,
finalize_id=finalise.id,
)
prepare = BriefPreparer(id="prepare_brief", agent_id=writer.id)
# Wire the workflow DAG. Edges mirror the numbered steps described in the
# module docstring. Because `WorkflowBuilder` is declarative, reading these
# edges is often the quickest way to understand execution order.
builder = (
WorkflowBuilder(max_iterations=6)
.set_start_executor(prepare)
.add_edge(prepare, writer)
.add_edge(writer, gateway)
.add_edge(gateway, review)
.add_edge(review, gateway) # human resumes loop
.add_edge(gateway, writer) # revisions
.add_edge(gateway, finalise)
)
# Opt-in to persistence when the caller provides storage. The workflow
# object itself is identical whether or not checkpointing is enabled.
if checkpoint_storage:
builder = builder.with_checkpointing(checkpoint_storage=checkpoint_storage)
return builder.build()
def _render_checkpoint_summary(checkpoints: list["WorkflowCheckpoint"]) -> None:
"""Pretty-print saved checkpoints with the new framework summaries."""
print("\nCheckpoint summary:")
for summary in [
RequestInfoExecutor.checkpoint_summary(cp) for cp in sorted(checkpoints, key=lambda c: c.timestamp)
]:
# Compose a single line per checkpoint so the user can scan the output
# and pick the resume point that still has outstanding human work.
line = (
f"- {summary.checkpoint_id} | iter={summary.iteration_count} "
f"| targets={summary.targets} | states={summary.executor_states}"
)
if summary.status:
line += f" | status={summary.status}"
if summary.draft_preview:
line += f" | draft_preview={summary.draft_preview}"
if summary.pending_requests:
line += f" | pending_request_id={summary.pending_requests[0].request_id}"
print(line)
def _print_events(events: list[Any]) -> tuple[WorkflowCompletedEvent | None, list[tuple[str, HumanApprovalRequest]]]:
"""Echo workflow events to the console and collect outstanding requests."""
completed: WorkflowCompletedEvent | None = None
requests: list[tuple[str, HumanApprovalRequest]] = []
for event in events:
print(f"Event: {event}")
if isinstance(event, WorkflowCompletedEvent):
completed = event
elif isinstance(event, RequestInfoEvent) and isinstance(event.data, HumanApprovalRequest):
# Capture pending human approvals so the caller can ask the user for
# input after the current batch of events is processed.
requests.append((event.request_id, event.data))
elif isinstance(event, WorkflowStatusEvent) and event.state in {
WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
}:
print(f"Workflow state: {event.state.name}")
return completed, requests
def _prompt_for_responses(requests: list[tuple[str, HumanApprovalRequest]]) -> dict[str, str] | None:
"""Interactive CLI prompt for any live RequestInfo requests."""
if not requests:
return None
answers: dict[str, str] = {}
for request_id, request in requests:
# Keep the prompt conversational so testers can use the script without
# memorising the workflow APIs.
print("\n=== Human approval needed ===")
print(f"request_id: {request_id}")
if request.iteration:
print(f"Iteration: {request.iteration}")
print(request.prompt)
print("Draft: \n---\n" + request.draft + "\n---")
answer = input("Type 'approve' or enter revision guidance (or 'exit' to quit): ").strip() # noqa: ASYNC250
if answer.lower() == "exit":
raise SystemExit("Stopped by user.")
answers[request_id] = answer
return answers
def _maybe_pre_supply_responses(cp: "WorkflowCheckpoint") -> dict[str, str] | None:
"""Offer to collect responses before resuming a checkpoint."""
pending = RequestInfoExecutor.pending_requests_from_checkpoint(cp)
if not pending:
return None
print(
"This checkpoint still has pending human input. Provide the responses now so the resume step "
"applies them immediately and does not re-emit the original RequestInfo event."
)
choice = input("Pre-supply responses for this checkpoint? [y/N]: ").strip().lower() # noqa: ASYNC250
if choice not in {"y", "yes"}:
return None
answers: dict[str, str] = {}
for item in pending:
iteration = item.iteration or 0
print(f"\nPending draft (iteration {iteration} | request_id={item.request_id}):")
draft_text = (item.draft or "").strip()
if draft_text:
# The shortened preview in the summary may truncate text; here we
# show the full draft so the reviewer can make an informed choice.
print("Draft:\n---\n" + draft_text + "\n---")
else:
print("Draft: [not captured in checkpoint payload - refer to your notes/log]")
prompt_text = (item.prompt or "Review the draft").strip()
print(prompt_text)
answer = input("Response ('approve' or guidance, 'exit' to abort): ").strip() # noqa: ASYNC250
if answer.lower() == "exit":
raise SystemExit("Resume aborted by user.")
answers[item.request_id] = answer
return answers
async def _consume(stream: AsyncIterable[Any]) -> list[Any]:
"""Materialise an async event stream into a list."""
return [event async for event in stream]
async def run_interactive_session(workflow: "Workflow", initial_message: str) -> WorkflowCompletedEvent | None:
"""Run the workflow until it either finishes or pauses for human input."""
pending_responses: dict[str, str] | None = None
completed: WorkflowCompletedEvent | None = None
first = True
while completed is None:
if first:
# Kick off the workflow with the initial brief. The returned events
# include RequestInfo events when the agent produces a draft.
events = await _consume(workflow.run_stream(initial_message))
first = False
elif pending_responses:
# Feed any answers the user just typed back into the workflow.
events = await _consume(workflow.send_responses_streaming(pending_responses))
else:
break
completed, requests = _print_events(events)
pending_responses = _prompt_for_responses(requests)
return completed
async def resume_from_checkpoint(
workflow: "Workflow",
checkpoint_id: str,
storage: FileCheckpointStorage,
pre_supplied: dict[str, str] | None,
) -> None:
"""Resume a stored checkpoint and continue until completion or another pause."""
print(f"\nResuming from checkpoint: {checkpoint_id}")
events = await _consume(
workflow.run_stream_from_checkpoint(
checkpoint_id,
checkpoint_storage=storage,
responses=pre_supplied,
)
)
completed, requests = _print_events(events)
if pre_supplied and not requests and completed is None:
# When the checkpoint only needed the provided answers we let the user
# know the workflow is waiting for the next superstep (usually another
# agent response).
print("Pre-supplied responses applied automatically; workflow is now waiting for the next step.")
pending = _prompt_for_responses(requests)
while completed is None and pending:
events = await _consume(workflow.send_responses_streaming(pending))
completed, requests = _print_events(events)
pending = _prompt_for_responses(requests)
if completed:
print(f"Workflow completed with: {completed.data}")
async def main() -> None:
"""Entry point used by both the initial run and subsequent resumes."""
for file in TEMP_DIR.glob("*.json"):
# Start each execution with a clean slate so the demonstration is
# deterministic even if the directory had stale checkpoints.
file.unlink()
storage = FileCheckpointStorage(storage_path=TEMP_DIR)
workflow = create_workflow(checkpoint_storage=storage)
brief = (
"Introduce our limited edition smart coffee grinder. Mention the $249 price, highlight the "
"sensor that auto-adjusts the grind, and invite customers to pre-order on the website."
)
print("Running workflow (human approval required)...")
completed = await run_interactive_session(workflow, initial_message=brief)
if completed:
print(f"Initial run completed with final copy: {completed.data}")
else:
print("Initial run paused for human input.")
checkpoints = await storage.list_checkpoints()
if not checkpoints:
print("No checkpoints recorded.")
return
# Show the user what is available before we prompt for the index. The
# summary helper keeps this output consistent with other tooling.
_render_checkpoint_summary(checkpoints)
sorted_cps = sorted(checkpoints, key=lambda c: c.timestamp)
print("\nAvailable checkpoints:")
for idx, cp in enumerate(sorted_cps):
print(f" [{idx}] id={cp.checkpoint_id} iter={cp.iteration_count}")
# For the pause/resume demo we typically pick the latest checkpoint whose summary
# status reads "awaiting human response" - that is the saved state that proves the
# workflow can rehydrate, collect the pending answer, and continue after a break.
selection = input("\nResume from which checkpoint? (press Enter to skip): ").strip() # noqa: ASYNC250
if not selection:
print("No resume selected. Exiting.")
return
try:
idx = int(selection)
except ValueError:
print("Invalid input; exiting.")
return
if not 0 <= idx < len(sorted_cps):
print("Index out of range; exiting.")
return
chosen = sorted_cps[idx]
summary = RequestInfoExecutor.checkpoint_summary(chosen)
if summary.status == "completed":
print("Selected checkpoint already reflects a completed workflow; nothing to resume.")
return
# If the user wants, capture their decisions now so the resume call can
# push them into the workflow and avoid re-prompting.
pre_responses = _maybe_pre_supply_responses(chosen)
resumed_workflow = create_workflow()
# Resume with a fresh workflow instance. The checkpoint carries the
# persistent state while this object holds the runtime wiring.
await resume_from_checkpoint(resumed_workflow, chosen.checkpoint_id, storage, pre_responses)
if __name__ == "__main__":
asyncio.run(main())
@@ -3,7 +3,7 @@
import asyncio
import os
from pathlib import Path
from typing import Any
from typing import TYPE_CHECKING, Any
from agent_framework import (
AgentExecutor,
@@ -12,6 +12,7 @@ from agent_framework import (
ChatMessage,
Executor,
FileCheckpointStorage,
RequestInfoExecutor,
Role,
WorkflowBuilder,
WorkflowCompletedEvent,
@@ -21,6 +22,10 @@ from agent_framework import (
from agent_framework.azure import AzureChatClient
from azure.identity import AzureCliCredential
if TYPE_CHECKING:
from agent_framework import Workflow
from agent_framework._workflow._checkpoint import WorkflowCheckpoint
"""
Sample: Checkpointing and Resuming a Workflow (with an Agent stage)
@@ -87,7 +92,7 @@ class UpperCaseExecutor(Executor):
class SubmitToLowerAgent(Executor):
"""Builds an AgentExecutorRequest to send to the lowercasing agent while keeping shared-state visibility."""
def __init__(self, agent_id: str, id: str | None = None):
def __init__(self, id: str, agent_id: str):
super().__init__(id=id)
self._agent_id = agent_id
@@ -132,10 +137,6 @@ class FinalizeFromAgent(Executor):
class ReverseTextExecutor(Executor):
"""Reverses the input text and persists local state."""
def __init__(self, id: str):
"""Initialize the executor with an ID."""
super().__init__(id=id)
@handler
async def reverse_text(self, text: str, ctx: WorkflowContext[str]) -> None:
result = text[::-1]
@@ -154,15 +155,10 @@ class ReverseTextExecutor(Executor):
await ctx.send_message(result)
async def main():
# Clear existing checkpoints in this sample directory for a clean run.
checkpoint_dir = Path(TEMP_DIR)
for file in checkpoint_dir.glob("*.json"):
file.unlink()
def create_workflow(checkpoint_storage: FileCheckpointStorage) -> "Workflow":
# Instantiate the pipeline executors.
upper_case_executor = UpperCaseExecutor(id="upper_case_executor")
reverse_text_executor = ReverseTextExecutor(id="reverse_text_executor")
upper_case_executor = UpperCaseExecutor(id="upper-case")
reverse_text_executor = ReverseTextExecutor(id="reverse-text")
# Configure the agent stage that lowercases the text.
chat_client = AzureChatClient(credential=AzureCliCredential())
@@ -174,14 +170,11 @@ async def main():
)
# Bridge to the agent and terminalization stage.
submit_lower = SubmitToLowerAgent(agent_id=lower_agent.id, id="submit_lower")
submit_lower = SubmitToLowerAgent(id="submit_lower", agent_id=lower_agent.id)
finalize = FinalizeFromAgent(id="finalize")
# Backing store for checkpoints written by with_checkpointing.
checkpoint_storage = FileCheckpointStorage(storage_path=TEMP_DIR)
# Build the workflow with checkpointing enabled.
workflow = (
return (
WorkflowBuilder(max_iterations=5)
.add_edge(upper_case_executor, reverse_text_executor) # Uppercase -> Reverse
.add_edge(reverse_text_executor, submit_lower) # Reverse -> Build Agent request
@@ -192,6 +185,40 @@ async def main():
.build()
)
def _render_checkpoint_summary(checkpoints: list["WorkflowCheckpoint"]) -> None:
"""Display human-friendly checkpoint metadata using framework summaries."""
if not checkpoints:
return
print("\nCheckpoint summary:")
for cp in sorted(checkpoints, key=lambda c: c.timestamp):
summary = RequestInfoExecutor.checkpoint_summary(cp)
msg_count = sum(len(v) for v in cp.messages.values())
state_keys = sorted(cp.executor_states.keys())
orig = cp.shared_state.get("original_input")
upper = cp.shared_state.get("upper_output")
line = (
f"- {summary.checkpoint_id} | iter={summary.iteration_count} | messages={msg_count} | states={state_keys}"
)
if summary.status:
line += f" | status={summary.status}"
line += f" | shared_state: original_input='{orig}', upper_output='{upper}'"
print(line)
async def main():
# Clear existing checkpoints in this sample directory for a clean run.
checkpoint_dir = Path(TEMP_DIR)
for file in checkpoint_dir.glob("*.json"):
file.unlink()
# Backing store for checkpoints written by with_checkpointing.
checkpoint_storage = FileCheckpointStorage(storage_path=TEMP_DIR)
workflow = create_workflow(checkpoint_storage=checkpoint_storage)
# Run the full workflow once and observe events as they stream.
print("Running workflow with initial message...")
async for event in workflow.run_stream(message="hello world"):
@@ -206,26 +233,20 @@ async def main():
# All checkpoints created by this run share the same workflow_id.
workflow_id = all_checkpoints[0].workflow_id
# Dump a quick summary including shared_state keys to illustrate what persisted.
print("\nCheckpoint summary:")
for cp in sorted(all_checkpoints, key=lambda c: c.timestamp):
msg_count = sum(len(v) for v in cp.messages.values())
state_keys = sorted(list(cp.executor_states.keys())) if hasattr(cp, "executor_states") else []
orig = cp.shared_state.get("original_input") if hasattr(cp, "shared_state") else None
upper = cp.shared_state.get("upper_output") if hasattr(cp, "shared_state") else None
print(
f"- {cp.checkpoint_id} | "
f"iter={cp.iteration_count} | messages={msg_count} | states={state_keys} | "
f"shared_state: original_input='{orig}', upper_output='{upper}'"
)
_render_checkpoint_summary(all_checkpoints)
# Offer an interactive selection of checkpoints to resume from.
sorted_cps = sorted([cp for cp in all_checkpoints if cp.workflow_id == workflow_id], key=lambda c: c.timestamp)
print("\nAvailable checkpoints to resume from:")
for idx, cp in enumerate(sorted_cps):
summary = RequestInfoExecutor.checkpoint_summary(cp)
line = f" [{idx}] id={summary.checkpoint_id} iter={summary.iteration_count}"
if summary.status:
line += f" status={summary.status}"
msg_count = sum(len(v) for v in cp.messages.values())
print(f" [{idx}] id={cp.checkpoint_id} iter={cp.iteration_count} messages={msg_count}")
line += f" messages={msg_count}"
print(line)
user_input = input(
"\nEnter checkpoint index (or paste checkpoint id) to resume from, or press Enter to skip resume: "
@@ -256,15 +277,7 @@ async def main():
# You can reuse the same workflow graph definition and resume from a prior checkpoint.
# This second workflow instance does not enable checkpointing to show that resumption
# reads from stored state but need not write new checkpoints.
new_workflow = (
WorkflowBuilder(max_iterations=5)
.add_edge(upper_case_executor, reverse_text_executor)
.add_edge(reverse_text_executor, submit_lower)
.add_edge(submit_lower, lower_agent)
.add_edge(lower_agent, finalize)
.set_start_executor(upper_case_executor)
.build()
)
new_workflow = create_workflow(checkpoint_storage=checkpoint_storage)
print(f"\nResuming from checkpoint: {chosen_cp_id}")
async for event in new_workflow.run_stream_from_checkpoint(chosen_cp_id, checkpoint_storage=checkpoint_storage):
@@ -275,15 +288,15 @@ async def main():
Running workflow with initial message...
UpperCaseExecutor: 'hello world' -> 'HELLO WORLD'
Event: ExecutorInvokedEvent(executor_id=upper_case_executor)
Event: ExecutorInvokeEvent(executor_id=upper_case_executor)
Event: ExecutorCompletedEvent(executor_id=upper_case_executor)
ReverseTextExecutor: 'HELLO WORLD' -> 'DLROW OLLEH'
Event: ExecutorInvokedEvent(executor_id=reverse_text_executor)
Event: ExecutorInvokeEvent(executor_id=reverse_text_executor)
Event: ExecutorCompletedEvent(executor_id=reverse_text_executor)
LowerAgent (shared_state): original_input='hello world', upper_output='HELLO WORLD'
Event: ExecutorInvokedEvent(executor_id=submit_lower)
Event: ExecutorInvokedEvent(executor_id=lower_agent)
Event: ExecutorInvokedEvent(executor_id=finalize)
Event: ExecutorInvokeEvent(executor_id=submit_lower)
Event: ExecutorInvokeEvent(executor_id=lower_agent)
Event: ExecutorInvokeEvent(executor_id=finalize)
Event: WorkflowCompletedEvent(data=dlrow olleh)
Checkpoint summary:
@@ -300,9 +313,9 @@ async def main():
Resuming from checkpoint: a78c345a-e5d9-45ba-82c0-cb725452d91b
LowerAgent (shared_state): original_input='hello world', upper_output='HELLO WORLD'
Resumed Event: ExecutorInvokedEvent(executor_id=submit_lower)
Resumed Event: ExecutorInvokedEvent(executor_id=lower_agent)
Resumed Event: ExecutorInvokedEvent(executor_id=finalize)
Resumed Event: ExecutorInvokeEvent(executor_id=submit_lower)
Resumed Event: ExecutorInvokeEvent(executor_id=lower_agent)
Resumed Event: ExecutorInvokeEvent(executor_id=finalize)
Resumed Event: WorkflowCompletedEvent(data=dlrow olleh)
""" # noqa: E501
@@ -50,7 +50,7 @@ class GuessNumberExecutor(Executor):
def __init__(self, bound: tuple[int, int], id: str | None = None):
"""Initialize the executor with a target number."""
super().__init__(id=id)
super().__init__(id=id or "guess_number")
self._lower = bound[0]
self._upper = bound[1]
@@ -83,7 +83,7 @@ class SubmitToJudgeAgent(Executor):
"""Send the numeric guess to a judge agent which replies ABOVE/BELOW/MATCHED."""
def __init__(self, judge_agent_id: str, target: int, id: str | None = None):
super().__init__(id=id)
super().__init__(id=id or "submit_to_judge")
self._judge_agent_id = judge_agent_id
self._target = target
@@ -87,7 +87,7 @@ class TurnManager(Executor):
"""
def __init__(self, id: str | None = None):
super().__init__(id=id)
super().__init__(id=id or "turn_manager")
@handler
async def start(self, _: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
@@ -43,7 +43,7 @@ class DispatchToExperts(Executor):
"""Dispatches the incoming prompt to all expert agent executors for parallel processing (fan out)."""
def __init__(self, expert_ids: list[str], id: str | None = None):
super().__init__(id)
super().__init__(id=id or "dispatch_to_experts")
self._expert_ids = expert_ids
@handler
@@ -71,7 +71,7 @@ class AggregateInsights(Executor):
"""Aggregates expert agent responses into a single consolidated result (fan in)."""
def __init__(self, expert_ids: list[str], id: str | None = None):
super().__init__(id)
super().__init__(id=id or "aggregate_insights")
self._expert_ids = expert_ids
@handler
@@ -61,7 +61,7 @@ class Split(Executor):
def __init__(self, map_executor_ids: list[str], id: str | None = None):
"""Store mapper ids so we can assign non overlapping ranges per mapper."""
super().__init__(id)
super().__init__(id=id or "split")
self._map_executor_ids = map_executor_ids
@handler
@@ -145,7 +145,7 @@ class Shuffle(Executor):
def __init__(self, reducer_ids: list[str], id: str | None = None):
"""Remember reducer ids so we can partition work deterministically."""
super().__init__(id)
super().__init__(id=id or "shuffle")
self._reducer_ids = reducer_ids
@handler
@@ -40,7 +40,7 @@ class DispatchToExperts(Executor):
"""Dispatches the incoming prompt to all expert agent executors (fan-out)."""
def __init__(self, expert_ids: list[str], id: str | None = None):
super().__init__(id)
super().__init__(id=id or "dispatch_to_experts")
self._expert_ids = expert_ids
@handler
@@ -67,7 +67,7 @@ class AggregateInsights(Executor):
"""Aggregates expert agent responses into a single consolidated result (fan-in)."""
def __init__(self, expert_ids: list[str], id: str | None = None):
super().__init__(id)
super().__init__(id=id or "aggregate_insights")
self._expert_ids = expert_ids
@handler