chore: import upstream snapshot with attribution

python/samples/03-workflows/orchestrations/magentic_checkpoint.py

@@ -0,0 +1,319 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import json
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import cast
+
+from agent_framework import (
+    Agent,
+    FileCheckpointStorage,
+    Message,
+    WorkflowCheckpoint,
+    WorkflowEvent,
+    WorkflowRunState,
+)
+from agent_framework.foundry import FoundryChatClient
+from agent_framework.orchestrations import MagenticBuilder, MagenticPlanReviewRequest
+from azure.identity import AzureCliCredential
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+"""
+Sample: Magentic Orchestration + Checkpointing
+
+The goal of this sample is to show the exact mechanics needed to pause a Magentic
+workflow that requires human plan review, persist the outstanding request via a
+checkpoint, and later resume the workflow by feeding in the saved response.
+
+Concepts highlighted here:
+1. **Deterministic executor IDs** - the orchestrator and plan-review request executor
+   must keep stable IDs so the checkpoint state aligns when we rebuild the graph.
+2. **Executor snapshotting** - checkpoints capture the pending plan-review request
+   map, at superstep boundaries.
+3. **Resume with responses** - `Workflow.run(responses=...)` accepts a
+   `responses` mapping so we can inject the stored human reply during restoration.
+
+Prerequisites:
+- FOUNDRY_PROJECT_ENDPOINT must be your Azure AI Foundry Agent Service (V2) project endpoint.
+- FOUNDRY_MODEL must be set to your Azure OpenAI model deployment name.
+- Authentication via azure-identity. Use AzureCliCredential and run az login before executing the sample.
+"""
+
+TASK = (
+    "Draft a concise internal brief describing how our research and implementation teams should collaborate "
+    "to launch a beta feature for data-driven email summarization. Highlight the key milestones, "
+    "risks, and communication cadence."
+)
+
+# Dedicated folder for captured checkpoints. Keeping it under the sample directory
+# makes it easy to inspect the JSON blobs produced by each run.
+CHECKPOINT_DIR = Path(__file__).parent / "tmp" / "magentic_checkpoints"
+
+
+def build_workflow(checkpoint_storage: FileCheckpointStorage):
+    """Construct the Magentic workflow graph with checkpointing enabled."""
+
+    # Two vanilla ChatAgents act as participants in the orchestration. They do not need
+    # extra state handling because their inputs/outputs are fully described by chat messages.
+    researcher = Agent(
+        name="ResearcherAgent",
+        description="Collects background facts and references for the project.",
+        instructions=("You are the research lead. Gather crisp bullet points the team should know."),
+        client=FoundryChatClient(
+            project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+            model=os.environ["FOUNDRY_MODEL"],
+            credential=AzureCliCredential(),
+        ),
+    )
+
+    writer = Agent(
+        name="WriterAgent",
+        description="Synthesizes the final brief for stakeholders.",
+        instructions=("You convert the research notes into a structured brief with milestones and risks."),
+        client=FoundryChatClient(
+            project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+            model=os.environ["FOUNDRY_MODEL"],
+            credential=AzureCliCredential(),
+        ),
+    )
+
+    # Create a manager agent for orchestration
+    manager_agent = Agent(
+        name="MagenticManager",
+        description="Orchestrator that coordinates the research and writing workflow",
+        instructions="You coordinate a team to complete complex tasks efficiently.",
+        client=FoundryChatClient(
+            project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
+            model=os.environ["FOUNDRY_MODEL"],
+            credential=AzureCliCredential(),
+        ),
+    )
+
+    # The builder wires in the Magentic orchestrator, sets the plan review path, and
+    # stores the checkpoint backend so the runtime knows where to persist snapshots.
+    return MagenticBuilder(
+        participants=[researcher, writer],
+        enable_plan_review=True,
+        checkpoint_storage=checkpoint_storage,
+        manager_agent=manager_agent,
+        max_round_count=10,
+        max_stall_count=3,
+    ).build()
+
+
+async def main() -> None:
+    # Stage 0: make sure the checkpoint folder is empty so we inspect only checkpoints
+    # written by this invocation. This prevents stale files from previous runs from
+    # confusing the analysis.
+    CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
+    for file in CHECKPOINT_DIR.glob("*.json"):
+        file.unlink()
+
+    checkpoint_storage = FileCheckpointStorage(CHECKPOINT_DIR)
+
+    print("\n=== Stage 1: run until plan review request (checkpointing active) ===")
+    workflow = build_workflow(checkpoint_storage)
+
+    # Run the workflow until the first  is surfaced. The event carries the
+    # request_id we must reuse on resume. In a real system this is where the UI would present
+    # the plan for human review.
+    plan_review_request: MagenticPlanReviewRequest | None = None
+    async for event in workflow.run(TASK, stream=True):
+        if event.type == "request_info" and event.request_type is MagenticPlanReviewRequest:
+            plan_review_request = event.data
+            print(f"Captured plan review request: {event.request_id}")
+
+        if event.type == "status" and event.state is WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
+            break
+
+    if plan_review_request is None:
+        print("No plan review request emitted; nothing to resume.")
+        return
+
+    resume_checkpoint = await checkpoint_storage.get_latest(workflow_name=workflow.name)
+    if not resume_checkpoint:
+        print("No checkpoints persisted.")
+        return
+
+    print(f"Using checkpoint {resume_checkpoint.checkpoint_id} at iteration {resume_checkpoint.iteration_count}")
+
+    # Show that the checkpoint JSON indeed contains the pending plan-review request record.
+    checkpoint_path = checkpoint_storage.storage_path / f"{resume_checkpoint.checkpoint_id}.json"
+    if checkpoint_path.exists():
+        with checkpoint_path.open() as f:
+            snapshot = json.load(f)
+        request_map = snapshot.get("pending_request_info_events", {})
+        print(f"Pending plan-review requests persisted in checkpoint: {list(request_map.keys())}")
+
+    print("\n=== Stage 2: resume from checkpoint and approve plan ===")
+    resumed_workflow = build_workflow(checkpoint_storage)
+
+    # Construct an approval reply to supply when the plan review request is re-emitted.
+    approval = plan_review_request.approve()
+
+    # Resume execution and capture the re-emitted plan review request.
+    request_info_event: WorkflowEvent | None = None
+    async for event in resumed_workflow.run(checkpoint_id=resume_checkpoint.checkpoint_id, stream=True):
+        if event.type == "request_info" and isinstance(event.data, MagenticPlanReviewRequest):
+            request_info_event = event
+
+    if request_info_event is None:
+        print("No plan review request re-emitted on resume; cannot approve.")
+        return
+    print(f"Resumed plan review request: {request_info_event.request_id}")
+
+    # Supply the approval and continue to run to completion.
+    final_event: WorkflowEvent | None = None
+    async for event in resumed_workflow.run(stream=True, responses={request_info_event.request_id: approval}):
+        if event.type == "output":
+            final_event = event
+
+    if final_event is None:
+        print("Workflow did not complete after resume.")
+        return
+
+    # Final sanity check: display the assistant's answer as proof the orchestration reached
+    # a natural completion after resuming from the checkpoint.
+    result = final_event.data
+    if not result:
+        print("No result data from workflow.")
+        return
+    output_messages = cast(list[Message], result)
+    print("\n=== Final Answer ===")
+    # The output of the Magentic workflow is a list of ChatMessages with only one final message
+    # generated by the orchestrator.
+    print(output_messages[-1].text)
+
+    # ------------------------------------------------------------------
+    # Stage 3: demonstrate resuming from a later checkpoint (post-plan)
+    # ------------------------------------------------------------------
+
+    def _pending_message_count(cp: WorkflowCheckpoint) -> int:
+        return sum(len(msg_list) for msg_list in cp.messages.values() if isinstance(msg_list, list))
+
+    all_checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=resume_checkpoint.workflow_name)
+    later_checkpoints_with_messages = [
+        cp
+        for cp in all_checkpoints
+        if cp.iteration_count > resume_checkpoint.iteration_count and _pending_message_count(cp) > 0
+    ]
+
+    if later_checkpoints_with_messages:
+        post_plan_checkpoint = max(later_checkpoints_with_messages, key=lambda cp: datetime.fromisoformat(cp.timestamp))
+    else:
+        later_checkpoints = [cp for cp in all_checkpoints if cp.iteration_count > resume_checkpoint.iteration_count]
+
+        if not later_checkpoints:
+            print("\nNo additional checkpoints recorded beyond plan approval; sample complete.")
+            return
+
+        post_plan_checkpoint = max(later_checkpoints, key=lambda cp: datetime.fromisoformat(cp.timestamp))
+    print("\n=== Stage 3: resume from post-plan checkpoint ===")
+    pending_messages = _pending_message_count(post_plan_checkpoint)
+    print(
+        f"Resuming from checkpoint {post_plan_checkpoint.checkpoint_id} at iteration "
+        f"{post_plan_checkpoint.iteration_count} (pending messages: {pending_messages})"
+    )
+    if pending_messages == 0:
+        print("Checkpoint has no pending messages; no additional work expected on resume.")
+
+    final_event_post: WorkflowEvent | None = None
+    post_emitted_events = False
+    post_plan_workflow = build_workflow(checkpoint_storage)
+    async for event in post_plan_workflow.run(checkpoint_id=post_plan_checkpoint.checkpoint_id, stream=True):
+        post_emitted_events = True
+        if event.type == "output":
+            final_event_post = event
+
+    if final_event_post is None:
+        if not post_emitted_events:
+            print("No new events were emitted; checkpoint already captured a completed run.")
+            print("\n=== Final Answer (post-plan resume) ===")
+            print(output_messages[-1].text)
+            return
+        print("Workflow did not complete after post-plan resume.")
+        return
+
+    post_result = final_event_post.data
+    if not post_result:
+        print("No result data from post-plan resume.")
+        return
+
+    output_messages = cast(list[Message], post_result)
+    print("\n=== Final Answer (post-plan resume) ===")
+    # The output of the Magentic workflow is a list of ChatMessages with only one final message
+    # generated by the orchestrator.
+    print(output_messages[-1].text)
+
+    """
+    Sample Output:
+
+    === Stage 1: run until plan review request (checkpointing active) ===
+    Captured plan review request: 3a1a4a09-4ed1-4c90-9cf6-9ac488d452c0
+    Using checkpoint 4c76d77a-6ff8-4d2b-84f6-824771ffac7e at iteration 1
+    Pending plan-review requests persisted in checkpoint: ['3a1a4a09-4ed1-4c90-9cf6-9ac488d452c0']
+
+    === Stage 2: resume from checkpoint and approve plan ===
+
+    === Final Answer ===
+    Certainly! Here's your concise internal brief on how the research and implementation teams should collaborate for
+    the beta launch of the data-driven email summarization feature:
+
+    ---
+
+    **Internal Brief: Collaboration Plan for Data-driven Email Summarization Beta Launch**
+
+    **Collaboration Approach**
+    - **Joint Kickoff:** Research and Implementation teams hold a project kickoff to align on objectives, requirements,
+        and success metrics.
+    - **Ongoing Coordination:** Teams collaborate closely; researchers share model developments and insights, while
+        implementation ensures smooth integration and user experience.
+    - **Real-time Feedback Loop:** Implementation provides early feedback on technical integration and UX, while
+        Research evaluates initial performance and user engagement signals post-integration.
+
+    **Key Milestones**
+    1. **Requirement Finalization & Scoping** - Define MVP feature set and success criteria.
+    2. **Model Prototyping & Evaluation** - Researchers develop and validate summarization models with agreed metrics.
+    3. **Integration & Internal Testing** - Implementation team integrates the model; internal alpha testing and
+        compliance checks.
+    4. **Beta User Onboarding** - Recruit a select cohort of beta users and guide them through onboarding.
+    5. **Beta Launch & Monitoring** - Soft-launch for beta group, with active monitoring of usage, feedback,
+      and performance.
+    6. **Iterative Improvements** - Address issues, refine features, and prepare for possible broader rollout.
+
+    **Top Risks**
+    - **Data Privacy & Compliance:** Strict protocols and compliance reviews to prevent data leakage.
+    - **Model Quality (Bias, Hallucination):** Careful monitoring of summary accuracy; rapid iterations if critical
+        errors occur.
+    - **User Adoption:** Ensuring the beta solves genuine user needs, collecting actionable feedback early.
+    - **Feedback Quality & Quantity:** Proactively schedule user outreach to ensure substantive beta feedback.
+
+    **Communication Cadence**
+    - **Weekly Team Syncs:** Short all-hands progress and blockers meeting.
+    - **Bi-Weekly Stakeholder Check-ins:** Leadership and project leads address escalations and strategic decisions.
+    - **Dedicated Slack Channel:** For real-time queries and updates.
+    - **Documentation Hub:** Up-to-date project docs and FAQs on a shared internal wiki.
+    - **Post-Milestone Retrospectives:** After critical phases (e.g., alpha, beta), reviewing what worked and what needs
+        improvement.
+
+    **Summary**
+    Clear alignment, consistent communication, and iterative feedback are key to a successful beta. All team members are
+        expected to surface issues quickly and keep documentation current as we drive toward launch.
+    ---
+
+    === Stage 3: resume from post-plan checkpoint ===
+    Resuming from checkpoint 9a3b... at iteration 3 (pending messages: 0)
+    No new events were emitted; checkpoint already captured a completed run.
+
+    === Final Answer (post-plan resume) ===
+    (same brief as above)
+    """
+
+
+if __name__ == "__main__":
+    asyncio.run(main())