Update rizzchart sample

.github/workflows/python_a2ui_agent_build_and_test.yml

@@ -46,10 +46,14 @@ jobs:
         working-directory: a2a_agents/python/a2ui_agent
         run: uv run pyink --check .
 
+      - name: Run unit tests
+        working-directory: a2a_agents/python/a2ui_agent
+        run: uv run --with pytest pytest tests/
+
       - name: Build the python SDK
         working-directory: a2a_agents/python/a2ui_agent
         run: uv build .
 
-      - name: Run unit tests
+      - name: Run validation scripts on assets packing
         working-directory: a2a_agents/python/a2ui_agent
-        run: uv run --with pytest pytest tests/
+        run: uv run python tests/integration/verify_load_real.py

a2a_agents/python/a2ui_agent/.gitignore

@@ -0,0 +1 @@
+src/a2ui/assets/**/*.json

a2a_agents/python/a2ui_agent/pack_specs_hook.py

@@ -0,0 +1,88 @@
+# 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.
+
+import importlib.util
+import os
+import shutil
+from hatchling.builders.hooks.plugin.interface import BuildHookInterface
+
+
+def load_constants(project_root):
+  """Loads the shared constants module directly from its path in src/."""
+  constants_path = os.path.join(
+      project_root, "src", "a2ui", "inference", "schema", "constants.py"
+  )
+  if not os.path.exists(constants_path):
+    raise RuntimeError(f"Could not find shared constants at {constants_path}")
+
+  spec = importlib.util.spec_from_file_location("_constants_load", constants_path)
+  if spec and spec.loader:
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+  raise RuntimeError(f"Could not load shared constants from {constants_path}")
+
+
+class PackSpecsBuildHook(BuildHookInterface):
+
+  def initialize(self, version, build_data):
+    project_root = self.root
+
+    # Load constants dynamically from src/a2ui/inference/schema/constants.py
+    a2ui_constants = load_constants(project_root)
+
+    spec_version_map = a2ui_constants.SPEC_VERSION_MAP
+    a2ui_asset_package = a2ui_constants.A2UI_ASSET_PACKAGE
+    specification_dir = a2ui_constants.SPECIFICATION_DIR
+
+    # project root is in a2a_agents/python/a2ui_agent
+    # Dynamically find repo root by looking for specification_dir
+    repo_root = a2ui_constants.find_repo_root(project_root)
+    if not repo_root:
+      # Check for PKG-INFO which implies a packaged state (sdist).
+      # If PKG-INFO is present, trust the bundled assets.
+      if os.path.exists(os.path.join(project_root, "PKG-INFO")):
+        print("Repository root not found, but PKG-INFO present (sdist). Skipping copy.")
+        return
+
+      raise RuntimeError(
+          f"Could not find repository root (looked for '{specification_dir}'"
+          " directory)."
+      )
+
+    # Target directory: src/a2ui/assets
+    target_base = os.path.join(
+        project_root, "src", a2ui_asset_package.replace(".", os.sep)
+    )
+
+    for ver, schema_map in spec_version_map.items():
+      target_dir = os.path.join(target_base, ver)
+      os.makedirs(target_dir, exist_ok=True)
+
+      for _schema_key, source_rel_path in schema_map.items():
+        source_path = os.path.join(repo_root, source_rel_path)
+
+        if not os.path.exists(source_path):
+          print(
+              f"WARNING: Source schema file not found at {source_path}. Build"
+              " might produce incomplete wheel if not running from monorepo"
+              " root."
+          )
+          continue
+
+        filename = os.path.basename(source_path)
+        dst_file = os.path.join(target_dir, filename)
+
+        print(f"Copying {source_path} -> {dst_file}")
+        shutil.copy2(source_path, dst_file)

a2a_agents/python/a2ui_agent/pyproject.toml

@@ -12,11 +12,18 @@ dependencies = [
 ]
 
 [build-system]
-requires = ["hatchling"]
+requires = ["hatchling", "jsonschema"]
 build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]
 packages = ["src/a2ui"]
+artifacts = ["src/a2ui/assets/**"]
+
+[tool.hatch.build.targets.sdist]
+artifacts = ["src/a2ui/assets/**"]
+
+[tool.hatch.build.hooks.custom]
+path = "pack_specs_hook.py"
 
 [[tool.uv.index]]
 url = "https://pypi.org/simple"

a2a_agents/python/a2ui_agent/src/a2ui/extension/send_a2ui_to_client_toolset.py

@@ -39,16 +39,16 @@
 
     ```python
     # Simple boolean and dict
-    toolset = SendA2uiToClientToolset(a2ui_enabled=True, a2ui_schema=MY_SCHEMA)
+    toolset = SendA2uiToClientToolset(a2ui_enabled=True, a2ui_catalog=MY_CATALOG)
 
     # 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)
+    async def get_catalog(ctx: ReadonlyContext) -> A2uiCatalog:
+      return await fetch_catalog(ctx)
 
-    toolset = SendA2uiToClientToolset(a2ui_enabled=check_enabled, a2ui_schema=get_schema)
+    toolset = SendA2uiToClientToolset(a2ui_enabled=check_enabled, a2ui_catalog=get_catalog)
     ```
 
   2. Integration with Agent:
@@ -60,7 +60,7 @@ async def get_schema(ctx: ReadonlyContext) -> dict[str, Any]:
         tools=[
             SendA2uiToClientToolset(
                 a2ui_enabled=True,
-                a2ui_schema=MY_SCHEMA
+                a2ui_catalog=MY_CATALOG
             )
         ]
     )
@@ -86,7 +86,7 @@ async def get_schema(ctx: ReadonlyContext) -> dict[str, Any]:
 
 from a2a import types as a2a_types
 from a2ui.extension.a2ui_extension import create_a2ui_part
-from a2ui.extension.a2ui_schema_utils import wrap_as_json_array
+from a2ui.inference.schema.catalog import A2uiCatalog
 from google.adk.a2a.converters import part_converter
 from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.models import LlmRequest
@@ -101,8 +101,11 @@ async def get_schema(ctx: ReadonlyContext) -> dict[str, Any]:
 A2uiEnabledProvider: TypeAlias = Callable[
     [ReadonlyContext], Union[bool, Awaitable[bool]]
 ]
-A2uiSchemaProvider: TypeAlias = Callable[
-    [ReadonlyContext], Union[dict[str, Any], Awaitable[dict[str, Any]]]
+A2uiCatalogProvider: TypeAlias = Callable[
+    [ReadonlyContext], Union[A2uiCatalog, Awaitable[A2uiCatalog]]
+]
+A2uiExamplesProvider: TypeAlias = Callable[
+    [ReadonlyContext], Union[str, Awaitable[str]]
 ]
 
 
@@ -113,11 +116,12 @@ class SendA2uiToClientToolset(base_toolset.BaseToolset):
   def __init__(
       self,
       a2ui_enabled: Union[bool, A2uiEnabledProvider],
-      a2ui_schema: Union[dict[str, Any], A2uiSchemaProvider],
+      a2ui_catalog: Union[A2uiCatalog, A2uiCatalogProvider],
+      a2ui_examples: Union[str, A2uiExamplesProvider],
   ):
     super().__init__()
     self._a2ui_enabled = a2ui_enabled
-    self._ui_tools = [self._SendA2uiJsonToClientTool(a2ui_schema)]
+    self._ui_tools = [self._SendA2uiJsonToClientTool(a2ui_catalog, a2ui_examples)]
 
   async def _resolve_a2ui_enabled(self, ctx: ReadonlyContext) -> bool:
     """The resolved self.a2ui_enabled field to construct instruction for this agent.
@@ -164,8 +168,13 @@ class _SendA2uiJsonToClientTool(BaseTool):
     A2UI_JSON_ARG_NAME = "a2ui_json"
     TOOL_ERROR_KEY = "error"
 
-    def __init__(self, a2ui_schema: Union[dict[str, Any], A2uiSchemaProvider]):
-      self._a2ui_schema = a2ui_schema
+    def __init__(
+        self,
+        a2ui_catalog: Union[A2uiCatalog, A2uiCatalogProvider],
+        a2ui_examples: Union[str, A2uiExamplesProvider],
+    ):
+      self._a2ui_catalog = a2ui_catalog
+      self._a2ui_examples = a2ui_examples
       super().__init__(
           name=self.TOOL_NAME,
           description=(
@@ -195,34 +204,39 @@ def _get_declaration(self) -> genai_types.FunctionDeclaration | None:
           ),
       )
 
-    async def _resolve_a2ui_schema(self, ctx: ReadonlyContext) -> dict[str, Any]:
-      """The resolved self.a2ui_schema field to construct instruction for this agent.
+    async def _resolve_a2ui_examples(self, ctx: ReadonlyContext) -> str:
+      """The resolved self.a2ui_examples 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.
+          The A2UI examples string.
       """
-      if isinstance(self._a2ui_schema, dict):
-        return self._a2ui_schema
+      if isinstance(self._a2ui_examples, str):
+        return self._a2ui_examples
       else:
-        a2ui_schema = self._a2ui_schema(ctx)
-        if inspect.isawaitable(a2ui_schema):
-          a2ui_schema = await a2ui_schema
-        return a2ui_schema
+        a2ui_examples = self._a2ui_examples(ctx)
+        if inspect.isawaitable(a2ui_examples):
+          a2ui_examples = await a2ui_examples
+        return a2ui_examples
 
-    async def get_a2ui_schema(self, ctx: ReadonlyContext) -> dict[str, Any]:
-      """Retrieves and wraps the A2UI schema.
+    async def _resolve_a2ui_catalog(self, ctx: ReadonlyContext) -> A2uiCatalog:
+      """The resolved self.a2ui_catalog field to construct instruction for this agent.
 
       Args:
-          ctx: The ReadonlyContext for resolving the schema.
+          ctx: The ReadonlyContext to resolve the provider with.
 
       Returns:
-          The wrapped A2UI schema.
+          The A2UI catalog object.
       """
-      a2ui_schema = await self._resolve_a2ui_schema(ctx)
-      return wrap_as_json_array(a2ui_schema)
+      if isinstance(self._a2ui_catalog, A2uiCatalog):
+        return self._a2ui_catalog
+      else:
+        a2ui_catalog = self._a2ui_catalog(ctx)
+        if inspect.isawaitable(a2ui_catalog):
+          a2ui_catalog = await a2ui_catalog
+        return a2ui_catalog
 
     async def process_llm_request(
         self, *, tool_context: ToolContext, llm_request: LlmRequest
@@ -231,15 +245,14 @@ async def process_llm_request(
           tool_context=tool_context, llm_request=llm_request
       )
 
-      a2ui_schema = await self.get_a2ui_schema(tool_context)
+      a2ui_catalog = await self._resolve_a2ui_catalog(tool_context)
+
+      instruction = a2ui_catalog.render_as_llm_instructions()
+      examples = await self._resolve_a2ui_examples(tool_context)
 
-      llm_request.append_instructions([f"""
----BEGIN A2UI JSON SCHEMA---
-{json.dumps(a2ui_schema)}
----END A2UI JSON SCHEMA---
-"""])
+      llm_request.append_instructions([instruction, examples])
 
-      logger.info("Added a2ui_schema to system instructions")
+      logger.info("Added A2UI schema and examples to system instructions")
 
     async def run_async(
         self, *, args: dict[str, Any], tool_context: ToolContext
@@ -261,8 +274,9 @@ async def run_async(
           )
           a2ui_json_payload = [a2ui_json_payload]
 
-        a2ui_schema = await self.get_a2ui_schema(tool_context)
-        jsonschema.validate(instance=a2ui_json_payload, schema=a2ui_schema)
+        a2ui_catalog = await self._resolve_a2ui_catalog(tool_context)
+
+        a2ui_catalog.validator.validate(a2ui_json_payload)
 
         logger.info(
             f"Validated call to tool {self.TOOL_NAME} with {self.A2UI_JSON_ARG_NAME}"

a2a_agents/python/a2ui_agent/src/a2ui/inference/__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.

a2a_agents/python/a2ui_agent/src/a2ui/inference/inference_strategy.py

@@ -0,0 +1,49 @@
+# 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.
+
+from abc import ABC, abstractmethod
+from typing import List, Optional, Any
+
+
+class InferenceStrategy(ABC):
+
+  @abstractmethod
+  def generate_system_prompt(
+      self,
+      role_description: str,
+      workflow_description: str = "",
+      ui_description: str = "",
+      supported_catalog_ids: List[str] = [],
+      allowed_components: List[str] = [],
+      include_schema: bool = False,
+      include_examples: bool = False,
+      validate_examples: bool = False,
+  ) -> str:
+    """
+    Generates a system prompt for all LLM requests.
+
+    Args:
+      role_description: Description of the agent's role.
+      workflow_description: Description of the workflow.
+      ui_description: Description of the UI.
+      supported_catalog_ids: List of supported catalog IDs.
+      allowed_components: List of allowed components.
+      include_schema: Whether to include the schema.
+      include_examples: Whether to include examples.
+      validate_examples: Whether to validate examples.
+
+    Returns:
+      The system prompt.
+    """
+    pass

a2a_agents/python/a2ui_agent/src/a2ui/inference/schema/__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.