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

* 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 07:47:20 +00:00
committed by GitHub
parent 09f59b21ad
commit 0f3f4dbcaf
127 changed files with 1646 additions and 1703 deletions
@@ -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