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,76 @@
+# 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 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.constants 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,47 @@
+# 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,
+  ) -> 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.
+
+    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.