Merge remote-tracking branch 'fork/doc/helper-utils' into doc/helper-utils

Author: ahoblitz

src/a2a/grpc/a2a_pb2.py

@@ -488,14 +488,12 @@ class SendMessageRequest(_message.Message):
     def __init__(self, request: _Optional[_Union[Message, _Mapping]] = ..., configuration: _Optional[_Union[SendMessageConfiguration, _Mapping]] = ..., metadata: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...
 
 class GetTaskRequest(_message.Message):
-    __slots__ = ("name", "history_length", "metadata")
+    __slots__ = ("name", "history_length")
     NAME_FIELD_NUMBER: _ClassVar[int]
     HISTORY_LENGTH_FIELD_NUMBER: _ClassVar[int]
-    METADATA_FIELD_NUMBER: _ClassVar[int]
     name: str
     history_length: int
-    metadata: _struct_pb2.Struct
-    def __init__(self, name: _Optional[str] = ..., history_length: _Optional[int] = ..., metadata: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...
+    def __init__(self, name: _Optional[str] = ..., history_length: _Optional[int] = ...) -> None: ...
 
 class CancelTaskRequest(_message.Message):
     __slots__ = ("name",)

src/a2a/grpc/a2a_pb2.pyi

@@ -46,14 +46,7 @@ def metadata(
     ) -> struct_pb2.Struct | None:
         if metadata is None:
             return None
-        return struct_pb2.Struct(
-            # TODO: Add support for other types.
-            fields={
-                key: struct_pb2.Value(string_value=value)
-                for key, value in metadata.items()
-                if isinstance(value, str)
-            }
-        )
+        return dict_to_struct(metadata)
 
     @classmethod
     def part(cls, part: types.Part) -> a2a_pb2.Part:
@@ -324,6 +317,23 @@ def capabilities(
         return a2a_pb2.AgentCapabilities(
             streaming=bool(capabilities.streaming),
             push_notifications=bool(capabilities.push_notifications),
+            extensions=[
+                cls.extension(x) for x in capabilities.extensions or []
+            ],
+        )
+
+    @classmethod
+    def extension(
+        cls,
+        extension: types.AgentExtension,
+    ) -> a2a_pb2.AgentExtension:
+        return a2a_pb2.AgentExtension(
+            uri=extension.uri,
+            description=extension.description,
+            params=dict_to_struct(extension.params)
+            if extension.params
+            else None,
+            required=extension.required,
         )
 
     @classmethod
@@ -477,11 +487,9 @@ def message(cls, message: a2a_pb2.Message) -> types.Message:
 
     @classmethod
     def metadata(cls, metadata: struct_pb2.Struct) -> dict[str, Any]:
-        return {
-            key: value.string_value
-            for key, value in metadata.fields.items()
-            if value.string_value
-        }
+        if not metadata.fields:
+            return {}
+        return json_format.MessageToDict(metadata)
 
     @classmethod
     def part(cls, part: a2a_pb2.Part) -> types.Part:
@@ -767,7 +775,7 @@ def task_query_params(
             if request.history_length
             else None,
             id=m.group(1),
-            metadata=request.metadata,
+            metadata=None,
         )
 
     @classmethod
@@ -777,6 +785,21 @@ def capabilities(
         return types.AgentCapabilities(
             streaming=capabilities.streaming,
             push_notifications=capabilities.push_notifications,
+            extensions=[
+                cls.agent_extension(x) for x in capabilities.extensions
+            ],
+        )
+
+    @classmethod
+    def agent_extension(
+        cls,
+        extension: a2a_pb2.AgentExtension,
+    ) -> types.AgentExtension:
+        return types.AgentExtension(
+            uri=extension.uri,
+            description=extension.description,
+            params=json_format.MessageToDict(extension.params),
+            required=extension.required,
         )
 
     @classmethod
@@ -916,3 +939,25 @@ def role(cls, role: a2a_pb2.Role) -> types.Role:
                 return types.Role.agent
             case _:
                 return types.Role.agent
+
+
+def dict_to_struct(dictionary: dict[str, Any]) -> struct_pb2.Struct:
+    """Converts a Python dict to a Struct proto.
+
+    Unfortunately, using `json_format.ParseDict` does not work because this
+    wants the dictionary to be an exact match of the Struct proto with fields
+    and keys and values, not the traditional Python dict structure.
+
+    Args:
+      dictionary: The Python dict to convert.
+
+    Returns:
+      The Struct proto.
+    """
+    struct = struct_pb2.Struct()
+    for key, val in dictionary.items():
+        if isinstance(val, dict):
+            struct[key] = dict_to_struct(val)
+        else:
+            struct[key] = val
+    return struct

src/a2a/utils/proto_utils.py

@@ -28,22 +28,10 @@ def new_task(request: Message) -> Task:
         if isinstance(part.root, TextPart) and not part.root.text:
             raise ValueError('TextPart content cannot be empty')
 
-    context_id_str = request.context_id
-    if context_id_str is not None:
-        try:
-            uuid.UUID(context_id_str)
-            context_id = context_id_str
-        except (ValueError, AttributeError, TypeError) as e:
-            raise ValueError(
-                f"Invalid context_id: '{context_id_str}' is not a valid UUID."
-            ) from e
-    else:
-        context_id = str(uuid.uuid4())
-
     return Task(
         status=TaskStatus(state=TaskState.submitted),
-        id=(request.task_id if request.task_id else str(uuid.uuid4())),
-        context_id=context_id,
+        id=request.task_id or str(uuid.uuid4()),
+        context_id=request.context_id or str(uuid.uuid4()),
         history=[request],
     )
 

src/a2a/utils/task.py

@@ -188,24 +188,6 @@ def test_completed_task_invalid_artifact_type(self):
                 history=[],
             )
 
-    def test_new_task_with_invalid_context_id(self):
-        """Test that new_task raises a ValueError for various invalid context_id formats."""
-        invalid_ids = ['not-a-uuid', '']
-        for invalid_id in invalid_ids:
-            with self.subTest(invalid_id=invalid_id):
-                with pytest.raises(
-                    ValueError,
-                    match=f"Invalid context_id: '{invalid_id}' is not a valid UUID.",
-                ):
-                    new_task(
-                        Message(
-                            role=Role.user,
-                            parts=[Part(root=TextPart(text='test message'))],
-                            message_id=str(uuid.uuid4()),
-                            context_id=invalid_id,
-                        )
-                    )
-
 
 if __name__ == '__main__':
     unittest.main()