Merge branch 'main' into kthota/typechanges

Author: kthota-g

src/a2a/client/client.py

@@ -24,6 +24,7 @@
     SetTaskPushNotificationConfigRequest,
     SetTaskPushNotificationConfigResponse,
 )
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
 class A2ACardResolver:
@@ -59,6 +60,7 @@ async def get_agent_card(
             ) from e
 
 
+@trace_class(kind=SpanKind.CLIENT)
 class A2AClient:
     """A2A Client."""
 
@@ -111,7 +113,7 @@ async def send_message_streaming(
         request: SendStreamingMessageRequest,
         *,
         http_kwargs: dict[str, Any] | None = None,
-    ) -> AsyncGenerator[SendStreamingMessageResponse, None]:
+    ) -> AsyncGenerator[SendStreamingMessageResponse]:
         if not request.id:
             request.id = str(uuid4())
 

src/a2a/server/agent_execution/context.py

@@ -4,6 +4,7 @@
     InvalidParamsError,
     Message,
     MessageSendParams,
+    MessageSendConfiguration,
     Task,
 )
 from a2a.utils import get_message_text
@@ -79,6 +80,10 @@ def task_id(self) -> str | None:
     def context_id(self) -> str | None:
         return self._context_id
 
+    @property
+    def configuration(self) -> MessageSendConfiguration | None:
+        return self._params.configuration
+
     def _check_or_generate_task_id(self) -> None:
         if not self._params:
             return

src/a2a/server/events/event_consumer.py

@@ -12,11 +12,13 @@
     TaskStatusUpdateEvent,
 )
 from a2a.utils.errors import ServerError
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
 logger = logging.getLogger(__name__)
 
 
+@trace_class(kind=SpanKind.SERVER)
 class EventConsumer:
     """Consumer to read events from the agent event queue."""
 

src/a2a/server/events/event_queue.py

@@ -11,6 +11,7 @@
     TaskArtifactUpdateEvent,
     TaskStatusUpdateEvent,
 )
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
 logger = logging.getLogger(__name__)
@@ -26,6 +27,7 @@
 )
 
 
+@trace_class(kind=SpanKind.SERVER)
 class EventQueue:
     """Event queue for A2A responses from agent."""
 

src/a2a/server/events/in_memory_queue_manager.py

@@ -6,8 +6,10 @@
     QueueManager,
     TaskQueueExists,
 )
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
+@trace_class(kind=SpanKind.SERVER)
 class InMemoryQueueManager(QueueManager):
     """InMemoryQueueManager is used for a single binary management.
 

src/a2a/server/request_handlers/default_request_handler.py

@@ -27,11 +27,13 @@
     UnsupportedOperationError,
 )
 from a2a.utils.errors import ServerError
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
 logger = logging.getLogger(__name__)
 
 
+@trace_class(kind=SpanKind.SERVER)
 class DefaultRequestHandler(RequestHandler):
     """Default request handler for all incoming requests."""
 
@@ -128,7 +130,7 @@ async def on_message_send(
         # agents.
         queue = await self._queue_manager.create_or_tap(task_id)
         result_aggregator = ResultAggregator(task_manager)
-        # TODO to manage the non-blocking flows.
+        # TODO: to manage the non-blocking flows.
         producer_task = asyncio.create_task(
             self._run_event_stream(
                 request_context,

src/a2a/server/request_handlers/jsonrpc_handler.py

@@ -36,13 +36,15 @@
 )
 from a2a.utils.errors import ServerError
 from a2a.utils.helpers import validate
+from a2a.utils.telemetry import SpanKind, trace_class
 
 
 logger = logging.getLogger(__name__)
 
 
+@trace_class(kind=SpanKind.SERVER)
 class JSONRPCHandler:
-    """A handler that maps the JSONRPC Objects to the request handler and back."""
+    """Maps the JSONRPC Objects to the request handler and back."""
 
     def __init__(
         self,
@@ -53,7 +55,7 @@ def __init__(
 
         Args:
             agent_card: The AgentCard describing the agent's capabilities.
-            request_handler: The handler instance responsible for processing A2A requests.
+            request_handler: The handler instance to process A2A requests.
         """
         self.agent_card = agent_card
         self.request_handler = request_handler

src/a2a/server/tasks/task_updater.py

@@ -42,6 +42,7 @@ def add_artifact(
         parts: list[Part],
         artifact_id=str(uuid.uuid4()),
         name: str | None = None,
+        metadata: dict[str, any] | None = None,
     ):
         """Add an artifact to the task."""
         self.event_queue.enqueue_event(
@@ -52,6 +53,7 @@ def add_artifact(
                     artifactId=artifact_id,
                     name=name,
                     parts=parts,
+                    metadata=metadata,
                 ),
             )
         )
@@ -64,6 +66,14 @@ def complete(self, message: Message | None = None):
             final=True,
         )
 
+    def failed(self, message: Message | None = None):
+        """Mark the task as failed."""
+        self.update_status(
+            TaskState.failed,
+            message=message,
+            final=True
+        )
+
     def submit(self, message: Message | None = None):
         """Mark the task as submitted."""
         self.update_status(

src/a2a/utils/__init__.py

@@ -1,15 +1,24 @@
-from a2a.utils.artifact import new_text_artifact
+from a2a.utils.artifact import (
+    new_artifact,
+    new_data_artifact,
+    new_text_artifact,
+)
 from a2a.utils.helpers import (
     append_artifact_to_task,
     build_text_artifact,
     create_task_obj,
+    are_modalities_compatible,
 )
 from a2a.utils.message import (
     get_message_text,
     get_text_parts,
     new_agent_text_message,
+    new_agent_parts_message,
+)
+from a2a.utils.task import (
+    new_task,
+    completed_task,
 )
-from a2a.utils.task import new_task
 
 
 __all__ = [
@@ -21,4 +30,9 @@
     'new_agent_text_message',
     'new_task',
     'new_text_artifact',
+    'new_agent_parts_message',
+    'completed_task',
+    'new_artifact',
+    'new_data_artifact',
+    'are_modalities_compatible',
 ]

src/a2a/utils/artifact.py

@@ -3,14 +3,30 @@
 from a2a.types import Artifact, Part, TextPart
 
 
-def new_text_artifact(
-    name: str,
-    text: str,
-    description: str = '',
+def new_artifact(
+    parts: list[Part], name: str, description: str = ''
 ) -> Artifact:
     return Artifact(
         artifactId=str(uuid.uuid4()),
-        parts=[Part(root=TextPart(text=text))],
+        parts=parts,
         name=name,
         description=description,
     )
+
+def new_text_artifact(
+    name: str,
+    text: str,
+    description: str = '',
+) -> Artifact:
+    return new_artifact(
+        [Part(root=TextPart(text=text))], name, description,
+    )
+
+def new_data_artifact(
+    name: str,
+    data: dict[str, any],
+    description: str = '',
+):
+    return new_artifact(
+        [Part(root=DataPart(data=data))], name, description,
+    )