Repository navigation
Expand file tree
/
Copy pathcheckpoint_with_human_in_the_loop.py
More file actions
339 lines (279 loc) · 12.7 KB
/
Copy pathcheckpoint_with_human_in_the_loop.py
File metadata and controls
339 lines (279 loc) · 12.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import os
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any
from agent_framework import (
Agent,
AgentExecutor,
AgentExecutorRequest,
AgentExecutorResponse,
Executor,
FileCheckpointStorage,
Message,
Workflow,
WorkflowBuilder,
WorkflowContext,
handler,
register_checkpoint_type,
response_handler,
)
from agent_framework.foundry import FoundryChatClient
from azure.identity import AzureCliCredential
from dotenv import load_dotenv
if sys.version_info >= (3, 12):
from typing import override # pragma: no cover
else:
from typing_extensions import override # pragma: no cover
# Load environment variables from .env file
load_dotenv()
"""
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 sends a request for approval for every draft.
4. An output executor emits the approved draft as the terminal workflow output.
5. The workflow records checkpoints between each superstep so you can stop the
program and restart later.
Key concepts demonstrated
-------------------------
- Minimal executor pipeline with checkpoint persistence.
- Human-in-the-loop pause/resume with checkpoint restoration.
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 and select that checkpoint. The workflow restores
and re-emits the pending request so the human can answer it.
"""
# 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, str]) -> 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 workflow state so downstream executors and
# future checkpoints can recover the original intent.
ctx.set_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=[Message("user", contents=[prompt])], should_respond=True),
target_id=self._agent_id,
)
@dataclass
class HumanApprovalRequest:
"""Request sent to the human reviewer."""
# These fields are intentionally simple because they are serialised into
# checkpoints and reconstructed when the workflow resumes.
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, writer_id: str) -> None:
super().__init__(id=id)
self._writer_id = writer_id
self._iteration = 0
@handler
async def on_agent_response(self, response: AgentExecutorResponse, ctx: WorkflowContext) -> None:
# Capture the agent output so we can surface it to the reviewer and persist iterations.
self._iteration += 1
# Emit a human approval request.
await ctx.request_info(
request_data=HumanApprovalRequest(
prompt="Review the draft. Reply 'approve' or provide edit instructions.",
draft=response.agent_response.text,
iteration=self._iteration,
),
response_type=str,
)
@response_handler
async def on_human_feedback(
self,
original_request: HumanApprovalRequest,
feedback: str,
ctx: WorkflowContext[AgentExecutorRequest | str, str],
) -> None:
# The `original_request` is the request we sent earlier that is now being answered.
reply = feedback.strip()
if len(reply) == 0 or reply.lower() == "approve":
# Workflow is completed when the human approves.
await ctx.yield_output(original_request.draft)
return
# Any other response loops us back to the writer with fresh guidance.
prompt = (
"Revise the launch note. Respond with the new copy only.\n\n"
f"Previous draft:\n{original_request.draft}\n\n"
f"Human guidance: {reply}"
)
await ctx.send_message(
AgentExecutorRequest(messages=[Message("user", contents=[prompt])], should_respond=True),
target_id=self._writer_id,
)
@override
async def on_checkpoint_save(self) -> dict[str, Any]:
# Save the current iteration count in executor state for checkpointing.
return {"iteration": self._iteration}
@override
async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
# Restore the iteration count from executor state during checkpoint recovery.
self._iteration = state.get("iteration", 0)
def create_workflow(checkpoint_storage: FileCheckpointStorage) -> Workflow:
"""Assemble the workflow graph used by both the initial run and resume."""
# 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.
writer_agent = Agent(
client=FoundryChatClient(
project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
model=os.environ["FOUNDRY_MODEL"],
credential=AzureCliCredential(),
),
instructions="Write concise, warm release notes that sound human and helpful.",
name="writer",
)
writer = AgentExecutor(writer_agent)
review_gateway = ReviewGateway(id="review_gateway", writer_id="writer")
prepare_brief = BriefPreparer(id="prepare_brief", agent_id="writer")
workflow_builder = (
WorkflowBuilder(
max_iterations=6,
start_executor=prepare_brief,
checkpoint_storage=checkpoint_storage,
output_from=[review_gateway],
)
.add_edge(prepare_brief, writer)
.add_edge(writer, review_gateway)
.add_edge(review_gateway, writer) # revisions loop
)
return workflow_builder.build()
def prompt_for_responses(requests: dict[str, HumanApprovalRequest]) -> dict[str, str]:
"""Interactive CLI prompt for any live RequestInfo requests."""
responses: dict[str, str] = {}
for request_id, request in requests.items():
print("\n=== Human approval needed ===")
print(f"request_id: {request_id}")
print(f"Iteration: {request.iteration}")
print(request.prompt)
print("Draft: \n---\n" + request.draft + "\n---")
response = input("Type 'approve' or enter revision guidance (or 'exit' to quit): ").strip()
if response.lower() == "exit":
raise SystemExit("Stopped by user.")
responses[request_id] = response
return responses
async def run_interactive_session(
workflow: Workflow,
initial_message: str | None = None,
checkpoint_id: str | None = None,
) -> str:
"""Run the workflow until it either finishes or pauses for human input."""
requests: dict[str, HumanApprovalRequest] = {}
responses: dict[str, str] | None = None
completed_output: str | None = None
while True:
if responses:
event_stream = workflow.run(stream=True, responses=responses)
requests.clear()
responses = None
else:
if initial_message:
print(f"\nStarting workflow with brief: {initial_message}\n")
event_stream = workflow.run(message=initial_message, stream=True)
elif checkpoint_id:
print("\nStarting workflow from checkpoint...\n")
event_stream = workflow.run(checkpoint_id=checkpoint_id, stream=True)
else:
raise ValueError("Either initial_message or checkpoint_id must be provided")
async for event in event_stream:
if event.type == "status":
print(event)
if event.type == "output":
completed_output = event.data
if event.type == "request_info":
if isinstance(event.data, HumanApprovalRequest):
requests[event.request_id] = event.data
else:
raise ValueError("Unexpected request data type")
if completed_output:
break
if requests:
responses = prompt_for_responses(requests)
continue
raise RuntimeError("Workflow stopped without completing or requesting input")
return completed_output
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()
# Register the application-defined request type so file storage can reconstruct it when loading checkpoints.
# Alternatively, scope permission to this storage instance:
# allowed_types = [f"{HumanApprovalRequest.__module__}:{HumanApprovalRequest.__qualname__}"]
# storage = FileCheckpointStorage(storage_path=TEMP_DIR, allowed_checkpoint_types=allowed_types)
register_checkpoint_type(HumanApprovalRequest)
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)...")
result = await run_interactive_session(workflow, initial_message=brief)
print(f"Workflow completed with: {result}")
checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
if not checkpoints:
print("No checkpoints recorded.")
return
sorted_cps = sorted(checkpoints, key=lambda cp: datetime.fromisoformat(cp.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]
new_workflow = create_workflow(checkpoint_storage=storage)
# Resume with a fresh workflow instance. The checkpoint carries the
# persistent state while this object holds the runtime wiring.
result = await run_interactive_session(new_workflow, checkpoint_id=chosen.checkpoint_id)
print(f"Workflow completed with: {result}")
if __name__ == "__main__":
asyncio.run(main())