Move example A2UI tools to SDK

Author: jonsharkey

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

@@ -0,0 +1,255 @@
+# 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.
+
+import json
+import jsonschema
+import logging
+import inspect
+from typing import Any, List, Optional, TypeAlias, Callable, Union, Awaitable
+
+from a2a import types as a2a_types
+from google.genai import types as genai_types
+
+from google.adk.a2a.converters import part_converter
+from google.adk.models import LlmRequest
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools import base_toolset
+from google.adk.tools.tool_context import ToolContext
+from google.adk.agents.readonly_context import ReadonlyContext
+from google.adk.utils.feature_decorator import experimental
+
+from a2ui.a2ui_extension import create_a2ui_part
+
+logger = logging.getLogger(__name__)
+
+
+@experimental
+class A2uiToolset(base_toolset.BaseToolset):
+    """A toolset that provides A2UI Tools and can be enabled/disabled."""
+
+    A2uiEnabledProvider: TypeAlias = Callable[
+        [ReadonlyContext], Union[bool, Awaitable[bool]]
+    ]
+    A2uiSchemaProvider: TypeAlias = Callable[
+        [ReadonlyContext], Union[dict[str, Any], Awaitable[dict[str, Any]]]
+    ]
+
+    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 = [SendA2uiJsonToClientTool(self)]
+
+    async def canonical_a2ui_enabled(self) -> bool:
+        """The resolved self.instruction field to construct instruction for this agent.
+
+        This method is only for use by Agent Development Kit.
+
+        Args:
+        ctx: The context to retrieve the session state.
+
+        Returns:
+        If a2ui is enabled, return True. Otherwise, return False.
+        """
+        if isinstance(self._a2ui_enabled, bool):
+            return self._a2ui_enabled
+        else:
+            a2ui_enabled = self._a2ui_enabled()
+            if inspect.isawaitable(a2ui_enabled):
+                a2ui_enabled = await a2ui_enabled
+            return a2ui_enabled
+
+    async def canonical_a2ui_schema(self) -> dict[str, Any]:
+        """The resolved self.instruction field to construct instruction for this agent.
+
+        This method is only for use by Agent Development Kit.
+
+        Args:
+        ctx: The context to retrieve the session state.
+
+        Returns:
+        The A2UI schema to send to the client.
+        """
+        if isinstance(self._a2ui_schema, dict):
+            return self._a2ui_schema
+        else:
+            a2ui_schema = self._a2ui_schema()
+            if inspect.isawaitable(a2ui_schema):
+                a2ui_schema = await a2ui_schema
+            return a2ui_schema
+
+    async def get_tools(
+        self,
+        readonly_context: Optional[ReadonlyContext] = None,
+    ) -> List[BaseTool]:
+        use_ui = self._a2ui_enabled
+        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 []
+
+@experimental
+class SendA2uiJsonToClientTool(BaseTool):
+    TOOL_NAME = "send_a2ui_json_to_client"
+    A2UI_JSON_ARG_NAME = "a2ui_json"
+
+    def __init__(self, toolset: A2uiToolset):
+        self._toolset = toolset
+        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"    {self.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={
+                    self.A2UI_JSON_ARG_NAME: genai_types.Schema(
+                        type=genai_types.Type.STRING,
+                        description="valid A2UI JSON Schema to send to the client.",
+                    ),
+                },
+                required=[self.A2UI_JSON_ARG_NAME],
+            ),
+        )
+
+    async def get_a2ui_schema(self) -> dict[str, Any]:
+        a2ui_schema = await self._toolset.canonical_a2ui_schema()
+        print(f"!!!!! A2UI Schema: {a2ui_schema}")
+        if not a2ui_schema:
+            raise ValueError("A2UI schema is empty")
+        a2ui_schema_object = {"type": "array", "items": a2ui_schema} # Make a list since we support multiple parts in this tool call
+        return a2ui_schema_object 
+
+    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()
+
+        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(self.A2UI_JSON_ARG_NAME)
+            if not a2ui_json:                
+                raise ValueError(
+                    f"Failed to call tool {self.TOOL_NAME} because missing required arg {self.A2UI_JSON_ARG_NAME} "
+                )
+
+            a2ui_json_payload = json.loads(a2ui_json)
+            a2ui_schema = await self.get_a2ui_schema()
+            jsonschema.validate(
+                instance=a2ui_json_payload, schema=a2ui_schema
+            )
+
+            logger.info(
+                f"Validated call to tool {self.TOOL_NAME} with {self.A2UI_JSON_ARG_NAME}"
+            )
+
+            # Don't do a second LLM inference call for the None response
+            tool_context.actions.skip_summarization = True
+
+            return None
+        except Exception as e:
+            err = f"Failed to call A2UI tool {self.TOOL_NAME}: {e}"
+            logger.error(err)
+
+            return {"error": err}
+
+@experimental
+class A2uiPartConverter:
+
+  def __init__(self):
+      self._a2ui_schema = None
+
+  def set_a2ui_schema(self, a2ui_schema: dict[str, Any]):
+      self._a2ui_schema = a2ui_schema    
+      
+  def convert_genai_part_to_a2a_part(self, part: genai_types.Part) -> List[a2a_types.Part]:
+      if (function_call := part.function_call) and function_call.name == SendA2uiJsonToClientTool.TOOL_NAME:
+          if self._a2ui_schema is None:
+              logger.error("A2UI schema is not set in part converter")
+              return []
+          
+          try:
+            a2ui_json = function_call.args.get(SendA2uiJsonToClientTool.A2UI_JSON_ARG_NAME)
+            if a2ui_json is None:
+                raise ValueError(f"Failed to convert A2UI function call because required arg {SendA2uiJsonToClientTool.A2UI_JSON_ARG_NAME} not found in {str(part)}")
+            if not a2ui_json.strip():
+                logger.info("Empty a2ui_json, skipping")
+                return []
+            
+            logger.info(f"Converting a2ui json: {a2ui_json}")
+
+            json_data = json.loads(a2ui_json)
+            if not isinstance(json_data, list):
+                logger.info("Received a single JSON object, wrapping in a list for validation.")
+                json_data = [json_data]
+            
+            a2ui_schema_object = {"type": "array", "items": self._a2ui_schema} # Make a list since we support multiple parts in this tool call
+            jsonschema.validate(
+                  instance=json_data, schema=a2ui_schema_object
+              )          
+
+            final_parts = []
+            if isinstance(json_data, list):
+                logger.info( f"Found {len(json_data)} messages. Creating individual DataParts." )
+                for message in json_data:
+                  final_parts.append(create_a2ui_part(message))
+            else:
+                # Handle the case where a single JSON object is returned
+                logger.info("Received a single JSON object. Creating a DataPart." )
+                final_parts.append(create_a2ui_part(json_data))
+
+            return final_parts
+          except Exception as e:
+              logger.error(f"Error converting A2UI function call to A2A parts: {str(e)}")
+              return []
+          
+      # Don't send a2ui tool responses
+      elif (function_response := part.function_response) and function_response.name == SendA2uiJsonToClientTool.TOOL_NAME:    
+          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 []