Python: [BREAKING] Simplify API: ChatAgent -> Agent, ChatMessage -> Message (#3747)

* [BREAKING] Rename ChatAgent -> Agent, ChatMessage -> Message, ChatClientProtocol -> SupportsChatGetResponse

Simplify the public API by removing redundant 'Chat' prefix from core types:
- ChatAgent -> Agent
- RawChatAgent -> RawAgent
- ChatMessage -> Message
- ChatClientProtocol -> SupportsChatGetResponse

Also renamed internal WorkflowMessage (was Message in _runner_context) to avoid collision.

No backward compatibility aliases - this is a clean breaking change.

* [BREAKING] Rename Agent chat_client parameter to client

* Fix rebase issues: WorkflowMessage references and broken markdown links

* Fix formatting and lint issues from code quality checks

* Fix import ordering in workflow sample files

* fixed rebase

* Fix test failures: use WorkflowMessage and A2AMessage after ChatMessage→Message rename

- Replace Message(data=..., source_id=...) with WorkflowMessage(...) in workflow tests
- Fix isinstance check in A2A agent to use A2AMessage instead of Message
- Fix import in test_workflow_observability.py (Message→WorkflowMessage)

* Fix lint, fmt, and sample errors after ChatMessage→Message rename

- Auto-fix 70+ ruff lint issues across samples (ChatMessage→Message refs)
- Fix HostedVectorStoreContent→Content.from_hosted_vector_store in file search sample
- Fix _normalize_messages→normalize_messages in custom agent sample
- Fix context.terminate→raise MiddlewareTermination in middleware samples
- Fix with_update_hook→with_transform_hook in override middleware sample
- Add TOptions_co import back to custom_chat_client sample
- Add noqa for FastAPI File() default in chatkit sample
- Fix B023 loop variable capture in weather agent sample

* fix: update Agent constructor calls from chat_client to client in declaration-only tool tests

* fix: add register_cleanup to devui lazy-loading proxy and type stub

* fixed tests and updated new pieces

* fix agui typevar

* fix merge errors

* fix merge conflicts

* fiux merge

* Remove unused links

---------

Co-authored-by: Evan Mattson <evan.mattson@microsoft.com>
This commit is contained in:
Eduard van Valkenburg
2026-02-11 00:04:32 +01:00
committed by GitHub
Unverified
parent a4c9e43afb
commit 0521f5bed8
418 changed files with 5385 additions and 5389 deletions
@@ -7,8 +7,8 @@ from agent_framework import (
AgentExecutorRequest,
AgentExecutorResponse,
AgentResponse,
ChatMessage,
Executor,
Message,
WorkflowContext,
WorkflowRunState,
handler,
@@ -32,7 +32,7 @@ class _FakeAgentExec(Executor):
@handler
async def run(self, request: AgentExecutorRequest, ctx: WorkflowContext[AgentExecutorResponse]) -> None:
response = AgentResponse(messages=ChatMessage(role="assistant", text=self._reply_text))
response = AgentResponse(messages=Message(role="assistant", text=self._reply_text))
full_conversation = list(request.messages) + list(response.messages)
await ctx.send_message(AgentExecutorResponse(self.id, response, full_conversation=full_conversation))
@@ -58,18 +58,18 @@ async def test_concurrent_default_aggregator_emits_single_user_and_assistants()
wf = ConcurrentBuilder(participants=[e1, e2, e3]).build()
completed = False
output: list[ChatMessage] | None = None
output: list[Message] | None = None
async for ev in wf.run("prompt: hello world", stream=True):
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
elif ev.type == "output":
output = cast(list[ChatMessage], ev.data)
output = cast(list[Message], ev.data)
if completed and output is not None:
break
assert completed
assert output is not None
messages: list[ChatMessage] = output
messages: list[Message] = output
# Expect one user message + one assistant message per participant
assert len(messages) == 1 + 3
@@ -89,7 +89,7 @@ async def test_concurrent_custom_aggregator_callback_is_used() -> None:
async def summarize(results: list[AgentExecutorResponse]) -> str:
texts: list[str] = []
for r in results:
msgs: list[ChatMessage] = r.agent_response.messages
msgs: list[Message] = r.agent_response.messages
texts.append(msgs[-1].text if msgs else "")
return " | ".join(sorted(texts))
@@ -120,7 +120,7 @@ async def test_concurrent_custom_aggregator_sync_callback_is_used() -> None:
def summarize_sync(results: list[AgentExecutorResponse], _ctx: WorkflowContext[Any]) -> str: # type: ignore[unused-argument]
texts: list[str] = []
for r in results:
msgs: list[ChatMessage] = r.agent_response.messages
msgs: list[Message] = r.agent_response.messages
texts.append(msgs[-1].text if msgs else "")
return " | ".join(sorted(texts))
@@ -164,7 +164,7 @@ async def test_concurrent_with_aggregator_executor_instance() -> None:
async def aggregate(self, results: list[AgentExecutorResponse], ctx: WorkflowContext[Never, str]) -> None:
texts: list[str] = []
for r in results:
msgs: list[ChatMessage] = r.agent_response.messages
msgs: list[Message] = r.agent_response.messages
texts.append(msgs[-1].text if msgs else "")
await ctx.yield_output(" & ".join(sorted(texts)))
@@ -215,7 +215,7 @@ async def test_concurrent_checkpoint_resume_round_trip() -> None:
wf = ConcurrentBuilder(participants=list(participants), checkpoint_storage=storage).build()
baseline_output: list[ChatMessage] | None = None
baseline_output: list[Message] | None = None
async for ev in wf.run("checkpoint concurrent", stream=True):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
@@ -239,7 +239,7 @@ async def test_concurrent_checkpoint_resume_round_trip() -> None:
)
wf_resume = ConcurrentBuilder(participants=list(resumed_participants), checkpoint_storage=storage).build()
resumed_output: list[ChatMessage] | None = None
resumed_output: list[Message] | None = None
async for ev in wf_resume.run(checkpoint_id=resume_checkpoint.checkpoint_id, stream=True):
if ev.type == "output":
resumed_output = ev.data # type: ignore[assignment]
@@ -261,7 +261,7 @@ async def test_concurrent_checkpoint_runtime_only() -> None:
agents = [_FakeAgentExec(id="agent1", reply_text="A1"), _FakeAgentExec(id="agent2", reply_text="A2")]
wf = ConcurrentBuilder(participants=agents).build()
baseline_output: list[ChatMessage] | None = None
baseline_output: list[Message] | None = None
async for ev in wf.run("runtime checkpoint test", checkpoint_storage=storage, stream=True):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
@@ -282,7 +282,7 @@ async def test_concurrent_checkpoint_runtime_only() -> None:
resumed_agents = [_FakeAgentExec(id="agent1", reply_text="A1"), _FakeAgentExec(id="agent2", reply_text="A2")]
wf_resume = ConcurrentBuilder(participants=resumed_agents).build()
resumed_output: list[ChatMessage] | None = None
resumed_output: list[Message] | None = None
async for ev in wf_resume.run(
checkpoint_id=resume_checkpoint.checkpoint_id, checkpoint_storage=storage, stream=True
):
@@ -311,7 +311,7 @@ async def test_concurrent_checkpoint_runtime_overrides_buildtime() -> None:
agents = [_FakeAgentExec(id="agent1", reply_text="A1"), _FakeAgentExec(id="agent2", reply_text="A2")]
wf = ConcurrentBuilder(participants=agents, checkpoint_storage=buildtime_storage).build()
baseline_output: list[ChatMessage] | None = None
baseline_output: list[Message] | None = None
async for ev in wf.run("override test", checkpoint_storage=runtime_storage, stream=True):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
@@ -5,16 +5,16 @@ from typing import Any, cast
import pytest
from agent_framework import (
Agent,
AgentExecutorResponse,
AgentResponse,
AgentResponseUpdate,
AgentThread,
BaseAgent,
ChatAgent,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Content,
Message,
WorkflowEvent,
WorkflowRunState,
)
@@ -38,7 +38,7 @@ class StubAgent(BaseAgent):
def run( # type: ignore[override]
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -49,7 +49,7 @@ class StubAgent(BaseAgent):
return self._run_impl()
async def _run_impl(self) -> AgentResponse:
response = ChatMessage(role="assistant", text=self._reply_text, author_name=self.name)
response = Message(role="assistant", text=self._reply_text, author_name=self.name)
return AgentResponse(messages=[response])
async def _run_stream_impl(self) -> AsyncIterable[AgentResponseUpdate]:
@@ -69,14 +69,14 @@ class MockChatClient:
raise NotImplementedError
class StubManagerAgent(ChatAgent):
class StubManagerAgent(Agent):
def __init__(self) -> None:
super().__init__(chat_client=MockChatClient(), name="manager_agent", description="Stub manager")
super().__init__(client=MockChatClient(), name="manager_agent", description="Stub manager")
self._call_count = 0
async def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -87,7 +87,7 @@ class StubManagerAgent(ChatAgent):
payload = {"terminate": False, "reason": "Selecting agent", "next_speaker": "agent", "final_message": None}
return AgentResponse(
messages=[
ChatMessage(
Message(
role="assistant",
text=(
'{"terminate": false, "reason": "Selecting agent", '
@@ -108,7 +108,7 @@ class StubManagerAgent(ChatAgent):
}
return AgentResponse(
messages=[
ChatMessage(
Message(
role="assistant",
text=(
'{"terminate": true, "reason": "Task complete", '
@@ -143,10 +143,10 @@ class StubMagenticManager(MagenticManagerBase):
super().__init__(max_stall_count=3, max_round_count=5)
self._round = 0
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role="assistant", text="plan", author_name="magentic_manager")
async def plan(self, magentic_context: MagenticContext) -> Message:
return Message(role="assistant", text="plan", author_name="magentic_manager")
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
async def replan(self, magentic_context: MagenticContext) -> Message:
return await self.plan(magentic_context)
async def create_progress_ledger(self, magentic_context: MagenticContext) -> MagenticProgressLedger:
@@ -169,8 +169,8 @@ class StubMagenticManager(MagenticManagerBase):
instruction_or_question=MagenticProgressLedgerItem(reason="", answer=""),
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role="assistant", text="final", author_name="magentic_manager")
async def prepare_final_answer(self, magentic_context: MagenticContext) -> Message:
return Message(role="assistant", text="final", author_name="magentic_manager")
async def test_group_chat_builder_basic_flow() -> None:
@@ -185,12 +185,12 @@ async def test_group_chat_builder_basic_flow() -> None:
orchestrator_name="manager",
).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run("coordinate task", stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
assert len(outputs) == 1
assert len(outputs[0]) >= 1
@@ -213,8 +213,8 @@ async def test_group_chat_as_agent_accepts_conversation() -> None:
agent = workflow.as_agent(name="group-chat-agent")
conversation = [
ChatMessage(role="user", text="kickoff", author_name="user"),
ChatMessage(role="assistant", text="noted", author_name="alpha"),
Message(role="user", text="kickoff", author_name="user"),
Message(role="assistant", text="noted", author_name="alpha"),
]
response = await agent.run(conversation)
@@ -324,12 +324,12 @@ class TestGroupChatWorkflow:
selection_func=selector,
).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run("test task", stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
# Should have terminated due to max_rounds, expect at least one output
assert len(outputs) >= 1
@@ -345,7 +345,7 @@ class TestGroupChatWorkflow:
def selector(state: GroupChatState) -> str:
return "agent"
def termination_condition(conversation: list[ChatMessage]) -> bool:
def termination_condition(conversation: list[Message]) -> bool:
replies = [msg for msg in conversation if msg.role == "assistant" and msg.author_name == "agent"]
return len(replies) >= 2
@@ -357,12 +357,12 @@ class TestGroupChatWorkflow:
selection_func=selector,
).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run("test task", stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
assert outputs, "Expected termination to yield output"
conversation = outputs[-1]
@@ -383,12 +383,12 @@ class TestGroupChatWorkflow:
orchestrator_agent=manager,
).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run("test task", stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
assert outputs, "Expected termination to yield output"
conversation = outputs[-1]
@@ -429,12 +429,12 @@ class TestCheckpointing:
selection_func=selector,
).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run("test task", stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
assert len(outputs) == 1 # Should complete normally
@@ -452,12 +452,12 @@ class TestConversationHandling:
workflow = GroupChatBuilder(participants=[agent], max_rounds=1, selection_func=selector).build()
with pytest.raises(ValueError, match="At least one ChatMessage is required to start the group chat workflow."):
with pytest.raises(ValueError, match="At least one Message is required to start the group chat workflow."):
async for _ in workflow.run([], stream=True):
pass
async def test_handle_string_input(self) -> None:
"""Test handling string input creates proper ChatMessage."""
"""Test handling string input creates proper Message."""
def selector(state: GroupChatState) -> str:
# Verify the conversation has the user message
@@ -470,18 +470,18 @@ class TestConversationHandling:
workflow = GroupChatBuilder(participants=[agent], max_rounds=1, selection_func=selector).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run("test string", stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
assert len(outputs) == 1
async def test_handle_chat_message_input(self) -> None:
"""Test handling ChatMessage input directly."""
task_message = ChatMessage(role="user", text="test message")
"""Test handling Message input directly."""
task_message = Message(role="user", text="test message")
def selector(state: GroupChatState) -> str:
# Verify the task message was preserved in conversation
@@ -493,20 +493,20 @@ class TestConversationHandling:
workflow = GroupChatBuilder(participants=[agent], max_rounds=1, selection_func=selector).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run(task_message, stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
assert len(outputs) == 1
async def test_handle_conversation_list_input(self) -> None:
"""Test handling conversation list preserves context."""
conversation = [
ChatMessage(role="system", text="system message"),
ChatMessage(role="user", text="user message"),
Message(role="system", text="system message"),
Message(role="user", text="user message"),
]
def selector(state: GroupChatState) -> str:
@@ -519,12 +519,12 @@ class TestConversationHandling:
workflow = GroupChatBuilder(participants=[agent], max_rounds=1, selection_func=selector).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run(conversation, stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
assert len(outputs) == 1
@@ -549,12 +549,12 @@ class TestRoundLimitEnforcement:
selection_func=selector,
).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run("test", stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
# Should have at least one output (the round limit message)
assert len(outputs) >= 1
@@ -582,12 +582,12 @@ class TestRoundLimitEnforcement:
selection_func=selector,
).build()
outputs: list[list[ChatMessage]] = []
outputs: list[list[Message]] = []
async for event in workflow.run("test", stream=True):
if event.type == "output":
data = event.data
if isinstance(data, list):
outputs.append(cast(list[ChatMessage], data))
outputs.append(cast(list[Message], data))
# Should have at least one output (the round limit message)
assert len(outputs) >= 1
@@ -608,10 +608,10 @@ async def test_group_chat_checkpoint_runtime_only() -> None:
wf = GroupChatBuilder(participants=[agent_a, agent_b], max_rounds=2, selection_func=selector).build()
baseline_output: list[ChatMessage] | None = None
baseline_output: list[Message] | None = None
async for ev in wf.run("runtime checkpoint test", checkpoint_storage=storage, stream=True):
if ev.type == "output":
baseline_output = cast(list[ChatMessage], ev.data) if isinstance(ev.data, list) else None # type: ignore
baseline_output = cast(list[Message], ev.data) if isinstance(ev.data, list) else None # type: ignore
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
@@ -644,10 +644,10 @@ async def test_group_chat_checkpoint_runtime_overrides_buildtime() -> None:
checkpoint_storage=buildtime_storage,
selection_func=selector,
).build()
baseline_output: list[ChatMessage] | None = None
baseline_output: list[Message] | None = None
async for ev in wf.run("override test", checkpoint_storage=runtime_storage, stream=True):
if ev.type == "output":
baseline_output = cast(list[ChatMessage], ev.data) if isinstance(ev.data, list) else None # type: ignore
baseline_output = cast(list[Message], ev.data) if isinstance(ev.data, list) else None # type: ignore
if ev.type == "status" and ev.state in (
WorkflowRunState.IDLE,
WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
@@ -781,8 +781,8 @@ def test_group_chat_builder_rejects_multiple_orchestrator_configurations():
def selector(state: GroupChatState) -> str:
return list(state.participants.keys())[0]
def agent_factory() -> ChatAgent:
return cast(ChatAgent, StubManagerAgent())
def agent_factory() -> Agent:
return cast(Agent, StubManagerAgent())
agent = StubAgent("test", "response")
@@ -801,8 +801,8 @@ def test_group_chat_builder_requires_exactly_one_orchestrator_option():
def selector(state: GroupChatState) -> str:
return list(state.participants.keys())[0]
def agent_factory() -> ChatAgent:
return cast(ChatAgent, StubManagerAgent())
def agent_factory() -> Agent:
return cast(Agent, StubManagerAgent())
agent = StubAgent("test", "response")
@@ -816,19 +816,19 @@ def test_group_chat_builder_requires_exactly_one_orchestrator_option():
async def test_group_chat_with_orchestrator_factory_returning_chat_agent():
"""Test workflow creation using orchestrator_factory that returns ChatAgent."""
"""Test workflow creation using orchestrator_factory that returns Agent."""
factory_call_count = 0
class DynamicManagerAgent(ChatAgent):
class DynamicManagerAgent(Agent):
"""Manager agent that dynamically selects from available participants."""
def __init__(self) -> None:
super().__init__(chat_client=MockChatClient(), name="dynamic_manager", description="Dynamic manager")
super().__init__(client=MockChatClient(), name="dynamic_manager", description="Dynamic manager")
self._call_count = 0
async def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -843,7 +843,7 @@ async def test_group_chat_with_orchestrator_factory_returning_chat_agent():
}
return AgentResponse(
messages=[
ChatMessage(
Message(
role="assistant",
text=(
'{"terminate": false, "reason": "Selecting alpha", '
@@ -863,7 +863,7 @@ async def test_group_chat_with_orchestrator_factory_returning_chat_agent():
}
return AgentResponse(
messages=[
ChatMessage(
Message(
role="assistant",
text=(
'{"terminate": true, "reason": "Task complete", '
@@ -875,10 +875,10 @@ async def test_group_chat_with_orchestrator_factory_returning_chat_agent():
value=payload,
)
def agent_factory() -> ChatAgent:
def agent_factory() -> Agent:
nonlocal factory_call_count
factory_call_count += 1
return cast(ChatAgent, DynamicManagerAgent())
return cast(Agent, DynamicManagerAgent())
alpha = StubAgent("alpha", "reply from alpha")
beta = StubAgent("beta", "reply from beta")
@@ -899,7 +899,7 @@ async def test_group_chat_with_orchestrator_factory_returning_chat_agent():
assert isinstance(final_messages, list)
assert any(
msg.text == "dynamic manager final"
for msg in cast(list[ChatMessage], final_messages)
for msg in cast(list[Message], final_messages)
if msg.author_name == "dynamic_manager"
)
@@ -939,10 +939,10 @@ async def test_group_chat_orchestrator_factory_reusable_builder():
"""Test that the builder can be reused to build multiple workflows with orchestrator factory."""
factory_call_count = 0
def agent_factory() -> ChatAgent:
def agent_factory() -> Agent:
nonlocal factory_call_count
factory_call_count += 1
return cast(ChatAgent, StubManagerAgent())
return cast(Agent, StubManagerAgent())
alpha = StubAgent("alpha", "reply from alpha")
beta = StubAgent("beta", "reply from beta")
@@ -971,13 +971,13 @@ def test_group_chat_orchestrator_factory_invalid_return_type():
with pytest.raises(
TypeError,
match=r"Orchestrator factory must return ChatAgent or BaseGroupChatOrchestrator instance",
match=r"Orchestrator factory must return Agent or BaseGroupChatOrchestrator instance",
):
GroupChatBuilder(participants=[alpha], orchestrator=invalid_factory).build()
with pytest.raises(
TypeError,
match=r"Orchestrator factory must return ChatAgent or BaseGroupChatOrchestrator instance",
match=r"Orchestrator factory must return Agent or BaseGroupChatOrchestrator instance",
):
GroupChatBuilder(participants=[alpha], orchestrator_agent=invalid_factory).build()
@@ -6,13 +6,13 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from agent_framework import (
ChatAgent,
ChatMessage,
Agent,
ChatResponse,
ChatResponseUpdate,
Content,
Context,
ContextProvider,
Message,
ResponseStream,
WorkflowEvent,
resolve_agent_id,
@@ -50,7 +50,7 @@ class MockChatClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], Bas
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
@@ -60,7 +60,7 @@ class MockChatClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], Bas
async def _get() -> ChatResponse:
contents = _build_reply_contents(self._name, self._handoff_to, self._next_call_id())
reply = ChatMessage(
reply = Message(
role="assistant",
contents=contents,
)
@@ -105,7 +105,7 @@ def _build_reply_contents(
return contents
class MockHandoffAgent(ChatAgent):
class MockHandoffAgent(Agent):
"""Mock agent that can hand off to another agent."""
def __init__(
@@ -121,7 +121,7 @@ class MockHandoffAgent(ChatAgent):
handoff_to: The name of the agent to hand off to, or None for no handoff.
This is hardcoded for testing purposes so that the agent always attempts to hand off.
"""
super().__init__(chat_client=MockChatClient(name=name, handoff_to=handoff_to), name=name, id=name)
super().__init__(client=MockChatClient(name=name, handoff_to=handoff_to), name=name, id=name)
async def _drain(stream: AsyncIterable[WorkflowEvent]) -> list[WorkflowEvent]:
@@ -196,7 +196,7 @@ async def test_autonomous_mode_yields_output_without_user_request():
final_conversation = outputs[-1].data
assert isinstance(final_conversation, list)
conversation_list = cast(list[ChatMessage], final_conversation)
conversation_list = cast(list[Message], final_conversation)
assert any(msg.role == "assistant" and (msg.text or "").startswith("specialist reply") for msg in conversation_list)
@@ -237,7 +237,7 @@ async def test_handoff_async_termination_condition() -> None:
"""Test that async termination conditions work correctly."""
termination_call_count = 0
async def async_termination(conv: list[ChatMessage]) -> bool:
async def async_termination(conv: list[Message]) -> bool:
nonlocal termination_call_count
termination_call_count += 1
user_count = sum(1 for msg in conv if msg.role == "user")
@@ -258,7 +258,7 @@ async def test_handoff_async_termination_condition() -> None:
events = await _drain(
workflow.run(
stream=True, responses={requests[-1].request_id: [ChatMessage(role="user", text="Second user message")]}
stream=True, responses={requests[-1].request_id: [Message(role="user", text="Second user message")]}
)
)
outputs = [ev for ev in events if ev.type == "output"]
@@ -266,7 +266,7 @@ async def test_handoff_async_termination_condition() -> None:
final_conversation = outputs[0].data
assert isinstance(final_conversation, list)
final_conv_list = cast(list[ChatMessage], final_conversation)
final_conv_list = cast(list[Message], final_conversation)
user_messages = [msg for msg in final_conv_list if msg.role == "user"]
assert len(user_messages) == 2
assert termination_call_count > 0
@@ -281,7 +281,7 @@ async def test_tool_choice_preserved_from_agent_config():
if options:
recorded_tool_choices.append(options.get("tool_choice"))
return ChatResponse(
messages=[ChatMessage(role="assistant", text="Response")],
messages=[Message(role="assistant", text="Response")],
response_id="test_response",
)
@@ -289,8 +289,8 @@ async def test_tool_choice_preserved_from_agent_config():
mock_client.get_response = AsyncMock(side_effect=mock_get_response)
# Create agent with specific tool_choice configuration via default_options
agent = ChatAgent(
chat_client=mock_client,
agent = Agent(
client=mock_client,
name="test_agent",
default_options={"tool_choice": {"mode": "required"}}, # type: ignore
)
@@ -313,7 +313,7 @@ async def test_context_provider_preserved_during_handoff():
class TestContextProvider(ContextProvider):
"""A test context provider that tracks its invocations."""
async def invoking(self, messages: Sequence[ChatMessage], **kwargs: Any) -> Context:
async def invoking(self, messages: Sequence[Message], **kwargs: Any) -> Context:
provider_calls.append("invoking")
return Context(instructions="Test context from provider.")
@@ -324,8 +324,8 @@ async def test_context_provider_preserved_during_handoff():
mock_client = MockChatClient(name="test_agent")
# Create agent with context provider using proper constructor
agent = ChatAgent(
chat_client=mock_client,
agent = Agent(
client=mock_client,
name="test_agent",
id="test_agent",
context_provider=context_provider,
@@ -11,9 +11,9 @@ from agent_framework import (
AgentResponseUpdate,
AgentThread,
BaseAgent,
ChatMessage,
Content,
Executor,
Message,
SupportsAgentRun,
Workflow,
WorkflowCheckpoint,
@@ -48,7 +48,7 @@ def test_magentic_context_reset_behavior():
participant_descriptions={"Alice": "Researcher"},
)
# seed context state
ctx.chat_history.append(ChatMessage("assistant", ["draft"]))
ctx.chat_history.append(Message("assistant", ["draft"]))
ctx.stall_count = 2
prev_reset = ctx.reset_count
@@ -61,8 +61,8 @@ def test_magentic_context_reset_behavior():
@dataclass
class _SimpleLedger:
facts: ChatMessage
plan: ChatMessage
facts: Message
plan: Message
class FakeManager(MagenticManagerBase):
@@ -108,25 +108,25 @@ class FakeManager(MagenticManagerBase):
plan_payload = cast(dict[str, Any] | None, ledger_dict.get("plan"))
if facts_payload is not None and plan_payload is not None:
try:
facts = ChatMessage.from_dict(facts_payload)
plan = ChatMessage.from_dict(plan_payload)
facts = Message.from_dict(facts_payload)
plan = Message.from_dict(plan_payload)
self.task_ledger = _SimpleLedger(facts=facts, plan=plan)
except Exception: # pragma: no cover - defensive
pass
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
facts = ChatMessage("assistant", ["GIVEN OR VERIFIED FACTS\n- A\n"])
plan = ChatMessage("assistant", ["- Do X\n- Do Y\n"])
async def plan(self, magentic_context: MagenticContext) -> Message:
facts = Message("assistant", ["GIVEN OR VERIFIED FACTS\n- A\n"])
plan = Message("assistant", ["- Do X\n- Do Y\n"])
self.task_ledger = _SimpleLedger(facts=facts, plan=plan)
combined = f"Task: {magentic_context.task}\n\nFacts:\n{facts.text}\n\nPlan:\n{plan.text}"
return ChatMessage("assistant", [combined], author_name=self.name)
return Message("assistant", [combined], author_name=self.name)
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
facts = ChatMessage("assistant", ["GIVEN OR VERIFIED FACTS\n- A2\n"])
plan = ChatMessage("assistant", ["- Do Z\n"])
async def replan(self, magentic_context: MagenticContext) -> Message:
facts = Message("assistant", ["GIVEN OR VERIFIED FACTS\n- A2\n"])
plan = Message("assistant", ["- Do Z\n"])
self.task_ledger = _SimpleLedger(facts=facts, plan=plan)
combined = f"Task: {magentic_context.task}\n\nFacts:\n{facts.text}\n\nPlan:\n{plan.text}"
return ChatMessage("assistant", [combined], author_name=self.name)
return Message("assistant", [combined], author_name=self.name)
async def create_progress_ledger(self, magentic_context: MagenticContext) -> MagenticProgressLedger:
# At least two messages in chat history means request is satisfied for testing
@@ -139,8 +139,8 @@ class FakeManager(MagenticManagerBase):
instruction_or_question=MagenticProgressLedgerItem(reason="test", answer=self.instruction_text),
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage("assistant", [self.FINAL_ANSWER], author_name=self.name)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> Message:
return Message("assistant", [self.FINAL_ANSWER], author_name=self.name)
class StubAgent(BaseAgent):
@@ -150,7 +150,7 @@ class StubAgent(BaseAgent):
def run( # type: ignore[override]
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -160,7 +160,7 @@ class StubAgent(BaseAgent):
return self._run_stream()
async def _run() -> AgentResponse:
response = ChatMessage("assistant", [self._reply_text], author_name=self.name)
response = Message("assistant", [self._reply_text], author_name=self.name)
return AgentResponse(messages=[response])
return _run()
@@ -177,7 +177,7 @@ class DummyExec(Executor):
@handler
async def _noop(
self, message: GroupChatRequestMessage, ctx: WorkflowContext[ChatMessage]
self, message: GroupChatRequestMessage, ctx: WorkflowContext[Message]
) -> None: # pragma: no cover - not called
pass
@@ -190,13 +190,13 @@ async def test_magentic_builder_returns_workflow_and_runs() -> None:
assert isinstance(workflow, Workflow)
outputs: list[ChatMessage] = []
outputs: list[Message] = []
orchestrator_event_count = 0
async for event in workflow.run("compose summary", stream=True):
if event.type == "output":
msg = event.data
if isinstance(msg, list):
outputs.extend(cast(list[ChatMessage], msg))
outputs.extend(cast(list[Message], msg))
elif event.type == "magentic_orchestrator":
orchestrator_event_count += 1
@@ -216,8 +216,8 @@ async def test_magentic_as_agent_does_not_accept_conversation() -> None:
agent = workflow.as_agent(name="magentic-agent")
conversation = [
ChatMessage("system", ["Guidelines"], author_name="system"),
ChatMessage("user", ["Summarize the findings"], author_name="requester"),
Message("system", ["Guidelines"], author_name="system"),
Message("user", ["Summarize the findings"], author_name="requester"),
]
with pytest.raises(ValueError, match="Magentic only support a single task message to start the workflow."):
await agent.run(conversation)
@@ -250,7 +250,7 @@ async def test_magentic_workflow_plan_review_approval_to_completion():
assert isinstance(req_event.data, MagenticPlanReviewRequest)
completed = False
output: list[ChatMessage] | None = None
output: list[Message] | None = None
async for ev in wf.run(stream=True, responses={req_event.request_id: req_event.data.approve()}):
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
@@ -262,7 +262,7 @@ async def test_magentic_workflow_plan_review_approval_to_completion():
assert completed
assert output is not None
assert isinstance(output, list)
assert all(isinstance(msg, ChatMessage) for msg in output)
assert all(isinstance(msg, Message) for msg in output)
async def test_magentic_plan_review_with_revise():
@@ -273,7 +273,7 @@ async def test_magentic_plan_review_with_revise():
def __init__(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def]
super().__init__(*args, **kwargs)
async def replan(self, magentic_context: MagenticContext) -> ChatMessage: # type: ignore[override]
async def replan(self, magentic_context: MagenticContext) -> Message: # type: ignore[override]
self.replan_count += 1
return await super().replan(magentic_context)
@@ -340,7 +340,7 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result():
assert isinstance(data, list)
assert len(data) > 0 # type: ignore
assert data[-1].role == "assistant" # type: ignore
assert all(isinstance(msg, ChatMessage) for msg in data) # type: ignore
assert all(isinstance(msg, Message) for msg in data) # type: ignore
async def test_magentic_checkpoint_resume_round_trip():
@@ -406,7 +406,7 @@ class StubManagerAgent(BaseAgent):
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: Any = None,
@@ -416,22 +416,22 @@ class StubManagerAgent(BaseAgent):
return self._run_stream()
async def _run() -> AgentResponse:
return AgentResponse(messages=[ChatMessage("assistant", ["ok"])])
return AgentResponse(messages=[Message("assistant", ["ok"])])
return _run()
async def _run_stream(self) -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(message_deltas=[ChatMessage("assistant", ["ok"])])
yield AgentResponseUpdate(message_deltas=[Message("assistant", ["ok"])])
async def test_standard_manager_plan_and_replan_via_complete_monkeypatch():
mgr = StandardMagenticManager(StubManagerAgent())
async def fake_complete_plan(messages: list[ChatMessage], **kwargs: Any) -> ChatMessage:
async def fake_complete_plan(messages: list[Message], **kwargs: Any) -> Message:
# Return a different response depending on call order length
if any("FACTS" in (m.text or "") for m in messages):
return ChatMessage("assistant", ["- step A\n- step B"])
return ChatMessage("assistant", ["GIVEN OR VERIFIED FACTS\n- fact1"])
return Message("assistant", ["- step A\n- step B"])
return Message("assistant", ["GIVEN OR VERIFIED FACTS\n- fact1"])
# First, patch to produce facts then plan
mgr._complete = fake_complete_plan # type: ignore[attr-defined]
@@ -444,10 +444,10 @@ async def test_standard_manager_plan_and_replan_via_complete_monkeypatch():
assert any(t in combined.text for t in ("- step A", "- step B", "- step"))
# Now replan with new outputs
async def fake_complete_replan(messages: list[ChatMessage], **kwargs: Any) -> ChatMessage:
async def fake_complete_replan(messages: list[Message], **kwargs: Any) -> Message:
if any("Please briefly explain" in (m.text or "") for m in messages):
return ChatMessage("assistant", ["- new step"])
return ChatMessage("assistant", ["GIVEN OR VERIFIED FACTS\n- updated"])
return Message("assistant", ["- new step"])
return Message("assistant", ["GIVEN OR VERIFIED FACTS\n- updated"])
mgr._complete = fake_complete_replan # type: ignore[attr-defined]
combined2 = await mgr.replan(ctx.clone())
@@ -459,7 +459,7 @@ async def test_standard_manager_progress_ledger_success_and_error():
ctx = MagenticContext(task="task", participant_descriptions={"alice": "desc"})
# Success path: valid JSON
async def fake_complete_ok(messages: list[ChatMessage], **kwargs: Any) -> ChatMessage:
async def fake_complete_ok(messages: list[Message], **kwargs: Any) -> Message:
json_text = (
'{"is_request_satisfied": {"reason": "r", "answer": false}, '
'"is_in_loop": {"reason": "r", "answer": false}, '
@@ -467,15 +467,15 @@ async def test_standard_manager_progress_ledger_success_and_error():
'"next_speaker": {"reason": "r", "answer": "alice"}, '
'"instruction_or_question": {"reason": "r", "answer": "do"}}'
)
return ChatMessage("assistant", [json_text])
return Message("assistant", [json_text])
mgr._complete = fake_complete_ok # type: ignore[attr-defined]
ledger = await mgr.create_progress_ledger(ctx.clone())
assert ledger.next_speaker.answer == "alice"
# Error path: invalid JSON now raises to avoid emitting planner-oriented instructions to agents
async def fake_complete_bad(messages: list[ChatMessage], **kwargs: Any) -> ChatMessage:
return ChatMessage("assistant", ["not-json"])
async def fake_complete_bad(messages: list[Message], **kwargs: Any) -> Message:
return Message("assistant", ["not-json"])
mgr._complete = fake_complete_bad # type: ignore[attr-defined]
with pytest.raises(RuntimeError):
@@ -487,11 +487,11 @@ class InvokeOnceManager(MagenticManagerBase):
super().__init__(max_round_count=5, max_stall_count=3, max_reset_count=2)
self._invoked = False
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage("assistant", ["ledger"])
async def plan(self, magentic_context: MagenticContext) -> Message:
return Message("assistant", ["ledger"])
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage("assistant", ["re-ledger"])
async def replan(self, magentic_context: MagenticContext) -> Message:
return Message("assistant", ["re-ledger"])
async def create_progress_ledger(self, magentic_context: MagenticContext) -> MagenticProgressLedger:
if not self._invoked:
@@ -513,8 +513,8 @@ class InvokeOnceManager(MagenticManagerBase):
instruction_or_question=MagenticProgressLedgerItem(reason="r", answer="done"),
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage("assistant", ["final"])
async def prepare_final_answer(self, magentic_context: MagenticContext) -> Message:
return Message("assistant", ["final"])
class StubThreadAgent(BaseAgent):
@@ -526,7 +526,7 @@ class StubThreadAgent(BaseAgent):
return self._run_stream()
async def _run():
return AgentResponse(messages=[ChatMessage("assistant", ["thread-ok"], author_name=self.name)])
return AgentResponse(messages=[Message("assistant", ["thread-ok"], author_name=self.name)])
return _run()
@@ -543,18 +543,18 @@ class StubAssistantsClient:
class StubAssistantsAgent(BaseAgent):
chat_client: object | None = None # allow assignment via Pydantic field
client: object | None = None # allow assignment via Pydantic field
def __init__(self) -> None:
super().__init__(name="agentA")
self.chat_client = StubAssistantsClient() # type name contains 'AssistantsClient'
self.client = StubAssistantsClient() # type name contains 'AssistantsClient'
def run(self, messages=None, *, stream: bool = False, thread=None, **kwargs): # type: ignore[override]
if stream:
return self._run_stream()
async def _run():
return AgentResponse(messages=[ChatMessage("assistant", ["assistants-ok"], author_name=self.name)])
return AgentResponse(messages=[Message("assistant", ["assistants-ok"], author_name=self.name)])
return _run()
@@ -566,8 +566,8 @@ class StubAssistantsAgent(BaseAgent):
)
async def _collect_agent_responses_setup(participant: SupportsAgentRun) -> list[ChatMessage]:
captured: list[ChatMessage] = []
async def _collect_agent_responses_setup(participant: SupportsAgentRun) -> list[Message]:
captured: list[Message] = []
wf = MagenticBuilder(participants=[participant], intermediate_outputs=True, manager=InvokeOnceManager()).build()
@@ -578,7 +578,7 @@ async def _collect_agent_responses_setup(participant: SupportsAgentRun) -> list[
# Capture streaming updates (type="output" with AgentResponseUpdate data)
if ev.type == "output" and isinstance(ev.data, AgentResponseUpdate):
captured.append(
ChatMessage(
Message(
role=ev.data.role or "assistant",
text=ev.data.text or "",
author_name=ev.data.author_name,
@@ -711,11 +711,11 @@ class NotProgressingManager(MagenticManagerBase):
A manager that never marks progress being made, to test stall/reset limits.
"""
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage("assistant", ["ledger"])
async def plan(self, magentic_context: MagenticContext) -> Message:
return Message("assistant", ["ledger"])
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage("assistant", ["re-ledger"])
async def replan(self, magentic_context: MagenticContext) -> Message:
return Message("assistant", ["re-ledger"])
async def create_progress_ledger(self, magentic_context: MagenticContext) -> MagenticProgressLedger:
return MagenticProgressLedger(
@@ -726,8 +726,8 @@ class NotProgressingManager(MagenticManagerBase):
instruction_or_question=MagenticProgressLedgerItem(reason="r", answer="done"),
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage("assistant", ["final"])
async def prepare_final_answer(self, magentic_context: MagenticContext) -> Message:
return Message("assistant", ["final"])
async def test_magentic_stall_and_reset_reach_limits():
@@ -747,7 +747,7 @@ async def test_magentic_stall_and_reset_reach_limits():
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
assert all(isinstance(msg, Message) for msg in output_event.data) # type: ignore
assert len(output_event.data) > 0 # type: ignore
assert output_event.data[-1].text is not None # type: ignore
assert output_event.data[-1].text == "Workflow terminated due to reaching maximum reset count." # type: ignore
@@ -760,7 +760,7 @@ async def test_magentic_checkpoint_runtime_only() -> None:
manager = FakeManager(max_round_count=10)
wf = MagenticBuilder(participants=[DummyExec("agentA")], manager=manager).build()
baseline_output: ChatMessage | None = None
baseline_output: Message | None = None
async for ev in wf.run("runtime checkpoint test", checkpoint_storage=storage, stream=True):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
@@ -794,7 +794,7 @@ async def test_magentic_checkpoint_runtime_overrides_buildtime() -> None:
participants=[DummyExec("agentA")], checkpoint_storage=buildtime_storage, manager=manager
).build()
baseline_output: ChatMessage | None = None
baseline_output: Message | None = None
async for ev in wf.run("override test", checkpoint_storage=runtime_storage, stream=True):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
@@ -821,8 +821,8 @@ async def test_magentic_context_no_duplicate_on_reset():
ctx = MagenticContext(task="task", participant_descriptions={"Alice": "Researcher"})
# Add some history
ctx.chat_history.append(ChatMessage("assistant", ["response1"]))
ctx.chat_history.append(ChatMessage("assistant", ["response2"]))
ctx.chat_history.append(Message("assistant", ["response1"]))
ctx.chat_history.append(Message("assistant", ["response2"]))
assert len(ctx.chat_history) == 2
# Reset
@@ -832,7 +832,7 @@ async def test_magentic_context_no_duplicate_on_reset():
assert len(ctx.chat_history) == 0, "chat_history should be empty after reset"
# Add new history
ctx.chat_history.append(ChatMessage("assistant", ["new_response"]))
ctx.chat_history.append(Message("assistant", ["new_response"]))
assert len(ctx.chat_history) == 1, "Should have exactly 1 message after adding to reset context"
@@ -844,8 +844,8 @@ async def test_magentic_checkpoint_restore_no_duplicate_history():
wf = MagenticBuilder(participants=[DummyExec("agentA")], checkpoint_storage=storage, manager=manager).build()
# Run with conversation history to create initial checkpoint
conversation: list[ChatMessage] = [
ChatMessage("user", ["task_msg"]),
conversation: list[Message] = [
Message("user", ["task_msg"]),
]
async for event in wf.run(conversation, stream=True):
@@ -1022,8 +1022,8 @@ def test_magentic_agent_factory_with_standard_manager_options():
from agent_framework_orchestrations._magentic import _MagenticTaskLedger # type: ignore
custom_task_ledger = _MagenticTaskLedger(
facts=ChatMessage("assistant", ["Custom facts"]),
plan=ChatMessage("assistant", ["Custom plan"]),
facts=Message("assistant", ["Custom facts"]),
plan=Message("assistant", ["Custom plan"]),
)
participant = StubAgent("agentA", "reply from agentA")
@@ -11,7 +11,7 @@ from agent_framework import (
AgentResponse,
AgentResponseUpdate,
AgentThread,
ChatMessage,
Message,
SupportsAgentRun,
)
from agent_framework._workflows._agent_executor import AgentExecutorRequest, AgentExecutorResponse
@@ -72,16 +72,16 @@ class TestAgentRequestInfoResponse:
def test_create_response_with_messages(self):
"""Test creating an AgentRequestInfoResponse with messages."""
messages = [ChatMessage(role="user", text="Additional info")]
messages = [Message(role="user", text="Additional info")]
response = AgentRequestInfoResponse(messages=messages)
assert response.messages == messages
def test_from_messages_factory(self):
"""Test creating response from ChatMessage list."""
"""Test creating response from Message list."""
messages = [
ChatMessage(role="user", text="Message 1"),
ChatMessage(role="user", text="Message 2"),
Message(role="user", text="Message 1"),
Message(role="user", text="Message 2"),
]
response = AgentRequestInfoResponse.from_messages(messages)
@@ -113,7 +113,7 @@ class TestAgentRequestInfoExecutor:
"""Test that request_info handler calls ctx.request_info."""
executor = AgentRequestInfoExecutor(id="test_executor")
agent_response = AgentResponse(messages=[ChatMessage(role="assistant", text="Agent response")])
agent_response = AgentResponse(messages=[Message(role="assistant", text="Agent response")])
agent_response = AgentExecutorResponse(
executor_id="test_agent",
agent_response=agent_response,
@@ -131,7 +131,7 @@ class TestAgentRequestInfoExecutor:
"""Test response handler when user provides additional messages."""
executor = AgentRequestInfoExecutor(id="test_executor")
agent_response = AgentResponse(messages=[ChatMessage(role="assistant", text="Original")])
agent_response = AgentResponse(messages=[Message(role="assistant", text="Original")])
original_request = AgentExecutorResponse(
executor_id="test_agent",
agent_response=agent_response,
@@ -157,7 +157,7 @@ class TestAgentRequestInfoExecutor:
"""Test response handler when user approves (no additional messages)."""
executor = AgentRequestInfoExecutor(id="test_executor")
agent_response = AgentResponse(messages=[ChatMessage(role="assistant", text="Original")])
agent_response = AgentResponse(messages=[Message(role="assistant", text="Original")])
original_request = AgentExecutorResponse(
executor_id="test_agent",
agent_response=agent_response,
@@ -200,7 +200,7 @@ class _TestAgent:
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -209,10 +209,10 @@ class _TestAgent:
"""Dummy run method."""
if stream:
return self._run_stream_impl()
return AgentResponse(messages=[ChatMessage(role="assistant", text="Test response")])
return AgentResponse(messages=[Message(role="assistant", text="Test response")])
async def _run_stream_impl(self) -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(messages=[ChatMessage(role="assistant", text="Test response stream")])
yield AgentResponseUpdate(messages=[Message(role="assistant", text="Test response stream")])
def get_new_thread(self, **kwargs: Any) -> AgentThread:
"""Creates a new conversation thread for the agent."""
@@ -10,9 +10,9 @@ from agent_framework import (
AgentResponseUpdate,
AgentThread,
BaseAgent,
ChatMessage,
Content,
Executor,
Message,
TypeCompatibilityError,
WorkflowContext,
WorkflowRunState,
@@ -27,7 +27,7 @@ class _EchoAgent(BaseAgent):
def run( # type: ignore[override]
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -37,7 +37,7 @@ class _EchoAgent(BaseAgent):
return self._run_stream()
async def _run() -> AgentResponse:
return AgentResponse(messages=[ChatMessage("assistant", [f"{self.name} reply"])])
return AgentResponse(messages=[Message("assistant", [f"{self.name} reply"])])
return _run()
@@ -50,11 +50,11 @@ class _SummarizerExec(Executor):
"""Custom executor that summarizes by appending a short assistant message."""
@handler
async def summarize(self, agent_response: AgentExecutorResponse, ctx: WorkflowContext[list[ChatMessage]]) -> None:
async def summarize(self, agent_response: AgentExecutorResponse, ctx: WorkflowContext[list[Message]]) -> None:
conversation = agent_response.full_conversation or []
user_texts = [m.text for m in conversation if m.role == "user"]
agents = [m.author_name or m.role for m in conversation if m.role == "assistant"]
summary = ChatMessage("assistant", [f"Summary of users:{len(user_texts)} agents:{len(agents)}"])
summary = Message("assistant", [f"Summary of users:{len(user_texts)} agents:{len(agents)}"])
await ctx.send_message(list(conversation) + [summary])
@@ -62,7 +62,7 @@ class _InvalidExecutor(Executor):
"""Invalid executor that does not have a handler that accepts a list of chat messages"""
@handler
async def summarize(self, conversation: list[str], ctx: WorkflowContext[list[ChatMessage]]) -> None:
async def summarize(self, conversation: list[str], ctx: WorkflowContext[list[Message]]) -> None:
pass
@@ -84,7 +84,7 @@ async def test_sequential_agents_append_to_context() -> None:
wf = SequentialBuilder(participants=[a1, a2]).build()
completed = False
output: list[ChatMessage] | None = None
output: list[Message] | None = None
async for ev in wf.run("hello sequential", stream=True):
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
@@ -96,7 +96,7 @@ async def test_sequential_agents_append_to_context() -> None:
assert completed
assert output is not None
assert isinstance(output, list)
msgs: list[ChatMessage] = output
msgs: list[Message] = output
assert len(msgs) == 3
assert msgs[0].role == "user" and "hello sequential" in msgs[0].text
assert msgs[1].role == "assistant" and (msgs[1].author_name == "A1" or True)
@@ -112,7 +112,7 @@ async def test_sequential_with_custom_executor_summary() -> None:
wf = SequentialBuilder(participants=[a1, summarizer]).build()
completed = False
output: list[ChatMessage] | None = None
output: list[Message] | None = None
async for ev in wf.run("topic X", stream=True):
if ev.type == "status" and ev.state == WorkflowRunState.IDLE:
completed = True
@@ -123,7 +123,7 @@ async def test_sequential_with_custom_executor_summary() -> None:
assert completed
assert output is not None
msgs: list[ChatMessage] = output
msgs: list[Message] = output
# Expect: [user, A1 reply, summary]
assert len(msgs) == 3
assert msgs[0].role == "user"
@@ -137,7 +137,7 @@ async def test_sequential_checkpoint_resume_round_trip() -> None:
initial_agents = (_EchoAgent(id="agent1", name="A1"), _EchoAgent(id="agent2", name="A2"))
wf = SequentialBuilder(participants=list(initial_agents), checkpoint_storage=storage).build()
baseline_output: list[ChatMessage] | None = None
baseline_output: list[Message] | None = None
async for ev in wf.run("checkpoint sequential", stream=True):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
@@ -158,7 +158,7 @@ async def test_sequential_checkpoint_resume_round_trip() -> None:
resumed_agents = (_EchoAgent(id="agent1", name="A1"), _EchoAgent(id="agent2", name="A2"))
wf_resume = SequentialBuilder(participants=list(resumed_agents), checkpoint_storage=storage).build()
resumed_output: list[ChatMessage] | None = None
resumed_output: list[Message] | None = None
async for ev in wf_resume.run(checkpoint_id=resume_checkpoint.checkpoint_id, stream=True):
if ev.type == "output":
resumed_output = ev.data # type: ignore[assignment]
@@ -180,7 +180,7 @@ async def test_sequential_checkpoint_runtime_only() -> None:
agents = (_EchoAgent(id="agent1", name="A1"), _EchoAgent(id="agent2", name="A2"))
wf = SequentialBuilder(participants=list(agents)).build()
baseline_output: list[ChatMessage] | None = None
baseline_output: list[Message] | None = None
async for ev in wf.run("runtime checkpoint test", checkpoint_storage=storage, stream=True):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]
@@ -201,7 +201,7 @@ async def test_sequential_checkpoint_runtime_only() -> None:
resumed_agents = (_EchoAgent(id="agent1", name="A1"), _EchoAgent(id="agent2", name="A2"))
wf_resume = SequentialBuilder(participants=list(resumed_agents)).build()
resumed_output: list[ChatMessage] | None = None
resumed_output: list[Message] | None = None
async for ev in wf_resume.run(
checkpoint_id=resume_checkpoint.checkpoint_id, checkpoint_storage=storage, stream=True
):
@@ -231,7 +231,7 @@ async def test_sequential_checkpoint_runtime_overrides_buildtime() -> None:
agents = (_EchoAgent(id="agent1", name="A1"), _EchoAgent(id="agent2", name="A2"))
wf = SequentialBuilder(participants=list(agents), checkpoint_storage=buildtime_storage).build()
baseline_output: list[ChatMessage] | None = None
baseline_output: list[Message] | None = None
async for ev in wf.run("override test", checkpoint_storage=runtime_storage, stream=True):
if ev.type == "output":
baseline_output = ev.data # type: ignore[assignment]