move to toolsets

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/_cli/_web.py

@@ -4,22 +4,20 @@
 
 import sys
 from pathlib import Path
+from typing import Any
 
-from pydantic import BaseModel, ImportString, ValidationError
+from pydantic import ImportString, TypeAdapter, ValidationError
 from rich.console import Console
 
 from pydantic_ai import Agent
 from pydantic_ai.builtin_tools import AbstractBuiltinTool, get_builtin_tool_cls
+from pydantic_ai.mcp import MCPServerSSE, MCPServerStdio, MCPServerStreamableHTTP, load_mcp_servers
 from pydantic_ai.models import infer_model
-from pydantic_ai.ui._web import create_web_app, load_mcp_server_tools
+from pydantic_ai.ui._web import create_web_app
 
 __all__ = ['_run_web_command']
 
-
-class _AgentLoader(BaseModel):
-    """Helper model for loading agents using Pydantic ImportString."""
-
-    agent: ImportString  # type: ignore[valid-type]
+_import_string_adapter: TypeAdapter[Any] = TypeAdapter(ImportString)
 
 
 def _load_agent(agent_path: str) -> Agent | None:
@@ -33,11 +31,10 @@ def _load_agent(agent_path: str) -> Agent | None:
     """
     sys.path.insert(0, str(Path.cwd()))
     try:
-        loader = _AgentLoader(agent=agent_path)
-        agent = loader.agent  # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType]
-        if not isinstance(agent, Agent):
+        obj = _import_string_adapter.validate_python(agent_path)
+        if not isinstance(obj, Agent):
             return None
-        return agent  # pyright: ignore[reportUnknownVariableType]
+        return obj  # pyright: ignore[reportUnknownVariableType]
     except ValidationError:
         return None
 
@@ -70,13 +67,7 @@ def _run_web_command(  # noqa: C901
             console.print(f'[red]Error: Could not load agent from {agent_path}[/red]')
             return 1
     else:
-        agent = Agent()
-
-        if instructions:
-
-            @agent.system_prompt
-            def system_prompt() -> str:  # pyright: ignore[reportUnusedFunction]
-                return instructions  # pragma: no cover
+        agent = Agent(instructions=instructions)
 
     if agent.model is None and not models:
         console.print('[red]Error: At least one model (-m) is required when agent has no model[/red]')
@@ -105,12 +96,12 @@ def system_prompt() -> str:  # pyright: ignore[reportUnusedFunction]
                 continue
             all_tool_instances.append(tool_cls())
 
-    # Load MCP server tools if specified
+    # Load MCP servers as toolsets if specified
+    mcp_servers: list[MCPServerStdio | MCPServerStreamableHTTP | MCPServerSSE] = []
     if mcp:
         try:
-            mcp_tools = load_mcp_server_tools(mcp)
-            all_tool_instances.extend(mcp_tools)
-            console.print(f'[dim]Loaded {len(mcp_tools)} MCP server(s) from {mcp}[/dim]')
+            mcp_servers = load_mcp_servers(mcp)
+            console.print(f'[dim]Loaded {len(mcp_servers)} MCP server(s) from {mcp}[/dim]')
         except FileNotFoundError as e:
             console.print(f'[red]Error: {e}[/red]')
             return 1
@@ -125,6 +116,7 @@ def system_prompt() -> str:  # pyright: ignore[reportUnusedFunction]
         agent,
         models=models,
         builtin_tools=all_tool_instances if all_tool_instances else None,
+        toolsets=mcp_servers if mcp_servers else None,
     )
 
     agent_desc = agent_path if agent_path else 'generic agent'

pydantic_ai_slim/pydantic_ai/builtin_tools.py

@@ -417,7 +417,7 @@ def _tool_discriminator(tool_data: dict[str, Any] | AbstractBuiltinTool) -> str:
         return tool_data.kind
 
 
-def get_builtin_tool_ids() -> frozenset[str]:
+def get_builtin_tool_ids() -> frozenset[str]:  # pragma: no cover
     """Get the set of all builtin tool IDs (excluding deprecated tools like url_context)."""
     return frozenset(_BUILTIN_TOOL_TYPES.keys() - {'url_context'})
 

pydantic_ai_slim/pydantic_ai/models/test.py

@@ -96,6 +96,7 @@ def __init__(
         custom_output_text: str | None = None,
         custom_output_args: Any | None = None,
         seed: int = 0,
+        model_name: str = 'test',
         profile: ModelProfileSpec | None = None,
         settings: ModelSettings | None = None,
     ):
@@ -105,7 +106,7 @@ def __init__(
         self.custom_output_args = custom_output_args
         self.seed = seed
         self.last_model_request_parameters = None
-        self._model_name = 'test'
+        self._model_name = model_name
         self._system = 'test'
         super().__init__(settings=settings, profile=profile)
 

pydantic_ai_slim/pydantic_ai/ui/_web/api.py

@@ -1,5 +1,8 @@
 """API routes for the web chat UI."""
 
+from collections.abc import Sequence
+from typing import Any
+
 from pydantic import BaseModel
 from pydantic.alias_generators import to_camel
 from starlette.applications import Starlette
@@ -8,6 +11,7 @@
 
 from pydantic_ai import Agent
 from pydantic_ai.builtin_tools import AbstractBuiltinTool
+from pydantic_ai.toolsets import AbstractToolset
 from pydantic_ai.ui.vercel_ai import VercelAIAdapter
 
 
@@ -52,17 +56,20 @@ def add_api_routes(
     app: Starlette,
     models: list[ModelInfo] | None = None,
     builtin_tools: list[AbstractBuiltinTool] | None = None,
+    toolsets: Sequence[AbstractToolset[Any]] | None = None,
 ) -> None:
     """Add API routes to a Starlette app.
 
     Args:
         app: The Starlette app to add routes to.
         models: Optional list of AI models. If not provided, the UI will have no model options.
         builtin_tools: Optional list of builtin tools. If not provided, no tools will be available.
+        toolsets: Optional list of toolsets (e.g., MCP servers). These provide additional tools.
     """
     _models = models or []
     _model_ids = {m.id for m in _models}
     _builtin_tools = builtin_tools or []
+    _toolsets = list(toolsets) if toolsets else None
     _tools_by_id: dict[str, AbstractBuiltinTool] = {tool.unique_id: tool for tool in _builtin_tools}
 
     async def options_chat(request: Request) -> Response:
@@ -102,6 +109,7 @@ async def post_chat(request: Request) -> Response:
             agent=agent,
             model=extra_data.model,
             builtin_tools=request_builtin_tools,
+            toolsets=_toolsets,
         )
         return streaming_response
 

pydantic_ai_slim/pydantic_ai/ui/_web/app.py

@@ -1,5 +1,3 @@
-"""Factory function for creating a web chat app for a Pydantic AI agent."""
-
 from __future__ import annotations
 
 import os