-
Notifications
You must be signed in to change notification settings - Fork 1.4k
Add FileSearchTool with support for OpenAI and Google
#3396
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 37 commits
4376b96
6cec96f
4c3fe56
3c8decf
2343679
666a1bb
7365e20
2ee21c9
deef1ec
18b4b86
11654ed
1542f5c
7d683b7
d8ef07d
6acbd76
380e25c
c83f125
8eba82d
b3a8930
19f32f9
00ea1ed
c6ed56c
9b5bb54
c2765ac
8286cd7
3011e05
bc278e8
5f694c9
8dc7c17
8216f31
ffcb21f
bc3ac7a
977ab53
68bafb6
29f8da0
eef4526
8cc3d60
b77d857
db475c6
fd62e29
13abf31
129dacd
8d3f359
065c711
50ad873
ff22a6d
1f249aa
c7798d1
9696c01
2fe8791
1c0589b
122be13
925a909
5cc9fb0
8d6a6e7
68738fd
99af405
0c7c582
b2d95e5
1f6f4c9
810fc25
7610c5a
10960b6
74f8078
2faa581
a5402fc
8c8c949
057b0dd
fe1cde9
5984c0d
17eb139
3f757d6
188085d
9e1c707
2d86b0d
2c3fcd5
1c8ccc6
fa338fc
87caa1d
b199b8b
250b3df
608b977
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 |
|---|---|---|
|
|
@@ -12,6 +12,7 @@ Pydantic AI supports the following built-in tools: | |
| - **[`UrlContextTool`][pydantic_ai.builtin_tools.UrlContextTool]**: Enables agents to pull URL contents into their context | ||
| - **[`MemoryTool`][pydantic_ai.builtin_tools.MemoryTool]**: Enables agents to use memory | ||
| - **[`MCPServerTool`][pydantic_ai.builtin_tools.MCPServerTool]**: Enables agents to use remote MCP servers with communication handled by the model provider | ||
| - **[`FileSearchTool`][pydantic_ai.builtin_tools.FileSearchTool]**: Enables agents to search through uploaded files using vector search (RAG) | ||
|
|
||
| These tools are passed to the agent via the `builtin_tools` parameter and are executed by the model provider's infrastructure. | ||
|
|
||
|
|
@@ -566,6 +567,65 @@ _(This example is complete, it can be run "as is")_ | |
| | `description` | ✅ | ❌ | | ||
| | `headers` | ✅ | ❌ | | ||
|
|
||
| ## File Search Tool | ||
|
|
||
| The [`FileSearchTool`][pydantic_ai.builtin_tools.FileSearchTool] enables your agent to search through uploaded files using vector search, providing a fully managed Retrieval-Augmented Generation (RAG) system. This tool handles file storage, chunking, embedding generation, and context injection into prompts. | ||
|
|
||
| ### Provider Support | ||
|
|
||
| | Provider | Supported | Notes | | ||
| |----------|-----------|-------| | ||
| | OpenAI Responses | ✅ | Full feature support. Requires files to be uploaded to vector stores via the [OpenAI Files API](https://platform.openai.com/docs/api-reference/files). Vector stores must be created and file IDs added before using the tool. | | ||
| | Google (Gemini) | ✅ | Requires files to be uploaded via the [Gemini Files API](https://ai.google.dev/gemini-api/docs/files). Files are automatically deleted after 48 hours. Supports up to 2 GB per file and 20 GB per project. Using built-in tools and function tools (including [output tools](output.md#tool-output)) at the same time is not supported; to use structured output, use [`PromptedOutput`](output.md#prompted-output) instead. | | ||
| | Anthropic | ❌ | Not supported | | ||
| | Groq | ❌ | Not supported | | ||
| | OpenAI Chat Completions | ❌ | Not supported | | ||
| | Bedrock | ❌ | Not supported | | ||
| | Mistral | ❌ | Not supported | | ||
| | Cohere | ❌ | Not supported | | ||
| | HuggingFace | ❌ | Not supported | | ||
| | Outlines | ❌ | Not supported | | ||
|
|
||
| ### Usage | ||
|
|
||
| #### OpenAI Responses | ||
|
|
||
| With OpenAI, you need to first [upload files to a vector store](https://platform.openai.com/docs/assistants/tools/file-search), then reference the vector store IDs when using the `FileSearchTool`: | ||
|
|
||
| ```py {title="file_search_openai.py" test="skip"} | ||
| from pydantic_ai import Agent, FileSearchTool | ||
|
|
||
| agent = Agent( | ||
| 'openai-responses:gpt-5', | ||
| builtin_tools=[FileSearchTool(vector_store_ids=['vs_abc123'])] # (1) | ||
| ) | ||
|
|
||
| result = agent.run_sync('What information is in my documents about pydantic?') | ||
| print(result.output) | ||
| #> Based on your documents, Pydantic is a data validation library for Python... | ||
| ``` | ||
|
|
||
| 1. Replace `vs_abc123` with your actual vector store ID from the OpenAI API. | ||
|
|
||
| #### Google (Gemini) | ||
|
|
||
| With Gemini, you need to first [create a file search store via the Files API](https://ai.google.dev/gemini-api/docs/files), then reference the file search store names: | ||
|
||
|
|
||
| ```py {title="file_search_google.py" test="skip"} | ||
| from pydantic_ai import Agent, FileSearchTool | ||
|
|
||
| agent = Agent( | ||
| 'google-gla:gemini-2.5-flash', | ||
| builtin_tools=[FileSearchTool(vector_store_ids=['files/abc123'])] # (1) | ||
| ) | ||
|
|
||
| result = agent.run_sync('Summarize the key points from my uploaded documents.') | ||
| print(result.output) | ||
| #> The documents discuss the following key points: ... | ||
| ``` | ||
|
|
||
| 1. Replace `files/abc123` with your actual file search store name from the Gemini Files API. | ||
|
|
||
| ## API Reference | ||
|
|
||
| For complete API documentation, see the [API Reference](api/builtin_tools.md). | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -17,6 +17,7 @@ | |
| 'ImageGenerationTool', | ||
| 'MemoryTool', | ||
| 'MCPServerTool', | ||
| 'FileSearchTool', | ||
| ) | ||
|
|
||
| _BUILTIN_TOOL_TYPES: dict[str, type[AbstractBuiltinTool]] = {} | ||
|
|
@@ -334,6 +335,30 @@ def unique_id(self) -> str: | |
| return ':'.join([self.kind, self.id]) | ||
|
|
||
|
|
||
| @dataclass(kw_only=True) | ||
| class FileSearchTool(AbstractBuiltinTool): | ||
| """A builtin tool that allows your agent to search through uploaded files using vector search. | ||
|
|
||
| This tool provides a fully managed Retrieval-Augmented Generation (RAG) system that handles | ||
| file storage, chunking, embedding generation, and context injection into prompts. | ||
|
|
||
| Supported by: | ||
|
|
||
| * OpenAI Responses | ||
| * Google (Gemini) | ||
DouweM marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| """ | ||
|
|
||
| vector_store_ids: list[str] | ||
|
||
| """List of vector store IDs to search through. | ||
|
|
||
| For OpenAI, these are the IDs of vector stores created via the OpenAI API. | ||
| For Google, these are file resource names that have been uploaded and processed. | ||
| """ | ||
|
|
||
| kind: str = 'file_search' | ||
| """The kind of tool.""" | ||
|
|
||
|
|
||
| def _tool_discriminator(tool_data: dict[str, Any] | AbstractBuiltinTool) -> str: | ||
| if isinstance(tool_data, dict): | ||
| return tool_data.get('kind', AbstractBuiltinTool.kind) | ||
|
|
||
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.
It would be awesome if you could have example code for the upload step as well, using the
OpenAIResponsesModel.clientThere 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.
Added! The example now shows the complete workflow using model.client