agent/agent-framework
2
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