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/src/a2ui/a2ui_extension.py

@@ -123,3 +123,4 @@ def try_activate_a2ui_extension(context: RequestContext) -> bool:
         context.add_activated_extension(A2UI_EXTENSION_URI)
         return True
     return False
+

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

@@ -0,0 +1,332 @@
+# 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.
+
+"""Module for the A2UI Toolset and Part Converter.
+
+This module provides the necessary components to enable an agent to send A2UI (Agent-to-User Interface)
+JSON payloads to a client. It includes a toolset for managing A2UI tools, a specific tool for the LLM
+to send JSON, and a part converter to translate the LLM's tool calls into A2A (Agent-to-Agent) parts.
+
+Key Components:
+  * `SendA2uiToClientToolset`: The main entry point. It accepts providers for determining
+    if A2UI is enabled and for fetching the A2UI schema. It manages the lifecycle of the
+    `_SendA2uiJsonToClientTool`.
+  * `_SendA2uiJsonToClientTool`: A tool exposed to the LLM. It allows the LLM to "call" a function
+    that effectively sends a JSON payload to the client. This tool validates the JSON against
+    the provided schema. It automatically wraps the provided schema in an array structure,
+    instructing the LLM that it can send a list of UI items.
+  * `SendA2uiToClientPartConverter`: A utility class that intercepts the `send_a2ui_json_to_client`
+    tool calls from the LLM and converts them into `a2a_types.Part` objects, which are then
+    processed by the A2A system.
+
+Usage Examples:
+
+  1. Defining Providers:
+    You can use simple values or callables (sync or async) for enablement and schema.
+
+    ```python
+    # Simple boolean and dict
+    toolset = SendA2uiToClientToolset(a2ui_enabled=True, a2ui_schema=MY_SCHEMA)
+
+    # Async providers
+    async def check_enabled(ctx: ReadonlyContext) -> bool:
+      return await some_condition(ctx)
+
+    async def get_schema(ctx: ReadonlyContext) -> dict[str, Any]:
+      return await fetch_schema(ctx)
+
+    toolset = SendA2uiToClientToolset(a2ui_enabled=check_enabled, a2ui_schema=get_schema)
+    ```
+
+  2. Integration with Agent:
+    Typically used when initializing an agent's toolset.
+
+    ```python
+    # In your agent initialization
+    self.a2ui_toolset = SendA2uiToClientToolset(
+        a2ui_enabled=self._is_a2ui_enabled,
+        a2ui_schema=self._get_a2ui_schema
+    )
+    ```
+
+  3. Part Conversion:
+    The converter needs to be aware of the schema to validate incoming tool calls.
+
+    ```python
+    converter = SendA2uiToClientToolset.get_part_converter()
+    parts = converter.convert_genai_part_to_a2a_part(llm_part)
+    ```
+"""
+
+import inspect
+import json
+import logging
+from typing import Any, Awaitable, Callable, Optional, TypeAlias, Union
+
+import jsonschema
+
+from a2a import types as a2a_types
+from a2ui.a2ui_extension import create_a2ui_part
+from google.adk.a2a.converters import part_converter
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.models import LlmRequest
+from google.adk.tools import base_toolset
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+from google.adk.utils.feature_decorator import experimental
+from google.genai import types as genai_types
+
+logger = logging.getLogger(__name__)
+
+A2uiEnabledProvider: TypeAlias = Callable[
+    [ReadonlyContext], Union[bool, Awaitable[bool]]
+]
+A2uiSchemaProvider: TypeAlias = Callable[
+    [ReadonlyContext], Union[dict[str, Any], Awaitable[dict[str, Any]]]
+]
+
+async def resolve_a2ui_enabled(a2ui_enabled: Union[bool, A2uiEnabledProvider], ctx: ReadonlyContext) -> bool:
+    """The resolved self.a2ui_enabled field to construct instruction for this agent.
+
+    Args:
+        ctx: The ReadonlyContext to resolve the provider with.
+
+    Returns:
+        If A2UI is enabled, return True. Otherwise, return False.
+    """
+    if isinstance(a2ui_enabled, bool):
+        return a2ui_enabled
+    else:
+        a2ui_enabled = a2ui_enabled(ctx)
+        if inspect.isawaitable(a2ui_enabled):
+            a2ui_enabled = await a2ui_enabled
+        return a2ui_enabled
+
+async def resolve_a2ui_schema(a2ui_schema: Union[dict[str, Any], A2uiSchemaProvider], ctx: ReadonlyContext) -> dict[str, Any]:
+    """The resolved self.a2ui_schema field to construct instruction for this agent.
+
+    Args:
+        ctx: The ReadonlyContext to resolve the provider with.
+
+    Returns:
+        The A2UI schema to send to the client.
+    """
+    if isinstance(a2ui_schema, dict):
+        return a2ui_schema
+    else:
+        a2ui_schema = a2ui_schema(ctx)
+        if inspect.isawaitable(a2ui_schema):
+            a2ui_schema = await a2ui_schema
+        return a2ui_schema
+
+@experimental
+class SendA2uiToClientPartConverter:
+
+    def __init__(self, toolname: str):
+        self._toolname = toolname
+
+    def convert_genai_part_to_a2a_part(self, part: genai_types.Part) -> list[a2a_types.Part]:
+        if (function_response := part.function_response) and function_response.name == self._toolname:
+            if "error" in function_response.response:
+                logger.warning(f"A2UI tool call failed: {function_response.response['error']}")
+                return []
+            
+            # The tool returns the list of messages directly on success
+            json_data = function_response.response.get("result")
+            if not json_data:
+                logger.info("No result in A2UI tool response")
+                return []
+
+            final_parts = []
+            for message in json_data:
+                logger.info(f"Found {len(json_data)} messages. Creating individual DataParts.")
+                final_parts.append(create_a2ui_part(message))
+
+            return final_parts
+
+        # Don't send a2ui tool call to client
+        elif (function_call := part.function_call) and function_call.name == self._toolname:
+            return []
+
+        # Use default part converter for other types (images, etc)
+        converted_part = part_converter.convert_genai_part_to_a2a_part(part)
+
+        logger.info(f"Returning converted part: {converted_part}")
+        return [converted_part] if converted_part else []
+
+@experimental
+class SendA2uiToClientToolset(base_toolset.BaseToolset):
+    """A toolset that provides A2UI Tools and can be enabled/disabled."""
+
+    A2UI_JSON_ARG_NAME = "a2ui_json"
+
+    def __init__(
+        self,
+        a2ui_enabled: Union[bool, A2uiEnabledProvider],
+        a2ui_schema: Union[dict[str, Any], A2uiSchemaProvider],
+    ):
+        super().__init__()
+        self._a2ui_enabled = a2ui_enabled
+        self._a2ui_schema = a2ui_schema
+        self._ui_tools = [self._SendA2uiJsonToClientTool(self._a2ui_schema)]
+
+    async def _resolve_a2ui_enabled(self, ctx: ReadonlyContext) -> bool:
+        return await resolve_a2ui_enabled(self._a2ui_enabled, ctx)
+
+    async def _resolve_a2ui_schema(self, ctx: ReadonlyContext) -> dict[str, Any]:
+        return await resolve_a2ui_schema(self._a2ui_schema, ctx)
+
+    @classmethod
+    def get_part_converter(cls) -> SendA2uiToClientPartConverter:
+        return SendA2uiToClientPartConverter(cls._SendA2uiJsonToClientTool.TOOL_NAME)
+
+    async def get_tools(
+        self,
+        readonly_context: Optional[ReadonlyContext] = None,
+    ) -> list[BaseTool]:
+        """Returns the list of tools provided by this toolset.
+
+        Args:
+            readonly_context: The ReadonlyContext for resolving tool enablement.
+
+        Returns:
+            A list of tools.
+        """
+        use_ui = False
+        if readonly_context is not None:
+            use_ui = await self._resolve_a2ui_enabled(readonly_context)
+        if use_ui:
+            logger.info("A2UI is ENABLED, adding ui tools")
+            return self._ui_tools
+        else:
+            logger.info("A2UI is DISABLED, not adding ui tools")
+            return []
+
+    @staticmethod
+    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}
+
+    class _SendA2uiJsonToClientTool(BaseTool):
+        TOOL_NAME = "send_a2ui_json_to_client"
+
+        def __init__(self, a2ui_schema: Union[dict[str, Any], A2uiSchemaProvider]):
+            self._a2ui_schema = a2ui_schema
+            super().__init__(
+                name=self.TOOL_NAME,
+                description="Sends A2UI JSON to the client to render rich UI for the user. This tool can be called multiple times in the same call to render multiple UI surfaces."
+                "Args:"
+                f"    {SendA2uiToClientToolset.A2UI_JSON_ARG_NAME}: Valid A2UI JSON Schema to send to the client. The A2UI JSON Schema definition is between ---BEGIN A2UI JSON SCHEMA--- and ---END A2UI JSON SCHEMA--- in the system instructions.",
+            )
+
+        def _get_declaration(self) -> genai_types.FunctionDeclaration | None:
+            return genai_types.FunctionDeclaration(
+                name=self.name,
+                description=self.description,
+                parameters=genai_types.Schema(
+                    type=genai_types.Type.OBJECT,
+                    properties={
+                        SendA2uiToClientToolset.A2UI_JSON_ARG_NAME: genai_types.Schema(
+                            type=genai_types.Type.STRING,
+                            description="valid A2UI JSON Schema to send to the client.",
+                        ),
+                    },
+                    required=[SendA2uiToClientToolset.A2UI_JSON_ARG_NAME],
+                ),
+            )
+
+        async def get_a2ui_schema(self, ctx: ReadonlyContext) -> dict[str, Any]:
+            """Retrieves and wraps the A2UI schema.
+
+            Args:
+                ctx: The ReadonlyContext for resolving the schema.
+
+            Returns:
+                The wrapped A2UI schema.
+            """
+            a2ui_schema = await resolve_a2ui_schema(self._a2ui_schema, ctx)
+            return SendA2uiToClientToolset.wrap_as_json_array(a2ui_schema)
+
+        async def process_llm_request(
+            self, *, tool_context: ToolContext, llm_request: LlmRequest
+        ) -> None:
+            await super().process_llm_request(
+                tool_context=tool_context, llm_request=llm_request
+            )
+
+            a2ui_schema = await self.get_a2ui_schema(tool_context)
+
+            llm_request.append_instructions(
+                [
+                    f"""
+---BEGIN A2UI JSON SCHEMA---
+{json.dumps(a2ui_schema)}
+---END A2UI JSON SCHEMA---
+"""
+                ]
+            )
+
+            logger.info("Added a2ui_schema to system instructions")
+
+        async def run_async(
+            self, *, args: dict[str, Any], tool_context: ToolContext
+        ) -> Any:
+            try:
+                a2ui_json = args.get(SendA2uiToClientToolset.A2UI_JSON_ARG_NAME)
+                if not a2ui_json:                
+                    raise ValueError(
+                        f"Failed to call tool {self.TOOL_NAME} because missing required arg {SendA2uiToClientToolset.A2UI_JSON_ARG_NAME} "
+                    )
+
+                a2ui_json_payload = json.loads(a2ui_json)
+                
+                # Auto-wrap single object in list
+                if not isinstance(a2ui_json_payload, list):
+                     logger.info("Received a single JSON object, wrapping in a list for validation.")
+                     a2ui_json_payload = [a2ui_json_payload]
+
+                a2ui_schema = await self.get_a2ui_schema(tool_context)
+                jsonschema.validate(
+                    instance=a2ui_json_payload, schema=a2ui_schema
+                )
+
+                logger.info(
+                    f"Validated call to tool {self.TOOL_NAME} with {SendA2uiToClientToolset.A2UI_JSON_ARG_NAME}"
+                )
+
+                # Don't do a second LLM inference call for the JSON response
+                tool_context.actions.skip_summarization = True
+
+                # Return the validated JSON so the converter can use it.
+                # We return it in a dict under "result" key for consistent JSON structure.
+                return {"result": a2ui_json_payload}
+
+            except Exception as e:
+                err = f"Failed to call A2UI tool {self.TOOL_NAME}: {e}"
+                logger.error(err)
+
+                return {"error": err}
+

…n/a2ui_extension/tests/test_extension.py …i_extension/tests/test_a2ui_extension.pya2a_agents/python/a2ui_extension/tests/test_extension.py renamed to a2a_agents/python/a2ui_extension/tests/test_a2ui_extension.py a2a_agents/python/a2ui_extension/tests/test_extension.py renamed to a2a_agents/python/a2ui_extension/tests/test_a2ui_extension.py

@@ -100,3 +100,5 @@ def test_try_activate_a2ui_extension_not_requested():
 
     assert not a2ui_extension.try_activate_a2ui_extension(context)
     context.add_activated_extension.assert_not_called()
+
+