google-cloud-db-context-engineering 0.6.0__tar.gz → 0.7.0__tar.gz

google_cloud_db_context_engineering-0.7.0/PKG-INFO

@@ -0,0 +1,110 @@
+Metadata-Version: 2.4
+Name: google-cloud-db-context-engineering
+Version: 0.7.0
+Summary: A FastMCP server for generating natural language to SQL templates from database schemas.
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: fastmcp==3.3.1
+Requires-Dist: google-genai==1.46.0
+Requires-Dist: toolbox-core==0.5.2
+Requires-Dist: google-cloud-geminidataanalytics==0.11.0
+Requires-Dist: google-auth==2.41.1
+Requires-Dist: requests==2.33.0
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-asyncio; extra == "test"
+Dynamic: license-file
+
+This is not an officially supported Google product. This project is not eligible for the [Google Open Source Software Vulnerability Rewards Program](https://bughunters.google.com/open-source-security), [Google Cloud Platform/SecOps Terms of Service](https://cloud.google.com/terms), [How Gemini for Google Cloud uses your data](https://cloud.google.com/gemini/docs/discover/data-governance). This tool is provided "as is" without warranty of any kind. Users are solely responsible for understanding and managing the tool's interaction with their databases. Use of this tool constitutes acceptance of all risks associated with database access, reading, usage, and modifications.
+
+# Context Engineering Agent
+
+The **Context Engineering Agent** is an AI coding agent plugin designed to run in developer agent harnesses (such as Claude Code, Antigravity, or Gemini CLI). It generates, evaluates, and iteratively tunes tailored context artifacts (`ContextSets` comprising `Templates`, `Facets`, and `Value Searches`) to enrich database schemas for **Gemini Data Analytics's data agent developer platform tools**, such as QueryData ([AlloyDB](https://docs.cloud.google.com/gemini/data-agents/querydata/alloydb/data-agent-overview) | Cloud SQL: [PostgreSQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-postgres/data-agent-overview) / [MySQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-mysql/data-agent-overview) | [Spanner (GoogleSQL)](https://docs.cloud.google.com/gemini/data-agents/querydata/spanner/data-agent-overview)).
+
+---
+
+## Why Context Engineering?
+
+When building data agents and natural language analytics interfaces, accurately translating user intent into database queries is critical. 
+
+As outlined in **Build Context with Context Engineering Agent** ([AlloyDB](https://docs.cloud.google.com/gemini/data-agents/querydata/alloydb/build-context-gemini-cli) | Cloud SQL: [PostgreSQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-postgres/build-context-gemini-cli) / [MySQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-mysql/build-context-gemini-cli) | [Spanner (GoogleSQL)](https://docs.cloud.google.com/gemini/data-agents/querydata/spanner/build-context-gemini-cli)), by optimizing a `ContextSet` to match your application's expected query stream, the **QueryData API** acts as a data agent tool capable of achieving **~100% NL-to-SQL translation accuracy with low latency**.
+
+---
+
+## Core Concepts
+
+A `ContextSet` is the central artifact generated and managed by the agent, containing structured knowledge in three primary forms:
+
+* **Templates**: Link natural language query patterns to complete query statements.
+* **Facets**: Reusable, modular query fragments (e.g., parameterized `WHERE` clauses or specialized join filters) linked to domain vocabulary.
+* **Value Searches**: Specialized mapping queries that dynamically resolve user-supplied values (e.g., *"Lndn"*) to database records (*"London"*) via the capabilities of the underlying database, such as embedding search, AI operators, or simple trigram search.
+
+For full schema details, structure specifications, and dialect-specific JSON representations of `ContextSets`, see the official **Context Sets Overview** ([AlloyDB](https://docs.cloud.google.com/gemini/data-agents/querydata/alloydb/context-sets-overview) | Cloud SQL: [PostgreSQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-postgres/context-sets-overview) / [MySQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-mysql/context-sets-overview) | [Spanner (GoogleSQL)](https://docs.cloud.google.com/gemini/data-agents/querydata/spanner/context-sets-overview)).
+
+---
+
+## Prerequisites & Environment Setup
+
+Before getting started, prepare your GCP environment, required APIs (Data Analytics API, Gemini for Google Cloud API, Dataplex Universal Catalog API), IAM permissions, and database Data API settings.
+
+Follow the step-by-step setup guide in the official documentation:
+👉 **Prepare Your Environment**: ([AlloyDB](https://docs.cloud.google.com/gemini/data-agents/querydata/alloydb/build-context-gemini-cli#prepare-your-environment) | Cloud SQL: [PostgreSQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-postgres/build-context-gemini-cli#prepare-your-environment) / [MySQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-mysql/build-context-gemini-cli#prepare-your-environment) | [Spanner (GoogleSQL)](https://docs.cloud.google.com/gemini/data-agents/querydata/spanner/build-context-gemini-cli#prepare-your-environment))
+
+---
+
+## Primary Workflow Phases
+
+The agent enables you to craft an optimized context for QueryData API through three primary phases:
+
+### Phase 1: Artifact Ingestion
+*Why it matters: Without broader context on the application's goals and scope, AI models generate sterile queries based solely on database column names, missing how your users actually ask for information.*
+
+1. **Connect & Discover**: Inspect any resources shared with the agent that may be relevant to your application and its domain.
+2. **Domain Concept Extraction**: Extract key terminology, domain rules, jargon, and common query patterns.
+3. **Schema Mapping**: Map extracted concepts directly to their corresponding underlying database tables and columns.
+
+### Phase 2: Dataset Creation
+*Why it matters: A realistic golden dataset establishes your benchmark for accuracy, ensuring your data agent is evaluated against messy human language rather than artificially simple database queries.*
+
+1. **Dataset Planning**: Set goals for your dataset, deciding on question difficulty, business topics to cover, and total dataset size.
+2. **Initial Question Generation**: Create a core set of natural language questions grounded in your business documents, verifying that each question's corresponding query runs accurately on your database.
+3. **Dataset Expansion**: Scale up the question set by adding real-world variations, such as different human phrasing, complex combinations of joins and measures, and realistic values—to cover the ambiguity of natural language.
+4. **Dataset Validation**: Audit the full dataset against your original plan and ask for your approval before proceeding. 
+
+### Phase 3: Context Optimization via Recursive Hill-Climbing
+*Why it matters: Iterative evaluation and gap analysis systematically fix query translation failures, driving accuracy toward ~100% while ensuring fixes don't break previously working queries.*
+
+The optimization loop creates an initial `ContextSet` and then iteratively refines it using evaluation feedback:
+
+1. **Bootstrap**: Generate an initial baseline context.
+2. **Evaluate**: Measure context effectiveness against a golden dataset.
+3. **Hill-Climbing**: Perform gap analysis on failures and generate automated fixes.
+4. **Iterate**: Apply the improved context and re-run evaluation to continuously improve quality.
+5. **Final Validation** (Optional): Verify mutations against a separated test set to ensure generalization and prevent overfitting.
+
+*Note: While there is a typical ordering for these CUJs, the agent is flexible in how you want to execute. You can run the full pipeline end-to-end, trigger any individual phase, or ask for targeted changes to the `ContextSet`.*
+
+---
+
+## Bring Artifacts via Filesystem or MCP
+
+To prevent the AI from generating trivial schema-only questions (e.g., *"What is the count of users?"*), the agent bridges the semantic gap by grounding context generation in real-world business documents—such as **product glossaries**, **business wikis**, **SOPs**, **ORM data models**, **emails**, or **application or database logs**.
+
+You can provide business artifacts directly to the agent from local filesystems or via **Model Context Protocol (MCP) Servers**.
+
+---
+
+## How to Use
+
+Launch your agent harness (Gemini CLI, Claude Code, or Antigravity) in your workspace directory and interact in natural language:
+
+### Example Natural Language Prompts
+* **End-to-End**: (From the app directory) *"Optimize context for my app."*
+* **Dataset Curation**: *"Expand my dataset with app changes in `<PULL_REQUEST_LINK>`."*
+* **Ad-hoc Evaluation**: *"Evaluate accuracy of QueryData with `ContextSet` `<context_set_id>` on dataset.json."*
+* **Targeted Authoring**: *"Add a facet for active premium subscriptions."*
+
+---
+
+> 🛠️ **Developer Note**: For developer setup instructions, local CLI linking, unit testing (`pytest`), linting (`ruff`), and fork release testing, see [docs/development.md](docs/development.md).

google_cloud_db_context_engineering-0.7.0/README.md

@@ -0,0 +1,92 @@
+This is not an officially supported Google product. This project is not eligible for the [Google Open Source Software Vulnerability Rewards Program](https://bughunters.google.com/open-source-security), [Google Cloud Platform/SecOps Terms of Service](https://cloud.google.com/terms), [How Gemini for Google Cloud uses your data](https://cloud.google.com/gemini/docs/discover/data-governance). This tool is provided "as is" without warranty of any kind. Users are solely responsible for understanding and managing the tool's interaction with their databases. Use of this tool constitutes acceptance of all risks associated with database access, reading, usage, and modifications.
+
+# Context Engineering Agent
+
+The **Context Engineering Agent** is an AI coding agent plugin designed to run in developer agent harnesses (such as Claude Code, Antigravity, or Gemini CLI). It generates, evaluates, and iteratively tunes tailored context artifacts (`ContextSets` comprising `Templates`, `Facets`, and `Value Searches`) to enrich database schemas for **Gemini Data Analytics's data agent developer platform tools**, such as QueryData ([AlloyDB](https://docs.cloud.google.com/gemini/data-agents/querydata/alloydb/data-agent-overview) | Cloud SQL: [PostgreSQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-postgres/data-agent-overview) / [MySQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-mysql/data-agent-overview) | [Spanner (GoogleSQL)](https://docs.cloud.google.com/gemini/data-agents/querydata/spanner/data-agent-overview)).
+
+---
+
+## Why Context Engineering?
+
+When building data agents and natural language analytics interfaces, accurately translating user intent into database queries is critical. 
+
+As outlined in **Build Context with Context Engineering Agent** ([AlloyDB](https://docs.cloud.google.com/gemini/data-agents/querydata/alloydb/build-context-gemini-cli) | Cloud SQL: [PostgreSQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-postgres/build-context-gemini-cli) / [MySQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-mysql/build-context-gemini-cli) | [Spanner (GoogleSQL)](https://docs.cloud.google.com/gemini/data-agents/querydata/spanner/build-context-gemini-cli)), by optimizing a `ContextSet` to match your application's expected query stream, the **QueryData API** acts as a data agent tool capable of achieving **~100% NL-to-SQL translation accuracy with low latency**.
+
+---
+
+## Core Concepts
+
+A `ContextSet` is the central artifact generated and managed by the agent, containing structured knowledge in three primary forms:
+
+* **Templates**: Link natural language query patterns to complete query statements.
+* **Facets**: Reusable, modular query fragments (e.g., parameterized `WHERE` clauses or specialized join filters) linked to domain vocabulary.
+* **Value Searches**: Specialized mapping queries that dynamically resolve user-supplied values (e.g., *"Lndn"*) to database records (*"London"*) via the capabilities of the underlying database, such as embedding search, AI operators, or simple trigram search.
+
+For full schema details, structure specifications, and dialect-specific JSON representations of `ContextSets`, see the official **Context Sets Overview** ([AlloyDB](https://docs.cloud.google.com/gemini/data-agents/querydata/alloydb/context-sets-overview) | Cloud SQL: [PostgreSQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-postgres/context-sets-overview) / [MySQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-mysql/context-sets-overview) | [Spanner (GoogleSQL)](https://docs.cloud.google.com/gemini/data-agents/querydata/spanner/context-sets-overview)).
+
+---
+
+## Prerequisites & Environment Setup
+
+Before getting started, prepare your GCP environment, required APIs (Data Analytics API, Gemini for Google Cloud API, Dataplex Universal Catalog API), IAM permissions, and database Data API settings.
+
+Follow the step-by-step setup guide in the official documentation:
+👉 **Prepare Your Environment**: ([AlloyDB](https://docs.cloud.google.com/gemini/data-agents/querydata/alloydb/build-context-gemini-cli#prepare-your-environment) | Cloud SQL: [PostgreSQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-postgres/build-context-gemini-cli#prepare-your-environment) / [MySQL](https://docs.cloud.google.com/gemini/data-agents/querydata/sql-mysql/build-context-gemini-cli#prepare-your-environment) | [Spanner (GoogleSQL)](https://docs.cloud.google.com/gemini/data-agents/querydata/spanner/build-context-gemini-cli#prepare-your-environment))
+
+---
+
+## Primary Workflow Phases
+
+The agent enables you to craft an optimized context for QueryData API through three primary phases:
+
+### Phase 1: Artifact Ingestion
+*Why it matters: Without broader context on the application's goals and scope, AI models generate sterile queries based solely on database column names, missing how your users actually ask for information.*
+
+1. **Connect & Discover**: Inspect any resources shared with the agent that may be relevant to your application and its domain.
+2. **Domain Concept Extraction**: Extract key terminology, domain rules, jargon, and common query patterns.
+3. **Schema Mapping**: Map extracted concepts directly to their corresponding underlying database tables and columns.
+
+### Phase 2: Dataset Creation
+*Why it matters: A realistic golden dataset establishes your benchmark for accuracy, ensuring your data agent is evaluated against messy human language rather than artificially simple database queries.*
+
+1. **Dataset Planning**: Set goals for your dataset, deciding on question difficulty, business topics to cover, and total dataset size.
+2. **Initial Question Generation**: Create a core set of natural language questions grounded in your business documents, verifying that each question's corresponding query runs accurately on your database.
+3. **Dataset Expansion**: Scale up the question set by adding real-world variations, such as different human phrasing, complex combinations of joins and measures, and realistic values—to cover the ambiguity of natural language.
+4. **Dataset Validation**: Audit the full dataset against your original plan and ask for your approval before proceeding. 
+
+### Phase 3: Context Optimization via Recursive Hill-Climbing
+*Why it matters: Iterative evaluation and gap analysis systematically fix query translation failures, driving accuracy toward ~100% while ensuring fixes don't break previously working queries.*
+
+The optimization loop creates an initial `ContextSet` and then iteratively refines it using evaluation feedback:
+
+1. **Bootstrap**: Generate an initial baseline context.
+2. **Evaluate**: Measure context effectiveness against a golden dataset.
+3. **Hill-Climbing**: Perform gap analysis on failures and generate automated fixes.
+4. **Iterate**: Apply the improved context and re-run evaluation to continuously improve quality.
+5. **Final Validation** (Optional): Verify mutations against a separated test set to ensure generalization and prevent overfitting.
+
+*Note: While there is a typical ordering for these CUJs, the agent is flexible in how you want to execute. You can run the full pipeline end-to-end, trigger any individual phase, or ask for targeted changes to the `ContextSet`.*
+
+---
+
+## Bring Artifacts via Filesystem or MCP
+
+To prevent the AI from generating trivial schema-only questions (e.g., *"What is the count of users?"*), the agent bridges the semantic gap by grounding context generation in real-world business documents—such as **product glossaries**, **business wikis**, **SOPs**, **ORM data models**, **emails**, or **application or database logs**.
+
+You can provide business artifacts directly to the agent from local filesystems or via **Model Context Protocol (MCP) Servers**.
+
+---
+
+## How to Use
+
+Launch your agent harness (Gemini CLI, Claude Code, or Antigravity) in your workspace directory and interact in natural language:
+
+### Example Natural Language Prompts
+* **End-to-End**: (From the app directory) *"Optimize context for my app."*
+* **Dataset Curation**: *"Expand my dataset with app changes in `<PULL_REQUEST_LINK>`."*
+* **Ad-hoc Evaluation**: *"Evaluate accuracy of QueryData with `ContextSet` `<context_set_id>` on dataset.json."*
+* **Targeted Authoring**: *"Add a facet for active premium subscriptions."*
+
+---
+
+> 🛠️ **Developer Note**: For developer setup instructions, local CLI linking, unit testing (`pytest`), linting (`ruff`), and fork release testing, see [docs/development.md](docs/development.md).

{google_cloud_db_context_engineering-0.6.0 → google_cloud_db_context_engineering-0.7.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "google-cloud-db-context-engineering"
-version = "0.6.0"
+version = "0.7.0"
 description = "A FastMCP server for generating natural language to SQL templates from database schemas."
 readme = "README.md"
 requires-python = ">=3.12"
@@ -10,6 +10,8 @@ dependencies = [
     "google-genai==1.46.0",
     "toolbox-core==0.5.2",
     "google-cloud-geminidataanalytics==0.11.0",
+    "google-auth==2.41.1",
+    "requests==2.33.0",
 ]
 
 [project.scripts]

google_cloud_db_context_engineering-0.7.0/src/google/cloud/db_context_enrichment/common/context_store_client.py

@@ -0,0 +1,190 @@
+"""REST client for the Context Store API.
+
+Hand-rolled because the service is GOOGLE_INTERNAL and has no public SDK.
+All calls target the autopush sandbox endpoint pinned in this module.
+"""
+
+import json
+import time
+from typing import Any
+
+import google.auth
+import google.auth.exceptions
+import requests
+from google.auth.transport import requests as auth_requests
+
+from google.cloud.db_context_enrichment.model import context
+
+CONTEXT_STORE_ENDPOINT = "https://autopush-dataplex.sandbox.googleapis.com"
+CONTEXT_STORE_LOCATION = "us-central1"
+API_VERSION = "v1"
+DEFAULT_OAUTH_SCOPES = ("https://www.googleapis.com/auth/cloud-platform",)
+
+# Per-request HTTP timeout + LRO poll schedule. Transient-error retry is
+# intentionally NOT done here — the MCP agent driving these tools retries
+# failed tool calls on its own, and the Cloud SDK (once available) will bring
+# proper retry back. Only the per-call timeout and op polling stay in this
+# hand-rolled client.
+_LRO_POLL_INTERVALS_SECONDS = (2.0, 4.0, 8.0, 16.0, 16.0, 16.0)
+_REQUEST_TIMEOUT_SECONDS = 30
+
+
+class ContextStoreClient:
+    """REST client for the Context Store API (autopush sandbox).
+
+    The client is stateless w.r.t. project. Operations that build a resource
+    path from IDs (`ensure_*`) take `project_id` as an explicit argument;
+    operations that take a fully-qualified resource name don't need it.
+
+    The quota project (billing / quota) is always taken from ADC's
+    `quota_project_id` and sent as `X-Goog-User-Project` on every request.
+    """
+
+    def __init__(self):
+        try:
+            credentials, _ = google.auth.default(scopes=DEFAULT_OAUTH_SCOPES)
+        except google.auth.exceptions.DefaultCredentialsError as e:
+            raise RuntimeError(
+                "No Application Default Credentials found. Run "
+                "'gcloud auth application-default login' first."
+            ) from e
+        if not credentials.quota_project_id:
+            raise RuntimeError(
+                "No quota project set. Run "
+                "'gcloud auth application-default set-quota-project <PROJECT_ID>'."
+            )
+        self._session = auth_requests.AuthorizedSession(credentials)
+        self._session.headers["X-Goog-User-Project"] = credentials.quota_project_id
+
+    def ensure_context_set_group(self, project_id: str, csg_id: str) -> str:
+        """Return a CSG's full resource name, creating it if absent.
+
+        Idempotent: a 409 (already exists, whether pre-existing or from a
+        concurrent create) is treated as success.
+        """
+        parent = f"projects/{project_id}/locations/{CONTEXT_STORE_LOCATION}"
+        create_url = (
+            f"{CONTEXT_STORE_ENDPOINT}/{API_VERSION}/{parent}/contextSetGroups"
+            f"?context_set_group_id={csg_id}"
+        )
+        try:
+            self._request_lro("POST", create_url, json_body={})
+        except requests.HTTPError as e:
+            if e.response.status_code != 409:
+                raise
+        return f"{parent}/contextSetGroups/{csg_id}"
+
+    def delete_context_set_group(self, csg_resource_name: str) -> None:
+        """Delete a ContextSetGroup by full resource name. Blocks on LRO.
+
+        Idempotent: a 404 (already deleted or never existed) is treated as
+        success. Cascades: all ContextSets inside the group are also deleted.
+        """
+        try:
+            self._request_lro(
+                "DELETE", f"{CONTEXT_STORE_ENDPOINT}/{API_VERSION}/{csg_resource_name}"
+            )
+        except requests.HTTPError as e:
+            if e.response.status_code != 404:
+                raise
+
+    def ensure_context_set(
+        self, project_id: str, csg_id: str, cs_id: str, version: str
+    ) -> str:
+        """Return a CS's full resource name, creating it (and the parent CSG)
+        if absent.
+
+        Idempotent: a 409 on the CS create is treated as success. Body is not
+        touched — the caller still needs `upload_context_set` to overwrite it.
+        """
+        csg_resource_name = self.ensure_context_set_group(project_id, csg_id)
+        url = (
+            f"{CONTEXT_STORE_ENDPOINT}/{API_VERSION}/{csg_resource_name}/contextSets"
+            f"?context_set_id={cs_id}"
+        )
+        body = {"context_set_id": cs_id, "version": version}
+        try:
+            self._request("POST", url, json_body=body)
+        except requests.HTTPError as e:
+            if e.response.status_code != 409:
+                raise
+        return f"{csg_resource_name}/contextSets/{cs_id}@{version}"
+
+    def upload_context_set(
+        self, cs_resource_name: str, ctx: context.ContextSet
+    ) -> None:
+        """Populate an existing ContextSet's body.
+
+        The resource must already exist (call `ensure_context_set` first).
+        """
+        url = f"{CONTEXT_STORE_ENDPOINT}/{API_VERSION}/{cs_resource_name}:upload"
+        body = {"context_json": ctx.model_dump_json(exclude_none=True)}
+        self._request("POST", url, json_body=body)
+
+    def download_context_set(self, cs_resource_name: str) -> context.ContextSet:
+        """Download and parse a ContextSet by full resource name."""
+        url = f"{CONTEXT_STORE_ENDPOINT}/{API_VERSION}/{cs_resource_name}:download"
+        resp = self._request("POST", url, json_body={})
+        raw = resp.get("contextJson")
+        if raw is None:
+            raise ValueError(f"DownloadContextSet response missing contextJson: {resp}")
+        # ContextSet models accept camelCase aliases too (see
+        # `_BaseContextModel`), so the server's mixed casing validates
+        # without conversion.
+        return context.ContextSet.model_validate(json.loads(raw))
+
+    def _request(
+        self,
+        method: str,
+        url: str,
+        json_body: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Send one HTTP request and return parsed JSON.
+
+        Returns an empty dict if the response body is empty (eg. 204 No
+        Content). Raises `requests.HTTPError` on 4xx/5xx,
+        `requests.RequestException` on network failure.
+        """
+        response = self._session.request(
+            method, url, json=json_body, timeout=_REQUEST_TIMEOUT_SECONDS
+        )
+        response.raise_for_status()
+        return response.json() if response.content else {}
+
+    def _request_lro(
+        self,
+        method: str,
+        url: str,
+        json_body: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Issue a request that returns an LRO and poll until it completes.
+
+        - HTTP layer (`_request`): each call returns 200; raises on 4xx/5xx.
+        - Op layer: `done` / `response` / `error` live in the JSON body, not
+          the HTTP status.
+        - Sleeps per `_LRO_POLL_INTERVALS_SECONDS`; raises `TimeoutError`
+          on budget exhaustion or `RuntimeError` on op-reported error.
+        """
+        op = self._request(method, url, json_body=json_body)
+        op_name = op.get("name")
+        if not op_name:
+            raise RuntimeError(f"{method} {url} returned no operation name: {op}")
+
+        # Skip the poll cycle if the initial op is already done (some LROs
+        # complete synchronously). Otherwise poll until it does or the
+        # budget is exhausted.
+        poll_url = f"{CONTEXT_STORE_ENDPOINT}/{API_VERSION}/{op_name}"
+        for delay in _LRO_POLL_INTERVALS_SECONDS:
+            if op.get("done"):
+                break
+            time.sleep(delay)
+            op = self._request("GET", poll_url)
+        if not op.get("done"):
+            raise TimeoutError(f"LRO {op_name} did not complete within the poll budget")
+        if "error" in op:
+            err = op["error"]
+            raise RuntimeError(
+                f"LRO {op_name} failed: [{err.get('code')}] "
+                f"{err.get('message', 'unknown LRO error')}"
+            )
+        return op.get("response", {})

{google_cloud_db_context_engineering-0.6.0 → google_cloud_db_context_engineering-0.7.0}/src/google/cloud/db_context_enrichment/evaluate/evaluate_generator.py

@@ -23,7 +23,7 @@ GOLDEN_QUERIES_NAME = "golden_queries.json"
 
 
 def generate_evalbench_configs(
-    experiment_name: str,
+    output_dir: str,
     dataset_path: str,
     context_set_id: str,
     toolbox_config_path: str,
@@ -32,6 +32,10 @@ def generate_evalbench_configs(
     """
     Main entrypoint: Generates Evalbench-compatible YAML configurations natively using
     private DB format converters and the google-cloud-geminidataanalytics API validations.
+
+    All output files land under `<output_dir>/eval_configs/`. The generated
+    `run_config.yaml` also points evalbench at `<output_dir>/eval_reports/`
+    for its results.
     """
     params = _extract_toolbox_params(toolbox_config_path, toolbox_source_name)
     generator = _get_db_generator(params)
@@ -39,13 +43,13 @@ def generate_evalbench_configs(
     db_config_yaml = generator.generate_db_config()
     model_config_yaml = generator.generate_model_config(context_set_id)
     llmrater_config_yaml = _generate_llmrater_config(params.get("project"))
-    run_config_yaml = _generate_run_config(experiment_name, generator.DIALECT)
+    run_config_yaml = _generate_run_config(output_dir, generator.DIALECT)
 
     # Convert simplified dataset to EvalBench standard format
     golden_queries_json = _convert_dataset(dataset_path, generator.DIALECT)
 
     # Write all files directly
-    eval_configs_dir = f"autoctx/experiments/{experiment_name}/eval_configs"
+    eval_configs_dir = os.path.join(output_dir, "eval_configs")
     os.makedirs(eval_configs_dir, exist_ok=True)
 
     with open(os.path.join(eval_configs_dir, DB_CONFIG_NAME), "w") as f:
@@ -140,10 +144,17 @@ def _get_db_generator(params: dict[str, Any]) -> BaseDBConfigGenerator:
     return generators[source_type](params)
 
 
-def _generate_run_config(experiment_name: str, dialect: str) -> str:
-    """Generates the main EvalBench Run Experiment scaffolding."""
-    configs_dir = f"autoctx/experiments/{experiment_name}/eval_configs"
-    reports_dir = f"autoctx/experiments/{experiment_name}/eval_reports"
+def _generate_run_config(output_dir: str, dialect: str) -> str:
+    """Generates the main EvalBench Run Experiment scaffolding.
+
+    Path values embedded in the YAML are normalized to POSIX separators so
+    the generated file is consistent across platforms (avoids mixed
+    `C:\\out\\eval_configs/db_config.yaml` on Windows and keeps string
+    assertions in tests portable).
+    """
+    output_dir_posix = output_dir.replace(os.sep, "/")
+    configs_dir = f"{output_dir_posix}/eval_configs"
+    reports_dir = f"{output_dir_posix}/eval_reports"
 
     return textwrap.dedent(f"""\
         ############################################################

{google_cloud_db_context_engineering-0.6.0 → google_cloud_db_context_engineering-0.7.0}/src/google/cloud/db_context_enrichment/main.py

@@ -1,13 +1,18 @@
 import json
+import pathlib
 
 from fastmcp import FastMCP
 
-from google.cloud.db_context_enrichment.common import context_mutator
+from google.cloud.db_context_enrichment.common import (
+    context_mutator,
+    context_store_client,
+)
 from google.cloud.db_context_enrichment.dataset import dataset_generator
 from google.cloud.db_context_enrichment.evaluate import (
     evaluate_generator,
     result_reader,
 )
+from google.cloud.db_context_enrichment.model import context
 
 mcp = FastMCP("Context Engineering Agent MCP")
 
@@ -18,7 +23,15 @@ async def generate_dataset(
     output_file_path: str,
 ) -> str:
     """
-    Validates a list of evaluation dataset entries and saves them to a JSON file.
+    Validates a list of evaluation dataset entries and saves them to a JSON file. This is the REQUIRED tool for generating 'golden' datasets.
+
+    CRITICAL: Any request to generate an evaluation dataset MUST follow the "context-engineering-workflow" skill.
+    Do NOT use generic file-writing tools to save the golden dataset. Using this tool is the final step of the dataset generation workflow.
+
+    The workflow REQUIRES the following steps BEFORE calling this tool:
+    1. Activate the `context-engineering-workflow` skill.
+    2. Read the `context-engineering-dataset-generation` skill's `SKILL.md` for the full procedure.
+    3. Generate mandatory audit reports (evalset_environment_inputs.md, evalset_gen_plan.md, evalset_report_pair_level.md, evalset_report_dataset_level.md) as specified in the workflow. These files are required for the verification process to pass.
 
     Args:
         dataset_entries_json: A JSON string representing a list of dataset items.
@@ -36,7 +49,7 @@ async def generate_dataset(
 
 @mcp.tool
 def generate_evalbench_configs(
-    experiment_name: str,
+    output_dir: str,
     dataset_path: str,
     context_set_id: str,
     toolbox_config_path: str,
@@ -45,17 +58,20 @@ def generate_evalbench_configs(
     """
     Generates Evalbench YAML configurations and converts the user-facing golden dataset to be compatible for evaluation, saving all files directly to disk.
 
-    This tool writes the following files inside `experiments/<experiment_name>/eval_configs/`:
+    This tool writes the following files inside `<output_dir>/eval_configs/`:
     - `db_config.yaml`
     - `model_config.yaml`
     - `run_config.yaml`
     - `llmrater_config.yaml`
     - `golden_queries.json` (converted to EvalBench internal format)
 
+    The generated `run_config.yaml` also points evalbench at
+    `<output_dir>/eval_reports/` for its results.
+
     Args:
-        experiment_name: The name of the target experiment folder.
+        output_dir: Directory (absolute or workspace-relative) where the eval configs and reports should live. Created if missing.
         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: The specific context_set_id inside the experiment.
+        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).
 
@@ -63,13 +79,13 @@ def generate_evalbench_configs(
         A message indicating that the configuration files were successfully created.
     """
     evaluate_generator.generate_evalbench_configs(
-        experiment_name,
+        output_dir,
         dataset_path,
         context_set_id,
         toolbox_config_path,
         toolbox_source_name,
     )
-    return f"Successfully generated all configs for evaluation in experiments/{experiment_name}/eval_configs/"
+    return f"Successfully generated all configs for evaluation in {output_dir}/eval_configs/"
 
 
 @mcp.tool
@@ -117,6 +133,72 @@ def generate_upload_url(
         return "Error: Invalid db_engine. Must be one of 'alloydb', 'cloudsql', or 'spanner'."
 
 
+# NOTE: `@mcp.tool` is intentionally NOT applied to upload_context_set /
+# download_context_set. The Context Store client library ships in this
+# release, but the MCP tool wrappers are held back until the Context Store
+# API is stable in production. Re-add the decorator to expose these to
+# agents when ready.
+def upload_context_set(
+    local_file_path: str,
+    project_id: str,
+    csg_id: str,
+    cs_id: str,
+    version: str,
+) -> str:
+    """
+    Upload a local ContextSet JSON file to the Context Store.
+
+    Resource hierarchy: a ContextSetGroup (CSG) is a logical container that
+    holds versioned ContextSets — typically one CSG per experiment, one
+    cs_id per lineage (eg. "autoctx"), and versions like "v0", "v1", "v2".
+
+    The CSG and the (cs_id, version) ContextSet resource are created if they
+    don't already exist; then the file contents are written as the
+    ContextSet body. Re-uploading the same (csg_id, cs_id, version)
+    overwrites the body.
+
+    Args:
+        local_file_path: Absolute path to a ContextSet JSON file.
+        project_id: GCP project where the CSG / CS should live. Typically
+            the same project the target DB lives in.
+        csg_id: ContextSetGroup ID (eg. an experiment name).
+        cs_id: ContextSet ID (eg. "autoctx"). Stable across versions.
+        version: Version label (eg. "baseline", "v1").
+
+    Returns:
+        Full ContextSet resource name, eg.
+        `projects/<p>/locations/<l>/contextSetGroups/<csg_id>/contextSets/<cs_id>@<version>`.
+    """
+    text = pathlib.Path(local_file_path).read_text()
+    ctx = context.ContextSet.model_validate_json(text)
+    client = context_store_client.ContextStoreClient()
+    cs_resource_name = client.ensure_context_set(project_id, csg_id, cs_id, version)
+    client.upload_context_set(cs_resource_name, ctx)
+    return cs_resource_name
+
+
+def download_context_set(cs_resource_name: str, output_file_path: str) -> str:
+    """
+    Download a ContextSet from the Context Store and write it to a local
+    JSON file. Parent directories are created as needed; the output file
+    is overwritten if it exists.
+
+    Args:
+        cs_resource_name: Full ContextSet resource name, eg.
+            `projects/<p>/locations/<l>/contextSetGroups/<csg_id>/contextSets/<cs_id>@<version>`.
+        output_file_path: Absolute path where the JSON file should be written.
+
+    Returns:
+        The output file path.
+    """
+    client = context_store_client.ContextStoreClient()
+    ctx = client.download_context_set(cs_resource_name)
+    out = pathlib.Path(output_file_path)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(ctx.model_dump_json(exclude_none=True, indent=2))
+    return output_file_path
+
+
 @mcp.tool
 def mutate_context_set(
     file_path: str,