Move example A2UI tools to SDK

Author: jonsharkey

a2a_agents/python/a2ui_extension/README.md

@@ -1,6 +1,7 @@
 # A2UI Extension Implementation
 
-This is the Python implementation of the a2ui extension.
+a2ui_extension.py is the Python implementation of the a2ui extension.
+send_a2ui_to_client_toolset.py is an example Python implementation of using ADK toolcalls to implement A2UI.
 
 ## Running Tests
 
@@ -13,7 +14,7 @@ This is the Python implementation of the a2ui extension.
 2. Run the tests
 
    ```bash
-   uv run --with pytest pytest tests/test_extension.py
+   uv run --with pytest pytest tests/*.py
    ```
 
 ## Disclaimer

a2a_agents/python/a2ui_extension/pyproject.toml

@@ -13,3 +13,19 @@ build-backend = "hatchling.build"
 [[tool.uv.index]]
 url = "https://pypi.org/simple"
 default = true
+
+[tool.pyink]
+unstable = true
+target-version = []
+pyink-indentation = 2
+pyink-use-majority-quotes = true
+pyink-annotation-pragmas = [
+  "noqa",
+  "pylint:",
+  "type: ignore",
+  "pytype:",
+  "mypy:",
+  "pyright:",
+  "pyre-",
+]
+

a2a_agents/python/a2ui_extension/src/a2ui/a2ui_extension.py

@@ -31,95 +31,99 @@
 
 STANDARD_CATALOG_ID = "https://github.com/google/A2UI/blob/main/specification/0.8/json/standard_catalog_definition.json"
 
+
 def create_a2ui_part(a2ui_data: dict[str, Any]) -> Part:
-    """Creates an A2A Part containing A2UI data.
-
-    Args:
-        a2ui_data: The A2UI data dictionary.
-
-    Returns:
-        An A2A Part with a DataPart containing the A2UI data.
-    """
-    return Part(
-        root=DataPart(
-            data=a2ui_data,
-            metadata={
-                MIME_TYPE_KEY: A2UI_MIME_TYPE,
-            },
-        )
-    )
+  """Creates an A2A Part containing A2UI data.
+
+  Args:
+      a2ui_data: The A2UI data dictionary.
+
+  Returns:
+      An A2A Part with a DataPart containing the A2UI data.
+  """
+  return Part(
+      root=DataPart(
+          data=a2ui_data,
+          metadata={
+              MIME_TYPE_KEY: A2UI_MIME_TYPE,
+          },
+      )
+  )
 
 
 def is_a2ui_part(part: Part) -> bool:
-    """Checks if an A2A Part contains A2UI data.
-
-    Args:
-        part: The A2A Part to check.
-
-    Returns:
-        True if the part contains A2UI data, False otherwise.
-    """
-    return (
-        isinstance(part.root, DataPart)
-        and part.root.metadata
-        and part.root.metadata.get(MIME_TYPE_KEY) == A2UI_MIME_TYPE
-    )
+  """Checks if an A2A Part contains A2UI data.
+
+  Args:
+      part: The A2A Part to check.
+
+  Returns:
+      True if the part contains A2UI data, False otherwise.
+  """
+  return (
+      isinstance(part.root, DataPart)
+      and part.root.metadata
+      and part.root.metadata.get(MIME_TYPE_KEY) == A2UI_MIME_TYPE
+  )
 
 
 def get_a2ui_datapart(part: Part) -> Optional[DataPart]:
-    """Extracts the DataPart containing A2UI data from an A2A Part, if present.
+  """Extracts the DataPart containing A2UI data from an A2A Part, if present.
 
-    Args:
-        part: The A2A Part to extract A2UI data from.
+  Args:
+      part: The A2A Part to extract A2UI data from.
 
-    Returns:
-        The DataPart containing A2UI data if present, None otherwise.
-    """
-    if is_a2ui_part(part):
-        return part.root
-    return None
+  Returns:
+      The DataPart containing A2UI data if present, None otherwise.
+  """
+  if is_a2ui_part(part):
+    return part.root
+  return None
 
 
 AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY = "supportedCatalogIds"
 AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY = "acceptsInlineCatalogs"
 
+
 def get_a2ui_agent_extension(
     accepts_inline_catalogs: bool = False,
     supported_catalog_ids: List[str] = [],
 ) -> AgentExtension:
-    """Creates the A2UI AgentExtension configuration.
-
-    Args:
-        accepts_inline_catalogs: Whether the agent accepts inline custom catalogs.
-        supported_catalog_ids: All pre-defined catalogs the agent is known to support.
-
-    Returns:
-        The configured A2UI AgentExtension.
-    """
-    params = {}
-    if accepts_inline_catalogs:
-        params[AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY] = True  # Only set if not default of False
-
-    if supported_catalog_ids:
-        params[AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY] = supported_catalog_ids
-
-    return AgentExtension(
-        uri=A2UI_EXTENSION_URI,
-        description="Provides agent driven UI using the A2UI JSON format.",
-        params=params if params else None,
+  """Creates the A2UI AgentExtension configuration.
+
+  Args:
+      accepts_inline_catalogs: Whether the agent accepts inline custom catalogs.
+      supported_catalog_ids: All pre-defined catalogs the agent is known to support.
+
+  Returns:
+      The configured A2UI AgentExtension.
+  """
+  params = {}
+  if accepts_inline_catalogs:
+    params[AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY] = (
+        True  # Only set if not default of False
     )
 
+  if supported_catalog_ids:
+    params[AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY] = supported_catalog_ids
+
+  return AgentExtension(
+      uri=A2UI_EXTENSION_URI,
+      description="Provides agent driven UI using the A2UI JSON format.",
+      params=params if params else None,
+  )
+
 
 def try_activate_a2ui_extension(context: RequestContext) -> bool:
-    """Activates the A2UI extension if requested.
-
-    Args:
-        context: The request context to check.
-
-    Returns:
-        True if activated, False otherwise.
-    """
-    if A2UI_EXTENSION_URI in context.requested_extensions:
-        context.add_activated_extension(A2UI_EXTENSION_URI)
-        return True
-    return False
+  """Activates the A2UI extension if requested.
+
+  Args:
+      context: The request context to check.
+
+  Returns:
+      True if activated, False otherwise.
+  """
+  if A2UI_EXTENSION_URI in context.requested_extensions:
+    context.add_activated_extension(A2UI_EXTENSION_URI)
+    return True
+  return False

a2a_agents/python/a2ui_extension/src/a2ui/a2ui_schema_utils.py

@@ -0,0 +1,34 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for A2UI Schema manipulation."""
+
+from typing import Any
+
+
+def wrap_as_json_array(a2ui_schema: dict[str, Any]) -> dict[str, Any]:
+  """Wraps the A2UI schema in an array object to support multiple parts.
+
+  Args:
+      a2ui_schema: The A2UI schema to wrap.
+
+  Returns:
+      The wrapped A2UI schema object.
+
+  Raises:
+      ValueError: If the A2UI schema is empty.
+  """
+  if not a2ui_schema:
+    raise ValueError("A2UI schema is empty")
+  return {"type": "array", "items": a2ui_schema}