Skip to content

用于 ADK 的 Files Retrieval 工具

Supported in ADKPython

FilesRetrieval 工具让你的 ADK 智能体使用检索增强生成（RAG）对本地文档建立索引并进行查询。它使用 Google 的 Gemini 嵌入模型，在你指定的目录上构建 LlamaIndex VectorStoreIndex。随后，你的智能体就能从本地文本文件、Markdown 文档和源代码文件中检索相关片段，使答案以项目专属的上下文为依据。

用例

  • 代码库与文档搜索：从本地仓库中检索相关函数、设计说明和文档，以回答技术问题。
  • 本地知识库基础信息获取 (Grounding)：直接从你自己的文件系统中索引内部 Markdown 文件、技术规范和指南，无需先把它们加载到托管文档存储中。文档内容会被发送到所配置的嵌入模型进行索引，因此在索引敏感材料之前，请先查看 Google AI Studio 或 Agent Platform 的数据处理条款。
  • 上下文增强的辅助：从报告、日志或文本文件中检索相关的领域特定数据，使智能体的响应以经过验证的源材料为依据。

前提条件

FilesRetrieval 工具使用 LlamaIndex 来索引文档，而 ADK 默认不会安装它。请安装提供该能力的扩展依赖：

pip install "google-adk[extensions]"

然后为 Google AI Studio 或 Agent Platform 配置凭据：

在 Google AI Studio 中生成一个 API key，并设置环境变量：

export GOOGLE_API_KEY="your-api-key"

使用你的 Google Cloud 凭据配置 Agent Platform 访问权限：

export GOOGLE_GENAI_USE_ENTERPRISE=TRUE
export GOOGLE_CLOUD_PROJECT="your-project-id"
export GOOGLE_CLOUD_LOCATION="<global | us | eu>"

Note

在生产环境中，请通过 embedding_model=GoogleGenAIEmbedding(model_name="gemini-embedding-2", embed_batch_size=1) 显式传入 GA 模型。有关详情，请参阅 Gemini Embedding。有关将你的 ADK 智能体连接到 Google Cloud 资源和服务的详情，请参阅 Google Cloud 连接指南。

与智能体一起使用

本示例为一个本地数据目录配置 FilesRetrieval，并把它挂载到一个 ADK 智能体上。运行之前，请在你的智能体模块旁创建一个 data/ 目录，并加入你想索引的 .txt 或 .md 文件：FilesRetrieval 在构造时会加载并嵌入整个目录，因此该目录必须已经存在，并且每次导入智能体模块时都会重新索引其中的内容。

import os
from google.adk.agents import Agent
from google.adk.tools.retrieval.files_retrieval import FilesRetrieval

# 源文档所在目录的路径
DATA_DIR = os.path.join(os.path.dirname(__file__), "data")

# 初始化 FilesRetrieval 工具
files_retrieval = FilesRetrieval(
    name="search_documents",
    description=(
        "搜索本地文档文件以找到相关信息。"
        "当用户询问有关架构、项目结构或工具的问题时，"
        "请使用此工具。"
    ),
    input_dir=DATA_DIR,
)

# 创建配备了检索工具的智能体
root_agent = Agent(
    model="gemini-flash-latest",
    name="files_retrieval_agent",
    instruction=(
        "你是一个有用的助手，会基于本地文档文件回答问题。"
        "在生成答案之前，务必先使用 search_documents 工具"
        "检索相关上下文。"
    ),
    tools=[files_retrieval],
)

可用工具

FilesRetrieval 类是一个工具。通过 tools=[...] 挂载后，智能体会看到一个函数，它只接受一个 query 字符串参数：

工具 描述
search_documents 对已索引目录中的文档执行语义向量搜索，并针对给定的自然语言查询返回最相关的内容片段。可用 name 参数为其重命名。

配置

FilesRetrieval 构造函数接受以下参数：

参数 类型 必填 默认值 描述
name str 是 — 工具的唯一标识符，供模型用于函数调用。
description str 是 — 说明智能体应在何时以及如何调用该检索工具。
input_dir str 是 — 包含待加载和索引文档的本地文件系统目录路径。
embedding_model Optional[BaseEmbedding] 否 None 自定义的 LlamaIndex BaseEmbedding 实例。省略时，默认为 GoogleGenAIEmbedding(model_name="gemini-embedding-2-preview", embed_batch_size=1)。

自定义嵌入模型

你可以通过传入符合 LlamaIndex BaseEmbedding 接口的实例来定制嵌入模型：

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="搜索本地知识库文件。",
    input_dir=os.path.join(os.path.dirname(__file__), "data"),
    embedding_model=custom_embedding,
)

其他资源