Test QueryData for Cloud SQL for PostgreSQL in Cloud SQL Studio

This document describes how to test QueryData and update the context set file. You can test QueryData's ability to generate SQL queries from natural language questions. If a generated query is not accurate, you can update the context set file.

To learn about context sets, see Context sets overview.

Before you begin

Make sure that you have created a context set in Knowledge Catalog and uploaded your context set file to the context set. For more information, see Manage context sets in Cloud SQL Studio.

To download, update, and upload context set files, make sure that the IAM user has the Dataplex Catalog Editor (roles/dataplex.catalogEditor) or Dataplex Catalog Admin (roles/dataplex.catalogAdmin) role on the project. For more information, see Context sets permissions.

Test QueryData

To test a QueryData, perform the following steps:

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

    Go to Cloud SQL

  2. Select an instance from the list.

  3. In the navigation menu, click Cloud SQL Studio.

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

  5. In the Explorer pane, click View actions next to the context set you're using.

  6. Click Test context set.

  7. In the query editor, click Generate SQL using QueryData with: CONTEXT_SET_NAME to open the query generation panel.

  8. Enter a natural language question in the query generation panel to generate a SQL query, and click Generate.

  9. Review the SQL query for accuracy. Click Insert to add the query to the editor.

Download and update a context set

If you aren't satisfied with the generated SQL query for a natural language question, download the existing context set file. You can then review and update the query template, and upload the updated context set file to the context set in Knowledge Catalog.

To download and update a context set, perform the following steps:

  1. In the Explorer pane, click View actions next to the context set you're using.
  2. Click Download context file.
  3. Follow steps in Build context sets using the context engineering agent to update context with additional query pairs.
  4. In the Explorer pane, click View actions next to the context set you're using.
  5. Click Edit context set.
  6. Click Browse in the Upload context file section, and select the updated context set file.
  7. Click Save to update the context set.

After you are satisfied with the accuracy of your responses, you can pass the context set ID to the stateless QueryData endpoint in requests from your application.

Find the context set ID

To reference your authored context set stored in Knowledge Catalog when calling the stateless QueryData endpoint from your application, you need the context set ID.

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

    Go to Cloud SQL

  2. Select an instance from the list.

  3. In the navigation menu, click Cloud SQL Studio.

  4. Sign into Studio using Identity and Access Management authentication.

  5. In the Explorer pane, click View actions next to the context set you're using.

  6. Click Edit context set.

  7. Note the context set resource name in Context set ID. The QueryData endpoint supports both of the following context set resource name formats in Knowledge Catalog:

    • 6-segment format: projects/PROJECT_ID/locations/LOCATION/contextSets/CONTEXT_SET_ID, such as projects/data-agents-project/locations/us-central1/contextSets/bdf_pg_all_templates
    • 8-segment format: projects/PROJECT_ID/locations/LOCATION/contextSetGroups/GROUP_ID/contextSets/CONTEXT_SET_ID@VERSION (for example, @v1)

Connect QueryData to your application

Set the context set ID in the QueryData method call to provide authored context from Knowledge Catalog for database data sources such as AlloyDB, Spanner, Cloud SQL for MySQL, Cloud SQL for PostgreSQL, Firestore, and Bigtable. For more information, see Define data agent context for database data sources.

After testing the context set, you can reference the database data source and the context set in your stateless QueryData call.

Example QueryData request with authored context

The following example sends a stateless QueryData request using the cloud_sql_reference database data source. The agent_context_reference.context_set_id field links the request to your pre-authored context set stored in Knowledge Catalog, not in your database:

curl -X POST \
  -H "Authorization: Bearer $(gcloud auth application-default print-access-token)" \
  -H "x-goog-user-project: PROJECT_ID" \
  -H "Content-Type: application/json" \
  https://geminidataanalytics.googleapis.com/v1beta/projects/PROJECT_ID/locations/LOCATION:queryData \
  -d '{
  "prompt": "How many accounts in the Prague region are eligible for loans? A3 contains the data of region.",
  "context": {
    "datasource_references": {
      "cloud_sql_reference": {
        "database_reference": {
          "engine": "POSTGRESQL",
          "project_id": "PROJECT_ID",
          "region": "LOCATION",
          "instance_id": "INSTANCE_ID",
          "database_id": "DATABASE_ID"
        },
        "agent_context_reference": {
          "context_set_id": "projects/PROJECT_ID/locations/LOCATION/contextSets/CONTEXT_SET_ID"
        }
      }
    }
  },
  "generation_options": {
    "generate_query_result": true,
    "generate_natural_language_answer": true,
    "generate_disambiguation_question": true,
    "generate_explanation": true
  }
}'

Replace the following:

  • PROJECT_ID: the ID of the Google Cloud project used for the request.
  • LOCATION: the region for the request, such as us-central1.
  • INSTANCE_ID: the instance ID of the Cloud SQL instance.
  • DATABASE_ID: the ID of the database.
  • CONTEXT_SET_ID: the ID of your context set.

The request body contains the following fields:

  • prompt: The natural language question from the end user.
  • context: Contains information about the data sources.
    • datasource_references: Specifies the data source type.
      • cloud_sql_reference: Required when querying the database. This field changes based on the database you are querying.
        • database_reference: Specifies information related to your database instance.
          • engine: The database engine. Set to POSTGRESQL for Cloud SQL instances.
          • project_id: The project ID of the database instance.
          • region: The region of the Cloud SQL instance.
          • instance_id: The instance ID of the Cloud SQL instance.
          • database_id: The ID of the database.
        • agent_context_reference: Links the stateless QueryData request to your authored context set stored in Knowledge Catalog, not in your database.
          • context_set_id: The full resource name of the context set stored in Knowledge Catalog. You can specify either the 6-segment format, projects/PROJECT_ID/locations/LOCATION/contextSets/CONTEXT_SET_ID, or the 8-segment format, projects/PROJECT_ID/locations/LOCATION/contextSetGroups/GROUP_ID/contextSets/CONTEXT_SET_ID@VERSION (for example, @v1).
  • generation_options: Configures the type of output to generate.
    • generate_query_result: Set to true to generate and return the query results.
    • generate_natural_language_answer: Optional. If set to true, generates a natural language answer.
    • generate_explanation: Optional. If set to true, generates an explanation of the SQL query.
    • generate_disambiguation_question: Optional. If set to true, generates disambiguation questions if the query is ambiguous.

Example QueryData response

Here is an example of a successful response from a QueryData call:

{
  "generatedQuery": "-- Count the number of accounts in Prague that are eligible for loans\nSELECT\n  COUNT(DISTINCT \"loans\".\"account_id\")\nFROM \"loans\"\nJOIN \"district\" -- Join based on district ID\n  ON \"loans\".\"district_id\" = \"district\".\"district_id\"\nWHERE\n  \"district\".\"A3\" = 'Prague'; -- Filter for the Prague region",
  "intentExplanation": "The question asks for the number of accounts eligible for loans in the Prague region. I need to join the `district` table with the `loans` table to filter by region and count the distinct accounts. The `A3` column in the `district` table contains the region information, and I'll filter for 'Prague'. The `loans` table contains information about loans, including the `account_id` and `district_id`. I will join these two tables on their respective district IDs.",
  "queryResult": {
    "columns": [
      {
        "name": "count"
      }
    ],
    "rows": [
      {
        "values": [
          {
            "value": "2"
          }
        ]
      }
    ],
    "totalRowCount": "1"
  },
  "naturalLanguageAnswer": "There are 2 accounts in Prague that are eligible for loans.",
  "tokenUsage": {
    "inputTokenCount": 312,
    "outputTokenCount": 148
  }
}

What's next