Open Telemetry for A2A-python-sdk, enable trace

Author: rajeshvelicheti

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."""
 

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

@@ -28,11 +28,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."""
 
@@ -112,7 +114,7 @@ async def on_message_send(
         else:
             queue = EventQueue()
         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(
                 RequestContext(
@@ -128,7 +130,8 @@ async def on_message_send(
         try:
             consumer = EventConsumer(queue)
 
-            # TODO - register the queue for the task upon the first sign it is a task.
+            # TODO: - register the queue for the task upon
+            # the first sign it is a task.
             result = await result_aggregator.consume_all(consumer)
             if not result:
                 raise ServerError(error=InternalError())

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/utils/__init__.py

@@ -1,3 +1,6 @@
+# type: ignore
+from opentelemetry.trace import SpanKind as OTelSpanKind
+
 from a2a.utils.artifact import new_text_artifact
 from a2a.utils.helpers import (
     append_artifact_to_task,
@@ -12,7 +15,9 @@
 from a2a.utils.task import new_task
 
 
+SpanKind = OTelSpanKind
 __all__ = [
+    'SpanKind',
     'append_artifact_to_task',
     'build_text_artifact',
     'create_task_obj',

src/a2a/utils/helpers.py

@@ -13,11 +13,13 @@
     TextPart,
 )
 from a2a.utils.errors import ServerError, UnsupportedOperationError
+from a2a.utils.telemetry import trace_function
 
 
 logger = logging.getLogger(__name__)
 
 
+@trace_function()
 def create_task_obj(message_send_params: MessageSendParams) -> Task:
     """Create a new task object from message send params."""
     if not message_send_params.message.contextId:
@@ -31,6 +33,7 @@ def create_task_obj(message_send_params: MessageSendParams) -> Task:
     )
 
 
+@trace_function()
 def append_artifact_to_task(task: Task, event: TaskArtifactUpdateEvent) -> None:
     """Helper method for updating Task with new artifact data."""
     if not task.artifacts:

src/a2a/utils/telemetry.py

@@ -210,7 +210,11 @@ def sync_wrapper(*args, **kwargs):
     return async_wrapper if is_async_func else sync_wrapper
 
 
-def trace_class(include_list: list[str] = None, exclude_list: list[str] = None):
+def trace_class(
+    include_list: list[str] = None,
+    exclude_list: list[str] = None,
+    kind=SpanKind.INTERNAL,
+):
     """A class decorator to automatically trace specified methods of a class.
 
     This decorator iterates over the methods of a class and applies the
@@ -278,7 +282,11 @@ def decorator(cls):
             all_methods[name] = method
             span_name = f'{cls.__module__}.{cls.__name__}.{name}'
             # Set the decorator on the method.
-            setattr(cls, name, trace_function(span_name=span_name)(method))
+            setattr(
+                cls,
+                name,
+                trace_function(span_name=span_name, kind=kind)(method),
+            )
         return cls
 
     return decorator