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

Python: [BREAKING] Types API Review improvements (#3647)

Browse Source 
* Replace Role and FinishReason classes with NewType + Literal

- Remove EnumLike metaclass from _types.py
- Replace Role class with NewType('Role', str) + RoleLiteral
- Replace FinishReason class with NewType('FinishReason', str) + FinishReasonLiteral
- Update all usages across codebase to use string literals
- Remove .value access patterns (direct string comparison now works)
- Add backward compatibility for legacy dict serialization format
- Update tests to reflect new string-based types

Addresses #3591, #3615

* Simplify ChatResponse and AgentResponse type hints (#3592)

- Remove overloads from ChatResponse.__init__
- Remove text parameter from ChatResponse.__init__
- Remove | dict[str, Any] from finish_reason and usage_details params
- Remove **kwargs from AgentResponse.__init__
- Both now accept ChatMessage | Sequence[ChatMessage] | None for messages
- Update docstrings and examples to reflect changes
- Fix tests that were using removed kwargs
- Fix Role type hint usage in ag-ui utils

* Remove text parameter from ChatResponseUpdate and AgentResponseUpdate (#3597)

- Remove text parameter from ChatResponseUpdate.__init__
- Remove text parameter from AgentResponseUpdate.__init__
- Remove **kwargs from both update classes
- Simplify contents parameter type to Sequence[Content] | None
- Update all usages to use contents=[Content.from_text(...)] pattern
- Fix imports in test files
- Update docstrings and examples

* Rename from_chat_response_updates to from_updates (#3593)

- ChatResponse.from_chat_response_updates → ChatResponse.from_updates
- ChatResponse.from_chat_response_generator → ChatResponse.from_update_generator
- AgentResponse.from_agent_run_response_updates → AgentResponse.from_updates

* Remove try_parse_value method from ChatResponse and AgentResponse (#3595)

- Remove try_parse_value method from ChatResponse
- Remove try_parse_value method from AgentResponse
- Remove try_parse_value calls from from_updates and from_update_generator methods
- Update samples to use try/except with response.value instead
- Update tests to use response.value pattern
- Users should now use response.value with try/except for safe parsing

* Add agent_id to AgentResponse and clarify author_name documentation (#3596)

- Add agent_id parameter to AgentResponse class
- Document that author_name is on ChatMessage objects, not responses
- Update ChatResponse docstring with author_name note
- Update AgentResponse docstring with author_name note

* Simplify ChatMessage.__init__ signature (#3618)

- Make contents a positional argument accepting Sequence[Content | str]
- Auto-convert strings in contents to TextContent
- Remove overloads, keep text kwarg for backward compatibility with serialization
- Update _parse_content_list to handle string items
- Update all usages across codebase to use new format: ChatMessage("role", ["text"])

* Allow Content as input on run and get_response

- Update prepare_messages and normalize_messages to accept Content
- Update type signatures in _agents.py and _clients.py
- Add tests for Content input handling

* Fix ChatMessage usage across packages and samples

Update all remaining ChatMessage(role=..., text=...) to use new
ChatMessage('role', ['text']) signature.

* Fix Role string usage and response format parsing

- Fix redis provider: remove .value access on string literals
- Fix durabletask ensure_response_format: set _response_format before accessing .value

* Fix ollama .value and ai_model_id issues, handle None in content list

- Fix ollama _chat_client: remove .value on string literals
- Fix ollama _chat_client: rename ai_model_id to model_id
- Fix _parse_content_list: skip None values gracefully

* Fix A2AAgent type signature to include Content

* Fix Role/FinishReason NewType dict annotations and improve test coverage to 95%

* Fix mypy errors for Role/FinishReason NewType usage

* Fix Role.TOOL and Role.ASSISTANT usage in _orchestrator_helpers.py

* Fix Role NewType usage in durabletask _models.py
This commit is contained in:
Eduard van Valkenburg
2026-02-04 10:13:23 +00:00
committed by GitHub
parent ef798629e5
commit 838a7fd61d
341 changed files with 3766 additions and 3228 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/tests/azure/test_azure_assistants_client.py
+6 -6
View File
@@ -277,7 +277,7 @@ async def test_azure_assistants_client_get_response() -> None:
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
messages.append(ChatMessage("user", ["What's the weather like today?"]))
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
@@ -295,7 +295,7 @@ async def test_azure_assistants_client_get_response_tools() -> None:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
messages.append(ChatMessage("user", ["What's the weather like in Seattle?"]))
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(
@@ -323,7 +323,7 @@ async def test_azure_assistants_client_streaming() -> None:
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
messages.append(ChatMessage("user", ["What's the weather like today?"]))
# Test that the client can be used to get a response
response = azure_assistants_client.get_streaming_response(messages=messages)
@@ -347,7 +347,7 @@ async def test_azure_assistants_client_streaming_tools() -> None:
assert isinstance(azure_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
messages.append(ChatMessage("user", ["What's the weather like in Seattle?"]))
# Test that the client can be used to get a response
response = azure_assistants_client.get_streaming_response(
@@ -372,7 +372,7 @@ async def test_azure_assistants_client_with_existing_assistant() -> None:
# First create an assistant to use in the test
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as temp_client:
# Get the assistant ID by triggering assistant creation
messages = [ChatMessage(role="user", text="Hello")]
messages = [ChatMessage("user", ["Hello"])]
await temp_client.get_response(messages=messages)
assistant_id = temp_client.assistant_id
@@ -383,7 +383,7 @@ async def test_azure_assistants_client_with_existing_assistant() -> None:
assert isinstance(azure_assistants_client, ChatClientProtocol)
assert azure_assistants_client.assistant_id == assistant_id
messages = [ChatMessage(role="user", text="What can you do?")]
messages = [ChatMessage("user", ["What can you do?"])]
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
python/packages/core/tests/azure/test_azure_chat_client.py
+4 -4
View File
@@ -665,7 +665,7 @@ async def test_azure_openai_chat_client_response() -> None:
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
messages.append(ChatMessage("user", ["who are Emily and David?"]))
# Test that the client can be used to get a response
response = await azure_chat_client.get_response(messages=messages)
@@ -686,7 +686,7 @@ async def test_azure_openai_chat_client_response_tools() -> None:
assert isinstance(azure_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
messages.append(ChatMessage("user", ["who are Emily and David?"]))
# Test that the client can be used to get a response
response = await azure_chat_client.get_response(
@@ -716,7 +716,7 @@ async def test_azure_openai_chat_client_streaming() -> None:
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
messages.append(ChatMessage("user", ["who are Emily and David?"]))
# Test that the client can be used to get a response
response = azure_chat_client.get_streaming_response(messages=messages)
@@ -742,7 +742,7 @@ async def test_azure_openai_chat_client_streaming_tools() -> None:
assert isinstance(azure_chat_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
messages.append(ChatMessage("user", ["who are Emily and David?"]))
# Test that the client can be used to get a response
response = azure_chat_client.get_streaming_response(
python/packages/core/tests/azure/test_azure_responses_client.py
+8 -8
View File
@@ -221,14 +221,14 @@ async def test_integration_options(
# Prepare test message
if option_name == "tools" or option_name == "tool_choice":
# Use weather-related prompt for tool tests
messages = [ChatMessage(role="user", text="What is the weather in Seattle?")]
messages = [ChatMessage("user", ["What is the weather in Seattle?"])]
elif option_name == "response_format":
# Use prompt that works well with structured output
messages = [ChatMessage(role="user", text="The weather in Seattle is sunny")]
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
messages = [ChatMessage("user", ["The weather in Seattle is sunny"])]
messages.append(ChatMessage("user", ["What is the weather in Seattle?"]))
else:
# Generic prompt for simple options
messages = [ChatMessage(role="user", text="Say 'Hello World' briefly.")]
messages = [ChatMessage("user", ["Say 'Hello World' briefly."])]
# Build options dict
options: dict[str, Any] = {option_name: option_value}
@@ -245,7 +245,7 @@ async def test_integration_options(
)
output_format = option_value if option_name == "response_format" else None
response = await ChatResponse.from_chat_response_generator(response_gen, output_format_type=output_format)
response = await ChatResponse.from_update_generator(response_gen, output_format_type=output_format)
else:
# Test non-streaming mode
response = await client.get_response(
@@ -293,7 +293,7 @@ async def test_integration_web_search() -> None:
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
response = await ChatResponse.from_update_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
@@ -318,7 +318,7 @@ async def test_integration_web_search() -> None:
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
response = await ChatResponse.from_update_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
assert response.text is not None
@@ -367,7 +367,7 @@ async def test_integration_client_file_search_streaming() -> None:
)
assert response is not None
full_response = await ChatResponse.from_chat_response_generator(response)
full_response = await ChatResponse.from_update_generator(response)
assert "sunny" in full_response.text.lower()
assert "75" in full_response.text
finally:
python/packages/core/tests/core/conftest.py
+10 -7
View File
@@ -21,7 +21,6 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
Role,
ToolProtocol,
tool,
use_chat_middleware,
@@ -95,7 +94,7 @@ class MockChatClient:
self.call_count += 1
if self.responses:
return self.responses.pop(0)
return ChatResponse(messages=ChatMessage(role="assistant", text="test response"))
return ChatResponse(messages=ChatMessage("assistant", ["test response"]))
async def get_streaming_response(
self,
@@ -108,7 +107,7 @@ class MockChatClient:
for update in self.streaming_responses.pop(0):
yield update
else:
yield ChatResponseUpdate(text=Content.from_text(text="test streaming response "), role="assistant")
yield ChatResponseUpdate(contents=[Content.from_text(text="test streaming response ")], role="assistant")
yield ChatResponseUpdate(contents=[Content.from_text(text="another update")], role="assistant")
@@ -143,7 +142,7 @@ class MockBaseChatClient(BaseChatClient[TOptions_co], Generic[TOptions_co]):
logger.debug(f"Running base chat client inner, with: {messages=}, {options=}, {kwargs=}")
self.call_count += 1
if not self.run_responses:
return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[-1].text}"))
return ChatResponse(messages=ChatMessage("assistant", [f"test response - {messages[-1].text}"]))
response = self.run_responses.pop(0)
@@ -168,10 +167,14 @@ class MockBaseChatClient(BaseChatClient[TOptions_co], Generic[TOptions_co]):
) -> AsyncIterable[ChatResponseUpdate]:
logger.debug(f"Running base chat client inner stream, with: {messages=}, {options=}, {kwargs=}")
if not self.streaming_responses:
yield ChatResponseUpdate(text=f"update - {messages[0].text}", role="assistant")
yield ChatResponseUpdate(
contents=[Content.from_text(text=f"update - {messages[0].text}")], role="assistant"
)
return
if options.get("tool_choice") == "none":
yield ChatResponseUpdate(text="I broke out of the function invocation loop...", role="assistant")
yield ChatResponseUpdate(
contents=[Content.from_text(text="I broke out of the function invocation loop...")], role="assistant"
)
return
response = self.streaming_responses.pop(0)
for update in response:
@@ -233,7 +236,7 @@ class MockAgent(AgentProtocol):
**kwargs: Any,
) -> AgentResponse:
logger.debug(f"Running mock agent, with: {messages=}, {thread=}, {kwargs=}")
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, contents=[Content.from_text("Response")])])
return AgentResponse(messages=[ChatMessage("assistant", [Content.from_text("Response")])])
async def run_stream(
self,
python/packages/core/tests/core/test_agents.py
+20 -25
View File
@@ -24,7 +24,6 @@ from agent_framework import (
Context,
ContextProvider,
HostedCodeInterpreterTool,
Role,
ToolProtocol,
tool,
)
@@ -43,7 +42,7 @@ def test_agent_type(agent: AgentProtocol) -> None:
async def test_agent_run(agent: AgentProtocol) -> None:
response = await agent.run("test")
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert response.messages[0].text == "Response"
@@ -104,12 +103,12 @@ async def test_chat_client_agent_get_new_thread(chat_client: ChatClientProtocol)
async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatClientProtocol) -> None:
agent = ChatAgent(chat_client=chat_client)
message = ChatMessage(role=Role.USER, text="Hello")
message = ChatMessage("user", ["Hello"])
thread = AgentThread(message_store=ChatMessageStore(messages=[message]))
_, _, result_messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=thread,
input_messages=[ChatMessage(role=Role.USER, text="Test")],
input_messages=[ChatMessage("user", ["Test"])],
)
assert len(result_messages) == 2
@@ -127,7 +126,7 @@ async def test_prepare_thread_does_not_mutate_agent_chat_options(chat_client: Ch
_, prepared_chat_options, _ = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=thread,
input_messages=[ChatMessage(role=Role.USER, text="Test")],
input_messages=[ChatMessage("user", ["Test"])],
)
assert prepared_chat_options.get("tools") is not None
@@ -139,7 +138,7 @@ async def test_prepare_thread_does_not_mutate_agent_chat_options(chat_client: Ch
async def test_chat_client_agent_update_thread_id(chat_client_base: ChatClientProtocol) -> None:
mock_response = ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, contents=[Content.from_text("test response")])],
messages=[ChatMessage("assistant", [Content.from_text("test response")])],
conversation_id="123",
)
chat_client_base.run_responses = [mock_response]
@@ -202,11 +201,7 @@ async def test_chat_client_agent_author_name_as_agent_name(chat_client: ChatClie
async def test_chat_client_agent_author_name_is_used_from_response(chat_client_base: ChatClientProtocol) -> None:
chat_client_base.run_responses = [
ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT, contents=[Content.from_text("test response")], author_name="TestAuthor"
)
]
messages=[ChatMessage("assistant", [Content.from_text("test response")], author_name="TestAuthor")]
)
]
@@ -256,7 +251,7 @@ class MockContextProvider(ContextProvider):
async def test_chat_agent_context_providers_model_invoking(chat_client: ChatClientProtocol) -> None:
"""Test that context providers' invoking is called during agent run."""
mock_provider = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Test context instructions")])
mock_provider = MockContextProvider(messages=[ChatMessage("system", ["Test context instructions"])])
agent = ChatAgent(chat_client=chat_client, context_provider=mock_provider)
await agent.run("Hello")
@@ -269,7 +264,7 @@ async def test_chat_agent_context_providers_thread_created(chat_client_base: Cha
mock_provider = MockContextProvider()
chat_client_base.run_responses = [
ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, contents=[Content.from_text("test response")])],
messages=[ChatMessage("assistant", [Content.from_text("test response")])],
conversation_id="test-thread-id",
)
]
@@ -296,19 +291,19 @@ async def test_chat_agent_context_providers_messages_adding(chat_client: ChatCli
async def test_chat_agent_context_instructions_in_messages(chat_client: ChatClientProtocol) -> None:
"""Test that AI context instructions are included in messages."""
mock_provider = MockContextProvider(messages=[ChatMessage(role="system", text="Context-specific instructions")])
mock_provider = MockContextProvider(messages=[ChatMessage("system", ["Context-specific instructions"])])
agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions", context_provider=mock_provider)
# We need to test the _prepare_thread_and_messages method directly
_, _, messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=None, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
thread=None, input_messages=[ChatMessage("user", ["Hello"])]
)
# Should have context instructions, and user message
assert len(messages) == 2
assert messages[0].role == Role.SYSTEM
assert messages[0].role == "system"
assert messages[0].text == "Context-specific instructions"
assert messages[1].role == Role.USER
assert messages[1].role == "user"
assert messages[1].text == "Hello"
# instructions system message is added by a chat_client
@@ -319,18 +314,18 @@ async def test_chat_agent_no_context_instructions(chat_client: ChatClientProtoco
agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions", context_provider=mock_provider)
_, _, messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=None, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
thread=None, input_messages=[ChatMessage("user", ["Hello"])]
)
# Should have agent instructions and user message only
assert len(messages) == 1
assert messages[0].role == Role.USER
assert messages[0].role == "user"
assert messages[0].text == "Hello"
async def test_chat_agent_run_stream_context_providers(chat_client: ChatClientProtocol) -> None:
"""Test that context providers work with run_stream method."""
mock_provider = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Stream context instructions")])
mock_provider = MockContextProvider(messages=[ChatMessage("system", ["Stream context instructions"])])
agent = ChatAgent(chat_client=chat_client, context_provider=mock_provider)
# Collect all stream updates
@@ -350,7 +345,7 @@ async def test_chat_agent_context_providers_with_thread_service_id(chat_client_b
mock_provider = MockContextProvider()
chat_client_base.run_responses = [
ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, contents=[Content.from_text("test response")])],
messages=[ChatMessage("assistant", [Content.from_text("test response")])],
conversation_id="service-thread-123",
)
]
@@ -585,7 +580,7 @@ async def test_agent_tool_receives_thread_in_kwargs(chat_client_base: Any) -> No
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
agent = ChatAgent(
@@ -928,7 +923,7 @@ async def test_chat_agent_context_provider_adds_tools_when_agent_has_none(chat_c
# Run the agent and verify context tools are added
_, options, _ = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=None, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
thread=None, input_messages=[ChatMessage("user", ["Hello"])]
)
# The context tools should now be in the options
@@ -952,7 +947,7 @@ async def test_chat_agent_context_provider_adds_instructions_when_agent_has_none
# Run the agent and verify context instructions are available
_, options, _ = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=None, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
thread=None, input_messages=[ChatMessage("user", ["Hello"])]
)
# The context instructions should now be in the options
@@ -972,7 +967,7 @@ async def test_chat_agent_raises_on_conversation_id_mismatch(chat_client_base: C
with pytest.raises(AgentExecutionException, match="conversation_id set on the agent is different"):
await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
thread=thread, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
thread=thread, input_messages=[ChatMessage("user", ["Hello"])]
)
python/packages/core/tests/core/test_as_tool_kwargs_propagation.py
+10 -10
View File
@@ -28,7 +28,7 @@ class TestAsToolKwargsPropagation:
# Setup mock response
chat_client.responses = [
ChatResponse(messages=[ChatMessage(role="assistant", text="Response from sub-agent")]),
ChatResponse(messages=[ChatMessage("assistant", ["Response from sub-agent"])]),
]
# Create sub-agent with middleware
@@ -70,7 +70,7 @@ class TestAsToolKwargsPropagation:
# Setup mock response
chat_client.responses = [
ChatResponse(messages=[ChatMessage(role="assistant", text="Response from sub-agent")]),
ChatResponse(messages=[ChatMessage("assistant", ["Response from sub-agent"])]),
]
sub_agent = ChatAgent(
@@ -122,8 +122,8 @@ class TestAsToolKwargsPropagation:
)
]
),
ChatResponse(messages=[ChatMessage(role="assistant", text="Response from agent_c")]),
ChatResponse(messages=[ChatMessage(role="assistant", text="Response from agent_b")]),
ChatResponse(messages=[ChatMessage("assistant", ["Response from agent_c"])]),
ChatResponse(messages=[ChatMessage("assistant", ["Response from agent_b"])]),
]
# Create agent C (bottom level)
@@ -173,7 +173,7 @@ class TestAsToolKwargsPropagation:
from agent_framework import ChatResponseUpdate
chat_client.streaming_responses = [
[ChatResponseUpdate(text=Content.from_text(text="Streaming response"), role="assistant")],
[ChatResponseUpdate(contents=[Content.from_text(text="Streaming response")], role="assistant")],
]
sub_agent = ChatAgent(
@@ -204,7 +204,7 @@ class TestAsToolKwargsPropagation:
"""Test that as_tool works correctly when no extra kwargs are provided."""
# Setup mock response
chat_client.responses = [
ChatResponse(messages=[ChatMessage(role="assistant", text="Response from agent")]),
ChatResponse(messages=[ChatMessage("assistant", ["Response from agent"])]),
]
sub_agent = ChatAgent(
@@ -233,7 +233,7 @@ class TestAsToolKwargsPropagation:
# Setup mock response
chat_client.responses = [
ChatResponse(messages=[ChatMessage(role="assistant", text="Response with options")]),
ChatResponse(messages=[ChatMessage("assistant", ["Response with options"])]),
]
sub_agent = ChatAgent(
@@ -280,8 +280,8 @@ class TestAsToolKwargsPropagation:
# Setup mock responses for both calls
chat_client.responses = [
ChatResponse(messages=[ChatMessage(role="assistant", text="First response")]),
ChatResponse(messages=[ChatMessage(role="assistant", text="Second response")]),
ChatResponse(messages=[ChatMessage("assistant", ["First response"])]),
ChatResponse(messages=[ChatMessage("assistant", ["Second response"])]),
]
sub_agent = ChatAgent(
@@ -327,7 +327,7 @@ class TestAsToolKwargsPropagation:
# Setup mock response
chat_client.responses = [
ChatResponse(messages=[ChatMessage(role="assistant", text="Response from sub-agent")]),
ChatResponse(messages=[ChatMessage("assistant", ["Response from sub-agent"])]),
]
sub_agent = ChatAgent(
python/packages/core/tests/core/test_clients.py
+11 -12
View File
@@ -7,7 +7,6 @@ from agent_framework import (
BaseChatClient,
ChatClientProtocol,
ChatMessage,
Role,
)
@@ -16,15 +15,15 @@ def test_chat_client_type(chat_client: ChatClientProtocol):
async def test_chat_client_get_response(chat_client: ChatClientProtocol):
response = await chat_client.get_response(ChatMessage(role="user", text="Hello"))
response = await chat_client.get_response(ChatMessage("user", ["Hello"]))
assert response.text == "test response"
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
async def test_chat_client_get_streaming_response(chat_client: ChatClientProtocol):
async for update in chat_client.get_streaming_response(ChatMessage(role="user", text="Hello")):
async for update in chat_client.get_streaming_response(ChatMessage("user", ["Hello"])):
assert update.text == "test streaming response " or update.text == "another update"
assert update.role == Role.ASSISTANT
assert update.role == "assistant"
def test_base_client(chat_client_base: ChatClientProtocol):
@@ -33,13 +32,13 @@ def test_base_client(chat_client_base: ChatClientProtocol):
async def test_base_client_get_response(chat_client_base: ChatClientProtocol):
response = await chat_client_base.get_response(ChatMessage(role="user", text="Hello"))
assert response.messages[0].role == Role.ASSISTANT
response = await chat_client_base.get_response(ChatMessage("user", ["Hello"]))
assert response.messages[0].role == "assistant"
assert response.messages[0].text == "test response - Hello"
async def test_base_client_get_streaming_response(chat_client_base: ChatClientProtocol):
async for update in chat_client_base.get_streaming_response(ChatMessage(role="user", text="Hello")):
async for update in chat_client_base.get_streaming_response(ChatMessage("user", ["Hello"])):
assert update.text == "update - Hello" or update.text == "another update"
@@ -54,17 +53,17 @@ async def test_chat_client_instructions_handling(chat_client_base: ChatClientPro
_, kwargs = mock_inner_get_response.call_args
messages = kwargs.get("messages", [])
assert len(messages) == 1
assert messages[0].role == Role.USER
assert messages[0].role == "user"
assert messages[0].text == "hello"
from agent_framework._types import prepend_instructions_to_messages
appended_messages = prepend_instructions_to_messages(
[ChatMessage(role=Role.USER, text="hello")],
[ChatMessage("user", ["hello"])],
instructions,
)
assert len(appended_messages) == 2
assert appended_messages[0].role == Role.SYSTEM
assert appended_messages[0].role == "system"
assert appended_messages[0].text == "You are a helpful assistant."
assert appended_messages[1].role == Role.USER
assert appended_messages[1].role == "user"
assert appended_messages[1].text == "hello"
python/packages/core/tests/core/test_function_invocation_logic.py
+62 -65
View File
@@ -13,7 +13,6 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
Role,
tool,
)
from agent_framework._middleware import FunctionInvocationContext, FunctionMiddleware
@@ -37,21 +36,21 @@ async def test_base_client_with_function_calling(chat_client_base: ChatClientPro
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]})
assert exec_counter == 1
assert len(response.messages) == 3
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert response.messages[0].contents[0].type == "function_call"
assert response.messages[0].contents[0].name == "test_function"
assert response.messages[0].contents[0].arguments == '{"arg1": "value1"}'
assert response.messages[0].contents[0].call_id == "1"
assert response.messages[1].role == Role.TOOL
assert response.messages[1].role == "tool"
assert response.messages[1].contents[0].type == "function_result"
assert response.messages[1].contents[0].call_id == "1"
assert response.messages[1].contents[0].result == "Processed value1"
assert response.messages[2].role == Role.ASSISTANT
assert response.messages[2].role == "assistant"
assert response.messages[2].text == "done"
@@ -81,16 +80,16 @@ async def test_base_client_with_function_calling_resets(chat_client_base: ChatCl
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]})
assert exec_counter == 2
assert len(response.messages) == 5
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[1].role == Role.TOOL
assert response.messages[2].role == Role.ASSISTANT
assert response.messages[3].role == Role.TOOL
assert response.messages[4].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert response.messages[1].role == "tool"
assert response.messages[2].role == "assistant"
assert response.messages[3].role == "tool"
assert response.messages[4].role == "assistant"
assert response.messages[0].contents[0].type == "function_call"
assert response.messages[1].contents[0].type == "function_result"
assert response.messages[2].contents[0].type == "function_call"
@@ -162,7 +161,7 @@ async def test_function_invocation_inside_aiohttp_server(chat_client_base: ChatC
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
agent = ChatAgent(chat_client=chat_client_base, tools=[ai_func])
@@ -219,7 +218,7 @@ async def test_function_invocation_in_threaded_aiohttp_app(chat_client_base: Cha
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
agent = ChatAgent(chat_client=chat_client_base, tools=[ai_func])
@@ -339,11 +338,11 @@ async def test_function_invocation_scenarios(
# Single function call content
func_call = Content.from_function_call(call_id="1", name=function_name, arguments='{"arg1": "value1"}')
completion = ChatMessage(role="assistant", text="done")
completion = ChatMessage("assistant", ["done"])
chat_client_base.run_responses = [
ChatResponse(messages=ChatMessage(role="assistant", contents=[func_call]))
] + ([] if approval_required else [ChatResponse(messages=completion)])
chat_client_base.run_responses = [ChatResponse(messages=ChatMessage("assistant", [func_call]))] + (
[] if approval_required else [ChatResponse(messages=completion)]
)
chat_client_base.streaming_responses = [
[
@@ -371,7 +370,7 @@ async def test_function_invocation_scenarios(
Content.from_function_call(call_id="2", name="approval_func", arguments='{"arg1": "value2"}'),
]
chat_client_base.run_responses = [ChatResponse(messages=ChatMessage(role="assistant", contents=func_calls))]
chat_client_base.run_responses = [ChatResponse(messages=ChatMessage("assistant", func_calls))]
chat_client_base.streaming_responses = [
[
@@ -432,7 +431,7 @@ async def test_function_invocation_scenarios(
assert messages[0].contents[0].type == "function_call"
assert messages[1].contents[0].type == "function_result"
assert messages[1].contents[0].result == "Processed value1"
assert messages[2].role == Role.ASSISTANT
assert messages[2].role == "assistant"
assert messages[2].text == "done"
assert exec_counter == 1
else:
@@ -497,7 +496,7 @@ async def test_rejected_approval(chat_client_base: ChatClientProtocol):
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Get the response with approval requests
@@ -527,7 +526,7 @@ async def test_rejected_approval(chat_client_base: ChatClientProtocol):
)
# Continue conversation with one approved and one rejected
all_messages = response.messages + [ChatMessage(role="user", contents=[approved_response, rejected_response])]
all_messages = response.messages + [ChatMessage("user", [approved_response, rejected_response])]
# Call get_response which will process the approvals
await chat_client_base.get_response(
@@ -561,9 +560,7 @@ async def test_rejected_approval(chat_client_base: ChatClientProtocol):
for msg in all_messages:
for content in msg.contents:
if content.type == "function_result":
assert msg.role == Role.TOOL, (
f"Message with FunctionResultContent must have role='tool', got '{msg.role}'"
)
assert msg.role == "tool", f"Message with FunctionResultContent must have role='tool', got '{msg.role}'"
async def test_approval_requests_in_assistant_message(chat_client_base: ChatClientProtocol):
@@ -593,7 +590,7 @@ async def test_approval_requests_in_assistant_message(chat_client_base: ChatClie
# Should have one assistant message containing both the call and approval request
assert len(response.messages) == 1
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert len(response.messages[0].contents) == 2
assert response.messages[0].contents[0].type == "function_call"
assert response.messages[0].contents[1].type == "function_approval_request"
@@ -620,7 +617,7 @@ async def test_persisted_approval_messages_replay_correctly(chat_client_base: Ch
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Get approval request
@@ -630,7 +627,7 @@ async def test_persisted_approval_messages_replay_correctly(chat_client_base: Ch
# Store messages (like a thread would)
persisted_messages = [
ChatMessage(role="user", contents=[Content.from_text(text="hello")]),
ChatMessage("user", [Content.from_text(text="hello")]),
*response1.messages,
]
@@ -641,7 +638,7 @@ async def test_persisted_approval_messages_replay_correctly(chat_client_base: Ch
function_call=approval_req.function_call,
approved=True,
)
persisted_messages.append(ChatMessage(role="user", contents=[approval_response]))
persisted_messages.append(ChatMessage("user", [approval_response]))
# Continue with all persisted messages
response2 = await chat_client_base.get_response(
@@ -670,7 +667,7 @@ async def test_no_duplicate_function_calls_after_approval_processing(chat_client
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response1 = await chat_client_base.get_response(
@@ -684,7 +681,7 @@ async def test_no_duplicate_function_calls_after_approval_processing(chat_client
approved=True,
)
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
all_messages = response1.messages + [ChatMessage("user", [approval_response])]
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [func_with_approval]})
# Count function calls with the same call_id
@@ -714,7 +711,7 @@ async def test_rejection_result_uses_function_call_id(chat_client_base: ChatClie
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response1 = await chat_client_base.get_response(
@@ -728,7 +725,7 @@ async def test_rejection_result_uses_function_call_id(chat_client_base: ChatClie
approved=False,
)
all_messages = response1.messages + [ChatMessage(role="user", contents=[rejection_response])]
all_messages = response1.messages + [ChatMessage("user", [rejection_response])]
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [func_with_approval]})
# Find the rejection result
@@ -771,7 +768,7 @@ async def test_max_iterations_limit(chat_client_base: ChatClientProtocol):
)
),
# Failsafe response when tool_choice is set to "none"
ChatResponse(messages=ChatMessage(role="assistant", text="giving up on tools")),
ChatResponse(messages=ChatMessage("assistant", ["giving up on tools"])),
]
# Set max_iterations to 1 in additional_properties
@@ -798,7 +795,7 @@ async def test_function_invocation_config_enabled_false(chat_client_base: ChatCl
return f"Processed {arg1}"
chat_client_base.run_responses = [
ChatResponse(messages=ChatMessage(role="assistant", text="response without function calling")),
ChatResponse(messages=ChatMessage("assistant", ["response without function calling"])),
]
# Disable function invocation
@@ -853,7 +850,7 @@ async def test_function_invocation_config_max_consecutive_errors(chat_client_bas
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="final response")),
ChatResponse(messages=ChatMessage("assistant", ["final response"])),
]
# Set max_consecutive_errors to 2
@@ -898,7 +895,7 @@ async def test_function_invocation_config_terminate_on_unknown_calls_false(chat_
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Set terminate_on_unknown_calls to False (default)
@@ -971,7 +968,7 @@ async def test_function_invocation_config_additional_tools(chat_client_base: Cha
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Add hidden_func to additional_tools
@@ -1010,7 +1007,7 @@ async def test_function_invocation_config_include_detailed_errors_false(chat_cli
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Set include_detailed_errors to False (default)
@@ -1044,7 +1041,7 @@ async def test_function_invocation_config_include_detailed_errors_true(chat_clie
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Set include_detailed_errors to True
@@ -1114,7 +1111,7 @@ async def test_argument_validation_error_with_detailed_errors(chat_client_base:
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Set include_detailed_errors to True
@@ -1148,7 +1145,7 @@ async def test_argument_validation_error_without_detailed_errors(chat_client_bas
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Set include_detailed_errors to False (default)
@@ -1184,12 +1181,12 @@ async def test_hosted_tool_approval_response(chat_client_base: ChatClientProtoco
)
chat_client_base.run_responses = [
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Send the approval response
response = await chat_client_base.get_response(
[ChatMessage(role="user", contents=[approval_response])],
[ChatMessage("user", [approval_response])],
tool_choice="auto",
tools=[local_func],
)
@@ -1215,7 +1212,7 @@ async def test_unapproved_tool_execution_raises_exception(chat_client_base: Chat
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Get approval request
@@ -1231,7 +1228,7 @@ async def test_unapproved_tool_execution_raises_exception(chat_client_base: Chat
)
# Continue conversation with rejection
all_messages = response1.messages + [ChatMessage(role="user", contents=[rejection_response])]
all_messages = response1.messages + [ChatMessage("user", [rejection_response])]
# This should handle the rejection gracefully (not raise ToolException to user)
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [test_func]})
@@ -1270,7 +1267,7 @@ async def test_approved_function_call_with_error_without_detailed_errors(chat_cl
contents=[Content.from_function_call(call_id="1", name="error_func", arguments='{"arg1": "value1"}')],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Set include_detailed_errors to False (default)
@@ -1288,7 +1285,7 @@ async def test_approved_function_call_with_error_without_detailed_errors(chat_cl
approved=True,
)
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
all_messages = response1.messages + [ChatMessage("user", [approval_response])]
# Execute the approved function (which will error)
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [error_func]})
@@ -1333,7 +1330,7 @@ async def test_approved_function_call_with_error_with_detailed_errors(chat_clien
contents=[Content.from_function_call(call_id="1", name="error_func", arguments='{"arg1": "value1"}')],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Set include_detailed_errors to True
@@ -1351,7 +1348,7 @@ async def test_approved_function_call_with_error_with_detailed_errors(chat_clien
approved=True,
)
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
all_messages = response1.messages + [ChatMessage("user", [approval_response])]
# Execute the approved function (which will error)
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [error_func]})
@@ -1396,7 +1393,7 @@ async def test_approved_function_call_with_validation_error(chat_client_base: Ch
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Set include_detailed_errors to True to see validation details
@@ -1414,7 +1411,7 @@ async def test_approved_function_call_with_validation_error(chat_client_base: Ch
approved=True,
)
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
all_messages = response1.messages + [ChatMessage("user", [approval_response])]
# Execute the approved function (which will fail validation)
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [typed_func]})
@@ -1455,7 +1452,7 @@ async def test_approved_function_call_successful_execution(chat_client_base: Cha
contents=[Content.from_function_call(call_id="1", name="success_func", arguments='{"arg1": "value1"}')],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Get approval request
@@ -1470,7 +1467,7 @@ async def test_approved_function_call_successful_execution(chat_client_base: Cha
approved=True,
)
all_messages = response1.messages + [ChatMessage(role="user", contents=[approval_response])]
all_messages = response1.messages + [ChatMessage("user", [approval_response])]
# Execute the approved function
await chat_client_base.get_response(all_messages, options={"tool_choice": "auto", "tools": [success_func]})
@@ -1516,7 +1513,7 @@ async def test_declaration_only_tool(chat_client_base: ChatClientProtocol):
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response = await chat_client_base.get_response(
@@ -1572,7 +1569,7 @@ async def test_multiple_function_calls_parallel_execution(chat_client_base: Chat
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [func1, func2]})
@@ -1608,7 +1605,7 @@ async def test_callable_function_converted_to_tool(chat_client_base: ChatClientP
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
# Pass plain function (will be auto-converted)
@@ -1639,7 +1636,7 @@ async def test_conversation_id_handling(chat_client_base: ChatClientProtocol):
conversation_id="conv_123", # Simulate service-side thread
),
ChatResponse(
messages=ChatMessage(role="assistant", text="done"),
messages=ChatMessage("assistant", ["done"]),
conversation_id="conv_123",
),
]
@@ -1668,7 +1665,7 @@ async def test_function_result_appended_to_existing_assistant_message(chat_clien
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [test_func]})
@@ -1712,7 +1709,7 @@ async def test_error_recovery_resets_counter(chat_client_base: ChatClientProtoco
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [sometimes_fails]})
@@ -2324,7 +2321,7 @@ async def test_terminate_loop_single_function_call(chat_client_base: ChatClientP
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response = await chat_client_base.get_response(
@@ -2339,9 +2336,9 @@ async def test_terminate_loop_single_function_call(chat_client_base: ChatClientP
# There should be 2 messages: assistant with function call, tool result from middleware
# The loop should NOT have continued to call the LLM again
assert len(response.messages) == 2
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert response.messages[0].contents[0].type == "function_call"
assert response.messages[1].role == Role.TOOL
assert response.messages[1].role == "tool"
assert response.messages[1].contents[0].type == "function_result"
assert response.messages[1].contents[0].result == "terminated by middleware"
@@ -2393,7 +2390,7 @@ async def test_terminate_loop_multiple_function_calls_one_terminates(chat_client
],
)
),
ChatResponse(messages=ChatMessage(role="assistant", text="done")),
ChatResponse(messages=ChatMessage("assistant", ["done"])),
]
response = await chat_client_base.get_response(
@@ -2410,9 +2407,9 @@ async def test_terminate_loop_multiple_function_calls_one_terminates(chat_client
# There should be 2 messages: assistant with function calls, tool results
# The loop should NOT have continued to call the LLM again
assert len(response.messages) == 2
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert len(response.messages[0].contents) == 2
assert response.messages[1].role == Role.TOOL
assert response.messages[1].role == "tool"
# Both function results should be present
assert len(response.messages[1].contents) == 2
python/packages/core/tests/core/test_kwargs_propagation_to_ai_function.py
+4 -7
View File
@@ -49,7 +49,7 @@ class TestKwargsPropagationToFunctionTool:
]
)
# Second call: return final response
return ChatResponse(messages=[ChatMessage(role="assistant", text="Done!")])
return ChatResponse(messages=[ChatMessage("assistant", ["Done!"])])
# Wrap the function with function invocation decorator
wrapped = _handle_function_calls_response(mock_get_response)
@@ -101,7 +101,7 @@ class TestKwargsPropagationToFunctionTool:
)
]
)
return ChatResponse(messages=[ChatMessage(role="assistant", text="Completed!")])
return ChatResponse(messages=[ChatMessage("assistant", ["Completed!"])])
wrapped = _handle_function_calls_response(mock_get_response)
@@ -149,7 +149,7 @@ class TestKwargsPropagationToFunctionTool:
)
]
)
return ChatResponse(messages=[ChatMessage(role="assistant", text="All done!")])
return ChatResponse(messages=[ChatMessage("assistant", ["All done!"])])
wrapped = _handle_function_calls_response(mock_get_response)
@@ -196,13 +196,10 @@ class TestKwargsPropagationToFunctionTool:
arguments='{"value": "streaming-test"}',
)
],
is_finished=True,
)
else:
# Second call: return final response
yield ChatResponseUpdate(
text=Content.from_text(text="Stream complete!"), role="assistant", is_finished=True
)
yield ChatResponseUpdate(contents=[Content.from_text(text="Stream complete!")], role="assistant")
wrapped = _handle_function_calls_streaming_response(mock_get_streaming_response)
python/packages/core/tests/core/test_mcp.py
+4 -5
View File
@@ -18,7 +18,6 @@ from agent_framework import (
MCPStdioTool,
MCPStreamableHTTPTool,
MCPWebsocketTool,
Role,
ToolProtocol,
)
from agent_framework._mcp import (
@@ -63,7 +62,7 @@ def test_mcp_prompt_message_to_ai_content():
ai_content = _parse_message_from_mcp(mcp_message)
assert isinstance(ai_content, ChatMessage)
assert ai_content.role.value == "user"
assert ai_content.role == "user"
assert len(ai_content.contents) == 1
assert ai_content.contents[0].type == "text"
assert ai_content.contents[0].text == "Hello, world!"
@@ -1056,7 +1055,7 @@ async def test_local_mcp_server_prompt_execution():
assert len(result) == 1
assert isinstance(result[0], ChatMessage)
assert result[0].role == Role.USER
assert result[0].role == "user"
assert len(result[0].contents) == 1
assert result[0].contents[0].text == "Test message"
@@ -1414,7 +1413,7 @@ async def test_mcp_tool_sampling_callback_chat_client_exception():
async def test_mcp_tool_sampling_callback_no_valid_content():
"""Test sampling callback when response has no valid content types."""
from agent_framework import ChatMessage, Role
from agent_framework import ChatMessage
tool = MCPStdioTool(name="test_tool", command="python")
@@ -1423,7 +1422,7 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
mock_response = Mock()
mock_response.messages = [
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_uri(
uri="data:application/json;base64,e30K",
python/packages/core/tests/core/test_memory.py
+6 -6
View File
@@ -4,7 +4,7 @@ import sys
from collections.abc import MutableSequence
from typing import Any
from agent_framework import ChatMessage, Role
from agent_framework import ChatMessage
from agent_framework._memory import Context, ContextProvider
@@ -69,7 +69,7 @@ class TestContext:
def test_context_with_values(self) -> None:
"""Test Context can be initialized with values."""
messages = [ChatMessage(role=Role.USER, text="Test message")]
messages = [ChatMessage("user", ["Test message"])]
context = Context(instructions="Test instructions", messages=messages)
assert context.instructions == "Test instructions"
assert len(context.messages) == 1
@@ -89,15 +89,15 @@ class TestContextProvider:
async def test_invoked(self) -> None:
"""Test invoked is called."""
provider = MockContextProvider()
message = ChatMessage(role=Role.USER, text="Test message")
message = ChatMessage("user", ["Test message"])
await provider.invoked(message)
assert provider.invoked_called
assert provider.new_messages == message
async def test_invoking(self) -> None:
"""Test invoking is called and returns context."""
provider = MockContextProvider(messages=[ChatMessage(role=Role.USER, text="Context message")])
message = ChatMessage(role=Role.USER, text="Test message")
provider = MockContextProvider(messages=[ChatMessage("user", ["Context message"])])
message = ChatMessage("user", ["Test message"])
context = await provider.invoking(message)
assert provider.invoking_called
assert provider.model_invoking_messages == message
@@ -114,7 +114,7 @@ class TestContextProvider:
async def test_base_invoked_does_nothing(self) -> None:
"""Test that base ContextProvider.invoked does nothing by default."""
provider = MinimalContextProvider()
message = ChatMessage(role=Role.USER, text="Test")
message = ChatMessage("user", ["Test"])
await provider.invoked(message)
await provider.invoked(message, response_messages=message)
await provider.invoked(message, invoke_exception=Exception("test"))
python/packages/core/tests/core/test_middleware.py
+67 -68
View File
@@ -15,7 +15,6 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
Role,
)
from agent_framework._middleware import (
AgentMiddleware,
@@ -36,7 +35,7 @@ class TestAgentRunContext:
def test_init_with_defaults(self, mock_agent: AgentProtocol) -> None:
"""Test AgentRunContext initialization with default values."""
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
assert context.agent is mock_agent
@@ -46,7 +45,7 @@ class TestAgentRunContext:
def test_init_with_custom_values(self, mock_agent: AgentProtocol) -> None:
"""Test AgentRunContext initialization with custom values."""
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
metadata = {"key": "value"}
context = AgentRunContext(agent=mock_agent, messages=messages, is_streaming=True, metadata=metadata)
@@ -59,7 +58,7 @@ class TestAgentRunContext:
"""Test AgentRunContext initialization with thread parameter."""
from agent_framework import AgentThread
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
thread = AgentThread()
context = AgentRunContext(agent=mock_agent, messages=messages, thread=thread)
@@ -98,7 +97,7 @@ class TestChatContext:
def test_init_with_defaults(self, mock_chat_client: Any) -> None:
"""Test ChatContext initialization with default values."""
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
@@ -112,7 +111,7 @@ class TestChatContext:
def test_init_with_custom_values(self, mock_chat_client: Any) -> None:
"""Test ChatContext initialization with custom values."""
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {"temperature": 0.5}
metadata = {"key": "value"}
@@ -169,10 +168,10 @@ class TestAgentMiddlewarePipeline:
async def test_execute_no_middleware(self, mock_agent: AgentProtocol) -> None:
"""Test pipeline execution with no middleware."""
pipeline = AgentMiddlewarePipeline()
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
expected_response = AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
expected_response = AgentResponse(messages=[ChatMessage("assistant", ["response"])])
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
return expected_response
@@ -197,10 +196,10 @@ class TestAgentMiddlewarePipeline:
middleware = OrderTrackingMiddleware("test")
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
expected_response = AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
expected_response = AgentResponse(messages=[ChatMessage("assistant", ["response"])])
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
execution_order.append("handler")
@@ -213,7 +212,7 @@ class TestAgentMiddlewarePipeline:
async def test_execute_stream_no_middleware(self, mock_agent: AgentProtocol) -> None:
"""Test pipeline streaming execution with no middleware."""
pipeline = AgentMiddlewarePipeline()
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentResponseUpdate]:
@@ -245,7 +244,7 @@ class TestAgentMiddlewarePipeline:
middleware = StreamOrderTrackingMiddleware("test")
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentResponseUpdate]:
@@ -267,14 +266,14 @@ class TestAgentMiddlewarePipeline:
"""Test pipeline execution with termination before next()."""
middleware = self.PreNextTerminateMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
execution_order: list[str] = []
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
# Handler should not be executed when terminated before next()
execution_order.append("handler")
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return AgentResponse(messages=[ChatMessage("assistant", ["response"])])
response = await pipeline.execute(mock_agent, messages, context, final_handler)
assert response is not None
@@ -287,13 +286,13 @@ class TestAgentMiddlewarePipeline:
"""Test pipeline execution with termination after next()."""
middleware = self.PostNextTerminateMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
execution_order: list[str] = []
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
execution_order.append("handler")
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return AgentResponse(messages=[ChatMessage("assistant", ["response"])])
response = await pipeline.execute(mock_agent, messages, context, final_handler)
assert response is not None
@@ -306,7 +305,7 @@ class TestAgentMiddlewarePipeline:
"""Test pipeline streaming execution with termination before next()."""
middleware = self.PreNextTerminateMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
execution_order: list[str] = []
@@ -330,7 +329,7 @@ class TestAgentMiddlewarePipeline:
"""Test pipeline streaming execution with termination after next()."""
middleware = self.PostNextTerminateMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
execution_order: list[str] = []
@@ -366,11 +365,11 @@ class TestAgentMiddlewarePipeline:
middleware = ThreadCapturingMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
thread = AgentThread()
context = AgentRunContext(agent=mock_agent, messages=messages, thread=thread)
expected_response = AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
expected_response = AgentResponse(messages=[ChatMessage("assistant", ["response"])])
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
return expected_response
@@ -393,10 +392,10 @@ class TestAgentMiddlewarePipeline:
middleware = ThreadCapturingMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages, thread=None)
expected_response = AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
expected_response = AgentResponse(messages=[ChatMessage("assistant", ["response"])])
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
return expected_response
@@ -560,11 +559,11 @@ class TestChatMiddlewarePipeline:
async def test_execute_no_middleware(self, mock_chat_client: Any) -> None:
"""Test pipeline execution with no middleware."""
pipeline = ChatMiddlewarePipeline()
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
expected_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
expected_response = ChatResponse(messages=[ChatMessage("assistant", ["response"])])
async def final_handler(ctx: ChatContext) -> ChatResponse:
return expected_response
@@ -587,11 +586,11 @@ class TestChatMiddlewarePipeline:
middleware = OrderTrackingChatMiddleware("test")
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
expected_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
expected_response = ChatResponse(messages=[ChatMessage("assistant", ["response"])])
async def final_handler(ctx: ChatContext) -> ChatResponse:
execution_order.append("handler")
@@ -604,7 +603,7 @@ class TestChatMiddlewarePipeline:
async def test_execute_stream_no_middleware(self, mock_chat_client: Any) -> None:
"""Test pipeline streaming execution with no middleware."""
pipeline = ChatMiddlewarePipeline()
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
@@ -635,7 +634,7 @@ class TestChatMiddlewarePipeline:
middleware = StreamOrderTrackingChatMiddleware("test")
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
@@ -658,7 +657,7 @@ class TestChatMiddlewarePipeline:
"""Test pipeline execution with termination before next()."""
middleware = self.PreNextTerminateChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
execution_order: list[str] = []
@@ -666,7 +665,7 @@ class TestChatMiddlewarePipeline:
async def final_handler(ctx: ChatContext) -> ChatResponse:
# Handler should not be executed when terminated before next()
execution_order.append("handler")
return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return ChatResponse(messages=[ChatMessage("assistant", ["response"])])
response = await pipeline.execute(mock_chat_client, messages, chat_options, context, final_handler)
assert response is None
@@ -678,14 +677,14 @@ class TestChatMiddlewarePipeline:
"""Test pipeline execution with termination after next()."""
middleware = self.PostNextTerminateChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
execution_order: list[str] = []
async def final_handler(ctx: ChatContext) -> ChatResponse:
execution_order.append("handler")
return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return ChatResponse(messages=[ChatMessage("assistant", ["response"])])
response = await pipeline.execute(mock_chat_client, messages, chat_options, context, final_handler)
assert response is not None
@@ -698,7 +697,7 @@ class TestChatMiddlewarePipeline:
"""Test pipeline streaming execution with termination before next()."""
middleware = self.PreNextTerminateChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
execution_order: list[str] = []
@@ -723,7 +722,7 @@ class TestChatMiddlewarePipeline:
"""Test pipeline streaming execution with termination after next()."""
middleware = self.PostNextTerminateChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
execution_order: list[str] = []
@@ -764,12 +763,12 @@ class TestClassBasedMiddleware:
middleware = MetadataAgentMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
metadata_updates.append("handler")
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return AgentResponse(messages=[ChatMessage("assistant", ["response"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
@@ -827,12 +826,12 @@ class TestFunctionBasedMiddleware:
execution_order.append("function_after")
pipeline = AgentMiddlewarePipeline([test_agent_middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
execution_order.append("handler")
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return AgentResponse(messages=[ChatMessage("assistant", ["response"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
@@ -890,12 +889,12 @@ class TestMixedMiddleware:
execution_order.append("function_after")
pipeline = AgentMiddlewarePipeline([ClassMiddleware(), function_middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
execution_order.append("handler")
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return AgentResponse(messages=[ChatMessage("assistant", ["response"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
@@ -954,13 +953,13 @@ class TestMixedMiddleware:
execution_order.append("function_after")
pipeline = ChatMiddlewarePipeline([ClassChatMiddleware(), function_chat_middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
async def final_handler(ctx: ChatContext) -> ChatResponse:
execution_order.append("handler")
return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return ChatResponse(messages=[ChatMessage("assistant", ["response"])])
result = await pipeline.execute(mock_chat_client, messages, chat_options, context, final_handler)
@@ -1001,12 +1000,12 @@ class TestMultipleMiddlewareOrdering:
middleware = [FirstMiddleware(), SecondMiddleware(), ThirdMiddleware()]
pipeline = AgentMiddlewarePipeline(middleware) # type: ignore
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
execution_order.append("handler")
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return AgentResponse(messages=[ChatMessage("assistant", ["response"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
@@ -1085,13 +1084,13 @@ class TestMultipleMiddlewareOrdering:
middleware = [FirstChatMiddleware(), SecondChatMiddleware(), ThirdChatMiddleware()]
pipeline = ChatMiddlewarePipeline(middleware) # type: ignore
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
async def final_handler(ctx: ChatContext) -> ChatResponse:
execution_order.append("handler")
return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return ChatResponse(messages=[ChatMessage("assistant", ["response"])])
result = await pipeline.execute(mock_chat_client, messages, chat_options, context, final_handler)
@@ -1127,7 +1126,7 @@ class TestContextContentValidation:
# Verify context content
assert context.agent is mock_agent
assert len(context.messages) == 1
assert context.messages[0].role == Role.USER
assert context.messages[0].role == "user"
assert context.messages[0].text == "test"
assert context.is_streaming is False
assert isinstance(context.metadata, dict)
@@ -1139,13 +1138,13 @@ class TestContextContentValidation:
middleware = ContextValidationMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
# Verify metadata was set by middleware
assert ctx.metadata.get("validated") is True
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return AgentResponse(messages=[ChatMessage("assistant", ["response"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
assert result is not None
@@ -1205,7 +1204,7 @@ class TestContextContentValidation:
# Verify context content
assert context.chat_client is mock_chat_client
assert len(context.messages) == 1
assert context.messages[0].role == Role.USER
assert context.messages[0].role == "user"
assert context.messages[0].text == "test"
assert context.is_streaming is False
assert isinstance(context.metadata, dict)
@@ -1219,14 +1218,14 @@ class TestContextContentValidation:
middleware = ChatContextValidationMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {"temperature": 0.5}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
async def final_handler(ctx: ChatContext) -> ChatResponse:
# Verify metadata was set by middleware
assert ctx.metadata.get("validated") is True
return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return ChatResponse(messages=[ChatMessage("assistant", ["response"])])
result = await pipeline.execute(mock_chat_client, messages, chat_options, context, final_handler)
assert result is not None
@@ -1248,14 +1247,14 @@ class TestStreamingScenarios:
middleware = StreamingFlagMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
# Test non-streaming
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
streaming_flags.append(ctx.is_streaming)
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return AgentResponse(messages=[ChatMessage("assistant", ["response"])])
await pipeline.execute(mock_agent, messages, context, final_handler)
@@ -1287,7 +1286,7 @@ class TestStreamingScenarios:
middleware = StreamProcessingMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_stream_handler(ctx: AgentRunContext) -> AsyncIterable[AgentResponseUpdate]:
@@ -1323,7 +1322,7 @@ class TestStreamingScenarios:
middleware = ChatStreamingFlagMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
# Test non-streaming
@@ -1331,7 +1330,7 @@ class TestStreamingScenarios:
async def final_handler(ctx: ChatContext) -> ChatResponse:
streaming_flags.append(ctx.is_streaming)
return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return ChatResponse(messages=[ChatMessage("assistant", ["response"])])
await pipeline.execute(mock_chat_client, messages, chat_options, context, final_handler)
@@ -1365,7 +1364,7 @@ class TestStreamingScenarios:
middleware = ChatStreamProcessingMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
@@ -1447,7 +1446,7 @@ class TestMiddlewareExecutionControl:
middleware = NoNextMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
handler_called = False
@@ -1455,7 +1454,7 @@ class TestMiddlewareExecutionControl:
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
nonlocal handler_called
handler_called = True
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="should not execute")])
return AgentResponse(messages=[ChatMessage("assistant", ["should not execute"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
@@ -1478,7 +1477,7 @@ class TestMiddlewareExecutionControl:
middleware = NoNextStreamingMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
handler_called = False
@@ -1551,7 +1550,7 @@ class TestMiddlewareExecutionControl:
await next(context)
pipeline = AgentMiddlewarePipeline([FirstMiddleware(), SecondMiddleware()])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
handler_called = False
@@ -1559,7 +1558,7 @@ class TestMiddlewareExecutionControl:
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
nonlocal handler_called
handler_called = True
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="should not execute")])
return AgentResponse(messages=[ChatMessage("assistant", ["should not execute"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
@@ -1580,7 +1579,7 @@ class TestMiddlewareExecutionControl:
middleware = NoNextChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
@@ -1589,7 +1588,7 @@ class TestMiddlewareExecutionControl:
async def final_handler(ctx: ChatContext) -> ChatResponse:
nonlocal handler_called
handler_called = True
return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="should not execute")])
return ChatResponse(messages=[ChatMessage("assistant", ["should not execute"])])
result = await pipeline.execute(mock_chat_client, messages, chat_options, context, final_handler)
@@ -1608,7 +1607,7 @@ class TestMiddlewareExecutionControl:
middleware = NoNextStreamingChatMiddleware()
pipeline = ChatMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options, is_streaming=True)
@@ -1644,7 +1643,7 @@ class TestMiddlewareExecutionControl:
await next(context)
pipeline = ChatMiddlewarePipeline([FirstChatMiddleware(), SecondChatMiddleware()])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
chat_options: dict[str, Any] = {}
context = ChatContext(chat_client=mock_chat_client, messages=messages, options=chat_options)
@@ -1653,7 +1652,7 @@ class TestMiddlewareExecutionControl:
async def final_handler(ctx: ChatContext) -> ChatResponse:
nonlocal handler_called
handler_called = True
return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="should not execute")])
return ChatResponse(messages=[ChatMessage("assistant", ["should not execute"])])
result = await pipeline.execute(mock_chat_client, messages, chat_options, context, final_handler)
python/packages/core/tests/core/test_middleware_context_result.py
+17 -20
View File
@@ -14,7 +14,6 @@ from agent_framework import (
ChatAgent,
ChatMessage,
Content,
Role,
)
from agent_framework._middleware import (
AgentMiddleware,
@@ -40,7 +39,7 @@ class TestResultOverrideMiddleware:
async def test_agent_middleware_response_override_non_streaming(self, mock_agent: AgentProtocol) -> None:
"""Test that agent middleware can override response for non-streaming execution."""
override_response = AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="overridden response")])
override_response = AgentResponse(messages=[ChatMessage("assistant", ["overridden response"])])
class ResponseOverrideMiddleware(AgentMiddleware):
async def process(
@@ -52,7 +51,7 @@ class TestResultOverrideMiddleware:
middleware = ResponseOverrideMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
handler_called = False
@@ -60,7 +59,7 @@ class TestResultOverrideMiddleware:
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
nonlocal handler_called
handler_called = True
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="original response")])
return AgentResponse(messages=[ChatMessage("assistant", ["original response"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
@@ -88,7 +87,7 @@ class TestResultOverrideMiddleware:
middleware = StreamResponseOverrideMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentResponseUpdate]:
@@ -149,7 +148,7 @@ class TestResultOverrideMiddleware:
# Then conditionally override based on content
if any("special" in msg.text for msg in context.messages if msg.text):
context.result = AgentResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Special response from middleware!")]
messages=[ChatMessage("assistant", ["Special response from middleware!"])]
)
# Create ChatAgent with override middleware
@@ -157,14 +156,14 @@ class TestResultOverrideMiddleware:
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
# Test override case
override_messages = [ChatMessage(role=Role.USER, text="Give me a special response")]
override_messages = [ChatMessage("user", ["Give me a special response"])]
override_response = await agent.run(override_messages)
assert override_response.messages[0].text == "Special response from middleware!"
# Verify chat client was called since middleware called next()
assert mock_chat_client.call_count == 1
# Test normal case
normal_messages = [ChatMessage(role=Role.USER, text="Normal request")]
normal_messages = [ChatMessage("user", ["Normal request"])]
normal_response = await agent.run(normal_messages)
assert normal_response.messages[0].text == "test response"
# Verify chat client was called for normal case
@@ -194,7 +193,7 @@ class TestResultOverrideMiddleware:
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
# Test streaming override case
override_messages = [ChatMessage(role=Role.USER, text="Give me a custom stream")]
override_messages = [ChatMessage("user", ["Give me a custom stream"])]
override_updates: list[AgentResponseUpdate] = []
async for update in agent.run_stream(override_messages):
override_updates.append(update)
@@ -205,7 +204,7 @@ class TestResultOverrideMiddleware:
assert override_updates[2].text == " response!"
# Test normal streaming case
normal_messages = [ChatMessage(role=Role.USER, text="Normal streaming request")]
normal_messages = [ChatMessage("user", ["Normal streaming request"])]
normal_updates: list[AgentResponseUpdate] = []
async for update in agent.run_stream(normal_messages):
normal_updates.append(update)
@@ -234,10 +233,10 @@ class TestResultOverrideMiddleware:
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
nonlocal handler_called
handler_called = True
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="executed response")])
return AgentResponse(messages=[ChatMessage("assistant", ["executed response"])])
# Test case where next() is NOT called
no_execute_messages = [ChatMessage(role=Role.USER, text="Don't run this")]
no_execute_messages = [ChatMessage("user", ["Don't run this"])]
no_execute_context = AgentRunContext(agent=mock_agent, messages=no_execute_messages)
no_execute_result = await pipeline.execute(mock_agent, no_execute_messages, no_execute_context, final_handler)
@@ -252,7 +251,7 @@ class TestResultOverrideMiddleware:
handler_called = False
# Test case where next() IS called
execute_messages = [ChatMessage(role=Role.USER, text="Please execute this")]
execute_messages = [ChatMessage("user", ["Please execute this"])]
execute_context = AgentRunContext(agent=mock_agent, messages=execute_messages)
execute_result = await pipeline.execute(mock_agent, execute_messages, execute_context, final_handler)
@@ -332,11 +331,11 @@ class TestResultObservability:
middleware = ObservabilityMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="executed response")])
return AgentResponse(messages=[ChatMessage("assistant", ["executed response"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
@@ -396,17 +395,15 @@ class TestResultObservability:
if "modify" in context.result.messages[0].text:
# Override after observing
context.result = AgentResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="modified after execution")]
)
context.result = AgentResponse(messages=[ChatMessage("assistant", ["modified after execution"])])
middleware = PostExecutionOverrideMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
messages = [ChatMessage("user", ["test"])]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentResponse:
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response to modify")])
return AgentResponse(messages=[ChatMessage("assistant", ["response to modify"])])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
python/packages/core/tests/core/test_middleware_with_agent.py
+75 -80
View File
@@ -15,7 +15,6 @@ from agent_framework import (
ChatResponseUpdate,
Content,
FunctionTool,
Role,
agent_middleware,
chat_middleware,
function_middleware,
@@ -58,13 +57,13 @@ class TestChatAgentClassBasedMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
# Note: conftest "MockChatClient" returns different text format
assert "test response" in response.messages[0].text
@@ -93,7 +92,7 @@ class TestChatAgentClassBasedMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
@@ -128,8 +127,8 @@ class TestChatAgentFunctionBasedMiddleware:
# Execute the agent with multiple messages
messages = [
ChatMessage(role=Role.USER, text="message1"),
ChatMessage(role=Role.USER, text="message2"), # This should not be processed due to termination
ChatMessage("user", ["message1"]),
ChatMessage("user", ["message2"]), # This should not be processed due to termination
]
response = await agent.run(messages)
@@ -158,15 +157,15 @@ class TestChatAgentFunctionBasedMiddleware:
# Execute the agent with multiple messages
messages = [
ChatMessage(role=Role.USER, text="message1"),
ChatMessage(role=Role.USER, text="message2"),
ChatMessage("user", ["message1"]),
ChatMessage("user", ["message2"]),
]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) == 1
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert "test response" in response.messages[0].text
# Verify middleware execution order
@@ -190,7 +189,7 @@ class TestChatAgentFunctionBasedMiddleware:
execution_order.append("middleware_after")
# Create a message to start the conversation
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
# Set up chat client to return a function call, then a final response
# If terminate works correctly, only the first response should be consumed
@@ -198,7 +197,7 @@ class TestChatAgentFunctionBasedMiddleware:
ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="test_call", name="test_function", arguments={"text": "test"}
@@ -207,7 +206,7 @@ class TestChatAgentFunctionBasedMiddleware:
)
]
),
ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="this should not be consumed")]),
ChatResponse(messages=[ChatMessage("assistant", ["this should not be consumed"])]),
]
# Create the test function with the expected signature
@@ -251,7 +250,7 @@ class TestChatAgentFunctionBasedMiddleware:
context.terminate = True
# Create a message to start the conversation
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
# Set up chat client to return a function call, then a final response
# If terminate works correctly, only the first response should be consumed
@@ -259,7 +258,7 @@ class TestChatAgentFunctionBasedMiddleware:
ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="test_call", name="test_function", arguments={"text": "test"}
@@ -268,7 +267,7 @@ class TestChatAgentFunctionBasedMiddleware:
)
]
),
ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="this should not be consumed")]),
ChatResponse(messages=[ChatMessage("assistant", ["this should not be consumed"])]),
]
# Create the test function with the expected signature
@@ -312,13 +311,13 @@ class TestChatAgentFunctionBasedMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[tracking_agent_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert response.messages[0].text == "test response"
assert chat_client.call_count == 1
@@ -340,7 +339,7 @@ class TestChatAgentFunctionBasedMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[tracking_function_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
@@ -376,13 +375,13 @@ class TestChatAgentStreamingMiddleware:
# Set up mock streaming responses
chat_client.streaming_responses = [
[
ChatResponseUpdate(contents=[Content.from_text(text="Streaming")], role=Role.ASSISTANT),
ChatResponseUpdate(contents=[Content.from_text(text=" response")], role=Role.ASSISTANT),
ChatResponseUpdate(contents=[Content.from_text(text="Streaming")], role="assistant"),
ChatResponseUpdate(contents=[Content.from_text(text=" response")], role="assistant"),
]
]
# Execute streaming
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
updates: list[AgentResponseUpdate] = []
async for update in agent.run_stream(messages):
updates.append(update)
@@ -411,7 +410,7 @@ class TestChatAgentStreamingMiddleware:
# Create ChatAgent with middleware
middleware = FlagTrackingMiddleware()
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
# Test non-streaming execution
response = await agent.run(messages)
@@ -452,7 +451,7 @@ class TestChatAgentMultipleMiddlewareOrdering:
agent = ChatAgent(chat_client=chat_client, middleware=[middleware1, middleware2, middleware3])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
@@ -511,7 +510,7 @@ class TestChatAgentMultipleMiddlewareOrdering:
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
@@ -567,7 +566,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="call_123",
@@ -578,7 +577,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Final response"])])
chat_client.responses = [function_call_response, final_response]
@@ -591,7 +590,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for Seattle")]
messages = [ChatMessage("user", ["Get weather for Seattle"])]
response = await agent.run(messages)
# Verify response
@@ -627,7 +626,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="call_456",
@@ -638,7 +637,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Final response"])])
chat_client.responses = [function_call_response, final_response]
@@ -650,7 +649,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for San Francisco")]
messages = [ChatMessage("user", ["Get weather for San Francisco"])]
response = await agent.run(messages)
# Verify response
@@ -699,7 +698,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="call_789",
@@ -710,7 +709,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Final response"])])
chat_client.responses = [function_call_response, final_response]
@@ -722,7 +721,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for New York")]
messages = [ChatMessage("user", ["Get weather for New York"])]
response = await agent.run(messages)
# Verify response
@@ -786,7 +785,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="test_call", name="sample_tool_function", arguments={"location": "Seattle"}
@@ -795,16 +794,14 @@ class TestChatAgentFunctionMiddlewareWithTools:
)
]
),
ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, contents=[Content.from_text("Function completed")])]
),
ChatResponse(messages=[ChatMessage("assistant", [Content.from_text("Function completed")])]),
]
# Create ChatAgent with function middleware
agent = ChatAgent(chat_client=chat_client, middleware=[kwargs_middleware], tools=[sample_tool_function])
# Execute the agent with custom parameters passed as kwargs
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages, custom_param="test_value")
# Verify response
@@ -1068,7 +1065,7 @@ class TestRunLevelMiddleware:
# Verify response is correct
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert "test response" in response.messages[0].text
# Verify middleware was executed
@@ -1097,8 +1094,8 @@ class TestRunLevelMiddleware:
# Set up mock streaming responses
chat_client.streaming_responses = [
[
ChatResponseUpdate(contents=[Content.from_text(text="Stream")], role=Role.ASSISTANT),
ChatResponseUpdate(contents=[Content.from_text(text=" response")], role=Role.ASSISTANT),
ChatResponseUpdate(contents=[Content.from_text(text="Stream")], role="assistant"),
ChatResponseUpdate(contents=[Content.from_text(text=" response")], role="assistant"),
]
]
@@ -1182,7 +1179,7 @@ class TestRunLevelMiddleware:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="test_call",
@@ -1193,7 +1190,7 @@ class TestRunLevelMiddleware:
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Final response"])])
chat_client.responses = [function_call_response, final_response]
# Create agent with agent-level middleware
@@ -1275,7 +1272,7 @@ class TestMiddlewareDecoratorLogic:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="test_call",
@@ -1286,7 +1283,7 @@ class TestMiddlewareDecoratorLogic:
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Final response"])])
chat_client.responses = [function_call_response, final_response]
# Should work without errors
@@ -1296,7 +1293,7 @@ class TestMiddlewareDecoratorLogic:
tools=[custom_tool_wrapped],
)
response = await agent.run([ChatMessage(role=Role.USER, text="test")])
response = await agent.run([ChatMessage("user", ["test"])])
assert response is not None
assert "decorator_type_match_agent" in execution_order
@@ -1317,7 +1314,7 @@ class TestMiddlewareDecoratorLogic:
await next(context)
agent = ChatAgent(chat_client=chat_client, middleware=[mismatched_middleware])
await agent.run([ChatMessage(role=Role.USER, text="test")])
await agent.run([ChatMessage("user", ["test"])])
async def test_only_decorator_specified(self, chat_client: Any) -> None:
"""Only decorator specified - rely on decorator."""
@@ -1346,7 +1343,7 @@ class TestMiddlewareDecoratorLogic:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="test_call",
@@ -1357,7 +1354,7 @@ class TestMiddlewareDecoratorLogic:
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Final response"])])
chat_client.responses = [function_call_response, final_response]
# Should work - relies on decorator
@@ -1367,7 +1364,7 @@ class TestMiddlewareDecoratorLogic:
tools=[custom_tool_wrapped],
)
response = await agent.run([ChatMessage(role=Role.USER, text="test")])
response = await agent.run([ChatMessage("user", ["test"])])
assert response is not None
assert "decorator_only_agent" in execution_order
@@ -1402,7 +1399,7 @@ class TestMiddlewareDecoratorLogic:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="test_call",
@@ -1413,7 +1410,7 @@ class TestMiddlewareDecoratorLogic:
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Final response"])])
chat_client.responses = [function_call_response, final_response]
# Should work - relies on type annotations
@@ -1421,7 +1418,7 @@ class TestMiddlewareDecoratorLogic:
chat_client=chat_client, middleware=[type_only_agent, type_only_function], tools=[custom_tool_wrapped]
)
response = await agent.run([ChatMessage(role=Role.USER, text="test")])
response = await agent.run([ChatMessage("user", ["test"])])
assert response is not None
assert "type_only_agent" in execution_order
@@ -1436,7 +1433,7 @@ class TestMiddlewareDecoratorLogic:
# Should raise MiddlewareException
with pytest.raises(MiddlewareException, match="Cannot determine middleware type"):
agent = ChatAgent(chat_client=chat_client, middleware=[no_info_middleware])
await agent.run([ChatMessage(role=Role.USER, text="test")])
await agent.run([ChatMessage("user", ["test"])])
async def test_insufficient_parameters_error(self, chat_client: Any) -> None:
"""Test that middleware with insufficient parameters raises an error."""
@@ -1450,7 +1447,7 @@ class TestMiddlewareDecoratorLogic:
pass
agent = ChatAgent(chat_client=chat_client, middleware=[insufficient_params_middleware])
await agent.run([ChatMessage(role=Role.USER, text="test")])
await agent.run([ChatMessage("user", ["test"])])
async def test_decorator_markers_preserved(self) -> None:
"""Test that decorator markers are properly set on functions."""
@@ -1523,7 +1520,7 @@ class TestChatAgentThreadBehavior:
thread = agent.get_new_thread()
# First run
first_messages = [ChatMessage(role=Role.USER, text="first message")]
first_messages = [ChatMessage("user", ["first message"])]
first_response = await agent.run(first_messages, thread=thread)
# Verify first response
@@ -1531,7 +1528,7 @@ class TestChatAgentThreadBehavior:
assert len(first_response.messages) > 0
# Second run - use the same thread
second_messages = [ChatMessage(role=Role.USER, text="second message")]
second_messages = [ChatMessage("user", ["second message"])]
second_response = await agent.run(second_messages, thread=thread)
# Verify second response
@@ -1603,13 +1600,13 @@ class TestChatAgentChatMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert "test response" in response.messages[0].text
assert execution_order == ["chat_middleware_before", "chat_middleware_after"]
@@ -1629,13 +1626,13 @@ class TestChatAgentChatMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[tracking_chat_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
assert "test response" in response.messages[0].text
assert execution_order == ["chat_middleware_before", "chat_middleware_after"]
@@ -1649,10 +1646,10 @@ class TestChatAgentChatMiddleware:
# Modify the first message by adding a prefix
if context.messages:
for idx, msg in enumerate(context.messages):
if msg.role.value == "system":
if msg.role == "system":
continue
original_text = msg.text or ""
context.messages[idx] = ChatMessage(role=msg.role, text=f"MODIFIED: {original_text}")
context.messages[idx] = ChatMessage(msg.role, [f"MODIFIED: {original_text}"])
break
await next(context)
@@ -1661,7 +1658,7 @@ class TestChatAgentChatMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[message_modifier_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify that the message was modified (MockBaseChatClient echoes back the input)
@@ -1677,7 +1674,7 @@ class TestChatAgentChatMiddleware:
) -> None:
# Override the response without calling next()
context.result = ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Middleware overridden response")],
messages=[ChatMessage("assistant", ["Middleware overridden response"])],
response_id="middleware-response-123",
)
context.terminate = True
@@ -1687,7 +1684,7 @@ class TestChatAgentChatMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[response_override_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify that the response was overridden
@@ -1717,7 +1714,7 @@ class TestChatAgentChatMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[first_middleware, second_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
@@ -1743,13 +1740,13 @@ class TestChatAgentChatMiddleware:
# Set up mock streaming responses
chat_client.streaming_responses = [
[
ChatResponseUpdate(contents=[Content.from_text(text="Stream")], role=Role.ASSISTANT),
ChatResponseUpdate(contents=[Content.from_text(text=" response")], role=Role.ASSISTANT),
ChatResponseUpdate(contents=[Content.from_text(text="Stream")], role="assistant"),
ChatResponseUpdate(contents=[Content.from_text(text=" response")], role="assistant"),
]
]
# Execute streaming
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
updates: list[AgentResponseUpdate] = []
async for update in agent.run_stream(messages):
updates.append(update)
@@ -1770,9 +1767,7 @@ class TestChatAgentChatMiddleware:
execution_order.append("middleware_before")
context.terminate = True
# Set a custom response since we're terminating
context.result = ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Terminated by middleware")]
)
context.result = ChatResponse(messages=[ChatMessage("assistant", ["Terminated by middleware"])])
# We call next() but since terminate=True, execution should stop
await next(context)
execution_order.append("middleware_after")
@@ -1782,7 +1777,7 @@ class TestChatAgentChatMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[PreTerminationChatMiddleware()])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response was from middleware
@@ -1807,7 +1802,7 @@ class TestChatAgentChatMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[PostTerminationChatMiddleware()])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response is from actual execution
@@ -1843,7 +1838,7 @@ class TestChatAgentChatMiddleware:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="call_456",
@@ -1854,7 +1849,7 @@ class TestChatAgentChatMiddleware:
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Final response"])])
chat_client = use_function_invocation(MockBaseChatClient)()
chat_client.run_responses = [function_call_response, final_response]
@@ -1867,7 +1862,7 @@ class TestChatAgentChatMiddleware:
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for San Francisco")]
messages = [ChatMessage("user", ["Get weather for San Francisco"])]
response = await agent.run(messages)
# Verify response
@@ -1924,7 +1919,7 @@ class TestChatAgentChatMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[kwargs_middleware])
# Execute the agent with custom parameters
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages, temperature=0.7, max_tokens=100, custom_param="test_value")
# Verify response
@@ -1973,7 +1968,7 @@ class TestMiddlewareWithProtocolOnlyAgent:
self.middleware = [TrackingMiddleware()]
async def run(self, messages=None, *, thread=None, **kwargs) -> AgentResponse:
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
return AgentResponse(messages=[ChatMessage("assistant", ["response"])])
def run_stream(self, messages=None, *, thread=None, **kwargs) -> AsyncIterable[AgentResponseUpdate]:
async def _stream():
python/packages/core/tests/core/test_middleware_with_chat.py
+23 -26
View File
@@ -12,7 +12,6 @@ from agent_framework import (
Content,
FunctionInvocationContext,
FunctionTool,
Role,
chat_middleware,
function_middleware,
use_chat_middleware,
@@ -43,13 +42,13 @@ class TestChatMiddleware:
chat_client_base.middleware = [LoggingChatMiddleware()]
# Execute chat client directly
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await chat_client_base.get_response(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
# Verify middleware execution order
assert execution_order == ["chat_middleware_before", "chat_middleware_after"]
@@ -68,13 +67,13 @@ class TestChatMiddleware:
chat_client_base.middleware = [logging_chat_middleware]
# Execute chat client directly
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await chat_client_base.get_response(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
# Verify middleware execution order
assert execution_order == ["function_middleware_before", "function_middleware_after"]
@@ -89,14 +88,14 @@ class TestChatMiddleware:
# Modify the first message by adding a prefix
if context.messages and len(context.messages) > 0:
original_text = context.messages[0].text or ""
context.messages[0] = ChatMessage(role=context.messages[0].role, text=f"MODIFIED: {original_text}")
context.messages[0] = ChatMessage(context.messages[0].role, [f"MODIFIED: {original_text}"])
await next(context)
# Add middleware to chat client
chat_client_base.middleware = [message_modifier_middleware]
# Execute chat client
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await chat_client_base.get_response(messages)
# Verify that the message was modified (MockChatClient echoes back the input)
@@ -114,7 +113,7 @@ class TestChatMiddleware:
) -> None:
# Override the response without calling next()
context.result = ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Middleware overridden response")],
messages=[ChatMessage("assistant", ["Middleware overridden response"])],
response_id="middleware-response-123",
)
context.terminate = True
@@ -123,7 +122,7 @@ class TestChatMiddleware:
chat_client_base.middleware = [response_override_middleware]
# Execute chat client
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await chat_client_base.get_response(messages)
# Verify that the response was overridden
@@ -152,7 +151,7 @@ class TestChatMiddleware:
chat_client_base.middleware = [first_middleware, second_middleware]
# Execute chat client
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await chat_client_base.get_response(messages)
# Verify response
@@ -180,13 +179,13 @@ class TestChatMiddleware:
agent = ChatAgent(chat_client=chat_client, middleware=[agent_level_chat_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].role == "assistant"
# Verify middleware execution order
assert execution_order == ["agent_chat_middleware_before", "agent_chat_middleware_after"]
@@ -211,7 +210,7 @@ class TestChatMiddleware:
agent = ChatAgent(chat_client=chat_client_base, middleware=[first_middleware, second_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await agent.run(messages)
# Verify response
@@ -237,7 +236,7 @@ class TestChatMiddleware:
chat_client_base.middleware = [streaming_middleware]
# Execute streaming response
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
updates: list[object] = []
async for update in chat_client_base.get_streaming_response(messages):
updates.append(update)
@@ -258,19 +257,19 @@ class TestChatMiddleware:
await next(context)
# First call with run-level middleware
messages = [ChatMessage(role=Role.USER, text="first message")]
messages = [ChatMessage("user", ["first message"])]
response1 = await chat_client_base.get_response(messages, middleware=[counting_middleware])
assert response1 is not None
assert execution_count["count"] == 1
# Second call WITHOUT run-level middleware - should not execute the middleware
messages = [ChatMessage(role=Role.USER, text="second message")]
messages = [ChatMessage("user", ["second message"])]
response2 = await chat_client_base.get_response(messages)
assert response2 is not None
assert execution_count["count"] == 1 # Should still be 1, not 2
# Third call with run-level middleware again - should execute
messages = [ChatMessage(role=Role.USER, text="third message")]
messages = [ChatMessage("user", ["third message"])]
response3 = await chat_client_base.get_response(messages, middleware=[counting_middleware])
assert response3 is not None
assert execution_count["count"] == 2 # Should be 2 now
@@ -301,7 +300,7 @@ class TestChatMiddleware:
chat_client_base.middleware = [kwargs_middleware]
# Execute chat client with custom parameters
messages = [ChatMessage(role=Role.USER, text="test message")]
messages = [ChatMessage("user", ["test message"])]
response = await chat_client_base.get_response(
messages, temperature=0.7, max_tokens=100, custom_param="test_value"
)
@@ -355,7 +354,7 @@ class TestChatMiddleware:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="call_1",
@@ -366,14 +365,12 @@ class TestChatMiddleware:
)
]
)
final_response = ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Based on the weather data, it's sunny!")]
)
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Based on the weather data, it's sunny!"])])
chat_client.run_responses = [function_call_response, final_response]
# Execute the chat client directly with tools - this should trigger function invocation and middleware
messages = [ChatMessage(role=Role.USER, text="What's the weather in San Francisco?")]
messages = [ChatMessage("user", ["What's the weather in San Francisco?"])]
response = await chat_client.get_response(messages, options={"tools": [sample_tool_wrapped]})
# Verify response
@@ -418,7 +415,7 @@ class TestChatMiddleware:
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[
Content.from_function_call(
call_id="call_2",
@@ -430,13 +427,13 @@ class TestChatMiddleware:
]
)
final_response = ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="The weather information has been retrieved!")]
messages=[ChatMessage("assistant", ["The weather information has been retrieved!"])]
)
chat_client.run_responses = [function_call_response, final_response]
# Execute the chat client directly with run-level middleware and tools
messages = [ChatMessage(role=Role.USER, text="What's the weather in New York?")]
messages = [ChatMessage("user", ["What's the weather in New York?"])]
response = await chat_client.get_response(
messages, options={"tools": [sample_tool_wrapped]}, middleware=[run_level_function_middleware]
)
python/packages/core/tests/core/test_observability.py
+40 -52
View File
@@ -14,12 +14,13 @@ from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentProtocol,
AgentResponse,
AgentResponseUpdate,
AgentThread,
BaseChatClient,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Role,
Content,
UsageDetails,
prepend_agent_framework_to_user_agent,
tool,
@@ -217,7 +218,7 @@ def mock_chat_client():
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
):
return ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")],
messages=[ChatMessage("assistant", ["Test response"])],
usage_details=UsageDetails(input_token_count=10, output_token_count=20),
finish_reason=None,
)
@@ -225,8 +226,8 @@ def mock_chat_client():
async def _inner_get_streaming_response(
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
):
yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield ChatResponseUpdate(text=" world", role=Role.ASSISTANT)
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")], role="assistant")
yield ChatResponseUpdate(contents=[Content.from_text(text=" world")], role="assistant")
return MockChatClient
@@ -236,7 +237,7 @@ async def test_chat_client_observability(mock_chat_client, span_exporter: InMemo
"""Test that when diagnostics are enabled, telemetry is applied."""
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test message")]
messages = [ChatMessage("user", ["Test message"])]
span_exporter.clear()
response = await client.get_response(messages=messages, model_id="Test")
assert response is not None
@@ -259,7 +260,7 @@ async def test_chat_client_streaming_observability(
):
"""Test streaming telemetry through the use_instrumentation decorator."""
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
span_exporter.clear()
# Collect all yielded updates
updates = []
@@ -288,7 +289,7 @@ async def test_chat_client_observability_with_instructions(
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test message")]
messages = [ChatMessage("user", ["Test message"])]
options = {"model_id": "Test", "instructions": "You are a helpful assistant."}
span_exporter.clear()
response = await client.get_response(messages=messages, options=options)
@@ -317,7 +318,7 @@ async def test_chat_client_streaming_observability_with_instructions(
import json
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
options = {"model_id": "Test", "instructions": "You are a helpful assistant."}
span_exporter.clear()
@@ -344,7 +345,7 @@ async def test_chat_client_observability_without_instructions(
"""Test that system_instructions attribute is not set when instructions are not provided."""
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test message")]
messages = [ChatMessage("user", ["Test message"])]
options = {"model_id": "Test"} # No instructions
span_exporter.clear()
response = await client.get_response(messages=messages, options=options)
@@ -365,7 +366,7 @@ async def test_chat_client_observability_with_empty_instructions(
"""Test that system_instructions attribute is not set when instructions is an empty string."""
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test message")]
messages = [ChatMessage("user", ["Test message"])]
options = {"model_id": "Test", "instructions": ""} # Empty string
span_exporter.clear()
response = await client.get_response(messages=messages, options=options)
@@ -388,7 +389,7 @@ async def test_chat_client_observability_with_list_instructions(
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test message")]
messages = [ChatMessage("user", ["Test message"])]
options = {"model_id": "Test", "instructions": ["Instruction 1", "Instruction 2"]}
span_exporter.clear()
response = await client.get_response(messages=messages, options=options)
@@ -409,7 +410,7 @@ async def test_chat_client_observability_with_list_instructions(
async def test_chat_client_without_model_id_observability(mock_chat_client, span_exporter: InMemorySpanExporter):
"""Test telemetry shouldn't fail when the model_id is not provided for unknown reason."""
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
span_exporter.clear()
response = await client.get_response(messages=messages)
@@ -428,7 +429,7 @@ async def test_chat_client_streaming_without_model_id_observability(
):
"""Test streaming telemetry shouldn't fail when the model_id is not provided for unknown reason."""
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
span_exporter.clear()
# Collect all yielded updates
updates = []
@@ -535,17 +536,16 @@ def mock_chat_agent():
async def run(self, messages=None, *, thread=None, **kwargs):
return AgentResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Agent response")],
messages=[ChatMessage("assistant", ["Agent response"])],
usage_details=UsageDetails(input_token_count=15, output_token_count=25),
response_id="test_response_id",
raw_representation=Mock(finish_reason=Mock(value="stop")),
)
async def run_stream(self, messages=None, *, thread=None, **kwargs):
from agent_framework import AgentResponseUpdate
yield AgentResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield AgentResponseUpdate(text=" from agent", role=Role.ASSISTANT)
yield AgentResponseUpdate(contents=[Content.from_text(text="Hello")], role="assistant")
yield AgentResponseUpdate(contents=[Content.from_text(text=" from agent")], role="assistant")
return MockChatClientAgent
@@ -1338,7 +1338,7 @@ async def test_chat_client_observability_exception(mock_chat_client, span_export
raise ValueError("Test error")
client = use_instrumentation(FailingChatClient)()
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
span_exporter.clear()
with pytest.raises(ValueError, match="Test error"):
@@ -1356,11 +1356,11 @@ async def test_chat_client_streaming_observability_exception(mock_chat_client, s
class FailingStreamingChatClient(mock_chat_client):
async def _inner_get_streaming_response(self, *, messages, options, **kwargs):
yield ChatResponseUpdate(text="Hello", role=Role.ASSISTANT)
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")], role="assistant")
raise ValueError("Streaming error")
client = use_instrumentation(FailingStreamingChatClient)()
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
span_exporter.clear()
with pytest.raises(ValueError, match="Streaming error"):
@@ -1431,12 +1431,11 @@ def test_get_response_attributes_with_finish_reason():
"""Test _get_response_attributes includes finish_reason."""
from unittest.mock import Mock
from agent_framework import FinishReason
from agent_framework.observability import OtelAttr, _get_response_attributes
response = Mock()
response.response_id = None
response.finish_reason = FinishReason.STOP
response.finish_reason = "stop"
response.raw_representation = None
response.usage_details = None
@@ -1608,11 +1607,10 @@ def test_get_response_attributes_finish_reason_from_raw():
"""Test _get_response_attributes gets finish_reason from raw_representation."""
from unittest.mock import Mock
from agent_framework import FinishReason
from agent_framework.observability import OtelAttr, _get_response_attributes
raw_rep = Mock()
raw_rep.finish_reason = FinishReason.LENGTH
raw_rep.finish_reason = "length"
response = Mock()
response.response_id = None
@@ -1668,8 +1666,7 @@ async def test_agent_observability(span_exporter: InMemorySpanExporter, enable_s
**kwargs,
):
return AgentResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")],
thread=thread,
messages=[ChatMessage("assistant", ["Test response"])],
)
async def run_stream(
@@ -1679,9 +1676,8 @@ async def test_agent_observability(span_exporter: InMemorySpanExporter, enable_s
thread=None,
**kwargs,
):
from agent_framework import AgentResponseUpdate
yield AgentResponseUpdate(text="Test", role=Role.ASSISTANT)
yield AgentResponseUpdate(contents=[Content.from_text(text="Test")], role="assistant")
decorated_agent = use_agent_instrumentation(MockAgent)
agent = decorated_agent()
@@ -1697,7 +1693,6 @@ async def test_agent_observability(span_exporter: InMemorySpanExporter, enable_s
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
async def test_agent_observability_with_exception(span_exporter: InMemorySpanExporter, enable_sensitive_data):
"""Test agent instrumentation captures exceptions."""
from agent_framework import AgentResponseUpdate
from agent_framework.observability import use_agent_instrumentation
class FailingAgent(AgentProtocol):
@@ -1730,7 +1725,7 @@ async def test_agent_observability_with_exception(span_exporter: InMemorySpanExp
async def run_stream(self, messages=None, *, thread=None, **kwargs):
# yield before raise to make this an async generator
yield AgentResponseUpdate(text="", role=Role.ASSISTANT)
yield AgentResponseUpdate(contents=[Content.from_text(text="")], role="assistant")
raise RuntimeError("Agent failed")
decorated_agent = use_agent_instrumentation(FailingAgent)
@@ -1751,7 +1746,6 @@ async def test_agent_observability_with_exception(span_exporter: InMemorySpanExp
@pytest.mark.parametrize("enable_sensitive_data", [True, False], indirect=True)
async def test_agent_streaming_observability(span_exporter: InMemorySpanExporter, enable_sensitive_data):
"""Test agent streaming instrumentation."""
from agent_framework import AgentResponseUpdate
from agent_framework.observability import use_agent_instrumentation
class StreamingAgent(AgentProtocol):
@@ -1781,13 +1775,12 @@ async def test_agent_streaming_observability(span_exporter: InMemorySpanExporter
async def run(self, messages=None, *, thread=None, **kwargs):
return AgentResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Test")],
thread=thread,
messages=[ChatMessage("assistant", ["Test"])],
)
async def run_stream(self, messages=None, *, thread=None, **kwargs):
yield AgentResponseUpdate(text="Hello ", role=Role.ASSISTANT)
yield AgentResponseUpdate(text="World", role=Role.ASSISTANT)
yield AgentResponseUpdate(contents=[Content.from_text(text="Hello ")], role="assistant")
yield AgentResponseUpdate(contents=[Content.from_text(text="World")], role="assistant")
decorated_agent = use_agent_instrumentation(StreamingAgent)
agent = decorated_agent()
@@ -1836,24 +1829,22 @@ async def test_capture_messages_with_finish_reason(mock_chat_client, span_export
"""Test that finish_reason is captured in output messages."""
import json
from agent_framework import FinishReason
class ClientWithFinishReason(mock_chat_client):
async def _inner_get_response(self, *, messages, options, **kwargs):
return ChatResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Done")],
messages=[ChatMessage("assistant", ["Done"])],
usage_details=UsageDetails(input_token_count=5, output_token_count=10),
finish_reason=FinishReason.STOP,
finish_reason="stop",
)
client = use_instrumentation(ClientWithFinishReason)()
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
span_exporter.clear()
response = await client.get_response(messages=messages, model_id="Test")
assert response is not None
assert response.finish_reason == FinishReason.STOP
assert response.finish_reason == "stop"
spans = span_exporter.get_finished_spans()
assert len(spans) == 1
span = spans[0]
@@ -1869,7 +1860,6 @@ async def test_capture_messages_with_finish_reason(mock_chat_client, span_export
@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
async def test_agent_streaming_exception(span_exporter: InMemorySpanExporter, enable_sensitive_data):
"""Test agent streaming captures exceptions."""
from agent_framework import AgentResponseUpdate
from agent_framework.observability import use_agent_instrumentation
class FailingStreamingAgent(AgentProtocol):
@@ -1898,10 +1888,10 @@ async def test_agent_streaming_exception(span_exporter: InMemorySpanExporter, en
return self._default_options
async def run(self, messages=None, *, thread=None, **kwargs):
return AgentResponse(messages=[], thread=thread)
return AgentResponse(messages=[])
async def run_stream(self, messages=None, *, thread=None, **kwargs):
yield AgentResponseUpdate(text="Starting", role=Role.ASSISTANT)
yield AgentResponseUpdate(contents=[Content.from_text(text="Starting")], role="assistant")
raise RuntimeError("Stream failed")
decorated_agent = use_agent_instrumentation(FailingStreamingAgent)
@@ -1924,7 +1914,7 @@ async def test_agent_streaming_exception(span_exporter: InMemorySpanExporter, en
async def test_chat_client_when_disabled(mock_chat_client, span_exporter: InMemorySpanExporter):
"""Test that no spans are created when instrumentation is disabled."""
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
span_exporter.clear()
response = await client.get_response(messages=messages, model_id="Test")
@@ -1939,7 +1929,7 @@ async def test_chat_client_when_disabled(mock_chat_client, span_exporter: InMemo
async def test_chat_client_streaming_when_disabled(mock_chat_client, span_exporter: InMemorySpanExporter):
"""Test streaming creates no spans when instrumentation is disabled."""
client = use_instrumentation(mock_chat_client)()
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
span_exporter.clear()
updates = []
@@ -1982,12 +1972,11 @@ async def test_agent_when_disabled(span_exporter: InMemorySpanExporter):
return self._default_options
async def run(self, messages=None, *, thread=None, **kwargs):
return AgentResponse(messages=[], thread=thread)
return AgentResponse(messages=[])
async def run_stream(self, messages=None, *, thread=None, **kwargs):
from agent_framework import AgentResponseUpdate
yield AgentResponseUpdate(text="test", role=Role.ASSISTANT)
yield AgentResponseUpdate(contents=[Content.from_text(text="test")], role="assistant")
decorated = use_agent_instrumentation(TestAgent)
agent = decorated()
@@ -2002,7 +1991,6 @@ async def test_agent_when_disabled(span_exporter: InMemorySpanExporter):
@pytest.mark.parametrize("enable_instrumentation", [False], indirect=True)
async def test_agent_streaming_when_disabled(span_exporter: InMemorySpanExporter):
"""Test agent streaming creates no spans when disabled."""
from agent_framework import AgentResponseUpdate
from agent_framework.observability import use_agent_instrumentation
class TestAgent(AgentProtocol):
@@ -2031,10 +2019,10 @@ async def test_agent_streaming_when_disabled(span_exporter: InMemorySpanExporter
return self._default_options
async def run(self, messages=None, *, thread=None, **kwargs):
return AgentResponse(messages=[], thread=thread)
return AgentResponse(messages=[])
async def run_stream(self, messages=None, *, thread=None, **kwargs):
yield AgentResponseUpdate(text="test", role=Role.ASSISTANT)
yield AgentResponseUpdate(contents=[Content.from_text(text="test")], role="assistant")
decorated = use_agent_instrumentation(TestAgent)
agent = decorated()
python/packages/core/tests/core/test_threads.py
+8 -8
View File
@@ -5,7 +5,7 @@ from typing import Any
import pytest
from agent_framework import AgentThread, ChatMessage, ChatMessageStore, Role
from agent_framework import AgentThread, ChatMessage, ChatMessageStore
from agent_framework._threads import AgentThreadState, ChatMessageStoreState
from agent_framework.exceptions import AgentThreadException
@@ -44,16 +44,16 @@ class MockChatMessageStore:
def sample_messages() -> list[ChatMessage]:
"""Fixture providing sample chat messages for testing."""
return [
ChatMessage(role=Role.USER, text="Hello", message_id="msg1"),
ChatMessage(role=Role.ASSISTANT, text="Hi there!", message_id="msg2"),
ChatMessage(role=Role.USER, text="How are you?", message_id="msg3"),
ChatMessage("user", ["Hello"], message_id="msg1"),
ChatMessage("assistant", ["Hi there!"], message_id="msg2"),
ChatMessage("user", ["How are you?"], message_id="msg3"),
]
@pytest.fixture
def sample_message() -> ChatMessage:
"""Fixture providing a single sample chat message for testing."""
return ChatMessage(role=Role.USER, text="Test message", message_id="test1")
return ChatMessage("user", ["Test message"], message_id="test1")
class TestAgentThread:
@@ -178,7 +178,7 @@ class TestAgentThread:
async def test_on_new_messages_with_existing_store(self, sample_message: ChatMessage) -> None:
"""Test _on_new_messages adds to existing message store."""
initial_messages = [ChatMessage(role=Role.USER, text="Initial", message_id="init1")]
initial_messages = [ChatMessage("user", ["Initial"], message_id="init1")]
store = ChatMessageStore(initial_messages)
thread = AgentThread(message_store=store)
@@ -226,7 +226,7 @@ class TestAgentThread:
thread = AgentThread(message_store=store)
serialized_data: dict[str, Any] = {
"service_thread_id": None,
"chat_message_store_state": {"messages": [ChatMessage(role="user", text="test")]},
"chat_message_store_state": {"messages": [ChatMessage("user", ["test"])]},
}
await thread.update_from_thread_state(serialized_data)
@@ -449,7 +449,7 @@ class TestThreadState:
def test_init_with_chat_message_store_state_object(self) -> None:
"""Test AgentThreadState initialization with ChatMessageStoreState object."""
store_state = ChatMessageStoreState(messages=[ChatMessage(role=Role.USER, text="test")])
store_state = ChatMessageStoreState(messages=[ChatMessage("user", ["test"])])
state = AgentThreadState(chat_message_store_state=store_state)
assert state.service_thread_id is None
python/packages/core/tests/core/test_tools.py
+13 -16
View File
@@ -959,7 +959,7 @@ def mock_chat_client():
return response
# Default response
return ChatResponse(
messages=[ChatMessage(role="assistant", contents=["Default response"])],
messages=[ChatMessage("assistant", ["Default response"])],
)
async def get_streaming_response(self, messages, **kwargs):
@@ -973,7 +973,7 @@ def mock_chat_client():
yield ChatResponseUpdate(contents=[content], role=msg.role)
else:
# Default response
yield ChatResponseUpdate(text="Default response", role="assistant")
yield ChatResponseUpdate(contents=[Content.from_text(text="Default response")], role="assistant")
return MockChatClient()
@@ -1015,7 +1015,7 @@ async def test_non_streaming_single_function_no_approval():
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role="assistant", text="The result is 10")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["The result is 10"])])
call_count = [0]
responses = [initial_response, final_response]
@@ -1100,7 +1100,7 @@ async def test_non_streaming_two_functions_both_no_approval():
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role="assistant", text="Both tools executed successfully")])
final_response = ChatResponse(messages=[ChatMessage("assistant", ["Both tools executed successfully"])])
call_count = [0]
responses = [initial_response, final_response]
@@ -1227,7 +1227,7 @@ async def test_streaming_single_function_no_approval():
role="assistant",
)
]
final_updates = [ChatResponseUpdate(text="The result is 10", role="assistant")]
final_updates = [ChatResponseUpdate(contents=[Content.from_text(text="The result is 10")], role="assistant")]
call_count = [0]
updates_list = [initial_updates, final_updates]
@@ -1246,13 +1246,12 @@ async def test_streaming_single_function_no_approval():
updates.append(update)
# Verify: should have function call update, tool result update (injected), and final update
from agent_framework import Role
assert len(updates) >= 3
# First update is the function call
assert updates[0].contents[0].type == "function_call"
# Second update should be the tool result (injected by the wrapper)
assert updates[1].role == Role.TOOL
assert updates[1].role == "tool"
assert updates[1].contents[0].type == "function_result"
assert updates[1].contents[0].result == 10 # 5 * 2
# Last update is the final message
@@ -1294,11 +1293,10 @@ async def test_streaming_single_function_requires_approval():
updates.append(update)
# Verify: should yield function call and then approval request
from agent_framework import Role
assert len(updates) == 2
assert updates[0].contents[0].type == "function_call"
assert updates[1].role == Role.ASSISTANT
assert updates[1].role == "assistant"
assert updates[1].contents[0].type == "function_approval_request"
@@ -1319,7 +1317,9 @@ async def test_streaming_two_functions_both_no_approval():
role="assistant",
),
]
final_updates = [ChatResponseUpdate(text="Both tools executed successfully", role="assistant")]
final_updates = [
ChatResponseUpdate(contents=[Content.from_text(text="Both tools executed successfully")], role="assistant")
]
call_count = [0]
updates_list = [initial_updates, final_updates]
@@ -1338,7 +1338,6 @@ async def test_streaming_two_functions_both_no_approval():
updates.append(update)
# Verify: should have both function calls, one tool result update with both results, and final message
from agent_framework import Role
assert len(updates) >= 2
# First update has both function calls
@@ -1346,7 +1345,7 @@ async def test_streaming_two_functions_both_no_approval():
assert updates[0].contents[0].type == "function_call"
assert updates[0].contents[1].type == "function_call"
# Should have a tool result update with both results
tool_updates = [u for u in updates if u.role == Role.TOOL]
tool_updates = [u for u in updates if u.role == "tool"]
assert len(tool_updates) == 1
assert len(tool_updates[0].contents) == 2
assert all(c.type == "function_result" for c in tool_updates[0].contents)
@@ -1392,13 +1391,12 @@ async def test_streaming_two_functions_both_require_approval():
updates.append(update)
# Verify: should yield both function calls and then approval requests
from agent_framework import Role
assert len(updates) == 3
assert updates[0].contents[0].type == "function_call"
assert updates[1].contents[0].type == "function_call"
# Assistant update with both approval requests
assert updates[2].role == Role.ASSISTANT
assert updates[2].role == "assistant"
assert len(updates[2].contents) == 2
assert all(c.type == "function_approval_request" for c in updates[2].contents)
@@ -1443,13 +1441,12 @@ async def test_streaming_two_functions_mixed_approval():
updates.append(update)
# Verify: should yield both function calls and then approval requests (when one needs approval, all wait)
from agent_framework import Role
assert len(updates) == 3
assert updates[0].contents[0].type == "function_call"
assert updates[1].contents[0].type == "function_call"
# Assistant update with both approval requests
assert updates[2].role == Role.ASSISTANT
assert updates[2].role == "assistant"
assert len(updates[2].contents) == 2
assert all(c.type == "function_approval_request" for c in updates[2].contents)
python/packages/core/tests/core/test_types.py
+1139 -125
View File
File diff suppressed because it is too large Load Diff
python/packages/core/tests/openai/test_openai_assistants_client.py
+23 -24
View File
@@ -22,7 +22,6 @@ from agent_framework import (
Content,
HostedCodeInterpreterTool,
HostedFileSearchTool,
Role,
tool,
)
from agent_framework.exceptions import ServiceInitializationError
@@ -405,7 +404,7 @@ async def test_process_stream_events_thread_run_created(mock_async_openai: Magic
update = updates[0]
assert isinstance(update, ChatResponseUpdate)
assert update.conversation_id == thread_id
assert update.role == Role.ASSISTANT
assert update.role == "assistant"
assert update.contents == []
assert update.raw_representation == mock_response.data
@@ -449,7 +448,7 @@ async def test_process_stream_events_message_delta_text(mock_async_openai: Magic
update = updates[0]
assert isinstance(update, ChatResponseUpdate)
assert update.conversation_id == thread_id
assert update.role == Role.ASSISTANT
assert update.role == "assistant"
assert update.text == "Hello from assistant"
assert update.raw_representation == mock_message_delta
@@ -488,7 +487,7 @@ async def test_process_stream_events_requires_action(mock_async_openai: MagicMoc
update = updates[0]
assert isinstance(update, ChatResponseUpdate)
assert update.conversation_id == thread_id
assert update.role == Role.ASSISTANT
assert update.role == "assistant"
assert len(update.contents) == 1
assert update.contents[0] == test_function_content
assert update.raw_representation == mock_run
@@ -568,7 +567,7 @@ async def test_process_stream_events_run_completed_with_usage(
update = updates[0]
assert isinstance(update, ChatResponseUpdate)
assert update.conversation_id == thread_id
assert update.role == Role.ASSISTANT
assert update.role == "assistant"
assert len(update.contents) == 1
# Check the usage content
@@ -696,7 +695,7 @@ def test_prepare_options_basic(mock_async_openai: MagicMock) -> None:
"top_p": 0.9,
}
messages = [ChatMessage(role=Role.USER, text="Hello")]
messages = [ChatMessage("user", ["Hello"])]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
@@ -725,7 +724,7 @@ def test_prepare_options_with_tool_tool(mock_async_openai: MagicMock) -> None:
"tool_choice": "auto",
}
messages = [ChatMessage(role=Role.USER, text="Hello")]
messages = [ChatMessage("user", ["Hello"])]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
@@ -750,7 +749,7 @@ def test_prepare_options_with_code_interpreter(mock_async_openai: MagicMock) ->
"tool_choice": "auto",
}
messages = [ChatMessage(role=Role.USER, text="Calculate something")]
messages = [ChatMessage("user", ["Calculate something"])]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
@@ -770,7 +769,7 @@ def test_prepare_options_tool_choice_none(mock_async_openai: MagicMock) -> None:
"tool_choice": "none",
}
messages = [ChatMessage(role=Role.USER, text="Hello")]
messages = [ChatMessage("user", ["Hello"])]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
@@ -791,7 +790,7 @@ def test_prepare_options_required_function(mock_async_openai: MagicMock) -> None
"tool_choice": tool_choice,
}
messages = [ChatMessage(role=Role.USER, text="Hello")]
messages = [ChatMessage("user", ["Hello"])]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
@@ -817,7 +816,7 @@ def test_prepare_options_with_file_search_tool(mock_async_openai: MagicMock) ->
"tool_choice": "auto",
}
messages = [ChatMessage(role=Role.USER, text="Search for information")]
messages = [ChatMessage("user", ["Search for information"])]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
@@ -842,7 +841,7 @@ def test_prepare_options_with_mapping_tool(mock_async_openai: MagicMock) -> None
"tool_choice": "auto",
}
messages = [ChatMessage(role=Role.USER, text="Use custom tool")]
messages = [ChatMessage("user", ["Use custom tool"])]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, options) # type: ignore
@@ -864,7 +863,7 @@ def test_prepare_options_with_pydantic_response_format(mock_async_openai: MagicM
model_config = ConfigDict(extra="forbid")
chat_client = create_test_openai_assistants_client(mock_async_openai)
messages = [ChatMessage(role=Role.USER, text="Test")]
messages = [ChatMessage("user", ["Test"])]
options = {"response_format": TestResponse}
run_options, _ = chat_client._prepare_options(messages, options) # type: ignore
@@ -880,8 +879,8 @@ def test_prepare_options_with_system_message(mock_async_openai: MagicMock) -> No
chat_client = create_test_openai_assistants_client(mock_async_openai)
messages = [
ChatMessage(role=Role.SYSTEM, text="You are a helpful assistant."),
ChatMessage(role=Role.USER, text="Hello"),
ChatMessage("system", ["You are a helpful assistant."]),
ChatMessage("user", ["Hello"]),
]
# Call the method
@@ -901,7 +900,7 @@ def test_prepare_options_with_image_content(mock_async_openai: MagicMock) -> Non
# Create message with image content
image_content = Content.from_uri(uri="https://example.com/image.jpg", media_type="image/jpeg")
messages = [ChatMessage(role=Role.USER, contents=[image_content])]
messages = [ChatMessage("user", [image_content])]
# Call the method
run_options, tool_results = chat_client._prepare_options(messages, {}) # type: ignore
@@ -1021,7 +1020,7 @@ async def test_get_response() -> None:
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
messages.append(ChatMessage("user", ["What's the weather like today?"]))
# Test that the client can be used to get a response
response = await openai_assistants_client.get_response(messages=messages)
@@ -1039,7 +1038,7 @@ async def test_get_response_tools() -> None:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
messages.append(ChatMessage("user", ["What's the weather like in Seattle?"]))
# Test that the client can be used to get a response
response = await openai_assistants_client.get_response(
@@ -1067,7 +1066,7 @@ async def test_streaming() -> None:
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
messages.append(ChatMessage("user", ["What's the weather like today?"]))
# Test that the client can be used to get a response
response = openai_assistants_client.get_streaming_response(messages=messages)
@@ -1091,7 +1090,7 @@ async def test_streaming_tools() -> None:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
messages.append(ChatMessage("user", ["What's the weather like in Seattle?"]))
# Test that the client can be used to get a response
response = openai_assistants_client.get_streaming_response(
@@ -1119,7 +1118,7 @@ async def test_with_existing_assistant() -> None:
# First create an assistant to use in the test
async with OpenAIAssistantsClient(model_id=INTEGRATION_TEST_MODEL) as temp_client:
# Get the assistant ID by triggering assistant creation
messages = [ChatMessage(role="user", text="Hello")]
messages = [ChatMessage("user", ["Hello"])]
await temp_client.get_response(messages=messages)
assistant_id = temp_client.assistant_id
@@ -1130,7 +1129,7 @@ async def test_with_existing_assistant() -> None:
assert isinstance(openai_assistants_client, ChatClientProtocol)
assert openai_assistants_client.assistant_id == assistant_id
messages = [ChatMessage(role="user", text="What can you do?")]
messages = [ChatMessage("user", ["What can you do?"])]
# Test that the client can be used to get a response
response = await openai_assistants_client.get_response(messages=messages)
@@ -1149,7 +1148,7 @@ async def test_file_search() -> None:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
messages.append(ChatMessage("user", ["What's the weather like today?"]))
file_id, vector_store = await create_vector_store(openai_assistants_client)
response = await openai_assistants_client.get_response(
@@ -1175,7 +1174,7 @@ async def test_file_search_streaming() -> None:
assert isinstance(openai_assistants_client, ChatClientProtocol)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
messages.append(ChatMessage("user", ["What's the weather like today?"]))
file_id, vector_store = await create_vector_store(openai_assistants_client)
response = openai_assistants_client.get_streaming_response(
python/packages/core/tests/openai/test_openai_chat_client.py
+17 -17
View File
@@ -154,7 +154,7 @@ def test_serialize_with_org_id(openai_unit_test_env: dict[str, str]) -> None:
async def test_content_filter_exception_handling(openai_unit_test_env: dict[str, str]) -> None:
"""Test that content filter errors are properly handled."""
client = OpenAIChatClient()
messages = [ChatMessage(role="user", text="test message")]
messages = [ChatMessage("user", ["test message"])]
# Create a mock BadRequestError with content_filter code
mock_response = MagicMock()
@@ -209,7 +209,7 @@ def get_weather(location: str) -> str:
async def test_exception_message_includes_original_error_details() -> None:
"""Test that exception messages include original error details in the new format."""
client = OpenAIChatClient(model_id="test-model", api_key="test-key")
messages = [ChatMessage(role="user", text="test message")]
messages = [ChatMessage("user", ["test message"])]
mock_response = MagicMock()
original_error_message = "Invalid API request format"
@@ -652,12 +652,12 @@ def test_function_approval_content_is_skipped_in_preparation(openai_unit_test_en
)
# Test that approval request is skipped
message_with_request = ChatMessage(role="assistant", contents=[approval_request])
message_with_request = ChatMessage("assistant", [approval_request])
prepared_request = client._prepare_message_for_openai(message_with_request)
assert len(prepared_request) == 0 # Should be empty - approval content is skipped
# Test that approval response is skipped
message_with_response = ChatMessage(role="user", contents=[approval_response])
message_with_response = ChatMessage("user", [approval_response])
prepared_response = client._prepare_message_for_openai(message_with_response)
assert len(prepared_response) == 0 # Should be empty - approval content is skipped
@@ -752,7 +752,7 @@ def test_prepare_options_without_model_id(openai_unit_test_env: dict[str, str])
client = OpenAIChatClient()
client.model_id = None # Remove model_id
messages = [ChatMessage(role="user", text="test")]
messages = [ChatMessage("user", ["test"])]
with pytest.raises(ValueError, match="model_id must be a non-empty string"):
client._prepare_options(messages, {})
@@ -786,7 +786,7 @@ def test_prepare_options_with_instructions(openai_unit_test_env: dict[str, str])
"""Test that instructions are prepended as system message."""
client = OpenAIChatClient()
messages = [ChatMessage(role="user", text="Hello")]
messages = [ChatMessage("user", ["Hello"])]
options = {"instructions": "You are a helpful assistant."}
prepared_options = client._prepare_options(messages, options)
@@ -836,7 +836,7 @@ def test_tool_choice_required_with_function_name(openai_unit_test_env: dict[str,
"""Test that tool_choice with required mode and function name is correctly prepared."""
client = OpenAIChatClient()
messages = [ChatMessage(role="user", text="test")]
messages = [ChatMessage("user", ["test"])]
options = {
"tools": [get_weather],
"tool_choice": {"mode": "required", "required_function_name": "get_weather"},
@@ -854,7 +854,7 @@ def test_response_format_dict_passthrough(openai_unit_test_env: dict[str, str])
"""Test that response_format as dict is passed through directly."""
client = OpenAIChatClient()
messages = [ChatMessage(role="user", text="test")]
messages = [ChatMessage("user", ["test"])]
custom_format = {
"type": "json_schema",
"json_schema": {"name": "Test", "schema": {"type": "object"}},
@@ -894,7 +894,7 @@ def test_prepare_options_removes_parallel_tool_calls_when_no_tools(openai_unit_t
"""Test that parallel_tool_calls is removed when no tools are present."""
client = OpenAIChatClient()
messages = [ChatMessage(role="user", text="test")]
messages = [ChatMessage("user", ["test"])]
options = {"allow_multiple_tool_calls": True}
prepared_options = client._prepare_options(messages, options)
@@ -906,7 +906,7 @@ def test_prepare_options_removes_parallel_tool_calls_when_no_tools(openai_unit_t
async def test_streaming_exception_handling(openai_unit_test_env: dict[str, str]) -> None:
"""Test that streaming errors are properly handled."""
client = OpenAIChatClient()
messages = [ChatMessage(role="user", text="test")]
messages = [ChatMessage("user", ["test"])]
# Create a mock error during streaming
mock_error = Exception("Streaming error")
@@ -1008,14 +1008,14 @@ async def test_integration_options(
# Prepare test message
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [ChatMessage(role="user", text="What is the weather in Seattle?")]
messages = [ChatMessage("user", ["What is the weather in Seattle?"])]
elif option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [ChatMessage(role="user", text="The weather in Seattle is sunny")]
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
messages = [ChatMessage("user", ["The weather in Seattle is sunny"])]
messages.append(ChatMessage("user", ["What is the weather in Seattle?"]))
else:
# Generic prompt for simple options
messages = [ChatMessage(role="user", text="Say 'Hello World' briefly.")]
messages = [ChatMessage("user", ["Say 'Hello World' briefly."])]
# Build options dict
options: dict[str, Any] = {option_name: option_value}
@@ -1032,7 +1032,7 @@ async def test_integration_options(
)
output_format = option_value if option_name.startswith("response_format") else None
response = await ChatResponse.from_chat_response_generator(response_gen, output_format_type=output_format)
response = await ChatResponse.from_update_generator(response_gen, output_format_type=output_format)
else:
# Test non-streaming mode
response = await client.get_response(
@@ -1080,7 +1080,7 @@ async def test_integration_web_search() -> None:
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
response = await ChatResponse.from_update_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
@@ -1105,7 +1105,7 @@ async def test_integration_web_search() -> None:
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
response = await ChatResponse.from_update_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
assert response.text is not None
python/packages/core/tests/openai/test_openai_chat_client_base.py
+12 -12
View File
@@ -69,7 +69,7 @@ async def test_cmc(
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
openai_chat_completion = OpenAIChatClient()
await openai_chat_completion.get_response(messages=chat_history)
@@ -88,7 +88,7 @@ async def test_cmc_chat_options(
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
openai_chat_completion = OpenAIChatClient()
await openai_chat_completion.get_response(
@@ -109,7 +109,7 @@ async def test_cmc_no_fcc_in_response(
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
orig_chat_history = deepcopy(chat_history)
openai_chat_completion = OpenAIChatClient()
@@ -131,7 +131,7 @@ async def test_cmc_structured_output_no_fcc(
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
# Define a mock response format
class Test(BaseModel):
@@ -153,7 +153,7 @@ async def test_scmc_chat_options(
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_streaming_chat_completion_response
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
openai_chat_completion = OpenAIChatClient()
async for msg in openai_chat_completion.get_streaming_response(
@@ -178,7 +178,7 @@ async def test_cmc_general_exception(
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
openai_chat_completion = OpenAIChatClient()
with pytest.raises(ServiceResponseException):
@@ -195,7 +195,7 @@ async def test_cmc_additional_properties(
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
openai_chat_completion = OpenAIChatClient()
await openai_chat_completion.get_response(messages=chat_history, options={"reasoning_effort": "low"})
@@ -233,7 +233,7 @@ async def test_get_streaming(
stream = MagicMock(spec=AsyncStream)
stream.__aiter__.return_value = [content1, content2]
mock_create.return_value = stream
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
orig_chat_history = deepcopy(chat_history)
openai_chat_completion = OpenAIChatClient()
@@ -272,7 +272,7 @@ async def test_get_streaming_singular(
stream = MagicMock(spec=AsyncStream)
stream.__aiter__.return_value = [content1, content2]
mock_create.return_value = stream
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
orig_chat_history = deepcopy(chat_history)
openai_chat_completion = OpenAIChatClient()
@@ -311,7 +311,7 @@ async def test_get_streaming_structured_output_no_fcc(
stream = MagicMock(spec=AsyncStream)
stream.__aiter__.return_value = [content1, content2]
mock_create.return_value = stream
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
# Define a mock response format
class Test(BaseModel):
@@ -334,7 +334,7 @@ async def test_get_streaming_no_fcc_in_response(
openai_unit_test_env: dict[str, str],
):
mock_create.return_value = mock_streaming_chat_completion_response
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
orig_chat_history = deepcopy(chat_history)
openai_chat_completion = OpenAIChatClient()
@@ -360,7 +360,7 @@ async def test_get_streaming_no_stream(
mock_chat_completion_response: ChatCompletion, # AsyncStream[ChatCompletionChunk]?
):
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(role="user", text="hello world"))
chat_history.append(ChatMessage("user", ["hello world"]))
openai_chat_completion = OpenAIChatClient()
with pytest.raises(ServiceResponseException):
python/packages/core/tests/openai/test_openai_responses_client.py
+44 -45
View File
@@ -39,7 +39,6 @@ from agent_framework import (
HostedImageGenerationTool,
HostedMCPTool,
HostedWebSearchTool,
Role,
tool,
)
from agent_framework.exceptions import (
@@ -215,7 +214,7 @@ def test_get_response_with_all_parameters() -> None:
with pytest.raises(ServiceResponseException):
asyncio.run(
client.get_response(
messages=[ChatMessage(role="user", text="Test message")],
messages=[ChatMessage("user", ["Test message"])],
options={
"include": ["message.output_text.logprobs"],
"instructions": "You are a helpful assistant",
@@ -261,7 +260,7 @@ def test_web_search_tool_with_location() -> None:
with pytest.raises(ServiceResponseException):
asyncio.run(
client.get_response(
messages=[ChatMessage(role="user", text="What's the weather?")],
messages=[ChatMessage("user", ["What's the weather?"])],
options={"tools": [web_search_tool], "tool_choice": "auto"},
)
)
@@ -278,7 +277,7 @@ def test_file_search_tool_with_invalid_inputs() -> None:
with pytest.raises(ValueError, match="HostedFileSearchTool requires inputs to be of type"):
asyncio.run(
client.get_response(
messages=[ChatMessage(role="user", text="Search files")],
messages=[ChatMessage("user", ["Search files"])],
options={"tools": [file_search_tool]},
)
)
@@ -294,7 +293,7 @@ def test_code_interpreter_tool_variations() -> None:
with pytest.raises(ServiceResponseException):
asyncio.run(
client.get_response(
messages=[ChatMessage(role="user", text="Run some code")],
messages=[ChatMessage("user", ["Run some code"])],
options={"tools": [code_tool_empty]},
)
)
@@ -307,7 +306,7 @@ def test_code_interpreter_tool_variations() -> None:
with pytest.raises(ServiceResponseException):
asyncio.run(
client.get_response(
messages=[ChatMessage(role="user", text="Process these files")],
messages=[ChatMessage("user", ["Process these files"])],
options={"tools": [code_tool_with_files]},
)
)
@@ -327,7 +326,7 @@ def test_content_filter_exception() -> None:
with patch.object(client.client.responses, "create", side_effect=mock_error):
with pytest.raises(OpenAIContentFilterException) as exc_info:
asyncio.run(client.get_response(messages=[ChatMessage(role="user", text="Test message")]))
asyncio.run(client.get_response(messages=[ChatMessage("user", ["Test message"])]))
assert "content error" in str(exc_info.value)
@@ -343,7 +342,7 @@ def test_hosted_file_search_tool_validation() -> None:
with pytest.raises((ValueError, ServiceInvalidRequestError)):
asyncio.run(
client.get_response(
messages=[ChatMessage(role="user", text="Test")],
messages=[ChatMessage("user", ["Test"])],
options={"tools": [empty_file_search_tool]},
)
)
@@ -364,9 +363,9 @@ def test_chat_message_parsing_with_function_calls() -> None:
function_result = Content.from_function_result(call_id="test-call-id", result="Function executed successfully")
messages = [
ChatMessage(role="user", text="Call a function"),
ChatMessage(role="assistant", contents=[function_call]),
ChatMessage(role="tool", contents=[function_result]),
ChatMessage("user", ["Call a function"]),
ChatMessage("assistant", [function_call]),
ChatMessage("tool", [function_result]),
]
# This should exercise the message parsing logic - will fail due to invalid API key
@@ -392,7 +391,7 @@ async def test_response_format_parse_path() -> None:
with patch.object(client.client.responses, "parse", return_value=mock_parsed_response):
response = await client.get_response(
messages=[ChatMessage(role="user", text="Test message")],
messages=[ChatMessage("user", ["Test message"])],
options={"response_format": OutputStruct, "store": True},
)
assert response.response_id == "parsed_response_123"
@@ -419,7 +418,7 @@ async def test_response_format_parse_path_with_conversation_id() -> None:
with patch.object(client.client.responses, "parse", return_value=mock_parsed_response):
response = await client.get_response(
messages=[ChatMessage(role="user", text="Test message")],
messages=[ChatMessage("user", ["Test message"])],
options={"response_format": OutputStruct, "store": True},
)
assert response.response_id == "parsed_response_123"
@@ -442,7 +441,7 @@ async def test_bad_request_error_non_content_filter() -> None:
with patch.object(client.client.responses, "parse", side_effect=mock_error):
with pytest.raises(ServiceResponseException) as exc_info:
await client.get_response(
messages=[ChatMessage(role="user", text="Test message")],
messages=[ChatMessage("user", ["Test message"])],
options={"response_format": OutputStruct},
)
@@ -463,7 +462,7 @@ async def test_streaming_content_filter_exception_handling() -> None:
mock_create.side_effect.code = "content_filter"
with pytest.raises(OpenAIContentFilterException, match="service encountered a content error"):
response_stream = client.get_streaming_response(messages=[ChatMessage(role="user", text="Test")])
response_stream = client.get_streaming_response(messages=[ChatMessage("user", ["Test"])])
async for _ in response_stream:
break
@@ -658,7 +657,7 @@ def test_prepare_content_for_opentool_approval_response() -> None:
function_call=function_call,
)
result = client._prepare_content_for_openai(Role.ASSISTANT, approval_response, {})
result = client._prepare_content_for_openai("assistant", approval_response, {})
assert result["type"] == "mcp_approval_response"
assert result["approval_request_id"] == "approval_001"
@@ -675,7 +674,7 @@ def test_prepare_content_for_openai_error_content() -> None:
error_details="Invalid parameter",
)
result = client._prepare_content_for_openai(Role.ASSISTANT, error_content, {})
result = client._prepare_content_for_openai("assistant", error_content, {})
# ErrorContent should return empty dict (logged but not sent)
assert result == {}
@@ -693,7 +692,7 @@ def test_prepare_content_for_openai_usage_content() -> None:
}
)
result = client._prepare_content_for_openai(Role.ASSISTANT, usage_content, {})
result = client._prepare_content_for_openai("assistant", usage_content, {})
# UsageContent should return empty dict (logged but not sent)
assert result == {}
@@ -707,7 +706,7 @@ def test_prepare_content_for_openai_hosted_vector_store_content() -> None:
vector_store_id="vs_123",
)
result = client._prepare_content_for_openai(Role.ASSISTANT, vector_store_content, {})
result = client._prepare_content_for_openai("assistant", vector_store_content, {})
# HostedVectorStoreContent should return empty dict (logged but not sent)
assert result == {}
@@ -807,7 +806,7 @@ def test_prepare_message_for_openai_with_function_approval_response() -> None:
function_call=function_call,
)
message = ChatMessage(role="user", contents=[approval_response])
message = ChatMessage("user", [approval_response])
call_id_to_id: dict[str, str] = {}
result = client._prepare_message_for_openai(message, call_id_to_id)
@@ -829,7 +828,7 @@ def test_chat_message_with_error_content() -> None:
error_code="TEST_ERR",
)
message = ChatMessage(role="assistant", contents=[error_content])
message = ChatMessage("assistant", [error_content])
call_id_to_id: dict[str, str] = {}
result = client._prepare_message_for_openai(message, call_id_to_id)
@@ -854,7 +853,7 @@ def test_chat_message_with_usage_content() -> None:
}
)
message = ChatMessage(role="assistant", contents=[usage_content])
message = ChatMessage("assistant", [usage_content])
call_id_to_id: dict[str, str] = {}
result = client._prepare_message_for_openai(message, call_id_to_id)
@@ -877,7 +876,7 @@ def test_hosted_file_content_preparation() -> None:
name="document.pdf",
)
result = client._prepare_content_for_openai(Role.USER, hosted_file, {})
result = client._prepare_content_for_openai("user", hosted_file, {})
assert result["type"] == "input_file"
assert result["file_id"] == "file_abc123"
@@ -900,7 +899,7 @@ def test_function_approval_response_with_mcp_tool_call() -> None:
function_call=mcp_call,
)
result = client._prepare_content_for_openai(Role.ASSISTANT, approval_response, {})
result = client._prepare_content_for_openai("assistant", approval_response, {})
assert result["type"] == "mcp_approval_response"
assert result["approval_request_id"] == "approval_mcp_001"
@@ -1358,14 +1357,14 @@ async def test_end_to_end_mcp_approval_flow(span_exporter) -> None:
# Patch the create call to return the two mocked responses in sequence
with patch.object(client.client.responses, "create", side_effect=[mock_response1, mock_response2]) as mock_create:
# First call: get the approval request
response = await client.get_response(messages=[ChatMessage(role="user", text="Trigger approval")])
response = await client.get_response(messages=[ChatMessage("user", ["Trigger approval"])])
assert response.messages[0].contents[0].type == "function_approval_request"
req = response.messages[0].contents[0]
assert req.id == "approval-1"
# Build a user approval and send it (include required function_call)
approval = Content.from_function_approval_response(approved=True, id=req.id, function_call=req.function_call)
approval_message = ChatMessage(role="user", contents=[approval])
approval_message = ChatMessage("user", [approval])
_ = await client.get_response(messages=[approval_message])
# Ensure two calls were made and the second includes the mcp_approval_response
@@ -1469,7 +1468,7 @@ def test_streaming_response_basic_structure() -> None:
# Should get a valid ChatResponseUpdate structure
assert isinstance(response, ChatResponseUpdate)
assert response.role == Role.ASSISTANT
assert response.role == "assistant"
assert response.model_id == "test-model"
assert isinstance(response.contents, list)
assert response.raw_representation is mock_event
@@ -1620,7 +1619,7 @@ def test_streaming_annotation_added_with_unknown_type() -> None:
def test_service_response_exception_includes_original_error_details() -> None:
"""Test that ServiceResponseException messages include original error details in the new format."""
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
messages = [ChatMessage(role="user", text="test message")]
messages = [ChatMessage("user", ["test message"])]
mock_response = MagicMock()
original_error_message = "Request rate limit exceeded"
@@ -1645,7 +1644,7 @@ def test_service_response_exception_includes_original_error_details() -> None:
def test_get_streaming_response_with_response_format() -> None:
"""Test get_streaming_response with response_format."""
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
messages = [ChatMessage(role="user", text="Test streaming with format")]
messages = [ChatMessage("user", ["Test streaming with format"])]
# It will fail due to invalid API key, but exercises the code path
with pytest.raises(ServiceResponseException):
@@ -1667,7 +1666,7 @@ def test_prepare_content_for_openai_image_content() -> None:
media_type="image/jpeg",
additional_properties={"detail": "high", "file_id": "file_123"},
)
result = client._prepare_content_for_openai(Role.USER, image_content_with_detail, {}) # type: ignore
result = client._prepare_content_for_openai("user", image_content_with_detail, {}) # type: ignore
assert result["type"] == "input_image"
assert result["image_url"] == "https://example.com/image.jpg"
assert result["detail"] == "high"
@@ -1675,7 +1674,7 @@ def test_prepare_content_for_openai_image_content() -> None:
# Test image content without additional properties (defaults)
image_content_basic = Content.from_uri(uri="https://example.com/basic.png", media_type="image/png")
result = client._prepare_content_for_openai(Role.USER, image_content_basic, {}) # type: ignore
result = client._prepare_content_for_openai("user", image_content_basic, {}) # type: ignore
assert result["type"] == "input_image"
assert result["detail"] == "auto"
assert result["file_id"] is None
@@ -1687,14 +1686,14 @@ def test_prepare_content_for_openai_audio_content() -> None:
# Test WAV audio content
wav_content = Content.from_uri(uri="data:audio/wav;base64,abc123", media_type="audio/wav")
result = client._prepare_content_for_openai(Role.USER, wav_content, {}) # type: ignore
result = client._prepare_content_for_openai("user", wav_content, {}) # type: ignore
assert result["type"] == "input_audio"
assert result["input_audio"]["data"] == "data:audio/wav;base64,abc123"
assert result["input_audio"]["format"] == "wav"
# Test MP3 audio content
mp3_content = Content.from_uri(uri="data:audio/mp3;base64,def456", media_type="audio/mp3")
result = client._prepare_content_for_openai(Role.USER, mp3_content, {}) # type: ignore
result = client._prepare_content_for_openai("user", mp3_content, {}) # type: ignore
assert result["type"] == "input_audio"
assert result["input_audio"]["format"] == "mp3"
@@ -1705,12 +1704,12 @@ def test_prepare_content_for_openai_unsupported_content() -> None:
# Test unsupported audio format
unsupported_audio = Content.from_uri(uri="data:audio/ogg;base64,ghi789", media_type="audio/ogg")
result = client._prepare_content_for_openai(Role.USER, unsupported_audio, {}) # type: ignore
result = client._prepare_content_for_openai("user", unsupported_audio, {}) # type: ignore
assert result == {}
# Test non-media content
text_uri_content = Content.from_uri(uri="https://example.com/document.txt", media_type="text/plain")
result = client._prepare_content_for_openai(Role.USER, text_uri_content, {}) # type: ignore
result = client._prepare_content_for_openai("user", text_uri_content, {}) # type: ignore
assert result == {}
@@ -1775,7 +1774,7 @@ def test_prepare_content_for_openai_text_reasoning_comprehensive() -> None:
"encrypted_content": "secure_data_456",
},
)
result = client._prepare_content_for_openai(Role.ASSISTANT, comprehensive_reasoning, {}) # type: ignore
result = client._prepare_content_for_openai("assistant", comprehensive_reasoning, {}) # type: ignore
assert result["type"] == "reasoning"
assert result["summary"]["text"] == "Comprehensive reasoning summary"
assert result["status"] == "in_progress"
@@ -2091,7 +2090,7 @@ def test_parse_response_from_openai_image_generation_fallback():
async def test_prepare_options_store_parameter_handling() -> None:
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
messages = [ChatMessage(role="user", text="Test message")]
messages = [ChatMessage("user", ["Test message"])]
test_conversation_id = "test-conversation-123"
chat_options = ChatOptions(store=True, conversation_id=test_conversation_id)
@@ -2117,7 +2116,7 @@ async def test_prepare_options_store_parameter_handling() -> None:
async def test_conversation_id_precedence_kwargs_over_options() -> None:
"""When both kwargs and options contain conversation_id, kwargs wins."""
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
messages = [ChatMessage(role="user", text="Hello")]
messages = [ChatMessage("user", ["Hello"])]
# options has a stale response id, kwargs carries the freshest one
opts = {"conversation_id": "resp_old_123"}
@@ -2224,14 +2223,14 @@ async def test_integration_options(
# Prepare test message
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [ChatMessage(role="user", text="What is the weather in Seattle?")]
messages = [ChatMessage("user", ["What is the weather in Seattle?"])]
elif option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [ChatMessage(role="user", text="The weather in Seattle is sunny")]
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
messages = [ChatMessage("user", ["The weather in Seattle is sunny"])]
messages.append(ChatMessage("user", ["What is the weather in Seattle?"]))
else:
# Generic prompt for simple options
messages = [ChatMessage(role="user", text="Say 'Hello World' briefly.")]
messages = [ChatMessage("user", ["Say 'Hello World' briefly."])]
# Build options dict
options: dict[str, Any] = {option_name: option_value}
@@ -2248,7 +2247,7 @@ async def test_integration_options(
)
output_format = option_value if option_name.startswith("response_format") else None
response = await ChatResponse.from_chat_response_generator(response_gen, output_format_type=output_format)
response = await ChatResponse.from_update_generator(response_gen, output_format_type=output_format)
else:
# Test non-streaming mode
response = await openai_responses_client.get_response(
@@ -2296,7 +2295,7 @@ async def test_integration_web_search() -> None:
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
response = await ChatResponse.from_update_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
@@ -2321,7 +2320,7 @@ async def test_integration_web_search() -> None:
},
}
if streaming:
response = await ChatResponse.from_chat_response_generator(client.get_streaming_response(**content))
response = await ChatResponse.from_update_generator(client.get_streaming_response(**content))
else:
response = await client.get_response(**content)
assert response.text is not None
python/packages/core/tests/workflow/test_agent_executor.py
+8 -11
View File
@@ -12,7 +12,6 @@ from agent_framework import (
ChatMessage,
ChatMessageStore,
Content,
Role,
SequentialBuilder,
WorkflowOutputEvent,
WorkflowRunState,
@@ -37,9 +36,7 @@ class _CountingAgent(BaseAgent):
**kwargs: Any,
) -> AgentResponse:
self.call_count += 1
return AgentResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text=f"Response #{self.call_count}: {self.name}")]
)
return AgentResponse(messages=[ChatMessage("assistant", [f"Response #{self.call_count}: {self.name}"])])
async def run_stream( # type: ignore[override]
self,
@@ -62,8 +59,8 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
# Add some initial messages to the thread to verify thread state persistence
initial_messages = [
ChatMessage(role=Role.USER, text="Initial message 1"),
ChatMessage(role=Role.ASSISTANT, text="Initial response 1"),
ChatMessage("user", ["Initial message 1"]),
ChatMessage("assistant", ["Initial response 1"]),
]
await initial_thread.on_new_messages(initial_messages)
@@ -166,9 +163,9 @@ async def test_agent_executor_save_and_restore_state_directly() -> None:
# Add messages to thread
thread_messages = [
ChatMessage(role=Role.USER, text="Message in thread 1"),
ChatMessage(role=Role.ASSISTANT, text="Thread response 1"),
ChatMessage(role=Role.USER, text="Message in thread 2"),
ChatMessage("user", ["Message in thread 1"]),
ChatMessage("assistant", ["Thread response 1"]),
ChatMessage("user", ["Message in thread 2"]),
]
await thread.on_new_messages(thread_messages)
@@ -176,8 +173,8 @@ async def test_agent_executor_save_and_restore_state_directly() -> None:
# Add messages to executor cache
cache_messages = [
ChatMessage(role=Role.USER, text="Cached user message"),
ChatMessage(role=Role.ASSISTANT, text="Cached assistant response"),
ChatMessage("user", ["Cached user message"]),
ChatMessage("assistant", ["Cached assistant response"]),
]
executor._cache = list(cache_messages) # type: ignore[reportPrivateUsage]
python/packages/core/tests/workflow/test_agent_executor_tool_calls.py
+7 -8
View File
@@ -21,7 +21,6 @@ from agent_framework import (
ChatResponseUpdate,
Content,
RequestInfoEvent,
Role,
WorkflowBuilder,
WorkflowContext,
WorkflowOutputEvent,
@@ -45,7 +44,7 @@ class _ToolCallingAgent(BaseAgent):
**kwargs: Any,
) -> AgentResponse:
"""Non-streaming run - not used in this test."""
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="done")])
return AgentResponse(messages=[ChatMessage("assistant", ["done"])])
async def run_stream(
self,
@@ -58,7 +57,7 @@ class _ToolCallingAgent(BaseAgent):
# First update: some text
yield AgentResponseUpdate(
contents=[Content.from_text(text="Let me search for that...")],
role=Role.ASSISTANT,
role="assistant",
)
# Second update: tool call (no text!)
@@ -70,7 +69,7 @@ class _ToolCallingAgent(BaseAgent):
arguments={"query": "weather"},
)
],
role=Role.ASSISTANT,
role="assistant",
)
# Third update: tool result (no text!)
@@ -81,13 +80,13 @@ class _ToolCallingAgent(BaseAgent):
result={"temperature": 72, "condition": "sunny"},
)
],
role=Role.TOOL,
role="tool",
)
# Fourth update: final text response
yield AgentResponseUpdate(
contents=[Content.from_text(text="The weather is sunny, 72°F.")],
role=Role.ASSISTANT,
role="assistant",
)
@@ -179,7 +178,7 @@ class MockChatClient:
)
)
else:
response = ChatResponse(messages=ChatMessage(role="assistant", text="Tool executed successfully."))
response = ChatResponse(messages=ChatMessage("assistant", ["Tool executed successfully."]))
self._iteration += 1
return response
@@ -212,7 +211,7 @@ class MockChatClient:
role="assistant",
)
else:
yield ChatResponseUpdate(text=Content.from_text(text="Tool executed "), role="assistant")
yield ChatResponseUpdate(contents=[Content.from_text(text="Tool executed ")], role="assistant")
yield ChatResponseUpdate(contents=[Content.from_text(text="successfully.")], role="assistant")
self._iteration += 1
python/packages/core/tests/workflow/test_agent_run_event_typing.py
+2 -2
View File
@@ -2,13 +2,13 @@
"""Tests for AgentRunEvent and AgentRunUpdateEvent type annotations."""
from agent_framework import AgentResponse, AgentResponseUpdate, ChatMessage, Role
from agent_framework import AgentResponse, AgentResponseUpdate, ChatMessage
from agent_framework._workflows._events import AgentRunEvent, AgentRunUpdateEvent
def test_agent_run_event_data_type() -> None:
"""Verify AgentRunEvent.data is typed as AgentResponse | None."""
response = AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Hello")])
response = AgentResponse(messages=[ChatMessage("assistant", ["Hello"])])
event = AgentRunEvent(executor_id="test", data=response)
# This assignment should pass type checking without a cast
python/packages/core/tests/workflow/test_concurrent.py
+5 -6
View File
@@ -12,7 +12,6 @@ from agent_framework import (
ChatMessage,
ConcurrentBuilder,
Executor,
Role,
WorkflowContext,
WorkflowOutputEvent,
WorkflowRunState,
@@ -36,7 +35,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=ChatMessage("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))
@@ -126,12 +125,12 @@ async def test_concurrent_default_aggregator_emits_single_user_and_assistants()
# Expect one user message + one assistant message per participant
assert len(messages) == 1 + 3
assert messages[0].role == Role.USER
assert messages[0].role == "user"
assert "hello world" in messages[0].text
assistant_texts = {m.text for m in messages[1:]}
assert assistant_texts == {"Alpha", "Beta", "Gamma"}
assert all(m.role == Role.ASSISTANT for m in messages[1:])
assert all(m.role == "assistant" for m in messages[1:])
async def test_concurrent_custom_aggregator_callback_is_used() -> None:
@@ -543,9 +542,9 @@ async def test_concurrent_with_register_participants() -> None:
# Expect one user message + one assistant message per participant
assert len(messages) == 1 + 3
assert messages[0].role == Role.USER
assert messages[0].role == "user"
assert "test prompt" in messages[0].text
assistant_texts = {m.text for m in messages[1:]}
assert assistant_texts == {"Alpha", "Beta", "Gamma"}
assert all(m.role == Role.ASSISTANT for m in messages[1:])
assert all(m.role == "assistant" for m in messages[1:])
python/packages/core/tests/workflow/test_executor.py
+2 -2
View File
@@ -537,7 +537,7 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
async def mutator(messages: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None:
# The handler mutates the input list by appending new messages
original_len = len(messages)
messages.append(ChatMessage(role="assistant", text="Added by executor"))
messages.append(ChatMessage("assistant", ["Added by executor"]))
await ctx.send_message(messages)
# Verify mutation happened
assert len(messages) == original_len + 1
@@ -545,7 +545,7 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
workflow = WorkflowBuilder().set_start_executor(mutator).build()
# Run with a single user message
input_messages = [ChatMessage(role="user", text="hello")]
input_messages = [ChatMessage("user", ["hello"])]
events = await workflow.run(input_messages)
# Find the invoked event for the Mutator executor
python/packages/core/tests/workflow/test_full_conversation.py
+8 -9
View File
@@ -16,7 +16,6 @@ from agent_framework import (
ChatMessage,
Content,
Executor,
Role,
SequentialBuilder,
WorkflowBuilder,
WorkflowContext,
@@ -40,7 +39,7 @@ class _SimpleAgent(BaseAgent):
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text=self._reply_text)])
return AgentResponse(messages=[ChatMessage("assistant", [self._reply_text])])
async def run_stream( # type: ignore[override]
self,
@@ -89,8 +88,8 @@ async def test_agent_executor_populates_full_conversation_non_streaming() -> Non
# Assert: full_conversation contains [user("hello world"), assistant("agent-reply")]
assert isinstance(payload, dict)
assert payload["length"] == 2
assert payload["roles"][0] == Role.USER and "hello world" in (payload["texts"][0] or "")
assert payload["roles"][1] == Role.ASSISTANT and "agent-reply" in (payload["texts"][1] or "")
assert payload["roles"][0] == "user" and "hello world" in (payload["texts"][0] or "")
assert payload["roles"][1] == "assistant" and "agent-reply" in (payload["texts"][1] or "")
class _CaptureAgent(BaseAgent):
@@ -116,9 +115,9 @@ class _CaptureAgent(BaseAgent):
if isinstance(m, ChatMessage):
norm.append(m)
elif isinstance(m, str):
norm.append(ChatMessage(role=Role.USER, text=m))
norm.append(ChatMessage("user", [m]))
self._last_messages = norm
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text=self._reply_text)])
return AgentResponse(messages=[ChatMessage("assistant", [self._reply_text])])
async def run_stream( # type: ignore[override]
self,
@@ -134,7 +133,7 @@ class _CaptureAgent(BaseAgent):
if isinstance(m, ChatMessage):
norm.append(m)
elif isinstance(m, str):
norm.append(ChatMessage(role=Role.USER, text=m))
norm.append(ChatMessage("user", [m]))
self._last_messages = norm
yield AgentResponseUpdate(contents=[Content.from_text(text=self._reply_text)])
@@ -154,5 +153,5 @@ async def test_sequential_adapter_uses_full_conversation() -> None:
# Assert: second agent should have seen the user prompt and A1's assistant reply
seen = a2._last_messages # pyright: ignore[reportPrivateUsage]
assert len(seen) == 2
assert seen[0].role == Role.USER and "hello seq" in (seen[0].text or "")
assert seen[1].role == Role.ASSISTANT and "A1 reply" in (seen[1].text or "")
assert seen[0].role == "user" and "hello seq" in (seen[0].text or "")
assert seen[1].role == "assistant" and "A1 reply" in (seen[1].text or "")
python/packages/core/tests/workflow/test_group_chat.py
+18 -19
View File
@@ -25,7 +25,6 @@ from agent_framework import (
MagenticProgressLedger,
MagenticProgressLedgerItem,
RequestInfoEvent,
Role,
WorkflowOutputEvent,
WorkflowRunState,
WorkflowStatusEvent,
@@ -45,7 +44,7 @@ class StubAgent(BaseAgent):
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
response = ChatMessage(role=Role.ASSISTANT, text=self._reply_text, author_name=self.name)
response = ChatMessage("assistant", [self._reply_text], author_name=self.name)
return AgentResponse(messages=[response])
def run_stream( # type: ignore[override]
@@ -57,7 +56,7 @@ class StubAgent(BaseAgent):
) -> AsyncIterable[AgentResponseUpdate]:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(
contents=[Content.from_text(text=self._reply_text)], role=Role.ASSISTANT, author_name=self.name
contents=[Content.from_text(text=self._reply_text)], role="assistant", author_name=self.name
)
return _stream()
@@ -94,7 +93,7 @@ class StubManagerAgent(ChatAgent):
return AgentResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
text=(
'{"terminate": false, "reason": "Selecting agent", '
'"next_speaker": "agent", "final_message": null}'
@@ -115,7 +114,7 @@ class StubManagerAgent(ChatAgent):
return AgentResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
text=(
'{"terminate": true, "reason": "Task complete", '
'"next_speaker": null, "final_message": "agent manager final"}'
@@ -146,7 +145,7 @@ class StubManagerAgent(ChatAgent):
)
)
],
role=Role.ASSISTANT,
role="assistant",
author_name=self.name,
)
@@ -162,7 +161,7 @@ class StubManagerAgent(ChatAgent):
)
)
],
role=Role.ASSISTANT,
role="assistant",
author_name=self.name,
)
@@ -192,7 +191,7 @@ class StubMagenticManager(MagenticManagerBase):
self._round = 0
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="plan", author_name="magentic_manager")
return ChatMessage("assistant", ["plan"], author_name="magentic_manager")
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
return await self.plan(magentic_context)
@@ -218,7 +217,7 @@ class StubMagenticManager(MagenticManagerBase):
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="final", author_name="magentic_manager")
return ChatMessage("assistant", ["final"], author_name="magentic_manager")
async def test_group_chat_builder_basic_flow() -> None:
@@ -263,8 +262,8 @@ async def test_group_chat_as_agent_accepts_conversation() -> None:
agent = workflow.as_agent(name="group-chat-agent")
conversation = [
ChatMessage(role=Role.USER, text="kickoff", author_name="user"),
ChatMessage(role=Role.ASSISTANT, text="noted", author_name="alpha"),
ChatMessage("user", ["kickoff"], author_name="user"),
ChatMessage("assistant", ["noted"], author_name="alpha"),
]
response = await agent.run(conversation)
@@ -425,7 +424,7 @@ class TestGroupChatWorkflow:
return "agent"
def termination_condition(conversation: list[ChatMessage]) -> bool:
replies = [msg for msg in conversation if msg.role == Role.ASSISTANT and msg.author_name == "agent"]
replies = [msg for msg in conversation if msg.role == "assistant" and msg.author_name == "agent"]
return len(replies) >= 2
agent = StubAgent("agent", "response")
@@ -447,7 +446,7 @@ class TestGroupChatWorkflow:
assert outputs, "Expected termination to yield output"
conversation = outputs[-1]
agent_replies = [msg for msg in conversation if msg.author_name == "agent" and msg.role == Role.ASSISTANT]
agent_replies = [msg for msg in conversation if msg.author_name == "agent" and msg.role == "assistant"]
assert len(agent_replies) == 2
final_output = conversation[-1]
# The orchestrator uses its ID as author_name by default
@@ -553,7 +552,7 @@ class TestConversationHandling:
def selector(state: GroupChatState) -> str:
# Verify the conversation has the user message
assert len(state.conversation) > 0
assert state.conversation[0].role == Role.USER
assert state.conversation[0].role == "user"
assert state.conversation[0].text == "test string"
return "agent"
@@ -578,7 +577,7 @@ class TestConversationHandling:
async def test_handle_chat_message_input(self) -> None:
"""Test handling ChatMessage input directly."""
task_message = ChatMessage(role=Role.USER, text="test message")
task_message = ChatMessage("user", ["test message"])
def selector(state: GroupChatState) -> str:
# Verify the task message was preserved in conversation
@@ -608,8 +607,8 @@ class TestConversationHandling:
async def test_handle_conversation_list_input(self) -> None:
"""Test handling conversation list preserves context."""
conversation = [
ChatMessage(role=Role.SYSTEM, text="system message"),
ChatMessage(role=Role.USER, text="user message"),
ChatMessage("system", ["system message"]),
ChatMessage("user", ["user message"]),
]
def selector(state: GroupChatState) -> str:
@@ -1118,7 +1117,7 @@ async def test_group_chat_with_orchestrator_factory_returning_chat_agent():
return AgentResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
text=(
'{"terminate": false, "reason": "Selecting alpha", '
'"next_speaker": "alpha", "final_message": null}'
@@ -1138,7 +1137,7 @@ async def test_group_chat_with_orchestrator_factory_returning_chat_agent():
return AgentResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
text=(
'{"terminate": true, "reason": "Task complete", '
'"next_speaker": null, "final_message": "dynamic manager final"}'
python/packages/core/tests/workflow/test_handoff.py
+14 -21
View File
@@ -15,7 +15,6 @@ from agent_framework import (
HandoffAgentUserRequest,
HandoffBuilder,
RequestInfoEvent,
Role,
WorkflowEvent,
WorkflowOutputEvent,
resolve_agent_id,
@@ -49,7 +48,7 @@ class MockChatClient:
async def get_response(self, messages: Any, **kwargs: Any) -> ChatResponse:
contents = _build_reply_contents(self._name, self._handoff_to, self._next_call_id())
reply = ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=contents,
)
return ChatResponse(messages=reply, response_id="mock_response")
@@ -57,7 +56,7 @@ class MockChatClient:
def get_streaming_response(self, messages: Any, **kwargs: Any) -> AsyncIterable[ChatResponseUpdate]:
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
contents = _build_reply_contents(self._name, self._handoff_to, self._next_call_id())
yield ChatResponseUpdate(contents=contents, role=Role.ASSISTANT)
yield ChatResponseUpdate(contents=contents, role="assistant")
return _stream()
@@ -123,7 +122,7 @@ async def test_handoff():
workflow = (
HandoffBuilder(participants=[triage, specialist, escalation])
.with_start_agent(triage)
.with_termination_condition(lambda conv: sum(1 for m in conv if m.role == Role.USER) >= 2)
.with_termination_condition(lambda conv: sum(1 for m in conv if m.role == "user") >= 2)
.build()
)
@@ -174,9 +173,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)
assert any(
msg.role == Role.ASSISTANT and (msg.text or "").startswith("specialist reply") for msg in conversation_list
)
assert any(msg.role == "assistant" and (msg.text or "").startswith("specialist reply") for msg in conversation_list)
async def test_autonomous_mode_resumes_user_input_on_turn_limit():
@@ -222,7 +219,7 @@ async def test_handoff_async_termination_condition() -> None:
async def async_termination(conv: list[ChatMessage]) -> bool:
nonlocal termination_call_count
termination_call_count += 1
user_count = sum(1 for msg in conv if msg.role == Role.USER)
user_count = sum(1 for msg in conv if msg.role == "user")
return user_count >= 2
coordinator = MockHandoffAgent(name="coordinator", handoff_to="worker")
@@ -240,9 +237,7 @@ async def test_handoff_async_termination_condition() -> None:
assert requests
events = await _drain(
workflow.send_responses_streaming({
requests[-1].request_id: [ChatMessage(role=Role.USER, text="Second user message")]
})
workflow.send_responses_streaming({requests[-1].request_id: [ChatMessage("user", ["Second user message"])]})
)
outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)]
assert len(outputs) == 1
@@ -250,7 +245,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)
user_messages = [msg for msg in final_conv_list if msg.role == Role.USER]
user_messages = [msg for msg in final_conv_list if msg.role == "user"]
assert len(user_messages) == 2
assert termination_call_count > 0
@@ -264,7 +259,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=Role.ASSISTANT, text="Response")],
messages=[ChatMessage("assistant", ["Response"])],
response_id="test_response",
)
@@ -480,7 +475,7 @@ async def test_handoff_with_participant_factories():
workflow = (
HandoffBuilder(participant_factories={"triage": create_triage, "specialist": create_specialist})
.with_start_agent("triage")
.with_termination_condition(lambda conv: sum(1 for m in conv if m.role == Role.USER) >= 2)
.with_termination_condition(lambda conv: sum(1 for m in conv if m.role == "user") >= 2)
.build()
)
@@ -493,7 +488,7 @@ async def test_handoff_with_participant_factories():
# Follow-up message
events = await _drain(
workflow.send_responses_streaming({requests[-1].request_id: [ChatMessage(role=Role.USER, text="More details")]})
workflow.send_responses_streaming({requests[-1].request_id: [ChatMessage("user", ["More details"])]})
)
outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)]
assert outputs
@@ -553,7 +548,7 @@ async def test_handoff_with_participant_factories_and_add_handoff():
.with_start_agent("triage")
.add_handoff("triage", ["specialist_a", "specialist_b"])
.add_handoff("specialist_a", ["specialist_b"])
.with_termination_condition(lambda conv: sum(1 for m in conv if m.role == Role.USER) >= 3)
.with_termination_condition(lambda conv: sum(1 for m in conv if m.role == "user") >= 3)
.build()
)
@@ -567,9 +562,7 @@ async def test_handoff_with_participant_factories_and_add_handoff():
# Second user message - specialist_a hands off to specialist_b
events = await _drain(
workflow.send_responses_streaming({
requests[-1].request_id: [ChatMessage(role=Role.USER, text="Need escalation")]
})
workflow.send_responses_streaming({requests[-1].request_id: [ChatMessage("user", ["Need escalation"])]})
)
requests = [ev for ev in events if isinstance(ev, RequestInfoEvent)]
assert requests
@@ -594,7 +587,7 @@ async def test_handoff_participant_factories_with_checkpointing():
HandoffBuilder(participant_factories={"triage": create_triage, "specialist": create_specialist})
.with_start_agent("triage")
.with_checkpointing(storage)
.with_termination_condition(lambda conv: sum(1 for m in conv if m.role == Role.USER) >= 2)
.with_termination_condition(lambda conv: sum(1 for m in conv if m.role == "user") >= 2)
.build()
)
@@ -604,7 +597,7 @@ async def test_handoff_participant_factories_with_checkpointing():
assert requests
events = await _drain(
workflow.send_responses_streaming({requests[-1].request_id: [ChatMessage(role=Role.USER, text="follow up")]})
workflow.send_responses_streaming({requests[-1].request_id: [ChatMessage("user", ["follow up"])]})
)
outputs = [ev for ev in events if isinstance(ev, WorkflowOutputEvent)]
assert outputs, "Should have workflow output after termination condition is met"
python/packages/core/tests/workflow/test_magentic.py
+39 -40
View File
@@ -27,7 +27,6 @@ from agent_framework import (
MagenticProgressLedger,
MagenticProgressLedgerItem,
RequestInfoEvent,
Role,
StandardMagenticManager,
Workflow,
WorkflowCheckpoint,
@@ -53,7 +52,7 @@ def test_magentic_context_reset_behavior():
participant_descriptions={"Alice": "Researcher"},
)
# seed context state
ctx.chat_history.append(ChatMessage(role=Role.ASSISTANT, text="draft"))
ctx.chat_history.append(ChatMessage("assistant", ["draft"]))
ctx.stall_count = 2
prev_reset = ctx.reset_count
@@ -120,18 +119,18 @@ class FakeManager(MagenticManagerBase):
pass
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
facts = ChatMessage(role=Role.ASSISTANT, text="GIVEN OR VERIFIED FACTS\n- A\n")
plan = ChatMessage(role=Role.ASSISTANT, text="- Do X\n- Do Y\n")
facts = ChatMessage("assistant", ["GIVEN OR VERIFIED FACTS\n- A\n"])
plan = ChatMessage("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(role=Role.ASSISTANT, text=combined, author_name=self.name)
return ChatMessage("assistant", [combined], author_name=self.name)
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
facts = ChatMessage(role=Role.ASSISTANT, text="GIVEN OR VERIFIED FACTS\n- A2\n")
plan = ChatMessage(role=Role.ASSISTANT, text="- Do Z\n")
facts = ChatMessage("assistant", ["GIVEN OR VERIFIED FACTS\n- A2\n"])
plan = ChatMessage("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(role=Role.ASSISTANT, text=combined, author_name=self.name)
return ChatMessage("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
@@ -145,7 +144,7 @@ class FakeManager(MagenticManagerBase):
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text=self.FINAL_ANSWER, author_name=self.name)
return ChatMessage("assistant", [self.FINAL_ANSWER], author_name=self.name)
class StubAgent(BaseAgent):
@@ -160,7 +159,7 @@ class StubAgent(BaseAgent):
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
response = ChatMessage(role=Role.ASSISTANT, text=self._reply_text, author_name=self.name)
response = ChatMessage("assistant", [self._reply_text], author_name=self.name)
return AgentResponse(messages=[response])
def run_stream( # type: ignore[override]
@@ -172,7 +171,7 @@ class StubAgent(BaseAgent):
) -> AsyncIterable[AgentResponseUpdate]:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(
contents=[Content.from_text(text=self._reply_text)], role=Role.ASSISTANT, author_name=self.name
contents=[Content.from_text(text=self._reply_text)], role="assistant", author_name=self.name
)
return _stream()
@@ -223,8 +222,8 @@ async def test_magentic_as_agent_does_not_accept_conversation() -> None:
agent = workflow.as_agent(name="magentic-agent")
conversation = [
ChatMessage(role=Role.SYSTEM, text="Guidelines", author_name="system"),
ChatMessage(role=Role.USER, text="Summarize the findings", author_name="requester"),
ChatMessage("system", ["Guidelines"], author_name="system"),
ChatMessage("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)
@@ -238,7 +237,7 @@ async def test_standard_manager_plan_and_replan_combined_ledger():
)
first = await manager.plan(ctx.clone())
assert first.role == Role.ASSISTANT and "Facts:" in first.text and "Plan:" in first.text
assert first.role == "assistant" and "Facts:" in first.text and "Plan:" in first.text
assert manager.task_ledger is not None
replanned = await manager.replan(ctx.clone())
@@ -352,7 +351,7 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result():
data = output_event.data
assert isinstance(data, list)
assert len(data) > 0 # type: ignore
assert data[-1].role == Role.ASSISTANT # type: ignore
assert data[-1].role == "assistant" # type: ignore
assert all(isinstance(msg, ChatMessage) for msg in data) # type: ignore
@@ -427,7 +426,7 @@ class StubManagerAgent(BaseAgent):
thread: Any = None,
**kwargs: Any,
) -> AgentResponse:
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="ok")])
return AgentResponse(messages=[ChatMessage("assistant", ["ok"])])
def run_stream(
self,
@@ -437,7 +436,7 @@ class StubManagerAgent(BaseAgent):
**kwargs: Any,
) -> AsyncIterable[AgentResponseUpdate]:
async def _gen() -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(message_deltas=[ChatMessage(role=Role.ASSISTANT, text="ok")])
yield AgentResponseUpdate(message_deltas=[ChatMessage("assistant", ["ok"])])
return _gen()
@@ -448,8 +447,8 @@ async def test_standard_manager_plan_and_replan_via_complete_monkeypatch():
async def fake_complete_plan(messages: list[ChatMessage], **kwargs: Any) -> ChatMessage:
# Return a different response depending on call order length
if any("FACTS" in (m.text or "") for m in messages):
return ChatMessage(role=Role.ASSISTANT, text="- step A\n- step B")
return ChatMessage(role=Role.ASSISTANT, text="GIVEN OR VERIFIED FACTS\n- fact1")
return ChatMessage("assistant", ["- step A\n- step B"])
return ChatMessage("assistant", ["GIVEN OR VERIFIED FACTS\n- fact1"])
# First, patch to produce facts then plan
mgr._complete = fake_complete_plan # type: ignore[attr-defined]
@@ -464,8 +463,8 @@ async def test_standard_manager_plan_and_replan_via_complete_monkeypatch():
# Now replan with new outputs
async def fake_complete_replan(messages: list[ChatMessage], **kwargs: Any) -> ChatMessage:
if any("Please briefly explain" in (m.text or "") for m in messages):
return ChatMessage(role=Role.ASSISTANT, text="- new step")
return ChatMessage(role=Role.ASSISTANT, text="GIVEN OR VERIFIED FACTS\n- updated")
return ChatMessage("assistant", ["- new step"])
return ChatMessage("assistant", ["GIVEN OR VERIFIED FACTS\n- updated"])
mgr._complete = fake_complete_replan # type: ignore[attr-defined]
combined2 = await mgr.replan(ctx.clone())
@@ -485,7 +484,7 @@ 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(role=Role.ASSISTANT, text=json_text)
return ChatMessage("assistant", [json_text])
mgr._complete = fake_complete_ok # type: ignore[attr-defined]
ledger = await mgr.create_progress_ledger(ctx.clone())
@@ -493,7 +492,7 @@ async def test_standard_manager_progress_ledger_success_and_error():
# 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(role=Role.ASSISTANT, text="not-json")
return ChatMessage("assistant", ["not-json"])
mgr._complete = fake_complete_bad # type: ignore[attr-defined]
with pytest.raises(RuntimeError):
@@ -506,10 +505,10 @@ class InvokeOnceManager(MagenticManagerBase):
self._invoked = False
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="ledger")
return ChatMessage("assistant", ["ledger"])
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="re-ledger")
return ChatMessage("assistant", ["re-ledger"])
async def create_progress_ledger(self, magentic_context: MagenticContext) -> MagenticProgressLedger:
if not self._invoked:
@@ -532,7 +531,7 @@ class InvokeOnceManager(MagenticManagerBase):
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="final")
return ChatMessage("assistant", ["final"])
class StubThreadAgent(BaseAgent):
@@ -543,11 +542,11 @@ class StubThreadAgent(BaseAgent):
yield AgentResponseUpdate(
contents=[Content.from_text(text="thread-ok")],
author_name=self.name,
role=Role.ASSISTANT,
role="assistant",
)
async def run(self, messages=None, *, thread=None, **kwargs): # type: ignore[override]
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="thread-ok", author_name=self.name)])
return AgentResponse(messages=[ChatMessage("assistant", ["thread-ok"], author_name=self.name)])
class StubAssistantsClient:
@@ -565,11 +564,11 @@ class StubAssistantsAgent(BaseAgent):
yield AgentResponseUpdate(
contents=[Content.from_text(text="assistants-ok")],
author_name=self.name,
role=Role.ASSISTANT,
role="assistant",
)
async def run(self, messages=None, *, thread=None, **kwargs): # type: ignore[override]
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="assistants-ok", author_name=self.name)])
return AgentResponse(messages=[ChatMessage("assistant", ["assistants-ok"], author_name=self.name)])
async def _collect_agent_responses_setup(participant: AgentProtocol) -> list[ChatMessage]:
@@ -586,7 +585,7 @@ async def _collect_agent_responses_setup(participant: AgentProtocol) -> list[Cha
if isinstance(ev, AgentRunUpdateEvent):
captured.append(
ChatMessage(
role=ev.data.role or Role.ASSISTANT,
role=ev.data.role or "assistant",
text=ev.data.text or "",
author_name=ev.data.author_name,
)
@@ -738,10 +737,10 @@ class NotProgressingManager(MagenticManagerBase):
"""
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="ledger")
return ChatMessage("assistant", ["ledger"])
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="re-ledger")
return ChatMessage("assistant", ["re-ledger"])
async def create_progress_ledger(self, magentic_context: MagenticContext) -> MagenticProgressLedger:
return MagenticProgressLedger(
@@ -753,7 +752,7 @@ class NotProgressingManager(MagenticManagerBase):
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="final")
return ChatMessage("assistant", ["final"])
async def test_magentic_stall_and_reset_reach_limits():
@@ -851,8 +850,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(role=Role.ASSISTANT, text="response1"))
ctx.chat_history.append(ChatMessage(role=Role.ASSISTANT, text="response2"))
ctx.chat_history.append(ChatMessage("assistant", ["response1"]))
ctx.chat_history.append(ChatMessage("assistant", ["response2"]))
assert len(ctx.chat_history) == 2
# Reset
@@ -862,7 +861,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(role=Role.ASSISTANT, text="new_response"))
ctx.chat_history.append(ChatMessage("assistant", ["new_response"]))
assert len(ctx.chat_history) == 1, "Should have exactly 1 message after adding to reset context"
@@ -881,7 +880,7 @@ async def test_magentic_checkpoint_restore_no_duplicate_history():
# Run with conversation history to create initial checkpoint
conversation: list[ChatMessage] = [
ChatMessage(role=Role.USER, text="task_msg"),
ChatMessage("user", ["task_msg"]),
]
async for event in wf.run_stream(conversation):
@@ -1248,8 +1247,8 @@ def test_magentic_agent_factory_with_standard_manager_options():
from agent_framework._workflows._magentic import _MagenticTaskLedger # type: ignore
custom_task_ledger = _MagenticTaskLedger(
facts=ChatMessage(role=Role.ASSISTANT, text="Custom facts"),
plan=ChatMessage(role=Role.ASSISTANT, text="Custom plan"),
facts=ChatMessage("assistant", ["Custom facts"]),
plan=ChatMessage("assistant", ["Custom plan"]),
)
participant = StubAgent("agentA", "reply from agentA")
python/packages/core/tests/workflow/test_orchestration_request_info.py
+10 -11
View File
@@ -14,7 +14,6 @@ from agent_framework import (
AgentResponseUpdate,
AgentThread,
ChatMessage,
Role,
)
from agent_framework._workflows._agent_executor import AgentExecutorRequest, AgentExecutorResponse
from agent_framework._workflows._orchestration_request_info import (
@@ -73,7 +72,7 @@ class TestAgentRequestInfoResponse:
def test_create_response_with_messages(self):
"""Test creating an AgentRequestInfoResponse with messages."""
messages = [ChatMessage(role=Role.USER, text="Additional info")]
messages = [ChatMessage("user", ["Additional info"])]
response = AgentRequestInfoResponse(messages=messages)
assert response.messages == messages
@@ -81,8 +80,8 @@ class TestAgentRequestInfoResponse:
def test_from_messages_factory(self):
"""Test creating response from ChatMessage list."""
messages = [
ChatMessage(role=Role.USER, text="Message 1"),
ChatMessage(role=Role.USER, text="Message 2"),
ChatMessage("user", ["Message 1"]),
ChatMessage("user", ["Message 2"]),
]
response = AgentRequestInfoResponse.from_messages(messages)
@@ -94,9 +93,9 @@ class TestAgentRequestInfoResponse:
response = AgentRequestInfoResponse.from_strings(texts)
assert len(response.messages) == 2
assert response.messages[0].role == Role.USER
assert response.messages[0].role == "user"
assert response.messages[0].text == "First message"
assert response.messages[1].role == Role.USER
assert response.messages[1].role == "user"
assert response.messages[1].text == "Second message"
def test_approve_factory(self):
@@ -114,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=Role.ASSISTANT, text="Agent response")])
agent_response = AgentResponse(messages=[ChatMessage("assistant", ["Agent response"])])
agent_response = AgentExecutorResponse(
executor_id="test_agent",
agent_response=agent_response,
@@ -132,7 +131,7 @@ class TestAgentRequestInfoExecutor:
"""Test response handler when user provides additional messages."""
executor = AgentRequestInfoExecutor(id="test_executor")
agent_response = AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Original")])
agent_response = AgentResponse(messages=[ChatMessage("assistant", ["Original"])])
original_request = AgentExecutorResponse(
executor_id="test_agent",
agent_response=agent_response,
@@ -158,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=Role.ASSISTANT, text="Original")])
agent_response = AgentResponse(messages=[ChatMessage("assistant", ["Original"])])
original_request = AgentExecutorResponse(
executor_id="test_agent",
agent_response=agent_response,
@@ -207,7 +206,7 @@ class _TestAgent:
**kwargs: Any,
) -> AgentResponse:
"""Dummy run method."""
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Test response")])
return AgentResponse(messages=[ChatMessage("assistant", ["Test response"])])
def run_stream(
self,
@@ -219,7 +218,7 @@ class _TestAgent:
"""Dummy run_stream method."""
async def generator():
yield AgentResponseUpdate(messages=[ChatMessage(role=Role.ASSISTANT, text="Test response stream")])
yield AgentResponseUpdate(messages=[ChatMessage("assistant", ["Test response stream"])])
return generator()
python/packages/core/tests/workflow/test_sequential.py
+16 -17
View File
@@ -14,7 +14,6 @@ from agent_framework import (
ChatMessage,
Content,
Executor,
Role,
SequentialBuilder,
TypeCompatibilityError,
WorkflowContext,
@@ -36,7 +35,7 @@ class _EchoAgent(BaseAgent):
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentResponse:
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text=f"{self.name} reply")])
return AgentResponse(messages=[ChatMessage("assistant", [f"{self.name} reply"])])
async def run_stream( # type: ignore[override]
self,
@@ -55,9 +54,9 @@ class _SummarizerExec(Executor):
@handler
async def summarize(self, agent_response: AgentExecutorResponse, ctx: WorkflowContext[list[ChatMessage]]) -> None:
conversation = agent_response.full_conversation or []
user_texts = [m.text for m in conversation if m.role == Role.USER]
agents = [m.author_name or m.role for m in conversation if m.role == Role.ASSISTANT]
summary = ChatMessage(role=Role.ASSISTANT, text=f"Summary of users:{len(user_texts)} agents:{len(agents)}")
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)}"])
await ctx.send_message(list(conversation) + [summary])
@@ -119,9 +118,9 @@ async def test_sequential_agents_append_to_context() -> None:
assert isinstance(output, list)
msgs: list[ChatMessage] = output
assert len(msgs) == 3
assert msgs[0].role == Role.USER and "hello sequential" in msgs[0].text
assert msgs[1].role == Role.ASSISTANT and (msgs[1].author_name == "A1" or True)
assert msgs[2].role == Role.ASSISTANT and (msgs[2].author_name == "A2" or True)
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)
assert msgs[2].role == "assistant" and (msgs[2].author_name == "A2" or True)
assert "A1 reply" in msgs[1].text
assert "A2 reply" in msgs[2].text
@@ -152,9 +151,9 @@ async def test_sequential_register_participants_with_agent_factories() -> None:
assert isinstance(output, list)
msgs: list[ChatMessage] = output
assert len(msgs) == 3
assert msgs[0].role == Role.USER and "hello factories" in msgs[0].text
assert msgs[1].role == Role.ASSISTANT and "A1 reply" in msgs[1].text
assert msgs[2].role == Role.ASSISTANT and "A2 reply" in msgs[2].text
assert msgs[0].role == "user" and "hello factories" in msgs[0].text
assert msgs[1].role == "assistant" and "A1 reply" in msgs[1].text
assert msgs[2].role == "assistant" and "A2 reply" in msgs[2].text
async def test_sequential_with_custom_executor_summary() -> None:
@@ -178,9 +177,9 @@ async def test_sequential_with_custom_executor_summary() -> None:
msgs: list[ChatMessage] = output
# Expect: [user, A1 reply, summary]
assert len(msgs) == 3
assert msgs[0].role == Role.USER
assert msgs[1].role == Role.ASSISTANT and "A1 reply" in msgs[1].text
assert msgs[2].role == Role.ASSISTANT and msgs[2].text.startswith("Summary of users:")
assert msgs[0].role == "user"
assert msgs[1].role == "assistant" and "A1 reply" in msgs[1].text
assert msgs[2].role == "assistant" and msgs[2].text.startswith("Summary of users:")
async def test_sequential_register_participants_mixed_agents_and_executors() -> None:
@@ -209,9 +208,9 @@ async def test_sequential_register_participants_mixed_agents_and_executors() ->
msgs: list[ChatMessage] = output
# Expect: [user, A1 reply, summary]
assert len(msgs) == 3
assert msgs[0].role == Role.USER and "topic Y" in msgs[0].text
assert msgs[1].role == Role.ASSISTANT and "A1 reply" in msgs[1].text
assert msgs[2].role == Role.ASSISTANT and msgs[2].text.startswith("Summary of users:")
assert msgs[0].role == "user" and "topic Y" in msgs[0].text
assert msgs[1].role == "assistant" and "A1 reply" in msgs[1].text
assert msgs[2].role == "assistant" and msgs[2].text.startswith("Summary of users:")
async def test_sequential_checkpoint_resume_round_trip() -> None:
python/packages/core/tests/workflow/test_workflow.py
+1 -2
View File
@@ -23,7 +23,6 @@ from agent_framework import (
FileCheckpointStorage,
Message,
RequestInfoEvent,
Role,
WorkflowBuilder,
WorkflowCheckpointException,
WorkflowContext,
@@ -869,7 +868,7 @@ class _StreamingTestAgent(BaseAgent):
**kwargs: Any,
) -> AgentResponse:
"""Non-streaming run - returns complete response."""
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text=self._reply_text)])
return AgentResponse(messages=[ChatMessage("assistant", [self._reply_text])])
async def run_stream(
self,
python/packages/core/tests/workflow/test_workflow_agent.py
+47 -48
View File
@@ -16,7 +16,6 @@ from agent_framework import (
ChatMessageStore,
Content,
Executor,
Role,
UsageDetails,
WorkflowAgent,
WorkflowBuilder,
@@ -41,11 +40,11 @@ class SimpleExecutor(Executor):
response_text = f"{self.response_text}: {input_text}"
# Create response message for both streaming and non-streaming cases
response_message = ChatMessage(role=Role.ASSISTANT, contents=[Content.from_text(text=response_text)])
response_message = ChatMessage("assistant", [Content.from_text(text=response_text)])
# Emit update event.
streaming_update = AgentResponseUpdate(
contents=[Content.from_text(text=response_text)], role=Role.ASSISTANT, message_id=str(uuid.uuid4())
contents=[Content.from_text(text=response_text)], role="assistant", message_id=str(uuid.uuid4())
)
await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=streaming_update))
@@ -68,7 +67,7 @@ class RequestingExecutor(Executor):
# Handle the response and emit completion response
update = AgentResponseUpdate(
contents=[Content.from_text(text="Request completed successfully")],
role=Role.ASSISTANT,
role="assistant",
message_id=str(uuid.uuid4()),
)
await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=update))
@@ -90,10 +89,10 @@ class ConversationHistoryCapturingExecutor(Executor):
message_count = len(messages)
response_text = f"Received {message_count} messages"
response_message = ChatMessage(role=Role.ASSISTANT, contents=[Content.from_text(text=response_text)])
response_message = ChatMessage("assistant", [Content.from_text(text=response_text)])
streaming_update = AgentResponseUpdate(
contents=[Content.from_text(text=response_text)], role=Role.ASSISTANT, message_id=str(uuid.uuid4())
contents=[Content.from_text(text=response_text)], 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])
@@ -232,7 +231,7 @@ class TestWorkflowAgent:
),
)
response_message = ChatMessage(role=Role.USER, contents=[approval_response])
response_message = ChatMessage("user", [approval_response])
# Continue the workflow with the response
continuation_result = await agent.run(response_message)
@@ -295,7 +294,7 @@ class TestWorkflowAgent:
workflow = WorkflowBuilder().set_start_executor(yielding_executor).build()
# Run directly - should return WorkflowOutputEvent in result
direct_result = await workflow.run([ChatMessage(role=Role.USER, contents=[Content.from_text(text="hello")])])
direct_result = await workflow.run([ChatMessage("user", [Content.from_text(text="hello")])])
direct_outputs = direct_result.get_outputs()
assert len(direct_outputs) == 1
assert direct_outputs[0] == "processed: hello"
@@ -362,7 +361,7 @@ class TestWorkflowAgent:
@executor
async def chat_message_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
msg = ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[Content.from_text(text="response text")],
author_name="custom-author",
)
@@ -374,7 +373,7 @@ class TestWorkflowAgent:
result = await agent.run("test")
assert len(result.messages) == 1
assert result.messages[0].role == Role.ASSISTANT
assert result.messages[0].role == "assistant"
assert result.messages[0].text == "response text"
assert result.messages[0].author_name == "custom-author"
@@ -425,10 +424,10 @@ class TestWorkflowAgent:
async def list_yielding_executor(messages: list[ChatMessage], ctx: WorkflowContext) -> None:
# Yield a list of ChatMessages (as SequentialBuilder does)
msg_list = [
ChatMessage(role=Role.USER, contents=[Content.from_text(text="first message")]),
ChatMessage(role=Role.ASSISTANT, contents=[Content.from_text(text="second message")]),
ChatMessage("user", [Content.from_text(text="first message")]),
ChatMessage("assistant", [Content.from_text(text="second message")]),
ChatMessage(
role=Role.ASSISTANT,
role="assistant",
contents=[Content.from_text(text="third"), Content.from_text(text="fourth")],
),
]
@@ -469,8 +468,8 @@ class TestWorkflowAgent:
# 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"),
ChatMessage("user", ["Previous user message"]),
ChatMessage("assistant", ["Previous assistant response"]),
]
message_store = ChatMessageStore(messages=history_messages)
thread = AgentThread(message_store=message_store)
@@ -499,9 +498,9 @@ class TestWorkflowAgent:
# 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!"),
ChatMessage("system", ["You are a helpful assistant"]),
ChatMessage("user", ["Hello"]),
ChatMessage("assistant", ["Hi there!"]),
]
message_store = ChatMessageStore(messages=history_messages)
thread = AgentThread(message_store=message_store)
@@ -579,7 +578,7 @@ class TestWorkflowAgent:
async def run(self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any) -> AgentResponse:
return AgentResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text=self._response_text)],
messages=[ChatMessage("assistant", [self._response_text])],
text=self._response_text,
)
@@ -589,7 +588,7 @@ class TestWorkflowAgent:
for word in self._response_text.split():
yield AgentResponseUpdate(
contents=[Content.from_text(text=word + " ")],
role=Role.ASSISTANT,
role="assistant",
author_name=self._name,
)
@@ -653,7 +652,7 @@ class TestWorkflowAgent:
async def run(self, messages: Any, *, thread: AgentThread | None = None, **kwargs: Any) -> AgentResponse:
return AgentResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text=self._response_text)],
messages=[ChatMessage("assistant", [self._response_text])],
text=self._response_text,
)
@@ -662,7 +661,7 @@ class TestWorkflowAgent:
) -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(
contents=[Content.from_text(text=self._response_text)],
role=Role.ASSISTANT,
role="assistant",
author_name=self._name,
)
@@ -728,7 +727,7 @@ class TestWorkflowAgentAuthorName:
# Emit update with explicit author_name
update = AgentResponseUpdate(
contents=[Content.from_text(text="Response with author")],
role=Role.ASSISTANT,
role="assistant",
author_name="custom_author_name", # Explicitly set
message_id=str(uuid.uuid4()),
)
@@ -780,7 +779,7 @@ class TestWorkflowAgentMergeUpdates:
# Response B, Message 2 (latest in resp B)
AgentResponseUpdate(
contents=[Content.from_text(text="RespB-Msg2")],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-b",
message_id="msg-2",
created_at="2024-01-01T12:02:00Z",
@@ -788,7 +787,7 @@ class TestWorkflowAgentMergeUpdates:
# Response A, Message 1 (earliest overall)
AgentResponseUpdate(
contents=[Content.from_text(text="RespA-Msg1")],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-a",
message_id="msg-1",
created_at="2024-01-01T12:00:00Z",
@@ -796,7 +795,7 @@ class TestWorkflowAgentMergeUpdates:
# Response B, Message 1 (earlier in resp B)
AgentResponseUpdate(
contents=[Content.from_text(text="RespB-Msg1")],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-b",
message_id="msg-1",
created_at="2024-01-01T12:01:00Z",
@@ -804,7 +803,7 @@ class TestWorkflowAgentMergeUpdates:
# Response A, Message 2 (later in resp A)
AgentResponseUpdate(
contents=[Content.from_text(text="RespA-Msg2")],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-a",
message_id="msg-2",
created_at="2024-01-01T12:00:30Z",
@@ -812,7 +811,7 @@ class TestWorkflowAgentMergeUpdates:
# Global dangling update (no response_id) - should go at end
AgentResponseUpdate(
contents=[Content.from_text(text="Global-Dangling")],
role=Role.ASSISTANT,
role="assistant",
response_id=None,
message_id="msg-global",
created_at="2024-01-01T11:59:00Z", # Earliest timestamp but should be last
@@ -886,7 +885,7 @@ class TestWorkflowAgentMergeUpdates:
usage_details={"input_token_count": 10, "output_token_count": 5, "total_token_count": 15}
),
],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-1",
message_id="msg-1",
created_at="2024-01-01T12:00:00Z",
@@ -899,7 +898,7 @@ class TestWorkflowAgentMergeUpdates:
usage_details={"input_token_count": 20, "output_token_count": 8, "total_token_count": 28}
),
],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-2",
message_id="msg-2",
created_at="2024-01-01T12:01:00Z", # Later timestamp
@@ -912,7 +911,7 @@ class TestWorkflowAgentMergeUpdates:
usage_details={"input_token_count": 5, "output_token_count": 3, "total_token_count": 8}
),
],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-1", # Same response_id as first
message_id="msg-3",
created_at="2024-01-01T11:59:00Z", # Earlier timestamp
@@ -975,7 +974,7 @@ class TestWorkflowAgentMergeUpdates:
# User question
AgentResponseUpdate(
contents=[Content.from_text(text="What is the weather?")],
role=Role.USER,
role="user",
response_id="resp-1",
message_id="msg-1",
created_at="2024-01-01T12:00:00Z",
@@ -985,7 +984,7 @@ class TestWorkflowAgentMergeUpdates:
contents=[
Content.from_function_call(call_id=call_id, name="get_weather", arguments='{"location": "NYC"}')
],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-1",
message_id="msg-2",
created_at="2024-01-01T12:00:01Z",
@@ -994,7 +993,7 @@ class TestWorkflowAgentMergeUpdates:
# and be placed at the end (the bug); fix now correctly associates via call_id
AgentResponseUpdate(
contents=[Content.from_function_result(call_id=call_id, result="Sunny, 72F")],
role=Role.TOOL,
role="tool",
response_id=None,
message_id="msg-3",
created_at="2024-01-01T12:00:02Z",
@@ -1002,7 +1001,7 @@ class TestWorkflowAgentMergeUpdates:
# Final assistant answer
AgentResponseUpdate(
contents=[Content.from_text(text="The weather in NYC is sunny and 72F.")],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-1",
message_id="msg-4",
created_at="2024-01-01T12:00:03Z",
@@ -1026,10 +1025,10 @@ class TestWorkflowAgentMergeUpdates:
# Verify correct ordering: user -> function_call -> function_result -> assistant_answer
expected_sequence = [
("text", Role.USER),
("function_call", Role.ASSISTANT),
("function_result", Role.TOOL),
("text", Role.ASSISTANT),
("text", "user"),
("function_call", "assistant"),
("function_result", "tool"),
("text", "assistant"),
]
assert content_sequence == expected_sequence, (
@@ -1073,7 +1072,7 @@ class TestWorkflowAgentMergeUpdates:
# User question
AgentResponseUpdate(
contents=[Content.from_text(text="What's the weather and time?")],
role=Role.USER,
role="user",
response_id="resp-1",
message_id="msg-1",
created_at="2024-01-01T12:00:00Z",
@@ -1083,7 +1082,7 @@ class TestWorkflowAgentMergeUpdates:
contents=[
Content.from_function_call(call_id=call_id_1, name="get_weather", arguments='{"location": "NYC"}')
],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-1",
message_id="msg-2",
created_at="2024-01-01T12:00:01Z",
@@ -1093,7 +1092,7 @@ class TestWorkflowAgentMergeUpdates:
contents=[
Content.from_function_call(call_id=call_id_2, name="get_time", arguments='{"timezone": "EST"}')
],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-1",
message_id="msg-3",
created_at="2024-01-01T12:00:02Z",
@@ -1101,7 +1100,7 @@ class TestWorkflowAgentMergeUpdates:
# Second function result arrives first (no response_id)
AgentResponseUpdate(
contents=[Content.from_function_result(call_id=call_id_2, result="3:00 PM EST")],
role=Role.TOOL,
role="tool",
response_id=None,
message_id="msg-4",
created_at="2024-01-01T12:00:03Z",
@@ -1109,7 +1108,7 @@ class TestWorkflowAgentMergeUpdates:
# First function result arrives second (no response_id)
AgentResponseUpdate(
contents=[Content.from_function_result(call_id=call_id_1, result="Sunny, 72F")],
role=Role.TOOL,
role="tool",
response_id=None,
message_id="msg-5",
created_at="2024-01-01T12:00:04Z",
@@ -1117,7 +1116,7 @@ class TestWorkflowAgentMergeUpdates:
# Final assistant answer
AgentResponseUpdate(
contents=[Content.from_text(text="It's sunny (72F) and 3 PM in NYC.")],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-1",
message_id="msg-6",
created_at="2024-01-01T12:00:05Z",
@@ -1168,7 +1167,7 @@ class TestWorkflowAgentMergeUpdates:
updates = [
AgentResponseUpdate(
contents=[Content.from_text(text="Hello")],
role=Role.USER,
role="user",
response_id="resp-1",
message_id="msg-1",
created_at="2024-01-01T12:00:00Z",
@@ -1176,14 +1175,14 @@ class TestWorkflowAgentMergeUpdates:
# Function result with no matching call
AgentResponseUpdate(
contents=[Content.from_function_result(call_id="orphan_call_id", result="orphan result")],
role=Role.TOOL,
role="tool",
response_id=None,
message_id="msg-2",
created_at="2024-01-01T12:00:01Z",
),
AgentResponseUpdate(
contents=[Content.from_text(text="Goodbye")],
role=Role.ASSISTANT,
role="assistant",
response_id="resp-1",
message_id="msg-3",
created_at="2024-01-01T12:00:02Z",
python/packages/core/tests/workflow/test_workflow_builder.py
+1 -2
View File
@@ -13,7 +13,6 @@ from agent_framework import (
BaseAgent,
ChatMessage,
Executor,
Role,
WorkflowBuilder,
WorkflowContext,
handler,
@@ -28,7 +27,7 @@ class DummyAgent(BaseAgent):
if isinstance(m, ChatMessage):
norm.append(m)
elif isinstance(m, str):
norm.append(ChatMessage(role=Role.USER, text=m))
norm.append(ChatMessage("user", [m]))
return AgentResponse(messages=norm)
async def run_stream(self, messages=None, *, thread: AgentThread | None = None, **kwargs): # type: ignore[override]
python/packages/core/tests/workflow/test_workflow_kwargs.py
+7 -8
View File
@@ -16,7 +16,6 @@ from agent_framework import (
GroupChatBuilder,
GroupChatState,
HandoffBuilder,
Role,
SequentialBuilder,
WorkflowRunState,
WorkflowStatusEvent,
@@ -57,7 +56,7 @@ class _KwargsCapturingAgent(BaseAgent):
**kwargs: Any,
) -> AgentResponse:
self.captured_kwargs.append(dict(kwargs))
return AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text=f"{self.name} response")])
return AgentResponse(messages=[ChatMessage("assistant", [f"{self.name} response"])])
async def run_stream(
self,
@@ -387,10 +386,10 @@ async def test_magentic_kwargs_flow_to_agents() -> None:
self.task_ledger = None
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="Plan: Test task", author_name="manager")
return ChatMessage("assistant", ["Plan: Test task"], author_name="manager")
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="Replan: Test task", author_name="manager")
return ChatMessage("assistant", ["Replan: Test task"], author_name="manager")
async def create_progress_ledger(self, magentic_context: MagenticContext) -> MagenticProgressLedger:
# Return completed on first call
@@ -403,7 +402,7 @@ async def test_magentic_kwargs_flow_to_agents() -> None:
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="Final answer", author_name="manager")
return ChatMessage("assistant", ["Final answer"], author_name="manager")
agent = _KwargsCapturingAgent(name="agent1")
manager = _MockManager()
@@ -437,10 +436,10 @@ async def test_magentic_kwargs_stored_in_shared_state() -> None:
self.task_ledger = None
async def plan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="Plan", author_name="manager")
return ChatMessage("assistant", ["Plan"], author_name="manager")
async def replan(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="Replan", author_name="manager")
return ChatMessage("assistant", ["Replan"], author_name="manager")
async def create_progress_ledger(self, magentic_context: MagenticContext) -> MagenticProgressLedger:
return MagenticProgressLedger(
@@ -452,7 +451,7 @@ async def test_magentic_kwargs_stored_in_shared_state() -> None:
)
async def prepare_final_answer(self, magentic_context: MagenticContext) -> ChatMessage:
return ChatMessage(role=Role.ASSISTANT, text="Final", author_name="manager")
return ChatMessage("assistant", ["Final"], author_name="manager")
agent = _KwargsCapturingAgent(name="agent1")
manager = _MockManager()