mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: basic python a2a support (#906)
* basic python a2a support * fixes * small fixes --------- Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
fe6492ce60
commit
c102706146
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import importlib.metadata
|
||||
|
||||
from ._agent import A2AAgent
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0" # Fallback for development mode
|
||||
|
||||
__all__ = [
|
||||
"A2AAgent",
|
||||
"__version__",
|
||||
]
|
||||
@@ -0,0 +1,383 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import AsyncIterable, Sequence
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from a2a.client import Client, ClientConfig, ClientFactory, minimal_agent_card
|
||||
from a2a.types import (
|
||||
AgentCard,
|
||||
Artifact,
|
||||
FilePart,
|
||||
FileWithBytes,
|
||||
FileWithUri,
|
||||
Message,
|
||||
Task,
|
||||
TaskState,
|
||||
TextPart,
|
||||
TransportProtocol,
|
||||
)
|
||||
from a2a.types import Message as A2AMessage
|
||||
from a2a.types import Part as A2APart
|
||||
from a2a.types import Role as A2ARole
|
||||
from agent_framework import (
|
||||
AgentRunResponse,
|
||||
AgentRunResponseUpdate,
|
||||
AgentThread,
|
||||
BaseAgent,
|
||||
ChatMessage,
|
||||
Contents,
|
||||
DataContent,
|
||||
Role,
|
||||
TextContent,
|
||||
UriContent,
|
||||
prepend_agent_framework_to_user_agent,
|
||||
)
|
||||
|
||||
__all__ = ["A2AAgent"]
|
||||
|
||||
URI_PATTERN = re.compile(r"^data:(?P<media_type>[^;]+);base64,(?P<base64_data>[A-Za-z0-9+/=]+)$")
|
||||
TERMINAL_TASK_STATES = [
|
||||
TaskState.completed,
|
||||
TaskState.failed,
|
||||
TaskState.canceled,
|
||||
TaskState.rejected,
|
||||
]
|
||||
|
||||
|
||||
def _get_uri_data(uri: str) -> str:
|
||||
match = URI_PATTERN.match(uri)
|
||||
if not match:
|
||||
raise ValueError(f"Invalid data URI format: {uri}")
|
||||
|
||||
return match.group("base64_data")
|
||||
|
||||
|
||||
class A2AAgent(BaseAgent):
|
||||
"""Agent-to-Agent (A2A) protocol implementation.
|
||||
|
||||
Wraps an A2A Client to connect the Agent Framework with external A2A-compliant agents
|
||||
via HTTP/JSON-RPC. Converts framework ChatMessages to A2A Messages on send, and converts
|
||||
A2A responses (Messages/Tasks) back to framework types. Inherits BaseAgent capabilities
|
||||
while managing the underlying A2A protocol communication.
|
||||
|
||||
Can be initialized with a URL, AgentCard, or existing A2A Client instance.
|
||||
"""
|
||||
|
||||
client: Client
|
||||
_http_client: httpx.AsyncClient | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
name: str | None = None,
|
||||
id: str | None = None,
|
||||
description: str | None = None,
|
||||
agent_card: AgentCard | None = None,
|
||||
url: str | None = None,
|
||||
client: Client | None = None,
|
||||
http_client: httpx.AsyncClient | None = None,
|
||||
) -> None:
|
||||
"""Initialize the A2AAgent.
|
||||
|
||||
Args:
|
||||
name: The name of the agent.
|
||||
id: The unique identifier for the agent, will be created automatically if not provided.
|
||||
description: A brief description of the agent's purpose.
|
||||
agent_card: The agent card for the agent.
|
||||
url: The URL for the A2A server.
|
||||
client: The A2A client for the agent.
|
||||
http_client: Optional httpx.AsyncClient to use.
|
||||
"""
|
||||
if client is None:
|
||||
if agent_card is None:
|
||||
if url is None:
|
||||
raise ValueError("Either agent_card or url must be provided")
|
||||
# Create minimal agent card from URL
|
||||
agent_card = minimal_agent_card(url, [TransportProtocol.jsonrpc])
|
||||
|
||||
# Create or use provided httpx client
|
||||
if http_client is None:
|
||||
timeout = httpx.Timeout(
|
||||
connect=10.0, # 10 seconds to establish connection
|
||||
read=60.0, # 60 seconds to read response (A2A operations can take time)
|
||||
write=10.0, # 10 seconds to send request
|
||||
pool=5.0, # 5 seconds to get connection from pool
|
||||
)
|
||||
headers = prepend_agent_framework_to_user_agent()
|
||||
http_client = httpx.AsyncClient(timeout=timeout, headers=headers)
|
||||
self._http_client = http_client # Store for cleanup
|
||||
|
||||
# Create A2A client using factory
|
||||
config = ClientConfig(
|
||||
httpx_client=http_client,
|
||||
supported_transports=[TransportProtocol.jsonrpc],
|
||||
)
|
||||
factory = ClientFactory(config)
|
||||
client = factory.create(agent_card)
|
||||
|
||||
args: dict[str, Any] = {"client": client}
|
||||
if name:
|
||||
args["name"] = name
|
||||
if id:
|
||||
args["id"] = id
|
||||
if description:
|
||||
args["description"] = description
|
||||
super().__init__(**args)
|
||||
|
||||
async def __aenter__(self) -> "A2AAgent":
|
||||
"""Async context manager entry."""
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: Any,
|
||||
) -> None:
|
||||
"""Async context manager exit with httpx client cleanup."""
|
||||
# Close our httpx client if we created it
|
||||
if self._http_client is not None:
|
||||
await self._http_client.aclose()
|
||||
|
||||
async def run(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentRunResponse:
|
||||
"""Get a response from the agent.
|
||||
|
||||
This method returns the final result of the agent's execution
|
||||
as a single AgentRunResponse object. The caller is blocked until
|
||||
the final result is available.
|
||||
|
||||
Args:
|
||||
messages: The message(s) to send to the agent.
|
||||
thread: The conversation thread associated with the message(s).
|
||||
kwargs: Additional keyword arguments.
|
||||
|
||||
Returns:
|
||||
An agent response item.
|
||||
"""
|
||||
# Collect all updates and use framework to consolidate updates into response
|
||||
updates = [update async for update in self.run_stream(messages, thread=thread, **kwargs)]
|
||||
return AgentRunResponse.from_agent_run_response_updates(updates)
|
||||
|
||||
async def run_stream(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
"""Run the agent as a stream.
|
||||
|
||||
This method will return the intermediate steps and final results of the
|
||||
agent's execution as a stream of AgentRunResponseUpdate objects to the caller.
|
||||
|
||||
Args:
|
||||
messages: The message(s) to send to the agent.
|
||||
thread: The conversation thread associated with the message(s).
|
||||
kwargs: Additional keyword arguments.
|
||||
|
||||
Yields:
|
||||
An agent response item.
|
||||
"""
|
||||
messages = self._normalize_messages(messages)
|
||||
a2a_message = self._chat_message_to_a2a_message(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)
|
||||
yield AgentRunResponseUpdate(
|
||||
contents=contents,
|
||||
role=Role.ASSISTANT if item.role == A2ARole.agent else Role.USER,
|
||||
response_id=str(getattr(item, "message_id", uuid.uuid4())),
|
||||
raw_representation=item,
|
||||
)
|
||||
elif isinstance(item, tuple) and len(item) == 2: # ClientEvent = (Task, UpdateEvent)
|
||||
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)
|
||||
if task_messages:
|
||||
for message in task_messages:
|
||||
# Use the artifact's ID from raw_representation as message_id for unique identification
|
||||
artifact_id = getattr(message.raw_representation, "artifact_id", None)
|
||||
yield AgentRunResponseUpdate(
|
||||
contents=message.contents,
|
||||
role=message.role,
|
||||
response_id=task.id,
|
||||
message_id=artifact_id,
|
||||
raw_representation=task,
|
||||
)
|
||||
else:
|
||||
# Empty task
|
||||
yield AgentRunResponseUpdate(
|
||||
contents=[],
|
||||
role=Role.ASSISTANT,
|
||||
response_id=task.id,
|
||||
raw_representation=task,
|
||||
)
|
||||
else:
|
||||
# Unknown response type
|
||||
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.
|
||||
|
||||
Transforms Agent Framework ChatMessage objects into A2A protocol Messages by:
|
||||
- Converting all message contents to appropriate A2A Part types
|
||||
- Mapping text content to TextPart objects
|
||||
- Converting file references (URI/data/hosted_file) to FilePart objects
|
||||
- Preserving metadata and additional properties from the original message
|
||||
- Setting the role to 'user' as framework messages are treated as user input
|
||||
"""
|
||||
parts: list[A2APart] = []
|
||||
if not message.contents:
|
||||
raise ValueError("ChatMessage.contents is empty; cannot convert to A2AMessage.")
|
||||
|
||||
# Process ALL contents
|
||||
for content in message.contents:
|
||||
match content.type:
|
||||
case "text":
|
||||
parts.append(
|
||||
A2APart(
|
||||
root=TextPart(
|
||||
text=content.text,
|
||||
metadata=content.additional_properties,
|
||||
)
|
||||
)
|
||||
)
|
||||
case "error":
|
||||
parts.append(
|
||||
A2APart(
|
||||
root=TextPart(
|
||||
text=content.message or "An error occurred.",
|
||||
metadata=content.additional_properties,
|
||||
)
|
||||
)
|
||||
)
|
||||
case "uri":
|
||||
parts.append(
|
||||
A2APart(
|
||||
root=FilePart(
|
||||
file=FileWithUri(
|
||||
uri=content.uri,
|
||||
mime_type=content.media_type,
|
||||
),
|
||||
metadata=content.additional_properties,
|
||||
)
|
||||
)
|
||||
)
|
||||
case "data":
|
||||
parts.append(
|
||||
A2APart(
|
||||
root=FilePart(
|
||||
file=FileWithBytes(
|
||||
bytes=_get_uri_data(content.uri),
|
||||
mime_type=content.media_type,
|
||||
),
|
||||
metadata=content.additional_properties,
|
||||
)
|
||||
)
|
||||
)
|
||||
case "hosted_file":
|
||||
parts.append(
|
||||
A2APart(
|
||||
root=FilePart(
|
||||
file=FileWithUri(
|
||||
uri=content.file_id,
|
||||
mime_type=None, # HostedFileContent doesn't specify media_type
|
||||
),
|
||||
metadata=content.additional_properties,
|
||||
)
|
||||
)
|
||||
)
|
||||
case _:
|
||||
raise ValueError(f"Unknown content type: {content.type}")
|
||||
|
||||
return A2AMessage(
|
||||
role=A2ARole("user"),
|
||||
parts=parts,
|
||||
message_id=message.message_id or uuid.uuid4().hex,
|
||||
metadata=message.additional_properties,
|
||||
)
|
||||
|
||||
def _a2a_parts_to_contents(self, parts: Sequence[A2APart]) -> list[Contents]:
|
||||
"""Convert A2A Parts to Agent Framework Contents.
|
||||
|
||||
Transforms A2A protocol Parts into framework-native Content objects,
|
||||
handling text, file (URI/bytes), and data parts with metadata preservation.
|
||||
"""
|
||||
contents: list[Contents] = []
|
||||
for part in parts:
|
||||
inner_part = part.root
|
||||
match inner_part.kind:
|
||||
case "text":
|
||||
contents.append(
|
||||
TextContent(
|
||||
text=inner_part.text,
|
||||
additional_properties=inner_part.metadata,
|
||||
raw_representation=inner_part,
|
||||
)
|
||||
)
|
||||
case "file":
|
||||
if isinstance(inner_part.file, FileWithUri):
|
||||
contents.append(
|
||||
UriContent(
|
||||
uri=inner_part.file.uri,
|
||||
media_type=inner_part.file.mime_type or "",
|
||||
additional_properties=inner_part.metadata,
|
||||
raw_representation=inner_part,
|
||||
)
|
||||
)
|
||||
elif isinstance(inner_part.file, FileWithBytes):
|
||||
contents.append(
|
||||
DataContent(
|
||||
data=base64.b64decode(inner_part.file.bytes),
|
||||
media_type=inner_part.file.mime_type or "",
|
||||
additional_properties=inner_part.metadata,
|
||||
raw_representation=inner_part,
|
||||
)
|
||||
)
|
||||
case "data":
|
||||
contents.append(
|
||||
TextContent(
|
||||
text=json.dumps(inner_part.data),
|
||||
additional_properties=inner_part.metadata,
|
||||
raw_representation=inner_part,
|
||||
)
|
||||
)
|
||||
case _:
|
||||
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."""
|
||||
messages: list[ChatMessage] = []
|
||||
|
||||
if task.artifacts is not None:
|
||||
for artifact in task.artifacts:
|
||||
messages.append(self._artifact_to_chat_message(artifact))
|
||||
|
||||
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)
|
||||
return ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
contents=contents,
|
||||
raw_representation=artifact,
|
||||
)
|
||||
Reference in New Issue
Block a user