feat(python-sdk): Implement A2UI Templates and Sandbox Prototype#1761
feat(python-sdk): Implement A2UI Templates and Sandbox Prototype#1761jacobsimionato wants to merge 3 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a prototype implementation of A2UI Templates in the Python Agent SDK, featuring a template expansion engine (TemplateProcessor), pre-baked JSON templates, a FastAPI demo server integrating Gemini, a React frontend client, and associated tests. The reviewer identified several critical issues that need to be addressed: broken imports and references to stream_queue in the integration tests due to a refactoring to active_clients, a FastAPI startup RuntimeError caused by an invalid CORS configuration, a logical bug in expand_template where ID mapping is performed before parameter substitution, and a limitation in the map_surface_id helper which does not recursively traverse dictionary values.
| import pytest | ||
| from fastapi.testclient import TestClient | ||
|
|
||
| from a2ui.experimental.templates.demo.server import app, stream_queue |
There was a problem hiding this comment.
The import of stream_queue is broken because server.py was refactored to use active_clients for multi-client pub/sub. This will cause an immediate ImportError when running the tests. Import active_clients and asyncio instead.
| from a2ui.experimental.templates.demo.server import app, stream_queue | |
| from a2ui.experimental.templates.demo.server import app, active_clients | |
| import asyncio |
| # Empty any stale queue elements | ||
| while not stream_queue.empty(): | ||
| stream_queue.get_nowait() |
There was a problem hiding this comment.
Replace the non-existent stream_queue with a dummy asyncio.Queue registered in active_clients to capture the broadcasted messages. Also, remember to update stream_queue to client_queue on lines 82 and 84.
| # Empty any stale queue elements | |
| while not stream_queue.empty(): | |
| stream_queue.get_nowait() | |
| client_queue = asyncio.Queue() | |
| active_clients.clear() | |
| active_clients.append(client_queue) |
| while not stream_queue.empty(): | ||
| stream_queue.get_nowait() | ||
|
|
There was a problem hiding this comment.
Initialize a dummy client_queue and register it in active_clients instead of using the non-existent stream_queue. Also, remember to update stream_queue to client_queue on lines 129 and 133.
| while not stream_queue.empty(): | |
| stream_queue.get_nowait() | |
| client_queue = asyncio.Queue() | |
| active_clients.clear() | |
| active_clients.append(client_queue) |
| app.add_middleware( | ||
| CORSMiddleware, | ||
| allow_origins=["*"], | ||
| allow_credentials=True, | ||
| allow_headers=["*"], | ||
| allow_methods=["*"], | ||
| ) |
There was a problem hiding this comment.
In FastAPI/Starlette, configuring CORSMiddleware with allow_origins=["*"] and allow_credentials=True will raise a RuntimeError at startup or during request processing. If credentials are not required for the sandbox, set allow_credentials=False. Otherwise, specify explicit origins.
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_headers=["*"], | |
| allow_methods=["*"], | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=False, | |
| allow_headers=["*"], | |
| allow_methods=["*"], | |
| ) |
| for comp in template.components: | ||
| comp_copy = dict(comp) | ||
|
|
||
| # Map the component's own ID | ||
| comp_copy["id"] = map_id(comp_copy["id"]) | ||
|
|
||
| # Map parent-child structural references | ||
| if "child" in comp_copy: | ||
| comp_copy["child"] = map_id(comp_copy["child"]) | ||
| if "children" in comp_copy: | ||
| # Handle static unrolling if children contains a static loop template | ||
| if isinstance(comp_copy["children"], dict) and "param" in comp_copy["children"] and "template" in comp_copy["children"]: | ||
| children_meta = comp_copy["children"] | ||
| param_name = children_meta["param"] | ||
| item_template_id = children_meta["template"] | ||
|
|
||
| array_data = params.get(param_name, []) | ||
| if not isinstance(array_data, list): | ||
| raise ValueError(f"Parameter '{param_name}' must be an array/list for static loop unrolling.") | ||
|
|
||
| unrolled_child_ids = [] | ||
| for idx, item in enumerate(array_data): | ||
| sub_instance_id = f"{comp_copy['id']}_item_{idx}" | ||
| unrolled_child_ids.append(sub_instance_id) | ||
|
|
||
| sub_params = item if isinstance(item, dict) else {"value": item} | ||
| # Recursively expand the child item template | ||
| sub_expanded = self.expand_template(sub_instance_id, item_template_id, sub_params) | ||
| expanded_components.extend(sub_expanded) | ||
|
|
||
| comp_copy["children"] = unrolled_child_ids | ||
| else: | ||
| comp_copy["children"] = map_child_list(comp_copy["children"]) | ||
|
|
||
| # Substitute parameters dynamically | ||
| comp_final = substitute_params(comp_copy, params) | ||
| expanded_components.append(comp_final) |
There was a problem hiding this comment.
In expand_template, ID mapping and namespacing are performed before parameter substitution. If any structural fields (such as id, child, or componentId in children) are defined dynamically via parameter references (e.g., {"param": "..."}), map_id will be called with a dictionary instead of a string. This causes incorrect string serialization of the dictionary and prevents substitute_params from recognizing and substituting the parameter later. Performing parameter substitution before ID mapping resolves this issue cleanly, ensuring all parameters are fully resolved to strings before namespacing is applied. Additionally, avoid using defensive .get() with fallback values when accessing params for the parameter name, as its presence is a strict invariant; raising a KeyError is preferred to fail fast.
# 4. First pass: Substitute parameters, map internal structure & check static array unrolling
for comp in template.components:
# Substitute parameters dynamically first to resolve any parameter-driven IDs/structures
comp_resolved = substitute_params(dict(comp), params)
# Map the component's own ID
comp_resolved["id"] = map_id(comp_resolved["id"])
# Map parent-child structural references
if "child" in comp_resolved:
comp_resolved["child"] = map_id(comp_resolved["child"])
if "children" in comp_resolved:
# Handle static unrolling if children contains a static loop template
if isinstance(comp_resolved["children"], dict) and "param" in comp_resolved["children"] and "template" in comp_resolved["children"]:
children_meta = comp_resolved["children"]
param_name = children_meta["param"]
item_template_id = children_meta["template"]
array_data = params[param_name]
if not isinstance(array_data, list):
raise ValueError(f"Parameter '{param_name}' must be an array/list for static loop unrolling.")
unrolled_child_ids = []
for idx, item in enumerate(array_data):
sub_instance_id = f"{comp_resolved['id']}_item_{idx}"
unrolled_child_ids.append(sub_instance_id)
sub_params = item if isinstance(item, dict) else {"value": item}
# Recursively expand the child item template
sub_expanded = self.expand_template(sub_instance_id, item_template_id, sub_params)
expanded_components.extend(sub_expanded)
comp_resolved["children"] = unrolled_child_ids
else:
comp_resolved["children"] = map_child_list(comp_resolved["children"])
expanded_components.append(comp_resolved)References
- Do not use defensive '.get()' or fallback values when accessing dictionary or metadata keys if the presence of the key is a strict invariant. Raising a 'KeyError' is preferred to fail fast and signal invariant violations.
| def map_surface_id(m: Any): | ||
| if isinstance(m, dict): | ||
| for k in ["createSurface", "updateComponents", "deleteSurface", "updateDataModel"]: | ||
| if k in m and isinstance(m[k], dict): | ||
| m[k]["surfaceId"] = "live_surface" | ||
| elif isinstance(m, list): | ||
| for x in m: | ||
| map_surface_id(x) |
There was a problem hiding this comment.
The map_surface_id helper function only recurses into lists, but does not recursively traverse dictionary values. If the LLM wraps the envelopes in a top-level dictionary (e.g., {"response": [...]} or similar metadata wrappers), the surfaceId mapping will be skipped. Making the traversal fully recursive for both dictionaries and lists ensures robustness against LLM output variations.
| def map_surface_id(m: Any): | |
| if isinstance(m, dict): | |
| for k in ["createSurface", "updateComponents", "deleteSurface", "updateDataModel"]: | |
| if k in m and isinstance(m[k], dict): | |
| m[k]["surfaceId"] = "live_surface" | |
| elif isinstance(m, list): | |
| for x in m: | |
| map_surface_id(x) | |
| def map_surface_id(m: Any): | |
| if isinstance(m, dict): | |
| for k in ["createSurface", "updateComponents", "deleteSurface", "updateDataModel"]: | |
| if k in m and isinstance(m[k], dict): | |
| m[k]["surfaceId"] = "live_surface" | |
| for v in m.values(): | |
| map_surface_id(v) | |
| elif isinstance(m, list): | |
| for x in m: | |
| map_surface_id(x) |
1 participant
Description of Changes
This PR introduces a complete, robust prototype implementation of the A2UI Templates concept inside the Python Agent SDK and a React frontend, allowing agents to generate UIs using high-level semantic templates which are then dynamically unwrapped on the server side into standard A2UI primitive layouts.
Key implementations include:
server.py):A2uiSchemaManagerin-memory using anA2uiCatalogobject.schema_manager.generate_system_prompt(include_schema=True).processor.py):'root'to guarantee client consistency.aioasync Gemini calls (gemini-3.5-flashmodel) to prevent event-loop blocking.App.jsx):examples/):TeamGoalList,TeamMemberKnowledgePanel, andTeamFeedbackBoard.a2ui_templates.md):Rationale
Testing/Running Instructions
A comprehensive test suite containing 5 unit and integration tests is fully green and verifies all unrolling, namespacing, list payloads, and SSE streams.
Run Python Tests
cd agent_sdks/python/a2ui_agent uv run pytest tests/experimental/test_templates.pyLaunch the Sandbox Prototype
To run both backend and frontend servers in parallel from the root of the repository:
Open
http://localhost:5173in your browser.