google-cloud-db-context-engineering 0.7.3__tar.gz → 0.7.4__tar.gz

{google_cloud_db_context_engineering-0.7.3/src/google_cloud_db_context_engineering.egg-info → google_cloud_db_context_engineering-0.7.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-cloud-db-context-engineering
-Version: 0.7.3
+Version: 0.7.4
 Summary: A FastMCP server for generating natural language to SQL templates from database schemas.
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown

{google_cloud_db_context_engineering-0.7.3 → google_cloud_db_context_engineering-0.7.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "google-cloud-db-context-engineering"
-version = "0.7.3"
+version = "0.7.4"
 description = "A FastMCP server for generating natural language to SQL templates from database schemas."
 readme = "README.md"
 requires-python = ">=3.12"
@@ -45,7 +45,7 @@ dev = [
 ]
 
 [tool.db-context-engineering]
-toolbox_version = "1.4.0"
+toolbox_version = "1.10.0"
 evalbench_version = "1.17.0"
 
 [tool.ruff]

google_cloud_db_context_engineering-0.7.4/src/google/cloud/db_context_enrichment/evaluate/db_generators/__init__.py

@@ -0,0 +1,15 @@
+from .alloydb import AlloyDBConfigGenerator
+from .base import BaseDBConfigGenerator
+from .custom import CustomDBConfigGenerator
+from .mysql import MySQLConfigGenerator
+from .postgres import PostgresConfigGenerator
+from .spanner import SpannerConfigGenerator
+
+__all__ = [
+    "BaseDBConfigGenerator",
+    "AlloyDBConfigGenerator",
+    "CustomDBConfigGenerator",
+    "MySQLConfigGenerator",
+    "PostgresConfigGenerator",
+    "SpannerConfigGenerator",
+]

google_cloud_db_context_engineering-0.7.4/src/google/cloud/db_context_enrichment/evaluate/db_generators/bigtable.py

@@ -0,0 +1,44 @@
+from typing import Any
+
+import yaml
+
+from .base import BaseDBConfigGenerator
+
+
+class BigtableConfigGenerator(BaseDBConfigGenerator):
+    """
+    Dedicated generator mapping properties to explicit Bigtable configuration
+    topologies. Bypasses the GDA Python SDK to construct the model configuration
+    directly, as the SDK may lack native BigtableReference definitions.
+    """
+
+    SOURCE_TYPE = "bigtable"
+    DIALECT = "bigtable"
+    REQUIRED_FIELDS = {"project", "instance"}
+
+    def generate_db_config(self) -> str:
+        db_config = {
+            "db_type": "bigtable",
+            "dialect": self.DIALECT,
+            "database_name": self.params.get("instance"),
+            "database_path": f"projects/{self.params.get('project')}/instances/{self.params.get('instance')}",
+            "instance_id": self.params.get("instance"),
+            "gcp_project_id": self.params.get("project"),
+            "max_executions_per_minute": 100,
+        }
+        return yaml.safe_dump(
+            db_config, sort_keys=False, default_flow_style=False
+        ).strip()
+
+    def build_datasource_reference(self, context_set_id: str) -> dict[str, Any]:
+        return {
+            "bigtable_reference": {
+                "database_reference": {
+                    "project_id": self.params.get("project"),
+                    "instance_id": self.params.get("instance"),
+                },
+                "agent_context_reference": {
+                    "context_set_id": context_set_id,
+                },
+            }
+        }

google_cloud_db_context_engineering-0.7.4/src/google/cloud/db_context_enrichment/evaluate/db_generators/custom.py

@@ -0,0 +1,104 @@
+import os
+from typing import Any
+
+import yaml
+
+from .base import BaseDBConfigGenerator
+
+
+class CustomDBConfigGenerator(BaseDBConfigGenerator):
+    """
+    Pluggable generator mapping properties to custom evaluation configurations.
+    Enables arbitrary third-party or proprietary internal engines to be plugged
+    into the evaluation framework via dynamic connector and generator SPIs
+    without exposing engine-specific code or configuration schemas.
+    """
+
+    SOURCE_TYPE = "custom"
+    DIALECT = "sql"
+    REQUIRED_FIELDS = set()
+
+    def __init__(self, params: dict[str, Any]):
+        self.params = params
+        if "project" not in self.params:
+            project_id = os.environ.get("GOOGLE_CLOUD_PROJECT") or os.environ.get(
+                "GCP_PROJECT"
+            )
+            if project_id:
+                self.params["project"] = project_id
+        self.connector_class = params.get("connector_class", "")
+        self.generator_class = params.get("generator_class", "")
+        self.DIALECT = params.get("dialect", self.DIALECT)
+        self.validate()
+
+    def validate(self) -> None:
+        """
+        Validates that either connector_class or generator_class is specified.
+        """
+        if not self.connector_class and not self.generator_class:
+            raise ValueError(
+                "Custom source configuration must specify at least 'connector_class' or 'generator_class'."
+            )
+
+    def generate_db_config(self) -> str:
+        """
+        Generates the db_config.yaml payload for custom connectors.
+        """
+        db_config: dict[str, Any] = {
+            "db_type": self.params.get("db_type", "custom"),
+            "dialect": self.DIALECT,
+        }
+        if self.connector_class:
+            db_config["connector_class"] = self.connector_class
+
+        # Forward all non-meta parameters from tools.yaml source block
+        excluded_keys = {
+            "kind",
+            "name",
+            "type",
+            "connector_class",
+            "generator_class",
+            "dialect",
+            "db_type",
+        }
+        for key, value in self.params.items():
+            if key not in excluded_keys:
+                db_config[key] = value
+
+        return yaml.safe_dump(
+            db_config, sort_keys=False, default_flow_style=False
+        ).strip()
+
+    def build_datasource_reference(self, context_set_id: str) -> dict[str, Any]:
+        """
+        Datasource reference dictionary. Custom engines manage their own
+        schema references or context binding.
+        """
+        return {}
+
+    def generate_model_config(self, context_set_id: str) -> str:
+        """
+        Generates model_config.yaml. If generator_class is specified, produces
+        a custom generator configuration for dynamic instantiation.
+        """
+        if self.generator_class:
+            model_config: dict[str, Any] = {
+                "generator": "custom",
+                "generator_class": self.generator_class,
+                "context_set_id": context_set_id,
+            }
+            excluded_keys = {
+                "kind",
+                "name",
+                "type",
+                "connector_class",
+                "generator_class",
+            }
+            for key, value in self.params.items():
+                if key not in excluded_keys:
+                    model_config[key] = value
+            return yaml.safe_dump(
+                model_config, sort_keys=False, default_flow_style=False
+            ).strip()
+
+        return super().generate_model_config(context_set_id)

google_cloud_db_context_engineering-0.7.4/src/google/cloud/db_context_enrichment/evaluate/db_generators/firestore.py

@@ -0,0 +1,60 @@
+from typing import Any
+
+import yaml
+
+from .base import BaseDBConfigGenerator
+
+
+class FirestoreConfigGenerator(BaseDBConfigGenerator):
+    """
+    Dedicated generator mapping properties to Firestore (MongoDB query dialect)
+    topologies utilized by both EvalBench binaries and GDA REST API.
+    """
+
+    SOURCE_TYPE = "firestore"
+    DIALECT = "mongodb"
+    REQUIRED_FIELDS = BaseDBConfigGenerator.REQUIRED_FIELDS | {
+        "project",
+        "database",
+    }
+
+    def __init__(self, params: dict[str, Any]):
+        super().__init__(params)
+        self.project = params.get("project")
+        self.database = params.get("database")
+        self.connection_string = params.get("connection_string")
+        self.collection_ids = params.get("collection_ids") or params.get("table_ids")
+
+    def generate_db_config(self) -> str:
+        db_type = "mongodb"
+
+        db_config = {
+            "db_type": db_type,
+            "dialect": self.DIALECT,
+            "database_name": self.database,
+            "database_path": "",
+            "firestore_database": f"projects/{self.project}/databases/{self.database}",
+            "max_executions_per_minute": 120,
+        }
+        if self.connection_string:
+            db_config["connection_string"] = self.connection_string
+
+        return yaml.safe_dump(
+            db_config, sort_keys=False, default_flow_style=False
+        ).strip()
+
+    def build_datasource_reference(self, context_set_id: str) -> dict[str, Any]:
+        db_ref: dict[str, Any] = {
+            "project_id": self.project,
+            "database_id": self.database,
+        }
+        if self.collection_ids:
+            db_ref["collection_ids"] = self.collection_ids
+
+        ref: dict[str, Any] = {"firestore_reference": {"database_reference": db_ref}}
+        if context_set_id:
+            ref["firestore_reference"]["agent_context_reference"] = {
+                "context_set_id": context_set_id
+            }
+
+        return ref

{google_cloud_db_context_engineering-0.7.3 → google_cloud_db_context_engineering-0.7.4}/src/google/cloud/db_context_enrichment/evaluate/evaluate_generator.py

@@ -1,7 +1,11 @@
+import importlib
 import json
+import logging
 import os
 import re
+import sys
 import textwrap
+import threading
 from typing import Any
 
 import yaml
@@ -10,10 +14,15 @@ from google.cloud.db_context_enrichment.common import config
 
 from .db_generators.alloydb import AlloyDBConfigGenerator
 from .db_generators.base import BaseDBConfigGenerator
+from .db_generators.bigtable import BigtableConfigGenerator
+from .db_generators.custom import CustomDBConfigGenerator
+from .db_generators.firestore import FirestoreConfigGenerator
 from .db_generators.mysql import MySQLConfigGenerator
 from .db_generators.postgres import PostgresConfigGenerator
 from .db_generators.spanner import SpannerConfigGenerator
 
+logger = logging.getLogger(__name__)
+
 # Constants for EvalBench configuration filenames
 DB_CONFIG_NAME = "db_config.yaml"
 MODEL_CONFIG_NAME = "model_config.yaml"
@@ -203,19 +212,93 @@ def _interpolate_env_vars(raw_yaml: str) -> str:
     return pattern.sub(replacer, raw_yaml)
 
 
+_SYS_PATH_LOCK = threading.Lock()
+
+
+def _import_with_cwd_fallback(mod_name: str):
+    """Imports mod_name, falling back to appending os.getcwd() to sys.path."""
+    try:
+        return importlib.import_module(mod_name)
+    except ModuleNotFoundError as e:
+        if e.name is None or not (
+            mod_name == e.name or mod_name.startswith(e.name + ".")
+        ):
+            raise
+        cwd = os.path.abspath(os.getcwd())
+        with _SYS_PATH_LOCK:
+            if not any(os.path.abspath(p or cwd) == cwd for p in sys.path):
+                # Keep cwd at the end of sys.path so lazy runtime imports work
+                # without shadowing standard library or installed packages.
+                sys.path.append(cwd)
+                importlib.invalidate_caches()
+        return importlib.import_module(mod_name)
+
+
 def _get_db_generator(params: dict[str, Any]) -> BaseDBConfigGenerator:
     """Factory function to build the correct Evaluation Generator."""
     source_type = params.get("type", "").lower()
 
     generators = {
         AlloyDBConfigGenerator.SOURCE_TYPE: AlloyDBConfigGenerator,
+        BigtableConfigGenerator.SOURCE_TYPE: BigtableConfigGenerator,
         PostgresConfigGenerator.SOURCE_TYPE: PostgresConfigGenerator,
         MySQLConfigGenerator.SOURCE_TYPE: MySQLConfigGenerator,
         SpannerConfigGenerator.SOURCE_TYPE: SpannerConfigGenerator,
         "spanner-postgres": SpannerConfigGenerator,
         "spanner-pg": SpannerConfigGenerator,
+        FirestoreConfigGenerator.SOURCE_TYPE: FirestoreConfigGenerator,
+        CustomDBConfigGenerator.SOURCE_TYPE: CustomDBConfigGenerator,
     }
 
+    # Dynamically register external custom database configuration generators.
+    # AUTOCTX_CUSTOM_GENERATORS holds a dotted Python module import path
+    # (e.g., "my_package.custom_generators") that exposes a
+    # CUSTOM_GENERATORS: dict[str, type[BaseDBConfigGenerator]] mapping
+    # custom tools.yaml source types to BaseDBConfigGenerator subclasses.
+    custom_gens = {}
+    custom_plugin = os.environ.get("AUTOCTX_CUSTOM_GENERATORS")
+    if custom_plugin:
+        try:
+            mod = _import_with_cwd_fallback(custom_plugin)
+            custom_gens = getattr(mod, "CUSTOM_GENERATORS", None)
+            if custom_gens is None:
+                raise AttributeError(
+                    f"Custom generator module '{custom_plugin}' must define 'CUSTOM_GENERATORS' dict."
+                )
+            if not isinstance(custom_gens, dict):
+                raise TypeError(
+                    f"CUSTOM_GENERATORS in '{custom_plugin}' must be a dictionary, "
+                    f"got {type(custom_gens).__name__}."
+                )
+            generators.update(custom_gens)
+        except Exception as e:
+            logger.error(
+                "Failed to load custom generators from plugin module '%s': %s",
+                custom_plugin,
+                e,
+            )
+            raise RuntimeError(
+                f"Failed to load custom generators from plugin module '{custom_plugin}': {e}"
+            ) from e
+
+    # If the source type is not handled by an external AUTOCTX_CUSTOM_GENERATORS
+    # plugin, allow inline connector_class / generator_class in tools.yaml to
+    # route directly to CustomDBConfigGenerator.
+    if source_type not in custom_gens and (
+        "connector_class" in params or "generator_class" in params
+    ):
+        if (
+            source_type in generators
+            and source_type != CustomDBConfigGenerator.SOURCE_TYPE
+        ):
+            logger.warning(
+                "Source type '%s' is being overridden by custom connector_class '%s' / generator_class '%s'.",
+                source_type,
+                params.get("connector_class"),
+                params.get("generator_class"),
+            )
+        return CustomDBConfigGenerator(params)
+
     if source_type not in generators:
         supported = ", ".join(generators.keys())
         raise ValueError(

{google_cloud_db_context_engineering-0.7.3 → google_cloud_db_context_engineering-0.7.4}/src/google/cloud/db_context_enrichment/main.py

@@ -74,7 +74,7 @@ def generate_evalbench_configs(
         dataset_path: The absolute path to the golden dataset file in the simplified user-facing format (JSON list of objects with keys: "id", "database", "nlq", "golden_sql").
         context_set_id: Full ContextSet resource name to evaluate against.
         toolbox_config_path: The absolute path to the tools.yaml configuration file.
-        toolbox_source_name: The name of the database source to use inside tools.yaml. The underlying source block must use a supported 'type' (cloud-sql-postgres, cloud-sql-mysql, spanner, alloydb-postgres).
+        toolbox_source_name: The name of the database source to use inside tools.yaml. The underlying source block must use a supported 'type' (cloud-sql-postgres, cloud-sql-mysql, spanner, alloydb-postgres, or custom).
 
     Returns:
         A message indicating that the configuration files were successfully created.
@@ -103,13 +103,14 @@ def generate_upload_url(
 
     Args:
         db_engine: The database engine. Accepted values are 'alloydb',
-                 'cloudsql', or 'spanner'. This can be derived from the 'kind'
-                 field in the tools.yaml file. For example, 'alloydb-postgres'
-                 becomes 'alloydb', and 'cloud-sql-postgres' becomes 'cloudsql'.
+                 'cloudsql', 'spanner', or 'bigtable'. This can be derived from
+                 the 'kind' field in the tools.yaml file. For example,
+                 'alloydb-postgres' becomes 'alloydb', 'cloud-sql-postgres'
+                 becomes 'cloudsql', and 'bigtable' becomes 'bigtable'.
         project_id: The Google Cloud project ID.
         location: The location of the AlloyDB cluster.
         cluster_id: The ID of the AlloyDB cluster.
-        instance_id: The ID of the Cloud SQL or Spanner instance.
+        instance_id: The ID of the Cloud SQL, Spanner, or Bigtable instance.
         database_id: The ID of the Spanner database.
 
     Returns:
@@ -130,8 +131,13 @@ def generate_upload_url(
             return f"https://console.cloud.google.com/spanner/instances/{instance_id}/databases/{database_id}/details/query?project={project_id}"
         else:
             return "Error: Missing instance_id, database_id, or project_id for spanner."
+    elif db_engine == "bigtable":
+        if instance_id and project_id:
+            return f"https://console.cloud.google.com/bigtable/instances/{instance_id}/overview?project={project_id}"
+        else:
+            return "Error: Missing instance_id or project_id for bigtable."
     else:
-        return "Error: Invalid db_engine. Must be one of 'alloydb', 'cloudsql', or 'spanner'."
+        return "Error: Invalid db_engine. Must be one of 'alloydb', 'cloudsql', 'spanner', or 'bigtable'."
 
 
 # NOTE: `@mcp.tool` is intentionally NOT applied to upload_context_set /

{google_cloud_db_context_engineering-0.7.3 → google_cloud_db_context_engineering-0.7.4/src/google_cloud_db_context_engineering.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-cloud-db-context-engineering
-Version: 0.7.3
+Version: 0.7.4
 Summary: A FastMCP server for generating natural language to SQL templates from database schemas.
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown

{google_cloud_db_context_engineering-0.7.3 → google_cloud_db_context_engineering-0.7.4}/src/google_cloud_db_context_engineering.egg-info/SOURCES.txt

@@ -16,6 +16,9 @@ src/google/cloud/db_context_enrichment/evaluate/result_reader.py
 src/google/cloud/db_context_enrichment/evaluate/db_generators/__init__.py
 src/google/cloud/db_context_enrichment/evaluate/db_generators/alloydb.py
 src/google/cloud/db_context_enrichment/evaluate/db_generators/base.py
+src/google/cloud/db_context_enrichment/evaluate/db_generators/bigtable.py
+src/google/cloud/db_context_enrichment/evaluate/db_generators/custom.py
+src/google/cloud/db_context_enrichment/evaluate/db_generators/firestore.py
 src/google/cloud/db_context_enrichment/evaluate/db_generators/mysql.py
 src/google/cloud/db_context_enrichment/evaluate/db_generators/postgres.py
 src/google/cloud/db_context_enrichment/evaluate/db_generators/spanner.py