"content": "\"\"\"Agentic Generative UI feature.\"\"\"\n\nfrom __future__ import annotations\n\nfrom textwrap import dedent\nfrom typing import Any, Literal\n\nfrom pydantic import BaseModel, Field\n\nfrom ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent\nfrom pydantic_ai import Agent\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\nagent = Agent(\n 'openai:gpt-4o-mini',\n instructions=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)\n\n\
[email protected]_plain\nasync def create_plan(steps: list[str]) -> StateSnapshotEvent:\n \"\"\"Create a plan with multiple steps.\n\n Args:\n steps: List of step descriptions to create the plan.\n\n Returns:\n StateSnapshotEvent 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\
[email protected]_plain\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: The index of the step to update.\n description: The new description for the step.\n status: The new status for the step.\n\n Returns:\n StateDeltaEvent 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\napp = agent.to_ag_ui()\n",
0 commit comments