a2ui_dynamic.py

"""
MS Agent Framework agent for the Declarative Generative UI (A2UI — Dynamic Schema) demo.

Pattern (ported from the LangGraph reference
`showcase/integrations/langgraph-python/src/agents/a2ui_dynamic.py`):

- The agent binds an explicit `generate_a2ui` tool. When called, it invokes a
  secondary LLM bound to `_design_a2ui_surface` (tool_choice forced) and returns the
  resulting `a2ui_operations` container.
- The runtime (see `src/app/api/copilotkit-declarative-gen-ui/route.ts`) uses
  `injectA2UITool: false` because the tool binding is owned by the agent here
  (double-injection would duplicate the tool slot).
"""

from __future__ import annotations

import json
from textwrap import dedent
from typing import Annotated, Any

from agent_framework import Agent, BaseChatClient, tool
from agent_framework_ag_ui import AgentFrameworkAgent
from pydantic import Field

from tools import build_a2ui_operations_from_tool_call

CUSTOM_CATALOG_ID = "declarative-gen-ui-catalog"


@tool(
    name="generate_a2ui",
    description=(
        "Generate dynamic A2UI components based on the conversation. "
        "A secondary LLM designs the UI schema and data."
    ),
)
def generate_a2ui(
    context: Annotated[
        str,
        # Default to empty so the primary LLM can call generate_a2ui() with
        # no args (aimock fixtures return `arguments: "{}"`); pydantic rejects
        # missing-context calls with "Argument parsing failed" before the
        # function body runs. Same pattern as
        # `src/agents/beautiful_chat.py::generate_a2ui`.
        Field(default="", description="Conversation context to generate UI from."),
    ] = "",
    session: Any = None,
) -> str:
    """Generate dynamic A2UI dashboard from conversation context."""
    from openai import OpenAI

    # Pull the latest user message from the active agent session so the
    # secondary LLM call sees what the user actually asked for. Without this,
    # aimock's substring matcher can't distinguish between "KPI dashboard",
    # "pie chart", and "bar chart" pills — they'd all hit the first matching
    # fixture. `session` is the AgentSession injected by agent_framework
    # (see `_tools.py:1483`). When unavailable (e.g. direct calls in tests),
    # we fall back to the caller-supplied `context` string.
    latest_user_message = ""
    if session is not None:
        try:
            messages = list(getattr(session, "input_messages", []) or [])
            for msg in reversed(messages):
                if getattr(msg, "role", None) == "user":
                    text = getattr(msg, "text", None) or str(
                        getattr(msg, "content", "") or ""
                    )
                    if text:
                        latest_user_message = text
                        break
        except Exception:
            latest_user_message = ""

    client = OpenAI()
    tool_schema = {
        "type": "function",
        "function": {
            "name": "_design_a2ui_surface",
            "description": "Render a dynamic A2UI v0.9 surface.",
            "parameters": {
                "type": "object",
                "properties": {
                    "surfaceId": {"type": "string"},
                    "catalogId": {"type": "string"},
                    "components": {"type": "array", "items": {"type": "object"}},
                    "data": {"type": "object"},
                },
                "required": ["surfaceId", "catalogId", "components"],
            },
        },
    }

    # Build the secondary-LLM user message. Priority:
    #   1. `latest_user_message` from the active AgentSession (preferred —
    #      lets aimock's substring matcher pick the right fixture per pill).
    #   2. The caller-supplied `context` arg (LangGraph-style summary).
    #   3. A generic catch-all that contains all four d5 demo keywords so
    #      direct/test invocations still match SOME fixture instead of
    #      falling through to the real-OpenAI proxy.
    user_content = (
        latest_user_message
        or context
        or "KPI dashboard with 3-4 metrics, pie chart sales by region, "
        "bar chart quarterly revenue, status report."
    )
    response = client.chat.completions.create(
        model="gpt-4.1",
        messages=[
            {
                "role": "system",
                "content": (
                    f"Generate a useful dashboard UI. Use catalogId='{CUSTOM_CATALOG_ID}'."
                ),
            },
            {"role": "user", "content": user_content},
        ],
        tools=[tool_schema],
        tool_choice={"type": "function", "function": {"name": "_design_a2ui_surface"}},
    )

    if not response.choices[0].message.tool_calls:
        return json.dumps({"error": "LLM did not call _design_a2ui_surface"})

    tool_call = response.choices[0].message.tool_calls[0]
    args = json.loads(tool_call.function.arguments)
    # Default the catalog to the dynamic-gen-ui catalog if the LLM omitted it.
    args.setdefault("catalogId", CUSTOM_CATALOG_ID)
    result = build_a2ui_operations_from_tool_call(args)
    return json.dumps(result)


SYSTEM_PROMPT = dedent(
    """
    You are a demo assistant for Declarative Generative UI (A2UI — Dynamic
    Schema). Whenever a response would benefit from a rich visual — a
    dashboard, status report, KPI summary, card layout, info grid, a
    pie/donut chart of part-of-whole breakdowns, a bar chart comparing
    values across categories, or anything more structured than plain text —
    call `generate_a2ui` to draw it. The registered catalog includes
    `Card`, `StatusBadge`, `Metric`, `InfoRow`, `PrimaryButton`, `PieChart`,
    and `BarChart` (in addition to the basic A2UI primitives). Prefer
    `PieChart` for part-of-whole breakdowns (sales by region, traffic
    sources, portfolio allocation) and `BarChart` for comparisons across
    categories (quarterly revenue, headcount by team, signups per month).
    `generate_a2ui` takes a `context` string summarising the user's request
    and handles the rendering automatically. Keep chat replies to one short
    sentence; let the UI do the talking.
    """
).strip()


def create_agent(chat_client: BaseChatClient) -> AgentFrameworkAgent:
    """Instantiate the MS-Agent-backed declarative-gen-ui agent."""
    base_agent = Agent(
        client=chat_client,
        name="declarative_gen_ui_agent",
        instructions=SYSTEM_PROMPT,
        tools=[generate_a2ui],
    )

    return AgentFrameworkAgent(
        agent=base_agent,
        name="CopilotKitMicrosoftAgentFrameworkAgent",
        description="Dynamic A2UI generator that designs rich UI surfaces on demand.",
        require_confirmation=False,
    )