main.py

"""Strands AG-UI Integration Example - Proverbs Agent.

This example demonstrates a Strands agent integrated with AG-UI, featuring:
- Shared state management between agent and UI
- Backend tool execution (get_weather, update_proverbs)
- Frontend tools (set_theme_color)
- Generative UI rendering
"""

import csv
import json
import os
from pathlib import Path
from typing import Any, Dict, List
from uuid import uuid4

from ag_ui_strands import (
    PredictStateMapping,
    StrandsAgent,
    StrandsAgentConfig,
    ToolBehavior,
    create_strands_app,
)
from copilotkit import a2ui
from dotenv import load_dotenv
from langchain_core.messages import SystemMessage
from langchain_core.tools import tool as lc_tool
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
from strands import Agent, tool
from strands.models.openai import OpenAIModel

# ---------------------------------------------------------------------------
# Env loading (shared demo root pattern used by the other integration demos)
# ---------------------------------------------------------------------------
_demo_root = Path(__file__).parent.parent
for env_path in (_demo_root / ".env", Path(".env")):
    if env_path.is_file():
        load_dotenv(env_path)
        break
else:
    load_dotenv(Path(__file__).resolve().parent.parent / ".env")
load_dotenv()

# ---------------------------------------------------------------------------
# Shared state schema: todos
# ---------------------------------------------------------------------------
# Strands "state" is a free-form dict carried on the AG-UI input. We keep
# the same todos shape as the reference demo so the frontend renders the
# same canvas.


class Todo(BaseModel):
    id: str = ""
    title: str
    description: str
    emoji: str
    status: str = "pending"  # "pending" | "completed"


# ---------------------------------------------------------------------------
# Tools — same names and contracts as langgraph-python
# ---------------------------------------------------------------------------


@tool
def manage_todos(todos: List[Todo]) -> str:
    """Manage the current todos.

    IMPORTANT: Always pass the full todo list, not just new items. Each todo
    should have a title, description, emoji, and status (pending/completed).
    """
    # Strands @tool validates with pydantic but passes ``model_dump()`` output
    # to the function body — so list elements arrive as plain dicts, not
    # ``Todo`` instances. Rehydrate before touching attributes.
    todos = [Todo.model_validate(t) for t in todos]
    # Ensure every todo has a stable id. The state emission callback
    # (state_from_args below) re-reads the tool arguments and sends the
    # final list to the UI, so id injection here is enough.
    for todo in todos:
        if not todo.id:
            todo.id = str(uuid4())
    return "Successfully updated todos"


@tool
def get_todos() -> str:
    """Get the current todos.

    Returns a JSON string of the current todos list. The list is injected
    into the prompt via the state context builder, but this tool is still
    useful when the model wants to re-confirm state.
    """
    # Strands tools don't get a runtime handle, so we rely on the state
    # context builder to surface the list. Returning a marker string tells
    # the model to read state from the prompt it already has.
    return "See the current todos list already provided in the conversation context."


_CSV_PATH = Path(__file__).parent / "src" / "db.csv"
with open(_CSV_PATH) as _f:
    _CACHED_DATA = list(csv.DictReader(_f))


@tool
def query_data(query: str) -> str:
    """Query the database with a natural-language query.

    Always call this before rendering a chart so the UI has data to plot.
    """
    return json.dumps(_CACHED_DATA)


# ---------------------------------------------------------------------------
# A2UI tools (framework-agnostic — use copilotkit.a2ui helpers directly)
# ---------------------------------------------------------------------------
CATALOG_ID = "copilotkit://app-dashboard-catalog"
FLIGHT_SURFACE_ID = "flight-search-results"
FLIGHT_SCHEMA = a2ui.load_schema(
    Path(__file__).parent / "src" / "a2ui" / "schemas" / "flight_schema.json"
)


class Flight(BaseModel):
    id: str
    airline: str
    airlineLogo: str
    flightNumber: str
    origin: str
    destination: str
    date: str
    departureTime: str
    arrivalTime: str
    duration: str
    status: str
    statusIcon: str
    price: str


class FlightList(BaseModel):
    flights: List[Flight]


@tool
def search_flights(flight_list: FlightList) -> str:
    """Search for flights and display the results as rich cards.

    Return exactly 2 flights. Each flight must have: id, airline, airlineLogo
    (Google favicon API URL for the airline domain), flightNumber, origin,
    destination, date (e.g. "Tue, Mar 18" — use near-future dates),
    departureTime, arrivalTime, duration (e.g. "4h 25m"), status (e.g.
    "On Time" or "Delayed"), statusIcon (colored dot URL:
    https://placehold.co/12/22c55e/22c55e.png for On Time,
    https://placehold.co/12/eab308/eab308.png for Delayed,
    https://placehold.co/12/ef4444/ef4444.png for Cancelled), and price
    (e.g. "$289").
    """
    # Strands @tool passes plain dicts (model_dump output) — ``flight_list``
    # is a dict, ``flight_list["flights"]`` is a list of dicts. Validate
    # back to Pydantic to enforce the schema, then dump for a2ui rendering.
    parsed = FlightList.model_validate(flight_list)
    flights_payload = [f.model_dump() for f in parsed.flights]
    return a2ui.render(
        operations=[
            a2ui.create_surface(FLIGHT_SURFACE_ID, catalog_id=CATALOG_ID),
            a2ui.update_components(FLIGHT_SURFACE_ID, FLIGHT_SCHEMA),
            a2ui.update_data_model(FLIGHT_SURFACE_ID, {"flights": flights_payload}),
        ],
    )


@lc_tool
def render_a2ui(
    surfaceId: str,
    catalogId: str,
    components: List[Dict[str, Any]],
    data: Dict[str, Any] | None = None,
) -> str:
    """Render a dynamic A2UI v0.9 surface.

    Args:
        surfaceId: Unique surface identifier.
        catalogId: The catalog ID (use "copilotkit://app-dashboard-catalog").
        components: A2UI v0.9 component array (flat format). The root
            component must have id "root".
        data: Optional initial data model for the surface (e.g. form values,
            list items for data-bound components).
    """
    return "rendered"


@tool
def generate_a2ui(user_intent: str, agent) -> str:
    """Generate dynamic A2UI components based on the conversation.

    A secondary LLM designs the UI schema and data. The result is returned
    as an a2ui_operations container for the middleware to detect and render.

    Seed the secondary LLM with the catalog + component schema entries
    that CopilotKit's runtime middleware injects into
    ``RunAgentInput.context``. The ag_ui_strands adapter forwards those
    entries onto ``agent.state`` under the ``agui_context`` key.
    """
    context_entries = []
    try:
        context_entries = agent.state.get("agui_context") or []
    except Exception:
        context_entries = []

    context_text = "\n\n".join(
        e.get("value", "")
        for e in context_entries
        if isinstance(e, dict) and e.get("value")
    )
    prompt = f"{context_text}\n\n{user_intent}" if context_text else user_intent

    model = ChatOpenAI(model="gpt-4.1")
    model_with_tool = model.bind_tools(
        [render_a2ui],
        tool_choice="render_a2ui",
    )
    try:
        response = model_with_tool.invoke(
            [SystemMessage(content=prompt)],
        )
    except Exception as exc:  # pragma: no cover — surface LLM/network failures
        return json.dumps({"error": f"dynamic-a2ui LLM call failed: {exc}"})

    if not response.tool_calls:
        return json.dumps({"error": "LLM did not call render_a2ui"})

    tool_call = response.tool_calls[0]
    args = tool_call["args"]

    surface_id = args.get("surfaceId", "dynamic-surface")
    catalog_id = args.get("catalogId", CATALOG_ID)
    components = args.get("components", []) or []
    data = args.get("data") or {}

    ops = [
        a2ui.create_surface(surface_id, catalog_id=catalog_id),
        a2ui.update_components(surface_id, components),
    ]
    if data:
        ops.append(a2ui.update_data_model(surface_id, data))

    return a2ui.render(operations=ops)


# ---------------------------------------------------------------------------
# Shared-state config: inject todos into the prompt, stream state back on
# every manage_todos tool call.
# ---------------------------------------------------------------------------


def build_todos_prompt(input_data, user_message: str) -> str:
    """Inject the current todos state into the prompt."""
    state_dict = getattr(input_data, "state", None)
    if isinstance(state_dict, dict) and "todos" in state_dict:
        todos_json = json.dumps(state_dict.get("todos", []), indent=2)
        return f"Current todos list:\n{todos_json}\n\nUser request: {user_message}"
    return user_message


async def todos_state_from_args(context):
    """Snapshot state for the UI after a manage_todos call.

    Strands calls this with the tool's parsed arguments. We return the
    `todos` list so the AG-UI layer can emit a STATE_SNAPSHOT event.
    """
    try:
        tool_input = context.tool_input
        if isinstance(tool_input, str):
            tool_input = json.loads(tool_input)
        todos = tool_input.get("todos", [])
        return {"todos": todos}
    except Exception:
        return None


shared_state_config = StrandsAgentConfig(
    state_context_builder=build_todos_prompt,
    tool_behaviors={
        "manage_todos": ToolBehavior(
            state_from_args=todos_state_from_args,
            predict_state=[
                PredictStateMapping(
                    state_key="todos",
                    tool="manage_todos",
                    tool_argument="todos",
                )
            ],
        )
    },
)


# ---------------------------------------------------------------------------
# Agent wiring
# ---------------------------------------------------------------------------
api_key = os.getenv("OPENAI_API_KEY", "")
model = OpenAIModel(
    client_args={"api_key": api_key},
    model_id="gpt-5.4",
    params={"parallel_tool_calls": False},
)

system_prompt = (
    "You are a polished, professional demo assistant. Keep responses to 1-2 sentences.\n\n"
    "Tool guidance:\n"
    "- Flights: call search_flights to show flight cards with a pre-built schema.\n"
    "- Dashboards & rich UI: call generate_a2ui to create dashboard UIs with metrics,\n"
    "  charts, tables, and cards. It handles rendering automatically.\n"
    "- Charts: call query_data first, then render with the chart component.\n"
    "- Todos: enable app mode first, then manage todos.\n"
    "- Diagrams (Excalidraw): when MCP Excalidraw tools are exposed (e.g. create_view),\n"
    "  call create_view ONCE with 3-5 elements (shapes + arrows + optional title text).\n"
    "  Include ONE cameraUpdate at the end to frame the diagram. Do NOT call read_me\n"
    "  even if it appears in the toolset — you already know the basic shape API.\n"
    '- A2UI actions: when you see a log_a2ui_event result (e.g. "view_details"),\n'
    "  respond with a brief confirmation. The UI already updated on the frontend."
)

strands_agent = Agent(
    model=model,
    system_prompt=system_prompt,
    tools=[manage_todos, get_todos, query_data, generate_a2ui, search_flights],
)

agui_agent = StrandsAgent(
    agent=strands_agent,
    name="todo_demo_agent",
    description=(
        "A polished demo assistant matching the canonical langgraph-python "
        "todo / charts / a2ui / flights showcase, running on Strands."
    ),
    config=shared_state_config,
)

agent_path = os.getenv("AGENT_PATH", "/")
app = create_strands_app(agui_agent, agent_path)


@app.get("/health")
async def health():
    return {"status": "ok"}


if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("AGENT_PORT", 8000))
    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=True)