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