Skip to content

Commit c4fbada

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
fix: update agent engine utils to support python-a2a sdk 1.0
PiperOrigin-RevId: 913139425
1 parent a99f340 commit c4fbada

1 file changed

Lines changed: 158 additions & 10 deletions

File tree

‎vertexai/_genai/_agent_engines_utils.py‎

Lines changed: 158 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,26 @@
110110

111111

112112
try:
113+
from a2a.utils.constants import TransportProtocol as _A2aVersionTest # noqa: F401
114+
115+
_A2A_SDK_VERSION: Optional[str] = "1.0"
116+
except ImportError:
117+
try:
118+
from a2a.types import TransportProtocol as _A2aVersionTest # noqa: F401
119+
120+
_A2A_SDK_VERSION = "0.3"
121+
except ImportError:
122+
_A2A_SDK_VERSION = None
123+
124+
if _A2A_SDK_VERSION == "1.0":
125+
from a2a.types import (
126+
AgentCard,
127+
Message,
128+
)
129+
from a2a.client import ClientConfig, ClientFactory
130+
from a2a.utils.constants import TransportProtocol
131+
from a2a.compat.v0_3.types import TaskIdParams, TaskQueryParams
132+
elif _A2A_SDK_VERSION == "0.3":
113133
from a2a.types import (
114134
AgentCard,
115135
TransportProtocol,
@@ -118,22 +138,18 @@
118138
TaskQueryParams,
119139
)
120140
from a2a.client import ClientConfig, ClientFactory
121-
122-
AgentCard = AgentCard
123-
TransportProtocol = TransportProtocol
124-
Message = Message
125-
ClientConfig = ClientConfig
126-
ClientFactory = ClientFactory
127-
TaskIdParams = TaskIdParams
128-
TaskQueryParams = TaskQueryParams
129-
except (ImportError, AttributeError):
141+
else:
130142
AgentCard = None
131143
TransportProtocol = None
132144
Message = None
133145
ClientConfig = None
134146
ClientFactory = None
135147
TaskIdParams = None
136148
TaskQueryParams = None
149+
SendMessageRequest = None
150+
GetTaskRequest = None
151+
CancelTaskRequest = None
152+
GetExtendedAgentCardRequest = None
137153

138154
_ACTIONS_KEY = "actions"
139155
_ACTION_APPEND = "append"
@@ -1737,7 +1753,9 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]
17371753
return _method
17381754

17391755

1740-
def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list[Any]]:
1756+
def _wrap_a2a_operation_v03(
1757+
method_name: str, agent_card: str
1758+
) -> Callable[..., list[Any]]:
17411759
"""Wraps an Agent Engine method, creating a callable for A2A API.
17421760
17431761
Args:
@@ -1854,6 +1872,136 @@ async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def]
18541872
return _method # type: ignore[return-value]
18551873

18561874

1875+
def _wrap_a2a_operation_v10(method_name: str, agent_card: str) -> Callable[..., list[Any]]:
1876+
"""Wraps an Agent Engine method, creating a callable for A2A API (v1.0.0+).
1877+
1878+
Args:
1879+
method_name: The name of the Agent Engine method to call.
1880+
agent_card: The agent card JSON string to use for the A2A API call.
1881+
Example:
1882+
{
1883+
'name': 'Sample Agent',
1884+
'description': (
1885+
'A helpful assistant agent that can answer questions.'
1886+
),
1887+
'supportedInterfaces': [{
1888+
'url': 'http://localhost:8080/a2a/rest/',
1889+
'protocolBinding': 'HTTP+JSON',
1890+
'protocolVersion': '1.0',
1891+
}],
1892+
'version': '1.0.0',
1893+
'capabilities': {
1894+
'streaming': True,
1895+
'pushNotifications': False,
1896+
'extendedAgentCard': True,
1897+
},
1898+
'defaultInputModes': ['text'],
1899+
'defaultOutputModes': ['text'],
1900+
'skills': [{
1901+
'id': 'question_answer',
1902+
'name': 'Q&A Agent',
1903+
'description': (
1904+
'A helpful assistant agent that can answer questions.'
1905+
),
1906+
'tags': ['Question-Answer'],
1907+
'examples': [
1908+
'Who is leading 2025 F1 Standings?',
1909+
'Where can i find an active volcano?',
1910+
],
1911+
'inputModes': ['text'],
1912+
'outputModes': ['text'],
1913+
}],
1914+
}
1915+
1916+
Returns:
1917+
A callable object that executes the method on the Agent Engine via
1918+
the A2A API.
1919+
"""
1920+
1921+
async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def]
1922+
if not self.api_client:
1923+
raise ValueError("api_client is not initialized.")
1924+
if not self.api_resource:
1925+
raise ValueError("api_resource is not initialized.")
1926+
1927+
a2a_agent_card = AgentCard()
1928+
json_format.ParseDict(
1929+
json.loads(agent_card), a2a_agent_card, ignore_unknown_fields=True
1930+
)
1931+
1932+
if a2a_agent_card.supported_interfaces:
1933+
interface = a2a_agent_card.supported_interfaces[0]
1934+
if interface.protocol_binding != TransportProtocol.HTTP_JSON:
1935+
raise ValueError(
1936+
"Only HTTP+JSON is supported for preferred transport on agent card"
1937+
)
1938+
else:
1939+
raise ValueError("Agent card does not define any supported interfaces.")
1940+
1941+
base_url = self.api_client._api_client._http_options.base_url.rstrip("/")
1942+
api_version = self.api_client._api_client._http_options.api_version
1943+
a2a_agent_card.supported_interfaces[0].url = (
1944+
f"{base_url}/{api_version}/{self.api_resource.name}/a2a"
1945+
)
1946+
1947+
config = ClientConfig(
1948+
supported_protocol_bindings=[
1949+
TransportProtocol.HTTP_JSON,
1950+
],
1951+
use_client_preference=True,
1952+
httpx_client=httpx.AsyncClient(
1953+
headers={
1954+
"Authorization": (
1955+
f"Bearer {self.api_client._api_client._credentials.token}"
1956+
)
1957+
},
1958+
timeout=(
1959+
self.api_client._api_client._http_options.timeout / 1000.0
1960+
if self.api_client._api_client._http_options.timeout
1961+
else None
1962+
),
1963+
),
1964+
)
1965+
factory = ClientFactory(config)
1966+
client = factory.create(a2a_agent_card)
1967+
1968+
context = kwargs.pop("context", None)
1969+
if context is not None:
1970+
from a2a.client.client import ClientCallContext
1971+
1972+
if not isinstance(context, ClientCallContext):
1973+
actual_context = ClientCallContext()
1974+
if hasattr(context, "state"):
1975+
actual_context.state = context.state
1976+
elif isinstance(context, dict):
1977+
actual_context.state = context
1978+
context = actual_context
1979+
1980+
req = kwargs["request"]
1981+
if method_name == "on_message_send":
1982+
response = client.send_message(req, context=context)
1983+
chunks = []
1984+
async for chunk in response:
1985+
chunks.append(chunk)
1986+
return chunks
1987+
elif method_name == "on_get_task":
1988+
return await client.get_task(req, context=context)
1989+
elif method_name == "on_cancel_task":
1990+
return await client.cancel_task(req, context=context)
1991+
elif method_name == "on_get_extended_agent_card":
1992+
return await client.get_extended_agent_card(req, context=context)
1993+
else:
1994+
raise ValueError(f"Unknown method name: {method_name}")
1995+
1996+
return _method # type: ignore[return-value]
1997+
1998+
1999+
if _A2A_SDK_VERSION == "1.0":
2000+
_wrap_a2a_operation = _wrap_a2a_operation_v10
2001+
else:
2002+
_wrap_a2a_operation = _wrap_a2a_operation_v03
2003+
2004+
18572005
def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]:
18582006
"""Converts the body of the HTTP Response message to JSON format.
18592007

0 commit comments

Comments
 (0)