Skip to content

用于 ADK 的 BigQuery 工具

Supported in ADKPython v1.1.0

这是一组旨在提供与 BigQuery 集成的工具,包括:

  • list_dataset_ids:获取 GCP 项目中存在的 BigQuery 数据集 ID。
  • get_dataset_info:获取 BigQuery 数据集的元数据。
  • list_table_ids:获取 BigQuery 数据集中的表 ID。
  • get_table_info:获取 BigQuery 表的元数据。
  • get_job_info:获取 BigQuery 作业的元数据信息(槽位使用情况、配置、统计数据、状态等)。
  • execute_sql:在 BigQuery 中运行 SQL 查询并获取结果。
  • forecast:使用 AI.FORECAST 函数运行 BigQuery AI 时间序列预测。
  • analyze_contribution:执行 BigQuery ML 贡献分析,以了解驱动指标变化的因素。
  • detect_anomalies:训练 ARIMA_PLUS 模型并检测时间序列数据中的异常。
  • ask_data_insights:使用自然语言回答有关 BigQuery 表中数据的问题。
  • search_catalog:通过 Dataplex 使用自然语言语义搜索查找 BigQuery 数据集和表面。

这些工具都打包在 BigQueryToolset 工具集中。

身份验证

BigQueryToolset 通过 BigQueryCredentialsConfig 支持多种身份验证机制。

应用默认凭据 (ADC)

你应该在本地开发以及在 Cloud Run 和 GKE 等 Google Cloud 服务上运行时使用此方法。

import google.auth
from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig

# 加载应用默认凭据
credentials, project_id = google.auth.default()

# 配置工具集
credentials_config = BigQueryCredentialsConfig(credentials=credentials)
bigquery_toolset = BigQueryToolset(credentials_config=credentials_config)

服务账号

你可以显式提供服务账号文件或信息。

from google.oauth2 import service_account
from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig

# 加载服务账号凭据
credentials = service_account.Credentials.from_service_account_file('path/to/key.json')

# 配置工具集
credentials_config = BigQueryCredentialsConfig(credentials=credentials)
bigquery_toolset = BigQueryToolset(credentials_config=credentials_config)

外部访问令牌

对于需要代表最终用户操作的应用程序,你可以传递直接从访问令牌实例化的用户凭据,例如来自 OAuth2 流程或外部 IDP 的令牌。

from google.oauth2.credentials import Credentials
from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig

# 假设 'user_token' 是通过外部 OAuth 流程获得的
credentials = Credentials(token=user_token)

# 配置工具集
credentials_config = BigQueryCredentialsConfig(credentials=credentials)
bigquery_toolset = BigQueryToolset(credentials_config=credentials_config)

外部身份验证提供程序

如果你正与由平台管理令牌的外部身份验证提供程序集成(例如 Gemini Enterprise),请使用 external_access_token_key

from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig

# 用于在会话状态中查找访问令牌的键
credentials_config = BigQueryCredentialsConfig(
    external_access_token_key="YOUR_AUTH_ID"
)
bigquery_toolset = BigQueryToolset(credentials_config=credentials_config)

交互式身份验证 (ADK Web)

在交互式会话中使用 adk web 界面时,你可以提供 OAuth 2.0 客户端凭据以触发登录流程。此机制适用于本地开发,也适用于将 ADK 智能体部署到 Cloud Run 等环境时。

from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig

# 提供 OAuth 2.0 Client ID 和 Secret
credentials_config = BigQueryCredentialsConfig(
    client_id="YOUR_CLIENT_ID",
    client_secret="YOUR_CLIENT_SECRET"
)
bigquery_toolset = BigQueryToolset(credentials_config=credentials_config)

示例代码

以下示例代码演示了如何在 ADK 智能体中使用应用默认凭据 (ADC) 来调用 BigQueryToolset

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio

from google.adk.agents import Agent
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.adk.tools.bigquery import BigQueryCredentialsConfig
from google.adk.tools.bigquery import BigQueryToolset
from google.adk.tools.bigquery.config import BigQueryToolConfig
from google.adk.tools.bigquery.config import WriteMode
from google.genai import types
import google.auth

# Define constants for this example agent
AGENT_NAME = "bigquery_agent"
APP_NAME = "bigquery_app"
USER_ID = "user1234"
SESSION_ID = "1234"
GEMINI_MODEL = "gemini-2.0-flash"

# Define a tool configuration to block any write operations
tool_config = BigQueryToolConfig(write_mode=WriteMode.BLOCKED)

# Use Application Default Credentials (ADC) for BigQuery authentication
# https://cloud.google.com/docs/authentication/provide-credentials-adc
application_default_credentials, _ = google.auth.default()
credentials_config = BigQueryCredentialsConfig(
    credentials=application_default_credentials
)

# Instantiate a BigQuery toolset
bigquery_toolset = BigQueryToolset(
    credentials_config=credentials_config, bigquery_tool_config=tool_config
)

# Agent Definition
bigquery_agent = Agent(
    model=GEMINI_MODEL,
    name=AGENT_NAME,
    description=(
        "Agent to answer questions about BigQuery data and models and execute"
        " SQL queries."
    ),
    instruction="""\
        You are a data science agent with access to several BigQuery tools.
        Make use of those tools to answer the user's questions.
    """,
    tools=[bigquery_toolset],
)

# Session and Runner
session_service = InMemorySessionService()
session = asyncio.run(
    session_service.create_session(
        app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID
    )
)
runner = Runner(
    agent=bigquery_agent, app_name=APP_NAME, session_service=session_service
)


# Agent Interaction
def call_agent(query):
    """
    Helper function to call the agent with a query.
    """
    content = types.Content(role="user", parts=[types.Part(text=query)])
    events = runner.run(user_id=USER_ID, session_id=SESSION_ID, new_message=content)

    print("USER:", query)
    for event in events:
        if event.is_final_response():
            final_response = event.content.parts[0].text
            print("AGENT:", final_response)


call_agent("Are there any ml datasets in bigquery-public-data project?")
call_agent("Tell me more about ml_datasets.")
call_agent("Which all tables does it have?")
call_agent("Tell me more about the census_adult_income table.")
call_agent("How many rows are there per income bracket?")
call_agent(
    "What is the statistical correlation between education_num, age, and the income_bracket?"
)

示例智能体

有关使用 BigQuery 驱动的智能体的完整、可运行示例(包含详细的身份验证示例),请参阅 GitHub 上的 BigQuery 示例智能体

注意:如果你想将 BigQuery 数据智能体作为工具访问,请参阅 用于 ADK 的数据智能体工具