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: cleanup and refactoring of chat clients (#2937)

Browse Source 
* refactoring and unifying naming schemes of internal methods of chat clients

* set tool_choice to auto

* fix for mypy

* added note on naming and fix #2951

* fix responses

* fixes in azure ai agents client
This commit is contained in:
Eduard van Valkenburg
2025-12-18 13:02:23 +01:00
committed by GitHub
Unverified
parent a71f768331
commit e5c11d38d6
26 changed files with 1128 additions and 1068 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/a2a/agent_framework_a2a/_agent.py
+14 -14
View File
@@ -237,14 +237,14 @@ class A2AAgent(BaseAgent):
An agent response item.
"""
messages = self._normalize_messages(messages)
a2a_message = self._chat_message_to_a2a_message(messages[-1])
a2a_message = self._prepare_message_for_a2a(messages[-1])
response_stream = self.client.send_message(a2a_message)
async for item in response_stream:
if isinstance(item, Message):
# Process A2A Message
contents = self._a2a_parts_to_contents(item.parts)
contents = self._parse_contents_from_a2a(item.parts)
yield AgentRunResponseUpdate(
contents=contents,
role=Role.ASSISTANT if item.role == A2ARole.agent else Role.USER,
@@ -255,7 +255,7 @@ class A2AAgent(BaseAgent):
task, _update_event = item
if isinstance(task, Task) and task.status.state in TERMINAL_TASK_STATES:
# Convert Task artifacts to ChatMessages and yield as separate updates
task_messages = self._task_to_chat_messages(task)
task_messages = self._parse_messages_from_task(task)
if task_messages:
for message in task_messages:
# Use the artifact's ID from raw_representation as message_id for unique identification
@@ -280,8 +280,8 @@ class A2AAgent(BaseAgent):
msg = f"Only Message and Task responses are supported from A2A agents. Received: {type(item)}"
raise NotImplementedError(msg)
def _chat_message_to_a2a_message(self, message: ChatMessage) -> A2AMessage:
"""Convert a ChatMessage to an A2A Message.
def _prepare_message_for_a2a(self, message: ChatMessage) -> A2AMessage:
"""Prepare a ChatMessage for the A2A protocol.
Transforms Agent Framework ChatMessage objects into A2A protocol Messages by:
- Converting all message contents to appropriate A2A Part types
@@ -361,8 +361,8 @@ class A2AAgent(BaseAgent):
metadata=cast(dict[str, Any], message.additional_properties),
)
def _a2a_parts_to_contents(self, parts: Sequence[A2APart]) -> list[Contents]:
"""Convert A2A Parts to Agent Framework Contents.
def _parse_contents_from_a2a(self, parts: Sequence[A2APart]) -> list[Contents]:
"""Parse A2A Parts into Agent Framework Contents.
Transforms A2A protocol Parts into framework-native Content objects,
handling text, file (URI/bytes), and data parts with metadata preservation.
@@ -410,17 +410,17 @@ class A2AAgent(BaseAgent):
raise ValueError(f"Unknown Part kind: {inner_part.kind}")
return contents
def _task_to_chat_messages(self, task: Task) -> list[ChatMessage]:
"""Convert A2A Task artifacts to ChatMessages with ASSISTANT role."""
def _parse_messages_from_task(self, task: Task) -> list[ChatMessage]:
"""Parse A2A Task artifacts into ChatMessages with ASSISTANT role."""
messages: list[ChatMessage] = []
if task.artifacts is not None:
for artifact in task.artifacts:
messages.append(self._artifact_to_chat_message(artifact))
messages.append(self._parse_message_from_artifact(artifact))
elif task.history is not None and len(task.history) > 0:
# Include the last history item as the agent response
history_item = task.history[-1]
contents = self._a2a_parts_to_contents(history_item.parts)
contents = self._parse_contents_from_a2a(history_item.parts)
messages.append(
ChatMessage(
role=Role.ASSISTANT if history_item.role == A2ARole.agent else Role.USER,
@@ -431,9 +431,9 @@ class A2AAgent(BaseAgent):
return messages
def _artifact_to_chat_message(self, artifact: Artifact) -> ChatMessage:
"""Convert A2A Artifact to ChatMessage using part contents."""
contents = self._a2a_parts_to_contents(artifact.parts)
def _parse_message_from_artifact(self, artifact: Artifact) -> ChatMessage:
"""Parse A2A Artifact into ChatMessage using part contents."""
contents = self._parse_contents_from_a2a(artifact.parts)
return ChatMessage(
role=Role.ASSISTANT,
contents=contents,
python/packages/a2a/tests/test_a2a_agent.py
+33 -33
View File
@@ -197,18 +197,18 @@ async def test_run_with_unknown_response_type_raises_error(a2a_agent: A2AAgent,
await a2a_agent.run("Test message")
def test_task_to_chat_messages_empty_artifacts(a2a_agent: A2AAgent) -> None:
"""Test _task_to_chat_messages with task containing no artifacts."""
def test_parse_messages_from_task_empty_artifacts(a2a_agent: A2AAgent) -> None:
"""Test _parse_messages_from_task with task containing no artifacts."""
task = MagicMock()
task.artifacts = None
result = a2a_agent._task_to_chat_messages(task)
result = a2a_agent._parse_messages_from_task(task)
assert len(result) == 0
def test_task_to_chat_messages_with_artifacts(a2a_agent: A2AAgent) -> None:
"""Test _task_to_chat_messages with task containing artifacts."""
def test_parse_messages_from_task_with_artifacts(a2a_agent: A2AAgent) -> None:
"""Test _parse_messages_from_task with task containing artifacts."""
task = MagicMock()
# Create mock artifacts
@@ -232,7 +232,7 @@ def test_task_to_chat_messages_with_artifacts(a2a_agent: A2AAgent) -> None:
task.artifacts = [artifact1, artifact2]
result = a2a_agent._task_to_chat_messages(task)
result = a2a_agent._parse_messages_from_task(task)
assert len(result) == 2
assert result[0].text == "Content 1"
@@ -240,8 +240,8 @@ def test_task_to_chat_messages_with_artifacts(a2a_agent: A2AAgent) -> None:
assert all(msg.role == Role.ASSISTANT for msg in result)
def test_artifact_to_chat_message(a2a_agent: A2AAgent) -> None:
"""Test _artifact_to_chat_message conversion."""
def test_parse_message_from_artifact(a2a_agent: A2AAgent) -> None:
"""Test _parse_message_from_artifact conversion."""
artifact = MagicMock()
artifact.artifact_id = "test-artifact"
@@ -253,7 +253,7 @@ def test_artifact_to_chat_message(a2a_agent: A2AAgent) -> None:
artifact.parts = [text_part]
result = a2a_agent._artifact_to_chat_message(artifact)
result = a2a_agent._parse_message_from_artifact(artifact)
assert isinstance(result, ChatMessage)
assert result.role == Role.ASSISTANT
@@ -276,7 +276,7 @@ def test_get_uri_data_invalid_uri() -> None:
_get_uri_data("not-a-valid-data-uri")
def test_a2a_parts_to_contents_conversion(a2a_agent: A2AAgent) -> None:
def test_parse_contents_from_a2a_conversion(a2a_agent: A2AAgent) -> None:
"""Test A2A parts to contents conversion."""
agent = A2AAgent(name="Test Agent", client=MockA2AClient(), _http_client=None)
@@ -285,7 +285,7 @@ def test_a2a_parts_to_contents_conversion(a2a_agent: A2AAgent) -> None:
parts = [Part(root=TextPart(text="First part")), Part(root=TextPart(text="Second part"))]
# Convert to contents
contents = agent._a2a_parts_to_contents(parts)
contents = agent._parse_contents_from_a2a(parts)
# Verify conversion
assert len(contents) == 2
@@ -295,30 +295,30 @@ def test_a2a_parts_to_contents_conversion(a2a_agent: A2AAgent) -> None:
assert contents[1].text == "Second part"
def test_chat_message_to_a2a_message_with_error_content(a2a_agent: A2AAgent) -> None:
"""Test _chat_message_to_a2a_message with ErrorContent."""
def test_prepare_message_for_a2a_with_error_content(a2a_agent: A2AAgent) -> None:
"""Test _prepare_message_for_a2a with ErrorContent."""
# Create ChatMessage with ErrorContent
error_content = ErrorContent(message="Test error message")
message = ChatMessage(role=Role.USER, contents=[error_content])
# Convert to A2A message
a2a_message = a2a_agent._chat_message_to_a2a_message(message)
a2a_message = a2a_agent._prepare_message_for_a2a(message)
# Verify conversion
assert len(a2a_message.parts) == 1
assert a2a_message.parts[0].root.text == "Test error message"
def test_chat_message_to_a2a_message_with_uri_content(a2a_agent: A2AAgent) -> None:
"""Test _chat_message_to_a2a_message with UriContent."""
def test_prepare_message_for_a2a_with_uri_content(a2a_agent: A2AAgent) -> None:
"""Test _prepare_message_for_a2a with UriContent."""
# Create ChatMessage with UriContent
uri_content = UriContent(uri="http://example.com/file.pdf", media_type="application/pdf")
message = ChatMessage(role=Role.USER, contents=[uri_content])
# Convert to A2A message
a2a_message = a2a_agent._chat_message_to_a2a_message(message)
a2a_message = a2a_agent._prepare_message_for_a2a(message)
# Verify conversion
assert len(a2a_message.parts) == 1
@@ -326,15 +326,15 @@ def test_chat_message_to_a2a_message_with_uri_content(a2a_agent: A2AAgent) -> No
assert a2a_message.parts[0].root.file.mime_type == "application/pdf"
def test_chat_message_to_a2a_message_with_data_content(a2a_agent: A2AAgent) -> None:
"""Test _chat_message_to_a2a_message with DataContent."""
def test_prepare_message_for_a2a_with_data_content(a2a_agent: A2AAgent) -> None:
"""Test _prepare_message_for_a2a with DataContent."""
# Create ChatMessage with DataContent (base64 data URI)
data_content = DataContent(uri="data:text/plain;base64,SGVsbG8gV29ybGQ=", media_type="text/plain")
message = ChatMessage(role=Role.USER, contents=[data_content])
# Convert to A2A message
a2a_message = a2a_agent._chat_message_to_a2a_message(message)
a2a_message = a2a_agent._prepare_message_for_a2a(message)
# Verify conversion
assert len(a2a_message.parts) == 1
@@ -342,14 +342,14 @@ def test_chat_message_to_a2a_message_with_data_content(a2a_agent: A2AAgent) -> N
assert a2a_message.parts[0].root.file.mime_type == "text/plain"
def test_chat_message_to_a2a_message_empty_contents_raises_error(a2a_agent: A2AAgent) -> None:
"""Test _chat_message_to_a2a_message with empty contents raises ValueError."""
def test_prepare_message_for_a2a_empty_contents_raises_error(a2a_agent: A2AAgent) -> None:
"""Test _prepare_message_for_a2a with empty contents raises ValueError."""
# Create ChatMessage with no contents
message = ChatMessage(role=Role.USER, contents=[])
# Should raise ValueError for empty contents
with raises(ValueError, match="ChatMessage.contents is empty"):
a2a_agent._chat_message_to_a2a_message(message)
a2a_agent._prepare_message_for_a2a(message)
async def test_run_stream_with_message_response(a2a_agent: A2AAgent, mock_a2a_client: MockA2AClient) -> None:
@@ -405,7 +405,7 @@ async def test_context_manager_no_cleanup_when_no_http_client() -> None:
pass
def test_chat_message_to_a2a_message_with_multiple_contents() -> None:
def test_prepare_message_for_a2a_with_multiple_contents() -> None:
"""Test conversion of ChatMessage with multiple contents."""
agent = A2AAgent(client=MagicMock(), _http_client=None)
@@ -421,7 +421,7 @@ def test_chat_message_to_a2a_message_with_multiple_contents() -> None:
],
)
result = agent._chat_message_to_a2a_message(message)
result = agent._prepare_message_for_a2a(message)
# Should have converted all 4 contents to parts
assert len(result.parts) == 4
@@ -433,7 +433,7 @@ def test_chat_message_to_a2a_message_with_multiple_contents() -> None:
assert result.parts[3].root.kind == "text" # JSON text remains as text (no parsing)
def test_a2a_parts_to_contents_with_data_part() -> None:
def test_parse_contents_from_a2a_with_data_part() -> None:
"""Test conversion of A2A DataPart."""
agent = A2AAgent(client=MagicMock(), _http_client=None)
@@ -441,7 +441,7 @@ def test_a2a_parts_to_contents_with_data_part() -> None:
# Create DataPart
data_part = Part(root=DataPart(data={"key": "value", "number": 42}, metadata={"source": "test"}))
contents = agent._a2a_parts_to_contents([data_part])
contents = agent._parse_contents_from_a2a([data_part])
assert len(contents) == 1
@@ -450,7 +450,7 @@ def test_a2a_parts_to_contents_with_data_part() -> None:
assert contents[0].additional_properties == {"source": "test"}
def test_a2a_parts_to_contents_unknown_part_kind() -> None:
def test_parse_contents_from_a2a_unknown_part_kind() -> None:
"""Test error handling for unknown A2A part kind."""
agent = A2AAgent(client=MagicMock(), _http_client=None)
@@ -459,10 +459,10 @@ def test_a2a_parts_to_contents_unknown_part_kind() -> None:
mock_part.root.kind = "unknown_kind"
with raises(ValueError, match="Unknown Part kind: unknown_kind"):
agent._a2a_parts_to_contents([mock_part])
agent._parse_contents_from_a2a([mock_part])
def test_chat_message_to_a2a_message_with_hosted_file() -> None:
def test_prepare_message_for_a2a_with_hosted_file() -> None:
"""Test conversion of ChatMessage with HostedFileContent to A2A message."""
agent = A2AAgent(client=MagicMock(), _http_client=None)
@@ -473,7 +473,7 @@ def test_chat_message_to_a2a_message_with_hosted_file() -> None:
contents=[HostedFileContent(file_id="hosted://storage/document.pdf")],
)
result = agent._chat_message_to_a2a_message(message) # noqa: SLF001
result = agent._prepare_message_for_a2a(message) # noqa: SLF001
# Verify the conversion
assert len(result.parts) == 1
@@ -488,7 +488,7 @@ def test_chat_message_to_a2a_message_with_hosted_file() -> None:
assert part.root.file.mime_type is None # HostedFileContent doesn't specify media_type
def test_a2a_parts_to_contents_with_hosted_file_uri() -> None:
def test_parse_contents_from_a2a_with_hosted_file_uri() -> None:
"""Test conversion of A2A FilePart with hosted file URI back to UriContent."""
agent = A2AAgent(client=MagicMock(), _http_client=None)
@@ -503,7 +503,7 @@ def test_a2a_parts_to_contents_with_hosted_file_uri() -> None:
)
)
contents = agent._a2a_parts_to_contents([file_part]) # noqa: SLF001
contents = agent._parse_contents_from_a2a([file_part]) # noqa: SLF001
assert len(contents) == 1