feat: traceability extension - more follow up changes - docs update, agent card (#330)

* feat: traceability extension - more follow up changes - docs update, agent card

* some small fixes

* removed unused import

* small linter fix

Author: mindpower

extensions/traceability/v1/spec.md

@@ -0,0 +1,123 @@
+# Traceability Extension
+
+## Overview
+
+This extension defines how to add traceability information to `Message` and `Artifact`
+objects.
+
+## Extension URI
+
+The URI of this extension is `https://github.com/a2aproject/a2a-samples/extensions/traceability/v1`.
+
+This is the only URI accepted for this extension.
+
+## Traceability Format
+
+
+
+## Message/Artifact Metadata Field
+
+Traceability information MUST be stored in the metadata for a Message or Artifact, under a
+field with the key `github.com/a2aproject/a2a-samples/extensions/traceability/v1/traceability`,
+or an addtional artifact in the returned completed response.
+
+```proto
+// A Trace message that contains a collection of spans.
+message ResponseTrace {
+ // A unique identifier for the trace.
+ string trace_id = 1;
+
+
+ // The list of steps that make up this trace.
+ repeated Step steps = 2;
+}
+
+
+enum CallTypeEnum {
+ AGENT = 1;
+ TOOL = 2;
+}
+
+
+message StepAction {
+ oneof action {
+   ToolInvocation tool_invocation = 1;
+   AgentInvocation agent_invocation = 2;
+ }
+}
+
+
+message ToolInvocation {
+ string tool_name = 1;
+ google.protobuf.Struct parameters = 2;
+}
+
+
+message AgentInvocation {
+ // The URL of the agent that was invoked.
+ string agent_url = 1;
+ // the agent name
+ string agent_name = 2;
+ // The request message sent to the agent.
+ google.protobuf.Struct requests = 3;
+ // intenral response trace for this specific steps, if the callee also
+ // supports the traceability extension and the caller requests traceability.
+ ResponseTrace response_trace = 4;
+}
+
+
+// A Step message that represents a single operation within a trace.
+message Step {
+ // A unique identifier for this step.
+ string step_id = 1;
+
+
+ // The trace_id of the trace this step belongs to.
+ string trace_id = 2;
+
+
+ // The step_id of the parent step. Empty if this is a root step.
+ string parent_step_id = 3;
+
+
+ // The type of the operation this step represents.
+ CallTypeEnum call_type = 4;
+
+
+ // Detailed invocation about the step
+ StepAction step_action = 5;
+
+
+ // The cost of the operation this step represents.
+ int64 cost = 6;
+
+
+ // The token of the operation this step represents.
+ int64 total_tokens = 7;
+
+
+ // A set of key-value attributes with additional details about the step.
+ map<string, string> additional_attributes = 8;
+
+
+ // The latency of the operation this step represents.
+ int64 latency = 9;
+
+
+ // The start time of the operation.
+ google.protobuf.Timestamp start_time = 10;
+
+
+ // The end time of the operation.
+ google.protobuf.Timestamp end_time = 11;
+}
+
+```
+
+## Extension Activation
+
+Clients indicate their desire to receive traceability on response by specifying
+the [Extension URI](#extension-uri) via the transport-defined extension
+activation mechanism. For JSON-RPC and HTTP transports, this is indicated via
+the `X-A2A-Extensions` HTTP header. For gRPC, this is indicated via the
+`X-A2A-Extensions` metadata value.

samples/python/extensions/traceability/src/traceability_ext/__init__.py

@@ -5,6 +5,8 @@
 from enum import Enum
 from typing import Any
 
+from a2a.types import AgentExtension
+
 
 _CORE_PATH = 'github.com/a2aproject/a2a-samples/extensions/traceability/v1'
 TRACEABILITY_EXTENSION_URI = f'https://{_CORE_PATH}'
@@ -238,7 +240,7 @@ def __exit__(
         self,
         exc_type: type[BaseException] | None,
         exc_val: BaseException | None,
-        traceback: types.TracebackType | None,
+        exc_traceback: types.TracebackType | None,
     ) -> bool:
         """Context manager exit point that finalizes the trace step.
 
@@ -253,10 +255,25 @@ def __exit__(
         error_msg = None
         if exc_type:
             error_msg = ''.join(
-                traceback.format_exception(exc_type, exc_val, traceback)
+                exc_traceback.format_exception(exc_type, exc_val, exc_traceback)
             )
         self.step.end_step(error=error_msg)
         if self.response_trace:
             self.response_trace.add_step(self.step)
         # Do not suppress exceptions
         return False
+
+
+class TraceabilityExtension:
+    """An implementation of the Traceability extension.
+
+    This extension implementation illustrates a simple way for an extension to
+    provide functionality to agent developers.
+    """
+
+    def agent_extension(self) -> AgentExtension:
+        """Get the AgentExtension representing this extension."""
+        return AgentExtension(
+            uri=TRACEABILITY_EXTENSION_URI,
+            description='Adds traceability information to artifacts.',
+        )

samples/python/hosts/a2a_multiagent_host/README.md

@@ -14,7 +14,7 @@ The application utilizes a multi-agent architecture where a host A2A server dele
 
 ![architecture](assets/A2A_multi_agents.jpg)
 
-### App UI
+### screenshot for CLI tool run with traceability information returned
 
 ![screenshot](assets/cli_trace_screenshot.png)
 

samples/python/hosts/a2a_multiagent_host/__main__.py

@@ -23,6 +23,7 @@
 from routing_agent import (
     root_agent,
 )
+from traceability_ext import TraceabilityExtension
 
 
 load_dotenv()
@@ -54,14 +55,22 @@ def main(host: str = DEFAULT_HOST, port: int = DEFAULT_PORT):
 
     app_url = os.environ.get('APP_URL', f'http://{host}:{port}')
 
+    traceability_ext = TraceabilityExtension()
+    capabilities = AgentCapabilities(
+        streaming=True,
+        extensions=[
+            traceability_ext.agent_extension(),
+        ],
+    )
+
     agent_card = AgentCard(
         name='Host A2A Agent',
         description='A2A server that helps with weather and airbnb',
         url=app_url,
         version='1.0.0',
         default_input_modes=['text'],
         default_output_modes=['text'],
-        capabilities=AgentCapabilities(streaming=True),
+        capabilities=capabilities,
         skills=[skill],
     )
 

samples/python/hosts/a2a_multiagent_host/host_agent_executor.py

@@ -1,7 +1,6 @@
 import json
 import logging
 
-from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING
 
 from a2a.server.agent_execution import AgentExecutor
@@ -65,12 +64,11 @@ async def _process_request(
         self._active_sessions.add(session_id)
 
         try:
-            event_iterator: AsyncIterator[Event] = self.runner.run_async(
+            async for event in self.runner.run_async(
                 user_id=DEFAULT_USER_ID,
                 session_id=session_id,
                 new_message=new_message,
-            )
-            async for event in event_iterator:
+            ):
                 logger.debug(
                     '### Event received: %s',
                     event.model_dump_json(exclude_none=True, indent=2),