+ "content": "\"\"\"Agentic Generative UI feature.\"\"\"\n\nfrom __future__ import annotations\n\nfrom textwrap import dedent\nfrom typing import Any, Literal\n\nfrom fastapi import FastAPI\nfrom pydantic import BaseModel, Field\n\nfrom ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent\nfrom ag_ui_adk import ADKAgent, add_adk_fastapi_endpoint\nfrom google.adk.agents import LlmAgent\n\nStepStatus = Literal['pending', 'completed']\n\n\nclass Step(BaseModel):\n \"\"\"Represents a step in a plan.\"\"\"\n\n description: str = Field(description='The description of the step')\n status: StepStatus = Field(\n default='pending',\n description='The status of the step (e.g., pending, completed)',\n )\n\n\nclass Plan(BaseModel):\n \"\"\"Represents a plan with multiple steps.\"\"\"\n\n steps: list[Step] = Field(default_factory=list, description='The steps in the plan')\n\n\nclass JSONPatchOp(BaseModel):\n \"\"\"A class representing a JSON Patch operation (RFC 6902).\"\"\"\n\n op: Literal['add', 'remove', 'replace', 'move', 'copy', 'test'] = Field(\n description='The operation to perform: add, remove, replace, move, copy, or test',\n )\n path: str = Field(description='JSON Pointer (RFC 6901) to the target location')\n value: Any = Field(\n default=None,\n description='The value to apply (for add, replace operations)',\n )\n from_: str | None = Field(\n default=None,\n alias='from',\n description='Source path (for move, copy operations)',\n )\n\n\nasync def create_plan(steps: list[str]) -> StateSnapshotEvent:\n \"\"\"Create a plan with multiple steps.\n\n Args:\n steps (list[str]): List of step descriptions to create the plan.\n\n Returns:\n StateSnapshotEvent: Event containing the initial state of the steps.\n \"\"\"\n plan: Plan = Plan(\n steps=[Step(description=step) for step in steps],\n )\n return StateSnapshotEvent(\n type=EventType.STATE_SNAPSHOT,\n snapshot=plan.model_dump(),\n )\n\n\nasync def update_plan_step(\n index: int, description: str | None = None, status: StepStatus | None = None\n) -> StateDeltaEvent:\n \"\"\"Update the plan with new steps or changes.\n\n Args:\n index (int): The index of the step to update.\n description (str | None): The new description for the step.\n status (StepStatus | None): The new status for the step.\n\n Returns:\n StateDeltaEvent: Event containing the changes made to the plan.\n \"\"\"\n changes: list[JSONPatchOp] = []\n if description is not None:\n changes.append(\n JSONPatchOp(\n op='replace', path=f'/steps/{index}/description', value=description\n )\n )\n if status is not None:\n changes.append(\n JSONPatchOp(op='replace', path=f'/steps/{index}/status', value=status)\n )\n return StateDeltaEvent(\n type=EventType.STATE_DELTA,\n delta=changes,\n )\n\n\n# Create the ADK agent\nagent = LlmAgent(\n name=\"planner\",\n model=\"gemini-2.0-flash\",\n instruction=dedent(\n \"\"\"\n When planning use tools only, without any other messages.\n IMPORTANT:\n - Use the `create_plan` tool to set the initial state of the steps\n - Use the `update_plan_step` tool to update the status of each step\n - Do NOT repeat the plan or summarise it in a message\n - Do NOT confirm the creation or updates in a message\n - Do NOT ask the user for additional information or next steps\n - Do NOT leave a plan hanging, always complete the plan via `update_plan_step` if one is ongoing.\n\n Only one plan can be active at a time, so do not call the `create_plan` tool\n again until all the steps in current plan are completed.\n \"\"\"\n ),\n tools=[create_plan, update_plan_step],\n)\n\n# Create ADK middleware agent instance\nadk_agent = ADKAgent(\n adk_agent=agent,\n app_name=\"demo_app\",\n user_id=\"demo_user\",\n session_timeout_seconds=3600,\n use_in_memory_services=True,\n)\n\n# Create FastAPI app\napp = FastAPI(title=\"ADK Middleware Agentic Generative UI\")\n\n# Add the ADK endpoint\nadd_adk_fastapi_endpoint(app, adk_agent, path=\"/\")\n",
0 commit comments