Refactor: extract basic catalog logic into its own package

Author: nan-yu

a2a_agents/python/a2ui_agent/pack_specs_hook.py

@@ -14,41 +14,58 @@
 
 import importlib.util
 import os
+import sys
 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}")
+def load_module(project_root, rel_path, filename, module_name):
+  """Loads a module directly from its path in src/."""
+  path = os.path.join(project_root, "src", rel_path.replace(".", os.sep), filename)
+  if not os.path.exists(path):
+    raise RuntimeError(f"Could not find module at {path}")
 
-  spec = importlib.util.spec_from_file_location("_constants_load", constants_path)
+  # Add src to sys.path so absolute imports work
+  src_path = os.path.abspath(os.path.join(project_root, "src"))
+  if src_path not in sys.path:
+    sys.path.insert(0, src_path)
+
+  spec = importlib.util.spec_from_file_location(module_name, path)
   if spec and spec.loader:
     module = importlib.util.module_from_spec(spec)
+    # Set the package context to allow relative imports if any
+    module.__package__ = rel_path
+    sys.modules[module_name] = module
     spec.loader.exec_module(module)
     return module
-  raise RuntimeError(f"Could not load shared constants from {constants_path}")
+  raise RuntimeError(f"Could not load module from {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)
+    # Load constants and utils dynamically from src/
+    schema_path = "a2ui.inference.schema"
+    a2ui_constants = load_module(
+        project_root, schema_path, "constants.py", "_constants_load"
+    )
+    a2ui_utils = load_module(project_root, schema_path, "utils.py", "_utils_load")
+
+    basic_catalog_constants = load_module(
+        project_root,
+        "a2ui.inference.basic_catalog",
+        "constants.py",
+        "_basic_catalog_constants_load",
+    )
 
     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)
+    repo_root = a2ui_utils.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.
@@ -66,23 +83,40 @@ def initialize(self, version, build_data):
         project_root, "src", a2ui_asset_package.replace(".", os.sep)
     )
 
-    for ver, schema_map in spec_version_map.items():
+    self._pack_schemas(repo_root, spec_version_map, target_base)
+    self._pack_basic_catalogs(
+        repo_root, basic_catalog_constants.BASIC_CATALOG_PATHS, target_base
+    )
+
+  def _pack_schemas(self, repo_root, spec_map, target_base):
+    for ver, schema_map in spec_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)
+        self._copy_schema(repo_root, source_rel_path, target_dir)
 
-        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
+  def _pack_basic_catalogs(self, repo_root, catalog_paths, target_base):
+    for ver, path_map in catalog_paths.items():
+      target_dir = os.path.join(target_base, ver)
+      os.makedirs(target_dir, exist_ok=True)
+
+      for _key, source_rel_path in path_map.items():
+        self._copy_schema(repo_root, source_rel_path, target_dir)
+
+  def _copy_schema(self, repo_root, source_rel_path, target_dir):
+    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."
+      )
+      return
 
-        filename = os.path.basename(source_path)
-        dst_file = os.path.join(target_dir, filename)
+    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)
+    print(f"Copying {source_path} -> {dst_file}")
+    shutil.copy2(source_path, dst_file)

a2a_agents/python/a2ui_agent/src/a2ui/inference/basic_catalog/__init__.py

@@ -0,0 +1,15 @@
+# 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 .provider import BasicCatalog

a2a_agents/python/a2ui_agent/src/a2ui/inference/basic_catalog/constants.py

@@ -0,0 +1,25 @@
+# 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 ..schema.constants import CATALOG_SCHEMA_KEY, VERSION_0_8, VERSION_0_9
+
+BASIC_CATALOG_NAME = "basic"
+
+# Maps version to the relative path of the basic catalog schema in the source repo
+BASIC_CATALOG_PATHS = {
+    VERSION_0_8: {
+        CATALOG_SCHEMA_KEY: "specification/v0_8/json/standard_catalog_definition.json"
+    },
+    VERSION_0_9: {CATALOG_SCHEMA_KEY: "specification/v0_9/json/basic_catalog.json"},
+}

a2a_agents/python/a2ui_agent/src/a2ui/inference/basic_catalog/provider.py

@@ -0,0 +1,61 @@
+# 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 typing import Any, Dict, Optional
+
+from ..schema.catalog import CatalogConfig
+from ..schema.catalog_provider import A2uiCatalogProvider
+from ..schema.utils import load_from_bundled_resource
+from ..schema.constants import BASE_SCHEMA_URL, CATALOG_ID_KEY, CATALOG_SCHEMA_KEY
+from .constants import BASIC_CATALOG_NAME, BASIC_CATALOG_PATHS
+
+
+class BundledCatalogProvider(A2uiCatalogProvider):
+  """Loads schemas from bundled package resources with fallbacks."""
+
+  def __init__(self, version: str):
+    self.version = version
+
+  def load(self) -> Dict[str, Any]:
+    # Use load_from_bundled_resource but with the specialized basic catalog paths
+    resource = load_from_bundled_resource(
+        self.version, CATALOG_SCHEMA_KEY, BASIC_CATALOG_PATHS
+    )
+
+    # Post-load processing for catalogs
+    if CATALOG_ID_KEY not in resource:
+      spec_map = BASIC_CATALOG_PATHS.get(self.version)
+      if spec_map and CATALOG_SCHEMA_KEY in spec_map:
+        rel_path = spec_map[CATALOG_SCHEMA_KEY]
+        # Strip the `json/` part from the catalog file path for the ID.
+        catalog_file = rel_path.replace("/json/", "/")
+        resource[CATALOG_ID_KEY] = BASE_SCHEMA_URL + catalog_file
+
+    if "$schema" not in resource:
+      resource["$schema"] = "https://json-schema.org/draft/2020-12/schema"
+
+    return resource
+
+
+class BasicCatalog:
+  """Helper for accessing the basic A2UI catalog."""
+
+  @staticmethod
+  def get_config(version: str, examples_path: Optional[str] = None) -> CatalogConfig:
+    """Returns a CatalogConfig for the basic bundled catalog."""
+    return CatalogConfig(
+        name=BASIC_CATALOG_NAME,
+        provider=BundledCatalogProvider(version),
+        examples_path=examples_path,
+    )

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

@@ -19,8 +19,8 @@
 from dataclasses import dataclass, field, replace
 from typing import Any, Dict, List, Optional, TYPE_CHECKING
 
-from .catalog_provider import A2uiCatalogProvider, FileSystemCatalogProvider, BundledCatalogProvider
-from .constants import CATALOG_COMPONENTS_KEY, CATALOG_ID_KEY, BASIC_CATALOG_NAME
+from .catalog_provider import A2uiCatalogProvider, FileSystemCatalogProvider
+from .constants import CATALOG_COMPONENTS_KEY, CATALOG_ID_KEY
 
 
 @dataclass
@@ -41,17 +41,6 @@ class CatalogConfig:
   provider: A2uiCatalogProvider
   examples_path: Optional[str] = None
 
-  @classmethod
-  def bundled(
-      cls, version: str, examples_path: Optional[str] = None
-  ) -> "CatalogConfig":
-    """Returns a CatalogConfig for the basic bundled catalog."""
-    return cls(
-        name=BASIC_CATALOG_NAME,
-        provider=BundledCatalogProvider(version),
-        examples_path=examples_path,
-    )
-
   @classmethod
   def from_path(
       cls, name: str, catalog_path: str, examples_path: Optional[str] = None

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

@@ -15,23 +15,10 @@
 """Module for providing A2UI catalog schemas and resources."""
 
 import json
-import logging
-import os
-import importlib.resources
 from abc import ABC, abstractmethod
 from json.decoder import JSONDecodeError
-from typing import Any, Dict, Optional
-
-from .constants import (
-    A2UI_ASSET_PACKAGE,
-    BASE_SCHEMA_URL,
-    CATALOG_ID_KEY,
-    CATALOG_SCHEMA_KEY,
-    SPEC_VERSION_MAP,
-    find_repo_root,
-)
-
-ENCODING = "utf-8"
+from typing import Any, Dict
+from .constants import ENCODING
 
 
 class A2uiCatalogProvider(ABC):
@@ -59,84 +46,3 @@ def load(self) -> Dict[str, Any]:
         return json.load(f)
     except (FileNotFoundError, JSONDecodeError) as e:
       raise IOError(f"Could not load schema from {self.path}: {e}") from e
-
-
-def load_from_bundled_resource(version: str, resource_key: str) -> Dict[str, Any]:
-  """Loads a schema resource from bundled package resources."""
-  spec_map = SPEC_VERSION_MAP.get(version)
-  if not spec_map:
-    raise ValueError(f"Unknown A2UI version: {version}")
-
-  if resource_key not in spec_map:
-    return None
-
-  rel_path = spec_map[resource_key]
-  filename = os.path.basename(rel_path)
-
-  # 1. Try to load from the bundled package resources.
-  try:
-    traversable = importlib.resources.files(A2UI_ASSET_PACKAGE)
-    traversable = traversable.joinpath(version).joinpath(filename)
-    with traversable.open("r", encoding=ENCODING) as f:
-      return json.load(f)
-  except Exception as e:
-    logging.debug("Could not load '%s' from package resources: %s", filename, e)
-
-  # 2. Fallback to local assets
-  # This handles cases where assets might be present in src but not installed
-  try:
-    potential_path = os.path.abspath(
-        os.path.join(
-            os.path.dirname(__file__),
-            "..",
-            "assets",
-            version,
-            filename,
-        )
-    )
-    if os.path.exists(potential_path):
-      provider = FileSystemCatalogProvider(potential_path)
-      return provider.load()
-  except Exception as e:
-    logging.debug("Could not load schema '%s' from local assets: %s", filename, e)
-
-  # 3. Fallback: Source Repository (specification/...)
-  # This handles cases where we are running directly from source tree
-  # And assets are not yet copied to src/a2ui/assets
-  # manager.py is at a2a_agents/python/a2ui_agent/src/a2ui/inference/schema/manager.py
-  # Dynamically find repo root by looking for "specification" directory
-  try:
-    repo_root = find_repo_root(os.path.dirname(__file__))
-    if repo_root:
-      source_path = os.path.join(repo_root, rel_path)
-      if os.path.exists(source_path):
-        provider = FileSystemCatalogProvider(source_path)
-        return provider.load()
-  except Exception as e:
-    logging.debug("Could not load schema from source repo: %s", e)
-
-  raise IOError(f"Could not load schema {filename} for version {version}")
-
-
-class BundledCatalogProvider(A2uiCatalogProvider):
-  """Loads schemas from bundled package resources with fallbacks."""
-
-  def __init__(self, version: str):
-    self.version = version
-
-  def load(self) -> Dict[str, Any]:
-
-    resource = load_from_bundled_resource(self.version, CATALOG_SCHEMA_KEY)
-
-    # Post-load processing for catalogs
-    if CATALOG_ID_KEY not in resource:
-      spec_map = SPEC_VERSION_MAP.get(self.version)
-      rel_path = spec_map[CATALOG_SCHEMA_KEY]
-      # Strip the `json/` part from the catalog file path for the ID.
-      catalog_file = rel_path.replace("/json/", "/")
-      resource[CATALOG_ID_KEY] = BASE_SCHEMA_URL + catalog_file
-
-    if "$schema" not in resource:
-      resource["$schema"] = "https://json-schema.org/draft/2020-12/schema"
-
-    return resource