feat: implement schema packaging and runtime loading

- Add a base InferenceStrategy class
- Add PackSpecsBuildHook to copy JSON schemas into the package during
  build time.
- Update pyproject.toml to include assets and configure the build hook.
- Implement A2uiSchemaManager for robust schema loading,  pruning, and
  prompt generation.

Author: nan-yu

.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,62 @@
+import os
+import shutil
+import sys
+from hatchling.builders.hooks.plugin.interface import BuildHookInterface
+
+
+class PackSpecsBuildHook(BuildHookInterface):
+
+  def initialize(self, version, build_data):
+    project_root = self.root
+
+    # Add src to sys.path to import the constant
+    src_path = os.path.join(project_root, "src")
+    if src_path not in sys.path:
+      sys.path.insert(0, src_path)
+
+    from a2ui.inference.schema.manager import (
+        SPEC_VERSION_MAP,
+        A2UI_ASSET_PACKAGE,
+        SPECIFICATION_DIR,
+        find_repo_root,
+    )
+
+    # project root is in a2a_agents/python/a2ui_agent
+    # Dynamically find repo root by looking for SPECIFICATION_DIR
+    repo_root = 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

@@ -17,6 +17,13 @@ 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/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,33 @@
+# 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
+
+
+class InferenceStrategy(ABC):
+
+  @abstractmethod
+  def generate_system_prompt(
+      self,
+      role_description: str,
+      workflow_description: str = "",
+      ui_description: str = "",
+      selected_components: List[str] = [],
+      include_examples: bool = False,
+  ) -> str:
+    """
+    Abstract method to be implemented by subclasses.
+    """
+    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.

a2a_agents/python/a2ui_agent/src/a2ui/inference/schema/loader.py

@@ -0,0 +1,69 @@
+# 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 json
+import os
+import importlib.resources
+from typing import List, Dict, Any
+
+from abc import ABC, abstractmethod
+
+ENCODING = "utf-8"
+
+
+class A2uiSchemaLoader(ABC):
+  """Abstract base class for loading schema files."""
+
+  @abstractmethod
+  def load(self, filename: str) -> Any:
+    """Loads a JSON file."""
+    pass
+
+
+class FileSystemLoader(A2uiSchemaLoader):
+  """Loads schema files from the local filesystem.
+
+  This loader assumes that all referenced schema files are located in the
+  same flat directory structure.
+  """
+
+  def __init__(self, base_dir: str):
+    self.base_dir = base_dir
+
+  def load(self, filename: str) -> Any:
+    path = os.path.join(self.base_dir, filename)
+    with open(path, "r", encoding=ENCODING) as f:
+      return json.load(f)
+
+
+class PackageLoader(A2uiSchemaLoader):
+  """Loads schema files from package resources.
+
+  This loader assumes that all referenced schema files are located in the
+  same flat package structure.
+  """
+
+  def __init__(self, package_path: str):
+    self.package_path = package_path
+
+  def load(self, filename: str) -> Any:
+    try:
+      traversable = importlib.resources.files(self.package_path)
+      resource_path = traversable.joinpath(filename)
+      with resource_path.open("r", encoding=ENCODING) as f:
+        return json.load(f)
+    except Exception as e:
+      raise IOError(
+          f"Could not load package resource {filename} in {self.package_path}: {e}"
+      ) from e