Python: Use generic for WorkflowContext and use its type parameters to indicate executor's output types (#444)

* Use generic for WorkflowContext and use its type parameters to indicate executor's output types

* Update

* Fix type errors and add in-line comments

* fix test

* type

* Fix executor type issues

python/samples/getting_started/workflow/step_00_foundation_patterns.py

@@ -21,8 +21,8 @@ The samples here use numbers and simple arithmetic operations to demonstrate the
 class AddOneExecutor(Executor):
     """An executor that processes a number by adding one."""
 
-    @handler(output_types=[int])
-    async def add_one(self, number: int, ctx: WorkflowContext) -> None:
+    @handler
+    async def add_one(self, number: int, ctx: WorkflowContext[int]) -> None:
         """Execute the task by adding one to the input number."""
         result = number + 1
 
@@ -35,8 +35,8 @@ class AddOneExecutor(Executor):
 class MultiplyByTwoExecutor(Executor):
     """An executor that processes a number by multiplying it by two."""
 
-    @handler(output_types=[int])
-    async def multiply_by_two(self, number: int, ctx: WorkflowContext) -> None:
+    @handler
+    async def multiply_by_two(self, number: int, ctx: WorkflowContext[int]) -> None:
         """Execute the task by multiplying the input number by two."""
         result = number * 2
 
@@ -49,8 +49,8 @@ class MultiplyByTwoExecutor(Executor):
 class DivideByTwoExecutor(Executor):
     """An executor that processes a number by dividing it by two."""
 
-    @handler(output_types=[float])
-    async def divide_by_two(self, number: int, ctx: WorkflowContext) -> None:
+    @handler
+    async def divide_by_two(self, number: int, ctx: WorkflowContext[float]) -> None:
         """Execute the task by dividing the input number by two."""
         result = number / 2
 
@@ -64,7 +64,7 @@ class AggregateResultExecutor(Executor):
     """An executor that receives results and prints them."""
 
     @handler
-    async def aggregate_results(self, results: Any, ctx: WorkflowContext) -> None:
+    async def aggregate_results(self, results: Any, ctx: WorkflowContext[None]) -> None:
         """Print whatever results are received."""
         print("Aggregating results:", results)
 

python/samples/getting_started/workflow/step_01_simple_workflow_sequential.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+from typing import Any
 
 from agent_framework.workflow import (
     Executor,
@@ -20,8 +21,8 @@ input string to uppercase, and the second executor reverses the string.
 class UpperCaseExecutor(Executor):
     """An executor that converts text to uppercase."""
 
-    @handler(output_types=[str])
-    async def to_upper_case(self, text: str, ctx: WorkflowContext) -> None:
+    @handler
+    async def to_upper_case(self, text: str, ctx: WorkflowContext[str]) -> None:
         """Execute the task by converting the input string to uppercase."""
         result = text.upper()
 
@@ -33,7 +34,7 @@ class ReverseTextExecutor(Executor):
     """An executor that reverses text."""
 
     @handler
-    async def reverse_text(self, text: str, ctx: WorkflowContext) -> None:
+    async def reverse_text(self, text: str, ctx: WorkflowContext[Any]) -> None:
         """Execute the task by reversing the input string."""
         result = text[::-1]
 

python/samples/getting_started/workflow/step_01a_simple_workflow_sequential_non_stream.py

@@ -20,8 +20,8 @@ input string to uppercase, and the second executor reverses the string.
 class UpperCaseExecutor(Executor):
     """An executor that converts text to uppercase."""
 
-    @handler(output_types=[str])
-    async def to_upper_case(self, text: str, ctx: WorkflowContext) -> None:
+    @handler
+    async def to_upper_case(self, text: str, ctx: WorkflowContext[str]) -> None:
         """Execute the task by converting the input string to uppercase."""
         result = text.upper()
 
@@ -33,7 +33,7 @@ class ReverseTextExecutor(Executor):
     """An executor that reverses text."""
 
     @handler
-    async def reverse_text(self, text: str, ctx: WorkflowContext) -> None:
+    async def reverse_text(self, text: str, ctx: WorkflowContext[str]) -> None:
         """Execute the task by reversing the input string."""
         result = text[::-1]
 

python/samples/getting_started/workflow/step_02_simple_workflow_condition.py

@@ -37,8 +37,8 @@ class SpamDetector(Executor):
         super().__init__(id=id)
         self._spam_keywords = spam_keywords
 
-    @handler(output_types=[SpamDetectorResponse])
-    async def handle_email(self, email: str, ctx: WorkflowContext) -> None:
+    @handler
+    async def handle_email(self, email: str, ctx: WorkflowContext[SpamDetectorResponse]) -> None:
         """Determine if the input string is spam."""
         result = any(keyword in email.lower() for keyword in self._spam_keywords)
 
@@ -52,7 +52,7 @@ class SendResponse(Executor):
     async def handle_detector_response(
         self,
         spam_detector_response: SpamDetectorResponse,
-        ctx: WorkflowContext,
+        ctx: WorkflowContext[None],
     ) -> None:
         """Respond with a message based on whether the input is spam."""
         if spam_detector_response.is_spam:
@@ -72,7 +72,7 @@ class RemoveSpam(Executor):
     async def handle_detector_response(
         self,
         spam_detector_response: SpamDetectorResponse,
-        ctx: WorkflowContext,
+        ctx: WorkflowContext[None],
     ) -> None:
         """Remove the spam message."""
         if spam_detector_response.is_spam is False:

python/samples/getting_started/workflow/step_03_simple_workflow_loop.py

@@ -41,8 +41,8 @@ class GuessNumberExecutor(Executor):
         self._lower = bound[0]
         self._upper = bound[1]
 
-    @handler(output_types=[int])
-    async def guess_number(self, feedback: NumberSignal, ctx: WorkflowContext) -> None:
+    @handler
+    async def guess_number(self, feedback: NumberSignal, ctx: WorkflowContext[int]) -> None:
         """Execute the task by guessing a number."""
         if feedback == NumberSignal.INIT:
             self._guess = (self._lower + self._upper) // 2
@@ -74,8 +74,8 @@ class JudgeExecutor(Executor):
         super().__init__(id=id)
         self._target = target
 
-    @handler(output_types=[NumberSignal])
-    async def judge(self, number: int, ctx: WorkflowContext) -> None:
+    @handler
+    async def judge(self, number: int, ctx: WorkflowContext[NumberSignal]) -> None:
         """Judge the guessed number."""
         if number == self._target:
             result = NumberSignal.MATCHED

python/samples/getting_started/workflow/step_04_simple_group_chat.py

@@ -33,8 +33,8 @@ class RoundRobinGroupChatManager(Executor):
         self._max_round = max_round
         self._current_round = 0
 
-    @handler(output_types=[AgentExecutorRequest])
-    async def start(self, task: str, ctx: WorkflowContext) -> None:
+    @handler
+    async def start(self, task: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
         """Execute the task by sending messages to the next executor in the round-robin sequence."""
         initial_message = ChatMessage(ChatRole.USER, text=task)
 
@@ -53,8 +53,10 @@ class RoundRobinGroupChatManager(Executor):
             target_id=self._get_next_member(),
         )
 
-    @handler(output_types=[AgentExecutorRequest])
-    async def handle_agent_response(self, response: AgentExecutorResponse, ctx: WorkflowContext) -> None:
+    @handler
+    async def handle_agent_response(
+        self, response: AgentExecutorResponse, ctx: WorkflowContext[AgentExecutorRequest]
+    ) -> None:
         """Execute the task by sending messages to the next executor in the round-robin sequence."""
         # Send the response to the other members
         await asyncio.gather(*[

python/samples/getting_started/workflow/step_05_simple_group_chat_with_hil.py

@@ -35,13 +35,19 @@ class CriticGroupChatManager(Executor):
         self._current_round = 0
         self._chat_history: list[ChatMessage] = []
 
-    @handler(output_types=[AgentExecutorRequest])
-    async def start(self, task: str, ctx: WorkflowContext) -> None:
+    @handler
+    async def start(self, task: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
         """Handler that starts the group chat with an initial task."""
         initial_message = ChatMessage(ChatRole.USER, text=task)
 
         # Send the initial message to the members
-        await self._broadcast_message([initial_message], ctx)
+        await asyncio.gather(*[
+            ctx.send_message(
+                AgentExecutorRequest(messages=[initial_message], should_respond=False),
+                target_id=member_id,
+            )
+            for member_id in self._members
+        ])
 
         # Invoke the first member to start the round-robin chat
         await ctx.send_message(
@@ -52,14 +58,25 @@ class CriticGroupChatManager(Executor):
         # Update the cache with the initial message
         self._chat_history.append(initial_message)
 
-    @handler(output_types=[AgentExecutorRequest, RequestInfoMessage])
-    async def handle_agent_response(self, response: AgentExecutorResponse, ctx: WorkflowContext) -> None:
+    @handler
+    async def handle_agent_response(
+        self,
+        response: AgentExecutorResponse,
+        ctx: WorkflowContext[RequestInfoMessage | AgentExecutorRequest],
+    ) -> None:
         """Handler that processes the response from the agent."""
         # Update the chat history with the response
         self._chat_history.extend(response.agent_run_response.messages)
 
         # Send the response to the other members
-        await self._broadcast_message(response.agent_run_response.messages, ctx, exclude_id=response.executor_id)
+        await asyncio.gather(*[
+            ctx.send_message(
+                AgentExecutorRequest(messages=response.agent_run_response.messages, should_respond=False),
+                target_id=member_id,
+            )
+            for member_id in self._members
+            if member_id != response.executor_id
+        ])
 
         # Check if we need to request additional information
         if self._should_request_info():
@@ -75,14 +92,22 @@ class CriticGroupChatManager(Executor):
         selection = self._get_next_member()
         await ctx.send_message(AgentExecutorRequest(messages=[], should_respond=True), target_id=selection)
 
-    @handler(output_types=[AgentExecutorRequest])
-    async def handle_request_response(self, response: list[ChatMessage], ctx: WorkflowContext) -> None:
+    @handler
+    async def handle_request_response(
+        self, response: list[ChatMessage], ctx: WorkflowContext[AgentExecutorRequest]
+    ) -> None:
         """Handler that processes the response from the RequestInfoExecutor."""
         # Update the chat history with the response
         self._chat_history.extend(response)
 
         # Send the response to the other members
-        await self._broadcast_message(response, ctx)
+        await asyncio.gather(*[
+            ctx.send_message(
+                AgentExecutorRequest(messages=response, should_respond=False),
+                target_id=member_id,
+            )
+            for member_id in self._members
+        ])
 
         # Check for termination condition
         if self._should_terminate():
@@ -93,22 +118,6 @@ class CriticGroupChatManager(Executor):
         selection = self._get_next_member()
         await ctx.send_message(AgentExecutorRequest(messages=[], should_respond=True), target_id=selection)
 
-    async def _broadcast_message(
-        self,
-        messages: list[ChatMessage],
-        ctx: WorkflowContext,
-        exclude_id: str | None = None,
-    ) -> None:
-        """Broadcast messages to all members."""
-        await asyncio.gather(*[
-            ctx.send_message(
-                AgentExecutorRequest(messages=messages, should_respond=False),
-                target_id=member_id,
-            )
-            for member_id in self._members
-            if member_id != exclude_id
-        ])
-
     def _should_terminate(self) -> bool:
         """Determine if the group chat should terminate based on the last message."""
         if len(self._chat_history) == 0:

python/samples/getting_started/workflow/step_06_map_reduce_and_visualization.py

@@ -5,6 +5,7 @@ import asyncio
 import os
 from collections import defaultdict
 from dataclasses import dataclass
+from typing import Any
 
 import aiofiles
 from agent_framework.workflow import (
@@ -52,8 +53,8 @@ class Split(Executor):
         super().__init__(id)
         self._map_executor_ids = map_executor_ids
 
-    @handler(output_types=[SplitCompleted])
-    async def split(self, data: str, ctx: WorkflowContext) -> None:
+    @handler
+    async def split(self, data: str, ctx: WorkflowContext[SplitCompleted]) -> None:
         """Execute the task by splitting the data into chunks.
 
         Args:
@@ -107,8 +108,8 @@ class MapCompleted:
 class Map(Executor):
     """An executor that applies a function to each item in the data and save the result to a file."""
 
-    @handler(output_types=[MapCompleted])
-    async def map(self, _: SplitCompleted, ctx: WorkflowContext) -> None:
+    @handler
+    async def map(self, _: SplitCompleted, ctx: WorkflowContext[MapCompleted]) -> None:
         """Execute the task by applying a function to each item and same result to a file.
 
         Args:
@@ -144,8 +145,8 @@ class Shuffle(Executor):
         super().__init__(id)
         self._reducer_ids = reducer_ids
 
-    @handler(output_types=[ShuffleCompleted])
-    async def shuffle(self, data: list[MapCompleted], ctx: WorkflowContext) -> None:
+    @handler
+    async def shuffle(self, data: list[MapCompleted], ctx: WorkflowContext[ShuffleCompleted]) -> None:
         """Execute the task by aggregating the results.
 
         Args:
@@ -219,8 +220,8 @@ class ReduceCompleted:
 class Reduce(Executor):
     """An executor that reduces the results from the ShuffleExecutor."""
 
-    @handler(output_types=[ReduceCompleted])
-    async def _execute(self, data: ShuffleCompleted, ctx: WorkflowContext) -> None:
+    @handler
+    async def _execute(self, data: ShuffleCompleted, ctx: WorkflowContext[ReduceCompleted]) -> None:
         """Execute the task by reducing the results.
 
         Args:
@@ -253,7 +254,7 @@ class CompletionExecutor(Executor):
     """An executor that completes the workflow by aggregating the results from the ReduceExecutors."""
 
     @handler
-    async def complete(self, data: list[ReduceCompleted], ctx: WorkflowContext) -> None:
+    async def complete(self, data: list[ReduceCompleted], ctx: WorkflowContext[Any]) -> None:
         """Execute the task by aggregating the results.
 
         Args:

python/samples/getting_started/workflow/step_07_checkpoint_resume.py

@@ -3,6 +3,7 @@
 import asyncio
 import os
 from pathlib import Path
+from typing import Any
 
 from agent_framework.workflow import (
     Executor,
@@ -38,8 +39,8 @@ os.makedirs(TEMP_DIR, exist_ok=True)
 
 
 class UpperCaseExecutor(Executor):
-    @handler(output_types=[str])
-    async def to_upper_case(self, text: str, ctx: WorkflowContext) -> None:
+    @handler
+    async def to_upper_case(self, text: str, ctx: WorkflowContext[str]) -> None:
         result = text.upper()
         print(f"UpperCaseExecutor: '{text}' -> '{result}'")
         # Persist executor state into checkpointable context
@@ -57,8 +58,8 @@ class UpperCaseExecutor(Executor):
 
 
 class LowerCaseExecutor(Executor):
-    @handler(output_types=[str])
-    async def to_lower_case(self, text: str, ctx: WorkflowContext) -> None:
+    @handler
+    async def to_lower_case(self, text: str, ctx: WorkflowContext[Any]) -> None:
         result = text.lower()
         print(f"LowerCaseExecutor: '{text}' -> '{result}'")
         # Read from shared_state written by UpperCaseExecutor
@@ -82,8 +83,8 @@ class ReverseTextExecutor(Executor):
         """Initialize the executor with an ID."""
         super().__init__(id=id)
 
-    @handler(output_types=[str])
-    async def reverse_text(self, text: str, ctx: WorkflowContext) -> None:
+    @handler
+    async def reverse_text(self, text: str, ctx: WorkflowContext[str]) -> None:
         result = text[::-1]
         print(f"ReverseTextExecutor: '{text}' -> '{result}'")
         # Persist executor state into checkpointable context