Python: Clean up and formatting (#487)

* Clean up and formatting

* Fix mypy

* Bug fix
This commit is contained in:
Tao Chen
2025-08-26 13:21:32 -07:00
committed by GitHub
Unverified
parent 5df0bd6cb0
commit beb5218838
6 changed files with 196 additions and 187 deletions
@@ -77,6 +77,11 @@ class Edge(AFBaseModel):
return self._condition(data)
def _default_edge_list() -> list[Edge]:
"""Get the default list of edges for the group."""
return []
class EdgeGroup(AFBaseModel):
"""Represents a group of edges that share some common properties and can be triggered together."""
@@ -84,7 +89,7 @@ class EdgeGroup(AFBaseModel):
default_factory=lambda: f"EdgeGroup/{uuid.uuid4()}", description="Unique identifier for the edge group"
)
type: str = Field(description="The type of edge group, corresponding to the class name")
edges: list[Edge] = Field(default_factory=list, description="List of edges in this group")
edges: list[Edge] = Field(default_factory=_default_edge_list, description="List of edges in this group")
def __init__(self, **kwargs: Any) -> None:
"""Initialize the edge group."""
@@ -97,24 +102,12 @@ class EdgeGroup(AFBaseModel):
@property
def source_executor_ids(self) -> list[str]:
"""Get the source executor IDs of the edges in the group."""
seen = set()
result = []
for edge in self.edges:
if edge.source_id not in seen:
result.append(edge.source_id)
seen.add(edge.source_id)
return result
return list(dict.fromkeys(edge.source_id for edge in self.edges))
@property
def target_executor_ids(self) -> list[str]:
"""Get the target executor IDs of the edges in the group."""
seen = set()
result = []
for edge in self.edges:
if edge.target_id not in seen:
result.append(edge.target_id)
seen.add(edge.target_id)
return result
return list(dict.fromkeys(edge.target_id for edge in self.edges))
class SingleEdgeGroup(EdgeGroup):
@@ -275,6 +268,11 @@ class SwitchCaseEdgeGroupDefault(AFBaseModel):
type: str = Field(default="Default", description="The type of the case")
def _default_case_list() -> list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault]:
"""Get the default list of cases for the group."""
return []
class SwitchCaseEdgeGroup(FanOutEdgeGroup):
"""Represents a group of edges that assemble a conditional routing pattern.
@@ -299,7 +297,8 @@ class SwitchCaseEdgeGroup(FanOutEdgeGroup):
"""
cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault] = Field(
default_factory=list, description="List of conditional cases for this switch-case group"
default_factory=_default_case_list,
description="List of conditional cases for this switch-case group",
)
def __init__(
@@ -1,9 +1,12 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
from typing import TYPE_CHECKING, Any
from agent_framework import AgentRunResponse, AgentRunResponseUpdate
if TYPE_CHECKING:
from ._executor import RequestInfoMessage
class WorkflowEvent:
"""Base class for workflow events."""
@@ -61,7 +64,7 @@ class RequestInfoEvent(WorkflowEvent):
request_id: str,
source_executor_id: str,
request_type: type,
request_data: Any,
request_data: "RequestInfoMessage",
):
"""Initialize the request info event.
@@ -188,8 +188,10 @@ class Executor(AFBaseModel):
# Check if interceptor handled it or needs to forward
if isinstance(response, RequestResponse):
# Add automatic correlation info to the response
correlated_response = RequestResponse._with_correlation(
response, request.data, request.request_id
correlated_response = RequestResponse[RequestInfoMessage, Any].with_correlation(
response, # pyright: ignore[reportUnknownArgumentType]
request.data,
request.request_id,
)
if correlated_response.is_handled:
@@ -257,23 +259,8 @@ class Executor(AFBaseModel):
ExecutorT = TypeVar("ExecutorT", bound="Executor")
@overload
def handler(
func: Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]],
) -> Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]: ...
@overload
def handler(
func: None = None,
) -> Callable[
[Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]],
Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]],
]: ...
def handler(
func: Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]] | None = None,
) -> (
Callable[[ExecutorT, Any, WorkflowContext[Any]], Awaitable[Any]]
| Callable[
@@ -382,14 +369,24 @@ def handler(
return wrapper
if func is None:
return decorator
return decorator(func)
# endregion: Handler Decorator
# region Request/Response Types
@dataclass
class RequestInfoMessage:
"""Base class for all request messages in workflows.
Any message that should be routed to the RequestInfoExecutor for external
handling must inherit from this class. This ensures type safety and makes
the request/response pattern explicit.
"""
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
TRequest = TypeVar("TRequest", bound="RequestInfoMessage")
TResponse = TypeVar("TResponse")
@@ -411,7 +408,7 @@ class RequestResponse(Generic[TRequest, TResponse]):
request_id: str | None = None # Added for tracking
@classmethod
def handled(cls, data: TResponse) -> "RequestResponse[Any, TResponse]":
def handled(cls, data: TResponse) -> "RequestResponse[TRequest, TResponse]":
"""Create a response indicating the request was handled.
Correlation info (original_request, request_id) will be added automatically
@@ -420,13 +417,15 @@ class RequestResponse(Generic[TRequest, TResponse]):
return cls(is_handled=True, data=data)
@classmethod
def forward(cls, modified_request: Any = None) -> "RequestResponse[Any, Any]":
def forward(cls, modified_request: Any = None) -> "RequestResponse[TRequest, TResponse]":
"""Create a response indicating the request should be forwarded."""
return cls(is_handled=False, forward_request=modified_request)
@staticmethod
def _with_correlation(
original_response: "RequestResponse[Any, TResponse]", original_request: TRequest, request_id: str
def with_correlation(
original_response: "RequestResponse[TRequest, TResponse]",
original_request: TRequest,
request_id: str,
) -> "RequestResponse[TRequest, TResponse]":
"""Internal method to add correlation info to a response.
@@ -451,7 +450,7 @@ class SubWorkflowRequestInfo:
request_id: str # Original request ID from sub-workflow
sub_workflow_id: str # ID of the WorkflowExecutor that sent this
data: Any # The actual request data
data: RequestInfoMessage # The actual request data
@dataclass
@@ -595,110 +594,9 @@ def intercepts_request(
# endregion: Intercepts Request Decorator
# region Agent Executor
@dataclass
class AgentExecutorRequest:
"""A request to an agent executor.
Attributes:
messages: A list of chat messages to be processed by the agent.
should_respond: A flag indicating whether the agent should respond to the messages.
If False, the messages will be saved to the executor's cache but not sent to the agent.
"""
messages: list[ChatMessage]
should_respond: bool = True
@dataclass
class AgentExecutorResponse:
"""A response from an agent executor.
Attributes:
executor_id: The ID of the executor that generated the response.
response: The agent run response containing the messages generated by the agent.
"""
executor_id: str
agent_run_response: AgentRunResponse
class AgentExecutor(Executor):
"""built-in executor that wraps an agent for handling messages."""
def __init__(
self,
agent: AIAgent,
*,
agent_thread: AgentThread | None = None,
streaming: bool = False,
id: str | None = None,
):
"""Initialize the executor with a unique identifier.
Args:
agent: The agent to be wrapped by this executor.
agent_thread: The thread to use for running the agent. If None, a new thread will be created.
streaming: Whether to enable streaming for the agent. If enabled, the executor will emit
AgentRunStreamingEvent updates instead of a single AgentRunEvent.
id: A unique identifier for the executor. If None, a new UUID will be generated.
"""
super().__init__(id or agent.id)
self._agent = agent
self._agent_thread = agent_thread or self._agent.get_new_thread()
self._streaming = streaming
self._cache: list[ChatMessage] = []
@handler
async def run(self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorResponse]) -> None:
"""Run the agent executor with the given request."""
self._cache.extend(request.messages)
if request.should_respond:
if self._streaming:
updates: list[AgentRunResponseUpdate] = []
async for update in self._agent.run_streaming(
self._cache,
thread=self._agent_thread,
):
updates.append(update)
await ctx.add_event(AgentRunStreamingEvent(self.id, update))
response = AgentRunResponse.from_agent_run_response_updates(updates)
else:
response = await self._agent.run(
self._cache,
thread=self._agent_thread,
)
await ctx.add_event(AgentRunEvent(self.id, response))
await ctx.send_message(AgentExecutorResponse(self.id, response))
self._cache.clear()
# endregion: Agent Executor
# region Request Info Executor
@dataclass
class RequestInfoMessage:
"""Base class for all request messages in workflows.
Any message that should be routed to the RequestInfoExecutor for external
handling must inherit from this class. This ensures type safety and makes
the request/response pattern explicit.
"""
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
# Note: SubWorkflowRequestInfo, SubWorkflowResponse, and RequestResponse
# have been moved before intercepts_request decorator
class RequestInfoExecutor(Executor):
"""Built-in executor that handles request/response patterns in workflows.
@@ -798,14 +696,105 @@ class RequestInfoExecutor(Executor):
else:
# Regular response - send directly back to source
# Create a correlated response that includes both the response data and original request
correlated_response = RequestResponse.handled(response_data)
correlated_response = RequestResponse._with_correlation(correlated_response, event.data, request_id)
if not isinstance(event.data, RequestInfoMessage):
raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}")
correlated_response = RequestResponse[RequestInfoMessage, Any].handled(response_data)
correlated_response = RequestResponse[RequestInfoMessage, Any].with_correlation(
correlated_response,
event.data,
request_id,
)
await ctx.send_message(correlated_response, target_id=event.source_executor_id)
# endregion: Request Info Executor
# region Agent Executor
@dataclass
class AgentExecutorRequest:
"""A request to an agent executor.
Attributes:
messages: A list of chat messages to be processed by the agent.
should_respond: A flag indicating whether the agent should respond to the messages.
If False, the messages will be saved to the executor's cache but not sent to the agent.
"""
messages: list[ChatMessage]
should_respond: bool = True
@dataclass
class AgentExecutorResponse:
"""A response from an agent executor.
Attributes:
executor_id: The ID of the executor that generated the response.
response: The agent run response containing the messages generated by the agent.
"""
executor_id: str
agent_run_response: AgentRunResponse
class AgentExecutor(Executor):
"""built-in executor that wraps an agent for handling messages."""
def __init__(
self,
agent: AIAgent,
*,
agent_thread: AgentThread | None = None,
streaming: bool = False,
id: str | None = None,
):
"""Initialize the executor with a unique identifier.
Args:
agent: The agent to be wrapped by this executor.
agent_thread: The thread to use for running the agent. If None, a new thread will be created.
streaming: Whether to enable streaming for the agent. If enabled, the executor will emit
AgentRunStreamingEvent updates instead of a single AgentRunEvent.
id: A unique identifier for the executor. If None, a new UUID will be generated.
"""
super().__init__(id or agent.id)
self._agent = agent
self._agent_thread = agent_thread or self._agent.get_new_thread()
self._streaming = streaming
self._cache: list[ChatMessage] = []
@handler
async def run(self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorResponse]) -> None:
"""Run the agent executor with the given request."""
self._cache.extend(request.messages)
if request.should_respond:
if self._streaming:
updates: list[AgentRunResponseUpdate] = []
async for update in self._agent.run_streaming(
self._cache,
thread=self._agent_thread,
):
updates.append(update)
await ctx.add_event(AgentRunStreamingEvent(self.id, update))
response = AgentRunResponse.from_agent_run_response_updates(updates)
else:
response = await self._agent.run(
self._cache,
thread=self._agent_thread,
)
await ctx.add_event(AgentRunEvent(self.id, response))
await ctx.send_message(AgentExecutorResponse(self.id, response))
self._cache.clear()
# endregion: Agent Executor
# region Workflow Executor
@@ -889,6 +878,8 @@ class WorkflowExecutor(Executor):
self._pending_requests[event.request_id] = event.data
# Wrap request with routing context and send to parent
if not isinstance(event.data, RequestInfoMessage):
raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}")
wrapped_request = SubWorkflowRequestInfo(
request_id=event.request_id,
sub_workflow_id=self.id,
@@ -957,6 +948,8 @@ class WorkflowExecutor(Executor):
self._pending_requests[event.request_id] = event.data
# Send the new request to parent
if not isinstance(event.data, RequestInfoMessage):
raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}")
wrapped_request = SubWorkflowRequestInfo(
request_id=event.request_id,
sub_workflow_id=self.id,