Python: A2UI (Agent-to-UI) support for the AG-UI adapter

python/packages/ag-ui/agent_framework_ag_ui/__init__.py

@@ -3,6 +3,7 @@
 """AG-UI protocol integration for Agent Framework."""
 
 import importlib.metadata
+from typing import TYPE_CHECKING, Any
 
 from ._agent import AgentFrameworkAgent
 from ._client import AGUIChatClient
@@ -22,6 +23,9 @@
 from ._types import AgentState, AGUIChatOptions, AGUIRequest, PredictStateConfig, RunMetadata
 from ._workflow import AgentFrameworkWorkflow, WorkflowFactory
 
+if TYPE_CHECKING:
+    from ._a2ui import A2UIAgent, AGUIContextAgent, enable_a2ui, plan_a2ui_injection
+
 try:
     __version__ = importlib.metadata.version(__name__)
 except importlib.metadata.PackageNotFoundError:
@@ -53,4 +57,21 @@
     "DEFAULT_TAGS",
     "state_update",
     "__version__",
+    # A2UI (lazy — require ag-ui-a2ui-toolkit)
+    "A2UIAgent",
+    "AGUIContextAgent",
+    "enable_a2ui",
+    "plan_a2ui_injection",
 ]
+
+# A2UI symbols are loaded lazily so importing this package does not require
+# ag-ui-a2ui-toolkit unless A2UI is actually used.
+_A2UI_LAZY = {"A2UIAgent", "AGUIContextAgent", "enable_a2ui", "plan_a2ui_injection"}
+
+
+def __getattr__(name: str) -> Any:
+    if name in _A2UI_LAZY:
+        import importlib
+
+        return getattr(importlib.import_module("._a2ui", __name__), name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

python/packages/ag-ui/agent_framework_ag_ui/_a2ui/__init__.py

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""A2UI (agent-generated UI) support for the AG-UI Agent Framework adapter.
+
+Toolkit-dependent exports (``AGUIContextAgent`` and, later, ``A2UIAgent`` /
+``get_a2ui_tools`` / ``plan_a2ui_injection``) are loaded lazily via PEP 562 so this
+package — and the toolkit-free context-plumbing helpers in :mod:`._state` it
+re-exports — can be imported by the hosting loop even when ``ag-ui-a2ui-toolkit`` is
+not installed.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+# Toolkit-free helpers — safe to import eagerly (no ag_ui_a2ui_toolkit dependency).
+from ._state import (
+    A2UI_CONTEXT_KEY,
+    A2UI_SCHEMA_CONTEXT_DESCRIPTION,
+    build_ag_ui_context_slice,
+    read_agent_state,
+    read_inject_a2ui_flag,
+    stamp_context_slice,
+)
+
+if TYPE_CHECKING:
+    from ._agent import A2UIAgent
+    from ._context_agent import AGUIContextAgent
+    from ._factory import enable_a2ui, plan_a2ui_injection
+
+__all__ = [
+    "A2UI_CONTEXT_KEY",
+    "A2UI_SCHEMA_CONTEXT_DESCRIPTION",
+    "A2UIAgent",
+    "AGUIContextAgent",
+    "build_ag_ui_context_slice",
+    "enable_a2ui",
+    "plan_a2ui_injection",
+    "read_agent_state",
+    "read_inject_a2ui_flag",
+    "stamp_context_slice",
+]
+
+# Lazy (toolkit-dependent) exports: module -> symbols defined there.
+_LAZY: dict[str, tuple[str, ...]] = {
+    "_agent": ("A2UIAgent",),
+    "_context_agent": ("AGUIContextAgent",),
+    "_factory": ("enable_a2ui", "plan_a2ui_injection"),
+}
+
+
+def __getattr__(name: str) -> Any:
+    for module, symbols in _LAZY.items():
+        if name in symbols:
+            import importlib
+
+            mod = importlib.import_module(f".{module}", __name__)
+            return getattr(mod, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")