mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
[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:
committed by
GitHub
Unverified
parent
907654a489
commit
a971d24f1e
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user