chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,353 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user