From 8324b6d30441062a52291483930c4195d709ecc0 Mon Sep 17 00:00:00 2001 From: Jeremy Wortz Date: Wed, 20 May 2026 00:56:13 -0500 Subject: [PATCH] feat(agent_sdk): add A2UIOutputMode enum for unified prompt generation Introduces A2UIOutputMode.TEXT and A2UIOutputMode.TOOL to generate_system_prompt(), preventing accidental double-injection of schemas when using SendA2uiToClientToolset. Addresses #1289 --- .../python/src/a2ui/schema/constants.py | 12 ++ agent_sdks/python/src/a2ui/schema/manager.py | 44 +++++- .../python/src/a2ui/schema/output_mode.py | 35 +++++ .../tests/schema/test_schema_manager.py | 142 ++++++++++++++++++ 4 files changed, 231 insertions(+), 2 deletions(-) create mode 100644 agent_sdks/python/src/a2ui/schema/output_mode.py diff --git a/agent_sdks/python/src/a2ui/schema/constants.py b/agent_sdks/python/src/a2ui/schema/constants.py index ed65b3b07d..476d5209e6 100644 --- a/agent_sdks/python/src/a2ui/schema/constants.py +++ b/agent_sdks/python/src/a2ui/schema/constants.py @@ -64,3 +64,15 @@ - Parent components MUST appear before their child components. This specific ordering allows the streaming parser to yield and render the UI incrementally as it arrives. """ + +TOOL_WORKFLOW_RULES = """ +The generated response MUST follow these rules: +- To send UI to the user, call the `send_a2ui_json_to_client` tool with valid A2UI JSON. +- Each response SHOULD include at least one tool call to render UI components. +- You may include brief conversational text alongside tool calls for context. +- The JSON payload passed to the tool MUST be a single, raw JSON object (usually a list of A2UI messages) and MUST validate against the provided A2UI JSON SCHEMA. +- Top-Down Component Ordering: Within the `components` list of a message: + - The 'root' component MUST be the FIRST element. + - Parent components MUST appear before their child components. + This specific ordering allows the streaming parser to yield and render the UI incrementally as it arrives. +""" diff --git a/agent_sdks/python/src/a2ui/schema/manager.py b/agent_sdks/python/src/a2ui/schema/manager.py index 8e8ea07f9b..faae46ed19 100644 --- a/agent_sdks/python/src/a2ui/schema/manager.py +++ b/agent_sdks/python/src/a2ui/schema/manager.py @@ -22,8 +22,11 @@ from .utils import load_from_bundled_resource from ..inference_strategy import InferenceStrategy from .constants import * +from .output_mode import A2UIOutputMode from .catalog import CatalogConfig, A2uiCatalog +logger = logging.getLogger(__name__) + class A2uiSchemaManager(InferenceStrategy): """Manages A2UI schema levels and prompt injection.""" @@ -209,11 +212,48 @@ def generate_system_prompt( include_schema: bool = False, include_examples: bool = False, validate_examples: bool = False, + output_mode: A2UIOutputMode = A2UIOutputMode.TEXT, ) -> str: - """Assembles the final system instruction for the LLM.""" + """Assembles the final system instruction for the LLM. + + Args: + role_description: A description of the agent's role. + workflow_description: Additional workflow rules appended to the base + rules. + ui_description: A description of the UI the agent should generate. + client_ui_capabilities: A dictionary of client UI capabilities, used + for catalog selection. + allowed_components: An optional list of component names to include in + the prompt. If None, all components are included. + allowed_messages: An optional list of message names to include in the + prompt. If None, all messages are included. + include_schema: Whether to include the JSON schema in the prompt. + include_examples: Whether to include examples in the prompt. + validate_examples: Whether to validate examples against the schema. + output_mode: Controls how A2UI instructions are delivered. TEXT (default) + injects the schema into the system prompt. TOOL skips schema and + example injection (handled by SendA2uiToClientToolset) and uses + tool-oriented workflow rules. + + Returns: + The assembled system prompt string. + """ parts = [role_description] - workflow = DEFAULT_WORKFLOW_RULES + if output_mode == A2UIOutputMode.TOOL: + if include_schema or include_examples: + logger.warning( + "output_mode=TOOL overrides include_schema and include_examples " + "to False. Schema and examples are injected by " + "SendA2uiToClientToolset.process_llm_request() instead." + ) + include_schema = False + include_examples = False + base_rules = TOOL_WORKFLOW_RULES + else: + base_rules = DEFAULT_WORKFLOW_RULES + + workflow = base_rules if workflow_description: workflow += f"\n{workflow_description}" parts.append(f"## Workflow Description:\n{workflow}") diff --git a/agent_sdks/python/src/a2ui/schema/output_mode.py b/agent_sdks/python/src/a2ui/schema/output_mode.py new file mode 100644 index 0000000000..aadb993bff --- /dev/null +++ b/agent_sdks/python/src/a2ui/schema/output_mode.py @@ -0,0 +1,35 @@ +# 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. + +"""Output mode configuration for A2UI system prompt generation.""" + +from enum import Enum + + +class A2UIOutputMode(Enum): + """Controls how A2UI schema and instructions are delivered to the LLM. + + Attributes: + TEXT: Schema is injected directly into the system prompt via + generate_system_prompt(). The LLM outputs A2UI JSON wrapped in + tags as part of its text response. This is the default. + TOOL: Schema is provided via SendA2uiToClientToolset's + process_llm_request(). The LLM calls the send_a2ui_json_to_client + tool to deliver A2UI JSON. When using this mode, + generate_system_prompt() automatically skips schema and example + injection to prevent double-injection. + """ + + TEXT = "text" + TOOL = "tool" diff --git a/agent_sdks/python/tests/schema/test_schema_manager.py b/agent_sdks/python/tests/schema/test_schema_manager.py index 6b7b2fdbb3..1c6040c581 100644 --- a/agent_sdks/python/tests/schema/test_schema_manager.py +++ b/agent_sdks/python/tests/schema/test_schema_manager.py @@ -19,6 +19,7 @@ from a2ui.basic_catalog.constants import BASIC_CATALOG_NAME from a2ui.schema.constants import ( DEFAULT_WORKFLOW_RULES, + TOOL_WORKFLOW_RULES, INLINE_CATALOG_NAME, VERSION_0_8, VERSION_0_9, @@ -29,6 +30,7 @@ INLINE_CATALOGS_KEY, SUPPORTED_CATALOG_IDS_KEY, ) +from a2ui.schema.output_mode import A2UIOutputMode @pytest.fixture @@ -125,3 +127,143 @@ def open_side_effect(path, *args, **kwargs): assert len(manager._supported_catalogs) >= 1 catalog = manager._supported_catalogs[0] assert "LocalText" in catalog.catalog_schema["components"] + + +# --- Tests for output_mode parameter --- + + +class TestOutputMode: + """Tests for the output_mode parameter on generate_system_prompt().""" + + @pytest.fixture + def manager(self): + return A2uiSchemaManager( + VERSION_0_8, + catalogs=[BasicCatalog.get_config(VERSION_0_8)], + ) + + def test_default_mode_is_text(self, manager): + """Default output_mode is TEXT, uses DEFAULT_WORKFLOW_RULES.""" + prompt = manager.generate_system_prompt( + role_description="Test agent", + ) + assert "send_a2ui_json_to_client" not in prompt + + def test_text_mode_includes_schema(self, manager): + """TEXT mode includes schema when include_schema=True.""" + prompt = manager.generate_system_prompt( + role_description="Test agent", + include_schema=True, + output_mode=A2UIOutputMode.TEXT, + ) + assert "BEGIN A2UI" in prompt or "A2UI JSON SCHEMA" in prompt + + def test_tool_mode_excludes_schema(self, manager): + """TOOL mode forces include_schema=False even when requested.""" + prompt = manager.generate_system_prompt( + role_description="Test agent", + include_schema=True, + output_mode=A2UIOutputMode.TOOL, + ) + assert "BEGIN A2UI" not in prompt + assert "A2UI JSON SCHEMA" not in prompt + + def test_tool_mode_excludes_examples(self, manager): + """TOOL mode forces include_examples=False even when requested.""" + prompt = manager.generate_system_prompt( + role_description="Test agent", + include_examples=True, + output_mode=A2UIOutputMode.TOOL, + ) + assert "Examples:" not in prompt + + def test_tool_mode_uses_tool_rules(self, manager): + """TOOL mode uses TOOL_WORKFLOW_RULES.""" + prompt = manager.generate_system_prompt( + role_description="Test agent", + output_mode=A2UIOutputMode.TOOL, + ) + assert "send_a2ui_json_to_client" in prompt + assert "call the" in prompt + + def test_text_mode_uses_default_rules(self, manager): + """TEXT mode uses DEFAULT_WORKFLOW_RULES.""" + prompt = manager.generate_system_prompt( + role_description="Test agent", + output_mode=A2UIOutputMode.TEXT, + ) + assert "send_a2ui_json_to_client" not in prompt + + def test_tool_mode_preserves_role_description(self, manager): + """TOOL mode still includes the role description.""" + role = "You are a data analytics dashboard agent." + prompt = manager.generate_system_prompt( + role_description=role, + output_mode=A2UIOutputMode.TOOL, + ) + assert role in prompt + + def test_tool_mode_preserves_ui_description(self, manager): + """TOOL mode still includes ui_description.""" + ui_desc = "Render KPI metrics as Card grids." + prompt = manager.generate_system_prompt( + role_description="Test agent", + ui_description=ui_desc, + output_mode=A2UIOutputMode.TOOL, + ) + assert ui_desc in prompt + + def test_tool_mode_preserves_workflow_description(self, manager): + """TOOL mode appends workflow_description after tool rules.""" + workflow = "Always show last 7 days of data." + prompt = manager.generate_system_prompt( + role_description="Test agent", + workflow_description=workflow, + output_mode=A2UIOutputMode.TOOL, + ) + assert workflow in prompt + assert "send_a2ui_json_to_client" in prompt + + def test_backward_compatible_no_output_mode(self, manager): + """Not passing output_mode produces identical output to before.""" + prompt_before = manager.generate_system_prompt( + role_description="Test agent", + ) + prompt_explicit = manager.generate_system_prompt( + role_description="Test agent", + output_mode=A2UIOutputMode.TEXT, + ) + assert prompt_before == prompt_explicit + + def test_tool_mode_differs_from_text(self, manager): + """TOOL mode produces different output than TEXT mode.""" + text_prompt = manager.generate_system_prompt( + role_description="Test agent", + output_mode=A2UIOutputMode.TEXT, + ) + tool_prompt = manager.generate_system_prompt( + role_description="Test agent", + output_mode=A2UIOutputMode.TOOL, + ) + assert text_prompt != tool_prompt + + def test_tool_rules_has_top_down_ordering(self): + """TOOL_WORKFLOW_RULES preserves the top-down ordering requirement.""" + assert "root" in TOOL_WORKFLOW_RULES + assert ( + "Parent components MUST appear before their child" in TOOL_WORKFLOW_RULES + ) + + def test_tool_mode_with_v09(self): + """TOOL mode works with v0.9 catalogs.""" + manager = A2uiSchemaManager( + VERSION_0_9, + catalogs=[BasicCatalog.get_config(VERSION_0_9)], + ) + prompt = manager.generate_system_prompt( + role_description="Test agent", + include_schema=True, + output_mode=A2UIOutputMode.TOOL, + ) + assert "send_a2ui_json_to_client" in prompt + assert "BEGIN A2UI" not in prompt