diff --git a/agent_sdks/python/a2ui_agent/agent_development.md b/agent_sdks/python/a2ui_agent/agent_development.md index bf9d10fc3..b5634714a 100644 --- a/agent_sdks/python/a2ui_agent/agent_development.md +++ b/agent_sdks/python/a2ui_agent/agent_development.md @@ -28,7 +28,8 @@ The first step in any A2UI-enabled agent is initializing the ```python from a2ui.schema.constants import VERSION_0_9 -from a2ui.schema.manager import A2uiSchemaManager, CatalogConfig +from a2ui.strategies.schema import A2uiSchemaManager +from a2ui.schema.catalog import CatalogConfig from a2ui.basic_catalog.provider import BasicCatalog # Define your catalogs (basic or bring your own) with optional examples diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/a2a/parts.py b/agent_sdks/python/a2ui_agent/src/a2ui/a2a/parts.py index e6598a89c..fa32a6958 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/a2a/parts.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/a2a/parts.py @@ -14,9 +14,10 @@ import logging from typing import Any, Optional, List, AsyncIterable, TYPE_CHECKING +from a2ui.parser.parser import Parser if TYPE_CHECKING: - from a2ui.parser.streaming import A2uiStreamParser + from a2ui.inference_formats.transport.streaming import TransportStreamParser from a2a.types import ( Part, DataPart, @@ -85,28 +86,72 @@ def get_a2ui_datapart(part: Part) -> Optional[DataPart]: return None -def parse_response_to_parts( +def parse_content_to_parts( content: str, - validator: Optional[Any] = None, + parser: Parser, fallback_text: Optional[str] = None, version: Optional[str] = None, ) -> List[Part]: - """Helper to parse LLM response content into A2A Parts, with optional validation. + """Helper to parse LLM response content into A2A Parts using a Parser instance. Args: content: The LLM response content, potentially containing A2UI delimiters. - validator: Optional validator to run against extracted JSON payloads. + parser: The Parser instance used to extract and compile format parts. fallback_text: Optional text to return if no parts are successfully created. version: Optional version string. Returns: A list of A2A Part objects (TextPart and/or DataPart). """ - from a2ui.parser.parser import parse_response + parts = [] + try: + response_parts = parser.parse_response(content) + + for part in response_parts: + if part.text: + parts.append(Part(root=TextPart(text=part.text))) + + if part.a2ui_json: + json_data = part.a2ui_json + if isinstance(json_data, list): + for message in json_data: + parts.append(create_a2ui_part(message, version=version)) + else: + parts.append(create_a2ui_part(json_data, version=version)) + + except Exception as e: + logger.warning(f"Failed to parse A2UI response: {e}") + + if not parts and fallback_text: + parts.append(Part(root=TextPart(text=fallback_text))) + + return parts + + +def parse_response_to_parts( + content: str, + validator: Optional[Any] = None, + fallback_text: Optional[str] = None, + version: Optional[str] = None, +) -> List[Part]: + """Deprecated compatibility wrapper around parse_response_to_parts. + + Please use parse_content_to_parts instead, providing a Parser instance. + """ + import warnings + + warnings.warn( + "parse_response_to_parts is deprecated. Please use parse_content_to_parts(...) " + "providing a Parser instance instead.", + DeprecationWarning, + stacklevel=2, + ) + + from a2ui.parser.parser import parse_response as legacy_parse_response parts = [] try: - response_parts = parse_response(content) + response_parts = legacy_parse_response(content) for part in response_parts: if part.text: @@ -114,7 +159,7 @@ def parse_response_to_parts( if part.a2ui_json: json_data = part.a2ui_json - if validator: + if validator is not None: validator.validate(json_data) if isinstance(json_data, list): @@ -124,7 +169,7 @@ def parse_response_to_parts( parts.append(create_a2ui_part(json_data, version=version)) except Exception as e: - logger.warning(f"Failed to parse or validate A2UI response: {e}") + logger.warning(f"Failed to parse legacy A2UI response: {e}") if not parts and fallback_text: parts.append(Part(root=TextPart(text=fallback_text))) @@ -133,14 +178,14 @@ def parse_response_to_parts( async def stream_response_to_parts( - parser: "A2uiStreamParser", + parser: "TransportStreamParser", token_stream: AsyncIterable[str], version: Optional[str] = None, ) -> AsyncIterable[Part]: """Helper to parse a stream of LLM tokens into A2A Parts incrementally. Args: - parser: A2uiStreamParser instance to process the stream. + parser: TransportStreamParser instance to process the stream. token_stream: An async iterable of strings (tokens). version: Optional version string. diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/adk/a2a/part_converter.py b/agent_sdks/python/a2ui_agent/src/a2ui/adk/a2a/part_converter.py index 7b88be20c..135f3f6be 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/adk/a2a/part_converter.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/adk/a2a/part_converter.py @@ -37,8 +37,9 @@ from a2a import types as a2a_types -from a2ui.a2a.parts import create_a2ui_part, parse_response_to_parts -from a2ui.parser.parser import has_a2ui_parts +from a2ui.a2a.parts import create_a2ui_part, parse_content_to_parts +from a2ui.parser.parser import Parser +from a2ui.inference_formats.transport.parser import TransportParser from a2ui.schema import constants from a2ui.schema.catalog import A2uiCatalog from google.adk.a2a.converters import part_converter @@ -69,11 +70,15 @@ def __init__( bypass_tool_check: bool = False, fallback_text: Optional[str] = None, version: str = constants.VERSION_0_8, + parser: Optional[Parser] = None, ): self._catalog = a2ui_catalog self._bypass_tool_check = bypass_tool_check self._fallback_text = fallback_text self._version = version + self._parser = parser or TransportParser( + a2ui_catalog, validator=a2ui_catalog.validator + ) def convert(self, part: genai_types.Part) -> list[a2a_types.Part]: """Converts a GenAI part to A2A parts, with A2UI validation. @@ -117,10 +122,10 @@ def convert(self, part: genai_types.Part) -> list[a2a_types.Part]: if function_response.response: result = function_response.response.get("result") - if isinstance(result, str) and has_a2ui_parts(result): - return parse_response_to_parts( + if isinstance(result, str) and self._parser.has_format_content(result): + return parse_content_to_parts( result, - validator=self._catalog.validator, + parser=self._parser, fallback_text=self._fallback_text, version=self._version, ) @@ -133,10 +138,10 @@ def convert(self, part: genai_types.Part) -> list[a2a_types.Part]: # 3. Handle Text-based A2UI (TextPart) if text := part.text: - if has_a2ui_parts(text): - return parse_response_to_parts( + if self._parser.has_format_content(text): + return parse_content_to_parts( text, - validator=self._catalog.validator, + parser=self._parser, fallback_text=self._fallback_text, version=self._version, ) diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/adk/orchestration/a2ui_subagent_map.py b/agent_sdks/python/a2ui_agent/src/a2ui/adk/orchestration/a2ui_subagent_map.py index df93a0eee..1f5a5b341 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/adk/orchestration/a2ui_subagent_map.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/adk/orchestration/a2ui_subagent_map.py @@ -88,11 +88,9 @@ A2UI_DELETE_SURFACE_KEY, A2UI_ACTIONS_KEY, A2UI_ERROR_KEY, - A2UI_CLIENT_DATA_MODEL_KEY, A2UI_CLIENT_DATA_MODEL_SURFACES_KEY, ) from a2ui.a2a.parts import is_a2ui_part -from a2a.server.events import Event as A2AEvent from a2a.types import Part, DataPart diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/adk/send_a2ui_to_client_toolset.py b/agent_sdks/python/a2ui_agent/src/a2ui/adk/send_a2ui_to_client_toolset.py index 58d529f4e..62dc933ef 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/adk/send_a2ui_to_client_toolset.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/adk/send_a2ui_to_client_toolset.py @@ -96,26 +96,18 @@ async def get_examples(ctx: ReadonlyContext) -> str: from google.adk.utils.feature_decorator import experimental from google.genai import types as genai_types -from a2ui.adk.a2a.event_converter import A2uiEventConverter from a2ui.adk.a2a.part_converter import A2uiPartConverter from a2ui.parser.payload_fixer import parse_and_fix from a2ui.schema import catalog from a2ui.core import A2uiValidationError from a2ui.schema import constants -from a2ui.schema.catalog import A2uiCatalog from a2ui.schema.constants import ( - A2UI_SCHEMA_BLOCK_END, - A2UI_SCHEMA_BLOCK_START, A2UI_TOOL_ERROR_KEY, A2UI_TOOL_NAME, A2UI_VALIDATED_JSON_KEY, ) -if TYPE_CHECKING: - from google.adk.agents.invocation_context import InvocationContext - from google.adk.events.event import Event - logger = logging.getLogger(__name__) A2uiEnabledProvider: TypeAlias = Callable[ diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/compiler.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/compiler.py index 0d566b6a6..17216b64f 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/compiler.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/compiler.py @@ -513,7 +513,7 @@ def _compile_value(self, val: Any, is_action: bool = False) -> Any: v, is_action ) - fn_schema = self.helper.functions[fn_name] + self.helper.functions[fn_name] if is_action: return { diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/decompiler.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/decompiler.py index ddf3b648d..c00dd810b 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/decompiler.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/decompiler.py @@ -20,7 +20,7 @@ import html import json import re -from typing import Any, Dict, List, Optional, Union +from typing import Any, Optional, Union from a2ui.core.catalog import Catalog from a2ui.schema.catalog import A2uiCatalog from a2ui.experimental.express.schema_helper import CatalogSchemaHelper @@ -254,7 +254,6 @@ def _render_component( if k not in ["id", "component"] and k not in all_props: all_props.append(k) - has_text_content = False text_content = "" for prop_name in all_props: diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/expression_parser.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/expression_parser.py index 0a3c00165..498cc287d 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/expression_parser.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/expression_parser.py @@ -18,7 +18,7 @@ including array literals, object literals, and path bindings prefixed with '$'. """ -from typing import Any, Dict, List, Union +from typing import Any, Dict, List from a2ui.core.basic_catalog.expression_parser import ExpressionParser, Scanner diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/parser.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/parser.py index afab83f5e..6ea8a1d25 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/parser.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/parser.py @@ -15,7 +15,7 @@ """Parser utilities to extract and compile A2UI Elemental HTML from LLM responses.""" import re -from typing import Any, Dict, List, Optional, Union +from typing import Any, List, Union from a2ui.core.catalog import Catalog from a2ui.schema.catalog import A2uiCatalog from a2ui.parser.response_part import ResponsePart @@ -94,7 +94,7 @@ def parse_elemental_response( a2ui_json=[compiled_json], ) ) - except Exception as e: + except Exception: # Graceful fallback: treat malformed/unparseable blocks as plain text fallback_text = html_content full_text = f"{text_part}\n{fallback_text}" if text_part else fallback_text diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/prompt_generator.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/prompt_generator.py index 11339beaf..356eea177 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/prompt_generator.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/elemental/prompt_generator.py @@ -20,7 +20,7 @@ import json import re -from typing import Any, Dict, List, Optional, Union +from typing import Any, Union from a2ui.core.catalog import Catalog from a2ui.schema.catalog import A2uiCatalog from a2ui.experimental.express.schema_helper import CatalogSchemaHelper diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/__init__.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/__init__.py index dcbe48147..ce8d1848b 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/__init__.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/__init__.py @@ -18,14 +18,6 @@ into standard A2UI v1.0 wire JSON messages and vice-versa. """ -import os - -if os.environ.get("A2UI_EXPRESS_ENABLED", "").lower() not in ("true", "1", "yes"): - raise ImportError( - "A2UI Express is an experimental proposal extension and is disabled by default." - " To enable it, set the environment variable A2UI_EXPRESS_ENABLED=true." - ) - from .compiler import ExpressCompiler from .decompiler import ExpressDecompiler from .prompt_generator import ExpressPromptGenerator diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/compiler.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/compiler.py index 33187cda3..91d16ce8d 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/compiler.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/compiler.py @@ -20,8 +20,7 @@ The grammar for A2UI Express is defined in Express.g4. """ -import re -from typing import Any, Dict, Optional, Union +from typing import Any, Optional, Union from antlr4 import InputStream, CommonTokenStream from a2ui.core.catalog import Catalog from a2ui.schema.catalog import A2uiCatalog diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/decompiler.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/decompiler.py index 0321a5f13..9c278bc17 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/decompiler.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/decompiler.py @@ -18,7 +18,7 @@ tailored for prompt tokens compression. """ -from typing import Any, Dict, Optional, Union +from typing import Any, Union from a2ui.core.catalog import Catalog from a2ui.schema.catalog import A2uiCatalog from .schema_helper import CatalogSchemaHelper diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/generated/express_lexer.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/generated/express_lexer.py index f35fe0070..c087819e8 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/generated/express_lexer.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/generated/express_lexer.py @@ -1,6 +1,5 @@ # Generated from Express.g4 by ANTLR 4.13.2 from antlr4 import * -from io import StringIO import sys if sys.version_info[1] > 5: from typing import TextIO diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/generated/express_parser.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/generated/express_parser.py index f50554267..d46c44543 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/generated/express_parser.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/generated/express_parser.py @@ -1,7 +1,6 @@ # Generated from Express.g4 by ANTLR 4.13.2 # encoding: utf-8 from antlr4 import * -from io import StringIO import sys if sys.version_info[1] > 5: from typing import TextIO diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/parser.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/parser.py index 1a58fb29a..d092ba0e3 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/parser.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/parser.py @@ -15,7 +15,7 @@ """Parser utilities to extract and compile A2UI Express DSL from LLM responses.""" import re -from typing import Any, Dict, List, Optional, Union +from typing import Any, List, Union from a2ui.core.catalog import Catalog from a2ui.schema.catalog import A2uiCatalog from a2ui.parser.response_part import ResponsePart diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/prompt_generator.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/prompt_generator.py index 4d0300b31..89f326318 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/prompt_generator.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/prompt_generator.py @@ -20,7 +20,7 @@ import json import re -from typing import Any, Dict, Optional, Union +from typing import Any, Optional, Union from a2ui.core.catalog import Catalog from a2ui.schema.catalog import A2uiCatalog from .decompiler import ExpressDecompiler diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/schema_helper.py b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/schema_helper.py index 518b70609..48305046e 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/schema_helper.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/experimental/express/schema_helper.py @@ -18,8 +18,7 @@ signatures, and requirements directly from standard catalog JSON schemas. """ -import json -from typing import Any, Dict, Optional, Union +from typing import Any, Optional, Union from a2ui.core.catalog import Catalog from a2ui.schema.catalog import A2uiCatalog diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/inference_format.py b/agent_sdks/python/a2ui_agent/src/a2ui/inference_format.py new file mode 100644 index 000000000..c2ec4de7b --- /dev/null +++ b/agent_sdks/python/a2ui_agent/src/a2ui/inference_format.py @@ -0,0 +1,83 @@ +# 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. + +"""Unified interface coordinating prompt generation and parsing of LLM response payloads.""" + +import warnings +from abc import ABC, abstractmethod +from typing import Any, Optional +from a2ui.prompt.generator import PromptGenerator +from a2ui.parser.parser import Parser + + +class InferenceFormat(ABC): + """Interface coordinating system prompt generation and response parsing.""" + + @property + @abstractmethod + def prompt_generator(self) -> PromptGenerator: + """The PromptGenerator instance associated with this inference format.""" + pass + + @property + @abstractmethod + def parser(self) -> Parser: + """The Parser instance associated with this inference format.""" + pass + + def generate_system_prompt( + self, + role_description: str, + workflow_description: str = "", + ui_description: str = "", + client_ui_capabilities: Optional[dict[str, Any]] = None, + allowed_components: Optional[list[str]] = None, + allowed_messages: Optional[list[str]] = None, + include_schema: bool = False, + include_examples: bool = False, + validate_examples: bool = False, + ) -> str: + """Generates a system prompt for all LLM requests (deprecated compatibility helper). + + Args: + role_description: Description of the agent's role. + workflow_description: Optional description of the task workflow. + ui_description: Optional UI context or rules. + client_ui_capabilities: Optional client UI capability details. + allowed_components: Optional list of component tags the LLM may use. + allowed_messages: Optional list of message types allowed. + include_schema: Whether to include component schemas in the prompt. + include_examples: Whether to include few-shot examples. + validate_examples: Whether to validate few-shot examples on generation. + + Returns: + The complete system prompt string. + """ + warnings.warn( + "generate_system_prompt is deprecated. Use prompt_generator.generate(...)" + " instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.prompt_generator.generate( + role_description=role_description, + workflow_description=workflow_description, + ui_description=ui_description, + client_ui_capabilities=client_ui_capabilities, + allowed_components=allowed_components, + allowed_messages=allowed_messages, + include_schema=include_schema, + include_examples=include_examples, + validate_examples=validate_examples, + ) diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/__init__.py b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/__init__.py new file mode 100644 index 000000000..b3b408135 --- /dev/null +++ b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/__init__.py @@ -0,0 +1,24 @@ +# 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 .format import TransportFormat +from .parser import TransportParser +from .streaming import TransportStreamParser, A2uiStreamParser + +__all__ = [ + "TransportFormat", + "TransportParser", + "TransportStreamParser", + "A2uiStreamParser", # Deprecated. Use TransportStreamParser instead. +] diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/format.py b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/format.py new file mode 100644 index 000000000..bb547e278 --- /dev/null +++ b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/format.py @@ -0,0 +1,274 @@ +# 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. + +"""Standard A2UI transport inference format coordination.""" + +import copy +from typing import Any, Optional, Callable, Union + +from a2ui.schema.utils import load_from_bundled_resource +from a2ui.inference_format import InferenceFormat +from a2ui.schema.constants import ( + SERVER_TO_CLIENT_SCHEMA_KEY, + COMMON_TYPES_SCHEMA_KEY, + SPEC_VERSION_MAP, + INLINE_CATALOGS_KEY, + SUPPORTED_CATALOG_IDS_KEY, + CATALOG_COMPONENTS_KEY, + INLINE_CATALOG_NAME, +) +from a2ui.schema.catalog import CatalogConfig, A2uiCatalog +from a2ui.core import A2uiCatalogError +from a2ui.inference_formats.transport.parser import TransportParser +from a2ui.inference_formats.transport.prompt_generator import TransportPromptGenerator + + +class TransportFormat(InferenceFormat): + """Manages standard A2UI JSON schema responses and prompt injection (Transport Format).""" + + def __init__( + self, + version: str, + catalogs: Optional[list[CatalogConfig]] = None, + accepts_inline_catalogs: bool = False, + schema_modifiers: Optional[ + list[Callable[[dict[str, Any]], dict[str, Any]]] + ] = None, + experiments: Optional[Union[set[str], frozenset[str]]] = None, + ): + """Initializes the TransportFormat with schemas and catalogs. + + Args: + version: The A2UI protocol specification version (e.g. "0.9"). + catalogs: Optional list of catalog configurations. + accepts_inline_catalogs: Whether inline catalog definitions are allowed. + schema_modifiers: Optional schema modifier functions to post-process schemas. + experiments: Optional set of enabled experimental feature flags. + """ + self._version = version + self._accepts_inline_catalogs = accepts_inline_catalogs + self.experiments = frozenset(experiments) if experiments else frozenset() + + self._server_to_client_schema: dict[str, Any] = {} + self._common_types_schema: dict[str, Any] = {} + self._supported_catalogs: list[A2uiCatalog] = [] + self._catalog_example_paths: dict[str, str] = {} + self._schema_modifiers = schema_modifiers or [] + self._parser: Optional[TransportParser] = None + self._prompt_generator: Optional[TransportPromptGenerator] = None + self._load_schemas(version, catalogs or []) + + @property + def prompt_generator(self) -> TransportPromptGenerator: + """Returns the PromptGenerator instance for this format.""" + if self._prompt_generator is None: + self._prompt_generator = TransportPromptGenerator(self) + return self._prompt_generator + + @property + def parser(self) -> TransportParser: + """Returns the Parser instance for this format.""" + if self._parser is None: + if not self._supported_catalogs: + raise ValueError( + "No supported catalogs configured for the transport format." + ) + default_catalog = self._supported_catalogs[0] + self._parser = TransportParser( + default_catalog, + default_catalog.validator, + ) + return self._parser + + @property + def accepts_inline_catalogs(self) -> bool: + """Whether this format accepts inline catalog definitions.""" + return self._accepts_inline_catalogs + + @property + def supported_catalog_ids(self) -> list[str]: + """A list of catalog IDs supported by this format.""" + return [c.catalog_id for c in 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, + catalogs: Optional[list[CatalogConfig]] = None, + ) -> None: + """Loads separate schema components and processes catalogs.""" + catalogs = catalogs or [] + if version not in SPEC_VERSION_MAP: + raise A2uiCatalogError( + f"Unknown A2UI specification version: {version}. Supported:" + f" {list(SPEC_VERSION_MAP.keys())}" + ) + + # Load server-to-client and common types schemas + self._server_to_client_schema = self._apply_modifiers( + load_from_bundled_resource( + version, SERVER_TO_CLIENT_SCHEMA_KEY, SPEC_VERSION_MAP + ) + ) + self._common_types_schema = self._apply_modifiers( + load_from_bundled_resource( + version, COMMON_TYPES_SCHEMA_KEY, SPEC_VERSION_MAP + ) + ) + + # Process catalogs + for config in catalogs: + catalog_schema = config.provider.load() + 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, + custom_cuttable_keys=config.custom_cuttable_keys, + experiments=self.experiments, + ) + self._supported_catalogs.append(catalog) + if config.examples_path: + self._catalog_example_paths[catalog.catalog_id] = config.examples_path + + def _select_catalog( + self, client_ui_capabilities: Optional[dict[str, Any]] = None + ) -> A2uiCatalog: + """Selects the component catalog for the prompt based on client capabilities. + + Selection priority: + 1. If inline catalogs are provided (and accepted by the agent), their + components are merged on top of a base catalog. The base is determined + by supportedCatalogIds (if also provided) or the agent's default catalog. + 2. If only supportedCatalogIds is provided, pick the first mutually + supported catalog. + 3. Fallback to the first agent-supported catalog (usually the bundled catalog). + + Args: + client_ui_capabilities: A dictionary of client UI capabilities, containing + inline catalogs and client-supported catalog IDs. + + Returns: + The resolved A2uiCatalog. + Raises: + ValueError: If inline catalogs are sent but not accepted, or if no + mutually supported catalog is found. + """ + if not self._supported_catalogs: + raise A2uiCatalogError( + "No supported catalogs found." + ) # This should not happen. + + if not client_ui_capabilities or not isinstance(client_ui_capabilities, dict): + return self._supported_catalogs[0] + + inline_catalogs: list[dict[str, Any]] = client_ui_capabilities.get( + INLINE_CATALOGS_KEY, [] + ) + client_supported_catalog_ids: list[str] = client_ui_capabilities.get( + SUPPORTED_CATALOG_IDS_KEY, [] + ) + + if not self._accepts_inline_catalogs and inline_catalogs: + raise A2uiCatalogError( + f"Inline catalog '{INLINE_CATALOGS_KEY}' is provided in client UI" + " capabilities. However, the agent does not accept inline catalogs." + ) + + if inline_catalogs: + # Determine the base catalog: use supportedCatalogIds if provided, + # otherwise fall back to the agent's default catalog. + base_catalog = self._supported_catalogs[0] + if client_supported_catalog_ids: + 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: + base_catalog = agent_supported_catalogs[cscid] + break + + merged_schema = copy.deepcopy(base_catalog.catalog_schema) + + for inline_catalog_schema in inline_catalogs: + inline_catalog_schema = self._apply_modifiers(inline_catalog_schema) + inline_components = inline_catalog_schema.get( + CATALOG_COMPONENTS_KEY, {} + ) + merged_schema[CATALOG_COMPONENTS_KEY].update(inline_components) + + return A2uiCatalog( + version=self._version, + name=INLINE_CATALOG_NAME, + catalog_schema=merged_schema, + s2c_schema=self._server_to_client_schema, + common_types_schema=self._common_types_schema, + experiments=self.experiments, + ) + + if not client_supported_catalog_ids: + return self._supported_catalogs[0] + + 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 A2uiCatalogError( + "No client-supported catalog found on the agent side. Agent-supported" + f" catalogs are: {[c.catalog_id for c in self._supported_catalogs]}" + ) + + def get_selected_catalog( + self, + client_ui_capabilities: Optional[dict[str, Any]] = None, + allowed_components: Optional[list[str]] = None, + allowed_messages: Optional[list[str]] = None, + ) -> A2uiCatalog: + """Selects and prunes the catalog according to client capabilities and restrictions. + + Args: + client_ui_capabilities: Optional client UI capability details. + allowed_components: Optional list of component tags allowed. + allowed_messages: Optional list of message types allowed. + + Returns: + The selected and pruned A2uiCatalog instance. + """ + catalog = self._select_catalog(client_ui_capabilities) + pruned_catalog = catalog.with_pruning(allowed_components, allowed_messages) + return pruned_catalog + + def load_examples(self, catalog: A2uiCatalog, validate: bool = False) -> str: + """Loads and optionally validates few-shot examples for the specified catalog. + + Args: + catalog: The A2uiCatalog to load examples for. + validate: Whether to validate the examples on load. + + Returns: + The examples text block, or an empty string. + """ + if catalog.catalog_id in self._catalog_example_paths: + return catalog.load_examples( + self._catalog_example_paths[catalog.catalog_id], validate=validate + ) + return "" diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/parser.py b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/parser.py new file mode 100644 index 000000000..feb618a21 --- /dev/null +++ b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/parser.py @@ -0,0 +1,155 @@ +# 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. + +"""Parser and compiler implementation for standard A2UI JSON schema responses.""" + +import re +from typing import List, Optional, Any +from a2ui.parser.parser import Parser +from a2ui.parser.response_part import ResponsePart +from a2ui.schema.catalog import A2uiCatalog +from a2ui.validation.validator import A2uiValidator +from a2ui.schema.constants import A2UI_OPEN_TAG, A2UI_CLOSE_TAG +from a2ui.core import A2uiParseError +from a2ui.parser.payload_fixer import parse_and_fix + +_A2UI_BLOCK_PATTERN = re.compile( + f"{re.escape(A2UI_OPEN_TAG)}(.*?){re.escape(A2UI_CLOSE_TAG)}", re.DOTALL +) + + +def _sanitize_json_string(json_string: str) -> str: + """Sanitizes the JSON string by removing markdown code blocks. + + Args: + json_string: The raw JSON string. + + Returns: + The sanitized JSON string. + """ + json_string = json_string.strip() + if json_string.startswith("```json"): + json_string = json_string[len("```json") :] + elif json_string.startswith("```"): + json_string = json_string[len("```") :] + if json_string.endswith("```"): + json_string = json_string[: -len("```")] + json_string = json_string.strip() + return json_string + + +def unwrap_response(content: str) -> List[ResponsePart]: + """Tokenizes the LLM response into a list of ResponsePart objects, extracting raw format content. + + Args: + content: The raw LLM response. + + Returns: + A list of ResponsePart objects. + """ + matches = list(_A2UI_BLOCK_PATTERN.finditer(content)) + + if not matches: + raise A2uiParseError( + f"A2UI tags '{A2UI_OPEN_TAG}' and '{A2UI_CLOSE_TAG}' not found in response." + ) + + response_parts = [] + last_end = 0 + + for match in matches: + start, end = match.span() + text_part = content[last_end:start].strip() + + json_string = match.group(1) + json_string_cleaned = _sanitize_json_string(json_string) + if not json_string_cleaned: + raise A2uiParseError("A2UI JSON part is empty.") + + response_parts.append( + ResponsePart(text=text_part, a2ui_raw=json_string_cleaned) + ) + last_end = end + + trailing_text = content[last_end:].strip() + if trailing_text: + response_parts.append(ResponsePart(text=trailing_text, a2ui_raw=None)) + + return response_parts + + +class TransportParser(Parser): + """Concrete parser implementation for standard A2UI JSON schema responses (Transport Format).""" + + def __init__( + self, + catalog: A2uiCatalog, + validator: Optional[A2uiValidator] = None, + ): + """Initializes the TransportParser. + + Args: + catalog: The A2uiCatalog mapping schema identifiers. + validator: Optional validator for payload verification. + """ + self._catalog = catalog + self._validator = validator + self._stream_parser: Optional[Any] = None + + def has_format_content(self, content: str, *, complete: bool = False) -> bool: + from a2ui.schema.constants import A2UI_OPEN_TAG, A2UI_CLOSE_TAG + + if complete: + return A2UI_OPEN_TAG in content and A2UI_CLOSE_TAG in content + return A2UI_OPEN_TAG in content + + def unwrap(self, content: str) -> List[ResponsePart]: + """Tokenizes response content into raw format-content parts. + + Args: + content: The raw response content. + + Returns: + A list of unwrapped ResponsePart objects. + """ + return unwrap_response(content) + + def compile(self, format_content: str) -> List[dict[str, Any]]: + """Validates and compiles raw A2UI JSON schema content. + + Args: + format_content: The raw A2UI JSON string. + + Returns: + A list of compiled A2UI message dictionaries. + """ + json_data = parse_and_fix(format_content) + if self._validator: + self._validator.validate(json_data) + return json_data + + def process_chunk(self, chunk: str) -> List[ResponsePart]: + """Processes streamed token chunks incrementally. + + Args: + chunk: The next token text chunk. + + Returns: + A list of parsed or completed ResponsePart objects. + """ + from a2ui.inference_formats.transport.streaming import TransportStreamParser + + if not self._stream_parser: + self._stream_parser = TransportStreamParser(self._catalog) + return self._stream_parser.process_chunk(chunk) diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/prompt_generator.py b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/prompt_generator.py new file mode 100644 index 000000000..ef2fb0755 --- /dev/null +++ b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/prompt_generator.py @@ -0,0 +1,93 @@ +# 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. + +"""Generator for standard A2UI JSON schema system prompt instructions.""" + +from typing import Optional, Any, TYPE_CHECKING +from a2ui.prompt.generator import PromptGenerator + +if TYPE_CHECKING: + from a2ui.inference_formats.transport.format import TransportFormat + + +class TransportPromptGenerator(PromptGenerator): + """Formats standard JSON schema system prompt instructions (Transport Format).""" + + def __init__(self, format_inst: "TransportFormat"): + """Initializes the prompt generator with a TransportFormat context. + + Args: + format_inst: The TransportFormat instance. + """ + self._format = format_inst + + def generate( + self, + role_description: str, + workflow_description: str = "", + ui_description: str = "", + client_ui_capabilities: Optional[dict[str, Any]] = None, + allowed_components: Optional[list[str]] = None, + allowed_messages: Optional[list[str]] = None, + include_schema: bool = False, + include_examples: bool = False, + validate_examples: bool = False, + ) -> str: + """Assembles prompt instructions contract for standard JSON. + + Args: + role_description: Description of the agent's role. + workflow_description: Optional description of the task workflow. + ui_description: Optional UI context or rules. + client_ui_capabilities: Optional client UI capability details. + allowed_components: Optional list of component tags the LLM may use. + allowed_messages: Optional list of A2UI message types allowed. + include_schema: Whether to include component schemas in the prompt. + include_examples: Whether to include few-shot examples. + validate_examples: Whether to validate few-shot examples on generation. + + Returns: + The complete generated prompt system instruction. + """ + selected_catalog = self._format.get_selected_catalog( + client_ui_capabilities, allowed_components, allowed_messages + ) + + examples_str = "" + if include_examples: + examples_str = self._format.load_examples( + selected_catalog, validate=validate_examples + ) + + parts = [role_description] + + from a2ui.schema.constants import DEFAULT_WORKFLOW_RULES + + rules = DEFAULT_WORKFLOW_RULES + if workflow_description: + rules += f"\n{workflow_description}" + parts.append(f"## Workflow Description:\n{rules}") + + if ui_description: + parts.append(f"## UI Description:\n{ui_description}") + + if include_schema: + instructions = selected_catalog.render_as_llm_instructions() + if instructions: + parts.append(instructions) + + if examples_str: + parts.append(f"### Examples:\n{examples_str}") + + return "\n\n".join(parts) diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/streaming.py b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/streaming.py new file mode 100644 index 000000000..171452c1a --- /dev/null +++ b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/streaming.py @@ -0,0 +1,1148 @@ +# 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 __future__ import annotations + +import copy +import json +import logging +import re +from typing import Any, List, Dict, Optional, Set, TYPE_CHECKING, Union + +from a2ui.parser.constants import * +from a2ui.schema.constants import ( + VERSION_0_9, + VERSION_0_8, + A2UI_OPEN_TAG, + A2UI_CLOSE_TAG, + SURFACE_ID_KEY, + CATALOG_COMPONENTS_KEY, +) +from a2ui.validation.validator import ( + extract_component_ref_fields, + extract_component_required_fields, +) +from a2ui.validation.validator import A2uiValidator +from a2ui.core.validating import analyze_topology +from a2ui.parser.response_part import ResponsePart +from a2ui.core.validating.validator import RELAXED_VALIDATION, STRICT_VALIDATION, ValidationConfig +from a2ui.core import A2uiParseError, A2uiIntegrityError + + +if TYPE_CHECKING: + from a2ui.schema.catalog import A2uiCatalog + +logger = logging.getLogger(__name__) + + +class TransportStreamParser: + """Parses a stream of text for A2UI JSON messages with fine-grained component yielding. + + This class acts as a factory that returns a version-specific parser instance + (V08 or V09) depending on the catalog version. + """ + + def __new__(cls, catalog: A2uiCatalog) -> TransportStreamParser: + if cls is TransportStreamParser: + version = catalog.version + # Lazy import inside __new__ to prevent circular import errors, as the + # version-specific subclass modules import TransportStreamParser from this module. + if version == VERSION_0_8: + from .streaming_v08 import TransportStreamParserV08 + + return TransportStreamParserV08(catalog=catalog) + else: + from .streaming_v09 import TransportStreamParserV09 + + return TransportStreamParserV09(catalog=catalog) + return super().__new__(cls) + + 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 = "" + self._json_buffer = "" + self._brace_stack: list[tuple[str, int]] = [] + self._brace_count = 0 + self._in_top_level_list = False + self._in_string = False + self._string_escaped = False + + self._seen_components: Dict[str, Dict[str, Any]] = {} + + # Track data model for path resolution + self._yielded_data_model: Dict[str, Any] = {} + self._deleted_surfaces: Set[str] = set() + + # Set of unique component IDs yielded per surface to prevent duplicate yielding + # surfaceId -> set of cids + self._yielded_ids: Dict[str, Set[str]] = {} + # (surfaceId, cid) -> hash of content for change detection + self._yielded_contents: Dict[Any, str] = {} + + self._root_ids: Dict[str, str] = {} # The root component IDs mapped per surface + self._default_root_id: Optional[str] = ( + None # Base default root ID for the protocol + ) + self._unbound_root_id: Optional[str] = ( + None # Temporary holding variable for when root arrives before surfaceId + ) + self._surface_id: Optional[str] = ( + None # The active surface ID tracking the context + ) + self._msg_types: List[str] = ( + [] + ) # Running list of message types seen in the block + + # A set of surface ids for which we have already yielded a start message + # Tracks if beginRendering or createSurface was emitted + self._yielded_start_messages: Set[str] = set() + + # The current active message type for component grouping + self._active_msg_type: Optional[str] = None + + # State for buffering updates until surface is ready + self._pending_messages: Dict[str, List[Dict[str, Any]]] = ( + {} + ) # surfaceId -> list of msgs delayed until start message arrives + self._buffered_start_message: Optional[Dict[str, Any]] = ( + None # The start message to yield before any components + ) + self._topology_dirty = False # Set to true if components are added out of order + self._found_valid_json_in_block = False + + @property + def _placeholder_component(self) -> Dict[str, Any]: + """Returns the version-specific placeholder component. + + This is used when a component references a child component that hasn't yet + streamed in. The placeholder component is added to the components list and + the reference is updated to point to the placeholder component. + """ + raise NotImplementedError("Subclasses must implement _placeholder_component") + + @property + def surface_id(self) -> Optional[str]: + return self._surface_id + + @surface_id.setter + def surface_id(self, value: Optional[str]) -> None: + self._surface_id = value + if value is not None and self._unbound_root_id is not None: + self._root_ids[value] = self._unbound_root_id + self._unbound_root_id = None + + @property + def root_id(self) -> Optional[str]: + if self._surface_id: + return self._root_ids.get(self._surface_id, self._default_root_id) + # Return unbound root ID if explicitly sniffed, otherwise use protocol default + return ( + self._unbound_root_id + if self._unbound_root_id is not None + else self._default_root_id + ) + + @root_id.setter + def root_id(self, value: Optional[str]) -> None: + if self._surface_id: + if value is not None: + self._root_ids[self._surface_id] = value + else: + self._root_ids.pop(self._surface_id, None) + else: + self._unbound_root_id = value + + @property + def msg_types(self) -> List[str]: + return self._msg_types + + def add_msg_type(self, msg_type: str) -> None: + if msg_type not in self._msg_types: + self._msg_types.append(msg_type) + if msg_type in ( + MSG_TYPE_SURFACE_UPDATE, + MSG_TYPE_UPDATE_COMPONENTS, + MSG_TYPE_CREATE_SURFACE, + ): + self._active_msg_type = msg_type + + @property + def _yielded_surfaces_set(self) -> Set[str]: + """Provides access to version-specific yielded surfaces set.""" + raise NotImplementedError("Subclasses must implement _yielded_surfaces_set") + + def is_protocol_msg(self, obj: Dict[str, Any]) -> bool: + """Checks if the object is a recognized A2UI message for this version.""" + raise NotImplementedError("Subclasses must implement is_protocol_msg") + + @property + def _data_model_msg_type(self) -> str: + """Returns the message type identifier for data model updates.""" + raise NotImplementedError("Subclasses must implement _data_model_msg_type") + + def _get_active_msg_type_for_components(self) -> Optional[str]: + """Determines which msg_type to use when wrapping component updates.""" + raise NotImplementedError( + "Subclasses must implement _get_active_msg_type_for_components" + ) + + def _construct_partial_message( + self, components: List[Dict[str, Any]], active_msg_type: str + ) -> Dict[str, Any]: + """Constructs a partial message of the correct type. Subclasses must implement.""" + raise NotImplementedError( + "Subclasses must implement _construct_partial_message" + ) + + def _deduplicate_data_model(self, m: Dict[str, Any]) -> bool: + """Returns True if message should be yielded, False if skipped.""" + return True + + def _yield_messages( + self, + messages_to_yield: List[Dict[str, Any]], + messages: List[ResponsePart], + config: ValidationConfig = STRICT_VALIDATION, + ) -> None: + """Validates and appends messages to the final output list.""" + for m in messages_to_yield: + if not self._deduplicate_data_model(m): + continue + + # Each surface update message must specify a surfaceId and satisfy catalog validation. + if self._validator: + try: + self._validator.validate(m, root_id=self.root_id, config=config) + except ValueError as e: + if config == STRICT_VALIDATION: + raise e + else: + logger.debug( + f"Validation failed for partial/sniffed message: {e}" + ) + continue + + # Consolidated appending logic + if messages and messages[-1].a2ui_json is None: + messages[-1].a2ui_json = [m] + elif messages and isinstance(messages[-1].a2ui_json, list): + messages[-1].a2ui_json.append(m) + else: + messages.append(ResponsePart(a2ui_json=[m])) + + def _delete_surface(self, sid: str) -> None: + """Clears all state related to a specific surface.""" + self._pending_messages.pop(sid, None) + self._yielded_ids.pop(sid, None) + + # Clear contents for this surface + self._yielded_contents = { + k: v for k, v in self._yielded_contents.items() if k[0] != sid + } + self._yielded_surfaces_set.discard(sid) + self._yielded_start_messages.discard(sid) + + self._deleted_surfaces.add(sid) + + def process_chunk(self, chunk: str) -> List[ResponsePart]: + """Processes a chunk of text and returns any complete A2UI messages found. + + This is the primary entry point for the streaming parser. It handles the + initial "tag hunt" and then delegates JSON fragment processing to + `_process_json_chunk`. It supports multiple A2UI blocks in a single stream. + + Args: + chunk: The chunk of raw text (e.g., from an LLM stream) to process. + + Returns: + A list of parsed A2UI message dictionaries. + """ + messages = [] + self._buffer += chunk + + while True: + if not self._found_delimiter: + # Looking for + if A2UI_OPEN_TAG in self._buffer: + parts = self._buffer.split(A2UI_OPEN_TAG, 1) + if parts[0]: + messages.append(ResponsePart(text=parts[0])) + self._found_delimiter = True + self._buffer = parts[1] + # Continue to process the content after the open tag + else: + # Yield conversational text while avoiding split tags + keep_len = 0 + for i in range(len(A2UI_OPEN_TAG) - 1, 0, -1): + if self._buffer.endswith(A2UI_OPEN_TAG[:i]): + keep_len = i + break + + if len(self._buffer) > keep_len: + safe_to_yield = len(self._buffer) - keep_len + text_to_yield = self._buffer[:safe_to_yield] + messages.append(ResponsePart(text=text_to_yield)) + self._buffer = self._buffer[safe_to_yield:] + break + + if self._found_delimiter: + # Looking for + if A2UI_CLOSE_TAG in self._buffer: + parts = self._buffer.split(A2UI_CLOSE_TAG, 1) + json_fragment = parts[0] + self._process_json_chunk(json_fragment, messages) + if not self._found_valid_json_in_block: + raise A2uiParseError( + "Failed to parse JSON: No valid JSON object found in A2UI" + " block." + ) + + # End of block: reset JSON state but keep seen_components + self._found_delimiter = False + self._reset_json_state() + + self._buffer = parts[1] + # Continue loop to look for next A2UI_OPEN_TAG in remaining buffer + else: + # Find if the buffer ends with a prefix of A2UI_CLOSE_TAG + # To avoid split-tag issues, we only delay processing if it looks like a close tag is starting + keep_len = 0 + for i in range(1, len(A2UI_CLOSE_TAG)): + if self._buffer.endswith(A2UI_CLOSE_TAG[:i]): + keep_len = i + + if keep_len < len(self._buffer): + to_process = self._buffer[: len(self._buffer) - keep_len] + self._buffer = self._buffer[len(self._buffer) - keep_len :] + self._process_json_chunk(to_process, messages) + break + + # Deduplicate surfaceUpdate messages to avoid over-yielding in a single chunk + for part in messages: + if not part.a2ui_json: + continue + + deduped_msgs = [] + seen_su = set() + # Iterate backwards to keep only the last (most complete) surfaceUpdate for each surface + for m in reversed(part.a2ui_json): + is_su = False + sid = None + if isinstance(m, dict) and MSG_TYPE_SURFACE_UPDATE in m: + is_su = True + sid = m[MSG_TYPE_SURFACE_UPDATE].get(SURFACE_ID_KEY) + + if is_su and sid: + if sid not in seen_su: + deduped_msgs.append(m) + seen_su.add(sid) + else: + deduped_msgs.append(m) + + deduped_msgs.reverse() + part.a2ui_json = deduped_msgs + + if messages: + logger.debug( + f"DEBUG: process_chunk returning {len(messages)} messages: {messages}" + ) + return messages + + def _reset_json_state(self) -> None: + """Resets the JSON-specific parsing state (e.g., at the end of a block).""" + self._json_buffer = "" + self._brace_stack = [] + self._brace_count = 0 + self._in_top_level_list = False + self._in_string = False + self._string_escaped = False + self._msg_types = [] + self._found_valid_json_in_block = False + # Note: we do NOT reset _active_msg_type or _yielded_contents here + + # so re-yielding works between blocks + + def _fix_json(self, fragment: str) -> str: + """Attempts to fix a partial JSON fragment by adding missing closing delimiters.""" + fixed = fragment.rstrip() + if not fixed: + return "" + + stack = [] + in_string = False + escaped = False + last_quote_idx = -1 + + # Single pass to track strings and braces + for i, char in enumerate(fixed): + if escaped: + escaped = False + continue + if char == "\\": + escaped = True + continue + if char == '"': + in_string = not in_string + if in_string: + last_quote_idx = i + elif not in_string: + if char in ("{", "["): + stack.append(char) + elif char in ("}", "]"): + if stack: + stack.pop() + + # 1. Close open strings (healing) + if in_string: + # We only auto-close strings for safe keys (CUTTABLE_KEYS) + prefix = fixed[:last_quote_idx].rstrip() + if prefix.endswith(":"): + key_match = re.findall(r'"([^"]+)"\s*:\s*$', prefix) + if key_match: + key = key_match[0] + if key not in self._cuttable_keys: + return "" + + # Special case: don't cut URL bindings, as partial URLs break images/links + if key == "valueString": + string_val = fixed[last_quote_idx + 1 :] + if ( + string_val.startswith("http://") + or string_val.startswith("https://") + or string_val.startswith("data:") + or string_val.startswith("/") + ): + return "" + + # Check if this value belongs to a URL-like key in the data model + # Look backwards in the prefix (max 200 chars) for the "key" assignment + prev_key_matches = re.findall( + r'"key"\s*:\s*"([^"]+)"', prefix[-200:] + ) + if prev_key_matches: + data_key = prev_key_matches[-1].lower() + if any( + k in data_key + for k in ("url", "link", "src", "href", "image") + ): + return "" + fixed += '"' + + # 2. Clean up trailing comma + fixed = fixed.rstrip() + if fixed.endswith(","): + fixed = fixed[:-1].rstrip() + + # 3. Close braces and brackets + while stack: + opening = stack.pop() + fixed += "}" if opening == "{" else "]" + + return fixed + + def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]) -> None: + for char in chunk: + char_handled = False + if self._brace_count == 0: + if char == "[": + self._in_top_level_list = True + elif char == "{": + pass + else: + continue + + # Track string state to avoid miscounting braces inside strings + if not char_handled and self._in_string: + if self._string_escaped: + self._string_escaped = False + if self._brace_count > 0: + self._json_buffer += char + elif char == "\\": + self._string_escaped = True + if self._brace_count > 0: + self._json_buffer += char + elif char == '"': + self._in_string = False + if self._brace_count > 0: + self._json_buffer += char + else: + if self._brace_count > 0: + self._json_buffer += char + char_handled = True + + if not char_handled: + if char == '"': + self._in_string = True + self._string_escaped = False + if self._brace_count > 0: + self._json_buffer += char + elif char == "{": + if self._brace_count == 0: + self._msg_types = [] + # Store (type, index) on stack + self._brace_stack.append(("{", len(self._json_buffer))) + self._json_buffer += "{" + self._brace_count += 1 + elif char == "}": + # Trigger object recognition + # In v0.8 streaming, we might be nested inside surfaceUpdate/components list + # So we check if it looks like a component even if brace_count > 1 + if self._brace_stack: # Ensure there's an opening brace to pop + # Pop the typed entry + b_type, start_idx = self._brace_stack.pop() + # If we popped a bracket while looking for a brace, we have a mismatch + # but we'll be resilient and just continue + self._json_buffer += "}" + self._brace_count -= 1 + + if ( + self._brace_count >= 0 + ): # Allow processing even if not top-level object + # The `i` here is the index within the current `chunk`. + # We need to get the full buffer from `start_idx` in `_json_buffer` + # up to the current point where `char` (which is '}') was added. + # The `_json_buffer` already has `char` appended. + obj_buffer = self._json_buffer[start_idx:] + if obj_buffer.startswith("{") and obj_buffer.endswith("}"): + try: + obj = json.loads(obj_buffer) + if isinstance(obj, dict): + self._found_valid_json_in_block = True + logger.debug( + f"[Parsed Dict] Keys: {list(obj.keys())}," + " protocol check follows..." + ) + + is_protocol = ( + self._in_top_level_list + and self.is_protocol_msg(obj) + ) + is_comp = obj.get("id") and obj.get("component") + # Process objects at top-level OR items in top-level list + # When in a list, we are top-level if the ONLY thing on the stack is the list opener + is_top_level = ( + len(self._brace_stack) == 0 + ) or ( + self._in_top_level_list + and len(self._brace_stack) == 1 + and self._brace_stack[0][0] == "[" + ) + if is_comp: + self._handle_partial_component( + obj, messages + ) + elif is_top_level or is_protocol: + if not self._handle_complete_object( + obj, self.surface_id, messages + ): + # Not a recognized message type. Validate to catch schema errors. + self._yield_messages([obj], messages) + + if self._brace_count == 0 or ( + self._in_top_level_list + and len(self._brace_stack) == 1 + ): + # Aggressively clear processed objects from the buffer to prevent slowdown. + if ( + len(self._brace_stack) == 1 + and self._brace_stack[0][0] == "[" + ): + # Keep '[' and remove the object after it + self._json_buffer = ( + self._json_buffer[:start_idx] + + self._json_buffer[ + start_idx + len(obj_buffer) : + ] + ) + else: + self._json_buffer = self._json_buffer[ + len(obj_buffer) : + ] + if self._brace_stack: + shift = len(obj_buffer) + self._brace_stack = [ + (b_t, i - shift) + for b_t, i in self._brace_stack + ] + + except json.JSONDecodeError as e: + logger.debug(f"Object recognition failed: {e}") + + elif char == "[": + self._brace_stack.append(("[", len(self._json_buffer))) + self._json_buffer += "[" + self._brace_count += 1 + elif char == "]": + if self._brace_stack and self._brace_stack[-1][0] == "[": + # Pop the typed entry + b_type, start_idx = self._brace_stack.pop() + self._json_buffer += "]" + self._brace_count -= 1 + if self._brace_count == 0: + self._in_top_level_list = False + else: + if self._brace_count > 0: + self._json_buffer += char + + # Sniff for metadata reactively on key delimiters to catch identifiers early + if self._brace_count > 0 and char in ('"', ":", ",", "}", "]"): + self._sniff_metadata() + + # Sniff for partial components at the end of the chunk + if self._brace_count >= 1 and self._json_buffer: + self._sniff_partial_component(messages) + self._sniff_partial_data_model(messages) + + if self._topology_dirty: + self.yield_reachable(messages, check_root=False, raise_on_orphans=False) + self._topology_dirty = False + + def _construct_sniffed_data_model_message( + self, active_msg_type: str, delta_msg_payload: Dict[str, Any] + ) -> Dict[str, Any]: + """Returns the message to yield for a partial data model update.""" + return {active_msg_type: delta_msg_payload} + + def _sniff_partial_data_model(self, messages: List[ResponsePart]) -> None: + msg_type = self._data_model_msg_type + if f'"{msg_type}"' not in self._json_buffer: + return + # Look through the brace stack for objects that might contain data model updates + for b_type, start_idx in reversed(self._brace_stack): + if b_type != "{": + continue + raw_fragment = self._json_buffer[start_idx:] + if not raw_fragment: + continue + + fixed_fragment = self._fix_json(raw_fragment) + obj = None + try: + obj = json.loads(fixed_fragment) + except json.JSONDecodeError: + # Fallback: iteratively strip from the last comma + # This handles cases where _fix_json produces invalid JSON + # from an incomplete trailing element (e.g. `{"key"}` from `{"ke`) + trimmed = raw_fragment + while "," in trimmed: + trimmed = trimmed.rsplit(",", 1)[0] + try: + fixed_trimmed = self._fix_json(trimmed) + if fixed_trimmed: + obj = json.loads(fixed_trimmed) + break + except json.JSONDecodeError: + continue + + if obj and isinstance(obj, dict): + active_msg_type = None + msg_type = self._data_model_msg_type + if msg_type in obj: + active_msg_type = msg_type + + if active_msg_type: + dm_obj = obj[active_msg_type] + if isinstance(dm_obj, dict) and "contents" in dm_obj: + + raw_contents = dm_obj["contents"] + contents_dict = self._parse_contents_to_dict(raw_contents) + + if contents_dict: + delta = {} + for k, v in contents_dict.items(): + if self._yielded_data_model.get(k) != v: + delta[k] = v + + if delta: + sid = ( + dm_obj.get(SURFACE_ID_KEY) + or self._surface_id + or "default" + ) + # Deduplicate delta_contents by only keeping the LATEST entry for each dirty key + delta_contents: Union[ + List[Dict[str, Any]], Dict[str, Any] + ] + if isinstance(raw_contents, list): + list_contents: List[Dict[str, Any]] = [] + seen_keys = set() + for entry in reversed(raw_contents): + if not isinstance(entry, dict): + continue + entry_key = entry.get("key") + # Only include entries that have a valid parsed key (cumulative) + if ( + entry_key + and entry_key in contents_dict + and entry_key not in seen_keys + ): + list_contents.insert(0, entry) + seen_keys.add(entry_key) + delta_contents = ( + self._prune_incomplete_datamodel_entries( + list_contents + ) + ) + else: + delta_contents = delta + + delta_msg_payload = { + SURFACE_ID_KEY: sid, + "contents": delta_contents, + } + if "path" in dm_obj: + delta_msg_payload["path"] = dm_obj["path"] + + delta_msg = self._construct_sniffed_data_model_message( + active_msg_type, delta_msg_payload + ) + self._yield_messages( + [delta_msg], messages, config=RELAXED_VALIDATION + ) + + self._yielded_data_model.update(contents_dict) + # Update internal model for path resolution + self.update_data_model(dm_obj, messages) + + def _sniff_partial_component(self, messages: List[ResponsePart]) -> None: + """Attempts to parse a partial component from the current buffer.""" + # We only care about components if we are inside a "components" array + if f'"{CATALOG_COMPONENTS_KEY}"' not in self._json_buffer: + return + # Try parsing from inner to outer to find the smallest complete component + for b_type, start_idx in reversed(self._brace_stack): + if b_type != "{": + continue + raw_fragment = self._json_buffer[start_idx:] + if not raw_fragment: + continue + + fixed_fragment = self._fix_json(raw_fragment) + try: + obj = json.loads(fixed_fragment) + if isinstance(obj, dict) and obj.get("id") and obj.get("component"): + if isinstance(obj["component"], str): + # Flat style (v0.9+): component type is a string + self._handle_partial_component(obj, messages) + elif ( + isinstance(obj["component"], dict) and len(obj["component"]) > 0 + ): + # Structured style (v0.8): Ignore components that are effectively empty (no type keys) + self._handle_partial_component(obj, messages) + + except Exception: + continue + + def _sniff_metadata(self) -> None: + """Sniffs for surfaceId, root, and msg_types in the current json_buffer.""" + raise NotImplementedError("Subclasses must implement _sniff_metadata") + + def _prune_incomplete_datamodel_entries(self, entries: Any) -> Any: + """Recursively removes data model entries that only contain 'key' and no valid values.""" + if not isinstance(entries, list): + return entries + + pruned = [] + for entry in entries: + if not isinstance(entry, dict): + pruned.append(entry) + continue + + has_val = False + for vkey in ("value", "valueString", "valueNumber", "valueBoolean"): + if vkey in entry: + has_val = True + break + + if "valueMap" in entry: + pruned_map = self._prune_incomplete_datamodel_entries(entry["valueMap"]) + # valueMap is considered valid even if empty, meaning map was explicitly empty + if isinstance(pruned_map, list): + if not pruned_map and len(entry["valueMap"]) > 0: + # If it was non-empty and became empty, it only had incomplete elements. Discard map. + del entry["valueMap"] + else: + entry["valueMap"] = pruned_map + has_val = True + + if has_val and "key" in entry: + pruned.append(entry) + + return pruned + + def _handle_partial_component( + self, comp: Dict[str, Any], messages: List[ResponsePart] + ) -> None: + """Handles a component discovered before its parent message is finished. + + When the parser sees a full JSON object that looks like a component + (contains `id` and `component` keys) within a larger message, it caches + the component and attempts to yield it immediately if it's reachable. + + Args: + comp: The parsed component dictionary. + messages: The list to append any renderable partial messages to. + """ + comp_id = comp.get("id") + if not comp_id: + return + + # Skip caching this component if it has empty dictionaries for complex properties. + # Elements like `children`, `text`, `url`, etc., violate A2UI schema if empty + # and will crash the client. We want the parent to yield a loading placeholder instead. + def _has_empty_dict(obj: Any) -> bool: + if isinstance(obj, dict): + if not obj: + return True + return any(_has_empty_dict(v) for v in obj.values()) + elif isinstance(obj, list): + return any(_has_empty_dict(v) for v in obj) + return False + + component_def = comp.get("component") + if isinstance(component_def, str): + # v0.9 flat style: check the whole component object for empty dicts + if _has_empty_dict(comp): + return + elif _has_empty_dict(component_def): + # v0.8 nested style: check properties inside component + return + + if isinstance(component_def, dict) and hasattr(self, "_required_fields_map"): + comp_type = next(iter(component_def.keys())) if component_def else None + if comp_type: + props = component_def.get(comp_type, {}) + if isinstance(props, dict): + required_fields = self._required_fields_map.get(comp_type, set()) + for req in required_fields: + if req not in props: + return + + self._seen_components[comp_id] = comp + self._topology_dirty = True + + def _parse_contents_to_dict(self, raw_contents: Any) -> Dict[str, Any]: + """Recursively parses a list of A2UI contents into a flat dictionary.""" + if isinstance(raw_contents, dict): + return raw_contents + if not isinstance(raw_contents, list): + return {} + + res = {} + for entry in raw_contents: + if not isinstance(entry, dict): + continue + key = entry.get("key") + val = None + for vkey in ["value", "valueString", "valueNumber", "valueBoolean"]: + if vkey in entry: + val = entry[vkey] + break + + if val is None and "valueMap" in entry: + val = self._parse_contents_to_dict(entry["valueMap"]) + + if key and val is not None: + res[key] = val + return res + + def update_data_model( + self, update: Dict[str, Any], messages: List[ResponsePart] + ) -> None: + """Updates the internal data model and marks affected components as dirty.""" + # Hook method to be overridden by subclasses to handle completed data model updates. + pass + + def _handle_complete_object( + self, + obj: Dict[str, Any], + sid: Optional[str], + messages: List[ResponsePart], + ) -> bool: + """Handles an object that has been fully parsed. To be implemented by subclasses.""" + raise NotImplementedError("Subclasses must implement _handle_complete_object") + + def yield_reachable( + self, + messages: List[ResponsePart], + check_root: bool = False, + raise_on_orphans: bool = False, + ) -> None: + """Yields a partial message containing all reachable and seen components. + + This is the core of the streaming logic. Instead of waiting for a UI message + which could contain dozens of components, we yield "partial" updates as soon + as we have enough components to build a renderable sub-tree from the root. + + Args: + messages: The list to which partial messages will be appended. + check_root: If True, raises an error if the root component isn't seen yet. + raise_on_orphans: If True, uses strict topology analysis to catch loops. + """ + active_msg_type = self._get_active_msg_type_for_components() + if not self.root_id or not active_msg_type: + return + + # Buffer components until we have a beginRendering or createSurface for a known surface. + if not self.surface_id: + return + + sid = self.surface_id + if sid not in self._yielded_surfaces_set and not self._buffered_start_message: + return + + try: + # Analyze topology of current seen components + components_to_analyze = list(self._seen_components.values()) + + if check_root and self.root_id not in self._seen_components: + raise A2uiIntegrityError( + f"No root component (id='{self.root_id}') found in" + f" {active_msg_type}" + ) + + reachable_ids = analyze_topology( + components_to_analyze, + self._ref_fields_map, + root_id=self.root_id, + allow_orphan_components=not raise_on_orphans, + ) + + # We only yield components we actually have in our "seen" cache + available_reachable = reachable_ids & set(self._seen_components.keys()) + + if check_root and not available_reachable: + raise A2uiIntegrityError( + f"No root component (id='{self.root_id}') found in" + f" {active_msg_type}" + ) + + # 1. Process placeholders and partial children + processed_components: List[Dict[str, Any]] = [] + extra_components: List[Dict[str, Any]] = [] + surface_id = self.surface_id or "unknown" + yielded_for_surface = self._yielded_ids.get(surface_id, set()) + + for rid in sorted(list(available_reachable)): + comp = copy.deepcopy(self._seen_components[rid]) + # Apply path placeholders and prune unseen children in a single pass + self._process_component_topology(comp, extra_components) + processed_components.append(comp) + + # Add generated placeholders to the yield + processed_components.extend(extra_components) + + # 2. Check if we have NEW or UPDATED reachable components to yield for THIS surface + surface_id = self.surface_id + if not surface_id or surface_id in self._deleted_surfaces: + return + + should_yield = False + if available_reachable - yielded_for_surface: + should_yield = True + else: + # Check if any yielded component's content has changed for this surface + for comp in processed_components: + cid = comp["id"] + content_str = json.dumps(comp, sort_keys=True) + state_key = (surface_id, cid) + if self._yielded_contents.get(state_key) != content_str: + should_yield = True + break + + if should_yield: + current_sid = self.surface_id or "unknown" + if ( + self._buffered_start_message + and current_sid not in self._yielded_start_messages + ): + self._yield_messages([self._buffered_start_message], messages) + self._yielded_start_messages.add(current_sid) + self._yielded_surfaces_set.add(current_sid) + + # Construct a partial message of the correct type + partial_msg = self._construct_partial_message( + processed_components, active_msg_type + ) + + # Use RELAXED_VALIDATION for partial fragments yielded during streaming + self._yield_messages([partial_msg], messages, config=RELAXED_VALIDATION) + self._yielded_ids.setdefault(surface_id, set()).update( + available_reachable + ) + + # Update content/placeholder tracking + for comp in processed_components: + cid = comp["id"] + self._yielded_contents[(surface_id, cid)] = json.dumps( + comp, sort_keys=True + ) + + except ValueError as e: + if "Circular reference detected" in str(e): + raise e + # Other topology errors (like orphans) are ignored during streaming + # as dependencies might still be on the wire. + msg = str(e) + if ( + raise_on_orphans + or "Circular" in msg + or "Self-reference" in msg + or "recursion" in msg.lower() + or check_root + ): + logger.debug(f"yield_reachable error (strict={check_root}): {msg}") + raise e + + def _get_placeholder_id(self, child_id: str) -> str: + """Returns the ID to use for a missing child placeholder.""" + return f"loading_{child_id}" + + def _process_component_topology( + self, + comp: Dict[str, Any], + extra_components: List[Dict[str, Any]], + ) -> None: + """Recursively processes path placeholders and child pruning in one pass.""" + comp_id = comp.get("id", "unknown") + + # Start recursion from the component content + if isinstance(comp.get("component"), dict): + self._traverse_component_topology( + comp.get("component", {}), extra_components, comp_id + ) + else: + # Flat style properties are siblings to 'component' type key + self._traverse_component_topology(comp, extra_components, comp_id) + + def _traverse_component_topology( + self, + obj: Any, + extra_components: List[Dict[str, Any]], + comp_id: str, + parent_key: Optional[str] = None, + ) -> None: + """Internal recursive helper to traverse and process component properties.""" + if isinstance(obj, dict): + # 1. Handle Path Placeholders (from _apply_placeholders) + if ( + "path" in obj + and isinstance(obj["path"], str) + and obj["path"].startswith("/") + ): + path = obj["path"] + key = path.lstrip("/") + if self._version != VERSION_0_9: + if "componentId" not in obj: + obj.clear() + obj.update({"path": "/" + key}) + elif self._version != VERSION_0_9: + # If not in data model, still ensure path has leading slash if it's a bindable object (v0.8 only) + current_path = obj.get("path") + if current_path is not None: + if not isinstance(current_path, str) or not current_path.startswith( + "/" + ): + obj["path"] = "/" + str(current_path) + + # 2. Handle Child Pruning (from _prune_unseen_children) + for field in ( + "children", + "explicitList", + "child", + "contentChild", + "entryPointChild", + "componentId", + ): + if field in obj: + if isinstance(obj[field], list): + valid_children = [] + for child_id in obj[field]: + if child_id in self._seen_components: + valid_children.append(child_id) + else: + # Individual placeholder for missing child + placeholder_id = self._get_placeholder_id(child_id) + valid_children.append(placeholder_id) + placeholder_comp = { + "id": placeholder_id, + **self._placeholder_component, + } + # Avoid duplicates in extra_components + if not any( + ec["id"] == placeholder_id + for ec in extra_components + ): + extra_components.append(placeholder_comp) + + if not valid_children and field in ( + "children", + "explicitList", + ): + # If list is empty, check if it was partial in the buffer + # (meaning it's a sequence that started but hasn't yielded items yet) + term = f'"{field}"' + if term in self._json_buffer: + # Simple check: is there a [ after the field name in the buffer? + after_field = self._json_buffer.split(term)[-1] + if ( + "[" in after_field + and "]" not in after_field.split("[")[0] + ): + placeholder_id = f"loading_children_{comp_id}" + valid_children.append(placeholder_id) + placeholder_comp = { + "id": placeholder_id, + **self._placeholder_component, + } + if not any( + ec["id"] == placeholder_id + for ec in extra_components + ): + extra_components.append(placeholder_comp) + obj[field] = valid_children + elif isinstance(obj[field], str): + child_id = obj[field] + if child_id not in self._seen_components: + placeholder_id = self._get_placeholder_id(child_id) + obj[field] = placeholder_id + placeholder_comp = { + "id": placeholder_id, + **self._placeholder_component, + } + if not any( + ec["id"] == placeholder_id for ec in extra_components + ): + extra_components.append(placeholder_comp) + + # Continue traversal on values + for k, v in list(obj.items()): + self._traverse_component_topology( + v, extra_components, comp_id, parent_key=k + ) + elif isinstance(obj, list): + for item in obj: + self._traverse_component_topology( + item, extra_components, comp_id, parent_key + ) + + +# Deprecated redirect. +A2uiStreamParser = TransportStreamParser diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v08.py b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/streaming_v08.py similarity index 95% rename from agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v08.py rename to agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/streaming_v08.py index ff910c282..c410d5798 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v08.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/streaming_v08.py @@ -15,20 +15,19 @@ from __future__ import annotations import re -import json 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 -from a2ui.core.validating.validator import ValidationConfig, RELAXED_VALIDATION, STRICT_VALIDATION +from a2ui.inference_formats.transport.streaming import TransportStreamParser +from a2ui.parser.response_part import ResponsePart +from a2ui.parser.constants import * +from a2ui.schema.constants import SURFACE_ID_KEY, CATALOG_COMPONENTS_KEY +from a2ui.core.validating.validator import RELAXED_VALIDATION if TYPE_CHECKING: - from ..schema.catalog import A2uiCatalog + from a2ui.schema.catalog import A2uiCatalog -class A2uiStreamParserV08(A2uiStreamParser): +class TransportStreamParserV08(TransportStreamParser): """Streaming parser implementation for A2UI v0.8 specification.""" def __init__(self, catalog: A2uiCatalog): diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v09.py b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/streaming_v09.py similarity index 97% rename from agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v09.py rename to agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/streaming_v09.py index eb412dcff..07ebd3746 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming_v09.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/inference_formats/transport/streaming_v09.py @@ -18,17 +18,17 @@ import json 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 +from a2ui.inference_formats.transport.streaming import TransportStreamParser +from a2ui.parser.response_part import ResponsePart +from a2ui.parser.constants import * +from a2ui.schema.constants import SURFACE_ID_KEY, CATALOG_COMPONENTS_KEY from a2ui.core.validating.validator import RELAXED_VALIDATION if TYPE_CHECKING: - from ..schema.catalog import A2uiCatalog + from a2ui.schema.catalog import A2uiCatalog -class A2uiStreamParserV09(A2uiStreamParser): +class TransportStreamParserV09(TransportStreamParser): """Streaming parser implementation for A2UI v0.9 specification.""" def __init__(self, catalog: A2uiCatalog): diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/inference_strategy.py b/agent_sdks/python/a2ui_agent/src/a2ui/inference_strategy.py index 64ade4dd2..3206122d3 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/inference_strategy.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/inference_strategy.py @@ -12,40 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -from abc import ABC, abstractmethod -from typing import Optional, Any +"""Deprecated compatibility redirect for InferenceStrategy.""" +import warnings +from a2ui.inference_format import InferenceFormat as InferenceStrategy -class InferenceStrategy(ABC): +warnings.warn( + "a2ui.inference_strategy is deprecated and will be removed. " + "Please import from a2ui.inference_format instead.", + DeprecationWarning, + stacklevel=2, +) - @abstractmethod - def generate_system_prompt( - self, - role_description: str, - workflow_description: str = "", - ui_description: str = "", - client_ui_capabilities: Optional[dict[str, Any]] = None, - allowed_components: Optional[list[str]] = None, - allowed_messages: Optional[list[str]] = None, - include_schema: bool = False, - include_examples: bool = False, - validate_examples: bool = False, - ) -> str: - """ - Generates a system prompt for all LLM requests. - - Args: - role_description: Description of the agent's role. - workflow_description: Description of the workflow. - ui_description: Description of the UI. - client_ui_capabilities: Capabilities reported by the client for targeted schema pruning. - allowed_components: List of allowed catalog components. - allowed_messages: List of allowed messages. - include_schema: Whether to include the schema. - include_examples: Whether to include examples. - validate_examples: Whether to validate examples. - - Returns: - The system prompt. - """ - pass +__all__ = [ + "InferenceStrategy", +] diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/parser/parser.py b/agent_sdks/python/a2ui_agent/src/a2ui/parser/parser.py index ca435d067..7f991d2ab 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/parser/parser.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/parser/parser.py @@ -12,78 +12,123 @@ # See the License for the specific language governing permissions and # limitations under the License. -import re -from typing import List, Optional, Any +"""Abstract parser interface and legacy parsing compatibility helpers.""" + +import warnings +from abc import ABC, abstractmethod +from typing import List, Any from .response_part import ResponsePart -from ..schema.constants import A2UI_OPEN_TAG, A2UI_CLOSE_TAG -from .payload_fixer import parse_and_fix -from a2ui.core import A2uiParseError -_A2UI_BLOCK_PATTERN = re.compile( - f"{re.escape(A2UI_OPEN_TAG)}(.*?){re.escape(A2UI_CLOSE_TAG)}", re.DOTALL -) +class Parser(ABC): + """Abstract interface defining the response parser and compiler.""" + @abstractmethod + def has_format_content(self, content: str, *, complete: bool = False) -> bool: + """Checks if the content contains blocks belonging to this parser's format. -def has_a2ui_parts(content: str) -> bool: - """Checks if the content has A2UI parts.""" - return A2UI_OPEN_TAG in content and A2UI_CLOSE_TAG in content + Args: + content: The raw LLM response. + complete: If True, checks that the format block is closed/complete. + Returns: + True if the content contains blocks belonging to this format. + """ + pass -def _sanitize_json_string(json_string: str) -> str: - """Sanitizes the JSON string by removing markdown code blocks.""" - json_string = json_string.strip() - if json_string.startswith("```json"): - json_string = json_string[len("```json") :] - elif json_string.startswith("```"): - json_string = json_string[len("```") :] - if json_string.endswith("```"): - json_string = json_string[: -len("```")] - json_string = json_string.strip() - return json_string + def parse_response(self, content: str) -> List[ResponsePart]: + """Parses full response content into standard JSON payload parts by unwrapping and compiling. + Args: + content: The raw LLM response. -def parse_response(content: str) -> List[ResponsePart]: - """ - Parses the LLM response into a list of ResponsePart objects. + Returns: + A list of ResponsePart objects containing text and compiled JSON. + """ + parts = self.unwrap(content) + for part in parts: + if part.a2ui_raw is not None: + part.a2ui_json = self.compile(part.a2ui_raw) + return parts - Args: - content: The raw LLM response. + @abstractmethod + def unwrap(self, content: str) -> List[ResponsePart]: + """Tokenizes response content into raw format-content parts. - Returns: - A list of ResponsePart objects. + Args: + content: The raw LLM response. - Raises: - A2uiParseError: If no A2UI tags are found or if the JSON part is invalid. - """ - matches = list(_A2UI_BLOCK_PATTERN.finditer(content)) + Returns: + A list of ResponsePart objects with a2ui_raw populated. + """ + pass + + @abstractmethod + def compile(self, format_content: str) -> List[dict[str, Any]]: + """Compiles raw format-content (inference format string) to structured A2UI messages. + + Args: + format_content: The raw format-content extracted from response. - if not matches: - raise A2uiParseError( - f"A2UI tags '{A2UI_OPEN_TAG}' and '{A2UI_CLOSE_TAG}' not found in response." - ) + Returns: + A list of compiled A2UI message dictionaries. + """ + pass - response_parts = [] - last_end = 0 + @abstractmethod + def process_chunk(self, chunk: str) -> List[ResponsePart]: + """Processes a streamed token chunk (incremental parsing). - for match in matches: - start, end = match.span() - # Text preceding the JSON block - text_part = content[last_end:start].strip() + Args: + chunk: The next text chunk from the stream. - # The JSON content within the tags - json_string = match.group(1) - json_string_cleaned = _sanitize_json_string(json_string) - if not json_string_cleaned: - raise A2uiParseError("A2UI JSON part is empty.") + Returns: + A list of parsed or completed ResponsePart objects. + """ + pass - json_data = parse_and_fix(json_string_cleaned) - response_parts.append(ResponsePart(text=text_part, a2ui_json=json_data)) - last_end = end - # Trailing text after the last JSON block - trailing_text = content[last_end:].strip() - if trailing_text: - response_parts.append(ResponsePart(text=trailing_text, a2ui_json=None)) +def has_a2ui_parts(content: str) -> bool: + """Checks if the content has A2UI parts (legacy compatibility helper). + + Args: + content: The raw response text. - return response_parts + Returns: + Whether the content contains open and close A2UI tags. + """ + warnings.warn( + "has_a2ui_parts is deprecated. Please use" + " format.parser.has_format_content(content, complete=True) on your" + " InferenceFormat instance instead.", + DeprecationWarning, + stacklevel=2, + ) + from a2ui.schema.constants import A2UI_OPEN_TAG, A2UI_CLOSE_TAG + + return A2UI_OPEN_TAG in content and A2UI_CLOSE_TAG in content + + +def parse_response(content: str) -> List[ResponsePart]: + """Parses the LLM response into a list of ResponsePart objects (legacy). + + Args: + content: The raw LLM response. + + Returns: + A list of ResponsePart objects. + """ + warnings.warn( + "parse_response is deprecated. Please use format.parser.parse_response(...) " + "on your InferenceFormat instance instead.", + DeprecationWarning, + stacklevel=2, + ) + from a2ui.inference_formats.transport.parser import unwrap_response + from a2ui.parser.payload_fixer import parse_and_fix + + parts = unwrap_response(content) + for part in parts: + if part.a2ui_raw is not None: + part.a2ui_json = parse_and_fix(part.a2ui_raw) + return parts diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/parser/response_part.py b/agent_sdks/python/a2ui_agent/src/a2ui/parser/response_part.py index 517b4d0a0..73705c29a 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/parser/response_part.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/parser/response_part.py @@ -22,10 +22,13 @@ class ResponsePart: Attributes: text: The conversational text part. Can be an empty string. - a2ui_json: The parsed A2UI JSON data, always a list of dictionaries if - it contains A2UI messages. None if this part only contains trailing - text. + a2ui_raw: The raw uncompiled format content string (e.g. raw XML/DSL/JSON). + None if this part only contains conversational text. + a2ui_json: The parsed/compiled A2UI JSON data, always a list of + dictionaries if it contains A2UI messages. None if this part only + contains conversational text. """ text: str = "" + a2ui_raw: Optional[str] = None a2ui_json: Optional[Any] = None diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming.py b/agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming.py index 23a885ac5..110e7e592 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/parser/streaming.py @@ -12,1142 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations +"""Deprecated compatibility redirect for A2uiStreamParser.""" -import copy -import json -import logging -import re -from typing import Any, List, Dict, Optional, Set, TYPE_CHECKING, Union +import warnings +from a2ui.inference_formats.transport.streaming import A2uiStreamParser -from .constants import * -from ..schema.constants import ( - VERSION_0_9, - VERSION_0_8, - A2UI_OPEN_TAG, - A2UI_CLOSE_TAG, - SURFACE_ID_KEY, - CATALOG_COMPONENTS_KEY, +warnings.warn( + "a2ui.parser.streaming is deprecated. " + "Please import from a2ui.inference_formats.transport instead.", + DeprecationWarning, + stacklevel=2, ) -from ..schema.validator import ( - extract_component_ref_fields, - extract_component_required_fields, -) -from ..schema.validator import A2uiValidator -from a2ui.core.validating import analyze_topology -from .response_part import ResponsePart -from a2ui.core.validating.validator import RELAXED_VALIDATION, STRICT_VALIDATION, ValidationConfig -from a2ui.core import A2uiParseError, A2uiIntegrityError - - -if TYPE_CHECKING: - from ..schema.catalog import A2uiCatalog - -logger = logging.getLogger(__name__) - - -class A2uiStreamParser: - """Parses a stream of text for A2UI JSON messages with fine-grained component yielding. - - This class acts as a factory that returns a version-specific parser instance - (V08 or V09) depending on the catalog version. - """ - - def __new__(cls, catalog: A2uiCatalog) -> A2uiStreamParser: - if cls is A2uiStreamParser: - 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 A2uiStreamParserV08(catalog=catalog) - else: - from .streaming_v09 import A2uiStreamParserV09 - - return A2uiStreamParserV09(catalog=catalog) - return super().__new__(cls) - - 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 = "" - self._json_buffer = "" - self._brace_stack: list[tuple[str, int]] = [] - self._brace_count = 0 - self._in_top_level_list = False - self._in_string = False - self._string_escaped = False - - self._seen_components: Dict[str, Dict[str, Any]] = {} - - # Track data model for path resolution - self._yielded_data_model: Dict[str, Any] = {} - self._deleted_surfaces: Set[str] = set() - - # Set of unique component IDs yielded per surface to prevent duplicate yielding - # surfaceId -> set of cids - self._yielded_ids: Dict[str, Set[str]] = {} - # (surfaceId, cid) -> hash of content for change detection - self._yielded_contents: Dict[Any, str] = {} - - self._root_ids: Dict[str, str] = {} # The root component IDs mapped per surface - self._default_root_id: Optional[str] = ( - None # Base default root ID for the protocol - ) - self._unbound_root_id: Optional[str] = ( - None # Temporary holding variable for when root arrives before surfaceId - ) - self._surface_id: Optional[str] = ( - None # The active surface ID tracking the context - ) - self._msg_types: List[str] = ( - [] - ) # Running list of message types seen in the block - - # A set of surface ids for which we have already yielded a start message - # Tracks if beginRendering or createSurface was emitted - self._yielded_start_messages: Set[str] = set() - - # The current active message type for component grouping - self._active_msg_type: Optional[str] = None - - # State for buffering updates until surface is ready - self._pending_messages: Dict[str, List[Dict[str, Any]]] = ( - {} - ) # surfaceId -> list of msgs delayed until start message arrives - self._buffered_start_message: Optional[Dict[str, Any]] = ( - None # The start message to yield before any components - ) - self._topology_dirty = False # Set to true if components are added out of order - self._in_top_level_list = False - self._found_valid_json_in_block = False - - @property - def _placeholder_component(self) -> Dict[str, Any]: - """Returns the version-specific placeholder component. - - This is used when a component references a child component that hasn't yet - streamed in. The placeholder component is added to the components list and - the reference is updated to point to the placeholder component. - """ - raise NotImplementedError("Subclasses must implement _placeholder_component") - - @property - def surface_id(self) -> Optional[str]: - return self._surface_id - - @surface_id.setter - def surface_id(self, value: Optional[str]) -> None: - self._surface_id = value - if value is not None and self._unbound_root_id is not None: - self._root_ids[value] = self._unbound_root_id - self._unbound_root_id = None - - @property - def root_id(self) -> Optional[str]: - if self._surface_id: - return self._root_ids.get(self._surface_id, self._default_root_id) - # Return unbound root ID if explicitly sniffed, otherwise use protocol default - return ( - self._unbound_root_id - if self._unbound_root_id is not None - else self._default_root_id - ) - - @root_id.setter - def root_id(self, value: Optional[str]) -> None: - if self._surface_id: - if value is not None: - self._root_ids[self._surface_id] = value - else: - self._root_ids.pop(self._surface_id, None) - else: - self._unbound_root_id = value - - @property - def msg_types(self) -> List[str]: - return self._msg_types - - def add_msg_type(self, msg_type: str) -> None: - if msg_type not in self._msg_types: - self._msg_types.append(msg_type) - if msg_type in ( - MSG_TYPE_SURFACE_UPDATE, - MSG_TYPE_UPDATE_COMPONENTS, - MSG_TYPE_CREATE_SURFACE, - ): - self._active_msg_type = msg_type - - @property - def _yielded_surfaces_set(self) -> Set[str]: - """Provides access to version-specific yielded surfaces set.""" - raise NotImplementedError("Subclasses must implement _yielded_surfaces_set") - - def is_protocol_msg(self, obj: Dict[str, Any]) -> bool: - """Checks if the object is a recognized A2UI message for this version.""" - raise NotImplementedError("Subclasses must implement is_protocol_msg") - - @property - def _data_model_msg_type(self) -> str: - """Returns the message type identifier for data model updates.""" - raise NotImplementedError("Subclasses must implement _data_model_msg_type") - - def _get_active_msg_type_for_components(self) -> Optional[str]: - """Determines which msg_type to use when wrapping component updates.""" - raise NotImplementedError( - "Subclasses must implement _get_active_msg_type_for_components" - ) - - def _construct_partial_message( - self, components: List[Dict[str, Any]], active_msg_type: str - ) -> Dict[str, Any]: - """Constructs a partial message of the correct type. Subclasses must implement.""" - raise NotImplementedError( - "Subclasses must implement _construct_partial_message" - ) - - def _deduplicate_data_model(self, m: Dict[str, Any]) -> bool: - """Returns True if message should be yielded, False if skipped.""" - return True - - def _yield_messages( - self, - messages_to_yield: List[Dict[str, Any]], - messages: List[ResponsePart], - config: ValidationConfig = STRICT_VALIDATION, - ) -> None: - """Validates and appends messages to the final output list.""" - for m in messages_to_yield: - if not self._deduplicate_data_model(m): - continue - - # Each surface update message must specify a surfaceId and satisfy catalog validation. - if self._validator: - try: - self._validator.validate(m, root_id=self.root_id, config=config) - except ValueError as e: - if config == STRICT_VALIDATION: - raise e - else: - logger.debug( - f"Validation failed for partial/sniffed message: {e}" - ) - continue - - # Consolidated appending logic - if messages and messages[-1].a2ui_json is None: - messages[-1].a2ui_json = [m] - elif messages and isinstance(messages[-1].a2ui_json, list): - messages[-1].a2ui_json.append(m) - else: - messages.append(ResponsePart(a2ui_json=[m])) - - def _delete_surface(self, sid: str) -> None: - """Clears all state related to a specific surface.""" - self._pending_messages.pop(sid, None) - self._yielded_ids.pop(sid, None) - - # Clear contents for this surface - self._yielded_contents = { - k: v for k, v in self._yielded_contents.items() if k[0] != sid - } - self._yielded_surfaces_set.discard(sid) - self._yielded_start_messages.discard(sid) - - self._deleted_surfaces.add(sid) - - def process_chunk(self, chunk: str) -> List[ResponsePart]: - """Processes a chunk of text and returns any complete A2UI messages found. - - This is the primary entry point for the streaming parser. It handles the - initial "tag hunt" and then delegates JSON fragment processing to - `_process_json_chunk`. It supports multiple A2UI blocks in a single stream. - - Args: - chunk: The chunk of raw text (e.g., from an LLM stream) to process. - - Returns: - A list of parsed A2UI message dictionaries. - """ - messages = [] - self._buffer += chunk - - while True: - if not self._found_delimiter: - # Looking for - if A2UI_OPEN_TAG in self._buffer: - parts = self._buffer.split(A2UI_OPEN_TAG, 1) - if parts[0]: - messages.append(ResponsePart(text=parts[0])) - self._found_delimiter = True - self._buffer = parts[1] - # Continue to process the content after the open tag - else: - # Yield conversational text while avoiding split tags - keep_len = 0 - for i in range(len(A2UI_OPEN_TAG) - 1, 0, -1): - if self._buffer.endswith(A2UI_OPEN_TAG[:i]): - keep_len = i - break - - if len(self._buffer) > keep_len: - safe_to_yield = len(self._buffer) - keep_len - text_to_yield = self._buffer[:safe_to_yield] - messages.append(ResponsePart(text=text_to_yield)) - self._buffer = self._buffer[safe_to_yield:] - break - - if self._found_delimiter: - # Looking for - if A2UI_CLOSE_TAG in self._buffer: - parts = self._buffer.split(A2UI_CLOSE_TAG, 1) - json_fragment = parts[0] - self._process_json_chunk(json_fragment, messages) - if not self._found_valid_json_in_block: - raise A2uiParseError( - "Failed to parse JSON: No valid JSON object found in A2UI" - " block." - ) - - # End of block: reset JSON state but keep seen_components - self._found_delimiter = False - self._reset_json_state() - - self._buffer = parts[1] - # Continue loop to look for next A2UI_OPEN_TAG in remaining buffer - else: - # Find if the buffer ends with a prefix of A2UI_CLOSE_TAG - # To avoid split-tag issues, we only delay processing if it looks like a close tag is starting - keep_len = 0 - for i in range(1, len(A2UI_CLOSE_TAG)): - if self._buffer.endswith(A2UI_CLOSE_TAG[:i]): - keep_len = i - - if keep_len < len(self._buffer): - to_process = self._buffer[: len(self._buffer) - keep_len] - self._buffer = self._buffer[len(self._buffer) - keep_len :] - self._process_json_chunk(to_process, messages) - break - - # Deduplicate surfaceUpdate messages to avoid over-yielding in a single chunk - for part in messages: - if not part.a2ui_json: - continue - - deduped_msgs = [] - seen_su = set() - # Iterate backwards to keep only the last (most complete) surfaceUpdate for each surface - for m in reversed(part.a2ui_json): - is_su = False - sid = None - if isinstance(m, dict) and MSG_TYPE_SURFACE_UPDATE in m: - is_su = True - sid = m[MSG_TYPE_SURFACE_UPDATE].get(SURFACE_ID_KEY) - - if is_su and sid: - if sid not in seen_su: - deduped_msgs.append(m) - seen_su.add(sid) - else: - deduped_msgs.append(m) - - deduped_msgs.reverse() - part.a2ui_json = deduped_msgs - - if messages: - logger.debug( - f"DEBUG: process_chunk returning {len(messages)} messages: {messages}" - ) - return messages - - def _reset_json_state(self) -> None: - """Resets the JSON-specific parsing state (e.g., at the end of a block).""" - self._json_buffer = "" - self._brace_stack = [] - self._brace_count = 0 - self._in_top_level_list = False - self._in_string = False - self._string_escaped = False - self._msg_types = [] - self._found_valid_json_in_block = False - # Note: we do NOT reset _active_msg_type or _yielded_contents here - - # so re-yielding works between blocks - - def _fix_json(self, fragment: str) -> str: - """Attempts to fix a partial JSON fragment by adding missing closing delimiters.""" - fixed = fragment.rstrip() - if not fixed: - return "" - - stack = [] - in_string = False - escaped = False - last_quote_idx = -1 - - # Single pass to track strings and braces - for i, char in enumerate(fixed): - if escaped: - escaped = False - continue - if char == "\\": - escaped = True - continue - if char == '"': - in_string = not in_string - if in_string: - last_quote_idx = i - elif not in_string: - if char in ("{", "["): - stack.append(char) - elif char in ("}", "]"): - if stack: - stack.pop() - - # 1. Close open strings (healing) - if in_string: - # We only auto-close strings for safe keys (CUTTABLE_KEYS) - prefix = fixed[:last_quote_idx].rstrip() - if prefix.endswith(":"): - key_match = re.findall(r'"([^"]+)"\s*:\s*$', prefix) - if key_match: - key = key_match[0] - if key not in self._cuttable_keys: - return "" - - # Special case: don't cut URL bindings, as partial URLs break images/links - if key == "valueString": - string_val = fixed[last_quote_idx + 1 :] - if ( - string_val.startswith("http://") - or string_val.startswith("https://") - or string_val.startswith("data:") - or string_val.startswith("/") - ): - return "" - - # Check if this value belongs to a URL-like key in the data model - # Look backwards in the prefix (max 200 chars) for the "key" assignment - prev_key_matches = re.findall( - r'"key"\s*:\s*"([^"]+)"', prefix[-200:] - ) - if prev_key_matches: - data_key = prev_key_matches[-1].lower() - if any( - k in data_key - for k in ("url", "link", "src", "href", "image") - ): - return "" - fixed += '"' - - # 2. Clean up trailing comma - fixed = fixed.rstrip() - if fixed.endswith(","): - fixed = fixed[:-1].rstrip() - - # 3. Close braces and brackets - while stack: - opening = stack.pop() - fixed += "}" if opening == "{" else "]" - - return fixed - - def _process_json_chunk(self, chunk: str, messages: List[ResponsePart]) -> None: - for char in chunk: - char_handled = False - if self._brace_count == 0: - if char == "[": - self._in_top_level_list = True - elif char == "{": - pass - else: - continue - - # Track string state to avoid miscounting braces inside strings - if not char_handled and self._in_string: - if self._string_escaped: - self._string_escaped = False - if self._brace_count > 0: - self._json_buffer += char - elif char == "\\": - self._string_escaped = True - if self._brace_count > 0: - self._json_buffer += char - elif char == '"': - self._in_string = False - if self._brace_count > 0: - self._json_buffer += char - else: - if self._brace_count > 0: - self._json_buffer += char - char_handled = True - - if not char_handled: - if char == '"': - self._in_string = True - self._string_escaped = False - if self._brace_count > 0: - self._json_buffer += char - elif char == "{": - if self._brace_count == 0: - self._msg_types = [] - # Store (type, index) on stack - self._brace_stack.append(("{", len(self._json_buffer))) - self._json_buffer += "{" - self._brace_count += 1 - elif char == "}": - # Trigger object recognition - # In v0.8 streaming, we might be nested inside surfaceUpdate/components list - # So we check if it looks like a component even if brace_count > 1 - if self._brace_stack: # Ensure there's an opening brace to pop - # Pop the typed entry - b_type, start_idx = self._brace_stack.pop() - # If we popped a bracket while looking for a brace, we have a mismatch - # but we'll be resilient and just continue - self._json_buffer += "}" - self._brace_count -= 1 - - if ( - self._brace_count >= 0 - ): # Allow processing even if not top-level object - # The `i` here is the index within the current `chunk`. - # We need to get the full buffer from `start_idx` in `_json_buffer` - # up to the current point where `char` (which is '}') was added. - # The `_json_buffer` already has `char` appended. - obj_buffer = self._json_buffer[start_idx:] - if obj_buffer.startswith("{") and obj_buffer.endswith("}"): - try: - obj = json.loads(obj_buffer) - if isinstance(obj, dict): - self._found_valid_json_in_block = True - logger.debug( - f"[Parsed Dict] Keys: {list(obj.keys())}," - " protocol check follows..." - ) - - is_protocol = ( - self._in_top_level_list - and self.is_protocol_msg(obj) - ) - is_comp = obj.get("id") and obj.get("component") - # Process objects at top-level OR items in top-level list - # When in a list, we are top-level if the ONLY thing on the stack is the list opener - is_top_level = ( - len(self._brace_stack) == 0 - ) or ( - self._in_top_level_list - and len(self._brace_stack) == 1 - and self._brace_stack[0][0] == "[" - ) - if is_comp: - self._handle_partial_component( - obj, messages - ) - elif is_top_level or is_protocol: - if not self._handle_complete_object( - obj, self.surface_id, messages - ): - # Not a recognized message type. Validate to catch schema errors. - self._yield_messages([obj], messages) - - if self._brace_count == 0 or ( - self._in_top_level_list - and len(self._brace_stack) == 1 - ): - # Aggressively clear processed objects from the buffer to prevent slowdown. - if ( - len(self._brace_stack) == 1 - and self._brace_stack[0][0] == "[" - ): - # Keep '[' and remove the object after it - self._json_buffer = ( - self._json_buffer[:start_idx] - + self._json_buffer[ - start_idx + len(obj_buffer) : - ] - ) - else: - self._json_buffer = self._json_buffer[ - len(obj_buffer) : - ] - if self._brace_stack: - shift = len(obj_buffer) - self._brace_stack = [ - (b_t, i - shift) - for b_t, i in self._brace_stack - ] - - except json.JSONDecodeError as e: - logger.debug(f"Object recognition failed: {e}") - - elif char == "[": - self._brace_stack.append(("[", len(self._json_buffer))) - self._json_buffer += "[" - self._brace_count += 1 - elif char == "]": - if self._brace_stack and self._brace_stack[-1][0] == "[": - # Pop the typed entry - b_type, start_idx = self._brace_stack.pop() - self._json_buffer += "]" - self._brace_count -= 1 - if self._brace_count == 0: - self._in_top_level_list = False - else: - if self._brace_count > 0: - self._json_buffer += char - - # Sniff for metadata reactively on key delimiters to catch identifiers early - if self._brace_count > 0 and char in ('"', ":", ",", "}", "]"): - self._sniff_metadata() - - # Sniff for partial components at the end of the chunk - if self._brace_count >= 1 and self._json_buffer: - self._sniff_partial_component(messages) - self._sniff_partial_data_model(messages) - - if self._topology_dirty: - self.yield_reachable(messages, check_root=False, raise_on_orphans=False) - self._topology_dirty = False - - def _construct_sniffed_data_model_message( - self, active_msg_type: str, delta_msg_payload: Dict[str, Any] - ) -> Dict[str, Any]: - """Returns the message to yield for a partial data model update.""" - return {active_msg_type: delta_msg_payload} - - def _sniff_partial_data_model(self, messages: List[ResponsePart]) -> None: - msg_type = self._data_model_msg_type - if f'"{msg_type}"' not in self._json_buffer: - return - # Look through the brace stack for objects that might contain data model updates - for b_type, start_idx in reversed(self._brace_stack): - if b_type != "{": - continue - raw_fragment = self._json_buffer[start_idx:] - if not raw_fragment: - continue - - fixed_fragment = self._fix_json(raw_fragment) - obj = None - try: - obj = json.loads(fixed_fragment) - except json.JSONDecodeError: - # Fallback: iteratively strip from the last comma - # This handles cases where _fix_json produces invalid JSON - # from an incomplete trailing element (e.g. `{"key"}` from `{"ke`) - trimmed = raw_fragment - while "," in trimmed: - trimmed = trimmed.rsplit(",", 1)[0] - try: - fixed_trimmed = self._fix_json(trimmed) - if fixed_trimmed: - obj = json.loads(fixed_trimmed) - break - except json.JSONDecodeError: - continue - - if obj and isinstance(obj, dict): - active_msg_type = None - msg_type = self._data_model_msg_type - if msg_type in obj: - active_msg_type = msg_type - - if active_msg_type: - dm_obj = obj[active_msg_type] - if isinstance(dm_obj, dict) and "contents" in dm_obj: - - raw_contents = dm_obj["contents"] - contents_dict = self._parse_contents_to_dict(raw_contents) - - if contents_dict: - delta = {} - for k, v in contents_dict.items(): - if self._yielded_data_model.get(k) != v: - delta[k] = v - - if delta: - sid = ( - dm_obj.get(SURFACE_ID_KEY) - or self._surface_id - or "default" - ) - # Deduplicate delta_contents by only keeping the LATEST entry for each dirty key - delta_contents: Union[ - List[Dict[str, Any]], Dict[str, Any] - ] - if isinstance(raw_contents, list): - list_contents: List[Dict[str, Any]] = [] - seen_keys = set() - for entry in reversed(raw_contents): - if not isinstance(entry, dict): - continue - entry_key = entry.get("key") - # Only include entries that have a valid parsed key (cumulative) - if ( - entry_key - and entry_key in contents_dict - and entry_key not in seen_keys - ): - list_contents.insert(0, entry) - seen_keys.add(entry_key) - delta_contents = ( - self._prune_incomplete_datamodel_entries( - list_contents - ) - ) - else: - delta_contents = delta - - delta_msg_payload = { - SURFACE_ID_KEY: sid, - "contents": delta_contents, - } - if "path" in dm_obj: - delta_msg_payload["path"] = dm_obj["path"] - - delta_msg = self._construct_sniffed_data_model_message( - active_msg_type, delta_msg_payload - ) - self._yield_messages( - [delta_msg], messages, config=RELAXED_VALIDATION - ) - - self._yielded_data_model.update(contents_dict) - # Update internal model for path resolution - self.update_data_model(dm_obj, messages) - - def _sniff_partial_component(self, messages: List[ResponsePart]) -> None: - """Attempts to parse a partial component from the current buffer.""" - # We only care about components if we are inside a "components" array - if f'"{CATALOG_COMPONENTS_KEY}"' not in self._json_buffer: - return - # Try parsing from inner to outer to find the smallest complete component - for b_type, start_idx in reversed(self._brace_stack): - if b_type != "{": - continue - raw_fragment = self._json_buffer[start_idx:] - if not raw_fragment: - continue - - fixed_fragment = self._fix_json(raw_fragment) - try: - obj = json.loads(fixed_fragment) - if isinstance(obj, dict) and obj.get("id") and obj.get("component"): - if isinstance(obj["component"], str): - # Flat style (v0.9+): component type is a string - self._handle_partial_component(obj, messages) - elif ( - isinstance(obj["component"], dict) and len(obj["component"]) > 0 - ): - # Structured style (v0.8): Ignore components that are effectively empty (no type keys) - self._handle_partial_component(obj, messages) - - except Exception: - continue - - def _sniff_metadata(self) -> None: - """Sniffs for surfaceId, root, and msg_types in the current json_buffer.""" - raise NotImplementedError("Subclasses must implement _sniff_metadata") - - def _prune_incomplete_datamodel_entries(self, entries: Any) -> Any: - """Recursively removes data model entries that only contain 'key' and no valid values.""" - if not isinstance(entries, list): - return entries - - pruned = [] - for entry in entries: - if not isinstance(entry, dict): - pruned.append(entry) - continue - - has_val = False - for vkey in ("value", "valueString", "valueNumber", "valueBoolean"): - if vkey in entry: - has_val = True - break - - if "valueMap" in entry: - pruned_map = self._prune_incomplete_datamodel_entries(entry["valueMap"]) - # valueMap is considered valid even if empty, meaning map was explicitly empty - if isinstance(pruned_map, list): - if not pruned_map and len(entry["valueMap"]) > 0: - # If it was non-empty and became empty, it only had incomplete elements. Discard map. - del entry["valueMap"] - else: - entry["valueMap"] = pruned_map - has_val = True - - if has_val and "key" in entry: - pruned.append(entry) - - return pruned - - def _handle_partial_component( - self, comp: Dict[str, Any], messages: List[ResponsePart] - ) -> None: - """Handles a component discovered before its parent message is finished. - - When the parser sees a full JSON object that looks like a component - (contains `id` and `component` keys) within a larger message, it caches - the component and attempts to yield it immediately if it's reachable. - - Args: - comp: The parsed component dictionary. - messages: The list to append any renderable partial messages to. - """ - comp_id = comp.get("id") - if not comp_id: - return - - # Skip caching this component if it has empty dictionaries for complex properties. - # Elements like `children`, `text`, `url`, etc., violate A2UI schema if empty - # and will crash the client. We want the parent to yield a loading placeholder instead. - def _has_empty_dict(obj: Any) -> bool: - if isinstance(obj, dict): - if not obj: - return True - return any(_has_empty_dict(v) for v in obj.values()) - elif isinstance(obj, list): - return any(_has_empty_dict(v) for v in obj) - return False - - component_def = comp.get("component") - if isinstance(component_def, str): - # v0.9 flat style: check the whole component object for empty dicts - if _has_empty_dict(comp): - return - elif _has_empty_dict(component_def): - # v0.8 nested style: check properties inside component - return - - if isinstance(component_def, dict) and hasattr(self, "_required_fields_map"): - comp_type = next(iter(component_def.keys())) if component_def else None - if comp_type: - props = component_def.get(comp_type, {}) - if isinstance(props, dict): - required_fields = self._required_fields_map.get(comp_type, set()) - for req in required_fields: - if req not in props: - return - - self._seen_components[comp_id] = comp - self._topology_dirty = True - - def _parse_contents_to_dict(self, raw_contents: Any) -> Dict[str, Any]: - """Recursively parses a list of A2UI contents into a flat dictionary.""" - if isinstance(raw_contents, dict): - return raw_contents - if not isinstance(raw_contents, list): - return {} - - res = {} - for entry in raw_contents: - if not isinstance(entry, dict): - continue - key = entry.get("key") - val = None - for vkey in ["value", "valueString", "valueNumber", "valueBoolean"]: - if vkey in entry: - val = entry[vkey] - break - - if val is None and "valueMap" in entry: - val = self._parse_contents_to_dict(entry["valueMap"]) - - if key and val is not None: - res[key] = val - return res - - def update_data_model( - self, update: Dict[str, Any], messages: List[ResponsePart] - ) -> None: - """Updates the internal data model and marks affected components as dirty.""" - # Data model update can be v0.8 flat or v0.9+ contents list - raw_contents = update.get("contents") - if raw_contents is not None: - contents = self._parse_contents_to_dict(raw_contents) - else: - # Fallback for old v0.8 flat structure or other layouts - contents = { - k: v - for k, v in update.items() - if k not in (SURFACE_ID_KEY, "root", "contents") - } - - def _handle_complete_object( - self, - obj: Dict[str, Any], - sid: Optional[str], - messages: List[ResponsePart], - ) -> bool: - """Handles an object that has been fully parsed. To be implemented by subclasses.""" - raise NotImplementedError("Subclasses must implement _handle_complete_object") - - def yield_reachable( - self, - messages: List[ResponsePart], - check_root: bool = False, - raise_on_orphans: bool = False, - ) -> None: - """Yields a partial message containing all reachable and seen components. - - This is the core of the streaming logic. Instead of waiting for a UI message - which could contain dozens of components, we yield "partial" updates as soon - as we have enough components to build a renderable sub-tree from the root. - - Args: - messages: The list to which partial messages will be appended. - check_root: If True, raises an error if the root component isn't seen yet. - raise_on_orphans: If True, uses strict topology analysis to catch loops. - """ - active_msg_type = self._get_active_msg_type_for_components() - if not self.root_id or not active_msg_type: - return - - # Buffer components until we have a beginRendering or createSurface for a known surface. - if not self.surface_id: - return - - sid = self.surface_id - if sid not in self._yielded_surfaces_set and not self._buffered_start_message: - return - - try: - # Analyze topology of current seen components - components_to_analyze = list(self._seen_components.values()) - - if check_root and self.root_id not in self._seen_components: - raise A2uiIntegrityError( - f"No root component (id='{self.root_id}') found in" - f" {active_msg_type}" - ) - - reachable_ids = analyze_topology( - components_to_analyze, - self._ref_fields_map, - root_id=self.root_id, - allow_orphan_components=not raise_on_orphans, - ) - - # We only yield components we actually have in our "seen" cache - available_reachable = reachable_ids & set(self._seen_components.keys()) - - if check_root and not available_reachable: - raise A2uiIntegrityError( - f"No root component (id='{self.root_id}') found in" - f" {active_msg_type}" - ) - - # 1. Process placeholders and partial children - processed_components: List[Dict[str, Any]] = [] - extra_components: List[Dict[str, Any]] = [] - surface_id = self.surface_id or "unknown" - yielded_for_surface = self._yielded_ids.get(surface_id, set()) - - for rid in sorted(list(available_reachable)): - comp = copy.deepcopy(self._seen_components[rid]) - # Apply path placeholders and prune unseen children in a single pass - re_yielding = rid in yielded_for_surface - self._process_component_topology( - comp, extra_components, inline_resolved=re_yielding - ) - processed_components.append(comp) - - # Add generated placeholders to the yield - processed_components.extend(extra_components) - - # 2. Check if we have NEW or UPDATED reachable components to yield for THIS surface - surface_id = self.surface_id - if not surface_id or surface_id in self._deleted_surfaces: - return - - should_yield = False - if available_reachable - yielded_for_surface: - should_yield = True - else: - # Check if any yielded component's content has changed for this surface - for comp in processed_components: - cid = comp["id"] - content_str = json.dumps(comp, sort_keys=True) - state_key = (surface_id, cid) - if self._yielded_contents.get(state_key) != content_str: - should_yield = True - break - - if should_yield: - current_sid = self.surface_id or "unknown" - if ( - self._buffered_start_message - and current_sid not in self._yielded_start_messages - ): - self._yield_messages([self._buffered_start_message], messages) - self._yielded_start_messages.add(current_sid) - self._yielded_surfaces_set.add(current_sid) - - # Construct a partial message of the correct type - partial_msg = self._construct_partial_message( - processed_components, active_msg_type - ) - - # Use RELAXED_VALIDATION for partial fragments yielded during streaming - self._yield_messages([partial_msg], messages, config=RELAXED_VALIDATION) - self._yielded_ids.setdefault(surface_id, set()).update( - available_reachable - ) - - # Update content/placeholder tracking - for comp in processed_components: - cid = comp["id"] - self._yielded_contents[(surface_id, cid)] = json.dumps( - comp, sort_keys=True - ) - - except ValueError as e: - if "Circular reference detected" in str(e): - raise e - # Other topology errors (like orphans) are ignored during streaming - # as dependencies might still be on the wire. - msg = str(e) - if ( - raise_on_orphans - or "Circular" in msg - or "Self-reference" in msg - or "recursion" in msg.lower() - or check_root - ): - logger.debug(f"yield_reachable error (strict={check_root}): {msg}") - raise e - - def _get_placeholder_id(self, child_id: str) -> str: - """Returns the ID to use for a missing child placeholder.""" - return f"loading_{child_id}" - - def _process_component_topology( - self, - comp: Dict[str, Any], - extra_components: List[Dict[str, Any]], - inline_resolved: bool = False, - ) -> None: - """Recursively processes path placeholders and child pruning in one pass.""" - comp_id = comp.get("id", "unknown") - - # Deduce the component type for better placeholder typing - comp_type = ( - next(iter(comp.get("component", {}).keys())) - if comp.get("component") and isinstance(comp.get("component"), dict) - else "" - ) - - def traverse(obj: Any, parent_key: Optional[str] = None) -> None: - if isinstance(obj, dict): - # 1. Handle Path Placeholders (from _apply_placeholders) - if ( - "path" in obj - and isinstance(obj["path"], str) - and obj["path"].startswith("/") - ): - path = obj["path"] - key = path.lstrip("/") - if self._version != VERSION_0_9: - if "componentId" not in obj: - obj.clear() - obj.update({"path": "/" + key}) - elif self._version != VERSION_0_9: - # If not in data model, still ensure path has leading slash if it's a bindable object (v0.8 only) - current_path = obj.get("path") - if current_path is not None: - if not isinstance( - current_path, str - ) or not current_path.startswith("/"): - obj["path"] = "/" + str(current_path) - - # 2. Handle Child Pruning (from _prune_unseen_children) - for field in ( - "children", - "explicitList", - "child", - "contentChild", - "entryPointChild", - "componentId", - ): - if field in obj: - if isinstance(obj[field], list): - valid_children = [] - for child_id in obj[field]: - if child_id in self._seen_components: - valid_children.append(child_id) - else: - # Individual placeholder for missing child - placeholder_id = self._get_placeholder_id(child_id) - valid_children.append(placeholder_id) - placeholder_comp = { - "id": placeholder_id, - **self._placeholder_component, - } - # Avoid duplicates in extra_components - if not any( - ec["id"] == placeholder_id - for ec in extra_components - ): - extra_components.append(placeholder_comp) - - if not valid_children and field in ( - "children", - "explicitList", - ): - # If list is empty, check if it was partial in the buffer - # (meaning it's a sequence that started but hasn't yielded items yet) - term = f'"{field}"' - if term in self._json_buffer: - # Simple check: is there a [ after the field name in the buffer? - after_field = self._json_buffer.split(term)[-1] - if ( - "[" in after_field - and "]" not in after_field.split("[")[0] - ): - placeholder_id = f"loading_children_{comp_id}" - valid_children.append(placeholder_id) - placeholder_comp = { - "id": placeholder_id, - **self._placeholder_component, - } - if not any( - ec["id"] == placeholder_id - for ec in extra_components - ): - extra_components.append(placeholder_comp) - obj[field] = valid_children - elif isinstance(obj[field], str): - child_id = obj[field] - if child_id not in self._seen_components: - placeholder_id = self._get_placeholder_id(child_id) - obj[field] = placeholder_id - placeholder_comp = { - "id": placeholder_id, - **self._placeholder_component, - } - if not any( - ec["id"] == placeholder_id - for ec in extra_components - ): - extra_components.append(placeholder_comp) - - # Continue traversal on values - for k, v in list(obj.items()): - traverse(v, parent_key=k) - elif isinstance(obj, list): - for item in obj: - traverse(item, parent_key) - # Start recursion from the component content - if isinstance(comp.get("component"), dict): - traverse(comp.get("component", {})) - else: - # Flat style properties are siblings to 'component' type key - traverse(comp) +__all__ = [ + "A2uiStreamParser", +] diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/prompt/__init__.py b/agent_sdks/python/a2ui_agent/src/a2ui/prompt/__init__.py new file mode 100644 index 000000000..948d05365 --- /dev/null +++ b/agent_sdks/python/a2ui_agent/src/a2ui/prompt/__init__.py @@ -0,0 +1,19 @@ +# 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. + +"""Prompt generation contracts and helpers package.""" + +from .generator import PromptGenerator + +__all__ = ["PromptGenerator"] diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/prompt/generator.py b/agent_sdks/python/a2ui_agent/src/a2ui/prompt/generator.py new file mode 100644 index 000000000..c9959e2c4 --- /dev/null +++ b/agent_sdks/python/a2ui_agent/src/a2ui/prompt/generator.py @@ -0,0 +1,53 @@ +# 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. + +"""Abstract prompt generator interface for inference formats.""" + +from abc import ABC, abstractmethod +from typing import Any, Optional + + +class PromptGenerator(ABC): + """Abstract base class for inference format prompt generators.""" + + @abstractmethod + def generate( + self, + role_description: str, + workflow_description: str = "", + ui_description: str = "", + client_ui_capabilities: Optional[dict[str, Any]] = None, + allowed_components: Optional[list[str]] = None, + allowed_messages: Optional[list[str]] = None, + include_schema: bool = False, + include_examples: bool = False, + validate_examples: bool = False, + ) -> str: + """Assembles prompt instructions contract for standard JSON. + + Args: + role_description: Description of the agent's role. + workflow_description: Optional description of the task workflow. + ui_description: Optional UI context or rules. + client_ui_capabilities: Optional client UI capability details. + allowed_components: Optional list of component tags the LLM may use. + allowed_messages: Optional list of A2UI message types allowed. + include_schema: Whether to include component schemas in the prompt. + include_examples: Whether to include few-shot examples. + validate_examples: Whether to validate few-shot examples on generation. + + Returns: + The complete generated prompt system instruction. + """ + pass diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/schema/catalog.py b/agent_sdks/python/a2ui_agent/src/a2ui/schema/catalog.py index 21d27d20c..69d1e2688 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/schema/catalog.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/schema/catalog.py @@ -20,8 +20,8 @@ import json import logging import os -from dataclasses import dataclass, field, replace -from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING +from dataclasses import dataclass, replace +from typing import Any, Dict, List, Optional, TYPE_CHECKING from urllib.parse import urlparse from a2ui.core.catalog import Catalog from a2ui.core import A2uiCatalogError @@ -39,7 +39,7 @@ ) if TYPE_CHECKING: - from .validator import A2uiValidator + from a2ui.validation.validator import A2uiValidator @dataclass @@ -163,6 +163,7 @@ class A2uiCatalog: common_types_schema: Dict[str, Any] catalog_schema: Dict[str, Any] custom_cuttable_keys: Optional[frozenset[str]] = None + experiments: Optional[frozenset[str]] = None @property def cuttable_keys(self) -> frozenset[str]: @@ -181,9 +182,9 @@ def catalog_id(self) -> str: @property def validator(self) -> "A2uiValidator": - from .validator import A2uiValidator + from a2ui.validation.validator import A2uiValidator - return A2uiValidator(self) + return A2uiValidator(self, experiments=self.experiments) @property def core_catalog(self) -> Catalog[Any, Any]: diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/schema/manager.py b/agent_sdks/python/a2ui_agent/src/a2ui/schema/manager.py index e1362104a..ff212fed7 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/schema/manager.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/schema/manager.py @@ -12,22 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy -import json -import logging -import os -import importlib.resources -from typing import Any, Optional, Callable -from dataclasses import dataclass, field -from .utils import load_from_bundled_resource -from ..inference_strategy import InferenceStrategy -from .constants import * -from .catalog import CatalogConfig, A2uiCatalog -from a2ui.core import A2uiCatalogError +"""Deprecated compatibility redirect for A2uiSchemaManager.""" +import warnings +from typing import Any, Optional, Callable, Union +from a2ui.inference_formats.transport.format import TransportFormat +from a2ui.schema.catalog import CatalogConfig -class A2uiSchemaManager(InferenceStrategy): - """Manages A2UI schema levels and prompt injection.""" +warnings.warn( + "a2ui.schema.manager is deprecated and will be removed. " + "Import from a2ui.inference_formats.transport instead.", + DeprecationWarning, + stacklevel=2, +) + + +class A2uiSchemaManager(TransportFormat): + """Deprecated compatibility wrapper around TransportFormat.""" def __init__( self, @@ -37,213 +38,17 @@ def __init__( schema_modifiers: Optional[ list[Callable[[dict[str, Any]], dict[str, Any]]] ] = None, + experiments: Optional[Union[set[str], frozenset[str]]] = None, ): - self._version = version - self._accepts_inline_catalogs = accepts_inline_catalogs - - self._server_to_client_schema: dict[str, Any] = {} - self._common_types_schema: dict[str, Any] = {} - self._supported_catalogs: list[A2uiCatalog] = [] - self._catalog_example_paths: dict[str, str] = {} - 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_catalog_ids(self) -> list[str]: - return [c.catalog_id for c in 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, - catalogs: Optional[list[CatalogConfig]] = None, - ) -> None: - """Loads separate schema components and processes catalogs.""" - catalogs = catalogs or [] - if version not in SPEC_VERSION_MAP: - raise A2uiCatalogError( - f"Unknown A2UI specification version: {version}. Supported:" - f" {list(SPEC_VERSION_MAP.keys())}" - ) - - # Load server-to-client and common types schemas - self._server_to_client_schema = self._apply_modifiers( - load_from_bundled_resource( - version, SERVER_TO_CLIENT_SCHEMA_KEY, SPEC_VERSION_MAP - ) + warnings.warn( + "A2uiSchemaManager is deprecated. Please use TransportFormat instead.", + DeprecationWarning, + stacklevel=2, ) - self._common_types_schema = self._apply_modifiers( - load_from_bundled_resource( - version, COMMON_TYPES_SCHEMA_KEY, SPEC_VERSION_MAP - ) + super().__init__( + version=version, + catalogs=catalogs, + accepts_inline_catalogs=accepts_inline_catalogs, + schema_modifiers=schema_modifiers, + experiments=experiments, ) - - # Process catalogs - for config in catalogs: - catalog_schema = config.provider.load() - 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, - custom_cuttable_keys=config.custom_cuttable_keys, - ) - self._supported_catalogs.append(catalog) - if config.examples_path: - self._catalog_example_paths[catalog.catalog_id] = config.examples_path - - def _select_catalog( - self, client_ui_capabilities: Optional[dict[str, Any]] = None - ) -> A2uiCatalog: - """Selects the component catalog for the prompt based on client capabilities. - - Selection priority: - 1. If inline catalogs are provided (and accepted by the agent), their - components are merged on top of a base catalog. The base is determined - by supportedCatalogIds (if also provided) or the agent's default catalog. - 2. If only supportedCatalogIds is provided, pick the first mutually - supported catalog. - 3. Fallback to the first agent-supported catalog (usually the bundled catalog). - - Args: - client_ui_capabilities: A dictionary of client UI capabilities, containing - inline catalogs and client-supported catalog IDs. - - Returns: - The resolved A2uiCatalog. - Raises: - ValueError: If inline catalogs are sent but not accepted, or if no - mutually supported catalog is found. - """ - if not self._supported_catalogs: - raise A2uiCatalogError( - "No supported catalogs found." - ) # This should not happen. - - if not client_ui_capabilities or not isinstance(client_ui_capabilities, dict): - return self._supported_catalogs[0] - - inline_catalogs: list[dict[str, Any]] = client_ui_capabilities.get( - INLINE_CATALOGS_KEY, [] - ) - client_supported_catalog_ids: list[str] = client_ui_capabilities.get( - SUPPORTED_CATALOG_IDS_KEY, [] - ) - - if not self._accepts_inline_catalogs and inline_catalogs: - raise A2uiCatalogError( - f"Inline catalog '{INLINE_CATALOGS_KEY}' is provided in client UI" - " capabilities. However, the agent does not accept inline catalogs." - ) - - if inline_catalogs: - # Determine the base catalog: use supportedCatalogIds if provided, - # otherwise fall back to the agent's default catalog. - base_catalog = self._supported_catalogs[0] - if client_supported_catalog_ids: - 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: - base_catalog = agent_supported_catalogs[cscid] - break - - merged_schema = copy.deepcopy(base_catalog.catalog_schema) - - for inline_catalog_schema in inline_catalogs: - inline_catalog_schema = self._apply_modifiers(inline_catalog_schema) - inline_components = inline_catalog_schema.get( - CATALOG_COMPONENTS_KEY, {} - ) - merged_schema[CATALOG_COMPONENTS_KEY].update(inline_components) - - return A2uiCatalog( - version=self._version, - name=INLINE_CATALOG_NAME, - catalog_schema=merged_schema, - s2c_schema=self._server_to_client_schema, - common_types_schema=self._common_types_schema, - ) - - if not client_supported_catalog_ids: - return self._supported_catalogs[0] - - 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 A2uiCatalogError( - "No client-supported catalog found on the agent side. Agent-supported" - f" catalogs are: {[c.catalog_id for c in self._supported_catalogs]}" - ) - - def get_selected_catalog( - self, - client_ui_capabilities: Optional[dict[str, Any]] = None, - allowed_components: Optional[list[str]] = None, - allowed_messages: Optional[list[str]] = None, - ) -> A2uiCatalog: - """Gets the selected catalog after selection and component pruning.""" - catalog = self._select_catalog(client_ui_capabilities) - pruned_catalog = catalog.with_pruning(allowed_components, allowed_messages) - return pruned_catalog - - def load_examples(self, catalog: A2uiCatalog, validate: bool = False) -> str: - """Loads examples for a catalog.""" - if catalog.catalog_id in self._catalog_example_paths: - return catalog.load_examples( - self._catalog_example_paths[catalog.catalog_id], validate=validate - ) - return "" - - def generate_system_prompt( - self, - role_description: str, - workflow_description: str = "", - ui_description: str = "", - client_ui_capabilities: Optional[dict[str, Any]] = None, - allowed_components: Optional[list[str]] = None, - allowed_messages: Optional[list[str]] = None, - include_schema: bool = False, - include_examples: bool = False, - validate_examples: bool = False, - ) -> str: - """Assembles the final system instruction for the LLM.""" - parts = [role_description] - - workflow = DEFAULT_WORKFLOW_RULES - if workflow_description: - workflow += f"\n{workflow_description}" - parts.append(f"## Workflow Description:\n{workflow}") - - if ui_description: - parts.append(f"## UI Description:\n{ui_description}") - - selected_catalog = self.get_selected_catalog( - client_ui_capabilities, allowed_components, allowed_messages - ) - - if include_schema: - parts.append(selected_catalog.render_as_llm_instructions()) - - if include_examples: - examples_str = self.load_examples( - selected_catalog, validate=validate_examples - ) - if examples_str: - parts.append(f"### Examples:\n{examples_str}") - - return "\n\n".join(parts) diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/schema/validator.py b/agent_sdks/python/a2ui_agent/src/a2ui/schema/validator.py index 5a6b070f0..3a7619f21 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/schema/validator.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/schema/validator.py @@ -12,254 +12,28 @@ # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union, Mapping +"""Deprecated compatibility redirect for A2uiValidator.""" -from .constants import VERSION_0_8, VERSION_0_9, VERSION_0_9_1, VERSION_1_0 -from .validator_v08 import ( - LegacyA2uiValidatorV08, - extract_component_required_fields as v08_req, - extract_component_ref_fields as v08_ref, +import warnings +from a2ui.validation.validator import ( + A2uiValidator, + extract_component_ref_fields, + extract_component_required_fields, ) -from a2ui.core.validating import A2uiValidator as CoreValidator +from a2ui.core.validating import analyze_topology from a2ui.core.validating.integrity_checker import get_component_references -from a2ui.core.validating.topology_analyzer import analyze_topology -from a2ui.core.validating.validator import ValidationConfig, STRICT_VALIDATION -from a2ui.core.validating.catalog_schema_validator import CatalogSchemaValidator -from a2ui.core import A2uiValidationError, A2uiCatalogError +warnings.warn( + "a2ui.schema.validator is deprecated and will be removed. " + "Please import from a2ui.validation.validator instead.", + DeprecationWarning, + stacklevel=2, +) -if TYPE_CHECKING: - from .catalog import A2uiCatalog - - -def extract_component_required_fields(catalog: A2uiCatalog) -> Dict[str, Set[str]]: - if catalog.version == VERSION_0_8: - return v08_req(catalog) - cs = catalog.catalog_schema - all_components = cs.get("components", {}) if isinstance(cs, dict) else {} - req_map = {} - for comp_name, comp_schema in all_components.items(): - if isinstance(comp_schema, dict): - reqs = set(comp_schema.get("required", [])) - {"component"} - if reqs: - req_map[comp_name] = reqs - return req_map - - -def extract_component_ref_fields( - catalog: A2uiCatalog, -) -> Mapping[str, Tuple[Set[str], Set[str]]]: - if catalog.version == VERSION_0_8: - return v08_ref(catalog) - result = CatalogSchemaValidator( - catalog.core_catalog, - catalog.common_types_schema, - ).extract_ref_fields() - return result - - -class A2uiValidatorWrapper: - """Validates v0.9+ payloads using a2ui_core.""" - - def __init__(self, catalog: A2uiCatalog): - self._catalog = catalog - self._validator = CoreValidator() - - def validate( - self, - a2ui_json: Union[Dict[str, Any], List[Any]], - root_id: Optional[str] = None, - config: ValidationConfig = STRICT_VALIDATION, - ) -> None: - target_ver = f"v{self._catalog.version}" - updated_config = config.model_copy(update={"target_version": target_ver}) - self._validator.validate( - schema_validator=CatalogSchemaValidator( - self._catalog.core_catalog, - self._catalog.common_types_schema, - ), - a2ui_payload=a2ui_json, - config=updated_config, - ) - - -class A2uiValidatorWrapperV10: - """Validates dynamic payloads (such as v0.9.1 and v1.0) using jsonschema and core component integrity checks.""" - - def __init__(self, catalog: A2uiCatalog): - self._catalog = catalog - from urllib.parse import urljoin - from jsonschema import Draft202012Validator - from referencing import Registry, Resource - - s2c = catalog.s2c_schema - common = catalog.common_types_schema - cat = catalog.catalog_schema - - resources = [] - for schema in [s2c, common]: - if schema and "$id" in schema: - resources.append((schema["$id"], Resource.from_contents(schema))) - - if isinstance(cat, dict): - cat_copy = dict(cat) - if "$schema" not in cat_copy: - cat_copy["$schema"] = "https://json-schema.org/draft/2020-12/schema" - resources.append(("catalog.json", Resource.from_contents(cat_copy))) - s2c_id = s2c.get("$id", "") if s2c else "" - if s2c_id: - resolved_catalog_uri = urljoin(s2c_id, "catalog.json") - cat_copy_uri = dict(cat_copy) - cat_copy_uri["$id"] = resolved_catalog_uri - resources.append( - (resolved_catalog_uri, Resource.from_contents(cat_copy_uri)) - ) - if "$id" in cat: - resources.append((cat["$id"], Resource.from_contents(cat_copy))) - - self._registry = Registry().with_resources(resources) - self._wrapped_schema = { - "$schema": "https://json-schema.org/draft/2020-12/schema", - "type": "array", - "items": {"$ref": s2c["$id"]}, - } - self._schema_validator = Draft202012Validator( - self._wrapped_schema, registry=self._registry - ) - - def validate( - self, - a2ui_json: Union[Dict[str, Any], List[Any]], - root_id: Optional[str] = None, - config: ValidationConfig = STRICT_VALIDATION, - ) -> None: - messages = a2ui_json if isinstance(a2ui_json, list) else [a2ui_json] - - # 1. Run schema validation - errors = list(self._schema_validator.iter_errors(messages)) - if errors: - from a2ui.core import A2uiErrorDetail - - details = [] - for err in errors: - path_str = ".".join(map(str, err.path)) if err.path else "root" - err_validator = getattr(err, "validator", "") - if err_validator == "required": - code = "missing_field" - elif err_validator == "type": - code = "type_mismatch" - elif err_validator == "additionalProperties": - code = "extra_field" - else: - code = "invalid_value" - details.append(A2uiErrorDetail(path_str, code, err.message)) - if err.context: - for sub_error in err.context: - sub_path = ( - ".".join(map(str, sub_error.path)) - if sub_error.path - else path_str - ) - sub_validator = getattr(sub_error, "validator", "") - if sub_validator == "required": - sub_code = "missing_field" - elif sub_validator == "type": - sub_code = "type_mismatch" - elif sub_validator == "additionalProperties": - sub_code = "extra_field" - else: - sub_code = "invalid_value" - details.append( - A2uiErrorDetail(sub_path, sub_code, sub_error.message) - ) - - msg = f"Validation failed: {errors[0].message}" - if errors[0].context: - msg += "\nContext failures:" - for sub_error in errors[0].context: - msg += f"\n - {sub_error.message}" - raise A2uiValidationError(msg, details=details) - - # 2. Run component integrity validation - from a2ui.core.validating.integrity_checker import ( - validate_component_integrity, - validate_recursion_and_paths, - ) - - all_components: list[dict[str, Any]] = [] - for message in messages: - if not isinstance(message, dict): - continue - if "createSurface" in message and isinstance( - message["createSurface"], dict - ): - comps = message["createSurface"].get("components") - if isinstance(comps, list): - all_components.extend(comps) - elif "updateComponents" in message and isinstance( - message["updateComponents"], dict - ): - comps = message["updateComponents"].get("components") - if isinstance(comps, list): - all_components.extend(comps) - - if all_components: - ref_fields = CatalogSchemaValidator( - self._catalog.core_catalog, - self._catalog.common_types_schema, - ).extract_ref_fields() - - validate_component_integrity( - all_components, - ref_fields, - allow_dangling_references=config.allow_dangling_references, - allow_missing_root=config.allow_missing_root, - ) - - validate_recursion_and_paths(messages) - - -class A2uiValidator: - """Version-aware validation facade dispatching to v0.8 or v0.9+ engines.""" - - def __init__(self, catalog: A2uiCatalog): - ver = catalog.version - self.version = ver if isinstance(ver, str) else VERSION_0_8 - if self.version == VERSION_0_8: - self._delegator: Union[ - LegacyA2uiValidatorV08, A2uiValidatorWrapper, A2uiValidatorWrapperV10 - ] = LegacyA2uiValidatorV08(catalog) - # TODO(a2ui-project/A2UI#1936): The V10 validator dynamically uses the `catalog` spec to validate. This should all be consolidated. - elif self.version == VERSION_0_9_1: - self._delegator = A2uiValidatorWrapperV10(catalog) - elif self.version == VERSION_1_0: - import os - - v1_0_enabled = os.environ.get("A2UI_VERSION_1_0", "").lower() in ( - "true", - "1", - "yes", - ) - express_enabled = os.environ.get("A2UI_EXPRESS_ENABLED", "").lower() in ( - "true", - "1", - "yes", - ) - if v1_0_enabled or express_enabled: - self._delegator = A2uiValidatorWrapperV10(catalog) - else: - raise A2uiCatalogError( - "A2UI v1.0 validation is experimental and is disabled by default. " - "To enable it, set the environment variable A2UI_VERSION_1_0=true." - ) - else: - self._delegator = A2uiValidatorWrapper(catalog) - - def validate( - self, - a2ui_json: Union[Dict[str, Any], List[Any]], - root_id: Optional[str] = None, - config: ValidationConfig = STRICT_VALIDATION, - ) -> None: - self._delegator.validate(a2ui_json, root_id=root_id, config=config) +__all__ = [ + "A2uiValidator", + "extract_component_ref_fields", + "extract_component_required_fields", + "analyze_topology", + "get_component_references", +] diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/template/manager.py b/agent_sdks/python/a2ui_agent/src/a2ui/template/manager.py index ccbc9761e..9613d29cc 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/template/manager.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/template/manager.py @@ -12,11 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ..inference_strategy import InferenceStrategy from typing import Optional, Any +from a2ui.inference_format import InferenceFormat -class A2uiTemplateManager(InferenceStrategy): +class A2uiTemplateManager(InferenceFormat): + + @property + def parser(self) -> Any: + raise NotImplementedError("This method is not yet implemented.") def generate_system_prompt( self, diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/validation/validator.py b/agent_sdks/python/a2ui_agent/src/a2ui/validation/validator.py new file mode 100644 index 000000000..ad3c972ee --- /dev/null +++ b/agent_sdks/python/a2ui_agent/src/a2ui/validation/validator.py @@ -0,0 +1,275 @@ +# 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. + +"""Facade validator dispatching to version-specific validation engines.""" + +from __future__ import annotations +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union, Mapping + +from a2ui.schema.constants import VERSION_0_8, VERSION_0_9_1, VERSION_1_0 +from .validator_v08 import ( + LegacyA2uiValidatorV08, + extract_component_required_fields as v08_req, + extract_component_ref_fields as v08_ref, +) +from a2ui.core.validating import A2uiValidator as CoreValidator +from a2ui.core.validating.integrity_checker import get_component_references as get_component_references +from a2ui.core.validating.topology_analyzer import analyze_topology as analyze_topology +from a2ui.core.validating.validator import ValidationConfig, STRICT_VALIDATION +from a2ui.core.validating.catalog_schema_validator import CatalogSchemaValidator +from a2ui.core import A2uiValidationError, A2uiCatalogError + + +if TYPE_CHECKING: + from a2ui.schema.catalog import A2uiCatalog + + +def extract_component_required_fields(catalog: A2uiCatalog) -> Dict[str, Set[str]]: + if catalog.version == VERSION_0_8: + return v08_req(catalog) + cs = catalog.catalog_schema + all_components = cs.get("components", {}) if isinstance(cs, dict) else {} + req_map = {} + for comp_name, comp_schema in all_components.items(): + if isinstance(comp_schema, dict): + reqs = set(comp_schema.get("required", [])) - {"component"} + if reqs: + req_map[comp_name] = reqs + return req_map + + +def extract_component_ref_fields( + catalog: A2uiCatalog, +) -> Mapping[str, Tuple[Set[str], Set[str]]]: + if catalog.version == VERSION_0_8: + return v08_ref(catalog) + result = CatalogSchemaValidator( + catalog.core_catalog, + catalog.common_types_schema, + ).extract_ref_fields() + return result + + +class A2uiValidatorWrapper: + """Validates v0.9+ payloads using a2ui_core.""" + + def __init__(self, catalog: A2uiCatalog): + self._catalog = catalog + self._validator = CoreValidator() + + def validate( + self, + a2ui_json: Union[Dict[str, Any], List[Any]], + root_id: Optional[str] = None, + config: ValidationConfig = STRICT_VALIDATION, + ) -> None: + target_ver = f"v{self._catalog.version}" + updated_config = config.model_copy(update={"target_version": target_ver}) + self._validator.validate( + schema_validator=CatalogSchemaValidator( + self._catalog.core_catalog, + self._catalog.common_types_schema, + ), + a2ui_payload=a2ui_json, + config=updated_config, + ) + + +class A2uiValidatorWrapperV10: + """Validates dynamic payloads (such as v0.9.1 and v1.0) using jsonschema and core component integrity checks.""" + + def __init__(self, catalog: A2uiCatalog): + self._catalog = catalog + from urllib.parse import urljoin + from jsonschema import Draft202012Validator + from referencing import Registry, Resource + + s2c = catalog.s2c_schema + common = catalog.common_types_schema + cat = catalog.catalog_schema + + resources = [] + for schema in [s2c, common]: + if schema and "$id" in schema: + resources.append((schema["$id"], Resource.from_contents(schema))) + + if isinstance(cat, dict): + cat_copy = dict(cat) + if "$schema" not in cat_copy: + cat_copy["$schema"] = "https://json-schema.org/draft/2020-12/schema" + resources.append(("catalog.json", Resource.from_contents(cat_copy))) + s2c_id = s2c.get("$id", "") if s2c else "" + if s2c_id: + resolved_catalog_uri = urljoin(s2c_id, "catalog.json") + cat_copy_uri = dict(cat_copy) + cat_copy_uri["$id"] = resolved_catalog_uri + resources.append( + (resolved_catalog_uri, Resource.from_contents(cat_copy_uri)) + ) + if "$id" in cat: + resources.append((cat["$id"], Resource.from_contents(cat_copy))) + + self._registry = Registry().with_resources(resources) + self._wrapped_schema = { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "array", + "items": {"$ref": s2c["$id"]}, + } + self._schema_validator = Draft202012Validator( + self._wrapped_schema, registry=self._registry + ) + + def validate( + self, + a2ui_json: Union[Dict[str, Any], List[Any]], + root_id: Optional[str] = None, + config: ValidationConfig = STRICT_VALIDATION, + ) -> None: + messages = a2ui_json if isinstance(a2ui_json, list) else [a2ui_json] + + # 1. Run schema validation + errors = list(self._schema_validator.iter_errors(messages)) + if errors: + from a2ui.core import A2uiErrorDetail + + details = [] + for err in errors: + path_str = ".".join(map(str, err.path)) if err.path else "root" + err_validator = getattr(err, "validator", "") + if err_validator == "required": + code = "missing_field" + elif err_validator == "type": + code = "type_mismatch" + elif err_validator == "additionalProperties": + code = "extra_field" + else: + code = "invalid_value" + details.append(A2uiErrorDetail(path_str, code, err.message)) + if err.context: + for sub_error in err.context: + sub_path = ( + ".".join(map(str, sub_error.path)) + if sub_error.path + else path_str + ) + sub_validator = getattr(sub_error, "validator", "") + if sub_validator == "required": + sub_code = "missing_field" + elif sub_validator == "type": + sub_code = "type_mismatch" + elif sub_validator == "additionalProperties": + sub_code = "extra_field" + else: + sub_code = "invalid_value" + details.append( + A2uiErrorDetail(sub_path, sub_code, sub_error.message) + ) + + msg = f"Validation failed: {errors[0].message}" + if errors[0].context: + msg += "\nContext failures:" + for sub_error in errors[0].context: + msg += f"\n - {sub_error.message}" + raise A2uiValidationError(msg, details=details) + + # 2. Run component integrity validation + from a2ui.core.validating.integrity_checker import ( + validate_component_integrity, + validate_recursion_and_paths, + ) + + all_components: list[dict[str, Any]] = [] + for message in messages: + if not isinstance(message, dict): + continue + if "createSurface" in message and isinstance( + message["createSurface"], dict + ): + comps = message["createSurface"].get("components") + if isinstance(comps, list): + all_components.extend(comps) + elif "updateComponents" in message and isinstance( + message["updateComponents"], dict + ): + comps = message["updateComponents"].get("components") + if isinstance(comps, list): + all_components.extend(comps) + + if all_components: + ref_fields = CatalogSchemaValidator( + self._catalog.core_catalog, + self._catalog.common_types_schema, + ).extract_ref_fields() + + validate_component_integrity( + all_components, + ref_fields, + allow_dangling_references=config.allow_dangling_references, + allow_missing_root=config.allow_missing_root, + ) + + validate_recursion_and_paths(messages) + + +class A2uiValidator: + """Version-aware validation facade dispatching to v0.8 or v0.9+ engines.""" + + def __init__( + self, + catalog: A2uiCatalog, + experiments: Optional[Union[set[str], frozenset[str]]] = None, + ): + ver = catalog.version + self.version = ver if isinstance(ver, str) else VERSION_0_8 + + self.experiments = (set(experiments) if experiments else set()) | ( + set(catalog.experiments) if catalog and catalog.experiments else set() + ) + + if self.version == VERSION_0_8: + self._delegator: Union[ + LegacyA2uiValidatorV08, A2uiValidatorWrapper, A2uiValidatorWrapperV10 + ] = LegacyA2uiValidatorV08(catalog) + # TODO(a2ui-project/A2UI#1936): The V10 validator dynamically uses the `catalog` spec to validate. This should all be consolidated. + elif self.version == VERSION_0_9_1: + self._delegator = A2uiValidatorWrapperV10(catalog) + elif self.version == VERSION_1_0: + if "version_1_0" in self.experiments: + self._delegator = A2uiValidatorWrapperV10(catalog) + else: + raise A2uiCatalogError( + "A2UI v1.0 validation is experimental and is disabled by default." + " To enable it, pass the experiment name 'version_1_0' in the" + " experiments configuration." + ) + else: + self._delegator = A2uiValidatorWrapper(catalog) + + def validate( + self, + a2ui_json: Union[Dict[str, Any], List[Any]], + root_id: Optional[str] = None, + config: ValidationConfig = STRICT_VALIDATION, + ) -> None: + """Validates the structure, integrity, and topology of the A2UI message. + + Args: + a2ui_json: The A2UI message payload (object or list of objects). + root_id: The ID of the root component context. + config: The ValidationConfig constraint settings (STRICT or RELAXED). + + Raises: + A2uiValidationError: If any JSON schema, reference integrity, or circular path checks fail. + """ + self._delegator.validate(a2ui_json, root_id=root_id, config=config) diff --git a/agent_sdks/python/a2ui_agent/src/a2ui/schema/validator_v08.py b/agent_sdks/python/a2ui_agent/src/a2ui/validation/validator_v08.py similarity index 98% rename from agent_sdks/python/a2ui_agent/src/a2ui/schema/validator_v08.py rename to agent_sdks/python/a2ui_agent/src/a2ui/validation/validator_v08.py index 7919f8982..64adabb37 100644 --- a/agent_sdks/python/a2ui_agent/src/a2ui/schema/validator_v08.py +++ b/agent_sdks/python/a2ui_agent/src/a2ui/validation/validator_v08.py @@ -15,9 +15,8 @@ from __future__ import annotations import copy -import logging import re -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union, Iterator +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union from jsonschema import Draft202012Validator from a2ui.core.validating.validator import ValidationConfig, STRICT_VALIDATION @@ -28,15 +27,13 @@ from a2ui.core.validating.topology_analyzer import analyze_topology as core_analyze_topology from a2ui.core import A2uiValidationError -from .utils import wrap_as_json_array +from a2ui.schema.utils import wrap_as_json_array if TYPE_CHECKING: - from .catalog import A2uiCatalog + from a2ui.schema.catalog import A2uiCatalog -from .constants import ( - BASE_SCHEMA_URL, +from a2ui.schema.constants import ( CATALOG_COMPONENTS_KEY, - CATALOG_ID_KEY, CATALOG_STYLES_KEY, VERSION_0_8, ) diff --git a/agent_sdks/python/a2ui_agent/tests/adk/a2a/test_event_converter.py b/agent_sdks/python/a2ui_agent/tests/adk/a2a/test_event_converter.py index 58cbd0f6c..2ee1206f7 100644 --- a/agent_sdks/python/a2ui_agent/tests/adk/a2a/test_event_converter.py +++ b/agent_sdks/python/a2ui_agent/tests/adk/a2a/test_event_converter.py @@ -16,7 +16,6 @@ from unittest.mock import MagicMock, patch -import pytest from a2ui.adk.a2a.event_converter import A2uiEventConverter from a2ui.adk.a2a.part_converter import A2uiPartConverter diff --git a/agent_sdks/python/a2ui_agent/tests/adk/a2a/test_part_converter.py b/agent_sdks/python/a2ui_agent/tests/adk/a2a/test_part_converter.py index 5ffa878e7..bfd664ed3 100644 --- a/agent_sdks/python/a2ui_agent/tests/adk/a2a/test_part_converter.py +++ b/agent_sdks/python/a2ui_agent/tests/adk/a2a/test_part_converter.py @@ -17,7 +17,6 @@ import json from unittest.mock import MagicMock, patch -import pytest from a2a import types as a2a_types from a2ui.a2a.parts import create_a2ui_part diff --git a/agent_sdks/python/a2ui_agent/tests/adk/orchestration/test_a2ui_subagent_map.py b/agent_sdks/python/a2ui_agent/tests/adk/orchestration/test_a2ui_subagent_map.py index 25ac3b5c8..5939eb76b 100644 --- a/agent_sdks/python/a2ui_agent/tests/adk/orchestration/test_a2ui_subagent_map.py +++ b/agent_sdks/python/a2ui_agent/tests/adk/orchestration/test_a2ui_subagent_map.py @@ -15,7 +15,6 @@ import unittest from unittest.mock import AsyncMock, patch, MagicMock -from google.adk.sessions.state import State from google.adk.sessions.session import Session from a2ui.adk.orchestration.a2ui_subagent_map import A2uiSubagentMap, SurfaceIdAlreadyExistsError diff --git a/agent_sdks/python/a2ui_agent/tests/adk/test_send_a2ui_to_client_toolset.py b/agent_sdks/python/a2ui_agent/tests/adk/test_send_a2ui_to_client_toolset.py index 679a3054b..818fe1d40 100644 --- a/agent_sdks/python/a2ui_agent/tests/adk/test_send_a2ui_to_client_toolset.py +++ b/agent_sdks/python/a2ui_agent/tests/adk/test_send_a2ui_to_client_toolset.py @@ -21,7 +21,6 @@ from a2ui.schema.catalog import A2uiCatalog from google.adk.agents.readonly_context import ReadonlyContext from google.adk.tools.tool_context import ToolContext -from google.genai import types as genai_types # region SendA2uiToClientToolset Tests """Tests for the SendA2uiToClientToolset class.""" @@ -34,7 +33,7 @@ async def test_toolset_init_bool(): a2ui_enabled=True, a2ui_catalog=catalog_mock, a2ui_examples="examples" ) ctx = MagicMock(spec=ReadonlyContext) - assert await toolset._resolve_a2ui_enabled(ctx) == True + assert await toolset._resolve_a2ui_enabled(ctx) # Access the tool to check schema resolution tool = toolset._ui_tools[0] @@ -52,7 +51,7 @@ async def test_toolset_init_callable(): a2ui_examples=examples_mock, ) ctx = MagicMock(spec=ReadonlyContext) - assert await toolset._resolve_a2ui_enabled(ctx) == True + assert await toolset._resolve_a2ui_enabled(ctx) # Access the tool to check schema resolution tool = toolset._ui_tools[0] @@ -82,7 +81,7 @@ async def async_examples(_ctx): a2ui_examples=async_examples, ) ctx = MagicMock(spec=ReadonlyContext) - assert await toolset._resolve_a2ui_enabled(ctx) == True + assert await toolset._resolve_a2ui_enabled(ctx) # Access the tool to check schema resolution tool = toolset._ui_tools[0] @@ -204,7 +203,7 @@ async def test_send_tool_run_async_valid(): valid_a2ui ) } - assert tool_context_mock.actions.skip_summarization == True + assert tool_context_mock.actions.skip_summarization catalog_mock.validator.validate.assert_called_once_with(valid_a2ui) @@ -230,7 +229,7 @@ async def test_send_tool_run_async_valid_list(): valid_a2ui ) } - assert tool_context_mock.actions.skip_summarization == True + assert tool_context_mock.actions.skip_summarization catalog_mock.validator.validate.assert_called_once_with(valid_a2ui) diff --git a/agent_sdks/python/a2ui_agent/tests/conformance/test_a2a_integration.py b/agent_sdks/python/a2ui_agent/tests/conformance/test_a2a_integration.py index 24cd16865..e3bb5b38d 100644 --- a/agent_sdks/python/a2ui_agent/tests/conformance/test_a2a_integration.py +++ b/agent_sdks/python/a2ui_agent/tests/conformance/test_a2a_integration.py @@ -13,7 +13,6 @@ # limitations under the License. import os -import re import yaml import pytest diff --git a/agent_sdks/python/a2ui_agent/tests/conformance/test_conformance.py b/agent_sdks/python/a2ui_agent/tests/conformance/test_conformance.py index 5afea9312..27b52247f 100644 --- a/agent_sdks/python/a2ui_agent/tests/conformance/test_conformance.py +++ b/agent_sdks/python/a2ui_agent/tests/conformance/test_conformance.py @@ -18,9 +18,10 @@ from a2ui.basic_catalog import BasicCatalog from a2ui.schema.catalog import A2uiCatalog -from a2ui.parser.streaming import A2uiStreamParser -from a2ui.schema.validator import A2uiValidator -from a2ui.schema.manager import A2uiSchemaManager, CatalogConfig +from a2ui.inference_formats.transport.streaming import TransportStreamParser +from a2ui.validation.validator import A2uiValidator +from a2ui.inference_formats.transport import TransportFormat +from a2ui.schema.catalog import CatalogConfig from a2ui.schema.common_modifiers import remove_strict_validation from a2ui.schema.constants import VERSION_0_8, VERSION_0_9 from a2ui.core import ( @@ -182,7 +183,7 @@ def get_conformance_cases(filename): def test_parser_conformance(name, test_case): catalog_config = test_case["catalog"] catalog = setup_catalog(catalog_config) - parser = A2uiStreamParser(catalog=catalog) + parser = TransportStreamParser(catalog=catalog) if test_case.get("disable_validation"): parser._validator = None @@ -357,7 +358,7 @@ def test_schema_manager_conformance(name, test_case): ) ) - manager = A2uiSchemaManager( + transport_format = TransportFormat( version=VERSION_0_9, catalogs=configs, accepts_inline_catalogs=accepts_inline_catalogs, @@ -365,9 +366,9 @@ def test_schema_manager_conformance(name, test_case): if "expect_error" in test_case: with assert_raises(test_case["expect_error"]): - manager.get_selected_catalog(client_capabilities) + transport_format.get_selected_catalog(client_capabilities) else: - selected = manager.get_selected_catalog(client_capabilities) + selected = transport_format.get_selected_catalog(client_capabilities) if "expect_selected" in test_case: assert selected.catalog_id == test_case["expect_selected"] if "expect_catalog_schema" in test_case: @@ -387,17 +388,17 @@ def test_schema_manager_conformance(name, test_case): configs.append( CatalogConfig.from_path(name=cfg["name"], catalog_path=full_path) ) - manager = A2uiSchemaManager( + transport_format = TransportFormat( version=VERSION_0_8, catalogs=configs, schema_modifiers=schema_modifiers ) - selected = manager.get_selected_catalog() + selected = transport_format.get_selected_catalog() expected = test_case["expect"] if "catalog_schema" in expected: assert selected.catalog_schema == expected["catalog_schema"] if "supported_catalog_ids" in expected: - assert [c.catalog_id for c in manager._supported_catalogs] == expected[ - "supported_catalog_ids" - ] + assert [ + c.catalog_id for c in transport_format._supported_catalogs + ] == expected["supported_catalog_ids"] elif action == "generate_prompt": version = args.get("version", VERSION_0_8) @@ -419,13 +420,13 @@ def test_schema_manager_conformance(name, test_case): examples_path=examples_path, ) - manager = A2uiSchemaManager( + transport_format = TransportFormat( version=version, catalogs=[config], accepts_inline_catalogs=args.get("accepts_inline_catalogs", False), ) - output = manager.generate_system_prompt( + output = transport_format.generate_system_prompt( role_description=role, workflow_description=workflow, ui_description=ui_desc, diff --git a/agent_sdks/python/a2ui_agent/tests/elemental/test_compiler.py b/agent_sdks/python/a2ui_agent/tests/elemental/test_compiler.py index 71e3e3a3d..f875328f5 100644 --- a/agent_sdks/python/a2ui_agent/tests/elemental/test_compiler.py +++ b/agent_sdks/python/a2ui_agent/tests/elemental/test_compiler.py @@ -18,7 +18,6 @@ import os import unittest -os.environ["A2UI_EXPRESS_ENABLED"] = "true" from a2ui.core.catalog import Catalog from a2ui.experimental.elemental.compiler import ElementalCompiler diff --git a/agent_sdks/python/a2ui_agent/tests/elemental/test_decompiler.py b/agent_sdks/python/a2ui_agent/tests/elemental/test_decompiler.py index cb261acef..3434d6303 100644 --- a/agent_sdks/python/a2ui_agent/tests/elemental/test_decompiler.py +++ b/agent_sdks/python/a2ui_agent/tests/elemental/test_decompiler.py @@ -18,7 +18,6 @@ import os import unittest -os.environ["A2UI_EXPRESS_ENABLED"] = "true" from a2ui.core.catalog import Catalog from a2ui.experimental.elemental.decompiler import ElementalDecompiler diff --git a/agent_sdks/python/a2ui_agent/tests/express/test_cli_tools.py b/agent_sdks/python/a2ui_agent/tests/express/test_cli_tools.py index ae2e310c2..387cf3435 100644 --- a/agent_sdks/python/a2ui_agent/tests/express/test_cli_tools.py +++ b/agent_sdks/python/a2ui_agent/tests/express/test_cli_tools.py @@ -22,9 +22,6 @@ import tempfile from unittest.mock import patch, MagicMock -# Set up environment variables -os.environ["A2UI_EXPRESS_ENABLED"] = "true" - # Add express proposals directory to sys.path to import run_* scripts EXPRESS_DIR = os.path.abspath( os.path.join( diff --git a/agent_sdks/python/a2ui_agent/tests/express/test_compiler.py b/agent_sdks/python/a2ui_agent/tests/express/test_compiler.py index aacfaf179..ac04e3929 100644 --- a/agent_sdks/python/a2ui_agent/tests/express/test_compiler.py +++ b/agent_sdks/python/a2ui_agent/tests/express/test_compiler.py @@ -18,8 +18,6 @@ import os import unittest -os.environ["A2UI_EXPRESS_ENABLED"] = "true" - from a2ui.core.catalog import Catalog from a2ui.schema.catalog import A2uiCatalog, CatalogConfig from a2ui.experimental.express.prompt_generator import ExpressPromptGenerator @@ -453,49 +451,26 @@ def compile_worker(dsl: str, expected_id: str): def test_v10_validator_gating(self): """Verifies that A2uiValidator gates v1.0 validation behind flags.""" - from a2ui.schema.catalog import CatalogConfig - from a2ui.schema.manager import A2uiSchemaManager - from a2ui.schema.validator import A2uiValidator + from a2ui.inference_formats.transport.format import TransportFormat + from a2ui.validation.validator import A2uiValidator + from a2ui.core import A2uiCatalogError catalog_config = CatalogConfig.from_path("basic_catalog", self.catalog_path) - manager = A2uiSchemaManager(version="1.0", catalogs=[catalog_config]) - catalog = manager.get_selected_catalog() - - orig_express = os.environ.get("A2UI_EXPRESS_ENABLED") - orig_v1_0 = os.environ.get("A2UI_VERSION_1_0") + transport_format = TransportFormat(version="1.0", catalogs=[catalog_config]) + catalog = transport_format.get_selected_catalog() - if "A2UI_EXPRESS_ENABLED" in os.environ: - del os.environ["A2UI_EXPRESS_ENABLED"] - if "A2UI_VERSION_1_0" in os.environ: - del os.environ["A2UI_VERSION_1_0"] + from unittest.mock import patch + import os - try: - with self.assertRaises(ValueError) as context: + # By default, version 1.0 is disabled + with patch.dict(os.environ, {}, clear=True): + with self.assertRaises(A2uiCatalogError) as context: A2uiValidator(catalog) - self.assertIn( - "A2UI v1.0 validation is experimental", str(context.exception) - ) + self.assertIn("A2UI v1.0 validation is experimental", str(context.exception)) - os.environ["A2UI_VERSION_1_0"] = "true" - validator = A2uiValidator(catalog) - self.assertEqual(validator.version, "1.0") - - del os.environ["A2UI_VERSION_1_0"] - - os.environ["A2UI_EXPRESS_ENABLED"] = "true" - validator = A2uiValidator(catalog) - self.assertEqual(validator.version, "1.0") - - finally: - if orig_express is not None: - os.environ["A2UI_EXPRESS_ENABLED"] = orig_express - elif "A2UI_EXPRESS_ENABLED" in os.environ: - del os.environ["A2UI_EXPRESS_ENABLED"] - - if orig_v1_0 is not None: - os.environ["A2UI_VERSION_1_0"] = orig_v1_0 - elif "A2UI_VERSION_1_0" in os.environ: - del os.environ["A2UI_VERSION_1_0"] + # It can be enabled by passing 'version_1_0' in experiments + validator = A2uiValidator(catalog, experiments={"version_1_0"}) + self.assertEqual(validator.version, "1.0") def test_semicolons_and_trailing_commas_and_line_continuation(self): """Verifies that optional semicolons, trailing commas, and line continuations compile correctly.""" diff --git a/agent_sdks/python/a2ui_agent/tests/express/test_decompiler.py b/agent_sdks/python/a2ui_agent/tests/express/test_decompiler.py index e55c4bd5f..a09278372 100644 --- a/agent_sdks/python/a2ui_agent/tests/express/test_decompiler.py +++ b/agent_sdks/python/a2ui_agent/tests/express/test_decompiler.py @@ -19,7 +19,6 @@ import unittest from a2ui.core.catalog import Catalog -os.environ["A2UI_EXPRESS_ENABLED"] = "true" from a2ui.experimental.express.compiler import ExpressCompiler from a2ui.experimental.express.decompiler import ExpressDecompiler diff --git a/agent_sdks/python/a2ui_agent/tests/express/test_integration.py b/agent_sdks/python/a2ui_agent/tests/express/test_integration.py index 831236d4d..18dfaba97 100644 --- a/agent_sdks/python/a2ui_agent/tests/express/test_integration.py +++ b/agent_sdks/python/a2ui_agent/tests/express/test_integration.py @@ -20,7 +20,6 @@ import unittest from typing import Any -os.environ["A2UI_EXPRESS_ENABLED"] = "true" from a2ui.core.catalog import Catalog from a2ui.experimental.express.compiler import ExpressCompiler diff --git a/agent_sdks/python/a2ui_agent/tests/integration/verify_load_real.py b/agent_sdks/python/a2ui_agent/tests/integration/verify_load_real.py index 9907c9a72..f5929514e 100644 --- a/agent_sdks/python/a2ui_agent/tests/integration/verify_load_real.py +++ b/agent_sdks/python/a2ui_agent/tests/integration/verify_load_real.py @@ -14,21 +14,21 @@ import sys -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport.format import TransportFormat from a2ui.schema.constants import CATALOG_COMPONENTS_KEY, VERSION_0_8, VERSION_0_9 from a2ui.schema.common_modifiers import remove_strict_validation from a2ui.basic_catalog.provider import BasicCatalog def verify(): - print('Verifying A2uiSchemaManager...') + print('Verifying TransportFormat...') try: - manager = A2uiSchemaManager( + transport_format = TransportFormat( version=VERSION_0_8, catalogs=[BasicCatalog.get_config(VERSION_0_8)], schema_modifiers=[remove_strict_validation], ) - catalog = manager.get_selected_catalog() + catalog = transport_format.get_selected_catalog() catalog_components = catalog.catalog_schema[CATALOG_COMPONENTS_KEY] print( f'Successfully loaded {VERSION_0_8}: {len(catalog_components)} components' @@ -416,12 +416,12 @@ def verify(): sys.exit(1) try: - manager = A2uiSchemaManager( + transport_format = TransportFormat( version=VERSION_0_9, catalogs=[BasicCatalog.get_config(VERSION_0_9)], schema_modifiers=[remove_strict_validation], ) - catalog = manager.get_selected_catalog() + catalog = transport_format.get_selected_catalog() catalog_components = catalog.catalog_schema[CATALOG_COMPONENTS_KEY] print( f'Successfully loaded {VERSION_0_9}: {len(catalog_components)} components' diff --git a/agent_sdks/python/a2ui_agent/tests/parser/test_streaming_v08.py b/agent_sdks/python/a2ui_agent/tests/parser/test_streaming_v08.py index d8907ff7d..1a461ed2d 100644 --- a/agent_sdks/python/a2ui_agent/tests/parser/test_streaming_v08.py +++ b/agent_sdks/python/a2ui_agent/tests/parser/test_streaming_v08.py @@ -12,25 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json import copy -from unittest.mock import MagicMock import pytest from a2ui.schema.constants import ( A2UI_OPEN_TAG, A2UI_CLOSE_TAG, VERSION_0_8, - SURFACE_ID_KEY, CATALOG_COMPONENTS_KEY, ) from a2ui.parser.constants import ( MSG_TYPE_SURFACE_UPDATE, MSG_TYPE_BEGIN_RENDERING, - MSG_TYPE_DELETE_SURFACE, - MSG_TYPE_DATA_MODEL_UPDATE, ) from a2ui.schema.catalog import A2uiCatalog -from a2ui.parser.streaming import A2uiStreamParser +from a2ui.inference_formats.transport.streaming import TransportStreamParser from a2ui.parser.response_part import ResponsePart @@ -232,7 +227,7 @@ def assertResponseContainsMessages(response, expected_messages): def assertResponseContainsNoA2UI(response): - assert len(response) == 0 or response[0].a2ui_json == None + assert len(response) == 0 or response[0].a2ui_json is None def assertResponseContainsText(response, expected_text): @@ -244,7 +239,7 @@ def assertResponseContainsText(response, expected_text): def test_add_msg_type_deduplication(mock_catalog): - parser = A2uiStreamParser(catalog=mock_catalog) + parser = TransportStreamParser(catalog=mock_catalog) parser.add_msg_type(MSG_TYPE_SURFACE_UPDATE) parser.add_msg_type(MSG_TYPE_SURFACE_UPDATE) assert parser.msg_types == [MSG_TYPE_SURFACE_UPDATE] @@ -256,7 +251,7 @@ def test_add_msg_type_deduplication(mock_catalog): def test_streaming_msg_type_deduplication(mock_catalog): - parser = A2uiStreamParser(catalog=mock_catalog) + parser = TransportStreamParser(catalog=mock_catalog) # 1. Send partial chunk that triggers sniffing chunk1 = A2UI_OPEN_TAG + '[{"surfaceUpdate": {"surfaceId": "s1", "components": [' parser.process_chunk(chunk1) @@ -278,7 +273,7 @@ def test_streaming_msg_type_deduplication(mock_catalog): def test_v08_path_heuristic_adds_slash(mock_catalog): """Tests that v0.8 adds a leading slash to relative paths.""" - parser = A2uiStreamParser(catalog=mock_catalog) + parser = TransportStreamParser(catalog=mock_catalog) # Disable validation for simplicity parser._validator = None diff --git a/agent_sdks/python/a2ui_agent/tests/parser/test_streaming_v09.py b/agent_sdks/python/a2ui_agent/tests/parser/test_streaming_v09.py index 2c3c6e5ba..7c02db4f0 100644 --- a/agent_sdks/python/a2ui_agent/tests/parser/test_streaming_v09.py +++ b/agent_sdks/python/a2ui_agent/tests/parser/test_streaming_v09.py @@ -12,25 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json import copy -from unittest.mock import MagicMock import pytest from a2ui.schema.constants import ( A2UI_OPEN_TAG, A2UI_CLOSE_TAG, VERSION_0_9, - SURFACE_ID_KEY, CATALOG_COMPONENTS_KEY, ) from a2ui.parser.constants import ( MSG_TYPE_CREATE_SURFACE, MSG_TYPE_UPDATE_COMPONENTS, - MSG_TYPE_DELETE_SURFACE, - MSG_TYPE_DATA_MODEL_UPDATE, ) from a2ui.schema.catalog import A2uiCatalog -from a2ui.parser.streaming import A2uiStreamParser +from a2ui.inference_formats.transport.streaming import TransportStreamParser from a2ui.parser.response_part import ResponsePart @@ -344,7 +339,7 @@ def assertResponseContainsMessages(response, expected_messages): def assertResponseContainsNoA2UI(response): - assert len(response) == 0 or response[0].a2ui_json == None + assert len(response) == 0 or response[0].a2ui_json is None def assertResponseContainsText(response, expected_text): @@ -356,7 +351,7 @@ def assertResponseContainsText(response, expected_text): def test_add_msg_type_deduplication(mock_catalog): - parser = A2uiStreamParser(catalog=mock_catalog) + parser = TransportStreamParser(catalog=mock_catalog) parser.add_msg_type(MSG_TYPE_UPDATE_COMPONENTS) parser.add_msg_type(MSG_TYPE_UPDATE_COMPONENTS) assert parser.msg_types == [MSG_TYPE_UPDATE_COMPONENTS] @@ -368,7 +363,7 @@ def test_add_msg_type_deduplication(mock_catalog): def test_streaming_msg_type_deduplication(mock_catalog): - parser = A2uiStreamParser(catalog=mock_catalog) + parser = TransportStreamParser(catalog=mock_catalog) # 1. Send partial chunk that triggers sniffing chunk1 = ( A2UI_OPEN_TAG @@ -392,7 +387,7 @@ def test_streaming_msg_type_deduplication(mock_catalog): def test_v09_path_heuristic_relative_path(mock_catalog): """Tests that v0.9 allows relative paths (no leading slash).""" - parser = A2uiStreamParser(catalog=mock_catalog) + parser = TransportStreamParser(catalog=mock_catalog) # Disable validation to avoid needing full catalog for this test parser._validator = None @@ -426,7 +421,7 @@ def test_v09_path_heuristic_relative_path(mock_catalog): def test_v09_path_heuristic_absolute_path(mock_catalog): """Tests that v0.9 still supports absolute paths (leading slash).""" - parser = A2uiStreamParser(catalog=mock_catalog) + parser = TransportStreamParser(catalog=mock_catalog) parser._validator = None # 1. Create surface @@ -458,7 +453,7 @@ def test_v09_path_heuristic_absolute_path(mock_catalog): def test_v09_single_top_level_object(mock_catalog): """Tests that v0.9 supports a single top-level object without array wrapping.""" - parser = A2uiStreamParser(catalog=mock_catalog) + parser = TransportStreamParser(catalog=mock_catalog) chunk = ( A2UI_OPEN_TAG @@ -476,7 +471,7 @@ def test_v09_single_top_level_object(mock_catalog): def test_v09_multiple_top_level_objects(mock_catalog): """Tests that v0.9 supports multiple consecutive top-level objects without array wrapping.""" - parser = A2uiStreamParser(catalog=mock_catalog) + parser = TransportStreamParser(catalog=mock_catalog) chunk = ( A2UI_OPEN_TAG diff --git a/agent_sdks/python/a2ui_agent/tests/schema/test_catalog.py b/agent_sdks/python/a2ui_agent/tests/schema/test_catalog.py index 4a518a939..c6d362c6f 100644 --- a/agent_sdks/python/a2ui_agent/tests/schema/test_catalog.py +++ b/agent_sdks/python/a2ui_agent/tests/schema/test_catalog.py @@ -12,14 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json -import os import pytest -from typing import Any, Dict, List from a2ui.schema.catalog import A2uiCatalog from a2ui.schema.constants import ( - A2UI_SCHEMA_BLOCK_START, - A2UI_SCHEMA_BLOCK_END, VERSION_0_8, VERSION_0_9, ) @@ -93,7 +88,6 @@ def test_catalog_config_from_path_schemes(): def test_basic_catalog_get_config_examples_path(): from a2ui.basic_catalog.provider import BasicCatalog - from a2ui.schema.constants import VERSION_0_9 # Test get_config with file:// scheme examples path config = BasicCatalog.get_config( diff --git a/agent_sdks/python/a2ui_agent/tests/schema/test_schema_manager.py b/agent_sdks/python/a2ui_agent/tests/schema/test_transport_format.py similarity index 80% rename from agent_sdks/python/a2ui_agent/tests/schema/test_schema_manager.py rename to agent_sdks/python/a2ui_agent/tests/schema/test_transport_format.py index 50b04b6e6..90a4ab4fd 100644 --- a/agent_sdks/python/a2ui_agent/tests/schema/test_schema_manager.py +++ b/agent_sdks/python/a2ui_agent/tests/schema/test_transport_format.py @@ -11,23 +11,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. import io import pytest -import json -import os -from unittest.mock import patch, MagicMock, PropertyMock -from a2ui.schema.manager import A2uiSchemaManager, A2uiCatalog, CatalogConfig +from unittest.mock import patch, MagicMock +from a2ui.inference_formats.transport.format import TransportFormat from a2ui.basic_catalog import BasicCatalog -from a2ui.basic_catalog.constants import BASIC_CATALOG_NAME from a2ui.schema.constants import ( - DEFAULT_WORKFLOW_RULES, - INLINE_CATALOG_NAME, VERSION_0_8, - VERSION_0_9, -) -from a2ui.schema.constants import ( - A2UI_SCHEMA_BLOCK_START, - A2UI_SCHEMA_BLOCK_END, - INLINE_CATALOGS_KEY, - SUPPORTED_CATALOG_IDS_KEY, ) @@ -72,23 +60,23 @@ def joinpath_side_effect(path): mock_traversable.joinpath.side_effect = joinpath_side_effect - manager = A2uiSchemaManager( + transport_format = TransportFormat( VERSION_0_8, catalogs=[BasicCatalog.get_config(VERSION_0_8)] ) - assert manager._server_to_client_schema["defs"] == "server_defs" + assert transport_format._server_to_client_schema["defs"] == "server_defs" # Basic catalog might have a URI-based ID if not explicitly matched # So we check if any catalog exists - assert len(manager._supported_catalogs) >= 1 + assert len(transport_format._supported_catalogs) >= 1 # The first one should be the basic one - catalog = manager._supported_catalogs[0] + catalog = transport_format._supported_catalogs[0] assert catalog.catalog_schema["version"] == VERSION_0_8 assert "Text" in catalog.catalog_schema["components"] def test_schema_manager_init_invalid_version(): with pytest.raises(ValueError, match="Unknown A2UI specification version"): - A2uiSchemaManager("invalid_version") + TransportFormat("invalid_version") def test_schema_manager_fallback_local_assets(mock_importlib_resources): @@ -117,11 +105,11 @@ def open_side_effect(path, *args, **kwargs): mock_open.side_effect = open_side_effect - manager = A2uiSchemaManager( + transport_format = TransportFormat( VERSION_0_8, catalogs=[BasicCatalog.get_config(VERSION_0_8)] ) - assert manager._server_to_client_schema["defs"] == "local_server" - assert len(manager._supported_catalogs) >= 1 - catalog = manager._supported_catalogs[0] + assert transport_format._server_to_client_schema["defs"] == "local_server" + assert len(transport_format._supported_catalogs) >= 1 + catalog = transport_format._supported_catalogs[0] assert "LocalText" in catalog.catalog_schema["components"] diff --git a/agent_sdks/python/a2ui_agent/tests/schema/test_validator.py b/agent_sdks/python/a2ui_agent/tests/schema/test_validator.py index 7ea739014..2e0ef7537 100644 --- a/agent_sdks/python/a2ui_agent/tests/schema/test_validator.py +++ b/agent_sdks/python/a2ui_agent/tests/schema/test_validator.py @@ -12,19 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json -import copy import pytest -from unittest.mock import MagicMock -from a2ui.schema.manager import A2uiSchemaManager, A2uiCatalog, CatalogConfig -from a2ui.schema.common_modifiers import remove_strict_validation +from a2ui.schema.catalog import A2uiCatalog from a2ui.schema.constants import VERSION_0_8, VERSION_0_9 -from a2ui.schema.validator import ( +from a2ui.validation.validator import ( extract_component_ref_fields, analyze_topology, get_component_references, ) -from a2ui.schema.validator_v08 import _find_root_id as find_root_id +from a2ui.validation.validator_v08 import _find_root_id as find_root_id class TestValidator: diff --git a/agent_sdks/python/a2ui_agent/tests/test_formats.py b/agent_sdks/python/a2ui_agent/tests/test_formats.py new file mode 100644 index 000000000..e88e777c2 --- /dev/null +++ b/agent_sdks/python/a2ui_agent/tests/test_formats.py @@ -0,0 +1,97 @@ +# 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 a2ui.schema.catalog import A2uiCatalog +from a2ui.schema.constants import VERSION_0_9 +from a2ui.inference_formats.transport import TransportFormat, TransportParser +from a2ui.adk.a2a.part_converter import A2uiPartConverter +from google.genai import types as genai_types + + +@pytest.fixture +def test_catalog(): + return A2uiCatalog( + version=VERSION_0_9, + name="test_catalog", + s2c_schema={}, + common_types_schema={}, + catalog_schema={ + "catalogId": "https://a2ui.org/test_catalog", + "components": { + "Text": { + "properties": {"text": {"type": "string", "positionalIndex": 0}} + } + }, + "functions": { + "openUrl": { + "properties": {"url": {"type": "string", "positionalIndex": 0}} + } + }, + }, + ) + + +def test_schema_strategy_prompt_generation(test_catalog): + from a2ui.schema.catalog import CatalogConfig + from a2ui.schema.catalog_provider import A2uiCatalogProvider + + class MemoryCatalogProvider(A2uiCatalogProvider): + + def __init__(self, schema): + self.schema = schema + + def load(self): + return self.schema + + config = CatalogConfig( + name="test_catalog", provider=MemoryCatalogProvider(test_catalog.catalog_schema) + ) + + transport_format = TransportFormat(version=VERSION_0_9, catalogs=[config]) + prompt = transport_format.generate_system_prompt( + role_description="You are a helpful assistant.", + workflow_description="Please adhere to constraints.", + include_schema=True, + client_ui_capabilities={ + "supportedCatalogIds": ["https://a2ui.org/test_catalog"] + }, + ) + assert "You are a helpful assistant." in prompt + assert "Please adhere to constraints." in prompt + assert "### Catalog Schema:" in prompt + + +def test_schema_parser(test_catalog): + parser = TransportParser(test_catalog) + parsed = parser.parse_response( + '[{"createSurface": {"surfaceId": "main", "layout": {"component":' + ' "Text"}}}]' + ) + assert len(parsed) == 1 + assert parsed[0].a2ui_json is not None + + +def test_strategy_based_converters(test_catalog, monkeypatch): + monkeypatch.setenv("A2UI_VERSION_1_0", "true") + # Test JSON default (A2uiSchemaParser) + json_converter = A2uiPartConverter(a2ui_catalog=test_catalog) + part_json = genai_types.Part( + text=( + '[{"version": "v0.9", "createSurface": {"surfaceId": "main",' + ' "catalogId": "https://a2ui.org/test_catalog"}}]' + ) + ) + parts_json = json_converter.convert(part_json) + assert len(parts_json) == 1 diff --git a/docs/public/guides/a2ui_over_mcp.md b/docs/public/guides/a2ui_over_mcp.md index 9f9616c24..54cedd7dc 100644 --- a/docs/public/guides/a2ui_over_mcp.md +++ b/docs/public/guides/a2ui_over_mcp.md @@ -396,7 +396,7 @@ pip install a2ui-agent-sdk ``` ```python -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.strategies.schema import A2uiSchemaManager from a2ui.basic_catalog.provider import BasicCatalog # Initialize the schema manager with the basic catalog diff --git a/docs/public/guides/agent-development.md b/docs/public/guides/agent-development.md index b3c30b6fa..562de6baa 100644 --- a/docs/public/guides/agent-development.md +++ b/docs/public/guides/agent-development.md @@ -106,7 +106,7 @@ In your agent file (e.g., `agent.py`), import the necessary classes: ```python from a2ui.schema.constants import VERSION_0_8, VERSION_0_9 -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.strategies.schema import A2uiSchemaManager from a2ui.basic_catalog.provider import BasicCatalog ``` diff --git a/eval/a2ui_eval/scorers.py b/eval/a2ui_eval/scorers.py index 72da393db..7115a6a55 100644 --- a/eval/a2ui_eval/scorers.py +++ b/eval/a2ui_eval/scorers.py @@ -20,7 +20,7 @@ from inspect_ai.scorer import scorer, Score, Target, accuracy, model_graded_qa from inspect_ai.solver import TaskState from inspect_ai.model._model import sample_model_usage -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport.format import TransportFormat from a2ui.schema.catalog import CatalogConfig from a2ui.parser.parser import parse_response from .shared.utils import GIT_ROOT @@ -48,8 +48,8 @@ async def score(state: TaskState, target: Target) -> Score: # pylint: disable=u resolved_catalog_path = str(GIT_ROOT / catalog_path) catalog_config = CatalogConfig.from_path("basic_catalog", resolved_catalog_path) - manager = A2uiSchemaManager(version=version, catalogs=[catalog_config]) - catalog = manager.get_selected_catalog() + transport_format = TransportFormat(version=version, catalogs=[catalog_config]) + catalog = transport_format.get_selected_catalog() validator = catalog.validator answer_text = state.output.completion or "" diff --git a/eval/a2ui_eval/strategies/direct.py b/eval/a2ui_eval/strategies/direct.py index cd8cb9343..feace83e2 100644 --- a/eval/a2ui_eval/strategies/direct.py +++ b/eval/a2ui_eval/strategies/direct.py @@ -14,7 +14,7 @@ from inspect_ai.solver import Solver, solver, TaskState, Generate from inspect_ai.model import ChatMessageSystem -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport.format import TransportFormat from a2ui.schema.catalog import CatalogConfig from ..shared.utils import GIT_ROOT, measured_generate @@ -28,12 +28,12 @@ async def solve(state: TaskState, generate: Generate) -> TaskState: resolved_catalog_path = str(GIT_ROOT / catalog_path) catalog_config = CatalogConfig.from_path('basic_catalog', resolved_catalog_path) - manager = A2uiSchemaManager(version=version, catalogs=[catalog_config]) + transport_format = TransportFormat(version=version, catalogs=[catalog_config]) role_description = state.metadata['role_description'] workflow_description = state.metadata['workflow_description'] - prompt = manager.generate_system_prompt( + prompt = transport_format.generate_system_prompt( role_description=role_description, workflow_description=workflow_description, include_schema=True, diff --git a/eval/a2ui_eval/strategies/elemental.py b/eval/a2ui_eval/strategies/elemental.py index f47581376..ee92124ec 100644 --- a/eval/a2ui_eval/strategies/elemental.py +++ b/eval/a2ui_eval/strategies/elemental.py @@ -15,18 +15,15 @@ """Evaluation strategy for A2UI Elemental.""" import json -import os import re -# Enable experimental extensions -os.environ["A2UI_EXPRESS_ENABLED"] = "true" from inspect_ai.solver import Solver, solver, TaskState, Generate from inspect_ai.model import ChatMessageSystem, ModelOutput, ChatCompletionChoice, ChatMessageAssistant from a2ui.core.catalog import Catalog from a2ui.experimental.elemental.prompt_generator import ElementalPromptGenerator from a2ui.experimental.elemental.parser import parse_elemental_response -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport.format import TransportFormat from a2ui.schema.catalog import CatalogConfig from ..shared.utils import GIT_ROOT, measured_generate @@ -60,10 +57,13 @@ async def solve(state: TaskState, generate: Generate) -> TaskState: catalog_path = state.metadata["catalog"] resolved_catalog_path = str(GIT_ROOT / catalog_path) - # Initialize the catalog schema validator for parsing and validation catalog_config = CatalogConfig.from_path("basic_catalog", resolved_catalog_path) - manager = A2uiSchemaManager(version=version, catalogs=[catalog_config]) - catalog = manager.get_selected_catalog() + transport_format = TransportFormat( + version=version, + catalogs=[catalog_config], + experiments={"version_1_0"} if version == "1.0" else None, + ) + catalog = transport_format.get_selected_catalog() validator = catalog.validator completion = state.output.completion.strip() diff --git a/eval/a2ui_eval/strategies/express.py b/eval/a2ui_eval/strategies/express.py index 61824d8fe..b2b3c8636 100644 --- a/eval/a2ui_eval/strategies/express.py +++ b/eval/a2ui_eval/strategies/express.py @@ -13,17 +13,14 @@ # limitations under the License. import json -import os import re -os.environ["A2UI_EXPRESS_ENABLED"] = "true" from inspect_ai.solver import Solver, solver, TaskState, Generate -from inspect_ai.model import ChatMessageSystem, ModelOutput, ChatCompletionChoice, ChatMessageAssistant, ChatMessageUser +from inspect_ai.model import ChatMessageSystem, ModelOutput, ChatCompletionChoice, ChatMessageAssistant from a2ui.core.catalog import Catalog from a2ui.experimental.express.prompt_generator import ExpressPromptGenerator -from a2ui.experimental.express.compiler import ExpressCompiler from a2ui.experimental.express.parser import parse_express_response -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport.format import TransportFormat from a2ui.schema.catalog import CatalogConfig from ..shared.utils import GIT_ROOT, measured_generate @@ -57,10 +54,13 @@ async def solve(state: TaskState, generate: Generate) -> TaskState: catalog_path = state.metadata["catalog"] resolved_catalog_path = str(GIT_ROOT / catalog_path) - # Initialize the catalog schema validator for parsing catalog_config = CatalogConfig.from_path("basic_catalog", resolved_catalog_path) - manager = A2uiSchemaManager(version=version, catalogs=[catalog_config]) - catalog = manager.get_selected_catalog() + transport_format = TransportFormat( + version=version, + catalogs=[catalog_config], + experiments={"version_1_0"} if version == "1.0" else None, + ) + catalog = transport_format.get_selected_catalog() validator = catalog.validator completion = state.output.completion.strip() diff --git a/eval/a2ui_eval/strategies/subagent_tool.py b/eval/a2ui_eval/strategies/subagent_tool.py index 788bb0869..4aaa0488a 100644 --- a/eval/a2ui_eval/strategies/subagent_tool.py +++ b/eval/a2ui_eval/strategies/subagent_tool.py @@ -14,15 +14,14 @@ import json from inspect_ai.solver import Solver, solver, TaskState, Generate, use_tools, system_message -from inspect_ai.model import ChatMessageSystem, ChatMessageTool, get_model, ModelOutput, ChatCompletionChoice, ChatMessageAssistant, ChatMessageUser +from inspect_ai.model import ChatMessageSystem, get_model, ModelOutput, ChatCompletionChoice, ChatMessageAssistant, ChatMessageUser from inspect_ai.tool import tool, Tool from inspect_ai.util import store -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from a2ui.schema.catalog import CatalogConfig from a2ui.parser.parser import parse_response from ..shared.utils import GIT_ROOT, measured_generate -from .direct import a2ui_system_prompt PAYLOAD_STORE_KEY = "a2ui_payload" @@ -42,12 +41,12 @@ async def execute(input: str) -> str: resolved_catalog_path = str(GIT_ROOT / catalog_path) catalog_config = CatalogConfig.from_path("basic_catalog", resolved_catalog_path) - manager = A2uiSchemaManager(version=version, catalogs=[catalog_config]) + transport_format = TransportFormat(version=version, catalogs=[catalog_config]) role_description = store().get("role_description") workflow_description = store().get("workflow_description") - system_content = manager.generate_system_prompt( + system_content = transport_format.generate_system_prompt( role_description=role_description, workflow_description=workflow_description, include_schema=True, diff --git a/samples/agent/adk/custom-components-example/agent.py b/samples/agent/adk/custom-components-example/agent.py index 67a2cd3ec..f0d2d5904 100644 --- a/samples/agent/adk/custom-components-example/agent.py +++ b/samples/agent/adk/custom-components-example/agent.py @@ -45,9 +45,8 @@ from a2ui.schema.constants import VERSION_0_8, VERSION_0_9, A2UI_OPEN_TAG, A2UI_CLOSE_TAG from a2ui.schema.common_modifiers import remove_strict_validation -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat, TransportStreamParser from a2ui.parser.parser import parse_response, ResponsePart -from a2ui.parser.streaming import A2uiStreamParser from a2ui.basic_catalog.provider import BasicCatalog from a2ui.a2a.extension import get_a2ui_agent_extension from a2ui.a2a.parts import create_a2ui_part, parse_response_to_parts, stream_response_to_parts @@ -68,15 +67,15 @@ def __init__(self, base_url: str): self._build_llm_agent() ) - self._schema_managers: Dict[str, A2uiSchemaManager] = {} + self._inference_formats: Dict[str, TransportFormat] = {} self._ui_runners: Dict[str, Runner] = {} - self._parsers: OrderedDict[str, A2uiStreamParser] = OrderedDict() + self._parsers: OrderedDict[str, TransportStreamParser] = OrderedDict() self._max_parsers = 1000 # Max active sessions to keep in memory for version in [VERSION_0_8, VERSION_0_9]: - schema_manager = self._build_schema_manager(version) - self._schema_managers[version] = schema_manager - agent = self._build_llm_agent(schema_manager) + inference_format = self._build_inference_format(version) + self._inference_formats[version] = inference_format + agent = self._build_llm_agent(inference_format) self._ui_runners[version] = self._build_runner(agent) self._agent_card = self._build_agent_card() @@ -85,13 +84,13 @@ def __init__(self, base_url: str): def agent_card(self) -> AgentCard: return self._agent_card - def get_schema_manager(self, version: Optional[str]) -> Optional[A2uiSchemaManager]: + def get_inference_format(self, version: Optional[str]) -> Optional[TransportFormat]: if version is None: return None - return self._schema_managers[version] + return self._inference_formats[version] - def _build_schema_manager(self, version: str) -> A2uiSchemaManager: - return A2uiSchemaManager( + def _build_inference_format(self, version: str) -> TransportFormat: + return TransportFormat( version=version, catalogs=[ BasicCatalog.get_config( @@ -104,8 +103,8 @@ def _build_schema_manager(self, version: str) -> A2uiSchemaManager: def _build_agent_card(self) -> AgentCard: extensions = [] - if self._schema_managers: - for version, sm in self._schema_managers.items(): + if self._inference_formats: + for version, sm in self._inference_formats.items(): ext = get_a2ui_agent_extension( version, sm.accepts_inline_catalogs, @@ -157,13 +156,13 @@ def get_processing_message(self) -> str: return "Looking up contact information..." def _build_llm_agent( - self, schema_manager: Optional[A2uiSchemaManager] = None + self, inference_format: Optional[TransportFormat] = None ) -> LlmAgent: """Builds the LLM agent for the contact agent.""" LITELLM_MODEL = os.getenv("LITELLM_MODEL", "gemini/gemini-3-flash-preview") instruction = ( - schema_manager.generate_system_prompt( + inference_format.generate_system_prompt( role_description=ROLE_DESCRIPTION, workflow_description=WORKFLOW_DESCRIPTION, ui_description=UI_DESCRIPTION, @@ -171,7 +170,7 @@ def _build_llm_agent( include_schema=True, validate_examples=False, # Missing inline_catalogs for OrgChart and WebFrame validation ) - if schema_manager + if inference_format else get_text_prompt() ) @@ -298,15 +297,15 @@ async def stream( # Determine which runner to use based on whether the a2ui extension is active. if ui_version: runner = self._ui_runners[ui_version] - schema_manager = self._schema_managers[ui_version] + inference_format = self._inference_formats[ui_version] selected_catalog = ( - schema_manager.get_selected_catalog(client_ui_capabilities) - if schema_manager + inference_format.get_selected_catalog(client_ui_capabilities) + if inference_format else None ) else: runner = self._text_runner - schema_manager = None + inference_format = None selected_catalog = None session = await runner.session_service.get_session( @@ -433,7 +432,7 @@ async def token_stream(): if session_id in self._parsers: self._parsers.move_to_end(session_id) else: - self._parsers[session_id] = A2uiStreamParser( + self._parsers[session_id] = TransportStreamParser( catalog=selected_catalog ) if len(self._parsers) > self._max_parsers: diff --git a/samples/agent/adk/custom-components-example/prompt_builder.py b/samples/agent/adk/custom-components-example/prompt_builder.py index 362f1da7f..9cffe04c7 100644 --- a/samples/agent/adk/custom-components-example/prompt_builder.py +++ b/samples/agent/adk/custom-components-example/prompt_builder.py @@ -15,7 +15,8 @@ import json from a2ui.schema.constants import VERSION_0_8, VERSION_0_9, A2UI_OPEN_TAG, A2UI_CLOSE_TAG -from a2ui.schema.manager import A2uiSchemaManager, CatalogConfig +from a2ui.inference_formats.transport import TransportFormat +from a2ui.schema.catalog import CatalogConfig from a2ui.schema.common_modifiers import remove_strict_validation from a2ui.schema.catalog_provider import A2uiCatalogProvider, FileSystemCatalogProvider from a2ui.basic_catalog.provider import BasicCatalog @@ -71,7 +72,7 @@ def get_text_prompt() -> str: my_base_url = "http://localhost:8000" my_version = VERSION_0_9 inline_catalog_path = f"inline_catalog_{my_version}.json" - schema_manager = A2uiSchemaManager( + transport_format = TransportFormat( my_version, catalogs=[ CatalogConfig.from_path( @@ -83,7 +84,7 @@ def get_text_prompt() -> str: accepts_inline_catalogs=True, schema_modifiers=[remove_strict_validation], ) - contact_prompt = schema_manager.generate_system_prompt( + contact_prompt = transport_format.generate_system_prompt( role_description=ROLE_DESCRIPTION, workflow_description=WORKFLOW_DESCRIPTION, ui_description=UI_DESCRIPTION, @@ -100,7 +101,7 @@ def get_text_prompt() -> str: inline_catalog = json.load(f) client_ui_capabilities = {"inlineCatalogs": [inline_catalog]} - inline_catalog = schema_manager.get_selected_catalog( + inline_catalog = transport_format.get_selected_catalog( client_ui_capabilities=client_ui_capabilities, ) request_prompt = inline_catalog.render_as_llm_instructions() diff --git a/samples/agent/adk/restaurant_finder/agent.py b/samples/agent/adk/restaurant_finder/agent.py index abcbeb9e5..a2d82e074 100644 --- a/samples/agent/adk/restaurant_finder/agent.py +++ b/samples/agent/adk/restaurant_finder/agent.py @@ -48,7 +48,7 @@ A2UI_OPEN_TAG, A2UI_CLOSE_TAG, ) -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from a2ui.parser.parser import parse_response, ResponsePart from a2ui.basic_catalog.provider import BasicCatalog from a2ui.schema.common_modifiers import remove_strict_validation @@ -71,15 +71,15 @@ def __init__(self, base_url: str): self._build_llm_agent() ) - self._schema_managers: Dict[str, A2uiSchemaManager] = {} + self._inference_formats: Dict[str, TransportFormat] = {} self._ui_runners: Dict[str, Runner] = {} self._parsers = OrderedDict() self._max_parsers = 1000 # Max active sessions to keep in memory for version in [VERSION_0_8, VERSION_0_9]: - schema_manager = self._build_schema_manager(version) - self._schema_managers[version] = schema_manager - agent = self._build_llm_agent(schema_manager) + inference_format = self._build_inference_format(version) + self._inference_formats[version] = inference_format + agent = self._build_llm_agent(inference_format) self._ui_runners[version] = self._build_runner(agent) self._agent_card = self._build_agent_card() @@ -88,8 +88,8 @@ def __init__(self, base_url: str): def agent_card(self) -> AgentCard: return self._agent_card - def _build_schema_manager(self, version: str) -> A2uiSchemaManager: - return A2uiSchemaManager( + def _build_inference_format(self, version: str) -> TransportFormat: + return TransportFormat( version=version, catalogs=[ BasicCatalog.get_config( @@ -101,8 +101,8 @@ def _build_schema_manager(self, version: str) -> A2uiSchemaManager: def _build_agent_card(self) -> AgentCard: extensions = [] - if self._schema_managers: - for version, sm in self._schema_managers.items(): + if self._inference_formats: + for version, sm in self._inference_formats.items(): ext = get_a2ui_agent_extension( version, sm.accepts_inline_catalogs, @@ -149,7 +149,7 @@ def get_processing_message(self) -> str: return "Finding restaurants that match your criteria..." def _build_llm_agent( - self, schema_manager: Optional[A2uiSchemaManager] = None + self, inference_format: Optional[TransportFormat] = None ) -> LlmAgent: """Builds the LLM agent for the restaurant agent.""" model_env = ( @@ -160,14 +160,14 @@ def _build_llm_agent( model_name = model_env.split("/")[-1] instruction = ( - schema_manager.generate_system_prompt( + inference_format.generate_system_prompt( role_description=ROLE_DESCRIPTION, ui_description=UI_DESCRIPTION, include_schema=True, include_examples=True, validate_examples=True, ) - if schema_manager + if inference_format else get_text_prompt() ) @@ -191,13 +191,13 @@ async def stream( # Determine which runner to use based on whether the a2ui extension is active. if ui_version: runner = self._ui_runners[ui_version] - schema_manager = self._schema_managers[ui_version] + inference_format = self._inference_formats[ui_version] selected_catalog = ( - schema_manager.get_selected_catalog() if schema_manager else None + inference_format.get_selected_catalog() if inference_format else None ) else: runner = self._text_runner - schema_manager = None + inference_format = None selected_catalog = None session = await runner.session_service.get_session( @@ -275,12 +275,12 @@ async def token_stream(): yield p.text if selected_catalog: - from a2ui.parser.streaming import A2uiStreamParser + from a2ui.inference_formats.transport.streaming import TransportStreamParser if session_id in self._parsers: self._parsers.move_to_end(session_id) else: - self._parsers[session_id] = A2uiStreamParser( + self._parsers[session_id] = TransportStreamParser( catalog=selected_catalog ) if len(self._parsers) > self._max_parsers: diff --git a/samples/agent/adk/restaurant_finder/prompt_builder.py b/samples/agent/adk/restaurant_finder/prompt_builder.py index 5c758ca7b..5abf0ab7c 100644 --- a/samples/agent/adk/restaurant_finder/prompt_builder.py +++ b/samples/agent/adk/restaurant_finder/prompt_builder.py @@ -13,7 +13,7 @@ # limitations under the License. from a2ui.schema.constants import VERSION_0_9 -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from a2ui.basic_catalog.provider import BasicCatalog from a2ui.schema.common_modifiers import remove_strict_validation @@ -61,7 +61,7 @@ def get_text_prompt() -> str: # For a different agent (e.g., a flight booker), you would pass in # different examples but use the same `get_ui_prompt` function. version = VERSION_0_9 - restaurant_prompt = A2uiSchemaManager( + restaurant_prompt = TransportFormat( version, catalogs=[ BasicCatalog.get_config( diff --git a/samples/agent/adk/tests/test_examples_validation.py b/samples/agent/adk/tests/test_examples_validation.py index 0bddfb134..86d5978eb 100644 --- a/samples/agent/adk/tests/test_examples_validation.py +++ b/samples/agent/adk/tests/test_examples_validation.py @@ -19,7 +19,8 @@ import pytest from a2ui.schema.constants import VERSION_0_9 -from a2ui.schema.manager import A2uiSchemaManager, CatalogConfig +from a2ui.inference_formats.transport.format import TransportFormat +from a2ui.schema.catalog import CatalogConfig from a2ui.basic_catalog.provider import BasicCatalog from a2ui.schema.common_modifiers import remove_strict_validation from a2ui.schema.catalog_provider import A2uiCatalogProvider @@ -69,7 +70,7 @@ def test_sample_examples_validation(config): sample_path ) # Change to sample dir to resolve relative catalog paths if any - manager = A2uiSchemaManager( + transport_format = TransportFormat( VERSION_0_9, catalogs=config["catalogs"], accepts_inline_catalogs=True, @@ -77,8 +78,8 @@ def test_sample_examples_validation(config): ) # Iterate through each catalog and validate its examples - for catalog in manager._supported_catalogs: - examples_path = manager._catalog_example_paths.get(catalog.catalog_id) + for catalog in transport_format._supported_catalogs: + examples_path = transport_format._catalog_example_paths.get(catalog.catalog_id) if not examples_path: continue diff --git a/samples/community/agent/adk/gemini_enterprise/agent_engine/agent.py b/samples/community/agent/adk/gemini_enterprise/agent_engine/agent.py index 075346d9e..6780984a6 100644 --- a/samples/community/agent/adk/gemini_enterprise/agent_engine/agent.py +++ b/samples/community/agent/adk/gemini_enterprise/agent_engine/agent.py @@ -32,7 +32,7 @@ from a2ui.core.parser.parser import parse_response from a2ui.core.schema.common_modifiers import remove_strict_validation from a2ui.core.schema.constants import A2UI_CLOSE_TAG, A2UI_OPEN_TAG, VERSION_0_8 -from a2ui.core.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat import dotenv from google.adk.agents import run_config from google.adk.agents.llm_agent import LlmAgent @@ -61,14 +61,14 @@ def __init__(self, base_url: str): self._build_llm_agent() ) - self._schema_managers: Dict[str, A2uiSchemaManager] = {} + self._inference_formats: Dict[str, TransportFormat] = {} self._ui_runners: Dict[str, Runner] = {} # Gemini Enerprise only supports VERSION_0_8 for now. for version in [VERSION_0_8]: - schema_manager = self._build_schema_manager(version) - self._schema_managers[version] = schema_manager - agent = self._build_llm_agent(schema_manager) + inference_format = self._build_inference_format(version) + self._inference_formats[version] = inference_format + agent = self._build_llm_agent(inference_format) self._ui_runners[version] = self._build_runner(agent) self._agent_card = self._build_agent_card() @@ -77,9 +77,9 @@ def __init__(self, base_url: str): def agent_card(self) -> AgentCard: return self._agent_card - def _build_schema_manager(self, version: str) -> A2uiSchemaManager: + def _build_inference_format(self, version: str) -> TransportFormat: # Gemini Enerprise only supports VERSION_0_8 for now. - return A2uiSchemaManager( + return TransportFormat( version=version, catalogs=[ BasicCatalog.get_config( @@ -95,8 +95,8 @@ def _build_schema_manager(self, version: str) -> A2uiSchemaManager: def _build_agent_card(self) -> AgentCard: """Builds the AgentCard for this agent, describing its capabilities and skills.""" extensions = [] - if self._schema_managers: - for version, sm in self._schema_managers.items(): + if self._inference_formats: + for version, sm in self._inference_formats.items(): ext = get_a2ui_agent_extension( version, sm.accepts_inline_catalogs, @@ -150,12 +150,12 @@ def get_processing_message(self) -> str: return "Looking up contact information..." def _build_llm_agent( - self, schema_manager: Optional[A2uiSchemaManager] = None + self, inference_format: Optional[TransportFormat] = None ) -> LlmAgent: """Builds the LLM agent for the contact agent.""" instruction = ( - schema_manager.generate_system_prompt( + inference_format.generate_system_prompt( role_description=ROLE_DESCRIPTION, workflow_description=WORKFLOW_DESCRIPTION, ui_description=UI_DESCRIPTION, @@ -163,7 +163,7 @@ def _build_llm_agent( include_examples=True, validate_examples=True, ) - if schema_manager + if inference_format else get_text_prompt() ) @@ -186,9 +186,9 @@ async def fetch_response( # active. if ui_version: runner = self._ui_runners[ui_version] - schema_manager = self._schema_managers[ui_version] + inference_format = self._inference_formats[ui_version] selected_catalog = ( - schema_manager.get_selected_catalog() if schema_manager else None + inference_format.get_selected_catalog() if inference_format else None ) else: runner = self._text_runner diff --git a/samples/community/agent/adk/gemini_enterprise/cloud_run/agent.py b/samples/community/agent/adk/gemini_enterprise/cloud_run/agent.py index f00acb9c2..26a8c7017 100644 --- a/samples/community/agent/adk/gemini_enterprise/cloud_run/agent.py +++ b/samples/community/agent/adk/gemini_enterprise/cloud_run/agent.py @@ -33,7 +33,7 @@ from a2ui.core.parser.parser import parse_response from a2ui.core.schema.common_modifiers import remove_strict_validation from a2ui.core.schema.constants import A2UI_CLOSE_TAG, A2UI_OPEN_TAG, VERSION_0_8 -from a2ui.core.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat import dotenv from google.adk.agents import run_config from google.adk.agents.llm_agent import LlmAgent @@ -65,14 +65,14 @@ def __init__(self, base_url: str): self._build_llm_agent() ) - self._schema_managers: Dict[str, A2uiSchemaManager] = {} + self._inference_formats: Dict[str, TransportFormat] = {} self._ui_runners: Dict[str, Runner] = {} # Gemini Enerprise only supports VERSION_0_8 for now. for version in [VERSION_0_8]: - schema_manager = self._build_schema_manager(version) - self._schema_managers[version] = schema_manager - agent = self._build_llm_agent(schema_manager) + inference_format = self._build_inference_format(version) + self._inference_formats[version] = inference_format + agent = self._build_llm_agent(inference_format) self._ui_runners[version] = self._build_runner(agent) self._agent_card = self._build_agent_card() @@ -81,9 +81,9 @@ def __init__(self, base_url: str): def agent_card(self) -> AgentCard: return self._agent_card - def _build_schema_manager(self, version: str) -> A2uiSchemaManager: + def _build_inference_format(self, version: str) -> TransportFormat: # Gemini Enerprise only supports VERSION_0_8 for now. - return A2uiSchemaManager( + return TransportFormat( version=version, catalogs=[ BasicCatalog.get_config( @@ -99,8 +99,8 @@ def _build_schema_manager(self, version: str) -> A2uiSchemaManager: def _build_agent_card(self) -> AgentCard: """Builds the AgentCard for this agent, describing its capabilities and skills.""" extensions = [] - if self._schema_managers: - for version, sm in self._schema_managers.items(): + if self._inference_formats: + for version, sm in self._inference_formats.items(): ext = get_a2ui_agent_extension( version, sm.accepts_inline_catalogs, @@ -153,12 +153,12 @@ def get_processing_message(self) -> str: return "Looking up contact information..." def _build_llm_agent( - self, schema_manager: Optional[A2uiSchemaManager] = None + self, inference_format: Optional[TransportFormat] = None ) -> LlmAgent: """Builds the LLM agent for the contact agent.""" instruction = ( - schema_manager.generate_system_prompt( + inference_format.generate_system_prompt( role_description=ROLE_DESCRIPTION, workflow_description=WORKFLOW_DESCRIPTION, ui_description=UI_DESCRIPTION, @@ -166,7 +166,7 @@ def _build_llm_agent( include_examples=True, validate_examples=True, ) - if schema_manager + if inference_format else get_text_prompt() ) @@ -188,9 +188,9 @@ async def fetch_response( # Determine which runner to use based on whether the a2ui extension is active. if ui_version: runner = self._ui_runners[ui_version] - schema_manager = self._schema_managers[ui_version] + inference_format = self._inference_formats[ui_version] selected_catalog = ( - schema_manager.get_selected_catalog() if schema_manager else None + inference_format.get_selected_catalog() if inference_format else None ) else: runner = self._text_runner diff --git a/samples/community/agent/adk/mcp_app_proxy/__main__.py b/samples/community/agent/adk/mcp_app_proxy/__main__.py index e7bf56c8e..d8d6981bf 100644 --- a/samples/community/agent/adk/mcp_app_proxy/__main__.py +++ b/samples/community/agent/adk/mcp_app_proxy/__main__.py @@ -16,7 +16,8 @@ from a2a.server.request_handlers import DefaultRequestHandler from a2a.server.tasks import InMemoryTaskStore from a2ui.schema.constants import VERSION_0_8 -from a2ui.schema.manager import A2uiSchemaManager, CatalogConfig +from a2ui.inference_formats.transport import TransportFormat +from a2ui.schema.catalog import CatalogConfig from agent import McpAppProxyAgent from agent_executor import McpAppProxyAgentExecutor, get_a2ui_enabled, get_a2ui_catalog, get_a2ui_examples from dotenv import load_dotenv diff --git a/samples/community/agent/adk/mcp_app_proxy/agent.py b/samples/community/agent/adk/mcp_app_proxy/agent.py index 273323f98..4125fc53c 100644 --- a/samples/community/agent/adk/mcp_app_proxy/agent.py +++ b/samples/community/agent/adk/mcp_app_proxy/agent.py @@ -19,7 +19,9 @@ from a2a.types import AgentCapabilities, AgentCard, AgentSkill from a2ui.a2a.extension import get_a2ui_agent_extension from a2ui.adk.send_a2ui_to_client_toolset import A2uiEnabledProvider, A2uiCatalogProvider, A2uiExamplesProvider, SendA2uiToClientToolset -from a2ui.schema.manager import A2uiSchemaManager, VERSION_0_8, VERSION_0_9, CatalogConfig +from a2ui.inference_formats.transport import TransportFormat +from a2ui.schema.catalog import CatalogConfig +from a2ui.schema.constants import VERSION_0_8, VERSION_0_9 from google.adk.agents.llm_agent import LlmAgent from google.adk.artifacts import InMemoryArtifactService from google.adk.memory.in_memory_memory_service import InMemoryMemoryService @@ -86,13 +88,13 @@ def __init__( self._build_llm_agent() ) - self._schema_managers: Dict[str, A2uiSchemaManager] = {} + self._inference_formats: Dict[str, TransportFormat] = {} self._ui_runners: Dict[str, Runner] = {} for version in [VERSION_0_8, VERSION_0_9]: - schema_manager = self._build_schema_manager(version) - self._schema_managers[version] = schema_manager - agent = self._build_llm_agent(schema_manager) + inference_format = self._build_inference_format(version) + self._inference_formats[version] = inference_format + agent = self._build_llm_agent(inference_format) self._ui_runners[version] = self._build_runner(agent) self._agent_card = self._build_agent_card() @@ -106,13 +108,13 @@ def get_runner(self, version: Optional[str]) -> Runner: return self._text_runner return self._ui_runners[version] - def get_schema_manager(self, version: Optional[str]) -> Optional[A2uiSchemaManager]: + def get_inference_format(self, version: Optional[str]) -> Optional[TransportFormat]: if version is None: return None - return self._schema_managers[version] + return self._inference_formats[version] - def _build_schema_manager(self, version: str) -> A2uiSchemaManager: - return A2uiSchemaManager( + def _build_inference_format(self, version: str) -> TransportFormat: + return TransportFormat( version=version, catalogs=[ CatalogConfig.from_path( @@ -125,8 +127,8 @@ def _build_schema_manager(self, version: str) -> A2uiSchemaManager: def _build_agent_card(self) -> AgentCard: extensions = [] - if self._schema_managers: - for version, sm in self._schema_managers.items(): + if self._inference_formats: + for version, sm in self._inference_formats.items(): ext = get_a2ui_agent_extension( version, sm.accepts_inline_catalogs, @@ -178,11 +180,11 @@ def _build_runner(self, agent: LlmAgent) -> Runner: ) def _build_llm_agent( - self, schema_manager: Optional[A2uiSchemaManager] = None + self, inference_format: Optional[TransportFormat] = None ) -> LlmAgent: """Builds the LLM agent for the contact agent.""" instruction = ( - schema_manager.generate_system_prompt( + inference_format.generate_system_prompt( role_description=ROLE_DESCRIPTION, workflow_description=WORKFLOW_DESCRIPTION, ui_description=UI_DESCRIPTION, @@ -190,7 +192,7 @@ def _build_llm_agent( include_examples=False, validate_examples=False, ) - if schema_manager + if inference_format else "" ) diff --git a/samples/community/agent/adk/mcp_app_proxy/agent_executor.py b/samples/community/agent/adk/mcp_app_proxy/agent_executor.py index 29c300fdf..041bc1528 100644 --- a/samples/community/agent/adk/mcp_app_proxy/agent_executor.py +++ b/samples/community/agent/adk/mcp_app_proxy/agent_executor.py @@ -20,7 +20,7 @@ from a2ui.a2a.extension import try_activate_a2ui_extension from a2ui.adk.a2a.event_converter import A2uiEventConverter from a2ui.schema.constants import A2UI_CLIENT_CAPABILITIES_KEY -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from google.adk.a2a.converters.request_converter import AgentRunRequest from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutorConfig @@ -106,7 +106,7 @@ async def _prepare_session( active_ui_version = try_activate_a2ui_extension(context, self._agent.agent_card) runner = self._agent.get_runner(active_ui_version) - schema_manager = self._agent.get_schema_manager(active_ui_version) + inference_format = self._agent.get_inference_format(active_ui_version) session = await super()._prepare_session(context, run_request, runner) @@ -120,10 +120,10 @@ async def _prepare_session( else None ) a2ui_catalog = ( - schema_manager.get_selected_catalog( + inference_format.get_selected_catalog( client_ui_capabilities=client_capabilities ) - if schema_manager + if inference_format else None ) diff --git a/samples/community/agent/adk/orchestrator/subagent_front_desk.py b/samples/community/agent/adk/orchestrator/subagent_front_desk.py index d53009738..ec0b7c832 100644 --- a/samples/community/agent/adk/orchestrator/subagent_front_desk.py +++ b/samples/community/agent/adk/orchestrator/subagent_front_desk.py @@ -35,16 +35,16 @@ from a2ui.a2a.extension import get_a2ui_agent_extension from a2ui.basic_catalog.provider import BasicCatalog -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from a2ui.adk.a2a.part_converter import A2uiPartConverter from a2ui.schema.common_modifiers import remove_strict_validation -schema_manager = A2uiSchemaManager( +inference_format = TransportFormat( version=VERSION_0_9, catalogs=[BasicCatalog.get_config(version=VERSION_0_9)], schema_modifiers=[remove_strict_validation], ) -my_catalog = schema_manager.get_selected_catalog() +my_catalog = inference_format.get_selected_catalog() a2ui_converter = A2uiPartConverter(a2ui_catalog=my_catalog, version=VERSION_0_9) load_dotenv() diff --git a/samples/community/agent/adk/orchestrator/subagent_housekeeping.py b/samples/community/agent/adk/orchestrator/subagent_housekeeping.py index 8e9ac333c..31bd3a593 100644 --- a/samples/community/agent/adk/orchestrator/subagent_housekeeping.py +++ b/samples/community/agent/adk/orchestrator/subagent_housekeeping.py @@ -35,16 +35,16 @@ from a2ui.a2a.extension import get_a2ui_agent_extension from a2ui.basic_catalog.provider import BasicCatalog -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from a2ui.adk.a2a.part_converter import A2uiPartConverter from a2ui.schema.common_modifiers import remove_strict_validation -schema_manager = A2uiSchemaManager( +inference_format = TransportFormat( version=VERSION_0_9, catalogs=[BasicCatalog.get_config(version=VERSION_0_9)], schema_modifiers=[remove_strict_validation], ) -my_catalog = schema_manager.get_selected_catalog() +my_catalog = inference_format.get_selected_catalog() a2ui_converter = A2uiPartConverter(a2ui_catalog=my_catalog, version=VERSION_0_9) load_dotenv() diff --git a/samples/community/agent/adk/orchestrator/subagent_maintenance.py b/samples/community/agent/adk/orchestrator/subagent_maintenance.py index 652fcab7a..b3d933bdf 100644 --- a/samples/community/agent/adk/orchestrator/subagent_maintenance.py +++ b/samples/community/agent/adk/orchestrator/subagent_maintenance.py @@ -35,16 +35,16 @@ from a2ui.a2a.extension import get_a2ui_agent_extension from a2ui.basic_catalog.provider import BasicCatalog -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from a2ui.adk.a2a.part_converter import A2uiPartConverter from a2ui.schema.common_modifiers import remove_strict_validation -schema_manager = A2uiSchemaManager( +inference_format = TransportFormat( version=VERSION_0_9, catalogs=[BasicCatalog.get_config(version=VERSION_0_9)], schema_modifiers=[remove_strict_validation], ) -my_catalog = schema_manager.get_selected_catalog() +my_catalog = inference_format.get_selected_catalog() a2ui_converter = A2uiPartConverter(a2ui_catalog=my_catalog, version=VERSION_0_9) load_dotenv() diff --git a/samples/community/agent/adk/orchestrator/subagent_room_service.py b/samples/community/agent/adk/orchestrator/subagent_room_service.py index 4076e2d0b..a1454edbb 100644 --- a/samples/community/agent/adk/orchestrator/subagent_room_service.py +++ b/samples/community/agent/adk/orchestrator/subagent_room_service.py @@ -35,16 +35,16 @@ from a2ui.a2a.extension import get_a2ui_agent_extension from a2ui.basic_catalog.provider import BasicCatalog -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from a2ui.adk.a2a.part_converter import A2uiPartConverter from a2ui.schema.common_modifiers import remove_strict_validation -schema_manager = A2uiSchemaManager( +inference_format = TransportFormat( version=VERSION_0_9, catalogs=[BasicCatalog.get_config(version=VERSION_0_9)], schema_modifiers=[remove_strict_validation], ) -my_catalog = schema_manager.get_selected_catalog() +my_catalog = inference_format.get_selected_catalog() a2ui_converter = A2uiPartConverter(a2ui_catalog=my_catalog, version=VERSION_0_9) load_dotenv() diff --git a/samples/community/agent/adk/rizzcharts/python/agent.py b/samples/community/agent/adk/rizzcharts/python/agent.py index 242e512fd..307a2282a 100644 --- a/samples/community/agent/adk/rizzcharts/python/agent.py +++ b/samples/community/agent/adk/rizzcharts/python/agent.py @@ -20,7 +20,8 @@ from a2a.types import AgentCapabilities, AgentCard, AgentSkill from a2ui.a2a.extension import get_a2ui_agent_extension from a2ui.adk.send_a2ui_to_client_toolset import SendA2uiToClientToolset, A2uiEnabledProvider, A2uiCatalogProvider, A2uiExamplesProvider -from a2ui.schema.manager import A2uiSchemaManager, CatalogConfig +from a2ui.inference_formats.transport import TransportFormat +from a2ui.schema.catalog import CatalogConfig from a2ui.basic_catalog.provider import BasicCatalog from a2ui.schema.constants import VERSION_0_8, VERSION_0_9 from google.adk.agents.llm_agent import LlmAgent @@ -117,13 +118,13 @@ def __init__( self._build_llm_agent() ) - self._schema_managers: Dict[str, A2uiSchemaManager] = {} + self._inference_formats: Dict[str, TransportFormat] = {} self._ui_runners: Dict[str, Runner] = {} for version in [VERSION_0_8, VERSION_0_9]: - schema_manager = self._build_schema_manager(version) - self._schema_managers[version] = schema_manager - agent = self._build_llm_agent(schema_manager) + inference_format = self._build_inference_format(version) + self._inference_formats[version] = inference_format + agent = self._build_llm_agent(inference_format) self._ui_runners[version] = self._build_runner(agent) self._agent_card = self._build_agent_card() @@ -137,13 +138,13 @@ def get_runner(self, version: Optional[str]) -> Runner: return self._text_runner return self._ui_runners[version] - def get_schema_manager(self, version: Optional[str]) -> Optional[A2uiSchemaManager]: + def get_inference_format(self, version: Optional[str]) -> Optional[TransportFormat]: if version is None: return None - return self._schema_managers[version] + return self._inference_formats[version] - def _build_schema_manager(self, version: str) -> A2uiSchemaManager: - return A2uiSchemaManager( + def _build_inference_format(self, version: str) -> TransportFormat: + return TransportFormat( version=version, catalogs=[ CatalogConfig.from_path( @@ -172,8 +173,8 @@ def _build_agent_card(self) -> AgentCard: An AgentCard object. """ extensions = [] - if self._schema_managers: - for version, sm in self._schema_managers.items(): + if self._inference_formats: + for version, sm in self._inference_formats.items(): ext = get_a2ui_agent_extension( version, sm.accepts_inline_catalogs, @@ -244,11 +245,11 @@ def _build_runner(self, agent: LlmAgent) -> Runner: ) def _build_llm_agent( - self, schema_manager: Optional[A2uiSchemaManager] = None + self, inference_format: Optional[TransportFormat] = None ) -> LlmAgent: """Builds the LLM agent for the contact agent.""" instruction = ( - schema_manager.generate_system_prompt( + inference_format.generate_system_prompt( role_description=ROLE_DESCRIPTION, workflow_description=WORKFLOW_DESCRIPTION, ui_description=UI_DESCRIPTION, @@ -256,7 +257,7 @@ def _build_llm_agent( include_examples=False, validate_examples=False, ) - if schema_manager + if inference_format else "" ) diff --git a/samples/community/agent/adk/rizzcharts/python/agent_executor.py b/samples/community/agent/adk/rizzcharts/python/agent_executor.py index f5da170e7..02f64ac3f 100644 --- a/samples/community/agent/adk/rizzcharts/python/agent_executor.py +++ b/samples/community/agent/adk/rizzcharts/python/agent_executor.py @@ -22,7 +22,7 @@ from a2ui.adk.a2a.event_converter import A2uiEventConverter from a2ui.adk.send_a2ui_to_client_toolset import SendA2uiToClientToolset from a2ui.schema.constants import A2UI_CLIENT_CAPABILITIES_KEY -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from google.adk.a2a.converters.request_converter import AgentRunRequest from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor @@ -105,7 +105,7 @@ async def _prepare_session( active_ui_version = try_activate_a2ui_extension(context, self._agent.agent_card) runner = self._agent.get_runner(active_ui_version) - schema_manager = self._agent.get_schema_manager(active_ui_version) + inference_format = self._agent.get_inference_format(active_ui_version) session = await super()._prepare_session(context, run_request, runner) @@ -119,14 +119,16 @@ async def _prepare_session( else None ) a2ui_catalog = ( - schema_manager.get_selected_catalog(client_ui_capabilities=capabilities) - if schema_manager + inference_format.get_selected_catalog( + client_ui_capabilities=capabilities + ) + if inference_format else None ) examples = ( - schema_manager.load_examples(a2ui_catalog, validate=True) - if schema_manager + inference_format.load_examples(a2ui_catalog, validate=True) + if inference_format else None ) diff --git a/samples/community/agent/adk/rizzcharts/python/prompt_builder.py b/samples/community/agent/adk/rizzcharts/python/prompt_builder.py index 9eafd6c58..3b7b4774e 100644 --- a/samples/community/agent/adk/rizzcharts/python/prompt_builder.py +++ b/samples/community/agent/adk/rizzcharts/python/prompt_builder.py @@ -16,7 +16,8 @@ # pylint: disable=g-importing-member, line-too-long from a2ui.schema.constants import VERSION_0_9 -from a2ui.schema.manager import A2uiSchemaManager, CatalogConfig +from a2ui.inference_formats.transport import TransportFormat +from a2ui.schema.catalog import CatalogConfig from a2ui.basic_catalog.provider import BasicCatalog from a2ui.schema.common_modifiers import remove_strict_validation from agent import ROLE_DESCRIPTION, WORKFLOW_DESCRIPTION, UI_DESCRIPTION @@ -24,7 +25,7 @@ if __name__ == "__main__": version = VERSION_0_9 - schema_manager = A2uiSchemaManager( + inference_format = TransportFormat( version, catalogs=[ CatalogConfig.from_path( @@ -43,7 +44,7 @@ # Generate prompt for rizzcharts catalog print("Building prompt and validating rizzcharts examples...") - system_prompt = schema_manager.generate_system_prompt( + system_prompt = inference_format.generate_system_prompt( role_description=ROLE_DESCRIPTION, workflow_description=WORKFLOW_DESCRIPTION, ui_description=UI_DESCRIPTION, @@ -57,7 +58,7 @@ # Also validate standard catalog examples print("Validating standard catalog examples...") # We can trigger this by selecting the basic catalog - std_prompt = schema_manager.generate_system_prompt( + std_prompt = inference_format.generate_system_prompt( role_description=ROLE_DESCRIPTION, workflow_description=WORKFLOW_DESCRIPTION, ui_description=UI_DESCRIPTION, diff --git a/samples/community/mcp/a2ui-over-mcp-recipe/server.py b/samples/community/mcp/a2ui-over-mcp-recipe/server.py index ca25e7774..9801b2051 100644 --- a/samples/community/mcp/a2ui-over-mcp-recipe/server.py +++ b/samples/community/mcp/a2ui-over-mcp-recipe/server.py @@ -22,7 +22,7 @@ import mcp.types as types from a2ui.basic_catalog.provider import BasicCatalog from a2ui.schema.constants import VERSION_0_9 -from a2ui.schema.manager import A2uiSchemaManager +from a2ui.inference_formats.transport import TransportFormat from mcp.server.lowlevel import Server from mcp.server.lowlevel.helper_types import ReadResourceContents from starlette.requests import Request @@ -49,10 +49,10 @@ ) def main(port: int, transport: str, bypass_verification: bool) -> int: # Initialize schema manager and validate sample - schema_manager = A2uiSchemaManager( + inference_format = TransportFormat( version=VERSION_0_9, catalogs=[BasicCatalog.get_config(version=VERSION_0_9)] ) - selected_catalog = schema_manager.get_selected_catalog() + selected_catalog = inference_format.get_selected_catalog() recipe_a2ui_json = json.loads( (pathlib.Path(__file__).resolve().parent / "recipe_a2ui.json").read_text()