agent/agent-framework
1
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

Fix WorkflowAgent to include thread convo history. Enable checkpointing. (#2774)

Browse Source
This commit is contained in:
Evan Mattson
2025-12-12 13:04:31 +09:00
committed by GitHub
Unverified
parent d7434d59ce
commit 0fc7933a92
5 changed files with 508 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/agent_framework/_workflows/_agent.py
+50 -6
View File
@@ -23,6 +23,7 @@ from agent_framework import (
)
from ..exceptions import AgentExecutionException
from ._checkpoint import CheckpointStorage
from ._events import (
AgentRunUpdateEvent,
RequestInfoEvent,
@@ -117,6 +118,8 @@ class WorkflowAgent(BaseAgent):
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AgentRunResponse:
"""Get a response from the workflow agent (non-streaming).
@@ -124,10 +127,16 @@ class WorkflowAgent(BaseAgent):
This method collects all streaming updates and merges them into a single response.
Args:
messages: The message(s) to send to the workflow.
messages: The message(s) to send to the workflow. Required for new runs,
should be None when resuming from checkpoint.
Keyword Args:
thread: The conversation thread. If None, a new thread will be created.
checkpoint_id: ID of checkpoint to restore from. If provided, the workflow
resumes from this checkpoint instead of starting fresh.
checkpoint_storage: Runtime checkpoint storage. When provided with checkpoint_id,
used to load and restore the checkpoint. When provided without checkpoint_id,
enables checkpointing for this run.
**kwargs: Additional keyword arguments.
Returns:
@@ -139,7 +148,9 @@ class WorkflowAgent(BaseAgent):
thread = thread or self.get_new_thread()
response_id = str(uuid.uuid4())
async for update in self._run_stream_impl(input_messages, response_id):
async for update in self._run_stream_impl(
input_messages, response_id, thread, checkpoint_id, checkpoint_storage
):
response_updates.append(update)
# Convert updates to final response.
@@ -155,15 +166,23 @@ class WorkflowAgent(BaseAgent):
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Stream response updates from the workflow agent.
Args:
messages: The message(s) to send to the workflow.
messages: The message(s) to send to the workflow. Required for new runs,
should be None when resuming from checkpoint.
Keyword Args:
thread: The conversation thread. If None, a new thread will be created.
checkpoint_id: ID of checkpoint to restore from. If provided, the workflow
resumes from this checkpoint instead of starting fresh.
checkpoint_storage: Runtime checkpoint storage. When provided with checkpoint_id,
used to load and restore the checkpoint. When provided without checkpoint_id,
enables checkpointing for this run.
**kwargs: Additional keyword arguments.
Yields:
@@ -174,7 +193,9 @@ class WorkflowAgent(BaseAgent):
response_updates: list[AgentRunResponseUpdate] = []
response_id = str(uuid.uuid4())
async for update in self._run_stream_impl(input_messages, response_id):
async for update in self._run_stream_impl(
input_messages, response_id, thread, checkpoint_id, checkpoint_storage
):
response_updates.append(update)
yield update
@@ -188,12 +209,18 @@ class WorkflowAgent(BaseAgent):
self,
input_messages: list[ChatMessage],
response_id: str,
thread: AgentThread,
checkpoint_id: str | None = None,
checkpoint_storage: CheckpointStorage | None = None,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Internal implementation of 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.
Yields:
AgentRunResponseUpdate objects representing the workflow execution progress.
@@ -217,10 +244,27 @@ class WorkflowAgent(BaseAgent):
# 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,
)
else:
# Execute workflow with streaming (initial run or no function responses)
# Pass the new input messages directly to the workflow
event_stream = self.workflow.run_stream(input_messages)
# 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,
)
# Process events from the stream
async for event in event_stream:
python/packages/core/tests/workflow/test_workflow_agent.py
+126
View File
@@ -9,7 +9,9 @@ from agent_framework import (
AgentRunResponse,
AgentRunResponseUpdate,
AgentRunUpdateEvent,
AgentThread,
ChatMessage,
ChatMessageStore,
Executor,
FunctionApprovalRequestContent,
FunctionApprovalResponseContent,
@@ -75,6 +77,31 @@ class RequestingExecutor(Executor):
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):
super().__init__(id=id)
self.received_messages: list[ChatMessage] = []
@handler
async def handle_message(self, messages: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None:
# Capture all received messages
self.received_messages = list(messages)
# Count messages by role for the response
message_count = len(messages)
response_text = f"Received {message_count} messages"
response_message = ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text=response_text)])
streaming_update = AgentRunResponseUpdate(
contents=[TextContent(text=response_text)], role=Role.ASSISTANT, message_id=str(uuid.uuid4())
)
await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=streaming_update))
await ctx.send_message([response_message])
class TestWorkflowAgent:
"""Test cases for WorkflowAgent end-to-end functionality."""
@@ -257,6 +284,105 @@ class TestWorkflowAgent:
with pytest.raises(ValueError, match="Workflow's start executor cannot handle list\\[ChatMessage\\]"):
workflow.as_agent()
async def test_thread_conversation_history_included_in_workflow_run(self) -> None:
"""Test that conversation history from thread is included when running WorkflowAgent.
This verifies that when a thread with existing messages is provided to agent.run(),
the workflow receives the complete conversation history (thread history + new messages).
"""
# Create an executor that captures all received messages
capturing_executor = ConversationHistoryCapturingExecutor(id="capturing")
workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Thread History Test Agent")
# Create a thread with existing conversation history
history_messages = [
ChatMessage(role=Role.USER, text="Previous user message"),
ChatMessage(role=Role.ASSISTANT, text="Previous assistant response"),
]
message_store = ChatMessageStore(messages=history_messages)
thread = AgentThread(message_store=message_store)
# Run the agent with the thread and a new message
new_message = "New user question"
await agent.run(new_message, thread=thread)
# Verify the executor received both history AND new message
assert len(capturing_executor.received_messages) == 3
# Verify the order: history first, then new message
assert capturing_executor.received_messages[0].text == "Previous user message"
assert capturing_executor.received_messages[1].text == "Previous assistant response"
assert capturing_executor.received_messages[2].text == "New user question"
async def test_thread_conversation_history_included_in_workflow_stream(self) -> None:
"""Test that conversation history from thread is included when streaming WorkflowAgent.
This verifies that run_stream also includes thread history.
"""
# Create an executor that captures all received messages
capturing_executor = ConversationHistoryCapturingExecutor(id="capturing_stream")
workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Thread Stream Test Agent")
# Create a thread with existing conversation history
history_messages = [
ChatMessage(role=Role.SYSTEM, text="You are a helpful assistant"),
ChatMessage(role=Role.USER, text="Hello"),
ChatMessage(role=Role.ASSISTANT, text="Hi there!"),
]
message_store = ChatMessageStore(messages=history_messages)
thread = AgentThread(message_store=message_store)
# Stream from the agent with the thread and a new message
async for _ in agent.run_stream("How are you?", thread=thread):
pass
# Verify the executor received all messages (3 from history + 1 new)
assert len(capturing_executor.received_messages) == 4
# Verify the order
assert capturing_executor.received_messages[0].text == "You are a helpful assistant"
assert capturing_executor.received_messages[1].text == "Hello"
assert capturing_executor.received_messages[2].text == "Hi there!"
assert capturing_executor.received_messages[3].text == "How are you?"
async def test_empty_thread_works_correctly(self) -> None:
"""Test that an empty thread (no message store) works correctly."""
capturing_executor = ConversationHistoryCapturingExecutor(id="empty_thread_test")
workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Empty Thread Test Agent")
# Create an empty thread
thread = AgentThread()
# Run with the empty thread
await agent.run("Just a new message", thread=thread)
# Should only receive the new message
assert len(capturing_executor.received_messages) == 1
assert capturing_executor.received_messages[0].text == "Just a new message"
async def test_checkpoint_storage_passed_to_workflow(self) -> None:
"""Test that checkpoint_storage parameter is passed through to the workflow."""
from agent_framework import InMemoryCheckpointStorage
capturing_executor = ConversationHistoryCapturingExecutor(id="checkpoint_test")
workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Checkpoint Test Agent")
# Create checkpoint storage
checkpoint_storage = InMemoryCheckpointStorage()
# Run with checkpoint storage enabled
async for _ in agent.run_stream("Test message", checkpoint_storage=checkpoint_storage):
pass
# Drain workflow events to get checkpoint
# The workflow should have created checkpoints
checkpoints = await checkpoint_storage.list_checkpoints(workflow.id)
assert len(checkpoints) > 0, "Checkpoints should have been created when checkpoint_storage is provided"
class TestWorkflowAgentMergeUpdates:
"""Test cases specifically for the WorkflowAgent.merge_updates static method."""