generate_a2ui.py

"""Dynamic A2UI tool: LLM-generated UI from conversation context.

This module provides the data preparation for a secondary LLM call that
generates v0.9 A2UI components. The actual LLM call is made by the
framework-specific wrapper (LangGraph, CrewAI, etc.) since each framework
has its own way of invoking LLMs.
"""

from __future__ import annotations

import logging
from typing import Any, Optional

_logger = logging.getLogger(__name__)

CUSTOM_CATALOG_ID = "copilotkit://app-dashboard-catalog"

# The _design_a2ui_surface tool schema that the secondary LLM is bound to.
DESIGN_A2UI_SURFACE_TOOL_SCHEMA = {
    "name": "_design_a2ui_surface",
    "description": (
        "Render a dynamic A2UI v0.9 surface.\n\n"
        "Args:\n"
        "    surfaceId: Unique surface identifier.\n"
        '    catalogId: The catalog ID (use "copilotkit://app-dashboard-catalog").\n'
        "    components: A2UI v0.9 component array (flat format). "
        'The root component must have id "root".\n'
        "    data: Optional initial data model for the surface."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "surfaceId": {
                "type": "string",
                "description": "Unique surface identifier.",
            },
            "catalogId": {"type": "string", "description": "The catalog ID."},
            "components": {
                "type": "array",
                "items": {"type": "object"},
                "description": "A2UI v0.9 component array (flat format).",
            },
            "data": {
                "type": "object",
                "description": "Optional initial data model for the surface.",
            },
        },
        "required": ["surfaceId", "catalogId", "components"],
    },
}


def generate_a2ui_impl(
    messages: list[dict[str, Any]],
    context_entries: Optional[list[dict[str, Any]]] = None,
) -> dict[str, Any]:
    """Prepare inputs for a secondary LLM call that generates A2UI components.

    Returns a dict with:
      - system_prompt: The system prompt for the secondary LLM (built from context)
      - tool_schema: The _design_a2ui_surface tool schema to bind to the LLM
      - tool_choice: The tool name to force
      - messages: The conversation messages to pass through
      - catalog_id: The default catalog ID

    The framework wrapper should:
      1. Make an LLM call with these inputs
      2. Extract the tool call args (surfaceId, catalogId, components, data)
      3. Build a2ui_operations from the args and return them
    """
    context_text = ""
    if context_entries:
        context_text = "\n\n".join(
            entry.get("value", "")
            for entry in context_entries
            if isinstance(entry, dict) and entry.get("value")
        )

    return {
        "system_prompt": context_text,
        "tool_schema": DESIGN_A2UI_SURFACE_TOOL_SCHEMA,
        "tool_choice": "_design_a2ui_surface",
        "messages": messages,
        "catalog_id": CUSTOM_CATALOG_ID,
    }


def build_a2ui_operations_from_tool_call(args: dict[str, Any]) -> dict[str, Any]:
    """Build a2ui_operations dict from the secondary LLM's tool call args.

    Call this after the framework wrapper extracts the tool call arguments.

    The returned operations use the A2UI v0.9 NESTED shape (mirroring
    `copilotkit.a2ui.create_surface` / `update_components` / `update_data_model`
    from langgraph-python). The `@ag-ui/a2ui-middleware` extracts the
    surfaceId via `op.createSurface?.surfaceId ?? op.updateComponents?.surfaceId
    ?? ...` — a flat `{"type": "create_surface", "surfaceId": ...}` shape leaves
    every op under a "default" surface and the renderer never binds to the
    registered catalog.
    """
    surface_id = args.get("surfaceId", "dynamic-surface")
    catalog_id = args.get("catalogId", CUSTOM_CATALOG_ID)
    components = args.get("components", [])
    if not components:
        _logger.warning(
            "build_a2ui_operations_from_tool_call received empty components list"
        )
    data = args.get("data")

    ops: list[dict[str, Any]] = [
        {
            "version": "v0.9",
            "createSurface": {
                "surfaceId": surface_id,
                "catalogId": catalog_id,
            },
        },
        {
            "version": "v0.9",
            "updateComponents": {
                "surfaceId": surface_id,
                "components": components,
            },
        },
    ]
    if data:
        ops.append(
            {
                "version": "v0.9",
                "updateDataModel": {
                    "surfaceId": surface_id,
                    "path": "/",
                    "value": data,
                },
            }
        )

    return {"a2ui_operations": ops}