From fc12ab9feddf4d807423b147d13c6b64ab65ca31 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 13 Oct 2025 12:18:17 -0700 Subject: [PATCH] Python: Refactor `RequestInfoExecutor` (#1403) * Refactor RequestInfoExecutor * Update AI script * Fix formatting * Address comments * fix unit test --- .../agent_framework/_workflows/__init__.py | 5 + .../agent_framework/_workflows/__init__.pyi | 5 + .../_workflows/_checkpoint_summary.py | 191 +++++ .../_workflows/_request_info_executor.py | 650 +++++------------- .../agent_framework/_workflows/_runner.py | 23 + .../_workflows/_workflow_context.py | 7 +- .../test_request_info_executor_rehydrate.py | 89 +-- .../tests/workflow/test_workflow_states.py | 38 - python/samples/_run_all_samples.py | 235 +++++++ .../workflow_as_agent_human_in_the_loop.py | 24 +- .../checkpoint_with_human_in_the_loop.py | 11 +- .../checkpoint/checkpoint_with_resume.py | 6 +- 12 files changed, 705 insertions(+), 579 deletions(-) create mode 100644 python/packages/core/agent_framework/_workflows/_checkpoint_summary.py create mode 100644 python/samples/_run_all_samples.py diff --git a/python/packages/core/agent_framework/_workflows/__init__.py b/python/packages/core/agent_framework/_workflows/__init__.py index 5af4049128..94950e1948 100644 --- a/python/packages/core/agent_framework/_workflows/__init__.py +++ b/python/packages/core/agent_framework/_workflows/__init__.py @@ -12,6 +12,7 @@ from ._checkpoint import ( InMemoryCheckpointStorage, WorkflowCheckpoint, ) +from ._checkpoint_summary import WorkflowCheckpointSummary, get_checkpoint_summary from ._concurrent import ConcurrentBuilder from ._const import ( DEFAULT_MAX_ITERATIONS, @@ -74,6 +75,7 @@ from ._magentic import ( StandardMagenticManager, ) from ._request_info_executor import ( + PendingRequestDetails, RequestInfoExecutor, RequestInfoMessage, RequestResponse, @@ -146,6 +148,7 @@ __all__ = [ "MagenticResponseMessage", "MagenticStartMessage", "Message", + "PendingRequestDetails", "RequestInfoEvent", "RequestInfoExecutor", "RequestInfoMessage", @@ -165,6 +168,7 @@ __all__ = [ "WorkflowAgent", "WorkflowBuilder", "WorkflowCheckpoint", + "WorkflowCheckpointSummary", "WorkflowContext", "WorkflowErrorDetails", "WorkflowEvent", @@ -181,6 +185,7 @@ __all__ = [ "WorkflowViz", "create_edge_runner", "executor", + "get_checkpoint_summary", "handler", "validate_workflow_graph", ] diff --git a/python/packages/core/agent_framework/_workflows/__init__.pyi b/python/packages/core/agent_framework/_workflows/__init__.pyi index db8f87eb4e..d98829c56d 100644 --- a/python/packages/core/agent_framework/_workflows/__init__.pyi +++ b/python/packages/core/agent_framework/_workflows/__init__.pyi @@ -12,6 +12,7 @@ from ._checkpoint import ( InMemoryCheckpointStorage, WorkflowCheckpoint, ) +from ._checkpoint_summary import WorkflowCheckpointSummary, get_checkpoint_summary from ._concurrent import ConcurrentBuilder from ._const import DEFAULT_MAX_ITERATIONS from ._edge import ( @@ -72,6 +73,7 @@ from ._magentic import ( StandardMagenticManager, ) from ._request_info_executor import ( + PendingRequestDetails, RequestInfoExecutor, RequestInfoMessage, RequestResponse, @@ -144,6 +146,7 @@ __all__ = [ "MagenticResponseMessage", "MagenticStartMessage", "Message", + "PendingRequestDetails", "RequestInfoEvent", "RequestInfoExecutor", "RequestInfoMessage", @@ -163,6 +166,7 @@ __all__ = [ "WorkflowAgent", "WorkflowBuilder", "WorkflowCheckpoint", + "WorkflowCheckpointSummary", "WorkflowContext", "WorkflowErrorDetails", "WorkflowEvent", @@ -179,6 +183,7 @@ __all__ = [ "WorkflowViz", "create_edge_runner", "executor", + "get_checkpoint_summary", "handler", "validate_workflow_graph", ] diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_summary.py b/python/packages/core/agent_framework/_workflows/_checkpoint_summary.py new file mode 100644 index 0000000000..e42e05dd91 --- /dev/null +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_summary.py @@ -0,0 +1,191 @@ +# Copyright (c) Microsoft. All rights reserved. + +import logging +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from textwrap import shorten +from typing import Any + +from ._checkpoint import WorkflowCheckpoint +from ._request_info_executor import PendingRequestDetails, RequestInfoMessage, RequestResponse +from ._runner_context import _decode_checkpoint_value # type: ignore + +logger = logging.getLogger(__name__) + + +@dataclass +class WorkflowCheckpointSummary: + """Human-readable summary of a workflow checkpoint.""" + + checkpoint_id: str + iteration_count: int + targets: list[str] + executor_ids: list[str] + status: str + draft_preview: str | None + pending_requests: list[PendingRequestDetails] + + +def get_checkpoint_summary( + checkpoint: WorkflowCheckpoint, + *, + request_executor_ids: Iterable[str] | None = None, + preview_width: int = 70, +) -> WorkflowCheckpointSummary: + targets = sorted(checkpoint.messages.keys()) + executor_ids = sorted(checkpoint.executor_states.keys()) + pending = _pending_requests_from_checkpoint(checkpoint, request_executor_ids=request_executor_ids) + + draft_preview: str | None = None + for entry in pending: + if entry.draft: + draft_preview = shorten(entry.draft, width=preview_width, placeholder="…") + break + + status = "idle" + if pending: + status = "awaiting request response" + elif not checkpoint.messages and "finalise" in executor_ids: + status = "completed" + elif checkpoint.messages: + status = "awaiting next superstep" + elif request_executor_ids is not None and any(tid in targets for tid in request_executor_ids): + status = "awaiting request delivery" + + return WorkflowCheckpointSummary( + checkpoint_id=checkpoint.checkpoint_id, + iteration_count=checkpoint.iteration_count, + targets=targets, + executor_ids=executor_ids, + status=status, + draft_preview=draft_preview, + pending_requests=pending, + ) + + +def _pending_requests_from_checkpoint( + checkpoint: WorkflowCheckpoint, + *, + request_executor_ids: Iterable[str] | None = None, +) -> list[PendingRequestDetails]: + executor_filter: set[str] | None = None + if request_executor_ids is not None: + executor_filter = {str(value) for value in request_executor_ids} + + pending: dict[str, PendingRequestDetails] = {} + + for state in checkpoint.executor_states.values(): + if not isinstance(state, Mapping): + continue + inner = state.get("pending_requests") + if isinstance(inner, Mapping): + for request_id, snapshot in inner.items(): # type: ignore[attr-defined] + _merge_snapshot(pending, str(request_id), snapshot) # type: ignore[arg-type] + + for source_id, message_list in checkpoint.messages.items(): + if executor_filter is not None and source_id not in executor_filter: + continue + if not isinstance(message_list, list): + continue + for message in message_list: + if not isinstance(message, Mapping): + continue + payload = _decode_checkpoint_value(message.get("data")) + _merge_message_payload(pending, payload, message) + + return list(pending.values()) + + +def _merge_snapshot(pending: dict[str, PendingRequestDetails], request_id: str, snapshot: Any) -> None: + if not request_id or not isinstance(snapshot, Mapping): + return + + details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id)) + + _apply_update( + details, + prompt=snapshot.get("prompt"), # type: ignore[attr-defined] + draft=snapshot.get("draft"), # type: ignore[attr-defined] + iteration=snapshot.get("iteration"), # type: ignore[attr-defined] + source_executor_id=snapshot.get("source_executor_id"), # type: ignore[attr-defined] + ) + + extra = snapshot.get("details") # type: ignore[attr-defined] + if isinstance(extra, Mapping): + _apply_update( + details, + prompt=extra.get("prompt"), # type: ignore[attr-defined] + draft=extra.get("draft"), # type: ignore[attr-defined] + iteration=extra.get("iteration"), # type: ignore[attr-defined] + ) + + +def _merge_message_payload( + pending: dict[str, PendingRequestDetails], + payload: Any, + raw_message: Mapping[str, Any], +) -> None: + if isinstance(payload, RequestResponse): + request_id = payload.request_id or _get_field(payload.original_request, "request_id") # type: ignore[arg-type] + if not request_id: + return + details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id)) + _apply_update( + details, + prompt=_get_field(payload.original_request, "prompt"), # type: ignore[arg-type] + draft=_get_field(payload.original_request, "draft"), # type: ignore[arg-type] + iteration=_get_field(payload.original_request, "iteration"), # type: ignore[arg-type] + source_executor_id=raw_message.get("source_id"), + original_request=payload.original_request, # type: ignore[arg-type] + ) + elif isinstance(payload, RequestInfoMessage): + request_id = getattr(payload, "request_id", None) + if not request_id: + return + details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id)) + _apply_update( + details, + prompt=getattr(payload, "prompt", None), + draft=getattr(payload, "draft", None), + iteration=getattr(payload, "iteration", None), + source_executor_id=raw_message.get("source_id"), + original_request=payload, + ) + + +def _apply_update( + details: PendingRequestDetails, + *, + prompt: Any = None, + draft: Any = None, + iteration: Any = None, + source_executor_id: Any = None, + original_request: Any = None, +) -> None: + if prompt and not details.prompt: + details.prompt = str(prompt) + if draft and not details.draft: + details.draft = str(draft) + if iteration is not None and details.iteration is None: + coerced = _coerce_int(iteration) + if coerced is not None: + details.iteration = coerced + if source_executor_id and not details.source_executor_id: + details.source_executor_id = str(source_executor_id) + if original_request is not None and details.original_request is None: + details.original_request = original_request + + +def _get_field(obj: Any, key: str) -> Any: + if obj is None: + return None + if isinstance(obj, Mapping): + return obj.get(key) # type: ignore[attr-defined,return-value] + return getattr(obj, key, None) + + +def _coerce_int(value: Any) -> int | None: + try: + return int(value) + except (TypeError, ValueError): + return None diff --git a/python/packages/core/agent_framework/_workflows/_request_info_executor.py b/python/packages/core/agent_framework/_workflows/_request_info_executor.py index ef6eb3fc22..6b3c710f62 100644 --- a/python/packages/core/agent_framework/_workflows/_request_info_executor.py +++ b/python/packages/core/agent_framework/_workflows/_request_info_executor.py @@ -5,17 +5,14 @@ import importlib import json import logging import uuid -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass, field, fields, is_dataclass -from textwrap import shorten from typing import Any, ClassVar, Generic, TypeVar, cast -from ._checkpoint import WorkflowCheckpoint from ._events import ( RequestInfoEvent, # type: ignore[reportPrivateUsage] ) from ._executor import Executor, handler -from ._runner_context import _decode_checkpoint_value # type: ignore from ._workflow_context import WorkflowContext logger = logging.getLogger(__name__) @@ -34,16 +31,17 @@ class PendingRequestDetails: @dataclass -class WorkflowCheckpointSummary: - """Human-readable summary of a workflow checkpoint.""" +class PendingRequestSnapshot: + """Snapshot of a pending request for internal tracking. - checkpoint_id: str - iteration_count: int - targets: list[str] - executor_states: list[str] - status: str - draft_preview: str | None - pending_requests: list[PendingRequestDetails] + This snapshot should be JSON-serializable and contain enough + information to reconstruct the original request if needed. + """ + + request_id: str + source_executor_id: str + request_type: str + request_as_json_safe_dict: dict[str, Any] @dataclass @@ -108,8 +106,10 @@ class RequestInfoExecutor(Executor): super().__init__(id=id) self._request_events: dict[str, RequestInfoEvent] = {} + # region Public Methods + @handler - async def run(self, message: RequestInfoMessage, ctx: WorkflowContext) -> None: + async def handle_request(self, message: RequestInfoMessage, ctx: WorkflowContext) -> None: """Run the RequestInfoExecutor with the given message.""" # Use source_executor_id from message if available, otherwise fall back to context source_executor_id = message.source_executor_id or ctx.get_source_executor_id() @@ -121,7 +121,7 @@ class RequestInfoExecutor(Executor): request_data=message, ) self._request_events[message.request_id] = event - await self._record_pending_request_snapshot(message, source_executor_id, ctx) + await self._record_pending_request(message, source_executor_id, ctx) await ctx.add_event(event) async def handle_response( @@ -139,7 +139,7 @@ class RequestInfoExecutor(Executor): """ event = self._request_events.get(request_id) if event is None: - event = await self._rehydrate_request_event(request_id, ctx) + event = await self._rehydrate_request_event(request_id, cast(WorkflowContext, ctx)) if event is None: raise ValueError(f"No request found with ID: {request_id}") @@ -151,72 +151,20 @@ class RequestInfoExecutor(Executor): correlated_response = RequestResponse(data=response_data, original_request=event.data, request_id=request_id) await ctx.send_message(correlated_response, target_id=event.source_executor_id) - await self._clear_pending_request_snapshot(request_id, ctx) - - async def _record_pending_request_snapshot( - self, - request: RequestInfoMessage, - source_executor_id: str, - ctx: WorkflowContext[Any], - ) -> None: - snapshot = self._build_request_snapshot(request, source_executor_id) - - pending = await self._load_pending_request_state(ctx) - pending[request.request_id] = snapshot - await self._persist_pending_request_state(pending, ctx) - await self._write_executor_state(ctx, pending) - - async def _clear_pending_request_snapshot(self, request_id: str, ctx: WorkflowContext[Any]) -> None: - pending = await self._load_pending_request_state(ctx) - if request_id in pending: - pending.pop(request_id, None) - await self._persist_pending_request_state(pending, ctx) - await self._write_executor_state(ctx, pending) - - async def _load_pending_request_state(self, ctx: WorkflowContext[Any]) -> dict[str, Any]: - try: - existing = await ctx.get_shared_state(self._PENDING_SHARED_STATE_KEY) - except KeyError: - return {} - except Exception as exc: # pragma: no cover - transport specific - logger.warning(f"RequestInfoExecutor {self.id} failed to read pending request state: {exc}") - return {} - - if not isinstance(existing, dict): - if existing not in (None, {}): - logger.warning( - f"RequestInfoExecutor {self.id} encountered non-dict pending state " - f"({type(existing).__name__}); resetting." - ) - return {} - - return dict(existing) # type: ignore[arg-type] - - async def _persist_pending_request_state(self, pending: dict[str, Any], ctx: WorkflowContext[Any]) -> None: - await self._safe_set_shared_state(ctx, pending) - await self._safe_set_runner_state(ctx, pending) - - async def _safe_set_shared_state(self, ctx: WorkflowContext[Any], pending: dict[str, Any]) -> None: - try: - await ctx.set_shared_state(self._PENDING_SHARED_STATE_KEY, pending) - except Exception as exc: # pragma: no cover - transport specific - logger.warning(f"RequestInfoExecutor {self.id} failed to update shared pending state: {exc}") - - async def _safe_set_runner_state(self, ctx: WorkflowContext[Any], pending: dict[str, Any]) -> None: - try: - await ctx.set_state({"pending_requests": pending}) - except Exception as exc: # pragma: no cover - transport specific - logger.warning(f"RequestInfoExecutor {self.id} failed to update runner state with pending requests: {exc}") + await self._erase_pending_request(request_id, cast(WorkflowContext, ctx)) def snapshot_state(self) -> dict[str, Any]: """Serialize pending requests so checkpoint restoration can resume seamlessly.""" - def _encode_event(event: RequestInfoEvent) -> dict[str, Any]: - request_data = event.data - payload: dict[str, Any] - data_cls = request_data.__class__ if request_data is not None else type(None) + def _encode_event(event: RequestInfoEvent) -> dict[str, Any] | None: + if event.data is None or not isinstance(event.data, RequestInfoMessage): + logger.warning( + f"RequestInfoExecutor {self.id} encountered invalid event data for request ID {event.request_id}: " + f"{type(event.data).__name__}. This request will be skipped in the checkpoint." + ) + return None - payload = self._encode_request_payload(request_data, data_cls) + payload = self._encode_request_payload(event.data, event.data.__class__) return { "source_executor_id": event.source_executor_id, @@ -225,20 +173,125 @@ class RequestInfoExecutor(Executor): } return { - "request_events": {rid: _encode_event(event) for rid, event in self._request_events.items()}, + "request_events": { + rid: encoded + for rid, event in self._request_events.items() + if (encoded := _encode_event(event)) is not None + }, } - def _encode_request_payload(self, request_data: RequestInfoMessage | None, data_cls: type[Any]) -> dict[str, Any]: - if request_data is None or isinstance(request_data, (str, int, float, bool)): - return { - "kind": "raw", - "type": f"{data_cls.__module__}:{data_cls.__qualname__}", - "value": request_data, - } + def restore_state(self, state: dict[str, Any]) -> None: + """Restore pending request bookkeeping from checkpoint state.""" + self._request_events.clear() + stored_events = state.get("request_events", {}) + for request_id, payload in stored_events.items(): + request_type_qual = payload.get("request_type", "") + try: + request_type = _import_qualname(request_type_qual) + except Exception as exc: # pragma: no cover - defensive fallback + logger.debug( + "RequestInfoExecutor %s failed to import %s during restore: %s", + self.id, + request_type_qual, + exc, + ) + request_type = RequestInfoMessage + request_data_meta = payload.get("request_data", {}) + request_data = self._decode_request_data(request_data_meta) + event = RequestInfoEvent( + request_id=request_id, + source_executor_id=payload.get("source_executor_id", ""), + request_type=request_type, + request_data=request_data, + ) + self._request_events[request_id] = event + + async def has_pending_request(self, request_id: str, ctx: WorkflowContext) -> bool: + """Check if there is a pending request with the given ID. + + Args: + request_id: The ID of the request to check. + ctx: The workflow context for accessing state if needed. + + Returns: True if the request is pending, False otherwise. + """ + if request_id in self._request_events: + return True + + pending_requests = await self._retrieve_existing_pending_requests(ctx) + return request_id in pending_requests + + # endregion: Public Methods + + # region: Internal Methods + + async def _record_pending_request( + self, + message: RequestInfoMessage, + source_executor_id: str, + ctx: WorkflowContext, + ) -> None: + """Record a pending request to the executor's state for checkpointing purposes.""" + pending_request_snapshot = self._build_pending_request_snapshot(message, source_executor_id) + + existing_pending_requests = await self._retrieve_existing_pending_requests(ctx) + existing_pending_requests[message.request_id] = pending_request_snapshot + + await self._persist_to_executor_state(existing_pending_requests, ctx) + + async def _erase_pending_request(self, request_id: str, ctx: WorkflowContext) -> None: + """Erase a pending request from the executor's state after it has been handled for checkpointing purposes.""" + existing_pending_requests = await self._retrieve_existing_pending_requests(ctx) + if request_id in existing_pending_requests: + existing_pending_requests.pop(request_id) + await self._persist_to_executor_state(existing_pending_requests, ctx) + + async def _retrieve_existing_pending_requests(self, ctx: WorkflowContext) -> dict[str, PendingRequestSnapshot]: + """Retrieve existing pending requests from executor state.""" + executor_state = await ctx.get_state() + if executor_state is None: + return {} + + stored_requests = executor_state.get(self._PENDING_SHARED_STATE_KEY, {}) + if not isinstance(stored_requests, dict): + raise TypeError(f"Unexpected type for pending requests: {type(stored_requests).__name__}") + + # Validate contents + for key, value in stored_requests.items(): # type: ignore + if not isinstance(key, str) or not isinstance(value, PendingRequestSnapshot): + raise TypeError( + "Invalid pending request entry in executor state. " + "Key must be `str` and value must be `PendingRequestSnapshot`." + ) + + return stored_requests # type: ignore + + async def _persist_to_executor_state( + self, pending: dict[str, PendingRequestSnapshot], ctx: WorkflowContext + ) -> None: + """Persist the current pending requests to the executor's state.""" + executor_state = await ctx.get_state() or {} + executor_state[self._PENDING_SHARED_STATE_KEY] = pending + await ctx.set_state(executor_state) + + def _build_pending_request_snapshot( + self, request: RequestInfoMessage, source_executor_id: str + ) -> PendingRequestSnapshot: + """Build a snapshot of the pending request for checkpointing.""" + request_as_json_safe_dict = self._convert_request_to_json_safe_dict(request) + + return PendingRequestSnapshot( + request_id=request.request_id, + source_executor_id=source_executor_id, + request_type=f"{type(request).__module__}:{type(request).__name__}", + request_as_json_safe_dict=request_as_json_safe_dict, + ) + + def _encode_request_payload(self, request_data: RequestInfoMessage, data_cls: type[Any]) -> dict[str, Any]: if is_dataclass(request_data) and not isinstance(request_data, type): dataclass_instance = cast(Any, request_data) - safe_value = self._make_json_safe(asdict(dataclass_instance)) + safe_value = _make_json_safe(asdict(dataclass_instance)) return { "kind": "dataclass", "type": f"{data_cls.__module__}:{data_cls.__qualname__}", @@ -251,7 +304,7 @@ class RequestInfoExecutor(Executor): dumped = to_dict_fn() except TypeError: dumped = to_dict_fn() - safe_value = self._make_json_safe(dumped) + safe_value = _make_json_safe(dumped) return { "kind": "dict", "type": f"{data_cls.__module__}:{data_cls.__qualname__}", @@ -278,76 +331,26 @@ class RequestInfoExecutor(Executor): converted = json.loads(decoded) except Exception: converted = decoded - safe_value = self._make_json_safe(converted) + safe_value = _make_json_safe(converted) return { "kind": "dict" if isinstance(converted, dict) else "json", "type": f"{data_cls.__module__}:{data_cls.__qualname__}", "value": safe_value, } - details = self._serialise_request_details(request_data) - if details is not None: - safe_value = self._make_json_safe(details) - return { - "kind": "raw", - "type": f"{data_cls.__module__}:{data_cls.__qualname__}", - "value": safe_value, - } - - safe_value = self._make_json_safe(request_data) return { "kind": "raw", "type": f"{data_cls.__module__}:{data_cls.__qualname__}", - "value": safe_value, + "value": self._convert_request_to_json_safe_dict(request_data), } - def restore_state(self, state: dict[str, Any]) -> None: - """Restore pending request bookkeeping from checkpoint state.""" - self._request_events.clear() - stored_events = state.get("request_events", {}) - - for request_id, payload in stored_events.items(): - request_type_qual = payload.get("request_type", "") - try: - request_type = self._import_qualname(request_type_qual) - except Exception as exc: # pragma: no cover - defensive fallback - logger.debug( - "RequestInfoExecutor %s failed to import %s during restore: %s", - self.id, - request_type_qual, - exc, - ) - request_type = RequestInfoMessage - request_data_meta = payload.get("request_data", {}) - request_data = self._decode_request_data(request_data_meta) - event = RequestInfoEvent( - request_id=request_id, - source_executor_id=payload.get("source_executor_id", ""), - request_type=request_type, - request_data=request_data, - ) - self._request_events[request_id] = event - - @staticmethod - def _import_qualname(qualname: str) -> type[Any]: - module_name, _, type_name = qualname.partition(":") - if not module_name or not type_name: - raise ValueError(f"Invalid qualified name: {qualname}") - module = importlib.import_module(module_name) - attr: Any = module - for part in type_name.split("."): - attr = getattr(attr, part) - if not isinstance(attr, type): - raise TypeError(f"Resolved object is not a type: {qualname}") - return attr - def _decode_request_data(self, metadata: dict[str, Any]) -> RequestInfoMessage: kind = metadata.get("kind") type_name = metadata.get("type", "") value: Any = metadata.get("value", {}) if type_name: try: - imported = self._import_qualname(type_name) + imported = _import_qualname(type_name) except Exception as exc: # pragma: no cover - defensive fallback logger.debug( "RequestInfoExecutor %s failed to import %s during decode: %s", @@ -396,114 +399,22 @@ class RequestInfoExecutor(Executor): return target_cls() return RequestInfoMessage() - async def _write_executor_state(self, ctx: WorkflowContext[Any], pending: dict[str, Any]) -> None: - state = self.snapshot_state() - state["pending_requests"] = pending + def _convert_request_to_json_safe_dict(self, request: RequestInfoMessage) -> dict[str, Any]: try: - await ctx.set_state(state) - except Exception as exc: # pragma: no cover - transport specific - logger.warning(f"RequestInfoExecutor {self.id} failed to persist executor state: {exc}") - - def _build_request_snapshot( - self, - request: RequestInfoMessage, - source_executor_id: str, - ) -> dict[str, Any]: - snapshot: dict[str, Any] = { - "request_id": request.request_id, - "source_executor_id": source_executor_id, - "request_type": f"{type(request).__module__}:{type(request).__name__}", - "summary": repr(request), - } - - details = self._serialise_request_details(request) - if details: - snapshot["details"] = details - for key in ("prompt", "draft", "iteration"): - if key in details and key not in snapshot: - snapshot[key] = details[key] - - return snapshot - - def _serialise_request_details(self, request: RequestInfoMessage) -> dict[str, Any] | None: - if is_dataclass(request): - data = self._make_json_safe(asdict(request)) + data = _make_json_safe(asdict(request)) if isinstance(data, dict): return cast(dict[str, Any], data) - return None + raise ValueError(f"Failed to convert {type(request).__name__} to dict") + except Exception as exc: + logger.error(f"RequestInfoExecutor {self.id} failed to serialize request: {exc}") + raise RuntimeError( + f"Failed to serialize request `{type(request).__name__}`: {exc}\n" + "Make sure request is a dataclass and derive from `RequestInfoMessage`." + ) from exc - to_dict = getattr(request, "to_dict", None) - if callable(to_dict): - try: - dump = self._make_json_safe(to_dict()) - except TypeError: - dump = self._make_json_safe(to_dict()) - if isinstance(dump, dict): - return cast(dict[str, Any], dump) - return None - - to_json = getattr(request, "to_json", None) - if callable(to_json): - try: - raw = to_json() - except TypeError: - raw = to_json() - converted = raw - if isinstance(raw, (str, bytes, bytearray)): - decoded: str | bytes | bytearray - if isinstance(raw, (bytes, bytearray)): - try: - decoded = raw.decode() - except Exception: - decoded = raw - else: - decoded = raw - try: - converted = json.loads(decoded) - except Exception: - converted = decoded - dump = self._make_json_safe(converted) - if isinstance(dump, dict): - return cast(dict[str, Any], dump) - return None - - attrs = getattr(request, "__dict__", None) - if isinstance(attrs, dict): - cleaned = self._make_json_safe(attrs) - if isinstance(cleaned, dict): - return cast(dict[str, Any], cleaned) - - return None - - def _make_json_safe(self, value: Any) -> Any: - if value is None or isinstance(value, (str, int, float, bool)): - return value - if isinstance(value, Mapping): - safe_dict: dict[str, Any] = {} - for key, val in value.items(): # type: ignore[attr-defined] - safe_dict[str(key)] = self._make_json_safe(val) # type: ignore[arg-type] - return safe_dict - if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): - return [self._make_json_safe(item) for item in value] # type: ignore[misc] - return repr(value) - - async def has_pending_request(self, request_id: str, ctx: WorkflowContext[Any]) -> bool: - if request_id in self._request_events: - return True - snapshot = await self._get_pending_request_snapshot(request_id, ctx) - return snapshot is not None - - async def _rehydrate_request_event( - self, - request_id: str, - ctx: WorkflowContext[Any], - ) -> RequestInfoEvent | None: - snapshot = await self._get_pending_request_snapshot(request_id, ctx) - if snapshot is None: - return None - - source_executor_id = snapshot.get("source_executor_id") - if not isinstance(source_executor_id, str) or not source_executor_id: + async def _rehydrate_request_event(self, request_id: str, ctx: WorkflowContext) -> RequestInfoEvent | None: + pending_requests = await self._retrieve_existing_pending_requests(ctx) + if (snapshot := pending_requests.get(request_id)) is None: return None request = self._construct_request_from_snapshot(snapshot) @@ -512,57 +423,18 @@ class RequestInfoExecutor(Executor): event = RequestInfoEvent( request_id=request_id, - source_executor_id=source_executor_id, + source_executor_id=snapshot.source_executor_id, request_type=type(request), request_data=request, ) self._request_events[request_id] = event return event - async def _get_pending_request_snapshot(self, request_id: str, ctx: WorkflowContext[Any]) -> dict[str, Any] | None: - pending = await self._collect_pending_request_snapshots(ctx) - snapshot = pending.get(request_id) - if snapshot is None: - return None - return snapshot - - async def _collect_pending_request_snapshots(self, ctx: WorkflowContext[Any]) -> dict[str, dict[str, Any]]: - combined: dict[str, dict[str, Any]] = {} - - try: - shared_pending = await ctx.get_shared_state(self._PENDING_SHARED_STATE_KEY) - except KeyError: - shared_pending = None - except Exception as exc: # pragma: no cover - transport specific - logger.warning(f"RequestInfoExecutor {self.id} failed to read shared pending state during rehydrate: {exc}") - shared_pending = None - - if isinstance(shared_pending, dict): - for key, value in shared_pending.items(): # type: ignore[attr-defined] - if isinstance(key, str) and isinstance(value, dict): - combined[key] = cast(dict[str, Any], value) - - try: - state = await ctx.get_state() - except Exception as exc: # pragma: no cover - transport specific - logger.warning(f"RequestInfoExecutor {self.id} failed to read runner state during rehydrate: {exc}") - state = None - - if isinstance(state, dict): - state_pending = state.get("pending_requests") - if isinstance(state_pending, dict): - for key, value in state_pending.items(): # type: ignore[attr-defined] - if isinstance(key, str) and isinstance(value, dict) and key not in combined: - combined[key] = cast(dict[str, Any], value) - - return combined - - def _construct_request_from_snapshot(self, snapshot: dict[str, Any]) -> RequestInfoMessage | None: - details_raw = snapshot.get("details") - details: dict[str, Any] = cast(dict[str, Any], details_raw) if isinstance(details_raw, dict) else {} + def _construct_request_from_snapshot(self, snapshot: PendingRequestSnapshot) -> RequestInfoMessage | None: + json_safe_dict = snapshot.request_as_json_safe_dict request_cls: type[RequestInfoMessage] = RequestInfoMessage - request_type_str = snapshot.get("request_type") + request_type_str = snapshot.request_type if isinstance(request_type_str, str) and ":" in request_type_str: module_name, class_name = request_type_str.split(":", 1) try: @@ -574,19 +446,19 @@ class RequestInfoExecutor(Executor): logger.warning(f"RequestInfoExecutor {self.id} could not import {module_name}.{class_name}: {exc}") request_cls = RequestInfoMessage - request: RequestInfoMessage | None = self._instantiate_request(request_cls, details) + request: RequestInfoMessage | None = self._instantiate_request(request_cls, json_safe_dict) if request is None and request_cls is not RequestInfoMessage: - request = self._instantiate_request(RequestInfoMessage, details) + request = self._instantiate_request(RequestInfoMessage, json_safe_dict) if request is None: logger.warning( f"RequestInfoExecutor {self.id} could not reconstruct request " - f"{request_type_str or RequestInfoMessage.__name__} from snapshot keys {sorted(details.keys())}" + f"{request_type_str or RequestInfoMessage.__name__} from snapshot keys {sorted(json_safe_dict.keys())}" ) return None - for key, value in details.items(): + for key, value in json_safe_dict.items(): if key == "request_id": continue try: @@ -597,7 +469,7 @@ class RequestInfoExecutor(Executor): ) continue - snapshot_request_id = snapshot.get("request_id") + snapshot_request_id = snapshot.request_id if isinstance(snapshot_request_id, str) and snapshot_request_id: try: request.request_id = snapshot_request_id @@ -664,178 +536,38 @@ class RequestInfoExecutor(Executor): return instance - @staticmethod - def pending_requests_from_checkpoint( - checkpoint: WorkflowCheckpoint, - *, - request_executor_ids: Iterable[str] | None = None, - ) -> list[PendingRequestDetails]: - executor_filter: set[str] | None = None - if request_executor_ids is not None: - executor_filter = {str(value) for value in request_executor_ids} + # endregion: Internal Methods - pending: dict[str, PendingRequestDetails] = {} - shared_map = checkpoint.shared_state.get(RequestInfoExecutor._PENDING_SHARED_STATE_KEY) - if isinstance(shared_map, Mapping): - for request_id, snapshot in shared_map.items(): # type: ignore[attr-defined] - RequestInfoExecutor._merge_snapshot(pending, str(request_id), snapshot) # type: ignore[arg-type] +# region: Utility Functions - for state in checkpoint.executor_states.values(): - if not isinstance(state, Mapping): - continue - inner = state.get("pending_requests") - if isinstance(inner, Mapping): - for request_id, snapshot in inner.items(): # type: ignore[attr-defined] - RequestInfoExecutor._merge_snapshot(pending, str(request_id), snapshot) # type: ignore[arg-type] - for source_id, message_list in checkpoint.messages.items(): - if executor_filter is not None and source_id not in executor_filter: - continue - if not isinstance(message_list, list): - continue - for message in message_list: - if not isinstance(message, Mapping): - continue - payload = _decode_checkpoint_value(message.get("data")) - RequestInfoExecutor._merge_message_payload(pending, payload, message) +def _make_json_safe(value: Any) -> Any: + """Recursively convert a value to a JSON-safe representation.""" + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, Mapping): + safe_dict: dict[str, Any] = {} + for key, val in value.items(): # type: ignore[attr-defined] + safe_dict[str(key)] = _make_json_safe(val) # type: ignore[arg-type] + return safe_dict + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return [_make_json_safe(item) for item in value] # type: ignore[misc] + return repr(value) - return list(pending.values()) - @staticmethod - def checkpoint_summary( - checkpoint: WorkflowCheckpoint, - *, - request_executor_ids: Iterable[str] | None = None, - preview_width: int = 70, - ) -> WorkflowCheckpointSummary: - targets = sorted(checkpoint.messages.keys()) - executor_states = sorted(checkpoint.executor_states.keys()) - pending = RequestInfoExecutor.pending_requests_from_checkpoint( - checkpoint, request_executor_ids=request_executor_ids - ) +def _import_qualname(qualname: str) -> type[Any]: + """Import a type given its qualified name in the format 'module:TypeName'.""" + module_name, _, type_name = qualname.partition(":") + if not module_name or not type_name: + raise ValueError(f"Invalid qualified name: {qualname}") + module = importlib.import_module(module_name) + attr: Any = module + for part in type_name.split("."): + attr = getattr(attr, part) + if not isinstance(attr, type): + raise TypeError(f"Resolved object is not a type: {qualname}") + return attr - draft_preview: str | None = None - for entry in pending: - if entry.draft: - draft_preview = shorten(entry.draft, width=preview_width, placeholder="…") - break - status = "idle" - if pending: - status = "awaiting human response" - elif not checkpoint.messages and "finalise" in executor_states: - status = "completed" - elif checkpoint.messages: - status = "awaiting next superstep" - elif request_executor_ids is not None and any(tid in targets for tid in request_executor_ids): - status = "awaiting request delivery" - - return WorkflowCheckpointSummary( - checkpoint_id=checkpoint.checkpoint_id, - iteration_count=checkpoint.iteration_count, - targets=targets, - executor_states=executor_states, - status=status, - draft_preview=draft_preview, - pending_requests=pending, - ) - - @staticmethod - def _merge_snapshot( - pending: dict[str, PendingRequestDetails], - request_id: str, - snapshot: Any, - ) -> None: - if not request_id or not isinstance(snapshot, Mapping): - return - - details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id)) - - RequestInfoExecutor._apply_update( - details, - prompt=snapshot.get("prompt"), # type: ignore[attr-defined] - draft=snapshot.get("draft"), # type: ignore[attr-defined] - iteration=snapshot.get("iteration"), # type: ignore[attr-defined] - source_executor_id=snapshot.get("source_executor_id"), # type: ignore[attr-defined] - ) - - extra = snapshot.get("details") # type: ignore[attr-defined] - if isinstance(extra, Mapping): - RequestInfoExecutor._apply_update( - details, - prompt=extra.get("prompt"), # type: ignore[attr-defined] - draft=extra.get("draft"), # type: ignore[attr-defined] - iteration=extra.get("iteration"), # type: ignore[attr-defined] - ) - - @staticmethod - def _merge_message_payload( - pending: dict[str, PendingRequestDetails], - payload: Any, - raw_message: Mapping[str, Any], - ) -> None: - if isinstance(payload, RequestResponse): - request_id = payload.request_id or RequestInfoExecutor._get_field(payload.original_request, "request_id") # type: ignore[arg-type] - if not request_id: - return - details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id)) - RequestInfoExecutor._apply_update( - details, - prompt=RequestInfoExecutor._get_field(payload.original_request, "prompt"), # type: ignore[arg-type] - draft=RequestInfoExecutor._get_field(payload.original_request, "draft"), # type: ignore[arg-type] - iteration=RequestInfoExecutor._get_field(payload.original_request, "iteration"), # type: ignore[arg-type] - source_executor_id=raw_message.get("source_id"), - original_request=payload.original_request, # type: ignore[arg-type] - ) - elif isinstance(payload, RequestInfoMessage): - request_id = getattr(payload, "request_id", None) - if not request_id: - return - details = pending.setdefault(request_id, PendingRequestDetails(request_id=request_id)) - RequestInfoExecutor._apply_update( - details, - prompt=getattr(payload, "prompt", None), - draft=getattr(payload, "draft", None), - iteration=getattr(payload, "iteration", None), - source_executor_id=raw_message.get("source_id"), - original_request=payload, - ) - - @staticmethod - def _apply_update( - details: PendingRequestDetails, - *, - prompt: Any = None, - draft: Any = None, - iteration: Any = None, - source_executor_id: Any = None, - original_request: Any = None, - ) -> None: - if prompt and not details.prompt: - details.prompt = str(prompt) - if draft and not details.draft: - details.draft = str(draft) - if iteration is not None and details.iteration is None: - coerced = RequestInfoExecutor._coerce_int(iteration) - if coerced is not None: - details.iteration = coerced - if source_executor_id and not details.source_executor_id: - details.source_executor_id = str(source_executor_id) - if original_request is not None and details.original_request is None: - details.original_request = original_request - - @staticmethod - def _get_field(obj: Any, key: str) -> Any: - if obj is None: - return None - if isinstance(obj, Mapping): - return obj.get(key) # type: ignore[attr-defined,return-value] - return getattr(obj, key, None) - - @staticmethod - def _coerce_int(value: Any) -> int | None: - try: - return int(value) - except (TypeError, ValueError): - return None +# endregion: Utility Functions diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index 11a0499acf..9a0d0a7790 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -304,6 +304,8 @@ class Runner: checkpoint_id, ) + await self._restore_executor_states(checkpoint.executor_states) + state = self._checkpoint_to_state(checkpoint) await self._ctx.set_checkpoint_state(state) if checkpoint.workflow_id: @@ -323,6 +325,27 @@ class Runner: logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}") return False + async def _restore_executor_states(self, executor_states: dict[str, dict[str, Any]]) -> None: + for exec_id, state in executor_states.items(): + executor = self._executors.get(exec_id) + if not executor: + logger.debug(f"Executor {exec_id} not found during state restoration; skipping.") + continue + + restored = False + restore_method = getattr(executor, "restore_state", None) + try: + if callable(restore_method): + maybe = restore_method(state) + if asyncio.iscoroutine(maybe): # type: ignore[arg-type] + await maybe # type: ignore[arg-type] + restored = True + except Exception as ex: # pragma: no cover - defensive + logger.debug(f"Executor {exec_id} restore_state failed: {ex}") + + if not restored: + logger.debug(f"Executor {exec_id} does not support state restoration; skipping.") + async def _restore_shared_state_from_context(self) -> None: try: restored_state = await self._ctx.get_checkpoint_state() diff --git a/python/packages/core/agent_framework/_workflows/_workflow_context.py b/python/packages/core/agent_framework/_workflows/_workflow_context.py index b55966fb7a..b7e077424e 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_context.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_context.py @@ -441,14 +441,11 @@ class WorkflowContext(Generic[T_Out, T_W_Out]): Executors call this with a JSON-serializable dict capturing the minimal state needed to resume. It replaces any previously stored state. """ - if hasattr(self._runner_context, "set_state"): - await self._runner_context.set_state(self._executor_id, state) # type: ignore[arg-type] + await self._runner_context.set_state(self._executor_id, state) async def get_state(self) -> dict[str, Any] | None: """Retrieve previously persisted state for this executor, if any.""" - if hasattr(self._runner_context, "get_state"): - return await self._runner_context.get_state(self._executor_id) # type: ignore[return-value] - return None + return await self._runner_context.get_state(self._executor_id) def is_streaming(self) -> bool: """Check if the workflow is running in streaming mode. diff --git a/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py b/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py index 2b52a04f61..15d0d7d7af 100644 --- a/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py +++ b/python/packages/core/tests/workflow/test_request_info_executor_rehydrate.py @@ -6,17 +6,19 @@ from datetime import datetime, timezone from typing import Any from agent_framework._workflows._checkpoint import CheckpointStorage, WorkflowCheckpoint +from agent_framework._workflows._checkpoint_summary import get_checkpoint_summary from agent_framework._workflows._events import RequestInfoEvent, WorkflowEvent from agent_framework._workflows._request_info_executor import ( PendingRequestDetails, + PendingRequestSnapshot, RequestInfoExecutor, RequestInfoMessage, RequestResponse, ) -from agent_framework._workflows._runner_context import ( # type: ignore +from agent_framework._workflows._runner_context import ( CheckpointState, Message, - _encode_checkpoint_value, + _encode_checkpoint_value, # type: ignore ) from agent_framework._workflows._shared_state import SharedState from agent_framework._workflows._workflow_context import WorkflowContext @@ -85,6 +87,12 @@ class _StubRunnerContext: async def set_checkpoint_state(self, state: CheckpointState) -> None: # pragma: no cover - unused pass + def set_streaming(self, streaming: bool) -> None: # pragma: no cover - unused + pass + + def is_streaming(self) -> bool: # pragma: no cover - unused + return False + @dataclass(kw_only=True) class SimpleApproval(RequestInfoMessage): @@ -109,30 +117,18 @@ async def test_rehydrate_falls_back_when_request_type_missing() -> None: This simulates resuming a workflow where the HumanApprovalRequest class is unavailable in the current process (e.g., defined in __main__ during the original run). """ - request_id = "request-123" - snapshot = { - "request_id": request_id, - "source_executor_id": "review_gateway", - "request_type": "nonexistent.module:MissingRequest", - "summary": "...", - "details": { + snapshot = PendingRequestSnapshot( + request_id=request_id, + source_executor_id="review_gateway", + request_type="nonexistent.module:MissingRequest", + request_as_json_safe_dict={ "request_id": request_id, - "prompt": "Review draft", - "draft": "Draft text", - "iteration": 2, }, - } + ) - shared_state = SharedState() - async with shared_state.hold(): - await shared_state.set_within_hold( - PENDING_STATE_KEY, - {request_id: snapshot}, - ) - - runner_ctx = _StubRunnerContext({"pending_requests": {request_id: snapshot}}) - ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], shared_state, runner_ctx) + runner_ctx = _StubRunnerContext({PENDING_STATE_KEY: {request_id: snapshot}}) + ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), runner_ctx) executor = RequestInfoExecutor(id="request_info") @@ -141,31 +137,21 @@ async def test_rehydrate_falls_back_when_request_type_missing() -> None: assert event is not None assert event.request_id == request_id assert isinstance(event.data, RequestInfoMessage) - assert getattr(event.data, "prompt", None) == "Review draft" - assert getattr(event.data, "iteration", None) == 2 async def test_has_pending_request_detects_snapshot() -> None: - request_id = "req-pending" - snapshot = { - "request_id": request_id, - "source_executor_id": "review_gateway", - "details": { + request_id = "request-123" + snapshot = PendingRequestSnapshot( + request_id=request_id, + source_executor_id="review_gateway", + request_type="nonexistent.module:MissingRequest", + request_as_json_safe_dict={ "request_id": request_id, - "prompt": "Review", - "draft": "Draft", }, - } + ) - shared_state = SharedState() - async with shared_state.hold(): - await shared_state.set_within_hold( - PENDING_STATE_KEY, - {request_id: snapshot}, - ) - - runner_ctx = _StubRunnerContext({"pending_requests": {request_id: snapshot}}) - ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], shared_state, runner_ctx) + runner_ctx = _StubRunnerContext({PENDING_STATE_KEY: {request_id: snapshot}}) + ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), runner_ctx) executor = RequestInfoExecutor(id="request_info") @@ -221,7 +207,12 @@ def test_pending_requests_from_checkpoint_and_summary() -> None: iteration_count=1, ) - pending = RequestInfoExecutor.pending_requests_from_checkpoint(checkpoint) + summary = get_checkpoint_summary(checkpoint) + assert summary.checkpoint_id == "cp-1" + assert summary.status == "awaiting request response" + assert summary.pending_requests[0].request_id == "req-42" + + pending = summary.pending_requests assert len(pending) == 1 entry = pending[0] assert isinstance(entry, PendingRequestDetails) @@ -231,11 +222,6 @@ def test_pending_requests_from_checkpoint_and_summary() -> None: assert entry.iteration == 3 assert entry.original_request is not None - summary = RequestInfoExecutor.checkpoint_summary(checkpoint) - assert summary.checkpoint_id == "cp-1" - assert summary.status == "awaiting human response" - assert summary.pending_requests[0].request_id == "req-42" - def test_snapshot_state_serializes_non_json_payloads() -> None: executor = RequestInfoExecutor(id="request_info") @@ -305,13 +291,10 @@ async def test_run_persists_pending_requests_in_runner_state() -> None: await executor.execute(approval, ctx.source_executor_ids, shared_state, runner_ctx) # Runner state should include both pending snapshot and serialized request events - assert "pending_requests" in runner_ctx._state # pyright: ignore[reportPrivateUsage] - assert approval.request_id in runner_ctx._state["pending_requests"] # pyright: ignore[reportPrivateUsage] - assert "request_events" in runner_ctx._state # pyright: ignore[reportPrivateUsage] - assert approval.request_id in runner_ctx._state["request_events"] # pyright: ignore[reportPrivateUsage] + assert PENDING_STATE_KEY in runner_ctx._state # pyright: ignore[reportPrivateUsage] + assert approval.request_id in runner_ctx._state[PENDING_STATE_KEY] # pyright: ignore[reportPrivateUsage] response_ctx: WorkflowContext[None] = WorkflowContext("request_info", ["source"], shared_state, runner_ctx) await executor.handle_response("approved", approval.request_id, response_ctx) # type: ignore - assert runner_ctx._state["pending_requests"] == {} # pyright: ignore[reportPrivateUsage] - assert runner_ctx._state.get("request_events", {}).get(approval.request_id) is None # pyright: ignore[reportPrivateUsage] + assert runner_ctx._state[PENDING_STATE_KEY] == {} # pyright: ignore[reportPrivateUsage] diff --git a/python/packages/core/tests/workflow/test_workflow_states.py b/python/packages/core/tests/workflow/test_workflow_states.py index 1f6adba134..1c4ba561d5 100644 --- a/python/packages/core/tests/workflow/test_workflow_states.py +++ b/python/packages/core/tests/workflow/test_workflow_states.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. from dataclasses import dataclass -from typing import Any import pytest from typing_extensions import Never @@ -25,7 +24,6 @@ from agent_framework import ( WorkflowStatusEvent, handler, ) -from agent_framework import WorkflowContext as WFContext class FailingExecutor(Executor): @@ -182,39 +180,3 @@ class SnapshotRequester(Executor): @handler async def ask(self, _: str, ctx: WorkflowContext[SnapshotRequest]) -> None: # pragma: no cover - simple helper await ctx.send_message(SnapshotRequest(prompt=self._prompt, draft=self._draft, iteration=1)) - - -async def test_request_info_executor_tracks_pending_requests_via_shared_state(): - prompt = "Review the launch copy" - draft = "Limited edition grinder now $249" - requester = SnapshotRequester(id="snapshot_req", prompt=prompt, draft=draft) - request_info = RequestInfoExecutor(id="request_info") - - wf = WorkflowBuilder().set_start_executor(requester).add_edge(requester, request_info).build() - - events = [event async for event in wf.run_stream("start")] - assert any(isinstance(event, RequestInfoEvent) for event in events) - - pending_map: dict[str, Any] = await wf._shared_state.get(RequestInfoExecutor._PENDING_SHARED_STATE_KEY) # type: ignore[reportPrivateUsage] - assert isinstance(pending_map, dict) - assert len(pending_map) == 1 - snapshot: dict[str, Any] = next(iter(pending_map.values())) - assert snapshot["prompt"] == prompt - assert snapshot["draft"] == draft - assert snapshot.get("iteration") == 1 - - request_id: str = snapshot["request_id"] - - request_info_resume = RequestInfoExecutor(id="request_info_resume") - resume_context: WFContext[Any] = WFContext( - executor_id=request_info_resume.id, - source_executor_ids=[wf.__class__.__name__], - shared_state=wf._shared_state, # type: ignore[reportPrivateUsage] - runner_context=wf._runner_context, # type: ignore[reportPrivateUsage] - ) - - await request_info_resume.handle_response("approve", request_id, resume_context) - - updated_pending: dict[str, Any] = await wf._shared_state.get(RequestInfoExecutor._PENDING_SHARED_STATE_KEY) # type: ignore[reportPrivateUsage] - assert isinstance(updated_pending, dict) - assert request_id not in updated_pending diff --git a/python/samples/_run_all_samples.py b/python/samples/_run_all_samples.py new file mode 100644 index 0000000000..ba5f21d8a5 --- /dev/null +++ b/python/samples/_run_all_samples.py @@ -0,0 +1,235 @@ +# Copyright (c) Microsoft. All rights reserved. + +""" +Script to run all Python samples in the samples directory concurrently. +This script will run all samples and report results at the end. + +Note: This script is AI generated. This is for internal validation purposes only. + +Samples that require human interaction are known to fail. + +Usage: + python run_all_samples.py # Run all samples using uv run (concurrent) + python run_all_samples.py --direct # Run all samples directly (concurrent, + # assumes environment is set up) + python run_all_samples.py --subdir # Run samples only in specific subdirectory + python run_all_samples.py --subdir getting_started/workflows # Example: run only workflow samples +""" + +import argparse +import os +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + + +def find_python_samples(samples_dir: Path, subdir: str | None = None) -> list[Path]: + """Find all Python sample files in the samples directory or a subdirectory.""" + python_files: list[Path] = [] + + # Determine the search directory + if subdir: + search_dir = samples_dir / subdir + if not search_dir.exists(): + print(f"Warning: Subdirectory '{subdir}' does not exist in {samples_dir}") + return [] + print(f"Searching in subdirectory: {search_dir}") + else: + search_dir = samples_dir + print(f"Searching in all samples: {search_dir}") + + # Walk through all subdirectories and find .py files + for root, dirs, files in os.walk(search_dir): + # Skip __pycache__ directories + dirs[:] = [d for d in dirs if d != "__pycache__"] + + for file in files: + if file.endswith(".py") and not file.startswith("_") and file != "_run_all_samples.py": + python_files.append(Path(root) / file) + + # Sort files for consistent execution order + return sorted(python_files) + + +def run_sample( + sample_path: Path, + use_uv: bool = True, + python_root: Path | None = None, +) -> tuple[bool, str, str]: + """ + Run a single sample file using subprocess and return (success, output, error_info). + + Args: + sample_path: Path to the sample file + use_uv: Whether to use uv run + python_root: Root directory for uv run + + Returns: + Tuple of (success, output, error_info) + """ + if use_uv and python_root: + cmd = ["uv", "run", "python", str(sample_path)] + cwd = python_root + else: + cmd = [sys.executable, sample_path.name] + cwd = sample_path.parent + + try: + result = subprocess.run( + cmd, + cwd=cwd, + capture_output=True, + text=True, + timeout=60, # 60 second timeout + ) + + if result.returncode == 0: + output = result.stdout.strip() if result.stdout.strip() else "No output" + return True, output, "" + + error_info = f"Exit code: {result.returncode}" + if result.stderr.strip(): + error_info += f"\nSTDERR: {result.stderr}" + + return False, result.stdout.strip() if result.stdout.strip() else "", error_info + + except subprocess.TimeoutExpired: + return False, "", f"TIMEOUT: {sample_path.name} (exceeded 60 seconds)" + except Exception as e: + return False, "", f"ERROR: {sample_path.name} - Exception: {str(e)}" + + +def parse_arguments() -> argparse.Namespace: + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Run Python samples concurrently", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python run_all_samples.py # Run all samples + python run_all_samples.py --direct # Run all samples directly + python run_all_samples.py --subdir getting_started # Run only getting_started samples + python run_all_samples.py --subdir getting_started/workflows # Run only workflow samples + python run_all_samples.py --subdir semantic-kernel-migration # Run only SK migration samples + """, + ) + + parser.add_argument( + "--direct", action="store_true", help="Run samples directly with python instead of using uv run" + ) + + parser.add_argument( + "--subdir", type=str, help="Run samples only in the specified subdirectory (relative to samples/)" + ) + + parser.add_argument( + "--max-workers", type=int, default=16, help="Maximum number of concurrent workers (default: 16)" + ) + + return parser.parse_args() + + +def main() -> None: + """Main function to run all samples concurrently.""" + args = parse_arguments() + + # Get the samples directory (assuming this script is in the samples directory) + samples_dir = Path(__file__).parent + python_root = samples_dir.parent # Go up to the python/ directory + + print("Python samples runner") + print(f"Samples directory: {samples_dir}") + + if args.direct: + print("Running samples directly (assuming environment is set up)") + else: + print(f"Using uv run from: {python_root}") + + if args.subdir: + print(f"Filtering to subdirectory: {args.subdir}") + + print("🚀 Running samples concurrently...") + + # Find all Python sample files + sample_files = find_python_samples(samples_dir, args.subdir) + + if not sample_files: + print("No Python sample files found!") + return + + print(f"Found {len(sample_files)} Python sample files") + + # Run samples concurrently + results: list[tuple[Path, bool, str, str]] = [] + + with ThreadPoolExecutor(max_workers=args.max_workers) as executor: + # Submit all tasks + future_to_sample = { + executor.submit(run_sample, sample_path, not args.direct, python_root): sample_path + for sample_path in sample_files + } + + # Collect results as they complete + for future in as_completed(future_to_sample): + sample_path = future_to_sample[future] + try: + success, output, error_info = future.result() + results.append((sample_path, success, output, error_info)) + + # Print progress - show relative path from samples directory + relative_path = sample_path.relative_to(samples_dir) + if success: + print(f"✅ {relative_path}") + else: + print(f"❌ {relative_path} - {error_info.split(':', 1)[0]}") + + except Exception as e: + error_info = f"Future exception: {str(e)}" + results.append((sample_path, False, "", error_info)) + relative_path = sample_path.relative_to(samples_dir) + print(f"❌ {relative_path} - {error_info}") + + # Sort results by original file order for consistent reporting + sample_to_index = {path: i for i, path in enumerate(sample_files)} + results.sort(key=lambda x: sample_to_index[x[0]]) + + successful_runs = sum(1 for _, success, _, _ in results if success) + failed_runs = len(results) - successful_runs + + # Print detailed results + print(f"\n{'=' * 80}") + print("DETAILED RESULTS:") + print(f"{'=' * 80}") + + for sample_path, success, output, error_info in results: + relative_path = sample_path.relative_to(samples_dir) + if success: + print(f"✅ {relative_path}") + if output and output != "No output": + print(f" Output preview: {output[:100]}{'...' if len(output) > 100 else ''}") + else: + print(f"❌ {relative_path}") + print(f" Error: {error_info}") + + # Print summary + print(f"\n{'=' * 80}") + if failed_runs == 0: + print("🎉 ALL SAMPLES COMPLETED SUCCESSFULLY!") + else: + print(f"❌ {failed_runs} SAMPLE(S) FAILED!") + print(f"Successful runs: {successful_runs}") + print(f"Failed runs: {failed_runs}") + + if args.subdir: + print(f"Subdirectory filter: {args.subdir}") + + print(f"{'=' * 80}") + + # Exit with error code if any samples failed + if failed_runs > 0: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/python/samples/getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py b/python/samples/getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py index 3c869439c2..e4fba82495 100644 --- a/python/samples/getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py +++ b/python/samples/getting_started/workflows/agents/workflow_as_agent_human_in_the_loop.py @@ -5,7 +5,7 @@ import sys from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path -from typing import Any, cast +from typing import Any # Ensure local getting_started package can be imported when running as a script. _SAMPLES_ROOT = Path(__file__).resolve().parents[3] @@ -150,23 +150,17 @@ async def main() -> None: else: raise TypeError("Unexpected argument type for human review function call.") - request_payload_obj: Any = request.data - if not isinstance(request_payload_obj, Mapping): - raise ValueError("Human review request payload must be a mapping.") - request_payload = cast(Mapping[str, Any], request_payload_obj) + request_payload: Any = request.data + if not isinstance(request_payload, HumanReviewRequest): + raise ValueError("Human review request payload must be a HumanReviewRequest.") - agent_request_obj = request_payload.get("agent_request") - if not isinstance(agent_request_obj, Mapping): - raise ValueError("Human review request must include agent_request mapping data.") - agent_request_data = cast(Mapping[str, Any], agent_request_obj) - - request_id_obj = agent_request_data.get("request_id") - if not isinstance(request_id_obj, str): - raise ValueError("Human review request_id must be a string.") - request_id_value = request_id_obj + agent_request = request_payload.agent_request + if agent_request is None: + raise ValueError("Human review request must include agent_request.") + request_id = agent_request.request_id # Mock a human response approval for demonstration purposes. - human_response = ReviewResponse(request_id=request_id_value, feedback="Approved", approved=True) + human_response = ReviewResponse(request_id=request_id, feedback="Approved", approved=True) # Create the function call result object to send back to the agent. human_review_function_result = FunctionResultContent( 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 e397e83088..3dc80339bc 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 @@ -23,6 +23,7 @@ from agent_framework import ( WorkflowOutputEvent, WorkflowRunState, WorkflowStatusEvent, + get_checkpoint_summary, handler, ) from agent_framework.azure import AzureOpenAIChatClient @@ -246,14 +247,12 @@ def _render_checkpoint_summary(checkpoints: list["WorkflowCheckpoint"]) -> None: """Pretty-print saved checkpoints with the new framework summaries.""" print("\nCheckpoint summary:") - for summary in [ - RequestInfoExecutor.checkpoint_summary(cp) for cp in sorted(checkpoints, key=lambda c: c.timestamp) - ]: + for summary in [get_checkpoint_summary(cp) for cp in sorted(checkpoints, key=lambda c: c.timestamp)]: # Compose a single line per checkpoint so the user can scan the output # and pick the resume point that still has outstanding human work. line = ( f"- {summary.checkpoint_id} | iter={summary.iteration_count} " - f"| targets={summary.targets} | states={summary.executor_states}" + f"| targets={summary.targets} | states={summary.executor_ids}" ) if summary.status: line += f" | status={summary.status}" @@ -312,7 +311,7 @@ def _prompt_for_responses(requests: list[tuple[str, HumanApprovalRequest]]) -> d def _maybe_pre_supply_responses(cp: "WorkflowCheckpoint") -> dict[str, str] | None: """Offer to collect responses before resuming a checkpoint.""" - pending = RequestInfoExecutor.pending_requests_from_checkpoint(cp) + pending = get_checkpoint_summary(cp).pending_requests if not pending: return None @@ -468,7 +467,7 @@ async def main() -> None: return chosen = sorted_cps[idx] - summary = RequestInfoExecutor.checkpoint_summary(chosen) + summary = get_checkpoint_summary(chosen) if summary.status == "completed": print("Selected checkpoint already reflects a completed workflow; nothing to resume.") return 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 a3dacc4a35..c85815c911 100644 --- a/python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py +++ b/python/samples/getting_started/workflows/checkpoint/checkpoint_with_resume.py @@ -12,10 +12,10 @@ from agent_framework import ( ChatMessage, Executor, FileCheckpointStorage, - RequestInfoExecutor, Role, WorkflowBuilder, WorkflowContext, + get_checkpoint_summary, handler, ) from agent_framework.azure import AzureOpenAIChatClient @@ -194,7 +194,7 @@ def _render_checkpoint_summary(checkpoints: list["WorkflowCheckpoint"]) -> None: print("\nCheckpoint summary:") for cp in sorted(checkpoints, key=lambda c: c.timestamp): - summary = RequestInfoExecutor.checkpoint_summary(cp) + summary = get_checkpoint_summary(cp) msg_count = sum(len(v) for v in cp.messages.values()) state_keys = sorted(cp.executor_states.keys()) orig = cp.shared_state.get("original_input") @@ -241,7 +241,7 @@ async def main(): print("\nAvailable checkpoints to resume from:") for idx, cp in enumerate(sorted_cps): - summary = RequestInfoExecutor.checkpoint_summary(cp) + summary = get_checkpoint_summary(cp) line = f" [{idx}] id={summary.checkpoint_id} iter={summary.iteration_count}" if summary.status: line += f" status={summary.status}"