From ad73cf9a4a73736a5534079f8a250d5ce42961ee Mon Sep 17 00:00:00 2001 From: totoleon Date: Tue, 21 Jul 2026 17:22:26 +0000 Subject: [PATCH 1/3] Add BigQuery support to evaluation config generation Implements the BigQuery portion of the Wave2 FR: - Add BigQueryConfigGenerator mapping tools.yaml bigquery sources to Evalbench db_config and GDA BigQueryTableReferences model config - Register the generator in the evaluate_generator factory map - Add bigquery.md connection reference template for the init skill - Attach agent_context_reference dynamically since the field is not yet in the public GDA SDK (restricted-visibility rollout) --- .../references/bigquery.md | 36 +++++++++ .../evaluate/db_generators/bigquery.py | 80 ++++++++++++++++++ .../evaluate/evaluate_generator.py | 2 + .../evaluate/db_generators/bigquery_test.py | 81 +++++++++++++++++++ 4 files changed, 199 insertions(+) create mode 100644 plugin/skills/context-engineering-init/references/bigquery.md create mode 100644 src/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery.py create mode 100644 tests/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery_test.py diff --git a/plugin/skills/context-engineering-init/references/bigquery.md b/plugin/skills/context-engineering-init/references/bigquery.md new file mode 100644 index 00000000..c9e108b1 --- /dev/null +++ b/plugin/skills/context-engineering-init/references/bigquery.md @@ -0,0 +1,36 @@ +## BigQuery + +**Required Information:** +- Data Source Name (e.g., `my-bigquery-db`) +- Google Cloud Project ID +- Dataset ID + +**Template:** + +```yaml +kind: source +name: +type: bigquery +project: +dataset: +--- +kind: tool +name: -list-schemas +type: bigquery-sql +source: +description: | + Use this tool to list tables and their schemas in the dataset. + + Progressive Schema Discovery (Recommended): + 1) Fetch structure first, + 2) Go deep on specific parts if interested, + 3) Use batching if info is too large. +statement: | + SELECT table_name, column_name, data_type, description FROM ``.``.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS ORDER BY table_name, column_name +--- +kind: tool +name: -execute-sql +type: bigquery-execute-sql +source: +description: Use this tool to execute SQL statements against the BigQuery dataset. +``` diff --git a/src/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery.py b/src/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery.py new file mode 100644 index 00000000..1b44adb7 --- /dev/null +++ b/src/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery.py @@ -0,0 +1,80 @@ +from typing import Any + +import google.cloud.geminidataanalytics_v1beta as gda +import yaml + +from .base import BaseDBConfigGenerator + + +class BigQueryConfigGenerator(BaseDBConfigGenerator): + """ + Dedicated generator mapping properties to explicit BigQuery configuration + topologies utilized by both EvalBench binaries and GDA Context objects. + """ + + SOURCE_TYPE = "bigquery" + DIALECT = "googlesql" + REQUIRED_FIELDS = BaseDBConfigGenerator.REQUIRED_FIELDS | { + "project", + "dataset", + } + + def __init__(self, params: dict[str, Any]): + super().__init__(params) + self.project = params.get("project") + self.dataset = params.get("dataset") + self.location = params.get("location") + # Optional explicit table scoping; the public GDA proto references + # BigQuery at table granularity rather than dataset granularity. + self.tables = params.get("tables") or [] + + def generate_db_config(self) -> str: + db_type = "bigquery" + db_path = f"projects/{self.project}/datasets/{self.dataset}" + + db_config = { + "db_type": db_type, + "dialect": self.DIALECT, + "database_name": self.dataset, + "database_path": db_path, + "gcp_project_id": self.project, + "max_executions_per_minute": 100, + } + if self.location: + db_config["location"] = self.location + return yaml.safe_dump( + db_config, sort_keys=False, default_flow_style=False + ).strip() + + def build_datasource_reference( + self, context_set_id: str + ) -> gda.DatasourceReferences: + datasource_ref = gda.DatasourceReferences() + + table_references = [ + gda.BigQueryTableReference( + project_id=self.project, + dataset_id=self.dataset, + table_id=table_id, + ) + for table_id in self.tables + ] + + bq_references = gda.BigQueryTableReferences( + table_references=table_references + ) + + # The agent_context_reference field on BigQueryTableReferences is not + # yet available in the public google-cloud-geminidataanalytics SDK + # (restricted-visibility rollout; see the Wave2 FR). Attach it + # dynamically so this generator works with internal SDK builds and + # degrades gracefully on public ones. + if "agent_context_reference" in { + f.name for f in type(bq_references).pb(bq_references).DESCRIPTOR.fields + }: + bq_references.agent_context_reference = gda.AgentContextReference( + context_set_id=context_set_id + ) + + datasource_ref.bq = bq_references + return datasource_ref diff --git a/src/google/cloud/db_context_enrichment/evaluate/evaluate_generator.py b/src/google/cloud/db_context_enrichment/evaluate/evaluate_generator.py index a776f923..aa45647f 100644 --- a/src/google/cloud/db_context_enrichment/evaluate/evaluate_generator.py +++ b/src/google/cloud/db_context_enrichment/evaluate/evaluate_generator.py @@ -10,6 +10,7 @@ from .db_generators.alloydb import AlloyDBConfigGenerator from .db_generators.base import BaseDBConfigGenerator +from .db_generators.bigquery import BigQueryConfigGenerator from .db_generators.mysql import MySQLConfigGenerator from .db_generators.postgres import PostgresConfigGenerator from .db_generators.spanner import SpannerConfigGenerator @@ -130,6 +131,7 @@ def _get_db_generator(params: dict[str, Any]) -> BaseDBConfigGenerator: generators = { AlloyDBConfigGenerator.SOURCE_TYPE: AlloyDBConfigGenerator, + BigQueryConfigGenerator.SOURCE_TYPE: BigQueryConfigGenerator, PostgresConfigGenerator.SOURCE_TYPE: PostgresConfigGenerator, MySQLConfigGenerator.SOURCE_TYPE: MySQLConfigGenerator, SpannerConfigGenerator.SOURCE_TYPE: SpannerConfigGenerator, diff --git a/tests/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery_test.py b/tests/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery_test.py new file mode 100644 index 00000000..1824b1da --- /dev/null +++ b/tests/google/cloud/db_context_enrichment/evaluate/db_generators/bigquery_test.py @@ -0,0 +1,81 @@ +import pytest +import yaml + +from google.cloud.db_context_enrichment.evaluate.db_generators.bigquery import ( + BigQueryConfigGenerator, +) + + +@pytest.fixture +def mock_params(): + return { + "project": "test-project", + "dataset": "test-dataset", + } + + +def test_generate_db_config(mock_params): + gen = BigQueryConfigGenerator(mock_params) + db_config_yaml = gen.generate_db_config() + + assert gen.DIALECT == "googlesql" + + config = yaml.safe_load(db_config_yaml) + assert config == { + "db_type": "bigquery", + "dialect": "googlesql", + "database_name": "test-dataset", + "database_path": "projects/test-project/datasets/test-dataset", + "gcp_project_id": "test-project", + "max_executions_per_minute": 100, + } + + +def test_generate_db_config_with_location(mock_params): + gen = BigQueryConfigGenerator({**mock_params, "location": "US"}) + config = yaml.safe_load(gen.generate_db_config()) + assert config["location"] == "US" + + +def test_missing_required_fields(): + with pytest.raises(ValueError, match="dataset"): + BigQueryConfigGenerator({"project": "test-project"}) + + +def test_generate_model_config(mock_params): + gen = BigQueryConfigGenerator(mock_params) + model_config_yaml = gen.generate_model_config( + "projects/test-project/locations/us-west1/contextSets/my-context" + ) + m_config = yaml.safe_load(model_config_yaml) + + assert m_config["generator"] == "query_data_api" + assert m_config["project_id"] == "test-project" + assert m_config["location"] == "global" + # The public GDA SDK references BigQuery at table granularity; with no + # explicit tables configured the reference set is empty but present. + assert "bq" in m_config["context"]["datasource_references"] + + +def test_generate_model_config_with_tables(mock_params): + gen = BigQueryConfigGenerator({**mock_params, "tables": ["t1", "t2"]}) + model_config_yaml = gen.generate_model_config( + "projects/test-project/locations/us-west1/contextSets/my-context" + ) + m_config = yaml.safe_load(model_config_yaml) + + table_refs = m_config["context"]["datasource_references"]["bq"][ + "table_references" + ] + assert table_refs == [ + { + "project_id": "test-project", + "dataset_id": "test-dataset", + "table_id": "t1", + }, + { + "project_id": "test-project", + "dataset_id": "test-dataset", + "table_id": "t2", + }, + ] From bbea672a03ca4ff006ac7d0b91e6e3211946a802 Mon Sep 17 00:00:00 2001 From: totoleon Date: Thu, 23 Jul 2026 17:51:14 +0000 Subject: [PATCH 2/3] Surface BigQuery in skill instructions and tool descriptions The init skill's database-type list, the generate_evalbench_configs tool docstring, and the evaluate skill's per-database references all enumerate supported types explicitly, so the BigQuery generator was unreachable from the agent workflows without these updates. --- .../references/bigquery.md | 17 +++++++++++++++++ plugin/skills/context-engineering-init/SKILL.md | 1 + src/google/cloud/db_context_enrichment/main.py | 2 +- 3 files changed, 19 insertions(+), 1 deletion(-) create mode 100644 plugin/skills/context-engineering-evaluate/references/bigquery.md diff --git a/plugin/skills/context-engineering-evaluate/references/bigquery.md b/plugin/skills/context-engineering-evaluate/references/bigquery.md new file mode 100644 index 00000000..ae73a0f0 --- /dev/null +++ b/plugin/skills/context-engineering-evaluate/references/bigquery.md @@ -0,0 +1,17 @@ +## BigQuery + +**Required properties from the `kind: source` block in `tools.yaml`:** +- Source Type (`type: bigquery`) +- Google Cloud Project ID (`project`) +- Dataset ID (`dataset`) + +**EvalBench Database Config Spec (`db_config.yaml`):** + +```yaml +db_type: bigquery +dialect: googlesql +database_name: +database_path: projects//datasets/ +gcp_project_id: +max_executions_per_minute: 100 +``` diff --git a/plugin/skills/context-engineering-init/SKILL.md b/plugin/skills/context-engineering-init/SKILL.md index 6c545518..2e3b8d43 100644 --- a/plugin/skills/context-engineering-init/SKILL.md +++ b/plugin/skills/context-engineering-init/SKILL.md @@ -70,6 +70,7 @@ When collecting information from the user, inform the user that only Application - Cloud SQL MySQL - AlloyDB Postgres - Spanner + - BigQuery 2. **Collect Information:** Request all **Required Information** based on the templates inside this directory. Do NOT assume missing fields; ask the user for them explicitly. 3. **Generate Configuration:** Replace all placeholders with the user's provided values and generate the complete `tools.yaml` content. Save it to the target location (e.g., `autoctx/tools.yaml` for Autoctx workflows, or `tools.yaml` in the current directory for standalone use). 4. **Validate:** After saving, validate the new connection using the toolbox script, replacing `` with the actual path to the file: diff --git a/src/google/cloud/db_context_enrichment/main.py b/src/google/cloud/db_context_enrichment/main.py index e2f0c393..28104ca4 100644 --- a/src/google/cloud/db_context_enrichment/main.py +++ b/src/google/cloud/db_context_enrichment/main.py @@ -73,7 +73,7 @@ def generate_evalbench_configs( dataset_path: The absolute path to the golden dataset file in the simplified user-facing format (JSON list of objects with keys: "id", "database", "nlq", "golden_sql"). context_set_id: Full ContextSet resource name to evaluate against. toolbox_config_path: The absolute path to the tools.yaml configuration file. - toolbox_source_name: The name of the database source to use inside tools.yaml. The underlying source block must use a supported 'type' (cloud-sql-postgres, cloud-sql-mysql, spanner, alloydb-postgres). + toolbox_source_name: The name of the database source to use inside tools.yaml. The underlying source block must use a supported 'type' (cloud-sql-postgres, cloud-sql-mysql, spanner, alloydb-postgres, bigquery). Returns: A message indicating that the configuration files were successfully created. From ceb27cf7b35a77d29726731f2f28d665f65d9257 Mon Sep 17 00:00:00 2001 From: totoleon Date: Thu, 23 Jul 2026 19:54:30 +0000 Subject: [PATCH 3/3] Require fully-qualified table names in BigQuery context generation BigQuery connections are scoped to a project with no default dataset at query time, so SQL referencing bare table names fails with 'Table not found'. Add BigQuery dialect references (template, facet, value search) to the context-generation-guide mandating project.dataset.table qualification, and call the rule out in the dataset-generation drafting chain of thought. --- .../references/generation-cot.md | 2 +- .../skills/context-generation-guide/SKILL.md | 5 +- .../references/facet/bigquery.md | 50 +++++++++++++++++ .../references/template/bigquery.md | 53 +++++++++++++++++++ .../references/value_search/bigquery.md | 41 ++++++++++++++ 5 files changed, 149 insertions(+), 2 deletions(-) create mode 100644 plugin/skills/context-generation-guide/references/facet/bigquery.md create mode 100644 plugin/skills/context-generation-guide/references/template/bigquery.md create mode 100644 plugin/skills/context-generation-guide/references/value_search/bigquery.md diff --git a/plugin/skills/context-engineering-dataset-generation/references/generation-cot.md b/plugin/skills/context-engineering-dataset-generation/references/generation-cot.md index a1f16db8..4ef3ecb4 100644 --- a/plugin/skills/context-engineering-dataset-generation/references/generation-cot.md +++ b/plugin/skills/context-engineering-dataset-generation/references/generation-cot.md @@ -3,7 +3,7 @@ When tasked with creating new pairs, using the generation plan as the guideline and north star. For *every* generated pair, execute the following internal Chain of Thought: -1. **Draft SQL (Schema-First):** Write syntactically perfect, dialect-compliant SQL using prioritized tables and conditions. Make sure to adhere to the technical schema and generation plan. Avoid inventing columns or tables based on business documents alone. +1. **Draft SQL (Schema-First):** Write syntactically perfect, dialect-compliant SQL using prioritized tables and conditions. Make sure to adhere to the technical schema and generation plan. Avoid inventing columns or tables based on business documents alone. For **BigQuery** sources, every table reference must be fully qualified as `` `project`.`dataset`.`table` `` (taken from the `kind: source` block in `tools.yaml`) — BigQuery has no default dataset at query time, so bare table names fail with "Table not found". 2. **Literal Translation:** Translate the SQL literally into English to guarantee no logical constraints are missed. 3. **Humanize & Bridge:** Rewrite the literal translation into natural business language, applying the *Semantic Bridge* principle. 4. **Verify and Refine:** Apply `acceptance-criteria.md` to verify the generated pair. If the pair fails the acceptance criteria, backtrack and try again. diff --git a/plugin/skills/context-generation-guide/SKILL.md b/plugin/skills/context-generation-guide/SKILL.md index 2b30f479..c1781fe1 100644 --- a/plugin/skills/context-generation-guide/SKILL.md +++ b/plugin/skills/context-generation-guide/SKILL.md @@ -21,7 +21,7 @@ Context generation allows you to create specific, high-value items in three form When asked to generate context items: 1. **Identify the Type**: Determine if the user wants to create a Template, Facet, or Value Search. 2. **Gather Information**: Ensure you have all the required information for the chosen context type as described in the "Context Type Definitions" section below. If information is missing, try to explore the database to find it or ask the user for clarification. -3. **Select Dialect Reference**: Identify the target database dialect (PostgreSQL, GoogleSQL, or MySQL) and consult the corresponding file in `references/` for specific syntax and patterns. +3. **Select Dialect Reference**: Identify the target database dialect (PostgreSQL, GoogleSQL for Spanner, GoogleSQL for BigQuery, or MySQL) and consult the corresponding file in `references/` for specific syntax and patterns. 4. **Parameterize**: Follow the [Phrase Extraction and Parameterization Guidelines](references/phrase_extraction/guidelines.md) to generalize the values. 5. **Format Output**: Construct the final JSON object according to the examples in the reference files. 6. **Save Context**: Use the appropriate MCP tool (e.g., `mutate_context_set`) to save or update the context set. @@ -131,12 +131,15 @@ For specific SQL templates, examples, and performance recommendations, refer to * **Templates**: * [PostgreSQL](references/template/postgresql.md) * [Spanner (GoogleSQL)](references/template/googlesql.md) + * [BigQuery (GoogleSQL)](references/template/bigquery.md) * [MySQL](references/template/mysql.md) * **Facets**: * [PostgreSQL](references/facet/postgresql.md) * [Spanner (GoogleSQL)](references/facet/googlesql.md) + * [BigQuery (GoogleSQL)](references/facet/bigquery.md) * [MySQL](references/facet/mysql.md) * **Value Searches**: * [PostgreSQL](references/value_search/postgresql.md) * [Spanner (GoogleSQL)](references/value_search/googlesql.md) + * [BigQuery (GoogleSQL)](references/value_search/bigquery.md) * [MySQL](references/value_search/mysql.md) diff --git a/plugin/skills/context-generation-guide/references/facet/bigquery.md b/plugin/skills/context-generation-guide/references/facet/bigquery.md new file mode 100644 index 00000000..60f1e9f0 --- /dev/null +++ b/plugin/skills/context-generation-guide/references/facet/bigquery.md @@ -0,0 +1,50 @@ +# BigQuery (GoogleSQL) Facet Generation Reference + +This reference provides best practices and ideal output definitions for generating Facets in BigQuery (GoogleSQL). + +## Concepts + +Facets are reusable, modular SQL fragments (like a `WHERE` clause or specialized join). They are dynamically injected filters linked to specific vocabulary or terminology. + +## Fully-Qualified Column References + +Every column reference in a facet's SQL snippet **must** be qualified with its table name as `table.column` (e.g., `products.rating`). Facets are injected into larger queries that may join multiple tables, so unqualified columns risk ambiguity errors or silently binding to the wrong column. Never use table aliases — the surrounding query controls aliasing. + +Use `table.column` (NOT `project.dataset.table.column`) for column +references inside snippets — the surrounding query's `FROM` clause carries +the fully-qualified `` `project`.`dataset`.`table` `` reference. However, if +a facet snippet itself contains a `FROM` clause (e.g., an `EXISTS` +subquery), every table inside it must be fully qualified as +`` `project`.`dataset`.`table` `` because BigQuery has no default dataset +at query time. + +## Parameterization + +Values in the SQL snippet and the intent must be replaced with positional parameters represented by `?`, according to the [Phrase Extraction and Parameterization Guidelines](../phrase_extraction/guidelines.md). + +### Example + +**Input**: +* **SQL Snippet**: `products.rating > 4.5` +* **Intent**: "highly rated products (above 4.5)" + +**Generated Output** (Conceptual): +```json +{ + "sql_snippet": "products.rating > 4.5", + "intent": "highly rated products (above 4.5)", + "manifest": "highly rated products (above a given number)", + "parameterized": { + "parameterized_sql_snippet": "products.rating > ?", + "parameterized_intent": "highly rated products (above ?)" + } +} +``` + +## Best Practices + +* Provide clear and reusable SQL snippets. +* **Always qualify columns as `table.column`** in both the literal and parameterized SQL snippets. Never use bare column names or table aliases. +* **Fully qualify any table inside subquery `FROM` clauses** as `` `project`.`dataset`.`table` ``. +* Ensure the SQL snippet follows BigQuery (GoogleSQL) syntax. +* The intent should clearly describe the condition or filter. diff --git a/plugin/skills/context-generation-guide/references/template/bigquery.md b/plugin/skills/context-generation-guide/references/template/bigquery.md new file mode 100644 index 00000000..0186b54c --- /dev/null +++ b/plugin/skills/context-generation-guide/references/template/bigquery.md @@ -0,0 +1,53 @@ +# BigQuery (GoogleSQL) Template Generation Reference + +This reference provides best practices and ideal output definitions for generating Templates in BigQuery (GoogleSQL). + +## Concepts + +Templates map full natural language questions to full SQL queries. They are used to teach the system overarching operational logic. + +## Fully-Qualified Table References (CRITICAL) + +Unlike Postgres or Spanner connections, a BigQuery connection is scoped to a +project only — there is **no default dataset at query time**. Every table +reference in the SQL **must** be fully qualified as +`` ``.``.`` `` (e.g., +`` `my-project`.`sales_data`.`orders` ``). SQL that references a bare table +name (`FROM orders`) will fail with "Table not found" when executed. + +Take the project and dataset IDs from the `kind: source` block in +`tools.yaml`. Apply this rule to the `sql` field AND the +`parameterized_sql` field — never parameterize the project, dataset, or +table identifiers themselves; only parameterize filter values. + +## Parameterization + +Values in the SQL query and the intent must be replaced with positional parameters represented by `?`, according to the [Phrase Extraction and Parameterization Guidelines](../phrase_extraction/guidelines.md). + +### Example + +**Input**: +* **Question**: "How many accounts are in London?" +* **SQL**: ``SELECT count(*) FROM `my-project`.`finance`.`account` WHERE city = 'London'`` +* **Intent**: "How many accounts are in London?" + +**Generated Output** (Conceptual): +```json +{ + "nl_query": "How many accounts are in London?", + "sql": "SELECT count(*) FROM `my-project`.`finance`.`account` WHERE city = 'London'", + "intent": "How many accounts are in London?", + "manifest": "How many accounts are in a given city?", + "parameterized": { + "parameterized_sql": "SELECT count(*) FROM `my-project`.`finance`.`account` WHERE city = ?", + "parameterized_intent": "How many accounts are in ?" + } +} +``` + +## Best Practices + +* Provide complete, executable SQL queries. +* Ensure the SQL follows BigQuery (GoogleSQL) syntax. +* **Always fully qualify every table as `` `project`.`dataset`.`table` ``** in both `sql` and `parameterized_sql`. +* The intent should accurately describe what the query does. diff --git a/plugin/skills/context-generation-guide/references/value_search/bigquery.md b/plugin/skills/context-generation-guide/references/value_search/bigquery.md new file mode 100644 index 00000000..afa848a2 --- /dev/null +++ b/plugin/skills/context-generation-guide/references/value_search/bigquery.md @@ -0,0 +1,41 @@ +# BigQuery (GoogleSQL) Value Search Templates + +This reference provides the SQL templates and examples for Value Search in BigQuery (GoogleSQL). + +## Requirements + +* Every table reference **must** be fully qualified as `` `{project}`.`{dataset}`.`{table}` `` — BigQuery has no default dataset at query time. + +## Supported Match Functions + +### 1. EXACT_MATCH_STRINGS + +**Description**: Exact match for strings in BigQuery. +**Example**: Use for exact IDs or state codes. + +**Template**: +```sql +SELECT CAST($value AS STRING) AS value, '{column}' AS `columns`, +'{concept_type}' AS concept_type, 0 AS distance, +JSON '{}' AS context +FROM `{project}`.`{dataset}`.`{table}` AS T +WHERE CAST(T.`{column}` AS STRING) = CAST($value AS STRING) +``` + +### 2. EDIT_DISTANCE_MATCH + +**Description**: String similarity using BigQuery's built-in `EDIT_DISTANCE` function (Levenshtein distance). No index prerequisites. +**Example**: Use for typos/misspellings (e.g., "Lndn" → "London"). + +**Template**: +```sql +SELECT CAST(T.`{column}` AS STRING) AS value, '{column}' AS `columns`, +'{concept_type}' AS concept_type, +EDIT_DISTANCE(LOWER(CAST(T.`{column}` AS STRING)), LOWER(CAST($value AS STRING))) AS distance, +JSON '{}' AS context +FROM `{project}`.`{dataset}`.`{table}` AS T +WHERE EDIT_DISTANCE(LOWER(CAST(T.`{column}` AS STRING)), LOWER(CAST($value AS STRING))) <= 3 +``` + +**Performance Recommendations**: +* Value-search scans are full-table scans in BigQuery. Prefer running them against low-cardinality dimension tables, or pre-materialize a `SELECT DISTINCT {column}` lookup table to bound bytes scanned.