a2ui-project / a2ui-project/a2ui
feat(python): Phase 2 - bidirectional AST deserialization for Pydantic fluent builders
- Dominant language
- TypeScript
- Stars
- 16.4k
- Forks
- 1.3k
- Avg merge
- 2d 13h
- Merged PRs (30d)
- 134
Description
## Objective
Implement **Phase 2: Bidirectional AST Deserialization** for the Python Agent SDK's Pydantic fluent builder system, turning flat wire A2UI JSON payloads and message envelopes into strongly-typed, navigable, mutable `ComponentTree` instances.
This builds directly on the Phase 1 fluent builder foundation introduced in [#2425](https://github.com/a2ui-project/a2ui/pull/2425) and detailed in [`specification/proposals/macros/pydantic_fluent_builders_and_deserialization.md`](https://github.com/a2ui-project/a2ui/blob/feat/typescript-cli-proposals/specification/proposals/macros/pydantic_fluent_builders_and_deserialization.md).
---
## Background & Motivation
Phase 1 provides outward serialization: agents write `Card(child=Column(...))` and call `.to_components()` to produce flat A2UI dictionaries.
However, many agent workflows require **inbound inspection and mutation**:
1. **Agent-to-agent negotiation**: An agent receives a UI payload from a subagent or peer and needs to inspect, validate, or modify specific nodes before re-emitting.
2. **Interactive UI updating**: An agent receives user interaction context and updates an existing surface tree incrementally.
3. **Template/Macro re-hydration**: Re-hydrating wire messages back into AST builder objects to perform semantic transformations.
Currently, reconstructing an AST from flat wire components requires ad-hoc dictionary traversals. Phase 2 standardizes this via native Pydantic v2 validation.
---
## Technical Architecture & Design
### 1. Single-Pass Contextual Slot Resolution (`WrapValidator`)
In `agent_sdks/python/a2ui_agent/src/a2ui/builder/base.py`:
Upgrade `Slot` and `SlotList` to use Pydantic's `WrapValidator`. When deserializing flat JSON, the validator receives the wire string ID, retrieves the raw component dictionary from `info.context["components"]`, records visited IDs for cycle detection, and validates the child recursively:
```python
from typing import Annotated, Any, Sequence, TypeAlias
from pydantic import ValidationInfo, WrapValidator
from a2ui.builder.base import ComponentBuilderNode
def _resolve_slot(
val: Any, handler: Any, info: ValidationInfo
) -> ComponentBuilderNode:
"""Resolves wire string IDs into ComponentBuilderNode instances during validation."""
if isinstance(val, str) and info.context and "components" in info.context:
by_id = info.context["components"]
visited = info.context.setdefault("_visited", set())
if val in visited:
raise ValueError(
f"Circular reference detected in component hierarchy at ID '{val}'"
)
if val in by_id:
visited.add(val)
raw_child_dict = by_id[val]
return handler(raw_child_dict)
# Directly instantiated model instances pass through transparently
return handler(val)
Slot: TypeAlias = Annotated[ComponentBuilderNode, WrapValidator(_resolve_slot)]
SlotList: TypeAlias = Sequence[Slot]
```
*Non-breaking guarantee*: Direct authoring (`Card(child=Button(...))`) passes model instances straight through `handler(val)` without alteration.
---
### 2. Catalog Evolution & Unknown Component Fallback
Generated catalog builders (`a2ui.builder.catalogs.basic`) must emit an `UnknownComponent` fallback and a discriminated union `Component`:
```python
from typing import Annotated, Literal, Optional, Union
from pydantic import ConfigDict, Field
from a2ui.builder.base import ComponentBuilderNode
class UnknownComponent(ComponentBuilderNode):
"""Fallback node preserving unrecognized wire components and attributes."""
model_config = ConfigDict(extra="allow")
component: str
# Discriminated union enabling polymorphic type validation in Pydantic:
Component = Annotated[
Union[
Button,
Card,
Column,
Row,
Text,
# ... all generated catalog components ...
UnknownComponent,
],
Field(discriminator="component"),
]
```
*Resilience guarantee*: If an upstream client introduces a new component or extra fields, `UnknownComponent` captures the payload with `extra="allow"` in `__pydantic_extra__`, allowing safe round-trip serialization without losing data.
---
### 3. Container Abstraction (`ComponentTree`)
Rather than conflating in-memory trees with a rendering surface canvas, provide a standalone `ComponentTree` container in `a2ui.builder.base`:
```python
class ComponentTree:
"""An in-memory hierarchy of components, containing a primary root and unlinked subtrees."""
def __init__(
self,
root: ComponentBuilderNode,
unlinked_roots: Sequence[ComponentBuilderNode] | None = None,
surface_id: str | None = None,
):
self.root = root
self.unlinked_roots = list(unlinked_roots or [])
self.surface_id = surface_id
def to_components(self) -> list[dict[str, Any]]:
"""Serializes root and unlinked subtrees into flat component dicts."""
comps = self.root.to_components()
for sub_tree in self.unlinked_roots:
comps.extend(sub_tree.to_components())
return comps
def to_update(self, surface_id: str | None = None) -> dict[str, Any]:
"""Packages the tree into an updateComponents envelope."""
target_id = surface_id or self.surface_id or "main"
return {
"updateComponents": {
"surfaceId": target_id,
"components": self.to_components(),
}
}
def to_surface(
self, surface_id: str | None = None, catalog_id: str | None = None
) -> list[dict[str, Any]]:
"""Packages the tree into createSurface + updateComponents envelopes."""
target_id = surface_id or self.surface_id or "main"
create_env: dict[str, Any] = {"createSurface": {"surfaceId": target_id}}
if catalog_id:
create_env["createSurface"]["catalogId"] = catalog_id
return [create_env, self.to_update(target_id)]
def prune_unlinked(self) -> None:
"""Clears all unlinked subtrees."""
self.unlinked_roots.clear()
```
---
### 4. Deserialization Entrypoint (`deserialize`)
Implement `deserialize()` in `a2ui.builder.base` (or `a2ui.builder.deserializer`):
```python
def deserialize(
payload: Mapping[str, Any] | Sequence[Mapping[str, Any]] | str,
adapter: TypeAdapter[Any] = TypeAdapter(Component),
) -> ComponentTree:
"""Rebuilds a typed ComponentTree from an A2UI payload in a single pass."""
```
**Algorithm requirements**:
1. Ingest string JSON, envelope dict (`createSurface`, `updateComponents`), or flat component array.
2. Extract `surfaceId` if present.
3. Index components by `id`: `by_id = {c["id"]: c for c in components}`.
4. Determine root component:
- If `rootId` is explicitly provided, use it.
- Otherwise, calculate incoming reference count (in-degree) for all IDs across all component slot attributes (`child`, `children`). The node with 0 incoming references is the primary root candidate.
5. Recursively validate `root_id` with `adapter.validate_python(by_id[root_id], context=context)`.
6. Identify unlinked components: `unvisited_ids = set(by_id.keys()) - context["_visited"]`.
7. Reconstruct typed ASTs for unvisited components and append to `unlinked_roots`.
8. Return `ComponentTree(root=root_node, unlinked_roots=unlinked_roots, surface_id=surface_id)`.
---
## Edge Cases to Handle
1. **Circular References**: Two components referencing each other (`c1.child = "c2"`, `c2.child = "c1"`). Must raise a descriptive `ValueError`.
2. **Dangling Slot References**: A component slot references an ID not present in the payload. Must coerce to `ExternalComponentBuilderNode(id=missing_id)` without crashing.
3. **Disconnected / Orphan Components**: Elements that belong to a different container or an incremental patch. Must be collected into `tree.unlinked_roots` rather than dropped.
4. **Custom / Unknown Components**: Wire components with unknown types must deserialize into `UnknownComponent` with exact dictionary retention.
5. **Data Bindings vs String Literals**: Distinguish between `{"path": "/user/name"}` (coerced to `DataBinding`) and plain string literals.
---
## Code Generator Alignment (`@a2ui/cli` & Dart CLI)
Both CLI implementations must emit:
1. `UnknownComponent(ComponentBuilderNode)` model with `extra="allow"`.
2. `Component = Annotated[Union[..., UnknownComponent], Field(discriminator="component")]`.
3. `Slot = Annotated[ComponentBuilderNode, WrapValidator(_resolve_slot)]`.
* Files to update:
* `dart/a2ui_cli/lib/src/emitters/python/python_emitter.dart`
* `javascript/a2ui_cli/src/emitters/python/python-emitter.ts`
* Conformance tests: `conformance/cli/codegen.yaml`
---
## Implementation Checklist
- [ ] **Base Runtime (`a2ui_agent/src/a2ui/builder/base.py`)**:
- Implement `_resolve_slot` WrapValidator.
- Define `Slot` and `SlotList` with WrapValidator.
- Implement `ComponentTree`.
- Implement `deserialize()`.
- Implement envelope convenience functions (`create_surface()`, `update_components()`).
- [ ] **CLI Code Generators**:
- Update Dart CLI Python emitter (`python_emitter.dart`) to emit `UnknownComponent` and discriminated union `Component`.
- Update TS CLI Python emitter (`python-emitter.ts`) to emit matching models.
- Re-run CLI conformance suite (`conformance/cli/codegen.yaml`).
- [ ] **Basic Catalog Builders**:
- Regenerate `agent_sdks/python/a2ui_agent/src/a2ui/builder/catalogs/basic/basic.py`.
- [ ] **Unit & Conformance Tests**:
- Add `agent_sdks/python/a2ui_agent/tests/test_pydantic_deserialization.py`.
- Test simple single-pass tree deserialization.
- Test envelope extraction (`createSurface`, `updateComponents`).
- Test circular reference detection.
- Test unknown component fallback and round-trip preservation.
- Test unlinked subtree collection and pruning.
- [ ] **Documentation**:
- Update `src/a2ui/builder/README.md` with deserialization examples and API reference.
---
## References & Design Links
- Full Design Proposal: [`specification/proposals/macros/pydantic_fluent_builders_and_deserialization.md`](https://github.com/a2ui-project/a2ui/blob/feat/typescript-cli-proposals/specification/proposals/macros/pydantic_fluent_builders_and_deserialization.md)
- Phase 1 Implementation PR: [#2425](https://github.com/a2ui-project/a2ui/pull/2425)
- Macro Runtime PR: [#2519](https://github.com/a2ui-project/a2ui/pull/2519)
Contributor guide
Assessment
This issue has not been assessed yet.