Expose all `pydantic_ai.messages` types on `pydantic_ai` to aid IDE auto-import

[moritzwilksch]

Question

pydantic_ai/__init__.py exposes a seemingly random subset of types by setting the modules __all__ attribute. Since the list is not exhaustive of all elements, this messes with my type checker and IDE auto-imports. What is the rationale behind the selected items? Is the list outdated and could it be changed or expanded?

Usually, we'd use this for more convenient, less deep import chains, e.g. from pydantic_ai import ToolDefinition, but other dev tooling doesn't seem to like this exposing of a subset of types.

Example

ToolDefinition is exposed here, so a type checker can find it:

pydantic-ai/pydantic_ai_slim/pydantic_ai/__init__.py

Line 59 in c028211

'ToolDefinition',

ModelRequest is not exposed, so no auto-import available and some type checkers can complain:

Additional Context

No response

[DouweM]

@moritzwilksch The rationale is hinted at in #2736, the most recent time I updated the __all__ list:

Let almost all types used in docs examples be imported directly from pydantic_ai

Found using from pydantic_ai\.(?!models|mcp|providers|ag_ui|messages|toolsets|profiles|direct|ext|retries|durable_exec|common_tools) limited to docs dir. Those are the more specialized modules whose types don't make sense at the top level.

The idea is that everything we use in examples (which, assuming all features are documented, means all public types) should be importable from pydantic_ai, unless:

1. it depends on optional dependencies (models, mcp, providers, ag_ui, ext, retries, durable_exec, common_tools)
2. it exposes a lot of types (messages, profiles)
3. it's a specifically scoped feature (toolsets, direct)

I realize 2 and 3 could easily be argued with :) Making all the toolsets available at the top level would makes sense to me, and I see why ModelRequest/Response would be useful to have there as well, but if we add those, should we then also move all of the other parts, events, etc? I kind of like keeping the "I'm dealing directly with messages" stuff in its own module.

[moritzwilksch]

Ok, I see the intention. The issue is that the docs for direct mode don’t show all types in the examples. That makes sense, since there are many and it’s not really needed to cover every single one. But because direct mode is meant for users too (these aren't private classes), missing IDE auto imports ends up being a big DX pain point.
Maybe there's a typing trick to tell pyright that the .direct mode still exists even though it's not mentioned in __all__?

[DouweM]

@moritzwilksch Hmm, looks like auto-imports are handled by the Pylance VSCode plugin, and its docs point out this specific issue:

Problem	Cause	Solution
Pylance is only offering top-level symbol options when adding imports.	By default, only top-level modules are indexed (depth=1).For example, you may see import matplotlib as a suggestion, but not import matplotlib.pyplot by default.	Try increasing the depth to which Pylance can index your installed libraries through the python.analysis.packageIndexDepths. Check code analysis settings.

So you can fix this in your IDE, but it would be useful if users didn't have to deal with this to begin with. We can't completely eliminate the issue as models etc (point 1 in my list above) can't be made available at the top level because of dependencies, but making all the types from pydantic_ai.messages available under pydantic_ai would make life a bit easier.

Let me see what the team thinks.

[DouweM]

@moritzwilksch Feel free to submit a PR that makes all of the types under messages, toolsets, direct, and profiles available on pydantic_ai! And then update all the examples to import directly from pydantic_ai. You can do that with a regex replace, and then run make update-examples which will auto-sort the imports including merging multiple from pydantic_ai import lines.