google-cloud-db-context-engineering 0.5.1__py3-none-any.whl

google/cloud/db_context_enrichment/__init__.py

@@ -0,0 +1 @@
+# db_context_enrichment package

google/cloud/db_context_enrichment/bootstrap/__init__.py

@@ -0,0 +1 @@
+# Initialization script for bootstrap module

google/cloud/db_context_enrichment/bootstrap/bootstrap_generator.py

@@ -0,0 +1,74 @@
+import json
+
+from pydantic import ValidationError
+
+from google.cloud.db_context_enrichment.common.context_mutator import (
+    Mutation,
+    mutate_context_set,
+)
+from google.cloud.db_context_enrichment.facet import facet_generator
+from google.cloud.db_context_enrichment.model import context
+from google.cloud.db_context_enrichment.template import template_generator
+
+
+async def generate_context(
+    output_file_path: str,
+    sql_dialect: str,
+    template_inputs_json: str | None = None,
+    facet_inputs_json: str | None = None,
+) -> str:
+    """
+    Core logic for generating a single unified ContextSet from key information and saving it to a file.
+    """
+    final_templates = None
+    final_facets = None
+
+    if template_inputs_json:
+        res_str = await template_generator.generate_templates(
+            template_inputs_json, sql_dialect
+        )
+        if '"error":' in res_str:
+            raise RuntimeError(f"Error generating templates: {res_str}")
+        try:
+            res_dict = json.loads(res_str)
+            final_templates = [
+                context.Template(**t) for t in res_dict.get("templates", [])
+            ]
+        except (json.JSONDecodeError, ValidationError) as e:
+            raise ValueError(f"Error parsing generated templates: {e}") from e
+
+    if facet_inputs_json:
+        res_str = await facet_generator.generate_facets(facet_inputs_json, sql_dialect)
+        if '"error":' in res_str:
+            raise RuntimeError(f"Error generating facets: {res_str}")
+        try:
+            res_dict = json.loads(res_str)
+            final_facets = [context.Facet(**f) for f in res_dict.get("facets", [])]
+        except (json.JSONDecodeError, ValidationError) as e:
+            raise ValueError(f"Error parsing generated facets: {e}") from e
+
+    mutations: list[Mutation] = []
+
+    if final_templates:
+        for t in final_templates:
+            mutations.append(
+                Mutation(
+                    operation="add",
+                    type="template",
+                    value=t.model_dump(exclude_none=True),
+                )
+            )
+
+    if final_facets:
+        for f in final_facets:
+            mutations.append(
+                Mutation(
+                    operation="add", type="facet", value=f.model_dump(exclude_none=True)
+                )
+            )
+
+    if not mutations:
+        raise ValueError("No templates or facets were generated to save.")
+
+    mutate_context_set(output_file_path, mutations)
+    return output_file_path

google/cloud/db_context_enrichment/common/config.py

@@ -0,0 +1,10 @@
+MODEL_NAME = "gemini-3.1-flash-lite"
+
+
+def get_model_name() -> str:
+    """Returns the configured model name.
+
+    Centralized in this file to make it easy to update.
+    Requires rebuilding the binary to take effect.
+    """
+    return MODEL_NAME

google/cloud/db_context_enrichment/common/context_mutator.py

@@ -0,0 +1,132 @@
+import json
+import os
+from typing import Any, Literal
+
+from pydantic import BaseModel, ValidationError
+
+from google.cloud.db_context_enrichment.model import context
+
+
+class Mutation(BaseModel):
+    """Defines a strict schema for context modifications."""
+
+    operation: Literal["add", "delete", "update"]
+    type: Literal["template", "facet", "value_search"]
+    identifier: dict[str, Any] = {}
+    value: dict[str, Any] | None = None
+
+
+def mutate_context_set(file_path: str, mutations: list[Mutation]) -> None:
+    """
+    Internal function to mutate (add, delete, update) elements in an existing ContextSet JSON file.
+
+    The mutations is expected to be a list of Mutation objects.
+    Each object specifies an 'operation', 'type', 'identifier' (to route to the correct item for deletes/updates), and 'value' (for adding/updating), for instance:
+
+    [
+      {
+        "operation": "add",
+        "type": "template",
+        "value": {"nl_query": "...", "sql": "...", "intent": "...", "manifest": "...", "parameterized": {...}}
+      },
+      {
+        "operation": "delete",
+        "type": "template",
+        "identifier": {"nl_query": "What are all users?"}
+      },
+      {
+        "operation": "update",
+        "type": "facet",
+        "identifier": {"intent": "high price"},
+        "value": {"sql_snippet": "price > 2000", "intent": "very high price"}
+      }
+    ]
+    """
+
+    # 1. Load exiting ContextSet (or create an empty one)
+    if not os.path.exists(file_path) or os.path.getsize(file_path) == 0:
+        context_set = context.ContextSet()
+    else:
+        try:
+            with open(file_path) as f:
+                raw_data = json.load(f)
+            context_set = context.ContextSet.model_validate(raw_data)
+        except ValidationError as e:
+            raise ValueError(
+                f"Validation Error loading ContextSet from {file_path}: {e}"
+            ) from e
+        except (json.JSONDecodeError, OSError) as e:
+            raise RuntimeError(f"Error reading JSON from {file_path}: {e}") from e
+
+    # 2. Model mapping for tracking and validation
+    type_to_model = {
+        "template": context.Template,
+        "facet": context.Facet,
+        "value_search": context.ValueSearch,
+    }
+
+    type_to_attr = {
+        "template": "templates",
+        "facet": "facets",
+        "value_search": "value_searches",
+    }
+
+    # 3. Apply mutations
+    for i, mut in enumerate(mutations):
+        op = mut.operation
+        item_type = mut.type
+        identifier = mut.identifier
+        value_data = mut.value
+
+        if item_type not in type_to_attr:
+            continue
+
+        attr_name = type_to_attr[item_type]
+        model_class = type_to_model[item_type]
+
+        target_list = getattr(context_set, attr_name)
+        if target_list is None:
+            target_list = []
+            setattr(context_set, attr_name, target_list)
+
+        if op == "add":
+            if value_data:
+                try:
+                    new_item = model_class.model_validate(value_data)
+                    target_list.append(new_item)
+                except ValidationError as e:
+                    raise ValueError(
+                        f"Validation Error on mutation {i} during 'add': {e}"
+                    ) from e
+
+        elif op == "delete":
+            new_list = []
+            for item in target_list:
+                item_dict = item.model_dump()
+                match = all(item_dict.get(k) == v for k, v in identifier.items())
+                if not match:
+                    new_list.append(item)
+            setattr(context_set, attr_name, new_list)
+
+        elif op == "update":
+            for idx, item in enumerate(target_list):
+                item_dict = item.model_dump()
+                match = all(item_dict.get(k) == v for k, v in identifier.items())
+                if match and value_data:
+                    updated_dict = {**item_dict, **value_data}
+                    try:
+                        updated_item = model_class.model_validate(updated_dict)
+                        target_list[idx] = updated_item
+                        break  # Only update first match
+                    except ValidationError as e:
+                        raise ValueError(
+                            f"Validation Error on mutation {i} during 'update': {e}"
+                        ) from e
+
+    # 4. Save validated ContextSet
+    try:
+        os.makedirs(os.path.dirname(os.path.abspath(file_path)), exist_ok=True)
+        with open(file_path, "w") as f:
+            f.write(context_set.model_dump_json(indent=2, exclude_none=True))
+    except OSError as e:
+        raise RuntimeError(f"Error saving ContextSet to {file_path}: {e}") from e

google/cloud/db_context_enrichment/common/parameterizer.py

@@ -0,0 +1,192 @@
+import re
+import textwrap
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from google import genai
+from google.cloud.db_context_enrichment.common import config
+
+
+class SQLDialect(Enum):
+    """Enumeration for supported database dialects."""
+
+    POSTGRESQL = "postgresql"
+    MYSQL = "mysql"
+    GOOGLESQL = "googlesql"
+
+
+class ValuePhrasePair(BaseModel):
+    """A key-value pair for a named entity and its types."""
+
+    key: str = Field(..., description="The extracted named entity.")
+    value: list[str] = Field(
+        ..., description="A list of identified types for the entity."
+    )
+
+
+class ValuePhrasesList(BaseModel):
+    """A list of named entity pairs."""
+
+    value_phrases: list[ValuePhrasePair] = Field(
+        ...,
+        description="A list of key-value pairs, where each key is a named entity and the value is a list of its types.",
+    )
+
+
+async def extract_value_phrases(nl_query: str) -> dict[str, list[str]]:
+    """
+    Extracts potential value phrases from a natural language question using an LLM.
+
+    This function replicates the core logic of the `value_phrases_extractor`
+    and `get_value_phrases_extractor_template` functions in `choose.sql`.
+    It builds a prompt to perform named entity recognition (NER) and calls a
+    generative model to extract entities based on a predefined list of types.
+
+    Args:
+        nl_question: The natural language question to analyze.
+
+    Returns:
+        A dictionary containing the extracted phrases and their types.
+
+    Raises:
+        Exception: If the model call fails or returns an invalid response.
+    """
+    prompt = textwrap.dedent(
+        f"""
+        Please extract the named entity (a real-world object, such as a person,
+        location, organization, product, etc., that can be denoted with a proper name)
+        from the query literally based on the following types:
+
+        [Types]
+        - country
+        - city
+        - email_address
+        - language
+        - law
+        - organization
+        - person
+        - product
+        - sport or activity
+        - work of art
+        - date
+        - time
+        - number
+        - currency
+        - region
+
+        The output should be a JSON object containing a list of key-value pairs.
+        Each pair should have a "key" (the extracted named entity) and a "value" (a list of its identified types).
+        For example: {{"value_phrases": [{{"key": "entity1", "value": ["type1"]}}, {{"key": "entity2", "value": ["type2", "type3"]}}]}}
+        DO NOT perform any spell checking or correction. If no entities are
+        identified, return an empty list.
+
+        [Query]
+        {nl_query}
+    """
+    )
+
+    client = genai.Client()
+    try:
+        response = await client.aio.models.generate_content(
+            model=config.get_model_name(),
+            contents=prompt,
+            config={
+                "response_mime_type": "application/json",
+                "response_schema": ValuePhrasesList,
+            },
+        )
+        if response.text:
+            phrases_obj = ValuePhrasesList.model_validate_json(response.text)
+            # Convert the list of pairs back to a dictionary
+            return {pair.key: pair.value for pair in phrases_obj.value_phrases}
+        else:
+            # Return an empty dict if the model returns no text
+            return {}
+    except Exception as e:
+        # Re-raise the exception to be handled by the caller
+        raise Exception(f"An error occurred during value phrase extraction: {e}") from e
+    finally:
+        client.close()
+        await client.aio.aclose()
+
+
+def parameterize_sql_and_intent(
+    value_phrases: dict[str, Any],
+    sql: str,
+    intent: str,
+    db_dialect: SQLDialect = SQLDialect.POSTGRESQL,
+) -> dict[str, str]:
+    """
+    Replaces value phrases in a SQL query and an intent string with placeholders.
+
+    This function iterates through a dictionary of value phrases and replaces
+    their occurrences in both the SQL and intent strings with positional
+    parameters. The syntax of the parameters is determined by the specified
+    database dialect.
+
+    The phrases are processed in descending order of length to handle nested
+    phrases correctly (e.g., "New York" before "York").
+
+    The replacement logic handles both quoted and unquoted occurrences of the
+    phrases, and it avoids replacing phrases that are already part of a
+    placeholder.
+
+    Args:
+        value_phrases: A dictionary where keys are the string phrases to be
+                       replaced. Values are not used.
+        sql: The SQL query string to parameterize.
+        intent: The natural language intent string to parameterize.
+        db_dialect: The SQL dialect to use for parameterization.
+
+    Returns:
+        A dictionary containing the parameterized 'sql' and 'intent' strings.
+    """
+    psql = sql
+    pintent = intent
+    param_index = 1
+
+    # Sort keys by length in descending order to prioritize longer matches
+    sorted_phrases = sorted(value_phrases.keys(), key=len, reverse=True)
+
+    for value in sorted_phrases:
+        # Determine the placeholder based on the database dialect
+        if db_dialect == SQLDialect.POSTGRESQL:
+            placeholder = f"${param_index}"
+        else:  # For mysql, googlesql, etc.
+            placeholder = "?"
+
+        search_phrase_quoted = f"'{value}'"
+        replaced = False
+
+        # Patterns with negative lookbehind to avoid replacing existing params
+        # e.g., don't replace 'foo' in "$'foo'"
+        quoted_pattern = re.compile(r"(?<!\$)" + re.escape(search_phrase_quoted))
+        unquoted_pattern = re.compile(r"(?<!\$)" + r"\b" + re.escape(value) + r"\b")
+
+        # Condition 1: Quoted in SQL, Quoted in Intent
+        if quoted_pattern.search(psql) and quoted_pattern.search(pintent):
+            psql = quoted_pattern.sub(placeholder, psql)
+            pintent = quoted_pattern.sub(placeholder, pintent)
+            replaced = True
+        # Condition 2: Quoted in SQL, Unquoted in Intent
+        elif quoted_pattern.search(psql) and unquoted_pattern.search(pintent):
+            psql = quoted_pattern.sub(placeholder, psql)
+            pintent = unquoted_pattern.sub(placeholder, pintent)
+            replaced = True
+        # Condition 3: Unquoted in SQL, Quoted in Intent
+        elif unquoted_pattern.search(psql) and quoted_pattern.search(pintent):
+            psql = unquoted_pattern.sub(placeholder, psql)
+            pintent = quoted_pattern.sub(placeholder, pintent)
+            replaced = True
+        # Condition 4: Unquoted in SQL, Unquoted in Intent
+        elif unquoted_pattern.search(psql) and unquoted_pattern.search(pintent):
+            psql = unquoted_pattern.sub(placeholder, psql)
+            pintent = unquoted_pattern.sub(placeholder, pintent)
+            replaced = True
+
+        if replaced:
+            param_index += 1
+
+    return {"sql": psql, "intent": pintent}

google/cloud/db_context_enrichment/dataset/dataset_generator.py

@@ -0,0 +1,44 @@
+import json
+import os
+
+
+async def generate_dataset(
+    dataset_entries_json: str,
+    output_file_path: str,
+) -> str:
+    """
+    Validates a list of evaluation dataset entries and saves them to a JSON file.
+
+    Args:
+        dataset_entries_json: A JSON string representing a list of dataset items.
+                             Each item should have "id", "database", "nlq", and "golden_sql" keys.
+                             Example: '[{"id": "eval_001", "database": "my_db", "nlq": "Count users", "golden_sql": "SELECT COUNT(*) FROM users"}]'
+        output_file_path: The absolute path where the dataset JSON file should be saved.
+
+    Returns:
+        The absolute file path where the dataset was saved.
+    """
+    try:
+        data = json.loads(dataset_entries_json)
+        if not isinstance(data, list):
+            raise ValueError("Dataset entries must be a list of objects.")
+
+        # Simple validation of keys
+        for i, entry in enumerate(data):
+            if not isinstance(entry, dict):
+                raise ValueError(f"Entry at index {i} is not an object.")
+            missing_keys = {"id", "database", "nlq", "golden_sql"} - set(entry.keys())
+            if missing_keys:
+                raise ValueError(
+                    f"Entry at index {i} is missing required keys: {missing_keys}"
+                )
+
+        # Ensure directory exists
+        os.makedirs(os.path.dirname(os.path.abspath(output_file_path)), exist_ok=True)
+
+        with open(output_file_path, "w") as f:
+            json.dump(data, f, indent=2)
+
+        return f"Successfully saved dataset to {output_file_path}"
+    except (json.JSONDecodeError, ValueError, OSError) as e:
+        return f"Error saving dataset: {str(e)}"

google/cloud/db_context_enrichment/evaluate/__init__.py

@@ -0,0 +1,3 @@
+"""
+Evaluation workflow integration and configuration generation for Evalbench.
+"""

google/cloud/db_context_enrichment/evaluate/db_generators/alloydb.py

@@ -0,0 +1,72 @@
+from typing import Any
+
+import google.cloud.geminidataanalytics_v1beta as gda
+import yaml
+
+from .base import BaseDBConfigGenerator
+
+
+class AlloyDBConfigGenerator(BaseDBConfigGenerator):
+    """
+    Dedicated generator mapping properties to explicit AlloyDB configuration
+    topologies utilized by both EvalBench binaries and GDA Context objects.
+    """
+
+    SOURCE_TYPE = "alloydb-postgres"
+    DIALECT = "postgres"
+    REQUIRED_FIELDS = BaseDBConfigGenerator.REQUIRED_FIELDS | {
+        "project",
+        "region",
+        "cluster",
+        "instance",
+        "database",
+    }
+
+    def __init__(self, params: dict[str, Any]):
+        super().__init__(params)
+        self.project = params.get("project")
+        self.region = params.get("region")
+        self.cluster = params.get("cluster")
+        self.instance = params.get("instance")
+        self.database = params.get("database")
+        self.user = params.get("user")
+        self.password = params.get("password")
+
+    def generate_db_config(self) -> str:
+        db_type = "alloydb"
+        db_path = f"projects/{self.project}/locations/{self.region}/clusters/{self.cluster}/instances/{self.instance}"
+
+        db_config = {
+            "db_type": db_type,
+            "dialect": self.DIALECT,
+            "database_name": self.database,
+            "database_path": db_path,
+            "max_executions_per_minute": 180,
+            "nl_config": "",  # Required by evalbench schema
+        }
+        if self.user:
+            db_config["user_name"] = self.user
+        if self.password:
+            db_config["password"] = self.password
+        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()
+
+        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

google/cloud/db_context_enrichment/evaluate/db_generators/base.py

@@ -0,0 +1,76 @@
+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.
+    """
+
+    SOURCE_TYPE = "unknown"
+    DIALECT = "unknown"
+    REQUIRED_FIELDS = {"project"}
+
+    def __init__(self, params: dict[str, Any]):
+        self.params = params
+        self.validate()
+
+    @abstractmethod
+    def generate_db_config(self) -> str:
+        """
+        Generates the Evalbench db_config.yaml payload natively.
+        """
+        raise NotImplementedError("Subclasses must implement generate_db_config")
+
+    @abstractmethod
+    def build_datasource_reference(
+        self, context_set_id: str
+    ) -> gda.DatasourceReferences:
+        """
+        Constructs the strict Protocol Buffer DatasourceReference required by the QueryDataAPI
+        context generation flow.
+        """
+        raise NotImplementedError(
+            "Subclasses must implement build_datasource_reference"
+        )
+
+    def validate(self) -> None:
+        """
+        Validates that the provided tools.yaml source configuration block contains all
+        the mandatory fields required by the specific Evalbench topology.
+        """
+        missing = [f for f in self.REQUIRED_FIELDS if f not in self.params]
+        if missing:
+            raise ValueError(
+                f"Missing required fields in tools.yaml config for '{self.SOURCE_TYPE}': "
+                f"{', '.join(missing)}"
+            )
+
+    def generate_model_config(self, context_set_id: str) -> str:
+        """
+        Standardized Model Builder converting the strictly typed GDA object 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,
+        }
+
+        return yaml.safe_dump(
+            model_config, sort_keys=False, default_flow_style=False
+        ).strip()