mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Clean up and formatting (#487)
* Clean up and formatting * Fix mypy * Bug fix
This commit is contained in:
committed by
GitHub
Unverified
parent
5df0bd6cb0
commit
beb5218838
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user