diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index 4215400e63..0817439b9a 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -15,7 +15,7 @@ - + diff --git a/dotnet/samples/GettingStarted/FoundryAgents/README.md b/dotnet/samples/GettingStarted/FoundryAgents/README.md new file mode 100644 index 0000000000..e1007473b2 --- /dev/null +++ b/dotnet/samples/GettingStarted/FoundryAgents/README.md @@ -0,0 +1,83 @@ +# Getting started with Foundry Agents + +The getting started with Foundry Agents samples demonstrate the fundamental concepts and functionalities +of Azure Foundry Agents and can be used with Azure Foundry as the AI provider. + +These samples showcase how to work with agents managed through Azure Foundry, including agent creation, +versioning, multi-turn conversations, and advanced features like code interpretation and computer use. + +## Getting started with Foundry Agents prerequisites + +Before you begin, ensure you have the following prerequisites: + +- .NET 8.0 SDK or later +- Azure Foundry service endpoint and project configured +- Azure CLI installed and authenticated (for Azure credential authentication) + +**Note**: These samples use Azure Foundry Agents. For more information, see [Azure AI Foundry documentation](https://learn.microsoft.com/en-us/azure/ai-foundry/). + +**Note**: These samples use Azure CLI credentials for authentication. Make sure you're logged in with `az login` and have access to the Azure Foundry resource. For more information, see the [Azure CLI documentation](https://learn.microsoft.com/cli/azure/authenticate-azure-cli-interactively). + +## Samples + +|Sample|Description| +|---|---| +|[Basics](./FoundryAgents_Step01.1_Basics/)|This sample demonstrates how to create and manage AI agents with versioning| +|[Running a simple agent](./FoundryAgents_Step01.2_Running/)|This sample demonstrates how to create and run a basic Foundry agent| +|[Multi-turn conversation](./FoundryAgents_Step02_MultiturnConversation/)|This sample demonstrates how to implement a multi-turn conversation with a Foundry agent| +|[Using function tools](./FoundryAgents_Step03.1_UsingFunctionTools/)|This sample demonstrates how to use function tools with a Foundry agent| +|[Using OpenAPI function tools](./FoundryAgents_Step03.2_UsingFunctionTools_FromOpenAPI/)|This sample demonstrates how to create function tools from an OpenAPI spec and use them with a Foundry agent| +|[Using function tools with approvals](./FoundryAgents_Step04_UsingFunctionToolsWithApprovals/)|This sample demonstrates how to use function tools where approvals require human in the loop approvals before execution| +|[Structured output](./FoundryAgents_Step05_StructuredOutput/)|This sample demonstrates how to use structured output with a Foundry agent| +|[Persisted conversations](./FoundryAgents_Step06_PersistedConversations/)|This sample demonstrates how to persist conversations and reload them later| +|[Observability](./FoundryAgents_Step07_Observability/)|This sample demonstrates how to add telemetry to a Foundry agent| +|[Dependency injection](./FoundryAgents_Step08_DependencyInjection/)|This sample demonstrates how to add and resolve a Foundry agent with a dependency injection container| +|[Using MCP client as tools](./FoundryAgents_Step09_UsingMcpClientAsTools/)|This sample demonstrates how to use MCP clients as tools with a Foundry agent| +|[Using images](./FoundryAgents_Step10_UsingImages/)|This sample demonstrates how to use image multi-modality with a Foundry agent| +|[Exposing as a function tool](./FoundryAgents_Step11_AsFunctionTool/)|This sample demonstrates how to expose a Foundry agent as a function tool| +|[Using middleware](./FoundryAgents_Step12_Middleware/)|This sample demonstrates how to use middleware with a Foundry agent| +|[Using plugins](./FoundryAgents_Step13_Plugins/)|This sample demonstrates how to use plugins with a Foundry agent| +|[Code interpreter](./FoundryAgents_Step14_CodeInterpreter/)|This sample demonstrates how to use the code interpreter tool with a Foundry agent| +|[Computer use](./FoundryAgents_Step15_ComputerUse/)|This sample demonstrates how to use computer use capabilities with a Foundry agent| + +## Running the samples from the console + +To run the samples, navigate to the desired sample directory, e.g. + +```powershell +cd FoundryAgents_Step01.2_Running +``` + +Set the following environment variables: + +```powershell +$env:AZURE_FOUNDRY_PROJECT_ENDPOINT="https://your-foundry-service.services.ai.azure.com/api/projects/your-foundry-project" # Replace with your Azure Foundry resource endpoint +$env:AZURE_FOUNDRY_PROJECT_DEPLOYMENT_NAME="gpt-4o-mini" # Optional, defaults to gpt-4o-mini +``` + +If the variables are not set, you will be prompted for the values when running the samples. + +Execute the following command to build the sample: + +```powershell +dotnet build +``` + +Execute the following command to run the sample: + +```powershell +dotnet run --no-build +``` + +Or just build and run in one step: + +```powershell +dotnet run +``` + +## Running the samples from Visual Studio + +Open the solution in Visual Studio and set the desired sample project as the startup project. Then, run the project using the built-in debugger or by pressing `F5`. + +You will be prompted for any required environment variables if they are not already set. + diff --git a/dotnet/samples/GettingStarted/README.md b/dotnet/samples/GettingStarted/README.md index 4fdf0a3d7f..4e349c7742 100644 --- a/dotnet/samples/GettingStarted/README.md +++ b/dotnet/samples/GettingStarted/README.md @@ -8,6 +8,7 @@ of the agent framework. |Sample|Description| |---|---| |[Agents](./Agents/README.md)|Step by step instructions for getting started with agents| +|[Foundry Agents](./FoundryAgents/README.md)|Getting started with Azure Foundry Agents| |[Agent Providers](./AgentProviders/README.md)|Getting started with creating agents using various providers| |[Agents With Retrieval Augmented Generation (RAG)](./AgentWithRAG/README.md)|Adding Retrieval Augmented Generation (RAG) capabilities to your agents.| |[Agents With Memory](./AgentWithMemory/README.md)|Adding Memory capabilities to your agents.| diff --git a/python/packages/core/agent_framework/_workflows/__init__.py b/python/packages/core/agent_framework/_workflows/__init__.py index 6c4948f0ea..18dd674a92 100644 --- a/python/packages/core/agent_framework/_workflows/__init__.py +++ b/python/packages/core/agent_framework/_workflows/__init__.py @@ -37,6 +37,8 @@ from ._events import ( ExecutorFailedEvent, ExecutorInvokedEvent, RequestInfoEvent, + SuperStepCompletedEvent, + SuperStepStartedEvent, WorkflowErrorDetails, WorkflowEvent, WorkflowEventSource, @@ -152,6 +154,8 @@ __all__ = [ "StandardMagenticManager", "SubWorkflowRequestMessage", "SubWorkflowResponseMessage", + "SuperStepCompletedEvent", + "SuperStepStartedEvent", "SwitchCaseEdgeGroup", "SwitchCaseEdgeGroupCase", "SwitchCaseEdgeGroupDefault", diff --git a/python/packages/core/agent_framework/_workflows/__init__.pyi b/python/packages/core/agent_framework/_workflows/__init__.pyi index 44247d685c..c9f8c6cb62 100644 --- a/python/packages/core/agent_framework/_workflows/__init__.pyi +++ b/python/packages/core/agent_framework/_workflows/__init__.pyi @@ -35,6 +35,8 @@ from ._events import ( ExecutorFailedEvent, ExecutorInvokedEvent, RequestInfoEvent, + SuperStepCompletedEvent, + SuperStepStartedEvent, WorkflowErrorDetails, WorkflowEvent, WorkflowEventSource, @@ -148,6 +150,8 @@ __all__ = [ "StandardMagenticManager", "SubWorkflowRequestMessage", "SubWorkflowResponseMessage", + "SuperStepCompletedEvent", + "SuperStepStartedEvent", "SwitchCaseEdgeGroup", "SwitchCaseEdgeGroupCase", "SwitchCaseEdgeGroupDefault", diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 8fa85b7f84..358cee94dd 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. import logging +import sys from dataclasses import dataclass from typing import Any, cast @@ -20,6 +21,11 @@ from ._message_utils import normalize_messages_input from ._request_info_mixin import response_handler from ._workflow_context import WorkflowContext +if sys.version_info >= (3, 12): + from typing import override +else: + from typing_extensions import override + logger = logging.getLogger(__name__) @@ -179,7 +185,8 @@ class AgentExecutor(Executor): self._pending_responses_to_agent.clear() await self._run_agent_and_emit(ctx) - async def snapshot_state(self) -> dict[str, Any]: + @override + async def on_checkpoint_save(self) -> dict[str, Any]: """Capture current executor state for checkpointing. NOTE: if the thread storage is on the server side, the full thread state @@ -196,9 +203,6 @@ class AgentExecutor(Executor): client_module = self._agent.chat_client.__class__.__module__ if client_class_name == "AzureAIAgentClient" and "azure_ai" in client_module: - # TODO(TaoChenOSU): update this warning when we surface the hooks for - # custom executor checkpointing. - # https://github.com/microsoft/agent-framework/issues/1816 logger.warning( "Checkpointing an AgentExecutor with AzureAIAgentClient that uses server-side threads. " "Currently, checkpointing does not capture messages from server-side threads " @@ -217,7 +221,8 @@ class AgentExecutor(Executor): "pending_responses_to_agent": encode_checkpoint_value(self._pending_responses_to_agent), } - async def restore_state(self, state: dict[str, Any]) -> None: + @override + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: """Restore executor state from checkpoint. Args: diff --git a/python/packages/core/agent_framework/_workflows/_base_group_chat_orchestrator.py b/python/packages/core/agent_framework/_workflows/_base_group_chat_orchestrator.py index 5752febab5..8b49de740c 100644 --- a/python/packages/core/agent_framework/_workflows/_base_group_chat_orchestrator.py +++ b/python/packages/core/agent_framework/_workflows/_base_group_chat_orchestrator.py @@ -4,6 +4,7 @@ import inspect import logging +import sys from abc import ABC, abstractmethod from collections.abc import Awaitable, Callable, Sequence from typing import Any @@ -13,6 +14,12 @@ from ._executor import Executor from ._orchestrator_helpers import ParticipantRegistry from ._workflow_context import WorkflowContext +if sys.version_info >= (3, 12): + from typing import override +else: + from typing_extensions import override + + logger = logging.getLogger(__name__) @@ -210,11 +217,12 @@ class BaseGroupChatOrchestrator(Executor, ABC): # State persistence (shared across all patterns) - def snapshot_state(self) -> dict[str, Any]: + @override + async def on_checkpoint_save(self) -> dict[str, Any]: """Capture current orchestrator state for checkpointing. Default implementation uses OrchestrationState to serialize common state. - Subclasses should override _snapshot_pattern_metadata() to add pattern-specific data. + Subclasses can override this method or _snapshot_pattern_metadata() to add pattern-specific data. Returns: Serialized state dict @@ -238,11 +246,12 @@ class BaseGroupChatOrchestrator(Executor, ABC): """ return {} - def restore_state(self, state: dict[str, Any]) -> None: + @override + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: """Restore orchestrator state from checkpoint. Default implementation uses OrchestrationState to deserialize common state. - Subclasses should override _restore_pattern_metadata() to restore pattern-specific data. + Subclasses can override this method or _restore_pattern_metadata() to restore pattern-specific data. Args: state: Serialized state dict diff --git a/python/packages/core/agent_framework/_workflows/_conversation_history.py b/python/packages/core/agent_framework/_workflows/_conversation_history.py index 7e19671b27..52d7d99c74 100644 --- a/python/packages/core/agent_framework/_workflows/_conversation_history.py +++ b/python/packages/core/agent_framework/_workflows/_conversation_history.py @@ -6,9 +6,7 @@ These utilities operate on standard `list[ChatMessage]` collections and simple dictionary snapshots so orchestrators can share logic without new mixins. """ -import json -from collections.abc import Mapping, Sequence -from typing import Any +from collections.abc import Sequence from .._types import ChatMessage @@ -26,25 +24,3 @@ def ensure_author(message: ChatMessage, fallback: str) -> ChatMessage: """Attach `fallback` author if message is missing `author_name`.""" message.author_name = message.author_name or fallback return message - - -def snapshot_state(conversation: Sequence[ChatMessage]) -> dict[str, Any]: - """Build an immutable snapshot for checkpoint storage.""" - if hasattr(conversation, "to_dict"): - result = conversation.to_dict() # type: ignore[attr-defined] - if isinstance(result, dict): - return result # type: ignore[return-value] - if isinstance(result, Mapping): - return dict(result) # type: ignore[arg-type] - serialisable: list[dict[str, Any]] = [] - for message in conversation: - if hasattr(message, "to_dict") and callable(message.to_dict): # type: ignore[attr-defined] - msg_dict = message.to_dict() # type: ignore[attr-defined] - serialisable.append(dict(msg_dict) if isinstance(msg_dict, Mapping) else msg_dict) # type: ignore[arg-type] - elif hasattr(message, "to_json") and callable(message.to_json): # type: ignore[attr-defined] - json_payload = message.to_json() # type: ignore[attr-defined] - parsed = json.loads(json_payload) if isinstance(json_payload, str) else json_payload - serialisable.append(dict(parsed) if isinstance(parsed, Mapping) else parsed) # type: ignore[arg-type] - else: - serialisable.append(dict(getattr(message, "__dict__", {}))) # type: ignore[arg-type] - return {"messages": serialisable} diff --git a/python/packages/core/agent_framework/_workflows/_events.py b/python/packages/core/agent_framework/_workflows/_events.py index 76ae7f8a4f..b681544876 100644 --- a/python/packages/core/agent_framework/_workflows/_events.py +++ b/python/packages/core/agent_framework/_workflows/_events.py @@ -294,6 +294,36 @@ class WorkflowOutputEvent(WorkflowEvent): return f"{self.__class__.__name__}(data={self.data}, source_executor_id={self.source_executor_id})" +class SuperStepEvent(WorkflowEvent): + """Event triggered when a superstep starts or ends.""" + + def __init__(self, iteration: int, data: Any | None = None): + """Initialize the superstep event. + + Args: + iteration: The number of the superstep (1-based index). + data: Optional data associated with the superstep event. + """ + super().__init__(data) + self.iteration = iteration + + def __repr__(self) -> str: + """Return a string representation of the superstep event.""" + return f"{self.__class__.__name__}(iteration={self.iteration}, data={self.data})" + + +class SuperStepStartedEvent(SuperStepEvent): + """Event triggered when a superstep starts.""" + + ... + + +class SuperStepCompletedEvent(SuperStepEvent): + """Event triggered when a superstep ends.""" + + ... + + class ExecutorEvent(WorkflowEvent): """Base class for executor events.""" @@ -310,17 +340,13 @@ class ExecutorEvent(WorkflowEvent): class ExecutorInvokedEvent(ExecutorEvent): """Event triggered when an executor handler is invoked.""" - def __repr__(self) -> str: - """Return a string representation of the executor handler invoke event.""" - return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})" + ... class ExecutorCompletedEvent(ExecutorEvent): """Event triggered when an executor handler is completed.""" - def __repr__(self) -> str: - """Return a string representation of the executor handler complete event.""" - return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})" + ... class ExecutorFailedEvent(ExecutorEvent): diff --git a/python/packages/core/agent_framework/_workflows/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 1563dd7c53..80df16592b 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -155,6 +155,11 @@ class Executor(RequestInfoMixin, DictConvertible): that parent workflows can intercept. See WorkflowExecutor documentation for details on workflow composition patterns and request/response handling. + ## State Management + Executors can contain states that persist across workflow runs and checkpoints. Override the + `on_checkpoint_save` and `on_checkpoint_restore` methods to implement custom state + serialization and restoration logic. + ## Implementation Notes - Do not call `execute()` directly - it's invoked by the workflow engine - Do not override `execute()` - define handlers using decorators instead @@ -460,6 +465,32 @@ class Executor(RequestInfoMixin, DictConvertible): return self._handlers[message_type] raise RuntimeError(f"Executor {self.__class__.__name__} cannot handle message of type {type(message)}.") + async def on_checkpoint_save(self) -> dict[str, Any]: + """Hook called when the workflow is being saved to a checkpoint. + + Override this method in subclasses to implement custom logic that should + return state to be saved in the checkpoint. + + The returned state dictionary will be passed to `on_checkpoint_restore` + when the workflow is restored from the checkpoint. The dictionary should + only contain JSON-serializable data. + + Returns: + A state dictionary to be saved during checkpointing. + """ + return {} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + """Hook called when the workflow is restored from a checkpoint. + + Override this method in subclasses to implement custom logic that should + run when the workflow is restored from a checkpoint. + + Args: + state: The state dictionary that was saved during checkpointing. + """ + ... + # endregion: Executor diff --git a/python/packages/core/agent_framework/_workflows/_handoff.py b/python/packages/core/agent_framework/_workflows/_handoff.py index c29e3f55ad..d18bc59562 100644 --- a/python/packages/core/agent_framework/_workflows/_handoff.py +++ b/python/packages/core/agent_framework/_workflows/_handoff.py @@ -16,6 +16,7 @@ Key properties: import logging import re +import sys from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass, field from typing import Any @@ -50,6 +51,12 @@ from ._workflow import Workflow from ._workflow_builder import WorkflowBuilder from ._workflow_context import WorkflowContext +if sys.version_info >= (3, 12): + from typing import override +else: + from typing_extensions import override + + logger = logging.getLogger(__name__) @@ -307,15 +314,6 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator): ctx: WorkflowContext[AgentExecutorRequest | list[ChatMessage], list[ChatMessage] | _ConversationForUserInput], ) -> None: """Process an agent's response and determine whether to route, request input, or terminate.""" - # Hydrate coordinator state (and detect new run) using checkpointable executor state - state = await ctx.get_executor_state() - if not state: - self._clear_conversation() - elif not self._get_conversation(): - restored = self._restore_conversation_from_state(state) - if restored: - self._conversation = list(restored) - source = ctx.get_source_executor_id() is_starting_agent = source == self._starting_agent_id @@ -343,7 +341,7 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator): # Update current agent when handoff occurs self._current_agent_id = target logger.info(f"Handoff detected: {source} -> {target}. Routing control to specialist '{target}'.") - await self._persist_state(ctx) + # Clean tool-related content before sending to next agent cleaned = clean_conversation_for_handoff(conversation) request = AgentExecutorRequest(messages=cleaned, should_respond=True) @@ -360,7 +358,6 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator): f"Agent '{source}' responded without handoff. " f"Requesting user input. Return-to-previous: {self._return_to_previous}" ) - await self._persist_state(ctx) if await self._check_termination(): # Clean the output conversation for display @@ -388,7 +385,6 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator): """Receive full conversation with new user input from gateway, update history, trim for agent.""" # Update authoritative conversation self._conversation = list(message.full_conversation) - await self._persist_state(ctx) # Check termination before sending to agent if await self._check_termination(): @@ -473,11 +469,7 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator): ) return list(conversation) - async def _persist_state(self, ctx: WorkflowContext[Any, Any]) -> None: - """Store authoritative conversation snapshot without losing rich metadata.""" - state_payload = self.snapshot_state() - await ctx.set_executor_state(state_payload) - + @override def _snapshot_pattern_metadata(self) -> dict[str, Any]: """Serialize pattern-specific state. @@ -492,6 +484,7 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator): } return {} + @override def _restore_pattern_metadata(self, metadata: dict[str, Any]) -> None: """Restore pattern-specific state. @@ -503,17 +496,6 @@ class _HandoffCoordinator(BaseGroupChatOrchestrator): if self._return_to_previous and "current_agent_id" in metadata: self._current_agent_id = metadata["current_agent_id"] - def _restore_conversation_from_state(self, state: Mapping[str, Any]) -> list[ChatMessage]: - """Rehydrate the coordinator's conversation history from checkpointed state. - - DEPRECATED: Use restore_state() instead. Kept for backward compatibility. - """ - from ._orchestration_state import OrchestrationState - - orch_state_dict = {"conversation": state.get("full_conversation", state.get("conversation", []))} - temp_state = OrchestrationState.from_dict(orch_state_dict) - return list(temp_state.conversation) - def _apply_response_metadata(self, conversation: list[ChatMessage], agent_response: AgentRunResponse) -> None: """Merge top-level response metadata into the latest assistant message.""" if not agent_response.additional_properties: diff --git a/python/packages/core/agent_framework/_workflows/_magentic.py b/python/packages/core/agent_framework/_workflows/_magentic.py index 9d21391ad8..ea6fb259a6 100644 --- a/python/packages/core/agent_framework/_workflows/_magentic.py +++ b/python/packages/core/agent_framework/_workflows/_magentic.py @@ -45,9 +45,15 @@ from ._workflow import Workflow, WorkflowRunResult from ._workflow_context import WorkflowContext if sys.version_info >= (3, 11): - from typing import Self # pragma: no cover + from typing import Self else: - from typing_extensions import Self # pragma: no cover + from typing_extensions import Self + +if sys.version_info >= (3, 12): + from typing import override +else: + from typing_extensions import override + logger = logging.getLogger(__name__) @@ -673,11 +679,11 @@ class MagenticManagerBase(ABC): """Prepare the final answer.""" ... - def snapshot_state(self) -> dict[str, Any]: + def on_checkpoint_save(self) -> dict[str, Any]: """Serialize runtime state for checkpointing.""" return {} - def restore_state(self, state: dict[str, Any]) -> None: + def on_checkpoint_restore(self, state: dict[str, Any]) -> None: """Restore runtime state from checkpoint data.""" return @@ -695,22 +701,6 @@ class StandardMagenticManager(MagenticManagerBase): task_ledger: _MagenticTaskLedger | None - def snapshot_state(self) -> dict[str, Any]: - state = super().snapshot_state() - if self.task_ledger is not None: - state = dict(state) - state["task_ledger"] = self.task_ledger.to_dict() - return state - - def restore_state(self, state: dict[str, Any]) -> None: - super().restore_state(state) - ledger = state.get("task_ledger") - if ledger is not None: - try: - self.task_ledger = _MagenticTaskLedger.from_dict(ledger) - except Exception: # pragma: no cover - defensive - logger.warning("Failed to restore manager task ledger from checkpoint state") - def __init__( self, chat_client: ChatClientProtocol, @@ -940,6 +930,22 @@ class StandardMagenticManager(MagenticManagerBase): author_name=response.author_name or MAGENTIC_MANAGER_NAME, ) + @override + def on_checkpoint_save(self) -> dict[str, Any]: + state: dict[str, Any] = {} + if self.task_ledger is not None: + state["task_ledger"] = self.task_ledger.to_dict() + return state + + @override + def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + ledger = state.get("task_ledger") + if ledger is not None: + try: + self.task_ledger = _MagenticTaskLedger.from_dict(ledger) + except Exception: # pragma: no cover - defensive + logger.warning("Failed to restore manager task ledger from checkpoint state") + # endregion Magentic Manager @@ -997,7 +1003,6 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): # Terminal state marker to stop further processing after completion/limits self._terminated = False # Tracks whether checkpoint state has been applied for this run - self._state_restored = False def _get_author_name(self) -> str: """Get the magentic manager name for orchestrator-generated messages.""" @@ -1036,7 +1041,8 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): ) await ctx.add_event(event) - def snapshot_state(self) -> dict[str, Any]: + @override + async def on_checkpoint_save(self) -> dict[str, Any]: """Capture current orchestrator state for checkpointing. Uses OrchestrationState for structure but maintains Magentic's complex metadata @@ -1055,14 +1061,16 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): state["magentic_context"] = self._context.to_dict() if self._task_ledger is not None: state["task_ledger"] = _message_to_payload(self._task_ledger) - manager_state: dict[str, Any] | None = None - with contextlib.suppress(Exception): - manager_state = self._manager.snapshot_state() - if manager_state: - state["manager_state"] = manager_state + + try: + state["manager_state"] = self._manager.on_checkpoint_save() + except Exception as exc: + logger.warning("Failed to save manager state for checkpoint: %s\nSkipping...", exc) + return state - def restore_state(self, state: dict[str, Any]) -> None: + @override + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: """Restore orchestrator state from checkpoint. Maintains backward compatibility with existing Magentic checkpoints @@ -1112,7 +1120,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): manager_state = state.get("manager_state") if manager_state is not None: try: - self._manager.restore_state(manager_state) + self._manager.on_checkpoint_restore(manager_state) except Exception as exc: # pragma: no cover logger.warning("Failed to restore manager state: %s", exc) @@ -1142,49 +1150,6 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): for name, description in expected.items(): restored[name] = description - def _snapshot_pattern_metadata(self) -> dict[str, Any]: - """Serialize pattern-specific state. - - Magentic uses custom snapshot_state() instead of base class hooks. - This method exists to satisfy the base class contract. - - Returns: - Empty dict (Magentic manages its own state) - """ - return {} - - def _restore_pattern_metadata(self, metadata: dict[str, Any]) -> None: - """Restore pattern-specific state. - - Magentic uses custom restore_state() instead of base class hooks. - This method exists to satisfy the base class contract. - - Args: - metadata: Pattern-specific state dict (ignored) - """ - pass - - async def _ensure_state_restored( - self, - context: WorkflowContext[Any, Any], - ) -> None: - if self._state_restored and self._context is not None: - return - state = await context.get_executor_state() - if not state: - self._state_restored = True - return - if not isinstance(state, dict): - self._state_restored = True - return - try: - self.restore_state(state) - except Exception as exc: # pragma: no cover - logger.warning("Magentic Orchestrator: Failed to apply checkpoint state: %s", exc, exc_info=True) - raise - else: - self._state_restored = True - @handler async def handle_start_message( self, @@ -1204,7 +1169,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): ) if message.messages: self._context.chat_history.extend(message.messages) - self._state_restored = True + # Non-streaming callback for the orchestrator receipt of the task await self._emit_orchestrator_message(context, message.task, ORCH_MSG_KIND_USER_TASK) @@ -1269,7 +1234,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): """Handle responses from agents.""" if getattr(self, "_terminated", False): return - await self._ensure_state_restored(context) + if self._context is None: raise RuntimeError("Magentic Orchestrator: Received response but not initialized") @@ -1301,7 +1266,7 @@ class MagenticOrchestratorExecutor(BaseGroupChatOrchestrator): ) -> None: if getattr(self, "_terminated", False): return - await self._ensure_state_restored(context) + if self._context is None: return @@ -1636,9 +1601,9 @@ class MagenticAgentExecutor(Executor): self._agent = agent self._agent_id = agent_id self._chat_history: list[ChatMessage] = [] - self._state_restored = False - def snapshot_state(self) -> dict[str, Any]: + @override + async def on_checkpoint_save(self) -> dict[str, Any]: """Capture current executor state for checkpointing. Returns: @@ -1650,7 +1615,8 @@ class MagenticAgentExecutor(Executor): "chat_history": encode_chat_messages(self._chat_history), } - def restore_state(self, state: dict[str, Any]) -> None: + @override + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: """Restore executor state from checkpoint. Args: @@ -1668,24 +1634,6 @@ class MagenticAgentExecutor(Executor): else: self._chat_history = [] - async def _ensure_state_restored(self, context: WorkflowContext[Any, Any]) -> None: - if self._state_restored and self._chat_history: - return - state = await context.get_executor_state() - if not state: - self._state_restored = True - return - if not isinstance(state, dict): - self._state_restored = True - return - try: - self.restore_state(state) - except Exception as exc: # pragma: no cover - logger.warning("Agent %s: Failed to apply checkpoint state: %s", self._agent_id, exc, exc_info=True) - raise - else: - self._state_restored = True - @handler async def handle_response_message( self, message: _MagenticResponseMessage, context: WorkflowContext[_MagenticResponseMessage] @@ -1693,8 +1641,6 @@ class MagenticAgentExecutor(Executor): """Handle response message (task ledger broadcast).""" logger.debug("Agent %s: Received response message", self._agent_id) - await self._ensure_state_restored(context) - # Check if this message is intended for this agent if message.target_agent is not None and message.target_agent != self._agent_id and not message.broadcast: # Message is targeted to a different agent, ignore it @@ -1735,8 +1681,6 @@ class MagenticAgentExecutor(Executor): logger.info("Agent %s: Received request to respond", self._agent_id) - await self._ensure_state_restored(context) - # Add persona adoption message with appropriate role persona_role = self._get_persona_adoption_role() persona_msg = ChatMessage( @@ -1783,7 +1727,6 @@ class MagenticAgentExecutor(Executor): """Reset the internal chat history of the agent (internal operation).""" logger.debug("Agent %s: Resetting chat history", self._agent_id) self._chat_history.clear() - self._state_restored = True async def _emit_agent_delta_event( self, diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 51ff79a864..8cc01c23cf 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -11,7 +11,7 @@ from ._checkpoint_encoding import DATACLASS_MARKER, MODEL_MARKER, decode_checkpo from ._const import EXECUTOR_STATE_KEY from ._edge import EdgeGroup from ._edge_runner import EdgeRunner, create_edge_runner -from ._events import WorkflowEvent +from ._events import SuperStepCompletedEvent, SuperStepStartedEvent, WorkflowEvent from ._executor import Executor from ._runner_context import ( Message, @@ -92,6 +92,7 @@ class Runner: while self._iteration < self._max_iterations: logger.info(f"Starting superstep {self._iteration + 1}") + yield SuperStepStartedEvent(iteration=self._iteration + 1) # Run iteration concurrently with live event streaming: we poll # for new events while the iteration coroutine progresses. @@ -126,6 +127,9 @@ class Runner: # Create checkpoint after each superstep iteration await self._create_checkpoint_if_enabled(f"superstep_{self._iteration}") + yield SuperStepCompletedEvent(iteration=self._iteration) + + # Check for convergence: no more messages to process if not await self._ctx.has_messages(): break @@ -183,8 +187,8 @@ class Runner: return None try: - # Auto-snapshot executor states - await self._auto_snapshot_executor_states() + # Snapshot executor states + await self._save_executor_states() checkpoint_category = "initial" if checkpoint_type == "after_initial_execution" else "superstep" metadata = { "superstep": self._iteration, @@ -203,41 +207,6 @@ class Runner: logger.warning(f"Failed to create {checkpoint_type} checkpoint: {e}") return None - async def _auto_snapshot_executor_states(self) -> None: - """Populate executor state by calling snapshot hooks on executors if available. - - TODO(@taochen#1614): this method is potentially problematic if executors also call - set_executor_state on the context directly. We should clarify the intended usage - pattern for executor state management. - - Convention: - - If an executor defines an async or sync method `snapshot_state(self) -> dict`, use it. - - Else if it has a plain attribute `state` that is a dict, use that. - Only JSON-serializable dicts should be provided by executors. - """ - for exec_id, executor in self._executors.items(): - state_dict: dict[str, Any] | None = None - snapshot = getattr(executor, "snapshot_state", None) - try: - if callable(snapshot): - maybe = snapshot() - if asyncio.iscoroutine(maybe): # type: ignore[arg-type] - maybe = await maybe # type: ignore[assignment] - if isinstance(maybe, dict): - state_dict = maybe # type: ignore[assignment] - else: - state_attr = getattr(executor, "state", None) - if isinstance(state_attr, dict): - state_dict = state_attr # type: ignore[assignment] - except Exception as ex: # pragma: no cover - logger.debug(f"Executor {exec_id} snapshot_state failed: {ex}") - - if state_dict is not None: - try: - await self._set_executor_state(exec_id, state_dict) - except Exception as ex: # pragma: no cover - logger.debug(f"Failed to persist state for executor {exec_id}: {ex}") - async def restore_from_checkpoint( self, checkpoint_id: str, @@ -300,7 +269,65 @@ class Runner: logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}") return False + async def _save_executor_states(self) -> None: + """Populate executor state by calling checkpoint hooks on executors. + + Backward compatibility behavior: + - If an executor defines an async or sync method `snapshot_state(self) -> dict`, use it. + - Else if it has a plain attribute `state` that is a dict, use that. + + Updated behavior: + - Executors should implement `on_checkpoint_save(self) -> dict` to provide state. + + This method will try the backward compatibility behavior first; if that does not yield state, + it falls back to the updated behavior. + + Only JSON-serializable dicts should be provided by executors. + """ + for exec_id, executor in self._executors.items(): + state_dict: dict[str, Any] | None = None + # Try backward compatibility behavior first + # TODO(@taochen): Remove backward compatibility + snapshot = getattr(executor, "snapshot_state", None) + try: + if callable(snapshot): + maybe = snapshot() + if asyncio.iscoroutine(maybe): # type: ignore[arg-type] + maybe = await maybe # type: ignore[assignment] + if isinstance(maybe, dict): + state_dict = maybe # type: ignore[assignment] + else: + state_attr = getattr(executor, "state", None) + if isinstance(state_attr, dict): + state_dict = state_attr # type: ignore[assignment] + except Exception as ex: # pragma: no cover + logger.debug(f"Executor {exec_id} snapshot_state failed: {ex}") + + if state_dict is None: + # Try the updated behavior only if backward compatibility did not yield state + try: + state_dict = await executor.on_checkpoint_save() + except Exception as ex: # pragma: no cover + raise ValueError(f"Executor {exec_id} on_checkpoint_save failed: {ex}") from ex + + try: + await self._set_executor_state(exec_id, state_dict) + except Exception as ex: # pragma: no cover + logger.debug(f"Failed to persist state for executor {exec_id}: {ex}") + async def _restore_executor_states(self) -> None: + """Restore executor state by calling restore hooks on executors. + + Backward compatibility behavior: + - If an executor defines an async or sync method `restore_state(self, state: dict)`, use it. + - Else, skip restoration for that executor. + + Updated behavior: + - Executors should implement `on_checkpoint_restore(self, state: dict)` to restore state. + + This method will try the backward compatibility behavior first; if that does not restore state, + it falls back to the updated behavior. + """ has_executor_states = await self._shared_state.has(EXECUTOR_STATE_KEY) if not has_executor_states: return @@ -309,16 +336,18 @@ class Runner: if not isinstance(executor_states, dict): raise ValueError("Executor states in shared state is not a dictionary. Unable to restore.") - for executor_id, state in executor_states.items(): + for executor_id, state in executor_states.items(): # pyright: ignore[reportUnknownVariableType] if not isinstance(executor_id, str): raise ValueError("Executor ID in executor states is not a string. Unable to restore.") - if not isinstance(state, dict): - raise ValueError(f"Executor state for {executor_id} is not a dictionary. Unable to restore.") + if not isinstance(state, dict) or not all(isinstance(k, str) for k in state): # pyright: ignore[reportUnknownVariableType] + raise ValueError(f"Executor state for {executor_id} is not a dict[str, Any]. Unable to restore.") executor = self._executors.get(executor_id) if not executor: raise ValueError(f"Executor {executor_id} not found during state restoration.") + # Try backward compatibility behavior first + # TODO(@taochen): Remove backward compatibility restored = False restore_method = getattr(executor, "restore_state", None) try: @@ -330,6 +359,14 @@ class Runner: except Exception as ex: # pragma: no cover - defensive raise ValueError(f"Executor {executor_id} restore_state failed: {ex}") from ex + if not restored: + # Try the updated behavior only if backward compatibility did not restore + try: + await executor.on_checkpoint_restore(state) # pyright: ignore[reportUnknownArgumentType] + restored = True + except Exception as ex: # pragma: no cover - defensive + raise ValueError(f"Executor {executor_id} on_checkpoint_restore failed: {ex}") from ex + if not restored: logger.debug(f"Executor {executor_id} does not support state restoration; skipping.") diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index e5fd02a611..a14542b2a6 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -109,9 +109,9 @@ class Workflow(DictConvertible): """A graph-based execution engine that orchestrates connected executors. ## Overview - A workflow executes a directed graph of executors connected via edge groups using a Pregel-like model, - running in supersteps until the graph becomes idle. Workflows are created using the - WorkflowBuilder class - do not instantiate this class directly. + A workflow executes a directed graph of executors connected via edge groups using a + Pregel-like model, running in supersteps until the graph becomes idle. Workflows + are created using the WorkflowBuilder class - do not instantiate this class directly. ## Execution Model Executors run in synchronized supersteps where each executor: @@ -142,6 +142,10 @@ class Workflow(DictConvertible): - HIL continuation: Provide `responses` to continue after RequestInfoExecutor requests - Runtime checkpointing: Provide `checkpoint_storage` to enable/override checkpointing for this run + ## State Management + Workflow instances contain states and states are preserved across calls to `run` and `run_stream`. + To execute multiple independent runs, create separate Workflow instances via WorkflowBuilder. + ## External Input Requests Executors within a workflow can request external input using `ctx.request_info()`: 1. Executor calls `ctx.request_info()` to request input diff --git a/python/packages/core/agent_framework/_workflows/_workflow_context.py b/python/packages/core/agent_framework/_workflows/_workflow_context.py index d2a3648298..dcf6715d62 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_context.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_context.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any, Generic, Union, cast, get_args, get_origi from opentelemetry.propagate import inject from opentelemetry.trace import SpanKind -from typing_extensions import Never, TypeVar +from typing_extensions import Never, TypeVar, deprecated from ..observability import OtelAttr, create_workflow_span from ._const import EXECUTOR_STATE_KEY @@ -410,6 +410,11 @@ class WorkflowContext(Generic[T_Out, T_W_Out]): """Get the shared state.""" return self._shared_state + @deprecated( + "Override `on_checkpoint_save()` methods instead. " + "For cross-executor state sharing, use set_shared_state() instead. " + "This API will be removed after 12/01/2025." + ) async def set_executor_state(self, state: dict[str, Any]) -> None: """Store executor state in shared state under a reserved key. @@ -428,6 +433,11 @@ class WorkflowContext(Generic[T_Out, T_W_Out]): existing_states[self._executor_id] = state await self._shared_state.set(EXECUTOR_STATE_KEY, existing_states) + @deprecated( + "Override `on_checkpoint_restore()` methods instead. " + "For cross-executor state sharing, use get_shared_state() instead. " + "This API will be removed after 12/01/2025." + ) async def get_executor_state(self) -> dict[str, Any] | None: """Retrieve previously persisted state for this executor, if any.""" has_existing_states = await self._shared_state.has(EXECUTOR_STATE_KEY) diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 77acbc5a58..cc028f337c 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -1,8 +1,8 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio -import contextlib import logging +import sys import uuid from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -26,6 +26,12 @@ from ._typing_utils import is_instance_of from ._workflow import WorkflowRunResult from ._workflow_context import WorkflowContext +if sys.version_info >= (3, 12): + from typing import override +else: + from typing_extensions import override + + logger = logging.getLogger(__name__) @@ -181,8 +187,7 @@ class WorkflowExecutor(Executor): # Includes all sub-workflow output types # Plus SubWorkflowRequestMessage if sub-workflow can make requests - output_types = workflow.output_types + [SubWorkflowRequestMessage] # if applicable - ``` + output_types = workflow.output_types + [SubWorkflowRequestMessage] # if applicable ## Error Handling WorkflowExecutor propagates sub-workflow failures: @@ -221,23 +226,10 @@ class WorkflowExecutor(Executor): ### Important Considerations **Shared Workflow Instance**: All concurrent executions use the same underlying workflow instance. - For proper isolation, ensure that: - - The wrapped workflow and its executors are stateless - - Executors use WorkflowContext state management instead of instance variables - - Any shared state is managed through WorkflowContext.get_shared_state/set_shared_state + For proper isolation, ensure that the wrapped workflow and its executors are stateless. .. code-block:: python - # Good: Stateless executor using context state - class StatelessExecutor(Executor): - @handler - async def process(self, data: str, ctx: WorkflowContext[str]) -> None: - # Use context state instead of instance variables - state = await ctx.get_executor_state() or {} - state["processed"] = data - await ctx.set_executor_state(state) - - # Avoid: Stateful executor with instance variables class StatefulExecutor(Executor): def __init__(self): @@ -246,23 +238,23 @@ class WorkflowExecutor(Executor): ## Integration with Parent Workflows Parent workflows can intercept sub-workflow requests: - ```python - class ParentExecutor(Executor): - @handler - async def handle_subworkflow_request( - self, - request: SubWorkflowRequestMessage, - ctx: WorkflowContext[SubWorkflowResponseMessage], - ) -> None: - # Handle request locally or forward to external source - if self.can_handle_locally(request): - # Send response back to sub-workflow - response = request.create_response(data="local response data") - await ctx.send_message(response, target_id=request.source_executor_id) - else: - # Forward to external handler - await ctx.request_info(request.source_event, response_type=request.source_event.response_type) - ``` + + .. code-block:: python + class ParentExecutor(Executor): + @handler + async def handle_subworkflow_request( + self, + request: SubWorkflowRequestMessage, + ctx: WorkflowContext[SubWorkflowResponseMessage], + ) -> None: + # Handle request locally or forward to external source + if self.can_handle_locally(request): + # Send response back to sub-workflow + response = request.create_response(data="local response data") + await ctx.send_message(response, target_id=request.source_executor_id) + else: + # Forward to external handler + await ctx.request_info(request.source_event, response_type=request.source_event.response_type) ## Implementation Notes - Sub-workflows run to completion before processing their results @@ -296,7 +288,6 @@ class WorkflowExecutor(Executor): self._execution_contexts: dict[str, ExecutionContext] = {} # execution_id -> ExecutionContext # Map request_id to execution_id for response routing self._request_to_execution: dict[str, str] = {} # request_id -> execution_id - self._state_loaded: bool = False @property def input_types(self) -> list[type[Any]]: @@ -362,8 +353,6 @@ class WorkflowExecutor(Executor): input_data: The input data to send to the sub-workflow. ctx: The workflow context from the parent. """ - await self._ensure_state_loaded(ctx) - # Create execution context for this sub-workflow run execution_id = str(uuid.uuid4()) execution_context = ExecutionContext( @@ -405,8 +394,6 @@ class WorkflowExecutor(Executor): response: The response to a previous request. ctx: The workflow context. """ - await self._ensure_state_loaded(ctx) - # Find the execution context for this request original_request = response.source_event execution_id = self._request_to_execution.get(original_request.request_id) @@ -434,8 +421,6 @@ class WorkflowExecutor(Executor): # Accumulate the response in this execution's context execution_context.collected_responses[original_request.request_id] = response.data - await self._persist_execution_state(ctx) - # Check if we have all expected responses for this execution if len(execution_context.collected_responses) < execution_context.expected_response_count: logger.debug( @@ -459,25 +444,20 @@ class WorkflowExecutor(Executor): if not execution_context.pending_requests: del self._execution_contexts[execution_id] - async def _ensure_state_loaded(self, ctx: WorkflowContext[Any]) -> None: - if self._state_loaded: - return + @override + async def on_checkpoint_save(self) -> dict[str, Any]: + """Get the current state of the WorkflowExecutor for checkpointing purposes.""" + return { + "execution_contexts": { + execution_id: encode_checkpoint_value(execution_context) + for execution_id, execution_context in self._execution_contexts.items() + }, + "request_to_execution": dict(self._request_to_execution), + } - state: dict[str, Any] | None = None - try: - state = await ctx.get_executor_state() - except Exception: - state = None - - if isinstance(state, dict) and state: - with contextlib.suppress(Exception): - await self.restore_state(state) - self._state_loaded = True - else: - self._state_loaded = True - - async def restore_state(self, state: dict[str, Any]) -> None: - """Restore pending request bookkeeping from a checkpoint snapshot.""" + @override + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + """Restore the WorkflowExecutor state from a checkpoint snapshot.""" # Validate the state contains the right keys if "execution_contexts" not in state: raise KeyError("Missing 'execution_contexts' in WorkflowExecutor state.") @@ -529,23 +509,6 @@ class WorkflowExecutor(Executor): for event in request_info_events ]) - self._state_loaded = True - - async def _persist_execution_state(self, ctx: WorkflowContext) -> None: - """Persist the state of the WorkflowExecutor for checkpointing purposes.""" - state = { - "execution_contexts": { - execution_id: encode_checkpoint_value(execution_context) - for execution_id, execution_context in self._execution_contexts.items() - }, - "request_to_execution": dict(self._request_to_execution), - } - - try: - await ctx.set_executor_state(state) - except Exception as exc: # pragma: no cover - transport specific - logger.warning(f"WorkflowExecutor {self.id} failed to persist state: {exc}") - async def _process_workflow_result( self, result: WorkflowRunResult, @@ -635,5 +598,3 @@ class WorkflowExecutor(Executor): ) else: raise RuntimeError(f"Unexpected workflow run state: {workflow_run_state}") - - await self._persist_execution_state(ctx) diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py index 77fd969f12..2815c3152c 100644 --- a/python/packages/core/tests/workflow/test_agent_executor.py +++ b/python/packages/core/tests/workflow/test_agent_executor.py @@ -158,8 +158,8 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None: assert thread_messages[1].text == "Initial response 1" -async def test_agent_executor_snapshot_and_restore_state_directly() -> None: - """Test AgentExecutor's snapshot_state and restore_state methods directly.""" +async def test_agent_executor_save_and_restore_state_directly() -> None: + """Test AgentExecutor's on_checkpoint_save and on_checkpoint_restore methods directly.""" # Create agent with thread containing messages agent = _CountingAgent(id="direct_test_agent", name="DirectTestAgent") thread = AgentThread(message_store=ChatMessageStore()) @@ -182,7 +182,7 @@ async def test_agent_executor_snapshot_and_restore_state_directly() -> None: executor._cache = list(cache_messages) # type: ignore[reportPrivateUsage] # Snapshot the state - state = await executor.snapshot_state() # type: ignore[reportUnknownMemberType] + state = await executor.on_checkpoint_save() # Verify snapshot contains both cache and thread assert "cache" in state @@ -206,7 +206,7 @@ async def test_agent_executor_snapshot_and_restore_state_directly() -> None: assert len(initial_thread_msgs) == 0 # Restore state - await new_executor.restore_state(state) # type: ignore[reportUnknownMemberType] + await new_executor.on_checkpoint_restore(state) # Verify cache is restored restored_cache = new_executor._cache # type: ignore[reportPrivateUsage] diff --git a/python/packages/core/tests/workflow/test_handoff.py b/python/packages/core/tests/workflow/test_handoff.py index a799fb6f73..5dfd7522df 100644 --- a/python/packages/core/tests/workflow/test_handoff.py +++ b/python/packages/core/tests/workflow/test_handoff.py @@ -288,57 +288,6 @@ def test_build_fails_without_participants(): HandoffBuilder().build() -async def test_multiple_runs_dont_leak_conversation(): - """Verify that running the same workflow multiple times doesn't leak conversation history.""" - triage = _RecordingAgent(name="triage", handoff_to="specialist") - specialist = _RecordingAgent(name="specialist") - - workflow = ( - HandoffBuilder(participants=[triage, specialist]) - .set_coordinator("triage") - .with_termination_condition(lambda conv: sum(1 for m in conv if m.role == Role.USER) >= 2) - .build() - ) - - # First run - events = await _drain(workflow.run_stream("First run message")) - requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)] - assert requests - events = await _drain(workflow.send_responses_streaming({requests[-1].request_id: "Second message"})) - outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)] - assert outputs, "First run should emit output" - - first_run_conversation = outputs[-1].data - assert isinstance(first_run_conversation, list) - first_run_conv_list = cast(list[ChatMessage], first_run_conversation) - first_run_user_messages = [msg for msg in first_run_conv_list if msg.role == Role.USER] - assert len(first_run_user_messages) == 2 - assert any("First run message" in msg.text for msg in first_run_user_messages if msg.text) - - # Second run - should start fresh, not include first run's messages - triage.calls.clear() - specialist.calls.clear() - - events = await _drain(workflow.run_stream("Second run different message")) - requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)] - assert requests - events = await _drain(workflow.send_responses_streaming({requests[-1].request_id: "Another message"})) - outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)] - assert outputs, "Second run should emit output" - - second_run_conversation = outputs[-1].data - assert isinstance(second_run_conversation, list) - second_run_conv_list = cast(list[ChatMessage], second_run_conversation) - second_run_user_messages = [msg for msg in second_run_conv_list if msg.role == Role.USER] - assert len(second_run_user_messages) == 2, ( - "Second run should have exactly 2 user messages, not accumulate first run" - ) - assert any("Second run different message" in msg.text for msg in second_run_user_messages if msg.text) - assert not any("First run message" in msg.text for msg in second_run_user_messages if msg.text), ( - "Second run should NOT contain first run's messages" - ) - - async def test_handoff_async_termination_condition() -> None: """Test that async termination conditions work correctly.""" termination_call_count = 0 @@ -585,7 +534,7 @@ async def test_return_to_previous_state_serialization(): coordinator._current_agent_id = "specialist_a" # type: ignore[reportPrivateUsage] # Snapshot the state - state = coordinator.snapshot_state() + state = await coordinator.on_checkpoint_save() # Verify pattern metadata includes current_agent_id assert "metadata" in state @@ -603,7 +552,7 @@ async def test_return_to_previous_state_serialization(): ) # Restore state - coordinator2.restore_state(state) + await coordinator2.on_checkpoint_restore(state) # Verify current_agent_id was restored assert coordinator2._current_agent_id == "specialist_a", "Current agent should be restored from checkpoint" # type: ignore[reportPrivateUsage] diff --git a/python/packages/core/tests/workflow/test_magentic.py b/python/packages/core/tests/workflow/test_magentic.py index eda0675361..cc1e8ad132 100644 --- a/python/packages/core/tests/workflow/test_magentic.py +++ b/python/packages/core/tests/workflow/test_magentic.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import sys from collections.abc import AsyncIterable from dataclasses import dataclass from typing import Any, cast @@ -42,6 +43,11 @@ from agent_framework._workflows._magentic import ( # type: ignore[reportPrivate _MagenticStartMessage, # type: ignore ) +if sys.version_info >= (3, 12): + from typing import override +else: + from typing_extensions import override + def test_magentic_start_message_from_string(): msg = _MagenticStartMessage.from_string("Do the thing") @@ -101,8 +107,9 @@ class FakeManager(MagenticManagerBase): next_speaker_name: str = "agentA" instruction_text: str = "Proceed with step 1" - def snapshot_state(self) -> dict[str, Any]: - state = super().snapshot_state() + @override + def on_checkpoint_save(self) -> dict[str, Any]: + state = super().on_checkpoint_save() if self.task_ledger is not None: state = dict(state) state["task_ledger"] = { @@ -111,8 +118,9 @@ class FakeManager(MagenticManagerBase): } return state - def restore_state(self, state: dict[str, Any]) -> None: - super().restore_state(state) + @override + def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + super().on_checkpoint_restore(state) ledger_state = state.get("task_ledger") if isinstance(ledger_state, dict): ledger_dict = cast(dict[str, Any], ledger_state) @@ -185,7 +193,6 @@ async def test_standard_manager_progress_ledger_and_fallback(): assert ledger2.is_request_satisfied.answer is False -@pytest.mark.skip(reason="Response handling refactored - responses no longer passed to run_stream()") async def test_magentic_workflow_plan_review_approval_to_completion(): manager = FakeManager(max_round_count=10) wf = ( @@ -204,7 +211,7 @@ async def test_magentic_workflow_plan_review_approval_to_completion(): completed = False output: ChatMessage | None = None - async for ev in wf.run_stream( + async for ev in wf.send_responses_streaming( responses={req_event.request_id: MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.APPROVE)} ): if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: @@ -218,7 +225,6 @@ async def test_magentic_workflow_plan_review_approval_to_completion(): assert isinstance(output, ChatMessage) -@pytest.mark.skip(reason="Response handling refactored - responses no longer passed to run_stream()") async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds(): class CountingManager(FakeManager): # Declare as a model field so assignment is allowed under Pydantic @@ -250,7 +256,7 @@ async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds() # Reply APPROVE with comments (no edited text). Expect one replan and no second review round. saw_second_review = False completed = False - async for ev in wf.run_stream( + async for ev in wf.send_responses_streaming( responses={ req_event.request_id: MagenticPlanReviewReply( decision=MagenticPlanReviewDecision.APPROVE, @@ -298,7 +304,6 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result(): assert data.role == Role.ASSISTANT -@pytest.mark.skip(reason="Response handling refactored - send_responses_streaming no longer exists") async def test_magentic_checkpoint_resume_round_trip(): storage = InMemoryCheckpointStorage() @@ -369,7 +374,7 @@ class _DummyExec(Executor): pass -def test_magentic_agent_executor_snapshot_roundtrip(): +async def test_magentic_agent_executor_on_checkpoint_save_and_restore_roundtrip(): backing_executor = _DummyExec("backing") agent_exec = MagenticAgentExecutor(backing_executor, "agentA") agent_exec._chat_history.extend([ # type: ignore[reportPrivateUsage] @@ -377,10 +382,10 @@ def test_magentic_agent_executor_snapshot_roundtrip(): ChatMessage(role=Role.ASSISTANT, text="world", author_name="agentA"), ]) - state = agent_exec.snapshot_state() + state = await agent_exec.on_checkpoint_save() restored_executor = MagenticAgentExecutor(_DummyExec("backing2"), "agentA") - restored_executor.restore_state(state) + await restored_executor.on_checkpoint_restore(state) assert len(restored_executor._chat_history) == 2 # type: ignore[reportPrivateUsage] assert restored_executor._chat_history[0].text == "hello" # type: ignore[reportPrivateUsage] diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 7f1a7fdce6..059c94803e 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -199,7 +199,10 @@ async def test_fan_out(): # Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent # executor_b will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore) - assert len(events) == 7 + # Each superstep will emit also emit a WorkflowStartedEvent and WorkflowCompletedEvent + # This workflow will converge in 2 supersteps because executor_c will send one more message + # after executor_b completes + assert len(events) == 11 assert events.get_final_state() == WorkflowRunState.IDLE outputs = events.get_outputs() @@ -220,7 +223,9 @@ async def test_fan_out_multiple_completed_events(): # Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent # executor_b and executor_c will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore) - assert len(events) == 8 + # Each superstep will emit also emit a WorkflowStartedEvent and WorkflowCompletedEvent + # This workflow will converge in 1 superstep because executor_a and executor_b will not send further messages + assert len(events) == 10 # Multiple outputs are expected from both executors outputs = events.get_outputs() @@ -246,7 +251,8 @@ async def test_fan_in(): # Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent # aggregator will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore) - assert len(events) == 9 + # Each superstep will emit also emit a WorkflowStartedEvent and WorkflowCompletedEvent + assert len(events) == 13 assert events.get_final_state() == WorkflowRunState.IDLE outputs = events.get_outputs() diff --git a/python/packages/devui/tests/test_checkpoints.py b/python/packages/devui/tests/test_checkpoints.py index 473b4faac0..838195fdd3 100644 --- a/python/packages/devui/tests/test_checkpoints.py +++ b/python/packages/devui/tests/test_checkpoints.py @@ -37,10 +37,14 @@ class WorkflowHILRequest: class WorkflowTestExecutor(Executor): """Test executor with HIL.""" + def __init__(self, id: str) -> None: + super().__init__(id=id) + self._data_value: str | None = None + @handler async def process(self, data: WorkflowTestData, ctx: WorkflowContext) -> None: """Process data and request approval.""" - await ctx.set_executor_state({"data_value": data.value}) + self._data_value = data.value # Request HIL (checkpoint created here) await ctx.request_info(request_data=WorkflowHILRequest(question=f"Approve {data.value}?"), response_type=str) @@ -50,8 +54,7 @@ class WorkflowTestExecutor(Executor): self, original_request: WorkflowHILRequest, response: str, ctx: WorkflowContext[str] ) -> None: """Handle HIL response.""" - state = await ctx.get_executor_state() or {} - value = state.get("data_value", "") + value = self._data_value or "" await ctx.send_message(f"{value}_approved" if response.lower() == "yes" else f"{value}_rejected") diff --git a/python/samples/README.md b/python/samples/README.md index c790e60519..937c303c77 100644 --- a/python/samples/README.md +++ b/python/samples/README.md @@ -17,11 +17,11 @@ This directory contains samples demonstrating the capabilities of Microsoft Agen | [`getting_started/agents/anthropic/anthropic_basic.py`](./getting_started/agents/anthropic/anthropic_basic.py) | Agent with Anthropic Client | | [`getting_started/agents/anthropic/anthropic_advanced.py`](./getting_started/agents/anthropic/anthropic_advanced.py) | Advanced sample with `thinking` and hosted tools. | -### Azure AI +### Azure AI (based on `azure-ai-agents` V1 package) | File | Description | |------|-------------| -| [`getting_started/agents/azure_ai_agent/azure_ai_basic.py`](./getting_started/agents/azure_ai/azure_ai_basic.py) | Azure AI Agent Basic Example | +| [`getting_started/agents/azure_ai_agent/azure_ai_basic.py`](./getting_started/agents/azure_ai_agent/azure_ai_basic.py) | Azure AI Agent Basic Example | | [`getting_started/agents/azure_ai_agent/azure_ai_with_azure_ai_search.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_azure_ai_search.py) | Azure AI Agent with Azure AI Search Example | | [`getting_started/agents/azure_ai_agent/azure_ai_with_bing_grounding.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_bing_grounding.py) | Azure AI agent with Bing Grounding search for real-time web information | | [`getting_started/agents/azure_ai_agent/azure_ai_with_code_interpreter.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_code_interpreter.py) | Azure AI Agent with Code Interpreter Example | @@ -36,6 +36,28 @@ This directory contains samples demonstrating the capabilities of Microsoft Agen | [`getting_started/agents/azure_ai_agent/azure_ai_with_openapi_tools.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_openapi_tools.py) | Azure AI agent with OpenAPI tools | | [`getting_started/agents/azure_ai_agent/azure_ai_with_thread.py`](./getting_started/agents/azure_ai_agent/azure_ai_with_thread.py) | Azure AI Agent with Thread Management Example | +### Azure AI (based on `azure-ai-projects` V2 package) + +| File | Description | +|------|-------------| +| [`getting_started/agents/azure_ai/azure_ai_basic.py`](./getting_started/agents/azure_ai/azure_ai_basic.py) | Azure AI Agent Basic Example | +| [`getting_started/agents/azure_ai/azure_ai_use_latest_version.py`](./getting_started/agents/azure_ai/azure_ai_use_latest_version.py) | Azure AI Agent latest version reuse example | +| [`getting_started/agents/azure_ai/azure_ai_with_azure_ai_search.py`](./getting_started/agents/azure_ai/azure_ai_with_azure_ai_search.py) | Azure AI Agent with Azure AI Search Example | +| [`getting_started/agents/azure_ai/azure_ai_with_bing_grounding.py`](./getting_started/agents/azure_ai/azure_ai_with_bing_grounding.py) | Azure AI Agent with Bing Grounding Example | +| [`getting_started/agents/azure_ai/azure_ai_with_bing_custom_search.py`](./getting_started/agents/azure_ai/azure_ai_with_bing_custom_search.py) | Azure AI Agent with Bing Custom Search Example | +| [`getting_started/agents/azure_ai/azure_ai_with_browser_automation.py`](./getting_started/agents/azure_ai/azure_ai_with_browser_automation.py) | Azure AI Agent with Browser Automation Example | +| [`getting_started/agents/azure_ai/azure_ai_with_code_interpreter.py`](./getting_started/agents/azure_ai/azure_ai_with_code_interpreter.py) | Azure AI Agent with Code Interpreter Example | +| [`getting_started/agents/azure_ai/azure_ai_with_existing_agent.py`](./getting_started/agents/azure_ai/azure_ai_with_existing_agent.py) | Azure AI Agent with Existing Agent Example | +| [`getting_started/agents/azure_ai/azure_ai_with_existing_conversation.py`](./getting_started/agents/azure_ai/azure_ai_with_existing_conversation.py) | Azure AI Agent with Existing Conversation Example | +| [`getting_started/agents/azure_ai/azure_ai_with_explicit_settings.py`](./getting_started/agents/azure_ai/azure_ai_with_explicit_settings.py) | Azure AI Agent with Explicit Settings Example | +| [`getting_started/agents/azure_ai/azure_ai_with_file_search.py`](./getting_started/agents/azure_ai/azure_ai_with_file_search.py) | Azure AI Agent with File Search Example | +| [`getting_started/agents/azure_ai/azure_ai_with_hosted_mcp.py`](./getting_started/agents/azure_ai/azure_ai_with_hosted_mcp.py) | Azure AI Agent with Hosted MCP Example | +| [`getting_started/agents/azure_ai/azure_ai_with_response_format.py`](./getting_started/agents/azure_ai/azure_ai_with_response_format.py) | Azure AI Agent with Structured Output Example | +| [`getting_started/agents/azure_ai/azure_ai_with_thread.py`](./getting_started/agents/azure_ai/azure_ai_with_thread.py) | Azure AI Agent with Thread Management Example | +| [`getting_started/agents/azure_ai/azure_ai_with_image_generation.py`](./getting_started/agents/azure_ai/azure_ai_with_image_generation.py) | Azure AI Agent with Image Generation Example | +| [`getting_started/agents/azure_ai/azure_ai_with_microsoft_fabric.py`](./getting_started/agents/azure_ai/azure_ai_with_microsoft_fabric.py) | Azure AI Agent with Microsoft Fabric Example | +| [`getting_started/agents/azure_ai/azure_ai_with_web_search.py`](./getting_started/agents/azure_ai/azure_ai_with_web_search.py) | Azure AI Agent with Web Search Example | + ### Azure OpenAI | File | Description | diff --git a/python/samples/getting_started/agents/README.md b/python/samples/getting_started/agents/README.md index e8c5f42e9c..6f1601cdc4 100644 --- a/python/samples/getting_started/agents/README.md +++ b/python/samples/getting_started/agents/README.md @@ -8,7 +8,8 @@ This folder contains examples demonstrating how to create and use agents with di | Folder | Description | |--------|-------------| -| **[`azure_ai/`](azure_ai/)** | Create agents using Azure AI Foundry Agent Service with various tools including function tools, code interpreter, MCP integration, and thread management | +| **[`azure_ai_agent/`](azure_ai_agent/)** | Create agents using Azure AI Agent Service (based on `azure-ai-agents` V1 package) including function tools, code interpreter, MCP integration, thread management, and more. | +| **[`azure_ai/`](azure_ai/)** | Create agents using Azure AI Agent Service (based on `azure-ai-projects` [V2](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/ai/azure-ai-projects/CHANGELOG.md#200b1-2025-11-11) package) including function tools, code interpreter, MCP integration, thread management, and more. | ### Microsoft Copilot Studio Examples diff --git a/python/samples/getting_started/agents/azure_ai/README.md b/python/samples/getting_started/agents/azure_ai/README.md index 475dad5a1a..aac3a7d01e 100644 --- a/python/samples/getting_started/agents/azure_ai/README.md +++ b/python/samples/getting_started/agents/azure_ai/README.md @@ -1,6 +1,6 @@ # Azure AI Agent Examples -This folder contains examples demonstrating different ways to create and use agents with the Azure AI client from the `agent_framework.azure` package. +This folder contains examples demonstrating different ways to create and use agents with the Azure AI client from the `agent_framework.azure` package. These examples use the `AzureAIClient` with the `azure-ai-projects` 2.x (V2) API surface (see [changelog](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/ai/azure-ai-projects/CHANGELOG.md#200b1-2025-11-11)). For V1 (`azure-ai-agents` 1.x) samples using `AzureAIAgentClient`, see the [Azure AI V1 examples folder](../azure_ai_agent/). ## Examples diff --git a/python/samples/getting_started/agents/azure_ai_agent/README.md b/python/samples/getting_started/agents/azure_ai_agent/README.md index 19165b89ed..c7337bbe82 100644 --- a/python/samples/getting_started/agents/azure_ai_agent/README.md +++ b/python/samples/getting_started/agents/azure_ai_agent/README.md @@ -1,6 +1,6 @@ # Azure AI Agent Examples -This folder contains examples demonstrating different ways to create and use agents with the Azure AI chat client from the `agent_framework.azure` package. +This folder contains examples demonstrating different ways to create and use agents with the Azure AI chat client from the `agent_framework.azure` package. These examples use the `AzureAIAgentClient` with the `azure-ai-agents` 1.x (V1) API surface. For updated V2 (`azure-ai-projects` 2.x) samples, see the [Azure AI V2 examples folder](../azure_ai/). ## Examples diff --git a/python/samples/getting_started/devui/spam_workflow/workflow.py b/python/samples/getting_started/devui/spam_workflow/workflow.py index 020d22a4c7..73be349cc6 100644 --- a/python/samples/getting_started/devui/spam_workflow/workflow.py +++ b/python/samples/getting_started/devui/spam_workflow/workflow.py @@ -17,7 +17,7 @@ Workflow Steps: import asyncio import logging from dataclasses import dataclass -from typing import Literal, Annotated +from typing import Literal from agent_framework import ( Case, @@ -31,9 +31,11 @@ from agent_framework import ( from pydantic import BaseModel, Field from typing_extensions import Never + # Define response model with clear user guidance class SpamDecision(BaseModel): """User's decision on whether the email is spam.""" + decision: Literal["spam", "not spam"] = Field( description="Enter 'spam' to mark as spam, or 'not spam' to mark as legitimate" ) @@ -71,10 +73,11 @@ class SpamDetectorResponse: class SpamApprovalRequest: """Human-in-the-loop approval request for spam classification.""" - email_message: str = "" - detected_as_spam: bool = False - confidence: float = 0.0 - reasons: str = "" + email_message: str + detected_as_spam: bool + confidence: float + reasons: list[str] + full_email_content: EmailContent @dataclass @@ -128,8 +131,6 @@ class EmailPreprocessor(Executor): await ctx.send_message(result) - - class SpamDetector(Executor): """Step 2: An executor that analyzes content and determines if a message is spam.""" @@ -139,7 +140,9 @@ class SpamDetector(Executor): self._spam_keywords = spam_keywords @handler - async def handle_email_content(self, email_content: EmailContent, ctx: WorkflowContext[SpamApprovalRequest]) -> None: + async def handle_email_content( + self, email_content: EmailContent, ctx: WorkflowContext[SpamApprovalRequest] + ) -> None: """Analyze email content and determine if the message is spam, then request human approval.""" await asyncio.sleep(2.0) # Simulate analysis and detection time @@ -186,25 +189,13 @@ class SpamDetector(Executor): is_spam = spam_score >= 0.5 - # Store detection result in executor state for later use - # Store minimal data needed (not complex objects that don't serialize well) - await ctx.set_executor_state({ - "original_message": email_content.original_message, - "cleaned_message": email_content.cleaned_message, - "word_count": email_content.word_count, - "has_suspicious_patterns": email_content.has_suspicious_patterns, - "is_spam": is_spam, - "ai_original_classification": is_spam, # Store original AI decision - "confidence_score": spam_score, - "spam_reasons": spam_reasons - }) - # Request human approval before proceeding using new API approval_request = SpamApprovalRequest( email_message=email_text[:200], # First 200 chars detected_as_spam=is_spam, confidence=spam_score, - reasons=", ".join(spam_reasons) if spam_reasons else "no specific reasons" + reasons=spam_reasons, + full_email_content=email_content, ) await ctx.request_info( @@ -214,20 +205,15 @@ class SpamDetector(Executor): @response_handler async def handle_human_response( - self, - original_request: SpamApprovalRequest, - response: SpamDecision, - ctx: WorkflowContext[SpamDetectorResponse] + self, original_request: SpamApprovalRequest, response: SpamDecision, ctx: WorkflowContext[SpamDetectorResponse] ) -> None: """Process human approval response and continue workflow.""" print(f"[SpamDetector] handle_human_response called with response: {response}") # Get stored detection result - state = await ctx.get_executor_state() or {} - print(f"[SpamDetector] Retrieved state: {state}") - ai_original = state.get("ai_original_classification", False) - confidence_score = state.get("confidence_score", 0.0) - spam_reasons = state.get("spam_reasons", []) + ai_original = original_request.detected_as_spam + confidence_score = original_request.confidence + spam_reasons = original_request.reasons # Parse human decision from the response model human_decision = response.decision.strip().lower() @@ -241,27 +227,21 @@ class SpamDetector(Executor): # Default to AI decision if unclear is_spam = ai_original - # Reconstruct EmailContent from stored primitives - email_content = EmailContent( - original_message=state.get("original_message", ""), - cleaned_message=state.get("cleaned_message", ""), - word_count=state.get("word_count", 0), - has_suspicious_patterns=state.get("has_suspicious_patterns", False) - ) - result = SpamDetectorResponse( - email_content=email_content, + email_content=original_request.full_email_content, is_spam=is_spam, confidence_score=confidence_score, spam_reasons=spam_reasons, human_reviewed=True, human_decision=response.decision, - ai_original_classification=ai_original + ai_original_classification=ai_original, ) - print(f"[SpamDetector] Sending SpamDetectorResponse: is_spam={is_spam}, confidence={confidence_score}, human_reviewed=True") + print( + f"[SpamDetector] Sending SpamDetectorResponse: is_spam={is_spam}, confidence={confidence_score}, human_reviewed=True" + ) await ctx.send_message(result) - print(f"[SpamDetector] Message sent successfully") + print("[SpamDetector] Message sent successfully") class SpamHandler(Executor): @@ -427,7 +407,9 @@ workflow = ( spam_detector, [ Case(condition=lambda x: isinstance(x, SpamDetectorResponse) and x.is_spam, target=spam_handler), - Default(target=legitimate_message_handler), # Default handles non-spam and non-SpamDetectorResponse messages + Default( + target=legitimate_message_handler + ), # Default handles non-spam and non-SpamDetectorResponse messages ], ) .add_edge(spam_handler, final_processor) diff --git a/python/samples/getting_started/workflows/checkpoint/checkpoint_with_human_in_the_loop.py b/python/samples/getting_started/workflows/checkpoint/checkpoint_with_human_in_the_loop.py index 9fb870bf01..7effcdd28e 100644 --- a/python/samples/getting_started/workflows/checkpoint/checkpoint_with_human_in_the_loop.py +++ b/python/samples/getting_started/workflows/checkpoint/checkpoint_with_human_in_the_loop.py @@ -3,6 +3,7 @@ import asyncio from dataclasses import dataclass from pathlib import Path +from typing import Any, override # 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 @@ -116,19 +117,19 @@ class ReviewGateway(Executor): 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. - draft = response.agent_run_response.text or "" - iteration = int((await ctx.get_executor_state() or {}).get("iteration", 0)) + 1 - await ctx.set_executor_state({"iteration": iteration, "last_draft": draft}) + 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=draft, - iteration=iteration, + draft=response.agent_run_response.text, + iteration=self._iteration, ), response_type=str, ) @@ -142,28 +143,33 @@ class ReviewGateway(Executor): ) -> None: # The `original_request` is the request we sent earlier that is now being answered. reply = feedback.strip() - state = await ctx.get_executor_state() or {} - draft = state.get("last_draft") or (original_request.draft or "") - if reply.lower() == "approve": + if len(reply) == 0 or reply.lower() == "approve": # Workflow is completed when the human approves. - await ctx.yield_output(draft) + await ctx.yield_output(original_request.draft) 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_executor_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}" + f"Previous draft:\n{original_request.draft}\n\n" + f"Human guidance: {reply}" ) await ctx.send_message( AgentExecutorRequest(messages=[ChatMessage(Role.USER, text=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.""" @@ -247,10 +253,10 @@ async def run_interactive_session( else: if initial_message: print(f"\nStarting workflow with brief: {initial_message}\n") - event_stream = workflow.run_stream(initial_message) + event_stream = workflow.run_stream(message=initial_message) elif checkpoint_id: print("\nStarting workflow from checkpoint...\n") - event_stream = workflow.run_stream(checkpoint_id) + event_stream = workflow.run_stream(checkpoint_id=checkpoint_id) else: raise ValueError("Either initial_message or checkpoint_id must be provided") diff --git a/python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py b/python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py index cb0c7705c5..0b7e874a1a 100644 --- a/python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py +++ b/python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py @@ -1,322 +1,157 @@ # Copyright (c) Microsoft. All rights reserved. -import asyncio -import os -from pathlib import Path -from typing import TYPE_CHECKING, Any - -from agent_framework import ( - AgentExecutor, - AgentExecutorRequest, - AgentExecutorResponse, - ChatMessage, - Executor, - FileCheckpointStorage, - Role, - WorkflowBuilder, - WorkflowContext, - get_checkpoint_summary, - handler, -) -from agent_framework.azure import AzureOpenAIChatClient -from azure.identity import AzureCliCredential - -if TYPE_CHECKING: - from agent_framework import Workflow - from agent_framework._workflows._checkpoint import WorkflowCheckpoint - """ -Sample: Checkpointing and Resuming a Workflow (with an Agent stage) +Sample: Checkpointing and Resuming a Workflow Purpose: -This sample shows how to enable checkpointing at superstep boundaries, persist both -executor-local state and shared workflow state, and then resume execution from a specific -checkpoint. The workflow demonstrates a simple text-processing pipeline that includes -an LLM-backed AgentExecutor stage. - -Pipeline: -1) UpperCaseExecutor converts input to uppercase and records state. -2) ReverseTextExecutor reverses the string. -3) SubmitToLowerAgent prepares an AgentExecutorRequest for the lowercasing agent. -4) lower_agent (AgentExecutor) converts text to lowercase via Azure OpenAI. -5) FinalizeFromAgent yields the final result. +This sample shows how to enable checkpointing for a long-running workflow +that can be paused and resumed. What you learn: -- How to persist executor state using ctx.get_executor_state and ctx.set_executor_state. -- How to persist shared workflow state using ctx.set_shared_state for cross-executor visibility. -- How to configure FileCheckpointStorage and call with_checkpointing on WorkflowBuilder. -- How to list and inspect checkpoints programmatically. -- How to interactively choose a checkpoint to resume from (instead of always resuming - from the most recent or a hard-coded one) using run_stream. -- How workflows complete by yielding outputs when idle, not via explicit completion events. +- How to configure checkpointing storage (InMemoryCheckpointStorage for testing) +- How to resume a workflow from a checkpoint after interruption +- How to implement executor state management with checkpoint hooks +- How to handle workflow interruptions and automatic recovery + +Pipeline: +This sample shows a workflow that computes factor pairs for numbers up to a given limit: +1) A start executor that receives the upper limit and creates the initial task +2) A worker executor that processes each number to find its factor pairs +3) The worker uses checkpoint hooks to save/restore its internal state Prerequisites: -- Azure AI or Azure OpenAI available for AzureOpenAIChatClient. -- Authentication with azure-identity via AzureCliCredential. Run az login locally. -- Filesystem access for writing JSON checkpoint files in a temp directory. +- Basic understanding of workflow concepts, including executors, edges, events, etc. """ -# Define the temporary directory for storing checkpoints. -# These files allow the workflow to be resumed later. -DIR = os.path.dirname(__file__) -TEMP_DIR = os.path.join(DIR, "tmp", "checkpoints") -os.makedirs(TEMP_DIR, exist_ok=True) +import asyncio +from dataclasses import dataclass +from random import random +from typing import Any, override + +from agent_framework import ( + Executor, + InMemoryCheckpointStorage, + SuperStepCompletedEvent, + WorkflowBuilder, + WorkflowCheckpoint, + WorkflowContext, + WorkflowOutputEvent, + handler, +) -class UpperCaseExecutor(Executor): - """Uppercases the input text and persists both local and shared state.""" +@dataclass +class ComputeTask: + """Task containing the list of numbers remaining to be processed.""" + + remaining_numbers: list[int] + + +class StartExecutor(Executor): + """Initiates the workflow by providing the upper limit for factor pair computation.""" @handler - async def to_upper_case(self, text: str, ctx: WorkflowContext[str]) -> None: - result = text.upper() - print(f"UpperCaseExecutor: '{text}' -> '{result}'") - - # Persist executor-local state so it is captured in checkpoints - # and available after resume for observability or logic. - prev = await ctx.get_executor_state() or {} - count = int(prev.get("count", 0)) + 1 - await ctx.set_executor_state({ - "count": count, - "last_input": text, - "last_output": result, - }) - - # Write to shared_state so downstream executors and any resumed runs can read it. - await ctx.set_shared_state("original_input", text) - await ctx.set_shared_state("upper_output", result) - - # Send transformed text to the next executor. - await ctx.send_message(result) + async def start(self, upper_limit: int, ctx: WorkflowContext[ComputeTask]) -> None: + """Start the workflow with a list of numbers to process.""" + print(f"StartExecutor: Starting factor pair computation up to {upper_limit}") + await ctx.send_message(ComputeTask(remaining_numbers=list(range(1, upper_limit + 1)))) -class SubmitToLowerAgent(Executor): - """Builds an AgentExecutorRequest to send to the lowercasing agent while keeping shared-state visibility.""" +class WorkerExecutor(Executor): + """Processes numbers to compute their factor pairs and manages executor state for checkpointing.""" - def __init__(self, id: str, agent_id: str): + def __init__(self, id: str) -> None: super().__init__(id=id) - self._agent_id = agent_id + self._composite_number_pairs: dict[int, list[tuple[int, int]]] = {} @handler - async def submit(self, text: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None: - # Demonstrate reading shared_state written by UpperCaseExecutor. - # Shared state survives across checkpoints and is visible to all executors. - orig = await ctx.get_shared_state("original_input") - upper = await ctx.get_shared_state("upper_output") - print(f"LowerAgent (shared_state): original_input='{orig}', upper_output='{upper}'") + async def compute( + self, + task: ComputeTask, + ctx: WorkflowContext[ComputeTask, dict[int, list[tuple[int, int]]]], + ) -> None: + """Process the next number in the task, computing its factor pairs.""" + next_number = task.remaining_numbers.pop(0) - # Build a minimal, deterministic prompt for the AgentExecutor. - prompt = f"Convert the following text to lowercase. Return ONLY the transformed text.\n\nText: {text}" + print(f"WorkerExecutor: Computing factor pairs for {next_number}") + pairs: list[tuple[int, int]] = [] + for i in range(1, next_number): + if next_number % i == 0: + pairs.append((i, next_number // i)) + self._composite_number_pairs[next_number] = pairs - # Send to the AgentExecutor. should_respond=True instructs the agent to produce a reply. - await ctx.send_message( - AgentExecutorRequest(messages=[ChatMessage(Role.USER, text=prompt)], should_respond=True), - target_id=self._agent_id, - ) + if not task.remaining_numbers: + # All numbers processed - output the results + await ctx.yield_output(self._composite_number_pairs) + else: + # More numbers to process - continue with remaining task + await ctx.send_message(task) + @override + async def on_checkpoint_save(self) -> dict[str, Any]: + """Save the executor's internal state for checkpointing.""" + return {"composite_number_pairs": self._composite_number_pairs} -class FinalizeFromAgent(Executor): - """Consumes the AgentExecutorResponse and yields the final result.""" - - @handler - async def finalize(self, response: AgentExecutorResponse, ctx: WorkflowContext[Any, str]) -> None: - result = response.agent_run_response.text or "" - - # Persist executor-local state for auditability when inspecting checkpoints. - prev = await ctx.get_executor_state() or {} - count = int(prev.get("count", 0)) + 1 - await ctx.set_executor_state({ - "count": count, - "last_output": result, - "final": True, - }) - - # Yield the final result so external consumers see the final value. - await ctx.yield_output(result) - - -class ReverseTextExecutor(Executor): - """Reverses the input text and persists local state.""" - - @handler - async def reverse_text(self, text: str, ctx: WorkflowContext[str]) -> None: - result = text[::-1] - print(f"ReverseTextExecutor: '{text}' -> '{result}'") - - # Persist executor-local state so checkpoint inspection can reveal progress. - prev = await ctx.get_executor_state() or {} - count = int(prev.get("count", 0)) + 1 - await ctx.set_executor_state({ - "count": count, - "last_input": text, - "last_output": result, - }) - - # Forward the reversed string to the next stage. - await ctx.send_message(result) - - -def create_workflow(checkpoint_storage: FileCheckpointStorage) -> "Workflow": - # Instantiate the pipeline executors. - upper_case_executor = UpperCaseExecutor(id="upper-case") - reverse_text_executor = ReverseTextExecutor(id="reverse-text") - - # Configure the agent stage that lowercases the text. - chat_client = AzureOpenAIChatClient(credential=AzureCliCredential()) - lower_agent = AgentExecutor( - chat_client.create_agent( - instructions=("You transform text to lowercase. Reply with ONLY the transformed text.") - ), - id="lower_agent", - ) - - # Bridge to the agent and terminalization stage. - submit_lower = SubmitToLowerAgent(id="submit_lower", agent_id=lower_agent.id) - finalize = FinalizeFromAgent(id="finalize") - - # Build the workflow with checkpointing enabled. - 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 - .add_edge(submit_lower, lower_agent) # Submit to AgentExecutor - .add_edge(lower_agent, finalize) # Agent output -> Finalize - .set_start_executor(upper_case_executor) # Entry point - .with_checkpointing(checkpoint_storage=checkpoint_storage) # Enable persistence - .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 = get_checkpoint_summary(cp) - msg_count = sum(len(v) for v in cp.messages.values()) - state_keys = sorted(summary.executor_ids) - 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) + @override + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + """Restore the executor's internal state from a checkpoint.""" + self._composite_number_pairs = state.get("composite_number_pairs", {}) 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"): # noqa: ASYNC240 - file.unlink() + # Create workflow executors + start_executor = StartExecutor(id="start") + worker_executor = WorkerExecutor(id="worker") - # Backing store for checkpoints written by with_checkpointing. - checkpoint_storage = FileCheckpointStorage(storage_path=TEMP_DIR) + # Build workflow with checkpointing enabled + workflow_builder = ( + WorkflowBuilder() + .set_start_executor(start_executor) + .add_edge(start_executor, worker_executor) + .add_edge(worker_executor, worker_executor) # Self-loop for iterative processing + ) + checkpoint_storage = InMemoryCheckpointStorage() + workflow_builder = workflow_builder.with_checkpointing(checkpoint_storage=checkpoint_storage) - workflow = create_workflow(checkpoint_storage=checkpoint_storage) + # Run workflow with automatic checkpoint recovery + latest_checkpoint: WorkflowCheckpoint | None = None + while True: + workflow = workflow_builder.build() - # 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"): - print(f"Event: {event}") + # Start from checkpoint or fresh execution + print(f"\n** Workflow {workflow.id} started **") + event_stream = ( + workflow.run_stream(message=10) + if latest_checkpoint is None + else workflow.run_stream(checkpoint_id=latest_checkpoint.checkpoint_id) + ) - # Inspect checkpoints written during the run. - all_checkpoints = await checkpoint_storage.list_checkpoints() - if not all_checkpoints: - print("No checkpoints found!") - return - - # All checkpoints created by this run share the same workflow_id. - workflow_id = all_checkpoints[0].workflow_id - - _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 = get_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()) - line += f" messages={msg_count}" - print(line) - - user_input = input( # noqa: ASYNC250 - "\nEnter checkpoint index (or paste checkpoint id) to resume from, or press Enter to skip resume: " - ).strip() - - if not user_input: - print("No checkpoint selected. Exiting without resuming.") - return - - chosen_cp_id: str | None = None - - # Try as index first - if user_input.isdigit(): - idx = int(user_input) - if 0 <= idx < len(sorted_cps): - chosen_cp_id = sorted_cps[idx].checkpoint_id - # Fall back to direct id match - if chosen_cp_id is None: - for cp in sorted_cps: - if cp.checkpoint_id.startswith(user_input): # allow prefix match for convenience - chosen_cp_id = cp.checkpoint_id + output: str | None = None + async for event in event_stream: + if isinstance(event, WorkflowOutputEvent): + output = event.data + break + if isinstance(event, SuperStepCompletedEvent) and random() < 0.5: + # Randomly simulate system interruptions + # The `SuperStepCompletedEvent` ensures we only interrupt after + # the current super-step is fully complete and checkpointed. + # If we interrupt mid-step, the workflow may resume from an earlier point. + print("\n** Simulating workflow interruption. Stopping execution. **") break - if chosen_cp_id is None: - print("Input did not match any checkpoint. Exiting without resuming.") - return + # Find the latest checkpoint to resume from + all_checkpoints = await checkpoint_storage.list_checkpoints() + if not all_checkpoints: + raise RuntimeError("No checkpoints available to resume from.") + latest_checkpoint = all_checkpoints[-1] + print( + f"Checkpoint {latest_checkpoint.checkpoint_id}: " + f"(iter={latest_checkpoint.iteration_count}, messages={latest_checkpoint.messages})" + ) - # 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 = create_workflow(checkpoint_storage=checkpoint_storage) - - print(f"\nResuming from checkpoint: {chosen_cp_id}") - async for event in new_workflow.run_stream(checkpoint_id=chosen_cp_id, checkpoint_storage=checkpoint_storage): - print(f"Resumed Event: {event}") - - """ - Sample Output: - - Running workflow with initial message... - UpperCaseExecutor: 'hello world' -> 'HELLO WORLD' - Event: ExecutorInvokeEvent(executor_id=upper_case_executor) - Event: ExecutorCompletedEvent(executor_id=upper_case_executor) - ReverseTextExecutor: 'HELLO WORLD' -> 'DLROW OLLEH' - 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: ExecutorInvokeEvent(executor_id=submit_lower) - Event: ExecutorInvokeEvent(executor_id=lower_agent) - Event: ExecutorInvokeEvent(executor_id=finalize) - - Checkpoint summary: - - dfc63e72-8e8d-454f-9b6d-0d740b9062e6 | label='after_initial_execution' | iter=0 | messages=1 | states=['upper_case_executor'] | shared_state: original_input='hello world', upper_output='HELLO WORLD' - - a78c345a-e5d9-45ba-82c0-cb725452d91b | label='superstep_1' | iter=1 | messages=1 | states=['reverse_text_executor', 'upper_case_executor'] | shared_state: original_input='hello world', upper_output='HELLO WORLD' - - 637c1dbd-a525-4404-9583-da03980537a2 | label='superstep_2' | iter=2 | messages=0 | states=['finalize', 'lower_agent', 'reverse_text_executor', 'submit_lower', 'upper_case_executor'] | shared_state: original_input='hello world', upper_output='HELLO WORLD' - - Available checkpoints to resume from: - [0] id=dfc63e72-... iter=0 messages=1 label='after_initial_execution' - [1] id=a78c345a-... iter=1 messages=1 label='superstep_1' - [2] id=637c1dbd-... iter=2 messages=0 label='superstep_2' - - Enter checkpoint index (or paste checkpoint id) to resume from, or press Enter to skip resume: 1 - - Resuming from checkpoint: a78c345a-e5d9-45ba-82c0-cb725452d91b - LowerAgent (shared_state): original_input='hello world', upper_output='HELLO WORLD' - Resumed Event: ExecutorInvokeEvent(executor_id=submit_lower) - Resumed Event: ExecutorInvokeEvent(executor_id=lower_agent) - Resumed Event: ExecutorInvokeEvent(executor_id=finalize) - """ # noqa: E501 + if output is not None: + print(f"\nWorkflow completed successfully with output: {output}") + break if __name__ == "__main__": diff --git a/python/samples/getting_started/workflows/checkpoint/sub_workflow_checkpoint.py b/python/samples/getting_started/workflows/checkpoint/sub_workflow_checkpoint.py index 62c88bf49f..2d90db8c2c 100644 --- a/python/samples/getting_started/workflows/checkpoint/sub_workflow_checkpoint.py +++ b/python/samples/getting_started/workflows/checkpoint/sub_workflow_checkpoint.py @@ -7,6 +7,7 @@ import uuid from dataclasses import dataclass, field, replace from datetime import datetime, timedelta from pathlib import Path +from typing import Any, override from agent_framework import ( Executor, @@ -205,6 +206,8 @@ class LaunchCoordinator(Executor): def __init__(self) -> None: super().__init__(id="launch_coordinator") + # Track pending requests to match responses + self._pending_requests: dict[str, SubWorkflowRequestMessage] = {} @handler async def kick_off(self, topic: str, ctx: WorkflowContext[DraftTask]) -> None: @@ -244,11 +247,9 @@ class LaunchCoordinator(Executor): if not isinstance(request.source_event.data, ReviewRequest): raise TypeError(f"Expected 'ReviewRequest', got {type(request.source_event.data)}") - # Record the request to response matching + # Record the request for response matching review_request = request.source_event.data - executor_state = await ctx.get_executor_state() or {} - executor_state[review_request.id] = request - await ctx.set_executor_state(executor_state) + self._pending_requests[review_request.id] = request # Send the request without modification await ctx.request_info(request_data=review_request, response_type=str) @@ -265,17 +266,25 @@ class LaunchCoordinator(Executor): Note that the response must be sent back using SubWorkflowResponseMessage to route the response back to the sub-workflow. """ - executor_state = await ctx.get_executor_state() or {} - request_message = executor_state.pop(original_request.id, None) - - # Save the executor state back to the context - await ctx.set_executor_state(executor_state) + request_message = self._pending_requests.pop(original_request.id, None) if request_message is None: raise ValueError("No matching pending request found for the resource response") await ctx.send_message(request_message.create_response(response)) + @override + async def on_checkpoint_save(self) -> dict[str, Any]: + """Capture any additional state needed for checkpointing.""" + return { + "pending_requests": self._pending_requests, + } + + @override + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + """Restore any additional state needed from checkpointing.""" + self._pending_requests = state.get("pending_requests", {}) + # --------------------------------------------------------------------------- # Workflow construction helpers @@ -356,9 +365,7 @@ async def main() -> None: workflow2 = build_parent_workflow(storage) request_info_event: RequestInfoEvent | None = None - async for event in workflow2.run_stream( - resume_checkpoint.checkpoint_id, - ): + async for event in workflow2.run_stream(checkpoint_id=resume_checkpoint.checkpoint_id): if isinstance(event, RequestInfoEvent): request_info_event = event