Add FileSearchTool with support for OpenAI and Google

docs/builtin-tools.md

@@ -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,98 @@ _(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). To include search results on the [`BuiltinToolReturnPart`][pydantic_ai.messages.BuiltinToolReturnPart] available via [`ModelResponse.builtin_tool_calls`][pydantic_ai.messages.ModelResponse.builtin_tool_calls], enable the [`OpenAIResponsesModelSettings.openai_include_file_search_results`][pydantic_ai.models.openai.OpenAIResponsesModelSettings.openai_include_file_search_results] [model setting](agents.md#model-run-settings). |
+| 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. |
+|| Google (Vertex AI) | ❌ | Not supported |
+| 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_upload.py" test="skip"}
+import asyncio
+
+from pydantic_ai import Agent, FileSearchTool
+from pydantic_ai.models.openai import OpenAIResponsesModel
+
+async def main():
+    model = OpenAIResponsesModel('gpt-5')
+
+    with open('my_document.txt', 'rb') as f:
+        file = await model.client.files.create(file=f, purpose='assistants')
+
+    vector_store = await model.client.vector_stores.create(name='my-docs')
+    await model.client.vector_stores.files.create(
+        vector_store_id=vector_store.id,
+        file_id=file.id
+    )
+
+    agent = Agent(
+        model,
+        builtin_tools=[FileSearchTool(file_store_ids=[vector_store.id])]
+    )
+
+    result = await agent.run('What information is in my documents about pydantic?')
+    print(result.output)
+    #> Based on your documents, Pydantic is a data validation library for Python...
+
+asyncio.run(main())
+```
+
+#### 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_upload.py" test="skip"}
+import asyncio
+
+from pydantic_ai import Agent, FileSearchTool
+from pydantic_ai.models.google import GoogleModel
+
+async def main():
+    model = GoogleModel('gemini-2.5-flash')
+
+    store = await model.client.aio.file_search_stores.create(
+        config={'display_name': 'my-docs'}
+    )
+
+    with open('my_document.txt', 'rb') as f:
+        await model.client.aio.file_search_stores.upload_to_file_search_store(
+            file_search_store_name=store.name,
+            file=f,
+            config={'mime_type': 'text/plain'}
+        )
+
+    agent = Agent(
+        model,
+        builtin_tools=[FileSearchTool(file_store_ids=[store.name])]
+    )
+
+    result = await agent.run('Summarize the key points from my uploaded documents.')
+    print(result.output)
+    #> The documents discuss the following key points: ...
+
+asyncio.run(main())
+```
+
 ## API Reference
 
 For complete API documentation, see the [API Reference](api/builtin_tools.md).

docs/models/openai.md

@@ -148,7 +148,7 @@ model_settings = OpenAIResponsesModelSettings(
     openai_builtin_tools=[
         FileSearchToolParam(
             type='file_search',
-            vector_store_ids=['your-history-book-vector-store-id']
+            file_store_ids=['your-history-book-vector-store-id']
         )
     ],
 )

pydantic_ai_slim/pydantic_ai/__init__.py

@@ -11,6 +11,7 @@
 )
 from .builtin_tools import (
     CodeExecutionTool,
+    FileSearchTool,
     ImageGenerationTool,
     MCPServerTool,
     MemoryTool,
@@ -214,13 +215,14 @@
     'ToolsetTool',
     'WrapperToolset',
     # builtin_tools
-    'WebSearchTool',
-    'WebSearchUserLocation',
-    'UrlContextTool',
     'CodeExecutionTool',
+    'FileSearchTool',
     'ImageGenerationTool',
-    'MemoryTool',
     'MCPServerTool',
+    'MemoryTool',
+    'UrlContextTool',
+    'WebSearchTool',
+    'WebSearchUserLocation',
     # output
     'ToolOutput',
     'NativeOutput',

pydantic_ai_slim/pydantic_ai/builtin_tools.py

@@ -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)
+    """
+
+    file_store_ids: list[str]
+    """List of file store IDs to search through.
+
+    For OpenAI, these are the IDs of vector stores created via the OpenAI API.
+    For Google, these are file search store names that have been uploaded and processed via the Gemini Files API.
+    """
+
+    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)

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -540,7 +540,7 @@ def _add_builtin_tools(
                     mcp_server_url_definition_param['authorization_token'] = tool.authorization_token
                 mcp_servers.append(mcp_server_url_definition_param)
                 beta_features.append('mcp-client-2025-04-04')
-            else:  # pragma: no cover
+            else:
                 raise UserError(
                     f'`{tool.__class__.__name__}` is not supported by `AnthropicModel`. If it should be, please file an issue.'
                 )

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -13,7 +13,7 @@
 from .. import UnexpectedModelBehavior, _utils, usage
 from .._output import OutputObjectDefinition
 from .._run_context import RunContext
-from ..builtin_tools import CodeExecutionTool, ImageGenerationTool, UrlContextTool, WebSearchTool
+from ..builtin_tools import CodeExecutionTool, FileSearchTool, ImageGenerationTool, UrlContextTool, WebSearchTool
 from ..exceptions import ModelAPIError, ModelHTTPError, UserError
 from ..messages import (
     BinaryContent,
@@ -63,6 +63,7 @@
         ExecutableCode,
         ExecutableCodeDict,
         FileDataDict,
+        FileSearchDict,
         FinishReason as GoogleFinishReason,
         FunctionCallDict,
         FunctionCallingConfigDict,
@@ -93,6 +94,7 @@
         'you can use the `google` optional group — `pip install "pydantic-ai-slim[google]"`'
     ) from _import_error
 
+
 LatestGoogleModelNames = Literal[
     'gemini-flash-latest',
     'gemini-flash-lite-latest',
@@ -350,6 +352,9 @@ def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[T
                     tools.append(ToolDict(url_context=UrlContextDict()))
                 elif isinstance(tool, CodeExecutionTool):
                     tools.append(ToolDict(code_execution=ToolCodeExecutionDict()))
+                elif isinstance(tool, FileSearchTool):
+                    file_search_config = FileSearchDict(file_search_store_names=tool.file_store_ids)
+                    tools.append(ToolDict(file_search=file_search_config))
                 elif isinstance(tool, ImageGenerationTool):  # pragma: no branch
                     if not self.profile.supports_image_output:
                         raise UserError(
@@ -652,6 +657,7 @@ class GeminiStreamedResponse(StreamedResponse):
     _timestamp: datetime
     _provider_name: str
     _provider_url: str
+    _file_search_tool_call_id: str | None = field(default=None, init=False)
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         code_execution_tool_call_id: str | None = None
@@ -697,6 +703,26 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 continue  # pragma: no cover
 
             for part in parts:
+                if self._file_search_tool_call_id and candidate.grounding_metadata:
+                    grounding_chunks = candidate.grounding_metadata.grounding_chunks
+                    if grounding_chunks:
+                        retrieved_contexts = [
+                            chunk.retrieved_context.model_dump(mode='json')
+                            for chunk in grounding_chunks
+                            if chunk.retrieved_context
+                        ]
+                        if retrieved_contexts:
+                            yield self._parts_manager.handle_part(
+                                vendor_part_id=uuid4(),
+                                part=BuiltinToolReturnPart(
+                                    provider_name=self.provider_name,
+                                    tool_name=FileSearchTool.kind,
+                                    tool_call_id=self._file_search_tool_call_id,
+                                    content={'retrieved_contexts': retrieved_contexts},
+                                ),
+                            )
+                            self._file_search_tool_call_id = None
+
                 provider_details: dict[str, Any] | None = None
                 if part.thought_signature:
                     # Per https://ai.google.dev/gemini-api/docs/function-calling?example=meeting#thought-signatures:
@@ -739,10 +765,27 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                         part=FilePart(content=BinaryContent.narrow_type(content), provider_details=provider_details),
                     )
                 elif part.executable_code is not None:
-                    code_execution_tool_call_id = _utils.generate_tool_call_id()
-                    part = _map_executable_code(part.executable_code, self.provider_name, code_execution_tool_call_id)
-                    part.provider_details = provider_details
-                    yield self._parts_manager.handle_part(vendor_part_id=uuid4(), part=part)
+                    code = part.executable_code.code
+                    if code and (file_search_query := _extract_file_search_query(code)):
+                        self._file_search_tool_call_id = _utils.generate_tool_call_id()
+                        part_obj = BuiltinToolCallPart(
+                            provider_name=self.provider_name,
+                            tool_name=FileSearchTool.kind,
+                            tool_call_id=self._file_search_tool_call_id,
+                            args={'query': file_search_query},
+                        )
+                        part_obj.provider_details = provider_details
+                        yield self._parts_manager.handle_part(
+                            vendor_part_id=uuid4(),
+                            part=part_obj,
+                        )
+                    else:
+                        code_execution_tool_call_id = _utils.generate_tool_call_id()
+                        part_obj = _map_executable_code(
+                            part.executable_code, self.provider_name, code_execution_tool_call_id
+                        )
+                        part_obj.provider_details = provider_details
+                        yield self._parts_manager.handle_part(vendor_part_id=uuid4(), part=part_obj)
                 elif part.code_execution_result is not None:
                     assert code_execution_tool_call_id is not None
                     part = _map_code_execution_result(
@@ -856,6 +899,11 @@ def _process_response_from_parts(
         items.append(web_search_call)
         items.append(web_search_return)
 
+    file_search_call, file_search_return = _map_file_search_grounding_metadata(grounding_metadata, provider_name)
+    if file_search_call and file_search_return:
+        items.append(file_search_call)
+        items.append(file_search_return)
+
     item: ModelResponsePart | None = None
     code_execution_tool_call_id: str | None = None
     for part in parts:
@@ -1007,3 +1055,47 @@ def _map_grounding_metadata(
         )
     else:
         return None, None
+
+
+def _map_file_search_grounding_metadata(
+    grounding_metadata: GroundingMetadata | None, provider_name: str
+) -> tuple[BuiltinToolCallPart, BuiltinToolReturnPart] | tuple[None, None]:
+    if not grounding_metadata or not (grounding_chunks := grounding_metadata.grounding_chunks):
+        return None, None
+
+    retrieved_contexts = [
+        chunk.retrieved_context.model_dump(mode='json') for chunk in grounding_chunks if chunk.retrieved_context
+    ]
+
+    if not retrieved_contexts:
+        return None, None
+
+    tool_call_id = _utils.generate_tool_call_id()
+    return (
+        BuiltinToolCallPart(
+            provider_name=provider_name,
+            tool_name=FileSearchTool.kind,
+            tool_call_id=tool_call_id,
+            args={},
+        ),
+        BuiltinToolReturnPart(
+            provider_name=provider_name,
+            tool_name=FileSearchTool.kind,
+            tool_call_id=tool_call_id,
+            content={'retrieved_contexts': retrieved_contexts},
+        ),
+    )
+
+
+def _extract_file_search_query(code: str) -> str | None:
+    """Extract the query from file_search.query() executable code.
+
+    Example: 'print(file_search.query(query="what is the capital of France?"))'
+    Returns: 'what is the capital of France?'
+    """
+    import re
+
+    match = re.search(r'file_search\.query\(query=(["\'])(.+?)\1\)', code)
+    if match:
+        return match.group(2)
+    return None