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

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

@@ -0,0 +1,69 @@
+from typing import Any
+
+import google.cloud.geminidataanalytics_v1beta as gda
+import yaml
+
+from .base import BaseDBConfigGenerator
+
+
+class MySQLConfigGenerator(BaseDBConfigGenerator):
+    """
+    Dedicated generator mapping properties to explicit Cloud SQL MySQL configuration
+    topologies utilized by both EvalBench binaries and GDA Context objects.
+    """
+
+    SOURCE_TYPE = "cloud-sql-mysql"
+    DIALECT = "mysql"
+    REQUIRED_FIELDS = BaseDBConfigGenerator.REQUIRED_FIELDS | {
+        "project",
+        "region",
+        "instance",
+        "database",
+    }
+
+    def __init__(self, params: dict[str, Any]):
+        super().__init__(params)
+        self.project = params.get("project")
+        self.region = params.get("region")
+        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 = "mysql"
+        db_path = f"{self.project}:{self.region}:{self.instance}"
+
+        db_config = {
+            "db_type": db_type,
+            "dialect": self.DIALECT,
+            "database_name": self.database,
+            "database_path": db_path,
+            "max_executions_per_minute": 180,
+        }
+        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.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

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

@@ -0,0 +1,69 @@
+from typing import Any
+
+import google.cloud.geminidataanalytics_v1beta as gda
+import yaml
+
+from .base import BaseDBConfigGenerator
+
+
+class PostgresConfigGenerator(BaseDBConfigGenerator):
+    """
+    Dedicated generator mapping properties to explicit Cloud SQL Postgres configuration
+    topologies utilized by both EvalBench binaries and GDA Context objects.
+    """
+
+    SOURCE_TYPE = "cloud-sql-postgres"
+    DIALECT = "postgres"
+    REQUIRED_FIELDS = BaseDBConfigGenerator.REQUIRED_FIELDS | {
+        "project",
+        "region",
+        "instance",
+        "database",
+    }
+
+    def __init__(self, params: dict[str, Any]):
+        super().__init__(params)
+        self.project = params.get("project")
+        self.region = params.get("region")
+        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 = "postgres"
+        db_path = f"{self.project}:{self.region}:{self.instance}"
+
+        db_config = {
+            "db_type": db_type,
+            "dialect": self.DIALECT,
+            "database_name": self.database,
+            "database_path": db_path,
+            "max_executions_per_minute": 180,
+        }
+        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.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

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

@@ -0,0 +1,62 @@
+from typing import Any
+
+import google.cloud.geminidataanalytics_v1beta as gda
+import yaml
+
+from .base import BaseDBConfigGenerator
+
+
+class SpannerConfigGenerator(BaseDBConfigGenerator):
+    """
+    Dedicated generator mapping properties to explicit Spanner configuration
+    topologies utilized by both EvalBench binaries and GDA Context objects.
+    """
+
+    SOURCE_TYPE = "spanner"
+    DIALECT = "spanner_gsql"
+    REQUIRED_FIELDS = BaseDBConfigGenerator.REQUIRED_FIELDS | {
+        "project",
+        "instance",
+        "database",
+    }
+
+    def __init__(self, params: dict[str, Any]):
+        super().__init__(params)
+        self.project = params.get("project")
+        self.instance = params.get("instance")
+        self.database = params.get("database")
+
+    def generate_db_config(self) -> str:
+        db_type = "spanner"
+        db_path = f"projects/{self.project}/instances/{self.instance}/databases/{self.database}"
+
+        db_config = {
+            "db_type": db_type,
+            "dialect": self.DIALECT,
+            "database_name": self.database,
+            "database_path": db_path,
+            "instance_id": self.instance,
+            "gcp_project_id": self.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
+    ) -> gda.DatasourceReferences:
+        datasource_ref = gda.DatasourceReferences()
+
+        datasource_ref.spanner_reference = gda.SpannerReference(
+            database_reference=gda.SpannerDatabaseReference(
+                engine=gda.SpannerDatabaseReference.Engine.GOOGLE_SQL,
+                project_id=self.project,
+                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/evaluate_generator.py

@@ -0,0 +1,239 @@
+import json
+import os
+import re
+import textwrap
+from typing import Any
+
+import yaml
+
+from google.cloud.db_context_enrichment.common import config
+
+from .db_generators.alloydb import AlloyDBConfigGenerator
+from .db_generators.base import BaseDBConfigGenerator
+from .db_generators.mysql import MySQLConfigGenerator
+from .db_generators.postgres import PostgresConfigGenerator
+from .db_generators.spanner import SpannerConfigGenerator
+
+# Constants for EvalBench configuration filenames
+DB_CONFIG_NAME = "db_config.yaml"
+MODEL_CONFIG_NAME = "model_config.yaml"
+RUN_CONFIG_NAME = "run_config.yaml"
+LLMRATER_CONFIG_NAME = "llmrater_config.yaml"
+GOLDEN_QUERIES_NAME = "golden_queries.json"
+
+
+def generate_evalbench_configs(
+    experiment_name: str,
+    dataset_path: str,
+    context_set_id: str,
+    toolbox_config_path: str,
+    toolbox_source_name: str,
+) -> None:
+    """
+    Main entrypoint: Generates Evalbench-compatible YAML configurations natively using
+    private DB format converters and the google-cloud-geminidataanalytics API validations.
+    """
+    params = _extract_toolbox_params(toolbox_config_path, toolbox_source_name)
+    generator = _get_db_generator(params)
+
+    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)
+
+    # 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"
+    os.makedirs(eval_configs_dir, exist_ok=True)
+
+    with open(os.path.join(eval_configs_dir, DB_CONFIG_NAME), "w") as f:
+        f.write(db_config_yaml)
+
+    with open(os.path.join(eval_configs_dir, MODEL_CONFIG_NAME), "w") as f:
+        f.write(model_config_yaml)
+
+    with open(os.path.join(eval_configs_dir, RUN_CONFIG_NAME), "w") as f:
+        f.write(run_config_yaml)
+
+    with open(os.path.join(eval_configs_dir, LLMRATER_CONFIG_NAME), "w") as f:
+        f.write(llmrater_config_yaml)
+
+    with open(os.path.join(eval_configs_dir, GOLDEN_QUERIES_NAME), "w") as f:
+        f.write(golden_queries_json)
+
+
+def _extract_toolbox_params(
+    toolbox_config_path: str, toolbox_source_name: str
+) -> dict[str, Any]:
+    """Deterministically extracts connection parameters for a specific database source from tools.yaml."""
+    try:
+        with open(toolbox_config_path) as f:
+            content = f.read()
+            interpolated = _interpolate_env_vars(content)
+            docs = yaml.safe_load_all(interpolated)
+            for doc in docs:
+                if not doc:
+                    continue
+                if (
+                    doc.get("kind") == "source"
+                    and doc.get("name") == toolbox_source_name
+                ):
+                    if not doc.get("type"):
+                        raise ValueError(
+                            f"Selected source '{toolbox_source_name}' is missing the 'type' field."
+                        )
+                    return doc
+
+            raise ValueError(
+                f"Could not find a 'kind: source' named '{toolbox_source_name}' in {toolbox_config_path}"
+            )
+
+    except FileNotFoundError:
+        raise ValueError(f"Config file not found: {toolbox_config_path}")
+    except PermissionError:
+        raise ValueError(
+            f"Permission denied reading config file: {toolbox_config_path}"
+        )
+    except yaml.YAMLError as e:
+        raise ValueError(f"Failed to parse {toolbox_config_path} as YAML: {e}")
+
+
+def _interpolate_env_vars(raw_yaml: str) -> str:
+    """Replaces ${ENV_NAME} or ${ENV_NAME:default_value} with environment variables."""
+    # Matches ${VAR_NAME} or ${VAR_NAME:fallback}
+    pattern = re.compile(r"\$\{(\w+)(?::([^}]*))?\}")
+
+    def replacer(match):
+        var_name = match.group(1)
+        fallback = match.group(2)
+
+        if var_name in os.environ:
+            return os.environ[var_name]
+        if fallback is not None:
+            return fallback
+        raise ValueError(
+            f"Environment variable '{var_name}' not found and no default provided."
+        )
+
+    return pattern.sub(replacer, raw_yaml)
+
+
+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,
+        PostgresConfigGenerator.SOURCE_TYPE: PostgresConfigGenerator,
+        MySQLConfigGenerator.SOURCE_TYPE: MySQLConfigGenerator,
+        SpannerConfigGenerator.SOURCE_TYPE: SpannerConfigGenerator,
+    }
+
+    if source_type not in generators:
+        supported = ", ".join(generators.keys())
+        raise ValueError(
+            f"Unsupported evaluating toolbox source type: '{source_type}'. Must be one of: {supported}"
+        )
+
+    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"
+
+    return textwrap.dedent(f"""\
+        ############################################################
+        ### Dataset / Eval Items
+        ############################################################
+        dataset_config: {configs_dir}/{GOLDEN_QUERIES_NAME}
+        dataset_format: evalbench-standard-format
+        database_configs:
+         - {configs_dir}/{DB_CONFIG_NAME}
+        dialect: {dialect}    # DB connection mapping
+        query_types:
+         - dql
+
+        ############################################################
+        ### Prompt and Generation Modules
+        ############################################################
+        model_config: {configs_dir}/{MODEL_CONFIG_NAME}
+        prompt_generator: 'NOOPGenerator'
+
+        ############################################################
+        ### Evaluator Execution / Parallelism Tuning
+        ############################################################
+        runners:
+          eval_runners: 4
+          sqlgen_runners: 20
+
+        ############################################################
+        ### Scorer Related Configs
+        ############################################################
+        scorers:
+          llmrater:
+            model_config: {configs_dir}/{LLMRATER_CONFIG_NAME}
+
+        ############################################################
+        ### Reporting Related Configs
+        ############################################################
+        reporting:
+          csv:
+            output_directory: '{reports_dir}/'
+    """).strip()
+
+
+def _generate_llmrater_config(project_id: str) -> str:
+    """Generates a dedicated LLM rater model configuration mimicking standard text models."""
+    return textwrap.dedent(f"""\
+        generator: gcp_vertex_gemini
+        vertex_model: {config.get_model_name()}
+        gcp_project_id: {project_id}
+        gcp_region: global
+        base_prompt: ""
+        execs_per_minute: 20
+    """).strip()
+
+
+def _convert_dataset(dataset_path: str, dialect: str) -> str:
+    """Reads simplified dataset and converts to EvalBench standard format."""
+    try:
+        with open(dataset_path) as f:
+            data = json.load(f)
+
+        if not isinstance(data, list):
+            raise ValueError("Dataset must be a JSON list.")
+
+        required_keys = {"id", "nlq", "database", "golden_sql"}
+        for i, entry in enumerate(data):
+            if not isinstance(entry, dict):
+                raise ValueError(f"Dataset entry at index {i} is not a dictionary.")
+            missing = required_keys - set(entry.keys())
+            if missing:
+                raise ValueError(
+                    f"Dataset entry at index {i} is missing required keys: {missing}"
+                )
+
+        converted = []
+        for entry in data:
+            converted_entry = {
+                "id": entry.get("id"),
+                "nl_prompt": entry.get("nlq"),
+                "query_type": "DQL",
+                "database": entry.get("database"),
+                "dialects": [dialect],
+                "golden_sql": {dialect: [entry.get("golden_sql")]},
+                "eval_query": {},
+                "setup_sql": {},
+                "cleanup_sql": {},
+                "other": {},
+                "tags": [],
+            }
+            converted.append(converted_entry)
+
+        return json.dumps(converted, indent=2)
+    except Exception as e:
+        raise ValueError(f"Failed to convert dataset at {dataset_path}: {e}")

google/cloud/db_context_enrichment/evaluate/result_reader.py

@@ -0,0 +1,181 @@
+import csv
+import os
+import re
+import textwrap
+from dataclasses import dataclass
+
+
+@dataclass
+class ScoreRecord:
+    id: str
+    score: float
+    comparison_logs: str = "N/A"
+
+
+@dataclass
+class EvalRecord:
+    id: str
+    nl_prompt: str = "N/A"
+    golden_sql: str = "N/A"
+    generated_sql: str = "N/A"
+    sql_generator_error: str = "N/A"
+    generated_error: str = "N/A"
+    other: str = "N/A"
+
+
+def read_eval_results(
+    run_folder_path: str, offset: int = 0, batch_size: int = 10
+) -> str:
+    """
+    Reads evaluation results from a folder and produces a markdown summary.
+
+    Args:
+        run_folder_path: The absolute path to the evaluation run result folder, which ends with the eval run job id.
+        offset: Offset to start reading failure cases from (default: 0).
+        batch_size: Number of failure cases to show in the report (default: 10).
+
+    Returns:
+        A string in markdown format containing the summary and failure cases.
+    """
+    scores_path = os.path.join(run_folder_path, "scores.csv")
+    evals_path = os.path.join(run_folder_path, "evals.csv")
+
+    # Read Summary
+    summary_md = _read_summary(run_folder_path)
+
+    # Read Scores to find failures
+    with open(scores_path, encoding="utf-8") as f:
+        failures = sorted(
+            [
+                ScoreRecord(
+                    id=row["id"],
+                    score=score,
+                    comparison_logs=row.get("comparison_logs") or "N/A",
+                )
+                for row in csv.DictReader(f)
+                if (score := _get_score(row)) is not None and score < 100
+            ],
+            key=lambda x: _natural_sort_key(x.id),
+        )
+
+    if not failures:
+        return summary_md + "No failure cases found (all passed)."
+
+    # Apply batching
+    batched_failures = failures[offset : offset + batch_size]
+
+    # Add failures info to summary
+    summary_md += textwrap.dedent(
+        f"""\
+        ## All Failures
+        {", ".join([str(f.id) for f in failures])}
+
+        **Showing failures**: {offset + 1} to {min(offset + batch_size, len(failures))} of {len(failures)}
+
+        """
+    )
+
+    # Read Evals to get prompts and golden SQL
+    with open(evals_path, encoding="utf-8") as f:
+        evals_data = {
+            row["id"]: EvalRecord(
+                id=row["id"],
+                nl_prompt=row.get("nl_prompt") or "N/A",
+                golden_sql=row.get("golden_sql") or "N/A",
+                generated_sql=row.get("generated_sql") or "N/A",
+                sql_generator_error=row.get("sql_generator_error") or "N/A",
+                generated_error=row.get("generated_error") or "N/A",
+                other=row.get("other") or "N/A",
+            )
+            for row in csv.DictReader(f)
+        }
+
+    # Format Failures
+    failures_md = _format_failures(batched_failures, evals_data)
+
+    return summary_md + failures_md
+
+
+def _format_failures(
+    failures: list[ScoreRecord], evals_data: dict[str, EvalRecord]
+) -> str:
+    failures_md = "# Failure Cases\n\n"
+    for fail in failures:
+        fail_id = fail.id
+        eval_info = evals_data.get(fail_id, EvalRecord(id=fail_id))
+
+        failures_md += textwrap.dedent(
+            f"""\
+            ## Case ID: {fail_id} (Score: {fail.score})
+
+            **Prompt**:
+            {eval_info.nl_prompt}
+
+            **Golden SQL**:
+            ```sql
+            {eval_info.golden_sql}
+            ```
+
+            **Generated SQL**:
+            ```sql
+            {eval_info.generated_sql}
+            ```
+
+            **SQL Generator Error** (Errors during SQL generation):
+            ```
+            {eval_info.sql_generator_error}
+            ```
+
+            **Execution Error** (Errors when executing the generated SQL):
+            ```
+            {eval_info.generated_error}
+            ```
+
+            **Additional Output**:
+            ```
+            {eval_info.other}
+            ```
+
+            **Evaluation Details**:
+            {fail.comparison_logs}
+
+            ---
+
+            """
+        )
+    return failures_md
+
+
+def _read_summary(run_folder_path: str) -> str:
+    summary_path = os.path.join(run_folder_path, "summary.csv")
+    summary_md = "# Evaluation Summary\n\n"
+    with open(summary_path, encoding="utf-8") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            summary_md += textwrap.dedent(
+                f"""\
+                - **Metric**: {row.get("metric_name", "N/A")}
+                  - **Correct / Total**: {row.get("correct_results_count", "N/A")}/{row.get("total_results_count", "N/A")}
+                  - **Run Time**: {row.get("run_time", "N/A")}
+
+                """
+            )
+    return summary_md
+
+
+def _natural_sort_key(val):
+    """
+    Splits a string into a list of string and integer chunks.
+    Example: "user10" -> ["user", 10, ""]
+    """
+    return [
+        int(text) if text.isdigit() else text.lower()
+        for text in re.split(r"(\d+)", str(val))
+    ]
+
+
+def _get_score(row: dict[str, str]) -> float | None:
+    try:
+        return float(row.get("score", ""))
+    except ValueError:
+        return None

google/cloud/db_context_enrichment/facet/facet_generator.py

@@ -0,0 +1,68 @@
+import json
+
+from google.cloud.db_context_enrichment.common import parameterizer
+from google.cloud.db_context_enrichment.model import context
+
+
+async def generate_facets(
+    facet_inputs_json: str, sql_dialect: str = "postgresql"
+) -> str:
+    """
+    Generates the final, detailed facets based on user-approved items.
+    """
+    try:
+        # Convert the string to the Enum member
+        db_dialect = parameterizer.SQLDialect(sql_dialect)
+    except ValueError:
+        return f'{{"error": "Invalid database dialect specified: {sql_dialect}"}}'
+
+    try:
+        # The input is expected to be a direct list of items
+        item_list = json.loads(facet_inputs_json)
+        if not isinstance(item_list, list):
+            raise json.JSONDecodeError("Input is not a list.", facet_inputs_json, 0)
+    except json.JSONDecodeError:
+        return '{"error": "Invalid JSON format for approved items. Expected a JSON array."}'
+
+    final_facets = []
+
+    for item in item_list:
+        sql_snippet = item.get("sql_snippet")
+        if not sql_snippet:
+            return '{"error": "Each item must have a \'sql_snippet\' key."}'
+
+        intent = item.get("intent")
+        if not intent:
+            return '{"error": "Each item must have an \'intent\' key."}'
+
+        # 1. Extract value phrases from the intent (used as nl_query)
+        phrases = await parameterizer.extract_value_phrases(nl_query=intent)
+
+        # 2. Generate the manifest
+        manifest = intent
+        # Sort keys by length descending to replace longer phrases first
+        sorted_phrases = sorted(phrases.keys(), key=len, reverse=True)
+        for phrase in sorted_phrases:
+            # Use the first identified type for the manifest
+            phrase_type = phrases[phrase][0] if phrases[phrase] else "value"
+            manifest = manifest.replace(phrase, f"a given {phrase_type}")
+
+        # 3. Parameterize the SQL and Intent
+        parameterized_result = parameterizer.parameterize_sql_and_intent(
+            phrases, sql_snippet, intent, db_dialect=db_dialect
+        )
+
+        # 4. Assemble the final facet object
+        facet = context.Facet(
+            sql_snippet=sql_snippet,
+            intent=intent,
+            manifest=manifest,
+            parameterized=context.ParameterizedFacet(
+                parameterized_sql_snippet=parameterized_result["sql"],
+                parameterized_intent=parameterized_result["intent"],
+            ),
+        )
+        final_facets.append(facet)
+
+    context_set = context.ContextSet(templates=None, facets=final_facets)
+    return context_set.model_dump_json(indent=2, exclude_none=True)