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

[BREAKING] Python: Refactor workflow events to unified discriminated union pattern (#3690)

Browse Source 
* Refactor events

* Merge main

* Fixes

* Cleanup

* Update samples and tests

* Remove unused imports

* PR feedback

* Merge main. Add properties for events to help typing

* Formatting

* Cleanup

* use builtins.type to avoid shadowing by WorkflowEvent.type attribute

* Final improvements
This commit is contained in:
Evan Mattson
2026-02-06 16:47:20 +09:00
committed by GitHub
Unverified
parent 09f59b21ad
commit 0f3f4dbcaf
127 changed files with 1646 additions and 1703 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/orchestrations/agent_framework_orchestrations/_base_group_chat_orchestrator.py
+28 -48
View File
@@ -61,48 +61,22 @@ GroupChatWorkflowContextOutT: TypeAlias = AgentExecutorRequest | GroupChatReques
# region Group chat events
class GroupChatEvent(WorkflowEvent):
"""Base class for group chat workflow events."""
def __init__(self, round_index: int, data: Any | None = None) -> None:
"""Initialize group chat event.
Args:
round_index: Current round index
data: Optional event-specific data
"""
super().__init__(data)
self.round_index = round_index
class GroupChatResponseReceivedEvent(GroupChatEvent):
"""Event emitted when a participant response is received."""
@dataclass
class GroupChatRequestSentEvent:
"""Data payload for group_chat request sent events."""
def __init__(self, round_index: int, participant_name: str, data: Any | None = None) -> None:
"""Initialize response received event.
Args:
round_index: Current round index
participant_name: Name of the participant who sent the response
data: Optional event-specific data
"""
super().__init__(round_index, data)
self.participant_name = participant_name
round_index: int
participant_name: str
class GroupChatRequestSentEvent(GroupChatEvent):
"""Event emitted when a request is sent to a participant."""
@dataclass
class GroupChatResponseReceivedEvent:
"""Data payload for group_chat response received events."""
def __init__(self, round_index: int, participant_name: str, data: Any | None = None) -> None:
"""Initialize request sent event.
Args:
round_index: Current round index
participant_name: Name of the participant to whom the request was sent
data: Optional event-specific data
"""
super().__init__(round_index, data)
self.participant_name = participant_name
round_index: int
participant_name: str
# endregion
@@ -273,10 +247,12 @@ class BaseGroupChatOrchestrator(Executor, ABC):
ctx: Workflow context
"""
await ctx.add_event(
GroupChatResponseReceivedEvent(
round_index=self._round_index,
participant_name=ctx.source_executor_ids[0] if ctx.source_executor_ids else "unknown",
data=response,
WorkflowEvent(
"group_chat",
data=GroupChatResponseReceivedEvent(
round_index=self._round_index,
participant_name=ctx.source_executor_ids[0] if ctx.source_executor_ids else "unknown",
),
)
)
await self._handle_response(response, ctx)
@@ -469,10 +445,12 @@ class BaseGroupChatOrchestrator(Executor, ABC):
request = AgentExecutorRequest(messages=messages, should_respond=True)
await ctx.send_message(request, target_id=target)
await ctx.add_event(
GroupChatRequestSentEvent(
round_index=self._round_index,
participant_name=target,
data=request,
WorkflowEvent(
"group_chat",
data=GroupChatRequestSentEvent(
round_index=self._round_index,
participant_name=target,
),
)
)
else:
@@ -480,10 +458,12 @@ class BaseGroupChatOrchestrator(Executor, ABC):
request = GroupChatRequestMessage(additional_instruction=additional_instruction, metadata=metadata) # type: ignore[assignment]
await ctx.send_message(request, target_id=target)
await ctx.add_event(
GroupChatRequestSentEvent(
round_index=self._round_index,
participant_name=target,
data=request,
WorkflowEvent(
"group_chat",
data=GroupChatRequestSentEvent(
round_index=self._round_index,
participant_name=target,
),
)
)
python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py
+2 -2
View File
@@ -463,9 +463,9 @@ class ConcurrentBuilder:
) -> "ConcurrentBuilder":
"""Enable request info after agent participant responses.
This enables human-in-the-loop (HIL) scenarios for the sequential orchestration.
This enables human-in-the-loop (HIL) scenarios for the concurrent orchestration.
When enabled, the workflow pauses after each agent participant runs, emitting
a RequestInfoEvent that allows the caller to review the conversation and optionally
a request_info event (type='request_info') that allows the caller to review the conversation and optionally
inject guidance for the agent participant to iterate. The caller provides input via
the standard response_handler/request_info pattern.
python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py
+1 -1
View File
@@ -866,7 +866,7 @@ class GroupChatBuilder:
This enables human-in-the-loop (HIL) scenarios for the group chat orchestration.
When enabled, the workflow pauses after each agent participant runs, emitting
a RequestInfoEvent that allows the caller to review the conversation and optionally
a request_info event (type='request_info') that allows the caller to review the conversation and optionally
inject guidance for the agent participant to iterate. The caller provides input via
the standard response_handler/request_info pattern.
python/packages/orchestrations/agent_framework_orchestrations/_handoff.py
+9 -13
View File
@@ -64,20 +64,14 @@ logger = logging.getLogger(__name__)
# region Handoff events
class HandoffSentEvent(WorkflowEvent):
"""Base class for handoff workflow events."""
def __init__(self, source: str, target: str, data: Any | None = None) -> None:
"""Initialize handoff sent event.
Args:
source: Identifier of the source agent initiating the handoff
target: Identifier of the target agent receiving the handoff
data: Optional event-specific data
"""
super().__init__(data)
self.source = source
self.target = target
@dataclass
class HandoffSentEvent:
"""Data payload for handoff_sent events."""
source: str
target: str
# endregion
@@ -421,7 +415,9 @@ class HandoffAgentExecutor(AgentExecutor):
await cast(WorkflowContext[AgentExecutorRequest], ctx).send_message(
AgentExecutorRequest(messages=[], should_respond=True), target_id=handoff_target
)
await ctx.add_event(HandoffSentEvent(source=self.id, target=handoff_target))
await ctx.add_event(
WorkflowEvent("handoff_sent", data=HandoffSentEvent(source=self.id, target=handoff_target))
)
self._autonomous_mode_turns = 0 # Reset autonomous mode turn counter on handoff
return
python/packages/orchestrations/agent_framework_orchestrations/_magentic.py
+30 -27
View File
@@ -19,7 +19,7 @@ from agent_framework import (
)
from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
from agent_framework._workflows._checkpoint import CheckpointStorage
from agent_framework._workflows._events import ExecutorEvent
from agent_framework._workflows._events import WorkflowEvent
from agent_framework._workflows._executor import Executor, handler
from agent_framework._workflows._model_utils import DictConvertible, encode_value
from agent_framework._workflows._request_info_mixin import response_handler
@@ -771,20 +771,11 @@ class MagenticOrchestratorEventType(str, Enum):
@dataclass
class MagenticOrchestratorEvent(ExecutorEvent):
"""Base class for Magentic orchestrator events."""
class MagenticOrchestratorEvent:
"""Data payload for magentic_orchestrator events."""
def __init__(
self,
executor_id: str,
event_type: MagenticOrchestratorEventType,
data: ChatMessage | MagenticProgressLedger,
) -> None:
super().__init__(executor_id, data)
self.event_type = event_type
def __repr__(self) -> str:
return f"{self.__class__.__name__}(executor_id={self.executor_id}, event_type={self.event_type})"
event_type: MagenticOrchestratorEventType
content: ChatMessage | MagenticProgressLedger
# region Request info related types
@@ -928,10 +919,13 @@ class MagenticOrchestrator(BaseGroupChatOrchestrator):
# Initial planning using the manager with real model calls
self._task_ledger = await self._manager.plan(self._magentic_context.clone(deep=True))
await ctx.add_event(
MagenticOrchestratorEvent(
WorkflowEvent(
"magentic_orchestrator",
executor_id=self.id,
event_type=MagenticOrchestratorEventType.PLAN_CREATED,
data=self._task_ledger,
data=MagenticOrchestratorEvent(
event_type=MagenticOrchestratorEventType.PLAN_CREATED,
content=self._task_ledger,
),
)
)
@@ -1006,10 +1000,13 @@ class MagenticOrchestrator(BaseGroupChatOrchestrator):
self._magentic_context.chat_history.extend(response.review)
self._task_ledger = await self._manager.replan(self._magentic_context.clone(deep=True))
await ctx.add_event(
MagenticOrchestratorEvent(
WorkflowEvent(
"magentic_orchestrator",
executor_id=self.id,
event_type=MagenticOrchestratorEventType.REPLANNED,
data=self._task_ledger,
data=MagenticOrchestratorEvent(
event_type=MagenticOrchestratorEventType.REPLANNED,
content=self._task_ledger,
),
)
)
# Continue the review process by sending the new plan for review again until approved
@@ -1072,10 +1069,13 @@ class MagenticOrchestrator(BaseGroupChatOrchestrator):
return
await ctx.add_event(
MagenticOrchestratorEvent(
WorkflowEvent(
"magentic_orchestrator",
executor_id=self.id,
event_type=MagenticOrchestratorEventType.PROGRESS_LEDGER_UPDATED,
data=self._progress_ledger,
data=MagenticOrchestratorEvent(
event_type=MagenticOrchestratorEventType.PROGRESS_LEDGER_UPDATED,
content=self._progress_ledger,
),
)
)
@@ -1149,10 +1149,13 @@ class MagenticOrchestrator(BaseGroupChatOrchestrator):
# Replan
self._task_ledger = await self._manager.replan(self._magentic_context.clone(deep=True))
await ctx.add_event(
MagenticOrchestratorEvent(
WorkflowEvent(
"magentic_orchestrator",
executor_id=self.id,
event_type=MagenticOrchestratorEventType.REPLANNED,
data=self._task_ledger,
data=MagenticOrchestratorEvent(
event_type=MagenticOrchestratorEventType.REPLANNED,
content=self._task_ledger,
),
)
)
# If a human must sign off, ask now and return. The response handler will resume.
@@ -1515,7 +1518,7 @@ class MagenticBuilder:
# During execution, handle plan review
async for event in workflow.run("task", stream=True):
if isinstance(event, RequestInfoEvent):
if event.type == "request_info":
request = event.data
if isinstance(request, MagenticHumanInterventionRequest):
if request.kind == MagenticHumanInterventionKind.PLAN_REVIEW:
python/packages/orchestrations/agent_framework_orchestrations/_orchestration_state.py
+3 -1
View File
@@ -6,6 +6,8 @@ Provides OrchestrationState dataclass for standardized checkpoint serialization
across GroupChat, Handoff, and Magentic patterns.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
@@ -69,7 +71,7 @@ class OrchestrationState:
return result
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "OrchestrationState":
def from_dict(cls, data: dict[str, Any]) -> OrchestrationState:
"""Deserialize from checkpointed dict.
Args:
python/packages/orchestrations/agent_framework_orchestrations/_sequential.py
+1 -1
View File
@@ -219,7 +219,7 @@ class SequentialBuilder:
This enables human-in-the-loop (HIL) scenarios for the sequential orchestration.
When enabled, the workflow pauses after each agent participant runs, emitting
a RequestInfoEvent that allows the caller to review the conversation and optionally
a request_info event (type='request_info') that allows the caller to review the conversation and optionally
inject guidance for the agent participant to iterate. The caller provides input via
the standard response_handler/request_info pattern.
python/packages/orchestrations/tests/test_concurrent.py
+24 -26
View File
@@ -10,9 +10,7 @@ from agent_framework import (
ChatMessage,
Executor,
WorkflowContext,
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
@@ -111,9 +109,9 @@ async def test_concurrent_default_aggregator_emits_single_user_and_assistants()
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run("prompt: hello world", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = cast(list[ChatMessage], ev.data)
if completed and output is not None:
break
@@ -149,9 +147,9 @@ async def test_concurrent_custom_aggregator_callback_is_used() -> None:
completed = False
output: str | None = None
async for ev in wf.run("prompt: custom", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = cast(str, ev.data)
if completed and output is not None:
break
@@ -180,9 +178,9 @@ async def test_concurrent_custom_aggregator_sync_callback_is_used() -> None:
completed = False
output: str | None = None
async for ev in wf.run("prompt: custom sync", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = cast(str, ev.data)
if completed and output is not None:
break
@@ -228,9 +226,9 @@ async def test_concurrent_with_aggregator_executor_instance() -> None:
completed = False
output: str | None = None
async for ev in wf.run("prompt: instance test", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = cast(str, ev.data)
if completed and output is not None:
break
@@ -266,9 +264,9 @@ async def test_concurrent_with_aggregator_executor_factory() -> None:
completed = False
output: str | None = None
async for ev in wf.run("prompt: factory test", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = cast(str, ev.data)
if completed and output is not None:
break
@@ -302,9 +300,9 @@ async def test_concurrent_with_aggregator_executor_factory_with_default_id() ->
completed = False
output: str | None = None
async for ev in wf.run("prompt: factory test", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = cast(str, ev.data)
if completed and output is not None:
break
@@ -352,9 +350,9 @@ async def test_concurrent_checkpoint_resume_round_trip() -> None:
baseline_output: list[ChatMessage] | None = None
async for ev in wf.run("checkpoint concurrent", stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
assert baseline_output is not None
@@ -376,9 +374,9 @@ async def test_concurrent_checkpoint_resume_round_trip() -> None:
resumed_output: list[ChatMessage] | None = None
async for ev in wf_resume.run(checkpoint_id=resume_checkpoint.checkpoint_id, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
resumed_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state in (
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -398,9 +396,9 @@ async def test_concurrent_checkpoint_runtime_only() -> None:
baseline_output: list[ChatMessage] | None = None
async for ev in wf.run("runtime checkpoint test", checkpoint_storage=storage, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
assert baseline_output is not None
@@ -421,9 +419,9 @@ async def test_concurrent_checkpoint_runtime_only() -> None:
async for ev in wf_resume.run(
checkpoint_id=resume_checkpoint.checkpoint_id, checkpoint_storage=storage, stream=True
):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
resumed_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state in (
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -448,9 +446,9 @@ async def test_concurrent_checkpoint_runtime_overrides_buildtime() -> None:
baseline_output: list[ChatMessage] | None = None
async for ev in wf.run("override test", checkpoint_storage=runtime_storage, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
assert baseline_output is not None
@@ -530,9 +528,9 @@ async def test_concurrent_with_register_participants() -> None:
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run("test prompt", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = cast(list[ChatMessage], ev.data)
if completed and output is not None:
break
python/packages/orchestrations/tests/test_group_chat.py
+27 -29
View File
@@ -15,10 +15,8 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
RequestInfoEvent,
WorkflowOutputEvent,
WorkflowEvent,
WorkflowRunState,
WorkflowStatusEvent,
)
from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
from agent_framework.orchestrations import (
@@ -190,7 +188,7 @@ async def test_group_chat_builder_basic_flow() -> None:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run("coordinate task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -362,7 +360,7 @@ class TestGroupChatWorkflow:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run("test task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -397,7 +395,7 @@ class TestGroupChatWorkflow:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run("test task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -425,7 +423,7 @@ class TestGroupChatWorkflow:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run("test task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -473,7 +471,7 @@ class TestCheckpointing:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run("test task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -526,7 +524,7 @@ class TestConversationHandling:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run("test string", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -555,7 +553,7 @@ class TestConversationHandling:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run(task_message, stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -587,7 +585,7 @@ class TestConversationHandling:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run(conversation, stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -619,7 +617,7 @@ class TestRoundLimitEnforcement:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run("test", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -654,7 +652,7 @@ class TestRoundLimitEnforcement:
outputs: list[list[ChatMessage]] = []
async for event in workflow.run("test", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
@@ -686,9 +684,9 @@ async def test_group_chat_checkpoint_runtime_only() -> None:
baseline_output: list[ChatMessage] | None = None
async for ev in wf.run("runtime checkpoint test", checkpoint_storage=storage, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = cast(list[ChatMessage], ev.data) if isinstance(ev.data, list) else None # type: ignore
if isinstance(ev, WorkflowStatusEvent) and ev.state in (
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -724,9 +722,9 @@ async def test_group_chat_checkpoint_runtime_overrides_buildtime() -> None:
)
baseline_output: list[ChatMessage] | None = None
async for ev in wf.run("override test", checkpoint_storage=runtime_storage, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = cast(list[ChatMessage], ev.data) if isinstance(ev.data, list) else None # type: ignore
if isinstance(ev, WorkflowStatusEvent) and ev.state in (
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -770,9 +768,9 @@ async def test_group_chat_with_request_info_filtering():
)
# Run until we get a request info event (should be before beta, not alpha)
request_events: list[RequestInfoEvent] = []
request_events: list[WorkflowEvent] = []
async for event in workflow.run("test task", stream=True):
if isinstance(event, RequestInfoEvent) and isinstance(event.data, AgentExecutorResponse):
if event.type == "request_info" and isinstance(event.data, AgentExecutorResponse):
request_events.append(event)
# Don't break - let stream complete naturally when paused
@@ -785,11 +783,11 @@ async def test_group_chat_with_request_info_filtering():
assert request_event.source_executor_id == "beta"
# Continue the workflow with a response
outputs: list[WorkflowOutputEvent] = []
outputs: list[WorkflowEvent] = []
async for event in workflow.send_responses_streaming({
request_event.request_id: AgentRequestInfoResponse.approve()
}):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
outputs.append(event)
# Workflow should complete
@@ -822,9 +820,9 @@ async def test_group_chat_with_request_info_no_filter_pauses_all():
)
# Run until we get a request info event
request_events: list[RequestInfoEvent] = []
request_events: list[WorkflowEvent] = []
async for event in workflow.run("test task", stream=True):
if isinstance(event, RequestInfoEvent) and isinstance(event.data, AgentExecutorResponse):
if event.type == "request_info" and isinstance(event.data, AgentExecutorResponse):
request_events.append(event)
break
@@ -926,9 +924,9 @@ async def test_group_chat_with_participant_factories():
# Factories should be called during build
assert call_count == 2
outputs: list[WorkflowOutputEvent] = []
outputs: list[WorkflowEvent] = []
async for event in workflow.run("coordinate task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
outputs.append(event)
assert len(outputs) == 1
@@ -991,9 +989,9 @@ async def test_group_chat_participant_factories_with_checkpointing():
.build()
)
outputs: list[WorkflowOutputEvent] = []
outputs: list[WorkflowEvent] = []
async for event in workflow.run("checkpoint test", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
outputs.append(event)
assert outputs, "Should have workflow output"
@@ -1119,9 +1117,9 @@ async def test_group_chat_with_orchestrator_factory_returning_chat_agent():
# Factory should be called during build
assert factory_call_count == 1
outputs: list[WorkflowOutputEvent] = []
outputs: list[WorkflowEvent] = []
async for event in workflow.run("coordinate task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
outputs.append(event)
assert len(outputs) == 1
python/packages/orchestrations/tests/test_handoff.py
+13 -15
View File
@@ -11,10 +11,8 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
RequestInfoEvent,
ResponseStream,
WorkflowEvent,
WorkflowOutputEvent,
resolve_agent_id,
)
from agent_framework._clients import BaseChatClient
@@ -150,7 +148,7 @@ async def test_handoff():
# escalation won't trigger a handoff, so the response from it will become
# a request for user input because autonomous mode is not enabled by default.
events = await _drain(workflow.run("Need technical support", stream=True))
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
requests = [ev for ev in events if ev.type == "request_info"]
assert requests
assert len(requests) == 1
@@ -184,10 +182,10 @@ async def test_autonomous_mode_yields_output_without_user_request():
)
events = await _drain(workflow.run("Package arrived broken", stream=True))
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
requests = [ev for ev in events if ev.type == "request_info"]
assert not requests, "Autonomous mode should not request additional user input"
outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)]
outputs = [ev for ev in events if ev.type == "output"]
assert outputs, "Autonomous mode should yield a workflow output"
final_conversation = outputs[-1].data
@@ -210,7 +208,7 @@ async def test_autonomous_mode_resumes_user_input_on_turn_limit():
)
events = await _drain(workflow.run("Start", stream=True))
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
requests = [ev for ev in events if ev.type == "request_info"]
assert requests and len(requests) == 1, "Turn limit should force a user input request"
assert requests[0].source_executor_id == worker.name
@@ -253,7 +251,7 @@ async def test_handoff_async_termination_condition() -> None:
)
events = await _drain(workflow.run("First user message", stream=True))
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
requests = [ev for ev in events if ev.type == "request_info"]
assert requests
events = await _drain(
@@ -261,7 +259,7 @@ async def test_handoff_async_termination_condition() -> None:
requests[-1].request_id: [ChatMessage(role="user", text="Second user message")]
})
)
outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)]
outputs = [ev for ev in events if ev.type == "output"]
assert len(outputs) == 1
final_conversation = outputs[0].data
@@ -505,14 +503,14 @@ async def test_handoff_with_participant_factories():
assert call_count == 2
events = await _drain(workflow.run("Need help", stream=True))
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
requests = [ev for ev in events if ev.type == "request_info"]
assert requests
# Follow-up message
events = await _drain(
workflow.send_responses_streaming({requests[-1].request_id: [ChatMessage(role="user", text="More details")]})
)
outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)]
outputs = [ev for ev in events if ev.type == "output"]
assert outputs
@@ -576,7 +574,7 @@ async def test_handoff_with_participant_factories_and_add_handoff():
# Start conversation - triage hands off to specialist_a
events = await _drain(workflow.run("Initial request", stream=True))
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
requests = [ev for ev in events if ev.type == "request_info"]
assert requests
# Verify specialist_a executor exists and was called
@@ -586,7 +584,7 @@ async def test_handoff_with_participant_factories_and_add_handoff():
events = await _drain(
workflow.send_responses_streaming({requests[-1].request_id: [ChatMessage(role="user", text="Need escalation")]})
)
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
requests = [ev for ev in events if ev.type == "request_info"]
assert requests
# Verify specialist_b executor exists
@@ -615,13 +613,13 @@ async def test_handoff_participant_factories_with_checkpointing():
# Run workflow and capture output
events = await _drain(workflow.run("checkpoint test", stream=True))
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
requests = [ev for ev in events if ev.type == "request_info"]
assert requests
events = await _drain(
workflow.send_responses_streaming({requests[-1].request_id: [ChatMessage(role="user", text="follow up")]})
)
outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)]
outputs = [ev for ev in events if ev.type == "output"]
assert outputs, "Should have workflow output after termination condition is met"
# List checkpoints - just verify they were created
@@ -693,7 +691,7 @@ async def test_handoff_participant_factories_autonomous_mode():
)
events = await _drain(workflow.run("Issue", stream=True))
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
requests = [ev for ev in events if ev.type == "request_info"]
assert requests and len(requests) == 1
assert requests[0].source_executor_id == "specialist"
python/packages/orchestrations/tests/test_magentic.py
+44 -44
View File
@@ -15,15 +15,12 @@ from agent_framework import (
ChatMessage,
Content,
Executor,
RequestInfoEvent,
Workflow,
WorkflowCheckpoint,
WorkflowCheckpointException,
WorkflowContext,
WorkflowEvent,
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
@@ -33,7 +30,6 @@ from agent_framework.orchestrations import (
MagenticContext,
MagenticManagerBase,
MagenticOrchestrator,
MagenticOrchestratorEvent,
MagenticPlanReviewRequest,
MagenticProgressLedger,
MagenticProgressLedgerItem,
@@ -197,11 +193,11 @@ async def test_magentic_builder_returns_workflow_and_runs() -> None:
outputs: list[ChatMessage] = []
orchestrator_event_count = 0
async for event in workflow.run("compose summary", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
msg = event.data
if isinstance(msg, list):
outputs.extend(cast(list[ChatMessage], msg))
elif isinstance(event, MagenticOrchestratorEvent):
elif event.type == "magentic_orchestrator":
orchestrator_event_count += 1
assert outputs, "Expected a final output message"
@@ -246,9 +242,9 @@ async def test_magentic_workflow_plan_review_approval_to_completion():
manager = FakeManager()
wf = MagenticBuilder().participants([DummyExec("agentA")]).with_manager(manager=manager).with_plan_review().build()
req_event: RequestInfoEvent | None = None
req_event: WorkflowEvent | None = None
async for ev in wf.run("do work", stream=True):
if isinstance(ev, RequestInfoEvent) and ev.request_type is MagenticPlanReviewRequest:
if ev.type == "request_info" and ev.request_type is MagenticPlanReviewRequest:
req_event = ev
assert req_event is not None
assert isinstance(req_event.data, MagenticPlanReviewRequest)
@@ -256,9 +252,9 @@ async def test_magentic_workflow_plan_review_approval_to_completion():
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.send_responses_streaming(responses={req_event.request_id: req_event.data.approve()}):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = ev.data # type: ignore[assignment]
if completed and output is not None:
break
@@ -291,9 +287,9 @@ async def test_magentic_plan_review_with_revise():
)
# Wait for the initial plan review request
req_event: RequestInfoEvent | None = None
req_event: WorkflowEvent | None = None
async for ev in wf.run("do work", stream=True):
if isinstance(ev, RequestInfoEvent) and ev.request_type is MagenticPlanReviewRequest:
if ev.type == "request_info" and ev.request_type is MagenticPlanReviewRequest:
req_event = ev
assert req_event is not None
assert isinstance(req_event.data, MagenticPlanReviewRequest)
@@ -304,7 +300,7 @@ async def test_magentic_plan_review_with_revise():
async for ev in wf.send_responses_streaming(
responses={req_event.request_id: req_event.data.revise("Looks good; consider Z")}
):
if isinstance(ev, RequestInfoEvent) and ev.request_type is MagenticPlanReviewRequest:
if ev.type == "request_info" and ev.request_type is MagenticPlanReviewRequest:
saw_second_review = True
req_event = ev
@@ -312,7 +308,7 @@ async def test_magentic_plan_review_with_revise():
async for ev in wf.send_responses_streaming(
responses={req_event.request_id: req_event.data.approve()} # type: ignore[union-attr]
):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
break
@@ -339,12 +335,12 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result():
events.append(ev)
idle_status = next(
(e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE),
(e for e in events if e.type == "status" and e.state == WorkflowRunState.IDLE),
None,
)
assert idle_status is not None
# Check that we got workflow output via WorkflowOutputEvent
output_event = next((e for e in events if isinstance(e, WorkflowOutputEvent)), None)
# Check that we got workflow output via WorkflowEvent with type "output"
output_event = next((e for e in events if e.type == "output"), None)
assert output_event is not None
data = output_event.data
assert isinstance(data, list)
@@ -367,9 +363,9 @@ async def test_magentic_checkpoint_resume_round_trip():
)
task_text = "checkpoint task"
req_event: RequestInfoEvent | None = None
req_event: WorkflowEvent | None = None
async for ev in wf.run(task_text, stream=True):
if isinstance(ev, RequestInfoEvent) and ev.request_type is MagenticPlanReviewRequest:
if ev.type == "request_info" and ev.request_type is MagenticPlanReviewRequest:
req_event = ev
assert req_event is not None
assert isinstance(req_event.data, MagenticPlanReviewRequest)
@@ -389,20 +385,20 @@ async def test_magentic_checkpoint_resume_round_trip():
.build()
)
completed: WorkflowOutputEvent | None = None
completed: WorkflowEvent | None = None
req_event = None
async for event in wf_resume.run(
resume_checkpoint.checkpoint_id,
stream=True,
):
if isinstance(event, RequestInfoEvent) and event.request_type is MagenticPlanReviewRequest:
if event.type == "request_info" and event.request_type is MagenticPlanReviewRequest:
req_event = event
assert req_event is not None
assert isinstance(req_event.data, MagenticPlanReviewRequest)
responses = {req_event.request_id: req_event.data.approve()}
async for event in wf_resume.send_responses_streaming(responses=responses):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
completed = event
assert completed is not None
@@ -595,7 +591,8 @@ async def _collect_agent_responses_setup(participant: AgentProtocol) -> list[Cha
events: list[WorkflowEvent] = []
async for ev in wf.run("task", stream=True): # plan review disabled
events.append(ev)
if isinstance(ev, WorkflowOutputEvent) and isinstance(ev.data, AgentResponseUpdate):
# Capture streaming updates (type="output" with AgentResponseUpdate data)
if ev.type == "output" and isinstance(ev.data, AgentResponseUpdate):
captured.append(
ChatMessage(
role=ev.data.role or "assistant",
@@ -603,6 +600,9 @@ async def _collect_agent_responses_setup(participant: AgentProtocol) -> list[Cha
author_name=ev.data.author_name,
)
)
# Break on final AgentResponse output
elif ev.type == "output" and isinstance(ev.data, AgentResponse):
break
return captured
@@ -640,7 +640,7 @@ async def test_magentic_checkpoint_resume_inner_loop_superstep():
)
async for event in workflow.run("inner-loop task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
break
checkpoints = await _collect_checkpoints(storage)
@@ -654,9 +654,9 @@ async def test_magentic_checkpoint_resume_inner_loop_superstep():
.build()
)
completed: WorkflowOutputEvent | None = None
completed: WorkflowEvent | None = None
async for event in resumed.run(checkpoint_id=inner_loop_checkpoint.checkpoint_id, stream=True): # type: ignore[reportUnknownMemberType]
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
completed = event
assert completed is not None
@@ -678,7 +678,7 @@ async def test_magentic_checkpoint_resume_from_saved_state():
)
async for event in workflow.run("checkpoint resume task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
break
checkpoints = await _collect_checkpoints(storage)
@@ -694,9 +694,9 @@ async def test_magentic_checkpoint_resume_from_saved_state():
.build()
)
completed: WorkflowOutputEvent | None = None
completed: WorkflowEvent | None = None
async for event in resumed_workflow.run(checkpoint_id=resumed_state.checkpoint_id, stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
completed = event
assert completed is not None
@@ -716,9 +716,9 @@ async def test_magentic_checkpoint_resume_rejects_participant_renames():
.build()
)
req_event: RequestInfoEvent | None = None
req_event: WorkflowEvent | None = None
async for event in workflow.run("task", stream=True):
if isinstance(event, RequestInfoEvent) and event.request_type is MagenticPlanReviewRequest:
if event.type == "request_info" and event.request_type is MagenticPlanReviewRequest:
req_event = event
assert req_event is not None
@@ -778,11 +778,11 @@ async def test_magentic_stall_and_reset_reach_limits():
events.append(ev)
idle_status = next(
(e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE),
(e for e in events if e.type == "status" and e.state == WorkflowRunState.IDLE),
None,
)
assert idle_status is not None
output_event = next((e for e in events if isinstance(e, WorkflowOutputEvent)), None)
output_event = next((e for e in events if e.type == "output"), None)
assert output_event is not None
assert isinstance(output_event.data, list)
assert all(isinstance(msg, ChatMessage) for msg in output_event.data) # type: ignore
@@ -800,9 +800,9 @@ async def test_magentic_checkpoint_runtime_only() -> None:
baseline_output: ChatMessage | None = None
async for ev in wf.run("runtime checkpoint test", checkpoint_storage=storage, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state in (
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -838,9 +838,9 @@ async def test_magentic_checkpoint_runtime_overrides_buildtime() -> None:
baseline_output: ChatMessage | None = None
async for ev in wf.run("override test", checkpoint_storage=runtime_storage, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state in (
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -897,7 +897,7 @@ async def test_magentic_checkpoint_restore_no_duplicate_history():
]
async for event in wf.run(conversation, stream=True):
if isinstance(event, WorkflowStatusEvent) and event.state in (
if event.type == "status" and event.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -1005,9 +1005,9 @@ async def test_magentic_with_participant_factories():
# Factory should be called during build
assert call_count == 1
outputs: list[WorkflowOutputEvent] = []
outputs: list[WorkflowEvent] = []
async for event in workflow.run("test task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
outputs.append(event)
assert len(outputs) == 1
@@ -1052,9 +1052,9 @@ async def test_magentic_participant_factories_with_checkpointing():
.build()
)
outputs: list[WorkflowOutputEvent] = []
outputs: list[WorkflowEvent] = []
async for event in workflow.run("checkpoint test", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
outputs.append(event)
assert outputs, "Should have workflow output"
@@ -1109,9 +1109,9 @@ async def test_magentic_with_manager_factory():
# Factory should be called during build
assert factory_call_count == 1
outputs: list[WorkflowOutputEvent] = []
outputs: list[WorkflowEvent] = []
async for event in workflow.run("test task", stream=True):
if isinstance(event, WorkflowOutputEvent):
if event.type == "output":
outputs.append(event)
assert len(outputs) == 1
python/packages/orchestrations/tests/test_sequential.py
+24 -26
View File
@@ -15,9 +15,7 @@ from agent_framework import (
Executor,
TypeCompatibilityError,
WorkflowContext,
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
handler,
)
from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
@@ -106,9 +104,9 @@ async def test_sequential_agents_append_to_context() -> None:
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run("hello sequential", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = ev.data # type: ignore[assignment]
if completed and output is not None:
break
@@ -139,9 +137,9 @@ async def test_sequential_register_participants_with_agent_factories() -> None:
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run("hello factories", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = ev.data
if completed and output is not None:
break
@@ -165,9 +163,9 @@ async def test_sequential_with_custom_executor_summary() -> None:
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run("topic X", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = ev.data
if completed and output is not None:
break
@@ -196,9 +194,9 @@ async def test_sequential_register_participants_mixed_agents_and_executors() ->
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run("topic Y", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = ev.data
if completed and output is not None:
break
@@ -221,9 +219,9 @@ async def test_sequential_checkpoint_resume_round_trip() -> None:
baseline_output: list[ChatMessage] | None = None
async for ev in wf.run("checkpoint sequential", stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
assert baseline_output is not None
@@ -242,9 +240,9 @@ async def test_sequential_checkpoint_resume_round_trip() -> None:
resumed_output: list[ChatMessage] | None = None
async for ev in wf_resume.run(checkpoint_id=resume_checkpoint.checkpoint_id, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
resumed_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state in (
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -264,9 +262,9 @@ async def test_sequential_checkpoint_runtime_only() -> None:
baseline_output: list[ChatMessage] | None = None
async for ev in wf.run("runtime checkpoint test", checkpoint_storage=storage, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
assert baseline_output is not None
@@ -287,9 +285,9 @@ async def test_sequential_checkpoint_runtime_only() -> None:
async for ev in wf_resume.run(
checkpoint_id=resume_checkpoint.checkpoint_id, checkpoint_storage=storage, stream=True
):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
resumed_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state in (
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -315,9 +313,9 @@ async def test_sequential_checkpoint_runtime_overrides_buildtime() -> None:
baseline_output: list[ChatMessage] | None = None
async for ev in wf.run("override test", checkpoint_storage=runtime_storage, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
assert baseline_output is not None
@@ -343,9 +341,9 @@ async def test_sequential_register_participants_with_checkpointing() -> None:
baseline_output: list[ChatMessage] | None = None
async for ev in wf.run("checkpoint with factories", stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
baseline_output = ev.data
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
assert baseline_output is not None
@@ -365,9 +363,9 @@ async def test_sequential_register_participants_with_checkpointing() -> None:
resumed_output: list[ChatMessage] | None = None
async for ev in wf_resume.run(checkpoint_id=resume_checkpoint.checkpoint_id, stream=True):
if isinstance(ev, WorkflowOutputEvent):
if ev.type == "output":
resumed_output = ev.data
if isinstance(ev, WorkflowStatusEvent) and ev.state in (
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
):
@@ -401,9 +399,9 @@ async def test_sequential_register_participants_factories_called_on_build() -> N
completed = False
output: list[ChatMessage] | None = None
async for ev in wf.run("test factories timing", stream=True):
if isinstance(ev, WorkflowStatusEvent) and ev.state == WorkflowRunState.IDLE:
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif isinstance(ev, WorkflowOutputEvent):
elif ev.type == "output":
output = ev.data # type: ignore[assignment]
if completed and output is not None:
break