Integrate QueryData with an application

This tutorial describes how to set up and use QueryData in Firestore with MongoDB compatibility using the Google Cloud console and integrate it with your application. Learn how to build the context set file, create a context set that uses the context set file, use MCP Toolbox to call the QueryData method in the Conversational Analytics API to generate MongoDB-compatible query language (MQL) queries for natural language questions, and integrate it with your application.

For more information, see QueryData overview.

Objectives

  • Create databases and populate them with data.
  • Build a context set file with the context engineering agent.
  • Create a context set and upload a context set file.
  • Test QueryData and generate MQL queries in Firestore with MongoDB compatibility Studio.
  • Integrate QueryData with your application using Gemini Data Analytics QueryData tool in MCP Toolbox.
  • Add grounding to LLM responses using value search queries.

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.

To avoid continued billing, delete the resources you created when you finish the tasks in this document. For more information, see Clean up.

Before you begin

Complete the following prerequisites before creating a context set.

Enable required services

Enable the following services for your project:

Prepare a Firestore database

Make sure that you have a Firestore Enterprise edition database with MongoDB compatibility.

If you're trying the examples on this page, we recommend that you create a new database for this purpose. For more information, see Create and manage databases.

Required roles and permissions

  • Grant the roles/datastore.viewer and roles/geminidataanalytics.queryDataUser roles to the IAM user at the project level. For more information, see Add IAM policy binding for a project.

Populate the flights and airports collections

In this section, you populate the flights and airports collections using the provided MQL scripts.

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

    Go to Firestore

  2. Select a database from the list. For example, if you created a new database earlier, select it. The database must be an Enterprise edition database with MongoDB compatibility.

  3. In the navigation menu, click Firestore Studio.

  4. Click New MQL editor tab or New tab to open a new tab.

  5. Populate the airports collection by running the following script. Make sure the Read-only option isn't selected.

  6. Populate the flights collection by running the following script. Make sure the Read-only option isn't selected.

Create a context set in Firestore Studio

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

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

    Go to Firestore

  2. Select a database from the list. The database must be an Enterprise edition database in MongoDB compatibility mode.

  3. In the navigation menu, click Firestore Studio.

  4. In the Explorer pane, next to Context sets, click More actions ().

  5. Click Create context set.

  6. In Context set name, enter flights-assistant.

  7. Click Create.

Test QueryData in Studio

In this section, use the flights-assistant context set by asking natural language questions to generate a MQL query. Since the context set doesn't have any context set file uploaded to it, even after asking a question with context such as nighttime traffic, QueryData generates a sub-optimal query.

  1. In the Explorer pane, next to your context set, click More actions ().
  2. Click Test context set.
  3. In the query editor, click Generate MQL using QueryData with: flights-assistant.

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

    Find flights from SFO to JFK.

    Review the MQL query. Notice that QueryData generates the correct MQL for this unambiguous question.

      db.flights.find({"departure_airport": "SFO", "arrival_airport": "JFK"})

  5. In the Generate MQL using QueryData with: flights-assistant window, click Edit.

  6. Enter the following natural language question to generate a MQL 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 MQL 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.
db.flights.find({
  $and: [
    {
      "departure_airport": {"$regex": "New York", "$options": "i"}
    },
    {
      "$or": [
        {"departure_time": {"$exists": true}},
        {"arrival_time": {"$exists": true}}
      ]
    }
  ]
})

Generate context for the context set

In this section, you prepare the environment and create a context file that helps improve the context set's querying capabilities.

Set up your environment

Before you can start generating context, you must prepare your environment.

To set up your environment, perform the following steps:

  1. Install Antigravity CLI. For more information, see Get Started with Antigravity CLI.
  2. Install the Google Cloud CLI.

  3. Set up Application Default Credentials (ADC). Run the following commands in your terminal to authenticate and select your project:

    gcloud auth application-default login

  4. Install uv, a Python package and project manager that is a prerequisite for the context engineering agent. If you don't already have Python on your system, uv can install it for you.

  5. Install the context engineering agent:

    agy plugin install https://github.com/GoogleCloudPlatform/db-context-enrichment/tree/main/plugin

  6. In your terminal, start the Antigravity CLI.

    agy

  7. Set up a database connection for context generation using MCP Toolbox. The connection is defined in the tools.yaml configuration file.

    To create the tools.yaml configuration file in your current directory, enter a prompt such as: 

    Help me set up context generation for a Firestore database with MongoDB compatibility.

    Follow the instructions provided by the skill. For more information about the tools.yaml file, see MCP Toolbox documentation.

  8. To reload the configuration after the tools.yaml file is created, run the following command in the Antigravity CLI and follow the instruction to reload the toolbox MCP server:

    /mcp

    1. Select db-context-engineering_toolbox.

    2. Select Restart.

  9. Verify that the MCP toolbox and the context engineering agent plugin are connected and ready to use.

    /mcp

Generate template context

In this section, to address the issue from the previous section where QueryData didn't recognize the term nighttime traffic, define the term in the context set file as traffic occurring between 5:00 PM and 7:00 PM.

To generate template context, perform the following steps:

  1. Enter the following prompt into the Antigravity CLI:

    Generate a template for the question: "Tell me flights that can help me beat nighttime traffic if traveling from New York" with the following MQL query:

db.flights.aggregate([
  {
    "$lookup": {
      "from": "airports",
      "localField": "departure_airport",
      "foreignField": "iata",
      "as": "departure_airport_info",
      "pipeline": [{ "$project": { "_id": 0, "name": 1, "city": 1 } }]
    }
  },
  { "$unwind": { "path": "$departure_airport_info" } },
  {
    "$match": {
      "$expr": {
        "$and": [
          { "$eq": ["$departure_airport_info.city", "New York"] },
          {
            "$or": [
              { "$lte": [{ "$hour": "$departure_time" }, 17] },
              { "$gte": [{ "$hour": "$departure_time" }, 19] }
            ]
          }
        ]
      }
    }
  },
])

  2. The agent converts your input into a specific format that refines the context set's performance across a wide range of user queries. For more information, see Context sets overview.

  3. 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.

  4. Prompt the agent to save the template:

    Save it to context_set.json.

    The agent creates the file 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": "db.flights.aggregate([\n  {\n    \"$lookup\": {\n      \"from\": \"airports\",\n      \"localField\": \"departure_airport\",\n      \"foreignField\": \"iata\",\n      \"as\": \"departure_airport_info\",\n      \"pipeline\": [{ \"$project\": { \"_id\": 0, \"name\": 1, \"city\": 1 } }]\n    }\n  },\n  { \"$unwind\": { \"path\": \"$departure_airport_info\" } },\n  {\n    \"$match\": {\n      \"$expr\": {\n        \"$and\": [\n          { \"$eq\": [\"$departure_airport_info.city\", \"New York\"] },\n          {\n            \"$or\": [\n              { \"$lte\": [{ \"$hour\": \"$departure_time\" }, 17] },\n              { \"$gte\": [{ \"$hour\": \"$departure_time\" }, 19] }\n            ]\n          }\n        ]\n      }\n    }\n  }\n])",
      "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": "db.flights.aggregate([\n  {\n    \"$lookup\": {\n      \"from\": \"airports\",\n      \"localField\": \"departure_airport\",\n      \"foreignField\": \"iata\",\n      \"as\": \"departure_airport_info\",\n      \"pipeline\": [{ \"$project\": { \"_id\": 0, \"name\": 1, \"city\": 1 } }]\n    }\n  },\n  { \"$unwind\": { \"path\": \"$departure_airport_info\" } },\n  {\n    \"$match\": {\n      \"$expr\": {\n        \"$and\": [\n          { \"$eq\": [\"$departure_airport_info.city\", \"$1\"] },\n          {\n            \"$or\": [\n              { \"$lte\": [{ \"$hour\": \"$departure_time\" }, 17] },\n              { \"$gte\": [{ \"$hour\": \"$departure_time\" }, 19] }\n            ]\n          }\n        ]\n      }\n    }\n  }\n])",
        "parameterized_intent": "Tell me flights that can help me beat nighttime traffic if traveling from $1"
      }
    }
  ]
}

Upload the context set file to the QueryData

In this section, you upload the context set file to QueryData, so that it improves the QueryData's MQL generation capabilities on your database.

To upload the context, perform the following steps:

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

    Go to Firestore

  2. Select a database from the list. The database must be an Enterprise edition database in MongoDB compatibility mode.

  3. In the navigation menu, click Firestore Studio.

  4. In the Explorer pane, next to Context sets, click the Actions () icon.

  5. Click Edit context set.

  6. Click Browse in the Upload context set file section, and select the context set file generated earlier.

  7. Click Save.

Generate MQL query using QueryData

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

To generate MQL queries, perform the following steps:

  1. In the Explorer pane, next to your context set, click More actions ().
  2. Click Test context set.
  3. In the query editor, click Generate MQL using QueryData with: flights-assistant.

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

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

    The generated MQL query looks similar to the following:

    // Tell me flights that can help me beat nighttime traffic if traveling from New York
db.flights.aggregate([
  {
    "$lookup": {
      "from": "airports",
      "localField": "departure_airport",
      "foreignField": "iata",
      "as": "departure_airport_info",
      "pipeline": [{ "$project": { "_id": 0, "name": 1, "city": 1 } }]
    }
  },
  { "$unwind": { "path": "$departure_airport_info" } },
  {
    "$match": {
      "$expr": {
        "$and": [
          { "$eq": ["$departure_airport_info.city", "New York"] },
          {
            "$or": [
              { "$lte": [{ "$hour": "$departure_time" }, 17] },
              { "$gte": [{ "$hour": "$departure_time" }, 19] }
            ]
          }
        ]
      }
    }
  }
])

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

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

  5. In the Generate MQL using QueryData with: flights-assistant window, click Edit.

  6. Enter the following similar question to generate a MQL 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, QueryData provides a consistent answer to this question by applying the same interpretation.

    The generated MQL query looks similar to the following:

    // What are the flights that can help me avoid evening traffic if departing from Boston
db.flights.aggregate([
  {
    "$lookup": {
      "from": "airports",
      "localField": "departure_airport",
      "foreignField": "iata",
      "as": "departure_airport_info",
      "pipeline": [{ "$project": { "_id": 0, "name": 1, "city": 1 } }]
    }
  },
  { "$unwind": { "path": "$departure_airport_info" } },
  {
    "$match": {
      "$expr": {
        "$and": [
          { "$eq": ["$departure_airport_info.city", "Boston"] },
          {
            "$or": [
              { "$lte": [{ "$hour": "$departure_time" }, 17] },
              { "$gte": [{ "$hour": "$departure_time" }, 19] }
            ]
          }
        ]
      }
    }
  }
])

Integrate QueryData 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 tables you created earlier. It also explains how to create and integrate this data agent into your application using Agent Development Kit (ADK), the Gemini Data Analytics QueryData MCP tool, and a context set to improve quality of the responses.

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

  2. Set up Application Default Credentials (ADC).

    gcloud auth application-default login

  3. Find the ID of the context set you created earlier. For more information about how to find the context set ID, see Find the context set ID.

  4. Append the following to 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.

    kind: source
name: gda-api-source
type: cloud-gemini-data-analytics
projectId: "PROJECT_ID"
---
kind: tool
name: cloud_gda_query_tool
type: 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 MongoDB-compatible query language (MQL), natural language answers, and explanations.
location: "REGION_ID"
context:
  datasourceReferences:
    firestoreReference:
      databaseReference:
        projectId: "PROJECT_ID"
        databaseId: "DATABASE_ID"
      agentContextReference:
        contextSetId: "CONTEXT_SET_ID"
generationOptions:
  generateQueryResult: true
  generateNaturalLanguageAnswer: true
  generateExplanation: true
  generateDisambiguationQuestion: true

    Replace the following:

    • PROJECT_ID: Your Google Cloud project ID.
    • DATABASE_ID: The name of the database to connect to.
    • REGION_ID: The region of your database (for example us-central1).
    • CONTEXT_SET_ID: The context set ID. For more information about how to find the context set ID, see Find the context set ID.

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

    ./toolbox --config autoctx/tools.yaml

  6. 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.

      uv init
uv add google-adk toolbox-core google-genai

    4. Set up an ADK agent.

      1. Create an ADK agent.

        uv run adk create my_agent

      2. Select the gemini-3.8-flash model, or newer.

      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. (Alternatively, if you don't need a free trial, you can select Vertex AI and specify a Google Cloud project with billing enabled).

    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 MQL. 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-3.8-flash",  # Or newer.
    name="root_agent",
    instruction=INSTRUCTION,
    tools=cast(list[ToolUnion], [mcp_tool]),
)

  7. 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.

    uv run adk web --port 8000

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

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

  9. 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?

  10. 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 QueryData 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 MQL 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 context engineering agent to generate context, and then upload the updated context set file to the existing flights-assistant data agent. For more information, see Manage context sets in Firestore Studio. 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 data 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 context version by choosing the corresponding tool name.

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

kind: source
name: gda-api-source
type: cloud-gemini-data-analytics
projectId: PROJECT_ID
---
kind: tool
name: cloud_gda_query_tool_v1
type: cloud-gemini-data-analytics-query
source: gda-api-source
context:
  datasourceReferences:
    DB_SOURCE:
      databaseReference: ...
      agentContextReference:
        contextSetId: "V1_YOUR_CONTEXT_SET_ID"
generationOptions: ...
---
kind: tool
name: cloud_gda_query_tool_v2
type: cloud-gemini-data-analytics-query
source: gda-api-source
context:
  datasourceReferences:
    DB_SOURCE:
      databaseReference: ...
      agentContextReference:
        contextSetId: "V2_YOUR_CONTEXT_SET_ID"
generationOptions: ...

Replace the following:

  • PROJECT_ID: Your Google Cloud project ID.
  • V1_YOUR_CONTEXT_SET_ID: The context set ID for version 1.
  • V2_YOUR_CONTEXT_SET_ID: The context set ID for version 2.

Access control and security considerations

In the preview launch, QueryData for Firestore with MongoDB compatibility doesn't support fine-grained access controls. When integrating QueryData with your application, consider how data access is controlled:

  • Database-level granularity: Access granularity can be controlled using IAM conditions at the database level. However, there is no granular access control at the collection, document, or field level.
  • Full database read access: The data agent has read access to the entire database. Any query generated by the agent can read from all collections, documents, and fields within that database. The data agent isn't restricted by Firebase security rules.
  • Application user scoping: Because fine-grained access controls are not supported, you cannot restrict data access to specific collections, documents, or fields based on application user identities. Queries generated by QueryData execute with the full read permissions of the database user or service account used by your application.

Clean up

After you have finished following this tutorial, you should delete the resources that you created to avoid incurring ongoing costs. The following sections describe how to delete these resources.

Delete the context set

Before you delete the database, delete the context set that you created.

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

    Go to Firestore

  2. Select a database from the list. The database must be an Enterprise edition database in MongoDB compatibility mode.

  3. In the navigation menu, click Firestore Studio.

  4. In the Explorer pane, next to your context set, click More actions ().

  5. In the Delete context set window, enter flight-assistant in the confirmation box.

  6. Click Confirm.

Delete the database

When you delete the database 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 Firestore page.

    Go to Firestore

  2. Find the row for the name of the database that you want to delete.

  3. Click the Actions () icon, then click Delete.

  4. Confirm that you want to delete the database by entering the database name and clicking Delete.

What's next