Skip to content

Remove notion of "custom catalogs" from agent SDK #705

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nan-yu wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Remove notion of "custom catalogs" from agent SDK #705

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion a2a_agents/python/a2ui_agent/src/a2ui/extension/a2ui_extension.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ def get_a2ui_agent_extension(
"""Creates the A2UI AgentExtension configuration.

Args:
accepts_inline_catalogs: Whether the agent accepts inline custom catalogs.
accepts_inline_catalogs: Whether the agent accepts inline catalogs.
supported_catalog_ids: All pre-defined catalogs the agent is known to support.

Returns:
Expand Down
135 changes: 25 additions & 110 deletions a2a_agents/python/a2ui_agent/src/a2ui/inference/schema/catalog.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from dataclasses import dataclass, field, replace
from typing import Any, Dict, List, Optional, TYPE_CHECKING

from .constants import CATALOG_COMPONENTS_KEY, CATALOG_ID_KEY
from .constants import CATALOG_COMPONENTS_KEY, CATALOG_ID_KEY, BASIC_CATALOG_NAME
from referencing import Registry, Resource

if TYPE_CHECKING:
Expand All @@ -28,13 +28,35 @@


@dataclass
class CustomCatalogConfig:
"""Configuration for a custom component catalog."""
class CatalogConfig:
"""
Configuration for a catalog of components.

The catalog must be free standing and not reference any other catalogs, except
for the common types schema.

If a catalog references other catalogs, the references must be resolved before
loading the catalog by using `tools/build_catalog/build_catalog.py` script.

Attributes:
name: The name of the catalog.
catalog_path: The path to the catalog JSON file.
examples_path: The path to the examples directory.
"""

name: str
catalog_path: str
examples_path: Optional[str] = None

@classmethod
def for_basic_catalog(cls, examples_path: Optional[str] = None) -> "CatalogConfig":
Copy link
Collaborator

@jacobsimionato jacobsimionato Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we move this to a top-level function in a different file, perhaps even a different folder, e.g. inference/basic_catalog or something?

That makes the conceptual separation clear, and avoids us accidentally sprinkling basic catalog assumptions into the core library. Ideally, it would be possible for us to publish basic catalog support as a completely separate library which depends on the core inference library.

Copy link
Collaborator

@wrenj wrenj Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1 on this, then we can run the basic catalog through the sdk like any other catalog

"""Returns a CatalogConfig for the basic catalog."""
return cls(
name=BASIC_CATALOG_NAME,
catalog_path="", # basic catalog doesn't need a path from the user
Copy link
Collaborator

@jacobsimionato jacobsimionato Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why doesn't this need a path? Otherwise, the core code won't know where to load it from, right?

Copy link
Collaborator Author

@nan-yu nan-yu Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The basic catalog and other core schemas (server-to-client and common types) are bundled directly within the SDK. pack_specs_hook.py packs the specs into the package assets. The SDK then loads them as package resources. We only require a path when a user is loading a custom catalog from their own filesystem.

examples_path=examples_path,
Copy link
Collaborator

@wrenj wrenj Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe just examples_paths=

versus doing an ls of a directory? what if there .git file in that directory?

)


@dataclass(frozen=True)
class A2uiCatalog:
Expand Down Expand Up @@ -165,113 +187,6 @@ def load_examples(self, path: Optional[str], validate: bool = False) -> str:
return ""
return "\n\n".join(merged_examples)

@staticmethod
def resolve_schema(basic: Dict[str, Any], custom: Dict[str, Any]) -> Dict[str, Any]:
"""Resolves references in custom catalog schema against the basic catalog.

Args:
basic: The basic catalog schema.
custom: The custom catalog schema.

Returns:
A new dictionary with references resolved.
"""
result = copy.deepcopy(custom)

# Initialize registry with basic catalog and maybe others from basic's $id
registry = Registry()
if CATALOG_ID_KEY in basic:
basic_resource = Resource.from_contents(basic)
registry = registry.with_resource(basic[CATALOG_ID_KEY], basic_resource)

def resolve_ref(ref_uri: str) -> Any:
try:
resolver = registry.resolver()
resolved = resolver.lookup(ref_uri)
return resolved.contents
except Exception as e:
logging.warning("Could not resolve reference %s: %s", ref_uri, e)
return None

def merge_into(target: Dict[str, Any], source: Dict[str, Any]):
for key, value in source.items():
if key not in target:
target[key] = copy.deepcopy(value)

# Process components
if CATALOG_COMPONENTS_KEY in result:
comp_dict = result[CATALOG_COMPONENTS_KEY]
if "$ref" in comp_dict:
resolved = resolve_ref(comp_dict["$ref"])
if isinstance(resolved, dict):
merge_into(comp_dict, resolved)
del comp_dict["$ref"]

# Process functions
if "functions" in result:
func_dict = result["functions"]
if "$ref" in func_dict:
resolved = resolve_ref(func_dict["$ref"])
if isinstance(resolved, dict):
merge_into(func_dict, resolved)
del func_dict["$ref"]

# Process $defs
if "$defs" in result:
res_defs = result["$defs"]
if "$ref" in res_defs:
resolved = resolve_ref(res_defs["$ref"])
if isinstance(resolved, dict):
merge_into(res_defs, resolved)
del res_defs["$ref"]

for name in ["anyComponent", "anyFunction"]:
if name in res_defs:
target = res_defs[name]
one_of = target.get("oneOf", [])
new_one_of = []
for item in one_of:
if isinstance(item, dict) and "$ref" in item:
ref_uri = item["$ref"]
# Check if it points to basic collector
resolved = resolve_ref(ref_uri)
if isinstance(resolved, dict) and "oneOf" in resolved:
# Merge oneOf items and resolve transitive refs to components/functions
for sub_item in resolved["oneOf"]:
if sub_item not in new_one_of:
new_one_of.append(copy.deepcopy(sub_item))
# Transitive merge: if sub_item is a ref to a component/function
if isinstance(sub_item, dict) and "$ref" in sub_item:
sub_ref = sub_item["$ref"]
if (
sub_ref.startswith("#/components/")
and CATALOG_COMPONENTS_KEY in basic
):
comp_name = sub_ref.split("/")[-1]
if comp_name in basic[CATALOG_COMPONENTS_KEY]:
if CATALOG_COMPONENTS_KEY not in result:
result[CATALOG_COMPONENTS_KEY] = {}
if comp_name not in result[CATALOG_COMPONENTS_KEY]:
result[CATALOG_COMPONENTS_KEY][comp_name] = copy.deepcopy(
basic[CATALOG_COMPONENTS_KEY][comp_name]
)
elif sub_ref.startswith("#/functions/") and "functions" in basic:
func_name = sub_ref.split("/")[-1]
if func_name in basic["functions"]:
if "functions" not in result:
result["functions"] = {}
if func_name not in result["functions"]:
result["functions"][func_name] = copy.deepcopy(
basic["functions"][func_name]
)
else:
new_one_of.append(item)
else:
new_one_of.append(item)
target["oneOf"] = new_one_of

return result

def _validate_example(self, full_path: str, basename: str, content: str) -> bool:
try:
json_data = json.loads(content)
Expand Down
122 changes: 54 additions & 68 deletions a2a_agents/python/a2ui_agent/src/a2ui/inference/schema/manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@
BASIC_CATALOG_NAME,
find_repo_root,
)
from .catalog import CustomCatalogConfig, A2uiCatalog
from .catalog import CatalogConfig, A2uiCatalog
from ...extension.a2ui_extension import INLINE_CATALOGS_KEY, SUPPORTED_CATALOG_IDS_KEY, get_a2ui_agent_extension
from a2a.types import AgentExtension

Expand Down Expand Up @@ -118,30 +118,28 @@ class A2uiSchemaManager(InferenceStrategy):
def __init__(
self,
version: str,
basic_examples_path: Optional[str] = None,
custom_catalogs: Optional[List[CustomCatalogConfig]] = None,
exclude_basic_catalog: bool = False,
catalogs: Optional[List[CatalogConfig]] = None,
accepts_inline_catalogs: bool = False,
schema_modifiers: List[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
schema_modifiers: Optional[
List[Callable[[Dict[str, Any]], Dict[str, Any]]]
] = None,
):
self._version = version
self._exclude_basic_catalog = exclude_basic_catalog
self._accepts_inline_catalogs = accepts_inline_catalogs

self._server_to_client_schema = None
self._common_types_schema = None
self._supported_catalogs: Dict[str, A2uiCatalog] = {}
self._supported_catalogs: List[A2uiCatalog] = []
self._catalog_example_paths: Dict[str, str] = {}
self._basic_catalog = None
self._schema_modifiers = schema_modifiers
self._load_schemas(version, custom_catalogs, basic_examples_path)
self._schema_modifiers = schema_modifiers or []
self._load_schemas(version, catalogs or [])

@property
def accepts_inline_catalogs(self) -> bool:
return self._accepts_inline_catalogs

@property
def supported_catalogs(self) -> Dict[str, A2uiCatalog]:
def supported_catalogs(self) -> List[A2uiCatalog]:
return self._supported_catalogs

def _apply_modifiers(self, schema: Dict[str, Any]) -> Dict[str, Any]:
Expand All @@ -153,8 +151,7 @@ def _apply_modifiers(self, schema: Dict[str, Any]) -> Dict[str, Any]:
def _load_schemas(
self,
version: str,
custom_catalogs: Optional[List[CustomCatalogConfig]] = None,
basic_examples_path: Optional[str] = None,
catalogs: List[CatalogConfig] = [],
):
"""Loads separate schema components and processes catalogs."""
if version not in SPEC_VERSION_MAP:
Expand All @@ -172,11 +169,7 @@ def _load_schemas(
)

# Process basic catalog
basic_catalog_schema = self._apply_modifiers(
_load_basic_component(version, CATALOG_SCHEMA_KEY)
)
if not basic_catalog_schema:
basic_catalog_schema = {}
basic_catalog_schema = _load_basic_component(version, CATALOG_SCHEMA_KEY)
Copy link
Collaborator

@jacobsimionato jacobsimionato Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ideally, avoid having this special cased logic. I wonder if we can publish the basic catalog schema as a JSON file bundled with the library?

Or perhaps the parameter of CatalogConfig could actually just be a "CatalogProvider" which could either be a FileReadingCatalogProvider which reads a JSON file on disk, or just an InlineCatalogProvider which returns a catalog baked into Python code as a string constant (so it's easy for us to create for the basic catalog).


# Ensure catalog id and schema url are set in the basic catalog schema
if CATALOG_ID_KEY not in basic_catalog_schema:
Expand All @@ -190,44 +183,36 @@ def _load_schemas(
if "$schema" not in basic_catalog_schema:
basic_catalog_schema["$schema"] = "https://json-schema.org/draft/2020-12/schema"

self._basic_catalog = A2uiCatalog(
version=version,
name=BASIC_CATALOG_NAME,
catalog_schema=basic_catalog_schema,
s2c_schema=self._server_to_client_schema,
common_types_schema=self._common_types_schema,
)
if not self._exclude_basic_catalog:
self._supported_catalogs[self._basic_catalog.catalog_id] = self._basic_catalog
self._catalog_example_paths[self._basic_catalog.catalog_id] = basic_examples_path

# Process custom catalogs
if custom_catalogs:
for config in custom_catalogs:
custom_catalog_schema = self._apply_modifiers(
_load_from_path(config.catalog_path)
)
resolved_catalog_schema = A2uiCatalog.resolve_schema(
basic_catalog_schema, custom_catalog_schema
)
catalog = A2uiCatalog(
version=version,
name=config.name,
catalog_schema=self._apply_modifiers(resolved_catalog_schema),
s2c_schema=self._server_to_client_schema,
common_types_schema=self._common_types_schema,
)
self._supported_catalogs[catalog.catalog_id] = catalog
self._catalog_example_paths[catalog.catalog_id] = config.examples_path
# Process catalogs
if not catalogs:
# If no catalogs are provided, use the basic catalog
catalogs = [CatalogConfig.for_basic_catalog()]

for config in catalogs:
catalog_schema = (
basic_catalog_schema
if config.name == BASIC_CATALOG_NAME
else _load_from_path(config.catalog_path)
)
catalog_schema = self._apply_modifiers(catalog_schema)
catalog = A2uiCatalog(
version=version,
name=config.name,
catalog_schema=catalog_schema,
s2c_schema=self._server_to_client_schema,
common_types_schema=self._common_types_schema,
)
self._supported_catalogs.append(catalog)
self._catalog_example_paths[catalog.catalog_id] = config.examples_path

def _determine_catalog(
self, client_ui_capabilities: Optional[dict[str, Any]] = None
) -> A2uiCatalog:
"""Determines the catalog to use based on supported catalog IDs.

If neither inline catalogs nor supported catalog IDs are provided, the basic catalog is used.
If neither inline catalogs nor client-supported catalog IDs are provided, the first agent-supported catalog is used.
Copy link
Collaborator

@wrenj wrenj Feb 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is it first in the client_capabitilities or first in the supported_catalogs?

If inline catalogs are provided, the first inline catalog is used.
If supported catalog IDs are provided, the first supported catalog that is recognized is used.
If client-supported catalog IDs are provided, the first one that is supported by the agent is used.

Args:
client_ui_capabilities: A dictionary of client UI capabilities.
Expand All @@ -236,16 +221,19 @@ def _determine_catalog(
The A2uiCatalog to use to generate the schema string in the prompt.

Raises:
ValueError: If both inline catalogs and supported catalog IDs are provided,
ValueError: If both inline catalogs and client-supported catalog IDs are provided,
or if no supported catalog is recognized.
"""
if not self._supported_catalogs:
raise ValueError("No supported catalogs found.") # This should not happen.

if not client_ui_capabilities or not isinstance(client_ui_capabilities, dict):
return self._basic_catalog
return self._supported_catalogs[0]

inline_catalogs: List[dict[str, Any]] = client_ui_capabilities.get(
INLINE_CATALOGS_KEY, []
)
supported_catalog_ids: List[str] = client_ui_capabilities.get(
client_supported_catalog_ids: List[str] = client_ui_capabilities.get(
SUPPORTED_CATALOG_IDS_KEY, []
)

Expand All @@ -255,37 +243,35 @@ def _determine_catalog(
" capabilities. However, the agent does not accept inline catalogs."
)

if inline_catalogs and supported_catalog_ids:
if inline_catalogs and client_supported_catalog_ids:
raise ValueError(
f"Both '{INLINE_CATALOGS_KEY}' and '{SUPPORTED_CATALOG_IDS_KEY}' "
"are provided in client UI capabilities. Only one is allowed."
)

if inline_catalogs:
# Load the first custom inline catalog schema.
# Load the first inline catalog schema.
inline_catalog_schema = inline_catalogs[0]
resolved_catalog_schema = A2uiCatalog.resolve_schema(
self._basic_catalog.catalog_schema, inline_catalog_schema
)
inline_catalog_schema = self._apply_modifiers(inline_catalog_schema)
return A2uiCatalog(
version=self._version,
name=INLINE_CATALOG_NAME,
catalog_schema=resolved_catalog_schema,
catalog_schema=inline_catalog_schema,
s2c_schema=self._server_to_client_schema,
common_types_schema=self._common_types_schema,
)

if not supported_catalog_ids:
return self._basic_catalog
if not client_supported_catalog_ids:
return self._supported_catalogs[0]

for scid in supported_catalog_ids:
if scid in self._supported_catalogs:
# Return the first supported catalog.
return self._supported_catalogs[scid]
agent_supported_catalogs = {c.catalog_id: c for c in self._supported_catalogs}
for cscid in client_supported_catalog_ids:
if cscid in agent_supported_catalogs:
return agent_supported_catalogs[cscid]

raise ValueError(
"No supported catalog found on the agent side. Agent supported catalogs are:"
f" {list(self._supported_catalogs.keys())}"
"No client-supported catalog found on the agent side. Agent-supported catalogs"
f" are: {[c.catalog_id for c in self._supported_catalogs]}"
)

def get_effective_catalog(
Expand Down Expand Up @@ -339,5 +325,5 @@ def generate_system_prompt(
return "\n\n".join(parts)

def get_agent_extension(self) -> AgentExtension:
catalog_ids = self._supported_catalogs.keys()
return get_a2ui_agent_extension(supported_catalog_ids=list(catalog_ids))
catalog_ids = [c.catalog_id for c in self._supported_catalogs]
return get_a2ui_agent_extension(supported_catalog_ids=catalog_ids)
2 changes: 1 addition & 1 deletion a2a_agents/python/a2ui_agent/tests/inference/test_modifiers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ def test_manager_with_modifiers():
assert "additionalProperties" not in manager._common_types_schema

# basic catalog should also be modified
for catalog in manager._supported_catalogs.values():
for catalog in manager._supported_catalogs:
assert "additionalProperties" not in catalog.catalog_schema


Expand Down
Loading
Loading