-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Add Agent.to_web() method and web chat UI #3456
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 1 commit
1cfa48e
d83caa5
bfffd4b
055e120
6bc8b16
32f7e1d
cf0e177
e7f44eb
c4ffde3
0595c27
0d24941
e5b30c2
f2dd19a
8ca7a27
fa3bb5f
e45c93f
f90b570
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,16 @@ | ||
| """Web-based chat UI for Pydantic AI agents.""" | ||
|
|
||
| from .agent_options import AI_MODELS, BUILTIN_TOOLS, AIModel, AIModelID, BuiltinTool, BuiltinToolID | ||
| from .api import create_api_router | ||
| from .app import create_chat_app | ||
|
|
||
| __all__ = [ | ||
| 'create_chat_app', | ||
| 'create_api_router', | ||
| 'AI_MODELS', | ||
| 'BUILTIN_TOOLS', | ||
| 'AIModel', | ||
| 'AIModelID', | ||
| 'BuiltinTool', | ||
| 'BuiltinToolID', | ||
| ] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| """Model and builtin tool configurations for the web chat UI.""" | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The goal is for the developer to be able to pass a list of models and builtin tools to |
||
|
|
||
| from typing import Literal | ||
|
|
||
| from pydantic import BaseModel | ||
|
|
||
| from pydantic_ai.builtin_tools import ( | ||
| AbstractBuiltinTool, | ||
| CodeExecutionTool, | ||
| ImageGenerationTool, | ||
| WebSearchTool, | ||
| ) | ||
|
|
||
| AIModelID = Literal[ | ||
| 'anthropic:claude-sonnet-4-5', | ||
| 'openai-responses:gpt-5', | ||
| 'google-gla:gemini-2.5-pro', | ||
| ] | ||
| BuiltinToolID = Literal['web_search', 'image_generation', 'code_execution'] | ||
|
|
||
|
|
||
| class AIModel(BaseModel): | ||
| """Defines an AI model with its associated built-in tools.""" | ||
|
|
||
| id: AIModelID | ||
| name: str | ||
| builtin_tools: list[BuiltinToolID] | ||
|
|
||
|
|
||
| class BuiltinTool(BaseModel): | ||
| """Defines a built-in tool.""" | ||
|
|
||
| id: BuiltinToolID | ||
| name: str | ||
|
|
||
|
|
||
| BUILTIN_TOOL_DEFS: list[BuiltinTool] = [ | ||
dsfaccini marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| BuiltinTool(id='web_search', name='Web Search'), | ||
| BuiltinTool(id='code_execution', name='Code Execution'), | ||
| BuiltinTool(id='image_generation', name='Image Generation'), | ||
| ] | ||
|
|
||
| BUILTIN_TOOLS: dict[BuiltinToolID, AbstractBuiltinTool] = { | ||
| 'web_search': WebSearchTool(), | ||
| 'code_execution': CodeExecutionTool(), | ||
| 'image_generation': ImageGenerationTool(), | ||
| } | ||
|
|
||
| AI_MODELS: list[AIModel] = [ | ||
| AIModel( | ||
| id='anthropic:claude-sonnet-4-5', | ||
| name='Claude Sonnet 4.5', | ||
| builtin_tools=[ | ||
| 'web_search', | ||
| 'code_execution', | ||
| ], | ||
| ), | ||
| AIModel( | ||
| id='openai-responses:gpt-5', | ||
| name='GPT 5', | ||
| builtin_tools=[ | ||
| 'web_search', | ||
| 'code_execution', | ||
| 'image_generation', | ||
| ], | ||
| ), | ||
| AIModel( | ||
| id='google-gla:gemini-2.5-pro', | ||
| name='Gemini 2.5 Pro', | ||
| builtin_tools=[ | ||
| 'web_search', | ||
| 'code_execution', | ||
| ], | ||
| ), | ||
| ] | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,80 @@ | ||
| """API router for the web chat UI.""" | ||
|
|
||
| from typing import Annotated | ||
|
|
||
| from fastapi import APIRouter, Depends, Request, Response | ||
| from pydantic import BaseModel | ||
| from pydantic.alias_generators import to_camel | ||
|
|
||
| from pydantic_ai import Agent | ||
| from pydantic_ai.ui.vercel_ai._adapter import VercelAIAdapter | ||
|
|
||
| from .agent_options import ( | ||
| AI_MODELS, | ||
| BUILTIN_TOOL_DEFS, | ||
| BUILTIN_TOOLS, | ||
| AIModel, | ||
| AIModelID, | ||
| BuiltinTool, | ||
| BuiltinToolID, | ||
| ) | ||
|
|
||
|
|
||
| def get_agent(request: Request) -> Agent: | ||
| """Get the agent from app state.""" | ||
| agent = getattr(request.app.state, 'agent', None) | ||
| if agent is None: | ||
| raise RuntimeError('No agent configured. Server must be started with a valid agent.') | ||
| return agent | ||
|
|
||
|
|
||
| def create_api_router() -> APIRouter: | ||
| """Create the API router for chat endpoints.""" | ||
| router = APIRouter() | ||
|
|
||
| @router.options('/api/chat') | ||
| def options_chat(): # pyright: ignore[reportUnusedFunction] | ||
| """Handle CORS preflight requests.""" | ||
| pass | ||
|
|
||
| class ConfigureFrontend(BaseModel, alias_generator=to_camel, populate_by_name=True): | ||
| """Response model for frontend configuration.""" | ||
|
|
||
| models: list[AIModel] | ||
| builtin_tools: list[BuiltinTool] | ||
|
|
||
| @router.get('/api/configure') | ||
| async def configure_frontend() -> ConfigureFrontend: # pyright: ignore[reportUnusedFunction] | ||
| """Endpoint to configure the frontend with available models and tools.""" | ||
| return ConfigureFrontend( | ||
| models=AI_MODELS, | ||
| builtin_tools=BUILTIN_TOOL_DEFS, | ||
| ) | ||
|
|
||
| @router.get('/api/health') | ||
| async def health() -> dict[str, bool]: # pyright: ignore[reportUnusedFunction] | ||
| """Health check endpoint.""" | ||
| return {'ok': True} | ||
|
|
||
| class ChatRequestExtra(BaseModel, extra='ignore', alias_generator=to_camel): | ||
| """Extra data extracted from chat request.""" | ||
|
|
||
| model: AIModelID | None = None | ||
| builtin_tools: list[BuiltinToolID] = [] | ||
|
|
||
| @router.post('/api/chat') | ||
| async def post_chat( # pyright: ignore[reportUnusedFunction] | ||
| request: Request, agent: Annotated[Agent, Depends(get_agent)] | ||
| ) -> Response: | ||
| """Handle chat requests via Vercel AI Adapter.""" | ||
| adapter = await VercelAIAdapter.from_request(request, agent=agent) | ||
| extra_data = ChatRequestExtra.model_validate(adapter.run_input.__pydantic_extra__) | ||
| streaming_response = await VercelAIAdapter.dispatch_request( | ||
| request, | ||
| agent=agent, | ||
| model=extra_data.model, | ||
| builtin_tools=[BUILTIN_TOOLS[tool_id] for tool_id in extra_data.builtin_tools], | ||
| ) | ||
| return streaming_response | ||
|
|
||
| return router |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,47 @@ | ||
| """Factory function for creating a web chat app for a Pydantic AI agent.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| from typing import TypeVar | ||
|
|
||
| import fastapi | ||
| import httpx | ||
| from fastapi import Request | ||
| from fastapi.responses import HTMLResponse | ||
|
|
||
| from pydantic_ai import Agent | ||
|
|
||
| from .api import create_api_router | ||
|
|
||
| CDN_URL = 'https://cdn.jsdelivr.net/npm/@pydantic/[email protected]/dist/index.html' | ||
dsfaccini marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
|
|
||
| AgentDepsT = TypeVar('AgentDepsT') | ||
| OutputDataT = TypeVar('OutputDataT') | ||
|
|
||
|
|
||
| def create_chat_app( | ||
| agent: Agent[AgentDepsT, OutputDataT], | ||
| ) -> fastapi.FastAPI: | ||
| """Create a FastAPI app that serves a web chat UI for the given agent. | ||
| Args: | ||
| agent: The Pydantic AI agent to serve | ||
| Returns: | ||
| A configured FastAPI application ready to be served | ||
| """ | ||
| app = fastapi.FastAPI() | ||
|
|
||
| app.state.agent = agent | ||
|
|
||
| app.include_router(create_api_router()) | ||
|
|
||
| @app.get('/') | ||
| @app.get('/{id}') | ||
| async def index(request: Request): # pyright: ignore[reportUnusedFunction] | ||
| """Serve the chat UI from CDN.""" | ||
| async with httpx.AsyncClient() as client: | ||
| response = await client.get(CDN_URL) | ||
|
||
| return HTMLResponse(content=response.content, status_code=response.status_code) | ||
|
|
||
| return app | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,101 @@ | ||
| """Tests for the web chat UI module.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import pytest | ||
|
|
||
| from pydantic_ai import Agent | ||
| from pydantic_ai.ui.web import AI_MODELS, BUILTIN_TOOLS, create_chat_app | ||
|
|
||
| from .conftest import try_import | ||
|
|
||
| with try_import() as fastapi_import_successful: | ||
| from fastapi import FastAPI | ||
| from fastapi.testclient import TestClient | ||
|
|
||
| pytestmark = [ | ||
| pytest.mark.skipif(not fastapi_import_successful, reason='fastapi not installed'), | ||
| ] | ||
|
|
||
|
|
||
| def test_create_chat_app_basic(): | ||
| """Test creating a basic chat app.""" | ||
| agent = Agent('test') | ||
| app = create_chat_app(agent) | ||
|
|
||
| assert isinstance(app, FastAPI) | ||
| assert app.state.agent is agent | ||
|
|
||
|
|
||
| def test_agent_to_web(): | ||
| """Test the Agent.to_web() method.""" | ||
| agent = Agent('test') | ||
| app = agent.to_web() | ||
|
|
||
| assert isinstance(app, FastAPI) | ||
| assert app.state.agent is agent | ||
|
|
||
|
|
||
| def test_chat_app_health_endpoint(): | ||
| """Test the /api/health endpoint.""" | ||
| agent = Agent('test') | ||
| app = create_chat_app(agent) | ||
|
|
||
| with TestClient(app) as client: | ||
| response = client.get('/api/health') | ||
| assert response.status_code == 200 | ||
| assert response.json() == {'ok': True} | ||
|
|
||
|
|
||
| def test_chat_app_configure_endpoint(): | ||
| """Test the /api/configure endpoint.""" | ||
| agent = Agent('test') | ||
| app = create_chat_app(agent) | ||
|
|
||
| with TestClient(app) as client: | ||
| response = client.get('/api/configure') | ||
| assert response.status_code == 200 | ||
| data = response.json() | ||
| assert 'models' in data | ||
|
||
| assert 'builtinTools' in data # camelCase due to alias generator | ||
|
|
||
| # no snapshot bc we're checking against the actual model/tool definitions | ||
| assert len(data['models']) == len(AI_MODELS) | ||
| assert len(data['builtinTools']) == len(BUILTIN_TOOLS) | ||
|
|
||
|
|
||
| def test_chat_app_index_endpoint(): | ||
| """Test that the index endpoint serves the UI from CDN.""" | ||
| agent = Agent('test') | ||
| app = create_chat_app(agent) | ||
|
|
||
| with TestClient(app) as client: | ||
| response = client.get('/') | ||
| assert response.status_code == 200 | ||
| assert response.headers['content-type'] == 'text/html; charset=utf-8' | ||
| assert len(response.content) > 0 | ||
|
|
||
|
|
||
| def test_ai_models_configuration(): | ||
| """Test that AI models are configured correctly.""" | ||
| assert len(AI_MODELS) == 3 | ||
|
|
||
| model_ids = {model.id for model in AI_MODELS} | ||
| assert 'anthropic:claude-sonnet-4-5' in model_ids | ||
| assert 'openai-responses:gpt-5' in model_ids | ||
| assert 'google-gla:gemini-2.5-pro' in model_ids | ||
|
|
||
|
|
||
| def test_builtin_tools_configuration(): | ||
| """Test that builtin tools are configured correctly.""" | ||
| assert len(BUILTIN_TOOLS) == 3 | ||
|
|
||
| assert 'web_search' in BUILTIN_TOOLS | ||
| assert 'code_execution' in BUILTIN_TOOLS | ||
| assert 'image_generation' in BUILTIN_TOOLS | ||
|
|
||
| from pydantic_ai.builtin_tools import CodeExecutionTool, ImageGenerationTool, WebSearchTool | ||
|
|
||
| assert isinstance(BUILTIN_TOOLS['web_search'], WebSearchTool) | ||
| assert isinstance(BUILTIN_TOOLS['code_execution'], CodeExecutionTool) | ||
| assert isinstance(BUILTIN_TOOLS['image_generation'], ImageGenerationTool) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We're gonna need some args here -- have a look at the
to_a2aandto_ag_uimethods. Not saying we need all of those args, but some may be useful