Refactor a2ui_agent to depend on a2ui_core

Author: nan-yu

agent_sdks/python/a2ui_agent/pyproject.toml

@@ -17,14 +17,22 @@ name = "a2ui-agent-sdk"
 dynamic = ["version"]
 description = "A2UI Agent SDK"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.14"
 dependencies = [
   "a2a-sdk>=0.3.0",
+  "a2ui-core>=0.1.0",
   "google-adk>=1.28.1",
   "google-genai>=1.27.0",
   "jsonschema>=4.0.0"
 ]
 
+[project.scripts]
+pyink = "pyink:patched_main"
+
+[tool.uv.sources]
+
+a2ui-core = { path = "../a2ui_core", editable = true }
+
 [build-system]
 requires = ["hatchling", "jsonschema"]
 build-backend = "hatchling.build"

agent_sdks/python/a2ui_agent/src/a2ui/__init__.py

@@ -11,4 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)
+
 from .version import __version__

agent_sdks/python/a2ui_agent/src/a2ui/basic_catalog/provider.py

@@ -19,6 +19,7 @@
 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
+from a2ui.core.catalog.json_catalog import JSON_SCHEMA_DRAFT_2020_12
 
 
 class BundledCatalogProvider(A2uiCatalogProvider):
@@ -43,7 +44,7 @@ def load(self) -> Dict[str, Any]:
         resource[CATALOG_ID_KEY] = BASE_SCHEMA_URL + catalog_file
 
     if "$schema" not in resource:
-      resource["$schema"] = "https://json-schema.org/draft/2020-12/schema"
+      resource["$schema"] = JSON_SCHEMA_DRAFT_2020_12
 
     return resource
 

agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming.py

@@ -32,6 +32,7 @@
     extract_component_ref_fields,
     extract_component_required_fields,
 )
+from ..schema.validator import A2uiValidator
 from .response_part import ResponsePart
 
 
@@ -48,31 +49,27 @@ class A2uiStreamParser:
   (V08 or V09) depending on the catalog version.
   """
 
-  def __new__(cls, catalog: "A2uiCatalog" = None, *args, **kwargs):
+  def __new__(cls, catalog: A2uiCatalog):
     if cls is A2uiStreamParser:
-      version = getattr(catalog, "version", None) if catalog else None
-      if version == VERSION_0_9:
-        from .streaming_v09 import A2uiStreamParserV09
+      version = catalog.version
+      # Lazy import inside __new__ to prevent circular import errors, as the
+      # version-specific subclass modules import A2uiStreamParser from this module.
+      if version == VERSION_0_8:
+        from .streaming_v08 import A2uiStreamParserV08
 
-        return A2uiStreamParserV09(catalog=catalog, *args, **kwargs)
+        return A2uiStreamParserV08(catalog=catalog)
       else:
-        from .streaming_v08 import A2uiStreamParserV08
+        from .streaming_v09 import A2uiStreamParserV09
 
-        return A2uiStreamParserV08(catalog=catalog, *args, **kwargs)
+        return A2uiStreamParserV09(catalog=catalog)
     return super().__new__(cls)
 
-  def __init__(self, catalog: "A2uiCatalog" = None):
-    self._version = getattr(catalog, "version", None) if catalog else None
-    self._cuttable_keys = (
-        getattr(catalog, "cuttable_keys", None) if catalog else frozenset()
-    )
-    self._ref_fields_map = extract_component_ref_fields(catalog) if catalog else {}
-    self._required_fields_map = (
-        extract_component_required_fields(catalog) if catalog else {}
-    )
-    from ..schema.validator import A2uiValidator
-
-    self._validator = A2uiValidator(catalog) if catalog else None
+  def __init__(self, catalog: A2uiCatalog):
+    self._version = catalog.version
+    self._cuttable_keys = catalog.cuttable_keys
+    self._ref_fields_map = extract_component_ref_fields(catalog)
+    self._required_fields_map = extract_component_required_fields(catalog)
+    self._validator = A2uiValidator(catalog)
 
     self._found_delimiter = False
     self._buffer = ""
@@ -878,6 +875,8 @@ def yield_reachable(
           self.root_id,
           components_to_analyze,
           self._ref_fields_map,
+          single_ref_fields=set(),
+          list_ref_fields=set(),
           raise_on_orphans=raise_on_orphans,
       )
 

agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v08.py

@@ -14,18 +14,21 @@
 
 import re
 import json
-from typing import Any, List, Dict, Optional, Set
+from typing import Any, List, Dict, Optional, Set, TYPE_CHECKING
 
 from .streaming import A2uiStreamParser
 from .response_part import ResponsePart
 from .constants import *
 from ..schema.constants import VERSION_0_8, SURFACE_ID_KEY, CATALOG_COMPONENTS_KEY
 
+if TYPE_CHECKING:
+  from ..schema.catalog import A2uiCatalog
+
 
 class A2uiStreamParserV08(A2uiStreamParser):
   """Streaming parser implementation for A2UI v0.8 specification."""
 
-  def __init__(self, catalog=None):
+  def __init__(self, catalog: A2uiCatalog):
     super().__init__(catalog=catalog)
     self._yielded_begin_rendering_surfaces: Set[str] = set()
 

agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v09.py

@@ -14,18 +14,21 @@
 
 import re
 import json
-from typing import Any, List, Dict, Optional, Set
+from typing import Any, List, Dict, Optional, Set, TYPE_CHECKING
 
 from .streaming import A2uiStreamParser
 from .response_part import ResponsePart
 from .constants import *
 from ..schema.constants import VERSION_0_9, SURFACE_ID_KEY, CATALOG_COMPONENTS_KEY
 
+if TYPE_CHECKING:
+  from ..schema.catalog import A2uiCatalog
+
 
 class A2uiStreamParserV09(A2uiStreamParser):
   """Streaming parser implementation for A2UI v0.9 specification."""
 
-  def __init__(self, catalog=None):
+  def __init__(self, catalog: A2uiCatalog):
     super().__init__(catalog=catalog)
     # v0.9 default root is "root"
     self._default_root_id = DEFAULT_ROOT_ID

agent_sdks/python/a2ui_agent/src/a2ui/schema/catalog.py

@@ -21,6 +21,7 @@
 from dataclasses import dataclass, field, replace
 from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING
 from urllib.parse import urlparse
+from a2ui.core.catalog import JsonCatalog
 
 from .catalog_provider import A2uiCatalogProvider, FileSystemCatalogProvider
 from .constants import (
@@ -32,7 +33,9 @@
     VERSION_0_8,
     ENCODING,
 )
-from .validator import A2uiValidator
+
+if TYPE_CHECKING:
+  from .validator import A2uiValidator
 
 
 @dataclass
@@ -48,12 +51,16 @@ class CatalogConfig:
     provider: The provider to use to load the catalog schema.
     examples_path: The path or glob pattern to the examples.
     custom_cuttable_keys: The optional custom set of cuttable keys.
+    custom_single_refs: The optional list of refs to be treated as single items.
+    custom_list_refs: The optional list of refs to be treated as list items.
   """
 
   name: str
   provider: A2uiCatalogProvider
   examples_path: Optional[str] = None
   custom_cuttable_keys: Optional[frozenset[str]] = None
+  custom_single_refs: Optional[List[str]] = None
+  custom_list_refs: Optional[List[str]] = None
 
   @classmethod
   def from_path(
@@ -62,6 +69,8 @@ def from_path(
       catalog_path: str,
       examples_path: Optional[str] = None,
       custom_cuttable_keys: Optional[frozenset[str]] = None,
+      custom_single_refs: Optional[List[str]] = None,
+      custom_list_refs: Optional[List[str]] = None,
   ) -> "CatalogConfig":
     """Returns a CatalogConfig that loads from a local path or 'file://' URI."""
     parsed = urlparse(catalog_path)
@@ -79,6 +88,10 @@ def from_path(
     }
     if custom_cuttable_keys is not None:
       kwargs["custom_cuttable_keys"] = custom_cuttable_keys
+    if custom_single_refs is not None:
+      kwargs["custom_single_refs"] = custom_single_refs
+    if custom_list_refs is not None:
+      kwargs["custom_list_refs"] = custom_list_refs
 
     return cls(**kwargs)
 
@@ -151,6 +164,8 @@ class A2uiCatalog:
     catalog_schema: The catalog schema.
     custom_cuttable_keys: The optional set of keys whose string values can be safely auto-closed
       (healed) if fragmented in the stream. If None, the default set is used.
+    custom_single_refs: The optional list of refs to be treated as single items.
+    custom_list_refs: The optional list of refs to be treated as list items.
   """
 
   version: str
@@ -159,6 +174,8 @@ class A2uiCatalog:
   common_types_schema: Dict[str, Any]
   catalog_schema: Dict[str, Any]
   custom_cuttable_keys: Optional[frozenset[str]] = None
+  custom_single_refs: Optional[List[str]] = None
+  custom_list_refs: Optional[List[str]] = None
 
   @property
   def cuttable_keys(self) -> frozenset[str]:
@@ -173,9 +190,22 @@ def catalog_id(self) -> str:
     return self.catalog_schema[CATALOG_ID_KEY]
 
   @property
-  def validator(self) -> A2uiValidator:
+  def validator(self) -> "A2uiValidator":
+    from .validator import A2uiValidator
+
     return A2uiValidator(self)
 
+  @property
+  def core_catalog(self) -> JsonCatalog:
+    return JsonCatalog(
+        version=self.version,
+        catalog_schema=self.catalog_schema,
+        catalog_id=self.catalog_id,
+        common_types_schema=self.common_types_schema,
+        custom_single_refs=self.custom_single_refs,
+        custom_list_refs=self.custom_list_refs,
+    )
+
   def _with_pruned_components(self, allowed_components: List[str]) -> A2uiCatalog:
     """Returns a new catalog with only allowed components.
 

agent_sdks/python/a2ui_agent/src/a2ui/schema/manager.py

@@ -95,6 +95,8 @@ def _load_schemas(
           s2c_schema=self._server_to_client_schema,
           common_types_schema=self._common_types_schema,
           custom_cuttable_keys=config.custom_cuttable_keys,
+          custom_single_refs=config.custom_single_refs,
+          custom_list_refs=config.custom_list_refs,
       )
       self._supported_catalogs.append(catalog)
       self._catalog_example_paths[catalog.catalog_id] = config.examples_path