diff --git a/src/a2a/client/helpers.py b/src/a2a/client/helpers.py index 930c71e6..67a39ce7 100644 --- a/src/a2a/client/helpers.py +++ b/src/a2a/client/helpers.py @@ -1,22 +1,32 @@ """Helper functions for the A2A client.""" +from typing import Any from uuid import uuid4 from a2a.types import Message, Part, Role, TextPart def create_text_message_object( - role: Role = Role.user, content: str = '' + role: Role = Role.user, + content: str = '', + extensions: list[str] | None = None, + metadata: dict[str, Any] | None = None, ) -> Message: """Create a Message object containing a single TextPart. Args: role: The role of the message sender (user or agent). Defaults to Role.user. content: The text content of the message. Defaults to an empty string. + extensions: The extensions of the message. Defaults to an empty list. + metadata: The metadata of the message. Defaults to an empty dictionary. Returns: A `Message` object with a new UUID message_id. """ return Message( - role=role, parts=[Part(TextPart(text=content))], message_id=str(uuid4()) + role=role, + parts=[Part(TextPart(text=content))], + message_id=str(uuid4()), + extensions=extensions or [], + metadata=metadata or {}, ) diff --git a/tests/client/test_jsonrpc_client.py b/tests/client/test_jsonrpc_client.py index 58feec25..57600cf5 100644 --- a/tests/client/test_jsonrpc_client.py +++ b/tests/client/test_jsonrpc_client.py @@ -378,6 +378,64 @@ async def test_send_message_success( assert isinstance(response, Message) assert response.model_dump() == success_response.model_dump() + @pytest.mark.asyncio + async def test_send_message_success_with_extensions( + self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock + ): + client = JsonRpcTransport( + httpx_client=mock_httpx_client, agent_card=mock_agent_card + ) + params = MessageSendParams( + message=create_text_message_object( + content='Hello', extensions=['test'] + ) + ) + success_response = create_text_message_object( + role=Role.agent, content='Hi there!', extensions=['test'] + ) + rpc_response = SendMessageSuccessResponse( + id='123', jsonrpc='2.0', result=success_response + ) + response = httpx.Response( + 200, json=rpc_response.model_dump(mode='json') + ) + response.request = httpx.Request('POST', 'http://agent.example.com/api') + mock_httpx_client.post.return_value = response + + response = await client.send_message(request=params) + + assert isinstance(response, Message) + assert response.model_dump() == success_response.model_dump() + + @pytest.mark.asyncio + async def test_send_message_success_with_metadata( + self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock + ): + client = JsonRpcTransport( + httpx_client=mock_httpx_client, agent_card=mock_agent_card + ) + params = MessageSendParams( + message=create_text_message_object( + content='Hello', metadata={'test': 'test'} + ) + ) + success_response = create_text_message_object( + role=Role.agent, content='Hi there!', metadata={'test': 'test'} + ) + rpc_response = SendMessageSuccessResponse( + id='123', jsonrpc='2.0', result=success_response + ) + response = httpx.Response( + 200, json=rpc_response.model_dump(mode='json') + ) + response.request = httpx.Request('POST', 'http://agent.example.com/api') + mock_httpx_client.post.return_value = response + + response = await client.send_message(request=params) + + assert isinstance(response, Message) + assert response.model_dump() == success_response.model_dump() + @pytest.mark.asyncio async def test_send_message_error_response( self, mock_httpx_client: AsyncMock, mock_agent_card: MagicMock