Rename a2ui_extension to a2ui_agent (#516)

nan-yu · web-flow · commit df770ebc9a89 · 2026-01-21T15:49:42.000-08:00
This major refactoring renames the python agent library from `a2ui_extension` to `a2ui_agent` to better reflect its purpose as a core agent library rather than just an extension.

- Rename directory `a2a_agents/python/a2ui_extension` -&gt; `a2a_agents/python/a2ui_agent`.
- Refactor project structure:
  - Move extension source files to `src/a2ui/extension/` to support proper namespacing.
  - Move extension tests to `tests/extension/`.
- Update pyproject.toml:
  - Rename project to `a2ui-agent`.
  - Configure `hatchling` build target.
- Update README.md with new directory paths and testing instructions.
- Fix imports in source and test files to match the new `a2ui.extension` namespace.
- Update all sample agents (`contact_lookup`, `contact_multiple_surfaces`, `orchestrator`, `restaurant_finder`, `rizzcharts`) to use the new `a2ui-agent` package name and path.
- Update python imports in all samples from `a2ui.a2ui_extension` to `a2ui.extension.a2ui_extension`.

Verification:
- Manually run the orchestration agent and client
- All workflow are working as expected.
diff --git a/.gemini/GEMINI.md b/.gemini/GEMINI.md
@@ -21,8 +21,8 @@ The A2UI repository is organized into several key directories:
         -   `angular/`: Clients using Angular.
     -   `mcp/`: MCP server samples (e.g., `flight_booking`).
 -   `a2a_agents/`: Contains source code for A2A extension integrations.
-    -   `python/a2ui_extension/`: Python implementation of the A2UI A2A extension.
-    -   `java/`: Java implementation of the A2UI A2A extension.
+    -   `python/a2ui_agent/`: Python implementation of the A2UI agent library.
+    -   `java/`: Java implementation of the A2UI agent library.
 -   `renderers/`: Contains renderer libraries.
     -   `lit/`: The shared Lit renderer library used by the Lit clients.
 -   `tools/`: Helper tools for development.
diff --git a/.github/workflows/python_samples_build.yml b/.github/workflows/python_samples_build.yml
@@ -20,11 +20,11 @@ on:
       - main
     paths:
       - 'samples/agent/adk/**'
-      - 'a2a_agents/python/a2ui_extension/**'
+      - 'a2a_agents/python/a2ui_agent/**'
   pull_request:
     paths:
       - 'samples/agent/adk/**'
-      - 'a2a_agents/python/a2ui_extension/**'
+      - 'a2a_agents/python/a2ui_agent/**'
 
 jobs:
   build:
diff --git a/a2a_agents/python/a2ui_agent/README.md b/a2a_agents/python/a2ui_agent/README.md
@@ -1,20 +1,27 @@
-# A2UI Extension Implementation
+# A2UI Agent Implementation
 
-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.
+The `a2a_agents/python/a2ui_agent/` is the Python implementation of the a2ui agent library.
+
+### Extension Components (`src/a2ui/extension`)
+
+The `src/a2ui/extension` directory contains the core logic for the A2UI agent extension:
+
+*   **`a2ui_extension.py`**: Core utilities for extension management and A2UI part handling.
+*   **`a2ui_schema_utils.py`**: Schema manipulation helpers.
+*   **`send_a2ui_to_client_toolset.py`**: An example implementation of using ADK toolcalls to implement A2UI.
 
 ## Running Tests
 
-1. Navigate to the a2ui_extension dir:
+1. Navigate to the a2ui_agent dir:
 
    ```bash
-   cd a2a_agents/python/a2ui_extension
+   cd a2a_agents/python/a2ui_agent
    ```
 
 2. Run the tests
 
    ```bash
-   uv run --with pytest pytest tests/*.py
+   uv run --with pytest pytest tests/
    ```
 
 ## Disclaimer
diff --git a/a2a_agents/python/a2ui_agent/pyproject.toml b/a2a_agents/python/a2ui_agent/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
-name = "a2ui"
+name = "a2ui-agent"
 version = "0.1.0"
-description = "A2UI Extension"
+description = "A2UI Agent"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
@@ -15,6 +15,9 @@ dependencies = [
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
+[tool.hatch.build.targets.wheel]
+packages = ["src/a2ui"]
+
 [[tool.uv.index]]
 url = "https://pypi.org/simple"
 default = true
diff --git a/a2a_agents/python/a2ui_agent/src/a2ui/__init__.py b/a2a_agents/python/a2ui_agent/src/a2ui/__init__.py
diff --git a/a2a_agents/python/a2ui_agent/src/a2ui/extension/__init__.py b/a2a_agents/python/a2ui_agent/src/a2ui/extension/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2026 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.
diff --git a/a2a_agents/python/a2ui_agent/src/a2ui/extension/a2ui_extension.py b/a2a_agents/python/a2ui_agent/src/a2ui/extension/a2ui_extension.py
diff --git a/a2a_agents/python/a2ui_agent/src/a2ui/extension/a2ui_schema_utils.py b/a2a_agents/python/a2ui_agent/src/a2ui/extension/a2ui_schema_utils.py
diff --git a/a2a_agents/python/a2ui_agent/src/a2ui/extension/send_a2ui_to_client_toolset.py b/a2a_agents/python/a2ui_agent/src/a2ui/extension/send_a2ui_to_client_toolset.py
@@ -85,8 +85,8 @@ async def get_schema(ctx: ReadonlyContext) -> dict[str, Any]:
 import jsonschema
 
 from a2a import types as a2a_types
-from a2ui.a2ui_extension import create_a2ui_part
-from a2ui.a2ui_schema_utils import wrap_as_json_array
+from a2ui.extension.a2ui_extension import create_a2ui_part
+from a2ui.extension.a2ui_schema_utils import wrap_as_json_array
 from google.adk.a2a.converters import part_converter
 from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.models import LlmRequest
diff --git a/a2a_agents/python/a2ui_agent/tests/__init__.py b/a2a_agents/python/a2ui_agent/tests/__init__.py
diff --git a/a2a_agents/python/a2ui_agent/tests/extension/test_a2ui_extension.py b/a2a_agents/python/a2ui_agent/tests/extension/test_a2ui_extension.py
@@ -15,8 +15,8 @@
 
 from a2a.server.agent_execution import RequestContext
 from a2a.types import DataPart, TextPart, Part
-from a2ui import a2ui_extension
-from a2ui.a2ui_extension import AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY, AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY
+from a2ui.extension import a2ui_extension
+from a2ui.extension.a2ui_extension import AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY, AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY
 from unittest.mock import MagicMock
 
 
diff --git a/a2a_agents/python/a2ui_agent/tests/extension/test_a2ui_schema_utils.py b/a2a_agents/python/a2ui_agent/tests/extension/test_a2ui_schema_utils.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import pytest
-from a2ui.a2ui_schema_utils import wrap_as_json_array
+from a2ui.extension.a2ui_schema_utils import wrap_as_json_array
 
 
 def test_wrap_as_json_array():
diff --git a/a2a_agents/python/a2ui_agent/tests/extension/test_send_a2ui_to_client_toolset.py b/a2a_agents/python/a2ui_agent/tests/extension/test_send_a2ui_to_client_toolset.py
@@ -18,10 +18,10 @@
 import pytest
 
 from a2a import types as a2a_types
-from a2ui.a2ui_extension import create_a2ui_part
+from a2ui.extension.a2ui_extension import create_a2ui_part
 
-from a2ui.send_a2ui_to_client_toolset import convert_send_a2ui_to_client_genai_part_to_a2a_part
-from a2ui.send_a2ui_to_client_toolset import SendA2uiToClientToolset
+from a2ui.extension.send_a2ui_to_client_toolset import convert_send_a2ui_to_client_genai_part_to_a2a_part
+from a2ui.extension.send_a2ui_to_client_toolset import SendA2uiToClientToolset
 from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.tools.tool_context import ToolContext
 from google.genai import types as genai_types
diff --git a/a2a_agents/python/a2ui_agent/uv.lock b/a2a_agents/python/a2ui_agent/uv.lock
diff --git a/samples/agent/adk/contact_lookup/__main__.py b/samples/agent/adk/contact_lookup/__main__.py
@@ -20,7 +20,7 @@
 from a2a.server.request_handlers import DefaultRequestHandler
 from a2a.server.tasks import InMemoryTaskStore
 from a2a.types import AgentCapabilities, AgentCard, AgentSkill
-from a2ui.a2ui_extension import get_a2ui_agent_extension
+from a2ui.extension.a2ui_extension import get_a2ui_agent_extension
 from agent import ContactAgent
 from agent_executor import ContactAgentExecutor
 from dotenv import load_dotenv
diff --git a/samples/agent/adk/contact_lookup/agent_executor.py b/samples/agent/adk/contact_lookup/agent_executor.py
@@ -33,7 +33,7 @@
 )
 from a2a.utils.errors import ServerError
 from agent import ContactAgent
-from a2ui.a2ui_extension import create_a2ui_part, try_activate_a2ui_extension
+from a2ui.extension.a2ui_extension import create_a2ui_part, try_activate_a2ui_extension
 
 logger = logging.getLogger(__name__)
 
diff --git a/samples/agent/adk/contact_lookup/pyproject.toml b/samples/agent/adk/contact_lookup/pyproject.toml
@@ -12,7 +12,7 @@ dependencies = [
     "python-dotenv>=1.1.0",
     "litellm",
     "jsonschema>=4.0.0",
-    "a2ui",
+    "a2ui-agent",
 ]
 
 [tool.hatch.build.targets.wheel]
diff --git a/samples/agent/adk/contact_multiple_surfaces/README.md b/samples/agent/adk/contact_multiple_surfaces/README.md
@@ -13,7 +13,7 @@ This sample uses the Agent Development Kit (ADK) along with the A2A protocol to
 1. Navigate to the samples directory:
 
     ```bash
-    cd a2a_samples/a2ui_contact_lookup
+    cd samples/agent/adk/contact_multiple_surfaces
     ```
 
 2. Create an environment file with your API key:
diff --git a/samples/agent/adk/contact_multiple_surfaces/__main__.py b/samples/agent/adk/contact_multiple_surfaces/__main__.py
@@ -20,7 +20,7 @@
 from a2a.server.request_handlers import DefaultRequestHandler
 from a2a.server.tasks import InMemoryTaskStore
 from a2a.types import AgentCapabilities, AgentCard, AgentSkill
-from a2ui.a2ui_extension import get_a2ui_agent_extension
+from a2ui.extension.a2ui_extension import get_a2ui_agent_extension
 from agent import ContactAgent
 from agent_executor import ContactAgentExecutor
 from dotenv import load_dotenv
diff --git a/samples/agent/adk/contact_multiple_surfaces/agent_executor.py b/samples/agent/adk/contact_multiple_surfaces/agent_executor.py
@@ -34,7 +34,7 @@
 )
 from a2a.utils.errors import ServerError
 from agent import ContactAgent
-from a2ui.a2ui_extension import create_a2ui_part, try_activate_a2ui_extension
+from a2ui.extension.a2ui_extension import create_a2ui_part, try_activate_a2ui_extension
 
 logger = logging.getLogger(__name__)
 
diff --git a/samples/agent/adk/contact_multiple_surfaces/pyproject.toml b/samples/agent/adk/contact_multiple_surfaces/pyproject.toml
@@ -12,7 +12,7 @@ dependencies = [
     "python-dotenv>=1.1.0",
     "litellm",
     "jsonschema>=4.0.0",
-    "a2ui",
+    "a2ui-agent",
 ]
 
 [tool.hatch.build.targets.wheel]
diff --git a/samples/agent/adk/orchestrator/agent.py b/samples/agent/adk/orchestrator/agent.py
@@ -36,7 +36,7 @@
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 from a2a.client.client import ClientConfig as A2AClientConfig
 from a2a.client.client_factory import ClientFactory as A2AClientFactory
-from a2ui.a2ui_extension import is_a2ui_part, A2UI_CLIENT_CAPABILITIES_KEY, A2UI_EXTENSION_URI, AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY, AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY, get_a2ui_agent_extension
+from a2ui.extension.a2ui_extension import is_a2ui_part, A2UI_CLIENT_CAPABILITIES_KEY, A2UI_EXTENSION_URI, AGENT_EXTENSION_SUPPORTED_CATALOG_IDS_KEY, AGENT_EXTENSION_ACCEPTS_INLINE_CATALOGS_KEY, get_a2ui_agent_extension
 from a2a.types import AgentCapabilities, AgentCard, AgentExtension
 
 logger = logging.getLogger(__name__)
diff --git a/samples/agent/adk/orchestrator/agent_executor.py b/samples/agent/adk/orchestrator/agent_executor.py
@@ -31,7 +31,7 @@
     A2aAgentExecutorConfig,
     A2aAgentExecutor,
 )
-from a2ui.a2ui_extension import is_a2ui_part, try_activate_a2ui_extension, A2UI_EXTENSION_URI, STANDARD_CATALOG_ID, SUPPORTED_CATALOG_IDS_KEY, A2UI_CLIENT_CAPABILITIES_KEY
+from a2ui.extension.a2ui_extension import is_a2ui_part, try_activate_a2ui_extension, A2UI_EXTENSION_URI, STANDARD_CATALOG_ID, SUPPORTED_CATALOG_IDS_KEY, A2UI_CLIENT_CAPABILITIES_KEY
 from google.adk.a2a.converters import event_converter
 from a2a.server.events import Event as A2AEvent
 from google.adk.events.event import Event
diff --git a/samples/agent/adk/orchestrator/part_converters.py b/samples/agent/adk/orchestrator/part_converters.py
@@ -19,7 +19,7 @@
 from google.genai import types as genai_types
 
 from google.adk.a2a.converters import part_converter
-from a2ui.a2ui_extension import is_a2ui_part
+from a2ui.extension.a2ui_extension import is_a2ui_part
 
 import pydantic
 
diff --git a/samples/agent/adk/orchestrator/pyproject.toml b/samples/agent/adk/orchestrator/pyproject.toml
@@ -12,7 +12,7 @@ dependencies = [
     "python-dotenv>=1.1.0",
     "litellm",
     "jsonschema>=4.0.0",
-    "a2ui",
+    "a2ui-agent",
 ]
 
 [tool.hatch.build.targets.wheel]
diff --git a/samples/agent/adk/pyproject.toml b/samples/agent/adk/pyproject.toml
@@ -6,5 +6,5 @@ default = true
 members = ["contact_lookup", "contact_multiple_surfaces", "orchestrator", "restaurant_finder", "rizzcharts"]
 
 [tool.uv.sources]
-a2ui = { path = "../../../a2a_agents/python/a2ui_extension", editable = true }
+a2ui-agent = { path = "../../../a2a_agents/python/a2ui_agent", editable = true }
 
diff --git a/samples/agent/adk/restaurant_finder/__main__.py b/samples/agent/adk/restaurant_finder/__main__.py
@@ -20,7 +20,7 @@
 from a2a.server.request_handlers import DefaultRequestHandler
 from a2a.server.tasks import InMemoryTaskStore
 from a2a.types import AgentCapabilities, AgentCard, AgentSkill
-from a2ui.a2ui_extension import get_a2ui_agent_extension
+from a2ui.extension.a2ui_extension import get_a2ui_agent_extension
 from agent import RestaurantAgent
 from agent_executor import RestaurantAgentExecutor
 from dotenv import load_dotenv
diff --git a/samples/agent/adk/restaurant_finder/agent_executor.py b/samples/agent/adk/restaurant_finder/agent_executor.py
@@ -32,7 +32,7 @@
     new_task,
 )
 from a2a.utils.errors import ServerError
-from a2ui.a2ui_extension import create_a2ui_part, try_activate_a2ui_extension
+from a2ui.extension.a2ui_extension import create_a2ui_part, try_activate_a2ui_extension
 from agent import RestaurantAgent
 
 logger = logging.getLogger(__name__)
diff --git a/samples/agent/adk/restaurant_finder/pyproject.toml b/samples/agent/adk/restaurant_finder/pyproject.toml
@@ -12,7 +12,7 @@ dependencies = [
     "python-dotenv>=1.1.0",
     "litellm",
     "jsonschema>=4.0.0",
-    "a2ui",
+    "a2ui-agent",
 ]
 
 [tool.hatch.build.targets.wheel]
diff --git a/samples/agent/adk/rizzcharts/agent.py b/samples/agent/adk/rizzcharts/agent.py
@@ -18,9 +18,9 @@
 import pkgutil
 from typing import Any, ClassVar
 
-from a2ui.a2ui_extension import STANDARD_CATALOG_ID
-from a2ui.a2ui_schema_utils import wrap_as_json_array
-from a2ui.send_a2ui_to_client_toolset import SendA2uiToClientToolset, A2uiEnabledProvider, A2uiSchemaProvider
+from a2ui.extension.a2ui_extension import STANDARD_CATALOG_ID
+from a2ui.extension.a2ui_schema_utils import wrap_as_json_array
+from a2ui.extension.send_a2ui_to_client_toolset import SendA2uiToClientToolset, A2uiEnabledProvider, A2uiSchemaProvider
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.planners.built_in_planner import BuiltInPlanner
diff --git a/samples/agent/adk/rizzcharts/agent_executor.py b/samples/agent/adk/rizzcharts/agent_executor.py
@@ -18,12 +18,12 @@
 
 from a2a.server.agent_execution import RequestContext
 from a2a.types import AgentCapabilities, AgentCard, AgentExtension, AgentSkill
-from a2ui.a2ui_extension import A2UI_CLIENT_CAPABILITIES_KEY
-from a2ui.a2ui_extension import A2UI_EXTENSION_URI
-from a2ui.a2ui_extension import STANDARD_CATALOG_ID
-from a2ui.a2ui_extension import get_a2ui_agent_extension
-from a2ui.a2ui_extension import try_activate_a2ui_extension
-from a2ui.send_a2ui_to_client_toolset import convert_send_a2ui_to_client_genai_part_to_a2a_part
+from a2ui.extension.a2ui_extension import A2UI_CLIENT_CAPABILITIES_KEY
+from a2ui.extension.a2ui_extension import A2UI_EXTENSION_URI
+from a2ui.extension.a2ui_extension import STANDARD_CATALOG_ID
+from a2ui.extension.a2ui_extension import get_a2ui_agent_extension
+from a2ui.extension.a2ui_extension import try_activate_a2ui_extension
+from a2ui.extension.send_a2ui_to_client_toolset import convert_send_a2ui_to_client_genai_part_to_a2a_part
 try:
     from .agent import A2UI_CATALOG_URI_STATE_KEY  # pylint: disable=import-error
     from .agent import RIZZCHARTS_CATALOG_URI  # pylint: disable=import-error
diff --git a/samples/agent/adk/rizzcharts/component_catalog_builder.py b/samples/agent/adk/rizzcharts/component_catalog_builder.py
@@ -15,7 +15,7 @@
 import json
 import logging
 from typing import Any, List, Optional
-from a2ui.a2ui_extension import INLINE_CATALOGS_KEY, SUPPORTED_CATALOG_IDS_KEY
+from a2ui.extension.a2ui_extension import INLINE_CATALOGS_KEY, SUPPORTED_CATALOG_IDS_KEY
 try:
     from .agent import RIZZCHARTS_CATALOG_URI, STANDARD_CATALOG_ID
 except ImportError:
diff --git a/samples/agent/adk/rizzcharts/pyproject.toml b/samples/agent/adk/rizzcharts/pyproject.toml
@@ -12,7 +12,7 @@ dependencies = [
     "python-dotenv>=1.1.0",
     "litellm",
     "jsonschema>=4.0.0",
-    "a2ui",
+    "a2ui-agent",
 ]
 
 [tool.hatch.build.targets.wheel]
@@ -30,4 +30,4 @@ url = "https://pypi.org/simple"
 default = true
 
 [tool.uv.sources]
-a2ui = { path = "../../../../a2a_agents/python/a2ui_extension", editable = true }
+a2ui-agent = { path = "../../../../a2a_agents/python/a2ui_agent", editable = true }
diff --git a/samples/agent/adk/uv.lock b/samples/agent/adk/uv.lock
