Skip to content

Commit 214a6c7

Browse files
committed
feat(ag-ui): ASGI server
Refactor to_ag_ui so it now returns a ASGI compatible server based off starlette. This makes it easier for users setup apps with minimal code. Fix some invalid references missed in the package refactor for cli and examples. Made enums, and exceptions private to the package, so they are not exposed in the public API.
1 parent 3291892 commit 214a6c7

File tree

12 files changed

+533
-384
lines changed

12 files changed

+533
-384
lines changed

docs/ag-ui.md

Lines changed: 34 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@ The team at [Rocket Science](https://www.rocketscience.gg/), contributed the
1212
protocol with PydanticAI agents.
1313

1414
This also includes an [`Agent.to_ag_ui`][pydantic_ai.Agent.to_ag_ui] convenience
15-
method which simplifies the creation of [`Adapter`][pydantic_ai.ag_ui.Adapter]
16-
for PydanticAI agents, which can then be used by as part of a
17-
[fastapi](https://fastapi.tiangolo.com/) app.
15+
method which simplifies the creation of [`FastAGUI`][pydantic_ai.ag_ui.FastAGUI]
16+
for PydanticAI agents, which is built on top of [Starlette](https://www.starlette.io/),
17+
meaning it's fully compatible with any ASGI server.
1818

1919
## AG-UI Adapter
2020

@@ -27,8 +27,6 @@ for all aspects of spec including:
2727
- [State Management](https://docs.ag-ui.com/concepts/state)
2828
- [Tools](https://docs.ag-ui.com/concepts/tools)
2929

30-
Let's have a quick look at how to use it:
31-
3230
### Installation
3331

3432
The only dependencies are:
@@ -40,14 +38,14 @@ The only dependencies are:
4038

4139
To run the examples you'll also need:
4240

43-
- [fastapi](https://fastapi.tiangolo.com/): to provide ASGI compatible server
41+
- [uvicorn](https://www.uvicorn.org/) or another ASGI compatible server
4442

4543
```bash
46-
pip/uv-add 'fastapi'
44+
pip/uv-add 'uvicorn'
4745
```
4846

49-
You can install PydanticAI with the `ag-ui` extra to include
50-
[Adapter][pydantic_ai.ag_ui.Adapter] run:
47+
You can install PydanticAI with the `ag-ui` extra to ensure you have all the
48+
required AG-UI dependencies:
5149

5250
```bash
5351
pip/uv-add 'pydantic-ai-slim[ag-ui]'
@@ -60,30 +58,10 @@ pip/uv-add 'pydantic-ai-slim[ag-ui]'
6058

6159
from __future__ import annotations
6260

63-
from typing import TYPE_CHECKING, Annotated
64-
65-
from fastapi import FastAPI, Header
66-
from fastapi.responses import StreamingResponse
67-
from pydantic_ai.ag_ui import SSE_CONTENT_TYPE
68-
6961
from pydantic_ai import Agent
7062

71-
if TYPE_CHECKING:
72-
from ag_ui.core import RunAgentInput
73-
7463
agent = Agent('openai:gpt-4.1', instructions='Be fun!')
75-
adapter = agent.to_ag_ui()
76-
app = FastAPI(title='AG-UI Endpoint')
77-
78-
79-
@app.post('/')
80-
async def root(
81-
input_data: RunAgentInput, accept: Annotated[str, Header()] = SSE_CONTENT_TYPE
82-
) -> StreamingResponse:
83-
return StreamingResponse(
84-
adapter.run(input_data, accept),
85-
media_type=SSE_CONTENT_TYPE,
86-
)
64+
app = agent.to_ag_ui()
8765
```
8866

8967
You can run the example with:
@@ -109,13 +87,16 @@ streamed back to the caller as Server-Sent Events (SSE).
10987
A user request may require multiple round trips between client UI and PydanticAI
11088
server, depending on the tools and events needed.
11189

112-
[Adapter][pydantic_ai.ag_ui.Adapter] can be used with any ASGI server.
90+
In addition to the [Adapter][pydantic_ai.ag_ui.Adapter] there is also
91+
[FastAGUI][pydantic_ai.ag_ui.FastAGUI] which is slim wrapper around
92+
[Starlette](https://www.starlette.io/) providing easy access to run a PydanticAI
93+
server with AG-UI support with any ASGI server.
11394

11495
### Features
11596

11697
To expose a PydanticAI agent as an AG-UI server including state support, you can
117-
use the [`to_ag_ui`][pydantic_ai.agent.Agent.to_ag_ui] method in combination
118-
with [fastapi](https://fastapi.tiangolo.com/).
98+
use the [`to_ag_ui`][pydantic_ai.agent.Agent.to_ag_ui] method create an ASGI
99+
compatible server.
119100

120101
In the example below we have document state which is shared between the UI and
121102
server using the [`StateDeps`][pydantic_ai.ag_ui.StateDeps] which implements the
@@ -134,17 +115,10 @@ real-time synchronization between agents and frontend applications.
134115

135116
from __future__ import annotations
136117

137-
from typing import TYPE_CHECKING, Annotated
138-
139-
from fastapi import FastAPI, Header
140-
from fastapi.responses import StreamingResponse
141118
from pydantic import BaseModel
142-
from pydantic_ai.ag_ui import SSE_CONTENT_TYPE, StateDeps
143119

144120
from pydantic_ai import Agent
145-
146-
if TYPE_CHECKING:
147-
from ag_ui.core import RunAgentInput
121+
from pydantic_ai.ag_ui import StateDeps
148122

149123

150124
class DocumentState(BaseModel):
@@ -158,29 +132,27 @@ agent = Agent(
158132
instructions='Be fun!',
159133
deps_type=StateDeps[DocumentState],
160134
)
161-
adapter = agent.to_ag_ui()
162-
app = FastAPI(title='AG-UI Endpoint')
163-
164-
165-
@app.post('/')
166-
async def root(
167-
input_data: RunAgentInput, accept: Annotated[str, Header()] = SSE_CONTENT_TYPE
168-
) -> StreamingResponse:
169-
return StreamingResponse(
170-
adapter.run(input_data, accept, deps=StateDeps(state_type=DocumentState)),
171-
media_type=SSE_CONTENT_TYPE,
172-
)
135+
app = agent.to_ag_ui(deps=StateDeps(state_type=DocumentState))
173136
```
174137

175-
Since `app` is an ASGI application, it can be used with any ASGI server.
138+
Since `app` is an ASGI application, it can be used with any ASGI server e.g.
176139

177140
```bash
178141
uvicorn agent_to_ag_ui:app --host 0.0.0.0 --port 8000
179142
```
180143

181144
Since the goal of [`to_ag_ui`][pydantic_ai.agent.Agent.to_ag_ui] is to be a
182-
convenience method, it accepts the same arguments as the
183-
[`Adapter`][pydantic_ai.ag_ui.Adapter] constructor.
145+
convenience method, it accepts the same a combination of the arguments require
146+
for:
147+
148+
- [`Adapter`][pydantic_ai.ag_ui.Adapter] constructor
149+
- [`Agent.iter`][pydantic_ai.agent.Agent.iter] method
150+
151+
If you want more control you can either use
152+
[`agent_to_ag_ui`][pydantic_ai.ag_ui.agent_to_ag_ui] helper method or create
153+
and [`Agent`][pydantic_ai.ag_ui.Agent] directly which also provide
154+
the ability to customise [`Starlette`](https://www.starlette.io/applications/#starlette.applications.Starlette)
155+
options.
184156

185157
#### Tools
186158

@@ -200,18 +172,16 @@ for custom events and state updates.
200172

201173
from __future__ import annotations
202174

203-
from typing import TYPE_CHECKING, Annotated
175+
from typing import TYPE_CHECKING
204176

205177
from ag_ui.core import CustomEvent, EventType, StateSnapshotEvent
206-
from fastapi import FastAPI, Header
207-
from fastapi.responses import StreamingResponse
208178
from pydantic import BaseModel
209-
from pydantic_ai.ag_ui import SSE_CONTENT_TYPE, StateDeps
210179

211180
from pydantic_ai import Agent, RunContext
181+
from pydantic_ai.ag_ui import StateDeps
212182

213183
if TYPE_CHECKING:
214-
from ag_ui.core import RunAgentInput
184+
pass
215185

216186

217187
class DocumentState(BaseModel):
@@ -225,8 +195,7 @@ agent = Agent(
225195
instructions='Be fun!',
226196
deps_type=StateDeps[DocumentState],
227197
)
228-
adapter = agent.to_ag_ui()
229-
app = FastAPI(title='AG-UI Endpoint')
198+
app = agent.to_ag_ui(deps=StateDeps(state_type=DocumentState))
230199

231200

232201
@agent.tool
@@ -251,16 +220,6 @@ def custom_events() -> list[CustomEvent]:
251220
value=2,
252221
),
253222
]
254-
255-
256-
@app.post('/')
257-
async def root(
258-
input_data: RunAgentInput, accept: Annotated[str, Header()] = SSE_CONTENT_TYPE
259-
) -> StreamingResponse:
260-
return StreamingResponse(
261-
adapter.run(input_data, accept, deps=StateDeps(state_type=DocumentState)),
262-
media_type=SSE_CONTENT_TYPE,
263-
)
264223
```
265224

266225
### Examples
@@ -296,11 +255,11 @@ options:
296255
Run with adapter debug logging:
297256

298257
```shell
299-
python -m pydantic_ai.ag_ui_examples.dojo_server --log-level debug
258+
python -m pydantic_ai_ag_ui_examples.dojo_server --log-level debug
300259
```
301260

302261
Using uvicorn:
303262

304263
```shell
305-
uvicorn pydantic_ai.ag_ui_examples.dojo_server:app --port 9000
264+
uvicorn pydantic_ai_ag_ui_examples.dojo_server:app --port 9000
306265
```

examples/pydantic_ai_ag_ui_examples/api/agent.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,4 +40,4 @@ def __init__(
4040
instructions=instructions,
4141
deps_type=deps_type,
4242
)
43-
self.adapter = self.agent.to_ag_ui()
43+
self.adapter = Adapter(agent=self.agent)

examples/pydantic_ai_ag_ui_examples/basic.py

Lines changed: 2 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -2,44 +2,13 @@
22

33
from __future__ import annotations
44

5-
from typing import TYPE_CHECKING, Annotated
6-
7-
from fastapi import FastAPI, Header
8-
from fastapi.responses import StreamingResponse
9-
105
from pydantic_ai import Agent
11-
from pydantic_ai.ag_ui import SSE_CONTENT_TYPE, Adapter
12-
13-
if TYPE_CHECKING:
14-
from ag_ui.core import RunAgentInput
15-
16-
app = FastAPI(title='AG-UI Endpoint')
176

187
agent: Agent[None, str] = Agent(
198
'openai:gpt-4o-mini',
209
instructions='You are a helpful assistant.',
2110
)
22-
adapter: Adapter[None, str] = agent.to_ag_ui()
23-
24-
25-
@app.post('/agent')
26-
async def handler(
27-
input_data: RunAgentInput, accept: Annotated[str, Header()] = SSE_CONTENT_TYPE
28-
) -> StreamingResponse:
29-
"""Endpoint to handle AG-UI protocol requests and stream responses.
30-
31-
Args:
32-
input_data: The AG-UI run input.
33-
accept: The Accept header to specify the response format.
34-
35-
Returns:
36-
A streaming response with event-stream media type.
37-
"""
38-
return StreamingResponse(
39-
adapter.run(input_data, accept),
40-
media_type=SSE_CONTENT_TYPE,
41-
)
42-
11+
app = agent.to_ag_ui()
4312

4413
if __name__ == '__main__':
4514
import uvicorn
@@ -49,7 +18,7 @@ async def handler(
4918
args: Args = parse_args()
5019

5120
uvicorn.run(
52-
'pydantic_ai.ag_ui_examples.dojo_server:app',
21+
'pydantic_ai_ag_ui_examples.dojo_server:app',
5322
port=args.port,
5423
reload=args.reload,
5524
log_level=args.log_level,

examples/pydantic_ai_ag_ui_examples/cli/args.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ def parse_args() -> Args:
6565
'--loggers',
6666
nargs='*',
6767
default=[
68-
'pydantic_ai.ag_ui.adapter',
68+
'pydantic_ai.ag_ui',
6969
],
7070
help='Logger names to configure (default: adapter and model loggers)',
7171
)

examples/pydantic_ai_ag_ui_examples/dojo_server.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
args: Args = parse_args()
4242

4343
uvicorn.run(
44-
'pydantic_ai.ag_ui_examples.dojo_server:app',
44+
'pydantic_ai_ag_ui_examples.dojo_server:app',
4545
port=args.port,
4646
reload=args.reload,
4747
log_config=args.log_config(),

mkdocs.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ nav:
6363
- examples/question-graph.md
6464
- examples/slack-lead-qualifier.md
6565
- API Reference:
66+
- api/ag_ui.md
6667
- api/agent.md
6768
- api/tools.md
6869
- api/common_tools.md
@@ -104,7 +105,6 @@ nav:
104105
- api/pydantic_evals/otel.md
105106
- api/pydantic_evals/generation.md
106107
- api/fasta2a.md
107-
- api/ag_ui.md
108108

109109
extra:
110110
# hide the "Made with Material for MkDocs" message

0 commit comments

Comments
 (0)