remove get agent utility and make supported builtin tools a classmethod

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -509,11 +509,9 @@ def label(self) -> str:
         """Human-friendly display label for the model."""
         return _format_model_label(self.model_name)
 
-    def supported_builtin_tools(self, profile: ModelProfile) -> frozenset[type[AbstractBuiltinTool]]:
-        """Return the set of builtin tool types this model can handle.
-
-        Args:
-            profile: The resolved model profile (passed to avoid circular dependency with self.profile)
+    @classmethod
+    def supported_builtin_tools(cls) -> frozenset[type[AbstractBuiltinTool]]:
+        """Return the set of builtin tool types this model class can handle.
 
         Subclasses should override this to reflect their actual capabilities.
         Default is empty set - subclasses must explicitly declare support.
@@ -536,7 +534,7 @@ def profile(self) -> ModelProfile:
             _profile = DEFAULT_PROFILE
 
         # Compute intersection: profile's allowed tools & model's implemented tools
-        model_supported = self.supported_builtin_tools(_profile)
+        model_supported = self.__class__.supported_builtin_tools()
         profile_supported = _profile.supported_builtin_tools
         effective_tools = profile_supported & model_supported
 

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -44,7 +44,7 @@
     ToolReturnPart,
     UserPromptPart,
 )
-from ..profiles import ModelProfile, ModelProfileSpec
+from ..profiles import ModelProfileSpec
 from ..providers import Provider, infer_provider
 from ..providers.anthropic import AsyncAnthropicClient
 from ..settings import ModelSettings, merge_model_settings
@@ -261,7 +261,8 @@ def system(self) -> str:
         """The model provider."""
         return self._provider.name
 
-    def supported_builtin_tools(self, profile: ModelProfile) -> frozenset[type[AbstractBuiltinTool]]:
+    @classmethod
+    def supported_builtin_tools(cls) -> frozenset[type[AbstractBuiltinTool]]:
         """Return the set of builtin tool types this model can handle."""
         return frozenset({WebSearchTool, CodeExecutionTool, WebFetchTool, MemoryTool, MCPServerTool})
 

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -201,7 +201,8 @@ def system(self) -> str:
         """The system / model provider."""
         return self._system
 
-    def supported_builtin_tools(self, profile: ModelProfile) -> frozenset[type[AbstractBuiltinTool]]:
+    @classmethod
+    def supported_builtin_tools(cls) -> frozenset[type[AbstractBuiltinTool]]:
         """FunctionModel supports all builtin tools for testing flexibility."""
         from ..builtin_tools import get_builtin_tool_types
 

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -37,7 +37,7 @@
     UserPromptPart,
     VideoUrl,
 )
-from ..profiles import ModelProfile, ModelProfileSpec
+from ..profiles import ModelProfileSpec
 from ..profiles.google import GoogleModelProfile
 from ..providers import Provider, infer_provider
 from ..settings import ModelSettings
@@ -229,7 +229,8 @@ def system(self) -> str:
         """The model provider."""
         return self._provider.name
 
-    def supported_builtin_tools(self, profile: ModelProfile) -> frozenset[type[AbstractBuiltinTool]]:
+    @classmethod
+    def supported_builtin_tools(cls) -> frozenset[type[AbstractBuiltinTool]]:
         """Return the set of builtin tool types this model can handle."""
         return frozenset({WebSearchTool, CodeExecutionTool, WebFetchTool, ImageGenerationTool})
 

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -179,7 +179,8 @@ def system(self) -> str:
         """The model provider."""
         return self._provider.name
 
-    def supported_builtin_tools(self, profile: ModelProfile) -> frozenset[type[AbstractBuiltinTool]]:
+    @classmethod
+    def supported_builtin_tools(cls) -> frozenset[type[AbstractBuiltinTool]]:
         """Return the set of builtin tool types this model can handle."""
         return frozenset({WebSearchTool})
 

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -8,6 +8,7 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field, replace
 from datetime import datetime
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, Literal, cast, overload
 
 from pydantic import ValidationError
@@ -424,15 +425,23 @@ def system(self) -> str:
         """The model provider."""
         return self._provider.name
 
-    def supported_builtin_tools(self, profile: ModelProfile) -> frozenset[type[AbstractBuiltinTool]]:
-        """Return the set of builtin tool types this model can handle.
+    @classmethod
+    def supported_builtin_tools(cls) -> frozenset[type[AbstractBuiltinTool]]:
+        """Return the set of builtin tool types this model can handle."""
+        return frozenset({WebSearchTool})
+
+    @cached_property
+    def profile(self) -> ModelProfile:
+        """The model profile.
 
         WebSearchTool is only supported if openai_chat_supports_web_search is True.
         """
-        openai_profile = OpenAIModelProfile.from_profile(profile)
-        if openai_profile.openai_chat_supports_web_search:
-            return frozenset({WebSearchTool})
-        return frozenset()
+        _profile = super().profile
+        openai_profile = OpenAIModelProfile.from_profile(_profile)
+        if not openai_profile.openai_chat_supports_web_search:
+            new_tools = _profile.supported_builtin_tools - {WebSearchTool}
+            _profile = replace(_profile, supported_builtin_tools=new_tools)
+        return _profile
 
     @property
     @deprecated('Set the `system_prompt_role` in the `OpenAIModelProfile` instead.')
@@ -1092,7 +1101,8 @@ def system(self) -> str:
         """The model provider."""
         return self._provider.name
 
-    def supported_builtin_tools(self, profile: ModelProfile) -> frozenset[type[AbstractBuiltinTool]]:
+    @classmethod
+    def supported_builtin_tools(cls) -> frozenset[type[AbstractBuiltinTool]]:
         """Return the set of builtin tool types this model can handle."""
         return frozenset({WebSearchTool, CodeExecutionTool, MCPServerTool, ImageGenerationTool})
 

pydantic_ai_slim/pydantic_ai/models/test.py

@@ -30,7 +30,7 @@
     ToolCallPart,
     ToolReturnPart,
 )
-from ..profiles import ModelProfile, ModelProfileSpec
+from ..profiles import ModelProfileSpec
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from ..usage import RequestUsage
@@ -158,7 +158,8 @@ def system(self) -> str:
         """The model provider."""
         return self._system
 
-    def supported_builtin_tools(self, profile: ModelProfile) -> frozenset[type[AbstractBuiltinTool]]:
+    @classmethod
+    def supported_builtin_tools(cls) -> frozenset[type[AbstractBuiltinTool]]:
         """TestModel supports all builtin tools for testing flexibility."""
         from ..builtin_tools import get_builtin_tool_types
 

pydantic_ai_slim/pydantic_ai/ui/_web/api.py

@@ -23,14 +23,6 @@ class ModelInfo(BaseModel, alias_generator=to_camel, populate_by_name=True):
     builtin_tools: list[str]
 
 
-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
-
-
 class BuiltinToolInfo(BaseModel, alias_generator=to_camel, populate_by_name=True):
     """Serializable info about a builtin tool for frontend config."""
 
@@ -56,6 +48,7 @@ class _ChatRequestExtra(BaseModel, extra='ignore', alias_generator=to_camel):
 
 def add_api_routes(
     app: Starlette,
+    agent: Agent,
     models: list[ModelInfo] | None = None,
     builtin_tools: list[AbstractBuiltinTool] | None = None,
     toolsets: Sequence[AbstractToolset[Any]] | None = None,
@@ -64,6 +57,7 @@ def add_api_routes(
 
     Args:
         app: The Starlette app to add routes to.
+        agent: Agent instance.
         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.
@@ -92,7 +86,6 @@ async def health(request: Request) -> Response:
 
     async def post_chat(request: Request) -> Response:
         """Handle chat requests via Vercel AI Adapter."""
-        agent = _get_agent(request)
         adapter = await VercelAIAdapter.from_request(request, agent=agent)
         extra_data = _ChatRequestExtra.model_validate(adapter.run_input.__pydantic_extra__)
 

pydantic_ai_slim/pydantic_ai/ui/_web/app.py

@@ -127,9 +127,13 @@ def create_web_app(
 
     app = Starlette()
 
-    app.state.agent = agent
-
-    add_api_routes(app, models=resolved_models, builtin_tools=builtin_tools, toolsets=toolsets)
+    add_api_routes(
+        app,
+        agent=agent,  # pyright: ignore[reportArgumentType]
+        models=resolved_models,
+        builtin_tools=builtin_tools,
+        toolsets=toolsets,
+    )
 
     async def index(request: Request) -> Response:
         """Serve the chat UI from filesystem cache or CDN."""

tests/test_ui_web.py

@@ -169,22 +169,6 @@ def test_chat_app_index_caching():
         assert response2.status_code == 200
 
 
-def test_get_agent_missing():
-    """Test that _get_agent raises RuntimeError when agent is not configured."""
-    from pydantic_ai.ui._web.api import _get_agent  # pyright: ignore[reportPrivateUsage]
-
-    app = Starlette()
-
-    class FakeRequest:
-        def __init__(self, app: Starlette):
-            self.app = app
-
-    request = FakeRequest(app)
-
-    with pytest.raises(RuntimeError, match='No agent configured'):
-        _get_agent(request)  # pyright: ignore[reportArgumentType]
-
-
 @pytest.mark.anyio
 async def test_post_chat_endpoint():
     """Test the POST /api/chat endpoint."""