[BREAKING] Python: Add context mode to AgentExecutor (#4668)

* Add context mode to AgentExecutor

* Fix unit tests

* Address comments

* Address comments

* REvise context mode and add tests

* Add chain config to sequential builder

* Add sample

* Fix pipeline

* Address comments

* Address comments
This commit is contained in:
Tao Chen
2026-03-20 11:27:02 -07:00
committed by GitHub
Unverified
parent 88ea9d08c7
commit 51828abed4
11 changed files with 549 additions and 89 deletions
@@ -100,24 +100,21 @@ class _AggregateAgentConversations(Executor):
assistant_replies: list[Message] = []
for r in results:
resp_messages = list(getattr(r.agent_response, "messages", []) or [])
conv = r.full_conversation if r.full_conversation is not None else resp_messages
resp_messages = list(r.agent_response.messages)
logger.debug(
f"Aggregating executor {getattr(r, 'executor_id', '<unknown>')}: "
f"{len(resp_messages)} response msgs, {len(conv)} conversation msgs"
f"{len(resp_messages)} response msgs, {len(r.full_conversation)} conversation msgs"
)
# Capture a single user prompt (first encountered across any conversation)
if prompt_message is None:
found_user = next((m for m in conv if _is_role(m, "user")), None)
if found_user is not None:
prompt_message = found_user
prompt_message = next((m for m in r.full_conversation if _is_role(m, "user")), None)
# Pick the final assistant message from the response; fallback to conversation search
final_assistant = next((m for m in reversed(resp_messages) if _is_role(m, "assistant")), None)
if final_assistant is None:
final_assistant = next((m for m in reversed(conv) if _is_role(m, "assistant")), None)
final_assistant = next((m for m in reversed(r.full_conversation) if _is_role(m, "assistant")), None)
if final_assistant is not None:
assistant_replies.append(final_assistant)
@@ -1,6 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
from dataclasses import dataclass
from typing import Literal
from agent_framework._agents import SupportsAgentRun
from agent_framework._types import Message
@@ -117,18 +118,25 @@ class AgentApprovalExecutor(WorkflowExecutor):
agent's output or send the final response to down stream executors in the orchestration.
"""
def __init__(self, agent: SupportsAgentRun) -> None:
def __init__(
self,
agent: SupportsAgentRun,
context_mode: Literal["full", "last_agent", "custom"] | None = None,
) -> None:
"""Initialize the AgentApprovalExecutor.
Args:
agent: The agent protocol to use for generating responses.
context_mode: The mode for providing context to the agent.
"""
super().__init__(workflow=self._build_workflow(agent), id=resolve_agent_id(agent), propagate_request=True)
self._context_mode: Literal["full", "last_agent", "custom"] | None = context_mode
self._description = agent.description
super().__init__(workflow=self._build_workflow(agent), id=resolve_agent_id(agent), propagate_request=True)
def _build_workflow(self, agent: SupportsAgentRun) -> Workflow:
"""Build the internal workflow for the AgentApprovalExecutor."""
agent_executor = AgentExecutor(agent)
agent_executor = AgentExecutor(agent, context_mode=self._context_mode)
request_info_executor = AgentRequestInfoExecutor(id="agent_request_info_executor")
return (
@@ -38,7 +38,7 @@ confusion and to mirror how the concurrent builder uses explicit dispatcher/aggr
import logging
from collections.abc import Sequence
from typing import Any
from typing import Any, Literal
from agent_framework import Message, SupportsAgentRun
from agent_framework._workflows._agent_executor import (
@@ -143,6 +143,7 @@ class SequentialBuilder:
*,
participants: Sequence[SupportsAgentRun | Executor],
checkpoint_storage: CheckpointStorage | None = None,
chain_only_agent_responses: bool = False,
intermediate_outputs: bool = False,
) -> None:
"""Initialize the SequentialBuilder.
@@ -150,10 +151,14 @@ class SequentialBuilder:
Args:
participants: Sequence of agent or executor instances to run sequentially.
checkpoint_storage: Optional checkpoint storage for enabling workflow state persistence.
chain_only_agent_responses: If True, only agent responses are chained between agents.
By default, the full conversation context is passed to the next agent. This also applies
to Executor -> Agent transitions if the executor sends `AgentExecutorResponse`.
intermediate_outputs: If True, enables intermediate outputs from agent participants.
"""
self._participants: list[SupportsAgentRun | Executor] = []
self._checkpoint_storage: CheckpointStorage | None = checkpoint_storage
self._chain_only_agent_responses: bool = chain_only_agent_responses
self._request_info_enabled: bool = False
self._request_info_filter: set[str] | None = None
self._intermediate_outputs: bool = intermediate_outputs
@@ -225,6 +230,10 @@ class SequentialBuilder:
participants: list[Executor | SupportsAgentRun] = self._participants
context_mode: Literal["full", "last_agent", "custom"] | None = (
"last_agent" if self._chain_only_agent_responses else None
)
executors: list[Executor] = []
for p in participants:
if isinstance(p, Executor):
@@ -234,9 +243,9 @@ class SequentialBuilder:
not self._request_info_filter or resolve_agent_id(p) in self._request_info_filter
):
# Handle request info enabled agents
executors.append(AgentApprovalExecutor(p))
executors.append(AgentApprovalExecutor(p, context_mode=context_mode))
else:
executors.append(AgentExecutor(p))
executors.append(AgentExecutor(p, context_mode=context_mode))
else:
raise TypeError(f"Participants must be SupportsAgentRun or Executor instances. Got {type(p).__name__}.")
@@ -117,6 +117,7 @@ class TestAgentRequestInfoExecutor:
agent_response = AgentExecutorResponse(
executor_id="test_agent",
agent_response=agent_response,
full_conversation=agent_response.messages,
)
ctx = MagicMock(spec=WorkflowContext)
@@ -135,6 +136,7 @@ class TestAgentRequestInfoExecutor:
original_request = AgentExecutorResponse(
executor_id="test_agent",
agent_response=agent_response,
full_conversation=agent_response.messages,
)
response = AgentRequestInfoResponse.from_strings(["Additional input"])
@@ -161,6 +163,7 @@ class TestAgentRequestInfoExecutor:
original_request = AgentExecutorResponse(
executor_id="test_agent",
agent_response=agent_response,
full_conversation=agent_response.messages,
)
response = AgentRequestInfoResponse.approve()
@@ -1,18 +1,20 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import AsyncIterable, Awaitable
from typing import Any
from typing import Any, Literal, overload
import pytest
from agent_framework import (
AgentExecutorResponse,
AgentResponse,
AgentResponseUpdate,
AgentRunInputs,
AgentSession,
BaseAgent,
Content,
Executor,
Message,
ResponseStream,
TypeCompatibilityError,
WorkflowContext,
WorkflowRunState,
@@ -25,26 +27,45 @@ from agent_framework.orchestrations import SequentialBuilder
class _EchoAgent(BaseAgent):
"""Simple agent that appends a single assistant message with its name."""
def run( # type: ignore[override]
@overload
def run(
self,
messages: str | Message | list[str] | list[Message] | None = None,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | AsyncIterable[AgentResponseUpdate]:
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
if stream:
return self._run_stream()
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(contents=[Content.from_text(text=f"{self.name} reply")])
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
async def _run() -> AgentResponse:
return AgentResponse(messages=[Message("assistant", [f"{self.name} reply"])])
return _run()
async def _run_stream(self) -> AsyncIterable[AgentResponseUpdate]:
# Minimal async generator with one assistant update
yield AgentResponseUpdate(contents=[Content.from_text(text=f"{self.name} reply")])
class _SummarizerExec(Executor):
"""Custom executor that summarizes by appending a short assistant message."""
@@ -251,3 +272,121 @@ async def test_sequential_builder_reusable_after_build_with_participants() -> No
assert builder._participants[0] is a1 # type: ignore
assert builder._participants[1] is a2 # type: ignore
# ---------------------------------------------------------------------------
# chain_only_agent_responses tests
# ---------------------------------------------------------------------------
class _CapturingAgent(BaseAgent):
"""Agent that records the messages it received and returns a configurable reply."""
def __init__(self, *, reply_text: str = "reply", **kwargs: Any):
super().__init__(**kwargs)
self.reply_text = reply_text
self.last_messages: list[Message] = []
@overload
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
captured: list[Message] = []
if messages:
for m in messages: # type: ignore[union-attr]
if isinstance(m, Message):
captured.append(m)
elif isinstance(m, str):
captured.append(Message("user", [m]))
self.last_messages = captured
if stream:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(contents=[Content.from_text(text=self.reply_text)])
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
async def _run() -> AgentResponse:
return AgentResponse(messages=[Message("assistant", [self.reply_text])])
return _run()
async def test_chain_only_agent_responses_false_passes_full_conversation() -> None:
"""Default (chain_only_agent_responses=False) passes full conversation to the second agent."""
a1 = _CapturingAgent(id="agent1", name="A1", reply_text="A1 reply")
a2 = _CapturingAgent(id="agent2", name="A2", reply_text="A2 reply")
wf = SequentialBuilder(participants=[a1, a2], chain_only_agent_responses=False).build()
async for ev in wf.run("hello", stream=True):
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
# Second agent should see full conversation: [user("hello"), assistant("A1 reply")]
seen = a2.last_messages
assert len(seen) == 2
assert seen[0].role == "user" and "hello" in (seen[0].text or "")
assert seen[1].role == "assistant" and "A1 reply" in (seen[1].text or "")
async def test_chain_only_agent_responses_true_passes_only_agent_messages() -> None:
"""chain_only_agent_responses=True passes only the previous agent's response messages."""
a1 = _CapturingAgent(id="agent1", name="A1", reply_text="A1 reply")
a2 = _CapturingAgent(id="agent2", name="A2", reply_text="A2 reply")
wf = SequentialBuilder(participants=[a1, a2], chain_only_agent_responses=True).build()
async for ev in wf.run("hello", stream=True):
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
# Second agent should see only the assistant message: [assistant("A1 reply")]
seen = a2.last_messages
assert len(seen) == 1
assert seen[0].role == "assistant" and "A1 reply" in (seen[0].text or "")
async def test_chain_only_agent_responses_three_agents() -> None:
"""chain_only_agent_responses=True with three agents: each sees only the prior agent's reply."""
a1 = _CapturingAgent(id="agent1", name="A1", reply_text="A1 reply")
a2 = _CapturingAgent(id="agent2", name="A2", reply_text="A2 reply")
a3 = _CapturingAgent(id="agent3", name="A3", reply_text="A3 reply")
wf = SequentialBuilder(participants=[a1, a2, a3], chain_only_agent_responses=True).build()
async for ev in wf.run("hello", stream=True):
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
break
# a2 should see only A1's reply
assert len(a2.last_messages) == 1
assert a2.last_messages[0].role == "assistant" and "A1 reply" in (a2.last_messages[0].text or "")
# a3 should see only A2's reply
assert len(a3.last_messages) == 1
assert a3.last_messages[0].role == "assistant" and "A2 reply" in (a3.last_messages[0].text or "")