Skip to content

Commit bdd7e99

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 bdd7e99

1 file changed

Lines changed: 174 additions & 18 deletions

File tree

‎vertexai/_genai/_agent_engines_utils.py‎

Lines changed: 174 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -110,30 +110,54 @@
110110

111111

112112
try:
113-
from a2a.types import (
114-
AgentCard,
115-
TransportProtocol,
116-
Message,
117-
TaskIdParams,
118-
TaskQueryParams,
119-
)
120-
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):
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+
try:
126+
from a2a.types import (
127+
AgentCard,
128+
Message,
129+
)
130+
from a2a.client import ClientConfig, ClientFactory
131+
from a2a.utils.constants import TransportProtocol
132+
from a2a.compat.v0_3.types import TaskIdParams, TaskQueryParams
133+
except (ImportError, AttributeError):
134+
_A2A_SDK_VERSION = None
135+
136+
if _A2A_SDK_VERSION == "0.3":
137+
try:
138+
from a2a.types import (
139+
AgentCard,
140+
TransportProtocol,
141+
Message,
142+
TaskIdParams,
143+
TaskQueryParams,
144+
)
145+
from a2a.client import ClientConfig, ClientFactory
146+
except (ImportError, AttributeError):
147+
_A2A_SDK_VERSION = None
148+
149+
if _A2A_SDK_VERSION is None:
130150
AgentCard = None
131151
TransportProtocol = None
132152
Message = None
133153
ClientConfig = None
134154
ClientFactory = None
135155
TaskIdParams = None
136156
TaskQueryParams = None
157+
SendMessageRequest = None
158+
GetTaskRequest = None
159+
CancelTaskRequest = None
160+
GetExtendedAgentCardRequest = None
137161

138162
_ACTIONS_KEY = "actions"
139163
_ACTION_APPEND = "append"
@@ -1737,7 +1761,9 @@ async def _method(self: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]
17371761
return _method
17381762

17391763

1740-
def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list[Any]]:
1764+
def _wrap_a2a_operation_v03(
1765+
method_name: str, agent_card: str
1766+
) -> Callable[..., list[Any]]:
17411767
"""Wraps an Agent Engine method, creating a callable for A2A API.
17421768
17431769
Args:
@@ -1854,6 +1880,136 @@ async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def]
18541880
return _method # type: ignore[return-value]
18551881

18561882

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

0 commit comments

Comments
 (0)