diff --git a/pyproject.toml b/pyproject.toml index fa629a2f..4d3b81e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,7 @@ requires-python = ">=3.13" dependencies = [ "httpx>=0.28.1", "httpx-sse>=0.4.0", + "mypy>=1.15.0", "opentelemetry-api>=1.33.0", "opentelemetry-sdk>=1.33.0", "pydantic>=2.11.3", diff --git a/src/a2a/client/client.py b/src/a2a/client/client.py index 2f032707..3645cee4 100644 --- a/src/a2a/client/client.py +++ b/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.""" diff --git a/src/a2a/server/events/event_consumer.py b/src/a2a/server/events/event_consumer.py index dcfa9d98..297579bf 100644 --- a/src/a2a/server/events/event_consumer.py +++ b/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.""" diff --git a/src/a2a/server/events/event_queue.py b/src/a2a/server/events/event_queue.py index ec651761..fbe63822 100644 --- a/src/a2a/server/events/event_queue.py +++ b/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.""" diff --git a/src/a2a/server/events/in_memory_queue_manager.py b/src/a2a/server/events/in_memory_queue_manager.py index a0d95f8e..8783ac91 100644 --- a/src/a2a/server/events/in_memory_queue_manager.py +++ b/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. diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py index f656e414..ebb479be 100644 --- a/src/a2a/server/request_handlers/default_request_handler.py +++ b/src/a2a/server/request_handlers/default_request_handler.py @@ -28,10 +28,12 @@ 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, @@ -144,7 +146,7 @@ async def on_message_send( ( result, interrupted, - ) = await result_aggregator.consume_and_break_on_interrupt(consumer) + ) = await result_aggregator.consume_and_break_on_interrupt(consumer) if not result: raise ServerError(error=InternalError()) finally: diff --git a/src/a2a/server/request_handlers/jsonrpc_handler.py b/src/a2a/server/request_handlers/jsonrpc_handler.py index d3192f59..0475ab12 100644 --- a/src/a2a/server/request_handlers/jsonrpc_handler.py +++ b/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 diff --git a/src/a2a/utils/__init__.py b/src/a2a/utils/__init__.py index 42e5d37e..25aafeec 100644 --- a/src/a2a/utils/__init__.py +++ b/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', diff --git a/src/a2a/utils/helpers.py b/src/a2a/utils/helpers.py index c52497ec..691aa616 100644 --- a/src/a2a/utils/helpers.py +++ b/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: diff --git a/src/a2a/utils/telemetry.py b/src/a2a/utils/telemetry.py index 926e3771..cd33843e 100644 --- a/src/a2a/utils/telemetry.py +++ b/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 diff --git a/uv.lock b/uv.lock index e9e8e313..5d721297 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,4 @@ version = 1 -revision = 1 requires-python = ">=3.13" [manifest] @@ -18,6 +17,7 @@ source = { editable = "." } dependencies = [ { name = "httpx" }, { name = "httpx-sse" }, + { name = "mypy" }, { name = "opentelemetry-api" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, @@ -41,6 +41,7 @@ dev = [ requires-dist = [ { name = "httpx", specifier = ">=0.28.1" }, { name = "httpx-sse", specifier = ">=0.4.0" }, + { name = "mypy", specifier = ">=1.15.0" }, { name = "opentelemetry-api", specifier = ">=1.33.0" }, { name = "opentelemetry-sdk", specifier = ">=1.33.0" }, { name = "pydantic", specifier = ">=2.11.3" }, @@ -246,7 +247,7 @@ name = "click" version = "8.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "platform_system == 'Windows'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cd/0f/62ca20172d4f87d93cf89665fbaedcd560ac48b465bd1d92bfc7ea6b0a41/click-8.2.0.tar.gz", hash = "sha256:f5452aeddd9988eefa20f90f05ab66f17fce1ee2a36907fd30b05bbb5953814d", size = 235857 } wheels = [ @@ -1701,6 +1702,7 @@ dependencies = [ { name = "greenlet", marker = "(python_full_version < '3.14' and platform_machine == 'AMD64') or (python_full_version < '3.14' and platform_machine == 'WIN32') or (python_full_version < '3.14' and platform_machine == 'aarch64') or (python_full_version < '3.14' and platform_machine == 'amd64') or (python_full_version < '3.14' and platform_machine == 'ppc64le') or (python_full_version < '3.14' and platform_machine == 'win32') or (python_full_version < '3.14' and platform_machine == 'x86_64')" }, { name = "typing-extensions" }, ] + sdist = { url = "https://files.pythonhosted.org/packages/63/66/45b165c595ec89aa7dcc2c1cd222ab269bc753f1fc7a1e68f8481bd957bf/sqlalchemy-2.0.41.tar.gz", hash = "sha256:edba70118c4be3c2b1f90754d308d0b79c6fe2c0fdc52d8ddf603916f83f4db9", size = 9689424 } wheels = [ { url = "https://files.pythonhosted.org/packages/d3/ad/2e1c6d4f235a97eeef52d0200d8ddda16f6c4dd70ae5ad88c46963440480/sqlalchemy-2.0.41-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4eeb195cdedaf17aab6b247894ff2734dcead6c08f748e617bfe05bd5a218443", size = 2115491 }, @@ -1795,7 +1797,7 @@ name = "tzlocal" version = "5.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "tzdata", marker = "sys_platform == 'win32'" }, + { name = "tzdata", marker = "platform_system == 'Windows'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/8b/2e/c14812d3d4d9cd1773c6be938f89e5735a1f11a9f184ac3639b93cef35d5/tzlocal-5.3.1.tar.gz", hash = "sha256:cceffc7edecefea1f595541dbd6e990cb1ea3d19bf01b2809f362a03dd7921fd", size = 30761 } wheels = [