mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Add FunctionExecutor and @executor decorator (#589)
* Add FunctionExecutor and @executor decorator * refactor * add single argument function * fix test * update example code * add support for sync funciton
This commit is contained in:
committed by
GitHub
Unverified
parent
3286e8c8b1
commit
ed8461aa7d
@@ -7,6 +7,8 @@ PACKAGE_NAME = "agent_framework_workflow"
|
||||
PACKAGE_EXTRA = "workflow"
|
||||
_IMPORTS = [
|
||||
"Executor",
|
||||
"FunctionExecutor",
|
||||
"executor",
|
||||
"WorkflowContext",
|
||||
"__version__",
|
||||
"events",
|
||||
|
||||
@@ -14,6 +14,7 @@ from agent_framework_workflow import (
|
||||
ExecutorEvent,
|
||||
ExecutorInvokeEvent,
|
||||
FileCheckpointStorage,
|
||||
FunctionExecutor,
|
||||
InMemoryCheckpointStorage,
|
||||
MagenticBuilder,
|
||||
MagenticCallbackEvent,
|
||||
@@ -42,6 +43,7 @@ from agent_framework_workflow import (
|
||||
WorkflowStartedEvent,
|
||||
WorkflowViz,
|
||||
__version__,
|
||||
executor,
|
||||
handler,
|
||||
intercepts_request,
|
||||
)
|
||||
@@ -60,6 +62,7 @@ __all__ = [
|
||||
"ExecutorEvent",
|
||||
"ExecutorInvokeEvent",
|
||||
"FileCheckpointStorage",
|
||||
"FunctionExecutor",
|
||||
"InMemoryCheckpointStorage",
|
||||
"MagenticBuilder",
|
||||
"MagenticCallbackEvent",
|
||||
@@ -88,6 +91,7 @@ __all__ = [
|
||||
"WorkflowStartedEvent",
|
||||
"WorkflowViz",
|
||||
"__version__",
|
||||
"executor",
|
||||
"handler",
|
||||
"intercepts_request",
|
||||
]
|
||||
|
||||
@@ -38,6 +38,7 @@ from ._executor import (
|
||||
handler,
|
||||
intercepts_request,
|
||||
)
|
||||
from ._function_executor import FunctionExecutor, executor
|
||||
from ._magentic import (
|
||||
MagenticAgentDeltaEvent,
|
||||
MagenticAgentExecutor,
|
||||
@@ -99,6 +100,7 @@ __all__ = [
|
||||
"ExecutorEvent",
|
||||
"ExecutorInvokeEvent",
|
||||
"FileCheckpointStorage",
|
||||
"FunctionExecutor",
|
||||
"GraphConnectivityError",
|
||||
"InMemoryCheckpointStorage",
|
||||
"InProcRunnerContext",
|
||||
@@ -145,6 +147,7 @@ __all__ = [
|
||||
"WorkflowValidationError",
|
||||
"WorkflowViz",
|
||||
"__version__",
|
||||
"executor",
|
||||
"handler",
|
||||
"intercepts_request",
|
||||
"validate_workflow_graph",
|
||||
|
||||
@@ -59,6 +59,7 @@ class Executor(AFBaseModel):
|
||||
|
||||
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._discover_handlers()
|
||||
|
||||
if not self._handlers and not self._request_interceptors:
|
||||
@@ -254,6 +255,34 @@ class Executor(AFBaseModel):
|
||||
"""
|
||||
return any(is_instance_of(message, message_type) for message_type in self._handlers)
|
||||
|
||||
def register_instance_handler(
|
||||
self,
|
||||
name: str,
|
||||
func: Callable[[Any, WorkflowContext[Any]], Awaitable[Any]],
|
||||
message_type: type,
|
||||
ctx_annotation: Any,
|
||||
output_types: list[type],
|
||||
) -> None:
|
||||
"""Register a handler at instance level.
|
||||
|
||||
Args:
|
||||
name: Name of the handler function for error reporting
|
||||
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
|
||||
"""
|
||||
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({
|
||||
"name": name,
|
||||
"message_type": message_type,
|
||||
"ctx_annotation": ctx_annotation,
|
||||
"output_types": output_types,
|
||||
})
|
||||
|
||||
def can_handle_type(self, message_type: type[Any]) -> bool:
|
||||
"""Check if the executor can handle a given message type.
|
||||
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Function-based Executor and decorator utilities.
|
||||
|
||||
This module provides:
|
||||
- FunctionExecutor: an Executor subclass that wraps a user-defined function
|
||||
with signature (message) or (message, ctx: WorkflowContext[T]). Both sync and async functions are supported.
|
||||
Synchronous functions are executed in a thread pool using asyncio.to_thread() to avoid blocking the event loop.
|
||||
- executor decorator: converts such a function into a ready-to-use Executor instance
|
||||
with proper type validation and handler registration.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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 ._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]
|
||||
|
||||
|
||||
class FunctionExecutor(Executor):
|
||||
"""Executor that wraps a user-defined function.
|
||||
|
||||
This executor allows users to define simple functions (both sync and async) and use them
|
||||
as workflow executors without needing to create full executor classes.
|
||||
|
||||
Synchronous functions are executed in a thread pool using asyncio.to_thread() to avoid
|
||||
blocking the event loop.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _validate_function(func: Callable[..., Any]) -> None:
|
||||
"""Validate that the function has the correct signature for an executor.
|
||||
|
||||
Args:
|
||||
func: The function to validate (can be sync or async)
|
||||
|
||||
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."
|
||||
)
|
||||
|
||||
def __init__(self, func: Callable[..., Any], id: str | None = None):
|
||||
"""Initialize the FunctionExecutor with a user-defined function.
|
||||
|
||||
Args:
|
||||
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
|
||||
|
||||
# Determine if function has WorkflowContext parameter
|
||||
has_context = len(params) == 2
|
||||
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")
|
||||
kwargs = {"id": executor_id, "type": "FunctionExecutor"}
|
||||
|
||||
# Set up the base class attributes manually to avoid _discover_handlers
|
||||
from pydantic import BaseModel
|
||||
|
||||
BaseModel.__init__(self, **kwargs)
|
||||
|
||||
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]] = []
|
||||
|
||||
# Store the original function and whether it has context
|
||||
self._original_func = func
|
||||
self._has_context = has_context
|
||||
self._is_async = is_async
|
||||
|
||||
# Create a wrapper function that always accepts both message and context
|
||||
if has_context and is_async:
|
||||
# Async function with context - already has the right signature
|
||||
wrapped_func: Callable[[Any, WorkflowContext[Any]], Awaitable[Any]] = func # type: ignore
|
||||
elif has_context and not is_async:
|
||||
# Sync function with context - wrap to make async using thread pool
|
||||
async def wrapped_func(message: Any, ctx: WorkflowContext[Any]) -> Any:
|
||||
# Call the sync function with both parameters in a thread
|
||||
return await asyncio.to_thread(func, message, ctx) # type: ignore
|
||||
|
||||
elif not has_context and is_async:
|
||||
# Async function without context - wrap to ignore context
|
||||
async def wrapped_func(message: Any, ctx: WorkflowContext[Any]) -> Any:
|
||||
# Call the async function with just the message
|
||||
return await func(message) # type: ignore
|
||||
|
||||
else:
|
||||
# Sync function without context - wrap to make async and ignore context using thread pool
|
||||
async def wrapped_func(message: Any, ctx: WorkflowContext[Any]) -> Any:
|
||||
# Call the sync function with just the message in a thread
|
||||
return await asyncio.to_thread(func, message) # type: ignore
|
||||
|
||||
# Now register our instance handler
|
||||
self.register_instance_handler(
|
||||
name=func.__name__,
|
||||
func=wrapped_func,
|
||||
message_type=message_type,
|
||||
ctx_annotation=ctx_annotation,
|
||||
output_types=output_types,
|
||||
)
|
||||
|
||||
# Now we can safely call _discover_handlers (it won't find any class-level handlers)
|
||||
self._discover_handlers()
|
||||
|
||||
|
||||
@overload
|
||||
def executor(func: Callable[..., Any]) -> FunctionExecutor: ...
|
||||
|
||||
|
||||
@overload
|
||||
def executor(*, id: str | None = None) -> Callable[[Callable[..., Any]], FunctionExecutor]: ...
|
||||
|
||||
|
||||
def executor(
|
||||
func: Callable[..., Any] | None = None, *, id: str | None = None
|
||||
) -> Callable[[Callable[..., Any]], FunctionExecutor] | FunctionExecutor:
|
||||
"""Decorator that converts a function into a FunctionExecutor instance.
|
||||
|
||||
Supports both synchronous and asynchronous functions. Synchronous functions
|
||||
are executed in a thread pool to avoid blocking the event loop.
|
||||
|
||||
Usage:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# With arguments (async function):
|
||||
@executor(id="upper_case")
|
||||
async def to_upper(text: str, ctx: WorkflowContext[str]):
|
||||
await ctx.send_message(text.upper())
|
||||
|
||||
|
||||
# Without parentheses (sync function - runs in thread pool):
|
||||
@executor
|
||||
def process_data(data: str):
|
||||
# Process data without sending messages
|
||||
return data.upper()
|
||||
|
||||
|
||||
# Sync function with context (runs in thread pool):
|
||||
@executor
|
||||
def sync_with_context(data: int, ctx: WorkflowContext[int]):
|
||||
# Note: sync functions can still use context
|
||||
return data * 2
|
||||
|
||||
Returns:
|
||||
An Executor instance that can be wired into a Workflow.
|
||||
"""
|
||||
|
||||
def wrapper(func: Callable[..., Any]) -> FunctionExecutor:
|
||||
return FunctionExecutor(func, id=id)
|
||||
|
||||
# If func is provided, this means @executor was used without parentheses
|
||||
if func is not None:
|
||||
return wrapper(func)
|
||||
|
||||
# Otherwise, return the wrapper for @executor() or @executor(id="...")
|
||||
return wrapper
|
||||
@@ -14,6 +14,21 @@ from ._executor import Executor
|
||||
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."""
|
||||
@@ -252,15 +267,10 @@ class WorkflowGraphValidator:
|
||||
|
||||
# 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]
|
||||
type_items: list[Any]
|
||||
type_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)]
|
||||
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,
|
||||
@@ -268,6 +278,62 @@ class WorkflowGraphValidator:
|
||||
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",
|
||||
)
|
||||
|
||||
# endregion
|
||||
|
||||
# region Edge and Type Validation
|
||||
@@ -408,6 +474,12 @@ class WorkflowGraphValidator:
|
||||
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]]:
|
||||
|
||||
@@ -0,0 +1,426 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from agent_framework.workflow import (
|
||||
FunctionExecutor,
|
||||
WorkflowBuilder,
|
||||
WorkflowCompletedEvent,
|
||||
WorkflowContext,
|
||||
executor,
|
||||
)
|
||||
|
||||
|
||||
class TestFunctionExecutor:
|
||||
"""Test suite for FunctionExecutor and @executor decorator."""
|
||||
|
||||
def test_function_executor_basic(self):
|
||||
"""Test basic FunctionExecutor creation and validation."""
|
||||
|
||||
async def process_string(text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(text.upper())
|
||||
|
||||
func_exec = FunctionExecutor(process_string)
|
||||
|
||||
# Check that handler was registered
|
||||
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]
|
||||
assert spec["name"] == "process_string"
|
||||
assert spec["message_type"] is str
|
||||
assert spec["output_types"] == [str]
|
||||
|
||||
def test_executor_decorator(self):
|
||||
"""Test @executor decorator creates proper FunctionExecutor."""
|
||||
|
||||
@executor(id="test_executor")
|
||||
async def process_int(value: int, ctx: WorkflowContext[int]) -> None:
|
||||
await ctx.send_message(value * 2)
|
||||
|
||||
assert isinstance(process_int, FunctionExecutor)
|
||||
assert process_int.id == "test_executor"
|
||||
assert int in process_int._handlers
|
||||
|
||||
# Check spec
|
||||
spec = process_int._instance_handler_specs[0]
|
||||
assert spec["message_type"] is int
|
||||
assert spec["output_types"] == [int]
|
||||
|
||||
def test_executor_decorator_without_id(self):
|
||||
"""Test @executor decorator uses function name as default ID."""
|
||||
|
||||
@executor
|
||||
async def my_function(data: dict, ctx: WorkflowContext[Any]) -> None:
|
||||
await ctx.send_message(data)
|
||||
|
||||
assert my_function.id == "my_function"
|
||||
|
||||
def test_executor_decorator_without_parentheses(self):
|
||||
"""Test @executor decorator works without parentheses."""
|
||||
|
||||
@executor
|
||||
async def no_parens_function(data: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(data.upper())
|
||||
|
||||
assert isinstance(no_parens_function, FunctionExecutor)
|
||||
assert no_parens_function.id == "no_parens_function"
|
||||
assert str in no_parens_function._handlers
|
||||
|
||||
# Also test with single parameter function
|
||||
@executor
|
||||
async def simple_no_parens(value: int):
|
||||
return value * 2
|
||||
|
||||
assert isinstance(simple_no_parens, FunctionExecutor)
|
||||
assert simple_no_parens.id == "simple_no_parens"
|
||||
assert int in simple_no_parens._handlers
|
||||
|
||||
def test_union_output_types(self):
|
||||
"""Test that union output types are properly inferred."""
|
||||
|
||||
@executor
|
||||
async def multi_output(text: str, ctx: WorkflowContext[str | int]) -> None:
|
||||
if text.isdigit():
|
||||
await ctx.send_message(int(text))
|
||||
else:
|
||||
await ctx.send_message(text.upper())
|
||||
|
||||
spec = multi_output._instance_handler_specs[0]
|
||||
assert set(spec["output_types"]) == {str, int}
|
||||
|
||||
def test_none_output_type(self):
|
||||
"""Test WorkflowContext[None] produces empty output types."""
|
||||
|
||||
@executor
|
||||
async def no_output(data: Any, ctx: WorkflowContext[None]) -> None:
|
||||
# This executor doesn't send any messages
|
||||
pass
|
||||
|
||||
spec = no_output._instance_handler_specs[0]
|
||||
assert spec["output_types"] == []
|
||||
|
||||
def test_any_output_type(self):
|
||||
"""Test WorkflowContext[Any] produces empty 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"] == []
|
||||
|
||||
def test_validation_errors(self):
|
||||
"""Test various validation errors in function signatures."""
|
||||
|
||||
# Wrong number of parameters (now accepts 1 or 2, so 0 or 3+ should fail)
|
||||
async def no_params() -> None:
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="one or two parameters"):
|
||||
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"):
|
||||
FunctionExecutor(too_many_params) # type: ignore
|
||||
|
||||
# Missing message type annotation
|
||||
async def no_msg_type(data, ctx: WorkflowContext[str]) -> None: # type: ignore
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="type annotation for the message"):
|
||||
FunctionExecutor(no_msg_type) # type: ignore
|
||||
|
||||
# Missing ctx annotation (only for 2-parameter functions)
|
||||
async def no_ctx_type(data: str, ctx) -> None: # type: ignore
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="annotated as 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\\]"):
|
||||
FunctionExecutor(wrong_ctx_type) # type: ignore
|
||||
|
||||
# Unparameterized WorkflowContext
|
||||
async def unparameterized_ctx(data: str, ctx: WorkflowContext) -> None: # type: ignore
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="concrete T"):
|
||||
FunctionExecutor(unparameterized_ctx) # type: ignore
|
||||
|
||||
async def test_execution_in_workflow(self):
|
||||
"""Test that FunctionExecutor works properly in a workflow."""
|
||||
|
||||
@executor(id="upper")
|
||||
async def to_upper(text: str, ctx: WorkflowContext[str]) -> None:
|
||||
result = text.upper()
|
||||
await ctx.send_message(result)
|
||||
|
||||
@executor(id="reverse")
|
||||
async def reverse_text(text: str, ctx: WorkflowContext[Any]) -> None:
|
||||
result = text[::-1]
|
||||
await ctx.add_event(WorkflowCompletedEvent(result))
|
||||
|
||||
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()
|
||||
|
||||
assert completed is not None
|
||||
assert completed.data == "DLROW OLLEH"
|
||||
|
||||
def test_can_handle_method(self):
|
||||
"""Test that can_handle method works with instance handlers."""
|
||||
|
||||
@executor
|
||||
async def string_processor(text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(text)
|
||||
|
||||
assert string_processor.can_handle("hello")
|
||||
assert not string_processor.can_handle(123)
|
||||
assert not string_processor.can_handle([])
|
||||
|
||||
def test_duplicate_handler_registration(self):
|
||||
"""Test that registering duplicate handlers raises an error."""
|
||||
|
||||
async def first_handler(text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(text)
|
||||
|
||||
func_exec = FunctionExecutor(first_handler)
|
||||
|
||||
# Try to register another handler for the same type
|
||||
async def second_handler(message: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(message)
|
||||
|
||||
with pytest.raises(ValueError, match="Handler for type .* already registered"):
|
||||
func_exec.register_instance_handler(
|
||||
name="second",
|
||||
func=second_handler,
|
||||
message_type=str,
|
||||
ctx_annotation=WorkflowContext[str],
|
||||
output_types=[str],
|
||||
)
|
||||
|
||||
def test_complex_type_annotations(self):
|
||||
"""Test with complex type annotations like List[str], Dict[str, int], etc."""
|
||||
|
||||
@executor
|
||||
async def process_list(items: list[str], ctx: WorkflowContext[dict[str, int]]) -> None:
|
||||
result = {item: len(item) for item in items}
|
||||
await ctx.send_message(result)
|
||||
|
||||
spec = process_list._instance_handler_specs[0]
|
||||
assert spec["message_type"] == list[str]
|
||||
assert spec["output_types"] == [dict[str, int]]
|
||||
|
||||
def test_single_parameter_function(self):
|
||||
"""Test FunctionExecutor with single-parameter functions."""
|
||||
|
||||
@executor(id="simple_processor")
|
||||
async def process_simple(text: str):
|
||||
return text.upper()
|
||||
|
||||
assert isinstance(process_simple, FunctionExecutor)
|
||||
assert process_simple.id == "simple_processor"
|
||||
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]
|
||||
assert spec["message_type"] is str
|
||||
assert spec["output_types"] == []
|
||||
assert spec["ctx_annotation"] is None
|
||||
|
||||
def test_single_parameter_validation(self):
|
||||
"""Test validation for single-parameter functions."""
|
||||
|
||||
# Valid single-parameter function
|
||||
async def valid_single(data: int):
|
||||
return data * 2
|
||||
|
||||
func_exec = FunctionExecutor(valid_single)
|
||||
assert int in func_exec._handlers
|
||||
|
||||
# Single parameter with missing type annotation should still fail
|
||||
async def no_annotation(data): # type: ignore
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="type annotation for the message"):
|
||||
FunctionExecutor(no_annotation) # type: ignore
|
||||
|
||||
def test_single_parameter_can_handle(self):
|
||||
"""Test that single-parameter functions work with can_handle method."""
|
||||
|
||||
@executor
|
||||
async def int_processor(value: int):
|
||||
return value * 2
|
||||
|
||||
assert int_processor.can_handle(42)
|
||||
assert not int_processor.can_handle("hello")
|
||||
assert not int_processor.can_handle([])
|
||||
|
||||
async def test_single_parameter_execution(self):
|
||||
"""Test that single-parameter functions can be executed properly."""
|
||||
|
||||
@executor(id="double")
|
||||
async def double_value(value: int):
|
||||
return value * 2
|
||||
|
||||
# Since single-parameter functions can't send messages,
|
||||
# they're typically used as terminal nodes or for side effects
|
||||
WorkflowBuilder().set_start_executor(double_value).build()
|
||||
|
||||
# For testing purposes, we can check that the handler is registered correctly
|
||||
assert double_value.can_handle(5)
|
||||
assert int in double_value._handlers
|
||||
|
||||
def test_sync_function_basic(self):
|
||||
"""Test basic synchronous function support."""
|
||||
|
||||
@executor(id="sync_processor")
|
||||
def process_sync(text: str):
|
||||
return text.upper()
|
||||
|
||||
assert isinstance(process_sync, FunctionExecutor)
|
||||
assert process_sync.id == "sync_processor"
|
||||
assert str in process_sync._handlers
|
||||
|
||||
# Check spec - sync single parameter functions have no output types
|
||||
spec = process_sync._instance_handler_specs[0]
|
||||
assert spec["message_type"] is str
|
||||
assert spec["output_types"] == []
|
||||
assert spec["ctx_annotation"] is None
|
||||
|
||||
def test_sync_function_with_context(self):
|
||||
"""Test synchronous function with WorkflowContext."""
|
||||
|
||||
@executor
|
||||
def sync_with_ctx(value: int, ctx: WorkflowContext[int]):
|
||||
# Sync functions can still use context
|
||||
return value * 2
|
||||
|
||||
assert isinstance(sync_with_ctx, FunctionExecutor)
|
||||
assert sync_with_ctx.id == "sync_with_ctx"
|
||||
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]
|
||||
assert spec["message_type"] is int
|
||||
assert spec["output_types"] == [int]
|
||||
|
||||
def test_sync_function_can_handle(self):
|
||||
"""Test that sync functions work with can_handle method."""
|
||||
|
||||
@executor
|
||||
def string_handler(text: str):
|
||||
return text.strip()
|
||||
|
||||
assert string_handler.can_handle("hello")
|
||||
assert not string_handler.can_handle(123)
|
||||
assert not string_handler.can_handle([])
|
||||
|
||||
def test_sync_function_validation(self):
|
||||
"""Test validation for synchronous functions."""
|
||||
|
||||
# Valid sync function with one parameter
|
||||
def valid_sync(data: str):
|
||||
return data.upper()
|
||||
|
||||
func_exec = FunctionExecutor(valid_sync)
|
||||
assert str in func_exec._handlers
|
||||
|
||||
# Valid sync function with two parameters
|
||||
def valid_sync_with_ctx(data: int, ctx: WorkflowContext[str]):
|
||||
return str(data)
|
||||
|
||||
func_exec2 = FunctionExecutor(valid_sync_with_ctx)
|
||||
assert int in func_exec2._handlers
|
||||
|
||||
# Sync function with missing type annotation should still fail
|
||||
def no_annotation(data): # type: ignore
|
||||
return data
|
||||
|
||||
with pytest.raises(ValueError, match="type annotation for the message"):
|
||||
FunctionExecutor(no_annotation) # type: ignore
|
||||
|
||||
def test_mixed_sync_async_decorator(self):
|
||||
"""Test that both sync and async functions work with decorator."""
|
||||
|
||||
@executor
|
||||
def sync_func(data: str):
|
||||
return data.lower()
|
||||
|
||||
@executor
|
||||
async def async_func(data: str):
|
||||
return data.upper()
|
||||
|
||||
# Both should be FunctionExecutor instances
|
||||
assert isinstance(sync_func, FunctionExecutor)
|
||||
assert isinstance(async_func, FunctionExecutor)
|
||||
|
||||
# Both should handle strings
|
||||
assert sync_func.can_handle("test")
|
||||
assert async_func.can_handle("test")
|
||||
|
||||
# Both should be different instances
|
||||
assert sync_func is not async_func
|
||||
|
||||
async def test_sync_function_in_workflow(self):
|
||||
"""Test that sync functions work properly in a workflow context."""
|
||||
|
||||
@executor(id="sync_upper")
|
||||
def to_upper_sync(text: str, ctx: WorkflowContext[str]):
|
||||
return text.upper()
|
||||
# Note: For the test, we'll use a sync send mechanism
|
||||
# In practice, the wrapper handles the async conversion
|
||||
|
||||
@executor(id="async_reverse")
|
||||
async def reverse_async(text: str, ctx: WorkflowContext[Any]):
|
||||
result = text[::-1]
|
||||
await ctx.add_event(WorkflowCompletedEvent(result))
|
||||
|
||||
# Verify the executors can handle their input types
|
||||
assert to_upper_sync.can_handle("hello")
|
||||
assert reverse_async.can_handle("HELLO")
|
||||
|
||||
# For integration testing, we mainly verify that the handlers are properly registered
|
||||
# and the functions are wrapped correctly
|
||||
assert str in to_upper_sync._handlers
|
||||
assert str in reverse_async._handlers
|
||||
|
||||
async def test_sync_function_thread_execution(self):
|
||||
"""Test that sync functions run in thread pool and don't block the event loop."""
|
||||
import threading
|
||||
import time
|
||||
|
||||
_ = threading.get_ident()
|
||||
execution_thread_id = None
|
||||
|
||||
@executor
|
||||
def blocking_function(data: str):
|
||||
nonlocal execution_thread_id
|
||||
execution_thread_id = threading.get_ident()
|
||||
# Simulate some CPU-bound work
|
||||
time.sleep(0.01) # Small sleep to verify thread execution
|
||||
return data.upper()
|
||||
|
||||
# Verify the function is wrapped and registered
|
||||
assert str in blocking_function._handlers
|
||||
|
||||
# For a more complete test, we'd need to create a full workflow context,
|
||||
# but for now we can verify that the function was properly wrapped
|
||||
# and that sync functions store the correct metadata
|
||||
assert not blocking_function._is_async
|
||||
assert not blocking_function._has_context
|
||||
|
||||
# The actual thread execution test would require a full workflow setup,
|
||||
# but the important thing is that asyncio.to_thread is used in the wrapper
|
||||
Reference in New Issue
Block a user