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

google/cloud/db_context_enrichment/prompts/targeted_value_search.py

@@ -0,0 +1,93 @@
+import textwrap
+
+from google.cloud.db_context_enrichment.value_search.match_templates import (
+    MatchFunction,
+)
+
+_SUPPORTED_DB_ENGINE = ["alloydb", "postgres", "mysql", "spanner"]
+
+GENERATE_TARGETED_VALUE_SEARCH_PROMPT = textwrap.dedent(
+    """
+        **Workflow for Generating Targeted Value Search**
+
+        1.  **Database Configuration:**
+            - Ask the user for the **Database Engine (Product name) and optionally version**. Supported database engines: {supported_db_engines}
+            - **Understand the difference between Database Engine and Dialect:**
+              - **Database Engine** is the product name (e.g., `alloydb`, `postgres`, `mysql`, `spanner`). This is what the user provides.
+              - **Dialect** is the underlying SQL dialect/protocol used by tools (e.g., `postgresql` for `alloydb`/`postgres`, `googlesql` for `spanner`). You MUST infer this dialect from the engine.
+            - Using the database engine provided, infer the dialect:
+              - **alloydb** or **postgres** -> `postgresql`
+              - **mysql** -> `mysql`
+              - **spanner** -> `googlesql`
+            - Confirm the database engine provided is one of the supported engines. If not, notify the user this database engine is not supported yet and end the workflow.
+
+        2.  **Fetch Capabilities:**
+            - **Immediately after** receiving the Database Engine (and Version if provided), call the `list_match_functions` tool.
+            - **Important:** Pass the **inferred dialect** (as defined in step 1) to the `dialect` parameter of the tool.
+            - If the tool returns an error (e.g., unsupported version), present the error to the user (which includes the list of supported versions) and end the workflow.
+            - Otherwise, present the available match functions to the user, strictly including the function name, its Description, and an Example of when it should be used, using the information returned by the tool.
+
+        3.  **User Input Loop:**
+            - Ask the user to provide the following details for a value search:
+              - **Table Name**
+              - **Column Name**
+              - **Concept Type** (e.g., "City", "Product ID")
+              - **Match Function** (Must be one of the function names retrieved in Step 2)
+              - **Dialect Specific Parameters** (Ask for these if the chosen function and dialect require them):
+                - For **Spanner (googlesql)** + `{trigram_match}`: Ask for **Column Tokens** column name.
+                - For **MySQL** + `{semantic_match}`: Ask for **Column Embedding** column name.
+              - **Description** (optional): A description of the value search.
+            - After capturing the details, check if the input is valid, especially whether Match Function is a valid string from the returned values from list_match_functions, if not ask the user to do a correction, if valid, ask the user if they would like to add another one.
+            - Continue this loop until the user indicates they have no more value searches to add.
+
+        4.  **Review and Confirmation:**
+            - Present the complete list of user-provided value search definitions for confirmation.
+              - **Use the following format for each value search:**
+                **Index [Number]**
+                **Table:** [Table Name]
+                **Column:** [Column Name]
+                **Concept:** [Concept Type]
+                **Function:** [Match Function]
+                **Engine Specific Parameters:** [Column Tokens / Column Embedding, if provided]
+                **Description:** [Description]
+            - Ask if any modifications are needed. If so, work with the user to refine the list.
+
+        5.  **Final Generation:**
+            - Once approved, call the `generate_value_searches` tool with the list of value search definitions.
+            - Mapping for `generate_value_searches` input (JSON):
+              - `column_tokens` (for Column Tokens)
+              - `column_embedding` (for Column Embedding)
+            - **Important:** Pass the **inferred dialect** (from step 1) and `db_version` to the tool.
+            - Combine all generated Value Search configurations into a single JSON structure (ContextSet).
+
+        6.  **Save Value Search:**
+            - Ask the user to choose one of the following options:
+              1. Create a new context set file.
+              2. Append value search to an existing context set file.
+
+            - **If creating a new file:**
+              - You will need to ask the user for the database instance and database name to create the filename.
+              - Call the `save_context_set` tool. You will need to provide the database instance, database name, the JSON content from the previous step, and the root directory where the Gemini CLI is running.
+
+            - **If appending to an existing file:**
+              - Ask the user to provide the path to the existing context set file.
+              - Call the `attach_context_set` tool with the JSON content and the absolute file path.
+
+        7.  **Generate Upload URL (Optional):**
+            - After the file is saved, ask the user if they want to generate a URL to upload the context set file.
+            - If the user confirms, you must collect the necessary database context from them. This includes:
+              - **Database Type:** 'alloydb', 'cloudsql', or 'spanner'.
+              - **Project ID:** The Google Cloud project ID.
+              - **And depending on the database type:**
+                - For 'alloydb': Location and Cluster ID.
+                - For 'cloudsql': Instance ID.
+                - For 'spanner': Instance ID and Database ID.
+            - Once you have the required information, call the `generate_upload_url` tool to provide the upload URL to the user.
+
+        Start the workflow.
+        """
+).format(
+    supported_db_engines=_SUPPORTED_DB_ENGINE,
+    trigram_match=MatchFunction.TRIGRAM_STRING_MATCH.value,
+    semantic_match=MatchFunction.SEMANTIC_SIMILARITY_MATCH.value,
+)

google/cloud/db_context_enrichment/template/template_generator.py

@@ -0,0 +1,67 @@
+import json
+
+from google.cloud.db_context_enrichment.common import parameterizer
+from google.cloud.db_context_enrichment.model import context
+
+
+async def generate_templates(
+    template_inputs_json: str, sql_dialect: str = "postgresql"
+) -> str:
+    """
+    Generates the final, detailed templates 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 now expected to be a direct list of items
+        item_list = json.loads(template_inputs_json)
+        if not isinstance(item_list, list):
+            raise json.JSONDecodeError("Input is not a list.", template_inputs_json, 0)
+    except json.JSONDecodeError:
+        return '{"error": "Invalid JSON format for approved items. Expected a JSON array."}'
+
+    final_templates = []
+
+    for item in item_list:
+        question = item["question"]
+        sql = item["sql"]
+        intent = item.get(
+            "intent", question
+        )  # Use provided intent or fallback to question
+
+        # 1. Extract value phrases from the intent
+        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, intent, db_dialect=db_dialect
+        )
+
+        # 4. Assemble the final template object
+        template = context.Template(
+            nl_query=question,
+            sql=sql,
+            intent=intent,
+            manifest=manifest,
+            parameterized=context.ParameterizedTemplate(
+                parameterized_sql=parameterized_result["sql"],
+                parameterized_intent=parameterized_result["intent"],
+            ),
+        )
+        final_templates.append(template)
+
+    context_set = context.ContextSet(templates=final_templates, facets=None)
+    return context_set.model_dump_json(indent=2, exclude_none=True)

google/cloud/db_context_enrichment/value_search/generator.py

@@ -0,0 +1,91 @@
+import json
+from typing import Any
+
+from google.cloud.db_context_enrichment.model import context
+from google.cloud.db_context_enrichment.value_search import match_templates
+
+
+def generate_value_searches(
+    value_search_inputs_json: str,
+    dialect: str,
+    db_version: str | None = None,
+) -> str:
+    """
+    Generates a list of Value Search configurations based on a JSON input list.
+
+    Args:
+        value_search_inputs_json: A JSON string representing a list of dictionaries.
+            Each dictionary must contain:
+            - table_name (str)
+            - column_name (str)
+            - concept_type (str)
+            - match_function (str)
+            - description (str, optional)
+        dialect: The database dialect (e.g., 'postgresql').
+        db_version: The specific database version (optional).
+
+    Returns:
+        A JSON string representation of a ContextSet containing all generated value searches.
+    """
+    try:
+        inputs: list[dict[str, Any]] = json.loads(value_search_inputs_json)
+    except json.JSONDecodeError as e:
+        return json.dumps({"error": f"Invalid JSON format: {str(e)}"})
+
+    value_searches = []
+
+    for index, item in enumerate(inputs):
+        required_fields = [
+            "table_name",
+            "column_name",
+            "concept_type",
+            "match_function",
+        ]
+        for field in required_fields:
+            if not item.get(field):
+                return json.dumps(
+                    {"error": f"Field '{field}' is missing at index {index}"}
+                )
+
+        table_name = item.get("table_name")
+        column_name = item.get("column_name")
+        concept_type = item.get("concept_type")
+        match_function = item.get("match_function")
+        description = item.get("description")
+
+        try:
+            template_def = match_templates.get_match_template(
+                dialect=dialect,
+                function_name=match_function,
+                version=db_version,
+            )
+            raw_sql = template_def["sql_template"]
+
+            # Prepare formatting arguments with defaults for optional fields
+            format_args = {
+                "table": table_name,
+                "column": column_name,
+                "concept_type": concept_type,
+                "column_tokens": item.get("column_tokens", ""),
+                "column_embedding": item.get("column_embedding", ""),
+            }
+
+            value_search_query = raw_sql.format(**format_args)
+
+            vs = context.ValueSearch(
+                concept_type=concept_type,
+                query=value_search_query,
+                description=description,
+            )
+            value_searches.append(vs)
+
+        except ValueError as e:
+            return json.dumps(
+                {
+                    "error": f"Error while processing value search at index {index}: {str(e)}"
+                }
+            )
+
+    return context.ContextSet(value_searches=value_searches).model_dump_json(
+        indent=2, exclude_none=True
+    )

google/cloud/db_context_enrichment/value_search/match_templates.py

@@ -0,0 +1,267 @@
+from enum import StrEnum
+from typing import Any
+
+
+class Dialect(StrEnum):
+    """Supported database dialects."""
+
+    POSTGRESQL = "postgresql"
+    MYSQL = "mysql"
+    GOOGLE_SQL = "googlesql"
+
+
+class MatchFunction(StrEnum):
+    """Available match functions."""
+
+    EXACT_MATCH_STRINGS = "EXACT_MATCH_STRINGS"
+    TRIGRAM_STRING_MATCH = "TRIGRAM_STRING_MATCH"
+    SEMANTIC_SIMILARITY_MATCH = "SEMANTIC_SIMILARITY_MATCH"
+
+
+_MATCH_CONFIG: dict[Dialect, dict[str, Any]] = {
+    Dialect.POSTGRESQL: {
+        "min_version": "13",
+        # Default templates
+        "defaults": {
+            MatchFunction.EXACT_MATCH_STRINGS.value: {
+                "description": "Exact match for strings (Standard SQL).",
+                "example": "Use when finding a specific state code (e.g., 'CA'), order ID, or exact product name where precise spelling is required.",
+                "sql_template": (
+                    "SELECT $value as value, '{table}.{column}' as columns, "
+                    "'{concept_type}' as concept_type, 0 as distance, "
+                    '\'\' as context FROM "{table}" T WHERE T."{column}" = $value'
+                ),
+            },
+            MatchFunction.TRIGRAM_STRING_MATCH.value: {
+                "description": "Fuzzy text match using trigram similarity (Requires extension: pg_trgm).",
+                "example": "Use when searching for names, addresses, or plain text where users might have typos, misspellings, or partial matches.",
+                "sql_template": (
+                    "WITH TrigramMetrics AS ("
+                    '    SELECT T."{column}" AS original_value, '
+                    '    (T."{column}" <-> $value::text) AS normalized_dist '
+                    '    FROM "{table}" T '
+                    '    WHERE T."{column}" % $value::text'
+                    ") "
+                    "SELECT original_value AS value, '{table}.{column}' AS columns, "
+                    "'{concept_type}' AS concept_type, normalized_dist AS distance, "
+                    "''::text AS context FROM TrigramMetrics"
+                ),
+            },
+            MatchFunction.SEMANTIC_SIMILARITY_MATCH.value: {
+                "description": "Semantic similarity search using Gemini text embeddings (Requires extensions: vector, google_ml_integration).",
+                "example": "Use when searching for concepts, descriptions, themes, or abstract text where the exact words might differ but the underlying meaning is similar.",
+                "sql_template": (
+                    "WITH SemanticMetrics AS ("
+                    '    SELECT T."{column}" AS original_value, ('
+                    "        (google_ml.embedding('gemini-embedding-001', $value)::vector <=> "
+                    "         google_ml.embedding('gemini-embedding-001', T.\"{column}\")::vector) / 2.0"
+                    "    ) AS normalized_dist "
+                    '    FROM "{table}" T '
+                    '    WHERE T."{column}" IS NOT NULL'
+                    ") "
+                    "SELECT original_value AS value, '{table}.{column}' AS columns, "
+                    "'{concept_type}' AS concept_type, normalized_dist AS distance, "
+                    "''::text AS context FROM SemanticMetrics"
+                ),
+            },
+        },
+        "overrides": {},
+    },
+    Dialect.GOOGLE_SQL: {
+        "min_version": "1",
+        "defaults": {
+            MatchFunction.EXACT_MATCH_STRINGS.value: {
+                "description": "Exact match for strings in Spanner.",
+                "example": "Use for exact IDs or state codes in Spanner.",
+                "sql_template": (
+                    "SELECT CAST($value AS STRING) AS value, '{column}' AS `columns`, "
+                    "'{concept_type}' AS concept_type, 0 AS distance, "
+                    "JSON '{{}}' AS context "
+                    "FROM `{table}` AS T "
+                    "WHERE CAST(T.`{column}` AS STRING) = CAST($value AS STRING) "
+                ),
+            },
+            MatchFunction.TRIGRAM_STRING_MATCH.value: {
+                "description": "String similarity using Spanner Search Indexes.",
+                "example": "Use for typos/misspellings in Spanner using SEARCH_NGRAMS.",
+                "sql_template": (
+                    "SELECT CAST(T.`{column}` AS STRING) AS value, '{column}' AS `columns`, "
+                    "'{concept_type}' AS concept_type, "
+                    "1 - SCORE_NGRAMS(T.`{column_tokens}`, CAST($value AS STRING)) AS distance, "
+                    "JSON '{{}}' AS context "
+                    "FROM `{table}` AS T "
+                    "WHERE SEARCH_NGRAMS(T.`{column_tokens}`, CAST($value AS STRING)) "
+                ),
+            },
+        },
+        "overrides": {},
+    },
+    Dialect.MYSQL: {
+        "min_version": "8",
+        "defaults": {
+            MatchFunction.EXACT_MATCH_STRINGS.value: {
+                "description": "Exact match for strings in MySQL.",
+                "example": "Use for exact matching in MySQL.",
+                "sql_template": (
+                    "SELECT $value AS value, '{column}' AS `columns`, "
+                    "'{concept_type}' AS concept_type, 0 AS distance, "
+                    "JSON_OBJECT() AS context "
+                    "FROM `{table}` AS T WHERE T.`{column}` = $value"
+                ),
+            },
+            MatchFunction.TRIGRAM_STRING_MATCH.value: {
+                "description": "Trigram fuzzy match in MySQL using FULLTEXT index with score normalization.",
+                "example": "Use for fuzzy matching in MySQL (requires FULLTEXT index with ngram).",
+                "sql_template": (
+                    "SELECT * FROM ("
+                    "  WITH TrigramMetrics AS ("
+                    "    SELECT T.`{column}` AS original_value, "
+                    "    MATCH(T.`{column}`) AGAINST($value IN NATURAL LANGUAGE MODE) AS raw_score "
+                    "    FROM `{table}` AS T "
+                    "    WHERE MATCH(T.`{column}`) AGAINST($value IN NATURAL LANGUAGE MODE) > 0 "
+                    "    ORDER BY raw_score DESC LIMIT 10"
+                    "  ), "
+                    "  NormalizationParams AS ("
+                    "    SELECT MAX(raw_score) AS max_score "
+                    "    FROM TrigramMetrics"
+                    "  ) "
+                    "  SELECT original_value AS value, '{column}' AS `columns`, "
+                    "  '{concept_type}' AS concept_type, "
+                    "  (CASE WHEN n.max_score > 0 THEN (1 - (m.raw_score / n.max_score)) ELSE 0 END) AS distance, "
+                    "  JSON_OBJECT() AS context "
+                    "  FROM TrigramMetrics m, NormalizationParams n"
+                    ") AS wrapped_query "
+                ),
+            },
+            MatchFunction.SEMANTIC_SIMILARITY_MATCH.value: {
+                "description": "Semantic match in MySQL using Vertex AI embedding.",
+                "example": "Use for semantic matching (requires mysql.ml_embedding).",
+                "sql_template": (
+                    "SELECT * FROM ("
+                    "  WITH search_embedding AS ("
+                    "    SELECT mysql.ml_embedding('text-embedding-005', $value) AS val"
+                    "  ) "
+                    "  SELECT T.`{column}` AS value, '{column}' AS `columns`, "
+                    "  '{concept_type}' AS concept_type, "
+                    "  COSINE_DISTANCE(T.`{column_embedding}`, search_embedding.val) AS distance, "
+                    "  JSON_OBJECT() AS context "
+                    "  FROM `{table}` AS T, search_embedding "
+                    "  WHERE T.`{column_embedding}` IS NOT NULL"
+                    ") AS wrapped_query "
+                ),
+            },
+        },
+        "overrides": {},
+    },
+}
+
+
+def _is_version_supported(version: str, min_version: str) -> bool:
+    """Helper to compare version strings (e.g. '13.2' >= '13')."""
+
+    def parse(v: str):
+        return tuple(map(int, v.split(".")))
+
+    try:
+        return parse(version) >= parse(min_version)
+    except ValueError:
+        return False
+
+
+def get_match_template(
+    dialect: str, function_name: str, version: str | None = None
+) -> dict:
+    """
+    Retrieves a match template with a default-fallback strategy.
+
+    Args:
+        dialect: The database dialect string (e.g., 'postgresql').
+        function_name: The name of the match function.
+        version: The specific database version (optional).
+
+    Returns:
+        A dictionary containing the template definition.
+
+    Raises:
+        ValueError:
+            - If dialect is invalid.
+            - If version is provided but unsupported.
+            - If function_name is not found (lists available templates).
+    """
+    try:
+        dialect_enum = Dialect(dialect.lower())
+    except ValueError:
+        supported = [d.value for d in Dialect]
+        raise ValueError(
+            f"Dialect '{dialect}' not supported. Supported dialects: {supported}"
+        )
+
+    engine_config = _MATCH_CONFIG.get(dialect_enum)
+    if not engine_config:
+        raise ValueError(f"Dialect '{dialect}' has no configuration registered.")
+
+    defaults = engine_config.get("defaults", {})
+    min_version = engine_config.get("min_version")
+    overrides = engine_config.get("overrides", {})
+
+    if version and min_version:
+        version = str(version)
+        if not _is_version_supported(version, min_version):
+            raise ValueError(
+                f"Version '{version}' is not supported for dialect '{dialect}'. "
+                f"Minimum required version: {min_version}"
+            )
+
+    # Identify specific overrides for this version (if any)
+    # Note: Overrides currently use exact version matches in the keys
+    version_overrides = overrides.get(version, {}) if version else {}
+
+    effective_templates = defaults | version_overrides
+    template = effective_templates.get(function_name)
+
+    if not template:
+        supported_templates = list(defaults.keys())
+        raise ValueError(
+            f"Match function '{function_name}' not found. "
+            f"Supported match templates: {supported_templates}"
+        )
+
+    return template
+
+
+def get_available_functions(
+    dialect: str, version: str | None = None
+) -> dict[str, dict[str, str]]:
+    """
+    Returns a dictionary of available match function names with their descriptions and examples for a given dialect.
+    Validates both the dialect and the version (if provided).
+    """
+    try:
+        dialect_enum = Dialect(dialect.lower())
+    except ValueError:
+        supported = [d.value for d in Dialect]
+        raise ValueError(
+            f"Dialect '{dialect}' not supported. Supported engine: {supported}"
+        )
+
+    engine_config = _MATCH_CONFIG.get(dialect_enum, {})
+
+    if version:
+        min_version = engine_config.get("min_version")
+        version = str(version)
+        if min_version and not _is_version_supported(version, min_version):
+            raise ValueError(
+                f"Version '{version}' is not supported for dialect '{dialect}'. "
+                f"Minimum required version: {min_version}"
+            )
+
+    defaults = engine_config.get("defaults", {})
+    version_overrides = (
+        engine_config.get("overrides", {}).get(version, {}) if version else {}
+    )
+    effective_templates = defaults | version_overrides
+
+    return {
+        k: {"description": v.get("description", ""), "example": v.get("example", "")}
+        for k, v in effective_templates.items()
+    }