Integrate a data agent with an application

This tutorial describes how to set up and use a data agent in AlloyDB for PostgreSQL using the Google Cloud console and integrate it with your application. You learn how to build the agent context file, create a data agent that uses the context, use MCP Toolbox to call the QueryData API to generate SQL queries for natural language questions, and finally integrate it with your application.

For more information, see Data agents overview.

Objectives

  • Create and populate tables.
  • Build agent context with Gemini CLI and MCP toolbox.
  • Create a data agent and upload context.
  • Inspect the agent and generate SQL queries in Studio.
  • Integrate the agent with your application using Gemini Data Analytics QueryData tool in MCP Toolbox.

Costs

In this document, you use the following billable components of Google Cloud:

To generate a cost estimate based on your projected usage, use the pricing calculator.

New Google Cloud users might be eligible for a free trial.

When you finish the tasks described in this document, avoid continued billing by deleting the resources you created. For more information, see Clean up.

Before you begin

Complete the following prerequisites before creating an agent.

Enable required services

Enable the following services for your project:

Prepare an AlloyDB for PostgreSQL cluster and instance

Make sure that you have access to an existing AlloyDB cluster and instance or create a new one.

Required roles and permissions

  • Add an Identity and Access Management (IAM) user or service account to the cluster at the database level. For more information, see Manage database users.
  • Grant the alloydb.databaseUser role and the serviceusage.serviceUsageConsumer permission to the IAM user at the project level. For more information, see Add IAM policy binding for a project.


To perform steps in this tutorial, sign in to Google Cloud, and then authenticate to the database using IAM authentication.

Create the flights and airports schema and tables

In this section, you create the flights and airports database for this tutorial.

  1. In the Google Cloud console, go to the AlloyDB page.

    Go to AlloyDB

  2. Select a cluster from the list.

  3. In the navigation menu, click AlloyDB Studio.

  4. Sign in to Studio using Identity and Access Management (IAM) authentication.

  5. Click Authenticate. The Explorer pane displays a list of the objects in your database.

  6. Click New SQL editor tab or New tab to open a new tab.

  7. Create the airports table and schema:

    CREATE TABLE IF NOT EXISTS airports (
  id INT PRIMARY KEY,
  iata TEXT,
  name TEXT,
  city TEXT,
  country TEXT
  );

  8. Create the flights table and schema:

    CREATE TABLE IF NOT EXISTS flights (
  id INT PRIMARY KEY,
  airline VARCHAR(10),
  flight_number INT,
  departure_airport VARCHAR(5),
  arrival_airport VARCHAR(5),
  departure_time TIMESTAMP,
  arrival_time TIMESTAMP,
  departure_gate VARCHAR(10),
  arrival_gate VARCHAR(10)
);

Populate the flights and airport table

In this section, you populate the flights and airports table using the provided SQL scripts.

  1. Populate the airports table.

  2. Populate the flights table.

  3. Run the following query to verify that the tables are populated:

    SELECT * FROM "public"."flights" LIMIT 10;
SELECT * FROM "public"."airports" LIMIT 10;

Create a data agent

In this section, you create a data agent named flights-assistant. This agent doesn't include any agent context uploaded to it.

  1. In the Explorer pane, next to Data Agents, click View actions.
  2. Click Create agent.
  3. In Name your agent, enter flights-assistant.
  4. Click Create.

Inspect the agent in Studio

In this section, you ask the flights-assistant agent a natural language question and it generates a SQL query. Since the agent doesn't have any context, even after asking a question with context such as nighttime traffic, the agent generates a sub-optimal query.

  1. In the Explorer pane, next to your data agent, click View actions.
  2. Click Inspect agent.
  3. In the query editor, click Generate SQL using agent: flights-assistant.

  4. Enter the following natural language question to generate a SQL query, and click Generate.

    Find flights from SFO to JFK.

    Review the SQL query. Notice that the agent generates the correct SQL for this unambiguous question.

      SELECT
    *
  FROM
    "flights"
  WHERE
    "departure_airport" = 'SFO' AND "arrival_airport" = 'JFK';

  5. In the Generate SQL using agent: flights-assistant window, click Edit.

  6. Enter the following natural language question to generate a SQL query, and click Update.

    Tell me flights that can help me beat nighttime traffic if traveling from New York

    The database fails to understand the term nighttime traffic. This might prevent it from generating a SQL query or cause it to generate a query that ignores the term, as the following query shows.

    -- The database schema does not contain information about traffic.
-- Returning all flights departing from New York airports.
SELECT
  f.airline,
  f.flight_number,
  a.name AS departure_airport_name,
  f.departure_time,
  b.name AS arrival_airport_name,
  f.arrival_time
FROM
  flights AS f
JOIN
  airports AS a
  ON f.departure_airport = a.iata
JOIN
  airports AS b
  ON f.arrival_airport = b.iata
WHERE
  a.city = 'New York'
ORDER BY
  f.departure_time;

Generate context for the agent

In this section, you create a context file that helps improve the agent's querying capabilities. To address the issue from the previous section where the agent didn't recognize the term nighttime traffic, define the term in the agent context as traffic occurring between 5:00 PM and 7:00 PM.

To generate agent context, perform the following steps:

  1. In your local directory, install Gemini CLI. For more information, see Gemini CLI quickstart.
  2. Install the Google Cloud CLI and set up Application Default Credentials (ADC).

  3. Install the MCP Toolbox Gemini CLI extension, which connects to the database.

    gemini extensions install https://github.com/gemini-cli-extensions/mcp-toolbox

  4. Create a tools.yaml configuration file in the same directory where you install the MCP toolbox to configure the database connection:

    sources:
  flight-sql-source:
    kind: alloydb-postgres
    project: PROJECT_ID
    region: REGION_ID
    cluster: CLUSTER_ID
    instance: INSTANCE_ID
    database: DATABASE_ID
    user: USER_NAME
    password: PASSWORD

tools:
  # (Optional) Fetches database schemas for context generation in the bulk generation (/generate_bulk_templates) phase.
  list_flight_schemas_tool:
    kind: postgres-list-tables
    source: flight-sql-source
    description: Use this tool to list all tables and their schemas in the flight database.
  # (Optional) Executes generated SQL for validation in the bulk generation (/generate_bulk_templates) phase.
  execute_sql_tool:
    kind: postgres-execute-sql
    source: flight-sql-source
    description: Use this tool to execute SQL against the flight database.

    Replace the following:

    • PROJECT_ID: Your Google Cloud project ID.
    • REGION_ID: The region of your AlloyDB cluster, for example, us-central1.
    • CLUSTER_ID: The ID of your AlloyDB cluster.
    • INSTANCE_ID: The ID of your AlloyDB primary instance.
    • DATABASE_ID: The name of the database to connect to.
    • USER_NAME: The database user. Set this as an environment variable and not literal value. For more information about how to set this value, see Sources in MCP toolbox.
    • PASSWORD: The password for the database user. Set this as an environment variable and not literal value. For more information about how to set this value, see Sources in MCP toolbox.

    For more information, see Quickstart (MCP with AlloyDB).

  5. Install the uv Python package installer by following official installation guide and verify the installation succeeds by running:

    uv --version

  6. Install the DB Context Enrichment MCP Server, which includes workflows for context generation.

    gemini extensions install https://github.com/GoogleCloudPlatform/db-context-enrichment

  7. Export the Gemini API key as an environment variable. For more information about how to find the API key, see Using Gemini API keys.

    export GEMINI_API_KEY="YOUR_API_KEY"

    Replace YOUR_API_KEY with your Gemini API key.

  8. In the same directory where you created the tools.yaml file, start Gemini:

    gemini

  9. Complete the Gemini CLI Authentication Setup.

  10. Verify that the MCP toolbox and the database enrichment extension are connected and ready to use.

    /mcp list

  11. Run the /generate_targeted_templates command and follow the workflow:

    /generate_targeted_templates

  12. Provide the natural language query that you want to add to the query template in the terminal.

    Tell me flights that can help me beat nighttime traffic if traveling from New York

  13. Provide a corresponding SQL query that you want to add to the query template. This query template defines the term nighttime as occurring between 5:00 PM and 7:00 PM.

    SELECT
  f.airline,
  f.flight_number,
  a.name AS airport_name,
  f.departure_time
FROM
  flights f
JOIN
  airports a
  ON f.departure_airport = a.iata
WHERE
  a.city = 'New York'
  AND (
    EXTRACT(HOUR FROM f.departure_time) < 17
    OR EXTRACT(HOUR FROM f.departure_time) >= 19
  )
ORDER BY
  f.departure_time;

  14. Press Enter. Gemini converts your input into a specific format that refines the agent's performance across a wide range of user queries. For more information, see Agent context.

    Optionally, run the /generate_bulk_templates workflow to let Gemini CLI generate more context by scanning your database schema and suggesting related context. Ensure you add both the list_flight_schemas_tool and the execute_sql_tool to the tools.yaml configuration you created in step 4.

  15. Review the generated query template. You can either save the query template as a new agent context file or append it to an existing agent context file.

  16. Select the option to create a new agent context file. Gemini creates a filename INSTANCE_ID_DATABASE_ID_context_set_TIMESTAMP.json in the same directory, with the following content:

    {
  "templates": [
    {
      "nl_query": "Tell me flights that can help me beat nighttime traffic if traveling from New York",
      "sql": "SELECT f.airline, f.flight_number, a.name AS airport_name, f.departure_time FROM flights f JOIN airports a ON f.departure_airport = a.iata WHERE a.city = 'New York' AND (EXTRACT(HOUR FROM f.departure_time) < 17 OR EXTRACT(HOUR FROM f.departure_time) >= 19) ORDER BY f.departure_time;",
      "intent": "Tell me flights that can help me beat nighttime traffic if traveling from New York",
      "manifest": "Tell me flights that can help me beat nighttime traffic if traveling from a given city",
      "parameterized": {
        "parameterized_sql": "SELECT f.airline, f.flight_number, a.name AS airport_name, f.departure_time FROM flights f JOIN airports a ON f.departure_airport = a.iata WHERE a.city = ? AND (EXTRACT(HOUR FROM f.departure_time) < 17 OR EXTRACT(HOUR FROM f.departure_time) >= 19) ORDER BY f.departure_time;",
        "parameterized_intent": "Tell me flights that can help me beat nighttime traffic if traveling from ?"
      }
    }
  ]
}

Upload context to the agent

In this section, you upload the agent context file to the data agent, so that it improves the agent's SQL generation capabilities on your database.

To upload the context, perform the following steps:

  1. In the Google Cloud console, go to the AlloyDB page.

    Go to AlloyDB

  2. Select a cluster from the list.

  3. In the navigation menu, click AlloyDB Studio.

  4. Sign in to Studio using Identity and Access Management (IAM) authentication.

  5. In the Explorer pane, next to Data Agents, click View actions.

  6. Click Edit agent.

  7. Optional: Edit Agent description.

  8. Click Browse in the Upload agent context file section, and select the agent context file generated earlier.

  9. Click Save.

Generate SQL query using agent context

In this section, you use the agent context file you uploaded to ask natural language questions. This lets you verify that the agent correctly understands and applies definitions for terms like nighttime traffic and other related phrases.

To generate SQL queries, perform the following steps:

  1. In the Explorer pane, next to your data agent, click View actions.
  2. Click Inspect agent.
  3. In the query editor, click Generate SQL using agent: flights-assistant.

  4. Enter the following natural language question to generate a SQL query, and click Generate.

    Tell me flights that can help me beat nighttime traffic if traveling from New York

    The generated SQL query looks similar to the following:

    SELECT
  f.airline,
  f.flight_number,
  a.name AS airport_name,
  f.departure_time
FROM
  flights f
JOIN
  airports a ON f.departure_airport = a.iata
WHERE
  a.city = 'New York'
  AND (
    EXTRACT(HOUR FROM f.departure_time) < 17
    OR EXTRACT(HOUR FROM f.departure_time) >= 19
  )
ORDER BY
  f.departure_time;

    This is the same question you added to the data agent's context. Observe that the agent can now accurately interpret the term nighttime traffic.

    Although the context originates from one particular question, the agent uses it to enhance SQL generation for a wide range of similar questions.

  5. In the Generate SQL using agent: flights-assistant window, click Edit.

  6. Enter the following similar question to generate a SQL query, and click Update.

    What are the flights that can help me avoid evening traffic if departing from Boston

    Since the question replaces the term nighttime traffic with a similar term, evening traffic, the agent provides a consistent answer to this question by applying the same interpretation.

    The generated SQL query looks similar to the following:

    -- What are the flights that can help me avoid evening traffic if departing from Boston
SELECT
  f.airline,
  f.flight_number,
  a.name AS airport_name,
  f.departure_time
FROM
  flights f
JOIN
  airports a
ON
  f.departure_airport = a.iata
WHERE
  a.city = 'Boston'
  AND (
    EXTRACT(HOUR FROM f.departure_time) < 17
    OR EXTRACT(HOUR FROM f.departure_time) >= 19
  )
ORDER BY
  f.departure_time;

Integrate the agent with your application

In this section, you create a data agent for a flight-finding application. This data agent provides a conversational interface to the flights and airports table you created earlier. It also explains how to create and integrate this agent into your application using Agent Development Kit (ADK), the Gemini Data Analytics QueryData MCP tool, and agent context to improve quality of the responses.

  1. Download MCP Toolbox version 0.24.0 or later. MCP toolbox exposes the data agent as a tool for applications to connect with. The MCP toolbox differs from the MCP Toolbox Gemini CLI extension you installed earlier, which generates context.

  2. In the terminal, set the project that you are using.

    gcloud config set project [PROJECT_ID]

  3. Set up Application Default Credentials (ADC).

    gcloud auth application-default login

  4. Find the agent context ID. For more information about how to find the context set ID, see Find the agent context ID.

  5. Create the tools.yaml configuration to connect to the data agent using the MCP toolbox. For more information, see Gemini Data Analytics Source and Gemini Data Analytics QueryData Tool.

    sources:
  gda-api-source:
    kind: cloud-gemini-data-analytics
    projectId: "PROJECT_ID"

tools:
  cloud_gda_query_tool:
    kind: cloud-gemini-data-analytics-query
    source: gda-api-source
    description: Use this tool to send natural language queries to the Gemini Data Analytics API and receive SQL, natural language answers, and explanations.
    location: "REGION_ID"
    context:
      datasourceReferences:
        alloydb:
          databaseReference:
            projectId: "PROJECT_ID"
            region: "REGION_ID"
            clusterId: "CLUSTER_ID"
            instanceId: "INSTANCE_ID"
            databaseId: "DATABASE_ID"
          agentContextReference:
            contextSetId: "DATA_AGENT_CONTEXT_SET_ID"
    generationOptions:
      generateQueryResult: true
      generateNaturalLanguageAnswer: true
      generateExplanation: true
      generateDisambiguationQuestion: true

    Replace the following:

    • PROJECT_ID: Your Google Cloud project ID.
    • REGION_ID: The region of your AlloyDB cluster, for example, us-central1.
    • CLUSTER_ID: The ID of your AlloyDB cluster.
    • INSTANCE_ID: The ID of your AlloyDB primary instance.
    • DATABASE_ID: The name of the database to connect to.
    • DATA_AGENT_CONTEXT_SET_ID: The data agent context set ID.

  6. Run the MCP Toolbox server with the tools.yaml file.

    ./toolbox --tools-file "tools.yaml"

  7. Create an ADK application that invokes Gemini Data Analytics QueryData tool using the MCP Toolbox's Python SDK. For more information about how to use the MCP Toolbox's Python SDK, see the quickstart for Toolbox and for Python ADK, see the quickstart for ADK.

    1. Create a directory to store the application, for example flight-assistant-app.

    2. Change directory to the flight-assistant-app directory.

      mkdir flight-assistant-app
cd flight-assistant-app

    3. Run the following commands under the flight-assistant-app directory to create a virtual environment and install required components.

      python3 -m venv .venv
source .venv/bin/activate
pip install toolbox-core
pip install google-genai
pip install google-adk

    4. Set up an ADK agent.

      1. Create an ADK agent.

        adk create my_agent

      2. Select the gemini-2.5-flash model.

      3. Select Google AI, and enter your Gemini API key. For more information about how to find your API key, see Using Gemini API keys.

    5. Replace the contents of the agent.py file with the following Flight Data Assistant sample application code.

      from typing import cast

from google.adk.agents.llm_agent import Agent
from google.adk.agents.llm_agent import ToolUnion

from toolbox_core import ToolboxSyncClient

TOOLBOX_URL = "http://127.0.0.1:5000"

INSTRUCTION = """
# ROLE
You are a friendly and factual flight data assistant. Your goal is to help users find the best flights for their needs by providing accurate information with a helpful, professional tone.
- use the Query Data Tool to answer the user's question, if the tool fails to generate a valid query, ask the user to clarify their question.

# OPERATIONAL CONSTRAINTS
- TOOL LIMITATION: You only have access to the Query Data Tool. Do not claim to have capabilities beyond what this tool provides.
- TRANSPARENCY POLICY: Maintain a seamless user experience. Never mention that you are using a tool, querying a database, or generating SQL. Frame all responses as your own direct assistance.
- SCOPE MANAGEMENT: If a user asks for something beyond your capabilities, politely state that you cannot perform that specific task. Guide the user towards what you can help with.

# COMMUNICATION STYLE
- Be concise and scannable when listing answers.
- Maintain a helpful, professional persona.

=====

# QUERY DATA TOOL

Inputs:
1. query: A natural language formulation of a database query.

Outputs: (all optional)
1. disambiguation_question: Clarification questions or comments where the tool needs the users' input.
2. generated_query: The generated query for the user query.
3. intent_explanation: An explanation for why the tool produced `generated_query`.
4. query_result: The result of executing `generated_query`.
5. natural_language_answer: The natural language answer that summarizes the `query` and `query_result`.

Usage guidance:
1. If `disambiguation_question` is produced, then solicit the needed inputs from the user and try the tool with a new `query` that has the needed clarification.
2. If `natural_language_answer` is produced, use `intent_explanation` and `generated_query` to see if you need to clarify any assumptions for the user.
3. If the tool output indicates failure or empty results, explain that clearly using the provided reasoning.
"""

client = ToolboxSyncClient(TOOLBOX_URL)

mcp_tool = client.load_tool("cloud_gda_query_tool")

root_agent = Agent(
    model="gemini-2.5-flash",
    name="root_agent",
    instruction=INSTRUCTION,
    tools=cast(list[ToolUnion], [mcp_tool]),
)

  8. Run the following commands under the flight-assistant-app directory to start the application and access the ADK web server at http://127.0.0.1:8000.

    adk web --port 8000

  9. Enter any text, such as hello, to start interacting with the agent.

    The ADK agent answers general questions and calls the required MCP tools.

  10. Enter the following flight-related question.

    How many flights depart from the west side?

    The MCP tool is called to answer this question. However, since the term the west is ambiguous and doesn't specify any airports, the MCP tool returns a disambiguation question which the agent uses to construct a response.

    I cannot determine how many flights depart from the 'west side' as the database does not contain information about which airports are considered to be on the 'west side'. However, I can help you with questions like:

1. How many flights depart from a specific airport?

2. What are the departure airports for all flights?

3. How many flights depart from each airport? Would you like to rephrase your question based on these options?

  11. Enter a question similar to that in the query template generated for the agent.

    Help me find flights from San Francisco that avoid the evening rush hour.

    Based on the agent context added earlier, the MCP tool understands that evening traffic occurs between 5 PM and 7 PM. The MCP tool returns the associated data for the agent to use in constructing its response.

    Here are the flights departing from San Francisco that avoid the evening rush hour (defined as 5 PM to 7 PM):

* UA 1532 departing at 05:50:00
* UA 1158 departing at 05:57:00
* CY 922 departing at 06:38:00
* OO 5441 departing at 07:08:00
* UA 616 departing at 07:14:00
* AA 24 departing at 07:14:00
* B6 434 departing at 08:00:00
* AA 242 departing at 08:18:00
* UA 1739 departing at 08:22:00
* OO 6336 departing at 08:32:00
* US 1784 departing at 08:47:00
* DL 1631 departing at 09:00:00
* DL 1106 departing at 09:06:00
* OO 5427 departing at 09:06:00
* CY 352 departing at 09:25:00

Iterate agent performance

The ADK web UI lets you inspect the request and response from the Gemini Data Analytics QueryData MCP tool. You can use this response to observe the tool responses such as generated SQL query, result set, intent explanation, disambiguation question, and natural language answer, to help you confirm the correctness of your agent's responses.

For example, for the input text How many flights depart from the west side? you entered earlier, click the agent bubble. In the Event tab in the left navigation, expand the functionResponse to see the following response.

"{"disambiguationQuestion": ["[NOT_ENOUGH_INFO] The database schema does not
contain information about which airports are on the 'west side'. Therefore, I
cannot determine how many flights depart from the west side.Possible alternative
questions: 1. How many flights depart from a specific airport? 2. What are the
departure airports for all flights? 3. How many flights depart from each
airport?"]}"

Refine response accuracy

You can continuously refine the accuracy of responses from the Gemini Data Analytics QueryData tool by adding additional context. Use the Gemini CLI to generate context, and then upload the updated agent context to the existing flights-assistant agent. For more information, see Build contexts using Gemini CLI. The console immediately ingests new context after you upload it, enabling you to enhance the agent's accuracy without any application downtime.

Multiple agents

In your development environment, you can perform A/B testing on multiple agent contexts by assigning distinct names to tools in your tools.yaml file. For example, you can create unique tools.yaml configurations by defining two cloud-gemini-data-analytics-query tools with different names, such as cloud_gda_query_tool_v1 and cloud_gda_query_tool_v2. This setup lets you implement application logic that programmatically selects the required agent context version by choosing the corresponding tool name.

The following example tools.yaml shows how to set up multiple agents for a database source:

sources:
  gda-api-source:
    kind: cloud-gemini-data-analytics
    projectId: "<var>PROJECT_ID</var>"
tools:
  cloud_gda_query_tool_v1:
    kind: cloud-gemini-data-analytics-query
    source: gda-api-source
    context:
      datasourceReferences:
        <var>DB_SOURCE</var>:
          databaseReference: ...
          agentContextReference:
            contextSetId: "V1_YOUR_DATA_AGENT_CONTEXT_SET_ID"
    generationOptions: ...
  cloud_gda_query_tool_v2:
    kind: cloud-gemini-data-analytics-query
    source: gda-api-source
    context:
      datasourceReferences:
        <var>DB_SOURCE</var>:
          databaseReference: ...
          agentContextReference:
            contextSetId: "V2_YOUR_DATA_AGENT_CONTEXT_SET_ID"
    generationOptions: ...

Replace the following:

  • PROJECT_ID: Your Google Cloud project ID.
  • V1_YOUR_DATA_AGENT_CONTEXT_SET_ID: The data agent context set ID for version 1.
  • V2_YOUR_DATA_AGENT_CONTEXT_SET_ID: The data agent context set ID for version 2.

Clean up

The following sections describe how to delete these resources and objects.

Delete the agent

Before you delete the cluster, delete the agent that you created.

  1. In the Google Cloud console, go to the AlloyDB page.

    Go to AlloyDB

  2. Select a cluster from the list.

  3. In the navigation menu, click AlloyDB Studio.

  4. Sign in to Studio using Identity and Access Management (IAM) authentication.

  5. In the Explorer pane, next to your data agent, click View actions.

  6. In the Delete agent window, enter flight-assistant in the confirmation box.

  7. Click Confirm.

Delete the cluster

When you delete the cluster that you created in the before you begin section, you also delete all of the objects you created.

  1. In the Google Cloud console, go to the AlloyDB page.

    Go to AlloyDB

  2. Select a cluster from the list.

  3. Click Delete cluster.

  4. In Delete cluster, enter name of the cluster to confirm you want to delete your cluster.

What's next