From 2133043f11c5d1c86cc1c25ab6a1c7bc1a23ddbb Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Tue, 23 Sep 2025 13:52:53 -0700 Subject: [PATCH] Python: [Breaking] Remove `WorkflowCompletedEvent`, introduce workflow output and migrate to `ctx.yield_output()` + a huge refactoring (#845) * Introduce input and output types for executor and workflow * WorkflowOutputContext handles two types * Remove can_handle_types from Executor * Update validation * Move workflow executor * Move workflow executor * Fix issues in WorkflowExecutor * refactor executor * update execute signature to create workflow context within Executor * fix simple sub workflow test; fix validation * fix output types in WorkflowExecutor * fix issue in Executor handling of SubWorkflowRequestInfo * update tests to use proper workflow output * update orchestration patterns to use output * Update sample -- not finished * Update python/packages/main/tests/workflow/test_workflow_states.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/main/tests/workflow/test_concurrent.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * address comments * WorkflowOutputContext --> WorkflowContext * remove WorkflowCompletedEvent * update samples * Update doc string for important classes; update WorkflowExecutor to support concurrent execution * use Never instead of None for default type * Update usage of WorkflowContext[None to WorkflowContext[Never * address comments * remove filter for None * address comments, minor fixes * quality of life improvement on interceptor types --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../devui/samples/fanout_workflow/workflow.py | 22 +- .../devui/samples/spam_workflow/workflow.py | 6 +- .../agent_framework/_workflow/__init__.py | 8 +- .../agent_framework/_workflow/__init__.pyi | 8 +- .../main/agent_framework/_workflow/_agent.py | 2 +- .../agent_framework/_workflow/_concurrent.py | 40 +- .../agent_framework/_workflow/_edge_runner.py | 18 +- .../main/agent_framework/_workflow/_events.py | 41 +- .../agent_framework/_workflow/_executor.py | 858 ++++++++++-------- .../_workflow/_function_executor.py | 137 +-- .../agent_framework/_workflow/_magentic.py | 38 +- .../main/agent_framework/_workflow/_runner.py | 42 +- .../agent_framework/_workflow/_sequential.py | 22 +- .../agent_framework/_workflow/_validation.py | 300 +----- .../main/agent_framework/_workflow/_viz.py | 4 +- .../agent_framework/_workflow/_workflow.py | 150 ++- .../_workflow/_workflow_context.py | 295 +++++- .../_workflow/_workflow_executor.py | 437 +++++++++ .../workflow/test_checkpoint_validation.py | 9 +- .../main/tests/workflow/test_concurrent.py | 56 +- .../tests/workflow/test_full_conversation.py | 42 +- .../tests/workflow/test_function_executor.py | 117 ++- .../main/tests/workflow/test_magentic.py | 38 +- .../main/tests/workflow/test_runner.py | 48 +- .../main/tests/workflow/test_sequential.py | 36 +- .../main/tests/workflow/test_serialization.py | 2 +- .../workflow/test_simple_sub_workflow.py | 12 +- .../main/tests/workflow/test_sub_workflow.py | 121 ++- .../main/tests/workflow/test_validation.py | 60 +- .../packages/main/tests/workflow/test_viz.py | 4 +- .../main/tests/workflow/test_workflow.py | 127 +-- .../tests/workflow/test_workflow_context.py | 208 ++++- .../workflow/test_workflow_observability.py | 15 +- .../tests/workflow/test_workflow_states.py | 38 +- .../observability/04-workflow.py | 22 +- .../_start-here/step1_executors_and_edges.py | 56 +- .../_start-here/step2_agents_in_a_workflow.py | 19 +- .../workflow/_start-here/step3_streaming.py | 30 +- .../agents/azure_chat_agents_streaming.py | 21 +- .../workflow/agents/custom_agent_executors.py | 21 +- .../agents/foundry_chat_agents_streaming.py | 20 +- .../workflow_as_agent_human_in_the_loop.py | 2 +- .../workflow_as_agent_reflection_pattern.py | 2 +- .../checkpoint_with_human_in_the_loop.py | 58 +- .../checkpoint/checkpoint_with_resume.py | 13 +- .../composition/sub_workflow_basics.py | 16 +- .../sub_workflow_parallel_requests.py | 34 +- .../sub_workflow_request_interception.py | 19 +- .../workflow/control-flow/edge_condition.py | 35 +- .../multi_selection_edge_group.py | 35 +- .../control-flow/sequential_executors.py | 34 +- .../control-flow/sequential_streaming.py | 37 +- .../workflow/control-flow/simple_loop.py | 9 +- .../control-flow/switch_case_edge_group.py | 32 +- .../guessing_game_with_human_input.py | 30 +- .../workflow/observability/tracing_basics.py | 5 +- .../orchestration/concurrent_agents.py | 27 +- .../concurrent_custom_agent_executors.py | 12 +- .../concurrent_custom_aggregator.py | 17 +- .../workflow/orchestration/magentic.py | 19 +- .../magentic_human_plan_update.py | 91 +- .../orchestration/sequential_agents.py | 21 +- .../sequential_custom_executors.py | 22 +- .../parallelism/fan_out_fan_in_edges.py | 19 +- .../map_reduce_and_visualization.py | 23 +- .../shared_states_with_agents.py | 31 +- .../concurrent_with_visualization.py | 19 +- 67 files changed, 2564 insertions(+), 1648 deletions(-) create mode 100644 python/packages/main/agent_framework/_workflow/_workflow_executor.py diff --git a/python/packages/devui/samples/fanout_workflow/workflow.py b/python/packages/devui/samples/fanout_workflow/workflow.py index ab5c98dcc8..b5b43fbd66 100644 --- a/python/packages/devui/samples/fanout_workflow/workflow.py +++ b/python/packages/devui/samples/fanout_workflow/workflow.py @@ -23,11 +23,11 @@ from typing import Literal from agent_framework import ( Executor, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, handler, ) from pydantic import BaseModel, Field +from typing_extensions import Never class DataType(Enum): @@ -315,7 +315,9 @@ class ValidationAggregator(Executor): """Aggregates validation results and decides on next steps.""" @handler - async def aggregate_validations(self, reports: list[ValidationReport], ctx: WorkflowContext[DataBatch]) -> None: + async def aggregate_validations( + self, reports: list[ValidationReport], ctx: WorkflowContext[DataBatch, str] + ) -> None: """Aggregate all validation reports and make processing decision.""" if not reports: return @@ -345,11 +347,9 @@ class ValidationAggregator(Executor): ) reason = " and ".join(failure_reason) - await ctx.add_event( - WorkflowCompletedEvent( - f"Batch {batch_id} failed validation: {reason}. " - f"Total issues: {total_issues}, Quality score: {quality_score:.2f}" - ) + await ctx.yield_output( + f"Batch {batch_id} failed validation: {reason}. " + f"Total issues: {total_issues}, Quality score: {quality_score:.2f}" ) return @@ -584,10 +584,12 @@ class FinalProcessor(Executor): """Final processing stage that combines all results.""" @handler - async def process_final_results(self, assessments: list[QualityAssessment], ctx: WorkflowContext[None]) -> None: + async def process_final_results( + self, assessments: list[QualityAssessment], ctx: WorkflowContext[Never, str] + ) -> None: """Generate final processing summary and complete workflow.""" if not assessments: - await ctx.add_event(WorkflowCompletedEvent("No quality assessments received")) + await ctx.yield_output("No quality assessments received") return batch_id = assessments[0].batch_id @@ -618,7 +620,7 @@ class FinalProcessor(Executor): f"🎖️ Final Status: {final_status}" ) - await ctx.add_event(WorkflowCompletedEvent(completion_message)) + await ctx.yield_output(completion_message) # Workflow Builder Helper diff --git a/python/packages/devui/samples/spam_workflow/workflow.py b/python/packages/devui/samples/spam_workflow/workflow.py index 32b9ff0985..d922098252 100644 --- a/python/packages/devui/samples/spam_workflow/workflow.py +++ b/python/packages/devui/samples/spam_workflow/workflow.py @@ -24,11 +24,11 @@ from agent_framework import ( Default, Executor, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, handler, ) from pydantic import BaseModel, Field +from typing_extensions import Never @dataclass @@ -259,7 +259,7 @@ class FinalProcessor(Executor): async def handle_processing_result( self, result: ProcessingResult, - ctx: WorkflowContext[None], + ctx: WorkflowContext[Never, str], ) -> None: """Complete the workflow with final processing and logging.""" await asyncio.sleep(1.5) # Simulate final processing time @@ -278,7 +278,7 @@ class FinalProcessor(Executor): f"Total time: {total_time:.1f}s" ) - await ctx.add_event(WorkflowCompletedEvent(completion_message)) + await ctx.yield_output(completion_message) # Create the workflow instance that DevUI can discover diff --git a/python/packages/main/agent_framework/_workflow/__init__.py b/python/packages/main/agent_framework/_workflow/__init__.py index 18563a8d15..c97927c0d5 100644 --- a/python/packages/main/agent_framework/_workflow/__init__.py +++ b/python/packages/main/agent_framework/_workflow/__init__.py @@ -33,12 +33,12 @@ from ._events import ( ExecutorFailedEvent, ExecutorInvokedEvent, RequestInfoEvent, - WorkflowCompletedEvent, WorkflowErrorDetails, WorkflowEvent, WorkflowEventSource, WorkflowFailedEvent, WorkflowLifecycleEvent, + WorkflowOutputEvent, WorkflowRunState, WorkflowStartedEvent, WorkflowStatusEvent, @@ -53,7 +53,6 @@ from ._executor import ( RequestResponse, SubWorkflowRequestInfo, SubWorkflowResponse, - WorkflowExecutor, handler, intercepts_request, ) @@ -92,7 +91,6 @@ from ._validation import ( EdgeDuplicationError, ExecutorDuplicationError, GraphConnectivityError, - HandlerOutputAnnotationError, TypeCompatibilityError, ValidationTypeEnum, WorkflowValidationError, @@ -101,6 +99,7 @@ from ._validation import ( from ._viz import WorkflowViz from ._workflow import Workflow, WorkflowBuilder, WorkflowRunResult from ._workflow_context import WorkflowContext +from ._workflow_executor import WorkflowExecutor __all__ = [ "DEFAULT_MAX_ITERATIONS", @@ -126,7 +125,6 @@ __all__ = [ "FileCheckpointStorage", "FunctionExecutor", "GraphConnectivityError", - "HandlerOutputAnnotationError", "InMemoryCheckpointStorage", "InProcRunnerContext", "MagenticAgentDeltaEvent", @@ -170,7 +168,6 @@ __all__ = [ "WorkflowAgent", "WorkflowBuilder", "WorkflowCheckpoint", - "WorkflowCompletedEvent", "WorkflowContext", "WorkflowErrorDetails", "WorkflowEvent", @@ -178,6 +175,7 @@ __all__ = [ "WorkflowExecutor", "WorkflowFailedEvent", "WorkflowLifecycleEvent", + "WorkflowOutputEvent", "WorkflowRunResult", "WorkflowRunState", "WorkflowStartedEvent", diff --git a/python/packages/main/agent_framework/_workflow/__init__.pyi b/python/packages/main/agent_framework/_workflow/__init__.pyi index 7099dfa2fe..349e6463a3 100644 --- a/python/packages/main/agent_framework/_workflow/__init__.pyi +++ b/python/packages/main/agent_framework/_workflow/__init__.pyi @@ -29,12 +29,12 @@ from ._events import ( ExecutorFailedEvent, ExecutorInvokedEvent, RequestInfoEvent, - WorkflowCompletedEvent, WorkflowErrorDetails, WorkflowEvent, WorkflowEventSource, WorkflowFailedEvent, WorkflowLifecycleEvent, + WorkflowOutputEvent, WorkflowRunState, WorkflowStartedEvent, WorkflowStatusEvent, @@ -49,7 +49,6 @@ from ._executor import ( RequestResponse, SubWorkflowRequestInfo, SubWorkflowResponse, - WorkflowExecutor, handler, intercepts_request, ) @@ -88,7 +87,6 @@ from ._validation import ( EdgeDuplicationError, ExecutorDuplicationError, GraphConnectivityError, - HandlerOutputAnnotationError, TypeCompatibilityError, ValidationTypeEnum, WorkflowValidationError, @@ -97,6 +95,7 @@ from ._validation import ( from ._viz import WorkflowViz from ._workflow import Workflow, WorkflowBuilder, WorkflowRunResult from ._workflow_context import WorkflowContext +from ._workflow_executor import WorkflowExecutor __all__ = [ "DEFAULT_MAX_ITERATIONS", @@ -122,7 +121,6 @@ __all__ = [ "FileCheckpointStorage", "FunctionExecutor", "GraphConnectivityError", - "HandlerOutputAnnotationError", "InMemoryCheckpointStorage", "InProcRunnerContext", "MagenticAgentDeltaEvent", @@ -166,7 +164,6 @@ __all__ = [ "WorkflowAgent", "WorkflowBuilder", "WorkflowCheckpoint", - "WorkflowCompletedEvent", "WorkflowContext", "WorkflowErrorDetails", "WorkflowEvent", @@ -174,6 +171,7 @@ __all__ = [ "WorkflowExecutor", "WorkflowFailedEvent", "WorkflowLifecycleEvent", + "WorkflowOutputEvent", "WorkflowRunResult", "WorkflowRunState", "WorkflowStartedEvent", diff --git a/python/packages/main/agent_framework/_workflow/_agent.py b/python/packages/main/agent_framework/_workflow/_agent.py index 1210611b67..b210b20e95 100644 --- a/python/packages/main/agent_framework/_workflow/_agent.py +++ b/python/packages/main/agent_framework/_workflow/_agent.py @@ -79,7 +79,7 @@ class WorkflowAgent(BaseAgent): except KeyError as exc: # Defensive: workflow lacks a configured entry point raise ValueError("Workflow's start executor is not defined.") from exc - if not start_executor.can_handle_type(list[ChatMessage]): + if list[ChatMessage] not in start_executor.input_types: raise ValueError("Workflow's start executor cannot handle list[ChatMessage]") super().__init__(id=id, name=name, description=description, **kwargs) diff --git a/python/packages/main/agent_framework/_workflow/_concurrent.py b/python/packages/main/agent_framework/_workflow/_concurrent.py index 42b567aaab..a33f425fee 100644 --- a/python/packages/main/agent_framework/_workflow/_concurrent.py +++ b/python/packages/main/agent_framework/_workflow/_concurrent.py @@ -6,9 +6,10 @@ import logging from collections.abc import Callable, Sequence from typing import Any +from typing_extensions import Never + from agent_framework import AgentProtocol, ChatMessage, Role -from ._events import WorkflowCompletedEvent from ._executor import AgentExecutorRequest, AgentExecutorResponse, Executor, handler from ._workflow import Workflow, WorkflowBuilder from ._workflow_context import WorkflowContext @@ -25,12 +26,14 @@ parallel workflow with: Notes: - Participants should be AgentProtocol instances or Executors. - A custom aggregator can be provided as: - - an Executor instance (it should handle list[AgentExecutorResponse] and add a WorkflowCompletedEvent), or + - an Executor instance (it should handle list[AgentExecutorResponse], + yield output), or - a callback function with signature: def cb(results: list[AgentExecutorResponse]) -> Any | None - def cb(results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> Any | None - If the callback returns a non-None value, it is sent as the data of a WorkflowCompletedEvent. - If it returns None, the callback may have already emitted a completion event via ctx. + def cb(results: list[AgentExecutorResponse], ctx: WorkflowContext) -> Any | None + The callback is wrapped in _CallbackAggregator. + If the callback returns a non-None value, _CallbackAggregator yields that as output. + If it returns None, the callback may have already yielded an output via ctx, so no further action is taken. """ @@ -70,7 +73,9 @@ class _AggregateAgentConversations(Executor): """ @handler - async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> None: + async def aggregate( + self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, list[ChatMessage]] + ) -> None: if not results: logger.error("Concurrent aggregator received empty results list") raise ValueError("Aggregation failed: no results provided") @@ -128,7 +133,7 @@ class _AggregateAgentConversations(Executor): logger.warning("No user prompt found in any conversation; emitting assistants only") output.extend(assistant_replies) - await ctx.add_event(WorkflowCompletedEvent(data=output)) + await ctx.yield_output(output) class _CallbackAggregator(Executor): @@ -141,7 +146,7 @@ class _CallbackAggregator(Executor): Notes: - Async callbacks are awaited directly. - Sync callbacks are executed via asyncio.to_thread to avoid blocking the event loop. - - If the callback returns a non-None value, it is wrapped in a WorkflowCompletedEvent. + - If the callback returns a non-None value, it is yielded as an output. """ def __init__(self, callback: Callable[..., Any], id: str | None = None) -> None: @@ -153,7 +158,7 @@ class _CallbackAggregator(Executor): self._param_count = len(inspect.signature(callback).parameters) @handler - async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> None: + async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, Any]) -> None: # Call according to provided signature, always non-blocking for sync callbacks if self._param_count >= 2: if inspect.iscoroutinefunction(self._callback): @@ -168,7 +173,7 @@ class _CallbackAggregator(Executor): # If the callback returned a value, finalize the workflow with it if ret is not None: - await ctx.add_event(WorkflowCompletedEvent(ret)) + await ctx.yield_output(ret) class ConcurrentBuilder: @@ -187,8 +192,7 @@ class ConcurrentBuilder: # Custom aggregator via callback (sync or async). The callback receives - # list[AgentExecutorResponse] and its return value becomes - # WorkflowCompletedEvent.data + # list[AgentExecutorResponse] and its return value becomes the workflow's output. def summarize(results): return " | ".join(r.agent_run_response.messages[-1].text for r in results) @@ -245,13 +249,13 @@ class ConcurrentBuilder: def with_aggregator(self, aggregator: Executor | Callable[..., Any]) -> "ConcurrentBuilder": r"""Override the default aggregator with an Executor or a callback. - - Executor: must handle `list[AgentExecutorResponse]` and add a - `WorkflowCompletedEvent` to the context. + - Executor: must handle `list[AgentExecutorResponse]` and + yield output using `ctx.yield_output(...)` and add a + output and the workflow becomes idle. - Callback: sync or async callable with one of the signatures: `(results: list[AgentExecutorResponse]) -> Any | None` or - `(results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> Any | None`. - If the callback returns a non-None value, it becomes the - `WorkflowCompletedEvent.data`. + `(results: list[AgentExecutorResponse], ctx: WorkflowContext) -> Any | None`. + If the callback returns a non-None value, it becomes the workflow's output. Example: ```python @@ -277,7 +281,7 @@ class ConcurrentBuilder: Wiring pattern: - Dispatcher (internal) fans out the input to all `participants` - Fan-in aggregator collects `AgentExecutorResponse` objects - - Aggregator emits a `WorkflowCompletedEvent` with either: + - Aggregator yields output and the workflow becomes idle. The output is either: - list[ChatMessage] (default aggregator: one user + one assistant per agent) - custom payload from the provided callback/executor diff --git a/python/packages/main/agent_framework/_workflow/_edge_runner.py b/python/packages/main/agent_framework/_workflow/_edge_runner.py index 7069b10baa..e186a22b8a 100644 --- a/python/packages/main/agent_framework/_workflow/_edge_runner.py +++ b/python/packages/main/agent_framework/_workflow/_edge_runner.py @@ -11,7 +11,6 @@ from ._edge import Edge, EdgeGroup, FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeG from ._executor import Executor from ._runner_context import Message, RunnerContext from ._shared_state import SharedState -from ._workflow_context import WorkflowContext logger = logging.getLogger(__name__) @@ -64,19 +63,16 @@ class EdgeRunner(ABC): target_executor = self._executors[target_id] - # Create WorkflowContext with trace contexts from message - workflow_context: WorkflowContext[Any] = WorkflowContext( - target_id, - source_ids, - shared_state, - ctx, - trace_contexts=message.trace_contexts, # Pass trace contexts to WorkflowContext + # Execute with trace context parameters + await target_executor.execute( + message.data, + source_ids, # source_executor_ids + shared_state, # shared_state + ctx, # runner_context + trace_contexts=message.trace_contexts, # Pass trace contexts source_span_ids=message.source_span_ids, # Pass source span IDs for linking ) - # Execute with trace context in WorkflowContext - await target_executor.execute(message.data, workflow_context) - class SingleEdgeRunner(EdgeRunner): """Runner for single edge groups.""" diff --git a/python/packages/main/agent_framework/_workflow/_events.py b/python/packages/main/agent_framework/_workflow/_events.py index d8c16aeb4c..90a2f1f174 100644 --- a/python/packages/main/agent_framework/_workflow/_events.py +++ b/python/packages/main/agent_framework/_workflow/_events.py @@ -66,16 +66,6 @@ class WorkflowStartedEvent(WorkflowEvent): ... -class WorkflowCompletedEvent(WorkflowEvent): - """Built-in lifecycle event emitted when a workflow run completes successfully. - - Unlike the framework-only `WorkflowLifecycleEvent` union, this event can be - emitted by developer-provided executors to return final workflow output. - """ - - ... - - class WorkflowWarningEvent(WorkflowEvent): """Executor-origin event signaling a warning surfaced by user code.""" @@ -120,16 +110,14 @@ class WorkflowRunState(str, Enum): request-for-information operations are outstanding. New work may still be scheduled while requests are in flight. - - IDLE: The workflow is quiescent with no outstanding requests, but has - not yet emitted a terminal result. Rare in practice but provided for - orchestration integrations that distinguish a quiescent state. + - IDLE: The workflow is quiescent with no outstanding requests and no more + work to do. This is the normal terminal state for workflows that have + finished executing, potentially having produced outputs along the way. - IDLE_WITH_PENDING_REQUESTS: The workflow is paused awaiting external input (e.g., emitted a `RequestInfoEvent`). This is a non-terminal state; the workflow can resume when responses are supplied. - - COMPLETED: Normal terminal state indicating successful completion. - - FAILED: Terminal state indicating an error surfaced. Accompanied by a `WorkflowFailedEvent` with structured error details. @@ -143,7 +131,6 @@ class WorkflowRunState(str, Enum): IN_PROGRESS_PENDING_REQUESTS = "IN_PROGRESS_PENDING_REQUESTS" # Active execution with outstanding requests IDLE = "IDLE" # No active work and no outstanding requests IDLE_WITH_PENDING_REQUESTS = "IDLE_WITH_PENDING_REQUESTS" # Paused awaiting external responses - COMPLETED = "COMPLETED" # Finished successfully FAILED = "FAILED" # Finished with an error CANCELLED = "CANCELLED" # Finished due to cancellation @@ -250,6 +237,28 @@ class RequestInfoEvent(WorkflowEvent): ) +class WorkflowOutputEvent(WorkflowEvent): + """Event triggered when a workflow executor yields output.""" + + def __init__( + self, + data: Any, + source_executor_id: str, + ): + """Initialize the workflow output event. + + Args: + data: The output yielded by the executor. + source_executor_id: ID of the executor that yielded the output. + """ + super().__init__(data) + self.source_executor_id = source_executor_id + + def __repr__(self) -> str: + """Return a string representation of the workflow output event.""" + return f"{self.__class__.__name__}(data={self.data}, source_executor_id={self.source_executor_id})" + + class ExecutorEvent(WorkflowEvent): """Base class for executor events.""" diff --git a/python/packages/main/agent_framework/_workflow/_executor.py b/python/packages/main/agent_framework/_workflow/_executor.py index 71d416f4f3..a9a77c3c95 100644 --- a/python/packages/main/agent_framework/_workflow/_executor.py +++ b/python/packages/main/agent_framework/_workflow/_executor.py @@ -9,8 +9,7 @@ import uuid from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import asdict, dataclass, field, fields, is_dataclass from textwrap import shorten -from types import UnionType -from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, Union, cast, get_args, get_origin, overload +from typing import Any, ClassVar, Generic, TypeVar, cast, overload from pydantic import Field @@ -24,16 +23,16 @@ from ._events import ( AgentRunEvent, AgentRunUpdateEvent, ExecutorCompletedEvent, + ExecutorFailedEvent, ExecutorInvokedEvent, RequestInfoEvent, + WorkflowErrorDetails, _framework_event_origin, # type: ignore[reportPrivateUsage] ) -from ._runner_context import _decode_checkpoint_value # type: ignore[reportPrivateUsage] +from ._runner_context import Message, RunnerContext, _decode_checkpoint_value +from ._shared_state import SharedState from ._typing_utils import is_instance_of -from ._workflow_context import WorkflowContext - -if TYPE_CHECKING: - from ._workflow import Workflow +from ._workflow_context import WorkflowContext, validate_function_signature logger = logging.getLogger(__name__) # region Executor @@ -65,7 +64,141 @@ class WorkflowCheckpointSummary: class Executor(AFBaseModel): - """An executor is a component that processes messages in a workflow.""" + """Base class for all workflow executors that process messages and perform computations. + + ## Overview + Executors are the fundamental building blocks of workflows, representing individual processing + units that receive messages, perform operations, and produce outputs. Each executor is uniquely + identified and can handle specific message types through decorated handler methods. + + ## Type System + Executors have a rich type system that defines their capabilities: + + ### Input Types + The types of messages an executor can process, discovered from handler method signatures: + ```python + class MyExecutor(Executor): + @handler + async def handle_string(self, message: str, ctx: WorkflowContext) -> None: + # This executor can handle 'str' input types + ``` + Access via the `input_types` property. + + ### Output Types + The types of messages an executor can send to other executors via `ctx.send_message()`: + ```python + class MyExecutor(Executor): + @handler + async def handle_data(self, message: str, ctx: WorkflowContext[int | bool]) -> None: + # This executor can send 'int' or 'bool' messages + ``` + Access via the `output_types` property. + + ### Workflow Output Types + The types of data an executor can emit as workflow-level outputs via `ctx.yield_output()`: + ```python + class MyExecutor(Executor): + @handler + async def process(self, message: str, ctx: WorkflowContext[int, str]) -> None: + # Can send 'int' messages AND yield 'str' workflow outputs + ``` + Access via the `workflow_output_types` property. + + ### Request Types + The types of sub-workflow requests an executor can intercept via `@intercepts_request`: + ```python + class ParentExecutor(Executor): + @intercepts_request + async def handle_request( + self, + request: MyRequest, + ctx: WorkflowContext, + ) -> RequestResponse[MyRequest, str]: + # Can intercept 'MyRequest' from sub-workflows + ``` + Access via the `request_types` property. + + ## Handler Discovery + Executors discover their capabilities through decorated methods: + + ### @handler Decorator + Marks methods that process incoming messages: + ```python + class MyExecutor(Executor): + @handler + async def handle_text(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message.upper()) + ``` + + ### @intercepts_request Decorator + Marks methods that intercept sub-workflow requests: + ```python + class ParentExecutor(Executor): + @intercepts_request + async def check_domain( + self, request: DomainRequest, ctx: WorkflowContext + ) -> RequestResponse[DomainRequest, bool]: + if self.is_allowed(request.domain): + return RequestResponse.handled(True) + return RequestResponse.forward() + ``` + + ## Context Types + Handler methods receive different WorkflowContext variants based on their type annotations: + + ### WorkflowContext (no type parameters) + For handlers that only perform side effects without sending messages or yielding outputs: + ```python + class LoggingExecutor(Executor): + @handler + async def log_message(self, msg: str, ctx: WorkflowContext) -> None: + print(f"Received: {msg}") # Only logging, no outputs + ``` + + ### WorkflowContext[T_Out] + Enables sending messages of type T_Out via `ctx.send_message()`: + ```python + class ProcessorExecutor(Executor): + @handler + async def handler(self, msg: str, ctx: WorkflowContext[int]) -> None: + await ctx.send_message(42) # Can send int messages + ``` + + ### WorkflowContext[T_Out, T_W_Out] + Enables both sending messages (T_Out) and yielding workflow outputs (T_W_Out): + ```python + class DualOutputExecutor(Executor): + @handler + async def handler(self, msg: str, ctx: WorkflowContext[int, str]) -> None: + await ctx.send_message(42) # Send int message + await ctx.yield_output("done") # Yield str workflow output + ``` + + ## Function Executors + Simple functions can be converted to executors using the `@executor` decorator: + ```python + @executor + async def process_text(text: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(text.upper()) + + + # Or with custom ID: + @executor(id="text_processor") + def sync_process(text: str, ctx: WorkflowContext[str]) -> None: + ctx.send_message(text.lower()) # Sync functions run in thread pool + ``` + + ## Sub-workflow Composition + Executors can contain sub-workflows using WorkflowExecutor. Sub-workflows can make requests + that parent workflows can intercept. See WorkflowExecutor documentation for details on + workflow composition patterns and request/response handling. + + ## Implementation Notes + - Do not call `execute()` directly - it's invoked by the workflow engine + - Do not override `execute()` - define handlers using decorators instead + - Each executor must have at least one `@handler` or `@intercepts_request` method + - Handler method signatures are validated at initialization time + """ # Provide a default so static analyzers (e.g., pyright) don't require passing `id`. # Runtime still sets a concrete value in __init__. @@ -92,9 +225,9 @@ class Executor(AFBaseModel): super().__init__(**kwargs) - self._handlers: dict[type, Callable[[Any, WorkflowContext[Any]], Any]] = {} + self._handlers: dict[type, Callable[[Any, WorkflowContext[Any, Any]], Awaitable[None]]] = {} self._request_interceptors: dict[type | str, list[dict[str, Any]]] = {} - self._instance_handler_specs: list[dict[str, Any]] = [] + self._handler_specs: list[dict[str, Any]] = [] self._discover_handlers() if not self._handlers and not self._request_interceptors: @@ -104,24 +237,34 @@ class Executor(AFBaseModel): "or @intercepts_request decorator." ) - async def execute(self, message: Any, context: WorkflowContext[Any]) -> None: - """Execute the executor with a given message and context. + async def execute( + self, + message: Any, + source_executor_ids: list[str], + shared_state: SharedState, + runner_context: RunnerContext, + trace_contexts: list[dict[str, str]] | None = None, + source_span_ids: list[str] | None = None, + ) -> None: + """Execute the executor with a given message and context parameters. + + - Do not call this method directly - it is invoked by the workflow engine. + - Do not override this method. Instead, define handlers using @handler decorator. Args: message: The message to be processed by the executor. - context: The workflow context in which the executor operates. + source_executor_ids: The IDs of the source executors that sent messages to this executor. + shared_state: The shared state for the workflow. + runner_context: The runner context that provides methods to send messages and events. + trace_contexts: Optional trace contexts from multiple sources for OpenTelemetry propagation. + source_span_ids: Optional source span IDs from multiple sources for linking. Returns: An awaitable that resolves to the result of the execution. """ # Create processing span for tracing (gracefully handles disabled tracing) - source_trace_contexts = getattr(context, "_trace_contexts", None) - source_span_ids = getattr(context, "_source_span_ids", None) - # Handle case where Message wrapper is passed instead of raw data - from ._runner_context import Message - if isinstance(message, Message): message = message.data @@ -129,38 +272,37 @@ class Executor(AFBaseModel): self.id, self.__class__.__name__, type(message).__name__, - source_trace_contexts=source_trace_contexts, + source_trace_contexts=trace_contexts, source_span_ids=source_span_ids, ): - # Lazy registration for SubWorkflowRequestInfo if we have interceptors - if self._request_interceptors and message.__class__.__name__ == "SubWorkflowRequestInfo": - # Directly handle SubWorkflowRequestInfo - with _framework_event_origin(): - invoke_event = ExecutorInvokedEvent(self.id) - await context.add_event(invoke_event) - try: - await self._handle_sub_workflow_request(message, context) - except Exception as exc: - # Surface structured executor failure before propagating - from ._events import ExecutorFailedEvent, WorkflowErrorDetails - - with _framework_event_origin(): - failure_event = ExecutorFailedEvent(self.id, WorkflowErrorDetails.from_exception(exc)) - await context.add_event(failure_event) - raise - with _framework_event_origin(): - completed_event = ExecutorCompletedEvent(self.id) - await context.add_event(completed_event) - return - - handler: Callable[[Any, WorkflowContext[Any]], Any] | None = None + # Find the handler and handler spec that matches the message type. + # This includes the self._handle_sub_workflow_request handler for SubWorkflowRequestInfo. + handler: Callable[[Any, WorkflowContext[Any, Any]], Awaitable[None]] | None = None + ctx_annotation = None for message_type in self._handlers: if is_instance_of(message, message_type): handler = self._handlers[message_type] + # Find the corresponding handler spec for context annotation + for spec in self._handler_specs: + if spec.get("message_type") == message_type: + ctx_annotation = spec.get("ctx_annotation") + break break if handler is None: raise RuntimeError(f"Executor {self.__class__.__name__} cannot handle message of type {type(message)}.") + + # Create the appropriate WorkflowContext based on handler specs + context = self._create_context_for_handler( + source_executor_ids=source_executor_ids, + shared_state=shared_state, + runner_context=runner_context, + ctx_annotation=ctx_annotation, + trace_contexts=trace_contexts, + source_span_ids=source_span_ids, + ) + + # Invoke the handler with the message and context with _framework_event_origin(): invoke_event = ExecutorInvokedEvent(self.id) await context.add_event(invoke_event) @@ -168,8 +310,6 @@ class Executor(AFBaseModel): await handler(message, context) except Exception as exc: # Surface structured executor failure before propagating - from ._events import ExecutorFailedEvent, WorkflowErrorDetails - with _framework_event_origin(): failure_event = ExecutorFailedEvent(self.id, WorkflowErrorDetails.from_exception(exc)) await context.add_event(failure_event) @@ -178,6 +318,38 @@ class Executor(AFBaseModel): completed_event = ExecutorCompletedEvent(self.id) await context.add_event(completed_event) + def _create_context_for_handler( + self, + source_executor_ids: list[str], + shared_state: SharedState, + runner_context: RunnerContext, + ctx_annotation: Any, + trace_contexts: list[dict[str, str]] | None = None, + source_span_ids: list[str] | None = None, + ) -> WorkflowContext[Any]: + """Create the appropriate WorkflowContext based on the handler's context annotation. + + Args: + source_executor_ids: The IDs of the source executors that sent messages to this executor. + shared_state: The shared state for the workflow. + runner_context: The runner context that provides methods to send messages and events. + ctx_annotation: The context annotation from the handler spec to determine which context type to create. + trace_contexts: Optional trace contexts from multiple sources for OpenTelemetry propagation. + source_span_ids: Optional source span IDs from multiple sources for linking. + + Returns: + WorkflowContext[Any] based on the handler's context annotation. + """ + # Create WorkflowContext + return WorkflowContext( + executor_id=self.id, + source_executor_ids=source_executor_ids, + shared_state=shared_state, + runner_context=runner_context, + trace_contexts=trace_contexts, + source_span_ids=source_span_ids, + ) + def _discover_handlers(self) -> None: """Discover message handlers and request interceptors in the executor class.""" # Use __class__.__dict__ to avoid accessing pydantic's dynamic attributes @@ -195,10 +367,30 @@ class Executor(AFBaseModel): if self._handlers.get(message_type) is not None: raise ValueError(f"Duplicate handler for type {message_type} in {self.__class__.__name__}") + + # RequestInfoExecutor is allowed to handle SubWorkflowRequestInfo directly + # but other executors should use @intercepts_request + if message_type is SubWorkflowRequestInfo and not isinstance(self, RequestInfoExecutor): + raise ValueError( + f"Executor {self.__class__.__name__} cannot define a handler for " + "SubWorkflowRequestInfo directly. " + f"Use @intercepts_request decorator to intercept sub-workflow requests." + ) + # Get the bound method bound_method = getattr(self, attr_name) self._handlers[message_type] = bound_method + # Add to unified handler specs list + self._handler_specs.append({ + "name": handler_spec["name"], + "message_type": message_type, + "output_types": handler_spec.get("output_types", []), + "workflow_output_types": handler_spec.get("workflow_output_types", []), + "ctx_annotation": handler_spec.get("ctx_annotation"), + "source": "class_method", # Distinguish from instance handlers if needed + }) + # Discover @intercepts_request methods if hasattr(attr, "_intercepts_request"): # Get the bound method for interceptors @@ -212,98 +404,149 @@ class Executor(AFBaseModel): if request_type not in self._request_interceptors: self._request_interceptors[request_type] = [] self._request_interceptors[request_type].append(interceptor_info) + + # We need to register the handler for SubWorkflowRequestInfo + # if not already registered. + if SubWorkflowRequestInfo not in self._handlers: + self._handlers[SubWorkflowRequestInfo] = self._handle_sub_workflow_request + self._handler_specs.append({ + "name": attr_name, + "message_type": SubWorkflowRequestInfo, + "output_types": [SubWorkflowResponse, SubWorkflowRequestInfo, RequestInfoMessage], + "workflow_output_types": [], + "ctx_annotation": WorkflowContext[ + SubWorkflowRequestInfo | SubWorkflowResponse | RequestInfoMessage + ], + "source": "class_method", + }) except AttributeError: # Skip attributes that may not be accessible continue - def _register_sub_workflow_handler(self) -> None: - """Register automatic handler for SubWorkflowRequestInfo messages.""" - # We need to use a string reference until the class is defined - # This will be resolved later when the class is actually used - pass # Will be registered lazily when needed - async def _handle_sub_workflow_request( self, request: "SubWorkflowRequestInfo", - ctx: WorkflowContext[Any], + ctx: WorkflowContext["SubWorkflowResponse | SubWorkflowRequestInfo | RequestInfoMessage"], ) -> None: """Automatic routing to @intercepts_request methods. This is only active for executors that have @intercepts_request methods. """ - # Try to match against registered interceptors + # Try to find a matching interceptor for the request and execute it for request_type, interceptor_list in self._request_interceptors.items(): - matched = False - - # Check type matching - if isinstance(request_type, type) and is_instance_of(request.data, request_type): - matched = True - elif ( - isinstance(request_type, str) - and hasattr(request.data, "__class__") - and request.data.__class__.__name__ == request_type - ): - # String matching - could check against type name or other attributes - matched = True - - if matched: - # Check each interceptor in the list for this request type + if self._does_request_match_type(request.data, request_type): for interceptor_info in interceptor_list: - # Check workflow scope if specified - from_workflow = interceptor_info["from_workflow"] - if from_workflow and request.sub_workflow_id != from_workflow: - continue # Skip this interceptor, wrong workflow - - # Check additional condition - condition = interceptor_info["condition"] - if condition and not condition(request): - continue - - # Call the interceptor method - method = interceptor_info["method"] - response = await method(request.data, ctx) - - # Check if interceptor handled it or needs to forward - if isinstance(response, RequestResponse): - # Add automatic correlation info to the response - correlated_response = RequestResponse[RequestInfoMessage, Any].with_correlation( - response, - request.data, - request.request_id, + if self._does_interceptor_apply_to_request(request, interceptor_info): + logger.debug( + f"Executor {self.id} intercepting request {request.request_id} " + f"of type {type(request.data).__name__} from sub-workflow {request.sub_workflow_id}" ) - - if correlated_response.is_handled: - # Send response back to sub-workflow - await ctx.send_message( - SubWorkflowResponse( - request_id=request.request_id, - data=correlated_response.data, - ), - target_id=request.sub_workflow_id, - ) - else: - # Forward WITH CONTEXT PRESERVED - # Update the data if interceptor provided a modified request - if correlated_response.forward_request: - request.data = correlated_response.forward_request - - # Send the inner request to RequestInfoExecutor to create external request - await ctx.send_message(request) - else: - # Legacy support: direct return means handled - await ctx.send_message( - SubWorkflowResponse( - request_id=request.request_id, - data=response, - ), - target_id=request.sub_workflow_id, - ) - return - + await self._execute_interceptor(request, interceptor_info, ctx) + return # Only the first matching interceptor is executed + logger.debug( + f"Executor {self.id} has no matching interceptor for request {request.request_id} " + f"of type {type(request.data).__name__} from sub-workflow {request.sub_workflow_id}; forwarding " + "inner request to RequestInfoExecutor for external handling." + ) # No interceptor found - forward inner request to RequestInfoExecutor - # This sends the original request to RequestInfoExecutor await ctx.send_message(request.data) + def _does_request_match_type(self, request_data: Any, request_type: type | str) -> bool: + """Check if request data matches the expected type. + + Args: + request_data: The request data to check + request_type: The type to match against (can be a type or string) + + Returns: + True if the request matches the type, False otherwise. + """ + if isinstance(request_type, type): + return is_instance_of(request_data, request_type) + + return ( + isinstance(request_type, str) + and hasattr(request_data, "__class__") + and request_data.__class__.__name__ == request_type + ) + + def _does_interceptor_apply_to_request( + self, request: "SubWorkflowRequestInfo", interceptor_info: dict[str, Any] + ) -> bool: + """Check if an interceptor applies to the given request. + + Args: + request: The sub-workflow request + interceptor_info: Information about the interceptor + + Returns: + True if the interceptor should handle this request, False otherwise. + """ + # Check workflow scope if specified + from_workflow = interceptor_info["from_workflow"] + if from_workflow and request.sub_workflow_id != from_workflow: + return False + + # Check additional condition + condition = interceptor_info["condition"] + return not (condition and not condition(request)) + + async def _execute_interceptor( + self, + request: "SubWorkflowRequestInfo", + interceptor_info: dict[str, Any], + ctx: WorkflowContext["SubWorkflowResponse | SubWorkflowRequestInfo | RequestInfoMessage"], + ) -> None: + """Execute a single interceptor method. + + Args: + request: The sub-workflow request + interceptor_info: Information about the interceptor method + ctx: The workflow context + """ + method = interceptor_info["method"] + response = await method(request.data, ctx) + + if not isinstance(response, RequestResponse): + raise RuntimeError( + f"Interceptor method {method.__name__} must return RequestResponse, got {type(response)}" + ) + + # Add automatic correlation info to the response + correlated_response = RequestResponse[RequestInfoMessage, Any].with_correlation( + response, + request.data, + request.request_id, + ) + + if correlated_response.is_handled: + logger.debug( + f"Executor {self.id}'s interceptor handled request {request.request_id} " + f"of type {type(request.data).__name__} from sub-workflow {request.sub_workflow_id}. " + f"Sending response back to sub-workflow." + ) + # Send response back to sub-workflow that made the request + await ctx.send_message( + SubWorkflowResponse( + request_id=request.request_id, + data=correlated_response.data, + ), + target_id=request.sub_workflow_id, + ) + else: + logger.debug( + f"Executor {self.id}'s interceptor did not handle request {request.request_id} " + f"of type {type(request.data).__name__} from sub-workflow {request.sub_workflow_id}. " + "Forwarding to RequestInfoExecutor for external handling." + ) + # Forward the request with potentially modified data + # Update the data if interceptor provided a modified request + if correlated_response.forward_request: + request.data = correlated_response.forward_request + + # Send the inner request to RequestInfoExecutor to create external request + await ctx.send_message(request) + def can_handle(self, message: Any) -> bool: """Check if the executor can handle a given message type. @@ -315,13 +558,14 @@ class Executor(AFBaseModel): """ return any(is_instance_of(message, message_type) for message_type in self._handlers) - def register_instance_handler( + def _register_instance_handler( self, name: str, func: Callable[[Any, WorkflowContext[Any]], Awaitable[Any]], message_type: type, ctx_annotation: Any, output_types: list[type], + workflow_output_types: list[type], ) -> None: """Register a handler at instance level. @@ -330,29 +574,84 @@ class Executor(AFBaseModel): func: The async handler function to register message_type: Type of message this handler processes ctx_annotation: The WorkflowContext[T] annotation from the function - output_types: List of output types inferred from ctx_annotation + output_types: List of output types for send_message() + workflow_output_types: List of workflow output types for yield_output() """ + if message_type is SubWorkflowRequestInfo: + raise ValueError( + f"Executor {self.__class__.__name__} cannot define a handler for " + "SubWorkflowRequestInfo directly. " + f"Use @intercepts_request decorator to intercept sub-workflow requests." + ) + if message_type in self._handlers: raise ValueError(f"Handler for type {message_type} already registered in {self.__class__.__name__}") self._handlers[message_type] = func - self._instance_handler_specs.append({ + self._handler_specs.append({ "name": name, "message_type": message_type, "ctx_annotation": ctx_annotation, "output_types": output_types, + "workflow_output_types": workflow_output_types, + "source": "instance_method", # Distinguish from class handlers if needed }) - def can_handle_type(self, message_type: type[Any]) -> bool: - """Check if the executor can handle a given message type. - - Args: - message_type: The message type to check. + @property + def input_types(self) -> list[type[Any]]: + """Get the list of input types that this executor can handle. Returns: - True if the executor can handle the message type, False otherwise. + A list of the message types that this executor's handlers can process. """ - return message_type in self._handlers + return list(self._handlers.keys()) + + @property + def output_types(self) -> list[type[Any]]: + """Get the list of output types that this executor can produce via send_message(). + + Returns: + A list of the output types inferred from the handlers' WorkflowContext[T] annotations. + """ + output_types: set[type[Any]] = set() + + # Collect output types from all handlers + for handler_spec in self._handler_specs: + handler_output_types = handler_spec.get("output_types", []) + output_types.update(handler_output_types) + + return list(output_types) + + @property + def workflow_output_types(self) -> list[type[Any]]: + """Get the list of workflow output types that this executor can produce via yield_output(). + + Returns: + A list of the workflow output types inferred from handlers' WorkflowContext[T, U] annotations. + """ + output_types: set[type[Any]] = set() + + # Collect workflow output types from all handlers + for handler_spec in self._handler_specs: + handler_workflow_output_types = handler_spec.get("workflow_output_types", []) + output_types.update(handler_workflow_output_types) + + return list(output_types) + + @property + def request_types(self) -> list[type[Any]]: + """Get the list of request types that this executor can intercept via @intercepts_request. + + Returns: + A list of the request types that this executor's interceptors can handle. + """ + request_types: list[type[Any]] = [] + + for request_type in self._request_interceptors: + if isinstance(request_type, type): + request_types.append(request_type) + + return request_types # endregion: Executor @@ -361,15 +660,16 @@ class Executor(AFBaseModel): ExecutorT = TypeVar("ExecutorT", bound="Executor") +ContextT = TypeVar("ContextT", bound="WorkflowContext[Any, Any]") def handler( - func: Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]], + func: Callable[[ExecutorT, Any, ContextT], Awaitable[Any]], ) -> ( - Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]] + Callable[[ExecutorT, Any, ContextT], Awaitable[Any]] | Callable[ - [Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]], - Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]], + [Callable[[ExecutorT, Any, ContextT], Awaitable[Any]]], + Callable[[ExecutorT, Any, ContextT], Awaitable[Any]], ] ): """Decorator to register a handler for an executor. @@ -390,73 +690,19 @@ def handler( ... """ - def _infer_output_types_from_ctx_annotation(ctx_annotation: Any) -> list[type[Any]]: - """Infer output types list from the WorkflowContext generic parameter. - - Examples: - - WorkflowContext[str] -> [str] - - WorkflowContext[str | int] -> [str, int] - - WorkflowContext[Union[str, int]] -> [str, int] - - WorkflowContext -> [] (unknown) - """ - # If no annotation or not parameterized, return empty list - try: - origin = get_origin(ctx_annotation) - except Exception: - origin = None - - # If annotation is unsubscripted WorkflowContext, nothing to infer - if origin is None: - # Might be the class itself or Any; try simple check by name to avoid import cycles - return [] - - # Expecting WorkflowContext[T] - if origin is not WorkflowContext: - return [] - - args = get_args(ctx_annotation) - if not args: - return [] - - t = args[0] - # If t is a Union, flatten it - t_origin = get_origin(t) - # If Any, treat as unknown -> no output types inferred - if t is Any: - return [] - - if t_origin in (Union, UnionType): - # Return all union args as-is (may include generic aliases like list[str]) - return [arg for arg in get_args(t) if arg is not Any and arg is not type(None)] - - # Single concrete or generic alias type (e.g., str, int, list[str]) - if t is Any or t is type(None): - return [] - return [t] - def decorator( - func: Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]], - ) -> Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]: - # Extract the message type from a handler function. + func: Callable[[ExecutorT, Any, ContextT], Awaitable[Any]], + ) -> Callable[[ExecutorT, Any, ContextT], Awaitable[Any]]: + # Extract the message type and validate using unified validation + message_type, ctx_annotation, inferred_output_types, inferred_workflow_output_types = ( + validate_function_signature(func, "Handler method") + ) + + # Get signature for preservation sig = inspect.signature(func) - params = list(sig.parameters.values()) - - if len(params) != 3: # self, message, ctx - raise ValueError(f"Handler must have exactly 3 parameters, got {len(params)}") - - message_type = params[1].annotation - if message_type is inspect.Parameter.empty: - raise ValueError("Handler's second parameter must have a type annotation") - - ctx_annotation = params[2].annotation - if ctx_annotation is inspect.Parameter.empty: - # Allow missing ctx annotation, but we can't infer outputs - inferred_output_types: list[type[Any]] = [] - else: - inferred_output_types = _infer_output_types_from_ctx_annotation(ctx_annotation) @functools.wraps(func) - async def wrapper(self: ExecutorT, message: Any, ctx: WorkflowContext[Any]) -> Any: + async def wrapper(self: ExecutorT, message: Any, ctx: ContextT) -> Any: """Wrapper function to call the handler.""" return await func(self, message, ctx) @@ -467,8 +713,10 @@ def handler( wrapper._handler_spec = { # type: ignore "name": func.__name__, "message_type": message_type, - # Keep output_types in spec for validators, inferred from WorkflowContext[T] + # Keep output_types and workflow_output_types in spec for validators "output_types": inferred_output_types, + "workflow_output_types": inferred_workflow_output_types, + "ctx_annotation": ctx_annotation, } return wrapper @@ -576,16 +824,15 @@ class SubWorkflowResponse: # TypeVar for request type that must be a RequestInfoMessage subclass RequestInfoMessageT = TypeVar("RequestInfoMessageT", bound="RequestInfoMessage") -# Type alias for interceptor functions -InterceptorFunc = Callable[ - [Any, RequestInfoMessageT, WorkflowContext[Any]], Awaitable[RequestResponse[RequestInfoMessageT, Any]] -] - @overload def intercepts_request( - func: Callable[..., Any], -) -> Callable[..., Any]: ... + func: Callable[ + [Any, RequestInfoMessageT, WorkflowContext[Any, Any]], Awaitable[RequestResponse[RequestInfoMessageT, Any]] + ], +) -> Callable[ + [Any, RequestInfoMessageT, WorkflowContext[Any, Any]], Awaitable[RequestResponse[RequestInfoMessageT, Any]] +]: ... @overload @@ -593,7 +840,16 @@ def intercepts_request( *, from_workflow: str | None = None, condition: Callable[[Any], bool] | None = None, -) -> Callable[[Callable[..., Any]], Callable[..., Any]]: ... +) -> Callable[ + [ + Callable[ + [Any, RequestInfoMessageT, WorkflowContext[Any, Any]], Awaitable[RequestResponse[RequestInfoMessageT, Any]] + ] + ], + Callable[ + [Any, RequestInfoMessageT, WorkflowContext[Any, Any]], Awaitable[RequestResponse[RequestInfoMessageT, Any]] + ], +]: ... def intercepts_request( @@ -621,7 +877,7 @@ def intercepts_request( Example: @intercepts_request async def check_domain( - self, request: DomainCheckRequest, ctx: WorkflowContext[Any] + self, request: DomainCheckRequest, ctx: WorkflowContext ) -> RequestResponse[DomainCheckRequest, bool]: # Type automatically inferred as DomainCheckRequest from parameter annotation if request.domain in self.approved_domains: @@ -630,7 +886,7 @@ def intercepts_request( @intercepts_request(from_workflow="email_validator") async def handle_specific( - self, request: EmailRequest, ctx: WorkflowContext[Any] + self, request: EmailRequest, ctx: WorkflowContext ) -> RequestResponse[EmailRequest, str]: # Only intercepts EmailRequest from the "email_validator" workflow return RequestResponse.handled("handled by parent") @@ -678,7 +934,7 @@ def intercepts_request( pass @functools.wraps(func) - async def wrapper(self: Any, request: Any, ctx: WorkflowContext[Any]) -> Any: + async def wrapper(self: Any, request: RequestInfoMessage, ctx: WorkflowContext[Any, Any]) -> Any: return await func(self, request, ctx) # Add metadata for discovery - store the inferred type @@ -722,7 +978,7 @@ class RequestInfoExecutor(Executor): self._sub_workflow_contexts: dict[str, dict[str, str]] = {} @handler - async def run(self, message: RequestInfoMessage, ctx: WorkflowContext[None]) -> None: + async def run(self, message: RequestInfoMessage, ctx: WorkflowContext) -> None: """Run the RequestInfoExecutor with the given message.""" source_executor_id = ctx.get_source_executor_id() @@ -740,7 +996,7 @@ class RequestInfoExecutor(Executor): async def handle_sub_workflow_request( self, message: SubWorkflowRequestInfo, - ctx: WorkflowContext[None], + ctx: WorkflowContext, ) -> None: """Handle forwarded sub-workflow request. @@ -770,7 +1026,7 @@ class RequestInfoExecutor(Executor): self, response_data: Any, request_id: str, - ctx: WorkflowContext[Any], + ctx: WorkflowContext[SubWorkflowResponse | RequestResponse[RequestInfoMessage, Any]], ) -> None: """Handle a response to a request. @@ -815,6 +1071,17 @@ class RequestInfoExecutor(Executor): await self._clear_pending_request_snapshot(request_id, ctx) + def _register_instance_handler( + self, + name: str, + func: Callable[[Any, WorkflowContext[Any]], Awaitable[Any]], + message_type: type, + ctx_annotation: Any, + output_types: list[type], + workflow_output_types: list[type], + ) -> None: + raise NotImplementedError("Cannot register handlers on RequestInfoExecutor") + async def _record_pending_request_snapshot( self, request: RequestInfoMessage, @@ -1441,176 +1708,3 @@ class AgentExecutor(Executor): # endregion: Agent Executor - - -# region Workflow Executor - - -class WorkflowExecutor(Executor): - """An executor that runs another workflow as its execution logic. - - This executor wraps a workflow to make it behave as an executor, enabling - hierarchical workflow composition. Sub-workflows can send requests that - are intercepted by parent workflows. - """ - - workflow: "Workflow" = Field(description="The workflow to execute as a sub-workflow") - - def __init__(self, workflow: "Workflow", id: str, **kwargs: Any): - """Initialize the WorkflowExecutor. - - Args: - workflow: The workflow to execute as a sub-workflow. - id: Unique identifier for this executor. - **kwargs: Additional keyword arguments passed to the parent constructor. - """ - kwargs.update({"workflow": workflow}) - super().__init__(id, **kwargs) - - # Track pending external responses by request_id - self._pending_responses: dict[str, Any] = {} # request_id -> response_data - # Track workflow state for proper resumption - support multiple concurrent requests - self._pending_requests: dict[str, Any] = {} # request_id -> original request data - self._active_executions: int = 0 # Count of active sub-workflow executions - # Response accumulation for multiple concurrent responses - self._collected_responses: dict[str, Any] = {} # Accumulate responses - self._expected_response_count: int = 0 # Track how many responses we're waiting for - - @handler # No output_types - can send any completion data type - async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any]) -> None: - """Execute the sub-workflow with raw input data. - - This handler starts a new sub-workflow execution. When the sub-workflow - needs external information, it pauses and sends a request to the parent. - - Args: - input_data: The input data to send to the sub-workflow. - ctx: The workflow context from the parent. - """ - # Skip SubWorkflowResponse and SubWorkflowRequestInfo - they have specific handlers - if isinstance(input_data, (SubWorkflowResponse, SubWorkflowRequestInfo)): - return - - from ._events import RequestInfoEvent, WorkflowCompletedEvent - - # Track this execution - self._active_executions += 1 - - try: - # Run the sub-workflow and collect all events - events = [event async for event in self.workflow.run_stream(input_data)] - - # Count requests and initialize response tracking - request_count = 0 - for event in events: - if isinstance(event, RequestInfoEvent): - request_count += 1 - - # Initialize response accumulation for this execution - # For sequential workflows (like step_08), expect only current requests - # For parallel workflows (like step_09), expect all requests at once - self._expected_response_count = request_count - self._collected_responses = {} - - # If no requests in initial run, handle completion immediately - if request_count == 0: - self._expected_response_count = 0 - - # Process events to check for completion or requests - for event in events: - if isinstance(event, WorkflowCompletedEvent): - # Sub-workflow completed normally - send result to parent - await ctx.send_message(event.data) - self._active_executions -= 1 - return # Exit after completion - - if isinstance(event, RequestInfoEvent): - # Sub-workflow needs external information - # Track the pending request - self._pending_requests[event.request_id] = event.data - - # Wrap request with routing context and send to parent - if not isinstance(event.data, RequestInfoMessage): - raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}") - wrapped_request = SubWorkflowRequestInfo( - request_id=event.request_id, - sub_workflow_id=self.id, - data=event.data, - ) - - await ctx.send_message(wrapped_request) - # Continue processing remaining events (no return) - - except Exception as e: - from ._events import ExecutorEvent - - # Sub-workflow failed - create error event - error_event = ExecutorEvent(executor_id=self.id, data={"error": str(e), "type": "sub_workflow_error"}) - await ctx.add_event(error_event) - self._active_executions -= 1 - raise - - @handler - async def handle_response( - self, - response: SubWorkflowResponse, - ctx: WorkflowContext[Any], - ) -> None: - """Handle response from parent for a forwarded request. - - This handler accumulates responses and only resumes the sub-workflow - when all expected responses have been received. - - Args: - response: The response to a previous request. - ctx: The workflow context. - """ - # Check if we have this pending request - pending_requests = getattr(self, "_pending_requests", {}) - if response.request_id not in pending_requests: - return - - # Remove the request from pending list - pending_requests.pop(response.request_id, None) - - # Accumulate the response - self._collected_responses[response.request_id] = response.data - - # Check if we have all expected responses for current batch - if len(self._collected_responses) >= self._expected_response_count: - from ._events import RequestInfoEvent, WorkflowCompletedEvent - - # Send all collected responses to the sub-workflow - responses_to_send = dict(self._collected_responses) - self._collected_responses.clear() # Clear for next batch - - result_events = [event async for event in self.workflow.send_responses_streaming(responses_to_send)] - - # Process the result events - new_request_count = 0 - for event in result_events: - if isinstance(event, WorkflowCompletedEvent): - # Sub-workflow completed - send result to parent - await ctx.send_message(event.data) - self._active_executions -= 1 - return - if isinstance(event, RequestInfoEvent): - # Sub-workflow sent more requests - prepare for next batch - new_request_count += 1 - self._pending_requests[event.request_id] = event.data - - # Send the new request to parent - if not isinstance(event.data, RequestInfoMessage): - raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}") - wrapped_request = SubWorkflowRequestInfo( - request_id=event.request_id, - sub_workflow_id=self.id, - data=event.data, - ) - await ctx.send_message(wrapped_request) - - # Update expected count for next batch of requests - self._expected_response_count = new_request_count - - -# endregion: Workflow Executor diff --git a/python/packages/main/agent_framework/_workflow/_function_executor.py b/python/packages/main/agent_framework/_workflow/_function_executor.py index 05edc07fbd..d2e0752936 100644 --- a/python/packages/main/agent_framework/_workflow/_function_executor.py +++ b/python/packages/main/agent_framework/_workflow/_function_executor.py @@ -11,67 +11,11 @@ This module provides: """ import asyncio -import inspect from collections.abc import Awaitable, Callable -from types import UnionType -from typing import Any, Union, get_args, get_origin, overload +from typing import Any, overload from ._executor import Executor -from ._workflow_context import WorkflowContext - - -def _is_workflow_context_type(annotation: Any) -> bool: - """Check if an annotation represents WorkflowContext[T].""" - origin = get_origin(annotation) - if origin is WorkflowContext: - return True - # Also handle the case where the raw WorkflowContext class is used - return annotation is WorkflowContext - - -def _infer_output_types_from_ctx_annotation(ctx_annotation: Any) -> list[type]: - """Infer output types list from the WorkflowContext generic parameter. - - Examples: - - WorkflowContext[str] -> [str] - - WorkflowContext[str | int] -> [str, int] - - WorkflowContext[Union[str, int]] -> [str, int] - - WorkflowContext[Any] -> [] (unknown) - - WorkflowContext[None] -> [] - """ - # If no annotation or not parameterized, return empty list - try: - origin = get_origin(ctx_annotation) - except Exception: - origin = None - - # If annotation is unsubscripted WorkflowContext, nothing to infer - if origin is None: - return [] - - # Expecting WorkflowContext[T] - if origin is not WorkflowContext: - return [] - - args = get_args(ctx_annotation) - if not args: - return [] - - t = args[0] - # If t is a Union, flatten it - t_origin = get_origin(t) - # If Any, treat as unknown -> no output types inferred - if t is Any: - return [] - - if t_origin in (Union, UnionType): - # Return all union args as-is (may include generic aliases like list[str]) - return [arg for arg in get_args(t) if arg is not Any and arg is not type(None)] - - # Single concrete or generic alias type (e.g., str, int, list[str]) - if t is Any or t is type(None): - return [] - return [t] +from ._workflow_context import WorkflowContext, validate_function_signature class FunctionExecutor(Executor): @@ -85,61 +29,19 @@ class FunctionExecutor(Executor): """ @staticmethod - def _validate_function(func: Callable[..., Any]) -> None: + def _validate_function(func: Callable[..., Any]) -> tuple[type, Any, list[type[Any]], list[type[Any]]]: """Validate that the function has the correct signature for an executor. Args: func: The function to validate (can be sync or async) + Returns: + Tuple of (message_type, ctx_annotation, output_types, workflow_output_types) + Raises: ValueError: If the function signature is incorrect """ - signature = inspect.signature(func) - params = list(signature.parameters.values()) - - if len(params) not in (1, 2): - raise ValueError( - f"Function {func.__name__} must have one or two parameters: " - f"(message: T) or (message: T, ctx: WorkflowContext[U]). Got {len(params)} parameters." - ) - - message_param = params[0] - - # Check message parameter has type annotation - if message_param.annotation == inspect.Parameter.empty: - raise ValueError(f"Function {func.__name__} must have a type annotation for the message parameter") - - # If there's a second parameter, validate it's WorkflowContext[T] - if len(params) == 2: - ctx_param = params[1] - - # Check ctx parameter has proper type annotation - if ctx_param.annotation == inspect.Parameter.empty: - raise ValueError(f"Function {func.__name__} second parameter must be annotated as WorkflowContext[T]") - - # Validate that ctx parameter is WorkflowContext[T] - if not _is_workflow_context_type(ctx_param.annotation): - raise ValueError( - f"Function {func.__name__} second parameter must be annotated as WorkflowContext[T], " - f"got {ctx_param.annotation}" - ) - - # Check that WorkflowContext has a concrete type parameter - if ctx_param.annotation is WorkflowContext: - # This is unparameterized WorkflowContext - raise ValueError( - f"Function {func.__name__} WorkflowContext must be parameterized with a concrete T. " - f"Use WorkflowContext[str], WorkflowContext[int], etc." - ) - - if hasattr(ctx_param.annotation, "__args__") and ctx_param.annotation.__args__: - # This is WorkflowContext[T] with a concrete T - pass - else: - raise ValueError( - f"Function {func.__name__} WorkflowContext must be parameterized with a concrete T. " - f"Use WorkflowContext[str], WorkflowContext[int], etc." - ) + return validate_function_signature(func, "Function") def __init__(self, func: Callable[..., Any], id: str | None = None): """Initialize the FunctionExecutor with a user-defined function. @@ -148,27 +50,13 @@ class FunctionExecutor(Executor): func: The function to wrap as an executor (can be sync or async) id: Optional executor ID. If None, uses the function name. """ - # Validate function signature first - self._validate_function(func) - - # Extract types from function signature - signature = inspect.signature(func) - params = list(signature.parameters.values()) - - message_type = params[0].annotation + # Validate function signature and extract types + message_type, ctx_annotation, output_types, workflow_output_types = self._validate_function(func) # Determine if function has WorkflowContext parameter - has_context = len(params) == 2 + has_context = ctx_annotation is not None is_async = asyncio.iscoroutinefunction(func) - if has_context: - ctx_annotation = params[1].annotation - output_types = _infer_output_types_from_ctx_annotation(ctx_annotation) - else: - # For single-parameter functions, we can't infer output types - ctx_annotation = None - output_types = [] - # Initialize parent WITHOUT calling _discover_handlers yet # We'll manually set up the attributes first executor_id = id or getattr(func, "__name__", "FunctionExecutor") @@ -181,7 +69,7 @@ class FunctionExecutor(Executor): self._handlers: dict[type, Callable[[Any, WorkflowContext[Any]], Any]] = {} self._request_interceptors: dict[type | str, list[dict[str, Any]]] = {} - self._instance_handler_specs: list[dict[str, Any]] = [] + self._handler_specs: list[dict[str, Any]] = [] # Store the original function and whether it has context self._original_func = func @@ -211,12 +99,13 @@ class FunctionExecutor(Executor): return await asyncio.to_thread(func, message) # type: ignore # Now register our instance handler - self.register_instance_handler( + self._register_instance_handler( name=func.__name__, func=wrapped_func, message_type=message_type, ctx_annotation=ctx_annotation, output_types=output_types, + workflow_output_types=workflow_output_types, ) # Now we can safely call _discover_handlers (it won't find any class-level handlers) diff --git a/python/packages/main/agent_framework/_workflow/_magentic.py b/python/packages/main/agent_framework/_workflow/_magentic.py index 44a2adc9dd..a5813d61e8 100644 --- a/python/packages/main/agent_framework/_workflow/_magentic.py +++ b/python/packages/main/agent_framework/_workflow/_magentic.py @@ -28,7 +28,7 @@ from agent_framework import ( from agent_framework._agents import BaseAgent from agent_framework._pydantic import AFBaseModel -from ._events import WorkflowCompletedEvent, WorkflowEvent +from ._events import WorkflowEvent from ._executor import Executor, RequestInfoMessage, RequestResponse, handler from ._workflow import Workflow, WorkflowBuilder, WorkflowRunResult from ._workflow_context import WorkflowContext @@ -840,7 +840,9 @@ class MagenticOrchestratorExecutor(Executor): async def handle_start_message( self, message: MagenticStartMessage, - context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage | MagenticPlanReviewRequest], + context: WorkflowContext[ + MagenticResponseMessage | MagenticRequestMessage | MagenticPlanReviewRequest, ChatMessage + ], ) -> None: """Handle the initial start message to begin orchestration.""" if getattr(self, "_terminated", False): @@ -877,7 +879,7 @@ class MagenticOrchestratorExecutor(Executor): # Start the inner loop ctx2 = cast( - WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], context, ) await self._run_inner_loop(ctx2) @@ -886,7 +888,7 @@ class MagenticOrchestratorExecutor(Executor): async def handle_response_message( self, message: MagenticResponseMessage, - context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], ) -> None: """Handle responses from agents.""" if getattr(self, "_terminated", False): @@ -916,7 +918,7 @@ class MagenticOrchestratorExecutor(Executor): response: RequestResponse[MagenticPlanReviewRequest, MagenticPlanReviewReply], context: WorkflowContext[ # may broadcast ledger next, or ask for another round of review - MagenticResponseMessage | MagenticRequestMessage | MagenticPlanReviewRequest + MagenticResponseMessage | MagenticRequestMessage | MagenticPlanReviewRequest, ChatMessage ], ) -> None: if getattr(self, "_terminated", False): @@ -968,7 +970,7 @@ class MagenticOrchestratorExecutor(Executor): # Enter the normal coordination loop ctx2 = cast( - WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], context, ) await self._run_inner_loop(ctx2) @@ -997,7 +999,7 @@ class MagenticOrchestratorExecutor(Executor): self._context.chat_history.append(self._task_ledger) # No further review requests; proceed directly into coordination ctx2 = cast( - WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], context, ) await self._run_inner_loop(ctx2) @@ -1032,7 +1034,7 @@ class MagenticOrchestratorExecutor(Executor): async def _run_outer_loop( self, - context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], ) -> None: """Run the outer orchestration loop - planning phase.""" if self._context is None: @@ -1056,7 +1058,7 @@ class MagenticOrchestratorExecutor(Executor): async def _run_inner_loop( self, - context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], ) -> None: """Run the inner orchestration loop. Coordination phase. Serialized with a lock.""" if self._context is None or self._task_ledger is None: @@ -1066,7 +1068,7 @@ class MagenticOrchestratorExecutor(Executor): async def _run_inner_loop_locked( self, - context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], ) -> None: """Run inner loop with exclusive access.""" # Narrow optional context for the remainder of this method @@ -1154,7 +1156,7 @@ class MagenticOrchestratorExecutor(Executor): async def _reset_and_replan( self, - context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], ) -> None: """Reset context and replan.""" if self._context is None: @@ -1178,7 +1180,7 @@ class MagenticOrchestratorExecutor(Executor): async def _prepare_final_answer( self, - context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], ) -> None: """Prepare the final answer using the manager.""" if self._context is None: @@ -1188,14 +1190,14 @@ class MagenticOrchestratorExecutor(Executor): final_answer = await self._manager.prepare_final_answer(self._context.model_copy(deep=True)) # Emit a completed event for the workflow - await context.add_event(WorkflowCompletedEvent(final_answer)) + await context.yield_output(final_answer) if self._result_callback: await self._result_callback(final_answer) async def _check_within_limits_or_complete( self, - context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage], + context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage, ChatMessage], ) -> bool: """Check if orchestrator is within operational limits.""" if self._context is None: @@ -1221,8 +1223,8 @@ class MagenticOrchestratorExecutor(Executor): author_name=MAGENTIC_MANAGER_NAME, ) - # Emit a completed event with the partial result - await context.add_event(WorkflowCompletedEvent(partial_result)) + # Yield the partial result and signal completion + await context.yield_output(partial_result) if self._result_callback: await self._result_callback(partial_result) @@ -1232,7 +1234,9 @@ class MagenticOrchestratorExecutor(Executor): async def _send_plan_review_request( self, - context: WorkflowContext[MagenticResponseMessage | MagenticRequestMessage | MagenticPlanReviewRequest], + context: WorkflowContext[ + MagenticResponseMessage | MagenticRequestMessage | MagenticPlanReviewRequest, ChatMessage + ], ) -> None: """Emit a PlanReviewRequest via RequestInfoExecutor.""" # If plan sign-off is disabled (e.g., ran out of review rounds), do nothing diff --git a/python/packages/main/agent_framework/_workflow/_runner.py b/python/packages/main/agent_framework/_workflow/_runner.py index bfbbce5e9b..71eaa0452b 100644 --- a/python/packages/main/agent_framework/_workflow/_runner.py +++ b/python/packages/main/agent_framework/_workflow/_runner.py @@ -11,7 +11,7 @@ if TYPE_CHECKING: from ._edge import EdgeGroup from ._edge_runner import EdgeRunner, create_edge_runner -from ._events import WorkflowCompletedEvent, WorkflowEvent, _framework_event_origin +from ._events import WorkflowEvent, WorkflowOutputEvent, _framework_event_origin from ._executor import Executor from ._runner_context import ( _DATACLASS_MARKER, @@ -23,7 +23,6 @@ from ._runner_context import ( ) from ._shared_state import SharedState from ._typing_utils import is_instance_of -from ._workflow_context import WorkflowContext logger = logging.getLogger(__name__) @@ -204,16 +203,14 @@ class Runner: f"from sub-workflow '{sub_request.sub_workflow_id}' " f"to executor '{executor.id}' for interception." ) - # Create WorkflowContext with trace context from message - workflow_ctx: WorkflowContext[Any] = WorkflowContext( - executor.id, - [message.source_id], - self._shared_state, - self._ctx, + await executor.execute( + sub_request, + [message.source_id], # source_executor_ids + self._shared_state, # shared_state + self._ctx, # runner_context trace_contexts=[message.trace_context] if message.trace_context else None, source_span_ids=[message.source_span_id] if message.source_span_id else None, ) - await executor.execute(sub_request, workflow_ctx) interceptor_found = True break if interceptor_found: @@ -226,20 +223,19 @@ class Runner: request_info_executor = self._find_request_info_executor() if request_info_executor: - request_info_workflow_ctx: WorkflowContext[None] = WorkflowContext( - request_info_executor.id, - [message.source_id], - self._shared_state, - self._ctx, - trace_contexts=[message.trace_context] if message.trace_context else None, - source_span_ids=[message.source_span_id] if message.source_span_id else None, - ) logger.info( f"Sending sub-workflow request of type '{sub_request.data.__class__.__name__}' " f"from sub-workflow '{sub_request.sub_workflow_id}' to RequestInfoExecutor " f"'{request_info_executor.id}'" ) - await request_info_executor.execute(sub_request, request_info_workflow_ctx) + await request_info_executor.execute( + sub_request, + [message.source_id], # source_executor_ids + self._shared_state, # shared_state + self._ctx, # runner_context + trace_contexts=[message.trace_context] if message.trace_context else None, + source_span_ids=[message.source_span_id] if message.source_span_id else None, + ) else: logger.warning( f"Sub-workflow request of type '{sub_request.data.__class__.__name__}' " @@ -303,8 +299,9 @@ class Runner: final_messages = message.data.agent_run_response.messages final_text = final_messages[-1].text if final_messages else "(no content)" with _framework_event_origin(): - completion_event = WorkflowCompletedEvent(final_text) - await self._ctx.add_event(completion_event) + # TODO(moonbox3): does user expect this event to contain the final text? + output_event = WorkflowOutputEvent(data=final_text, source_executor_id="") + await self._ctx.add_event(output_event) continue # Terminal handled except Exception as exc: # pragma: no cover - defensive logger.debug("Suppressed exception during terminal message type check: %s", exc) @@ -326,8 +323,9 @@ class Runner: final_messages = message.data.agent_run_response.messages final_text = final_messages[-1].text if final_messages else "(no content)" with _framework_event_origin(): - completion_event = WorkflowCompletedEvent(final_text) - await self._ctx.add_event(completion_event) + # TODO(moonbox3): does user expect this event to contain the final text? + output_event = WorkflowOutputEvent(data=final_text, source_executor_id="") + await self._ctx.add_event(output_event) continue except Exception as exc: # pragma: no cover logger.debug("Terminal completion emission failed: %s", exc) diff --git a/python/packages/main/agent_framework/_workflow/_sequential.py b/python/packages/main/agent_framework/_workflow/_sequential.py index cc091601c4..0202f21b34 100644 --- a/python/packages/main/agent_framework/_workflow/_sequential.py +++ b/python/packages/main/agent_framework/_workflow/_sequential.py @@ -8,10 +8,10 @@ workflow where: - A shared conversation context (list[ChatMessage]) is passed along the chain - Agents append their assistant messages to the context - Custom executors can transform or summarize and return a refined context -- The workflow completes with the final context produced by the last participant +- The workflow finishes with the final context produced by the last participant Typical wiring: - input -> _InputToConversation -> participant1 -> (agent? -> _ResponseToConversation) -> ... -> participantN -> _CompleteWithConversation + input -> _InputToConversation -> participant1 -> (agent? -> _ResponseToConversation) -> ... -> participantN -> _EndWithConversation Notes: - Participants can mix AgentProtocol and Executor objects @@ -27,9 +27,8 @@ Why include the small internal adapter executors? - Agent response adaptation ("to-conversation:"): agents (via AgentExecutor) emit `AgentExecutorResponse`. The adapter converts that to a `list[ChatMessage]` using `full_conversation` so original prompts aren't lost when chaining. -- Explicit completion ("complete"): emits a `WorkflowCompletedEvent` with the final - conversation list, giving a consistent terminal payload shape for both agents and - custom executors. +- Result output ("end"): yields the final conversation list and the workflow becomes idle + giving a consistent terminal payload shape for both agents and custom executors. These adapters are first-class executors by design so they are type-checked at edges, observable (ExecutorInvoke/Completed events), and easily testable/reusable. Their IDs are @@ -43,7 +42,6 @@ from typing import Any from agent_framework import AgentProtocol, ChatMessage, Role -from ._events import WorkflowCompletedEvent from ._executor import ( AgentExecutor, AgentExecutorResponse, @@ -84,12 +82,12 @@ class _ResponseToConversation(Executor): await ctx.send_message(list(response.full_conversation)) -class _CompleteWithConversation(Executor): +class _EndWithConversation(Executor): """Terminates the workflow by emitting the final conversation context.""" @handler - async def complete(self, conversation: list[ChatMessage], ctx: WorkflowContext[Any]) -> None: - await ctx.add_event(WorkflowCompletedEvent(data=list(conversation))) + async def end(self, conversation: list[ChatMessage], ctx: WorkflowContext[Any, list[ChatMessage]]) -> None: + await ctx.yield_output(list(conversation)) class SequentialBuilder: @@ -148,14 +146,14 @@ class SequentialBuilder: - If Agent (or AgentExecutor): pass conversation to the agent, then convert response to conversation via _ResponseToConversation - Else (custom Executor): pass conversation directly to the executor - - _CompleteWithConversation emits WorkflowCompletedEvent with the final conversation + - _EndWithConversation yields the final conversation and the workflow becomes idle """ if not self._participants: raise ValueError("No participants provided. Call .participants([...]) first.") # Internal nodes input_conv = _InputToConversation(id="input-conversation") - complete = _CompleteWithConversation(id="complete") + end = _EndWithConversation(id="end") builder = WorkflowBuilder() builder.set_start_executor(input_conv) @@ -182,6 +180,6 @@ class SequentialBuilder: raise TypeError(f"Unsupported participant type: {type(p).__name__}") # Terminate with the final conversation - builder.add_edge(prior, complete) + builder.add_edge(prior, end) return builder.build() diff --git a/python/packages/main/agent_framework/_workflow/_validation.py b/python/packages/main/agent_framework/_workflow/_validation.py index f7da01b969..f2ece7e3c0 100644 --- a/python/packages/main/agent_framework/_workflow/_validation.py +++ b/python/packages/main/agent_framework/_workflow/_validation.py @@ -9,26 +9,12 @@ from types import UnionType from typing import Any, Union, get_args, get_origin from ._edge import Edge, EdgeGroup, FanInEdgeGroup -from ._executor import Executor +from ._executor import Executor, RequestInfoExecutor +from ._workflow_executor import WorkflowExecutor logger = logging.getLogger(__name__) -def _is_type_like(x: Any) -> bool: - """Check if a value is a type-like entity. - - A "type-like" entry is either a class/type or a typing alias - (e.g., list[str] has an origin and args). - - Args: - x: The value to check - - Returns: - True if the value is type-like, False otherwise - """ - return isinstance(x, type) or get_origin(x) is not None - - # region Enums and Base Classes class ValidationTypeEnum(Enum): """Enumeration of workflow validation types.""" @@ -108,23 +94,6 @@ class GraphConnectivityError(WorkflowValidationError): super().__init__(message, validation_type=ValidationTypeEnum.GRAPH_CONNECTIVITY) -class HandlerOutputAnnotationError(WorkflowValidationError): - """Exception raised when a handler's WorkflowContext output annotation is invalid or missing.""" - - def __init__(self, executor_id: str, handler_name: str, reason: str): - super().__init__( - message=( - "Invalid WorkflowContext output annotation in handler " - f"'{handler_name}' of executor '{executor_id}': {reason}. " - "Handlers must annotate their third parameter as WorkflowContext[T]. " - "Use WorkflowContext[None] if the handler emits no messages." - ), - validation_type=ValidationTypeEnum.HANDLER_OUTPUT_ANNOTATION, - ) - self.executor_id = executor_id - self.handler_name = handler_name - - class InterceptorConflictError(WorkflowValidationError): """Exception raised when multiple executors intercept the same request type from the same sub-workflow.""" @@ -216,7 +185,6 @@ class WorkflowGraphValidator: self._validate_type_compatibility() self._validate_graph_connectivity(start_executor_id) self._validate_self_loops() - self._validate_handler_ambiguity() self._validate_dead_ends() self._validate_cycles() self._validate_interceptor_uniqueness() @@ -224,158 +192,18 @@ class WorkflowGraphValidator: def _validate_handler_output_annotations(self) -> None: """Validate that each handler's ctx parameter is annotated with WorkflowContext[T]. - Requirements: - - WorkflowContext annotation must be present - - T_Out must be provided; if no outputs, it must be None - - T_Out elements must be valid types (class) or typing generics (e.g., list[str]); - values like int() or 123 are invalid + Note: This validation is now primarily handled at handler registration time + via the unified validation functions in _workflow_context.py when the @handler + decorator is applied. This method is kept minimal for any edge cases. """ - from ._workflow_context import WorkflowContext # Local import to avoid cycles - - # Iterate over all registered executors in the workflow graph - for executor_id, executor in self._executors.items(): - for attr_name in dir(executor.__class__): - if attr_name.startswith("_"): - continue - # Retrieve attributes without binding (so the first parameter remains 'self'). - # This ensures inspect.signature sees all three parameters: (self, message, ctx). - attr = None - from contextlib import suppress - - with suppress(Exception): - attr = inspect.getattr_static(executor.__class__, attr_name) - if attr is None: - continue - # Consider only callables that were decorated with @handler - if not callable(attr) or not hasattr(attr, "_handler_spec"): - continue - - handler_spec = attr._handler_spec # type: ignore[attr-defined] - handler_name = handler_spec.get("name", attr_name) - - try: - # Inspect the function signature of the unbound function - sig = inspect.signature(attr) - except (TypeError, ValueError): - continue - - params = list(sig.parameters.values()) - # Handlers must have exactly three parameters: (self, message, ctx) - if len(params) != 3: - continue - - ctx_param = params[2] - ctx_ann = ctx_param.annotation - - # If ctx lacks an annotation entirely, fail fast with a clear message - if ctx_ann is inspect.Parameter.empty: - raise HandlerOutputAnnotationError(executor_id, handler_name, "missing type annotation for ctx") - - # Validate that the ctx annotation is WorkflowContext[...] and is properly parameterized - ctx_origin = get_origin(ctx_ann) - if ctx_origin is None: - # If it's exactly the WorkflowContext class, T_Out is missing (e.g., WorkflowContext) - if ctx_ann is WorkflowContext: - raise HandlerOutputAnnotationError( - executor_id, - handler_name, - "T_Out is missing; use WorkflowContext[None] or specify concrete types", - ) - else: - # The annotation is parameterized, but must be for WorkflowContext - if ctx_origin is not WorkflowContext: - raise HandlerOutputAnnotationError( - executor_id, handler_name, f"ctx must be WorkflowContext[T], got {ctx_ann}" - ) - - # Extract and validate T_Out - type_args = get_args(ctx_ann) - if not type_args: - raise HandlerOutputAnnotationError( - executor_id, - handler_name, - "T_Out is missing; use WorkflowContext[None] or specify concrete types", - ) - - t_out = type_args[0] - - # Allow Any for T_Out (unspecified outputs). We accept this here and - # skip type compatibility later, but still enforce shape validity elsewhere. - if t_out is Any: - continue - - # Allow None (no outputs) explicitly declared - if t_out is type(None): - continue - - # If T_Out is a union, validate each member (e.g., str | int) - union_origin = get_origin(t_out) - type_items: list[Any] - type_items = list(get_args(t_out)) if union_origin in (Union, UnionType) else [t_out] - - invalid = [x for x in type_items if not _is_type_like(x) and x is not type(None)] - if invalid: - raise HandlerOutputAnnotationError( - executor_id, - handler_name, - f"T_Out contains invalid entries: {invalid}. Use proper types or typing generics", - ) - - # Also validate instance-level handler specs if present - if hasattr(executor, "_instance_handler_specs"): - for spec in executor._instance_handler_specs: - handler_name = spec.get("name", "unknown") - ctx_ann = spec.get("ctx_annotation") - - if ctx_ann is None: - continue # Skip if no annotation stored - - # Validate that the ctx annotation is WorkflowContext[...] and is properly parameterized - ctx_origin = get_origin(ctx_ann) - if ctx_origin is None: - if ctx_ann is WorkflowContext: - raise HandlerOutputAnnotationError( - executor_id, - handler_name, - "T_Out is missing; use WorkflowContext[None] or specify concrete types", - ) - else: - if ctx_origin is not WorkflowContext: - raise HandlerOutputAnnotationError( - executor_id, handler_name, f"ctx must be WorkflowContext[T], got {ctx_ann}" - ) - - # Extract and validate T_Out - type_args = get_args(ctx_ann) - if not type_args: - raise HandlerOutputAnnotationError( - executor_id, - handler_name, - "T_Out is missing; use WorkflowContext[None] or specify concrete types", - ) - - t_out = type_args[0] - - # Allow Any for T_Out (unspecified outputs) - if t_out is Any: - continue - - # Allow None (no outputs) explicitly declared - if t_out is type(None): - continue - - # If T_Out is a union, validate each member - union_origin = get_origin(t_out) - instance_type_items: list[Any] - instance_type_items = list(get_args(t_out)) if union_origin in (Union, UnionType) else [t_out] - - invalid = [x for x in instance_type_items if not _is_type_like(x) and x is not type(None)] - if invalid: - raise HandlerOutputAnnotationError( - executor_id, - handler_name, - f"T_Out contains invalid entries: {invalid}. Use proper types or typing generics", - ) + # The comprehensive validation is already done during handler registration: + # 1. @handler decorator calls validate_function_signature() + # 2. FunctionExecutor constructor calls validate_function_signature() + # 3. Both use validate_workflow_context_annotation() for WorkflowContext validation + # + # All executors in the workflow must have gone through one of these paths, + # so redundant validation here is unnecessary and has been removed. + pass # endregion @@ -444,29 +272,25 @@ class WorkflowGraphValidator: target_executor = self._executors[edge.target_id] # Get output types from source executor - source_output_types = self._get_executor_output_types(source_executor) + source_output_types = list(source_executor.output_types) + # Also include intercepted request types as potential outputs + # since @intercepts_request methods can forward requests + source_output_types.extend(source_executor.request_types) # Get input types from target executor - target_input_types = self._get_executor_input_types(target_executor) + target_input_types = target_executor.input_types # If either executor has no type information, log warning and skip validation # This allows for dynamic typing scenarios but warns about reduced validation coverage if not source_output_types or not target_input_types: - # Suppress warnings for built-in workflow components where dynamic typing is expected - try: - from ._executor import RequestInfoExecutor, WorkflowExecutor # local import to avoid cycles - - builtin_types = (RequestInfoExecutor, WorkflowExecutor) - except Exception: - builtin_types = tuple() # type: ignore[assignment] - - if not source_output_types and not isinstance(source_executor, builtin_types): + # Suppress warnings for RequestInfoExecutor where dynamic typing is expected + if not source_output_types and not isinstance(source_executor, RequestInfoExecutor): logger.warning( f"Executor '{source_executor.id}' has no output type annotations. " f"Type compatibility validation will be skipped for edges from this executor. " f"Consider adding WorkflowContext[T] generics in handlers for better validation." ) - if not target_input_types and not isinstance(target_executor, builtin_types): + if not target_input_types and not isinstance(target_executor, RequestInfoExecutor): logger.warning( f"Executor '{target_executor.id}' has no input type annotations. " f"Type compatibility validation will be skipped for edges to this executor. " @@ -506,62 +330,6 @@ class WorkflowGraphValidator: target_input_types, ) - def _get_executor_output_types(self, executor: Executor) -> list[type[Any]]: - """Extract output types from an executor's message handlers. - - Args: - executor: The executor to analyze - - Returns: - list of types that this executor can output - """ - output_types: list[type[Any]] = [] - - for attr_name in dir(executor.__class__): - if attr_name.startswith("_"): - continue - try: - attr = getattr(executor.__class__, attr_name) - if callable(attr) and hasattr(attr, "_handler_spec"): - handler_spec = attr._handler_spec # type: ignore - handler_output_types = handler_spec.get("output_types", []) - output_types.extend(handler_output_types) - except AttributeError: - # Skip attributes that may not be accessible - continue - - # Also include intercepted request types as potential outputs - # since @intercepts_request methods can forward requests - if hasattr(executor, "_request_interceptors"): - for request_type in executor._request_interceptors: - if isinstance(request_type, type): - output_types.append(request_type) - - # Include output types from instance-level handler specs - if hasattr(executor, "_instance_handler_specs"): - for spec in executor._instance_handler_specs: - handler_output_types = spec.get("output_types", []) - output_types.extend(handler_output_types) - - return output_types - - def _get_executor_input_types(self, executor: Executor) -> list[type[Any]]: - """Extract input types from an executor's message handlers. - - Args: - executor: The executor to analyze - - Returns: - list of types that this executor can handle as input - """ - input_types: list[type[Any]] = [] - - # Access the private _handlers attribute to get input types - if hasattr(executor, "_handlers"): - input_types.extend(executor._handlers.keys()) # type: ignore - - return input_types - # endregion # region Graph Connectivity Validation @@ -650,30 +418,6 @@ class WorkflowGraphValidator: f"This may cause infinite recursion if not properly handled with conditions." ) - def _validate_handler_ambiguity(self) -> None: - """Check for potential ambiguity in message handlers. - - Warns when executors have multiple handlers that could handle the same type, - which might lead to unexpected behavior. - """ - for executor_id, executor in self._executors.items(): - input_types = self._get_executor_input_types(executor) - - # Check for duplicate input types - seen_types: set[type[Any]] = set() - duplicate_types: set[type[Any]] = set() - - for input_type in input_types: - if input_type in seen_types: - duplicate_types.add(input_type) - seen_types.add(input_type) - - if duplicate_types: - logger.warning( - f"Executor '{executor_id}' has multiple handlers for the same input types: " - f"{[str(t) for t in duplicate_types]}. This may lead to ambiguous message routing." - ) - def _validate_dead_ends(self) -> None: """Identify executors that have no outgoing edges (potential dead ends). @@ -744,8 +488,6 @@ class WorkflowGraphValidator: This prevents non-deterministic behavior where multiple executors could intercept the same request type from the same sub-workflow. """ - from ._executor import WorkflowExecutor - # Find all WorkflowExecutor instances in the workflow workflow_executors: dict[str, WorkflowExecutor] = {} for executor_id, executor in self._executors.items(): diff --git a/python/packages/main/agent_framework/_workflow/_viz.py b/python/packages/main/agent_framework/_workflow/_viz.py index 37eafdd6d3..f650dc5f23 100644 --- a/python/packages/main/agent_framework/_workflow/_viz.py +++ b/python/packages/main/agent_framework/_workflow/_viz.py @@ -253,7 +253,7 @@ class WorkflowViz: """Emit DOT subgraphs for any WorkflowExecutor instances found in the workflow.""" # Lazy import to avoid any potential import cycles try: - from ._executor import WorkflowExecutor # type: ignore + from ._workflow_executor import WorkflowExecutor # type: ignore except ImportError: # pragma: no cover - best-effort; if unavailable, skip subgraphs return @@ -327,7 +327,7 @@ class WorkflowViz: def _emit_sub_workflows_mermaid(self, wf: Workflow, lines: list[str], indent: str) -> None: try: - from ._executor import WorkflowExecutor # type: ignore + from ._workflow_executor import WorkflowExecutor # type: ignore except ImportError: # pragma: no cover return diff --git a/python/packages/main/agent_framework/_workflow/_workflow.py b/python/packages/main/agent_framework/_workflow/_workflow.py index aa083ff92b..d4b21f6124 100644 --- a/python/packages/main/agent_framework/_workflow/_workflow.py +++ b/python/packages/main/agent_framework/_workflow/_workflow.py @@ -30,10 +30,10 @@ from ._edge import ( ) from ._events import ( RequestInfoEvent, - WorkflowCompletedEvent, WorkflowErrorDetails, WorkflowEvent, WorkflowFailedEvent, + WorkflowOutputEvent, WorkflowRunState, WorkflowStartedEvent, WorkflowStatusEvent, @@ -56,32 +56,36 @@ logger = logging.getLogger(__name__) class WorkflowRunResult(list[WorkflowEvent]): - """A list of events generated during the workflow execution in non-streaming mode. + """Container for events generated during non-streaming workflow execution. - Preserves the historical contract that the list contains data-plane events - only (executor invoke/complete, completed, requests), while exposing the - control-plane status timeline via accessors. + ## Overview + Represents the complete execution results of a workflow run, containing all events + generated from start to idle state. Workflows produce outputs incrementally through + ctx.yield_output() calls during execution. + + ## Event Structure + Maintains separation between data-plane and control-plane events: + - Data-plane events: Executor invocations, completions, outputs, and requests (in main list) + - Control-plane events: Status timeline accessible via status_timeline() method + + ## Key Methods + - get_outputs(): Extract all workflow outputs from the execution + - get_request_info_events(): Retrieve external input requests made during execution + - get_final_state(): Get the final workflow state (IDLE, IDLE_WITH_PENDING_REQUESTS, etc.) + - status_timeline(): Access the complete status event history """ def __init__(self, events: list[WorkflowEvent], status_events: list[WorkflowStatusEvent] | None = None) -> None: super().__init__(events) self._status_events: list[WorkflowStatusEvent] = status_events or [] - def get_completed_event(self) -> WorkflowCompletedEvent | None: - """Get the completed event from the workflow run result. + def get_outputs(self) -> list[Any]: + """Get all outputs from the workflow run result. Returns: - A completed WorkflowEvent instance if the workflow has a completed event, otherwise None. - - Raises: - ValueError: If there are multiple completed events in the workflow run result. + A list of outputs produced by the workflow during its execution. """ - completed_events = [event for event in self if isinstance(event, WorkflowCompletedEvent)] - if not completed_events: - return None - if len(completed_events) > 1: - raise ValueError("Multiple completed events found.") - return completed_events[0] + return [event.data for event in self if isinstance(event, WorkflowOutputEvent)] def get_request_info_events(self) -> list[RequestInfoEvent]: """Get all request info events from the workflow run result. @@ -113,10 +117,54 @@ class WorkflowRunResult(list[WorkflowEvent]): class Workflow(AFBaseModel): - """A class representing a workflow that can be executed. + """A graph-based execution engine that orchestrates connected executors. - This class is a placeholder for the workflow logic and does not implement any specific functionality. - It serves as a base class for more complex workflows that can be defined in subclasses. + ## 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. + + ## Execution Model + Executors run in synchronized supersteps where each executor: + - Is invoked when it receives messages from connected edge groups + - Can send messages to downstream executors via ctx.send_message() + - Can yield workflow-level outputs via ctx.yield_output() + - Can emit custom events via ctx.add_event() + + Messages between executors are delivered at the end of each superstep and are not + visible in the event stream. Only workflow-level events (outputs, custom events) + and status events are observable to callers. + + ## Input/Output Types + Workflow types are discovered at runtime by inspecting: + - Input types: From the start executor's input types + - Output types: Union of all executors' workflow output types + Access these via the input_types and output_types properties. + + ## Execution Methods + - run(): Execute to completion, returns WorkflowRunResult with all events + - run_stream(): Returns async generator yielding events as they occur + - run_from_checkpoint(): Resume from a saved checkpoint + - run_stream_from_checkpoint(): Resume from checkpoint with streaming + + ## External Input Requests + Workflows can request external input using a RequestInfoExecutor: + 1. Executor connects to RequestInfoExecutor via edge group and back to itself + 2. Executor sends RequestInfoMessage to RequestInfoExecutor + 3. RequestInfoExecutor emits RequestInfoEvent and workflow enters IDLE_WITH_PENDING_REQUESTS + 4. Caller handles requests and uses send_responses()/send_responses_streaming() to continue + + ## Checkpointing + When enabled, checkpoints are created at the end of each superstep, capturing: + - Executor states + - Messages in transit + - Shared state + Workflows can be paused and resumed across process restarts using checkpoint storage. + + ## Composition + Workflows can be nested using WorkflowExecutor, which wraps a child workflow as an executor. + The nested workflow's input/output types become part of the WorkflowExecutor's types. + When invoked, the WorkflowExecutor runs the nested workflow to completion and processes its outputs. """ edge_groups: list[EdgeGroup] = Field( @@ -202,7 +250,7 @@ class Workflow(AFBaseModel): # Get the original executor object and serialize its workflow original_executor = self.executors.get(executor_id) if original_executor and hasattr(original_executor, "workflow"): - from ._executor import WorkflowExecutor + from ._workflow_executor import WorkflowExecutor if isinstance(original_executor, WorkflowExecutor): executor_data["workflow"] = original_executor.workflow.model_dump(**kwargs) @@ -249,7 +297,6 @@ class Workflow(AFBaseModel): OtelAttr.WORKFLOW_ID: self.id, }, ) as span: - saw_completed = False saw_request = False emitted_in_progress_pending = False try: @@ -273,25 +320,19 @@ class Workflow(AFBaseModel): # All executor executions happen within workflow span async for event in self._runner.run_until_convergence(): - # Track terminal indicators while forwarding events - if isinstance(event, WorkflowCompletedEvent): - saw_completed = True - elif isinstance(event, RequestInfoEvent): + # Track request events for final status determination + if isinstance(event, RequestInfoEvent): saw_request = True yield event - if isinstance(event, RequestInfoEvent) and not emitted_in_progress_pending and not saw_completed: + if isinstance(event, RequestInfoEvent) and not emitted_in_progress_pending: emitted_in_progress_pending = True with _framework_event_origin(): pending_status = WorkflowStatusEvent(WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS) yield pending_status - # Success path: emit a final status based on observed terminal signals - if saw_completed: - with _framework_event_origin(): - terminal_status = WorkflowStatusEvent(WorkflowRunState.COMPLETED) - yield terminal_status - elif saw_request: + # Workflow runs until idle - emit final status based on whether requests are pending + if saw_request: with _framework_event_origin(): terminal_status = WorkflowStatusEvent(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS) yield terminal_status @@ -334,14 +375,11 @@ class Workflow(AFBaseModel): executor = self.get_start_executor() await executor.execute( message, - WorkflowContext( - executor.id, - [self.__class__.__name__], - self._shared_state, - self._runner.context, - trace_contexts=None, # No parent trace context for workflow start - source_span_ids=None, # No source span for workflow start - ), + [self.__class__.__name__], # source_executor_ids + self._shared_state, # shared_state + self._runner.context, # runner_context + trace_contexts=None, # No parent trace context for workflow start + source_span_ids=None, # No source span for workflow start ) async for event in self._run_workflow_with_tracing(initial_executor_fn=initial_execution, reset_context=True): @@ -774,6 +812,36 @@ class Workflow(AFBaseModel): def graph_signature_hash(self) -> str: return self._graph_signature_hash + @property + def input_types(self) -> list[type[Any]]: + """Get the input types of the workflow. + + The input types are the list of input types of the start executor. + + Returns: + A list of input types that the workflow can accept. + """ + start_executor = self.get_start_executor() + return start_executor.input_types + + @property + def output_types(self) -> list[type[Any]]: + """Get the output types of the workflow. + + The output types are the list of all workflow output types from executors + that have workflow output types. + + Returns: + A list of output types that the workflow can produce. + """ + output_types: set[type[Any]] = set() + + for executor in self.executors.values(): + workflow_output_types = executor.workflow_output_types + output_types.update(workflow_output_types) + + return list(output_types) + def as_agent(self, name: str | None = None) -> WorkflowAgent: """Create a WorkflowAgent that wraps this workflow. diff --git a/python/packages/main/agent_framework/_workflow/_workflow_context.py b/python/packages/main/agent_framework/_workflow/_workflow_context.py index 41012161fc..9bc5864934 100644 --- a/python/packages/main/agent_framework/_workflow/_workflow_context.py +++ b/python/packages/main/agent_framework/_workflow/_workflow_context.py @@ -1,10 +1,16 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + +import inspect import logging -from typing import Any, Generic, TypeVar, cast, get_args +from collections.abc import Callable +from types import UnionType +from typing import Any, Generic, Union, cast, get_args, get_origin from opentelemetry.propagate import inject from opentelemetry.trace import SpanKind +from typing_extensions import Never, TypeVar from ..observability import OtelAttr, create_workflow_span from ._events import ( @@ -12,19 +18,249 @@ from ._events import ( WorkflowEventSource, WorkflowFailedEvent, WorkflowLifecycleEvent, + WorkflowOutputEvent, WorkflowStartedEvent, WorkflowStatusEvent, WorkflowWarningEvent, + _framework_event_origin, ) from ._runner_context import Message, RunnerContext from ._shared_state import SharedState -T_Out = TypeVar("T_Out") +T_Out = TypeVar("T_Out", default=Never) +T_W_Out = TypeVar("T_W_Out", default=Never) logger = logging.getLogger(__name__) +def infer_output_types_from_ctx_annotation(ctx_annotation: Any) -> tuple[list[type[Any]], list[type[Any]]]: + """Infer message types and workflow output types from the WorkflowContext generic parameters. + + Examples: + - WorkflowContext -> ([], []) + - WorkflowContext[str] -> ([str], []) + - WorkflowContext[str, int] -> ([str], [int]) + - WorkflowContext[str | int, bool | int] -> ([str, int], [bool, int]) + - WorkflowContext[Union[str, int], Union[bool, int]] -> ([str, int], [bool, int]) + - WorkflowContext[Any] -> ([Any], []) + - WorkflowContext[Any, Any] -> ([Any], [Any]) + - WorkflowContext[Never, Never] -> ([], []) + - WorkflowContext[Never, int] -> ([], [int]) + + Returns: + Tuple of (message_types, workflow_output_types) + """ + # If no annotation or not parameterized, return empty lists + try: + origin = get_origin(ctx_annotation) + except Exception: + origin = None + + # If annotation is unsubscripted WorkflowContext, nothing to infer + if origin is None: + return [], [] + + # Expecting WorkflowContext[T_Out, T_W_Out] + if origin is not WorkflowContext: + return [], [] + + args = list(get_args(ctx_annotation)) + if not args: + return [], [] + + # WorkflowContext[T_Out] -> message_types from T_Out, no workflow output types + if len(args) == 1: + t = args[0] + t_origin = get_origin(t) + if t is Any: + return [cast(type[Any], Any)], [] + + if t_origin in (Union, UnionType): + message_types = [arg for arg in get_args(t) if arg is not Any and arg is not Never] + return message_types, [] + + if t is Never: + return [], [] + return [t], [] + + # WorkflowContext[T_Out, T_W_Out] -> message_types from T_Out, workflow_output_types from T_W_Out + t_out, t_w_out = args[:2] # Take first two args in case there are more + + # Process T_Out for message_types + message_types = [] + t_out_origin = get_origin(t_out) + if t_out is Any: + message_types = [cast(type[Any], Any)] + elif t_out is not Never: + if t_out_origin in (Union, UnionType): + message_types = [arg for arg in get_args(t_out) if arg is not Any and arg is not Never] + else: + message_types = [t_out] + + # Process T_W_Out for workflow_output_types + workflow_output_types = [] + t_w_out_origin = get_origin(t_w_out) + if t_w_out is Any: + workflow_output_types = [cast(type[Any], Any)] + elif t_w_out is not Never: + if t_w_out_origin in (Union, UnionType): + workflow_output_types = [arg for arg in get_args(t_w_out) if arg is not Any and arg is not Never] + else: + workflow_output_types = [t_w_out] + + return message_types, workflow_output_types + + +def _is_workflow_context_type(annotation: Any) -> bool: + """Check if an annotation represents WorkflowContext, WorkflowContext[T], or WorkflowContext[T, U].""" + origin = get_origin(annotation) + if origin is WorkflowContext: + return True + # Also handle the case where the raw class is used + return annotation is WorkflowContext + + +def validate_workflow_context_annotation( + annotation: Any, + parameter_name: str, + context_description: str, +) -> tuple[list[type[Any]], list[type[Any]]]: + """Validate a WorkflowContext annotation and return inferred types. + + Args: + annotation: The type annotation to validate + parameter_name: Name of the parameter (for error messages) + context_description: Description of the context (e.g., "Function func1", "Handler method") + + Returns: + Tuple of (output_types, workflow_output_types) + + Raises: + ValueError: If the annotation is invalid + """ + if annotation == inspect.Parameter.empty: + raise ValueError( + f"{context_description} {parameter_name} must have a WorkflowContext, " + f"WorkflowContext[T] or WorkflowContext[T, U] type annotation, " + f"where T is output message type and U is workflow output type" + ) + + if not _is_workflow_context_type(annotation): + raise ValueError( + f"{context_description} {parameter_name} must be annotated as " + f"WorkflowContext, WorkflowContext[T], or WorkflowContext[T, U], " + f"got {annotation}" + ) + + # Validate type arguments for WorkflowContext[T] or WorkflowContext[T, U] + type_args = get_args(annotation) + + if len(type_args) > 2: + raise ValueError( + f"{context_description} {parameter_name} must have at most 2 type arguments, " + "WorkflowContext, WorkflowContext[T], or WorkflowContext[T, U], " + f"got {len(type_args)} arguments" + ) + + if type_args: + # Helper function to check if a value is a valid type annotation + def _is_type_like(x: Any) -> bool: + """Check if a value is a type-like entity (class, type, or typing construct).""" + return isinstance(x, type) or get_origin(x) is not None or x is Never + + for i, type_arg in enumerate(type_args): + param_description = "T_Out" if i == 0 else "T_W_Out" + + # Allow Any explicitly + if type_arg is Any: + continue + + # Check if it's a union type and validate each member + union_origin = get_origin(type_arg) + if union_origin in (Union, UnionType): + union_members = get_args(type_arg) + invalid_members = [m for m in union_members if not _is_type_like(m) and m is not Any] + if invalid_members: + raise ValueError( + f"{context_description} {parameter_name} {param_description} " + f"contains invalid type entries: {invalid_members}. " + f"Use proper types or typing generics" + ) + else: + # Check if it's a valid type + if not _is_type_like(type_arg): + raise ValueError( + f"{context_description} {parameter_name} {param_description} " + f"contains invalid type entry: {type_arg}. " + f"Use proper types or typing generics" + ) + + return infer_output_types_from_ctx_annotation(annotation) + + +def validate_function_signature( + func: Callable[..., Any], context_description: str +) -> tuple[type, Any, list[type[Any]], list[type[Any]]]: + """Validate function signature for executor functions. + + Args: + func: The function to validate + context_description: Description for error messages (e.g., "Function", "Handler method") + + Returns: + Tuple of (message_type, ctx_annotation, output_types, workflow_output_types) + + Raises: + ValueError: If the function signature is invalid + """ + signature = inspect.signature(func) + params = list(signature.parameters.values()) + + # Determine expected parameter count based on context + expected_counts: tuple[int, ...] + if context_description.startswith("Function"): + # Function executor: (message) or (message, ctx) + expected_counts = (1, 2) + param_description = "(message: T) or (message: T, ctx: WorkflowContext[U])" + else: + # Handler method: (self, message, ctx) + expected_counts = (3,) + param_description = "(self, message: T, ctx: WorkflowContext[U])" + + if len(params) not in expected_counts: + raise ValueError( + f"{context_description} {func.__name__} must have {param_description}. Got {len(params)} parameters." + ) + + # Extract message parameter (index 0 for functions, index 1 for methods) + message_param_idx = 0 if context_description.startswith("Function") else 1 + message_param = params[message_param_idx] + + # Check message parameter has type annotation + if message_param.annotation == inspect.Parameter.empty: + raise ValueError(f"{context_description} {func.__name__} must have a type annotation for the message parameter") + + message_type = message_param.annotation + + # Check if there's a context parameter + ctx_param_idx = message_param_idx + 1 + if len(params) > ctx_param_idx: + ctx_param = params[ctx_param_idx] + output_types, workflow_output_types = validate_workflow_context_annotation( + ctx_param.annotation, f"parameter '{ctx_param.name}'", context_description + ) + ctx_annotation = ctx_param.annotation + else: + # No context parameter (only valid for function executors) + if not context_description.startswith("Function"): + raise ValueError(f"{context_description} {func.__name__} must have a WorkflowContext parameter") + output_types, workflow_output_types = [], [] + ctx_annotation = None + + return message_type, ctx_annotation, output_types, workflow_output_types + + _FRAMEWORK_LIFECYCLE_EVENT_TYPES: tuple[type[WorkflowEvent], ...] = cast( tuple[type[WorkflowEvent], ...], tuple(get_args(WorkflowLifecycleEvent)) @@ -36,11 +272,47 @@ _FRAMEWORK_LIFECYCLE_EVENT_TYPES: tuple[type[WorkflowEvent], ...] = cast( ) -class WorkflowContext(Generic[T_Out]): - """Context for executors in a workflow. +class WorkflowContext(Generic[T_Out, T_W_Out]): + """Execution context that enables executors to interact with workflows and other executors. - This class is used to provide a way for executors to interact with the workflow - context and shared state, while preventing direct access to the runtime context. + ## Overview + WorkflowContext provides a controlled interface for executors to send messages, yield outputs, + manage state, and interact with the broader workflow ecosystem. It enforces type safety through + generic parameters while preventing direct access to internal runtime components. + + ## Type Parameters + The context is parameterized to enforce type safety for different operations: + + ### WorkflowContext (no parameters) + For executors that only perform side effects without sending messages or yielding outputs: + ```python + async def log_handler(message: str, ctx: WorkflowContext) -> None: + print(f"Received: {message}") # Only side effects + ``` + + ### WorkflowContext[T_Out] + Enables sending messages of type T_Out to other executors: + ```python + async def processor(message: str, ctx: WorkflowContext[int]) -> None: + result = len(message) + await ctx.send_message(result) # Send int to downstream executors + ``` + + ### WorkflowContext[T_Out, T_W_Out] + Enables both sending messages (T_Out) and yielding workflow outputs (T_W_Out): + ```python + async def dual_output(message: str, ctx: WorkflowContext[int, str]) -> None: + await ctx.send_message(42) # Send int message + await ctx.yield_output("complete") # Yield str workflow output + ``` + + ### Union Types + Multiple types can be specified using union notation: + ```python + async def flexible(message: str, ctx: WorkflowContext[int | str, bool | dict]) -> None: + await ctx.send_message("text") # or send 42 + await ctx.yield_output(True) # or yield {"status": "done"} + ``` """ def __init__( @@ -105,6 +377,17 @@ class WorkflowContext(Generic[T_Out]): await self._runner_context.send_message(msg) + async def yield_output(self, output: T_W_Out) -> None: + """Set the output of the workflow. + + Args: + output: The output to yield. This must conform to the workflow output type(s) + declared on this context. + """ + with _framework_event_origin(): + event = WorkflowOutputEvent(data=output, source_executor_id=self._executor_id) + await self._runner_context.add_event(event) + async def add_event(self, event: WorkflowEvent) -> None: """Add an event to the workflow context.""" if event.origin == WorkflowEventSource.EXECUTOR and isinstance(event, _FRAMEWORK_LIFECYCLE_EVENT_TYPES): diff --git a/python/packages/main/agent_framework/_workflow/_workflow_executor.py b/python/packages/main/agent_framework/_workflow/_workflow_executor.py new file mode 100644 index 0000000000..3f9ded3681 --- /dev/null +++ b/python/packages/main/agent_framework/_workflow/_workflow_executor.py @@ -0,0 +1,437 @@ +# Copyright (c) Microsoft. All rights reserved. + +import logging +import uuid +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from ._workflow import Workflow + +from pydantic import Field + +from ._events import ( + WorkflowErrorEvent, + WorkflowFailedEvent, + WorkflowRunState, +) +from ._executor import ( + Executor, + RequestInfoExecutor, + RequestInfoMessage, + SubWorkflowRequestInfo, + SubWorkflowResponse, + handler, +) +from ._workflow_context import WorkflowContext + +logger = logging.getLogger(__name__) + + +@dataclass +class ExecutionContext: + """Context for tracking a single sub-workflow execution.""" + + execution_id: str + collected_responses: dict[str, Any] # request_id -> response_data + expected_response_count: int + pending_requests: dict[str, Any] # request_id -> original request data + + +class WorkflowExecutor(Executor): + """An executor that wraps a workflow to enable hierarchical workflow composition. + + ## Overview + WorkflowExecutor makes a workflow behave as a single executor within a parent workflow, + enabling nested workflow architectures. It handles the complete lifecycle of sub-workflow + execution including event processing, output forwarding, and request/response coordination + between parent and child workflows. + + ## Execution Model + When invoked, WorkflowExecutor: + 1. Starts the wrapped workflow with the input message + 2. Runs the sub-workflow to completion or until it needs external input + 3. Processes the sub-workflow's complete event stream after execution + 4. Forwards outputs to the parent workflow's event stream + 5. Handles external requests by routing them to the parent workflow + 6. Accumulates responses and resumes sub-workflow execution + + ## Event Stream Processing + WorkflowExecutor processes events after sub-workflow completion: + + ### Output Forwarding + All outputs from the sub-workflow are automatically forwarded to the parent: + ```python + # Sub-workflow yields outputs + await ctx.yield_output("sub-workflow result") + + # WorkflowExecutor forwards to parent via ctx.send_message() + # Parent receives the output as a regular message + ``` + + ### Request/Response Coordination + When sub-workflows need external information: + ```python + # Sub-workflow makes request + request = MyDataRequest(query="user info") + # RequestInfoExecutor emits RequestInfoEvent + + # WorkflowExecutor wraps and forwards to parent + wrapped = SubWorkflowRequestInfo(request_id="...", sub_workflow_id="child_workflow", data=request) + # Parent workflow can intercept via @intercepts_request + ``` + + ### State Management + WorkflowExecutor maintains execution state across request/response cycles: + - Tracks pending requests by request_id + - Accumulates responses until all expected responses are received + - Resumes sub-workflow execution with complete response batch + - Handles concurrent executions and multiple pending requests + + ## Type System Integration + WorkflowExecutor inherits its type signature from the wrapped workflow: + + ### Input Types + Matches the wrapped workflow's start executor input types: + ```python + # If sub-workflow accepts str, WorkflowExecutor accepts str + workflow_executor = WorkflowExecutor(my_workflow, id="wrapper") + assert workflow_executor.input_types == my_workflow.input_types + ``` + + ### Output Types + Combines sub-workflow outputs with request coordination types: + ```python + # Includes all sub-workflow output types + # Plus SubWorkflowRequestInfo if sub-workflow can make requests + output_types = workflow.output_types + [SubWorkflowRequestInfo] # if applicable + ``` + + ## Error Handling + WorkflowExecutor propagates sub-workflow failures: + - Captures WorkflowFailedEvent from sub-workflow + - Converts to WorkflowErrorEvent in parent context + - Provides detailed error information including sub-workflow ID + + ## Concurrent Execution Support + WorkflowExecutor fully supports multiple concurrent sub-workflow executions: + + ### Per-Execution State Isolation + Each sub-workflow invocation creates an isolated ExecutionContext: + ```python + # Multiple concurrent invocations are supported + workflow_executor = WorkflowExecutor(my_workflow, id="concurrent_executor") + + # Each invocation gets its own execution context + # Execution 1: processes input_1 independently + # Execution 2: processes input_2 independently + # No state interference between executions + ``` + + ### Request/Response Coordination + Responses are correctly routed to the originating execution: + - Each execution tracks its own pending requests and expected responses + - Request-to-execution mapping ensures responses reach the correct sub-workflow + - Response accumulation is isolated per execution + - Automatic cleanup when execution completes + + ### Memory Management + - Unlimited concurrent executions supported + - Each execution has unique UUID-based identification + - Cleanup of completed execution contexts + - Thread-safe state management for concurrent access + + ### 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 + ```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_state() or {} + state["processed"] = data + await ctx.set_state(state) + + + # Avoid: Stateful executor with instance variables + class StatefulExecutor(Executor): + def __init__(self): + super().__init__(id="stateful") + self.data = [] # This will be shared across concurrent executions! + ``` + + ## Integration with Parent Workflows + Parent workflows can intercept sub-workflow requests: + ```python + class ParentExecutor(Executor): + @intercepts_request + async def handle_child_request( + self, request: MyDataRequest, ctx: WorkflowContext[Any] + ) -> RequestResponse[MyDataRequest, str]: + # Handle request locally or forward to external source + if self.can_handle_locally(request): + return RequestResponse.handled("local result") + return RequestResponse.forward() # Send to external handler + ``` + + ## Implementation Notes + - Sub-workflows run to completion before processing their results + - Event processing is atomic - all outputs are forwarded before requests + - Response accumulation ensures sub-workflows receive complete response batches + - Execution state is maintained for proper resumption after external requests + - Concurrent executions are fully isolated and do not interfere with each other + """ + + workflow: "Workflow" = Field(description="The workflow to execute as a sub-workflow") + + def __init__(self, workflow: "Workflow", id: str, **kwargs: Any): + """Initialize the WorkflowExecutor. + + Args: + workflow: The workflow to execute as a sub-workflow. + id: Unique identifier for this executor. + **kwargs: Additional keyword arguments passed to the parent constructor. + """ + kwargs.update({"workflow": workflow}) + super().__init__(id, **kwargs) + + # Track execution contexts for concurrent sub-workflow executions + 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._active_executions: int = 0 # Count of active sub-workflow executions + + @property + def input_types(self) -> list[type[Any]]: + """Get the input types based on the underlying workflow's input types. + + Returns: + A list of input types that the underlying workflow can accept. + """ + return self.workflow.input_types + + @property + def output_types(self) -> list[type[Any]]: + """Get the output types based on the underlying workflow's output types. + + Returns: + A list of output types that the underlying workflow can produce. + Includes SubWorkflowRequestInfo if the sub-workflow contains RequestInfoExecutor. + """ + output_types = list(self.workflow.output_types) + + # Check if the sub-workflow contains a RequestInfoExecutor + # If so, this WorkflowExecutor can also output SubWorkflowRequestInfo messages + for executor in self.workflow.executors.values(): + if isinstance(executor, RequestInfoExecutor): + if SubWorkflowRequestInfo not in output_types: + output_types.append(SubWorkflowRequestInfo) + break + + return output_types + + @handler # No output_types - can send any completion data type + async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any]) -> None: + """Execute the sub-workflow with raw input data. + + This handler starts a new sub-workflow execution. When the sub-workflow + needs external information, it pauses and sends a request to the parent. + + Args: + input_data: The input data to send to the sub-workflow. + ctx: The workflow context from the parent. + """ + # Skip SubWorkflowResponse and SubWorkflowRequestInfo - they have specific handlers + if isinstance(input_data, (SubWorkflowResponse, SubWorkflowRequestInfo)): + logger.debug(f"WorkflowExecutor {self.id} ignoring input of type {type(input_data)}") + return + + # Create execution context for this sub-workflow run + execution_id = str(uuid.uuid4()) + execution_context = ExecutionContext( + execution_id=execution_id, + collected_responses={}, + expected_response_count=0, + pending_requests={}, + ) + self._execution_contexts[execution_id] = execution_context + + # Track this execution + self._active_executions += 1 + + logger.debug(f"WorkflowExecutor {self.id} starting sub-workflow {self.workflow.id} execution {execution_id}") + + try: + # Run the sub-workflow and collect all events + result = await self.workflow.run(input_data) + + logger.debug( + f"WorkflowExecutor {self.id} sub-workflow {self.workflow.id} " + f"execution {execution_id} completed with {len(result)} events" + ) + + # Process the workflow result using shared logic + await self._process_workflow_result(result, execution_context, ctx) + finally: + # Clean up execution context if it's completed (no pending requests) + if execution_id in self._execution_contexts: + exec_ctx = self._execution_contexts[execution_id] + if not exec_ctx.pending_requests: + del self._execution_contexts[execution_id] + self._active_executions -= 1 + + async def _process_workflow_result( + self, result: Any, execution_context: ExecutionContext, ctx: WorkflowContext[Any] + ) -> None: + """Process the result from a workflow execution. + + This method handles the common logic for processing outputs, request info events, + and final states that is shared between process_workflow and handle_response. + + Args: + result: The workflow execution result. + execution_context: The execution context for this sub-workflow run. + ctx: The workflow context. + """ + # Collect all events from the workflow + request_info_events = result.get_request_info_events() + outputs = result.get_outputs() + final_state = result.get_final_state() + logger.debug( + f"WorkflowExecutor {self.id} processing workflow result with " + f"{len(outputs)} outputs and {len(request_info_events)} request info events, " + f"final state: {final_state}" + ) + + # Process outputs + for output in outputs: + await ctx.send_message(output) + + # Process request info events + for event in request_info_events: + # Track the pending request in execution context + execution_context.pending_requests[event.request_id] = event.data + # Map request to execution for response routing + self._request_to_execution[event.request_id] = execution_context.execution_id + # Wrap request with routing context and send to parent + if not isinstance(event.data, RequestInfoMessage): + raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}") + wrapped_request = SubWorkflowRequestInfo( + request_id=event.request_id, + sub_workflow_id=self.id, + data=event.data, + ) + await ctx.send_message(wrapped_request) + + # Update expected response count for this execution + execution_context.expected_response_count = len(request_info_events) + + # Handle final state + if final_state == WorkflowRunState.FAILED: + # Find the WorkflowFailedEvent. + failed_events = [e for e in result if isinstance(e, WorkflowFailedEvent)] + if failed_events: + failed_event = failed_events[0] + error_type = failed_event.details.error_type + error_message = failed_event.details.message + exception = Exception( + f"Sub-workflow {self.workflow.id} failed with error: {error_type} - {error_message}" + ) + error_event = WorkflowErrorEvent( + data=exception, + ) + await ctx.add_event(error_event) + self._active_executions -= 1 + elif final_state == WorkflowRunState.IDLE: + # Sub-workflow is idle - nothing more to do now + logger.debug(f"Sub-workflow {self.workflow.id} is idle with {self._active_executions} active executions") + self._active_executions -= 1 # Treat idle as completion for now + elif final_state == WorkflowRunState.CANCELLED: + # Sub-workflow was cancelled - treat as completion + logger.debug( + f"Sub-workflow {self.workflow.id} was cancelled with {self._active_executions} active executions" + ) + self._active_executions -= 1 + elif final_state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS: + # Sub-workflow is still running with pending requests + logger.debug( + f"Sub-workflow {self.workflow.id} is still in progress with {len(request_info_events)} " + f"pending requests with {self._active_executions} active executions" + ) + elif final_state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS: + # Sub-workflow is idle but has pending requests + logger.debug( + f"Sub-workflow {self.workflow.id} is idle with pending requests: " + f"{len(request_info_events)} with {self._active_executions} active executions" + ) + else: + raise RuntimeError(f"Unexpected final state: {final_state}") + + @handler + async def handle_response( + self, + response: SubWorkflowResponse, + ctx: WorkflowContext[Any], + ) -> None: + """Handle response from parent for a forwarded request. + + This handler accumulates responses and only resumes the sub-workflow + when all expected responses have been received for that execution. + + Args: + response: The response to a previous request. + ctx: The workflow context. + """ + # Find the execution context for this request + execution_id = self._request_to_execution.get(response.request_id) + if not execution_id or execution_id not in self._execution_contexts: + logger.warning( + f"WorkflowExecutor {self.id} received response for unknown request_id: {response.request_id}, ignoring" + ) + return + + execution_context = self._execution_contexts[execution_id] + + # Check if we have this pending request in the execution context + if response.request_id not in execution_context.pending_requests: + logger.warning( + f"WorkflowExecutor {self.id} received response for unknown request_id: " + f"{response.request_id} in execution {execution_id}, ignoring" + ) + return + + # Remove the request from pending list and request mapping + execution_context.pending_requests.pop(response.request_id, None) + self._request_to_execution.pop(response.request_id, None) + + # Accumulate the response in this execution's context + execution_context.collected_responses[response.request_id] = response.data + + # Check if we have all expected responses for this execution + if len(execution_context.collected_responses) < execution_context.expected_response_count: + logger.debug( + f"WorkflowExecutor {self.id} execution {execution_id} waiting for more responses: " + f"{len(execution_context.collected_responses)}/{execution_context.expected_response_count} received" + ) + return # Wait for more responses + + # Send all collected responses to the sub-workflow + responses_to_send = dict(execution_context.collected_responses) + execution_context.collected_responses.clear() # Clear for next batch + + try: + # Resume the sub-workflow with all collected responses + result = await self.workflow.send_responses(responses_to_send) + + # Process the workflow result using shared logic + await self._process_workflow_result(result, execution_context, ctx) + finally: + # Clean up execution context if it's completed (no pending requests) + if not execution_context.pending_requests: + del self._execution_contexts[execution_id] + self._active_executions -= 1 diff --git a/python/packages/main/tests/workflow/test_checkpoint_validation.py b/python/packages/main/tests/workflow/test_checkpoint_validation.py index defdbb1717..f1b52f6ac6 100644 --- a/python/packages/main/tests/workflow/test_checkpoint_validation.py +++ b/python/packages/main/tests/workflow/test_checkpoint_validation.py @@ -1,8 +1,9 @@ # Copyright (c) Microsoft. All rights reserved. import pytest +from typing_extensions import Never -from agent_framework import WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler +from agent_framework import WorkflowBuilder, WorkflowContext, WorkflowRunState, WorkflowStatusEvent, handler from agent_framework._workflow._checkpoint import InMemoryCheckpointStorage from agent_framework._workflow._executor import Executor @@ -15,8 +16,8 @@ class StartExecutor(Executor): class FinishExecutor(Executor): @handler - async def finish(self, message: str, ctx: WorkflowContext[None]) -> None: - await ctx.add_event(WorkflowCompletedEvent(message)) + async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None: + await ctx.yield_output(message) def build_workflow(storage: InMemoryCheckpointStorage, finish_id: str = "finish"): @@ -70,4 +71,4 @@ async def test_resume_succeeds_when_graph_matches() -> None: ) ] - assert any(isinstance(event, WorkflowCompletedEvent) for event in events) + assert any(isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE for event in events) diff --git a/python/packages/main/tests/workflow/test_concurrent.py b/python/packages/main/tests/workflow/test_concurrent.py index bfcf8802da..ccb7b9ee3d 100644 --- a/python/packages/main/tests/workflow/test_concurrent.py +++ b/python/packages/main/tests/workflow/test_concurrent.py @@ -12,8 +12,10 @@ from agent_framework import ( ConcurrentBuilder, Executor, Role, - WorkflowCompletedEvent, WorkflowContext, + WorkflowOutputEvent, + WorkflowRunState, + WorkflowStatusEvent, handler, ) @@ -57,15 +59,19 @@ async def test_concurrent_default_aggregator_emits_single_user_and_assistants() wf = ConcurrentBuilder().participants([e1, e2, e3]).build() - completed: WorkflowCompletedEvent | None = None + completed = False + output: list[ChatMessage] | None = None async for ev in wf.run_stream("prompt: hello world"): - if isinstance(ev, WorkflowCompletedEvent): - completed = ev + if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: + completed = True + elif isinstance(ev, WorkflowOutputEvent): + output = cast(list[ChatMessage], ev.data) + if completed and output is not None: break - assert completed is not None - assert isinstance(completed.data, list) - messages: list[ChatMessage] = cast(list[ChatMessage], completed.data) # type: ignore + assert completed + assert output is not None + messages: list[ChatMessage] = output # Expect one user message + one assistant message per participant assert len(messages) == 1 + 3 @@ -91,16 +97,21 @@ async def test_concurrent_custom_aggregator_callback_is_used() -> None: wf = ConcurrentBuilder().participants([e1, e2]).with_aggregator(summarize).build() - completed: WorkflowCompletedEvent | None = None + completed = False + output: str | None = None async for ev in wf.run_stream("prompt: custom"): - if isinstance(ev, WorkflowCompletedEvent): - completed = ev + if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: + completed = True + elif isinstance(ev, WorkflowOutputEvent): + output = cast(str, ev.data) + if completed and output is not None: break - assert completed is not None + assert completed + assert output is not None # Custom aggregator returns a string payload - assert isinstance(completed.data, str) - assert completed.data == "One | Two" + assert isinstance(output, str) + assert output == "One | Two" async def test_concurrent_custom_aggregator_sync_callback_is_used() -> None: @@ -108,7 +119,7 @@ async def test_concurrent_custom_aggregator_sync_callback_is_used() -> None: e2 = _FakeAgentExec("agentB", "Two") # Sync callback with ctx parameter (should run via asyncio.to_thread) - def summarize_sync(results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> str: # type: ignore[unused-argument] + def summarize_sync(results: list[AgentExecutorResponse], _ctx: WorkflowContext[Any]) -> str: # type: ignore[unused-argument] texts: list[str] = [] for r in results: msgs: list[ChatMessage] = r.agent_run_response.messages @@ -117,15 +128,20 @@ async def test_concurrent_custom_aggregator_sync_callback_is_used() -> None: wf = ConcurrentBuilder().participants([e1, e2]).with_aggregator(summarize_sync).build() - completed: WorkflowCompletedEvent | None = None + completed = False + output: str | None = None async for ev in wf.run_stream("prompt: custom sync"): - if isinstance(ev, WorkflowCompletedEvent): - completed = ev + if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: + completed = True + elif isinstance(ev, WorkflowOutputEvent): + output = cast(str, ev.data) + if completed and output is not None: break - assert completed is not None - assert isinstance(completed.data, str) - assert completed.data == "One | Two" + assert completed + assert output is not None + assert isinstance(output, str) + assert output == "One | Two" def test_concurrent_custom_aggregator_uses_callback_name_for_id() -> None: diff --git a/python/packages/main/tests/workflow/test_full_conversation.py b/python/packages/main/tests/workflow/test_full_conversation.py index a56d05a7b1..be5fa6e333 100644 --- a/python/packages/main/tests/workflow/test_full_conversation.py +++ b/python/packages/main/tests/workflow/test_full_conversation.py @@ -4,6 +4,7 @@ from collections.abc import AsyncIterable from typing import Any from pydantic import PrivateAttr +from typing_extensions import Never from agent_framework import ( AgentExecutor, @@ -16,11 +17,13 @@ from agent_framework import ( SequentialBuilder, TextContent, WorkflowBuilder, - WorkflowCompletedEvent, - WorkflowContext, + WorkflowOutputEvent, + WorkflowRunState, + WorkflowStatusEvent, handler, ) from agent_framework._workflow._executor import AgentExecutorResponse, Executor +from agent_framework._workflow._workflow_context import WorkflowContext class _SimpleAgent(BaseAgent): @@ -54,19 +57,17 @@ class _CaptureFullConversation(Executor): """Captures AgentExecutorResponse.full_conversation and completes the workflow.""" @handler - async def capture(self, response: AgentExecutorResponse, ctx: WorkflowContext[None]) -> None: + async def capture(self, response: AgentExecutorResponse, ctx: WorkflowContext[Never, dict]) -> None: full = response.full_conversation # The AgentExecutor contract guarantees full_conversation is populated. assert full is not None - await ctx.add_event( - WorkflowCompletedEvent( - data={ - "length": len(full), - "roles": [m.role for m in full], - "texts": [m.text for m in full], - } - ) - ) + payload = { + "length": len(full), + "roles": [m.role for m in full], + "texts": [m.text for m in full], + } + await ctx.yield_output(payload) + pass async def test_agent_executor_populates_full_conversation_non_streaming() -> None: @@ -78,15 +79,20 @@ async def test_agent_executor_populates_full_conversation_non_streaming() -> Non wf = WorkflowBuilder().set_start_executor(agent_exec).add_edge(agent_exec, capturer).build() # Act: run with a simple user prompt - completed: WorkflowCompletedEvent | None = None + completed = False + output: dict | None = None async for ev in wf.run_stream("hello world"): - if isinstance(ev, WorkflowCompletedEvent): - completed = ev + if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: + completed = True + elif isinstance(ev, WorkflowOutputEvent): + output = ev.data # type: ignore[assignment] + if completed and output is not None: break # Assert: full_conversation contains [user("hello world"), assistant("agent-reply")] - assert completed is not None - payload = completed.data # type: ignore[assignment] + assert completed + assert output is not None + payload = output assert isinstance(payload, dict) assert payload["length"] == 2 assert payload["roles"][0] == Role.USER and "hello world" in (payload["texts"][0] or "") @@ -148,7 +154,7 @@ async def test_sequential_adapter_uses_full_conversation() -> None: # Act async for ev in wf.run_stream("hello seq"): - if isinstance(ev, WorkflowCompletedEvent): + if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: break # Assert: second agent should have seen the user prompt and A1's assistant reply diff --git a/python/packages/main/tests/workflow/test_function_executor.py b/python/packages/main/tests/workflow/test_function_executor.py index 1ca16e9c66..229e399b09 100644 --- a/python/packages/main/tests/workflow/test_function_executor.py +++ b/python/packages/main/tests/workflow/test_function_executor.py @@ -3,11 +3,11 @@ from typing import Any import pytest +from typing_extensions import Never from agent_framework import ( FunctionExecutor, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, executor, ) @@ -28,9 +28,9 @@ class TestFunctionExecutor: assert len(func_exec._handlers) == 1 assert str in func_exec._handlers - # Check instance handler spec was created - assert len(func_exec._instance_handler_specs) == 1 - spec = func_exec._instance_handler_specs[0] + # Check handler spec was created + assert len(func_exec._handler_specs) == 1 + spec = func_exec._handler_specs[0] assert spec["name"] == "process_string" assert spec["message_type"] is str assert spec["output_types"] == [str] @@ -47,7 +47,7 @@ class TestFunctionExecutor: assert int in process_int._handlers # Check spec - spec = process_int._instance_handler_specs[0] + spec = process_int._handler_specs[0] assert spec["message_type"] is int assert spec["output_types"] == [int] @@ -81,7 +81,7 @@ class TestFunctionExecutor: assert int in simple_no_parens._handlers def test_union_output_types(self): - """Test that union output types are properly inferred.""" + """Test that union output types are properly inferred for both messages and workflow outputs.""" @executor async def multi_output(text: str, ctx: WorkflowContext[str | int]) -> None: @@ -90,29 +90,56 @@ class TestFunctionExecutor: else: await ctx.send_message(text.upper()) - spec = multi_output._instance_handler_specs[0] + spec = multi_output._handler_specs[0] assert set(spec["output_types"]) == {str, int} + assert spec["workflow_output_types"] == [] # No workflow outputs defined + + # Test union types for workflow outputs too + @executor + async def multi_workflow_output(data: str, ctx: WorkflowContext[Never, str | int | bool]) -> None: + if data.isdigit(): + await ctx.yield_output(int(data)) + elif data.lower() in ("true", "false"): + await ctx.yield_output(data.lower() == "true") + else: + await ctx.yield_output(data.upper()) + + workflow_spec = multi_workflow_output._handler_specs[0] + assert workflow_spec["output_types"] == [] # None means no message outputs + assert set(workflow_spec["workflow_output_types"]) == {str, int, bool} def test_none_output_type(self): - """Test WorkflowContext[None] produces empty output types.""" + """Test WorkflowContext produces empty output types.""" @executor - async def no_output(data: Any, ctx: WorkflowContext[None]) -> None: + async def no_output(data: Any, ctx: WorkflowContext) -> None: # This executor doesn't send any messages pass - spec = no_output._instance_handler_specs[0] + spec = no_output._handler_specs[0] assert spec["output_types"] == [] + assert spec["workflow_output_types"] == [] # No workflow outputs defined def test_any_output_type(self): - """Test WorkflowContext[Any] produces empty output types.""" + """Test WorkflowContext[Any] and WorkflowContext[Any, Any] produce Any output types.""" @executor async def any_output(data: str, ctx: WorkflowContext[Any]) -> None: await ctx.send_message("result") - spec = any_output._instance_handler_specs[0] - assert spec["output_types"] == [] + spec = any_output._handler_specs[0] + assert spec["output_types"] == [Any] + assert spec["workflow_output_types"] == [] # No workflow outputs defined + + # Test both parameters as Any + @executor + async def any_both_output(data: str, ctx: WorkflowContext[Any, Any]) -> None: + await ctx.send_message("message") + await ctx.yield_output("workflow_output") + + both_spec = any_both_output._handler_specs[0] + assert both_spec["output_types"] == [Any] + assert both_spec["workflow_output_types"] == [Any] def test_validation_errors(self): """Test various validation errors in function signatures.""" @@ -121,13 +148,17 @@ class TestFunctionExecutor: async def no_params() -> None: pass - with pytest.raises(ValueError, match="one or two parameters"): + with pytest.raises( + ValueError, match="must have \\(message: T\\) or \\(message: T, ctx: WorkflowContext\\[U\\]\\)" + ): FunctionExecutor(no_params) # type: ignore async def too_many_params(data: str, ctx: WorkflowContext[str], extra: int) -> None: pass - with pytest.raises(ValueError, match="one or two parameters"): + with pytest.raises( + ValueError, match="must have \\(message: T\\) or \\(message: T, ctx: WorkflowContext\\[U\\]\\)" + ): FunctionExecutor(too_many_params) # type: ignore # Missing message type annotation @@ -141,22 +172,24 @@ class TestFunctionExecutor: async def no_ctx_type(data: str, ctx) -> None: # type: ignore pass - with pytest.raises(ValueError, match="annotated as WorkflowContext"): + with pytest.raises(ValueError, match="must have a WorkflowContext"): FunctionExecutor(no_ctx_type) # type: ignore # Wrong ctx type async def wrong_ctx_type(data: str, ctx: str) -> None: # type: ignore pass - with pytest.raises(ValueError, match="WorkflowContext\\[T\\]"): + with pytest.raises(ValueError, match="must be annotated as WorkflowContext"): FunctionExecutor(wrong_ctx_type) # type: ignore - # Unparameterized WorkflowContext + # Unparameterized WorkflowContext is now allowed async def unparameterized_ctx(data: str, ctx: WorkflowContext) -> None: # type: ignore pass - with pytest.raises(ValueError, match="concrete T"): - FunctionExecutor(unparameterized_ctx) # type: ignore + # This should now succeed since unparameterized WorkflowContext is allowed + executor = FunctionExecutor(unparameterized_ctx) + assert executor.output_types == [] # Unparameterized has no inferred types + assert executor.workflow_output_types == [] # No workflow output types async def test_execution_in_workflow(self): """Test that FunctionExecutor works properly in a workflow.""" @@ -167,18 +200,28 @@ class TestFunctionExecutor: await ctx.send_message(result) @executor(id="reverse") - async def reverse_text(text: str, ctx: WorkflowContext[Any]) -> None: + async def reverse_text(text: str, ctx: WorkflowContext[Any, str]) -> None: result = text[::-1] - await ctx.add_event(WorkflowCompletedEvent(result)) + await ctx.yield_output(result) + + # Verify type inference for both executors + upper_spec = to_upper._handler_specs[0] + assert upper_spec["output_types"] == [str] + assert upper_spec["workflow_output_types"] == [] # No workflow outputs + + reverse_spec = reverse_text._handler_specs[0] + assert reverse_spec["output_types"] == [Any] # First parameter is Any + assert reverse_spec["workflow_output_types"] == [str] # Second parameter is str workflow = WorkflowBuilder().add_edge(to_upper, reverse_text).set_start_executor(to_upper).build() # Run workflow events = await workflow.run("hello world") - completed = events.get_completed_event() + outputs = events.get_outputs() - assert completed is not None - assert completed.data == "DLROW OLLEH" + # Assert that we got the expected output + assert len(outputs) == 1 + assert outputs[0] == "DLROW OLLEH" def test_can_handle_method(self): """Test that can_handle method works with instance handlers.""" @@ -204,12 +247,13 @@ class TestFunctionExecutor: await ctx.send_message(message) with pytest.raises(ValueError, match="Handler for type .* already registered"): - func_exec.register_instance_handler( + func_exec._register_instance_handler( name="second", func=second_handler, message_type=str, ctx_annotation=WorkflowContext[str], output_types=[str], + workflow_output_types=[], ) def test_complex_type_annotations(self): @@ -220,7 +264,7 @@ class TestFunctionExecutor: result = {item: len(item) for item in items} await ctx.send_message(result) - spec = process_list._instance_handler_specs[0] + spec = process_list._handler_specs[0] assert spec["message_type"] == list[str] assert spec["output_types"] == [dict[str, int]] @@ -236,7 +280,7 @@ class TestFunctionExecutor: assert str in process_simple._handlers # Check spec - single parameter functions have no output types since they can't send messages - spec = process_simple._instance_handler_specs[0] + spec = process_simple._handler_specs[0] assert spec["message_type"] is str assert spec["output_types"] == [] assert spec["ctx_annotation"] is None @@ -296,7 +340,7 @@ class TestFunctionExecutor: assert str in process_sync._handlers # Check spec - sync single parameter functions have no output types - spec = process_sync._instance_handler_specs[0] + spec = process_sync._handler_specs[0] assert spec["message_type"] is str assert spec["output_types"] == [] assert spec["ctx_annotation"] is None @@ -314,7 +358,7 @@ class TestFunctionExecutor: assert int in sync_with_ctx._handlers # Check spec - sync functions with context can infer output types - spec = sync_with_ctx._instance_handler_specs[0] + spec = sync_with_ctx._handler_specs[0] assert spec["message_type"] is int assert spec["output_types"] == [int] @@ -385,9 +429,18 @@ class TestFunctionExecutor: # In practice, the wrapper handles the async conversion @executor(id="async_reverse") - async def reverse_async(text: str, ctx: WorkflowContext[Any]): + async def reverse_async(text: str, ctx: WorkflowContext[Any, str]): result = text[::-1] - await ctx.add_event(WorkflowCompletedEvent(result)) + await ctx.yield_output(result) + + # Verify type inference for sync and async functions + sync_spec = to_upper_sync._handler_specs[0] + assert sync_spec["output_types"] == [str] + assert sync_spec["workflow_output_types"] == [] # No workflow outputs + + async_spec = reverse_async._handler_specs[0] + assert async_spec["output_types"] == [Any] # First parameter is Any + assert async_spec["workflow_output_types"] == [str] # Second parameter is str # Verify the executors can handle their input types assert to_upper_sync.can_handle("hello") diff --git a/python/packages/main/tests/workflow/test_magentic.py b/python/packages/main/tests/workflow/test_magentic.py index 13e170cdc3..e603673aa8 100644 --- a/python/packages/main/tests/workflow/test_magentic.py +++ b/python/packages/main/tests/workflow/test_magentic.py @@ -23,9 +23,11 @@ from agent_framework import ( RequestInfoEvent, Role, TextContent, - WorkflowCompletedEvent, WorkflowContext, WorkflowEvent, # type: ignore # noqa: E402 + WorkflowOutputEvent, + WorkflowRunState, + WorkflowStatusEvent, handler, ) from agent_framework._agents import BaseAgent @@ -169,15 +171,20 @@ async def test_magentic_workflow_plan_review_approval_to_completion(): req_event = ev assert req_event is not None - completed: WorkflowCompletedEvent | None = None + completed = False + output: ChatMessage | None = None async for ev in wf.send_responses_streaming({ req_event.request_id: MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.APPROVE) }): - if isinstance(ev, WorkflowCompletedEvent): - completed = ev + if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: + completed = True + elif isinstance(ev, WorkflowOutputEvent): + output = ev.data # type: ignore[assignment] + if completed and output is not None: break - assert completed is not None - assert isinstance(getattr(completed, "data", None), ChatMessage) + assert completed + assert output is not None + assert isinstance(output, ChatMessage) async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds(): @@ -210,7 +217,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: WorkflowCompletedEvent | None = None + completed = False async for ev in wf.send_responses_streaming({ req_event.request_id: MagenticPlanReviewReply( decision=MagenticPlanReviewDecision.APPROVE, @@ -219,11 +226,11 @@ async def test_magentic_plan_review_approve_with_comments_replans_and_proceeds() }): if isinstance(ev, RequestInfoEvent) and ev.request_type is MagenticPlanReviewRequest: saw_second_review = True - if isinstance(ev, WorkflowCompletedEvent): - completed = ev + if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: + completed = True break - assert completed is not None + assert completed assert manager.replan_count >= 1 assert saw_second_review is False # Replan from FakeManager updates facts/plan to include A2 / Do Z @@ -245,9 +252,14 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result(): if len(events) > 50: break - completed = next((e for e in events if isinstance(e, WorkflowCompletedEvent)), None) - assert completed is not None - data = getattr(completed, "data", None) + idle_status = next( + (e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE), None + ) + assert idle_status is not None + # Check that we got workflow output via WorkflowOutputEvent + output_event = next((e for e in events if isinstance(e, WorkflowOutputEvent)), None) + assert output_event is not None + data = output_event.data assert isinstance(data, ChatMessage) assert data.role == Role.ASSISTANT diff --git a/python/packages/main/tests/workflow/test_runner.py b/python/packages/main/tests/workflow/test_runner.py index 722856134a..cc2486ef71 100644 --- a/python/packages/main/tests/workflow/test_runner.py +++ b/python/packages/main/tests/workflow/test_runner.py @@ -9,10 +9,11 @@ from agent_framework import ( AgentExecutorResponse, AgentRunResponse, Executor, - WorkflowCompletedEvent, WorkflowContext, WorkflowEvent, - WorkflowEventSource, + WorkflowOutputEvent, + WorkflowRunState, + WorkflowStatusEvent, handler, ) from agent_framework._workflow._edge import SingleEdgeGroup @@ -32,11 +33,12 @@ class MockExecutor(Executor): """A mock executor for testing purposes.""" @handler - async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage]) -> None: + async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None: if message.data < 10: await ctx.send_message(MockMessage(data=message.data + 1)) else: - await ctx.add_event(WorkflowCompletedEvent(data=message.data)) + await ctx.yield_output(message.data) + pass def test_create_runner(): @@ -77,18 +79,14 @@ async def test_runner_run_until_convergence(): result: int | None = None await executor_a.execute( MockMessage(data=0), - WorkflowContext( - executor_id=executor_a.id, - source_executor_ids=["START"], - shared_state=shared_state, - runner_context=ctx, - ), + ["START"], # source_executor_ids + shared_state, # shared_state + ctx, # runner_context ) async for event in runner.run_until_convergence(): assert isinstance(event, WorkflowEvent) - if isinstance(event, WorkflowCompletedEvent): + if isinstance(event, WorkflowOutputEvent): result = event.data - assert event.origin is WorkflowEventSource.EXECUTOR assert result is not None and result == 10 @@ -112,16 +110,13 @@ async def test_runner_run_until_convergence_not_completed(): await executor_a.execute( MockMessage(data=0), - WorkflowContext( - executor_id=executor_a.id, - source_executor_ids=["START"], - shared_state=shared_state, - runner_context=ctx, - ), + ["START"], # source_executor_ids + shared_state, # shared_state + ctx, # runner_context ) with pytest.raises(RuntimeError, match="Runner did not converge after 5 iterations."): async for event in runner.run_until_convergence(): - assert not isinstance(event, WorkflowCompletedEvent) + assert not isinstance(event, WorkflowStatusEvent) or event.state != WorkflowRunState.IDLE async def test_runner_already_running(): @@ -143,12 +138,9 @@ async def test_runner_already_running(): await executor_a.execute( MockMessage(data=0), - WorkflowContext( - executor_id=executor_a.id, - source_executor_ids=["START"], - shared_state=shared_state, - runner_context=ctx, - ), + ["START"], # source_executor_ids + shared_state, # shared_state + ctx, # runner_context ) with pytest.raises(RuntimeError, match="Runner is already running."): @@ -172,6 +164,6 @@ async def test_runner_emits_runner_completion_for_agent_response_without_targets ) events: list[WorkflowEvent] = [event async for event in runner.run_until_convergence()] - completions = [e for e in events if isinstance(e, WorkflowCompletedEvent)] - assert completions - assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in completions) + # The runner should complete without errors when handling AgentExecutorResponse without targets + # No specific events are expected since there are no executors to process the message + assert isinstance(events, list) # Just verify the runner completed without errors diff --git a/python/packages/main/tests/workflow/test_sequential.py b/python/packages/main/tests/workflow/test_sequential.py index 930ec4db76..6aafdaa004 100644 --- a/python/packages/main/tests/workflow/test_sequential.py +++ b/python/packages/main/tests/workflow/test_sequential.py @@ -15,8 +15,10 @@ from agent_framework import ( Role, SequentialBuilder, TextContent, - WorkflowCompletedEvent, WorkflowContext, + WorkflowOutputEvent, + WorkflowRunState, + WorkflowStatusEvent, handler, ) @@ -66,15 +68,20 @@ async def test_sequential_agents_append_to_context() -> None: wf = SequentialBuilder().participants([a1, a2]).build() - completed: WorkflowCompletedEvent | None = None + completed = False + output: list[ChatMessage] | None = None async for ev in wf.run_stream("hello sequential"): - if isinstance(ev, WorkflowCompletedEvent): - completed = ev + if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: + completed = True + elif isinstance(ev, WorkflowOutputEvent): + output = ev.data # type: ignore[assignment] + if completed and output is not None: break - assert completed is not None - assert isinstance(completed.data, list) - msgs: list[ChatMessage] = completed.data # type: ignore[assignment] + assert completed + assert output is not None + assert isinstance(output, list) + msgs: list[ChatMessage] = output assert len(msgs) == 3 assert msgs[0].role == Role.USER and "hello sequential" in msgs[0].text assert msgs[1].role == Role.ASSISTANT and (msgs[1].author_name == "A1" or True) @@ -89,14 +96,19 @@ async def test_sequential_with_custom_executor_summary() -> None: wf = SequentialBuilder().participants([a1, summarizer]).build() - completed: WorkflowCompletedEvent | None = None + completed = False + output: list[ChatMessage] | None = None async for ev in wf.run_stream("topic X"): - if isinstance(ev, WorkflowCompletedEvent): - completed = ev + if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE: + completed = True + elif isinstance(ev, WorkflowOutputEvent): + output = ev.data # type: ignore[assignment] + if completed and output is not None: break - assert completed is not None - msgs: list[ChatMessage] = completed.data # type: ignore[assignment] + assert completed + assert output is not None + msgs: list[ChatMessage] = output # Expect: [user, A1 reply, summary] assert len(msgs) == 3 assert msgs[0].role == Role.USER diff --git a/python/packages/main/tests/workflow/test_serialization.py b/python/packages/main/tests/workflow/test_serialization.py index 59a754c9ac..8a4fae5db4 100644 --- a/python/packages/main/tests/workflow/test_serialization.py +++ b/python/packages/main/tests/workflow/test_serialization.py @@ -17,7 +17,7 @@ from agent_framework._workflow._edge import ( SwitchCaseEdgeGroupCase, SwitchCaseEdgeGroupDefault, ) -from agent_framework._workflow._executor import ( +from agent_framework._workflow._workflow_executor import ( WorkflowExecutor, ) diff --git a/python/packages/main/tests/workflow/test_simple_sub_workflow.py b/python/packages/main/tests/workflow/test_simple_sub_workflow.py index 73585fafe9..f57203a1cb 100644 --- a/python/packages/main/tests/workflow/test_simple_sub_workflow.py +++ b/python/packages/main/tests/workflow/test_simple_sub_workflow.py @@ -3,6 +3,8 @@ import asyncio from dataclasses import dataclass +from typing_extensions import Never + from agent_framework import ( Executor, WorkflowBuilder, @@ -33,13 +35,11 @@ class SimpleSubExecutor(Executor): super().__init__(id="simple_sub") @handler - async def process(self, request: SimpleRequest, ctx: WorkflowContext[None]) -> None: + async def process(self, request: SimpleRequest, ctx: WorkflowContext[Never, SimpleResponse]) -> None: """Process a simple request.""" - from agent_framework import WorkflowCompletedEvent - # Just echo back with prefix and complete response = SimpleResponse(result=f"processed: {request.text}") - await ctx.add_event(WorkflowCompletedEvent(data=response)) + await ctx.yield_output(response) class SimpleParent(Executor): @@ -57,7 +57,7 @@ class SimpleParent(Executor): await ctx.send_message(request, target_id="sub_workflow") @handler - async def collect(self, response: SimpleResponse, ctx: WorkflowContext[None]) -> None: + async def collect(self, response: SimpleResponse, ctx: WorkflowContext) -> None: """Collect the result.""" self.result = response @@ -72,7 +72,7 @@ async def test_simple_sub_workflow(): super().__init__(id="dummy") @handler - async def process(self, message: object, ctx: WorkflowContext[None]) -> None: + async def process(self, message: object, ctx: WorkflowContext) -> None: pass # Do nothing dummy = DummyExecutor() diff --git a/python/packages/main/tests/workflow/test_sub_workflow.py b/python/packages/main/tests/workflow/test_sub_workflow.py index 5cc1781a2c..2067734734 100644 --- a/python/packages/main/tests/workflow/test_sub_workflow.py +++ b/python/packages/main/tests/workflow/test_sub_workflow.py @@ -5,6 +5,7 @@ from dataclasses import dataclass from typing import Any from pydantic import Field +from typing_extensions import Never from agent_framework import ( Executor, @@ -12,7 +13,6 @@ from agent_framework import ( RequestInfoMessage, RequestResponse, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, WorkflowExecutor, handler, @@ -54,7 +54,7 @@ class EmailValidator(Executor): @handler async def validate_request( - self, request: EmailValidationRequest, ctx: WorkflowContext[RequestInfoMessage | ValidationResult] + self, request: EmailValidationRequest, ctx: WorkflowContext[RequestInfoMessage, ValidationResult] ) -> None: """Validate an email address.""" # Extract domain and check if it's approved @@ -62,7 +62,7 @@ class EmailValidator(Executor): if not domain: result = ValidationResult(email=request.email, is_valid=False, reason="Invalid email format") - await ctx.add_event(WorkflowCompletedEvent(data=result)) + await ctx.yield_output(result) return # Request domain check from external source @@ -71,7 +71,7 @@ class EmailValidator(Executor): @handler async def handle_domain_response( - self, response: RequestResponse[DomainCheckRequest, bool], ctx: WorkflowContext[ValidationResult] + self, response: RequestResponse[DomainCheckRequest, bool], ctx: WorkflowContext[Never, ValidationResult] ) -> None: """Handle domain check response with correlation.""" # Use the original email from the correlated response @@ -80,7 +80,7 @@ class EmailValidator(Executor): is_valid=response.data or False, reason="Domain approved" if response.data else "Domain not approved", ) - await ctx.add_event(WorkflowCompletedEvent(data=result)) + await ctx.yield_output(result) class ParentOrchestrator(Executor): @@ -114,7 +114,7 @@ class ParentOrchestrator(Executor): return RequestResponse[DomainCheckRequest, bool].forward() @handler - async def collect_result(self, result: ValidationResult, ctx: WorkflowContext[None]) -> None: + async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None: """Collect validation results.""" self.results.append(result) @@ -146,11 +146,11 @@ async def test_basic_sub_workflow() -> None: await ctx.send_message(request, target_id="email_workflow") @handler - async def collect(self, result: ValidationResult, ctx: WorkflowContext[None]) -> None: + async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None: self.result = result parent = SimpleParent() - workflow_executor = WorkflowExecutor(validation_workflow, id="email_workflow") + workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow") main_request_info = RequestInfoExecutor(id="main_request_info") main_workflow = ( @@ -199,7 +199,7 @@ async def test_sub_workflow_with_interception(): # Create parent workflow with interception parent = ParentOrchestrator(approved_domains={"example.com", "internal.org"}) - workflow_executor = WorkflowExecutor(validation_workflow, id="email_workflow") + workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow") parent_request_info = RequestInfoExecutor(id="request_info") main_workflow = ( @@ -275,7 +275,7 @@ async def test_conditional_forwarding() -> None: return RequestResponse[DomainCheckRequest, bool].forward() @handler - async def collect(self, result: ValidationResult, ctx: WorkflowContext[None]) -> None: + async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None: self.result = result # Setup workflows @@ -291,7 +291,7 @@ async def test_conditional_forwarding() -> None: ) parent = ConditionalParent() - workflow_executor = WorkflowExecutor(validation_workflow, id="email_workflow") + workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow") parent_request_info = RequestInfoExecutor(id="request_info") main_workflow = ( @@ -358,7 +358,7 @@ async def test_workflow_scoped_interception() -> None: return RequestResponse[DomainCheckRequest, bool].forward() @handler - async def collect(self, result: ValidationResult, ctx: WorkflowContext[None]) -> None: + async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None: self.results[result.email] = result # Create two identical sub-workflows @@ -377,8 +377,8 @@ async def test_workflow_scoped_interception() -> None: workflow_b = create_validation_workflow() parent = MultiWorkflowParent() - executor_a = WorkflowExecutor(workflow_a, id="workflow_a") - executor_b = WorkflowExecutor(workflow_b, id="workflow_b") + executor_a = WorkflowExecutor(workflow_a, "workflow_a") + executor_b = WorkflowExecutor(workflow_b, "workflow_b") parent_request_info = RequestInfoExecutor(id="request_info") main_workflow = ( @@ -407,9 +407,102 @@ async def test_workflow_scoped_interception() -> None: assert parent.results["user@random.com"].is_valid is True +async def test_concurrent_sub_workflow_execution() -> None: + """Test that WorkflowExecutor can handle multiple concurrent invocations properly.""" + + class ConcurrentProcessor(Executor): + """Processor that sends multiple concurrent requests to the same sub-workflow.""" + + results: list[ValidationResult] = Field(default_factory=list) + + def __init__(self, **kwargs: Any): + super().__init__(id="concurrent_processor", **kwargs) + + @handler + async def start(self, emails: list[str], ctx: WorkflowContext[EmailValidationRequest]) -> None: + """Send multiple concurrent requests to the same sub-workflow.""" + # Send all requests concurrently to the same workflow executor + for email in emails: + request = EmailValidationRequest(email=email) + await ctx.send_message(request, target_id="email_workflow") + + @handler + async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None: + """Collect results from concurrent executions.""" + self.results.append(result) + + # Create sub-workflow for email validation + email_validator = EmailValidator() + email_request_info = RequestInfoExecutor(id="email_request_info") + + validation_workflow = ( + WorkflowBuilder() + .set_start_executor(email_validator) + .add_edge(email_validator, email_request_info) + .add_edge(email_request_info, email_validator) + .build() + ) + + # Create parent workflow + processor = ConcurrentProcessor() + workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow") + parent_request_info = RequestInfoExecutor(id="request_info") + + main_workflow = ( + WorkflowBuilder() + .set_start_executor(processor) + .add_edge(processor, workflow_executor) + .add_edge(workflow_executor, processor) + .add_edge(workflow_executor, parent_request_info) # For external requests + .add_edge(parent_request_info, workflow_executor) # For SubWorkflowResponse routing + .build() + ) + + # Test concurrent execution with multiple emails + emails = [ + "user1@domain1.com", + "user2@domain2.com", + "user3@domain3.com", + "user4@domain4.com", + "user5@domain5.com", + ] + + result = await main_workflow.run(emails) + + # Each email should generate one external request + request_events = result.get_request_info_events() + assert len(request_events) == len(emails) + + # Verify each request corresponds to the correct domain + domains_requested = {event.data.domain for event in request_events} # type: ignore[union-attr] + expected_domains = {f"domain{i}.com" for i in range(1, 6)} + assert domains_requested == expected_domains + + # Send responses for all requests (approve all domains) + responses = {event.request_id: True for event in request_events} + await main_workflow.send_responses(responses) + + # All results should be collected + assert len(processor.results) == len(emails) + + # Verify each email was processed correctly + result_emails = {result.email for result in processor.results} + expected_emails = set(emails) + assert result_emails == expected_emails + + # All should be valid since we approved all domains + for result_obj in processor.results: + assert result_obj.is_valid is True + assert result_obj.reason == "Domain approved" + + # Verify that concurrent executions were properly isolated + # (This is implicitly tested by the fact that we got correct results for all emails) + + if __name__ == "__main__": # Run tests asyncio.run(test_basic_sub_workflow()) asyncio.run(test_sub_workflow_with_interception()) asyncio.run(test_conditional_forwarding()) asyncio.run(test_workflow_scoped_interception()) + asyncio.run(test_concurrent_sub_workflow_execution()) diff --git a/python/packages/main/tests/workflow/test_validation.py b/python/packages/main/tests/workflow/test_validation.py index da41ff4858..c9b47de1b9 100644 --- a/python/packages/main/tests/workflow/test_validation.py +++ b/python/packages/main/tests/workflow/test_validation.py @@ -19,7 +19,6 @@ from agent_framework import ( validate_workflow_graph, ) from agent_framework._workflow._edge import SingleEdgeGroup -from agent_framework._workflow._validation import HandlerOutputAnnotationError class StringExecutor(Executor): @@ -51,8 +50,8 @@ class AnyExecutor(Executor): class NoOutputTypesExecutor(Executor): @handler - async def handle_message(self, message: str, ctx: WorkflowContext[Any]) -> None: - await ctx.send_message("processed") + async def handle_message(self, message: str, ctx: WorkflowContext) -> None: + await ctx.send_message("processed") # type: ignore[arg-type] class MultiTypeExecutor(Executor): @@ -575,58 +574,33 @@ def test_validation_enum_usage() -> None: def test_handler_ctx_missing_annotation_raises() -> None: - class BadExecutor(Executor): - @handler - async def handle(self, message: str, ctx) -> None: # type: ignore[no-untyped-def] - pass + # Validation now happens at handler registration time, not workflow build time + with pytest.raises(ValueError) as exc: - start = StringExecutor(id="s") - bad = BadExecutor(id="b") + class BadExecutor(Executor): + @handler + async def handle(self, message: str, ctx) -> None: # type: ignore[no-untyped-def] + pass - with pytest.raises(HandlerOutputAnnotationError) as exc: - WorkflowBuilder().add_edge(start, bad).set_start_executor(start).build() - - assert exc.value.validation_type == ValidationTypeEnum.HANDLER_OUTPUT_ANNOTATION - assert "missing type annotation" in str(exc.value) - - -def test_handler_ctx_unsubscripted_workflow_context_raises() -> None: - class BadExecutor(Executor): - @handler - async def handle(self, message: str, ctx: WorkflowContext) -> None: # type: ignore # missing T - pass - - start = StringExecutor(id="s") - bad = BadExecutor(id="b") - - with pytest.raises(HandlerOutputAnnotationError) as exc: - WorkflowBuilder().add_edge(start, bad).set_start_executor(start).build() - - assert exc.value.validation_type == ValidationTypeEnum.HANDLER_OUTPUT_ANNOTATION - # Message should mention missing T or WorkflowContext[None] - assert "WorkflowContext[None]" in str(exc.value) or "missing" in str(exc.value).lower() + assert "must have a WorkflowContext" in str(exc.value) def test_handler_ctx_invalid_t_out_entries_raises() -> None: - class BadExecutor(Executor): - @handler - async def handle(self, message: str, ctx: WorkflowContext[123]) -> None: # type: ignore[valid-type] - pass + # Validation now happens at handler registration time, not workflow build time + with pytest.raises(ValueError) as exc: - start = StringExecutor(id="s") - bad = BadExecutor(id="b") + class BadExecutor(Executor): + @handler + async def handle(self, message: str, ctx: WorkflowContext[123]) -> None: # type: ignore[valid-type] + pass - with pytest.raises(HandlerOutputAnnotationError) as exc: - WorkflowBuilder().add_edge(start, bad).set_start_executor(start).build() - - assert exc.value.validation_type == ValidationTypeEnum.HANDLER_OUTPUT_ANNOTATION - assert "invalid entries" in str(exc.value) + assert "invalid type entry" in str(exc.value) def test_handler_ctx_none_is_allowed() -> None: class NoneExecutor(Executor): @handler - async def handle(self, message: str, ctx: WorkflowContext[None]) -> None: + async def handle(self, message: str, ctx: WorkflowContext) -> None: # does not emit return None diff --git a/python/packages/main/tests/workflow/test_viz.py b/python/packages/main/tests/workflow/test_viz.py index 7c76eeeb03..3856a3c5de 100644 --- a/python/packages/main/tests/workflow/test_viz.py +++ b/python/packages/main/tests/workflow/test_viz.py @@ -11,7 +11,7 @@ class MockExecutor(Executor): """A mock executor for testing purposes.""" @handler - async def mock_handler(self, message: str, ctx: WorkflowContext[None]) -> None: + async def mock_handler(self, message: str, ctx: WorkflowContext) -> None: """A mock handler that does nothing.""" pass @@ -20,7 +20,7 @@ class ListStrTargetExecutor(Executor): """A mock executor that accepts a list of strings (for fan-in targets).""" @handler - async def handle(self, message: list[str], ctx: WorkflowContext[None]) -> None: # type: ignore[type-arg] + async def handle(self, message: list[str], ctx: WorkflowContext) -> None: pass diff --git a/python/packages/main/tests/workflow/test_workflow.py b/python/packages/main/tests/workflow/test_workflow.py index 5baf05ed80..73ef0950f1 100644 --- a/python/packages/main/tests/workflow/test_workflow.py +++ b/python/packages/main/tests/workflow/test_workflow.py @@ -15,9 +15,11 @@ from agent_framework import ( RequestInfoMessage, RequestResponse, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, WorkflowEvent, + WorkflowOutputEvent, + WorkflowRunState, + WorkflowStatusEvent, handler, ) @@ -36,20 +38,20 @@ class IncrementExecutor(Executor): increment: int = 1 @handler - async def mock_handler(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage]) -> None: + async def mock_handler(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage, int]) -> None: if message.data < self.limit: await ctx.send_message(NumberMessage(data=message.data + self.increment)) else: - await ctx.add_event(WorkflowCompletedEvent(data=message.data)) + await ctx.yield_output(message.data) class AggregatorExecutor(Executor): """A mock executor that aggregates results from multiple executors.""" @handler - async def mock_handler(self, messages: list[NumberMessage], ctx: WorkflowContext[Any]) -> None: - # This mock simply returns the data incremented by 1 - await ctx.add_event(WorkflowCompletedEvent(data=sum(msg.data for msg in messages))) + async def mock_handler(self, messages: list[NumberMessage], ctx: WorkflowContext[Any, int]) -> None: + # This mock simply returns the sum of the data + await ctx.yield_output(sum(msg.data for msg in messages)) @dataclass @@ -70,18 +72,21 @@ class MockExecutorRequestApproval(Executor): @handler async def mock_handler_b( - self, message: RequestResponse[RequestInfoMessage, ApprovalMessage], ctx: WorkflowContext[NumberMessage] + self, + message: RequestResponse[RequestInfoMessage, ApprovalMessage], + ctx: WorkflowContext[NumberMessage, int], ) -> None: """A mock handler that processes the approval response.""" data = await ctx.get_shared_state(self.id) + assert isinstance(data, int) assert isinstance(message.data, ApprovalMessage) if message.data.approved: - await ctx.add_event(WorkflowCompletedEvent(data=data)) + await ctx.yield_output(data) else: await ctx.send_message(NumberMessage(data=data)) -async def test_workflow_run_streaming(): +async def test_workflow_run_streaming() -> None: """Test the workflow run stream.""" executor_a = IncrementExecutor(id="executor_a") executor_b = IncrementExecutor(id="executor_b") @@ -97,7 +102,7 @@ async def test_workflow_run_streaming(): result: int | None = None async for event in workflow.run_stream(NumberMessage(data=0)): assert isinstance(event, WorkflowEvent) - if isinstance(event, WorkflowCompletedEvent): + if isinstance(event, WorkflowOutputEvent): result = event.data assert result is not None and result == 10 @@ -136,9 +141,9 @@ async def test_workflow_run(): ) events = await workflow.run(NumberMessage(data=0)) - completed_event = events.get_completed_event() - assert isinstance(completed_event, WorkflowCompletedEvent) - assert completed_event.data == 10 + assert events.get_final_state() == WorkflowRunState.IDLE + outputs = events.get_outputs() + assert outputs[0] == 10 async def test_workflow_run_not_completed(): @@ -182,13 +187,18 @@ async def test_workflow_send_responses_streaming(): assert request_info_event is not None result: int | None = None + completed = False async for event in workflow.send_responses_streaming({ request_info_event.request_id: ApprovalMessage(approved=True) }): - if isinstance(event, WorkflowCompletedEvent): + if isinstance(event, WorkflowOutputEvent): result = event.data + elif isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE: + completed = True - assert result is not None and result == 1 # The data should be incremented by 1 from the initial message + assert ( + completed and result is not None and result == 1 + ) # The data should be incremented by 1 from the initial message async def test_workflow_send_responses(): @@ -214,9 +224,9 @@ async def test_workflow_send_responses(): result = await workflow.send_responses({request_info_events[0].request_id: ApprovalMessage(approved=True)}) - completed_event = result.get_completed_event() - assert isinstance(completed_event, WorkflowCompletedEvent) - assert completed_event.data == 1 # The data should be incremented by 1 from the initial message + assert result.get_final_state() == WorkflowRunState.IDLE + outputs = result.get_outputs() + assert outputs[0] == 1 # The data should be incremented by 1 from the initial message async def test_fan_out(): @@ -232,11 +242,12 @@ async def test_fan_out(): events = await workflow.run(NumberMessage(data=0)) # Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent - # executor_b will also emit a WorkflowCompletedEvent + # executor_b will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore) assert len(events) == 7 - completed_event = events.get_completed_event() - assert completed_event is not None and completed_event.data == 1 + assert events.get_final_state() == WorkflowRunState.IDLE + outputs = events.get_outputs() + assert outputs[0] == 1 async def test_fan_out_multiple_completed_events(): @@ -252,11 +263,12 @@ async def test_fan_out_multiple_completed_events(): events = await workflow.run(NumberMessage(data=0)) # Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent - # executor_a and executor_b will also emit a WorkflowCompletedEvent + # executor_b and executor_c will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore) assert len(events) == 8 - with pytest.raises(ValueError): - events.get_completed_event() + # Multiple outputs are expected from both executors + outputs = events.get_outputs() + assert len(outputs) == 2 async def test_fan_in(): @@ -277,18 +289,19 @@ async def test_fan_in(): events = await workflow.run(NumberMessage(data=0)) # Each executor will emit two events: ExecutorInvokedEvent and ExecutorCompletedEvent - # aggregator will also emit a WorkflowCompletedEvent + # aggregator will also emit a WorkflowOutputEvent (no WorkflowCompletedEvent anymore) assert len(events) == 9 - completed_event = events.get_completed_event() - assert completed_event is not None and completed_event.data == 4 + assert events.get_final_state() == WorkflowRunState.IDLE + outputs = events.get_outputs() + assert outputs[0] == 4 # executor_a(0->1), both executor_b and executor_c(1->2), aggregator(2+2=4) @pytest.fixture def simple_executor() -> Executor: class SimpleExecutor(Executor): @handler - async def handle_message(self, message: str, context: WorkflowContext[None]) -> None: + async def handle_message(self, message: str, context: WorkflowContext) -> None: pass return SimpleExecutor(id="test_executor") @@ -442,7 +455,7 @@ async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Execu # Test non-streaming run_from_checkpoint method result = await workflow.run_from_checkpoint(checkpoint_id) assert isinstance(result, list) # Should return WorkflowRunResult which extends list - assert hasattr(result, "get_completed_event") # Should have WorkflowRunResult methods + assert hasattr(result, "get_outputs") # Should have WorkflowRunResult methods async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executor: Executor): @@ -498,7 +511,7 @@ class StateTrackingExecutor(Executor): """An executor that tracks state in shared state to test context reset behavior.""" @handler - async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext[Any]) -> None: + async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext[Any, list]) -> None: """Handle the message and track it in shared state.""" # Get existing messages from shared state try: @@ -513,8 +526,8 @@ class StateTrackingExecutor(Executor): # Update shared state await ctx.set_shared_state("processed_messages", existing_messages) - # Complete workflow with current shared state - await ctx.add_event(WorkflowCompletedEvent(data=existing_messages.copy())) # type: ignore + # Yield output + await ctx.yield_output(existing_messages.copy()) # type: ignore async def test_workflow_multiple_runs_no_state_collision(): @@ -536,27 +549,27 @@ async def test_workflow_multiple_runs_no_state_collision(): # Run 1: Should only see messages from run 1 result1 = await workflow.run(StateTrackingMessage(data="message1", run_id="run1")) - completed1 = result1.get_completed_event() - assert completed1 is not None - assert completed1.data == ["run1:message1"] + assert result1.get_final_state() == WorkflowRunState.IDLE + outputs1 = result1.get_outputs() + assert outputs1[0] == ["run1:message1"] # Run 2: Should only see messages from run 2, not run 1 result2 = await workflow.run(StateTrackingMessage(data="message2", run_id="run2")) - completed2 = result2.get_completed_event() - assert completed2 is not None - assert completed2.data == ["run2:message2"] # Should NOT contain run1 data + assert result2.get_final_state() == WorkflowRunState.IDLE + outputs2 = result2.get_outputs() + assert outputs2[0] == ["run2:message2"] # Should NOT contain run1 data # Run 3: Should only see messages from run 3 result3 = await workflow.run(StateTrackingMessage(data="message3", run_id="run3")) - completed3 = result3.get_completed_event() - assert completed3 is not None - assert completed3.data == ["run3:message3"] # Should NOT contain run1 or run2 data + assert result3.get_final_state() == WorkflowRunState.IDLE + outputs3 = result3.get_outputs() + assert outputs3[0] == ["run3:message3"] # Should NOT contain run1 or run2 data # Verify that each run only processed its own message # This confirms that the checkpointable context properly resets between runs - assert completed1.data != completed2.data - assert completed2.data != completed3.data - assert completed1.data != completed3.data + assert outputs1[0] != outputs2[0] + assert outputs2[0] != outputs3[0] + assert outputs1[0] != outputs3[0] async def test_comprehensive_edge_groups_workflow(): @@ -604,17 +617,17 @@ async def test_comprehensive_edge_groups_workflow(): # router(2->3) -> switch routes to proc_a -> proc_a(3->4) -> fanout_hub(4->5) # -> [parallel_1(5->8), parallel_2(5->10)] -> aggregator(8+10=18) events_small = await workflow.run(NumberMessage(data=2)) - completed_small = events_small.get_completed_event() - assert completed_small is not None - assert completed_small.data == 18 # Exact expected result: 8+10 from parallel processors + assert events_small.get_final_state() == WorkflowRunState.IDLE + outputs_small = events_small.get_outputs() + assert outputs_small[0] == 18 # Exact expected result: 8+10 from parallel processors # Test with large number (should go through processor_b) # router(8->9) -> switch routes to proc_b -> proc_b(9->11) -> fanout_hub(11->12) # -> [parallel_1(12->15), parallel_2(12->17)] -> aggregator(15+17=32) events_large = await workflow.run(NumberMessage(data=8)) - completed_large = events_large.get_completed_event() - assert completed_large is not None - assert completed_large.data == 32 # Exact expected result: 15+17 from parallel processors + assert events_large.get_final_state() == WorkflowRunState.IDLE + outputs_large = events_large.get_outputs() + assert outputs_large[0] == 32 # Exact expected result: 15+17 from parallel processors # The key verification is that we successfully executed a workflow using all three edge group types # and that both switch-case paths work (small vs large numbers) @@ -624,9 +637,9 @@ async def test_comprehensive_edge_groups_workflow(): assert len(events_large) >= 6 # Verify different paths were taken by checking exact results - assert completed_small.data == 18, f"Small number path should result in 18, got {completed_small.data}" - assert completed_large.data == 32, f"Large number path should result in 32, got {completed_large.data}" - assert completed_small.data != completed_large.data, "Different paths should produce different results" + assert outputs_small[0] == 18, f"Small number path should result in 18, got {outputs_small[0]}" + assert outputs_large[0] == 32, f"Large number path should result in 32, got {outputs_large[0]}" + assert outputs_small[0] != outputs_large[0], "Different paths should produce different results" # Both tests should complete successfully, proving all edge group types work @@ -660,11 +673,9 @@ async def test_workflow_with_simple_cycle_and_exit_condition(): # Test the cycle # Expected: exec_a(2->4) -> exec_b(4->5) -> exec_a(5->7, completes because 7 >= 6) events = await workflow.run(NumberMessage(data=2)) - completed_event = events.get_completed_event() - assert completed_event is not None - assert ( - completed_event.data is not None and completed_event.data >= 6 - ) # Should complete when executor_a reaches its limit + assert events.get_final_state() == WorkflowRunState.IDLE + outputs = events.get_outputs() + assert outputs[0] is not None and outputs[0] >= 6 # Should complete when executor_a reaches its limit # Verify cycling occurred (should have events from both executors) # Check for ExecutorInvokedEvent and ExecutorCompletedEvent types that have executor_id diff --git a/python/packages/main/tests/workflow/test_workflow_context.py b/python/packages/main/tests/workflow/test_workflow_context.py index b78de9af67..504230e293 100644 --- a/python/packages/main/tests/workflow/test_workflow_context.py +++ b/python/packages/main/tests/workflow/test_workflow_context.py @@ -3,14 +3,19 @@ import asyncio from collections.abc import AsyncIterator from contextlib import asynccontextmanager -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any + +from typing_extensions import Never from agent_framework import ( - WorkflowCompletedEvent, + Executor, + WorkflowBuilder, WorkflowContext, WorkflowEvent, WorkflowRunState, WorkflowStatusEvent, + executor, + handler, ) if TYPE_CHECKING: @@ -57,8 +62,203 @@ async def test_executor_cannot_emit_framework_lifecycle_event(caplog: "LogCaptur async def test_executor_emits_normal_event() -> None: async with make_context() as (ctx, runner_ctx): - await ctx.add_event(WorkflowCompletedEvent("done")) + # Create a normal event to test event emission + await ctx.add_event(_TestEvent()) events: list[WorkflowEvent] = await runner_ctx.drain_events() assert len(events) == 1 - assert isinstance(events[0], WorkflowCompletedEvent) + assert isinstance(events[0], _TestEvent) + + +class _TestEvent(WorkflowEvent): + pass + + +async def test_workflow_context_type_annotations_no_parameter() -> None: + # Test function-based executor + @executor(id="func1") + async def func1(text: str, ctx: WorkflowContext) -> None: + await ctx.add_event(_TestEvent()) + + wf = WorkflowBuilder().set_start_executor(func1).build() + events = await wf.run("hello") + test_events = [e for e in events if isinstance(e, _TestEvent)] + assert len(test_events) == 1 + + # Test class-based executor + class _exec1(Executor): + @handler + async def func1(self, text: str, ctx: WorkflowContext) -> None: + await ctx.add_event(_TestEvent()) + + executor1 = _exec1(id="exec1") + + assert executor1.input_types == [str] + assert executor1.output_types == [] + assert executor1.workflow_output_types == [] + + wf2 = WorkflowBuilder().set_start_executor(executor1).build() + events2 = await wf2.run("hello") + test_events2 = [e for e in events2 if isinstance(e, _TestEvent)] + assert len(test_events2) == 1 + + +async def test_workflow_context_type_annotations_message_type_parameter() -> None: + # Test function-based executor + @executor(id="func1") + async def func1(text: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message("world") + + @executor(id="func2") + async def func2(text: str, ctx: WorkflowContext) -> None: + await ctx.add_event(_TestEvent(data=text)) + + wf = WorkflowBuilder().add_edge(func1, func2).set_start_executor(func1).build() + events = await wf.run("hello") + test_events = [e for e in events if isinstance(e, _TestEvent)] + assert len(test_events) == 1 + assert test_events[0].data == "world" + + # Test class-based executor + class _exec1(Executor): + @handler + async def func1(self, text: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message("world") + + class _exec2(Executor): + @handler + async def func2(self, text: str, ctx: WorkflowContext) -> None: + await ctx.add_event(_TestEvent(data=text)) + + executor1 = _exec1(id="exec1") + executor2 = _exec2(id="exec2") + + assert executor1.input_types == [str] + assert executor1.output_types == [str] + assert executor1.workflow_output_types == [] + assert executor2.input_types == [str] + assert executor2.output_types == [] + assert executor2.workflow_output_types == [] + + wf2 = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + events2 = await wf2.run("hello") + test_events2 = [e for e in events2 if isinstance(e, _TestEvent)] + assert len(test_events2) == 1 + assert test_events2[0].data == "world" + + +async def test_workflow_context_type_annotations_message_and_output_type_parameters() -> None: + # Test function-based executor + @executor(id="func1") + async def func1(text: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message("world") + + @executor(id="func2") + async def func2(text: str, ctx: WorkflowContext[Never, str]) -> None: + await ctx.add_event(_TestEvent(data=text)) + await ctx.yield_output(text) + + wf = WorkflowBuilder().add_edge(func1, func2).set_start_executor(func1).build() + events = await wf.run("hello") + outputs = events.get_outputs() + assert len(outputs) == 1 + assert outputs[0] == "world" + + # Test class-based executor + class _exec1(Executor): + @handler + async def func1(self, text: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message("world") + + class _exec2(Executor): + @handler + async def func2(self, text: str, ctx: WorkflowContext[Never, str]) -> None: + await ctx.add_event(_TestEvent(data=text)) + await ctx.yield_output(text) + + executor1 = _exec1(id="exec1") + executor2 = _exec2(id="exec2") + + assert executor1.input_types == [str] + assert executor1.output_types == [str] + assert executor1.workflow_output_types == [] + assert executor2.input_types == [str] + assert executor2.output_types == [] + assert executor2.workflow_output_types == [str] + + wf2 = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + events2 = await wf2.run("hello") + outputs2 = events2.get_outputs() + assert len(outputs2) == 1 + assert outputs2[0] == "world" + + +async def test_workflow_context_type_annotations_any() -> None: + class _exec1(Executor): + @handler + async def func1(self, text: str, ctx: WorkflowContext[Any]) -> None: + await ctx.add_event(_TestEvent()) + await ctx.send_message(123) + + executor1 = _exec1(id="exec1") + assert executor1.input_types == [str] + assert executor1.output_types == [Any] + + class _exec2(Executor): + @handler + async def func2(self, number: int, ctx: WorkflowContext[Any, Any]) -> None: + await ctx.add_event(_TestEvent()) + await ctx.send_message(456) + await ctx.yield_output(3.14) + + executor2 = _exec2(id="exec2") + assert executor2.input_types == [int] + assert executor2.output_types == [Any] + assert executor2.workflow_output_types == [Any] + + +async def test_workflow_context_missing_annotation_error() -> None: + """Test that missing WorkflowContext annotation raises appropriate error.""" + import pytest + + # Test function-based executor with missing ctx annotation + with pytest.raises(ValueError, match="must have a WorkflowContext"): + + @executor(id="bad_func") + async def bad_func(text: str, ctx) -> None: # type: ignore[no-untyped-def] + pass + + # Test class-based executor with missing ctx annotation + with pytest.raises(ValueError, match="must have a WorkflowContext"): + + class _BadExecutor(Executor): + @handler + async def bad_handler(self, text: str, ctx) -> None: # type: ignore[no-untyped-def] + pass + + +async def test_workflow_context_invalid_type_parameter_error() -> None: + """Test that invalid type parameters like int values raise appropriate errors.""" + import pytest + + # Test function-based executor with invalid type parameter (int value instead of type) + with pytest.raises(ValueError, match="invalid type entry"): + + @executor(id="bad_func") + async def bad_func(text: str, ctx: WorkflowContext[123]) -> None: # type: ignore[valid-type] + pass + + # Test class-based executor with invalid type parameter + with pytest.raises(ValueError, match="invalid type entry"): + + class _BadExecutor(Executor): + @handler + async def bad_handler(self, text: str, ctx: WorkflowContext[456]) -> None: # type: ignore[valid-type] + pass + + # Test two-parameter WorkflowContext with invalid workflow output type + with pytest.raises(ValueError, match="invalid type entry"): + + @executor(id="bad_func2") + async def bad_func2(text: str, ctx: WorkflowContext[str, 789]) -> None: # type: ignore[valid-type] + pass diff --git a/python/packages/main/tests/workflow/test_workflow_observability.py b/python/packages/main/tests/workflow/test_workflow_observability.py index 1f2c890a67..2a218eb0b1 100644 --- a/python/packages/main/tests/workflow/test_workflow_observability.py +++ b/python/packages/main/tests/workflow/test_workflow_observability.py @@ -48,7 +48,7 @@ class SecondExecutor(Executor): self._processed_messages: list[str] = [] @handler - async def handle_message(self, message: str, ctx: WorkflowContext[None]) -> None: + async def handle_message(self, message: str, ctx: WorkflowContext) -> None: """Handle string messages.""" self._processed_messages.append(message) @@ -87,7 +87,7 @@ class FanInAggregator(Executor): self._processed_messages: list[Any] = [] @handler - async def handle_aggregated_data(self, messages: list[str], ctx: WorkflowContext[None]) -> None: + async def handle_aggregated_data(self, messages: list[str], ctx: WorkflowContext) -> None: # Process aggregated messages from fan-in aggregated = f"aggregated: {', '.join(messages)}" self._processed_messages.append(aggregated) @@ -196,7 +196,14 @@ async def test_trace_context_handling(span_exporter: InMemorySpanExporter) -> No assert message.source_span_id is not None # Test executor trace context handling - await executor.execute("test message", workflow_ctx) + await executor.execute( + "test message", + ["source"], # source_executor_ids + shared_state, # shared_state + ctx, # runner_context + trace_contexts=[{"traceparent": "00-12345678901234567890123456789012-1234567890123456-01"}], + source_span_ids=["1234567890123456"], + ) # Check that spans were created with proper attributes spans = span_exporter.get_finished_spans() @@ -372,7 +379,7 @@ async def test_workflow_error_handling_in_tracing(span_exporter: InMemorySpanExp super().__init__(id="failing_executor") @handler - async def handle_message(self, message: str, ctx: WorkflowContext[None]) -> None: + async def handle_message(self, message: str, ctx: WorkflowContext) -> None: raise ValueError("Test error") failing_executor = FailingExecutor() diff --git a/python/packages/main/tests/workflow/test_workflow_states.py b/python/packages/main/tests/workflow/test_workflow_states.py index bef26a4c98..1f6adba134 100644 --- a/python/packages/main/tests/workflow/test_workflow_states.py +++ b/python/packages/main/tests/workflow/test_workflow_states.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Any import pytest +from typing_extensions import Never from agent_framework import ( Executor, @@ -15,7 +16,6 @@ from agent_framework import ( SharedState, Workflow, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, WorkflowEventSource, WorkflowFailedEvent, @@ -32,7 +32,7 @@ class FailingExecutor(Executor): """Executor that raises at runtime to test failure signaling.""" @handler - async def fail(self, msg: int, ctx: WorkflowContext[None]) -> None: # pragma: no cover - invoked via workflow + async def fail(self, msg: int, ctx: WorkflowContext) -> None: # pragma: no cover - invoked via workflow raise RuntimeError("boom") @@ -57,14 +57,14 @@ async def test_executor_failed_and_workflow_failed_events_streaming(): async def test_executor_failed_event_emitted_on_direct_execute(): failing = FailingExecutor(id="f") ctx = InProcRunnerContext() - wf_ctx: WFContext[None] = WFContext( - executor_id=failing.id, - source_executor_ids=["START"], - shared_state=SharedState(), - runner_context=ctx, - ) + shared_state = SharedState() with pytest.raises(RuntimeError, match="boom"): - await failing.execute(0, wf_ctx) + await failing.execute( + 0, + ["START"], + shared_state, + ctx, + ) drained = await ctx.drain_events() failed = [e for e in drained if isinstance(e, ExecutorFailedEvent)] assert failed @@ -98,17 +98,17 @@ class Completer(Executor): """Executor that completes immediately with provided data for testing.""" @handler - async def run(self, msg: str, ctx: WorkflowContext[str]) -> None: # pragma: no cover - await ctx.add_event(WorkflowCompletedEvent(msg)) + async def run(self, msg: str, ctx: WorkflowContext[Never, str]) -> None: # pragma: no cover + await ctx.yield_output(msg) async def test_completed_status_streaming(): c = Completer(id="c") wf = WorkflowBuilder().set_start_executor(c).build() events = [ev async for ev in wf.run_stream("ok")] # no raise - # Last status should be COMPLETED + # Last status should be IDLE status = [e for e in events if isinstance(e, WorkflowStatusEvent)] - assert status and status[-1].state == WorkflowRunState.COMPLETED + assert status and status[-1].state == WorkflowRunState.IDLE assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in status) @@ -120,8 +120,12 @@ async def test_started_and_completed_event_origins(): started = next(e for e in events if isinstance(e, WorkflowStartedEvent)) assert started.origin is WorkflowEventSource.FRAMEWORK - completed = next(e for e in events if isinstance(e, WorkflowCompletedEvent)) - assert completed.origin is WorkflowEventSource.EXECUTOR + # Check for IDLE status indicating completion + idle_status = next( + (e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE), None + ) + assert idle_status is not None + assert idle_status.origin is WorkflowEventSource.FRAMEWORK async def test_non_streaming_final_state_helpers(): @@ -129,7 +133,7 @@ async def test_non_streaming_final_state_helpers(): c = Completer(id="c") wf1 = WorkflowBuilder().set_start_executor(c).build() result1: WorkflowRunResult = await wf1.run("done") - assert result1.get_final_state() == WorkflowRunState.COMPLETED + assert result1.get_final_state() == WorkflowRunState.IDLE # Idle-with-pending-request case req = Requester(id="req") @@ -145,7 +149,7 @@ async def test_run_includes_status_events_completed(): result: WorkflowRunResult = await wf.run("ok") timeline = result.status_timeline() assert timeline, "Expected status timeline in non-streaming run() results" - assert timeline[-1].state == WorkflowRunState.COMPLETED + assert timeline[-1].state == WorkflowRunState.IDLE async def test_run_includes_status_events_idle_with_requests(): diff --git a/python/samples/getting_started/observability/04-workflow.py b/python/samples/getting_started/observability/04-workflow.py index 420ee7551f..ca2f23685e 100644 --- a/python/samples/getting_started/observability/04-workflow.py +++ b/python/samples/getting_started/observability/04-workflow.py @@ -7,6 +7,7 @@ from agent_framework import ( Executor, WorkflowBuilder, WorkflowCompletedEvent, + WorkflowOutputEvent, WorkflowContext, handler, ) @@ -42,14 +43,15 @@ class ReverseTextExecutor(Executor): """An executor that reverses text.""" @handler - async def reverse_text(self, text: str, ctx: WorkflowContext[Any]) -> None: + async def reverse_text(self, text: str, ctx: WorkflowContext[Any, str]) -> None: """Execute the task by reversing the input string.""" print(f"ReverseTextExecutor: Processing '{text}'") result = text[::-1] print(f"ReverseTextExecutor: Result '{result}'") - # Send the result with a workflow completion event. - await ctx.add_event(WorkflowCompletedEvent(result)) + # Yield the output and signal workflow completion. + await ctx.yield_output(result) + await ctx.add_event(WorkflowCompletedEvent()) async def run_sequential_workflow() -> None: @@ -84,17 +86,17 @@ async def run_sequential_workflow() -> None: input_text = "hello world" print(f"Starting workflow with input: '{input_text}'") - completion_event = None + output_event = None async for event in workflow.run_stream(input_text): print(f"Event: {event}") - if isinstance(event, WorkflowCompletedEvent): - # The WorkflowCompletedEvent contains the final result. - completion_event = event + if isinstance(event, WorkflowOutputEvent): + # The WorkflowOutputEvent contains the final result. + output_event = event - if completion_event: - print(f"Workflow completed with result: '{completion_event.data}'") + if output_event: + print(f"Workflow completed with result: '{output_event.data}'") else: - print("Workflow completed without a completion event") + print("Workflow completed without an output event") except Exception as e: current_span.record_exception(e) diff --git a/python/samples/getting_started/workflow/_start-here/step1_executors_and_edges.py b/python/samples/getting_started/workflow/_start-here/step1_executors_and_edges.py index 7744018c75..aa1445b5bb 100644 --- a/python/samples/getting_started/workflow/_start-here/step1_executors_and_edges.py +++ b/python/samples/getting_started/workflow/_start-here/step1_executors_and_edges.py @@ -2,10 +2,11 @@ import asyncio +from typing_extensions import Never + from agent_framework import ( Executor, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, executor, handler, @@ -17,18 +18,28 @@ Step 1: Foundational patterns: Executors and edges What this example shows - Two ways to define a unit of work (an Executor node): 1) Custom class that subclasses Executor with an async method marked by @handler. - Signature: (text: str, ctx: WorkflowContext[str]) -> None. The typed ctx - advertises the type this node emits via ctx.send_message(...). + Possible handler signatures: + - (text: str, ctx: WorkflowContext) -> None, + - (text: str, ctx: WorkflowContext[str]) -> None, or + - (text: str, ctx: WorkflowContext[Never, str]) -> None. + The first parameter is the typed input to this node, the input type is str here. + The second parameter is a WorkflowContext[T_Out, T_W_Out]. + WorkflowContext[T_Out] is used for nodes that send messages to downstream nodes with ctx.send_message(T_Out). + WorkflowContext[T_Out, T_W_Out] is used for nodes that also yield workflow + output with ctx.yield_output(T_W_Out). + WorkflowContext without type parameters is equivalent to WorkflowContext[Never, Never], meaning this node + neither sends messages to downstream nodes nor yields workflow output. + 2) Standalone async function decorated with @executor using the same signature. - Simple steps can use this form; a terminal step can emit a - WorkflowCompletedEvent to end the workflow. + Simple steps can use this form; a terminal step can yield output + using ctx.yield_output() to provide workflow results. - Fluent WorkflowBuilder API: add_edge(A, B) to connect nodes, set_start_executor(A), then build() -> Workflow. - Running and results: - workflow.run(initial_input) executes the graph. The last node emits a - WorkflowCompletedEvent that carries the final result. + workflow.run(initial_input) executes the graph. Terminal nodes yield + outputs using ctx.yield_output(). The workflow runs until idle. Prerequisites - No external services required. @@ -43,8 +54,8 @@ Prerequisites # # Handler signature contract: # - First parameter is the typed input to this node (here: text: str) -# - Second parameter is a WorkflowContext[T], where T is the type of data this -# node will emit via ctx.send_message (here: T is str) +# - Second parameter is a WorkflowContext[T_Out], where T_Out is the type of data this +# node will emit via ctx.send_message (here: T_Out is str) # # Within a handler you typically: # - Compute a result @@ -70,22 +81,25 @@ class UpperCase(Executor): # ----------------------------------------------- # # For simple steps you can skip subclassing and define an async function with the -# same signature pattern (typed input + WorkflowContext[T]) and decorate it with +# same signature pattern (typed input + WorkflowContext[T_Out, T_W_Out]) and decorate it with # @executor. This creates a fully functional node that can be wired into a flow. @executor(id="reverse_text_executor") -async def reverse_text(text: str, ctx: WorkflowContext[str]) -> None: - """Reverse the input string and signal workflow completion. +async def reverse_text(text: str, ctx: WorkflowContext[Never, str]) -> None: + """Reverse the input string and yield the workflow output. - This node emits a terminal event using ctx.add_event(WorkflowCompletedEvent). - The data carried by the WorkflowCompletedEvent becomes the final result of - the workflow (returned by workflow.run(...)). + This node yields the final output using ctx.yield_output(result). + The workflow will complete when it becomes idle (no more work to do). + + The WorkflowContext is parameterized with two types: + - T_Out = Never: this node does not send messages to downstream nodes. + - T_W_Out = str: this node yields workflow output of type str. """ result = text[::-1] - # Send the result with a workflow completion event. - await ctx.add_event(WorkflowCompletedEvent(result)) + # Yield the output - the workflow will complete when idle + await ctx.yield_output(result) async def main(): @@ -100,17 +114,17 @@ async def main(): workflow = WorkflowBuilder().add_edge(upper_case, reverse_text).set_start_executor(upper_case).build() # Run the workflow by sending the initial message to the start node. - # The run(...) call returns an event collection; its get_completed_event() - # provides the WorkflowCompletedEvent emitted by the terminal node. + # The run(...) call returns an event collection; its get_outputs() method + # retrieves the outputs yielded by any terminal nodes. events = await workflow.run("hello world") - print(events.get_completed_event()) + print(events.get_outputs()) # Summarize the final run state (e.g., COMPLETED) print("Final state:", events.get_final_state()) """ Sample Output: - WorkflowCompletedEvent(data=DLROW OLLEH) + ['DLROW OLLEH'] Final state: WorkflowRunState.COMPLETED """ diff --git a/python/samples/getting_started/workflow/_start-here/step2_agents_in_a_workflow.py b/python/samples/getting_started/workflow/_start-here/step2_agents_in_a_workflow.py index 1c0a113965..303e85aed4 100644 --- a/python/samples/getting_started/workflow/_start-here/step2_agents_in_a_workflow.py +++ b/python/samples/getting_started/workflow/_start-here/step2_agents_in_a_workflow.py @@ -13,9 +13,9 @@ This sample uses two custom executors. A Writer agent creates or edits content, then hands the conversation to a Reviewer agent which evaluates and finalizes the result. Purpose: -Show how to wrap chat agents created by AzureChatClient inside workflow executors. Demonstrate the @handler pattern -with typed inputs and typed WorkflowContext[T] outputs, connect executors with the fluent WorkflowBuilder, and finish -by emitting a WorkflowCompletedEvent from the terminal node. +Show how to wrap chat agents created by AzureChatClient inside workflow executors. Demonstrate how agents +automatically yield outputs when they complete, removing the need for explicit completion events. +The workflow completes when it becomes idle. Prerequisites: - Azure OpenAI configured for AzureChatClient with required environment variables. @@ -51,14 +51,12 @@ async def main(): # Run the workflow with the user's initial message. # For foundational clarity, use run (non streaming) and print the terminal event. events = await workflow.run("Create a slogan for a new electric SUV that is affordable and fun to drive.") - # The terminal node emits a WorkflowCompletedEvent; print its contents. - - # Print interim-agent run events + # Print agent run events and final outputs for event in events: if isinstance(event, AgentRunEvent): print(f"{event.executor_id}: {event.data}") - print(f"{'=' * 60}\n{events.get_completed_event()}") + print(f"{'=' * 60}\nWorkflow Outputs: {events.get_outputs()}") # Summarize the final run state (e.g., COMPLETED) print("Final state:", events.get_final_state()) @@ -74,14 +72,13 @@ async def main(): - Consider specifying "SUV" for clarity in some uses. - Strong, upbeat tone suitable for marketing. ============================================================ - Workflow Completed Event: - WorkflowCompletedEvent(data=Slogan: "Plug Into Fun—Affordable Adventure, Electrified." + Workflow Outputs: ['Slogan: "Plug Into Fun—Affordable Adventure, Electrified." - **Feedback:**s + **Feedback:** - Clear focus on affordability and enjoyment. - "Plug into fun" connects emotionally and highlights electric nature. - Consider specifying "SUV" for clarity in some uses. - - Strong, upbeat tone suitable for marketing.) + - Strong, upbeat tone suitable for marketing.'] """ diff --git a/python/samples/getting_started/workflow/_start-here/step3_streaming.py b/python/samples/getting_started/workflow/_start-here/step3_streaming.py index f5d29e178a..a09c3bd1c8 100644 --- a/python/samples/getting_started/workflow/_start-here/step3_streaming.py +++ b/python/samples/getting_started/workflow/_start-here/step3_streaming.py @@ -2,19 +2,21 @@ import asyncio +from typing_extensions import Never + from agent_framework import ( ChatAgent, ChatMessage, Executor, ExecutorFailedEvent, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, WorkflowFailedEvent, WorkflowRunState, WorkflowStatusEvent, handler, ) +from agent_framework._workflow._events import WorkflowOutputEvent from agent_framework.azure import AzureChatClient from azure.identity import AzureCliCredential @@ -28,9 +30,9 @@ The workflow is invoked with run_stream so you can observe events as they occur. Purpose: Show how to wrap chat agents created by AzureChatClient inside workflow executors, wire them with WorkflowBuilder, and consume streaming events from the workflow. Demonstrate the @handler pattern with typed inputs and typed -WorkflowContext[T] outputs, and finish by emitting a WorkflowCompletedEvent from the terminal node while printing -intermediate events for observability. The streaming loop also surfaces WorkflowEvent.origin so you can -distinguish runner-generated lifecycle events from executor-generated data-plane events. +WorkflowContext[T_Out, T_W_Out] outputs. Agents automatically yield outputs when they complete. +The streaming loop also surfaces WorkflowEvent.origin so you can distinguish runner-generated lifecycle events +from executor-generated data-plane events. Prerequisites: - Azure OpenAI configured for AzureChatClient with required environment variables. @@ -96,14 +98,14 @@ class Reviewer(Executor): super().__init__(agent=agent, id=id) @handler - async def handle(self, messages: list[ChatMessage], ctx: WorkflowContext[str]) -> None: - """Review the full conversation transcript and complete with a final string. + async def handle(self, messages: list[ChatMessage], ctx: WorkflowContext[Never, str]) -> None: + """Review the full conversation transcript and yield the final output. This node consumes all messages so far. It uses its agent to produce the final text, - then signals completion by adding a WorkflowCompletedEvent to the event stream. + then yields the output. The workflow completes when it becomes idle. """ response = await self.agent.run(messages) - await ctx.add_event(WorkflowCompletedEvent(response.text)) + await ctx.yield_output(response.text) async def main(): @@ -119,7 +121,7 @@ async def main(): workflow = WorkflowBuilder().set_start_executor(writer).add_edge(writer, reviewer).build() # Run the workflow with the user's initial message and stream events as they occur. - # In addition to executor events and WorkflowCompletedEvent, this also surfaces run-state and errors. + # This surfaces executor events, workflow outputs, run-state changes, and errors. async for event in workflow.run_stream( ChatMessage(role="user", text="Create a slogan for a new electric SUV that is affordable and fun to drive.") ): @@ -127,8 +129,6 @@ async def main(): prefix = f"State ({event.origin.value}): " if event.state == WorkflowRunState.IN_PROGRESS: print(prefix + "IN_PROGRESS") - elif event.state == WorkflowRunState.COMPLETED: - print(prefix + "COMPLETED") elif event.state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS: print(prefix + "IN_PROGRESS_PENDING_REQUESTS (requests in flight)") elif event.state == WorkflowRunState.IDLE: @@ -137,8 +137,8 @@ async def main(): print(prefix + "IDLE_WITH_PENDING_REQUESTS (prompt user or UI now)") else: print(prefix + str(event.state)) - elif isinstance(event, WorkflowCompletedEvent): - print(f"Workflow completed ({event.origin.value}): {event.data}") + elif isinstance(event, WorkflowOutputEvent): + print(f"Workflow output ({event.origin.value}): {event.data}") elif isinstance(event, ExecutorFailedEvent): print( f"Executor failed ({event.origin.value}): " @@ -157,9 +157,9 @@ async def main(): ExecutorInvokeEvent (RUNNER): ExecutorInvokeEvent(executor_id=writer) ExecutorCompletedEvent (RUNNER): ExecutorCompletedEvent(executor_id=writer) ExecutorInvokeEvent (RUNNER): ExecutorInvokeEvent(executor_id=reviewer) - Workflow completed (EXECUTOR): Drive the Future. Affordable Adventure, Electrified. + Workflow output (EXECUTOR): Drive the Future. Affordable Adventure, Electrified. ExecutorCompletedEvent (RUNNER): ExecutorCompletedEvent(executor_id=reviewer) - State (RUNNER): COMPLETED + State (RUNNER): IDLE """ diff --git a/python/samples/getting_started/workflow/agents/azure_chat_agents_streaming.py b/python/samples/getting_started/workflow/agents/azure_chat_agents_streaming.py index 175ee30214..ba8410bb53 100644 --- a/python/samples/getting_started/workflow/agents/azure_chat_agents_streaming.py +++ b/python/samples/getting_started/workflow/agents/azure_chat_agents_streaming.py @@ -2,7 +2,7 @@ import asyncio -from agent_framework import AgentRunUpdateEvent, WorkflowBuilder, WorkflowCompletedEvent +from agent_framework import AgentRunUpdateEvent, WorkflowBuilder, WorkflowOutputEvent from agent_framework.azure import AzureChatClient from azure.identity import AzureCliCredential @@ -18,7 +18,7 @@ Show how to wire chat agents directly into a WorkflowBuilder pipeline where agen Demonstrate: - Automatic streaming of agent deltas via AgentRunUpdateEvent. - A simple console aggregator that groups updates by executor id and prints them as they arrive. -- A final WorkflowCompletedEvent that contains the reviewer outcome after both agents finish. +- The workflow completes when idle and outputs are available in events.get_outputs(). Prerequisites: - Azure OpenAI configured for AzureChatClient with required environment variables. @@ -54,12 +54,10 @@ async def main(): workflow = WorkflowBuilder().set_start_executor(writer_agent).add_edge(writer_agent, reviewer_agent).build() # Stream events from the workflow. We aggregate partial token updates per executor for readable output. - completed_event: WorkflowCompletedEvent | None = None last_executor_id = None - async for event in workflow.run_stream( - "Create a slogan for a new electric SUV that is affordable and fun to drive." - ): + events = workflow.run_stream("Create a slogan for a new electric SUV that is affordable and fun to drive.") + async for event in events: if isinstance(event, AgentRunUpdateEvent): # AgentRunUpdateEvent contains incremental text deltas from the underlying agent. # Print a prefix when the executor changes, then append updates on the same line. @@ -70,14 +68,9 @@ async def main(): print(f"{eid}:", end=" ", flush=True) last_executor_id = eid print(event.data, end="", flush=True) - elif isinstance(event, WorkflowCompletedEvent): - # Terminal event with the final reviewer output. - completed_event = event - - # Print the final consolidated reviewer result. - if completed_event: - print("\n===== Final Output =====") - print(completed_event.data) + elif isinstance(event, WorkflowOutputEvent): + print("===== Final Output =====") + print(event.data) """ Sample Output: diff --git a/python/samples/getting_started/workflow/agents/custom_agent_executors.py b/python/samples/getting_started/workflow/agents/custom_agent_executors.py index eb9497ae77..80f9871d6e 100644 --- a/python/samples/getting_started/workflow/agents/custom_agent_executors.py +++ b/python/samples/getting_started/workflow/agents/custom_agent_executors.py @@ -7,7 +7,6 @@ from agent_framework import ( ChatMessage, Executor, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, handler, ) @@ -23,7 +22,7 @@ then hands the conversation to a Reviewer agent which evaluates and finalizes th Purpose: Show how to wrap chat agents created by AzureChatClient inside workflow executors. Demonstrate the @handler pattern with typed inputs and typed WorkflowContext[T] outputs, connect executors with the fluent WorkflowBuilder, and finish -by emitting a WorkflowCompletedEvent from the terminal node. +by yielding outputs from the terminal node. Prerequisites: - Azure OpenAI configured for AzureChatClient with required environment variables. @@ -53,7 +52,7 @@ class Writer(Executor): super().__init__(agent=agent, id=id) @handler - async def handle(self, message: ChatMessage, ctx: WorkflowContext[list[ChatMessage]]) -> None: + async def handle(self, message: ChatMessage, ctx: WorkflowContext[list[ChatMessage], str]) -> None: """Generate content using the agent and forward the updated conversation. Contract for this handler: @@ -79,7 +78,7 @@ class Reviewer(Executor): This class demonstrates: - Consuming a typed payload produced upstream. - - Emitting a terminal WorkflowCompletedEvent with the final text outcome. + - Yielding the final text outcome to complete the workflow. """ agent: ChatAgent @@ -94,14 +93,14 @@ class Reviewer(Executor): super().__init__(agent=agent, id=id) @handler - async def handle(self, messages: list[ChatMessage], ctx: WorkflowContext[str]) -> None: + async def handle(self, messages: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage], str]) -> None: """Review the full conversation transcript and complete with a final string. This node consumes all messages so far. It uses its agent to produce the final text, - then signals completion by adding a WorkflowCompletedEvent to the event stream. + then signals completion by yielding the output. """ response = await self.agent.run(messages) - await ctx.add_event(WorkflowCompletedEvent(response.text)) + await ctx.yield_output(response.text) async def main(): @@ -118,12 +117,14 @@ async def main(): workflow = WorkflowBuilder().set_start_executor(writer).add_edge(writer, reviewer).build() # Run the workflow with the user's initial message. - # For foundational clarity, use run (non streaming) and print the terminal event. + # For foundational clarity, use run (non streaming) and print the workflow output. events = await workflow.run( ChatMessage(role="user", text="Create a slogan for a new electric SUV that is affordable and fun to drive.") ) - # The terminal node emits a WorkflowCompletedEvent; print its contents. - print(events.get_completed_event()) + # The terminal node yields output; print its contents. + outputs = events.get_outputs() + if outputs: + print(outputs[-1]) if __name__ == "__main__": diff --git a/python/samples/getting_started/workflow/agents/foundry_chat_agents_streaming.py b/python/samples/getting_started/workflow/agents/foundry_chat_agents_streaming.py index e7bf4b0ffa..ab600dea0b 100644 --- a/python/samples/getting_started/workflow/agents/foundry_chat_agents_streaming.py +++ b/python/samples/getting_started/workflow/agents/foundry_chat_agents_streaming.py @@ -5,7 +5,7 @@ from collections.abc import Awaitable, Callable from contextlib import AsyncExitStack from typing import Any -from agent_framework import AgentRunUpdateEvent, WorkflowBuilder, WorkflowCompletedEvent +from agent_framework import AgentRunUpdateEvent, WorkflowBuilder, WorkflowOutputEvent from agent_framework.foundry import FoundryChatClient from azure.identity.aio import AzureCliCredential @@ -21,7 +21,7 @@ Show how to wire chat agents directly into a WorkflowBuilder pipeline where agen Demonstrate: - Automatic streaming of agent deltas via AgentRunUpdateEvent. - A simple console aggregator that groups updates by executor id and prints them as they arrive. -- A final WorkflowCompletedEvent that contains the reviewer outcome after both agents finish. +- The workflow completes when idle and outputs are available in events.get_outputs(). Prerequisites: - Foundry Agent Service configured, along with the required environment variables. @@ -69,12 +69,10 @@ async def main() -> None: workflow = WorkflowBuilder().set_start_executor(writer).add_edge(writer, reviewer).build() - completed: WorkflowCompletedEvent | None = None last_executor_id: str | None = None - async for event in workflow.run_stream( - "Create a slogan for a new electric SUV that is affordable and fun to drive." - ): + events = workflow.run_stream("Create a slogan for a new electric SUV that is affordable and fun to drive.") + async for event in events: if isinstance(event, AgentRunUpdateEvent): eid = event.executor_id if eid != last_executor_id: @@ -83,13 +81,9 @@ async def main() -> None: print(f"{eid}:", end=" ", flush=True) last_executor_id = eid print(event.data, end="", flush=True) - elif isinstance(event, WorkflowCompletedEvent): - completed = event - - if completed: - print("\n===== Final Output =====") - print(completed.data) - + elif isinstance(event, WorkflowOutputEvent): + print("\n===== Final output =====") + print(event.data) finally: await close() diff --git a/python/samples/getting_started/workflow/agents/workflow_as_agent_human_in_the_loop.py b/python/samples/getting_started/workflow/agents/workflow_as_agent_human_in_the_loop.py index e80c181128..e76615becd 100644 --- a/python/samples/getting_started/workflow/agents/workflow_as_agent_human_in_the_loop.py +++ b/python/samples/getting_started/workflow/agents/workflow_as_agent_human_in_the_loop.py @@ -39,7 +39,7 @@ This sample demonstrates how to build a workflow agent that escalates uncertain decisions to a human manager. A Worker generates results, while a Reviewer evaluates them. When the Reviewer is not confident, it escalates the decision to a human via RequestInfoExecutor, receives the human response, and then -forwards that response back to the Worker. +forwards that response back to the Worker. The workflow completes when idle. Prerequisites: - OpenAI account configured and accessible for OpenAIChatClient. diff --git a/python/samples/getting_started/workflow/agents/workflow_as_agent_reflection_pattern.py b/python/samples/getting_started/workflow/agents/workflow_as_agent_reflection_pattern.py index f9d9e9792c..775c927af9 100644 --- a/python/samples/getting_started/workflow/agents/workflow_as_agent_reflection_pattern.py +++ b/python/samples/getting_started/workflow/agents/workflow_as_agent_reflection_pattern.py @@ -27,7 +27,7 @@ This sample demonstrates how to wrap a workflow as an agent using WorkflowAgent. It uses a reflection pattern where a Worker executor generates responses and a Reviewer executor evaluates them. If the response is not approved, the Worker regenerates the output based on feedback until the Reviewer approves it. Only -approved responses are emitted to the external consumer. +approved responses are emitted to the external consumer. The workflow completes when idle. Key Concepts Demonstrated: - WorkflowAgent: Wraps a workflow to behave like a regular agent. diff --git a/python/samples/getting_started/workflow/checkpoint/checkpoint_with_human_in_the_loop.py b/python/samples/getting_started/workflow/checkpoint/checkpoint_with_human_in_the_loop.py index 06b75f8148..528c986d6c 100644 --- a/python/samples/getting_started/workflow/checkpoint/checkpoint_with_human_in_the_loop.py +++ b/python/samples/getting_started/workflow/checkpoint/checkpoint_with_human_in_the_loop.py @@ -19,8 +19,8 @@ from agent_framework import ( RequestResponse, Role, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, + WorkflowOutputEvent, WorkflowRunState, WorkflowStatusEvent, handler, @@ -87,7 +87,7 @@ class BriefPreparer(Executor): self._agent_id = agent_id @handler - async def prepare(self, brief: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None: + async def prepare(self, brief: str, ctx: WorkflowContext[AgentExecutorRequest, str]) -> None: # Collapse errant whitespace so the prompt is stable between runs. normalized = " ".join(brief.split()).strip() if not normalized.endswith("."): @@ -133,7 +133,7 @@ class ReviewGateway(Executor): async def on_agent_response( self, response: AgentExecutorResponse, - ctx: WorkflowContext[HumanApprovalRequest], + ctx: WorkflowContext[HumanApprovalRequest, str], ) -> None: # Capture the agent output so we can surface it to the reviewer and # persist iterations. The `RequestInfoExecutor` relies on this state to @@ -157,7 +157,7 @@ class ReviewGateway(Executor): async def on_human_feedback( self, feedback: RequestResponse[HumanApprovalRequest, str], - ctx: WorkflowContext[AgentExecutorRequest | str], + ctx: WorkflowContext[AgentExecutorRequest | str, str], ) -> None: # The RequestResponse wrapper gives us both the human data and the # original request message, even when resuming from checkpoints. @@ -190,11 +190,11 @@ class FinaliseExecutor(Executor): """Publishes the approved text.""" @handler - async def publish(self, text: str, ctx: WorkflowContext[Any]) -> None: + async def publish(self, text: str, ctx: WorkflowContext[Any, str]) -> None: # Store the output so diagnostics or a UI could fetch the final copy. await ctx.set_state({"published_text": text}) - # Emit a workflow completion event so the runner stops cleanly. - await ctx.add_event(WorkflowCompletedEvent(text)) + # Yield the final output so the workflow completes cleanly. + await ctx.yield_output(text) def create_workflow(*, checkpoint_storage: FileCheckpointStorage | None = None) -> "Workflow": @@ -264,17 +264,17 @@ def _render_checkpoint_summary(checkpoints: list["WorkflowCheckpoint"]) -> None: print(line) -def _print_events(events: list[Any]) -> tuple[WorkflowCompletedEvent | None, list[tuple[str, HumanApprovalRequest]]]: +def _print_events(events: list[Any]) -> tuple[str | None, list[tuple[str, HumanApprovalRequest]]]: """Echo workflow events to the console and collect outstanding requests.""" - completed: WorkflowCompletedEvent | None = None + completed_output: str | None = None requests: list[tuple[str, HumanApprovalRequest]] = [] for event in events: print(f"Event: {event}") - if isinstance(event, WorkflowCompletedEvent): - completed = event - elif isinstance(event, RequestInfoEvent) and isinstance(event.data, HumanApprovalRequest): + if isinstance(event, WorkflowOutputEvent): + completed_output = event.data + if isinstance(event, RequestInfoEvent) and isinstance(event.data, HumanApprovalRequest): # Capture pending human approvals so the caller can ask the user for # input after the current batch of events is processed. requests.append((event.request_id, event.data)) @@ -284,7 +284,7 @@ def _print_events(events: list[Any]) -> tuple[WorkflowCompletedEvent | None, lis }: print(f"Workflow state: {event.state.name}") - return completed, requests + return completed_output, requests def _prompt_for_responses(requests: list[tuple[str, HumanApprovalRequest]]) -> dict[str, str] | None: @@ -350,14 +350,14 @@ async def _consume(stream: AsyncIterable[Any]) -> list[Any]: return [event async for event in stream] -async def run_interactive_session(workflow: "Workflow", initial_message: str) -> WorkflowCompletedEvent | None: +async def run_interactive_session(workflow: "Workflow", initial_message: str) -> str | None: """Run the workflow until it either finishes or pauses for human input.""" pending_responses: dict[str, str] | None = None - completed: WorkflowCompletedEvent | None = None + completed_output: str | None = None first = True - while completed is None: + while completed_output is None: if first: # Kick off the workflow with the initial brief. The returned events # include RequestInfo events when the agent produces a draft. @@ -369,10 +369,11 @@ async def run_interactive_session(workflow: "Workflow", initial_message: str) -> else: break - completed, requests = _print_events(events) - pending_responses = _prompt_for_responses(requests) + completed_output, requests = _print_events(events) + if completed_output is None: + pending_responses = _prompt_for_responses(requests) - return completed + return completed_output async def resume_from_checkpoint( @@ -391,21 +392,24 @@ async def resume_from_checkpoint( responses=pre_supplied, ) ) - completed, requests = _print_events(events) - if pre_supplied and not requests and completed is None: + completed_output, requests = _print_events(events) + if pre_supplied and not requests and completed_output is None: # When the checkpoint only needed the provided answers we let the user # know the workflow is waiting for the next superstep (usually another # agent response). print("Pre-supplied responses applied automatically; workflow is now waiting for the next step.") pending = _prompt_for_responses(requests) - while completed is None and pending: + while completed_output is None and pending: events = await _consume(workflow.send_responses_streaming(pending)) - completed, requests = _print_events(events) - pending = _prompt_for_responses(requests) + completed_output, requests = _print_events(events) + if completed_output is None: + pending = _prompt_for_responses(requests) + else: + break - if completed: - print(f"Workflow completed with: {completed.data}") + if completed_output: + print(f"Workflow completed with: {completed_output}") async def main() -> None: @@ -427,7 +431,7 @@ async def main() -> None: print("Running workflow (human approval required)...") completed = await run_interactive_session(workflow, initial_message=brief) if completed: - print(f"Initial run completed with final copy: {completed.data}") + print(f"Initial run completed with final copy: {completed}") else: print("Initial run paused for human input.") diff --git a/python/samples/getting_started/workflow/checkpoint/checkpoint_with_resume.py b/python/samples/getting_started/workflow/checkpoint/checkpoint_with_resume.py index c8f675147c..ab5f5d4053 100644 --- a/python/samples/getting_started/workflow/checkpoint/checkpoint_with_resume.py +++ b/python/samples/getting_started/workflow/checkpoint/checkpoint_with_resume.py @@ -15,7 +15,6 @@ from agent_framework import ( RequestInfoExecutor, Role, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, handler, ) @@ -49,6 +48,7 @@ What you learn: - 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_from_checkpoint. +- How workflows complete by yielding outputs when idle, not via explicit completion events. Prerequisites: - Azure AI or Azure OpenAI available for AzureChatClient. @@ -115,10 +115,10 @@ class SubmitToLowerAgent(Executor): class FinalizeFromAgent(Executor): - """Consumes the AgentExecutorResponse and emits the terminal WorkflowCompletedEvent.""" + """Consumes the AgentExecutorResponse and yields the final result.""" @handler - async def finalize(self, response: AgentExecutorResponse, ctx: WorkflowContext[Any]) -> None: + 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. @@ -130,8 +130,8 @@ class FinalizeFromAgent(Executor): "final": True, }) - # Emit a terminal event so external consumers see the final value. - await ctx.add_event(WorkflowCompletedEvent(result)) + # Yield the final result so external consumers see the final value. + await ctx.yield_output(result) class ReverseTextExecutor(Executor): @@ -185,6 +185,7 @@ def create_workflow(checkpoint_storage: FileCheckpointStorage) -> "Workflow": .build() ) + def _render_checkpoint_summary(checkpoints: list["WorkflowCheckpoint"]) -> None: """Display human-friendly checkpoint metadata using framework summaries.""" @@ -297,7 +298,6 @@ async def main(): Event: ExecutorInvokeEvent(executor_id=submit_lower) Event: ExecutorInvokeEvent(executor_id=lower_agent) Event: ExecutorInvokeEvent(executor_id=finalize) - Event: WorkflowCompletedEvent(data=dlrow olleh) 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' @@ -316,7 +316,6 @@ async def main(): Resumed Event: ExecutorInvokeEvent(executor_id=submit_lower) Resumed Event: ExecutorInvokeEvent(executor_id=lower_agent) Resumed Event: ExecutorInvokeEvent(executor_id=finalize) - Resumed Event: WorkflowCompletedEvent(data=dlrow olleh) """ # noqa: E501 diff --git a/python/samples/getting_started/workflow/composition/sub_workflow_basics.py b/python/samples/getting_started/workflow/composition/sub_workflow_basics.py index 1c8013bacf..80a6078898 100644 --- a/python/samples/getting_started/workflow/composition/sub_workflow_basics.py +++ b/python/samples/getting_started/workflow/composition/sub_workflow_basics.py @@ -4,10 +4,11 @@ import asyncio from dataclasses import dataclass from typing import Any +from typing_extensions import Never + from agent_framework import ( Executor, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, WorkflowEvent, WorkflowExecutor, @@ -20,6 +21,7 @@ Sample: Sub-Workflows (Basics) What it does: - Shows how a parent workflow invokes a sub-workflow via `WorkflowExecutor` and collects results. - Example: parent orchestrates multiple text processors that count words/characters. +- Demonstrates how sub-workflows complete by yielding outputs when processing is done. Prerequisites: - No external services required. @@ -60,7 +62,9 @@ class TextProcessor(Executor): super().__init__(id="text_processor") @handler - async def process_text(self, request: TextProcessingRequest, ctx: WorkflowContext[TextProcessingResult]) -> None: + async def process_text( + self, request: TextProcessingRequest, ctx: WorkflowContext[Never, TextProcessingResult] + ) -> None: """Process a text string and return statistics.""" text_preview = f"'{request.text[:50]}{'...' if len(request.text) > 50 else ''}'" print(f"🔍 Sub-workflow processing text (Task {request.task_id}): {text_preview}") @@ -80,8 +84,8 @@ class TextProcessor(Executor): ) print(f"✅ Sub-workflow completed task {request.task_id}") - # Signal completion - await ctx.add_event(WorkflowCompletedEvent(data=result)) + # Signal completion by yielding the result + await ctx.yield_output(result) # Parent workflow @@ -110,7 +114,7 @@ class TextProcessingOrchestrator(Executor): await ctx.send_message(request, target_id="text_processor_workflow") @handler - async def collect_result(self, result: TextProcessingResult, ctx: WorkflowContext[None]) -> None: + async def collect_result(self, result: TextProcessingResult, ctx: WorkflowContext) -> None: """Collect results from sub-workflows.""" print(f"📥 Collected result from {result.task_id}") self.results.append(result) @@ -173,7 +177,7 @@ async def main(): print("=" * 60) # Step 4: Run the workflow - result = await main_workflow.run(test_texts) + await main_workflow.run(test_texts) # Step 5: Display results print("\n📊 Processing Results:") diff --git a/python/samples/getting_started/workflow/composition/sub_workflow_parallel_requests.py b/python/samples/getting_started/workflow/composition/sub_workflow_parallel_requests.py index a5a5ef45bd..22f3f37b64 100644 --- a/python/samples/getting_started/workflow/composition/sub_workflow_parallel_requests.py +++ b/python/samples/getting_started/workflow/composition/sub_workflow_parallel_requests.py @@ -4,11 +4,12 @@ import asyncio from dataclasses import dataclass from typing import Any +from typing_extensions import Never + from agent_framework import ( Executor, RequestInfoExecutor, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, handler, ) @@ -138,6 +139,7 @@ class ResourceRequester(Executor): for req_data in requests: req_type = req_data.get("request_type", "resource") + request: ResourceRequest | PolicyCheckRequest if req_type == "resource": print(f" 📦 Requesting resource: {req_data.get('type', 'cpu')} x{req_data.get('amount', 1)}") request = ResourceRequest( @@ -164,7 +166,7 @@ class ResourceRequester(Executor): async def handle_resource_response( self, response: RequestResponse[ResourceRequest, ResourceResponse], - ctx: WorkflowContext[None], + ctx: WorkflowContext[Never, RequestFinished], ) -> None: """Handle resource allocation response.""" if response.data: @@ -174,12 +176,12 @@ class ResourceRequester(Executor): f"from {response.data.source}" ) if self._collect_results(): - # Emit completion event and send RequestFinished to the parent workflow. - await ctx.add_event(WorkflowCompletedEvent(RequestFinished())) + # Yield completion result to the parent workflow. + await ctx.yield_output(RequestFinished()) @handler async def handle_policy_response( - self, response: RequestResponse[PolicyCheckRequest, PolicyResponse], ctx: WorkflowContext[None] + self, response: RequestResponse[PolicyCheckRequest, PolicyResponse], ctx: WorkflowContext[Never, RequestFinished] ) -> None: """Handle policy check response.""" if response.data: @@ -189,8 +191,8 @@ class ResourceRequester(Executor): f"{response.data.approved} - {response.data.reason}" ) if self._collect_results(): - # Emit completion event and send RequestFinished to the parent workflow. - await ctx.add_event(WorkflowCompletedEvent(RequestFinished())) + # Yield completion result to the parent workflow. + await ctx.yield_output(RequestFinished()) def _collect_results(self) -> bool: """Collect and summarize results.""" @@ -213,7 +215,7 @@ class ResourceCache(Executor): @intercepts_request async def check_cache( - self, request: ResourceRequest, ctx: WorkflowContext[None] + self, request: ResourceRequest, ctx: WorkflowContext ) -> RequestResponse[ResourceRequest, ResourceResponse]: """Intercept RESOURCE requests and check cache first.""" print(f"🏪 CACHE interceptor checking: {request.amount} {request.resource_type}") @@ -234,7 +236,7 @@ class ResourceCache(Executor): @handler async def collect_result( - self, response: RequestResponse[ResourceRequest, ResourceResponse], ctx: WorkflowContext[None] + self, response: RequestResponse[ResourceRequest, ResourceResponse], ctx: WorkflowContext ) -> None: """Collect results from external requests that were forwarded.""" if response.data and response.data.source != "cache": # Don't double-count our own results @@ -263,7 +265,7 @@ class PolicyEngine(Executor): @intercepts_request async def check_policy( - self, request: PolicyCheckRequest, ctx: WorkflowContext[None] + self, request: PolicyCheckRequest, ctx: WorkflowContext ) -> RequestResponse[PolicyCheckRequest, PolicyResponse]: """Intercept POLICY requests and apply rules.""" print(f"🛡️ POLICY interceptor checking: {request.amount} {request.resource_type}, policy={request.policy_type}") @@ -286,7 +288,7 @@ class PolicyEngine(Executor): @handler async def collect_policy_result( - self, response: RequestResponse[PolicyCheckRequest, PolicyResponse], ctx: WorkflowContext[None] + self, response: RequestResponse[PolicyCheckRequest, PolicyResponse], ctx: WorkflowContext ) -> None: """Collect policy results from external requests that were forwarded.""" if response.data: @@ -299,15 +301,15 @@ class Coordinator(Executor): super().__init__(id="coordinator") @handler - async def start(self, requests: list[dict[str, Any]], ctx: WorkflowContext[object]) -> None: + async def start(self, requests: list[dict[str, Any]], ctx: WorkflowContext[list[dict[str, Any]]]) -> None: """Start the resource allocation process.""" await ctx.send_message(requests, target_id="resource_workflow") @handler - async def handle_completion(self, completion: RequestFinished, ctx: WorkflowContext[None]) -> None: + async def handle_completion(self, completion: RequestFinished, ctx: WorkflowContext) -> None: """Handle sub-workflow completion. - It comes from the sub-workflow emitted WorkflowCompletionEvent's data field. + It comes from the sub-workflow yielded output. """ print("🎯 Main workflow received completion.") @@ -377,10 +379,10 @@ async def main() -> None: # 8. Run the workflow print("🎬 Running workflow...") - result = await main_workflow.run(test_requests) + events = await main_workflow.run(test_requests) # 9. Handle any external requests that couldn't be intercepted - request_events = result.get_request_info_events() + request_events = events.get_request_info_events() if request_events: print(f"\n🌐 Handling {len(request_events)} external request(s)...") diff --git a/python/samples/getting_started/workflow/composition/sub_workflow_request_interception.py b/python/samples/getting_started/workflow/composition/sub_workflow_request_interception.py index 2cc2be9a01..dfa1dc57b6 100644 --- a/python/samples/getting_started/workflow/composition/sub_workflow_request_interception.py +++ b/python/samples/getting_started/workflow/composition/sub_workflow_request_interception.py @@ -2,7 +2,6 @@ import asyncio from dataclasses import dataclass -from typing import Any from agent_framework import ( Executor, @@ -10,7 +9,6 @@ from agent_framework import ( RequestInfoMessage, RequestResponse, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, WorkflowExecutor, handler, @@ -43,6 +41,7 @@ Key concepts demonstrated: - Concurrent processing: Multiple emails processed simultaneously without interference - External request routing: RequestInfoExecutor handles forwarded external requests - Sub-workflow isolation: Sub-workflows work normally without knowing they're nested +- Sub-workflows complete by yielding outputs when validation is finished Prerequisites: - No external services required (external calls are simulated via `RequestInfoExecutor`). @@ -101,7 +100,9 @@ class EmailValidator(Executor): @handler async def validate_request( - self, request: EmailValidationRequest, ctx: WorkflowContext[DomainCheckRequest | ValidationResult] + self, + request: EmailValidationRequest, + ctx: WorkflowContext[DomainCheckRequest | ValidationResult, ValidationResult], ) -> None: """Validate an email address.""" print(f"🔍 Sub-workflow validating email: {request.email}") @@ -112,7 +113,7 @@ class EmailValidator(Executor): if not domain: print(f"❌ Invalid email format: {request.email}") result = ValidationResult(email=request.email, is_valid=False, reason="Invalid email format") - await ctx.add_event(WorkflowCompletedEvent(data=result)) + await ctx.yield_output(result) return print(f"🌐 Sub-workflow requesting domain check for: {domain}") @@ -126,7 +127,7 @@ class EmailValidator(Executor): async def handle_domain_response( self, response: RequestResponse[DomainCheckRequest, bool], - ctx: WorkflowContext[ValidationResult], + ctx: WorkflowContext[ValidationResult, ValidationResult], ) -> None: """Handle domain check response from RequestInfo with correlation.""" approved = bool(response.data) @@ -151,7 +152,7 @@ class EmailValidator(Executor): reason="Domain approved" if approved else "Domain not approved", ) print(f"✅ Sub-workflow completing validation for: {email}") - await ctx.add_event(WorkflowCompletedEvent(data=result)) + await ctx.yield_output(result) # 3. Implement the parent workflow with request interception @@ -181,7 +182,7 @@ class SmartEmailOrchestrator(Executor): @intercepts_request async def check_domain( - self, request: DomainCheckRequest, ctx: WorkflowContext[Any] + self, request: DomainCheckRequest, ctx: WorkflowContext ) -> RequestResponse[DomainCheckRequest, bool]: """Intercept domain check requests from sub-workflows.""" print(f"🔍 Parent intercepting domain check for: {request.domain}") @@ -192,8 +193,8 @@ class SmartEmailOrchestrator(Executor): return RequestResponse[DomainCheckRequest, bool].forward() @handler - async def collect_result(self, result: ValidationResult, ctx: WorkflowContext[None]) -> None: - """Collect validation results. It comes from the sub-workflow emitted WorkflowCompletionEvent's data field.""" + async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None: + """Collect validation results. It comes from the sub-workflow yielded output.""" status_icon = "✅" if result.is_valid else "❌" print(f"📥 {status_icon} Validation result: {result.email} -> {result.reason}") self._results.append(result) diff --git a/python/samples/getting_started/workflow/control-flow/edge_condition.py b/python/samples/getting_started/workflow/control-flow/edge_condition.py index 405b237fe5..3b051148ad 100644 --- a/python/samples/getting_started/workflow/control-flow/edge_condition.py +++ b/python/samples/getting_started/workflow/control-flow/edge_condition.py @@ -4,6 +4,8 @@ import asyncio import os from typing import Any +from typing_extensions import Never + from agent_framework import ( # Core chat primitives used to build requests AgentExecutor, # Wraps an LLM agent that can be invoked inside a workflow AgentExecutorRequest, # Input message bundle for an AgentExecutor @@ -11,7 +13,6 @@ from agent_framework import ( # Core chat primitives used to build requests ChatMessage, Role, WorkflowBuilder, # Fluent builder for wiring executors and edges - WorkflowCompletedEvent, # Event we emit at the end to signal completion WorkflowContext, # Per-run context and event bus executor, # Decorator to declare a Python function as a workflow executor ) @@ -41,15 +42,16 @@ and have the Azure OpenAI environment variables set as documented in the getting High level flow: 1) spam_detection_agent reads an email and returns DetectionResult. 2) If not spam, we transform the detection output into a user message for email_assistant_agent, then finish by -sending the drafted reply. -3) If spam, we short circuit to a spam handler that emits a completion event. +yielding the drafted reply as workflow output. +3) If spam, we short circuit to a spam handler that yields a spam notice as workflow output. Output: -- The final WorkflowCompletedEvent is printed to stdout, either with a drafted reply or a spam notice. +- The final workflow output is printed to stdout, either with a drafted reply or a spam notice. Notes: - Conditions read the agent response text and validate it into DetectionResult for robust routing. - Executors are small and single purpose to keep control flow easy to follow. +- The workflow completes when it becomes idle, not via explicit completion events. """ @@ -96,18 +98,18 @@ def get_condition(expected_result: bool): @executor(id="send_email") -async def handle_email_response(response: AgentExecutorResponse, ctx: WorkflowContext[None]) -> None: - # Downstream of the email assistant. Parse a validated EmailResponse and emit a completion event. +async def handle_email_response(response: AgentExecutorResponse, ctx: WorkflowContext[Never, str]) -> None: + # Downstream of the email assistant. Parse a validated EmailResponse and yield the workflow output. email_response = EmailResponse.model_validate_json(response.agent_run_response.text) - await ctx.add_event(WorkflowCompletedEvent(f"Email sent:\n{email_response.response}")) + await ctx.yield_output(f"Email sent:\n{email_response.response}") @executor(id="handle_spam") -async def handle_spam_classifier_response(response: AgentExecutorResponse, ctx: WorkflowContext[None]) -> None: - # Spam path. Confirm the DetectionResult and finish with the reason. Guard against accidental non spam input. +async def handle_spam_classifier_response(response: AgentExecutorResponse, ctx: WorkflowContext[Never, str]) -> None: + # Spam path. Confirm the DetectionResult and yield the workflow output. Guard against accidental non spam input. detection = DetectionResult.model_validate_json(response.agent_run_response.text) if detection.is_spam: - await ctx.add_event(WorkflowCompletedEvent(f"Email marked as spam: {detection.reason}")) + await ctx.yield_output(f"Email marked as spam: {detection.reason}") else: # This indicates the routing predicate and executor contract are out of sync. raise RuntimeError("This executor should only handle spam messages.") @@ -184,11 +186,12 @@ async def main() -> None: email = email_file.read() # Execute the workflow. Since the start is an AgentExecutor, pass an AgentExecutorRequest. - # run_stream yields events as they occur. We watch for the terminal WorkflowCompletedEvent and print it. + # The workflow completes when it becomes idle (no more work to do). request = AgentExecutorRequest(messages=[ChatMessage(Role.USER, text=email)], should_respond=True) - async for event in workflow.run_stream(request): - if isinstance(event, WorkflowCompletedEvent): - print(f"{event}") + events = await workflow.run(request) + outputs = events.get_outputs() + if outputs: + print(f"Workflow output: {outputs[0]}") """ Sample Output: @@ -214,7 +217,7 @@ async def main() -> None: (555) 123-4567 ---------------------------------------- - WorkflowCompletedEvent(data=Email sent: +Workflow output: Email sent: Hi Alex, Thank you for the follow-up and for summarizing the action items from this morning's meeting. The points you listed accurately reflect our discussion, and I don't have any additional items to add at this time. @@ -224,7 +227,7 @@ async def main() -> None: Thank you again for outlining the next steps. Best regards, - Sarah) + Sarah """ # noqa: E501 diff --git a/python/samples/getting_started/workflow/control-flow/multi_selection_edge_group.py b/python/samples/getting_started/workflow/control-flow/multi_selection_edge_group.py index 547f0ed740..ec8265bd55 100644 --- a/python/samples/getting_started/workflow/control-flow/multi_selection_edge_group.py +++ b/python/samples/getting_started/workflow/control-flow/multi_selection_edge_group.py @@ -8,6 +8,8 @@ from dataclasses import dataclass from typing import Literal from uuid import uuid4 +from typing_extensions import Never + from agent_framework import ( AgentExecutor, AgentExecutorRequest, @@ -15,9 +17,9 @@ from agent_framework import ( ChatMessage, Role, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, WorkflowEvent, + WorkflowOutputEvent, executor, ) from agent_framework.azure import AzureChatClient @@ -30,7 +32,7 @@ Sample: Multi-Selection Edge Group for email triage and response. The workflow stores an email, classifies it as NotSpam, Spam, or Uncertain, and then routes to one or more branches. Non-spam emails are drafted into replies, long ones are also summarized, spam is blocked, and uncertain cases are -flagged. Each path ends with simulated database persistence. +flagged. Each path ends with simulated database persistence. The workflow completes when it becomes idle. Purpose: Demonstrate how to use a multi-selection edge group to fan out from one executor to multiple possible targets. @@ -123,9 +125,9 @@ async def submit_to_email_assistant(analysis: AnalysisResult, ctx: WorkflowConte @executor(id="finalize_and_send") -async def finalize_and_send(response: AgentExecutorResponse, ctx: WorkflowContext[None]) -> None: +async def finalize_and_send(response: AgentExecutorResponse, ctx: WorkflowContext[Never, str]) -> None: parsed = EmailResponse.model_validate_json(response.agent_run_response.text) - await ctx.add_event(WorkflowCompletedEvent(f"Email sent: {parsed.response}")) + await ctx.yield_output(f"Email sent: {parsed.response}") @executor(id="summarize_email") @@ -155,28 +157,26 @@ async def merge_summary(response: AgentExecutorResponse, ctx: WorkflowContext[An @executor(id="handle_spam") -async def handle_spam(analysis: AnalysisResult, ctx: WorkflowContext[None]) -> None: +async def handle_spam(analysis: AnalysisResult, ctx: WorkflowContext[Never, str]) -> None: if analysis.spam_decision == "Spam": - await ctx.add_event(WorkflowCompletedEvent(f"Email marked as spam: {analysis.reason}")) + await ctx.yield_output(f"Email marked as spam: {analysis.reason}") else: raise RuntimeError("This executor should only handle Spam messages.") @executor(id="handle_uncertain") -async def handle_uncertain(analysis: AnalysisResult, ctx: WorkflowContext[None]) -> None: +async def handle_uncertain(analysis: AnalysisResult, ctx: WorkflowContext[Never, str]) -> None: if analysis.spam_decision == "Uncertain": email: Email | None = await ctx.get_shared_state(f"{EMAIL_STATE_PREFIX}{analysis.email_id}") - await ctx.add_event( - WorkflowCompletedEvent( - f"Email marked as uncertain: {analysis.reason}. Email content: {getattr(email, 'email_content', '')}" - ) + await ctx.yield_output( + f"Email marked as uncertain: {analysis.reason}. Email content: {getattr(email, 'email_content', '')}" ) else: raise RuntimeError("This executor should only handle Uncertain messages.") @executor(id="database_access") -async def database_access(analysis: AnalysisResult, ctx: WorkflowContext[None]) -> None: +async def database_access(analysis: AnalysisResult, ctx: WorkflowContext[Never, str]) -> None: # Simulate DB writes for email and analysis (and summary if present) await asyncio.sleep(0.05) await ctx.add_event(DatabaseEvent(f"Email {analysis.email_id} saved to database.")) @@ -263,14 +263,18 @@ async def main() -> None: print("Unable to find resource file, using default text.") email = "Hello team, here are the updates for this week..." + # Print outputs and database events from streaming async for event in workflow.run_stream(email): - if isinstance(event, (WorkflowCompletedEvent, DatabaseEvent)): + if isinstance(event, DatabaseEvent): print(f"{event}") + elif isinstance(event, WorkflowOutputEvent): + print(f"Workflow output: {event.data}") """ Sample Output: - WorkflowCompletedEvent(data=Email sent: Hi Alex, + DatabaseEvent(data=Email 32021432-2d4e-4c54-b04c-f81b4120340c saved to database.) + Workflow output: Email sent: Hi Alex, Thank you for summarizing the action items from this morning's meeting. I have noted the three tasks and will begin working on them right away. @@ -281,8 +285,7 @@ async def main() -> None: If anything else comes up, please let me know. Best regards, - Sarah) - DatabaseEvent(data=Email 32021432-2d4e-4c54-b04c-f81b4120340c saved to database.) + Sarah """ # noqa: E501 diff --git a/python/samples/getting_started/workflow/control-flow/sequential_executors.py b/python/samples/getting_started/workflow/control-flow/sequential_executors.py index d4921e99fa..c9973c5296 100644 --- a/python/samples/getting_started/workflow/control-flow/sequential_executors.py +++ b/python/samples/getting_started/workflow/control-flow/sequential_executors.py @@ -1,13 +1,15 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio -from typing import Any +from typing import cast + +from typing_extensions import Never from agent_framework import ( Executor, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, + WorkflowOutputEvent, handler, ) @@ -19,8 +21,8 @@ the second reverses the text and completes the workflow. The run_stream loop pri Purpose: Show how to define explicit Executor classes with @handler methods, wire them in order with -WorkflowBuilder, and consume streaming events. Demonstrate typed WorkflowContext[T] for outputs, -ctx.send_message to pass intermediate values, and ctx.add_event to signal completion with a WorkflowCompletedEvent. +WorkflowBuilder, and consume streaming events. Demonstrate typed WorkflowContext[T_Out, T_W_Out] for outputs, +ctx.send_message to pass intermediate values, and ctx.yield_output to provide workflow outputs. Prerequisites: - No external services required. @@ -44,21 +46,21 @@ class UpperCaseExecutor(Executor): class ReverseTextExecutor(Executor): - """Reverses the incoming string and completes the workflow. + """Reverses the incoming string and yields workflow output. Concepts: - - Use ctx.add_event to publish a WorkflowCompletedEvent when the terminal result is ready. + - Use ctx.yield_output to provide workflow outputs when the terminal result is ready. - The terminal node does not forward messages further. """ @handler - async def reverse_text(self, text: str, ctx: WorkflowContext[Any]) -> None: - """Reverse the input string and emit a completion event.""" + async def reverse_text(self, text: str, ctx: WorkflowContext[Never, str]) -> None: + """Reverse the input string and yield the workflow output.""" result = text[::-1] - await ctx.add_event(WorkflowCompletedEvent(result)) + await ctx.yield_output(result) -async def main(): +async def main() -> None: """Build a two step sequential workflow and run it with streaming to observe events.""" # Step 1: Create executor instances. upper_case_executor = UpperCaseExecutor(id="upper_case_executor") @@ -74,15 +76,15 @@ async def main(): ) # Step 3: Stream events for a single input. - # The stream will include executor invoke and completion events, plus the final WorkflowCompletedEvent. - completion_event = None + # The stream will include executor invoke and completion events, plus workflow outputs. + outputs: list[str] = [] async for event in workflow.run_stream("hello world"): print(f"Event: {event}") - if isinstance(event, WorkflowCompletedEvent): - completion_event = event + if isinstance(event, WorkflowOutputEvent): + outputs.append(cast(str, event.data)) - if completion_event: - print(f"Workflow completed with result: {completion_event.data}") + if outputs: + print(f"Workflow outputs: {outputs}") if __name__ == "__main__": diff --git a/python/samples/getting_started/workflow/control-flow/sequential_streaming.py b/python/samples/getting_started/workflow/control-flow/sequential_streaming.py index fe6056cfa6..1c3b52a7ea 100644 --- a/python/samples/getting_started/workflow/control-flow/sequential_streaming.py +++ b/python/samples/getting_started/workflow/control-flow/sequential_streaming.py @@ -2,19 +2,20 @@ import asyncio -from agent_framework import WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, executor +from typing_extensions import Never + +from agent_framework import WorkflowBuilder, WorkflowContext, WorkflowOutputEvent, executor """ Sample: Foundational sequential workflow with streaming using function-style executors. Two lightweight steps run in order. The first converts text to uppercase. -The second reverses the text and completes the workflow. Events are printed as they arrive from run_stream. +The second reverses the text and yields the workflow output. Events are printed as they arrive from run_stream. Purpose: Show how to declare executors with the @executor decorator, connect them with WorkflowBuilder, -pass intermediate values using ctx.send_message, and signal completion with ctx.add_event by emitting a -WorkflowCompletedEvent. Demonstrate how streaming exposes ExecutorInvokedEvent and WorkflowCompletedEvent -for observability. +pass intermediate values using ctx.send_message, and yield final output using ctx.yield_output(). +Demonstrate how streaming exposes ExecutorInvokedEvent and ExecutorCompletedEvent for observability. Prerequisites: - No external services required. @@ -37,17 +38,17 @@ async def to_upper_case(text: str, ctx: WorkflowContext[str]) -> None: @executor(id="reverse_text_executor") -async def reverse_text(text: str, ctx: WorkflowContext[str]) -> None: - """Reverse the input and complete the workflow with the final result. +async def reverse_text(text: str, ctx: WorkflowContext[Never, str]) -> None: + """Reverse the input and yield the workflow output. Concepts: - - Terminal nodes publish a WorkflowCompletedEvent using ctx.add_event. - - No further messages are forwarded after completion. + - Terminal nodes yield output using ctx.yield_output(). + - The workflow completes when it becomes idle (no more work to do). """ result = text[::-1] - # Emit the terminal event that carries the final output for this run. - await ctx.add_event(WorkflowCompletedEvent(result)) + # Yield the final output for this workflow run. + await ctx.yield_output(result) async def main(): @@ -57,17 +58,11 @@ async def main(): workflow = WorkflowBuilder().add_edge(to_upper_case, reverse_text).set_start_executor(to_upper_case).build() # Step 3: Run the workflow and stream events in real time. - completion_event = None async for event in workflow.run_stream("hello world"): - # You will see executor invoke and completion events, and then the final WorkflowCompletedEvent. + # You will see executor invoke and completion events as the workflow progresses. print(f"Event: {event}") - if isinstance(event, WorkflowCompletedEvent): - # The WorkflowCompletedEvent contains the final result. - completion_event = event - - # Print the final result after the streaming loop concludes. - if completion_event: - print(f"Workflow completed with result: {completion_event.data}") + if isinstance(event, WorkflowOutputEvent): + print(f"Workflow completed with result: {event.data}") """ Sample Output: @@ -75,8 +70,8 @@ async def main(): Event: ExecutorInvokedEvent(executor_id=upper_case_executor) Event: ExecutorCompletedEvent(executor_id=upper_case_executor) Event: ExecutorInvokedEvent(executor_id=reverse_text_executor) - Event: WorkflowCompletedEvent(data=DLROW OLLEH) Event: ExecutorCompletedEvent(executor_id=reverse_text_executor) + Event: WorkflowOutputEvent(data='DLROW OLLEH', source_executor_id=reverse_text_executor) Workflow completed with result: DLROW OLLEH """ diff --git a/python/samples/getting_started/workflow/control-flow/simple_loop.py b/python/samples/getting_started/workflow/control-flow/simple_loop.py index 6a16a73ed9..8d48274830 100644 --- a/python/samples/getting_started/workflow/control-flow/simple_loop.py +++ b/python/samples/getting_started/workflow/control-flow/simple_loop.py @@ -12,8 +12,8 @@ from agent_framework import ( ExecutorCompletedEvent, Role, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, + WorkflowOutputEvent, handler, ) from agent_framework.azure import AzureChatClient @@ -25,6 +25,7 @@ Sample: Simple Loop (with an Agent Judge) What it does: - Guesser performs a binary search; judge is an agent that returns ABOVE/BELOW/MATCHED. - Demonstrates feedback loops in workflows with agent steps. +- The workflow completes when the correct number is guessed. Prerequisites: - Azure AI/ Azure OpenAI for `AzureChatClient` agent. @@ -55,14 +56,14 @@ class GuessNumberExecutor(Executor): self._upper = bound[1] @handler - async def guess_number(self, feedback: NumberSignal, ctx: WorkflowContext[int]) -> None: + async def guess_number(self, feedback: NumberSignal, ctx: WorkflowContext[int, str]) -> None: """Execute the task by guessing a number.""" if feedback == NumberSignal.INIT: self._guess = (self._lower + self._upper) // 2 await ctx.send_message(self._guess) elif feedback == NumberSignal.MATCHED: # The previous guess was correct. - await ctx.add_event(WorkflowCompletedEvent(f"Guessed the number: {self._guess}")) + await ctx.yield_output(f"Guessed the number: {self._guess}") elif feedback == NumberSignal.ABOVE: # The previous guess was too low. # Update the lower bound to the previous guess. @@ -150,6 +151,8 @@ async def main(): async for event in workflow.run_stream(NumberSignal.INIT): if isinstance(event, ExecutorCompletedEvent) and event.executor_id == guess_number_executor.id: iterations += 1 + elif isinstance(event, WorkflowOutputEvent): + print(f"Final result: {event.data}") print(f"Event: {event}") # This is essentially a binary search, so the number of iterations should be logarithmic. diff --git a/python/samples/getting_started/workflow/control-flow/switch_case_edge_group.py b/python/samples/getting_started/workflow/control-flow/switch_case_edge_group.py index 0833ff5b12..cc2193baeb 100644 --- a/python/samples/getting_started/workflow/control-flow/switch_case_edge_group.py +++ b/python/samples/getting_started/workflow/control-flow/switch_case_edge_group.py @@ -6,6 +6,8 @@ from dataclasses import dataclass from typing import Any, Literal from uuid import uuid4 +from typing_extensions import Never + from agent_framework import ( # Core chat primitives used to form LLM requests AgentExecutor, # Wraps an agent so it can run inside a workflow AgentExecutorRequest, # Message bundle sent to an AgentExecutor @@ -15,7 +17,6 @@ from agent_framework import ( # Core chat primitives used to form LLM requests Default, # Default branch when no cases match Role, WorkflowBuilder, # Fluent builder for assembling the graph - WorkflowCompletedEvent, # Terminal event for successful completion WorkflowContext, # Per-run context and event bus executor, # Decorator to turn a function into a workflow executor ) @@ -36,6 +37,7 @@ Demonstrate deterministic one of N routing with switch-case edges. Show how to: - Validate agent JSON with Pydantic models for robust parsing. - Keep executor responsibilities narrow. Transform model output to a typed DetectionResult, then route based on that type. +- Use ctx.yield_output() to provide workflow results - the workflow completes when idle with no pending work. Prerequisites: - Familiarity with WorkflowBuilder, executors, edges, and events. @@ -124,30 +126,28 @@ async def submit_to_email_assistant(detection: DetectionResult, ctx: WorkflowCon @executor(id="finalize_and_send") -async def finalize_and_send(response: AgentExecutorResponse, ctx: WorkflowContext[None]) -> None: - # Terminal step for the drafting branch. Emit a completion event with the reply. +async def finalize_and_send(response: AgentExecutorResponse, ctx: WorkflowContext[Never, str]) -> None: + # Terminal step for the drafting branch. Yield the email response as output. parsed = EmailResponse.model_validate_json(response.agent_run_response.text) - await ctx.add_event(WorkflowCompletedEvent(f"Email sent: {parsed.response}")) + await ctx.yield_output(f"Email sent: {parsed.response}") @executor(id="handle_spam") -async def handle_spam(detection: DetectionResult, ctx: WorkflowContext[None]) -> None: +async def handle_spam(detection: DetectionResult, ctx: WorkflowContext[Never, str]) -> None: # Spam path terminal. Include the detector's rationale. if detection.spam_decision == "Spam": - await ctx.add_event(WorkflowCompletedEvent(f"Email marked as spam: {detection.reason}")) + await ctx.yield_output(f"Email marked as spam: {detection.reason}") else: raise RuntimeError("This executor should only handle Spam messages.") @executor(id="handle_uncertain") -async def handle_uncertain(detection: DetectionResult, ctx: WorkflowContext[None]) -> None: +async def handle_uncertain(detection: DetectionResult, ctx: WorkflowContext[Never, str]) -> None: # Uncertain path terminal. Surface the original content to aid human review. if detection.spam_decision == "Uncertain": email: Email | None = await ctx.get_shared_state(f"{EMAIL_STATE_PREFIX}{detection.email_id}") - await ctx.add_event( - WorkflowCompletedEvent( - f"Email marked as uncertain: {detection.reason}. Email content: {getattr(email, 'email_content', '')}" - ) + await ctx.yield_output( + f"Email marked as uncertain: {detection.reason}. Email content: {getattr(email, 'email_content', '')}" ) else: raise RuntimeError("This executor should only handle Uncertain messages.") @@ -215,10 +215,12 @@ async def main(): "Let me know if you'd like more details." ) - # Run and print the terminal event for whichever branch completes. - async for event in workflow.run_stream(email): - if isinstance(event, WorkflowCompletedEvent): - print(f"{event}") + # Run and print the outputs from whichever branch completes. + events = await workflow.run(email) + outputs = events.get_outputs() + if outputs: + for output in outputs: + print(f"Workflow output: {output}") if __name__ == "__main__": diff --git a/python/samples/getting_started/workflow/human-in-the-loop/guessing_game_with_human_input.py b/python/samples/getting_started/workflow/human-in-the-loop/guessing_game_with_human_input.py index c9a9880b68..4eec70ebb4 100644 --- a/python/samples/getting_started/workflow/human-in-the-loop/guessing_game_with_human_input.py +++ b/python/samples/getting_started/workflow/human-in-the-loop/guessing_game_with_human_input.py @@ -15,8 +15,8 @@ from agent_framework import ( RequestResponse, # Correlates a human response with the original request Role, # Enum of chat roles (user, assistant, system) WorkflowBuilder, # Fluent builder for assembling the graph - WorkflowCompletedEvent, # Terminal event used to finish the workflow WorkflowContext, # Per run context and event bus + WorkflowOutputEvent, # Event emitted when workflow yields output WorkflowRunState, # Enum of workflow run states WorkflowStatusEvent, # Event emitted on run state changes handler, # Decorator to expose an Executor method as a step @@ -30,8 +30,7 @@ Sample: Human in the loop guessing game An agent guesses a number, then a human guides it with higher, lower, or correct via RequestInfoExecutor. The loop continues until the human confirms -correct, at which point the workflow -completes. +correct, at which point the workflow completes when idle with no pending work. Purpose: Show how to integrate a human step in the middle of an LLM workflow using RequestInfoExecutor and correlated @@ -132,7 +131,7 @@ class TurnManager(Executor): async def on_human_feedback( self, feedback: RequestResponse[HumanFeedbackRequest, str], - ctx: WorkflowContext[AgentExecutorRequest | WorkflowCompletedEvent], + ctx: WorkflowContext[AgentExecutorRequest, str], ) -> None: """Continue the game or finish based on human feedback. @@ -144,7 +143,7 @@ class TurnManager(Executor): last_guess = getattr(feedback.original_request, "guess", None) if reply == "correct": - await ctx.add_event(WorkflowCompletedEvent(f"Guessed correctly: {last_guess}")) + await ctx.yield_output(f"Guessed correctly: {last_guess}") return # Provide feedback to the agent to try again. @@ -195,7 +194,8 @@ async def main() -> None: # Human in the loop run: alternate between invoking the workflow and supplying collected responses. pending_responses: dict[str, str] | None = None - completed: WorkflowCompletedEvent | None = None + completed = False + workflow_output: str | None = None # User guidance printing: # If you want to instruct users up front, print a short banner before the loop. @@ -219,15 +219,16 @@ async def main() -> None: events = [event async for event in stream] pending_responses = None - # Collect human requests and the terminal completion if present. + # Collect human requests, workflow outputs, and check for completion. requests: list[tuple[str, str]] = [] # (request_id, prompt) for event in events: - if isinstance(event, WorkflowCompletedEvent): - completed = event - elif isinstance(event, RequestInfoEvent) and isinstance(event.data, HumanFeedbackRequest): + if isinstance(event, RequestInfoEvent) and isinstance(event.data, HumanFeedbackRequest): # RequestInfoEvent for our HumanFeedbackRequest. requests.append((event.request_id, event.data.prompt)) - # Other events are ignored for brevity. + elif isinstance(event, WorkflowOutputEvent): + # Capture workflow output as they're yielded + workflow_output = str(event.data) + completed = True # In this sample, we finish after one output. # Detect run state transitions for a better developer experience. pending_status = any( @@ -258,9 +259,8 @@ async def main() -> None: responses[req_id] = answer pending_responses = responses - # Show final result. - print(completed) - + # Show final result from workflow output captured during streaming. + print(f"Workflow output: {workflow_output}") """ Sample Output: @@ -272,7 +272,7 @@ async def main() -> None: Enter higher/lower/correct/exit: lower HITL> The agent guessed: 9. Type one of: higher (your number is higher than this guess), lower (your number is lower than this guess), correct, or exit. Enter higher/lower/correct/exit: correct - WorkflowCompletedEvent(data=Guessed correctly: 9) + Workflow output: Guessed correctly: 9 """ # noqa: E501 diff --git a/python/samples/getting_started/workflow/observability/tracing_basics.py b/python/samples/getting_started/workflow/observability/tracing_basics.py index 37d671fa28..fd1daedc9b 100644 --- a/python/samples/getting_started/workflow/observability/tracing_basics.py +++ b/python/samples/getting_started/workflow/observability/tracing_basics.py @@ -2,7 +2,6 @@ import asyncio import os -from typing import Any from agent_framework import Executor, WorkflowBuilder, WorkflowContext, get_logger, handler @@ -57,8 +56,8 @@ class StartExecutor(Executor): class EndExecutor(Executor): @handler # type: ignore[misc] - async def handle_final(self, message: str, ctx: WorkflowContext[Any]) -> None: - # Sink executor. The framework emits WorkflowCompletedEvent automatically after this handler returns. + async def handle_final(self, message: str, ctx: WorkflowContext) -> None: + # Sink executor. The workflow completes when idle with no pending work. print(f"Final result: {message}") diff --git a/python/samples/getting_started/workflow/orchestration/concurrent_agents.py b/python/samples/getting_started/workflow/orchestration/concurrent_agents.py index df3d212b4c..13f167582b 100644 --- a/python/samples/getting_started/workflow/orchestration/concurrent_agents.py +++ b/python/samples/getting_started/workflow/orchestration/concurrent_agents.py @@ -3,7 +3,7 @@ import asyncio from typing import Any -from agent_framework import ChatMessage, ConcurrentBuilder, WorkflowCompletedEvent +from agent_framework import ChatMessage, ConcurrentBuilder from agent_framework.azure import AzureChatClient from azure.identity import AzureCliCredential @@ -12,13 +12,13 @@ Sample: Concurrent fan-out/fan-in (agent-only API) with default aggregator Build a high-level concurrent workflow using ConcurrentBuilder and three domain agents. The default dispatcher fans out the same user prompt to all agents in parallel. -The default aggregator fans in their results and emits a WorkflowCompletedEvent whose -data is a list[ChatMessage] representing the concatenated conversations from all agents. +The default aggregator fans in their results and yields output containing +a list[ChatMessage] representing the concatenated conversations from all agents. Demonstrates: - Minimal wiring with ConcurrentBuilder().participants([...]).build() - Fan-out to multiple agents, fan-in aggregation of final ChatMessages -- Streaming of AgentRunEvent for simple progress visibility +- Workflow completion when idle with no pending work Prerequisites: - Azure OpenAI access configured for AzureChatClient (use az login + env vars) @@ -58,18 +58,17 @@ async def main() -> None: # Participants are either Agents (type of AgentProtocol) or Executors workflow = ConcurrentBuilder().participants([researcher, marketer, legal]).build() - # 3) Run with a single prompt, stream progress, and pretty-print the final combined messages - completion: WorkflowCompletedEvent | None = None - async for event in workflow.run_stream("We are launching a new budget-friendly electric bike for urban commuters."): - if isinstance(event, WorkflowCompletedEvent): - completion = event + # 3) Run with a single prompt and pretty-print the final combined messages + events = await workflow.run("We are launching a new budget-friendly electric bike for urban commuters.") + outputs = events.get_outputs() - if completion: + if outputs: print("===== Final Aggregated Conversation (messages) =====") - messages: list[ChatMessage] | Any = completion.data - for i, msg in enumerate(messages, start=1): - name = msg.author_name if msg.author_name else "user" - print(f"{'-' * 60}\n\n{i:02d} [{name}]:\n{msg.text}") + for output in outputs: + messages: list[ChatMessage] | Any = output + for i, msg in enumerate(messages, start=1): + name = msg.author_name if msg.author_name else "user" + print(f"{'-' * 60}\n\n{i:02d} [{name}]:\n{msg.text}") """ Sample Output: diff --git a/python/samples/getting_started/workflow/orchestration/concurrent_custom_agent_executors.py b/python/samples/getting_started/workflow/orchestration/concurrent_custom_agent_executors.py index e4883e6789..6a25a8a922 100644 --- a/python/samples/getting_started/workflow/orchestration/concurrent_custom_agent_executors.py +++ b/python/samples/getting_started/workflow/orchestration/concurrent_custom_agent_executors.py @@ -10,7 +10,6 @@ from agent_framework import ( ChatMessage, ConcurrentBuilder, Executor, - WorkflowCompletedEvent, WorkflowContext, handler, ) @@ -30,6 +29,7 @@ Demonstrates: - A @handler that converts AgentExecutorRequest -> AgentExecutorResponse - ConcurrentBuilder().participants([...]) to build fan-out/fan-in - Default aggregator returning list[ChatMessage] (one user + one assistant per agent) +- Workflow completion when all participants become idle Prerequisites: - Azure OpenAI configured for AzureChatClient (az login + required env vars) @@ -105,14 +105,12 @@ async def main() -> None: workflow = ConcurrentBuilder().participants([researcher, marketer, legal]).build() - completion: WorkflowCompletedEvent | None = None - async for event in workflow.run_stream("We are launching a new budget-friendly electric bike for urban commuters."): - if isinstance(event, WorkflowCompletedEvent): - completion = event + events = await workflow.run("We are launching a new budget-friendly electric bike for urban commuters.") + outputs = events.get_outputs() - if completion: + if outputs: print("===== Final Aggregated Conversation (messages) =====") - messages: list[ChatMessage] | Any = completion.data + messages: list[ChatMessage] | Any = outputs[0] # Get the first (and typically only) output for i, msg in enumerate(messages, start=1): name = msg.author_name if msg.author_name else "user" print(f"{'-' * 60}\n\n{i:02d} [{name}]:\n{msg.text}") diff --git a/python/samples/getting_started/workflow/orchestration/concurrent_custom_aggregator.py b/python/samples/getting_started/workflow/orchestration/concurrent_custom_aggregator.py index 02fa8c892a..15fe455464 100644 --- a/python/samples/getting_started/workflow/orchestration/concurrent_custom_aggregator.py +++ b/python/samples/getting_started/workflow/orchestration/concurrent_custom_aggregator.py @@ -3,7 +3,7 @@ import asyncio from typing import Any -from agent_framework import ChatMessage, ConcurrentBuilder, Role, WorkflowCompletedEvent +from agent_framework import ChatMessage, ConcurrentBuilder, Role from agent_framework.azure import AzureChatClient from azure.identity import AzureCliCredential @@ -14,12 +14,13 @@ Build a concurrent workflow with ConcurrentBuilder that fans out one prompt to multiple domain agents and fans in their responses. Override the default aggregator with a custom async callback that uses AzureChatClient.get_response() to synthesize a concise, consolidated summary from the experts' outputs. +The workflow completes when all participants become idle. Demonstrates: - ConcurrentBuilder().participants([...]).with_custom_aggregator(callback) - Fan-out to agents and fan-in at an aggregator - Aggregation implemented via an LLM call (chat_client.get_response) -- WorkflowCompletedEvent carrying the synthesized summary string +- Workflow output yielded with the synthesized summary string Prerequisites: - Azure OpenAI configured for AzureChatClient (az login + required env vars) @@ -82,20 +83,18 @@ async def main() -> None: # Each participant becomes a parallel branch (fan-out) from an internal dispatcher. # - with_aggregator(...) overrides the default aggregator: # • Default aggregator -> returns list[ChatMessage] (one user + one assistant per agent) - # • Custom callback -> return value becomes WorkflowCompletedEvent.data (string here) + # • Custom callback -> return value becomes workflow output (string here) # The callback can be sync or async; it receives list[AgentExecutorResponse]. workflow = ( ConcurrentBuilder().participants([researcher, marketer, legal]).with_aggregator(summarize_results).build() ) - completion: WorkflowCompletedEvent | None = None - async for event in workflow.run_stream("We are launching a new budget-friendly electric bike for urban commuters."): - if isinstance(event, WorkflowCompletedEvent): - completion = event + events = await workflow.run("We are launching a new budget-friendly electric bike for urban commuters.") + outputs = events.get_outputs() - if completion: + if outputs: print("===== Final Consolidated Output =====") - print(completion.data) + print(outputs[0]) # Get the first (and typically only) output """ Sample Output: diff --git a/python/samples/getting_started/workflow/orchestration/magentic.py b/python/samples/getting_started/workflow/orchestration/magentic.py index 52be43a418..14e00b0418 100644 --- a/python/samples/getting_started/workflow/orchestration/magentic.py +++ b/python/samples/getting_started/workflow/orchestration/magentic.py @@ -13,7 +13,7 @@ from agent_framework import ( MagenticCallbackMode, MagenticFinalResultEvent, MagenticOrchestratorMessageEvent, - WorkflowCompletedEvent, + WorkflowOutputEvent, ) from agent_framework.openai import OpenAIChatClient, OpenAIResponsesClient @@ -38,7 +38,7 @@ The workflow is configured with: When run, the script builds the workflow, submits a task about estimating the energy efficiency and CO2 emissions of several ML models, streams intermediate -events, and prints the final answer. +events, and prints the final answer. The workflow completes when idle. Prerequisites: - OpenAI credentials configured for `OpenAIChatClient` and `OpenAIResponsesClient`. @@ -132,17 +132,14 @@ async def main() -> None: print("\nStarting workflow execution...") try: - completion_event = None + output: str | None = None async for event in workflow.run_stream(task): - print(f"Event: {event}") + print(event) + if isinstance(event, WorkflowOutputEvent): + output = str(event.data) - if isinstance(event, WorkflowCompletedEvent): - completion_event = event - - if completion_event is not None: - data = getattr(completion_event, "data", None) - preview = getattr(data, "text", None) or (str(data) if data is not None else "") - print(f"Workflow completed with result:\n\n{preview}") + if output is not None: + print(f"Workflow completed with result:\n\n{output}") except Exception as e: print(f"Workflow execution failed: {e}") diff --git a/python/samples/getting_started/workflow/orchestration/magentic_human_plan_update.py b/python/samples/getting_started/workflow/orchestration/magentic_human_plan_update.py index aec0aaa334..fb45f94855 100644 --- a/python/samples/getting_started/workflow/orchestration/magentic_human_plan_update.py +++ b/python/samples/getting_started/workflow/orchestration/magentic_human_plan_update.py @@ -18,7 +18,7 @@ from agent_framework import ( MagenticPlanReviewReply, MagenticPlanReviewRequest, RequestInfoEvent, - WorkflowCompletedEvent, + WorkflowOutputEvent, ) from agent_framework.openai import OpenAIChatClient, OpenAIResponsesClient @@ -41,6 +41,7 @@ Key behaviors demonstrated: replies with PlanReviewReply (here we auto-approve, but you can edit/collect input) - Callbacks: on_agent_stream (incremental chunks), on_agent_response (final messages), on_result (final answer), and on_exception +- Workflow completion when idle Prerequisites: - OpenAI credentials configured for `OpenAIChatClient` and `OpenAIResponsesClient`. @@ -73,6 +74,9 @@ async def main() -> None: print(f"Exception occurred: {exception}") logger.exception("Workflow exception", exc_info=exception) + last_stream_agent_id: str | None = None + stream_line_open: bool = False + # Unified callback async def on_event(event: MagenticCallbackEvent) -> None: nonlocal last_stream_agent_id, stream_line_open @@ -105,9 +109,6 @@ async def main() -> None: print("\nBuilding Magentic Workflow...") - last_stream_agent_id: str | None = None - stream_line_open: bool = False - workflow = ( MagenticBuilder() .participants(researcher=researcher_agent, coder=coder_agent) @@ -136,51 +137,61 @@ async def main() -> None: print("\nStarting workflow execution...") try: - completion_event: WorkflowCompletedEvent | None = None pending_request: RequestInfoEvent | None = None + pending_responses: dict[str, MagenticPlanReviewReply] | None = None + completed = False + workflow_output: str | None = None - while True: - # Phase 1: run until either completion or a HIL request - if pending_request is None: - async for event in workflow.run_stream(task): - print(f"Event: {event}") + while not completed: + # Use streaming for both initial run and response sending + if pending_responses is not None: + stream = workflow.send_responses_streaming(pending_responses) + else: + stream = workflow.run_stream(task) - if isinstance(event, WorkflowCompletedEvent): - completion_event = event + # Collect events from the stream + events = [event async for event in stream] + pending_responses = None - if isinstance(event, RequestInfoEvent) and event.request_type is MagenticPlanReviewRequest: - pending_request = event - review_req = cast(MagenticPlanReviewRequest, event.data) - if review_req.plan_text: - print(f"\n=== PLAN REVIEW REQUEST ===\n{review_req.plan_text}\n") + # Process events to find request info events, outputs, and completion status + for event in events: + if isinstance(event, RequestInfoEvent) and event.request_type is MagenticPlanReviewRequest: + pending_request = event + review_req = cast(MagenticPlanReviewRequest, event.data) + if review_req.plan_text: + print(f"\n=== PLAN REVIEW REQUEST ===\n{review_req.plan_text}\n") + elif isinstance(event, WorkflowOutputEvent): + # Capture workflow output during streaming + workflow_output = str(event.data) + completed = True - # Break if completed - if completion_event is not None: - data = getattr(completion_event, "data", None) - preview = getattr(data, "text", None) or (str(data) if data is not None else "") - print(f"Workflow completed with result:\n\n{preview}") - - # Phase 2: respond to the pending plan review (HIL) request + # Handle pending plan review request if pending_request is not None: - # For demo purposes we approve as-is. Replace this with UI input - # to collect a human decision/comments/edited plan. - reply = MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.APPROVE) + # Get human input for plan review decision + print("Plan review options:") + print("1. approve - Approve the plan as-is") + print("2. revise - Request revision of the plan") + print("3. exit - Exit the workflow") - async for event in workflow.send_responses_streaming({pending_request.request_id: reply}): - print(f"Event: {event}") + while True: + choice = input("Enter your choice (approve/revise/exit): ").strip().lower() # noqa: ASYNC250 + if choice in ["approve", "1"]: + reply = MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.APPROVE) + break + if choice in ["revise", "2"]: + reply = MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.REVISE) + break + if choice in ["exit", "3"]: + print("Exiting workflow...") + return + print("Invalid choice. Please enter 'approve', 'revise', or 'exit'.") - if isinstance(event, WorkflowCompletedEvent): - completion_event = event + pending_responses = {pending_request.request_id: reply} + pending_request = None - if isinstance(event, RequestInfoEvent) and event.request_type is MagenticPlanReviewRequest: - # Another review cycle requested; keep pending - pending_request = event - review_req = cast(MagenticPlanReviewRequest, event.data) - if review_req.plan_text: - print(f"\n=== PLAN REVIEW REQUEST ===\n{review_req.plan_text}\n") - else: - # Clear pending if no immediate new request - pending_request = None + # Show final result from captured workflow output + if workflow_output: + print(f"Workflow completed with result:\n\n{workflow_output}") except Exception as e: print(f"Workflow execution failed: {e}") diff --git a/python/samples/getting_started/workflow/orchestration/sequential_agents.py b/python/samples/getting_started/workflow/orchestration/sequential_agents.py index dfc7f154a5..a444514b99 100644 --- a/python/samples/getting_started/workflow/orchestration/sequential_agents.py +++ b/python/samples/getting_started/workflow/orchestration/sequential_agents.py @@ -1,9 +1,9 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio -from typing import Any +from typing import cast -from agent_framework import ChatMessage, Role, SequentialBuilder, WorkflowCompletedEvent +from agent_framework import ChatMessage, Role, SequentialBuilder, WorkflowOutputEvent from agent_framework.azure import AzureChatClient from azure.identity import AzureCliCredential @@ -12,8 +12,8 @@ Sample: Sequential workflow (agent-focused API) with shared conversation context Build a high-level sequential workflow using SequentialBuilder and two domain agents. The shared conversation (list[ChatMessage]) flows through each participant. Each agent -appends its assistant message to the context. The final WorkflowCompletedEvent includes -the final conversation list. +appends its assistant message to the context. The workflow outputs the final conversation +list when complete. Note on internal adapters: - Sequential orchestration includes small adapter nodes for input normalization @@ -44,16 +44,15 @@ async def main() -> None: # 2) Build sequential workflow: writer -> reviewer workflow = SequentialBuilder().participants([writer, reviewer]).build() - # 3) Run and print final conversation - completion: WorkflowCompletedEvent | None = None + # 3) Run and collect outputs + outputs: list[list[ChatMessage]] = [] async for event in workflow.run_stream("Write a tagline for a budget-friendly eBike."): - if isinstance(event, WorkflowCompletedEvent): - completion = event + if isinstance(event, WorkflowOutputEvent): + outputs.append(cast(list[ChatMessage], event.data)) - if completion: + if outputs: print("===== Final Conversation =====") - messages: list[ChatMessage] | Any = completion.data - for i, msg in enumerate(messages, start=1): + for i, msg in enumerate(outputs[-1], start=1): name = msg.author_name or ("assistant" if msg.role == Role.ASSISTANT else "user") print(f"{'-' * 60}\n{i:02d} [{name}]\n{msg.text}") diff --git a/python/samples/getting_started/workflow/orchestration/sequential_custom_executors.py b/python/samples/getting_started/workflow/orchestration/sequential_custom_executors.py index 889833f634..452234bb42 100644 --- a/python/samples/getting_started/workflow/orchestration/sequential_custom_executors.py +++ b/python/samples/getting_started/workflow/orchestration/sequential_custom_executors.py @@ -3,12 +3,13 @@ import asyncio from typing import Any +from typing_extensions import Never + from agent_framework import ( ChatMessage, Executor, Role, SequentialBuilder, - WorkflowCompletedEvent, WorkflowContext, handler, ) @@ -20,8 +21,8 @@ Sample: Sequential workflow mixing agents and a custom summarizer executor This demonstrates how SequentialBuilder chains participants with a shared conversation context (list[ChatMessage]). An agent produces content; a custom -executor appends a compact summary to the conversation. The final WorkflowCompletedEvent -contains the complete conversation. +executor appends a compact summary to the conversation. The workflow completes +when idle, and the final output contains the complete conversation. Custom executor contract: - Provide at least one @handler accepting list[ChatMessage] and a WorkflowContext[list[ChatMessage]] @@ -42,11 +43,12 @@ class Summarizer(Executor): """Simple summarizer: consumes full conversation and appends an assistant summary.""" @handler - async def summarize(self, conversation: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None: + async def summarize(self, conversation: list[ChatMessage], ctx: WorkflowContext[Never, list[ChatMessage]]) -> None: users = sum(1 for m in conversation if m.role == Role.USER) assistants = sum(1 for m in conversation if m.role == Role.ASSISTANT) summary = ChatMessage(role=Role.ASSISTANT, text=f"Summary -> users:{users} assistants:{assistants}") - await ctx.send_message(list(conversation) + [summary]) + final_conversation = list(conversation) + [summary] + await ctx.yield_output(final_conversation) async def main() -> None: @@ -62,14 +64,12 @@ async def main() -> None: workflow = SequentialBuilder().participants([content, summarizer]).build() # 3) Run and print final conversation - completion: WorkflowCompletedEvent | None = None - async for event in workflow.run_stream("Explain the benefits of budget eBikes for commuters."): - if isinstance(event, WorkflowCompletedEvent): - completion = event + events = await workflow.run("Explain the benefits of budget eBikes for commuters.") + outputs = events.get_outputs() - if completion: + if outputs: print("===== Final Conversation =====") - messages: list[ChatMessage] | Any = completion.data + messages: list[ChatMessage] | Any = outputs[0] for i, msg in enumerate(messages, start=1): name = msg.author_name or ("assistant" if msg.role == Role.ASSISTANT else "user") print(f"{'-' * 60}\n{i:02d} [{name}]\n{msg.text}") diff --git a/python/samples/getting_started/workflow/parallelism/fan_out_fan_in_edges.py b/python/samples/getting_started/workflow/parallelism/fan_out_fan_in_edges.py index c2c95392c4..ff1cc1a73c 100644 --- a/python/samples/getting_started/workflow/parallelism/fan_out_fan_in_edges.py +++ b/python/samples/getting_started/workflow/parallelism/fan_out_fan_in_edges.py @@ -2,7 +2,8 @@ import asyncio from dataclasses import dataclass -from typing import Any + +from typing_extensions import Never from agent_framework import ( # Core chat primitives to build LLM requests AgentExecutor, # Wraps an LLM agent for use inside a workflow @@ -13,8 +14,8 @@ from agent_framework import ( # Core chat primitives to build LLM requests Executor, # Base class for custom Python executors Role, # Enum of chat roles (user, assistant, system) WorkflowBuilder, # Fluent builder for wiring the workflow graph - WorkflowCompletedEvent, # Terminal event carrying the final result WorkflowContext, # Per run context and event bus + WorkflowOutputEvent, # Event emitted when workflow yields output handler, # Decorator to mark an Executor method as invokable ) from agent_framework.azure import AzureChatClient # Client wrapper for Azure OpenAI chat models @@ -75,7 +76,7 @@ class AggregateInsights(Executor): self._expert_ids = expert_ids @handler - async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> None: + async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, str]) -> None: # Map responses to text by executor id for a simple, predictable demo. by_id: dict[str, str] = {} for r in results: @@ -101,7 +102,7 @@ class AggregateInsights(Executor): f"Legal/Compliance Notes:\n{aggregated.legal}\n" ) - await ctx.add_event(WorkflowCompletedEvent(data=consolidated)) + await ctx.yield_output(consolidated) async def main() -> None: @@ -151,17 +152,13 @@ async def main() -> None: ) # 3) Run with a single prompt and print progress plus the final consolidated output - completion: WorkflowCompletedEvent | None = None async for event in workflow.run_stream("We are launching a new budget-friendly electric bike for urban commuters."): if isinstance(event, AgentRunEvent): # Show which agent ran and what step completed for lightweight observability. print(event) - if isinstance(event, WorkflowCompletedEvent): - completion = event - - if completion: - print("===== Final Aggregated Output =====") - print(completion.data) + elif isinstance(event, WorkflowOutputEvent): + print("===== Final Aggregated Output =====") + print(event.data) if __name__ == "__main__": diff --git a/python/samples/getting_started/workflow/parallelism/map_reduce_and_visualization.py b/python/samples/getting_started/workflow/parallelism/map_reduce_and_visualization.py index 6098da7303..d82835897f 100644 --- a/python/samples/getting_started/workflow/parallelism/map_reduce_and_visualization.py +++ b/python/samples/getting_started/workflow/parallelism/map_reduce_and_visualization.py @@ -5,14 +5,15 @@ import asyncio import os from collections import defaultdict from dataclasses import dataclass -from typing import Any import aiofiles +from typing_extensions import Never + from agent_framework import ( Executor, # Base class for custom workflow steps - WorkflowBuilder, # Fluent graph builder for executors and edges - WorkflowCompletedEvent, # Terminal event that carries final output + WorkflowBuilder, # Fluent builder for executors and edges WorkflowContext, # Per run context with shared state and messaging + WorkflowOutputEvent, # Event emitted when workflow yields output WorkflowViz, # Utility to visualize a workflow graph handler, # Decorator to expose an Executor method as a step ) @@ -246,12 +247,12 @@ class Reduce(Executor): class CompletionExecutor(Executor): - """Joins all reducer outputs and emits the final completion event.""" + """Joins all reducer outputs and yields the final output.""" @handler - async def complete(self, data: list[ReduceCompleted], ctx: WorkflowContext[Any]) -> None: - """Collect reducer output file paths and publish a terminal event.""" - await ctx.add_event(WorkflowCompletedEvent(data=[result.file_path for result in data])) + async def complete(self, data: list[ReduceCompleted], ctx: WorkflowContext[Never, list[str]]) -> None: + """Collect reducer output file paths and yield final output.""" + await ctx.yield_output([result.file_path for result in data]) async def main(): @@ -303,14 +304,10 @@ async def main(): raw_text = await f.read() # Step 4: Run the workflow with the raw text as input. - completion_event = None async for event in workflow.run_stream(raw_text): print(f"Event: {event}") - if isinstance(event, WorkflowCompletedEvent): - completion_event = event - - if completion_event: - print(f"Completion Event: {completion_event}") + if isinstance(event, WorkflowOutputEvent): + print(f"Final Output: {event.data}") if __name__ == "__main__": diff --git a/python/samples/getting_started/workflow/state-management/shared_states_with_agents.py b/python/samples/getting_started/workflow/state-management/shared_states_with_agents.py index c5584bc0fc..b329049c27 100644 --- a/python/samples/getting_started/workflow/state-management/shared_states_with_agents.py +++ b/python/samples/getting_started/workflow/state-management/shared_states_with_agents.py @@ -6,13 +6,14 @@ from dataclasses import dataclass from typing import Any from uuid import uuid4 +from typing_extensions import Never + from agent_framework import ( AgentExecutorRequest, AgentExecutorResponse, ChatMessage, Role, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, executor, ) @@ -31,7 +32,7 @@ Show how to: - Use shared state to decouple large payloads from messages and pass around lightweight references. - Enforce structured agent outputs with Pydantic models via response_format for robust parsing. - Route using conditional edges based on a typed intermediate DetectionResult. -- Compose agent backed executors with function style executors and print a terminal WorkflowCompletedEvent. +- Compose agent backed executors with function style executors and yield the final output when the workflow completes. Prerequisites: - Azure OpenAI configured for AzureChatClient with required environment variables. @@ -139,17 +140,17 @@ async def submit_to_email_assistant(detection: DetectionResult, ctx: WorkflowCon @executor(id="finalize_and_send") -async def finalize_and_send(response: AgentExecutorResponse, ctx: WorkflowContext[None]) -> None: - """Validate the drafted reply and complete the workflow with a terminal event.""" +async def finalize_and_send(response: AgentExecutorResponse, ctx: WorkflowContext[Never, str]) -> None: + """Validate the drafted reply and yield the final output.""" parsed = EmailResponse.model_validate_json(response.agent_run_response.text) - await ctx.add_event(WorkflowCompletedEvent(f"Email sent: {parsed.response}")) + await ctx.yield_output(f"Email sent: {parsed.response}") @executor(id="handle_spam") -async def handle_spam(detection: DetectionResult, ctx: WorkflowContext[None]) -> None: - """Emit a completion event describing why the email was marked as spam.""" +async def handle_spam(detection: DetectionResult, ctx: WorkflowContext[Never, str]) -> None: + """Yield output describing why the email was marked as spam.""" if detection.is_spam: - await ctx.add_event(WorkflowCompletedEvent(f"Email marked as spam: {detection.reason}")) + await ctx.yield_output(f"Email marked as spam: {detection.reason}") else: raise RuntimeError("This executor should only handle spam messages.") @@ -206,18 +207,20 @@ async def main() -> None: print("Unable to find resource file, using default text.") email = "You are a WINNER! Click here for a free lottery offer!!!" - # Run and print the terminal result. Streaming surfaces intermediate execution events as well. - async for event in workflow.run_stream(email): - if isinstance(event, WorkflowCompletedEvent): - print(event) + # Run and print the final result. Streaming surfaces intermediate execution events as well. + events = await workflow.run(email) + outputs = events.get_outputs() + + if outputs: + print(f"Final result: {outputs[0]}") """ Sample Output: - WorkflowCompletedEvent(data=Email marked as spam: This email exhibits several common spam and scam characteristics: + Final result: Email marked as spam: This email exhibits several common spam and scam characteristics: unrealistic claims of large cash winnings, urgent time pressure, requests for sensitive personal and financial information, and a demand for a processing fee. The sender impersonates a generic lottery commission, and the - message contains a suspicious link. All these are typical of phishing and lottery scam emails.) + message contains a suspicious link. All these are typical of phishing and lottery scam emails. """ diff --git a/python/samples/getting_started/workflow/visualization/concurrent_with_visualization.py b/python/samples/getting_started/workflow/visualization/concurrent_with_visualization.py index 5132299f43..a07b390051 100644 --- a/python/samples/getting_started/workflow/visualization/concurrent_with_visualization.py +++ b/python/samples/getting_started/workflow/visualization/concurrent_with_visualization.py @@ -2,7 +2,8 @@ import asyncio from dataclasses import dataclass -from typing import Any + +from typing_extensions import Never from agent_framework import ( AgentExecutor, @@ -13,8 +14,8 @@ from agent_framework import ( Executor, Role, WorkflowBuilder, - WorkflowCompletedEvent, WorkflowContext, + WorkflowOutputEvent, WorkflowViz, handler, ) @@ -71,7 +72,7 @@ class AggregateInsights(Executor): self._expert_ids = expert_ids @handler - async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> None: + async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, str]) -> None: # Map responses to text by executor id for a simple, predictable demo. by_id: dict[str, str] = {} for r in results: @@ -97,7 +98,7 @@ class AggregateInsights(Executor): f"Legal/Compliance Notes:\n{aggregated.legal}\n" ) - await ctx.add_event(WorkflowCompletedEvent(data=consolidated)) + await ctx.yield_output(consolidated) async def main() -> None: @@ -165,17 +166,13 @@ async def main() -> None: print("Tip: Install 'viz' extra to export workflow visualization: pip install agent-framework[viz]") # 3) Run with a single prompt - completion: WorkflowCompletedEvent | None = None async for event in workflow.run_stream("We are launching a new budget-friendly electric bike for urban commuters."): if isinstance(event, AgentRunEvent): # Show which agent ran and what step completed. print(event) - if isinstance(event, WorkflowCompletedEvent): - completion = event - - if completion: - print("===== Final Aggregated Output =====") - print(completion.data) + elif isinstance(event, WorkflowOutputEvent): + print("===== Final Aggregated Output =====") + print(event.data) if __name__ == "__main__":