Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
# 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.


def remove_strict_validation(schema):
if isinstance(schema, dict):
new_schema = {k: remove_strict_validation(v) for k, v in schema.items()}
if (
'additionalProperties' in new_schema
and new_schema['additionalProperties'] is False
):
del new_schema['additionalProperties']
return new_schema
elif isinstance(schema, list):
return [remove_strict_validation(item) for item in schema]
return schema
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
import logging
import os
import importlib.resources
from typing import List, Dict, Any, Optional
from typing import List, Dict, Any, Optional, Callable
from dataclasses import dataclass, field
from .loader import A2uiSchemaLoader, PackageLoader, FileSystemLoader
from ..inference_strategy import InferenceStrategy
Expand Down Expand Up @@ -122,6 +122,7 @@ def __init__(
custom_catalogs: Optional[List[CustomCatalogConfig]] = None,
exclude_basic_catalog: bool = False,
accepts_inline_catalogs: bool = False,
schema_modifiers: List[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
):
self._version = version
self._exclude_basic_catalog = exclude_basic_catalog
Expand All @@ -132,6 +133,7 @@ def __init__(
self._supported_catalogs: Dict[str, 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)

@property
Expand All @@ -142,6 +144,12 @@ def accepts_inline_catalogs(self) -> bool:
def supported_catalogs(self) -> Dict[str, A2uiCatalog]:
return self._supported_catalogs

def _apply_modifiers(self, schema: Dict[str, Any]) -> Dict[str, Any]:
if self._schema_modifiers:
for modifier in self._schema_modifiers:
schema = modifier(schema)
return schema

def _load_schemas(
self,
version: str,
Expand All @@ -156,13 +164,17 @@ def _load_schemas(
)

# Load server-to-client and common types schemas
self._server_to_client_schema = _load_basic_component(
version, SERVER_TO_CLIENT_SCHEMA_KEY
self._server_to_client_schema = self._apply_modifiers(
_load_basic_component(version, SERVER_TO_CLIENT_SCHEMA_KEY)
)
self._common_types_schema = self._apply_modifiers(
_load_basic_component(version, COMMON_TYPES_SCHEMA_KEY)
)
self._common_types_schema = _load_basic_component(version, COMMON_TYPES_SCHEMA_KEY)

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

Expand Down Expand Up @@ -192,14 +204,16 @@ def _load_schemas(
# Process custom catalogs
if custom_catalogs:
for config in custom_catalogs:
custom_catalog_schema = _load_from_path(config.catalog_path)
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=resolved_catalog_schema,
catalog_schema=self._apply_modifiers(resolved_catalog_schema),
s2c_schema=self._server_to_client_schema,
common_types_schema=self._common_types_schema,
)
Expand Down
74 changes: 74 additions & 0 deletions a2a_agents/python/a2ui_agent/tests/inference/test_modifiers.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
# 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 pytest
from unittest.mock import patch
from a2ui.inference.schema.manager import A2uiSchemaManager
from a2ui.inference.schema.common_modifiers import remove_strict_validation


def test_remove_strict_validation():
"""Tests the remove_strict_validation modifier."""
schema = {
"type": "object",
"properties": {
"a": {"type": "string", "additionalProperties": False},
"b": {
"type": "array",
"items": {"type": "object", "additionalProperties": False},
},
},
"additionalProperties": False,
}

modified = remove_strict_validation(schema)

# Check that additionalProperties: False is removed
assert "additionalProperties" not in modified
assert "additionalProperties" not in modified["properties"]["a"]
assert "additionalProperties" not in modified["properties"]["b"]["items"]

# Check that it didn't mutate the original
assert schema["additionalProperties"] is False
assert schema["properties"]["a"]["additionalProperties"] is False


def test_manager_with_modifiers():
"""Tests that A2uiSchemaManager applies modifiers during loading."""
# Mock _load_basic_component to return a simple schema with strict validation
mock_schema = {"type": "object", "additionalProperties": False}
with patch(
"a2ui.inference.schema.manager._load_basic_component", return_value=mock_schema
):
manager = A2uiSchemaManager("0.8", schema_modifiers=[remove_strict_validation])

# Verify that loaded schemas have modifiers applied
assert "additionalProperties" not in manager._server_to_client_schema
assert "additionalProperties" not in manager._common_types_schema

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


def test_manager_no_modifiers():
"""Tests that A2uiSchemaManager works fine without modifiers."""
mock_schema = {"type": "object", "additionalProperties": False}
with patch(
"a2ui.inference.schema.manager._load_basic_component", return_value=mock_schema
):
manager = A2uiSchemaManager("0.8", schema_modifiers=None)

# Verify that schemas are NOT modified
assert manager._server_to_client_schema["additionalProperties"] is False
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,13 @@

from a2ui.inference.schema.manager import A2uiSchemaManager
from a2ui.inference.schema.constants import CATALOG_COMPONENTS_KEY
from a2ui.inference.schema.common_modifiers import remove_strict_validation


def verify():
print('Verifying A2uiSchemaManager...')
try:
manager = A2uiSchemaManager('0.8')
manager = A2uiSchemaManager('0.8', schema_modifiers=[remove_strict_validation])
catalog = manager.get_effective_catalog()
catalog_components = catalog.catalog_schema[CATALOG_COMPONENTS_KEY]
print(f'Successfully loaded 0.8: {len(catalog_components)} components')
Expand Down Expand Up @@ -364,6 +365,13 @@ def verify():
'key': 'imageUrl',
'valueString': 'http://localhost:10003/static/profile2.png',
},
{
'key': 'contacts',
'valueMap': [{
'key': 'contact1',
'valueMap': [{'key': 'name', 'valueString': 'Casey Smith'}],
}],
},
],
}
},
Expand All @@ -375,7 +383,7 @@ def verify():
sys.exit(1)

try:
manager = A2uiSchemaManager('0.9')
manager = A2uiSchemaManager('0.9', schema_modifiers=[remove_strict_validation])
catalog = manager.get_effective_catalog()
catalog_components = catalog.catalog_schema[CATALOG_COMPONENTS_KEY]
print(f'Successfully loaded 0.9: {len(catalog_components)} components')
Expand All @@ -389,6 +397,7 @@ def verify():
'catalogId': (
'https://a2ui.dev/specification/v0_9/standard_catalog.json'
),
'fakeProperty': 'should be allowed',
},
},
{
Expand Down
38 changes: 5 additions & 33 deletions samples/agent/adk/contact_multiple_surfaces/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,6 @@
from a2a.server.apps import A2AStarletteApplication
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
from a2a.types import AgentCapabilities, AgentCard, AgentSkill
from a2ui.extension.a2ui_extension import get_a2ui_agent_extension
from agent import ContactAgent
from agent_executor import ContactAgentExecutor
from dotenv import load_dotenv
Expand All @@ -46,48 +44,22 @@ def main(host, port):
if not os.getenv("GOOGLE_GENAI_USE_VERTEXAI") == "TRUE":
if not os.getenv("GEMINI_API_KEY"):
raise MissingAPIKeyError(
"GEMINI_API_KEY environment variable not set and GOOGLE_GENAI_USE_VERTEXAI"
" is not TRUE."
"GEMINI_API_KEY environment variable not set and"
" GOOGLE_GENAI_USE_VERTEXAI is not TRUE."
)

capabilities = AgentCapabilities(
streaming=True,
extensions=[get_a2ui_agent_extension()],
)
skill = AgentSkill(
id="find_contact",
name="Find Contact Tool",
description=(
"Helps find contact information for colleagues (e.g., email, location,"
" team)."
),
tags=["contact", "directory", "people", "finder"],
examples=["Who is David Chen in marketing?", "Find Sarah Lee from engineering"],
)

base_url = f"http://{host}:{port}"

agent_card = AgentCard(
name="Contact Lookup Agent",
description=(
"This agent helps find contact info for people in your organization."
),
url=base_url, # <-- Use base_url here
version="1.0.0",
default_input_modes=ContactAgent.SUPPORTED_CONTENT_TYPES,
default_output_modes=ContactAgent.SUPPORTED_CONTENT_TYPES,
capabilities=capabilities,
skills=[skill],
)
agent = ContactAgent(base_url=base_url, use_ui=True)

agent_executor = ContactAgentExecutor(base_url=base_url)
agent_executor = ContactAgentExecutor(agent=agent)

request_handler = DefaultRequestHandler(
agent_executor=agent_executor,
task_store=InMemoryTaskStore(),
)
server = A2AStarletteApplication(
agent_card=agent_card, http_handler=request_handler
agent_card=agent.get_agent_card(), http_handler=request_handler
)
import uvicorn

Expand Down
72 changes: 0 additions & 72 deletions samples/agent/adk/contact_multiple_surfaces/a2ui_examples.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
from pathlib import Path

import jsonschema
from a2ui_schema import A2UI_SCHEMA

logger = logging.getLogger(__name__)

Expand All @@ -35,77 +34,6 @@
FLOOR_PLAN_FILE = "floor_plan.json"


def load_examples(base_url: str = "http://localhost:10004") -> str:
"""
Loads, validates, and formats the UI examples from JSON files.

Args:
base_url: The base URL to replace placeholder URLs with.
(Currently examples have http://localhost:10004 hardcoded,
but we can make this dynamic if needed).

Returns:
A string containing all formatted examples for the prompt.
"""

# Pre-parse validator
try:
single_msg_schema = json.loads(A2UI_SCHEMA)
# Examples are typically lists of messages
list_schema = {"type": "array", "items": single_msg_schema}
except json.JSONDecodeError:
logger.error("Failed to parse A2UI_SCHEMA for validation")
list_schema = None

examples_dir = Path(os.path.dirname(__file__)) / "examples"
formatted_output = []

for curr_name, filename in EXAMPLE_FILES.items():
file_path = examples_dir / filename
try:
content = file_path.read_text(encoding="utf-8")

# basic replacement if we decide to template the URL in JSON files
# content = content.replace("{{BASE_URL}}", base_url)

# Validation
if list_schema:
try:
data = json.loads(content)
jsonschema.validate(instance=data, schema=list_schema)
except (json.JSONDecodeError, jsonschema.ValidationError) as e:
logger.warning(f"Example {filename} validation failed: {e}")

formatted_output.append(f"---BEGIN {curr_name}---")
# Handle examples that include user/model text
if curr_name == "ORG_CHART_EXAMPLE":
formatted_output.append("User: Show me the org chart for Casey Smith")
formatted_output.append("Model: Here is the organizational chart.")
formatted_output.append("---a2ui_JSON---")
elif curr_name == "MULTI_SURFACE_EXAMPLE":
formatted_output.append("User: Full profile for Casey Smith")
formatted_output.append(
"Model: Here is the full profile including contact details and org chart."
)
formatted_output.append("---a2ui_JSON---")
elif curr_name == "CHART_NODE_CLICK_EXAMPLE":
formatted_output.append(
'User: ACTION: chart_node_click (context: clickedNodeName="John Smith")'
" (from modal)"
)
formatted_output.append("Model: Here is the profile for John Smith.")
formatted_output.append("---a2ui_JSON---")

formatted_output.append(content.strip())
formatted_output.append(f"---END {curr_name}---")
formatted_output.append("") # Newline

except FileNotFoundError:
logger.error(f"Example file not found: {file_path}")

return "\n".join(formatted_output)


def load_floor_plan_example() -> str:
"""Loads the floor plan example specifically."""
examples_dir = Path(os.path.dirname(__file__)) / "examples"
Expand Down
Loading
Loading