Skip to content

Commit a3e7e1e

Browse files
feat: Update proto conversion utilities (#424)
# Proto Util improvements - Improves metadata data support in proto conversions (supports more than just `map<string, string>` now and supports full json object `map<string, any>` - Adds support to convert extension data on agent cards
1 parent 39c6b43 commit a3e7e1e

File tree

1 file changed

+58
-13
lines changed

1 file changed

+58
-13
lines changed

src/a2a/utils/proto_utils.py

Lines changed: 58 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -46,14 +46,7 @@ def metadata(
4646
) -> struct_pb2.Struct | None:
4747
if metadata is None:
4848
return None
49-
return struct_pb2.Struct(
50-
# TODO: Add support for other types.
51-
fields={
52-
key: struct_pb2.Value(string_value=value)
53-
for key, value in metadata.items()
54-
if isinstance(value, str)
55-
}
56-
)
49+
return dict_to_struct(metadata)
5750

5851
@classmethod
5952
def part(cls, part: types.Part) -> a2a_pb2.Part:
@@ -324,6 +317,23 @@ def capabilities(
324317
return a2a_pb2.AgentCapabilities(
325318
streaming=bool(capabilities.streaming),
326319
push_notifications=bool(capabilities.push_notifications),
320+
extensions=[
321+
cls.extension(x) for x in capabilities.extensions or []
322+
],
323+
)
324+
325+
@classmethod
326+
def extension(
327+
cls,
328+
extension: types.AgentExtension,
329+
) -> a2a_pb2.AgentExtension:
330+
return a2a_pb2.AgentExtension(
331+
uri=extension.uri,
332+
description=extension.description,
333+
params=dict_to_struct(extension.params)
334+
if extension.params
335+
else None,
336+
required=extension.required,
327337
)
328338

329339
@classmethod
@@ -477,11 +487,9 @@ def message(cls, message: a2a_pb2.Message) -> types.Message:
477487

478488
@classmethod
479489
def metadata(cls, metadata: struct_pb2.Struct) -> dict[str, Any]:
480-
return {
481-
key: value.string_value
482-
for key, value in metadata.fields.items()
483-
if value.string_value
484-
}
490+
if not metadata.fields:
491+
return {}
492+
return json_format.MessageToDict(metadata)
485493

486494
@classmethod
487495
def part(cls, part: a2a_pb2.Part) -> types.Part:
@@ -777,6 +785,21 @@ def capabilities(
777785
return types.AgentCapabilities(
778786
streaming=capabilities.streaming,
779787
push_notifications=capabilities.push_notifications,
788+
extensions=[
789+
cls.agent_extension(x) for x in capabilities.extensions
790+
],
791+
)
792+
793+
@classmethod
794+
def agent_extension(
795+
cls,
796+
extension: a2a_pb2.AgentExtension,
797+
) -> types.AgentExtension:
798+
return types.AgentExtension(
799+
uri=extension.uri,
800+
description=extension.description,
801+
params=json_format.MessageToDict(extension.params),
802+
required=extension.required,
780803
)
781804

782805
@classmethod
@@ -916,3 +939,25 @@ def role(cls, role: a2a_pb2.Role) -> types.Role:
916939
return types.Role.agent
917940
case _:
918941
return types.Role.agent
942+
943+
944+
def dict_to_struct(dictionary: dict[str, Any]) -> struct_pb2.Struct:
945+
"""Converts a Python dict to a Struct proto.
946+
947+
Unfortunately, using `json_format.ParseDict` does not work because this
948+
wants the dictionary to be an exact match of the Struct proto with fields
949+
and keys and values, not the traditional Python dict structure.
950+
951+
Args:
952+
dictionary: The Python dict to convert.
953+
954+
Returns:
955+
The Struct proto.
956+
"""
957+
struct = struct_pb2.Struct()
958+
for key, val in dictionary.items():
959+
if isinstance(val, dict):
960+
struct[key] = dict_to_struct(val)
961+
else:
962+
struct[key] = val
963+
return struct

0 commit comments

Comments
 (0)