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

google_cloud_db_context_engineering-0.7.1/PKG-INFO

@@ -0,0 +1,122 @@
+Metadata-Version: 2.4
+Name: google-cloud-db-context-engineering
+Version: 0.7.1
+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**.
+
+---
+
+## Installing via a compatible Agent Plugins client
+
+This repository ships a valid [Agent Plugins](https://github.com/agentplugins/agent-plugins-spec) (v1) plugin under [`plugin/`](plugin/), bundling both agent skills and its MCP servers. Any [Agent Plugins–compatible client](https://agent-plugins.org/compatible-clients) (VS Code, Cursor, GitHub Copilot, Kiro, and others) can install it directly using its own built-in plugin command — no extra tooling required — by pointing at this repository:
+
+```
+https://github.com/GoogleCloudPlatform/db-context-enrichment
+```
+
+See your agent's documentation for its exact install command.
+
+---
+
+## 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.1/README.md

@@ -0,0 +1,104 @@
+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**.
+
+---
+
+## Installing via a compatible Agent Plugins client
+
+This repository ships a valid [Agent Plugins](https://github.com/agentplugins/agent-plugins-spec) (v1) plugin under [`plugin/`](plugin/), bundling both agent skills and its MCP servers. Any [Agent Plugins–compatible client](https://agent-plugins.org/compatible-clients) (VS Code, Cursor, GitHub Copilot, Kiro, and others) can install it directly using its own built-in plugin command — no extra tooling required — by pointing at this repository:
+
+```
+https://github.com/GoogleCloudPlatform/db-context-enrichment
+```
+
+See your agent's documentation for its exact install command.
+
+---
+
+## 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.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "google-cloud-db-context-engineering"
-version = "0.6.0"
+version = "0.7.1"
 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]
@@ -44,7 +46,7 @@ dev = [
 
 [tool.db-context-engineering]
 toolbox_version = "1.4.0"
-evalbench_version = "1.9.0"
+evalbench_version = "1.12.0"
 
 [tool.ruff]
 line-length = 88

google_cloud_db_context_engineering-0.7.1/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.1}/src/google/cloud/db_context_enrichment/evaluate/db_generators/alloydb.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import google.cloud.geminidataanalytics_v1beta as gda
 import yaml
 
 from .base import BaseDBConfigGenerator
@@ -52,21 +51,20 @@ class AlloyDBConfigGenerator(BaseDBConfigGenerator):
             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()
-
-        datasource_ref.alloydb = gda.AlloyDbReference(
-            database_reference=gda.AlloyDbDatabaseReference(
-                project_id=self.project,
-                region=self.region,
-                cluster_id=self.cluster,
-                instance_id=self.instance,
-                database_id=self.database,
-            ),
-            agent_context_reference=gda.AgentContextReference(
-                context_set_id=context_set_id
-            ),
-        )
-        return datasource_ref
+    def build_datasource_reference(self, context_set_id: str) -> dict[str, Any]:
+        ref: dict[str, Any] = {
+            "alloydb": {
+                "database_reference": {
+                    "project_id": self.project,
+                    "region": self.region,
+                    "cluster_id": self.cluster,
+                    "instance_id": self.instance,
+                    "database_id": self.database,
+                }
+            }
+        }
+        if context_set_id:
+            ref["alloydb"]["agent_context_reference"] = {
+                "context_set_id": context_set_id
+            }
+        return ref

{google_cloud_db_context_engineering-0.6.0 → google_cloud_db_context_engineering-0.7.1}/src/google/cloud/db_context_enrichment/evaluate/db_generators/base.py

@@ -1,16 +1,14 @@
 from abc import ABC, abstractmethod
 from typing import Any
 
-import google.cloud.geminidataanalytics_v1beta as gda
 import yaml
-from google.protobuf.json_format import MessageToDict
 
 
 class BaseDBConfigGenerator(ABC):
     """
     Abstract Base Class enforcing the construction contract for Evalbench database topologies.
     Each distinct DB type (Spanner, Postgres, AlloyDB, MySQL) must inherit and implement
-    the mappings required by both the standard EvalBench framework and the GDA SDK model.
+    the mappings required by both the standard EvalBench framework and the GDA REST API.
     """
 
     SOURCE_TYPE = "unknown"
@@ -29,11 +27,9 @@ class BaseDBConfigGenerator(ABC):
         raise NotImplementedError("Subclasses must implement generate_db_config")
 
     @abstractmethod
-    def build_datasource_reference(
-        self, context_set_id: str
-    ) -> gda.DatasourceReferences:
+    def build_datasource_reference(self, context_set_id: str) -> dict[str, Any]:
         """
-        Constructs the strict Protocol Buffer DatasourceReference required by the QueryDataAPI
+        Constructs the REST datasource_references dictionary required by the QueryDataAPI
         context generation flow.
         """
         raise NotImplementedError(
@@ -54,23 +50,21 @@ class BaseDBConfigGenerator(ABC):
 
     def generate_model_config(self, context_set_id: str) -> str:
         """
-        Standardized Model Builder converting the strictly typed GDA object into an EvalBench model dict.
+        Standardized Model Builder converting the datasource reference dict into an EvalBench model dict.
         """
         datasource_ref = self.build_datasource_reference(context_set_id)
 
-        query_context = gda.QueryDataContext(datasource_references=datasource_ref)
-
-        query_context_dict = MessageToDict(
-            query_context._pb, preserving_proto_field_name=True
-        )
-
         model_config = {
             "generator": "query_data_api",
             "project_id": self.params.get("project"),
             "location": self.params.get("region") or "global",
-            "context": query_context_dict,
+            "use_rest_api": True,
+            "context": {"datasource_references": datasource_ref},
         }
 
+        if self.params.get("api_endpoint"):
+            model_config["api_endpoint"] = self.params.get("api_endpoint")
+
         return yaml.safe_dump(
             model_config, sort_keys=False, default_flow_style=False
         ).strip()

{google_cloud_db_context_engineering-0.6.0 → google_cloud_db_context_engineering-0.7.1}/src/google/cloud/db_context_enrichment/evaluate/db_generators/mysql.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import google.cloud.geminidataanalytics_v1beta as gda
 import yaml
 
 from .base import BaseDBConfigGenerator
@@ -49,21 +48,20 @@ class MySQLConfigGenerator(BaseDBConfigGenerator):
             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()
-
-        datasource_ref.cloud_sql_reference = gda.CloudSqlReference(
-            database_reference=gda.CloudSqlDatabaseReference(
-                engine=gda.CloudSqlDatabaseReference.Engine.MYSQL,
-                project_id=self.project,
-                region=self.region,
-                instance_id=self.instance,
-                database_id=self.database,
-            ),
-            agent_context_reference=gda.AgentContextReference(
-                context_set_id=context_set_id
-            ),
-        )
-        return datasource_ref
+    def build_datasource_reference(self, context_set_id: str) -> dict[str, Any]:
+        ref: dict[str, Any] = {
+            "cloud_sql_reference": {
+                "database_reference": {
+                    "engine": "MYSQL",
+                    "project_id": self.project,
+                    "region": self.region,
+                    "instance_id": self.instance,
+                    "database_id": self.database,
+                }
+            }
+        }
+        if context_set_id:
+            ref["cloud_sql_reference"]["agent_context_reference"] = {
+                "context_set_id": context_set_id
+            }
+        return ref

{google_cloud_db_context_engineering-0.6.0 → google_cloud_db_context_engineering-0.7.1}/src/google/cloud/db_context_enrichment/evaluate/db_generators/postgres.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-import google.cloud.geminidataanalytics_v1beta as gda
 import yaml
 
 from .base import BaseDBConfigGenerator
@@ -49,21 +48,20 @@ class PostgresConfigGenerator(BaseDBConfigGenerator):
             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()
-
-        datasource_ref.cloud_sql_reference = gda.CloudSqlReference(
-            database_reference=gda.CloudSqlDatabaseReference(
-                engine=gda.CloudSqlDatabaseReference.Engine.POSTGRESQL,
-                project_id=self.project,
-                region=self.region,
-                instance_id=self.instance,
-                database_id=self.database,
-            ),
-            agent_context_reference=gda.AgentContextReference(
-                context_set_id=context_set_id
-            ),
-        )
-        return datasource_ref
+    def build_datasource_reference(self, context_set_id: str) -> dict[str, Any]:
+        ref: dict[str, Any] = {
+            "cloud_sql_reference": {
+                "database_reference": {
+                    "engine": "POSTGRESQL",
+                    "project_id": self.project,
+                    "region": self.region,
+                    "instance_id": self.instance,
+                    "database_id": self.database,
+                }
+            }
+        }
+        if context_set_id:
+            ref["cloud_sql_reference"]["agent_context_reference"] = {
+                "context_set_id": context_set_id
+            }
+        return ref