docs(ag-ui): various documentation improvements

Author: stevenh

adapter_ag_ui/adapter_ag_ui/adapter.py

@@ -77,7 +77,6 @@
 
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
-_LOGGER.setLevel(logging.DEBUG)
 
 
 @dataclass(repr=False)
@@ -105,12 +104,9 @@ def new_message_id(self) -> str:
 class AdapterAGUI(Generic[AgentDepsT, OutputDataT]):
     """An agent adapter providing AG-UI protocol support for PydanticAI agents.
 
-    This class manages the agent runs, tool calls, and provides an adapter
-    for running agents with Server-Sent Event (SSE) streaming responses using
-    the AG-UI protocol.
-
-    It supports multiple threads and runs, allowing for concurrent agent
-    execution and tool calls.
+    This class manages the agent runs, tool calls, state storage and providing
+    an adapter for running agents with Server-Sent Event (SSE) streaming
+    responses using the AG-UI protocol.
 
     Example:
 
@@ -143,8 +139,8 @@ async def root(input_data: RunAgentInput, accept: Annotated[str, Header()] = SSE
                 media_type=SSE_CONTENT_TYPE,
             )
 
-    PydanticAI Tools which return AG-UI events will be sent to the client
-    as part of the event stream, single events and iterables of events are
+    PydanticAI tools which return AG-UI events will be sent to the client
+    as part of the event stream, single events and event iterables are
     supported.
 
     Examples:
@@ -296,8 +292,9 @@ async def _tool_events(
         Yields:
             AG-UI Server-Sent Events (SSE).
         """
-        # TODO(steve): Determine how to handle when we have multiple parts. Currently
-        # AG-UI only seems to support a single tool call per request.
+        # TODO(steve): Determine how to handle multiple parts. Currently
+        # AG-UI only supports a single tool call per request, but that
+        # may change in the future.
         part: ModelRequestPart
         for part in parts:
             if not isinstance(part, ToolReturnPart):
@@ -323,9 +320,9 @@ async def _tool_events(
     def _convert_tools(self, run_tools: list[ToolAGUI]) -> list[Tool[AgentDepsT]]:
         """Convert AG-UI tools to PydanticAI tools.
 
-        Creates placeholder Tool objects from AG-UI tool definitions. These tools
-        don't actually execute anything - they just provide the necessary tool
-        definitions and names for the PydanticAI agent.
+        Creates `Tool` objects from AG-UI tool definitions. These tools don't
+        actually execute anything, that is done by AG-UI client - they just
+        provide the necessary tool definitions to PydanticAI agent.
 
         Args:
             run_tools: List of AG-UI tool definitions to convert.
@@ -342,7 +339,7 @@ def _tool_call(self, tool: ToolAGUI) -> Tool[AgentDepsT]:
             tool: The AG-UI tool definition to convert.
 
         Returns:
-            A PydanticAI Tool object that calls the AG-UI tool.
+            A PydanticAI `Tool` object that calls the AG-UI tool.
         """
 
         def _tool_stub(*args: Any, **kwargs: Any) -> ToolResult:
@@ -352,7 +349,7 @@ def _tool_stub(*args: Any, **kwargs: Any) -> ToolResult:
                 Never returns as it always raises an exception.
 
             Raises:
-                UnexpectedToolCallError: Always raised since this is a stub.
+                UnexpectedToolCallError: Always raised since this should never be called.
             """
             raise UnexpectedToolCallError(tool_name=tool.name)  # pragma: no cover
 
@@ -521,13 +518,13 @@ async def _handle_agent_event(
 
 
 def _convert_history(messages: list[Message]) -> list[ModelMessage]:
-    """Convert a list of AG-UI history to a PydanticAI one.
+    """Convert a AG-UI history to a PydanticAI one.
 
     Args:
-        messages: List of `ag_ui.Message` to convert.
+        messages: List of AG-UI messages to convert.
 
     Returns:
-        List of `ModelMessage` compatible with PydanticAI.
+        List of PydanticAI model messages.
     """
     msg: Message
     result: list[ModelMessage] = []

adapter_ag_ui/adapter_ag_ui/deps.py

@@ -9,11 +9,21 @@
 from ._exceptions import InvalidStateError
 
 StateT = TypeVar('StateT', bound=BaseModel, contravariant=True)
+"""Type variable for the state type, which must be a subclass of `BaseModel`."""
 
 
 @dataclass(kw_only=True)
 class StateDeps(Generic[StateT]):
-    """Provides AG-UI state management."""
+    """Provides AG-UI state management.
+
+    This class is used to manage the state of an agent run. It allows setting
+    the state of the agent run with a specific type of state model, which must
+    be a subclass of `BaseModel`.
+
+    The state is set using the `set_state` when the run starts by the `AdapterAGUI`.
+
+    Implements the `StateHandler` protocol.
+    """
 
     state_type: type[StateT]
     state: StateT = field(init=False)

docs/ag-ui.md

@@ -3,16 +3,18 @@
 The [Agent User Interaction (AG-UI) Protocol](https://docs.ag-ui.com/introduction)
 is an open standard introduced by the
 [CopilotKit](https://webflow.copilotkit.ai/blog/introducing-ag-ui-the-protocol-where-agents-meet-users)
-team that standardises how front-end applications connect to AI agents through
+team that standardises how frontend applications connect to AI agents through
 an open protocol. Think of it as a universal translator for AI-driven systems
 no matter what language an agent speaks: AG-UI ensures fluent communication.
 
 The team at [Rocket Science](https://www.rocketscience.gg/), contributed the
 [adapter-ag-ui](#adapter-ag-ui) library to make it easy to implement the AG-UI
 protocol with PydanticAI agents.
 
-This also includes a convenience method that expose PydanticAI agents as AG-UI
-servers - let's have a quick look at how to use it:
+This also includes a convenience method that exposes PydanticAI agents as an
+AG-UI server, which can then be used by as part of a
+[fastapi](https://fastapi.tiangolo.com/) app. Let's have a quick look at how to
+use it:
 
 ```py {title="agent_to_ag_ui.py" py="3.10" hl_lines="17-27"}
 """Basic example for AG-UI with FastAPI and Pydantic AI."""
@@ -53,20 +55,27 @@ uvicorn agent_to_ag_ui:app --host 0.0.0.0 --port 8000
 This will expose the agent as an AG-UI server, and you can start sending
 requests to it.
 
-## Adapter AG UI
+## AG-UI Adapter
 
-[AdapterAGUI][adapter_ag_ui.AdapterAGUI]is an adapter between PydanticAI agents
-and the AG-UI protocol written in Python.
+The [AdapterAGUI][adapter_ag_ui.AdapterAGUI] class is an adapter between
+PydanticAI agents and the AG-UI protocol written in Python. It provides support
+for all aspects of spec including:
+
+- [Events](https://docs.ag-ui.com/concepts/events)
+- [Messages](https://docs.ag-ui.com/concepts/messages)
+- [State Management](https://docs.ag-ui.com/concepts/state)
+- [Tools](https://docs.ag-ui.com/concepts/tools)
 
 ### Design
 
 The adapter receives messages in the form of a
 [`RunAgentInput`](https://docs.ag-ui.com/sdk/js/core/types#runagentinput)
 which describes the details of a request being passed to the agent including
 messages and state. These are then converted to PydanticAI types, passed to the
-provided agent which then process the request. Results from the agent are
-converted from PydanticAI types to AG-UI events and streamed back to the caller
-as Server-Sent Events (SSE).
+agent which then process the request.
+
+Results from the agent are converted from PydanticAI types to AG-UI events and
+streamed back to the caller as Server-Sent Events (SSE).
 
 A user request may require multiple round trips between client UI and PydanticAI
 server, depending on the tools and events needed.