diff --git a/docs/integrations/bigquery.md b/docs/integrations/bigquery.md index 17f9cba157..bc11c1fde0 100644 --- a/docs/integrations/bigquery.md +++ b/docs/integrations/bigquery.md @@ -37,7 +37,7 @@ You should use this approach for local development and running on Google Cloud s ```python import google.auth -from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig +from google.adk.integrations.bigquery import BigQueryToolset, BigQueryCredentialsConfig # Load Application Default Credentials credentials, project_id = google.auth.default() @@ -53,7 +53,7 @@ You can explicitly provide a service account file or info. ```python from google.oauth2 import service_account -from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig +from google.adk.integrations.bigquery import BigQueryToolset, BigQueryCredentialsConfig # Load Service Account credentials credentials = service_account.Credentials.from_service_account_file('path/to/key.json') @@ -69,7 +69,7 @@ For applications that need to act on behalf of an end-user, you can pass user cr ```python from google.oauth2.credentials import Credentials -from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig +from google.adk.integrations.bigquery import BigQueryToolset, BigQueryCredentialsConfig # Assume 'user_token' is obtained via an external OAuth flow credentials = Credentials(token=user_token) @@ -84,7 +84,7 @@ bigquery_toolset = BigQueryToolset(credentials_config=credentials_config) If you are integrating with an external authentication provider where the token is managed by the platform, such as Gemini Enterprise, use `external_access_token_key`. ```python -from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig +from google.adk.integrations.bigquery import BigQueryToolset, BigQueryCredentialsConfig # The key used to look up the access token in the session state credentials_config = BigQueryCredentialsConfig( @@ -98,7 +98,7 @@ bigquery_toolset = BigQueryToolset(credentials_config=credentials_config) When using the `adk web` interface for interactive sessions, you can provide OAuth 2.0 client credentials to trigger a login flow. This mechanism works for both local development and when your ADK agent is deployed to environments like Cloud Run. ```python -from google.adk.tools.bigquery import BigQueryToolset, BigQueryCredentialsConfig +from google.adk.integrations.bigquery import BigQueryToolset, BigQueryCredentialsConfig # Provide OAuth 2.0 Client ID and Secret credentials_config = BigQueryCredentialsConfig( diff --git a/examples/python/snippets/tools/built-in-tools/bigquery.py b/examples/python/snippets/tools/built-in-tools/bigquery.py index 7be429ec76..c44abb6b52 100644 --- a/examples/python/snippets/tools/built-in-tools/bigquery.py +++ b/examples/python/snippets/tools/built-in-tools/bigquery.py @@ -15,12 +15,12 @@ import asyncio from google.adk.agents import Agent +from google.adk.integrations.bigquery import BigQueryCredentialsConfig +from google.adk.integrations.bigquery import BigQueryToolset +from google.adk.integrations.bigquery.config import BigQueryToolConfig +from google.adk.integrations.bigquery.config import WriteMode 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 @@ -62,37 +62,49 @@ ) # Session and Runner -session_service = InMemorySessionService() -session = asyncio.run( - session_service.create_session( +async def setup_session_and_runner(): + session_service = InMemorySessionService() + await 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 -) + return Runner( + agent=bigquery_agent, app_name=APP_NAME, session_service=session_service + ) # Agent Interaction -def call_agent(query): +async def call_agent_async(runner, 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) + events = runner.run_async( + user_id=USER_ID, session_id=SESSION_ID, new_message=content + ) print("USER:", query) - for event in events: + async 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?" -) +async def main(): + runner = await setup_session_and_runner() + await call_agent_async( + runner, "Are there any ml datasets in bigquery-public-data project?" + ) + await call_agent_async(runner, "Tell me more about ml_datasets.") + await call_agent_async(runner, "Which all tables does it have?") + await call_agent_async(runner, "Tell me more about the census_adult_income table.") + await call_agent_async(runner, "How many rows are there per income bracket?") + await call_agent_async( + runner, + "What is the statistical correlation between education_num, age, and the" + " income_bracket?", + ) + + +# Note: In Colab or another notebook, an event loop is already running, so call +# `await main()` directly instead of `asyncio.run(main())`. +asyncio.run(main())