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

Python: [BREAKING] Move workflow to main package (#767)

Browse Source 
* Move workflow to main package

* Remove workflow specific unit test config

* Remove workflow-specific version info

* Revert unintended telemetry changes

* Removed the obsolete packages/workflow/tests target

* Rename dir workflow to _workflow

* Fix test imports
This commit is contained in:
Evan Mattson
2025-09-16 21:04:07 +09:00
committed by GitHub
Unverified
parent 65dd48aa1d
commit 732d9f6cd7
87 changed files with 524 additions and 595 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/agent_framework/__init__.py
+1
View File
@@ -16,3 +16,4 @@ from ._memory import * # noqa: F403
from ._threads import * # noqa: F403
from ._tools import * # noqa: F403
from ._types import * # noqa: F403
from ._workflow import * # noqa: F403
python/packages/main/agent_framework/_workflow/README.md
+15
View File
@@ -0,0 +1,15 @@
# Get Started with Microsoft Agent Framework Workflow
Workflow capabilities now ship with the core `agent-framework` package.
```bash
pip install agent-framework
```
Optional visualization support is still available via the `viz` extra:
```bash
pip install agent-framework[viz]
```
See the [project README](https://github.com/microsoft/agent-framework/tree/main/python/README.md) for more information.
python/packages/main/agent_framework/_workflow/__init__.py
+197
View File
@@ -0,0 +1,197 @@
# Copyright (c) Microsoft. All rights reserved.
import contextlib
from ._agent import WorkflowAgent
from ._checkpoint import (
CheckpointStorage,
FileCheckpointStorage,
InMemoryCheckpointStorage,
WorkflowCheckpoint,
)
from ._concurrent import ConcurrentBuilder
from ._const import (
DEFAULT_MAX_ITERATIONS,
)
from ._edge import (
Case,
Default,
Edge,
FanInEdgeGroup,
FanOutEdgeGroup,
SingleEdgeGroup,
SwitchCaseEdgeGroup,
SwitchCaseEdgeGroupCase,
SwitchCaseEdgeGroupDefault,
)
from ._edge_runner import create_edge_runner
from ._events import (
AgentRunEvent,
AgentRunUpdateEvent,
ExecutorCompletedEvent,
ExecutorEvent,
ExecutorFailedEvent,
ExecutorInvokeEvent,
RequestInfoEvent,
WorkflowCompletedEvent,
WorkflowErrorDetails,
WorkflowEvent,
WorkflowFailedEvent,
WorkflowRunState,
WorkflowStartedEvent,
WorkflowStatusEvent,
)
from ._executor import (
AgentExecutor,
AgentExecutorRequest,
AgentExecutorResponse,
Executor,
RequestInfoExecutor,
RequestInfoMessage,
RequestResponse,
SubWorkflowRequestInfo,
SubWorkflowResponse,
WorkflowExecutor,
handler,
intercepts_request,
)
from ._function_executor import FunctionExecutor, executor
from ._magentic import (
MagenticAgentDeltaEvent,
MagenticAgentExecutor,
MagenticAgentMessageEvent,
MagenticBuilder,
MagenticCallbackEvent,
MagenticCallbackMode,
MagenticContext,
MagenticFinalResultEvent,
MagenticManagerBase,
MagenticOrchestratorExecutor,
MagenticOrchestratorMessageEvent,
MagenticPlanReviewDecision,
MagenticPlanReviewReply,
MagenticPlanReviewRequest,
MagenticProgressLedger,
MagenticProgressLedgerItem,
MagenticRequestMessage,
MagenticResponseMessage,
MagenticStartMessage,
StandardMagenticManager,
)
from ._runner import Runner
from ._runner_context import (
InProcRunnerContext,
Message,
RunnerContext,
)
from ._sequential import SequentialBuilder
from ._shared_state import SharedState
from ._telemetry import EdgeGroupDeliveryStatus, WorkflowTracer, workflow_tracer
from ._validation import (
EdgeDuplicationError,
GraphConnectivityError,
HandlerOutputAnnotationError,
TypeCompatibilityError,
ValidationTypeEnum,
WorkflowValidationError,
validate_workflow_graph,
)
from ._viz import WorkflowViz
from ._workflow import Workflow, WorkflowBuilder, WorkflowRunResult
from ._workflow_context import WorkflowContext
__all__ = [
"DEFAULT_MAX_ITERATIONS",
"AgentExecutor",
"AgentExecutorRequest",
"AgentExecutorResponse",
"AgentRunEvent",
"AgentRunUpdateEvent",
"Case",
"CheckpointStorage",
"ConcurrentBuilder",
"Default",
"Edge",
"EdgeDuplicationError",
"EdgeGroupDeliveryStatus",
"Executor",
"ExecutorCompletedEvent",
"ExecutorEvent",
"ExecutorFailedEvent",
"ExecutorInvokeEvent",
"FanInEdgeGroup",
"FanOutEdgeGroup",
"FileCheckpointStorage",
"FunctionExecutor",
"GraphConnectivityError",
"HandlerOutputAnnotationError",
"InMemoryCheckpointStorage",
"InProcRunnerContext",
"MagenticAgentDeltaEvent",
"MagenticAgentExecutor",
"MagenticAgentMessageEvent",
"MagenticBuilder",
"MagenticCallbackEvent",
"MagenticCallbackMode",
"MagenticContext",
"MagenticFinalResultEvent",
"MagenticManagerBase",
"MagenticOrchestratorExecutor",
"MagenticOrchestratorMessageEvent",
"MagenticPlanReviewDecision",
"MagenticPlanReviewReply",
"MagenticPlanReviewRequest",
"MagenticProgressLedger",
"MagenticProgressLedgerItem",
"MagenticRequestMessage",
"MagenticResponseMessage",
"MagenticStartMessage",
"Message",
"RequestInfoEvent",
"RequestInfoExecutor",
"RequestInfoMessage",
"RequestResponse",
"Runner",
"RunnerContext",
"SequentialBuilder",
"SharedState",
"SingleEdgeGroup",
"StandardMagenticManager",
"SubWorkflowRequestInfo",
"SubWorkflowResponse",
"SwitchCaseEdgeGroup",
"SwitchCaseEdgeGroupCase",
"SwitchCaseEdgeGroupDefault",
"TypeCompatibilityError",
"ValidationTypeEnum",
"Workflow",
"WorkflowAgent",
"WorkflowBuilder",
"WorkflowCheckpoint",
"WorkflowCompletedEvent",
"WorkflowContext",
"WorkflowErrorDetails",
"WorkflowEvent",
"WorkflowExecutor",
"WorkflowFailedEvent",
"WorkflowRunResult",
"WorkflowRunState",
"WorkflowStartedEvent",
"WorkflowStatusEvent",
"WorkflowTracer",
"WorkflowValidationError",
"WorkflowViz",
"create_edge_runner",
"executor",
"handler",
"intercepts_request",
"validate_workflow_graph",
"workflow_tracer",
]
# Rebuild models to resolve forward references after all imports are complete
with contextlib.suppress(AttributeError, TypeError, ValueError):
# Rebuild WorkflowExecutor to resolve Workflow forward reference
WorkflowExecutor.model_rebuild()
# Rebuild WorkflowAgent to resolve Workflow forward reference
WorkflowAgent.model_rebuild()
python/packages/main/agent_framework/_workflow/__init__.pyi
+186
View File
@@ -0,0 +1,186 @@
# Copyright (c) Microsoft. All rights reserved.
from ._agent import WorkflowAgent
from ._checkpoint import (
CheckpointStorage,
FileCheckpointStorage,
InMemoryCheckpointStorage,
WorkflowCheckpoint,
)
from ._concurrent import ConcurrentBuilder
from ._const import DEFAULT_MAX_ITERATIONS
from ._edge import (
Case,
Default,
Edge,
FanInEdgeGroup,
FanOutEdgeGroup,
SingleEdgeGroup,
SwitchCaseEdgeGroup,
SwitchCaseEdgeGroupCase,
SwitchCaseEdgeGroupDefault,
)
from ._edge_runner import create_edge_runner
from ._events import (
AgentRunEvent,
AgentRunUpdateEvent,
ExecutorCompletedEvent,
ExecutorEvent,
ExecutorFailedEvent,
ExecutorInvokeEvent,
RequestInfoEvent,
WorkflowCompletedEvent,
WorkflowErrorDetails,
WorkflowEvent,
WorkflowFailedEvent,
WorkflowRunState,
WorkflowStartedEvent,
WorkflowStatusEvent,
)
from ._executor import (
AgentExecutor,
AgentExecutorRequest,
AgentExecutorResponse,
Executor,
RequestInfoExecutor,
RequestInfoMessage,
RequestResponse,
SubWorkflowRequestInfo,
SubWorkflowResponse,
WorkflowExecutor,
handler,
intercepts_request,
)
from ._function_executor import FunctionExecutor, executor
from ._magentic import (
MagenticAgentDeltaEvent,
MagenticAgentExecutor,
MagenticAgentMessageEvent,
MagenticBuilder,
MagenticCallbackEvent,
MagenticCallbackMode,
MagenticContext,
MagenticFinalResultEvent,
MagenticManagerBase,
MagenticOrchestratorExecutor,
MagenticOrchestratorMessageEvent,
MagenticPlanReviewDecision,
MagenticPlanReviewReply,
MagenticPlanReviewRequest,
MagenticProgressLedger,
MagenticProgressLedgerItem,
MagenticRequestMessage,
MagenticResponseMessage,
MagenticStartMessage,
StandardMagenticManager,
)
from ._runner import Runner
from ._runner_context import (
InProcRunnerContext,
Message,
RunnerContext,
)
from ._sequential import SequentialBuilder
from ._shared_state import SharedState
from ._telemetry import EdgeGroupDeliveryStatus, WorkflowTracer, workflow_tracer
from ._validation import (
EdgeDuplicationError,
GraphConnectivityError,
HandlerOutputAnnotationError,
TypeCompatibilityError,
ValidationTypeEnum,
WorkflowValidationError,
validate_workflow_graph,
)
from ._viz import WorkflowViz
from ._workflow import Workflow, WorkflowBuilder, WorkflowRunResult
from ._workflow_context import WorkflowContext
__all__ = [
"DEFAULT_MAX_ITERATIONS",
"AgentExecutor",
"AgentExecutorRequest",
"AgentExecutorResponse",
"AgentRunEvent",
"AgentRunUpdateEvent",
"Case",
"CheckpointStorage",
"ConcurrentBuilder",
"Default",
"Edge",
"EdgeDuplicationError",
"EdgeGroupDeliveryStatus",
"Executor",
"ExecutorCompletedEvent",
"ExecutorEvent",
"ExecutorFailedEvent",
"ExecutorInvokeEvent",
"FanInEdgeGroup",
"FanOutEdgeGroup",
"FileCheckpointStorage",
"FunctionExecutor",
"GraphConnectivityError",
"HandlerOutputAnnotationError",
"InMemoryCheckpointStorage",
"InProcRunnerContext",
"MagenticAgentDeltaEvent",
"MagenticAgentExecutor",
"MagenticAgentMessageEvent",
"MagenticBuilder",
"MagenticCallbackEvent",
"MagenticCallbackMode",
"MagenticContext",
"MagenticFinalResultEvent",
"MagenticManagerBase",
"MagenticOrchestratorExecutor",
"MagenticOrchestratorMessageEvent",
"MagenticPlanReviewDecision",
"MagenticPlanReviewReply",
"MagenticPlanReviewRequest",
"MagenticProgressLedger",
"MagenticProgressLedgerItem",
"MagenticRequestMessage",
"MagenticResponseMessage",
"MagenticStartMessage",
"Message",
"RequestInfoEvent",
"RequestInfoExecutor",
"RequestInfoMessage",
"RequestResponse",
"Runner",
"RunnerContext",
"SequentialBuilder",
"SharedState",
"SingleEdgeGroup",
"StandardMagenticManager",
"SubWorkflowRequestInfo",
"SubWorkflowResponse",
"SwitchCaseEdgeGroup",
"SwitchCaseEdgeGroupCase",
"SwitchCaseEdgeGroupDefault",
"TypeCompatibilityError",
"ValidationTypeEnum",
"Workflow",
"WorkflowAgent",
"WorkflowBuilder",
"WorkflowCheckpoint",
"WorkflowCompletedEvent",
"WorkflowContext",
"WorkflowErrorDetails",
"WorkflowEvent",
"WorkflowExecutor",
"WorkflowFailedEvent",
"WorkflowRunResult",
"WorkflowRunState",
"WorkflowStartedEvent",
"WorkflowStatusEvent",
"WorkflowTracer",
"WorkflowValidationError",
"WorkflowViz",
"create_edge_runner",
"executor",
"handler",
"intercepts_request",
"validate_workflow_graph",
"workflow_tracer",
]
python/packages/main/agent_framework/_workflow/_agent.py
+454
View File
@@ -0,0 +1,454 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
import uuid
from collections.abc import AsyncIterable, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, cast
from pydantic import Field
from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
BaseAgent,
ChatMessage,
FunctionCallContent,
FunctionResultContent,
Role,
TextContent,
UsageDetails,
)
from agent_framework._pydantic import AFBaseModel
from agent_framework.exceptions import AgentExecutionException
from ._events import (
AgentRunUpdateEvent,
RequestInfoEvent,
WorkflowEvent,
)
if TYPE_CHECKING:
from ._workflow import Workflow
logger = logging.getLogger(__name__)
class WorkflowAgent(BaseAgent):
"""An `Agent` subclass that wraps a workflow and exposes it as an agent."""
# Class variable for the request info function name
REQUEST_INFO_FUNCTION_NAME: ClassVar[str] = "request_info"
class RequestInfoFunctionArgs(AFBaseModel):
request_id: str
data: Any
workflow: "Workflow" = Field(description="The workflow wrapped as an agent")
pending_requests: dict[str, RequestInfoEvent] = Field(
default_factory=dict, description="Pending request info events"
)
def __init__(
self,
workflow: "Workflow",
*,
id: str | None = None,
name: str | None = None,
description: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize the WorkflowAgent.
Args:
workflow: The workflow to wrap as an agent.
id: Unique identifier for the agent. If None, will be generated.
name: Optional name for the agent.
description: Optional description of the agent.
**kwargs: Additional keyword arguments passed to BaseAgent.
"""
if id is None:
id = f"WorkflowAgent_{uuid.uuid4().hex[:8]}"
# Initialize with standard BaseAgent parameters first
kwargs["workflow"] = workflow
# Validate the workflow's start executor can handle agent-facing message inputs
try:
start_executor = workflow.get_start_executor()
except KeyError as exc: # Defensive: workflow lacks a configured entry point
raise ValueError("Workflow's start executor is not defined.") from exc
if not start_executor.can_handle_type(list[ChatMessage]):
raise ValueError("Workflow's start executor cannot handle list[ChatMessage]")
super().__init__(id=id, name=name, description=description, **kwargs)
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentRunResponse:
"""Get a response from the workflow agent (non-streaming).
This method collects all streaming updates and merges them into a single response.
Args:
messages: The message(s) to send to the workflow.
thread: The conversation thread. If None, a new thread will be created.
**kwargs: Additional keyword arguments.
Returns:
The final workflow response as an AgentRunResponse.
"""
# Collect all streaming updates
response_updates: list[AgentRunResponseUpdate] = []
input_messages = self._normalize_messages(messages)
thread = thread or self.get_new_thread()
response_id = str(uuid.uuid4())
async for update in self._run_stream_impl(input_messages, response_id):
response_updates.append(update)
# Convert updates to final response.
response = self.merge_updates(response_updates, response_id)
# Notify thread of new messages (both input and response messages)
await self._notify_thread_of_new_messages(thread, input_messages)
await self._notify_thread_of_new_messages(thread, response.messages)
return response
async def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Stream response updates from the workflow agent.
Args:
messages: The message(s) to send to the workflow.
thread: The conversation thread. If None, a new thread will be created.
**kwargs: Additional keyword arguments.
Yields:
AgentRunResponseUpdate objects representing the workflow execution progress.
"""
input_messages = self._normalize_messages(messages)
thread = thread or self.get_new_thread()
response_updates: list[AgentRunResponseUpdate] = []
response_id = str(uuid.uuid4())
async for update in self._run_stream_impl(input_messages, response_id):
response_updates.append(update)
yield update
# Convert updates to final response.
response = self.merge_updates(response_updates, response_id)
# Notify thread of new messages (both input and response messages)
await self._notify_thread_of_new_messages(thread, input_messages)
await self._notify_thread_of_new_messages(thread, response.messages)
async def _run_stream_impl(
self,
input_messages: list[ChatMessage],
response_id: str,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Internal implementation of streaming execution.
Args:
input_messages: Normalized input messages to process.
response_id: The unique response ID for this workflow execution.
Yields:
AgentRunResponseUpdate objects representing the workflow execution progress.
"""
# Determine the event stream based on whether we have function responses
if bool(self.pending_requests):
# This is a continuation - use send_responses_streaming to send function responses back
logger.info(f"Continuing workflow to address {len(self.pending_requests)} requests")
# Extract function responses from input messages, and ensure that
# only function responses are present in messages if there is any
# pending request.
function_responses = self._extract_function_responses(input_messages)
# Pop pending requests if fulfilled.
for request_id in list(self.pending_requests.keys()):
if request_id in function_responses:
self.pending_requests.pop(request_id)
# NOTE: It is possible that some pending requests are not fulfilled,
# and we will let the workflow to handle this -- the agent does not
# have an opinion on this.
event_stream = self.workflow.send_responses_streaming(function_responses)
else:
# Execute workflow with streaming (initial run or no function responses)
# Pass the new input messages directly to the workflow
event_stream = self.workflow.run_stream(input_messages)
# Process events from the stream
async for event in event_stream:
# Convert workflow event to agent update
update = self._convert_workflow_event_to_agent_update(response_id, event)
if update:
yield update
def _normalize_messages(
self,
messages: str | ChatMessage | Sequence[str] | Sequence[ChatMessage] | None = None,
) -> list[ChatMessage]:
"""Normalize input messages to a list of ChatMessage objects."""
if messages is None:
return []
if isinstance(messages, str):
return [ChatMessage(role=Role.USER, contents=[TextContent(text=messages)])]
if isinstance(messages, ChatMessage):
return [messages]
normalized: list[ChatMessage] = []
for msg in messages:
if isinstance(msg, str):
normalized.append(ChatMessage(role=Role.USER, contents=[TextContent(text=msg)]))
elif isinstance(msg, ChatMessage):
normalized.append(msg)
return normalized
def _convert_workflow_event_to_agent_update(
self,
response_id: str,
event: WorkflowEvent,
) -> AgentRunResponseUpdate | None:
"""Convert a workflow event to an AgentRunResponseUpdate.
Only AgentRunUpdateEvent and RequestInfoEvent are processed and the rest
are not relevant. Returns None if the event is not relevant.
"""
match event:
case AgentRunUpdateEvent(data=update):
# Direct pass-through of update in an agent streaming event
if update:
return cast(AgentRunResponseUpdate, update)
return None
case RequestInfoEvent(request_id=request_id):
# Store the pending request for later correlation
self.pending_requests[request_id] = event
# Convert to function call content
# TODO(ekzhu): update this to FunctionApprovalRequestContent
# monitor: https://github.com/microsoft/agent-framework/issues/285
function_call = FunctionCallContent(
call_id=request_id,
name=self.REQUEST_INFO_FUNCTION_NAME,
arguments=self.RequestInfoFunctionArgs(request_id=request_id, data=event.data).model_dump(),
)
return AgentRunResponseUpdate(
contents=[function_call],
role=Role.ASSISTANT,
author_name=self.name,
response_id=response_id,
message_id=str(uuid.uuid4()),
created_at=datetime.now().strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
)
case _:
# Ignore non-agent workflow events
pass
# We only care about the above two events and discard the rest.
return None
def _extract_function_responses(self, input_messages: list[ChatMessage]) -> dict[str, Any]:
"""Extract function responses from input messages."""
function_responses: dict[str, Any] = {}
for message in input_messages:
for content in message.contents:
# TODO(ekzhu): update this to FunctionApprovalResponseContent
# monitor: https://github.com/microsoft/agent-framework/issues/285
if isinstance(content, FunctionResultContent):
request_id = content.call_id
# Check if we have a pending request for this call_id
if request_id in self.pending_requests:
response_data = content.result if hasattr(content, "result") else str(content)
function_responses[request_id] = response_data
elif bool(self.pending_requests):
# Function result for unknown request when we have pending requests - this is an error
raise AgentExecutionException(
"Only FunctionResultContent for pending requests is allowed in input messages "
"when there are pending requests."
)
else:
if bool(self.pending_requests):
# Non-function content when we have pending requests - this is an error
raise AgentExecutionException(
"Only FunctionResultContent is allowed in input messages when there are pending requests."
)
return function_responses
class _ResponseState(TypedDict):
"""State for grouping response updates by message_id."""
by_msg: dict[str, list[AgentRunResponseUpdate]]
dangling: list[AgentRunResponseUpdate]
@staticmethod
def merge_updates(updates: list[AgentRunResponseUpdate], response_id: str) -> AgentRunResponse:
"""Merge streaming updates into a single AgentRunResponse.
Behavior:
- Group updates by response_id; within each response_id, group by message_id and keep a dangling bucket for
updates without message_id.
- Convert each group (per message and dangling) into an intermediate AgentRunResponse via
AgentRunResponse.from_agent_run_response_updates, then sort by created_at and merge.
- Append messages from updates without any response_id at the end (global dangling), while aggregating metadata.
Args:
updates: The list of AgentRunResponseUpdate objects to merge.
response_id: The response identifier to set on the returned AgentRunResponse.
Returns:
An AgentRunResponse with messages in processing order and aggregated metadata.
"""
# PHASE 1: GROUP UPDATES BY RESPONSE_ID AND MESSAGE_ID
states: dict[str, WorkflowAgent._ResponseState] = {}
global_dangling: list[AgentRunResponseUpdate] = []
for u in updates:
if u.response_id:
state = states.setdefault(u.response_id, {"by_msg": {}, "dangling": []})
by_msg = state["by_msg"]
dangling = state["dangling"]
if u.message_id:
by_msg.setdefault(u.message_id, []).append(u)
else:
dangling.append(u)
else:
global_dangling.append(u)
# HELPER FUNCTIONS
def _parse_dt(value: str | None) -> tuple[int, datetime | str | None]:
if not value:
return (1, None)
v = value
if v.endswith("Z"):
v = v[:-1] + "+00:00"
try:
return (0, datetime.fromisoformat(v))
except Exception:
return (0, v)
def _sum_usage(a: UsageDetails | None, b: UsageDetails | None) -> UsageDetails | None:
if a is None:
return b
if b is None:
return a
return a + b
def _merge_responses(current: AgentRunResponse | None, incoming: AgentRunResponse) -> AgentRunResponse:
if current is None:
return incoming
raw_list: list[object] = []
def _add_raw(value: object) -> None:
if isinstance(value, list):
raw_list.extend(cast(list[object], value))
else:
raw_list.append(value)
if current.raw_representation is not None:
_add_raw(current.raw_representation)
if incoming.raw_representation is not None:
_add_raw(incoming.raw_representation)
return AgentRunResponse(
messages=(current.messages or []) + (incoming.messages or []),
response_id=current.response_id or incoming.response_id,
created_at=incoming.created_at or current.created_at,
usage_details=_sum_usage(current.usage_details, incoming.usage_details),
raw_representation=raw_list if raw_list else None,
additional_properties=incoming.additional_properties or current.additional_properties,
)
# PHASE 2: CONVERT GROUPED UPDATES TO RESPONSES AND MERGE
final_messages: list[ChatMessage] = []
merged_usage: UsageDetails | None = None
latest_created_at: str | None = None
merged_additional_properties: dict[str, Any] | None = None
raw_representations: list[object] = []
for grouped_response_id in states:
state = states[grouped_response_id]
by_msg = state["by_msg"]
dangling = state["dangling"]
per_message_responses: list[AgentRunResponse] = []
for _, msg_updates in by_msg.items():
if msg_updates:
per_message_responses.append(AgentRunResponse.from_agent_run_response_updates(msg_updates))
if dangling:
per_message_responses.append(AgentRunResponse.from_agent_run_response_updates(dangling))
per_message_responses.sort(key=lambda r: _parse_dt(r.created_at))
aggregated: AgentRunResponse | None = None
for resp in per_message_responses:
if resp.response_id and grouped_response_id and resp.response_id != grouped_response_id:
resp.response_id = grouped_response_id
aggregated = _merge_responses(aggregated, resp)
if aggregated:
final_messages.extend(aggregated.messages)
if aggregated.usage_details:
merged_usage = _sum_usage(merged_usage, aggregated.usage_details)
if aggregated.created_at and (
not latest_created_at or _parse_dt(aggregated.created_at) > _parse_dt(latest_created_at)
):
latest_created_at = aggregated.created_at
if aggregated.additional_properties:
if merged_additional_properties is None:
merged_additional_properties = {}
merged_additional_properties.update(aggregated.additional_properties)
raw_value = aggregated.raw_representation
if raw_value:
cast_value = cast(object | list[object], raw_value)
if isinstance(cast_value, list):
raw_representations.extend(cast(list[object], cast_value))
else:
raw_representations.append(cast_value)
# PHASE 3: HANDLE GLOBAL DANGLING UPDATES (NO RESPONSE_ID)
if global_dangling:
flattened = AgentRunResponse.from_agent_run_response_updates(global_dangling)
final_messages.extend(flattened.messages)
if flattened.usage_details:
merged_usage = _sum_usage(merged_usage, flattened.usage_details)
if flattened.created_at and (
not latest_created_at or _parse_dt(flattened.created_at) > _parse_dt(latest_created_at)
):
latest_created_at = flattened.created_at
if flattened.additional_properties:
if merged_additional_properties is None:
merged_additional_properties = {}
merged_additional_properties.update(flattened.additional_properties)
flat_raw = flattened.raw_representation
if flat_raw:
cast_flat = cast(object | list[object], flat_raw)
if isinstance(cast_flat, list):
raw_representations.extend(cast(list[object], cast_flat))
else:
raw_representations.append(cast_flat)
# PHASE 4: CONSTRUCT FINAL RESPONSE WITH INPUT RESPONSE_ID
return AgentRunResponse(
messages=final_messages,
response_id=response_id,
created_at=latest_created_at,
usage_details=merged_usage,
raw_representation=raw_representations if raw_representations else None,
additional_properties=merged_additional_properties,
)
python/packages/main/agent_framework/_workflow/_checkpoint.py
+199
View File
@@ -0,0 +1,199 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import json
import logging
import os
import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Protocol
from ._const import DEFAULT_MAX_ITERATIONS
logger = logging.getLogger(__name__)
@dataclass(slots=True)
class WorkflowCheckpoint:
"""Represents a complete checkpoint of workflow state."""
checkpoint_id: str = field(default_factory=lambda: str(uuid.uuid4()))
workflow_id: str = ""
timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
# Core workflow state
messages: dict[str, list[dict[str, Any]]] = field(default_factory=dict) # type: ignore[misc]
shared_state: dict[str, Any] = field(default_factory=dict) # type: ignore[misc]
executor_states: dict[str, dict[str, Any]] = field(default_factory=dict) # type: ignore[misc]
# Runtime state
iteration_count: int = 0
max_iterations: int = DEFAULT_MAX_ITERATIONS
# Metadata
metadata: dict[str, Any] = field(default_factory=dict) # type: ignore[misc]
version: str = "1.0"
def to_dict(self) -> dict[str, Any]:
return asdict(self)
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "WorkflowCheckpoint":
return cls(**data)
class CheckpointStorage(Protocol):
"""Protocol for checkpoint storage backends."""
async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str:
"""Save a checkpoint and return its ID."""
...
async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
"""Load a checkpoint by ID."""
...
async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]:
"""List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
...
async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]:
"""List checkpoint objects. If workflow_id is provided, filter by that workflow."""
...
async def delete_checkpoint(self, checkpoint_id: str) -> bool:
"""Delete a checkpoint by ID."""
...
class InMemoryCheckpointStorage:
"""In-memory checkpoint storage for testing and development."""
def __init__(self) -> None:
"""Initialize the memory storage."""
self._checkpoints: dict[str, WorkflowCheckpoint] = {}
async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str:
"""Save a checkpoint and return its ID."""
self._checkpoints[checkpoint.checkpoint_id] = checkpoint
logger.debug(f"Saved checkpoint {checkpoint.checkpoint_id} to memory")
return checkpoint.checkpoint_id
async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
"""Load a checkpoint by ID."""
checkpoint = self._checkpoints.get(checkpoint_id)
if checkpoint:
logger.debug(f"Loaded checkpoint {checkpoint_id} from memory")
return checkpoint
async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]:
"""List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
if workflow_id is None:
return list(self._checkpoints.keys())
return [cp.checkpoint_id for cp in self._checkpoints.values() if cp.workflow_id == workflow_id]
async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]:
"""List checkpoint objects. If workflow_id is provided, filter by that workflow."""
if workflow_id is None:
return list(self._checkpoints.values())
return [cp for cp in self._checkpoints.values() if cp.workflow_id == workflow_id]
async def delete_checkpoint(self, checkpoint_id: str) -> bool:
"""Delete a checkpoint by ID."""
if checkpoint_id in self._checkpoints:
del self._checkpoints[checkpoint_id]
logger.debug(f"Deleted checkpoint {checkpoint_id} from memory")
return True
return False
class FileCheckpointStorage:
"""File-based checkpoint storage for persistence."""
def __init__(self, storage_path: str | Path):
"""Initialize the file storage."""
self.storage_path = Path(storage_path)
self.storage_path.mkdir(parents=True, exist_ok=True)
logger.info(f"Initialized file checkpoint storage at {self.storage_path}")
async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str:
"""Save a checkpoint and return its ID."""
file_path = self.storage_path / f"{checkpoint.checkpoint_id}.json"
checkpoint_dict = asdict(checkpoint)
def _write_atomic() -> None:
tmp_path = file_path.with_suffix(".json.tmp")
with open(tmp_path, "w") as f:
json.dump(checkpoint_dict, f, indent=2, ensure_ascii=False)
os.replace(tmp_path, file_path)
await asyncio.to_thread(_write_atomic)
logger.info(f"Saved checkpoint {checkpoint.checkpoint_id} to {file_path}")
return checkpoint.checkpoint_id
async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
"""Load a checkpoint by ID."""
file_path = self.storage_path / f"{checkpoint_id}.json"
if not file_path.exists():
return None
def _read() -> dict[str, Any]:
with open(file_path) as f:
return json.load(f) # type: ignore[no-any-return]
checkpoint_dict = await asyncio.to_thread(_read)
checkpoint = WorkflowCheckpoint(**checkpoint_dict)
logger.info(f"Loaded checkpoint {checkpoint_id} from {file_path}")
return checkpoint
async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]:
"""List checkpoint IDs. If workflow_id is provided, filter by that workflow."""
def _list_ids() -> list[str]:
checkpoint_ids: list[str] = []
for file_path in self.storage_path.glob("*.json"):
try:
with open(file_path) as f:
data = json.load(f)
if workflow_id is None or data.get("workflow_id") == workflow_id:
checkpoint_ids.append(data.get("checkpoint_id", file_path.stem))
except Exception as e:
logger.warning(f"Failed to read checkpoint file {file_path}: {e}")
return checkpoint_ids
return await asyncio.to_thread(_list_ids)
async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]:
"""List checkpoint objects. If workflow_id is provided, filter by that workflow."""
def _list_checkpoints() -> list[WorkflowCheckpoint]:
checkpoints: list[WorkflowCheckpoint] = []
for file_path in self.storage_path.glob("*.json"):
try:
with open(file_path) as f:
data = json.load(f)
if workflow_id is None or data.get("workflow_id") == workflow_id:
checkpoints.append(WorkflowCheckpoint.from_dict(data))
except Exception as e:
logger.warning(f"Failed to read checkpoint file {file_path}: {e}")
return checkpoints
return await asyncio.to_thread(_list_checkpoints)
async def delete_checkpoint(self, checkpoint_id: str) -> bool:
"""Delete a checkpoint by ID."""
file_path = self.storage_path / f"{checkpoint_id}.json"
def _delete() -> bool:
if file_path.exists():
file_path.unlink()
logger.info(f"Deleted checkpoint {checkpoint_id} from {file_path}")
return True
return False
return await asyncio.to_thread(_delete)
python/packages/main/agent_framework/_workflow/_concurrent.py
+304
View File
@@ -0,0 +1,304 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import inspect
import logging
from collections.abc import Callable, Sequence
from typing import Any
from agent_framework import AgentProtocol, ChatMessage, Role
from ._events import WorkflowCompletedEvent
from ._executor import AgentExecutorRequest, AgentExecutorResponse, Executor, handler
from ._workflow import Workflow, WorkflowBuilder
from ._workflow_context import WorkflowContext
logger = logging.getLogger(__name__)
"""Concurrent builder for agent-only fan-out/fan-in workflows.
This module provides a high-level, agent-focused API to quickly assemble a
parallel workflow with:
- a default dispatcher that broadcasts the input to all agent participants
- a default aggregator that combines all agent conversations and completes the workflow
Notes:
- Participants should be AgentProtocol instances or Executors.
- A custom aggregator can be provided as:
- an Executor instance (it should handle list[AgentExecutorResponse] and add a WorkflowCompletedEvent), or
- a callback function with signature:
def cb(results: list[AgentExecutorResponse]) -> Any | None
def cb(results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> Any | None
If the callback returns a non-None value, it is sent as the data of a WorkflowCompletedEvent.
If it returns None, the callback may have already emitted a completion event via ctx.
"""
class _DispatchToAllParticipants(Executor):
"""Broadcasts input to all downstream participants (via fan-out edges)."""
@handler
async def from_request(self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
# No explicit target: edge routing delivers to all connected participants.
await ctx.send_message(request)
@handler
async def from_str(self, prompt: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
request = AgentExecutorRequest(messages=[ChatMessage(Role.USER, text=prompt)], should_respond=True)
await ctx.send_message(request)
@handler
async def from_message(self, message: ChatMessage, ctx: WorkflowContext[AgentExecutorRequest]) -> None: # type: ignore[name-defined]
request = AgentExecutorRequest(messages=[message], should_respond=True)
await ctx.send_message(request)
@handler
async def from_messages(self, messages: list[ChatMessage], ctx: WorkflowContext[AgentExecutorRequest]) -> None: # type: ignore[name-defined]
request = AgentExecutorRequest(messages=list(messages), should_respond=True)
await ctx.send_message(request)
class _AggregateAgentConversations(Executor):
"""Aggregates agent responses and completes with combined ChatMessages.
Emits a list[ChatMessage] shaped as:
[ single_user_prompt?, agent1_final_assistant, agent2_final_assistant, ... ]
- Extracts a single user prompt (first user message seen across results).
- For each result, selects the final assistant message (prefers agent_run_response.messages).
- Avoids duplicating the same user message per agent.
"""
@handler
async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> None:
if not results:
logger.error("Concurrent aggregator received empty results list")
raise ValueError("Aggregation failed: no results provided")
def _is_role(msg: Any, role: Role) -> bool:
r = getattr(msg, "role", None)
if r is None:
return False
# Normalize both r and role to lowercase strings for comparison
r_str = str(r).lower() if isinstance(r, str) or hasattr(r, "__str__") else r
role_str = getattr(role, "value", None)
if role_str is None:
role_str = str(role)
role_str = role_str.lower()
return r_str == role_str
prompt_message: ChatMessage | None = None
assistant_replies: list[ChatMessage] = []
for r in results:
resp_messages = list(getattr(r.agent_run_response, "messages", []) or [])
conv = r.full_conversation if r.full_conversation is not None else resp_messages
logger.debug(
f"Aggregating executor {getattr(r, 'executor_id', '<unknown>')}: "
f"{len(resp_messages)} response msgs, {len(conv)} conversation msgs"
)
# Capture a single user prompt (first encountered across any conversation)
if prompt_message is None:
found_user = next((m for m in conv if _is_role(m, Role.USER)), None)
if found_user is not None:
prompt_message = found_user
# Pick the final assistant message from the response; fallback to conversation search
final_assistant = next((m for m in reversed(resp_messages) if _is_role(m, Role.ASSISTANT)), None)
if final_assistant is None:
final_assistant = next((m for m in reversed(conv) if _is_role(m, Role.ASSISTANT)), None)
if final_assistant is not None:
assistant_replies.append(final_assistant)
else:
logger.warning(
f"No assistant reply found for executor {getattr(r, 'executor_id', '<unknown>')}; skipping"
)
if not assistant_replies:
logger.error(f"Aggregation failed: no assistant replies found across {len(results)} results")
raise RuntimeError("Aggregation failed: no assistant replies found")
output: list[ChatMessage] = []
if prompt_message is not None:
output.append(prompt_message)
else:
logger.warning("No user prompt found in any conversation; emitting assistants only")
output.extend(assistant_replies)
await ctx.add_event(WorkflowCompletedEvent(data=output))
class _CallbackAggregator(Executor):
"""Wraps a Python callback as an aggregator.
Accepts either an async or sync callback with one of the signatures:
- (results: list[AgentExecutorResponse]) -> Any | None
- (results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> Any | None
Notes:
- Async callbacks are awaited directly.
- Sync callbacks are executed via asyncio.to_thread to avoid blocking the event loop.
- If the callback returns a non-None value, it is wrapped in a WorkflowCompletedEvent.
"""
def __init__(self, callback: Callable[..., Any], id: str | None = None) -> None:
super().__init__(id)
self._callback = callback
self._param_count = len(inspect.signature(callback).parameters)
@handler
async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> None:
# Call according to provided signature, always non-blocking for sync callbacks
if self._param_count >= 2:
if inspect.iscoroutinefunction(self._callback):
ret = await self._callback(results, ctx) # type: ignore[misc]
else:
ret = await asyncio.to_thread(self._callback, results, ctx)
else:
if inspect.iscoroutinefunction(self._callback):
ret = await self._callback(results) # type: ignore[misc]
else:
ret = await asyncio.to_thread(self._callback, results)
# If the callback returned a value, finalize the workflow with it
if ret is not None:
await ctx.add_event(WorkflowCompletedEvent(ret))
class ConcurrentBuilder:
r"""High-level builder for concurrent agent workflows.
- `participants([...])` accepts a list of AgentProtocol (recommended) or Executor.
- `build()` wires: dispatcher -> fan-out -> participants -> fan-in -> aggregator.
- `with_custom_aggregator(...)` overrides the default aggregator with an Executor or callback.
Usage:
```python
from agent_framework import ConcurrentBuilder
# Minimal: use default aggregator (returns list[ChatMessage])
workflow = ConcurrentBuilder().participants([agent1, agent2, agent3]).build()
# Custom aggregator via callback (sync or async). The callback receives
# list[AgentExecutorResponse] and its return value becomes
# WorkflowCompletedEvent.data
def summarize(results):
return " | ".join(r.agent_run_response.messages[-1].text for r in results)
workflow = ConcurrentBuilder().participants([agent1, agent2, agent3]).with_custom_aggregator(summarize).build()
```
"""
def __init__(self) -> None:
self._participants: list[AgentProtocol | Executor] = []
self._aggregator: Executor | None = None
def participants(self, participants: Sequence[AgentProtocol | Executor]) -> "ConcurrentBuilder":
r"""Define the parallel participants for this concurrent workflow.
Accepts AgentProtocol instances (e.g., created by a chat client) or Executor
instances. Each participant is wired as a parallel branch using fan-out edges
from an internal dispatcher.
Raises:
ValueError: if `participants` is empty or contains duplicates
TypeError: if any entry is not AgentProtocol or Executor
Example:
```python
wf = ConcurrentBuilder().participants([researcher_agent, marketer_agent, legal_agent]).build()
# Mixing agent(s) and executor(s) is supported
wf2 = ConcurrentBuilder().participants([researcher_agent, my_custom_executor]).build()
```
"""
if not participants:
raise ValueError("participants cannot be empty")
# Defensive duplicate detection
seen_agent_ids: set[int] = set()
seen_executor_ids: set[str] = set()
for p in participants:
if isinstance(p, Executor):
if p.id in seen_executor_ids:
raise ValueError(f"Duplicate executor participant detected: id '{p.id}'")
seen_executor_ids.add(p.id)
elif isinstance(p, AgentProtocol):
pid = id(p)
if pid in seen_agent_ids:
raise ValueError("Duplicate agent participant detected (same agent instance provided twice)")
seen_agent_ids.add(pid)
else:
raise TypeError(f"participants must be AgentProtocol or Executor instances; got {type(p).__name__}")
self._participants = list(participants)
return self
def with_aggregator(self, aggregator: Executor | Callable[..., Any]) -> "ConcurrentBuilder":
r"""Override the default aggregator with an Executor or a callback.
- Executor: must handle `list[AgentExecutorResponse]` and add a
`WorkflowCompletedEvent` to the context.
- Callback: sync or async callable with one of the signatures:
`(results: list[AgentExecutorResponse]) -> Any | None` or
`(results: list[AgentExecutorResponse], ctx: WorkflowContext[Any]) -> Any | None`.
If the callback returns a non-None value, it becomes the
`WorkflowCompletedEvent.data`.
Example:
```python
# Callback-based aggregator (string result)
async def summarize(results):
return " | ".join(r.agent_run_response.messages[-1].text for r in results)
wf = ConcurrentBuilder().participants([a1, a2, a3]).with_custom_aggregator(summarize).build()
```
"""
if isinstance(aggregator, Executor):
self._aggregator = aggregator
elif callable(aggregator):
self._aggregator = _CallbackAggregator(aggregator)
else:
raise TypeError("aggregator must be an Executor or a callable")
return self
def build(self) -> Workflow:
r"""Build and validate the concurrent workflow.
Wiring pattern:
- Dispatcher (internal) fans out the input to all `participants`
- Fan-in aggregator collects `AgentExecutorResponse` objects
- Aggregator emits a `WorkflowCompletedEvent` with either:
- list[ChatMessage] (default aggregator: one user + one assistant per agent)
- custom payload from the provided callback/executor
Returns:
Workflow: a ready-to-run workflow instance
Raises:
ValueError: if no participants were defined
Example:
```python
workflow = ConcurrentBuilder().participants([agent1, agent2]).build()
```
"""
if not self._participants:
raise ValueError("No participants provided. Call .participants([...]) first.")
dispatcher = _DispatchToAllParticipants(id="dispatcher")
aggregator = self._aggregator or _AggregateAgentConversations(id="aggregator")
builder = WorkflowBuilder()
return (
builder.set_start_executor(dispatcher)
.add_fan_out_edges(dispatcher, list(self._participants))
.add_fan_in_edges(list(self._participants), aggregator)
.build()
)
python/packages/main/agent_framework/_workflow/_const.py
+3
View File
@@ -0,0 +1,3 @@
# Copyright (c) Microsoft. All rights reserved.
DEFAULT_MAX_ITERATIONS = 100 # Default maximum iterations for workflow execution.
python/packages/main/agent_framework/_workflow/_edge.py
+349
View File
@@ -0,0 +1,349 @@
# Copyright (c) Microsoft. All rights reserved.
import logging
import uuid
from collections.abc import Callable, Sequence
from dataclasses import dataclass
from typing import Any, ClassVar
from pydantic import Field
from agent_framework._pydantic import AFBaseModel
from ._executor import Executor
logger = logging.getLogger(__name__)
def _extract_function_name(func: Callable[..., Any]) -> str:
"""Extract the name of any callable function for serialization.
Args:
func: The function to extract the name from.
Returns:
The name of the function, or a placeholder for lambda functions.
"""
if hasattr(func, "__name__"):
name = func.__name__
# Check if it's a lambda function
if name == "<lambda>":
return "<lambda>"
return name
# Fallback for other callable objects
return "<callable>"
class Edge(AFBaseModel):
"""Represents a directed edge in a graph."""
ID_SEPARATOR: ClassVar[str] = "->"
source_id: str = Field(min_length=1, description="The ID of the source executor of the edge")
target_id: str = Field(min_length=1, description="The ID of the target executor of the edge")
condition_name: str | None = Field(default=None, description="The name of the condition function for serialization")
def __init__(
self,
source_id: str,
target_id: str,
condition: Callable[[Any], bool] | None = None,
**kwargs: Any,
) -> None:
"""Initialize the edge with a source and target node.
Args:
source_id (str): The ID of the source executor of the edge.
target_id (str): The ID of the target executor of the edge.
condition (Callable[[Any], bool], optional): A condition function that determines
if the edge can handle the data. If None, the edge can handle any data type.
Defaults to None.
kwargs: Additional keyword arguments. Unused in this implementation.
"""
condition_name = _extract_function_name(condition) if condition is not None else None
kwargs.update({"source_id": source_id, "target_id": target_id, "condition_name": condition_name})
super().__init__(**kwargs)
self._condition = condition
@property
def id(self) -> str:
"""Get the unique ID of the edge."""
return f"{self.source_id}{self.ID_SEPARATOR}{self.target_id}"
def should_route(self, data: Any) -> bool:
"""Determine if message should be routed through this edge based on the condition."""
if self._condition is None:
return True
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."""
id: str = Field(
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=_default_edge_list, description="List of edges in this group")
def __init__(self, **kwargs: Any) -> None:
"""Initialize the edge group."""
if "id" not in kwargs:
kwargs["id"] = f"{self.__class__.__name__}/{uuid.uuid4()}"
if "type" not in kwargs:
kwargs["type"] = self.__class__.__name__
super().__init__(**kwargs)
@property
def source_executor_ids(self) -> list[str]:
"""Get the source executor IDs of the edges in the group."""
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."""
return list(dict.fromkeys(edge.target_id for edge in self.edges))
class SingleEdgeGroup(EdgeGroup):
"""Represents a single edge group that contains only one edge.
A concrete implementation of EdgeGroup that represent a group containing exactly one edge.
"""
def __init__(
self, source_id: str, target_id: str, condition: Callable[[Any], bool] | None = None, **kwargs: Any
) -> None:
"""Initialize the single edge group with an edge.
Args:
source_id (str): The source executor ID.
target_id (str): The target executor ID that the source executor can send messages to.
condition (Callable[[Any], bool], optional): A condition function that determines
if the edge will pass the data to the target executor. If None, the edge will
always pass the data to the target executor.
kwargs: Additional keyword arguments. Unused in this implementation.
"""
edge = Edge(source_id=source_id, target_id=target_id, condition=condition)
kwargs["edges"] = [edge]
super().__init__(**kwargs)
class FanOutEdgeGroup(EdgeGroup):
"""Represents a group of edges that share the same source executor.
Assembles a Fan-out pattern where multiple edges share the same source executor
and send messages to their respective target executors.
"""
selection_func_name: str | None = Field(
default=None, description="The name of the selection function for serialization"
)
def __init__(
self,
source_id: str,
target_ids: Sequence[str],
selection_func: Callable[[Any, list[str]], list[str]] | None = None,
**kwargs: Any,
) -> None:
"""Initialize the fan-out edge group with a list of edges.
Args:
source_id (str): The source executor ID.
target_ids (Sequence[str]): A list of target executor IDs that the source executor can send messages to.
selection_func (Callable[[Any, list[str]], list[str]], optional): A function that selects which target
executors to send messages to. The function takes in the message data and a list of target executor
IDs, and returns a list of selected target executor IDs.
kwargs: Additional keyword arguments. Unused in this implementation.
"""
if len(target_ids) <= 1:
raise ValueError("FanOutEdgeGroup must contain at least two targets.")
# Extract selection function name for serialization
selection_func_name = None
if selection_func is not None:
selection_func_name = _extract_function_name(selection_func)
edges = [Edge(source_id=source_id, target_id=target_id) for target_id in target_ids]
kwargs.update({"edges": edges, "selection_func_name": selection_func_name})
super().__init__(**kwargs)
self._target_ids = list(target_ids)
self._selection_func = selection_func
@property
def target_ids(self) -> list[str]:
"""Get the target executor IDs for selection."""
return self._target_ids
@property
def selection_func(self) -> Callable[[Any, list[str]], list[str]] | None:
"""Get the selection function for this fan-out group."""
return self._selection_func
class FanInEdgeGroup(EdgeGroup):
"""Represents a group of edges that share the same target executor.
Assembles a Fan-in pattern where multiple edges send messages to a single target executor.
Messages are buffered until all edges in the group have data to send.
"""
def __init__(self, source_ids: Sequence[str], target_id: str, **kwargs: Any) -> None:
"""Initialize the fan-in edge group with a list of edges.
Args:
source_ids (Sequence[str]): A list of source executor IDs that can send messages to the target executor.
target_id (str): The target executor ID that receives a list of messages aggregated from all sources.
kwargs: Additional keyword arguments. Unused in this implementation.
"""
if len(source_ids) <= 1:
raise ValueError("FanInEdgeGroup must contain at least two sources.")
edges = [Edge(source_id=source_id, target_id=target_id) for source_id in source_ids]
kwargs["edges"] = edges
super().__init__(**kwargs)
@dataclass
class Case:
"""Represents a single case in the switch-case edge group.
Args:
condition (Callable[[Any], bool]): The condition function for the case.
target (Executor): The target executor for the case.
"""
condition: Callable[[Any], bool]
target: Executor
@dataclass
class Default:
"""Represents the default case in the switch-case edge group.
Args:
target (Executor): The target executor for the default case.
"""
target: Executor
class SwitchCaseEdgeGroupCase(AFBaseModel):
"""A single case in the SwitchCaseEdgeGroup. This is used internally."""
target_id: str = Field(description="The target executor ID for this case")
condition_name: str | None = Field(default=None, description="The name of the condition function for serialization")
type: str = Field(default="Case", description="The type of the case")
def __init__(self, condition: Callable[[Any], bool], target_id: str, **kwargs: Any) -> None:
"""Initialize the switch case with a condition and target.
Args:
condition: The condition function for the case.
target_id: The target executor ID for this case.
kwargs: Additional keyword arguments.
"""
condition_name = _extract_function_name(condition)
kwargs.update({"target_id": target_id, "condition_name": condition_name})
super().__init__(**kwargs)
self._condition = condition
@property
def condition(self) -> Callable[[Any], bool]:
"""Get the condition function for this case."""
return self._condition
class SwitchCaseEdgeGroupDefault(AFBaseModel):
"""The default case in the SwitchCaseEdgeGroup. This is used internally."""
target_id: str = Field(description="The target executor ID for the default case")
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.
This is similar to a switch-case construct:
switch(data):
case condition_1:
edge_1
break
case condition_2:
edge_2
break
default:
edge_3
break
Or equivalently an if-elif-else construct:
if condition_1:
edge_1
elif condition_2:
edge_2
else:
edge_4
"""
cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault] = Field(
default_factory=_default_case_list,
description="List of conditional cases for this switch-case group",
)
def __init__(
self,
source_id: str,
cases: Sequence[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault],
**kwargs: Any,
) -> None:
"""Initialize the switch-case edge group with a list of edges.
Args:
source_id (str): The source executor ID.
cases (Sequence[Case | Default]): A list of cases for the switch-case edge group.
There should be exactly one default case.
kwargs: Additional keyword arguments. Unused in this implementation.
"""
if len(cases) < 2:
raise ValueError("SwitchCaseEdgeGroup must contain at least two cases (including the default case).")
default_case = [isinstance(case, SwitchCaseEdgeGroupDefault) for case in cases]
if sum(default_case) != 1:
raise ValueError("SwitchCaseEdgeGroup must contain exactly one default case.")
if not isinstance(cases[-1], SwitchCaseEdgeGroupDefault):
logger.warning(
"Default case in the switch-case edge group is not the last case. "
"This will result in unexpected behavior."
)
def selection_func(data: Any, targets: list[str]) -> list[str]:
"""Select the target executor based on the conditions."""
for index, case in enumerate(cases):
if isinstance(case, SwitchCaseEdgeGroupDefault):
return [case.target_id]
if isinstance(case, SwitchCaseEdgeGroupCase):
try:
if case.condition(data):
return [case.target_id]
except Exception as e:
logger.warning(f"Error occurred while evaluating condition for case {index}: {e}")
raise RuntimeError("No matching case found in SwitchCaseEdgeGroup.")
target_ids = [case.target_id for case in cases]
kwargs.update({"cases": cases})
super().__init__(source_id, target_ids, selection_func=selection_func, **kwargs)
python/packages/main/agent_framework/_workflow/_edge_runner.py
+349
View File
@@ -0,0 +1,349 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import logging
from abc import ABC, abstractmethod
from collections import defaultdict
from typing import Any
from ._edge import Edge, EdgeGroup, FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeGroup, SwitchCaseEdgeGroup
from ._executor import Executor
from ._runner_context import Message, RunnerContext
from ._shared_state import SharedState
from ._telemetry import EdgeGroupDeliveryStatus, workflow_tracer
from ._workflow_context import WorkflowContext
logger = logging.getLogger(__name__)
class EdgeRunner(ABC):
"""Abstract base class for edge runners that handle message delivery."""
def __init__(self, edge_group: EdgeGroup, executors: dict[str, Executor]) -> None:
"""Initialize the edge runner with an edge group and executor map.
Args:
edge_group: The edge group to run.
executors: Map of executor IDs to executor instances.
"""
self._edge_group = edge_group
self._executors = executors
@abstractmethod
async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool:
"""Send a message through the edge group.
Args:
message: The message to send.
shared_state: The shared state to use for holding data.
ctx: The context for the runner.
Returns:
bool: True if the message was processed successfully,
False if the target executor cannot handle the message.
"""
raise NotImplementedError
def _can_handle(self, executor_id: str, message_data: Any) -> bool:
"""Check if an executor can handle the given message data."""
if executor_id not in self._executors:
return False
return self._executors[executor_id].can_handle(message_data)
async def _execute_on_target(
self,
target_id: str,
source_ids: list[str],
message: Message,
shared_state: SharedState,
ctx: RunnerContext,
) -> None:
"""Execute a message on a target executor with trace context."""
if target_id not in self._executors:
raise RuntimeError(f"Target executor {target_id} not found.")
target_executor = self._executors[target_id]
# Create WorkflowContext with trace contexts from message
workflow_context: WorkflowContext[Any] = WorkflowContext(
target_id,
source_ids,
shared_state,
ctx,
trace_contexts=message.trace_contexts, # Pass trace contexts to WorkflowContext
source_span_ids=message.source_span_ids, # Pass source span IDs for linking
)
# Execute with trace context in WorkflowContext
await target_executor.execute(message.data, workflow_context)
class SingleEdgeRunner(EdgeRunner):
"""Runner for single edge groups."""
def __init__(self, edge_group: SingleEdgeGroup, executors: dict[str, Executor]) -> None:
super().__init__(edge_group, executors)
self._edge = edge_group.edges[0]
async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool:
"""Send a message through the single edge."""
should_execute = False
target_id = None
source_id = None
with workflow_tracer.create_edge_group_processing_span(
self._edge_group.__class__.__name__,
edge_group_id=self._edge_group.id,
message_source_id=message.source_id,
message_target_id=message.target_id,
source_trace_contexts=message.trace_contexts,
source_span_ids=message.source_span_ids,
):
try:
if message.target_id and message.target_id != self._edge.target_id:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH
)
return False
if self._can_handle(self._edge.target_id, message.data):
if self._edge.should_route(message.data):
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED)
should_execute = True
target_id = self._edge.target_id
source_id = self._edge.source_id
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE
)
# Return True here because message was processed, just condition failed
return True
else:
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH)
return False
except Exception as e:
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION)
raise e
# Execute outside the span
if should_execute and target_id and source_id:
await self._execute_on_target(target_id, [source_id], message, shared_state, ctx)
return True
return False
class FanOutEdgeRunner(EdgeRunner):
"""Runner for fan-out edge groups."""
def __init__(self, edge_group: FanOutEdgeGroup, executors: dict[str, Executor]) -> None:
super().__init__(edge_group, executors)
self._edges = edge_group.edges
self._target_ids = edge_group.target_ids
self._target_map = {edge.target_id: edge for edge in self._edges}
self._selection_func = edge_group.selection_func
async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool:
"""Send a message through all edges in the fan-out edge group."""
deliverable_edges = []
single_target_edge = None
# Process routing logic within span
with workflow_tracer.create_edge_group_processing_span(
self._edge_group.__class__.__name__,
edge_group_id=self._edge_group.id,
message_source_id=message.source_id,
message_target_id=message.target_id,
source_trace_contexts=message.trace_contexts,
source_span_ids=message.source_span_ids,
):
try:
selection_results = (
self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids
)
if not self._validate_selection_result(selection_results):
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION)
raise RuntimeError(
f"Invalid selection result: {selection_results}. "
f"Expected selections to be a subset of valid target executor IDs: {self._target_ids}."
)
if message.target_id:
# If the target ID is specified and the selection result contains it, send the message to that edge
if message.target_id in selection_results:
edge = self._target_map.get(message.target_id)
if edge and self._can_handle(edge.target_id, message.data):
if edge.should_route(message.data):
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED)
single_target_edge = edge
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE
)
# For targeted messages with condition failure, return True (message was processed)
return True
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH
)
# For targeted messages that can't be handled, return False
return False
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH
)
# For targeted messages not in selection, return False
return False
else:
# If no target ID, send the message to the selected targets
for target_id in selection_results:
edge = self._target_map[target_id]
if self._can_handle(edge.target_id, message.data) and edge.should_route(message.data):
deliverable_edges.append(edge)
if len(deliverable_edges) > 0:
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED)
else:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH
)
except Exception as e:
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION)
raise e
# Execute outside the span
if single_target_edge:
await self._execute_on_target(
single_target_edge.target_id, [single_target_edge.source_id], message, shared_state, ctx
)
return True
if deliverable_edges:
async def send_to_edge(edge: Edge) -> bool:
await self._execute_on_target(edge.target_id, [edge.source_id], message, shared_state, ctx)
return True
tasks = [send_to_edge(edge) for edge in deliverable_edges]
results = await asyncio.gather(*tasks)
return any(results)
# If we get here, it's a broadcast message with no deliverable edges
return False
def _validate_selection_result(self, selection_results: list[str]) -> bool:
"""Validate the selection results to ensure all IDs are valid target executor IDs."""
return all(result in self._target_ids for result in selection_results)
class FanInEdgeRunner(EdgeRunner):
"""Runner for fan-in edge groups."""
def __init__(self, edge_group: FanInEdgeGroup, executors: dict[str, Executor]) -> None:
super().__init__(edge_group, executors)
self._edges = edge_group.edges
# Buffer to hold messages before sending them to the target executor
# Key is the source executor ID, value is a list of messages
self._buffer: dict[str, list[Message]] = defaultdict(list)
async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool:
"""Send a message through all edges in the fan-in edge group."""
execution_data: dict[str, Any] | None = None
with workflow_tracer.create_edge_group_processing_span(
self._edge_group.__class__.__name__,
edge_group_id=self._edge_group.id,
message_source_id=message.source_id,
message_target_id=message.target_id,
source_trace_contexts=message.trace_contexts,
source_span_ids=message.source_span_ids,
):
try:
if message.target_id and message.target_id != self._edges[0].target_id:
workflow_tracer.set_edge_group_span_attributes(
False, EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH
)
return False
# Check if target can handle list of message data (fan-in aggregates multiple messages)
if self._can_handle(self._edges[0].target_id, [message.data]):
# If the edge can handle the data, buffer the message
self._buffer[message.source_id].append(message)
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.BUFFERED)
else:
# If the edge cannot handle the data, return False
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH)
return False
if self._is_ready_to_send():
# If all edges in the group have data, prepare for execution
messages_to_send = [msg for edge in self._edges for msg in self._buffer[edge.source_id]]
self._buffer.clear()
# Send aggregated data to target
aggregated_data = [msg.data for msg in messages_to_send]
# Collect all trace contexts and source span IDs for fan-in linking
trace_contexts = [msg.trace_context for msg in messages_to_send if msg.trace_context]
source_span_ids = [msg.source_span_id for msg in messages_to_send if msg.source_span_id]
# Create a new Message object for the aggregated data
aggregated_message = Message(
data=aggregated_data,
source_id=self._edge_group.__class__.__name__, # This won't be used in self._execute_on_target.
trace_contexts=trace_contexts,
source_span_ids=source_span_ids,
)
workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED)
# Store execution data for later
execution_data = {
"target_id": self._edges[0].target_id,
"source_ids": [edge.source_id for edge in self._edges],
"message": aggregated_message,
}
except Exception as e:
workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION)
raise e
# Execute outside the span if needed
if execution_data:
await self._execute_on_target(
execution_data["target_id"], execution_data["source_ids"], execution_data["message"], shared_state, ctx
)
return True
return True # Return True for buffered messages (waiting for more)
def _is_ready_to_send(self) -> bool:
"""Check if all edges in the group have data to send."""
return all(self._buffer[edge.source_id] for edge in self._edges)
class SwitchCaseEdgeRunner(FanOutEdgeRunner):
"""Runner for switch-case edge groups (inherits from FanOutEdgeRunner)."""
def __init__(self, edge_group: SwitchCaseEdgeGroup, executors: dict[str, Executor]) -> None:
super().__init__(edge_group, executors)
def create_edge_runner(edge_group: EdgeGroup, executors: dict[str, Executor]) -> EdgeRunner:
"""Factory function to create the appropriate edge runner for an edge group.
Args:
edge_group: The edge group to create a runner for.
executors: Map of executor IDs to executor instances.
Returns:
The appropriate EdgeRunner instance.
"""
if isinstance(edge_group, SingleEdgeGroup):
return SingleEdgeRunner(edge_group, executors)
if isinstance(edge_group, SwitchCaseEdgeGroup):
return SwitchCaseEdgeRunner(edge_group, executors)
if isinstance(edge_group, FanOutEdgeGroup):
return FanOutEdgeRunner(edge_group, executors)
if isinstance(edge_group, FanInEdgeGroup):
return FanInEdgeRunner(edge_group, executors)
raise ValueError(f"Unsupported edge group type: {type(edge_group)}")
python/packages/main/agent_framework/_workflow/_events.py
+259
View File
@@ -0,0 +1,259 @@
# Copyright (c) Microsoft. All rights reserved.
import traceback as _traceback
from dataclasses import dataclass
from enum import Enum
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."""
def __init__(self, data: Any | None = None):
"""Initialize the workflow event with optional data."""
self.data = data
def __repr__(self) -> str:
"""Return a string representation of the workflow event."""
return f"{self.__class__.__name__}(data={self.data if self.data is not None else 'None'})"
class WorkflowStartedEvent(WorkflowEvent):
"""Event triggered when a workflow starts."""
...
class WorkflowCompletedEvent(WorkflowEvent):
"""Event triggered when a workflow completes."""
...
class WorkflowWarningEvent(WorkflowEvent):
"""Event triggered when a warning occurs in the workflow."""
def __init__(self, data: str):
"""Initialize the workflow warning event with optional data and warning message."""
super().__init__(data)
def __repr__(self) -> str:
"""Return a string representation of the workflow warning event."""
return f"{self.__class__.__name__}(message={self.data})"
class WorkflowErrorEvent(WorkflowEvent):
"""Event triggered when an error occurs in the workflow."""
def __init__(self, data: Exception):
"""Initialize the workflow error event with optional data and error message."""
super().__init__(data)
def __repr__(self) -> str:
"""Return a string representation of the workflow error event."""
return f"{self.__class__.__name__}(exception={self.data})"
class WorkflowRunState(str, Enum):
"""Run-level state of a workflow execution.
Semantics:
- STARTED: Run has been initiated and the workflow context has been created.
This is an initial state before any meaningful work is performed. In this
codebase we emit a dedicated `WorkflowStartedEvent` for telemetry, and
typically advance the status directly to `IN_PROGRESS`. Consumers may
still rely on `STARTED` for state machines that need an explicit pre-work
phase.
- IN_PROGRESS: The workflow is actively executing (e.g., the initial
message has been delivered to the start executor or a superstep is
running). This status is emitted at the beginning of a run and can be
followed by other statuses as the run progresses.
- IN_PROGRESS_PENDING_REQUESTS: Active execution while one or more
request-for-information operations are outstanding. New work may still
be scheduled while requests are in flight.
- IDLE: The workflow is quiescent with no outstanding requests, but has
not yet emitted a terminal result. Rare in practice but provided for
orchestration integrations that distinguish a quiescent state.
- IDLE_WITH_PENDING_REQUESTS: The workflow is paused awaiting external
input (e.g., emitted a `RequestInfoEvent`). This is a non-terminal
state; the workflow can resume when responses are supplied.
- COMPLETED: Normal terminal state indicating successful completion.
- FAILED: Terminal state indicating an error surfaced. Accompanied by a
`WorkflowFailedEvent` with structured error details.
- CANCELLED: Terminal state indicating the run was cancelled by a caller
or orchestrator. Not currently emitted by default runner paths but
included for integrators/orchestrators that support cancellation.
"""
STARTED = "STARTED" # Explicit pre-work phase (rarely emitted as status; see note above)
IN_PROGRESS = "IN_PROGRESS" # Active execution is underway
IN_PROGRESS_PENDING_REQUESTS = "IN_PROGRESS_PENDING_REQUESTS" # Active execution with outstanding requests
IDLE = "IDLE" # No active work and no outstanding requests
IDLE_WITH_PENDING_REQUESTS = "IDLE_WITH_PENDING_REQUESTS" # Paused awaiting external responses
COMPLETED = "COMPLETED" # Finished successfully
FAILED = "FAILED" # Finished with an error
CANCELLED = "CANCELLED" # Finished due to cancellation
class WorkflowStatusEvent(WorkflowEvent):
"""Event indicating a transition in the workflow run state."""
def __init__(self, state: WorkflowRunState, data: Any | None = None):
super().__init__(data)
self.state = state
def __repr__(self) -> str: # pragma: no cover - representation only
return f"{self.__class__.__name__}(state={self.state}, data={self.data!r})"
@dataclass
class WorkflowErrorDetails:
"""Structured error information to surface in error events/results."""
error_type: str
message: str
traceback: str | None = None
executor_id: str | None = None
extra: dict[str, Any] | None = None
@classmethod
def from_exception(
cls,
exc: BaseException,
*,
executor_id: str | None = None,
extra: dict[str, Any] | None = None,
) -> "WorkflowErrorDetails":
tb = None
try:
tb = "".join(_traceback.format_exception(type(exc), exc, exc.__traceback__))
except Exception:
tb = None
return cls(
error_type=exc.__class__.__name__,
message=str(exc),
traceback=tb,
executor_id=executor_id,
extra=extra,
)
class WorkflowFailedEvent(WorkflowEvent):
"""Terminal failure event for a workflow run."""
def __init__(self, details: WorkflowErrorDetails, data: Any | None = None):
super().__init__(data)
self.details = details
def __repr__(self) -> str: # pragma: no cover - representation only
return f"{self.__class__.__name__}(details={self.details}, data={self.data!r})"
class RequestInfoEvent(WorkflowEvent):
"""Event triggered when a workflow executor requests external information."""
def __init__(
self,
request_id: str,
source_executor_id: str,
request_type: type,
request_data: "RequestInfoMessage",
):
"""Initialize the request info event.
Args:
request_id: Unique identifier for the request.
source_executor_id: ID of the executor that made the request.
request_type: Type of the request (e.g., a specific data type).
request_data: The data associated with the request.
"""
super().__init__(request_data)
self.request_id = request_id
self.source_executor_id = source_executor_id
self.request_type = request_type
def __repr__(self) -> str:
"""Return a string representation of the request info event."""
return (
f"{self.__class__.__name__}("
f"request_id={self.request_id}, "
f"source_executor_id={self.source_executor_id}, "
f"request_type={self.request_type.__name__}, "
f"data={self.data})"
)
class ExecutorEvent(WorkflowEvent):
"""Base class for executor events."""
def __init__(self, executor_id: str, data: Any | None = None):
"""Initialize the executor event with an executor ID and optional data."""
super().__init__(data)
self.executor_id = executor_id
def __repr__(self) -> str:
"""Return a string representation of the executor event."""
return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})"
class ExecutorInvokeEvent(ExecutorEvent):
"""Event triggered when an executor handler is invoked."""
def __repr__(self) -> str:
"""Return a string representation of the executor handler invoke event."""
return f"{self.__class__.__name__}(executor_id={self.executor_id})"
class ExecutorCompletedEvent(ExecutorEvent):
"""Event triggered when an executor handler is completed."""
def __repr__(self) -> str:
"""Return a string representation of the executor handler complete event."""
return f"{self.__class__.__name__}(executor_id={self.executor_id})"
class ExecutorFailedEvent(ExecutorEvent):
"""Event triggered when an executor handler raises an error."""
def __init__(self, executor_id: str, details: WorkflowErrorDetails):
super().__init__(executor_id, details)
self.details = details
def __repr__(self) -> str: # pragma: no cover - representation only
return f"{self.__class__.__name__}(executor_id={self.executor_id}, details={self.details})"
class AgentRunUpdateEvent(ExecutorEvent):
"""Event triggered when an agent is streaming messages."""
def __init__(self, executor_id: str, data: AgentRunResponseUpdate | None = None):
"""Initialize the agent streaming event."""
super().__init__(executor_id, data)
def __repr__(self) -> str:
"""Return a string representation of the agent streaming event."""
return f"{self.__class__.__name__}(executor_id={self.executor_id}, messages={self.data})"
class AgentRunEvent(ExecutorEvent):
"""Event triggered when an agent run is completed."""
def __init__(self, executor_id: str, data: AgentRunResponse | None = None):
"""Initialize the agent run event."""
super().__init__(executor_id, data)
def __repr__(self) -> str:
"""Return a string representation of the agent run event."""
return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})"
python/packages/main/agent_framework/_workflow/_executor.py
+1099
View File
File diff suppressed because it is too large Load Diff
python/packages/main/agent_framework/_workflow/_function_executor.py
+277
View File
@@ -0,0 +1,277 @@
# 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.
"""
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
python/packages/main/agent_framework/_workflow/_magentic.py
+1850
View File
File diff suppressed because it is too large Load Diff
python/packages/main/agent_framework/_workflow/_runner.py
+481
View File
@@ -0,0 +1,481 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import logging
from collections import defaultdict
from collections.abc import AsyncGenerator, Sequence
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from ._executor import RequestInfoExecutor
from ._edge import EdgeGroup
from ._edge_runner import EdgeRunner, create_edge_runner
from ._events import WorkflowCompletedEvent, WorkflowEvent
from ._executor import Executor
from ._runner_context import Message, RunnerContext
from ._shared_state import SharedState
from ._typing_utils import is_instance_of
from ._workflow_context import WorkflowContext
logger = logging.getLogger(__name__)
class Runner:
"""A class to run a workflow in Pregel supersteps."""
def __init__(
self,
edge_groups: Sequence[EdgeGroup],
executors: dict[str, Executor],
shared_state: SharedState,
ctx: RunnerContext,
max_iterations: int = 100,
workflow_id: str | None = None,
) -> None:
"""Initialize the runner with edges, shared state, and context.
Args:
edge_groups: The edge groups of the workflow.
executors: Map of executor IDs to executor instances.
shared_state: The shared state for the workflow.
ctx: The runner context for the workflow.
max_iterations: The maximum number of iterations to run.
workflow_id: The workflow ID for checkpointing.
"""
self._executors = executors
self._edge_runners = [create_edge_runner(group, executors) for group in edge_groups]
self._edge_runner_map = self._parse_edge_runners(self._edge_runners)
self._ctx = ctx
self._iteration = 0
self._max_iterations = max_iterations
self._shared_state = shared_state
self._workflow_id = workflow_id
self._running = False
self._resumed_from_checkpoint = False # Track whether we resumed
# Set workflow ID in context if provided
if workflow_id:
self._ctx.set_workflow_id(workflow_id)
@property
def context(self) -> RunnerContext:
"""Get the workflow context."""
return self._ctx
def mark_resumed(self, iteration: int | None = None, max_iterations: int | None = None) -> None:
"""Mark the runner as having resumed from a checkpoint.
Optionally set the current iteration and max iterations.
"""
self._resumed_from_checkpoint = True
if iteration is not None:
self._iteration = iteration
if max_iterations is not None:
self._max_iterations = max_iterations
async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]:
"""Run the workflow until no more messages are sent."""
if self._running:
raise RuntimeError("Runner is already running.")
self._running = True
try:
# Emit any events already produced prior to entering loop
if await self._ctx.has_events():
logger.info("Yielding pre-loop events")
for event in await self._ctx.drain_events():
yield event
# Create first checkpoint if there are messages from initial execution
if await self._ctx.has_messages() and self._ctx.has_checkpointing():
if not self._resumed_from_checkpoint:
logger.info("Creating checkpoint after initial execution")
await self._create_checkpoint_if_enabled("after_initial_execution")
else:
logger.info("Skipping 'after_initial_execution' checkpoint because we resumed from a checkpoint")
# Initialize context with starting iteration state
await self._update_context_with_shared_state()
while self._iteration < self._max_iterations:
logger.info(f"Starting superstep {self._iteration + 1}")
# Run iteration concurrently with live event streaming: we poll
# for new events while the iteration coroutine progresses.
iteration_task = asyncio.create_task(self._run_iteration())
while not iteration_task.done():
try:
# Wait briefly for any new event; timeout allows progress checks
event = await asyncio.wait_for(self._ctx.next_event(), timeout=0.05)
yield event
except asyncio.TimeoutError:
# Periodically continue to let iteration advance
continue
# Propagate errors from iteration, but first surface any pending events
try:
await iteration_task
except Exception:
# Make sure failure-related events (like ExecutorFailedEvent) are surfaced
if await self._ctx.has_events():
for event in await self._ctx.drain_events():
yield event
raise
self._iteration += 1
# Drain any straggler events emitted at tail end
if await self._ctx.has_events():
for event in await self._ctx.drain_events():
yield event
# Update context with current iteration state immediately
await self._update_context_with_shared_state()
logger.info(f"Completed superstep {self._iteration}")
# Create checkpoint after each superstep iteration
await self._create_checkpoint_if_enabled(f"superstep_{self._iteration}")
if not await self._ctx.has_messages():
break
if self._iteration >= self._max_iterations and await self._ctx.has_messages():
raise RuntimeError(f"Runner did not converge after {self._max_iterations} iterations.")
logger.info(f"Workflow completed after {self._iteration} supersteps")
self._iteration = 0
self._resumed_from_checkpoint = False # Reset resume flag for next run
finally:
self._running = False
async def _run_iteration(self) -> None:
async def _deliver_messages(source_executor_id: str, messages: list[Message]) -> None:
"""Outer loop to concurrently deliver messages from all sources to their targets."""
# Special handling for SubWorkflowRequestInfo messages
async def _deliver_sub_workflow_requests(messages: list[Message]) -> None:
from ._executor import SubWorkflowRequestInfo
# Handle SubWorkflowRequestInfo messages - only process those not already targeted
sub_workflow_messages: list[Message] = []
for msg in messages:
# Skip messages sent directly to RequestInfoExecutor - they are already forwarded
if self._is_message_to_request_info_executor(msg):
continue
if isinstance(msg.data, SubWorkflowRequestInfo):
sub_workflow_messages.append(msg)
for message in sub_workflow_messages:
# message.data is guaranteed to be SubWorkflowRequestInfo via filtering above
sub_request = message.data # type: ignore[assignment]
# Find executor that can intercept the wrapped type
interceptor_found = False
for executor in self._executors.values():
interceptors = getattr(executor, "_request_interceptors", [])
if interceptors and executor.id != message.source_id:
for registered_type in interceptors: # type: ignore[assignment]
# Check type matching - handle both type and string cases
matched = False
if (
isinstance(registered_type, type)
and is_instance_of(sub_request.data, registered_type)
) or (
isinstance(registered_type, str)
and hasattr(sub_request.data, "__class__")
and sub_request.data.__class__.__name__ == registered_type
):
matched = True
if matched:
# Send directly to the intercepting executor
logger.info(
f"Sending sub-workflow request of type '{sub_request.data.__class__.__name__}' "
f"from sub-workflow '{sub_request.sub_workflow_id}' "
f"to executor '{executor.id}' for interception."
)
# Create WorkflowContext with trace context from message
workflow_ctx: WorkflowContext[Any] = WorkflowContext(
executor.id,
[message.source_id],
self._shared_state,
self._ctx,
trace_contexts=[message.trace_context] if message.trace_context else None,
source_span_ids=[message.source_span_id] if message.source_span_id else None,
)
await executor.execute(sub_request, workflow_ctx)
interceptor_found = True
break
if interceptor_found:
break
if not interceptor_found:
# No interceptor found - send directly to RequestInfoExecutor if available.
# Find the RequestInfoExecutor instance
request_info_executor = self._find_request_info_executor()
if request_info_executor:
request_info_workflow_ctx: WorkflowContext[None] = WorkflowContext(
request_info_executor.id,
[message.source_id],
self._shared_state,
self._ctx,
trace_contexts=[message.trace_context] if message.trace_context else None,
source_span_ids=[message.source_span_id] if message.source_span_id else None,
)
logger.info(
f"Sending sub-workflow request of type '{sub_request.data.__class__.__name__}' "
f"from sub-workflow '{sub_request.sub_workflow_id}' to RequestInfoExecutor "
f"'{request_info_executor.id}'"
)
await request_info_executor.execute(sub_request, request_info_workflow_ctx)
else:
logger.warning(
f"Sub-workflow request of type '{sub_request.data.__class__.__name__}' "
f"from sub-workflow '{sub_request.sub_workflow_id}' could not be handled: "
f"no RequestInfoExecutor found in the workflow. Add a RequestInfoExecutor "
f"to handle external requests or add an @intercepts_request handler."
)
async def _deliver_message_inner(edge_runner: EdgeRunner, message: Message) -> bool:
"""Inner loop to deliver a single message through an edge runner."""
return await edge_runner.send_message(message, self._shared_state, self._ctx)
# Handle SubWorkflowRequestInfo messages specially
await _deliver_sub_workflow_requests(messages)
# Filter out SubWorkflowRequestInfo messages from normal edge routing
# since they were handled specially
from ._executor import SubWorkflowRequestInfo
non_sub_workflow_messages: list[Message] = []
for msg in messages:
# Keep messages sent directly to RequestInfoExecutor (forwarded messages)
if self._is_message_to_request_info_executor(msg):
non_sub_workflow_messages.append(msg)
continue
# Skip SubWorkflowRequestInfo messages (handled by special routing)
if isinstance(msg.data, SubWorkflowRequestInfo):
continue
non_sub_workflow_messages.append(msg)
associated_edge_runners = self._edge_runner_map.get(source_executor_id, [])
for message in non_sub_workflow_messages:
# Deliver a message through all edge runners associated with the source executor concurrently.
tasks = [_deliver_message_inner(edge_runner, message) for edge_runner in associated_edge_runners]
if not tasks:
# No outgoing edges. If this is an AgentExecutorResponse, treat it as an
# intentional terminal emission and emit a WorkflowCompletedEvent here.
# (Previously this relied on the executor to emit, but AgentExecutor only
# sends an AgentExecutorResponse message; centralized completion keeps the
# contract consistent with other executors.)
try: # Local import to avoid circular dependencies at module import time.
from ._executor import AgentExecutorResponse # type: ignore
if isinstance(message.data, AgentExecutorResponse):
final_messages = message.data.agent_run_response.messages
final_text = final_messages[-1].text if final_messages else "(no content)"
await self._ctx.add_event(WorkflowCompletedEvent(final_text))
continue # Terminal handled
except Exception as exc: # pragma: no cover - defensive
logger.debug("Suppressed exception during terminal message type check: %s", exc)
# Otherwise keep prior behavior (emit warning for unexpected undelivered message).
logger.warning(
f"Message {message} could not be delivered (no outgoing edges). "
"Add a downstream executor or remove the send if this is unexpected."
)
continue
results = await asyncio.gather(*tasks)
if not any(results):
# Outgoing edges exist but none accepted the message. If this is an
# AgentExecutorResponse, treat as natural terminal and emit completion.
try:
from ._executor import AgentExecutorResponse # type: ignore
if isinstance(message.data, AgentExecutorResponse):
# Emit a single completion event with final text (best-effort extraction)
final_messages = message.data.agent_run_response.messages
final_text = final_messages[-1].text if final_messages else "(no content)"
await self._ctx.add_event(WorkflowCompletedEvent(final_text))
continue
except Exception as exc: # pragma: no cover
logger.debug("Terminal completion emission failed: %s", exc)
logger.warning(
f"Message {message} could not be delivered. "
"This may be due to type incompatibility or no matching targets."
)
messages = await self._ctx.drain_messages()
tasks = [_deliver_messages(source_executor_id, messages) for source_executor_id, messages in messages.items()]
await asyncio.gather(*tasks)
async def _create_checkpoint_if_enabled(self, checkpoint_type: str) -> str | None:
"""Create a checkpoint if checkpointing is enabled and attach a label and metadata."""
if not self._ctx.has_checkpointing():
return None
try:
# Auto-snapshot executor states
await self._auto_snapshot_executor_states()
await self._update_context_with_shared_state()
checkpoint_category = "initial" if checkpoint_type == "after_initial_execution" else "superstep"
metadata = {
"superstep": self._iteration,
"checkpoint_type": checkpoint_category,
}
checkpoint_id = await self._ctx.create_checkpoint(metadata=metadata)
logger.info(f"Created {checkpoint_type} checkpoint: {checkpoint_id}")
return checkpoint_id
except Exception as e:
logger.warning(f"Failed to create {checkpoint_type} checkpoint: {e}")
return None
async def _auto_snapshot_executor_states(self) -> None:
"""Populate executor state by calling snapshot hooks on executors if available.
Convention:
- If an executor defines an async or sync method `snapshot_state(self) -> dict`, use it.
- Else if it has a plain attribute `state` that is a dict, use that.
Only JSON-serializable dicts should be provided by executors.
"""
for exec_id, executor in self._executors.items():
state_dict: dict[str, Any] | None = None
snapshot = getattr(executor, "snapshot_state", None)
try:
if callable(snapshot):
maybe = snapshot()
if asyncio.iscoroutine(maybe): # type: ignore[arg-type]
maybe = await maybe # type: ignore[assignment]
if isinstance(maybe, dict):
state_dict = maybe # type: ignore[assignment]
else:
state_attr = getattr(executor, "state", None)
if isinstance(state_attr, dict):
state_dict = state_attr # type: ignore[assignment]
except Exception as ex: # pragma: no cover
logger.debug(f"Executor {exec_id} snapshot_state failed: {ex}")
if state_dict is not None:
try:
await self._ctx.set_state(exec_id, state_dict)
except Exception as ex: # pragma: no cover
logger.debug(f"Failed to persist state for executor {exec_id}: {ex}")
async def _update_context_with_shared_state(self) -> None:
if not self._ctx.has_checkpointing():
return
try:
current_state = await self._ctx.get_checkpoint_state()
shared_state_data = {}
async with self._shared_state.hold():
if hasattr(self._shared_state, "_state"):
shared_state_data = dict(self._shared_state._state) # type: ignore[attr-defined]
current_state["shared_state"] = shared_state_data
current_state["iteration_count"] = self._iteration
current_state["max_iterations"] = self._max_iterations
await self._ctx.set_checkpoint_state(current_state)
except Exception as e:
logger.warning(f"Failed to update context with shared state: {e}")
async def restore_from_checkpoint(self, checkpoint_id: str) -> bool:
"""Restore workflow state from a checkpoint.
Args:
checkpoint_id: The ID of the checkpoint to restore from
Returns:
True if restoration was successful, False otherwise
"""
if not self._ctx.has_checkpointing():
logger.warning("Context does not support checkpointing")
return False
try:
success = await self._ctx.restore_from_checkpoint(checkpoint_id)
if not success:
return False
await self._restore_shared_state_from_context()
self.mark_resumed() # mark resumed; iteration/max already restored from context
logger.info(f"Successfully restored workflow from checkpoint: {checkpoint_id}")
return True
except Exception as e:
logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}")
return False
async def _restore_shared_state_from_context(self) -> None:
if not self._ctx.has_checkpointing():
return
try:
restored_state = await self._ctx.get_checkpoint_state()
shared_state_data = restored_state.get("shared_state", {})
if shared_state_data and hasattr(self._shared_state, "_state"):
async with self._shared_state.hold():
self._shared_state._state.clear() # type: ignore[attr-defined]
self._shared_state._state.update(shared_state_data) # type: ignore[attr-defined]
self._iteration = restored_state.get("iteration_count", 0)
self._max_iterations = restored_state.get("max_iterations", self._max_iterations)
except Exception as e:
logger.warning(f"Failed to restore shared state from context: {e}")
def _parse_edge_runners(self, edge_runners: list[EdgeRunner]) -> dict[str, list[EdgeRunner]]:
"""Parse the edge runners of the workflow into a mapping where each source executor ID maps to its edge runners.
Args:
edge_runners: A list of edge runners in the workflow.
Returns:
A dictionary mapping each source executor ID to a list of edge runners.
"""
parsed: defaultdict[str, list[EdgeRunner]] = defaultdict(list)
for runner in edge_runners:
# Accessing protected attribute (_edge_group) intentionally for internal wiring.
for source_executor_id in runner._edge_group.source_executor_ids: # type: ignore[attr-defined]
parsed[source_executor_id].append(runner)
return parsed
def _find_request_info_executor(self) -> "RequestInfoExecutor | None":
"""Find the RequestInfoExecutor instance in this workflow.
Returns:
The RequestInfoExecutor instance if found, None otherwise.
"""
from ._executor import RequestInfoExecutor
for executor in self._executors.values():
if isinstance(executor, RequestInfoExecutor):
return executor
return None
def _is_message_to_request_info_executor(self, msg: "Message") -> bool:
"""Check if message targets any RequestInfoExecutor in this workflow.
Args:
msg: The message to check.
Returns:
True if the message targets a RequestInfoExecutor, False otherwise.
"""
from ._executor import RequestInfoExecutor
if not msg.target_id:
return False
# Check all executors to see if target_id matches a RequestInfoExecutor
for executor in self._executors.values():
if executor.id == msg.target_id and isinstance(executor, RequestInfoExecutor):
return True
return False
python/packages/main/agent_framework/_workflow/_runner_context.py
+559
View File
@@ -0,0 +1,559 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import importlib
import logging
import uuid
from collections import defaultdict
from dataclasses import dataclass, fields, is_dataclass
from typing import Any, Protocol, TypedDict, TypeVar, cast, runtime_checkable
from ._checkpoint import CheckpointStorage, WorkflowCheckpoint
from ._const import DEFAULT_MAX_ITERATIONS
from ._events import AgentRunUpdateEvent, WorkflowEvent
from ._shared_state import SharedState
logger = logging.getLogger(__name__)
T = TypeVar("T")
@dataclass
class Message:
"""A class representing a message in the workflow."""
data: Any
source_id: str
target_id: str | None = None
# OpenTelemetry trace context fields for message propagation
# These are plural to support fan-in scenarios where multiple messages are aggregated
trace_contexts: list[dict[str, str]] | None = None # W3C Trace Context headers from multiple sources
source_span_ids: list[str] | None = None # Publishing span IDs for linking from multiple sources
# Backward compatibility properties
@property
def trace_context(self) -> dict[str, str] | None:
"""Get the first trace context for backward compatibility."""
return self.trace_contexts[0] if self.trace_contexts else None
@property
def source_span_id(self) -> str | None:
"""Get the first source span ID for backward compatibility."""
return self.source_span_ids[0] if self.source_span_ids else None
class CheckpointState(TypedDict):
messages: dict[str, list[dict[str, Any]]]
shared_state: dict[str, Any]
executor_states: dict[str, dict[str, Any]]
iteration_count: int
max_iterations: int
# Checkpoint serialization helpers
_PYDANTIC_MARKER = "__af_pydantic_model__"
_DATACLASS_MARKER = "__af_dataclass__"
# Guards to prevent runaway recursion while encoding arbitrary user data
_MAX_ENCODE_DEPTH = 100
_CYCLE_SENTINEL = "<cycle>"
def _is_pydantic_model(obj: object) -> bool:
"""Best-effort check for Pydantic models (e.g., AFBaseModel).
We avoid hard dependencies by duck-typing on model_dump/model_validate.
"""
try:
obj_type: type[Any] = type(obj)
return hasattr(obj, "model_dump") and hasattr(obj_type, "model_validate")
except Exception:
return False
def _encode_checkpoint_value(value: Any) -> Any:
"""Recursively encode values into JSON-serializable structures.
- Pydantic models -> { _PYDANTIC_MARKER: "module:Class", value: model_dump(mode="json") }
- dataclass instances -> { _DATACLASS_MARKER: "module:Class", value: {field: encoded} }
- dict -> encode keys as str and values recursively
- list/tuple/set -> list of encoded items
- other -> returned as-is if already JSON-serializable
Includes cycle and depth protection to avoid infinite recursion.
"""
def _enc(v: Any, stack: set[int], depth: int) -> Any:
# Depth guard
if depth > _MAX_ENCODE_DEPTH:
logger.debug(f"Max encode depth reached at depth={depth} for type={type(v)}")
return "<max_depth>"
# Pydantic (AFBaseModel) handling
if _is_pydantic_model(v):
cls = cast(type[Any], type(v)) # type: ignore
try:
return {
_PYDANTIC_MARKER: f"{cls.__module__}:{cls.__name__}",
"value": v.model_dump(mode="json"),
}
except Exception as exc: # best-effort fallback
logger.debug("Pydantic model_dump failed for %s: %s", cls, exc)
return str(v)
# Dataclasses (instances only)
if is_dataclass(v) and not isinstance(v, type):
oid = id(v)
if oid in stack:
logger.debug("Cycle detected while encoding dataclass instance")
return _CYCLE_SENTINEL
stack.add(oid)
try:
# type(v) already narrows sufficiently; cast was redundant
dc_cls: type[Any] = type(v)
field_values: dict[str, Any] = {}
for f in fields(v): # type: ignore[arg-type]
field_values[f.name] = _enc(getattr(v, f.name), stack, depth + 1)
return {
_DATACLASS_MARKER: f"{dc_cls.__module__}:{dc_cls.__name__}",
"value": field_values,
}
finally:
stack.remove(oid)
# Collections
if isinstance(v, dict):
v_dict = cast("dict[object, object]", v)
oid = id(v_dict)
if oid in stack:
logger.debug("Cycle detected while encoding dict")
return _CYCLE_SENTINEL
stack.add(oid)
try:
json_dict: dict[str, Any] = {}
for k_any, val_any in v_dict.items(): # type: ignore[assignment]
k_str: str = str(k_any)
json_dict[k_str] = _enc(val_any, stack, depth + 1)
return json_dict
finally:
stack.remove(oid)
if isinstance(v, (list, tuple, set)):
iterable_v = cast("list[object] | tuple[object, ...] | set[object]", v)
oid = id(iterable_v)
if oid in stack:
logger.debug("Cycle detected while encoding iterable")
return _CYCLE_SENTINEL
stack.add(oid)
try:
seq: list[object] = list(iterable_v)
encoded_list: list[Any] = []
for item in seq:
encoded_list.append(_enc(item, stack, depth + 1))
return encoded_list
finally:
stack.remove(oid)
# Primitives (or unknown objects): ensure JSON-serializable
if isinstance(v, (str, int, float, bool)) or v is None:
return v
# Fallback: stringify unknown objects to avoid JSON serialization errors
try:
return str(v)
except Exception:
return f"<{type(v).__name__}>"
return _enc(value, set(), 0)
def _decode_checkpoint_value(value: Any) -> Any:
"""Recursively decode values previously encoded by _encode_checkpoint_value."""
if isinstance(value, dict):
value_dict = cast(dict[str, Any], value) # encoded form always uses string keys
# Pydantic marker handling
if _PYDANTIC_MARKER in value_dict and "value" in value_dict:
type_key: str | None = value_dict.get(_PYDANTIC_MARKER) # type: ignore[assignment]
raw: Any = value_dict.get("value")
if isinstance(type_key, str):
try:
module_name, class_name = type_key.split(":", 1)
module = importlib.import_module(module_name)
cls: Any = getattr(module, class_name)
if hasattr(cls, "model_validate"):
return cls.model_validate(raw)
except Exception as exc:
logger.debug(
"Failed to decode pydantic model %s: %s; returning raw value",
type_key,
exc,
)
# Dataclass marker handling
if _DATACLASS_MARKER in value_dict and "value" in value_dict:
type_key_dc: str | None = value_dict.get(_DATACLASS_MARKER) # type: ignore[assignment]
raw_dc: Any = value_dict.get("value")
if isinstance(type_key_dc, str):
try:
module_name, class_name = type_key_dc.split(":", 1)
module = importlib.import_module(module_name)
cls_dc: Any = getattr(module, class_name)
decoded_raw = _decode_checkpoint_value(raw_dc)
if isinstance(decoded_raw, dict):
return cls_dc(**decoded_raw)
except Exception as exc:
logger.debug(
"Failed to decode dataclass %s: %s; returning raw value",
type_key_dc,
exc,
)
# Fallback to decoded raw value
return _decode_checkpoint_value(raw_dc)
# Regular dict: decode recursively
decoded: dict[str, Any] = {}
for k_any, v_any in value_dict.items():
decoded[k_any] = _decode_checkpoint_value(v_any)
return decoded
if isinstance(value, list):
# After isinstance check, treat value as list[Any] for decoding
value_list: list[Any] = value # type: ignore[assignment]
return [_decode_checkpoint_value(v_any) for v_any in value_list]
return value
@runtime_checkable
class RunnerContext(Protocol):
"""Protocol for the execution context used by the runner.
A single context that supports messaging, events, and optional checkpointing.
If checkpoint storage is not configured, checkpoint methods may raise.
"""
async def send_message(self, message: Message) -> None:
"""Send a message from the executor to the context.
Args:
message: The message to be sent.
"""
...
async def drain_messages(self) -> dict[str, list[Message]]:
"""Drain all messages from the context.
Returns:
A dictionary mapping executor IDs to lists of messages.
"""
...
async def has_messages(self) -> bool:
"""Check if there are any messages in the context.
Returns:
True if there are messages, False otherwise.
"""
...
async def add_event(self, event: WorkflowEvent) -> None:
"""Add an event to the execution context.
Args:
event: The event to be added.
"""
...
async def drain_events(self) -> list[WorkflowEvent]:
"""Drain all events from the context.
Returns:
A list of events that were added to the context.
"""
...
async def has_events(self) -> bool:
"""Check if there are any events in the context.
Returns:
True if there are events, False otherwise.
"""
...
async def next_event(self) -> WorkflowEvent: # pragma: no cover - interface only
"""Wait for and return the next event emitted by the workflow run."""
...
async def set_state(self, executor_id: str, state: dict[str, Any]) -> None:
"""Set the state for a specific executor.
Args:
executor_id: The ID of the executor whose state is being set.
state: The state data to be set for the executor.
"""
...
async def get_state(self, executor_id: str) -> dict[str, Any] | None:
"""Get the state for a specific executor.
Args:
executor_id: The ID of the executor whose state is being retrieved.
Returns:
The state data for the executor, or None if not found.
"""
...
# Checkpointing capability
def has_checkpointing(self) -> bool:
"""Check if the context supports checkpointing.
Returns:
True if checkpointing is supported, False otherwise.
"""
...
# Checkpointing APIs (optional, enabled by storage)
def set_workflow_id(self, workflow_id: str) -> None:
"""Set the workflow ID for the context."""
...
def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None:
"""Reset the context for a new workflow run."""
...
async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str:
"""Create a checkpoint of the current workflow state.
Args:
metadata: Optional metadata to associate with the checkpoint.
"""
...
async def restore_from_checkpoint(self, checkpoint_id: str) -> bool:
"""Restore the context from a checkpoint.
Args:
checkpoint_id: The ID of the checkpoint to restore from.
Returns:
True if the restoration was successful, False otherwise.
"""
...
async def get_checkpoint_state(self) -> CheckpointState:
"""Get the current state of the context suitable for checkpointing."""
...
async def set_checkpoint_state(self, state: CheckpointState) -> None:
"""Set the state of the context from a checkpoint.
Args:
state: The state data to set for the context.
"""
...
class InProcRunnerContext:
"""In-process execution context for local execution and optional checkpointing."""
def __init__(self, checkpoint_storage: CheckpointStorage | None = None):
"""Initialize the in-process execution context.
Args:
checkpoint_storage: Optional storage to enable checkpointing.
"""
self._messages: defaultdict[str, list[Message]] = defaultdict(list)
# Event queue for immediate streaming of events (e.g., AgentRunUpdateEvent)
self._event_queue: asyncio.Queue[WorkflowEvent] = asyncio.Queue()
# Checkpointing configuration/state
self._checkpoint_storage = checkpoint_storage
self._workflow_id: str | None = None
self._shared_state: dict[str, Any] = {}
self._executor_states: dict[str, dict[str, Any]] = {}
self._iteration_count: int = 0
self._max_iterations: int = 100
async def send_message(self, message: Message) -> None:
self._messages[message.source_id].append(message)
async def drain_messages(self) -> dict[str, list[Message]]:
messages = dict(self._messages)
self._messages.clear()
return messages
async def has_messages(self) -> bool:
return bool(self._messages)
async def add_event(self, event: WorkflowEvent) -> None:
"""Add an event to the context immediately.
Events are enqueued so runners can stream them in real time instead of
waiting for superstep boundaries.
"""
# Filter out empty AgentRunUpdateEvent updates to avoid emitting None/empty chunks
try:
if isinstance(event, AgentRunUpdateEvent):
update = getattr(event, "data", None)
# Skip if no update payload
if not update:
return
# Robust emptiness check: allow either top-level text or any text-bearing content
text_val = getattr(update, "text", None)
contents = getattr(update, "contents", None)
has_text_content = False
if contents:
for c in contents:
if getattr(c, "text", None):
has_text_content = True
break
if not (text_val or has_text_content):
return
except Exception as exc: # pragma: no cover - defensive logging path
# Best-effort filtering only; never block event delivery on filtering errors
logger.debug("Error while filtering event %r: %s", event, exc, exc_info=True)
await self._event_queue.put(event)
async def drain_events(self) -> list[WorkflowEvent]:
"""Drain all currently queued events without blocking for new ones."""
events: list[WorkflowEvent] = []
while True:
try:
events.append(self._event_queue.get_nowait())
except asyncio.QueueEmpty: # type: ignore[attr-defined]
break
return events
async def has_events(self) -> bool:
return not self._event_queue.empty()
async def next_event(self) -> WorkflowEvent:
"""Wait for and return the next event.
Used by the runner to interleave event emission with ongoing iteration work.
"""
return await self._event_queue.get()
async def set_state(self, executor_id: str, state: dict[str, Any]) -> None:
self._executor_states[executor_id] = state
async def get_state(self, executor_id: str) -> dict[str, Any] | None:
return self._executor_states.get(executor_id)
def has_checkpointing(self) -> bool:
return self._checkpoint_storage is not None
def set_workflow_id(self, workflow_id: str) -> None:
self._workflow_id = workflow_id
def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None:
self._messages.clear()
# Clear any pending events (best-effort) by recreating the queue
self._event_queue = asyncio.Queue()
self._shared_state.clear()
self._executor_states.clear()
self._iteration_count = 0
if workflow_shared_state is not None and hasattr(workflow_shared_state, "_state"):
workflow_shared_state._state.clear() # type: ignore[attr-defined]
async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str:
if not self._checkpoint_storage:
raise ValueError("Checkpoint storage not configured")
wf_id = self._workflow_id or str(uuid.uuid4())
self._workflow_id = wf_id
state = await self.get_checkpoint_state()
checkpoint = WorkflowCheckpoint(
workflow_id=wf_id,
messages=state["messages"],
shared_state=state.get("shared_state", {}),
executor_states=state.get("executor_states", {}),
iteration_count=state.get("iteration_count", 0),
max_iterations=state.get("max_iterations", DEFAULT_MAX_ITERATIONS),
metadata=metadata or {},
)
checkpoint_id = await self._checkpoint_storage.save_checkpoint(checkpoint)
logger.info(f"Created checkpoint {checkpoint_id} for workflow {wf_id}'")
return checkpoint_id
async def restore_from_checkpoint(self, checkpoint_id: str) -> bool:
if not self._checkpoint_storage:
raise ValueError("Checkpoint storage not configured")
checkpoint = await self._checkpoint_storage.load_checkpoint(checkpoint_id)
if not checkpoint:
logger.error(f"Checkpoint {checkpoint_id} not found")
return False
state: CheckpointState = {
"messages": checkpoint.messages,
"shared_state": checkpoint.shared_state,
"executor_states": checkpoint.executor_states,
"iteration_count": checkpoint.iteration_count,
"max_iterations": checkpoint.max_iterations,
}
await self.set_checkpoint_state(state)
self._workflow_id = checkpoint.workflow_id
logger.info(f"Restored state from checkpoint {checkpoint_id}'")
return True
async def get_checkpoint_state(self) -> CheckpointState:
serializable_messages: dict[str, list[dict[str, Any]]] = {}
for source_id, message_list in self._messages.items():
serializable_messages[source_id] = [
{
"data": _encode_checkpoint_value(msg.data),
"source_id": msg.source_id,
"target_id": msg.target_id,
"trace_contexts": msg.trace_contexts,
"source_span_ids": msg.source_span_ids,
}
for msg in message_list
]
return {
"messages": serializable_messages,
"shared_state": _encode_checkpoint_value(self._shared_state),
"executor_states": _encode_checkpoint_value(self._executor_states),
"iteration_count": self._iteration_count,
"max_iterations": self._max_iterations,
}
async def set_checkpoint_state(self, state: CheckpointState) -> None:
self._messages.clear()
messages_data = state.get("messages", {})
for source_id, message_list in messages_data.items():
self._messages[source_id] = [
Message(
data=_decode_checkpoint_value(msg.get("data")),
source_id=msg.get("source_id", ""),
target_id=msg.get("target_id"),
trace_contexts=msg.get("trace_contexts"),
source_span_ids=msg.get("source_span_ids"),
)
for msg in message_list
]
# Restore shared_state
decoded_shared_raw = _decode_checkpoint_value(state.get("shared_state", {}))
if isinstance(decoded_shared_raw, dict):
self._shared_state = cast(dict[str, Any], decoded_shared_raw)
else: # fallback to empty dict if corrupted
self._shared_state = {}
# Restore executor_states ensuring value types are dicts
decoded_exec_raw = _decode_checkpoint_value(state.get("executor_states", {}))
if isinstance(decoded_exec_raw, dict):
typed_exec: dict[str, dict[str, Any]] = {}
for k_raw, v_raw in decoded_exec_raw.items(): # type: ignore[assignment]
if isinstance(k_raw, str) and isinstance(v_raw, dict):
# Filter inner dict to string keys only (best-effort)
inner: dict[str, Any] = {}
for inner_k, inner_v in v_raw.items(): # type: ignore[assignment]
if isinstance(inner_k, str):
inner[inner_k] = inner_v
typed_exec[k_raw] = inner
self._executor_states = typed_exec
else:
self._executor_states = {}
self._iteration_count = state.get("iteration_count", 0)
self._max_iterations = state.get("max_iterations", 100)
python/packages/main/agent_framework/_workflow/_sequential.py
+187
View File
@@ -0,0 +1,187 @@
# Copyright (c) Microsoft. All rights reserved.
"""Sequential builder for agent/executor workflows with shared conversation context.
This module provides a high-level, agent-focused API to assemble a sequential
workflow where:
- Participants are a sequence of AgentProtocol instances or Executors
- A shared conversation context (list[ChatMessage]) is passed along the chain
- Agents append their assistant messages to the context
- Custom executors can transform or summarize and return a refined context
- The workflow completes with the final context produced by the last participant
Typical wiring:
input -> _InputToConversation -> participant1 -> (agent? -> _ResponseToConversation) -> ... -> participantN -> _CompleteWithConversation
Notes:
- Participants can mix AgentProtocol and Executor objects
- Agents are auto-wrapped by WorkflowBuilder as AgentExecutor
- AgentExecutor produces AgentExecutorResponse; _ResponseToConversation converts this to list[ChatMessage]
- Non-agent executors must define a handler that consumes `list[ChatMessage]` and sends back
the updated `list[ChatMessage]` via their workflow context
Why include the small internal adapter executors?
- Input normalization ("input-conversation"): ensures the workflow always starts with a
`list[ChatMessage]` regardless of whether callers pass a `str`, a single `ChatMessage`,
or a list. This keeps the first hop strongly typed and avoids boilerplate in participants.
- Agent response adaptation ("to-conversation:<participant>"): agents (via AgentExecutor)
emit `AgentExecutorResponse`. The adapter converts that to a `list[ChatMessage]`
using `full_conversation` so original prompts aren't lost when chaining.
- Explicit completion ("complete"): emits a `WorkflowCompletedEvent` with the final
conversation list, giving a consistent terminal payload shape for both agents and
custom executors.
These adapters are first-class executors by design so they are type-checked at edges,
observable (ExecutorInvoke/Completed events), and easily testable/reusable. Their IDs are
deterministic and self-describing (for example, "to-conversation:writer") to reduce event-log
confusion and to mirror how the concurrent builder uses explicit dispatcher/aggregator nodes.
""" # noqa: E501
import logging
from collections.abc import Sequence
from typing import Any
from agent_framework import AgentProtocol, ChatMessage, Role
from ._events import WorkflowCompletedEvent
from ._executor import (
AgentExecutor,
AgentExecutorResponse,
Executor,
handler,
)
from ._workflow import Workflow, WorkflowBuilder
from ._workflow_context import WorkflowContext
logger = logging.getLogger(__name__)
class _InputToConversation(Executor):
"""Normalizes initial input into a list[ChatMessage] conversation."""
@handler
async def from_str(self, prompt: str, ctx: WorkflowContext[list[ChatMessage]]) -> None:
await ctx.send_message([ChatMessage(Role.USER, text=prompt)])
@handler
async def from_message(self, message: ChatMessage, ctx: WorkflowContext[list[ChatMessage]]) -> None: # type: ignore[name-defined]
await ctx.send_message([message])
@handler
async def from_messages(self, messages: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None: # type: ignore[name-defined]
# Make a copy to avoid mutation downstream
await ctx.send_message(list(messages))
class _ResponseToConversation(Executor):
"""Converts AgentExecutorResponse to list[ChatMessage] conversation for chaining."""
@handler
async def convert(self, response: AgentExecutorResponse, ctx: WorkflowContext[list[ChatMessage]]) -> None:
# Always use full_conversation; AgentExecutor guarantees it is populated.
if response.full_conversation is None: # Defensive: indicates a contract violation
raise RuntimeError("AgentExecutorResponse.full_conversation missing. AgentExecutor must populate it.")
await ctx.send_message(list(response.full_conversation))
class _CompleteWithConversation(Executor):
"""Terminates the workflow by emitting the final conversation context."""
@handler
async def complete(self, conversation: list[ChatMessage], ctx: WorkflowContext[Any]) -> None:
await ctx.add_event(WorkflowCompletedEvent(data=list(conversation)))
class SequentialBuilder:
r"""High-level builder for sequential agent/executor workflows with shared context.
- `participants([...])` accepts a list of AgentProtocol (recommended) or Executor
- The workflow wires participants in order, passing a list[ChatMessage] down the chain
- Agents append their assistant messages to the conversation
- Custom executors can transform/summarize and return a list[ChatMessage]
- The final output is the conversation produced by the last participant
Usage:
```python
from agent_framework import SequentialBuilder
workflow = SequentialBuilder().participants([agent1, agent2, summarizer_exec]).build()
```
"""
def __init__(self) -> None:
self._participants: list[AgentProtocol | Executor] = []
def participants(self, participants: Sequence[AgentProtocol | Executor]) -> "SequentialBuilder":
"""Define the ordered participants for this sequential workflow.
Accepts AgentProtocol instances (auto-wrapped as AgentExecutor) or Executor instances.
Raises if empty or duplicates are provided for clarity.
"""
if not participants:
raise ValueError("participants cannot be empty")
# Defensive duplicate detection
seen_agent_ids: set[int] = set()
seen_executor_ids: set[str] = set()
for p in participants:
if isinstance(p, Executor):
if p.id in seen_executor_ids:
raise ValueError(f"Duplicate executor participant detected: id '{p.id}'")
seen_executor_ids.add(p.id)
else:
# Treat non-Executor as agent-like (AgentProtocol). Structural checks can be brittle at runtime.
pid = id(p)
if pid in seen_agent_ids:
raise ValueError("Duplicate agent participant detected (same agent instance provided twice)")
seen_agent_ids.add(pid)
self._participants = list(participants)
return self
def build(self) -> Workflow:
"""Build and validate the sequential workflow.
Wiring pattern:
- _InputToConversation normalizes the initial input into list[ChatMessage]
- For each participant in order:
- If Agent (or AgentExecutor): pass conversation to the agent, then convert response
to conversation via _ResponseToConversation
- Else (custom Executor): pass conversation directly to the executor
- _CompleteWithConversation emits WorkflowCompletedEvent with the final conversation
"""
if not self._participants:
raise ValueError("No participants provided. Call .participants([...]) first.")
# Internal nodes
input_conv = _InputToConversation(id="input-conversation")
complete = _CompleteWithConversation(id="complete")
builder = WorkflowBuilder()
builder.set_start_executor(input_conv)
# Start of the chain is the input normalizer
prior: Executor | AgentProtocol = input_conv
for p in self._participants:
# Agent-like branch: either explicitly an AgentExecutor or any non-AgentExecutor
if not (isinstance(p, Executor) and not isinstance(p, AgentExecutor)):
# input conversation -> (agent) -> response -> conversation
builder.add_edge(prior, p)
# Give the adapter a deterministic, self-describing id
label: str
label = p.id if isinstance(p, Executor) else getattr(p, "name", None) or p.__class__.__name__
resp_to_conv = _ResponseToConversation(id=f"to-conversation:{label}")
builder.add_edge(p, resp_to_conv)
prior = resp_to_conv
elif isinstance(p, Executor):
# Custom executor operates on list[ChatMessage]
builder.add_edge(prior, p)
prior = p
else: # pragma: no cover - defensive
raise TypeError(f"Unsupported participant type: {type(p).__name__}")
# Terminate with the final conversation
builder.add_edge(prior, complete)
return builder.build()
python/packages/main/agent_framework/_workflow/_shared_state.py
+69
View File
@@ -0,0 +1,69 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any
class SharedState:
"""A class to manage shared state in a workflow."""
def __init__(self) -> None:
"""Initialize the shared state."""
self._state: dict[str, Any] = {}
self._shared_state_lock = asyncio.Lock()
async def set(self, key: str, value: Any) -> None:
"""Set a value in the shared state."""
async with self._shared_state_lock:
await self.set_within_hold(key, value)
async def get(self, key: str) -> Any:
"""Get a value from the shared state."""
async with self._shared_state_lock:
return await self.get_within_hold(key)
async def has(self, key: str) -> bool:
"""Check if a key exists in the shared state."""
async with self._shared_state_lock:
return await self.has_within_hold(key)
async def delete(self, key: str) -> None:
"""Delete a key from the shared state."""
async with self._shared_state_lock:
await self.delete_within_hold(key)
@asynccontextmanager
async def hold(self) -> AsyncIterator["SharedState"]:
"""Context manager to hold the shared state lock for multiple operations.
Usage:
async with shared_state.hold():
await shared_state.set_within_hold("key", value)
value = await shared_state.get_within_hold("key")
"""
async with self._shared_state_lock:
yield self
# Unsafe methods that don't acquire locks (for use within hold() context)
async def set_within_hold(self, key: str, value: Any) -> None:
"""Set a value without acquiring the lock (unsafe - use within hold() context)."""
self._state[key] = value
async def get_within_hold(self, key: str) -> Any:
"""Get a value without acquiring the lock (unsafe - use within hold() context)."""
if key not in self._state:
raise KeyError(f"Key '{key}' not found in shared state.")
return self._state[key]
async def has_within_hold(self, key: str) -> bool:
"""Check if a key exists without acquiring the lock (unsafe - use within hold() context)."""
return key in self._state
async def delete_within_hold(self, key: str) -> None:
"""Delete a key without acquiring the lock (unsafe - use within hold() context)."""
if key in self._state:
del self._state[key]
else:
raise KeyError(f"Key '{key}' not found in shared state.")
python/packages/main/agent_framework/_workflow/_telemetry.py
+310
View File
@@ -0,0 +1,310 @@
# Copyright (c) Microsoft. All rights reserved.
from enum import Enum
from typing import TYPE_CHECKING, Any, ClassVar
from opentelemetry.trace import Link, NoOpTracer, SpanKind, StatusCode, get_current_span, get_tracer
from opentelemetry.trace.span import SpanContext
from opentelemetry.util.types import Attributes
from agent_framework._pydantic import AFBaseSettings
if TYPE_CHECKING:
from ._workflow import Workflow
class EdgeGroupDeliveryStatus(Enum):
"""Enum for edge group delivery status values."""
DELIVERED = "delivered"
DROPPED_TYPE_MISMATCH = "dropped type mismatch"
DROPPED_TARGET_MISMATCH = "dropped target mismatch"
DROPPED_CONDITION_FALSE = "dropped condition evaluated to false"
EXCEPTION = "exception"
BUFFERED = "buffered"
# Span name constants
_WORKFLOW_BUILD_SPAN = "workflow.build"
_WORKFLOW_RUN_SPAN = "workflow.run"
_EXECUTOR_PROCESS_SPAN = "executor.process"
_MESSAGE_SEND_SPAN = "message.send"
_EDGE_GROUP_PROCESS_SPAN = "edge_group.process"
class WorkflowDiagnosticSettings(AFBaseSettings):
"""Settings for workflow tracing diagnostics."""
env_prefix: ClassVar[str] = "AGENT_FRAMEWORK_WORKFLOW_"
enable_otel: bool = False
@property
def ENABLED(self) -> bool:
return self.enable_otel
class WorkflowTracer:
"""Central tracing coordinator for workflow system.
Manages OpenTelemetry span creation and relationships for:
- Workflow build spans (workflow.build)
- Workflow execution spans (workflow.run)
- Executor processing spans (executor.process)
- Message sending spans (message.send)
- Edge group processing spans (edge_group.process)
Implements span linking for causality without unwanted nesting.
"""
def __init__(self) -> None:
self.settings = WorkflowDiagnosticSettings()
self.tracer = get_tracer("agent_framework") if self.settings.ENABLED else NoOpTracer()
@property
def enabled(self) -> bool:
return self.settings.ENABLED
def create_workflow_run_span(self, workflow: "Workflow") -> Any:
"""Create a workflow execution span."""
attributes: dict[str, str | int] = {
"workflow.id": workflow.id,
}
return self.tracer.start_as_current_span(_WORKFLOW_RUN_SPAN, kind=SpanKind.INTERNAL, attributes=attributes)
def create_workflow_build_span(self) -> Any:
"""Create a workflow build span."""
return self.tracer.start_as_current_span(_WORKFLOW_BUILD_SPAN, kind=SpanKind.INTERNAL)
def create_processing_span(
self,
executor_id: str,
executor_type: str,
message_type: str,
source_trace_contexts: list[dict[str, str]] | None = None,
source_span_ids: list[str] | None = None,
) -> Any:
"""Create an executor processing span with optional links to source spans.
Processing spans are created as children of the current workflow span and
linked (not nested) to the source publishing spans for causality tracking.
This supports multiple links for fan-in scenarios.
"""
# Create links to source spans for causality without nesting
links = []
if source_trace_contexts and source_span_ids:
# Create links for all source spans (supporting fan-in with multiple sources)
for trace_context, span_id in zip(source_trace_contexts, source_span_ids, strict=False):
try:
# Extract trace and span IDs from the trace context
# This is a simplified approach - in production you'd want more robust parsing
traceparent = trace_context.get("traceparent", "")
if traceparent:
# traceparent format: "00-{trace_id}-{parent_span_id}-{trace_flags}"
parts = traceparent.split("-")
if len(parts) >= 3:
trace_id_hex = parts[1]
# Use the source_span_id that was saved from the publishing span
# Create span context for linking
span_context = SpanContext(
trace_id=int(trace_id_hex, 16),
span_id=int(span_id, 16),
is_remote=True,
)
links.append(Link(span_context))
except (ValueError, TypeError, AttributeError):
# If linking fails, continue without link (graceful degradation)
pass
return self.tracer.start_as_current_span(
_EXECUTOR_PROCESS_SPAN,
kind=SpanKind.INTERNAL,
attributes={
"executor.id": executor_id,
"executor.type": executor_type,
"message.type": message_type,
},
links=links,
)
def create_sending_span(self, message_type: str, target_executor_id: str | None = None) -> Any:
"""Create a message send span.
Sending spans are created as children of the current processing span
to track message emission for distributed tracing.
"""
attributes: dict[str, str] = {
"message.type": message_type,
}
if target_executor_id is not None:
attributes["message.destination_executor_id"] = target_executor_id
return self.tracer.start_as_current_span(
_MESSAGE_SEND_SPAN,
kind=SpanKind.PRODUCER,
attributes=attributes,
)
def create_edge_group_processing_span(
self,
edge_group_type: str,
edge_group_id: str | None = None,
message_source_id: str | None = None,
message_target_id: str | None = None,
source_trace_contexts: list[dict[str, str]] | None = None,
source_span_ids: list[str] | None = None,
) -> Any:
"""Create an edge group processing span with optional links to source spans.
Edge group processing spans track the processing operations in edge runners
before message delivery, including condition checking and routing decisions.
Links to source spans provide causality tracking without unwanted nesting.
Args:
edge_group_type: The type of the edge group (class name).
edge_group_id: The unique ID of the edge group.
message_source_id: The source ID of the message being processed.
message_target_id: The target ID of the message being processed.
source_trace_contexts: Optional trace contexts from source spans for linking.
source_span_ids: Optional source span IDs for linking.
"""
attributes: dict[str, str] = {
"edge_group.type": edge_group_type,
}
if edge_group_id is not None:
attributes["edge_group.id"] = edge_group_id
if message_source_id is not None:
attributes["message.source_id"] = message_source_id
if message_target_id is not None:
attributes["message.target_id"] = message_target_id
# Create links to source spans for causality without nesting
links = []
if source_trace_contexts and source_span_ids:
# Create links for all source spans (supporting fan-in with multiple sources)
for trace_context, span_id in zip(source_trace_contexts, source_span_ids, strict=False):
try:
# Extract trace and span IDs from the trace context
# This is a simplified approach - in production you'd want more robust parsing
traceparent = trace_context.get("traceparent", "")
if traceparent:
# traceparent format: "00-{trace_id}-{parent_span_id}-{trace_flags}"
parts = traceparent.split("-")
if len(parts) >= 3:
trace_id_hex = parts[1]
# Use the source_span_id that was saved from the publishing span
# Create span context for linking
span_context = SpanContext(
trace_id=int(trace_id_hex, 16),
span_id=int(span_id, 16),
is_remote=True,
)
links.append(Link(span_context))
except (ValueError, TypeError, AttributeError):
# If linking fails, continue without link (graceful degradation)
pass
return self.tracer.start_as_current_span(
_EDGE_GROUP_PROCESS_SPAN,
kind=SpanKind.INTERNAL,
attributes=attributes,
links=links,
)
def set_edge_group_span_attributes(self, delivered: bool, delivery_status: EdgeGroupDeliveryStatus) -> None:
"""Set edge group span attributes for delivery status.
Args:
delivered: Whether the message was delivered.
delivery_status: The delivery status from EdgeGroupDeliveryStatus enum.
"""
span = get_current_span()
if span and span.is_recording():
span.set_attributes({
"edge_group.delivered": delivered,
"edge_group.delivery_status": delivery_status.value,
})
def add_workflow_event(self, event_name: str, attributes: Attributes | None = None) -> None:
"""Add an event to the current workflow span.
Args:
event_name: Name of the event (e.g., "workflow.started", "workflow.completed")
attributes: Optional attributes to attach to the event
"""
span = get_current_span()
if span and span.is_recording():
span.add_event(event_name, attributes)
def add_workflow_error_event(self, error: Exception, attributes: Attributes | None = None) -> None:
"""Add an error event to the current workflow span.
Args:
error: The exception that occurred
attributes: Optional additional attributes to attach to the event
"""
span = get_current_span()
if span and span.is_recording():
event_attributes: dict[str, str | bool | int | float] = {
"error.message": str(error),
"error.type": type(error).__name__,
}
if attributes:
# Safely merge attributes, ensuring type compatibility
for key, value in attributes.items():
if isinstance(value, (str, bool, int, float)):
event_attributes[key] = value
span.add_event("workflow.error", event_attributes)
span.set_status(StatusCode.ERROR, str(error))
def set_workflow_build_span_attributes(self, workflow: "Workflow") -> None:
"""Set workflow attributes on the current span.
Args:
workflow: The workflow instance to extract attributes from
"""
span = get_current_span()
if span and span.is_recording():
span.set_attributes({
"workflow.id": workflow.id,
"workflow.definition": workflow.model_dump_json(by_alias=True),
})
def add_build_event(self, event_name: str, attributes: Attributes | None = None) -> None:
"""Add an event to the current workflow build span.
Args:
event_name: Name of the build event (e.g., "build.started", "build.validation_completed")
attributes: Optional attributes to attach to the event
"""
span = get_current_span()
if span and span.is_recording():
span.add_event(event_name, attributes)
def add_build_error_event(self, error: Exception, attributes: Attributes | None = None) -> None:
"""Add an error event to the current workflow build span.
Args:
error: The exception that occurred during build
attributes: Optional additional attributes to attach to the event
"""
span = get_current_span()
if span and span.is_recording():
event_attributes: dict[str, str | bool | int | float] = {
"build.error.message": str(error),
"build.error.type": type(error).__name__,
}
if attributes:
# Safely merge attributes, ensuring type compatibility
for key, value in attributes.items():
if isinstance(value, (str, bool, int, float)):
event_attributes[key] = value
span.add_event("build.error", event_attributes)
span.set_status(StatusCode.ERROR, str(error))
# Global workflow tracer instance
workflow_tracer = WorkflowTracer()
python/packages/main/agent_framework/_workflow/_typing_utils.py
+78
View File
@@ -0,0 +1,78 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any, Union, get_args, get_origin
def is_instance_of(data: Any, target_type: type) -> bool:
"""Check if the data is an instance of the target type.
Args:
data (Any): The data to check.
target_type (type): The type to check against.
Returns:
bool: True if data is an instance of target_type, False otherwise.
"""
# Case 0: target_type is Any - always return True
if target_type is Any:
return True
origin = get_origin(target_type)
args = get_args(target_type)
# Case 1: origin is None, meaning target_type is not a generic type
if origin is None:
return isinstance(data, target_type)
# Case 2: target_type is Optional[T] or Union[T1, T2, ...]
# Optional[T] is really just as Union[T, None]
if origin is Union:
return any(is_instance_of(data, arg) for arg in args)
# Case 3: target_type is a generic type
if origin in [list, set]:
return isinstance(data, origin) and all(is_instance_of(item, args[0]) for item in data) # type: ignore
# Case 4: target_type is a tuple
if origin is tuple:
if len(args) == 1 and args[0] is Ellipsis: # Tuple[...] case
return isinstance(data, tuple)
return (
isinstance(data, tuple)
and len(data) == len(args) # type: ignore
and all(is_instance_of(item, arg) for item, arg in zip(data, args, strict=False)) # type: ignore
)
# Case 5: target_type is a dict
if origin is dict:
return isinstance(data, dict) and all(
is_instance_of(key, args[0]) and is_instance_of(value, args[1])
for key, value in data.items() # type: ignore
)
# Case 6: target_type is RequestResponse[T, U] - validate generic parameters
if origin and hasattr(origin, "__name__") and origin.__name__ == "RequestResponse":
if not isinstance(data, origin):
return False
# Validate generic parameters for RequestResponse[TRequest, TResponse]
if len(args) >= 2:
request_type, response_type = args[0], args[1]
# Check if the original_request matches TRequest and data matches TResponse
if (
hasattr(data, "original_request")
and data.original_request is not None
and not is_instance_of(data.original_request, request_type)
):
return False
if hasattr(data, "data") and data.data is not None and not is_instance_of(data.data, response_type):
return False
return True
# Case 7: Other custom generic classes - check origin type only
# For generic classes, we check if data is an instance of the origin type
# We don't validate the generic parameters at runtime since that's handled by type system
if origin and hasattr(origin, "__name__"):
return isinstance(data, origin)
# Fallback: if we reach here, we assume data is an instance of the target_type
return isinstance(data, target_type)
python/packages/main/agent_framework/_workflow/_validation.py
+809
View File
@@ -0,0 +1,809 @@
# Copyright (c) Microsoft. All rights reserved.
import inspect
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
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."""
EDGE_DUPLICATION = "EDGE_DUPLICATION"
TYPE_COMPATIBILITY = "TYPE_COMPATIBILITY"
GRAPH_CONNECTIVITY = "GRAPH_CONNECTIVITY"
HANDLER_OUTPUT_ANNOTATION = "HANDLER_OUTPUT_ANNOTATION"
INTERCEPTOR_CONFLICT = "INTERCEPTOR_CONFLICT"
class WorkflowValidationError(Exception):
"""Base exception for workflow validation errors."""
def __init__(self, message: str, validation_type: ValidationTypeEnum):
super().__init__(message)
self.message = message
self.validation_type = validation_type
def __str__(self) -> str:
return f"[{self.validation_type.value}] {self.message}"
class EdgeDuplicationError(WorkflowValidationError):
"""Exception raised when duplicate edges are detected in the workflow."""
def __init__(self, edge_id: str):
super().__init__(
message=f"Duplicate edge detected: {edge_id}. Each edge in the workflow must be unique.",
validation_type=ValidationTypeEnum.EDGE_DUPLICATION,
)
self.edge_id = edge_id
class TypeCompatibilityError(WorkflowValidationError):
"""Exception raised when type incompatibility is detected between connected executors."""
def __init__(
self,
source_executor_id: str,
target_executor_id: str,
source_types: list[type[Any]],
target_types: list[type[Any]],
):
# Use a placeholder for incompatible types - will be computed in WorkflowGraphValidator
super().__init__(
message=f"Type incompatibility between executors '{source_executor_id}' -> '{target_executor_id}'. "
f"Source executor outputs types {[str(t) for t in source_types]} but target executor "
f"can only handle types {[str(t) for t in target_types]}.",
validation_type=ValidationTypeEnum.TYPE_COMPATIBILITY,
)
self.source_executor_id = source_executor_id
self.target_executor_id = target_executor_id
self.source_types = source_types
self.target_types = target_types
class GraphConnectivityError(WorkflowValidationError):
"""Exception raised when graph connectivity issues are detected."""
def __init__(self, message: str):
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
class InterceptorConflictError(WorkflowValidationError):
"""Exception raised when multiple executors intercept the same request type from the same sub-workflow."""
def __init__(self, message: str):
super().__init__(message, validation_type=ValidationTypeEnum.INTERCEPTOR_CONFLICT)
# endregion
# region Workflow Graph Validator
class WorkflowGraphValidator:
"""Validator for workflow graphs.
This validator performs multiple validation checks:
1. Edge duplication validation
2. Type compatibility validation between connected executors
3. Graph connectivity validation
"""
def __init__(self) -> None:
self._edges: list[Edge] = []
self._executors: dict[str, Executor] = {}
# region Core Validation Methods
def validate_workflow(
self, edge_groups: Sequence[EdgeGroup], executors: dict[str, Executor], start_executor: Executor | str
) -> None:
"""Validate the entire workflow graph.
Args:
edge_groups: list of edge groups in the workflow
executors: Map of executor IDs to executor instances
start_executor: The starting executor (can be instance or ID)
Raises:
WorkflowValidationError: If any validation fails
"""
self._executors = executors
self._edges = [edge for group in edge_groups for edge in group.edges]
self._edge_groups = edge_groups
# If only the start executor exists, add it to the executor map
# Handle the special case where the workflow consists of only a single executor and no edges.
# In this scenario, the executor map will be empty because there are no edge groups to reference executors.
# Adding the start executor to the map ensures that single-executor workflows (without any edges) are supported,
# allowing validation and execution to proceed for workflows that do not require inter-executor communication.
if not self._executors and start_executor and isinstance(start_executor, Executor):
self._executors[start_executor.id] = start_executor
# Validate that start_executor exists in the graph
# It should because we check for it in the WorkflowBuilder
# but we do it here for completeness.
start_executor_id = start_executor.id if isinstance(start_executor, Executor) else start_executor
if start_executor_id not in self._executors:
raise GraphConnectivityError(f"Start executor '{start_executor_id}' is not present in the workflow graph")
# Additional presence verification:
# A start executor that is only injected via the builder (present in the executors map)
# but not referenced by any edge while other executors ARE referenced indicates a
# configuration error: the chosen start node is effectively disconnected / unknown to the
# defined graph topology. For single-node workflows (no edges) we allow the start executor
# to stand alone (handled above when we inject it into the map). We perform this refined
# check only when there is at least one edge group defined.
if self._edges: # Only evaluate when the workflow defines edges
edge_executor_ids: set[str] = set()
for _e in self._edges:
edge_executor_ids.add(_e.source_id)
edge_executor_ids.add(_e.target_id)
if start_executor_id not in edge_executor_ids:
raise GraphConnectivityError(
f"Start executor '{start_executor_id}' is not present in the workflow graph"
)
# 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()
self._validate_handler_ambiguity()
self._validate_dead_ends()
self._validate_cycles()
self._validate_interceptor_uniqueness()
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.__class__):
if attr_name.startswith("_"):
continue
# 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.__class__, 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)
type_items: list[Any]
type_items = list(get_args(t_out)) if union_origin in (Union, UnionType) else [t_out]
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,
handler_name,
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
def _validate_edge_duplication(self) -> None:
"""Validate that there are no duplicate edges in the workflow.
Raises:
EdgeDuplicationError: If duplicate edges are found
"""
seen_edge_ids: set[str] = set()
for edge in self._edges:
edge_id = edge.id
if edge_id in seen_edge_ids:
raise EdgeDuplicationError(edge_id)
seen_edge_ids.add(edge_id)
def _validate_type_compatibility(self) -> None:
"""Validate type compatibility between connected executors.
This checks that the output types of source executors are compatible
with the input types expected by target executors.
Raises:
TypeCompatibilityError: If type incompatibility is detected
"""
for edge_group in self._edge_groups:
for edge in edge_group.edges:
self._validate_edge_type_compatibility(edge, edge_group)
def _validate_edge_type_compatibility(self, edge: Edge, edge_group: EdgeGroup) -> None:
"""Validate type compatibility for a specific edge.
This checks that the output types of the source executor are compatible
with the input types expected by the target executor.
Args:
edge: The edge to validate
edge_group: The edge group containing this edge
Raises:
TypeCompatibilityError: If type incompatibility is detected
"""
source_executor = self._executors[edge.source_id]
target_executor = self._executors[edge.target_id]
# Get output types from source executor
source_output_types = self._get_executor_output_types(source_executor)
# Get input types from target executor
target_input_types = self._get_executor_input_types(target_executor)
# If either executor has no type information, log warning and skip validation
# This allows for dynamic typing scenarios but warns about reduced validation coverage
if not source_output_types or not target_input_types:
# Suppress warnings for built-in workflow components where dynamic typing is expected
try:
from ._executor import RequestInfoExecutor, WorkflowExecutor # local import to avoid cycles
builtin_types = (RequestInfoExecutor, WorkflowExecutor)
except Exception:
builtin_types = tuple() # type: ignore[assignment]
if not source_output_types and not isinstance(source_executor, builtin_types):
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 WorkflowContext[T] generics in handlers for better validation."
)
if not target_input_types and not isinstance(target_executor, builtin_types):
logger.warning(
f"Executor '{target_executor.id}' has no input type annotations. "
f"Type compatibility validation will be skipped for edges to this executor. "
f"Consider adding type annotations to message handler parameters for better validation."
)
return
# Check if any source output type is compatible with any target input type
compatible = False
compatible_pairs: list[tuple[type[Any], type[Any]]] = []
for source_type in source_output_types:
for target_type in target_input_types:
if isinstance(edge_group, FanInEdgeGroup):
# If the edge is part of an edge group, the target expects a list of data types
if self._is_type_compatible(list[source_type], target_type): # type: ignore[valid-type]
compatible = True
compatible_pairs.append((list[source_type], target_type)) # type: ignore[valid-type]
else:
if self._is_type_compatible(source_type, target_type):
compatible = True
compatible_pairs.append((source_type, target_type))
# Log successful type compatibility for debugging
if compatible:
logger.debug(
f"Type compatibility validated for edge '{source_executor.id}' -> '{target_executor.id}'. "
f"Compatible type pairs: {[(str(s), str(t)) for s, t in compatible_pairs]}"
)
if not compatible:
# Enhanced error with more detailed information
raise TypeCompatibilityError(
source_executor.id,
target_executor.id,
source_output_types,
target_input_types,
)
def _get_executor_output_types(self, executor: Executor) -> list[type[Any]]:
"""Extract output types from an executor's message handlers.
Args:
executor: The executor to analyze
Returns:
list of types that this executor can output
"""
output_types: list[type[Any]] = []
for attr_name in dir(executor.__class__):
if attr_name.startswith("_"):
continue
try:
attr = getattr(executor.__class__, attr_name)
if callable(attr) and hasattr(attr, "_handler_spec"):
handler_spec = attr._handler_spec # type: ignore
handler_output_types = handler_spec.get("output_types", [])
output_types.extend(handler_output_types)
except AttributeError:
# Skip attributes that may not be accessible
continue
# Also include intercepted request types as potential outputs
# since @intercepts_request methods can forward requests
if hasattr(executor, "_request_interceptors"):
for request_type in executor._request_interceptors:
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]]:
"""Extract input types from an executor's message handlers.
Args:
executor: The executor to analyze
Returns:
list of types that this executor can handle as input
"""
input_types: list[type[Any]] = []
# Access the private _handlers attribute to get input types
if hasattr(executor, "_handlers"):
input_types.extend(executor._handlers.keys()) # type: ignore
return input_types
# endregion
# region Graph Connectivity Validation
def _validate_graph_connectivity(self, start_executor_id: str) -> None:
"""Validate graph connectivity and detect potential issues.
This performs several checks:
- Detects unreachable executors from the start node
- Detects isolated executors (no incoming or outgoing edges)
- Warns about potential infinite loops
Args:
start_executor_id: The ID of the starting executor
Raises:
GraphConnectivityError: If connectivity issues are detected
"""
# Build adjacency list for the graph
graph: dict[str, list[str]] = defaultdict(list)
all_executors = set(self._executors.keys())
for edge in self._edges:
graph[edge.source_id].append(edge.target_id)
# Find reachable nodes from start
reachable = self._find_reachable_nodes(graph, start_executor_id)
# Check for unreachable executors
unreachable = all_executors - reachable
if unreachable:
raise GraphConnectivityError(
f"The following executors are unreachable from the start executor '{start_executor_id}': "
f"{sorted(unreachable)}. This may indicate a disconnected workflow graph."
)
# Check for isolated executors (no edges)
isolated_executors: list[str] = []
for executor_id in all_executors:
has_incoming = any(edge.target_id == executor_id for edge in self._edges)
has_outgoing = any(edge.source_id == executor_id for edge in self._edges)
if not has_incoming and not has_outgoing and executor_id != start_executor_id:
isolated_executors.append(executor_id)
if isolated_executors:
raise GraphConnectivityError(
f"The following executors are isolated (no incoming or outgoing edges): "
f"{sorted(isolated_executors)}. Isolated executors will never be executed."
)
def _find_reachable_nodes(self, graph: dict[str, list[str]], start: str) -> set[str]:
"""Find all nodes reachable from the start node using DFS.
Args:
graph: Adjacency list representation of the graph
start: Starting node ID
Returns:
Set of reachable node IDs
"""
visited: set[str] = set()
stack = [start]
while stack:
node = stack.pop()
if node not in visited:
visited.add(node)
stack.extend(graph[node])
return visited
# endregion
# region Additional Validation Scenarios
def _validate_self_loops(self) -> None:
"""Detect and log self-loops (edges from executor to itself).
Self-loops might indicate recursive processing which could be intentional
but should be highlighted for review.
"""
self_loops = [edge for edge in self._edges if edge.source_id == edge.target_id]
for edge in self_loops:
logger.warning(
f"Self-loop detected: Executor '{edge.source_id}' connects to itself. "
f"This may cause infinite recursion if not properly handled with conditions."
)
def _validate_handler_ambiguity(self) -> None:
"""Check for potential ambiguity in message handlers.
Warns when executors have multiple handlers that could handle the same type,
which might lead to unexpected behavior.
"""
for executor_id, executor in self._executors.items():
input_types = self._get_executor_input_types(executor)
# Check for duplicate input types
seen_types: set[type[Any]] = set()
duplicate_types: set[type[Any]] = set()
for input_type in input_types:
if input_type in seen_types:
duplicate_types.add(input_type)
seen_types.add(input_type)
if duplicate_types:
logger.warning(
f"Executor '{executor_id}' has multiple handlers for the same input types: "
f"{[str(t) for t in duplicate_types]}. This may lead to ambiguous message routing."
)
def _validate_dead_ends(self) -> None:
"""Identify executors that have no outgoing edges (potential dead ends).
These might be intentional final nodes or could indicate missing connections.
"""
executors_with_outgoing = {edge.source_id for edge in self._edges}
all_executor_ids = set(self._executors.keys())
dead_ends = all_executor_ids - executors_with_outgoing
if dead_ends:
logger.info(
f"Dead-end executors detected (no outgoing edges): {sorted(dead_ends)}. "
f"Verify these are intended as final nodes in the workflow."
)
def _validate_cycles(self) -> None:
"""Detect cycles in the workflow graph.
Cycles might be intentional for iterative processing but should be flagged
for review to ensure proper termination conditions exist.
"""
# Build adjacency list
graph: dict[str, list[str]] = defaultdict(list)
for edge in self._edges:
graph[edge.source_id].append(edge.target_id)
# Use DFS to detect cycles
white = set(self._executors.keys()) # Unvisited
gray: set[str] = set() # Currently being processed
black: set[str] = set() # Completely processed
def has_cycle(node: str) -> bool:
if node in gray: # Back edge found - cycle detected
return True
if node in black: # Already processed
return False
# Mark as being processed
white.discard(node)
gray.add(node)
# Visit neighbors
for neighbor in graph[node]:
if has_cycle(neighbor):
return True
# Mark as completely processed
gray.discard(node)
black.add(node)
return False
# Check for cycles starting from any unvisited node
cycle_detected = False
while white and not cycle_detected:
start_node = next(iter(white))
if has_cycle(start_node):
cycle_detected = True
if cycle_detected:
logger.warning(
"Cycle detected in the workflow graph. "
"Ensure proper termination conditions exist to prevent infinite loops."
)
def _validate_interceptor_uniqueness(self) -> None:
"""Validate that only one executor intercepts a given request type from a specific sub-workflow.
This prevents non-deterministic behavior where multiple executors could intercept
the same request type from the same sub-workflow.
"""
from ._executor import WorkflowExecutor
# Find all WorkflowExecutor instances in the workflow
workflow_executors: dict[str, WorkflowExecutor] = {}
for executor_id, executor in self._executors.items():
if isinstance(executor, WorkflowExecutor):
workflow_executors[executor_id] = executor
# For each WorkflowExecutor, check which executors can intercept its requests
for workflow_id, _workflow_executor in workflow_executors.items():
# Map of request_type -> list of intercepting executor IDs
interceptors_by_type: dict[type | str, list[str]] = {}
# Find all executors that have edges from this WorkflowExecutor
# These are potential interceptors
for edge in self._edges:
if edge.source_id == workflow_id:
target_executor = self._executors.get(edge.target_id)
if target_executor and hasattr(target_executor, "_request_interceptors"):
# Check what request types this executor intercepts
for request_type, interceptor_list in target_executor._request_interceptors.items():
# Check if any interceptor is scoped to this workflow or unscoped
for interceptor_info in interceptor_list:
from_workflow = interceptor_info.get("from_workflow")
# If unscoped or specifically scoped to this workflow
if from_workflow is None or from_workflow == workflow_id:
if request_type not in interceptors_by_type:
interceptors_by_type[request_type] = []
interceptors_by_type[request_type].append(edge.target_id)
# Check for duplicates
for request_type, executor_ids in interceptors_by_type.items():
unique_executors = list(set(executor_ids)) # Remove duplicates from same executor
if len(unique_executors) > 1:
type_name = request_type.__name__ if isinstance(request_type, type) else str(request_type)
raise InterceptorConflictError(
f"Multiple executors intercept the same request type '{type_name}' "
f"from sub-workflow '{workflow_id}': {', '.join(unique_executors)}. "
f"Only one executor should intercept a given request type from a specific sub-workflow "
f"to ensure deterministic behavior."
)
# endregion
# region Type Compatibility Utilities
@staticmethod
def _is_type_compatible(source_type: type[Any], target_type: type[Any]) -> bool:
"""Check if source_type is compatible with target_type."""
# Handle Any type
if source_type is Any or target_type is Any:
return True
# Handle exact match
if source_type == target_type:
return True
# Handle inheritance
try:
if inspect.isclass(source_type) and inspect.isclass(target_type):
return issubclass(source_type, target_type)
except TypeError:
# Handle generic types that can't be used with issubclass
pass
# Handle Union types
source_origin = get_origin(source_type)
target_origin = get_origin(target_type)
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 in (Union, UnionType):
source_args = get_args(source_type)
return all(WorkflowGraphValidator._is_type_compatible(arg, target_type) for arg in source_args)
# Handle generic types
if source_origin is not None and target_origin is not None and source_origin == target_origin:
source_args = get_args(source_type)
target_args = get_args(target_type)
if len(source_args) == len(target_args):
return all(
WorkflowGraphValidator._is_type_compatible(s_arg, t_arg)
for s_arg, t_arg in zip(source_args, target_args, strict=True)
)
return False
# endregion
# endregion
def validate_workflow_graph(
edge_groups: Sequence[EdgeGroup], executors: dict[str, Executor], start_executor: Executor | str
) -> None:
"""Convenience function to validate a workflow graph.
Args:
edge_groups: list of edge groups in the workflow
executors: Map of executor IDs to executor instances
start_executor: The starting executor (can be instance or ID)
Raises:
WorkflowValidationError: If any validation fails
"""
validator = WorkflowGraphValidator()
validator.validate_workflow(edge_groups, executors, start_executor)
python/packages/main/agent_framework/_workflow/_viz.py
+342
View File
@@ -0,0 +1,342 @@
# Copyright (c) Microsoft. All rights reserved.
import hashlib
import re
import tempfile
import uuid
from pathlib import Path
from typing import Literal
from ._edge import FanInEdgeGroup
from ._workflow import Workflow
# Import of WorkflowExecutor is performed lazily inside methods to avoid cycles
"""Workflow visualization module using graphviz."""
class WorkflowViz:
"""A class for visualizing workflows using graphviz."""
def __init__(self, workflow: Workflow):
"""Initialize the WorkflowViz with a workflow.
Args:
workflow: The workflow to visualize.
"""
self._workflow = workflow
def to_digraph(self) -> str:
"""Export the workflow as a DOT format digraph string.
Returns:
A string representation of the workflow in DOT format.
"""
lines = ["digraph Workflow {"]
lines.append(" rankdir=TD;") # Top to bottom layout
lines.append(" node [shape=box, style=filled, fillcolor=lightblue];")
lines.append(" edge [color=black, arrowhead=vee];")
lines.append("")
# Emit the top-level workflow nodes/edges
self._emit_workflow_digraph(self._workflow, lines, indent=" ")
# Emit sub-workflows hosted by WorkflowExecutor as nested clusters
self._emit_sub_workflows_digraph(self._workflow, lines, indent=" ")
lines.append("}")
return "\n".join(lines)
def export(self, format: Literal["svg", "png", "pdf", "dot"] = "svg", filename: str | None = None) -> str:
"""Export the workflow visualization to a file or return the file path.
Args:
format: The output format. Supported formats: 'svg', 'png', 'pdf', 'dot'.
filename: Optional filename to save the output. If None, creates a temporary file.
Returns:
The path to the saved file.
Raises:
ImportError: If graphviz is not installed.
ValueError: If an unsupported format is specified.
"""
# Validate format first
if format not in ["svg", "png", "pdf", "dot"]:
raise ValueError(f"Unsupported format: {format}. Supported formats: svg, png, pdf, dot")
if format == "dot":
content = self.to_digraph()
if filename:
with open(filename, "w", encoding="utf-8") as f:
f.write(content)
return filename
# Create temporary file for dot format
with tempfile.NamedTemporaryFile(mode="w", suffix=".dot", delete=False, encoding="utf-8") as temp_file:
temp_file.write(content)
return temp_file.name
try:
import graphviz # type: ignore
except ImportError as e:
raise ImportError(
"viz extra is required for export. Install it with: pip install agent-framework[viz]. "
"You also need to install graphviz separately. E.g., sudo apt-get install graphviz on Debian/Ubuntu "
"or brew install graphviz on macOS. See https://graphviz.org/download/ for details."
) from e
# Create a temporary graphviz Source object
dot_content = self.to_digraph()
source = graphviz.Source(dot_content)
if filename:
# Save to specified file
output_path = Path(filename)
if output_path.suffix and output_path.suffix[1:] != format:
raise ValueError(f"File extension {output_path.suffix} doesn't match format {format}")
# Remove extension if present since graphviz.render() adds it
base_name = str(output_path.with_suffix(""))
source.render(base_name, format=format, cleanup=True)
# Return the actual filename with extension
return f"{base_name}.{format}"
# Create temporary file
with tempfile.NamedTemporaryFile(suffix=f".{format}", delete=False) as temp_file:
temp_path = Path(temp_file.name)
base_name = str(temp_path.with_suffix(""))
source.render(base_name, format=format, cleanup=True)
return f"{base_name}.{format}"
def save_svg(self, filename: str) -> str:
"""Convenience method to save as SVG.
Args:
filename: The filename to save the SVG file.
Returns:
The path to the saved SVG file.
"""
return self.export(format="svg", filename=filename)
def save_png(self, filename: str) -> str:
"""Convenience method to save as PNG.
Args:
filename: The filename to save the PNG file.
Returns:
The path to the saved PNG file.
"""
return self.export(format="png", filename=filename)
def save_pdf(self, filename: str) -> str:
"""Convenience method to save as PDF.
Args:
filename: The filename to save the PDF file.
Returns:
The path to the saved PDF file.
"""
return self.export(format="pdf", filename=filename)
def to_mermaid(self) -> str:
"""Export the workflow as a Mermaid flowchart string.
Returns:
A string representation of the workflow in Mermaid flowchart syntax.
"""
def _san(s: str) -> str:
"""Sanitize an ID for Mermaid (alphanumeric and underscore, start with letter)."""
s2 = re.sub(r"[^0-9A-Za-z_]", "_", s)
if not s2 or not s2[0].isalpha():
s2 = f"n_{s2}"
return s2
lines: list[str] = ["flowchart TD"]
# Emit top-level workflow
self._emit_workflow_mermaid(self._workflow, lines, indent=" ")
# Emit sub-workflows as Mermaid subgraphs
self._emit_sub_workflows_mermaid(self._workflow, lines, indent=" ")
return "\n".join(lines)
# region Private helpers
def _fan_in_digest(self, target: str, sources: list[str]) -> str:
sources_sorted = sorted(sources)
return hashlib.sha256((target + "|" + "|".join(sources_sorted)).encode("utf-8")).hexdigest()[:8]
def _compute_fan_in_descriptors(self, wf: Workflow | None = None) -> list[tuple[str, list[str], str]]:
"""Return list of (node_id, sources, target) for fan-in groups.
node_id is DOT-oriented: fan_in::target::digest
"""
result: list[tuple[str, list[str], str]] = []
workflow = wf or self._workflow
for group in workflow.edge_groups:
if isinstance(group, FanInEdgeGroup):
target = group.target_executor_ids[0]
sources = list(group.source_executor_ids)
digest = self._fan_in_digest(target, sources)
node_id = f"fan_in::{target}::{digest}"
result.append((node_id, sorted(sources), target))
return result
def _compute_normal_edges(self, wf: Workflow | None = None) -> list[tuple[str, str, bool]]:
"""Return list of (source_id, target_id, is_conditional) for non-fan-in groups."""
edges: list[tuple[str, str, bool]] = []
workflow = wf or self._workflow
for group in workflow.edge_groups:
if isinstance(group, FanInEdgeGroup):
continue
for edge in group.edges:
is_cond = getattr(edge, "_condition", None) is not None
edges.append((edge.source_id, edge.target_id, is_cond))
return edges
# endregion
# region Internal emitters (DOT)
def _emit_workflow_digraph(self, wf: Workflow, lines: list[str], indent: str, ns: str | None = None) -> None:
"""Emit DOT nodes/edges for the given workflow.
If ns (namespace) is provided, node ids are prefixed with f"{ns}/" for uniqueness,
but labels remain the original executor ids.
"""
def map_id(x: str) -> str:
return f"{ns}/{x}" if ns else x
# Nodes
start_executor_id = wf.start_executor_id
lines.append(
f'{indent}"{map_id(start_executor_id)}" [fillcolor=lightgreen, label="{start_executor_id}\\n(Start)"];'
)
for executor_id in wf.executors:
if executor_id != start_executor_id:
lines.append(f'{indent}"{map_id(executor_id)}" [label="{executor_id}"];')
# Fan-in nodes
fan_in_nodes = self._compute_fan_in_descriptors(wf)
if fan_in_nodes:
lines.append("")
for node_id, _, _ in fan_in_nodes:
lines.append(f'{indent}"{map_id(node_id)}" [shape=ellipse, fillcolor=lightgoldenrod, label="fan-in"];')
# Fan-in edges
for node_id, sources, target in fan_in_nodes:
for src in sources:
lines.append(f'{indent}"{map_id(src)}" -> "{map_id(node_id)}";')
lines.append(f'{indent}"{map_id(node_id)}" -> "{map_id(target)}";')
# Normal edges
for src, tgt, is_cond in self._compute_normal_edges(wf):
edge_attr = ' [style=dashed, label="conditional"]' if is_cond else ""
lines.append(f'{indent}"{map_id(src)}" -> "{map_id(tgt)}"{edge_attr};')
def _emit_sub_workflows_digraph(self, wf: Workflow, lines: list[str], indent: str) -> None:
"""Emit DOT subgraphs for any WorkflowExecutor instances found in the workflow."""
# Lazy import to avoid any potential import cycles
try:
from ._executor import WorkflowExecutor # type: ignore
except ImportError: # pragma: no cover - best-effort; if unavailable, skip subgraphs
return
for exec_id, exec_obj in wf.executors.items():
if isinstance(exec_obj, WorkflowExecutor) and hasattr(exec_obj, "workflow") and exec_obj.workflow:
subgraph_id = f"cluster_{uuid.uuid5(uuid.NAMESPACE_OID, exec_id).hex[:8]}"
lines.append(f"{indent}subgraph {subgraph_id} {{")
lines.append(f'{indent} label="sub-workflow: {exec_id}";')
lines.append(f"{indent} style=dashed;")
# Emit the nested workflow inside this cluster using a namespace
ns = exec_id
self._emit_workflow_digraph(exec_obj.workflow, lines, indent=f"{indent} ", ns=ns)
# Recurse into deeper nested sub-workflows
self._emit_sub_workflows_digraph(exec_obj.workflow, lines, indent=f"{indent} ")
lines.append(f"{indent}}}")
# endregion
# region Internal emitters (Mermaid)
def _emit_workflow_mermaid(self, wf: Workflow, lines: list[str], indent: str, ns: str | None = None) -> None:
def _san(s: str) -> str:
s2 = re.sub(r"[^0-9A-Za-z_]", "_", s)
if not s2 or not s2[0].isalpha():
s2 = f"n_{s2}"
return s2
def map_id(x: str) -> str:
if ns:
return f"{_san(ns)}__{_san(x)}"
return _san(x)
# Nodes
start_executor_id = wf.start_executor_id
lines.append(f'{indent}{map_id(start_executor_id)}["{start_executor_id} (Start)"];')
for executor_id in wf.executors:
if executor_id == start_executor_id:
continue
lines.append(f'{indent}{map_id(executor_id)}["{executor_id}"];')
# Fan-in nodes
fan_in_nodes_dot = self._compute_fan_in_descriptors(wf)
fan_in_nodes: list[tuple[str, list[str], str]] = []
for dot_node_id, sources, target in fan_in_nodes_dot:
digest = dot_node_id.split("::")[-1]
base = f"{target}__{digest}"
fan_node_id = f"fan_in__{_san(ns) + '__' if ns else ''}{_san(base)}"
fan_in_nodes.append((fan_node_id, sources, target))
for fan_node_id, _, _ in fan_in_nodes:
# Keep this line without trailing semicolon to match existing tests
lines.append(f"{indent}{fan_node_id}((fan-in))")
# Fan-in edges
for fan_node_id, sources, target in fan_in_nodes:
for s in sources:
lines.append(f"{indent}{map_id(s)} --> {fan_node_id};")
lines.append(f"{indent}{fan_node_id} --> {map_id(target)};")
# Normal edges
for src, tgt, is_cond in self._compute_normal_edges(wf):
s = map_id(src)
t = map_id(tgt)
if is_cond:
lines.append(f"{indent}{s} -. conditional .-> {t};")
else:
lines.append(f"{indent}{s} --> {t};")
def _emit_sub_workflows_mermaid(self, wf: Workflow, lines: list[str], indent: str) -> None:
try:
from ._executor import WorkflowExecutor # type: ignore
except ImportError: # pragma: no cover
return
def _san(s: str) -> str:
s2 = re.sub(r"[^0-9A-Za-z_]", "_", s)
if not s2 or not s2[0].isalpha():
s2 = f"n_{s2}"
return s2
for exec_id, exec_obj in wf.executors.items():
if isinstance(exec_obj, WorkflowExecutor) and hasattr(exec_obj, "workflow") and exec_obj.workflow:
sg_id = _san(exec_id)
lines.append(f"{indent}subgraph {sg_id}")
# Render nested workflow within this subgraph using namespacing
self._emit_workflow_mermaid(exec_obj.workflow, lines, indent=f"{indent} ", ns=exec_id)
# Recurse into deeper sub-workflows
self._emit_sub_workflows_mermaid(exec_obj.workflow, lines, indent=f"{indent} ")
lines.append(f"{indent}end")
# endregion
python/packages/main/agent_framework/_workflow/_workflow.py
+1002
View File
File diff suppressed because it is too large Load Diff
python/packages/main/agent_framework/_workflow/_workflow_context.py
+126
View File
@@ -0,0 +1,126 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any, Generic, TypeVar
from opentelemetry.propagate import inject
from ._events import WorkflowEvent
from ._runner_context import Message, RunnerContext
from ._shared_state import SharedState
from ._telemetry import workflow_tracer
T_Out = TypeVar("T_Out")
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
context and shared state, while preventing direct access to the runtime context.
"""
def __init__(
self,
executor_id: str,
source_executor_ids: list[str],
shared_state: SharedState,
runner_context: RunnerContext,
trace_contexts: list[dict[str, str]] | None = None,
source_span_ids: list[str] | None = None,
):
"""Initialize the executor context with the given workflow context.
Args:
executor_id: The unique identifier of the executor that this context belongs to.
source_executor_ids: The IDs of the source executors that sent messages to this executor.
This is a list to support fan_in scenarios where multiple sources send aggregated
messages to the same executor.
shared_state: The shared state for the workflow.
runner_context: The runner context that provides methods to send messages and events.
trace_contexts: Optional trace contexts from multiple sources for OpenTelemetry propagation.
source_span_ids: Optional source span IDs from multiple sources for linking (not for nesting).
"""
self._executor_id = executor_id
self._source_executor_ids = source_executor_ids
self._runner_context = runner_context
self._shared_state = shared_state
# Store trace contexts and source span IDs for linking (supporting multiple sources)
self._trace_contexts = trace_contexts or []
self._source_span_ids = source_span_ids or []
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: T_Out, target_id: str | None = None) -> None:
"""Send a message to the workflow context.
Args:
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.
"""
# Create publishing span (inherits current trace context automatically)
with workflow_tracer.create_sending_span(type(message).__name__, target_id) as span:
# Create Message wrapper
msg = Message(data=message, source_id=self._executor_id, target_id=target_id)
# Inject current trace context if tracing enabled
if workflow_tracer.enabled and span and span.is_recording():
trace_context: dict[str, str] = {}
inject(trace_context) # Inject current trace context for message propagation
msg.trace_contexts = [trace_context]
msg.source_span_ids = [format(span.get_span_context().span_id, "016x")]
await self._runner_context.send_message(msg)
async def add_event(self, event: WorkflowEvent) -> None:
"""Add an event to the workflow context."""
await self._runner_context.add_event(event)
async def get_shared_state(self, key: str) -> Any:
"""Get a value from the shared state."""
return await self._shared_state.get(key)
async def set_shared_state(self, key: str, value: Any) -> None:
"""Set a value in the shared state."""
await self._shared_state.set(key, value)
def get_source_executor_id(self) -> str:
"""Get the ID of the source executor that sent the message to this executor.
Raises:
RuntimeError: If there are multiple source executors, this method raises an error.
"""
if len(self._source_executor_ids) > 1:
raise RuntimeError(
"Cannot get source executor ID when there are multiple source executors. "
"Access the full list via the source_executor_ids property instead."
)
return self._source_executor_ids[0]
@property
def source_executor_ids(self) -> list[str]:
"""Get the IDs of the source executors that sent messages to this executor."""
return self._source_executor_ids
@property
def shared_state(self) -> SharedState:
"""Get the shared state."""
return self._shared_state
async def set_state(self, state: dict[str, Any]) -> None:
"""Persist this executor's state into the checkpointable context.
Executors call this with a JSON-serializable dict capturing the minimal
state needed to resume. It replaces any previously stored state.
"""
if hasattr(self._runner_context, "set_state"):
await self._runner_context.set_state(self._executor_id, state) # type: ignore[arg-type]
async def get_state(self) -> dict[str, Any] | None:
"""Retrieve previously persisted state for this executor, if any."""
if hasattr(self._runner_context, "get_state"):
return await self._runner_context.get_state(self._executor_id) # type: ignore[return-value]
return None
python/packages/main/agent_framework/workflow/__init__.py
-82
View File
@@ -1,82 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import importlib
from typing import Any
PACKAGE_NAME = "agent_framework_workflow"
PACKAGE_EXTRA = "workflow"
_IMPORTS = [
"Executor",
"FunctionExecutor",
"executor",
"WorkflowContext",
"__version__",
"events",
"WorkflowBuilder",
"ExecutorCompletedEvent",
"ExecutorEvent",
"ExecutorInvokeEvent",
"RequestInfoEvent",
"WorkflowCompletedEvent",
"WorkflowEvent",
"WorkflowStartedEvent",
"AgentRunEvent",
"AgentRunUpdateEvent",
"handler",
"AgentExecutor",
"MagenticAgentDeltaEvent",
"MagenticCallbackEvent",
"MagenticCallbackMode",
"MagenticAgentMessageEvent",
"MagenticFinalResultEvent",
"MagenticManagerBase",
"MagenticOrchestratorMessageEvent",
"AgentExecutorRequest",
"AgentExecutorResponse",
"RequestInfoExecutor",
"RequestInfoMessage",
"WorkflowRunResult",
"Workflow",
"WorkflowAgent",
"WorkflowViz",
"FileCheckpointStorage",
"ExecutorFailedEvent",
"InMemoryCheckpointStorage",
"CheckpointStorage",
"WorkflowCheckpoint",
"Case",
"Default",
"RequestResponse",
"SubWorkflowRequestInfo",
"SubWorkflowResponse",
"WorkflowExecutor",
"intercepts_request",
"MagenticBuilder",
"PlanReviewDecision",
"PlanReviewReply",
"PlanReviewRequest",
"RequestInfoEvent",
"StandardMagenticManager",
"WorkflowStatusEvent",
"WorkflowRunState",
"WorkflowErrorDetails",
"WorkflowFailedEvent",
"ConcurrentBuilder",
"SequentialBuilder",
]
def __getattr__(name: str) -> Any:
if name in _IMPORTS:
try:
return getattr(importlib.import_module(PACKAGE_NAME), name)
except ModuleNotFoundError as exc:
raise ModuleNotFoundError(
f"The '{PACKAGE_EXTRA}' extra is not installed, "
f"please do `pip install agent-framework[{PACKAGE_EXTRA}]`"
) from exc
raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
def __dir__() -> list[str]:
return _IMPORTS
python/packages/main/agent_framework/workflow/__init__.pyi
-111
View File
@@ -1,111 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from agent_framework_workflow import (
AgentExecutor,
AgentExecutorRequest,
AgentExecutorResponse,
AgentRunEvent,
AgentRunUpdateEvent,
Case,
CheckpointStorage,
ConcurrentBuilder,
Default,
Executor,
ExecutorCompletedEvent,
ExecutorEvent,
ExecutorFailedEvent,
ExecutorInvokeEvent,
FileCheckpointStorage,
FunctionExecutor,
InMemoryCheckpointStorage,
MagenticBuilder,
MagenticCallbackEvent,
MagenticCallbackMode,
MagenticPlanReviewDecision,
MagenticPlanReviewReply,
MagenticPlanReviewRequest,
MagenticProgressLedger,
MagenticProgressLedgerItem,
RequestInfoEvent,
RequestInfoExecutor,
RequestInfoMessage,
RequestResponse,
SequentialBuilder,
StandardMagenticManager,
SubWorkflowRequestInfo,
SubWorkflowResponse,
Workflow,
WorkflowAgent,
WorkflowBuilder,
WorkflowCheckpoint,
WorkflowCompletedEvent,
WorkflowContext,
WorkflowErrorDetails,
WorkflowEvent,
WorkflowExecutor,
WorkflowFailedEvent,
WorkflowRunResult,
WorkflowRunState,
WorkflowStartedEvent,
WorkflowStatusEvent,
WorkflowViz,
__version__,
executor,
handler,
intercepts_request,
)
__all__ = [
"AgentExecutor",
"AgentExecutorRequest",
"AgentExecutorResponse",
"AgentRunEvent",
"AgentRunUpdateEvent",
"Case",
"CheckpointStorage",
"ConcurrentBuilder",
"Default",
"Executor",
"ExecutorCompletedEvent",
"ExecutorEvent",
"ExecutorFailedEvent",
"ExecutorInvokeEvent",
"FileCheckpointStorage",
"FunctionExecutor",
"InMemoryCheckpointStorage",
"MagenticBuilder",
"MagenticCallbackEvent",
"MagenticCallbackMode",
"MagenticPlanReviewDecision",
"MagenticPlanReviewReply",
"MagenticPlanReviewRequest",
"MagenticProgressLedger",
"MagenticProgressLedgerItem",
"RequestInfoEvent",
"RequestInfoExecutor",
"RequestInfoMessage",
"RequestResponse",
"SequentialBuilder",
"StandardMagenticManager",
"SubWorkflowRequestInfo",
"SubWorkflowResponse",
"Workflow",
"WorkflowAgent",
"WorkflowBuilder",
"WorkflowCheckpoint",
"WorkflowCompletedEvent",
"WorkflowContext",
"WorkflowErrorDetails",
"WorkflowEvent",
"WorkflowExecutor",
"WorkflowFailedEvent",
"WorkflowRunResult",
"WorkflowRunState",
"WorkflowStartedEvent",
"WorkflowStatusEvent",
"WorkflowViz",
"__version__",
"executor",
"handler",
"intercepts_request",
]