[BREAKING] Python: Fix workflow as agent streaming output (#3649)

* WIP: with_output_from

* Add with_output_from to other modules; next: workflow as agent

* WIP: remove agent run events

* orchestrations

* WIP: update samples; next start at guessing_game_With_human_input.py

* Update all samples

* WIP: consolidate workflow as agent streaming vs non-streaming

* Consolidate workflow as agent streaming vs non-streaming

* Move request info event processing to a share method

* Final pass on the samples

* Fix mypy

* Fix mypy

* Comments

---------

Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com>
This commit is contained in:
Tao Chen
2026-02-04 16:16:45 -08:00
committed by GitHub
Unverified
parent 907654a489
commit a971d24f1e
68 changed files with 2652 additions and 2247 deletions
@@ -38,8 +38,6 @@ from ._edge import (
)
from ._edge_runner import create_edge_runner
from ._events import (
AgentRunEvent,
AgentRunUpdateEvent,
ExecutorCompletedEvent,
ExecutorEvent,
ExecutorFailedEvent,
@@ -131,8 +129,6 @@ __all__ = [
"AgentExecutorRequest",
"AgentExecutorResponse",
"AgentRequestInfoResponse",
"AgentRunEvent",
"AgentRunUpdateEvent",
"BaseGroupChatOrchestrator",
"Case",
"CheckpointStorage",
@@ -21,16 +21,14 @@ from agent_framework import (
from .._types import add_usage_details
from ..exceptions import AgentExecutionException
from ._agent_executor import AgentExecutor
from ._checkpoint import CheckpointStorage
from ._events import (
AgentRunUpdateEvent,
RequestInfoEvent,
WorkflowEvent,
WorkflowOutputEvent,
)
from ._message_utils import normalize_messages_input
from ._typing_utils import is_type_compatible
from ._typing_utils import is_instance_of, is_type_compatible
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
@@ -93,6 +91,12 @@ class WorkflowAgent(BaseAgent):
name: Optional name for the agent.
description: Optional description of the agent.
**kwargs: Additional keyword arguments passed to BaseAgent.
Note:
Only WorkflowOutputEvents and RequestInfoEvents from the workflow are considered and
converted to agent responses of the WorkflowAgent. Other workflow events are ignored.
Use `with_output_from` in WorkflowBuilder to control which executors' outputs are surfaced
as agent responses.
"""
if id is None:
id = f"WorkflowAgent_{uuid.uuid4().hex[:8]}"
@@ -118,6 +122,8 @@ class WorkflowAgent(BaseAgent):
def pending_requests(self) -> dict[str, RequestInfoEvent]:
return self._pending_requests
# region Run Methods
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
@@ -129,7 +135,7 @@ class WorkflowAgent(BaseAgent):
) -> AgentResponse:
"""Get a response from the workflow agent (non-streaming).
This method collects all streaming updates and merges them into a single response.
This method runs the workflow in non-streaming mode.
Args:
messages: The message(s) to send to the workflow. Required for new runs,
@@ -146,21 +152,19 @@ class WorkflowAgent(BaseAgent):
and tool functions.
Returns:
The final workflow response as an AgentResponse.
An AgentResponse representing the workflow execution results. The response
includes all output events and requests emitted during the workflow run.
WorkflowOutputEvents will be converted to ChatMessages in the response.
RequestInfoEvents will be converted to function call and approval request contents
in the response.
"""
# Collect all streaming updates
response_updates: list[AgentResponseUpdate] = []
input_messages = normalize_messages_input(messages)
thread = thread or self.get_new_thread()
response_id = str(uuid.uuid4())
async for update in self._run_stream_impl(
response = await self._run_impl(
input_messages, response_id, thread, checkpoint_id, checkpoint_storage, **kwargs
):
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, response.messages)
@@ -194,6 +198,10 @@ class WorkflowAgent(BaseAgent):
Yields:
AgentResponseUpdate objects representing the workflow execution progress.
Updates include output events and requests emitted during the workflow run.
WorkflowOutputEvents will be converted to AgentResponseUpdate objects.
RequestInfoEvents will be converted to function call and approval request contents
in the updates.
"""
input_messages = normalize_messages_input(messages)
thread = thread or self.get_new_thread()
@@ -212,6 +220,38 @@ class WorkflowAgent(BaseAgent):
# Notify thread of new messages (both input and response messages)
await self._notify_thread_of_new_messages(thread, input_messages, response.messages)
async def _run_impl(
self,
input_messages: list[ChatMessage],
response_id: str,
thread: AgentThread,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AgentResponse:
"""Internal implementation of non-streaming execution.
Args:
input_messages: Normalized input messages to process.
response_id: The unique response ID for this workflow execution.
thread: The conversation thread containing message history.
checkpoint_id: ID of checkpoint to restore from.
checkpoint_storage: Runtime checkpoint storage.
**kwargs: Additional keyword arguments passed through to the underlying
workflow and tool functions.
Returns:
An AgentResponse representing the workflow execution results.
"""
output_events: list[WorkflowOutputEvent | RequestInfoEvent] = []
async for event in self._run_core(
input_messages, thread, checkpoint_id, checkpoint_storage, streaming=False, **kwargs
):
if isinstance(event, WorkflowOutputEvent | RequestInfoEvent):
output_events.append(event)
return self._convert_workflow_events_to_agent_response(response_id, output_events)
async def _run_stream_impl(
self,
input_messages: list[ChatMessage],
@@ -235,150 +275,325 @@ class WorkflowAgent(BaseAgent):
Yields:
AgentResponseUpdate objects representing the workflow execution progress.
"""
# Determine the event stream based on whether we have function responses
async for event in self._run_core(
input_messages, thread, checkpoint_id, checkpoint_storage, streaming=True, **kwargs
):
updates = self._convert_workflow_event_to_agent_response_update(response_id, event)
for update in updates:
yield update
async def _run_core(
self,
input_messages: list[ChatMessage],
thread: AgentThread,
checkpoint_id: str | None,
checkpoint_storage: CheckpointStorage | None,
streaming: bool,
**kwargs: Any,
) -> AsyncIterable[WorkflowEvent]:
"""Core implementation that yields workflow events for both streaming and non-streaming modes.
Args:
input_messages: Normalized input messages to process.
thread: The conversation thread containing message history.
checkpoint_id: ID of checkpoint to restore from.
checkpoint_storage: Runtime checkpoint storage.
streaming: Whether to use streaming workflow methods.
**kwargs: Additional keyword arguments passed through to the underlying
workflow and tool functions.
Yields:
WorkflowEvent objects from the workflow execution.
"""
# Determine the execution mode based on state
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")
# This is a continuation - send function responses back
function_responses = self._process_pending_requests(input_messages)
# 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)
if streaming:
async for event in self.workflow.send_responses_streaming(function_responses):
yield event
else:
workflow_result = await self.workflow.send_responses(function_responses)
for event in workflow_result:
yield event
# 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)
elif checkpoint_id is not None:
# Resume from checkpoint - don't prepend thread history since workflow state
# is being restored from the checkpoint
event_stream = self.workflow.run_stream(
message=None,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
**kwargs,
)
if streaming:
async for event in self.workflow.run_stream(
message=None,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
**kwargs,
):
yield event
else:
workflow_result = await self.workflow.run(
message=None,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
**kwargs,
)
for event in workflow_result:
yield event
else:
# Execute workflow with streaming (initial run or no function responses)
# Build the complete conversation by prepending thread history to input messages
conversation_messages: list[ChatMessage] = []
if thread.message_store:
history = await thread.message_store.list_messages()
if history:
conversation_messages.extend(history)
conversation_messages.extend(input_messages)
event_stream = self.workflow.run_stream(
message=conversation_messages,
checkpoint_storage=checkpoint_storage,
**kwargs,
)
# Initial run - build conversation from thread history
conversation_messages = await self._build_conversation_messages(thread, 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
if streaming:
async for event in self.workflow.run_stream(
message=conversation_messages,
checkpoint_storage=checkpoint_storage,
**kwargs,
):
yield event
else:
workflow_result = await self.workflow.run(
message=conversation_messages,
checkpoint_storage=checkpoint_storage,
**kwargs,
)
for event in workflow_result:
yield event
def _convert_workflow_event_to_agent_update(
# endregion Run Methods
async def _build_conversation_messages(
self,
thread: AgentThread,
input_messages: list[ChatMessage],
) -> list[ChatMessage]:
"""Build the complete conversation by prepending thread history to input messages.
Args:
thread: The conversation thread containing message history.
input_messages: The new input messages to append.
Returns:
A list of ChatMessage objects representing the full conversation.
"""
conversation_messages: list[ChatMessage] = []
if thread.message_store:
history = await thread.message_store.list_messages()
if history:
conversation_messages.extend(history)
conversation_messages.extend(input_messages)
return conversation_messages
def _process_pending_requests(self, input_messages: list[ChatMessage]) -> dict[str, Any]:
"""Process pending requests by extracting function responses and updating state.
Args:
input_messages: Input messages that may contain function responses.
Returns:
A dictionary mapping request IDs to their response data.
"""
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.
return function_responses
def _convert_workflow_events_to_agent_response(
self,
response_id: str,
output_events: list[WorkflowOutputEvent | RequestInfoEvent],
) -> AgentResponse:
"""Convert a list of workflow output events to an AgentResponse."""
messages: list[ChatMessage] = []
raw_representations: list[object] = []
merged_usage: UsageDetails | None = None
latest_created_at: str | None = None
for output_event in output_events:
if isinstance(output_event, RequestInfoEvent):
function_call, approval_request = self._process_request_info_event(output_event)
messages.append(
ChatMessage(
contents=[function_call, approval_request],
role="assistant",
author_name=output_event.source_executor_id,
message_id=str(uuid.uuid4()),
raw_representation=output_event,
)
)
raw_representations.append(output_event)
else:
data = output_event.data
if isinstance(data, AgentResponseUpdate):
# We cannot support AgentResponseUpdate in non-streaming mode. This is because the message
# sequence cannot be guaranteed when there are streaming updates in between non-streaming
# responses.
raise AgentExecutionException(
"WorkflowOutputEvent with AgentResponseUpdate data cannot be emitted in non-streaming mode. "
"Please ensure executors emit AgentResponse for non-streaming workflows."
)
if isinstance(data, AgentResponse):
messages.extend(data.messages)
raw_representations.append(data.raw_representation)
merged_usage = add_usage_details(merged_usage, data.usage_details)
latest_created_at = (
data.created_at
if not latest_created_at
else max(latest_created_at, data.created_at)
if data.created_at
else latest_created_at
)
elif isinstance(data, ChatMessage):
messages.append(data)
raw_representations.append(data.raw_representation)
elif is_instance_of(data, list[ChatMessage]):
chat_messages = cast(list[ChatMessage], data)
messages.extend(chat_messages)
raw_representations.append(data)
else:
contents = self._extract_contents(data)
if not contents:
continue
messages.append(
ChatMessage(
contents=contents,
role="assistant",
author_name=output_event.executor_id,
message_id=str(uuid.uuid4()),
raw_representation=data,
)
)
raw_representations.append(data)
return AgentResponse(
messages=messages,
response_id=response_id,
created_at=latest_created_at,
usage_details=merged_usage,
raw_representation=raw_representations,
)
def _convert_workflow_event_to_agent_response_update(
self,
response_id: str,
event: WorkflowEvent,
) -> AgentResponseUpdate | None:
) -> list[AgentResponseUpdate]:
"""Convert a workflow event to an AgentResponseUpdate.
AgentRunUpdateEvent, RequestInfoEvent, and WorkflowOutputEvent are processed.
Only WorkflowOutputEvent and RequestInfoEvent are processed.
Other workflow events are ignored as they are workflow-internal.
For AgentRunUpdateEvent from AgentExecutor instances, only events from executors
with output_response=True are converted to agent updates. This prevents agent
responses from executors that were not explicitly marked to surface their output.
Non-AgentExecutor executors that emit AgentRunUpdateEvent directly are allowed
through since they explicitly chose to emit the event.
"""
match event:
case AgentRunUpdateEvent(data=update, executor_id=executor_id):
# For AgentExecutor instances, only pass through if output_response=True.
# Non-AgentExecutor executors that emit AgentRunUpdateEvent are allowed through.
executor = self.workflow.executors.get(executor_id)
if isinstance(executor, AgentExecutor) and not executor.output_response:
return None
if update:
# Enrich with executor identity if author_name is not already set
if not update.author_name:
update.author_name = executor_id
return update
return None
# Convert workflow output to an agent response update.
case WorkflowOutputEvent(data=data, executor_id=executor_id):
# Convert workflow output to an agent response update.
# Handle different data types appropriately.
# Skip AgentResponse from AgentExecutor with output_response=True
# since streaming events already surfaced the content.
if isinstance(data, AgentResponse):
executor = self.workflow.executors.get(executor_id)
if isinstance(executor, AgentExecutor) and executor.output_response:
return None
return [
AgentResponseUpdate(
contents=[content for message in data.messages for content in message.contents],
role="assistant",
author_name=executor_id,
response_id=response_id,
created_at=data.created_at,
raw_representation=data,
)
]
if isinstance(data, AgentResponseUpdate):
return data
return [data]
if isinstance(data, ChatMessage):
return AgentResponseUpdate(
contents=list(data.contents),
role=data.role,
author_name=data.author_name or executor_id,
return [
AgentResponseUpdate(
contents=list(data.contents),
role=data.role,
author_name=data.author_name,
response_id=response_id,
message_id=data.message_id or str(uuid.uuid4()),
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
raw_representation=data,
)
]
if is_instance_of(data, list[ChatMessage]):
chat_messages = cast(list[ChatMessage], data)
return [
AgentResponseUpdate(
contents=list(msg.contents),
role=msg.role,
author_name=msg.author_name,
response_id=response_id,
message_id=msg.message_id or str(uuid.uuid4()),
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
raw_representation=msg,
)
for msg in chat_messages
]
contents = self._extract_contents(data)
if not contents:
return []
return [
AgentResponseUpdate(
contents=contents,
role="assistant",
author_name=executor_id,
response_id=response_id,
message_id=str(uuid.uuid4()),
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
raw_representation=data,
)
contents = self._extract_contents(data)
if not contents:
return None
return AgentResponseUpdate(
contents=contents,
role="assistant",
author_name=executor_id,
response_id=response_id,
message_id=str(uuid.uuid4()),
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
raw_representation=data,
)
]
case RequestInfoEvent(request_id=request_id):
# Store the pending request for later correlation
self.pending_requests[request_id] = event
args = self.RequestInfoFunctionArgs(request_id=request_id, data=event.data).to_dict()
function_call = Content.from_function_call(
call_id=request_id,
name=self.REQUEST_INFO_FUNCTION_NAME,
arguments=args,
)
approval_request = Content.from_function_approval_request(
id=request_id,
function_call=function_call,
additional_properties={"request_id": request_id},
)
return AgentResponseUpdate(
contents=[function_call, approval_request],
role="assistant",
author_name=self.name,
response_id=response_id,
message_id=str(uuid.uuid4()),
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
)
case RequestInfoEvent():
function_call, approval_request = self._process_request_info_event(event)
return [
AgentResponseUpdate(
contents=[function_call, approval_request],
role="assistant",
author_name=self.name,
response_id=response_id,
message_id=str(uuid.uuid4()),
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
)
]
case _:
# Ignore workflow-internal events
pass
return None
return []
def _process_request_info_event(self, event: RequestInfoEvent) -> tuple[Content, Content]:
"""Process a RequestInfoEvent by adding it to pending requests."""
# Store the pending request for later correlation
self.pending_requests[event.request_id] = event
args = self.RequestInfoFunctionArgs(request_id=event.request_id, data=event.data).to_dict()
function_call = Content.from_function_call(
call_id=event.request_id,
name=self.REQUEST_INFO_FUNCTION_NAME,
arguments=args,
)
approval_request = Content.from_function_approval_request(
id=event.request_id,
function_call=function_call,
additional_properties={"request_id": event.request_id},
)
return function_call, approval_request
def _extract_function_responses(self, input_messages: list[ChatMessage]) -> dict[str, Any]:
"""Extract function responses from input messages."""
@@ -428,8 +643,6 @@ class WorkflowAgent(BaseAgent):
def _extract_contents(self, data: Any) -> list[Content]:
"""Recursively extract Content from workflow output data."""
if isinstance(data, ChatMessage):
return list(data.contents)
if isinstance(data, list):
return [c for item in data for c in self._extract_contents(item)]
if isinstance(data, Content):
@@ -2,26 +2,24 @@
import logging
import sys
import types
from dataclasses import dataclass
from typing import Any, cast
from typing_extensions import Never
from agent_framework import Content
from .._agents import AgentProtocol, ChatAgent
from .._agents import AgentProtocol
from .._threads import AgentThread
from .._types import AgentResponse, AgentResponseUpdate, ChatMessage
from ._agent_utils import resolve_agent_id
from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value
from ._const import WORKFLOW_RUN_KWARGS_KEY
from ._conversation_state import encode_chat_messages
from ._events import (
AgentRunEvent,
AgentRunUpdateEvent,
)
from ._executor import Executor, handler
from ._message_utils import normalize_messages_input
from ._request_info_mixin import response_handler
from ._typing_utils import is_chat_agent
from ._workflow_context import WorkflowContext
if sys.version_info >= (3, 12):
@@ -55,7 +53,7 @@ class AgentExecutorResponse:
agent_response: The underlying agent run response (unaltered from client).
full_conversation: The full conversation context (prior inputs + all assistant/tool outputs) that
should be used when chaining to another AgentExecutor. This prevents downstream agents losing
user prompts while keeping the emitted AgentRunEvent text faithful to the raw agent output.
user prompts.
"""
executor_id: str
@@ -67,8 +65,15 @@ class AgentExecutor(Executor):
"""built-in executor that wraps an agent for handling messages.
AgentExecutor adapts its behavior based on the workflow execution mode:
- run_stream(): Emits incremental AgentRunUpdateEvent events as the agent produces tokens
- run(): Emits a single AgentRunEvent containing the complete response
- run_stream(): Emits incremental WorkflowOutputEvents as the agent produces tokens
- run(): Emits a single WorkflowOutputEvent containing the complete response
Use `with_output_from` in WorkflowBuilder to control whether the AgentResponse
or AgentResponseUpdate objects are yielded as workflow outputs.
Messages sent to downstream executors will always be the complete AgentResponse. In
streaming mode, incremental AgentResponseUpdates will be concatenated to form the full
response to be sent downstream.
The executor automatically detects the mode via WorkflowContext.is_streaming().
"""
@@ -78,7 +83,6 @@ class AgentExecutor(Executor):
agent: AgentProtocol,
*,
agent_thread: AgentThread | None = None,
output_response: bool = False,
id: str | None = None,
):
"""Initialize the executor with a unique identifier.
@@ -86,7 +90,6 @@ class AgentExecutor(Executor):
Args:
agent: The agent to be wrapped by this executor.
agent_thread: The thread to use for running the agent. If None, a new thread will be created.
output_response: Whether to yield an AgentResponse as a workflow output when the agent completes.
id: A unique identifier for the executor. If None, the agent's name will be used if available.
"""
# Prefer provided id; else use agent.name if present; else generate deterministic prefix
@@ -96,27 +99,15 @@ class AgentExecutor(Executor):
super().__init__(exec_id)
self._agent = agent
self._agent_thread = agent_thread or self._agent.get_new_thread()
self._pending_agent_requests: dict[str, Content] = {}
self._pending_responses_to_agent: list[Content] = []
self._output_response = output_response
# AgentExecutor maintains an internal cache of messages in between runs
self._cache: list[ChatMessage] = []
# This tracks the full conversation after each run
self._full_conversation: list[ChatMessage] = []
@property
def output_response(self) -> bool:
"""Whether this executor yields AgentResponse as workflow output when complete."""
return self._output_response
@property
def workflow_output_types(self) -> list[type[Any] | types.UnionType]:
# Override to declare AgentResponse as a possible output type only if enabled.
if self._output_response:
return [AgentResponse]
return []
@property
def description(self) -> str | None:
"""Get the description of the underlying agent."""
@@ -124,7 +115,9 @@ class AgentExecutor(Executor):
@handler
async def run(
self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse]
self,
request: AgentExecutorRequest,
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
) -> None:
"""Handle an AgentExecutorRequest (canonical input).
@@ -137,7 +130,9 @@ class AgentExecutor(Executor):
@handler
async def from_response(
self, prior: AgentExecutorResponse, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse]
self,
prior: AgentExecutorResponse,
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
) -> None:
"""Enable seamless chaining: accept a prior AgentExecutorResponse as input.
@@ -152,7 +147,9 @@ class AgentExecutor(Executor):
await self._run_agent_and_emit(ctx)
@handler
async def from_str(self, text: str, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse]) -> None:
async def from_str(
self, text: str, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate]
) -> None:
"""Accept a raw user prompt string and run the agent (one-shot)."""
self._cache = normalize_messages_input(text)
await self._run_agent_and_emit(ctx)
@@ -161,7 +158,7 @@ class AgentExecutor(Executor):
async def from_message(
self,
message: ChatMessage,
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse],
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
) -> None:
"""Accept a single ChatMessage as input."""
self._cache = normalize_messages_input(message)
@@ -171,7 +168,7 @@ class AgentExecutor(Executor):
async def from_messages(
self,
messages: list[str | ChatMessage],
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse],
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
) -> None:
"""Accept a list of chat inputs (strings or ChatMessage) as conversation context."""
self._cache = normalize_messages_input(messages)
@@ -182,7 +179,7 @@ class AgentExecutor(Executor):
self,
original_request: Content,
response: Content,
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse],
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
) -> None:
"""Handle user input responses for function approvals during agent execution.
@@ -215,7 +212,7 @@ class AgentExecutor(Executor):
Dict containing serialized cache and thread state
"""
# Check if using AzureAIAgentClient with server-side thread and warn about checkpointing limitations
if isinstance(self._agent, ChatAgent) and self._agent_thread.service_thread_id is not None:
if is_chat_agent(self._agent) and self._agent_thread.service_thread_id is not None:
client_class_name = self._agent.chat_client.__class__.__name__
client_module = self._agent.chat_client.__class__.__module__
@@ -293,23 +290,26 @@ class AgentExecutor(Executor):
logger.debug("AgentExecutor %s: Resetting cache", self.id)
self._cache.clear()
async def _run_agent_and_emit(self, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse]) -> None:
async def _run_agent_and_emit(
self,
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
) -> None:
"""Execute the underlying agent, emit events, and enqueue response.
Checks ctx.is_streaming() to determine whether to emit incremental AgentRunUpdateEvent
events (streaming mode) or a single AgentRunEvent (non-streaming mode).
Checks ctx.is_streaming() to determine whether to emit WorkflowOutputEvents
containing incremental updates (streaming mode) or a single WorkflowOutputEvent
containing the complete response (non-streaming mode).
"""
if ctx.is_streaming():
# Streaming mode: emit incremental updates
response = await self._run_agent_streaming(cast(WorkflowContext, ctx))
response = await self._run_agent_streaming(cast(WorkflowContext[Never, AgentResponseUpdate], ctx))
else:
# Non-streaming mode: use run() and emit single event
response = await self._run_agent(cast(WorkflowContext, ctx))
response = await self._run_agent(cast(WorkflowContext[Never, AgentResponse], ctx))
# Always extend full conversation with cached messages plus agent outputs
# (agent_response.messages) after each run. This is to avoid losing context
# when agent did not complete and the cache is cleared when responses come back.
# Do not mutate response.messages so AgentRunEvent remains faithful to the raw output.
self._full_conversation.extend(list(self._cache) + (list(response.messages) if response else []))
if response is None:
@@ -317,14 +317,11 @@ class AgentExecutor(Executor):
logger.info("AgentExecutor %s: Agent did not complete, awaiting user input", self.id)
return
if self._output_response:
await ctx.yield_output(response)
agent_response = AgentExecutorResponse(self.id, response, full_conversation=self._full_conversation)
await ctx.send_message(agent_response)
self._cache.clear()
async def _run_agent(self, ctx: WorkflowContext) -> AgentResponse | None:
async def _run_agent(self, ctx: WorkflowContext[Never, AgentResponse]) -> AgentResponse | None:
"""Execute the underlying agent in non-streaming mode.
Args:
@@ -340,7 +337,7 @@ class AgentExecutor(Executor):
thread=self._agent_thread,
**run_kwargs,
)
await ctx.add_event(AgentRunEvent(self.id, response))
await ctx.yield_output(response)
# Handle any user input requests
if response.user_input_requests:
@@ -351,7 +348,7 @@ class AgentExecutor(Executor):
return response
async def _run_agent_streaming(self, ctx: WorkflowContext) -> AgentResponse | None:
async def _run_agent_streaming(self, ctx: WorkflowContext[Never, AgentResponseUpdate]) -> AgentResponse | None:
"""Execute the underlying agent in streaming mode and collect the full response.
Args:
@@ -370,13 +367,13 @@ class AgentExecutor(Executor):
**run_kwargs,
):
updates.append(update)
await ctx.add_event(AgentRunUpdateEvent(self.id, update))
await ctx.yield_output(update)
if update.user_input_requests:
user_input_requests.extend(update.user_input_requests)
# Build the final AgentResponse from the collected updates
if isinstance(self._agent, ChatAgent):
if is_chat_agent(self._agent):
response_format = self._agent.default_options.get("response_format")
response = AgentResponse.from_updates(
updates,
@@ -246,6 +246,7 @@ class ConcurrentBuilder:
self._checkpoint_storage: CheckpointStorage | None = None
self._request_info_enabled: bool = False
self._request_info_filter: set[str] | None = None
self._intermediate_outputs: bool = False
def register_participants(
self,
@@ -489,6 +490,19 @@ class ConcurrentBuilder:
return self
def with_intermediate_outputs(self) -> "ConcurrentBuilder":
"""Enable intermediate outputs from agent participants before aggregation.
When enabled, the workflow returns each agent participant's response or yields
streaming updates as they become available. The output of the aggregator will
always be available as the final output of the workflow.
Returns:
Self for fluent chaining
"""
self._intermediate_outputs = True
return self
def _resolve_participants(self) -> list[Executor]:
"""Resolve participant instances into Executor objects."""
if not self._participants and not self._participant_factories:
@@ -568,6 +582,10 @@ class ConcurrentBuilder:
# Direct fan-in to aggregator
builder.add_fan_in_edges(participants, aggregator)
if not self._intermediate_outputs:
# Constrain output to aggregator only
builder = builder.with_output_from([aggregator])
if self._checkpoint_storage is not None:
builder = builder.with_checkpointing(self._checkpoint_storage)
@@ -8,8 +8,6 @@ from dataclasses import dataclass
from enum import Enum
from typing import Any, TypeAlias
from agent_framework import AgentResponse, AgentResponseUpdate
from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value
from ._typing_utils import deserialize_type, serialize_type
@@ -364,32 +362,4 @@ class ExecutorFailedEvent(ExecutorEvent):
return f"{self.__class__.__name__}(executor_id={self.executor_id}, details={self.details})"
class AgentRunUpdateEvent(ExecutorEvent):
"""Event triggered when an agent is streaming messages."""
data: AgentResponseUpdate
def __init__(self, executor_id: str, data: AgentResponseUpdate):
"""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."""
data: AgentResponse
def __init__(self, executor_id: str, data: AgentResponse):
"""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})"
WorkflowLifecycleEvent: TypeAlias = WorkflowStartedEvent | WorkflowStatusEvent | WorkflowFailedEvent
@@ -537,6 +537,9 @@ class GroupChatBuilder:
self._request_info_enabled: bool = False
self._request_info_filter: set[str] = set()
# Intermediate outputs
self._intermediate_outputs = False
@overload
def with_orchestrator(self, *, agent: ChatAgent | Callable[[], ChatAgent]) -> "GroupChatBuilder":
"""Set the orchestrator for this group chat workflow using a ChatAgent.
@@ -880,6 +883,19 @@ class GroupChatBuilder:
return self
def with_intermediate_outputs(self) -> "GroupChatBuilder":
"""Enable intermediate outputs from agent participants.
When enabled, the workflow returns each agent participant's response or yields
streaming updates as they become available. The output of the orchestrator will
always be available as the final output of the workflow.
Returns:
Self for fluent chaining
"""
self._intermediate_outputs = True
return self
def _resolve_orchestrator(self, participants: Sequence[Executor]) -> Executor:
"""Determine the orchestrator to use for the workflow.
@@ -986,6 +1002,11 @@ class GroupChatBuilder:
# Orchestrator and participant bi-directional edges
workflow_builder = workflow_builder.add_edge(orchestrator, participant)
workflow_builder = workflow_builder.add_edge(participant, orchestrator)
if not self._intermediate_outputs:
# Constrain output to orchestrator only
workflow_builder = workflow_builder.with_output_from([orchestrator])
if self._checkpoint_storage is not None:
workflow_builder = workflow_builder.with_checkpointing(self._checkpoint_storage)
@@ -42,7 +42,7 @@ from .._agents import AgentProtocol, ChatAgent
from .._middleware import FunctionInvocationContext, FunctionMiddleware
from .._threads import AgentThread
from .._tools import FunctionTool, tool
from .._types import AgentResponse, ChatMessage
from .._types import AgentResponse, AgentResponseUpdate, ChatMessage
from ._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
from ._agent_utils import resolve_agent_id
from ._base_group_chat_orchestrator import TerminationCondition
@@ -365,7 +365,9 @@ class HandoffAgentExecutor(AgentExecutor):
return _handoff_tool
@override
async def _run_agent_and_emit(self, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse]) -> None:
async def _run_agent_and_emit(
self, ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate]
) -> None:
"""Override to support handoff."""
# When the full conversation is empty, it means this is the first run.
# Broadcast the initial cache to all other agents. Subsequent runs won't
@@ -383,10 +385,10 @@ class HandoffAgentExecutor(AgentExecutor):
# Run the agent
if ctx.is_streaming():
# Streaming mode: emit incremental updates
response = await self._run_agent_streaming(cast(WorkflowContext, ctx))
response = await self._run_agent_streaming(cast(WorkflowContext[Never, AgentResponseUpdate], ctx))
else:
# Non-streaming mode: use run() and emit single event
response = await self._run_agent(cast(WorkflowContext, ctx))
response = await self._run_agent(cast(WorkflowContext[Never, AgentResponse], ctx))
# Clear the cache after running the agent
self._cache.clear()
@@ -466,7 +468,9 @@ class HandoffAgentExecutor(AgentExecutor):
# Append the user response messages to the cache
self._cache.extend(response)
await self._run_agent_and_emit(ctx)
await self._run_agent_and_emit(
cast(WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate], ctx)
)
async def _broadcast_messages(
self,
@@ -562,7 +566,10 @@ class HandoffBuilder:
The final conversation history as a list of ChatMessage once the group chat completes.
Note:
Agents in handoff workflows must be ChatAgent instances and support local tool calls.
1. Agents in handoff workflows must be ChatAgent instances and support local tool calls.
2. Handoff doesn't support intermediate outputs from agents. All outputs are returned as
they become available. This is because agents in handoff workflows are not considered
sub-agents of a central orchestrator, thus all outputs are directly emitted.
"""
def __init__(
@@ -1389,6 +1389,9 @@ class MagenticBuilder:
self._checkpoint_storage: CheckpointStorage | None = None
# Intermediate outputs
self._intermediate_outputs = False
def register_participants(
self,
participant_factories: Sequence[Callable[[], AgentProtocol | Executor]],
@@ -1904,6 +1907,19 @@ class MagenticBuilder:
return self
def with_intermediate_outputs(self) -> Self:
"""Enable intermediate outputs from agent participants before aggregation.
When enabled, the workflow returns each agent participant's response or yields
streaming updates as they become available. The output of the orchestrator will
always be available as the final output of the workflow.
Returns:
Self for fluent chaining
"""
self._intermediate_outputs = True
return self
def _resolve_orchestrator(self, participants: Sequence[Executor]) -> Executor:
"""Determine the orchestrator to use for the workflow.
@@ -1977,6 +1993,10 @@ class MagenticBuilder:
if self._checkpoint_storage is not None:
workflow_builder = workflow_builder.with_checkpointing(self._checkpoint_storage)
if not self._intermediate_outputs:
# Constrain output to orchestrator only
workflow_builder = workflow_builder.with_output_from([orchestrator])
return workflow_builder.build()
@@ -290,7 +290,7 @@ class InProcRunnerContext:
checkpoint_storage: Optional storage to enable checkpointing.
"""
self._messages: dict[str, list[Message]] = {}
# Event queue for immediate streaming of events (e.g., AgentRunUpdateEvent)
# Event queue for immediate streaming of events
self._event_queue: asyncio.Queue[WorkflowEvent] = asyncio.Queue()
# An additional storage for pending request info events
@@ -153,6 +153,7 @@ class SequentialBuilder:
self._checkpoint_storage: CheckpointStorage | None = None
self._request_info_enabled: bool = False
self._request_info_filter: set[str] | None = None
self._intermediate_outputs: bool = False
def register_participants(
self,
@@ -242,6 +243,19 @@ class SequentialBuilder:
return self
def with_intermediate_outputs(self) -> "SequentialBuilder":
"""Enable intermediate outputs from agent participants.
When enabled, the workflow returns each agent participant's response or yields
streaming updates as they become available. The output of the last participant
will always be available as the final output of the workflow.
Returns:
Self for fluent chaining
"""
self._intermediate_outputs = True
return self
def _resolve_participants(self) -> list[Executor]:
"""Resolve participant instances into Executor objects."""
if not self._participants and not self._participant_factories:
@@ -305,6 +319,10 @@ class SequentialBuilder:
# Terminate with the final conversation
builder.add_edge(prior, end)
if not self._intermediate_outputs:
# Constrain output to end only
builder = builder.with_output_from([end])
if self._checkpoint_storage is not None:
builder = builder.with_checkpointing(self._checkpoint_storage)
@@ -1,9 +1,21 @@
# Copyright (c) Microsoft. All rights reserved.
from types import UnionType
from typing import Any, TypeVar, Union, cast, get_args, get_origin
from typing import Any, TypeGuard, Union, cast, get_args, get_origin
T = TypeVar("T")
from .._agents import ChatAgent
def is_chat_agent(agent: Any) -> TypeGuard[ChatAgent]:
"""Check if the given agent is a ChatAgent.
Args:
agent (Any): The agent to check.
Returns:
TypeGuard[ChatAgent]: True if the agent is a ChatAgent, False otherwise.
"""
return isinstance(agent, ChatAgent)
def resolve_type_annotation(
@@ -23,7 +23,7 @@ class ValidationTypeEnum(Enum):
TYPE_COMPATIBILITY = "TYPE_COMPATIBILITY"
GRAPH_CONNECTIVITY = "GRAPH_CONNECTIVITY"
HANDLER_OUTPUT_ANNOTATION = "HANDLER_OUTPUT_ANNOTATION"
INTERCEPTOR_CONFLICT = "INTERCEPTOR_CONFLICT"
OUTPUT_VALIDATION = "OUTPUT_VALIDATION"
class WorkflowValidationError(Exception):
@@ -79,13 +79,6 @@ class GraphConnectivityError(WorkflowValidationError):
super().__init__(message, validation_type=ValidationTypeEnum.GRAPH_CONNECTIVITY)
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
@@ -109,6 +102,7 @@ class WorkflowGraphValidator:
edge_groups: Sequence[EdgeGroup],
executors: dict[str, Executor],
start_executor: Executor,
output_executors: list[str],
) -> None:
"""Validate the entire workflow graph.
@@ -116,6 +110,7 @@ class WorkflowGraphValidator:
edge_groups: list of edge groups in the workflow
executors: Map of executor IDs to executor instances
start_executor: The starting executor
output_executors: List of output executor IDs
Raises:
WorkflowValidationError: If any validation fails
@@ -162,6 +157,7 @@ class WorkflowGraphValidator:
self._validate_graph_connectivity(start_executor.id)
self._validate_self_loops()
self._validate_dead_ends()
self._output_validation(output_executors)
def _validate_handler_output_annotations(self) -> None:
"""Validate that each handler's ctx parameter is annotated with WorkflowContext[T].
@@ -357,6 +353,26 @@ class WorkflowGraphValidator:
# endregion
# region Output Validation
def _output_validation(self, output_executors: list[str]) -> None:
"""Validate that output executors exist in the workflow and have the correct workflow context annotations."""
for output_id in output_executors:
if output_id not in self._executors:
raise WorkflowValidationError(
f"Output executor '{output_id}' is not present in the workflow graph",
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
)
output_executor = self._executors[output_id]
if not output_executor.workflow_output_types:
raise WorkflowValidationError(
f"Output executor '{output_id}' must have output type annotations defined.",
validation_type=ValidationTypeEnum.OUTPUT_VALIDATION,
)
# endregion
# region Additional Validation Scenarios
def _validate_self_loops(self) -> None:
"""Detect and log self-loops (edges from executor to itself).
@@ -397,13 +413,15 @@ def validate_workflow_graph(
edge_groups: Sequence[EdgeGroup],
executors: dict[str, Executor],
start_executor: Executor,
output_executors: list[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)
start_executor: The starting executor instance
output_executors: List of output executor IDs
Raises:
WorkflowValidationError: If any validation fails
@@ -413,4 +431,5 @@ def validate_workflow_graph(
edge_groups,
executors,
start_executor,
output_executors,
)
@@ -180,6 +180,7 @@ class Workflow(DictConvertible):
max_iterations: int = DEFAULT_MAX_ITERATIONS,
name: str | None = None,
description: str | None = None,
output_executors: list[str] | None = None,
**kwargs: Any,
):
"""Initialize the workflow with a list of edges.
@@ -192,6 +193,8 @@ class Workflow(DictConvertible):
max_iterations: The maximum number of iterations the workflow will run for convergence.
name: Optional human-readable name for the workflow.
description: Optional description of what the workflow does.
output_executors: Optional list of executor IDs whose outputs will be considered workflow outputs.
If None or empty, all executor outputs are treated as workflow outputs.
kwargs: Additional keyword arguments. Unused in this implementation.
"""
self.edge_groups = list(edge_groups)
@@ -202,6 +205,10 @@ class Workflow(DictConvertible):
self.name = name
self.description = description
# `WorkflowOutputEvent`s from these executors are treated as workflow outputs.
# If None or empty, all executor outputs are considered workflow outputs.
self._output_executors = list(output_executors) if output_executors else list(self.executors.keys())
# Store non-serializable runtime objects as private attributes
self._runner_context = runner_context
self._shared_state = SharedState()
@@ -241,6 +248,7 @@ class Workflow(DictConvertible):
"max_iterations": self.max_iterations,
"edge_groups": [group.to_dict() for group in self.edge_groups],
"executors": {executor_id: executor.to_dict() for executor_id, executor in self.executors.items()},
"output_executors": self._output_executors,
}
# Add optional name and description if provided
@@ -277,6 +285,10 @@ class Workflow(DictConvertible):
"""
return self.executors[self.start_executor_id]
def get_output_executors(self) -> list[Executor]:
"""Get the list of output executors in the workflow."""
return [self.executors[executor_id] for executor_id in self._output_executors]
def get_executors_list(self) -> list[Executor]:
"""Get the list of executors in the workflow."""
return list(self.executors.values())
@@ -539,6 +551,8 @@ class Workflow(DictConvertible):
streaming=True,
run_kwargs=kwargs if kwargs else None,
):
if isinstance(event, WorkflowOutputEvent) and not self._should_yield_output_event(event):
continue
yield event
finally:
if checkpoint_storage is not None:
@@ -562,6 +576,8 @@ class Workflow(DictConvertible):
reset_context=False, # Don't reset context when sending responses
streaming=True,
):
if isinstance(event, WorkflowOutputEvent) and not self._should_yield_output_event(event):
continue
yield event
finally:
self._reset_running_flag()
@@ -687,6 +703,8 @@ class Workflow(DictConvertible):
if include_status_events:
filtered.append(ev)
continue
if isinstance(ev, WorkflowOutputEvent) and not self._should_yield_output_event(ev):
continue
filtered.append(ev)
return WorkflowRunResult(filtered, status_events)
@@ -710,7 +728,13 @@ class Workflow(DictConvertible):
)
]
status_events = [e for e in events if isinstance(e, WorkflowStatusEvent)]
filtered_events = [e for e in events if not isinstance(e, (WorkflowStatusEvent, WorkflowStartedEvent))]
filtered_events: list[WorkflowEvent] = []
for e in events:
if isinstance(e, WorkflowOutputEvent) and not self._should_yield_output_event(e):
continue
if isinstance(e, (WorkflowStatusEvent, WorkflowStartedEvent)):
continue
filtered_events.append(e)
return WorkflowRunResult(filtered_events, status_events)
finally:
self._reset_running_flag()
@@ -750,6 +774,22 @@ class Workflow(DictConvertible):
raise ValueError(f"Executor with ID {executor_id} not found.")
return self.executors[executor_id]
def _should_yield_output_event(self, event: WorkflowOutputEvent) -> bool:
"""Determine if a WorkflowOutputEvent should be yielded as a workflow output.
Args:
event: The WorkflowOutputEvent to evaluate.
Returns:
True if the event should be yielded as a workflow output, False otherwise.
"""
# If no specific output executors are defined, yield all outputs
if not self._output_executors:
return True
# Check if the event's source executor is in the list of output executors
return event.executor_id in self._output_executors
# Graph signature helpers
def _compute_graph_signature(self) -> dict[str, Any]:
@@ -6,12 +6,11 @@ from collections.abc import Callable, Sequence
from dataclasses import dataclass
from typing import Any
from typing_extensions import deprecated
from .._agents import AgentProtocol
from .._threads import AgentThread
from ..observability import OtelAttr, capture_exception, create_workflow_span
from ._agent_executor import AgentExecutor
from ._agent_utils import resolve_agent_id
from ._checkpoint import CheckpointStorage
from ._const import DEFAULT_MAX_ITERATIONS
from ._edge import (
@@ -173,10 +172,9 @@ class WorkflowBuilder:
self._name: str | None = name
self._description: str | None = description
# Maps underlying AgentProtocol object id -> wrapped Executor so we reuse the same wrapper
# across set_start_executor / add_edge calls. Without this, unnamed agents (which receive
# random UUID based executor ids) end up wrapped multiple times, giving different ids for
# the start node vs edge nodes and triggering a GraphConnectivityError during validation.
self._agent_wrappers: dict[int, Executor] = {}
# across set_start_executor / add_edge calls. This avoids multiple AgentExecutor instances
# being created for the same agent.
self._agent_wrappers: dict[str, Executor] = {}
# Registrations for lazy initialization of executors
self._edge_registry: list[
@@ -188,6 +186,9 @@ class WorkflowBuilder:
] = []
self._executor_registry: dict[str, Callable[[], Executor]] = {}
# Output executors filter; if set, only outputs from these executors are yielded
self._output_executors: list[Executor | AgentProtocol | str] = []
# Agents auto-wrapped by builder now always stream incremental updates.
def _add_executor(self, executor: Executor) -> str:
@@ -207,13 +208,7 @@ class WorkflowBuilder:
return executor.id
def _maybe_wrap_agent(
self,
candidate: Executor | AgentProtocol,
agent_thread: Any | None = None,
output_response: bool = False,
executor_id: str | None = None,
) -> Executor:
def _maybe_wrap_agent(self, candidate: Executor | AgentProtocol) -> Executor:
"""If the provided object implements AgentProtocol, wrap it in an AgentExecutor.
This allows fluent builder APIs to directly accept agents instead of
@@ -221,9 +216,9 @@ class WorkflowBuilder:
Args:
candidate: The executor or agent to wrap.
agent_thread: The thread to use for running the agent. If None, a new thread will be created.
output_response: Whether to yield an AgentResponse as a workflow output when the agent completes.
executor_id: A unique identifier for the executor. If None, the agent's name will be used if available.
Returns:
An Executor instance, wrapping the agent if necessary.
"""
try: # Local import to avoid hard dependency at import time
from agent_framework import AgentProtocol # type: ignore
@@ -234,28 +229,20 @@ class WorkflowBuilder:
return candidate
if isinstance(candidate, AgentProtocol): # type: ignore[arg-type]
# Reuse existing wrapper for the same agent instance if present
agent_instance_id = id(candidate)
agent_instance_id = str(id(candidate))
existing = self._agent_wrappers.get(agent_instance_id)
if existing is not None:
return existing
# Use agent name if available and unique among current executors
name = getattr(candidate, "name", None)
proposed_id: str | None = executor_id
if proposed_id is None and name:
proposed_id = str(name)
if proposed_id in self._executors:
raise ValueError(
f"Duplicate executor ID '{proposed_id}' from agent name. "
"Agent names must be unique within a workflow."
)
wrapper = AgentExecutor(
candidate,
agent_thread=agent_thread,
output_response=output_response,
id=proposed_id,
)
executor_id = resolve_agent_id(candidate)
if executor_id in self._executors:
raise ValueError(
f"Duplicate executor ID '{executor_id}' from agent. "
"Agent IDs or names must be unique within a workflow."
)
wrapper = AgentExecutor(candidate, id=executor_id)
self._agent_wrappers[agent_instance_id] = wrapper
return wrapper
raise TypeError(
f"WorkflowBuilder expected an Executor or AgentProtocol instance; got {type(candidate).__name__}."
)
@@ -337,7 +324,6 @@ class WorkflowBuilder:
factory_func: Callable[[], AgentProtocol],
name: str,
agent_thread: AgentThread | None = None,
output_response: bool = False,
) -> Self:
"""Register an agent factory function for lazy initialization.
@@ -351,7 +337,6 @@ class WorkflowBuilder:
the agent's internal name. But it must be unique within the workflow.
agent_thread: The thread to use for running the agent. If None, a new thread will be created when
the agent is instantiated.
output_response: Whether to yield an AgentResponse as a workflow output when the agent completes.
Example:
.. code-block:: python
@@ -382,69 +367,12 @@ class WorkflowBuilder:
return AgentExecutor(
agent,
agent_thread=agent_thread,
output_response=output_response,
)
self._executor_registry[name] = wrapped_factory
return self
@deprecated("Use register_agent() for lazy initialization instead.")
def add_agent(
self,
agent: AgentProtocol,
agent_thread: Any | None = None,
output_response: bool = False,
id: str | None = None,
) -> Self:
"""Add an agent to the workflow by wrapping it in an AgentExecutor.
This method creates an AgentExecutor that wraps the agent with the given parameters
and ensures that subsequent uses of the same agent instance in other builder methods
(like add_edge, set_start_executor, etc.) will reuse the same wrapped executor.
Note: Agents adapt their behavior based on how the workflow is executed:
- run_stream(): Agents emit incremental AgentRunUpdateEvent events as tokens are produced
- run(): Agents emit a single AgentRunEvent containing the complete response
Args:
agent: The agent to add to the workflow.
agent_thread: The thread to use for running the agent. If None, a new thread will be created.
output_response: Whether to yield an AgentResponse as a workflow output when the agent completes.
id: A unique identifier for the executor. If None, the agent's name will be used if available.
Returns:
Self: The WorkflowBuilder instance for method chaining.
Raises:
ValueError: If the provided id or agent name conflicts with an existing executor.
Example:
.. code-block:: python
from agent_framework import WorkflowBuilder
from agent_framework_anthropic import AnthropicAgent
# Create an agent
agent = AnthropicAgent(name="writer", model="claude-3-5-sonnet-20241022")
# Add the agent to a workflow
workflow = WorkflowBuilder().add_agent(agent, output_response=True).set_start_executor(agent).build()
"""
logger.warning(
"Adding an agent instance directly to WorkflowBuilder is not recommended, "
"because workflow instances created from the builder will share the same agent instance. "
"Consider using register_agent() for lazy initialization instead."
)
executor = self._maybe_wrap_agent(
agent,
agent_thread=agent_thread,
output_response=output_response,
executor_id=id,
)
self._add_executor(executor)
return self
def add_edge(
self,
source: Executor | AgentProtocol | str,
@@ -1139,10 +1067,35 @@ class WorkflowBuilder:
self._checkpoint_storage = checkpoint_storage
return self
def _resolve_edge_registry(
self,
) -> tuple[Executor, list[Executor], list[EdgeGroup]]:
"""Resolve deferred edge registrations into executors and edge groups."""
def with_output_from(self, executors: list[Executor | AgentProtocol | str]) -> Self:
"""Specify which executors' outputs should be collected as workflow outputs.
By default, outputs from all executors are collected. This method allows
filtering to only include outputs from specified executors.
Args:
executors: A list of executors or registered names of the executor factories
whose outputs should be collected.
Returns:
Self: The WorkflowBuilder instance for method chaining.
"""
self._output_executors = list(executors)
return self
def _resolve_edge_registry(self) -> tuple[Executor, dict[str, Executor], list[EdgeGroup]]:
"""Resolve deferred edge registrations into executors and edge groups.
Returns:
tuple: A tuple containing:
- The starting Executor instance.
- A dictionary mapping registered factory names to resolved Executor instances.
- A list of EdgeGroup instances representing the workflow edges composed of resolved executors.
Notes:
Non-factory executors (i.e., those added directly) are not included in the returned list,
as they are already part of the workflow builder's internal state.
"""
if not self._start_executor:
raise ValueError("Starting executor must be set using set_start_executor before building the workflow.")
@@ -1158,7 +1111,9 @@ class WorkflowBuilder:
for name, exec_factory in self._executor_registry.items():
instance = exec_factory()
if instance.id in executor_id_to_instance:
raise ValueError(f"Executor with ID '{instance.id}' has already been created.")
raise ValueError(f"Executor with ID '{instance.id}' has already been registered.")
if instance.id in self._executors:
raise ValueError(f"Executor ID collision: An executor with ID '{instance.id}' already exists.")
executor_id_to_instance[instance.id] = instance
if isinstance(self._start_executor, str) and name == self._start_executor:
@@ -1211,11 +1166,7 @@ class WorkflowBuilder:
if start_executor is None:
raise ValueError("Failed to resolve starting executor from registered factories.")
return (
start_executor,
list(executor_id_to_instance.values()),
deferred_edge_groups,
)
return (start_executor, factory_name_to_instance, deferred_edge_groups)
def build(self) -> Workflow:
"""Build and return the constructed workflow.
@@ -1271,14 +1222,24 @@ class WorkflowBuilder:
# Resolve lazy edge registrations
start_executor, deferred_executors, deferred_edge_groups = self._resolve_edge_registry()
executors = self._executors | {exe.id: exe for exe in deferred_executors}
executors = self._executors | {exe.id: exe for exe in deferred_executors.values()}
edge_groups = self._edge_groups + deferred_edge_groups
output_executors = (
[
deferred_executors[factory_name].id
for factory_name in self._output_executors
if isinstance(factory_name, str)
]
+ [ex.id for ex in self._output_executors if isinstance(ex, Executor)]
+ [resolve_agent_id(agent) for agent in self._output_executors if isinstance(agent, AgentProtocol)]
)
# Perform validation before creating the workflow
validate_workflow_graph(
edge_groups,
executors,
start_executor,
output_executors,
)
# Add validation completed event
@@ -1295,6 +1256,7 @@ class WorkflowBuilder:
self._max_iterations,
name=self._name,
description=self._description,
output_executors=output_executors,
)
build_attributes: dict[str, Any] = {
OtelAttr.WORKFLOW_ID: workflow.id,
@@ -2,7 +2,7 @@
"""Tests for AgentExecutor handling of tool calls and results in streaming mode."""
from collections.abc import AsyncIterable
from collections.abc import AsyncIterable, Sequence
from typing import Any
from typing_extensions import Never
@@ -12,7 +12,6 @@ from agent_framework import (
AgentExecutorResponse,
AgentResponse,
AgentResponseUpdate,
AgentRunUpdateEvent,
AgentThread,
BaseAgent,
ChatAgent,
@@ -38,7 +37,7 @@ class _ToolCallingAgent(BaseAgent):
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -48,7 +47,7 @@ class _ToolCallingAgent(BaseAgent):
async def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -99,9 +98,9 @@ async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None:
workflow = WorkflowBuilder().set_start_executor(agent_exec).build()
# Act: run in streaming mode
events: list[AgentRunUpdateEvent] = []
events: list[WorkflowOutputEvent] = []
async for event in workflow.run_stream("What's the weather?"):
if isinstance(event, AgentRunUpdateEvent):
if isinstance(event, WorkflowOutputEvent):
events.append(event)
# Assert: we should receive 4 events (text, function call, function result, text)
@@ -148,7 +147,7 @@ class MockChatClient:
async def get_response(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage],
messages: str | ChatMessage | Sequence[str | ChatMessage],
**kwargs: Any,
) -> ChatResponse:
if self._iteration == 0:
@@ -185,7 +184,7 @@ class MockChatClient:
async def get_streaming_response(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage],
messages: str | ChatMessage | Sequence[str | ChatMessage],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
if self._iteration == 0:
@@ -231,7 +230,13 @@ async def test_agent_executor_tool_call_with_approval() -> None:
tools=[mock_tool_requiring_approval],
)
workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build()
workflow = (
WorkflowBuilder()
.set_start_executor(agent)
.add_edge(agent, test_executor)
.with_output_from([test_executor])
.build()
)
# Act
events = await workflow.run("Invoke tool requiring approval")
@@ -300,7 +305,13 @@ async def test_agent_executor_parallel_tool_call_with_approval() -> None:
tools=[mock_tool_requiring_approval],
)
workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build()
workflow = (
WorkflowBuilder()
.set_start_executor(agent)
.add_edge(agent, test_executor)
.with_output_from([test_executor])
.build()
)
# Act
events = await workflow.run("Invoke tool requiring approval")
@@ -1,15 +1,15 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for AgentRunEvent and AgentRunUpdateEvent type annotations."""
"""Tests for agent run event typing."""
from agent_framework import AgentResponse, AgentResponseUpdate, ChatMessage
from agent_framework._workflows._events import AgentRunEvent, AgentRunUpdateEvent
from agent_framework._workflows._events import WorkflowOutputEvent
def test_agent_run_event_data_type() -> None:
"""Verify AgentRunEvent.data is typed as AgentResponse | None."""
response = AgentResponse(messages=[ChatMessage("assistant", ["Hello"])])
event = AgentRunEvent(executor_id="test", data=response)
"""Verify WorkflowOutputEvent.data is typed as AgentResponse | None."""
response = AgentResponse(messages=[ChatMessage(role="assistant", text="Hello")])
event = WorkflowOutputEvent(data=response, executor_id="test")
# This assignment should pass type checking without a cast
data: AgentResponse | None = event.data
@@ -18,9 +18,9 @@ def test_agent_run_event_data_type() -> None:
def test_agent_run_update_event_data_type() -> None:
"""Verify AgentRunUpdateEvent.data is typed as AgentResponseUpdate | None."""
"""Verify WorkflowOutputEvent.data is typed as AgentResponseUpdate | None."""
update = AgentResponseUpdate()
event = AgentRunUpdateEvent(executor_id="test", data=update)
event = WorkflowOutputEvent(data=update, executor_id="test")
# This assignment should pass type checking without a cast
data: AgentResponseUpdate | None = event.data
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import AsyncIterable
from collections.abc import AsyncIterable, Sequence
from typing import Any
from pydantic import PrivateAttr
@@ -34,7 +34,7 @@ class _SimpleAgent(BaseAgent):
async def run( # type: ignore[override]
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -43,7 +43,7 @@ class _SimpleAgent(BaseAgent):
async def run_stream( # type: ignore[override]
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -56,7 +56,7 @@ class _CaptureFullConversation(Executor):
"""Captures AgentExecutorResponse.full_conversation and completes the workflow."""
@handler
async def capture(self, response: AgentExecutorResponse, ctx: WorkflowContext[Never, dict]) -> None:
async def capture(self, response: AgentExecutorResponse, ctx: WorkflowContext[Never, dict[str, Any]]) -> None:
full = response.full_conversation
# The AgentExecutor contract guarantees full_conversation is populated.
assert full is not None
@@ -75,7 +75,13 @@ async def test_agent_executor_populates_full_conversation_non_streaming() -> Non
agent_exec = AgentExecutor(agent, id="agent1-exec")
capturer = _CaptureFullConversation(id="capture")
wf = WorkflowBuilder().set_start_executor(agent_exec).add_edge(agent_exec, capturer).build()
wf = (
WorkflowBuilder()
.set_start_executor(agent_exec)
.add_edge(agent_exec, capturer)
.with_output_from([capturer])
.build()
)
# Act: use run() instead of run_stream() to test non-streaming mode
result = await wf.run("hello world")
@@ -103,7 +109,7 @@ class _CaptureAgent(BaseAgent):
async def run( # type: ignore[override]
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -121,7 +127,7 @@ class _CaptureAgent(BaseAgent):
async def run_stream( # type: ignore[override]
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -11,7 +11,6 @@ from agent_framework import (
AgentProtocol,
AgentResponse,
AgentResponseUpdate,
AgentRunUpdateEvent,
AgentThread,
BaseAgent,
ChatMessage,
@@ -574,15 +573,19 @@ class StubAssistantsAgent(BaseAgent):
async def _collect_agent_responses_setup(participant: AgentProtocol) -> list[ChatMessage]:
captured: list[ChatMessage] = []
wf = MagenticBuilder().participants([participant]).with_manager(manager=InvokeOnceManager()).build()
wf = (
MagenticBuilder()
.participants([participant])
.with_manager(manager=InvokeOnceManager())
.with_intermediate_outputs()
.build()
)
# Run a bounded stream to allow one invoke and then completion
events: list[WorkflowEvent] = []
async for ev in wf.run_stream("task"): # plan review disabled
events.append(ev)
if isinstance(ev, WorkflowOutputEvent):
break
if isinstance(ev, AgentRunUpdateEvent):
if isinstance(ev, WorkflowOutputEvent) and isinstance(ev.data, AgentResponseUpdate):
captured.append(
ChatMessage(
role=ev.data.role or "assistant",
@@ -597,7 +600,6 @@ async def _collect_agent_responses_setup(participant: AgentProtocol) -> list[Cha
async def test_agent_executor_invoke_with_thread_chat_client():
agent = StubThreadAgent()
captured = await _collect_agent_responses_setup(agent)
# Should have at least one response from agentA via _MagenticAgentExecutor path
assert any((m.author_name == agent.name and "ok" in (m.text or "")) for m in captured)
@@ -177,7 +177,7 @@ def test_graph_connectivity_isolated_executors():
executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2, executor3.id: executor3}
with pytest.raises(GraphConnectivityError) as exc_info:
validate_workflow_graph(edge_groups, executors, executor1)
validate_workflow_graph(edge_groups, executors, executor1, [])
assert "unreachable" in str(exc_info.value).lower()
assert "executor3" in str(exc_info.value)
@@ -258,12 +258,12 @@ def test_direct_validation_function():
executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2}
# This should not raise any exceptions
validate_workflow_graph(edge_groups, executors, executor1)
validate_workflow_graph(edge_groups, executors, executor1, [])
# Test with invalid start executor
executor3 = StringExecutor(id="executor3")
with pytest.raises(GraphConnectivityError):
validate_workflow_graph(edge_groups, executors, executor3)
validate_workflow_graph(edge_groups, executors, executor3, [])
def test_fan_out_validation():
@@ -557,3 +557,155 @@ def test_handler_ctx_any_is_allowed_but_skips_type_checks(caplog: Any) -> None:
# Builds; later edges from this executor will skip type compatibility when outputs are unspecified
wf = WorkflowBuilder().add_edge(start, any_out).set_start_executor(start).build()
assert wf is not None
# region Output Validation Tests
class OutputExecutor(Executor):
@handler
async def handle_string(self, message: str, ctx: WorkflowContext[str, str]) -> None:
pass
def test_output_validation_with_valid_output_executors():
"""Test that output validation passes when output executors exist and have output types."""
executor1 = OutputExecutor(id="executor1")
executor2 = OutputExecutor(id="executor2")
# Build workflow with valid output executors
workflow = (
WorkflowBuilder()
.add_edge(executor1, executor2)
.set_start_executor(executor1)
.with_output_from([executor2])
.build()
)
assert workflow is not None
assert workflow._output_executors == ["executor2"]
def test_output_validation_with_multiple_valid_output_executors():
"""Test that output validation passes with multiple valid output executors."""
executor1 = OutputExecutor(id="executor1")
executor2 = OutputExecutor(id="executor2")
executor3 = OutputExecutor(id="executor3")
workflow = (
WorkflowBuilder()
.add_edge(executor1, executor2)
.add_edge(executor2, executor3)
.set_start_executor(executor1)
.with_output_from([executor1, executor3])
.build()
)
assert workflow is not None
assert set(workflow._output_executors) == {"executor1", "executor3"}
def test_output_validation_fails_for_nonexistent_executor():
"""Test that output validation fails when an output executor doesn't exist in the graph."""
executor1 = OutputExecutor(id="executor1")
executor2 = OutputExecutor(id="executor2")
edge_groups = [SingleEdgeGroup(executor1.id, executor2.id)]
executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2}
# Directly test validation with a nonexistent output executor
with pytest.raises(WorkflowValidationError) as exc_info:
validate_workflow_graph(edge_groups, executors, executor1, ["nonexistent_executor"])
assert "not present in the workflow graph" in str(exc_info.value)
assert "nonexistent_executor" in str(exc_info.value)
assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
def test_output_validation_fails_for_executor_without_output_types():
"""Test that output validation fails when an output executor has no output type annotations."""
executor1 = OutputExecutor(id="executor1")
no_output_executor = NoOutputTypesExecutor(id="no_output")
with pytest.raises(WorkflowValidationError) as exc_info:
(
WorkflowBuilder()
.add_edge(executor1, no_output_executor)
.set_start_executor(executor1)
.with_output_from([no_output_executor])
.build()
)
assert "must have output type annotations defined" in str(exc_info.value)
assert "no_output" in str(exc_info.value)
assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
def test_output_validation_empty_list_passes():
"""Test that output validation passes with an empty output executors list."""
executor1 = OutputExecutor(id="executor1")
executor2 = OutputExecutor(id="executor2")
workflow = (
WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).with_output_from([]).build()
)
assert workflow is not None
# All executors are outputs
assert workflow._output_executors == ["executor1", "executor2"] # type: ignore
def test_output_validation_with_direct_validate_workflow_graph():
"""Test _output_validation directly via validate_workflow_graph function."""
executor1 = OutputExecutor(id="executor1")
executor2 = OutputExecutor(id="executor2")
edge_groups = [SingleEdgeGroup(executor1.id, executor2.id)]
executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2}
# Valid output executors
validate_workflow_graph(edge_groups, executors, executor1, ["executor2"])
# Invalid output executor (doesn't exist)
with pytest.raises(WorkflowValidationError) as exc_info:
validate_workflow_graph(edge_groups, executors, executor1, ["nonexistent"])
assert "not present in the workflow graph" in str(exc_info.value)
assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
def test_output_validation_with_no_output_types_via_direct_validation():
"""Test _output_validation fails for executors without output types via direct validation."""
executor1 = OutputExecutor(id="executor1")
no_output_executor = NoOutputTypesExecutor(id="no_output")
edge_groups = [SingleEdgeGroup(executor1.id, no_output_executor.id)]
executors: dict[str, Executor] = {executor1.id: executor1, no_output_executor.id: no_output_executor}
# Should fail because no_output_executor has no output types
with pytest.raises(WorkflowValidationError) as exc_info:
validate_workflow_graph(edge_groups, executors, executor1, ["no_output"])
assert "must have output type annotations defined" in str(exc_info.value)
assert exc_info.value.validation_type == ValidationTypeEnum.OUTPUT_VALIDATION
def test_output_validation_partial_invalid_list():
"""Test that output validation fails if any executor in the list is invalid."""
executor1 = OutputExecutor(id="executor1")
executor2 = OutputExecutor(id="executor2")
edge_groups = [SingleEdgeGroup(executor1.id, executor2.id)]
executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2}
# First executor is valid, second doesn't exist - validation should fail
with pytest.raises(WorkflowValidationError) as exc_info:
validate_workflow_graph(edge_groups, executors, executor1, ["executor2", "nonexistent"])
assert "not present in the workflow graph" in str(exc_info.value)
assert "nonexistent" in str(exc_info.value)
def test_output_validation_type_enum_value():
"""Test that OUTPUT_VALIDATION is properly defined in ValidationTypeEnum."""
assert hasattr(ValidationTypeEnum, "OUTPUT_VALIDATION")
assert ValidationTypeEnum.OUTPUT_VALIDATION.value == "OUTPUT_VALIDATION"
# endregion
@@ -2,9 +2,9 @@
import asyncio
import tempfile
from collections.abc import AsyncIterable
from collections.abc import AsyncIterable, Sequence
from dataclasses import dataclass, field
from typing import Any
from typing import Any, cast
from uuid import uuid4
import pytest
@@ -13,8 +13,6 @@ from agent_framework import (
AgentExecutor,
AgentResponse,
AgentResponseUpdate,
AgentRunEvent,
AgentRunUpdateEvent,
AgentThread,
BaseAgent,
ChatMessage,
@@ -862,7 +860,7 @@ class _StreamingTestAgent(BaseAgent):
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -872,7 +870,7 @@ class _StreamingTestAgent(BaseAgent):
async def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -884,7 +882,7 @@ class _StreamingTestAgent(BaseAgent):
async def test_agent_streaming_vs_non_streaming() -> None:
"""Test that run() emits AgentRunEvent while run_stream() emits AgentRunUpdateEvent."""
"""Test that run() and run_stream() both emits WorkflowOutputEvents correctly with the right data types."""
agent = _StreamingTestAgent(id="test_agent", name="TestAgent", reply_text="Hello World")
agent_exec = AgentExecutor(agent, id="agent_exec")
@@ -894,15 +892,17 @@ async def test_agent_streaming_vs_non_streaming() -> None:
result = await workflow.run("test message")
# Filter for agent events (result is a list of events)
agent_run_events = [e for e in result if isinstance(e, AgentRunEvent)]
agent_update_events = [e for e in result if isinstance(e, AgentRunUpdateEvent)]
agent_response = [e for e in result if isinstance(e, WorkflowOutputEvent) and isinstance(e.data, AgentResponse)]
agent_response_updates = [
e for e in result if isinstance(e, WorkflowOutputEvent) and isinstance(e.data, AgentResponseUpdate)
]
# In non-streaming mode, should have AgentRunEvent, no AgentRunUpdateEvent
assert len(agent_run_events) == 1, "Expected exactly one AgentRunEvent in non-streaming mode"
assert len(agent_update_events) == 0, "Expected no AgentRunUpdateEvent in non-streaming mode"
assert agent_run_events[0].executor_id == "agent_exec"
assert agent_run_events[0].data is not None
assert agent_run_events[0].data.messages[0].text == "Hello World"
# In non-streaming mode, should have AgentResponse, no AgentResponseUpdate
assert len(agent_response) == 1, "Expected exactly one AgentResponse in non-streaming mode"
assert len(agent_response_updates) == 0, "Expected no AgentResponseUpdate in non-streaming mode"
assert agent_response[0].executor_id == "agent_exec"
assert agent_response[0].data is not None
assert agent_response[0].data.messages[0].text == "Hello World"
# Test streaming mode with run_stream()
stream_events: list[WorkflowEvent] = []
@@ -910,22 +910,31 @@ async def test_agent_streaming_vs_non_streaming() -> None:
stream_events.append(event)
# Filter for agent events
stream_agent_run_events = [e for e in stream_events if isinstance(e, AgentRunEvent)]
stream_agent_update_events = [e for e in stream_events if isinstance(e, AgentRunUpdateEvent)]
agent_response = [
cast(AgentResponse, e.data) # type: ignore
for e in stream_events
if isinstance(e, WorkflowOutputEvent) and isinstance(e.data, AgentResponse)
]
agent_response_updates = [
e.data for e in stream_events if isinstance(e, WorkflowOutputEvent) and isinstance(e.data, AgentResponseUpdate)
]
# In streaming mode, should have AgentRunUpdateEvent, no AgentRunEvent
assert len(stream_agent_run_events) == 0, "Expected no AgentRunEvent in streaming mode"
assert len(stream_agent_update_events) > 0, "Expected AgentRunUpdateEvent events in streaming mode"
# In streaming mode, should have AgentResponseUpdate, no AgentResponse
assert len(agent_response) == 0, "Expected no AgentResponse in streaming mode"
assert len(agent_response_updates) > 0, "Expected AgentResponseUpdate events in streaming mode"
# Verify we got incremental updates (one per character in "Hello World")
assert len(stream_agent_update_events) == len("Hello World"), "Expected one update per character"
assert len(agent_response_updates) == len("Hello World"), "Expected one update per character"
# Verify the updates build up to the full message
accumulated_text = "".join(
e.data.contents[0].text
for e in stream_agent_update_events
if e.data and e.data.contents and e.data.contents[0].text
)
accumulated_text = "".join([
e.contents[0].text
for e in agent_response_updates
if e.contents
and isinstance(e.contents[0], Content)
and e.contents[0].type == "text"
and e.contents[0].text is not None
])
assert accumulated_text == "Hello World", f"Expected 'Hello World', got '{accumulated_text}'"
@@ -974,3 +983,253 @@ async def test_workflow_run_stream_parameter_validation(
# Invalid combinations already tested in test_workflow_run_parameter_validation
# This test ensures streaming works correctly for valid parameters
# region Output executor filtering tests
class OutputProducerExecutor(Executor):
"""An executor that produces a unique output value for testing output filtering."""
def __init__(self, id: str, output_value: int) -> None:
super().__init__(id=id)
self.output_value = output_value
@handler
async def handle_message(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage, int]) -> None:
await ctx.yield_output(self.output_value)
class PassthroughExecutor(Executor):
"""An executor that passes through messages and produces an output."""
def __init__(self, id: str, output_value: int) -> None:
super().__init__(id=id)
self.output_value = output_value
@handler
async def handle_message(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage, int]) -> None:
await ctx.yield_output(self.output_value)
await ctx.send_message(message)
async def test_output_executors_empty_yields_all_outputs() -> None:
"""Test that when _output_executors is empty (default), all outputs are yielded."""
# Create executors that each produce different outputs
executor_a = PassthroughExecutor(id="executor_a", output_value=10)
executor_b = OutputProducerExecutor(id="executor_b", output_value=20)
# Build workflow with a -> b
workflow = WorkflowBuilder().set_start_executor(executor_a).add_edge(executor_a, executor_b).build()
result = await workflow.run(NumberMessage(data=0))
outputs = result.get_outputs()
# Both executors' outputs should be present
assert len(outputs) == 2
assert outputs == [10, 20]
output_events = [event for event in result if isinstance(event, WorkflowOutputEvent)]
assert len(output_events) == 2
assert output_events[0].executor_id == "executor_a"
assert output_events[1].executor_id == "executor_b"
async def test_output_executors_filters_outputs_non_streaming() -> None:
"""Test that only outputs from specified executors are yielded in non-streaming mode."""
# Create executors that each produce different outputs
executor_a = PassthroughExecutor(id="executor_a", output_value=10)
executor_b = OutputProducerExecutor(id="executor_b", output_value=20)
# Build workflow with a -> b
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
.add_edge(executor_a, executor_b)
.with_output_from([executor_b])
.build()
)
result = await workflow.run(NumberMessage(data=0))
outputs = result.get_outputs()
# Only executor_b's output should be present
assert len(outputs) == 1
assert outputs[0] == 20
output_events = [event for event in result if isinstance(event, WorkflowOutputEvent)]
assert len(output_events) == 1
assert output_events[0].executor_id == "executor_b"
async def test_output_executors_filters_outputs_streaming() -> None:
"""Test that only outputs from specified executors are yielded in streaming mode."""
# Create executors that each produce different outputs
executor_a = PassthroughExecutor(id="executor_a", output_value=100)
executor_b = OutputProducerExecutor(id="executor_b", output_value=200)
# Build workflow with a -> b
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
.add_edge(executor_a, executor_b)
.with_output_from([executor_a])
.build()
)
# Collect outputs from streaming
output_events: list[WorkflowOutputEvent] = []
async for event in workflow.run_stream(NumberMessage(data=0)):
if isinstance(event, WorkflowOutputEvent):
output_events.append(event)
# Only executor_a's output should be present
assert len(output_events) == 1
assert output_events[0].data == 100
assert output_events[0].executor_id == "executor_a"
async def test_output_executors_with_multiple_specified_executors() -> None:
"""Test filtering with multiple executors in the output list."""
# Create three executors with pass-through to reach all of them
executor_a = PassthroughExecutor(id="executor_a", output_value=1)
executor_b = PassthroughExecutor(id="executor_b", output_value=2)
executor_c = OutputProducerExecutor(id="executor_c", output_value=3)
# Build workflow with a -> b -> c
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_c)
.with_output_from([executor_a, executor_c])
.build()
)
result = await workflow.run(NumberMessage(data=0))
outputs = result.get_outputs()
# Only executor_a and executor_c outputs should be present
assert len(outputs) == 2
assert 1 in outputs # executor_a
assert 3 in outputs # executor_c
assert 2 not in outputs # executor_b should be filtered out
async def test_output_executors_with_nonexistent_executor_id() -> None:
"""Test that specifying a non-existent executor ID doesn't break the workflow."""
executor_a = OutputProducerExecutor(id="executor_a", output_value=42)
workflow = WorkflowBuilder().set_start_executor(executor_a).build()
# Set output_executors to an ID that doesn't exist
workflow._output_executors = ["nonexistent_executor"] # type: ignore
result = await workflow.run(NumberMessage(data=0))
outputs = result.get_outputs()
# No outputs should be yielded since the executor ID doesn't match
assert len(outputs) == 0
async def test_output_executors_filtering_with_fan_in() -> None:
"""Test output filtering in a fan-in workflow."""
class FanOutStartExecutor(Executor):
"""Executor that sends messages to fan-out targets."""
@handler
async def handle(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage, int]) -> None:
await ctx.yield_output(999) # This should be filtered out
await ctx.send_message(NumberMessage(data=5))
class FanOutTargetExecutor(Executor):
"""Executor that processes fan-out messages."""
def __init__(self, id: str, increment: int) -> None:
super().__init__(id=id)
self.increment = increment
@handler
async def handle(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage, int]) -> None:
await ctx.yield_output(888) # This should be filtered out
await ctx.send_message(NumberMessage(data=message.data + self.increment))
# Create executors for fan-in pattern
executor_start = FanOutStartExecutor(id="executor_start")
executor_a = FanOutTargetExecutor(id="executor_a", increment=10)
executor_b = FanOutTargetExecutor(id="executor_b", increment=20)
aggregator = AggregatorExecutor(id="aggregator")
# Build fan-in workflow: start -> [a, b] -> aggregator
workflow = (
WorkflowBuilder()
.set_start_executor(executor_start)
.add_fan_out_edges(executor_start, [executor_a, executor_b])
.add_fan_in_edges([executor_a, executor_b], aggregator)
.with_output_from([aggregator])
.build()
)
result = await workflow.run(NumberMessage(data=0))
outputs = result.get_outputs()
# Only aggregator output should be present
# executor_a sends 5+10=15, executor_b sends 5+20=25, aggregator sums: 15+25=40
assert len(outputs) == 1
assert outputs[0] == 40
async def test_output_executors_filtering_with_send_responses() -> None:
"""Test output filtering works correctly with send_responses method."""
executor = MockExecutorRequestApproval(id="approval_executor")
workflow = WorkflowBuilder().set_start_executor(executor).with_output_from([executor]).build()
# Run workflow which will request approval
result = await workflow.run(NumberMessage(data=42))
# Get request info events
request_events = result.get_request_info_events()
assert len(request_events) == 1
# Send approval response
responses = {request_events[0].request_id: ApprovalMessage(approved=True)}
response_result = await workflow.send_responses(responses)
outputs = response_result.get_outputs()
# Output should be yielded since approval_executor is in output_executors
assert len(outputs) == 1
assert outputs[0] == 42
async def test_output_executors_filtering_with_send_responses_streaming() -> None:
"""Test output filtering works correctly with send_responses_streaming method."""
executor = MockExecutorRequestApproval(id="approval_executor")
workflow = WorkflowBuilder().set_start_executor(executor).build()
# Run workflow which will request approval
events_list: list[WorkflowEvent] = []
async for event in workflow.run_stream(NumberMessage(data=99)):
events_list.append(event)
# Get request info events
request_events = [e for e in events_list if isinstance(e, RequestInfoEvent)]
assert len(request_events) == 1
# Set output_executors to exclude the approval executor
workflow._output_executors = ["other_executor"] # type: ignore
# Send approval response via streaming
responses = {request_events[0].request_id: ApprovalMessage(approved=True)}
output_events: list[WorkflowOutputEvent] = []
async for event in workflow.send_responses_streaming(responses):
if isinstance(event, WorkflowOutputEvent):
output_events.append(event)
# No outputs should be yielded since approval_executor is not in output_executors
assert len(output_events) == 0
# endregion
@@ -1,16 +1,17 @@
# Copyright (c) Microsoft. All rights reserved.
import uuid
from collections.abc import AsyncIterable
from collections.abc import AsyncIterable, Sequence
from typing import Any
import pytest
from typing_extensions import Never
from agent_framework import (
AgentExecutorRequest,
AgentProtocol,
AgentResponse,
AgentResponseUpdate,
AgentRunUpdateEvent,
AgentThread,
ChatMessage,
ChatMessageStore,
@@ -27,26 +28,34 @@ from agent_framework import (
class SimpleExecutor(Executor):
"""Simple executor that emits AgentRunEvent or AgentRunStreamingEvent."""
"""Simple executor that emits a response based on input."""
def __init__(self, id: str, response_text: str, emit_streaming: bool = False):
def __init__(self, id: str, response_text: str, streaming: bool = False):
super().__init__(id=id)
self.response_text = response_text
self.emit_streaming = emit_streaming
self.streaming = streaming
@handler
async def handle_message(self, message: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None:
async def handle_message(
self,
message: list[ChatMessage],
ctx: WorkflowContext[list[ChatMessage], AgentResponseUpdate | AgentResponse],
) -> None:
input_text = message[0].contents[0].text if message and message[0].contents[0].type == "text" else "no input"
response_text = f"{self.response_text}: {input_text}"
# Create response message for both streaming and non-streaming cases
response_message = ChatMessage("assistant", [Content.from_text(text=response_text)])
# Emit update event.
streaming_update = AgentResponseUpdate(
contents=[Content.from_text(text=response_text)], role="assistant", message_id=str(uuid.uuid4())
)
await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=streaming_update))
if self.streaming:
# Emit update event.
streaming_update = AgentResponseUpdate(
contents=[Content.from_text(text=response_text)], role="assistant", message_id=str(uuid.uuid4())
)
await ctx.yield_output(streaming_update)
else:
response = AgentResponse(messages=[response_message])
await ctx.yield_output(response)
# Pass message to next executor if any (for both streaming and non-streaming)
await ctx.send_message([response_message])
@@ -55,6 +64,10 @@ class SimpleExecutor(Executor):
class RequestingExecutor(Executor):
"""Executor that requests info."""
def __init__(self, id: str, streaming: bool = False):
super().__init__(id=id)
self.streaming = streaming
@handler
async def handle_message(self, _: list[ChatMessage], ctx: WorkflowContext) -> None:
# Send a RequestInfoMessage to trigger the request info process
@@ -62,26 +75,49 @@ class RequestingExecutor(Executor):
@response_handler
async def handle_request_response(
self, original_request: str, response: str, ctx: WorkflowContext[ChatMessage]
self,
original_request: str,
response: str,
ctx: WorkflowContext[ChatMessage, AgentResponseUpdate | AgentResponse],
) -> None:
# Handle the response and emit completion response
update = AgentResponseUpdate(
contents=[Content.from_text(text="Request completed successfully")],
role="assistant",
message_id=str(uuid.uuid4()),
content = Content.from_text(text=f"Request completed with response: {response}")
if self.streaming:
await ctx.yield_output(
AgentResponseUpdate(
contents=[content],
role="assistant",
message_id=str(uuid.uuid4()),
)
)
return
await ctx.yield_output(
AgentResponse(
messages=[
ChatMessage(
role="assistant",
contents=[content],
)
],
)
)
await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=update))
class ConversationHistoryCapturingExecutor(Executor):
"""Executor that captures the received conversation history for verification."""
def __init__(self, id: str):
def __init__(self, id: str, streaming: bool = False):
super().__init__(id=id)
self.received_messages: list[ChatMessage] = []
self.streaming = streaming
@handler
async def handle_message(self, messages: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None:
async def handle_message(
self,
messages: list[ChatMessage],
ctx: WorkflowContext[list[ChatMessage], AgentResponseUpdate | AgentResponse],
) -> None:
# Capture all received messages
self.received_messages = list(messages)
@@ -91,10 +127,16 @@ class ConversationHistoryCapturingExecutor(Executor):
response_message = ChatMessage("assistant", [Content.from_text(text=response_text)])
streaming_update = AgentResponseUpdate(
contents=[Content.from_text(text=response_text)], role="assistant", message_id=str(uuid.uuid4())
)
await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=streaming_update))
if self.streaming:
# Emit streaming update
streaming_update = AgentResponseUpdate(
contents=[Content.from_text(text=response_text)], role="assistant", message_id=str(uuid.uuid4())
)
await ctx.yield_output(streaming_update)
else:
response = AgentResponse(messages=[response_message])
await ctx.yield_output(response)
await ctx.send_message([response_message])
@@ -102,10 +144,10 @@ class TestWorkflowAgent:
"""Test cases for WorkflowAgent end-to-end functionality."""
async def test_end_to_end_basic_workflow(self):
"""Test basic end-to-end workflow execution with 2 executors emitting AgentRunEvent."""
"""Test basic end-to-end workflow execution with 2 executors emitting AgentResponse."""
# Create workflow with two executors
executor1 = SimpleExecutor(id="executor1", response_text="Step1", emit_streaming=False)
executor2 = SimpleExecutor(id="executor2", response_text="Step2", emit_streaming=False)
executor1 = SimpleExecutor(id="executor1", response_text="Step1", streaming=False)
executor2 = SimpleExecutor(id="executor2", response_text="Step2", streaming=False)
workflow = WorkflowBuilder().set_start_executor(executor1).add_edge(executor1, executor2).build()
@@ -126,6 +168,7 @@ class TestWorkflowAgent:
first_content = message.contents[0]
if first_content.type == "text":
text = first_content.text
assert text is not None
if text.startswith("Step1:"):
step1_messages.append(message)
elif text.startswith("Step2:"):
@@ -136,16 +179,18 @@ class TestWorkflowAgent:
assert len(step2_messages) >= 1, "Should have received message from Step2 executor"
# Verify the processing worked for both
step1_text: str = step1_messages[0].contents[0].text # type: ignore[attr-defined]
step2_text: str = step2_messages[0].contents[0].text # type: ignore[attr-defined]
step1_text = step1_messages[0].contents[0].text
step2_text = step2_messages[0].contents[0].text
assert step1_text is not None
assert step2_text is not None
assert "Step1: Hello World" in step1_text
assert "Step2: Step1: Hello World" in step2_text
async def test_end_to_end_basic_workflow_streaming(self):
"""Test end-to-end workflow with streaming executor that emits AgentRunStreamingEvent."""
# Create a single streaming executor
executor1 = SimpleExecutor(id="stream1", response_text="Streaming1", emit_streaming=True)
executor2 = SimpleExecutor(id="stream2", response_text="Streaming2", emit_streaming=True)
executor1 = SimpleExecutor(id="stream1", response_text="Streaming1")
executor2 = SimpleExecutor(id="stream2", response_text="Streaming2")
# Create workflow with just one executor
workflow = WorkflowBuilder().set_start_executor(executor1).add_edge(executor1, executor2).build()
@@ -165,15 +210,17 @@ class TestWorkflowAgent:
first_content: Content = updates[0].contents[0] # type: ignore[assignment]
second_content: Content = updates[1].contents[0] # type: ignore[assignment]
assert first_content.type == "text"
assert first_content.text is not None
assert "Streaming1: Test input" in first_content.text
assert second_content.type == "text"
assert second_content.text is not None
assert "Streaming2: Streaming1: Test input" in second_content.text
async def test_end_to_end_request_info_handling(self):
"""Test end-to-end workflow with RequestInfoEvent handling."""
# Create workflow with requesting executor -> request info executor (no cycle)
simple_executor = SimpleExecutor(id="simple", response_text="SimpleResponse", emit_streaming=False)
requesting_executor = RequestingExecutor(id="requester")
simple_executor = SimpleExecutor(id="simple", response_text="SimpleResponse", streaming=False)
requesting_executor = RequestingExecutor(id="requester", streaming=False)
workflow = (
WorkflowBuilder().set_start_executor(simple_executor).add_edge(simple_executor, requesting_executor).build()
@@ -208,6 +255,8 @@ class TestWorkflowAgent:
assert function_call.arguments.get("request_id") == approval_request.id
# Approval request should reference the same function call
assert approval_request.id is not None
assert approval_request.function_call is not None
assert approval_request.function_call.call_id == function_call.call_id
assert approval_request.function_call.name == function_call.name
@@ -245,7 +294,7 @@ class TestWorkflowAgent:
def test_workflow_as_agent_method(self) -> None:
"""Test that Workflow.as_agent() creates a properly configured WorkflowAgent."""
# Create a simple workflow
executor = SimpleExecutor(id="executor1", response_text="Response", emit_streaming=False)
executor = SimpleExecutor(id="executor1", response_text="Response")
workflow = WorkflowBuilder().set_start_executor(executor).build()
# Test as_agent with a name
@@ -286,7 +335,7 @@ class TestWorkflowAgent:
"""
@executor
async def yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
async def yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext[Never, str]) -> None:
# Extract text from input for demonstration
input_text = messages[0].text if messages else "no input"
await ctx.yield_output(f"processed: {input_text}")
@@ -311,7 +360,7 @@ class TestWorkflowAgent:
"""Test that ctx.yield_output() surfaces as AgentResponseUpdate when streaming."""
@executor
async def yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
async def yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext[Never, str]) -> None:
await ctx.yield_output("first output")
await ctx.yield_output("second output")
@@ -331,7 +380,7 @@ class TestWorkflowAgent:
"""Test that yield_output preserves different content types (Content, Content, etc.)."""
@executor
async def content_yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
async def content_yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext[Never, Content]) -> None:
# Yield different content types
await ctx.yield_output(Content.from_text(text="text content"))
await ctx.yield_output(Content.from_data(data=b"binary data", media_type="application/octet-stream"))
@@ -359,7 +408,7 @@ class TestWorkflowAgent:
"""Test that yield_output with ChatMessage preserves the message structure."""
@executor
async def chat_message_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
async def chat_message_executor(messages: list[ChatMessage], ctx: WorkflowContext[Never, ChatMessage]) -> None:
msg = ChatMessage(
role="assistant",
contents=[Content.from_text(text="response text")],
@@ -389,7 +438,9 @@ class TestWorkflowAgent:
return f"CustomData({self.value})"
@executor
async def raw_yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
async def raw_yielding_executor(
messages: list[ChatMessage], ctx: WorkflowContext[Never, Content | CustomData | str]
) -> None:
# Yield different types of data
await ctx.yield_output("simple string")
await ctx.yield_output(Content.from_text(text="text content"))
@@ -408,8 +459,11 @@ class TestWorkflowAgent:
# Verify raw_representation is set for each update
assert updates[0].raw_representation == "simple string"
assert isinstance(updates[1].raw_representation, Content)
assert updates[1].raw_representation.type == "text"
assert updates[1].raw_representation.text == "text content"
assert isinstance(updates[2].raw_representation, CustomData)
assert updates[2].raw_representation.value == 42
@@ -421,7 +475,9 @@ class TestWorkflowAgent:
"""
@executor
async def list_yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
async def list_yielding_executor(
messages: list[ChatMessage], ctx: WorkflowContext[Never, list[ChatMessage]]
) -> None:
# Yield a list of ChatMessages (as SequentialBuilder does)
msg_list = [
ChatMessage("user", [Content.from_text(text="first message")]),
@@ -441,19 +497,20 @@ class TestWorkflowAgent:
async for update in agent.run_stream("test"):
updates.append(update)
assert len(updates) == 1
assert len(updates[0].contents) == 4
texts = [c.text for c in updates[0].contents if c.type == "text"]
assert texts == ["first message", "second message", "third", "fourth"]
assert len(updates) == 3
full_response = AgentResponse.from_updates(updates)
assert len(full_response.messages) == 3
texts = [message.text for message in full_response.messages]
# Note: `from_agent_run_response_updates` coalesces multiple text contents into one content
assert texts == ["first message", "second message", "thirdfourth"]
# Verify run() coalesces text contents (expected behavior)
# Verify run()
result = await agent.run("test")
assert isinstance(result, AgentResponse)
assert len(result.messages) == 1
# Content items are coalesced into one
assert len(result.messages[0].contents) == 1
assert result.messages[0].text == "first messagesecond messagethirdfourth"
assert len(result.messages) == 3
texts = [message.text for message in result.messages]
assert texts == ["first message", "second message", "third fourth"]
async def test_thread_conversation_history_included_in_workflow_run(self) -> None:
"""Test that conversation history from thread is included when running WorkflowAgent.
@@ -462,7 +519,7 @@ class TestWorkflowAgent:
the workflow receives the complete conversation history (thread history + new messages).
"""
# Create an executor that captures all received messages
capturing_executor = ConversationHistoryCapturingExecutor(id="capturing")
capturing_executor = ConversationHistoryCapturingExecutor(id="capturing", streaming=False)
workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Thread History Test Agent")
@@ -561,41 +618,41 @@ class TestWorkflowAgent:
"""Mock agent for testing."""
def __init__(self, name: str, response_text: str) -> None:
self._name = name
self.id = str(uuid.uuid4())
self.name = name
self.description: str | None = None
self._response_text = response_text
self._description: str | None = None
@property
def name(self) -> str | None:
return self._name
@property
def description(self) -> str | None:
return self._description
def get_new_thread(self) -> AgentThread:
def get_new_thread(self, **kwargs: Any) -> AgentThread:
return AgentThread()
async def run(self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any) -> AgentResponse:
async def run(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
return AgentResponse(
messages=[ChatMessage("assistant", [self._response_text])],
text=self._response_text,
)
async def run_stream(
self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
for word in self._response_text.split():
yield AgentResponseUpdate(
contents=[Content.from_text(text=word + " ")],
role="assistant",
author_name=self._name,
author_name=self.name,
)
@executor
async def start_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
from agent_framework import AgentExecutorRequest
async def start_executor(messages: list[ChatMessage], ctx: WorkflowContext[AgentExecutorRequest, str]) -> None:
await ctx.yield_output("Start output")
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
@@ -604,12 +661,11 @@ class TestWorkflowAgent:
WorkflowBuilder()
.register_executor(lambda: start_executor, "start")
.register_agent(lambda: MockAgent("agent1", "Agent1 output - should NOT appear"), "agent1")
.register_agent(
lambda: MockAgent("agent2", "Agent2 output - SHOULD appear"), "agent2", output_response=True
)
.register_agent(lambda: MockAgent("agent2", "Agent2 output - SHOULD appear"), "agent2")
.set_start_executor("start")
.add_edge("start", "agent1")
.add_edge("agent1", "agent2")
.with_output_from(["start", "agent2"])
.build()
)
@@ -635,47 +691,45 @@ class TestWorkflowAgent:
"""Mock agent for testing."""
def __init__(self, name: str, response_text: str) -> None:
self._name = name
self.id = str(uuid.uuid4())
self.name = name
self.description: str | None = None
self._response_text = response_text
self._description: str | None = None
@property
def name(self) -> str | None:
return self._name
@property
def description(self) -> str | None:
return self._description
def get_new_thread(self) -> AgentThread:
def get_new_thread(self, **kwargs: Any) -> AgentThread:
return AgentThread()
async def run(self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any) -> AgentResponse:
return AgentResponse(
messages=[ChatMessage("assistant", [self._response_text])],
text=self._response_text,
)
async def run(
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
return AgentResponse(messages=[ChatMessage("assistant", [self._response_text])])
async def run_stream(
self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any
self,
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(
contents=[Content.from_text(text=self._response_text)],
role="assistant",
author_name=self._name,
author_name=self.name,
)
@executor
async def start_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
from agent_framework import AgentExecutorRequest
async def start_executor(messages: list[ChatMessage], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
# Build workflow with single agent that has output_response=True
# Build workflow with single agent
workflow = (
WorkflowBuilder()
.register_executor(lambda: start_executor, "start")
.register_agent(lambda: MockAgent("agent", "Unique response text"), "agent", output_response=True)
.register_agent(lambda: MockAgent("agent", "Unique response text"), "agent")
.set_start_executor("start")
.add_edge("start", "agent")
.build()
@@ -694,14 +748,14 @@ class TestWorkflowAgent:
class TestWorkflowAgentAuthorName:
"""Test cases for author_name enrichment in WorkflowAgent (GitHub issue #1331)."""
async def test_agent_run_update_event_gets_executor_id_as_author_name(self):
"""Test that AgentRunUpdateEvent gets executor_id as author_name when not already set.
async def test_agent_response_update_gets_executor_id_as_author_name(self):
"""Test that AgentResponseUpdate gets executor_id as author_name when not already set.
This validates the fix for GitHub issue #1331: agent responses should include
identification of which agent produced them in multi-agent workflows.
"""
# Create workflow with executor that emits AgentRunUpdateEvent without author_name
executor1 = SimpleExecutor(id="my_executor_id", response_text="Response", emit_streaming=False)
# Create workflow with executor that emits AgentResponseUpdate without author_name
executor1 = SimpleExecutor(id="my_executor_id", response_text="Response")
workflow = WorkflowBuilder().set_start_executor(executor1).build()
agent = WorkflowAgent(workflow=workflow, name="Test Agent")
@@ -716,14 +770,18 @@ class TestWorkflowAgentAuthorName:
# Verify author_name is set to executor_id
assert updates[0].author_name == "my_executor_id"
async def test_agent_run_update_event_preserves_existing_author_name(self):
async def test_agent_response_update_preserves_existing_author_name(self):
"""Test that existing author_name is preserved and not overwritten."""
class AuthorNameExecutor(Executor):
"""Executor that sets author_name explicitly."""
@handler
async def handle_message(self, message: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None:
async def handle_message(
self,
message: list[ChatMessage],
ctx: WorkflowContext[list[ChatMessage], AgentResponseUpdate],
) -> None:
# Emit update with explicit author_name
update = AgentResponseUpdate(
contents=[Content.from_text(text="Response with author")],
@@ -731,7 +789,7 @@ class TestWorkflowAgentAuthorName:
author_name="custom_author_name", # Explicitly set
message_id=str(uuid.uuid4()),
)
await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=update))
await ctx.yield_output(update)
executor = AuthorNameExecutor(id="executor_id")
workflow = WorkflowBuilder().set_start_executor(executor).build()
@@ -749,8 +807,8 @@ class TestWorkflowAgentAuthorName:
async def test_multiple_executors_have_distinct_author_names(self):
"""Test that multiple executors in a workflow have their own author_name."""
# Create workflow with two executors
executor1 = SimpleExecutor(id="first_executor", response_text="First", emit_streaming=False)
executor2 = SimpleExecutor(id="second_executor", response_text="Second", emit_streaming=False)
executor1 = SimpleExecutor(id="first_executor", response_text="First")
executor2 = SimpleExecutor(id="second_executor", response_text="Second")
workflow = WorkflowBuilder().set_start_executor(executor1).add_edge(executor1, executor2).build()
agent = WorkflowAgent(workflow=workflow, name="Multi-Executor Agent")
@@ -834,11 +892,11 @@ class TestWorkflowAgentMergeUpdates:
# The exact order depends on dict iteration order for response_ids,
# but within each response group, chronological order should be maintained
# and global dangling should be last
assert "Global-Dangling" in message_texts[-1] # Global dangling at end
assert "Global-Dangling" in message_texts[-1] # type: ignore # Global dangling at end
# Find positions of resp-a and resp-b messages
resp_a_positions = [i for i, text in enumerate(message_texts) if "RespA" in text]
resp_b_positions = [i for i, text in enumerate(message_texts) if "RespB" in text]
resp_a_positions = [i for i, text in enumerate(message_texts) if "RespA" in text] # type: ignore
resp_b_positions = [i for i, text in enumerate(message_texts) if "RespB" in text] # type: ignore
# Within resp-a group: Msg1 (earlier) should come before Msg2 (later)
resp_a_texts = [message_texts[i] for i in resp_a_positions]
@@ -1013,7 +1071,7 @@ class TestWorkflowAgentMergeUpdates:
assert len(result.messages) == 4
# Extract content types for verification
content_sequence = []
content_sequence: list[tuple[str, str]] = []
for msg in result.messages:
for content in msg.contents:
if content.type == "text":
@@ -1128,7 +1186,7 @@ class TestWorkflowAgentMergeUpdates:
assert len(result.messages) == 6
# Build a sequence of (content_type, call_id_if_applicable)
content_sequence = []
content_sequence: list[tuple[str, str | None]] = []
for msg in result.messages:
for content in msg.contents:
if content.type == "text":
@@ -1194,7 +1252,7 @@ class TestWorkflowAgentMergeUpdates:
assert len(result.messages) == 3
# Orphan function result should be at the end since it can't be matched
content_types = []
content_types: list[str] = []
for msg in result.messages:
for content in msg.contents:
if content.type == "text":
@@ -15,6 +15,7 @@ from agent_framework import (
Executor,
WorkflowBuilder,
WorkflowContext,
WorkflowValidationError,
handler,
)
@@ -57,7 +58,7 @@ class MockExecutor(Executor):
"""A mock executor for testing purposes."""
@handler
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage]) -> None:
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage, MockMessage]) -> None:
"""A mock handler that does nothing."""
pass
@@ -104,99 +105,23 @@ def test_workflow_builder_fluent_api():
assert len(workflow.executors) == 6
def test_add_agent_with_custom_parameters():
"""Test adding an agent with custom parameters."""
agent = DummyAgent(id="agent_custom", name="custom_agent")
builder = WorkflowBuilder()
# Add agent with custom parameters
with pytest.deprecated_call():
result = builder.add_agent(agent, output_response=True, id="my_custom_id")
# Verify that add_agent returns the builder for chaining
assert result is builder
# Build workflow and verify executor is present
workflow = builder.set_start_executor(agent).build()
assert "my_custom_id" in workflow.executors
# Verify the executor was created with correct parameters
executor = workflow.executors["my_custom_id"]
assert isinstance(executor, AgentExecutor)
assert executor.id == "my_custom_id"
assert getattr(executor, "_output_response", False) is True
def test_add_agent_reuses_same_wrapper():
"""Test that using the same agent instance multiple times reuses the same wrapper."""
agent = DummyAgent(id="agent_reuse", name="reuse_agent")
reuse_agent = DummyAgent(id="agent_reuse", name="reuse_agent")
agent_a = DummyAgent(id="agent_a", name="agent_a")
builder = WorkflowBuilder()
# Add agent with specific parameters
with pytest.deprecated_call():
builder.add_agent(agent, output_response=True, id="agent_exec")
# Use the same agent instance in add_edge - should reuse the same wrapper
builder.set_start_executor(agent)
builder.set_start_executor(reuse_agent)
builder.add_edge(reuse_agent, agent_a)
builder.add_edge(agent_a, reuse_agent)
workflow = builder.build()
# Verify only one executor exists for this agent
assert workflow.start_executor_id == "agent_exec"
assert "agent_exec" in workflow.executors
assert len([e for e in workflow.executors.values() if isinstance(e, AgentExecutor)]) == 1
# Verify the executor has the parameters from add_agent
start_executor = workflow.get_start_executor()
assert isinstance(start_executor, AgentExecutor)
assert getattr(start_executor, "_output_response", False) is True
def test_add_agent_then_use_in_edges():
"""Test that an agent added via add_agent can be used in edge definitions."""
agent1 = DummyAgent(id="agent1", name="first")
agent2 = DummyAgent(id="agent2", name="second")
builder = WorkflowBuilder()
# Add agents with specific settings
with pytest.deprecated_call():
builder.add_agent(agent1, output_response=False, id="exec1")
builder.add_agent(agent2, output_response=True, id="exec2")
# Use the same agent instances to create edges
workflow = builder.set_start_executor(agent1).add_edge(agent1, agent2).build()
# Verify the executors maintain their settings
assert workflow.start_executor_id == "exec1"
assert "exec1" in workflow.executors
assert "exec2" in workflow.executors
e1 = workflow.executors["exec1"]
e2 = workflow.executors["exec2"]
assert isinstance(e1, AgentExecutor)
assert isinstance(e2, AgentExecutor)
assert getattr(e1, "_output_response", True) is False
assert getattr(e2, "_output_response", False) is True
def test_add_agent_without_explicit_id_uses_agent_name():
"""Test that add_agent uses agent name as id when no explicit id is provided."""
agent = DummyAgent(id="agent_x", name="named_agent")
builder = WorkflowBuilder()
with pytest.deprecated_call():
result = builder.add_agent(agent)
# Verify that add_agent returns the builder for chaining
assert result is builder
workflow = builder.set_start_executor(agent).build()
assert "named_agent" in workflow.executors
# Verify the executor id matches the agent name
executor = workflow.executors["named_agent"]
assert executor.id == "named_agent"
assert workflow.start_executor_id == "reuse_agent"
assert "reuse_agent" in workflow.executors
assert len([e for e in workflow.executors.values() if isinstance(e, AgentExecutor)]) == 2
def test_add_agent_duplicate_id_raises_error():
@@ -205,13 +130,8 @@ def test_add_agent_duplicate_id_raises_error():
agent2 = DummyAgent(id="agent2", name="first") # Same name as agent1
builder = WorkflowBuilder()
# Add first agent
with pytest.deprecated_call():
builder.add_agent(agent1)
# Adding second agent with same name should raise ValueError
with pytest.deprecated_call(), pytest.raises(ValueError, match="Duplicate executor ID"):
builder.add_agent(agent2)
with pytest.raises(ValueError, match="Duplicate executor ID"):
builder.set_start_executor(agent1).add_edge(agent1, agent2).build()
# Tests for new executor registration patterns
@@ -303,7 +223,7 @@ def test_register_duplicate_id_raises_error():
builder.set_start_executor("MyExecutor1")
# Registering second executor with same ID should raise ValueError
with pytest.raises(ValueError, match="Executor with ID 'executor' has already been created."):
with pytest.raises(ValueError, match="Executor with ID 'executor' has already been registered."):
builder.build()
@@ -312,9 +232,7 @@ def test_register_agent_basic():
builder = WorkflowBuilder()
# Register an agent factory
result = builder.register_agent(
lambda: DummyAgent(id="agent_test", name="test_agent"), name="TestAgent", output_response=True
)
result = builder.register_agent(lambda: DummyAgent(id="agent_test", name="test_agent"), name="TestAgent")
# Verify that register_agent returns the builder for chaining
assert result is builder
@@ -323,7 +241,6 @@ def test_register_agent_basic():
workflow = builder.set_start_executor("TestAgent").build()
assert "test_agent" in workflow.executors
assert isinstance(workflow.executors["test_agent"], AgentExecutor)
assert workflow.executors["test_agent"]._output_response is True # type: ignore
def test_register_agent_with_thread():
@@ -336,7 +253,6 @@ def test_register_agent_with_thread():
lambda: DummyAgent(id="agent_with_thread", name="threaded_agent"),
name="ThreadedAgent",
agent_thread=custom_thread,
output_response=False,
)
# Build workflow and verify agent executor configuration
@@ -345,7 +261,6 @@ def test_register_agent_with_thread():
assert isinstance(executor, AgentExecutor)
assert executor.id == "threaded_agent"
assert executor._output_response is False # type: ignore
assert executor._agent_thread is custom_thread # type: ignore
@@ -549,3 +464,151 @@ def test_register_agent_creates_unique_instances():
# Verify that two different agent instances were created
assert len(instance_ids) == 2
assert instance_ids[0] != instance_ids[1]
# region with_output_from tests
def test_with_output_from_returns_builder():
"""Test that with_output_from returns the builder for method chaining."""
executor_a = MockExecutor(id="executor_a")
builder = WorkflowBuilder()
result = builder.with_output_from([executor_a])
assert result is builder
def test_with_output_from_with_executor_instances():
"""Test with_output_from with direct executor instances."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
.add_edge(executor_a, executor_b)
.with_output_from([executor_b])
.build()
)
# Verify that the workflow was built with the correct output executors
assert workflow._output_executors == ["executor_b"] # type: ignore
def test_with_output_from_with_agent_instances():
"""Test with_output_from with agent instances."""
agent_a = DummyAgent(id="agent_a", name="writer")
agent_b = DummyAgent(id="agent_b", name="reviewer")
workflow = (
WorkflowBuilder().set_start_executor(agent_a).add_edge(agent_a, agent_b).with_output_from([agent_b]).build()
)
# Verify that the workflow was built with the agent's name as output executor
assert workflow._output_executors == ["reviewer"] # type: ignore
def test_with_output_from_with_registered_names():
"""Test with_output_from with registered factory names (strings)."""
workflow = (
WorkflowBuilder()
.register_executor(lambda: MockExecutor(id="ExecutorA"), name="ExecutorAFactory")
.register_executor(lambda: MockExecutor(id="ExecutorB"), name="ExecutorBFactory")
.set_start_executor("ExecutorAFactory")
.add_edge("ExecutorAFactory", "ExecutorBFactory")
.with_output_from(["ExecutorBFactory"])
.build()
)
# Verify that the workflow was built with the correct output executors
assert workflow._output_executors == ["ExecutorB"] # type: ignore
def test_with_output_from_with_multiple_executors():
"""Test with_output_from with multiple executors."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
executor_c = MockExecutor(id="executor_c")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_c)
.with_output_from([executor_a, executor_c])
.build()
)
# Verify that the workflow was built with both output executors
assert set(workflow._output_executors) == {"executor_a", "executor_c"} # type: ignore
def test_with_output_from_can_be_called_multiple_times():
"""Test that calling with_output_from multiple times overwrites the previous setting."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
.add_edge(executor_a, executor_b)
.with_output_from([executor_a])
.with_output_from([executor_b]) # This should overwrite the previous setting
.build()
)
# Verify that only the last setting is applied
assert workflow._output_executors == ["executor_b"] # type: ignore
def test_with_output_from_with_registered_agents():
"""Test with_output_from with registered agent factory names."""
workflow = (
WorkflowBuilder()
.register_agent(lambda: DummyAgent(id="agent1", name="writer"), name="WriterAgent")
.register_agent(lambda: DummyAgent(id="agent2", name="reviewer"), name="ReviewerAgent")
.set_start_executor("WriterAgent")
.add_edge("WriterAgent", "ReviewerAgent")
.with_output_from(["ReviewerAgent"])
.build()
)
# Verify that the workflow was built with the agent's resolved name
assert workflow._output_executors == ["reviewer"] # type: ignore
def test_with_output_from_in_fluent_chain():
"""Test that with_output_from works correctly in a fluent builder chain."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
executor_c = MockExecutor(id="executor_c")
# Build workflow with with_output_from in the middle of the chain
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
.with_output_from([executor_c]) # Set early in the chain
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_c)
.build()
)
# Verify that the setting persists through the chain
assert workflow._output_executors == ["executor_c"] # type: ignore
def test_with_output_from_with_invalid_executor_raises_validation_error():
"""Test that with_output_from with an invalid executor raises an error."""
executor_a = MockExecutor(id="executor_a")
builder = WorkflowBuilder().set_start_executor(executor_a)
# Attempting to set output from an executor not in the workflow should raise an error
with pytest.raises(
WorkflowValidationError, match="Output executor 'executor_b' is not present in the workflow graph"
):
builder.with_output_from([MockExecutor(id="executor_b")]).build()
# endregion
@@ -12,7 +12,7 @@ from datetime import datetime
from typing import Any, Union
from uuid import uuid4
from agent_framework import ChatMessage, Content
from agent_framework import ChatMessage, Content, WorkflowOutputEvent
from openai.types.responses import (
Response,
ResponseContentPartAddedEvent,
@@ -179,11 +179,10 @@ class MessageMapper:
# Import Agent Framework types for proper isinstance checks
try:
from agent_framework import AgentResponse, AgentResponseUpdate, WorkflowEvent
from agent_framework._workflows._events import AgentRunUpdateEvent
# Handle AgentRunUpdateEvent - workflow event wrapping AgentResponseUpdate
# This must be checked BEFORE generic WorkflowEvent check
if isinstance(raw_event, AgentRunUpdateEvent):
if isinstance(raw_event, WorkflowOutputEvent):
# Extract the AgentResponseUpdate from the event's data attribute
if raw_event.data and isinstance(raw_event.data, AgentResponseUpdate):
# Preserve executor_id in context for proper output routing