a2ui-project / a2ui-project/a2ui

feat(python): Phase 2 - bidirectional AST deserialization for Pydantic fluent builders

Abierto
#2,524 0 comentarios 0 reacciones 0 asignados Ver en GitHub
component: agent_sdk P2 status: first-line-handled type: feature/enhancement
Lenguaje dominante
TypeScript
Estrellas
16.4k
Forks
1.3k
Merge medio
2 d 13 h
PR fusionados (30 d)
134

Descripción

## 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)

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.