Integrate parser into send_a2ui_to_client_toolset.py

The previous commit defines a default workflow rule, which defines the
LLM response with both a text part and an A2UI JSON part, separated by
a delimiter. The default workflow rule is injected into the system
prompt, which changes the rizzcharts response format.

This commit updates the send_a2ui_to_client_toolset.py to handle the new
response correctly by leveraging the parser.

It also updates the payload fixer to not depend on the A2UICatalog and
moves it outside of the schema package. Now the payload fixer only
parses the raw JSON string and tries to fix it by removing the trailing
comma. The parser can use the payload fixes to extract the JSON object.

Author: nan-yu

agent_sdks/agent_development.md agent_sdks/python/agent_development.mdagent_sdks/agent_development.md renamed to agent_sdks/python/agent_development.md

@@ -99,18 +99,16 @@ agent_executor = MyAgentExecutor(
 To ensure reliability, always validate the LLM's JSON output before returning it. The SDK's `A2uiCatalog` provides a validator that checks the payload against the A2UI schema. If the payload is invalid, the validator will attempt to fix it.
 
 ```python
+from a2ui.core.parser import parse_response
+
 # Get the catalog for the current request
 selected_catalog = schema_manager.get_selected_catalog()
 
-try:
-    # Parse the LLM's JSON part
-    parsed_json = json.loads(json_string)
-    
-    # Validate and fix against the schema
-    selected_catalog.payload_fixer.validate_and_fix(parsed_json)
-except Exception as e:
-    # Handle validation errors (e.g., log error or retry with correction prompt)
-    print(f"Validation failed: {e}")
+# Parse the LLM's JSON part with simple fixers like removing trailing commas
+text_part, json_data = parse_response(text)
+
+# Validate the JSON part against the schema
+selected_catalog.validator.validate(json_data)
 ```
 
 #### 4c. Stream the A2UI Payload
@@ -203,8 +201,9 @@ ui_toolset = SendA2uiToClientToolset(
 #### 2c. Runtime Validation
 
 When the LLM calls the UI tool, the toolset uses the dynamic catalog to:
-1.  **Generate Instructions**: Automatically inject the specific schema and examples into the LLM's system prompt for that turn.
-2.  **Validate and Fix Payloads**: Automatically validate and fix the LLM's generated JSON against the specific `A2uiCatalog` object's validator and auto-fixer.
+1.  **Generate Instructions**: Inject the specific schema and examples into the LLM's system prompt for that turn.
+2.  **Parse and Fix Payloads**: Parse and fix the LLM's generated JSON using the parser and payload-fixer.
+3.  **Validate Payloads**: Validate the LLM's generated JSON against the specific `A2uiCatalog` object's validator.
 
 ### 3. Orchestration and Delegation
 

agent_sdks/python/src/a2ui/adk/a2a_extension/send_a2ui_to_client_toolset.py

@@ -28,18 +28,20 @@
     that effectively sends a JSON payload to the client. This tool validates the JSON against
     the provided schema. It automatically wraps the provided schema in an array structure,
     instructing the LLM that it can send a list of UI items.
-  * `convert_send_a2ui_to_client_genai_part_to_a2a_part`: A utility function that intercepts the `send_a2ui_json_to_client`
-    tool calls from the LLM and converts them into `a2a_types.Part` objects, which are then
-    returned by the A2A Agent Executor.
+  * `A2uiEventConverter`: An event converter that automatically injects the A2UI catalog into part conversion.
 
 Usage Examples:
 
   1. Defining Providers:
-    You can use simple values or callables (sync or async) for enablement and schema.
+    You can use simple values or callables (sync or async) for enablement, catalog schema, and examples.
 
     ```python
     # Simple boolean and dict
-    toolset = SendA2uiToClientToolset(a2ui_enabled=True, a2ui_catalog=MY_CATALOG)
+    toolset = SendA2uiToClientToolset(
+        a2ui_enabled=True,
+        a2ui_catalog=MY_CATALOG,
+        a2ui_examples=MY_EXAMPLES,
+    )
 
     # Async providers
     async def check_enabled(ctx: ReadonlyContext) -> bool:
@@ -48,7 +50,14 @@ async def check_enabled(ctx: ReadonlyContext) -> bool:
     async def get_catalog(ctx: ReadonlyContext) -> A2uiCatalog:
       return await fetch_catalog(ctx)
 
-    toolset = SendA2uiToClientToolset(a2ui_enabled=check_enabled, a2ui_catalog=get_catalog)
+    async def get_examples(ctx: ReadonlyContext) -> str:
+      return await fetch_examples(ctx)
+
+    toolset = SendA2uiToClientToolset(
+        a2ui_enabled=check_enabled,
+        a2ui_catalog=get_catalog,
+        a2ui_examples=get_examples,
+    )
     ```
 
   2. Integration with Agent:
@@ -59,10 +68,11 @@ async def get_catalog(ctx: ReadonlyContext) -> A2uiCatalog:
     LlmAgent(
         tools=[
             SendA2uiToClientToolset(
-                a2ui_enabled=True,
-                a2ui_catalog=MY_CATALOG
-            )
-        ]
+                a2ui_enabled=check_enabled,
+                a2ui_catalog=get_catalog,
+                a2ui_examples=get_examples,
+            ),
+        ],
     )
     ```
 
@@ -71,7 +81,7 @@ async def get_catalog(ctx: ReadonlyContext) -> A2uiCatalog:
 
     ```python
     config = A2aAgentExecutorConfig(
-        genai_part_converter=convert_send_a2ui_to_client_genai_part_to_a2a_part
+        event_converter=A2uiEventConverter()
     )
     executor = A2aAgentExecutor(config)
     ```
@@ -81,13 +91,23 @@ async def get_catalog(ctx: ReadonlyContext) -> A2uiCatalog:
 import json
 import logging
 import re
-from typing import Any, Awaitable, Callable, Optional, TypeAlias, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Optional,
+    TypeAlias,
+    Union,
+)
 
 import jsonschema
 
 from a2a import types as a2a_types
 from ...a2a import create_a2ui_part
+from a2ui.core.parser import parse_response, parse_and_fix
 from a2ui.core.schema.catalog import A2uiCatalog
+from a2ui.core.schema.constants import A2UI_DELIMITER
 from google.adk.a2a.converters import part_converter
 from google.adk.agents.readonly_context import ReadonlyContext
 from google.adk.models import LlmRequest
@@ -97,6 +117,12 @@ async def get_catalog(ctx: ReadonlyContext) -> A2uiCatalog:
 from google.adk.utils.feature_decorator import experimental
 from google.genai import types as genai_types
 
+if TYPE_CHECKING:
+  from a2a.server.events import Event as A2AEvent
+  from google.adk.a2a.converters.part_converter import GenAIPartToA2APartConverter
+  from google.adk.agents.invocation_context import InvocationContext
+  from google.adk.events.event import Event
+
 logger = logging.getLogger(__name__)
 
 A2uiEnabledProvider: TypeAlias = Callable[
@@ -163,6 +189,18 @@ async def get_tools(
       logger.info("A2UI is DISABLED, not adding ui tools")
       return []
 
+  async def get_part_converter(self, ctx: ReadonlyContext) -> "A2uiPartConverter":
+    """Returns a configured A2uiPartConverter for the given context.
+
+    Args:
+        ctx: The ReadonlyContext to resolve the catalog with.
+
+    Returns:
+        A configured A2uiPartConverter.
+    """
+    catalog = await self._ui_tools[0]._resolve_a2ui_catalog(ctx)
+    return A2uiPartConverter(catalog)
+
   class _SendA2uiJsonToClientTool(BaseTool):
     TOOL_NAME = "send_a2ui_json_to_client"
     VALIDATED_A2UI_JSON_KEY = "validated_a2ui_json"
@@ -267,7 +305,8 @@ async def run_async(
           )
 
         a2ui_catalog = await self._resolve_a2ui_catalog(tool_context)
-        a2ui_json_payload = a2ui_catalog.payload_fixer.validate_and_fix(a2ui_json)
+        a2ui_json_payload = parse_and_fix(a2ui_json)
+        a2ui_catalog.validator.validate(a2ui_json_payload)
 
         logger.info(
             f"Validated call to tool {self.TOOL_NAME} with {self.A2UI_JSON_ARG_NAME}"
@@ -288,49 +327,133 @@ async def run_async(
 
 
 @experimental
-def convert_send_a2ui_to_client_genai_part_to_a2a_part(
-    part: genai_types.Part,
-) -> list[a2a_types.Part]:
-  if (
-      (function_response := part.function_response)
-      and function_response.name
-      == SendA2uiToClientToolset._SendA2uiJsonToClientTool.TOOL_NAME
-  ):
+class A2uiPartConverter:
+  """A catalog-aware GenAI to A2A part converter.
+
+  This converter handles both tool-based A2UI (via `send_a2ui_json_to_client`)
+  and text-based A2UI (via A2UI delimiter tags). It uses the provided
+  catalog to validate and fix JSON payloads.
+  """
+
+  def __init__(self, a2ui_catalog: A2uiCatalog):
+    self._catalog = a2ui_catalog
+
+  def convert(self, part: genai_types.Part) -> list[a2a_types.Part]:
+    """Converts a GenAI part to A2A parts, with A2UI validation.
+
+    Args:
+        part: The GenAI part to convert.
+
+    Returns:
+        A list of A2A parts.
+    """
+    # 1. Handle Tool Responses (FunctionResponse)
     if (
-        SendA2uiToClientToolset._SendA2uiJsonToClientTool.TOOL_ERROR_KEY
-        in function_response.response
+        (function_response := part.function_response)
+        and function_response.name
+        == SendA2uiToClientToolset._SendA2uiJsonToClientTool.TOOL_NAME
     ):
-      logger.warning(
-          "A2UI tool call failed:"
-          f" {function_response.response[SendA2uiToClientToolset._SendA2uiJsonToClientTool.TOOL_ERROR_KEY]}"
+      if (
+          SendA2uiToClientToolset._SendA2uiJsonToClientTool.TOOL_ERROR_KEY
+          in function_response.response
+      ):
+        logger.warning(
+            "A2UI tool call failed:"
+            f" {function_response.response[SendA2uiToClientToolset._SendA2uiJsonToClientTool.TOOL_ERROR_KEY]}"
+        )
+        return []
+
+      # The tool returns the list of messages directly on success
+      json_data = function_response.response.get(
+          SendA2uiToClientToolset._SendA2uiJsonToClientTool.VALIDATED_A2UI_JSON_KEY
       )
-      return []
+      if not json_data:
+        logger.info("No result in A2UI tool response")
+        return []
 
-    # The tool returns the list of messages directly on success
-    json_data = function_response.response.get(
-        SendA2uiToClientToolset._SendA2uiJsonToClientTool.VALIDATED_A2UI_JSON_KEY
-    )
-    if not json_data:
-      logger.info("No result in A2UI tool response")
+      return [create_a2ui_part(message) for message in json_data]
+
+    # 2. Handle Tool Calls (FunctionCall) - Skip sending to client
+    if (
+        (function_call := part.function_call)
+        and function_call.name
+        == SendA2uiToClientToolset._SendA2uiJsonToClientTool.TOOL_NAME
+    ):
       return []
 
-    final_parts = []
-    for message in json_data:
+    # 3. Handle Text-based A2UI (TextPart)
+    if text := part.text:
+      if A2UI_DELIMITER in text:
+        return self._convert_text_with_a2ui(text)
+
+    # 4. Default conversion for other parts
+    converted_part = part_converter.convert_genai_part_to_a2a_part(part)
+    return [converted_part] if converted_part else []
+
+  def _convert_text_with_a2ui(self, text: str) -> list[a2a_types.Part]:
+    """Helper to split text and extract/validate A2UI JSON."""
+    parts = []
+    try:
+      text_part, json_data = parse_response(text)
+      self._catalog.validator.validate(json_data)
+
+      if text_part:
+        parts.append(
+            a2a_types.Part(root=a2a_types.TextPart(kind="text", text=text_part))
+        )
+
       logger.info(f"Found {len(json_data)} messages. Creating individual DataParts.")
-      final_parts.append(create_a2ui_part(message))
+      for message in json_data:
+        parts.append(create_a2ui_part(message))
+
+    except Exception as e:
+      logger.error(f"Failed to parse or validate text-based A2UI JSON: {e}")
+      # Fallback: at least try to return the leading text part if we can split it
+      if not parts:
+        segments = text.split(A2UI_DELIMITER, 1)
+        if segments[0].strip():
+          parts.append(
+              a2a_types.Part(
+                  root=a2a_types.TextPart(kind="text", text=segments[0].strip())
+              )
+          )
 
-    return final_parts
+    return parts
 
-  # Don't send a2ui tool call to client
-  elif (
-      (function_call := part.function_call)
-      and function_call.name
-      == SendA2uiToClientToolset._SendA2uiJsonToClientTool.TOOL_NAME
-  ):
-    return []
 
-  # Use default part converter for other types (images, etc)
-  converted_part = part_converter.convert_genai_part_to_a2a_part(part)
+@experimental
+class A2uiEventConverter:
+  """An event converter that automatically injects the A2UI catalog into part conversion.
+
+  This allows text-based A2UI extraction and validation to work even when the
+  catalog is session-specific.
+  """
 
-  logger.info(f"Returning converted part: {converted_part}")
-  return [converted_part] if converted_part else []
+  def __init__(self, catalog_key: str = "system:a2ui_catalog"):
+    self._catalog_key = catalog_key
+
+  def __call__(
+      self,
+      event: "Event",
+      invocation_context: "InvocationContext",
+      task_id: Optional[str] = None,
+      context_id: Optional[str] = None,
+      part_converter_func: "GenAIPartToA2APartConverter" = part_converter.convert_genai_part_to_a2a_part,
+  ) -> list["A2AEvent"]:
+    """Converts an ADK event to A2A events, using the session catalog if available."""
+    from google.adk.a2a.converters.event_converter import convert_event_to_a2a_events
+
+    catalog = invocation_context.session.state.get(self._catalog_key)
+    if catalog:
+      # Use the catalog-aware part converter
+      effective_converter = A2uiPartConverter(catalog).convert
+    else:
+      effective_converter = part_converter_func
+
+    return convert_event_to_a2a_events(
+        event,
+        invocation_context,
+        task_id,
+        context_id,
+        effective_converter,
+    )

agent_sdks/python/src/a2ui/core/parser/__init__.py

@@ -13,5 +13,6 @@
 # limitations under the License.
 
 from .parser import parse_response
+from .payload_fixer import parse_and_fix
 
-__all__ = ["parse_response"]
+__all__ = ["parse_response", "parse_and_fix"]

agent_sdks/python/src/a2ui/core/parser/parser.py

@@ -15,6 +15,7 @@
 import json
 from typing import Tuple, Any
 from ..schema.constants import A2UI_DELIMITER
+from .payload_fixer import parse_and_fix
 
 
 def parse_response(content: str) -> Tuple[str, Any]:
@@ -47,10 +48,5 @@ def parse_response(content: str) -> Tuple[str, Any]:
   if not json_string_cleaned:
     raise ValueError("A2UI JSON part is empty.")
 
-  try:
-    parsed_json_data = json.loads(json_string_cleaned)
-    return text_part, parsed_json_data
-  except json.JSONDecodeError as e:
-    # Truncate string if it's too long to keep logs readable
-    snippet = (json_string[:100] + "...") if len(json_string) > 100 else json_string
-    raise ValueError(f"Failed to parse JSON: {snippet}") from e
+  json_data = parse_and_fix(json_string_cleaned)
+  return text_part, json_data