Let almost all types used in docs examples be imported directly from pydantic_ai (#2736)

Author: DouweM

docs/ag-ui.md

@@ -43,24 +43,24 @@ This example uses [`run_ag_ui()`][pydantic_ai.ag_ui.run_ag_ui] and performs its
 This can be modified to work with any web framework.
 
 ```py {title="run_ag_ui.py"}
+import json
+from http import HTTPStatus
+
 from ag_ui.core import RunAgentInput
 from fastapi import FastAPI
-from http import HTTPStatus
 from fastapi.requests import Request
 from fastapi.responses import Response, StreamingResponse
 from pydantic import ValidationError
-import json
 
 from pydantic_ai import Agent
-from pydantic_ai.ag_ui import run_ag_ui, SSE_CONTENT_TYPE
-
+from pydantic_ai.ag_ui import SSE_CONTENT_TYPE, run_ag_ui
 
 agent = Agent('openai:gpt-4.1', instructions='Be fun!')
 
 app = FastAPI()
 
 
-@app.post("/")
+@app.post('/')
 async def run_agent(request: Request) -> Response:
     accept = request.headers.get('accept', SSE_CONTENT_TYPE)
     try:
@@ -97,12 +97,11 @@ from starlette.responses import Response
 from pydantic_ai import Agent
 from pydantic_ai.ag_ui import handle_ag_ui_request
 
-
 agent = Agent('openai:gpt-4.1', instructions='Be fun!')
 
 app = FastAPI()
 
-@app.post("/")
+@app.post('/')
 async def run_agent(request: Request) -> Response:
     return await handle_ag_ui_request(agent, request)
 ```

docs/agents.md

@@ -116,7 +116,7 @@ import asyncio
 from collections.abc import AsyncIterable
 from datetime import date
 
-from pydantic_ai import Agent
+from pydantic_ai import Agent, RunContext
 from pydantic_ai.messages import (
     AgentStreamEvent,
     FinalResultEvent,
@@ -128,7 +128,6 @@ from pydantic_ai.messages import (
     ThinkingPartDelta,
     ToolCallPartDelta,
 )
-from pydantic_ai.tools import RunContext
 
 weather_agent = Agent(
     'openai:gpt-4o',
@@ -215,10 +214,10 @@ Unlike `run_stream()`, it always runs the agent graph to completion even if text
     To get the best of both worlds, at the expense of some additional complexity, you can use [`agent.iter()`][pydantic_ai.agent.AbstractAgent.iter] as described in the next section, which lets you [iterate over the agent graph](#iterating-over-an-agents-graph) and [stream both events and output](#streaming-all-events-and-output) at every step.
 
 ```python {title="run_events.py" requires="run_stream_events.py"}
-from run_stream_events import weather_agent, event_stream_handler, output_messages
-
 import asyncio
 
+from run_stream_events import event_stream_handler, output_messages, weather_agent
+
 
 async def main():
     user_prompt = 'What will the weather be like in Paris on Tuesday?'
@@ -395,7 +394,7 @@ import asyncio
 from dataclasses import dataclass
 from datetime import date
 
-from pydantic_ai import Agent
+from pydantic_ai import Agent, RunContext
 from pydantic_ai.messages import (
     FinalResultEvent,
     FunctionToolCallEvent,
@@ -406,7 +405,6 @@ from pydantic_ai.messages import (
     ThinkingPartDelta,
     ToolCallPartDelta,
 )
-from pydantic_ai.tools import RunContext
 
 
 @dataclass
@@ -548,9 +546,7 @@ You can apply these settings by passing the `usage_limits` argument to the `run{
 Consider the following example, where we limit the number of response tokens:
 
 ```py
-from pydantic_ai import Agent
-from pydantic_ai.exceptions import UsageLimitExceeded
-from pydantic_ai.usage import UsageLimits
+from pydantic_ai import Agent, UsageLimitExceeded, UsageLimits
 
 agent = Agent('anthropic:claude-3-5-sonnet-latest')
 
@@ -578,9 +574,7 @@ Restricting the number of requests can be useful in preventing infinite loops or
 ```py
 from typing_extensions import TypedDict
 
-from pydantic_ai import Agent, ModelRetry
-from pydantic_ai.exceptions import UsageLimitExceeded
-from pydantic_ai.usage import UsageLimits
+from pydantic_ai import Agent, ModelRetry, UsageLimitExceeded, UsageLimits
 
 
 class NeverOutputType(TypedDict):
@@ -636,9 +630,8 @@ For example, if you'd like to set the `temperature` setting to `0.0` to ensure l
 you can do the following:
 
 ```py
-from pydantic_ai import Agent
+from pydantic_ai import Agent, ModelSettings
 from pydantic_ai.models.openai import OpenAIChatModel
-from pydantic_ai.settings import ModelSettings
 
 # 1. Model-level defaults
 model = OpenAIChatModel(

docs/cli.md

@@ -115,7 +115,13 @@ Both `Agent.to_cli()` and `Agent.to_cli_sync()` support a `message_history` para
 
 ```python {title="agent_with_history.py" test="skip"}
 from pydantic_ai import Agent
-from pydantic_ai.messages import ModelMessage, ModelRequest, ModelResponse, UserPromptPart, TextPart
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    TextPart,
+    UserPromptPart,
+)
 
 agent = Agent('openai:gpt-4.1')
 

docs/common-tools.md

@@ -106,7 +106,7 @@ Here's an example of how you can use the Tavily search tool with an agent:
 ```py {title="tavily_search.py" test="skip"}
 import os
 
-from pydantic_ai.agent import Agent
+from pydantic_ai import Agent
 from pydantic_ai.common_tools.tavily import tavily_search_tool
 
 api_key = os.getenv('TAVILY_API_KEY')

docs/direct.md

@@ -40,13 +40,14 @@ You can also use the direct API to work with function/tool calling.
 Even here we can use Pydantic to generate the JSON schema for the tool:
 
 ```python
+from typing import Literal
+
 from pydantic import BaseModel
-from typing_extensions import Literal
 
+from pydantic_ai import ToolDefinition
 from pydantic_ai.direct import model_request
 from pydantic_ai.messages import ModelRequest
 from pydantic_ai.models import ModelRequestParameters
-from pydantic_ai.tools import ToolDefinition
 
 
 class Divide(BaseModel):

docs/evals.md

@@ -58,11 +58,11 @@ Pydantic Evals includes several built-in evaluators and allows you to create cus
 ```python {title="simple_eval_evaluator.py" requires="simple_eval_dataset.py"}
 from dataclasses import dataclass
 
-from simple_eval_dataset import dataset
-
 from pydantic_evals.evaluators import Evaluator, EvaluatorContext
 from pydantic_evals.evaluators.common import IsInstance
 
+from simple_eval_dataset import dataset
+
 dataset.add_evaluator(IsInstance(type_name='str'))  # (1)!
 
 
@@ -619,11 +619,11 @@ You can also write datasets as JSON files:
 ```python {title="generate_dataset_example_json.py" requires="generate_dataset_example.py"}
 from pathlib import Path
 
-from generate_dataset_example import AnswerOutput, MetadataType, QuestionInputs
-
 from pydantic_evals import Dataset
 from pydantic_evals.generation import generate_dataset
 
+from generate_dataset_example import AnswerOutput, MetadataType, QuestionInputs
+
 
 async def main():
     dataset = await generate_dataset(  # (1)!

docs/logfire.md

@@ -226,7 +226,7 @@ from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.trace import set_tracer_provider
 
-from pydantic_ai.agent import Agent
+from pydantic_ai import Agent
 
 os.environ['OTEL_EXPORTER_OTLP_ENDPOINT'] = 'http://localhost:4318'
 exporter = OTLPSpanExporter()
@@ -296,7 +296,7 @@ By default, the global `TracerProvider` and `EventLoggerProvider` are used. Thes
 from opentelemetry.sdk._events import EventLoggerProvider
 from opentelemetry.sdk.trace import TracerProvider
 
-from pydantic_ai.agent import Agent, InstrumentationSettings
+from pydantic_ai import Agent, InstrumentationSettings
 
 instrumentation_settings = InstrumentationSettings(
     tracer_provider=TracerProvider(),
@@ -322,7 +322,7 @@ agent = Agent(model)
 ### Excluding binary content
 
 ```python {title="excluding_binary_content.py"}
-from pydantic_ai.agent import Agent, InstrumentationSettings
+from pydantic_ai import Agent, InstrumentationSettings
 
 instrumentation_settings = InstrumentationSettings(include_binary_content=False)
 
@@ -338,7 +338,7 @@ For privacy and security reasons, you may want to monitor your agent's behavior
 When `include_content=False` is set, Pydantic AI will exclude sensitive content from OpenTelemetry events, including user prompts and model completions, tool call arguments and responses, and any other message content.
 
 ```python {title="excluding_sensitive_content.py"}
-from pydantic_ai.agent import Agent
+from pydantic_ai import Agent
 from pydantic_ai.models.instrumented import InstrumentationSettings
 
 instrumentation_settings = InstrumentationSettings(include_content=False)

docs/mcp/client.md

@@ -174,10 +174,9 @@ call needs.
 ```python {title="mcp_process_tool_call.py"}
 from typing import Any
 
-from pydantic_ai import Agent
+from pydantic_ai import Agent, RunContext
 from pydantic_ai.mcp import CallToolFunc, MCPServerStdio, ToolResult
 from pydantic_ai.models.test import TestModel
-from pydantic_ai.tools import RunContext
 
 
 async def process_tool_call(
@@ -245,29 +244,29 @@ parameter that lets you pass your own pre-configured
 [`httpx.AsyncClient`](https://www.python-httpx.org/async/).
 
 ```python {title="mcp_custom_tls_client.py"}
-import httpx
 import ssl
 
+import httpx
+
 from pydantic_ai import Agent
 from pydantic_ai.mcp import MCPServerSSE
 
-
 # Trust an internal / self-signed CA
-ssl_ctx = ssl.create_default_context(cafile="/etc/ssl/private/my_company_ca.pem")
+ssl_ctx = ssl.create_default_context(cafile='/etc/ssl/private/my_company_ca.pem')
 
 # OPTIONAL: if the server requires **mutual TLS** load your client certificate
-ssl_ctx.load_cert_chain(certfile="/etc/ssl/certs/client.crt", keyfile="/etc/ssl/private/client.key",)
+ssl_ctx.load_cert_chain(certfile='/etc/ssl/certs/client.crt', keyfile='/etc/ssl/private/client.key',)
 
 http_client = httpx.AsyncClient(
     verify=ssl_ctx,
     timeout=httpx.Timeout(10.0),
 )
 
 server = MCPServerSSE(
-    url="http://localhost:3001/sse",
+    url='http://localhost:3001/sse',
     http_client=http_client,  # (1)!
 )
-agent = Agent("openai:gpt-4o", toolsets=[server])
+agent = Agent('openai:gpt-4o', toolsets=[server])
 
 async def main():
     async with agent:

docs/mcp/server.md

@@ -102,7 +102,12 @@ from typing import Any
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
 from mcp.shared.context import RequestContext
-from mcp.types import CreateMessageRequestParams, CreateMessageResult, ErrorData, TextContent
+from mcp.types import (
+    CreateMessageRequestParams,
+    CreateMessageResult,
+    ErrorData,
+    TextContent,
+)
 
 
 async def sampling_callback(

docs/message-history.md

@@ -402,9 +402,8 @@ History processors can optionally accept a [`RunContext`][pydantic_ai.tools.RunC
 additional information about the current run, such as dependencies, model information, and usage statistics:
 
 ```python {title="context_aware_processor.py"}
-from pydantic_ai import Agent
+from pydantic_ai import Agent, RunContext
 from pydantic_ai.messages import ModelMessage
-from pydantic_ai.tools import RunContext
 
 
 def context_aware_processor(