agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
This commit is contained in:
Eric Zhu
2025-08-20 19:40:43 -07:00
committed by GitHub
Unverified
parent 123a0bca10
commit 65836ab125
18 changed files with 423 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/workflow/agent_framework_workflow/_executor.py
+82 -29
View File
@@ -1,11 +1,13 @@
# Copyright (c) Microsoft. All rights reserved.
import contextlib
import functools
import inspect
import uuid
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any, ClassVar, TypeVar, overload
from types import UnionType
from typing import Any, ClassVar, TypeVar, Union, get_args, get_origin, overload
from agent_framework import AgentRunResponse, AgentRunResponseUpdate, AgentThread, AIAgent, ChatMessage
@@ -33,7 +35,7 @@ class Executor:
"""
self._id = id or f"{self.__class__.__name__}/{uuid.uuid4()}"
self._handlers: dict[type, Callable[[Any, WorkflowContext], Any]] = {}
self._handlers: dict[type, Callable[[Any, WorkflowContext[Any]], Any]] = {}
self._discover_handlers()
if not self._handlers:
@@ -42,7 +44,7 @@ class Executor:
"Please define at least one handler using the @handler decorator."
)
async def execute(self, message: Any, context: WorkflowContext) -> None:
async def execute(self, message: Any, context: WorkflowContext[Any]) -> None:
"""Execute the executor with a given message and context.
Args:
@@ -52,7 +54,7 @@ class Executor:
Returns:
An awaitable that resolves to the result of the execution.
"""
handler: Callable[[Any, WorkflowContext], Any] | None = None
handler: Callable[[Any, WorkflowContext[Any]], Any] | None = None
for message_type in self._handlers:
if is_instance_of(message, message_type):
handler = self._handlers[message_type]
@@ -104,54 +106,93 @@ ExecutorT = TypeVar("ExecutorT", bound="Executor")
@overload
def handler(
func: Callable[[ExecutorT, Any, WorkflowContext], Awaitable[Any]],
) -> Callable[[ExecutorT, Any, WorkflowContext], Awaitable[Any]]: ...
func: Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]],
) -> Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]: ...
@overload
def handler(
func: None = None,
*,
output_types: list[type] | None = None,
) -> Callable[
[Callable[[ExecutorT, Any, WorkflowContext], Awaitable[Any]]],
Callable[[ExecutorT, Any, WorkflowContext], Awaitable[Any]],
[Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]],
Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]],
]: ...
def handler(
func: Callable[[ExecutorT, Any, WorkflowContext], Awaitable[Any]] | None = None,
*,
output_types: list[type] | None = None,
func: Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]] | None = None,
) -> (
Callable[[ExecutorT, Any, WorkflowContext], Awaitable[Any]]
Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]
| Callable[
[Callable[[ExecutorT, Any, WorkflowContext], Awaitable[Any]]],
Callable[[ExecutorT, Any, WorkflowContext], Awaitable[Any]],
[Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]],
Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]],
]
):
"""Decorator to register a handler for an executor.
Args:
func: The function to decorate. Can be None when using with parameters.
output_types: Optional list of message types this handler can emit.
func: The function to decorate. Can be None when used without parameters.
Returns:
The decorated function with handler metadata.
Example:
@handler
async def handle_string(self, message: str, ctx: WorkflowContext) -> None:
async def handle_string(self, message: str, ctx: WorkflowContext[str]) -> None:
...
@handler(output_types=[str, int])
async def handle_data(self, message: dict, ctx: WorkflowContext) -> None:
@handler
async def handle_data(self, message: dict, ctx: WorkflowContext[str | int]) -> None:
...
"""
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], Awaitable[Any]],
) -> Callable[[ExecutorT, Any, WorkflowContext], Awaitable[Any]]:
func: Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]],
) -> Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]:
# Extract the message type from a handler function.
sig = inspect.signature(func)
params = list(sig.parameters.values())
@@ -163,15 +204,27 @@ def handler(
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:
async def wrapper(self: ExecutorT, message: Any, ctx: WorkflowContext[Any]) -> Any:
"""Wrapper function to call the handler."""
return await func(self, message, ctx)
# Preserve the original function signature for introspection during validation
with contextlib.suppress(Exception):
wrapper.__signature__ = sig # type: ignore[attr-defined]
wrapper._handler_spec = { # type: ignore
"name": func.__name__,
"message_type": message_type,
"output_types": output_types or [],
# Keep output_types in spec for validators, inferred from WorkflowContext[T]
"output_types": inferred_output_types,
}
return wrapper
@@ -239,8 +292,8 @@ class AgentExecutor(Executor):
self._streaming = streaming
self._cache: list[ChatMessage] = []
@handler(output_types=[AgentExecutorResponse])
async def run(self, request: AgentExecutorRequest, ctx: WorkflowContext) -> None:
@handler
async def run(self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorResponse]) -> None:
"""Run the agent executor with the given request."""
self._cache.extend(request.messages)
@@ -300,7 +353,7 @@ class RequestInfoExecutor(Executor):
self._request_events: dict[str, RequestInfoEvent] = {}
@handler
async def run(self, message: RequestInfoMessage, ctx: WorkflowContext) -> None:
async def run(self, message: RequestInfoMessage, ctx: WorkflowContext[None]) -> None:
"""Run the RequestInfoExecutor with the given message."""
source_executor_id = ctx.get_source_executor_id()
@@ -317,7 +370,7 @@ class RequestInfoExecutor(Executor):
self,
response_data: Any,
request_id: str,
ctx: WorkflowContext,
ctx: WorkflowContext[Any],
) -> None:
"""Handle a response to a request.
python/packages/workflow/agent_framework_workflow/_validation.py
+126 -3
View File
@@ -5,6 +5,7 @@ import logging
from collections import defaultdict
from collections.abc import Sequence
from enum import Enum
from types import UnionType
from typing import Any, Union, get_args, get_origin
from ._edge import Edge, EdgeGroup, FanInEdgeGroup
@@ -20,6 +21,7 @@ class ValidationTypeEnum(Enum):
EDGE_DUPLICATION = "EDGE_DUPLICATION"
TYPE_COMPATIBILITY = "TYPE_COMPATIBILITY"
GRAPH_CONNECTIVITY = "GRAPH_CONNECTIVITY"
HANDLER_OUTPUT_ANNOTATION = "HANDLER_OUTPUT_ANNOTATION"
class WorkflowValidationError(Exception):
@@ -75,6 +77,23 @@ 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
# endregion
@@ -116,6 +135,7 @@ class WorkflowGraphValidator:
# Run all checks
self._validate_edge_duplication()
self._validate_handler_output_annotations()
self._validate_type_compatibility()
self._validate_graph_connectivity(start_executor_id)
self._validate_self_loops()
@@ -132,6 +152,109 @@ class WorkflowGraphValidator:
return executors
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
"""
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):
# 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, 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)
items: list[Any]
items = list(get_args(t_out)) if union_origin in (Union, UnionType) else [t_out]
def _is_type_like(x: Any) -> bool:
# A "type-like" entry is either a class/type or a typing alias
# (e.g., list[str] has an origin and args)
return isinstance(x, type) or get_origin(x) is not None
invalid = [x for x in 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",
)
# endregion
# region Edge and Type Validation
@@ -191,7 +314,7 @@ class WorkflowGraphValidator:
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 output_types to @handler decorators for better validation."
f"Consider adding WorkflowContext[T] generics in handlers for better validation."
)
if not target_input_types:
logger.warning(
@@ -472,11 +595,11 @@ class WorkflowGraphValidator:
source_origin = get_origin(source_type)
target_origin = get_origin(target_type)
if target_origin is Union:
if target_origin in (Union, UnionType):
target_args = get_args(target_type)
return any(WorkflowGraphValidator._is_type_compatible(source_type, arg) for arg in target_args)
if source_origin is Union:
if source_origin in (Union, UnionType):
source_args = get_args(source_type)
return all(WorkflowGraphValidator._is_type_compatible(arg, target_type) for arg in source_args)
python/packages/workflow/agent_framework_workflow/_workflow_context.py
+6 -4
View File
@@ -1,13 +1,15 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
from typing import Any, Generic, TypeVar
from ._events import WorkflowEvent
from ._runner_context import Message, RunnerContext
from ._shared_state import SharedState
T_Out = TypeVar("T_Out")
class WorkflowContext:
class WorkflowContext(Generic[T_Out]):
"""Context for executors in a workflow.
This class is used to provide a way for executors to interact with the workflow
@@ -39,11 +41,11 @@ class WorkflowContext:
if not self._source_executor_ids:
raise ValueError("source_executor_ids cannot be empty. At least one source executor ID is required.")
async def send_message(self, message: Any, target_id: str | None = None) -> None:
async def send_message(self, message: T_Out, target_id: str | None = None) -> None:
"""Send a message to the workflow context.
Args:
message: The message to send. This can be any data type that the target executor can handle.
message: The message to send. This must conform to the output type(s) declared on this context.
target_id: The ID of the target executor to send the message to.
If None, the message will be sent to all target executors.
"""
python/packages/workflow/tests/test_executor.py
+4 -4
View File
@@ -74,13 +74,13 @@ def test_executor_handlers_with_output_types():
class MockExecutorWithOutputTypes(Executor): # type: ignore
"""A mock executor with handlers that specify output types."""
@handler(output_types=[str])
async def handle_string(self, text: str, ctx: WorkflowContext) -> None: # type: ignore
@handler
async def handle_string(self, text: str, ctx: WorkflowContext[str]) -> None: # type: ignore
"""A mock handler that outputs a string."""
pass
@handler(output_types=[int])
async def handle_integer(self, number: int, ctx: WorkflowContext) -> None: # type: ignore
@handler
async def handle_integer(self, number: int, ctx: WorkflowContext[int]) -> None: # type: ignore
"""A mock handler that outputs an integer."""
pass
python/packages/workflow/tests/test_runner.py
+2 -2
View File
@@ -22,8 +22,8 @@ class MockMessage:
class MockExecutor(Executor):
"""A mock executor for testing purposes."""
@handler(output_types=[MockMessage])
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext) -> None:
@handler
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage]) -> None:
if message.data < 10:
await ctx.send_message(MockMessage(data=message.data + 1))
else:
python/packages/workflow/tests/test_validation.py
+106 -25
View File
@@ -18,48 +18,49 @@ from agent_framework_workflow import (
validate_workflow_graph,
)
from agent_framework_workflow._edge import SingleEdgeGroup
from agent_framework_workflow._validation import HandlerOutputAnnotationError
class StringExecutor(Executor):
@handler(output_types=[str])
async def handle_string(self, message: str, ctx: WorkflowContext) -> None:
@handler
async def handle_string(self, message: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message(message.upper())
class StringAggregator(Executor):
"""A mock executor that aggregates results from multiple executors."""
@handler(output_types=[str])
async def mock_handler(self, messages: list[str], ctx: WorkflowContext) -> None:
@handler
async def mock_handler(self, messages: list[str], ctx: WorkflowContext[str]) -> None:
# This mock simply returns the data incremented by 1
await ctx.send_message("Aggregated: " + ", ".join(messages))
class IntExecutor(Executor):
@handler(output_types=[int])
async def handle_int(self, message: int, ctx: WorkflowContext) -> None:
@handler
async def handle_int(self, message: int, ctx: WorkflowContext[int]) -> None:
await ctx.send_message(message * 2)
class AnyExecutor(Executor):
@handler
async def handle_any(self, message: Any, ctx: WorkflowContext) -> None:
async def handle_any(self, message: Any, ctx: WorkflowContext[Any]) -> None:
await ctx.send_message(f"Processed: {message}")
class NoOutputTypesExecutor(Executor):
@handler
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
async def handle_message(self, message: str, ctx: WorkflowContext[Any]) -> None:
await ctx.send_message("processed")
class MultiTypeExecutor(Executor):
@handler(output_types=[str])
async def handle_string(self, message: str, ctx: WorkflowContext) -> None:
@handler
async def handle_string(self, message: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message(f"String: {message}")
@handler(output_types=[int])
async def handle_int(self, message: int, ctx: WorkflowContext) -> None:
@handler
async def handle_int(self, message: int, ctx: WorkflowContext[str]) -> None:
await ctx.send_message(f"Int: {message}")
@@ -221,13 +222,13 @@ def test_complex_workflow_validation():
def test_type_compatibility_inheritance():
class BaseExecutor(Executor):
@handler(output_types=[str])
async def handle_base(self, message: str, ctx: WorkflowContext) -> None:
@handler
async def handle_base(self, message: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message("base")
class DerivedExecutor(Executor):
@handler(output_types=[str])
async def handle_derived(self, message: str, ctx: WorkflowContext) -> None:
@handler
async def handle_derived(self, message: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message("derived")
base_executor = BaseExecutor(id="base")
@@ -306,7 +307,7 @@ def test_logging_for_missing_output_types(caplog: Any) -> None:
assert workflow is not None
assert "has no output type annotations" in caplog.text
assert "Consider adding output_types to @handler decorators" in caplog.text
assert "Consider adding WorkflowContext[T] generics" in caplog.text
def test_logging_for_missing_input_types(caplog: Any) -> None:
@@ -504,13 +505,13 @@ def test_enhanced_type_compatibility_error_details():
def test_union_type_compatibility_validation() -> None:
class UnionOutputExecutor(Executor):
@handler(output_types=[str, int])
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
@handler
async def handle_message(self, message: str, ctx: WorkflowContext[str | int]) -> None:
await ctx.send_message("output")
class UnionInputExecutor(Executor):
@handler(output_types=[str])
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
@handler
async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message("processed")
union_output = UnionOutputExecutor(id="union_output")
@@ -524,13 +525,13 @@ def test_union_type_compatibility_validation() -> None:
def test_generic_type_compatibility() -> None:
class ListOutputExecutor(Executor):
@handler(output_types=[list[str]])
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
@handler
async def handle_message(self, message: str, ctx: WorkflowContext[list[str]]) -> None:
await ctx.send_message(["output"])
class ListInputExecutor(Executor):
@handler(output_types=[str])
async def handle_message(self, message: list[str], ctx: WorkflowContext) -> None:
@handler
async def handle_message(self, message: list[str], ctx: WorkflowContext[str]) -> None:
await ctx.send_message("processed")
list_output = ListOutputExecutor(id="list_output")
@@ -556,3 +557,83 @@ def test_validation_enum_usage() -> None:
# Test enum string representation
assert str(ValidationTypeEnum.EDGE_DUPLICATION) == "ValidationTypeEnum.EDGE_DUPLICATION"
assert ValidationTypeEnum.EDGE_DUPLICATION.value == "EDGE_DUPLICATION"
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
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
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: # 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()
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
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
assert "invalid entries" 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:
# does not emit
return None
start = StringExecutor(id="s")
none_exec = NoneExecutor(id="n")
# Should build successfully
wf = WorkflowBuilder().add_edge(start, none_exec).set_start_executor(start).build()
assert wf is not None
def test_handler_ctx_any_is_allowed_but_skips_type_checks(caplog: Any) -> None:
caplog.set_level(logging.WARNING)
class AnyOutExecutor(Executor):
@handler
async def handle(self, message: str, ctx: WorkflowContext[Any]) -> None:
return None
start = StringExecutor(id="s")
any_out = AnyOutExecutor(id="a")
# Builds; later edges from this executor will skip type compatibility when outputs are unspecified
wf = WorkflowBuilder().add_edge(start, any_out).set_start_executor(start).build()
assert wf is not None
python/packages/workflow/tests/test_viz.py
+3 -3
View File
@@ -9,8 +9,8 @@ from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowContext,
class MockExecutor(Executor):
"""A mock executor for testing purposes."""
@handler(output_types=[str])
async def mock_handler(self, message: str, ctx: WorkflowContext) -> None:
@handler
async def mock_handler(self, message: str, ctx: WorkflowContext[None]) -> None:
"""A mock handler that does nothing."""
pass
@@ -19,7 +19,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: # type: ignore[type-arg]
async def handle(self, message: list[str], ctx: WorkflowContext[None]) -> None: # type: ignore[type-arg]
pass
python/packages/workflow/tests/test_workflow.py
+11 -10
View File
@@ -2,6 +2,7 @@
import tempfile
from dataclasses import dataclass
from typing import Any
import pytest
from agent_framework.workflow import (
@@ -35,8 +36,8 @@ class MockExecutor(Executor):
super().__init__(id=id)
self.limit = limit
@handler(output_types=[MockMessage])
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext) -> None:
@handler
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage]) -> None:
if message.data < self.limit:
await ctx.send_message(MockMessage(data=message.data + 1))
else:
@@ -47,7 +48,7 @@ class MockAggregator(Executor):
"""A mock executor that aggregates results from multiple executors."""
@handler
async def mock_handler(self, messages: list[MockMessage], ctx: WorkflowContext) -> None:
async def mock_handler(self, messages: list[MockMessage], 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)))
@@ -62,14 +63,14 @@ class ApprovalMessage:
class MockExecutorRequestApproval(Executor):
"""A mock executor that simulates a request for approval."""
@handler(output_types=[RequestInfoMessage])
async def mock_handler_a(self, message: MockMessage, ctx: WorkflowContext) -> None:
@handler
async def mock_handler_a(self, message: MockMessage, ctx: WorkflowContext[RequestInfoMessage]) -> None:
"""A mock handler that requests approval."""
await ctx.set_shared_state(self.id, message.data)
await ctx.send_message(RequestInfoMessage())
@handler(output_types=[MockMessage])
async def mock_handler_b(self, message: ApprovalMessage, ctx: WorkflowContext) -> None:
@handler
async def mock_handler_b(self, message: ApprovalMessage, ctx: WorkflowContext[MockMessage]) -> None:
"""A mock handler that processes the approval response."""
data = await ctx.get_shared_state(self.id)
if message.approved:
@@ -285,7 +286,7 @@ async def test_fan_in():
def simple_executor() -> Executor:
class SimpleExecutor(Executor):
@handler
async def handle_message(self, message: Message, context: WorkflowContext) -> None:
async def handle_message(self, message: Message, context: WorkflowContext[None]) -> None:
pass
return SimpleExecutor("test_executor")
@@ -494,8 +495,8 @@ class StateTrackingMessage:
class StateTrackingExecutor(Executor):
"""An executor that tracks state in shared state to test context reset behavior."""
@handler(output_types=[])
async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext) -> None:
@handler
async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext[Any]) -> None:
"""Handle the message and track it in shared state."""
# Get existing messages from shared state
try:
python/packages/workflow/tests/test_workflow_builder.py
+4 -4
View File
@@ -17,8 +17,8 @@ class MockMessage:
class MockExecutor(Executor):
"""A mock executor for testing purposes."""
@handler(output_types=[MockMessage])
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext) -> None:
@handler
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage]) -> None:
"""A mock handler that does nothing."""
pass
@@ -26,8 +26,8 @@ class MockExecutor(Executor):
class MockAggregator(Executor):
"""A mock executor that aggregates results from multiple executors."""
@handler(output_types=[MockMessage])
async def mock_handler(self, messages: list[MockMessage], ctx: WorkflowContext) -> None:
@handler
async def mock_handler(self, messages: list[MockMessage], ctx: WorkflowContext[MockMessage]) -> None:
# This mock simply returns the data incremented by 1
pass
python/samples/getting_started/workflow/step_00_foundation_patterns.py
+7 -7
View File
@@ -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
+4 -3
View File
@@ -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
+3 -3
View File
@@ -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
+4 -4
View File
@@ -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
+4 -4
View File
@@ -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
+6 -4
View File
@@ -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
+34 -25
View File
@@ -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
+10 -9
View File
@@ -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
+7 -6
View File
@@ -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