Files Retrieval tool for ADK¶
The FilesRetrieval tool lets your ADK agent index and query local documents
using retrieval-augmented generation (RAG). It builds a LlamaIndex
VectorStoreIndex over a directory you specify, using Google's
Gemini embedding model. Your agent can then retrieve
relevant excerpts from local text files, Markdown documents, and source files
to ground its answers in project-specific context.
Use cases¶
- Codebase and Documentation Search: Retrieve relevant functions, design notes and documentation from a local repository to answer technical questions.
- Local Knowledge Base Grounding: Index internal markdown files, technical specifications, and guides straight from your own filesystem, without first loading them into a hosted document store. Document content is sent to the configured embedding model for indexing, so review the data handling terms for Google AI Studio or Agent Platform before indexing sensitive material.
- Context-Augmented Assistance: Retrieve relevant domain-specific data from reports, logs, or text files to ground agent responses in verified source material.
Prerequisites¶
The FilesRetrieval tool indexes documents with LlamaIndex, which ADK does not install by default. Install the extra that provides it:
Then configure credentials for either Google AI Studio or Agent Platform:
Generate an API key in Google AI Studio and set the environment variable:
Note
For production, pass the GA model explicitly with embedding_model=GoogleGenAIEmbedding(model_name="gemini-embedding-2", embed_batch_size=1). For more information, see Gemini Embedding. For more information on connecting your ADK agent to Google Cloud resources
and services, see the Google Cloud Connection Guide.
Use with agent¶
This example configures FilesRetrieval for a local data directory and
attaches it to an ADK agent. Before running it, create a data/ directory next to your agent module and add the .txt or .md files you want indexed: FilesRetrieval loads and embeds the entire directory when it is constructed, so the directory must already exist, and the content is re-indexed each time the agent module is imported.
import os
from google.adk.agents import Agent
from google.adk.tools.retrieval.files_retrieval import FilesRetrieval
# Path to the directory containing your source documents
DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
# Initialize the FilesRetrieval tool
files_retrieval = FilesRetrieval(
name="search_documents",
description=(
"Search through local documentation files to find relevant"
" information. Use this tool when the user asks questions about"
" architecture, project structure, or tools."
),
input_dir=DATA_DIR,
)
# Create an agent equipped with the retrieval tool
root_agent = Agent(
model="gemini-flash-latest",
name="files_retrieval_agent",
instruction=(
"You are a helpful assistant that answers questions based on local"
" documentation files. Always use the search_documents tool to retrieve"
" relevant context before generating your answer."
),
tools=[files_retrieval],
)
Available tools¶
The FilesRetrieval class is a tool. Once attached with tools=[...], the agent sees one function, which takes a single query string parameter:
| Tool | Description |
|---|---|
search_documents |
Performs semantic vector search over documents in the indexed directory and returns the most relevant content chunk for a given natural language query. Rename it with the name parameter. |
Configuration¶
The FilesRetrieval constructor accepts the following parameters:
| Parameter | Type | Required | Default | Description |
|---|---|---|---|---|
name |
str |
Yes | — | Unique identifier for the tool, used by the model for function calling. |
description |
str |
Yes | — | Explanation of when and how the agent should invoke the retrieval tool. |
input_dir |
str |
Yes | — | Local filesystem directory path containing the documents to load and index. |
embedding_model |
Optional[BaseEmbedding] |
No | None |
Custom LlamaIndex BaseEmbedding instance. When omitted, defaults to GoogleGenAIEmbedding(model_name="gemini-embedding-2-preview", embed_batch_size=1). |
Custom embedding models¶
You can customize the embedding model by passing an instance conforming to LlamaIndex's BaseEmbedding interface:
from google.adk.tools.retrieval.files_retrieval import FilesRetrieval
from llama_index.embeddings.google_genai import GoogleGenAIEmbedding
custom_embedding = GoogleGenAIEmbedding(
model_name="gemini-embedding-2",
embed_batch_size=1,
)
files_retrieval = FilesRetrieval(
name="search_documents",
description="Search local knowledge base files.",
input_dir=os.path.join(os.path.dirname(__file__), "data"),
embedding_model=custom_embedding,
)