diff --git a/agent_sdks/python/src/a2ui/core/schema/catalog.py b/agent_sdks/python/src/a2ui/core/schema/catalog.py index b845c170c..15898a3a6 100644 --- a/agent_sdks/python/src/a2ui/core/schema/catalog.py +++ b/agent_sdks/python/src/a2ui/core/schema/catalog.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import collections import copy import json import logging @@ -102,7 +103,7 @@ def with_pruned_components(self, allowed_components: List[str]) -> "A2uiCatalog" # Allow all components if no allowed components are specified if not allowed_components: - return self + return self._with_pruned_common_types() if CATALOG_COMPONENTS_KEY in schema_copy and isinstance( schema_copy[CATALOG_COMPONENTS_KEY], dict @@ -132,7 +133,58 @@ def with_pruned_components(self, allowed_components: List[str]) -> "A2uiCatalog" any_comp["oneOf"] = filtered_one_of - return replace(self, catalog_schema=schema_copy) + pruned_catalog = replace(self, catalog_schema=schema_copy) + return pruned_catalog._with_pruned_common_types() + + def _with_pruned_common_types(self) -> "A2uiCatalog": + """Returns a new catalog with unused common types pruned from the schema.""" + if not self.common_types_schema or "$defs" not in self.common_types_schema: + return self + + def _collect_refs(obj: Any) -> set[str]: + refs = set() + if isinstance(obj, dict): + for k, v in obj.items(): + if k == "$ref" and isinstance(v, str): + refs.add(v) + else: + refs.update(_collect_refs(v)) + elif isinstance(obj, list): + for item in obj: + refs.update(_collect_refs(item)) + return refs + + visited_defs = set() + internal_refs_queue = collections.deque() + + # Initialize queue with ONLY refs targeting common_types.json from external schemas + external_refs = _collect_refs(self.catalog_schema) + external_refs.update(_collect_refs(self.s2c_schema)) + + for ref in external_refs: + if ref.startswith("common_types.json#/$defs/"): + internal_refs_queue.append(ref.split("#/$defs/")[-1]) + + while internal_refs_queue: + def_name = internal_refs_queue.popleft() + if def_name in self.common_types_schema["$defs"] and def_name not in visited_defs: + visited_defs.add(def_name) + + # Collect internal references (which just use #/$defs/) + internal_refs = _collect_refs(self.common_types_schema["$defs"][def_name]) + for ref in internal_refs: + if ref.startswith("#/$defs/"): + # Note: This assumes a flat `$defs` namespace and no escaped + # slashes (~1) or tildes (~0) in the definition names as per RFC 6901. + internal_refs_queue.append(ref.split("#/$defs/")[-1]) + + new_common_types_schema = copy.deepcopy(self.common_types_schema) + all_defs = new_common_types_schema["$defs"] + new_common_types_schema["$defs"] = { + k: v for k, v in all_defs.items() if k in visited_defs + } + + return replace(self, common_types_schema=new_common_types_schema) def render_as_llm_instructions(self) -> str: """Renders the catalog and schema as LLM instructions.""" @@ -144,7 +196,11 @@ def render_as_llm_instructions(self) -> str: ) all_schemas.append(f"### Server To Client Schema:\n{server_client_str}") - if self.common_types_schema: + if ( + self.common_types_schema + and "$defs" in self.common_types_schema + and self.common_types_schema["$defs"] + ): common_str = json.dumps(self.common_types_schema, indent=2) all_schemas.append(f"### Common Types Schema:\n{common_str}") diff --git a/agent_sdks/python/tests/core/schema/test_catalog.py b/agent_sdks/python/tests/core/schema/test_catalog.py index f29d37a4d..8c773ac34 100644 --- a/agent_sdks/python/tests/core/schema/test_catalog.py +++ b/agent_sdks/python/tests/core/schema/test_catalog.py @@ -193,7 +193,7 @@ def test_render_as_llm_instructions(): version=VERSION_0_9, name=BASIC_CATALOG_NAME, s2c_schema={"s2c": "schema"}, - common_types_schema={"common": "types"}, + common_types_schema={"$defs": {"common": "types"}}, catalog_schema={ "$schema": "https://json-schema.org/draft/2020-12/schema", "catalog": "schema", @@ -204,8 +204,87 @@ def test_render_as_llm_instructions(): schema_str = catalog.render_as_llm_instructions() assert A2UI_SCHEMA_BLOCK_START in schema_str assert '### Server To Client Schema:\n{\n "s2c": "schema"\n}' in schema_str - assert '### Common Types Schema:\n{\n "common": "types"\n}' in schema_str + assert ( + '### Common Types Schema:\n{\n "$defs": {\n "common": "types"\n }\n}' + in schema_str + ) assert "### Catalog Schema:" in schema_str assert '"catalog": "schema"' in schema_str assert '"catalogId": "id_basic"' in schema_str assert A2UI_SCHEMA_BLOCK_END in schema_str + + +def test_render_as_llm_instructions_drops_empty_common_types(): + # Test with empty common_types_schema + catalog_empty = A2uiCatalog( + version=VERSION_0_9, + name=BASIC_CATALOG_NAME, + s2c_schema={"s2c": "schema"}, + common_types_schema={}, + catalog_schema={ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "catalog": "schema", + "catalogId": "id_basic", + }, + ) + schema_str_empty = catalog_empty.render_as_llm_instructions() + assert "### Common Types Schema:" not in schema_str_empty + + # Test with common_types_schema missing $defs + catalog_no_defs = A2uiCatalog( + version=VERSION_0_9, + name=BASIC_CATALOG_NAME, + s2c_schema={"s2c": "schema"}, + common_types_schema={"something": "else"}, + catalog_schema={ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "catalog": "schema", + "catalogId": "id_basic", + }, + ) + schema_str_no_defs = catalog_no_defs.render_as_llm_instructions() + assert "### Common Types Schema:" not in schema_str_no_defs + + # Test with common_types_schema having empty $defs + catalog_empty_defs = A2uiCatalog( + version=VERSION_0_9, + name=BASIC_CATALOG_NAME, + s2c_schema={"s2c": "schema"}, + common_types_schema={"$defs": {}}, + catalog_schema={ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "catalog": "schema", + "catalogId": "id_basic", + }, + ) + schema_str_empty_defs = catalog_empty_defs.render_as_llm_instructions() + assert "### Common Types Schema:" not in schema_str_empty_defs + + +def test_with_pruned_components_prunes_common_types(): + common_types = { + "$defs": { + "TypeForCompA": {"type": "string"}, + "TypeForCompB": {"type": "number"}, + } + } + catalog_schema = { + "catalogId": "basic", + "components": { + "CompA": {"$ref": "common_types.json#/$defs/TypeForCompA"}, + "CompB": {"$ref": "common_types.json#/$defs/TypeForCompB"}, + }, + } + catalog = A2uiCatalog( + version=VERSION_0_8, + name=BASIC_CATALOG_NAME, + s2c_schema={}, + common_types_schema=common_types, + catalog_schema=catalog_schema, + ) + + pruned_catalog = catalog.with_pruned_components(["CompA"]) + pruned_defs = pruned_catalog.common_types_schema["$defs"] + + assert "TypeForCompA" in pruned_defs + assert "TypeForCompB" not in pruned_defs diff --git a/agent_sdks/python/tests/core/schema/test_schema_manager.py b/agent_sdks/python/tests/core/schema/test_schema_manager.py index 9355fd776..777832b59 100644 --- a/agent_sdks/python/tests/core/schema/test_schema_manager.py +++ b/agent_sdks/python/tests/core/schema/test_schema_manager.py @@ -279,7 +279,7 @@ def joinpath_side_effect(path): content = '{"$schema": "https://json-schema.org/draft/2020-12/schema"}' if path == "common_types.json": content = ( - '{"$schema": "https://json-schema.org/draft/2020-12/schema", "types":' + '{"$schema": "https://json-schema.org/draft/2020-12/schema", "$defs":' ' {"Common": {}}}' ) elif "server_to_client" in path: @@ -290,9 +290,9 @@ def joinpath_side_effect(path): elif "catalog" in path: content = ( '{"$schema": "https://json-schema.org/draft/2020-12/schema", "catalogId":' - ' "basic", "components": {}}' + ' "basic", "components": {}, "$defs": {"test": {"$ref":' + ' "common_types.json#/$defs/Common"}}}' ) - mock_file.open.return_value.__enter__.return_value = io.StringIO(content) return mock_file @@ -306,7 +306,7 @@ def joinpath_side_effect(path): prompt = manager.generate_system_prompt("Role", include_schema=True) assert "### Common Types Schema:" in prompt - assert '"types":{"Common":{}}' in prompt.replace(" ", "").replace("\n", "") + assert '"$defs":{"Common":{}}' in prompt.replace(" ", "").replace("\n", "") def test_generate_system_prompt_minimal_args(mock_importlib_resources):