notebooks/(standalone)Run_OpenDataQnA.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div style=\"display: flex; align-items: left;\">\n",
        "    <a href=\"https://sites.google.com/corp/google.com/genai-solutions/home?authuser=0\">\n",
        "        <img src=\"https://storage.googleapis.com/miscfilespublic/Linkedin%20Banner%20%E2%80%93%202.png\" style=\"margin-right\">\n",
        "    </a>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DRyGcAepAPJ5"
      },
      "source": [
        "\n",
        "<h1 align=\"center\">Open Data QnA - Chat with your SQL Database (standalone notebook)</h1> \n",
        "\n",
        "---\n",
        "\n",
        "This notebook is a standalone version for quick testing of the solution without the need to clone the complete Github repository. \n",
        "Given the standalone nature, the notebook only supports BigQuery as the data source and for the Vector Store. \n",
        "\n",
        "The notebook first walks through the Vector Store Setup needed for running the Open Data QnA application. \n",
        "\n",
        "Currently supported Source DBs **from the complete solution** are: \n",
        "- PostgreSQL on Google Cloud SQL \n",
        "- BigQuery\n",
        "\n",
        "Furthermore, the following vector stores are supported \n",
        "- pgvector on PostgreSQL \n",
        "- BigQuery vector\n",
        "\n",
        "**If you want to use a different data source or a different vector store, make sure to clone the repository and run the standard (non-standalone) notebook instead!**\n",
        "\n",
        "The setup part covers the following steps: \n",
        "> 1. Configuration: Intial GCP project, IAM permissions, Environment  and Databases setup including logging on Bigquery for analytics\n",
        "\n",
        "> 2. Creation of Table, Column and Known Good Query Embeddings in the Vector Store  for Retreival Augmented Generation(RAG)\n",
        "\n",
        "\n",
        "Afterwards, you will be able to run the Open Data QnA Pipeline to generate SQL queries and answer questions over your data source. \n",
        "\n",
        "The pipeline run covers the following steps: \n",
        "\n",
        "> 1. Take user question and generate sql in the dialect corresponding to data source\n",
        "\n",
        "> 2. Execute the sql query and retreive the data\n",
        "\n",
        "> 3. Generate natural language respose and charts to display\n",
        "\n",
        "> 4. Clean Up resources\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jsWGZW_fUJjN"
      },
      "source": [
        "### 📒 Using this interactive notebook\n",
        "\n",
        "If you have not used this IDE with jupyter notebooks it will ask for installing Python + Jupyter extensions. Please go ahead install them\n",
        "\n",
        "Click the **run** icons ▶️  of each cell within this notebook.\n",
        "\n",
        "> 💡 Alternatively, you can run the currently selected cell with `Ctrl + Enter` (or `⌘ + Enter` on a Mac).\n",
        "\n",
        "> ⚠️ **To avoid any errors**, wait for each section to finish in their order before clicking the next “run” icon.\n",
        "\n",
        "This sample must be connected to a **Google Cloud project**, but nothing else is needed other than your Google Cloud project.\n",
        "\n",
        "You can use an existing project. Alternatively, you can create a new Cloud project [with cloud credits for free.](https://cloud.google.com/free/docs/gcp-free-tier)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RicDCkdI-hmp"
      },
      "source": [
        "# 🚧 **0. Prerequisites**\n",
        "\n",
        "Make sure that Google Cloud CLI is installed before moving to the next cell! You can refer to the link below for guidance\n",
        "\n",
        "Installation Guide: https://cloud.google.com/sdk/docs/install"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##  **0.1. Setup Poetry Environment and Setup GCP Project** "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 💻 **Install Code Dependencies**\n",
        "Install the dependencies by runnign the commands below. \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!pip install vertexai \n",
        "!pip install pandas \n",
        "!pip install db-dtypes\n",
        "!pip install google \n",
        "!pip install google-cloud-bigquery \n",
        "!pip install google-oauth \n",
        "!pip install tabulate\n",
        "!pip install google-cloud-bigquery-connection"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 📌 **Important Step: Activate your virtual environment and authenticate with Google Cloud CLI [Run all these on Terminal]**\n",
        "\n",
        "Once the venv created either in the local directory or in the cache directory. Open the terminal on the same machine where your notebooks are running and start running the below commands.\n",
        "\n",
        "\n",
        "```\n",
        "poetry shell #this command should activate your venv and you should see it enters into the venv\n",
        "\n",
        "##inside the activated venv shell\n",
        "\n",
        "gcloud auth login\n",
        "gcloud auth application-default login\n",
        "gcloud services enable \\\n",
        "    serviceusage.googleapis.com \\\n",
        "    cloudresourcemanager.googleapis.com --project <<Enter Project Id>>\n",
        "gcloud auth application-default set-quota-project <<Enter Project Id for using resources>>\n",
        "\n",
        "```\n",
        "For IDEs adding Juypter Extensions will automatically give you option to change the kernel. If not, manually select the python interpreter in your IDE (The exact is shown in the above cell. Path would look like e.g. /home/admin_/Talk2Data/.venv/bin/python or ~cache/user/Talk2Data/.venv/bin/python)\n",
        "\n",
        "**Extra Steps if you are running inside Jupyter Lab or Jupyter Environments on Workbench etc**\n",
        "\n",
        "We need to manually add venv as Kernel in the those instance where you don't have choice to select the path manually like above.\n",
        "\n",
        "Run the steps above and continue with below\n",
        "\n",
        "```\n",
        "##still inside the activated venv shell\n",
        "\n",
        "pip install jupyter\n",
        "\n",
        "ipython kernel install --name \"openqna-venv\" --user \n",
        "\n",
        "```\n",
        "\n",
        "Restart your kernel or close the exsiting notebook and open again, you should now see the \"openqna-venv\" in the kernel drop down\n",
        "\n",
        "**What did we do here?**\n",
        "\n",
        "* Created Application Default Credentials to use for the code\n",
        "* Added venv to kernel to select for runningt the notebooks (For standalone Jupyter setups like Workbench etc)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yygMe6rPWxHS"
      },
      "source": [
        "### 🔐 **Authenticate to Google Cloud** (Colab)\n",
        "Authenticate to Google Cloud as the IAM user logged into this notebook in order to access your Google Cloud Project.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\"\"\"Colab Auth\"\"\" \n",
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p4W6FPnrYEE8"
      },
      "source": [
        "### 🔗 **Connect Your Google Cloud Project**\n",
        "Time to connect your Google Cloud Project to this notebook. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os \n",
        "\n",
        "#@markdown Please fill in the value below with your GCP project ID and then run the cell.\n",
        "project_id = \"my_project\"\n",
        "\n",
        "# Quick input validations.\n",
        "assert project_id, \"⚠️ Please provide your Google Cloud Project ID\"\n",
        "\n",
        "# Configure gcloud.\n",
        "!gcloud config set project {project_id}\n",
        "print(f'Project has been set to {project_id}')\n",
        "\n",
        "os.environ['GOOGLE_CLOUD_QUOTA_PROJECT']=project_id\n",
        "os.environ['GOOGLE_CLOUD_PROJECT']=project_id"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### ⚙️ **Enable Required API Services in the GCP Project**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#Enable all the required APIs for the Open Data QnA solution\n",
        "\n",
        "!gcloud services enable \\\n",
        "  cloudapis.googleapis.com \\\n",
        "  compute.googleapis.com \\\n",
        "  iam.googleapis.com \\\n",
        "  run.googleapis.com \\\n",
        "  sqladmin.googleapis.com \\\n",
        "  aiplatform.googleapis.com \\\n",
        "  bigquery.googleapis.com \\"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### **Set Parameters**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Fill out the parameters and configuration settings below. \n",
        "These are the parameters for connecting to the source databases and setting configurations for the vector store tables to be created. \n",
        "Additionally, you can specify whether you have and want to use known-good-queries for the pipeline run and whether you want to enable logging.\n",
        "\n",
        "**Known good queries:** if you have known working user question <-> SQL query pairs, you can put them into the file `scripts/known_good_sql.csv`. This will be used as a caching layer and for in-context learning: If an exact match of the user question is found in the vector store, the pipeline will skip SQL Generation and output the cached SQL query. If the similarity score is between 90-100%, the known good queries will be used as few-shot examples by the SQL Generator Agent. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Data source details\n",
        "DATA_SOURCE = 'bigquery' # Options: 'bigquery' and 'cloudsql-pg' i.e, PostgreSQL database on Google Cloud SQL\n",
        "\n",
        "# Please specify what you would like to use as vector store for embeddings\n",
        "VECTOR_STORE = 'bigquery-vector' # Options: 'bigquery-vector' i.e, Bigquery vector and 'cloudsql-pgvector' i.e, pgvector on PostgreSQL\n",
        "\n",
        "\n",
        "# If you have chosen 'bigquery' as DATA_SOURCE; provide information below\n",
        "BQ_REGION = 'us-central1'\n",
        "BQ_DATASET_NAME = 'imdb'\n",
        "# you can specify the names of the bq tables you want to query over specifially. If left empty, Open Data QnA will parse through all the tables in the dataset.\n",
        "BQ_TABLE_LIST = None # either None or a list of table names in format ['reviews', 'ratings']\n",
        "\n",
        "\n",
        "# Specify if you have example question & known-good-query pairs you want to leverage \n",
        "EXAMPLES = False \n",
        "\n",
        "# If Logging is enabled OR you are using bigquery-vector as the data store, a Big Query dataset will be created to hold the tables. \n",
        "# You can rename the table below if you wish to do so. \n",
        "BQ_OPENDATAQNA_DATASET_NAME = 'opendataqna'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Quick input verifications below:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "# Input verification - Source\n",
        "assert DATA_SOURCE in {'bigquery', 'cloudsql-pg'}, \"⚠️ Invalid DATA_SOURCE. Must be 'bigquery' or 'cloudsql-pg'\"\n",
        "\n",
        "# Input verification - Vector Store\n",
        "assert VECTOR_STORE in {'bigquery-vector', 'cloudsql-pgvector'}, \"⚠️ Invalid VECTOR_STORE. Must be 'bigquery-vector' or 'cloudsql-pgvector'\"\n",
        "\n",
        "if DATA_SOURCE == 'bigquery':\n",
        "    assert BQ_REGION, \"⚠️ Please provide the Data Set Region\"\n",
        "    assert BQ_DATASET_NAME, \"⚠️ Please provide the name of the dataset on Bigquery\"\n",
        "\n",
        "    DATASET_REGION = BQ_REGION\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## **0.2. Implement Agent Classes and Helper Functions**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Big Query Connector\n",
        "\n",
        "The `BQConnector` is a class for interacting with Google BigQuery. It simplifies the process of connecting to BigQuery, retrieving data into Pandas DataFrames, and finding similar or exact matches for queries using vector embeddings. This connector also enables users to validate SQL queries through dry runs before execution, ensuring efficient and cost-effective data analysis. Whether you're exploring datasets, building data pipelines, or integrating BigQuery with other tools, the BQConnector provides a convenient and powerful interface for working with your data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "BigQuery Connector Class\n",
        "\"\"\"\n",
        "from google.cloud import bigquery\n",
        "from google.cloud import bigquery_connection_v1 as bq_connection\n",
        "from abc import ABC\n",
        "from datetime import datetime\n",
        "import google.auth\n",
        "import pandas as pd\n",
        "from google.cloud.exceptions import NotFound\n",
        "\n",
        "def get_auth_user():\n",
        "    credentials, project_id = google.auth.default()\n",
        "\n",
        "    if hasattr(credentials, 'service_account_email'):\n",
        "        return credentials.service_account_email\n",
        "    else:\n",
        "        return \"Not Determined\"\n",
        "\n",
        "def bq_specific_data_types(): \n",
        "    return '''\n",
        "    BigQuery offers a wide variety of datatypes to store different types of data effectively. Here's a breakdown of the available categories:\n",
        "    Numeric Types -\n",
        "    INTEGER (INT64): Stores whole numbers within the range of -9,223,372,036,854,775,808 to 9,223,372,036,854,775,807. Ideal for non-fractional values.\n",
        "    FLOAT (FLOAT64): Stores approximate floating-point numbers with a range of -1.7E+308 to 1.7E+308. Suitable for decimals with a degree of imprecision.\n",
        "    NUMERIC: Stores exact fixed-precision decimal numbers, with up to 38 digits of precision and 9 digits to the right of the decimal point. Useful for precise financial and accounting calculations.\n",
        "    BIGNUMERIC: Similar to NUMERIC but with even larger scale and precision. Designed for extreme precision in calculations.\n",
        "    \n",
        "    Character Types -\n",
        "    STRING: Stores variable-length Unicode character sequences. Enclosed using single, double, or triple quotes.\n",
        "    \n",
        "    Boolean Type -\n",
        "    BOOLEAN: Stores logical values of TRUE or FALSE (case-insensitive).\n",
        "    \n",
        "    Date and Time Types -\n",
        "    DATE: Stores dates without associated time information.\n",
        "    TIME: Stores time information independent of a specific date.\n",
        "    DATETIME: Stores both date and time information (without timezone information).\n",
        "    TIMESTAMP: Stores an exact moment in time with microsecond precision, including a timezone component for global accuracy.\n",
        "    \n",
        "    Other Types\n",
        "    BYTES: Stores variable-length binary data. Distinguished from strings by using 'B' or 'b' prefix in values.\n",
        "    GEOGRAPHY: Stores points, lines, and polygons representing locations on the Earth's surface.\n",
        "    ARRAY: Stores an ordered collection of zero or more elements of the same (non-ARRAY) data type.\n",
        "    STRUCT: Stores an ordered collection of fields, each with its own name and data type (can be nested).\n",
        "    \n",
        "    This list covers the most common datatypes in BigQuery.\n",
        "    '''\n",
        "\n",
        "class BQConnector(ABC):\n",
        "    \"\"\"\n",
        "    Instantiates a BigQuery Connector.\n",
        "    \"\"\"\n",
        "    connectorType: str = \"Base\"\n",
        "\n",
        "    def __init__(self,\n",
        "                 project_id:str,\n",
        "                 region:str,\n",
        "                 dataset_name:str,\n",
        "                 opendataqna_dataset:str):\n",
        "\n",
        "\n",
        "        self.project_id = project_id\n",
        "        self.region = region\n",
        "        self.dataset_name = dataset_name\n",
        "        self.opendataqna_dataset = opendataqna_dataset\n",
        "        self.client=self.getconn()\n",
        "\n",
        "    def getconn(self):\n",
        "        client = bigquery.Client(project=self.project_id)\n",
        "        return client\n",
        "    \n",
        "    def retrieve_df(self,query):\n",
        "        return self.client.query_and_wait(query).to_dataframe()\n",
        "   \n",
        "    \n",
        "    def retrieve_matches(self, mode, schema, qe, similarity_threshold, limit): \n",
        "        \"\"\"\n",
        "        This function retrieves the most similar table_schema and column_schema.\n",
        "        Modes can be either 'table', 'column', or 'example' \n",
        "        \"\"\"\n",
        "        matches = []\n",
        "\n",
        "        if mode == 'table':\n",
        "            sql = '''select base.content as tables_content from vector_search(TABLE `{}.table_details_embeddings`, \"embedding\", \n",
        "            (SELECT {} as qe), top_k=> {},distance_type=>\"COSINE\") where 1-distance > {} '''\n",
        "        \n",
        "        elif mode == 'column':\n",
        "            sql='''select base.content as columns_content from vector_search(TABLE `{}.tablecolumn_details_embeddings`, \"embedding\",\n",
        "            (SELECT {} as qe), top_k=> {}, distance_type=>\"COSINE\") where 1-distance > {} '''\n",
        "\n",
        "        elif mode == 'example': \n",
        "            sql='''select base.example_user_question, base.example_generated_sql from vector_search ( TABLE `{}.example_prompt_sql_embeddings`, \"embedding\",\n",
        "            (select {} as qe), top_k=> {}, distance_type=>\"COSINE\") where 1-distance > {} '''\n",
        "    \n",
        "        else: \n",
        "            ValueError(\"No valid mode. Must be either table, column, or example\")\n",
        "            name_txt = ''\n",
        "\n",
        "\n",
        "        results=self.client.query_and_wait(sql.format('{}.{}'.format(self.project_id,self.opendataqna_dataset),qe,limit,similarity_threshold)).to_dataframe()\n",
        "\n",
        "\n",
        "        # CHECK RESULTS \n",
        "        if len(results) == 0:\n",
        "            print(\"Did not find any results. Adjust the query parameters.\")\n",
        "\n",
        "        if mode == 'table': \n",
        "            name_txt = ''\n",
        "            for _ , r in results.iterrows():\n",
        "                name_txt=name_txt+r[\"tables_content\"]+\"\\n\"\n",
        "\n",
        "        elif mode == 'column': \n",
        "            name_txt = '' \n",
        "            for _ ,r in results.iterrows():\n",
        "                name_txt=name_txt+r[\"columns_content\"]+\"\\n\"\n",
        "\n",
        "        elif mode == 'example': \n",
        "            name_txt = ''\n",
        "            for _ , r in results.iterrows():\n",
        "                example_user_question=r[\"example_user_question\"]\n",
        "                example_sql=r[\"example_generated_sql\"]\n",
        "                name_txt = name_txt + \"\\n Example_question: \"+example_user_question+ \"; Example_SQL: \"+example_sql\n",
        "\n",
        "        else: \n",
        "            ValueError(\"No valid mode. Must be either table, column, or example\")\n",
        "            name_txt = ''\n",
        "\n",
        "        matches.append(name_txt)\n",
        "        \n",
        "\n",
        "        return matches\n",
        "\n",
        "    def getSimilarMatches(self, mode, schema, qe, num_matches, similarity_threshold):\n",
        "\n",
        "        if mode == 'table': \n",
        "            match_result= self.retrieve_matches(mode, schema, qe, similarity_threshold, num_matches)\n",
        "            match_result = match_result[0]\n",
        "            # print(match_result)\n",
        "\n",
        "        elif mode == 'column': \n",
        "            match_result= self.retrieve_matches(mode, schema, qe, similarity_threshold, num_matches)\n",
        "            match_result = match_result[0]\n",
        "        \n",
        "        elif mode == 'example': \n",
        "            match_result= self.retrieve_matches(mode, schema, qe, similarity_threshold, num_matches)\n",
        "            if len(match_result) == 0:\n",
        "                match_result = None\n",
        "            else:\n",
        "                match_result = match_result[0]\n",
        "\n",
        "        return match_result\n",
        "\n",
        "    def getExactMatches(self, query):\n",
        "        \"\"\"Checks if the exact question is already present in the example SQL set\"\"\"\n",
        "        check_history_sql=f\"\"\"SELECT example_user_question,example_generated_sql FROM {self.project_id}.{self.opendataqna_dataset}.example_prompt_sql_embeddings\n",
        "                          WHERE lower(example_user_question) = lower('{query}') LIMIT 1; \"\"\"\n",
        "\n",
        "        exact_sql_history = self.client.query_and_wait(check_history_sql).to_dataframe()\n",
        "\n",
        "\n",
        "        if exact_sql_history[exact_sql_history.columns[0]].count() != 0:\n",
        "            sql_example_txt = ''\n",
        "            exact_sql = ''\n",
        "            for index, row in exact_sql_history.iterrows():\n",
        "                example_user_question=row[\"example_user_question\"]\n",
        "                example_sql=row[\"example_generated_sql\"]\n",
        "                exact_sql=example_sql\n",
        "                sql_example_txt = sql_example_txt + \"\\n Example_question: \"+example_user_question+ \"; Example_SQL: \"+example_sql\n",
        "\n",
        "            print(\"Found a matching question from the history!\" + str(sql_example_txt))\n",
        "            final_sql=exact_sql\n",
        "\n",
        "        else: \n",
        "            print(\"No exact match found for the user prompt\")\n",
        "            final_sql = None\n",
        "\n",
        "        return final_sql\n",
        "\n",
        "    def test_sql_plan_execution(self, generated_sql):\n",
        "        try:\n",
        "\n",
        "            job_config=bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)\n",
        "            query_job = self.client.query(generated_sql,job_config=job_config)\n",
        "            # print(query_job)\n",
        "            exec_result_df=(\"This query will process {} bytes.\".format(query_job.total_bytes_processed))\n",
        "            correct_sql = True\n",
        "            print(exec_result_df)\n",
        "            return correct_sql, exec_result_df\n",
        "        except Exception as e:\n",
        "            return False,str(e)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "connector = BQConnector(project_id,BQ_REGION,BQ_DATASET_NAME,BQ_OPENDATAQNA_DATASET_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### **Initial Step: Specifying Agent Classes**\n",
        "Before moving forward with the standalone notebook, let's first implement the class functions for the agents we are going to use. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Base Class for Agents\n",
        "The `Agent` class is the foundation for building agents with Large Language Models (LLMs) on Vertex AI. It handles model initialization based on the specified model ID and provides a method to generate responses from the chosen LLM. This base class serves as a template for creating specialized agents within the Open Data QnA project, promoting code reusability and a structured approach to working with LLMs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "Provides the base class for all Agents \n",
        "\"\"\"\n",
        "\n",
        "from abc import ABC\n",
        "import vertexai\n",
        "from vertexai.language_models import TextGenerationModel\n",
        "from vertexai.language_models import CodeGenerationModel\n",
        "from vertexai.language_models import CodeChatModel\n",
        "from vertexai.generative_models import GenerativeModel\n",
        "from vertexai.generative_models import HarmCategory,HarmBlockThreshold\n",
        "from vertexai.generative_models import GenerationConfig\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class Agent(ABC):\n",
        "    \"\"\"\n",
        "    The core class for all Agents\n",
        "    \"\"\"\n",
        "\n",
        "    agentType: str = \"Agent\"\n",
        "\n",
        "    def __init__(self,\n",
        "                model_id:str):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            PROJECT_ID (str | None): GCP Project Id.\n",
        "            dataset_name (str): \n",
        "            TODO\n",
        "        \"\"\"\n",
        "\n",
        "        self.model_id = model_id \n",
        "\n",
        "        if model_id == 'code-bison-32k':\n",
        "            self.model = CodeGenerationModel.from_pretrained('code-bison-32k')\n",
        "        elif model_id == 'text-bison-32k':\n",
        "            self.model = TextGenerationModel.from_pretrained('text-bison-32k')\n",
        "        elif model_id == 'codechat-bison-32k':\n",
        "            self.model = CodeChatModel.from_pretrained(\"codechat-bison-32k\")\n",
        "        elif model_id == 'gemini-1.0-pro' or model_id == 'gemini-1.5-pro':\n",
        "            self.model = GenerativeModel(\"gemini-1.0-pro-001\")\n",
        "            self.safety_settings: Optional[dict] = {\n",
        "                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,\n",
        "                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,\n",
        "                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,\n",
        "                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,\n",
        "            }\n",
        "        else:\n",
        "            raise ValueError(\"Please specify a compatible model.\")\n",
        "\n",
        "    def generate_llm_response(self,prompt):\n",
        "        context_query = self.model.generate_content(prompt,safety_settings=self.safety_settings,stream=False)\n",
        "        return str(context_query.candidates[0].text).replace(\"```sql\", \"\").replace(\"```\", \"\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Agent class for generating embeddings\n",
        "The `EmbedderAgent` class is designed to generate text embeddings using Vertex AI's TextEmbeddingModel. It takes a string or a list of strings as input and returns the corresponding embeddings as a list of vectors. The agent's create method handles both single and multiple text inputs, ensuring that the embeddings are correctly retrieved and formatted."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from abc import ABC\n",
        "from vertexai.language_models import TextEmbeddingModel\n",
        "\n",
        "\n",
        "\n",
        "class EmbedderAgent(Agent, ABC): \n",
        "    \"\"\" \n",
        "    This Agent generates embeddings \n",
        "    \"\"\" \n",
        "\n",
        "    agentType: str = \"EmbedderAgent\"\n",
        "\n",
        "    def __init__(self, mode, embeddings_model='textembedding-gecko@002'): \n",
        "        if mode == 'vertex': \n",
        "            self.mode = mode \n",
        "            self.model = TextEmbeddingModel.from_pretrained(embeddings_model)\n",
        "\n",
        "        else: raise ValueError('EmbedderAgent mode must be vertex')\n",
        "\n",
        "\n",
        "\n",
        "    def create(self, question): \n",
        "        \"\"\"Text embedding with a Large Language Model.\"\"\"\n",
        "\n",
        "        if self.mode == 'vertex': \n",
        "            if isinstance(question, str): \n",
        "                embeddings = self.model.get_embeddings([question])\n",
        "                for embedding in embeddings:\n",
        "                    vector = embedding.values\n",
        "                return vector\n",
        "            \n",
        "            elif isinstance(question, list):  \n",
        "                vector = list() \n",
        "                for q in question: \n",
        "                    embeddings = self.model.get_embeddings([q])\n",
        "\n",
        "                    for embedding in embeddings:\n",
        "                        vector.append(embedding.values) \n",
        "                return vector\n",
        "            \n",
        "            else: raise ValueError('Input must be either str or list')\n",
        "        "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Agent class for building SQL\n",
        "\n",
        "The `BuildSQLAgent` class specializes in constructing SQL queries for BigQuery (and PostgreSQL, on the main repo). It leverages an LLM to analyze user questions, table schemas, and column details to generate syntactically and semantically correct BigQuery-compliant SQL queries. It adheres to specific guidelines like minimal table joins, safe casting, and accurate column referencing. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "class BuildSQLAgent(Agent, ABC):\n",
        "    \"\"\"\n",
        "    This Agent produces the SQL query \n",
        "    \"\"\"\n",
        "\n",
        "    agentType: str = \"BuildSQLAgent\"\n",
        "\n",
        "    def build_sql(self,source_type, user_question,tables_schema,tables_detailed_schema, similar_sql, max_output_tokens=2048, temperature=0.4, top_p=1, top_k=32): \n",
        "        context_prompt = f\"\"\"\n",
        "            You are a BigQuery SQL guru. Write a SQL comformant query for Bigquery that answers the following question while using the provided context to correctly refer to the BigQuery tables and the needed column names.\n",
        "\n",
        "            Guidelines:\n",
        "            - Join as minimal tables as possible.\n",
        "            - When joining tables ensure all join columns are the same data_type.\n",
        "            - Analyze the database and the table schema provided as parameters and undestand the relations (column and table relations).\n",
        "            - Use always SAFE_CAST. If performing a SAFE_CAST, use only Bigquery supported datatypes.\n",
        "            - Always SAFE_CAST and then use aggregate functions\n",
        "            - Don't include any comments in code.\n",
        "            - Remove ```sql and ``` from the output and generate the SQL in single line.\n",
        "            - Tables should be refered to using a fully qualified name with enclosed in ticks (`) e.g. `project_id.owner.table_name`.\n",
        "            - Use all the non-aggregated columns from the \"SELECT\" statement while framing \"GROUP BY\" block.\n",
        "            - Return syntactically and symantically correct SQL for BigQuery with proper relation mapping i.e project_id, owner, table and column relation.\n",
        "            - Use ONLY the column names (column_name) mentioned in Table Schema. DO NOT USE any other column names outside of this.\n",
        "            - Associate column_name mentioned in Table Schema only to the table_name specified under Table Schema.\n",
        "            - Use SQL 'AS' statement to assign a new name temporarily to a table column or even a table wherever needed.\n",
        "            - Table names are case sensitive. DO NOT uppercase or lowercase the table names.\n",
        "            - Always enclose subqueries and union queries in brackets.\n",
        "            - Refer to the examples provided i.e. {similar_sql}\n",
        "\n",
        "\n",
        "        Here are some examples of user-question and SQL queries:\n",
        "        {similar_sql}\n",
        "\n",
        "        question:\n",
        "        {user_question}\n",
        "\n",
        "        Table Schema:\n",
        "        {tables_schema}\n",
        "\n",
        "        Column Description:\n",
        "        {tables_detailed_schema}\n",
        "\n",
        "        \"\"\"\n",
        "        \n",
        "\n",
        "        if 'gemini' in self.model_id:\n",
        "            # Generation Config\n",
        "            config = GenerationConfig(\n",
        "                max_output_tokens=max_output_tokens, temperature=temperature, top_p=top_p, top_k=top_k\n",
        "            )\n",
        "\n",
        "            # Generate text\n",
        "            context_query = self.model.generate_content(context_prompt, generation_config=config, stream=False)\n",
        "            generated_sql = str(context_query.candidates[0].text)\n",
        "\n",
        "        else:\n",
        "            context_query = self.model.predict(context_prompt, max_output_tokens = max_output_tokens, temperature=temperature)\n",
        "            generated_sql = str(context_query.candidates[0])\n",
        "\n",
        "\n",
        "        return generated_sql"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Agent class for debugging loop\n",
        "\n",
        "The `DebugSQLAgent` is designed to troubleshoot and refine BigQuery SQL queries. It initiates a chat session with a language model (e.g., Gemini or CodeChat) and iteratively analyzes queries, identifying and addressing errors. If a query is invalid, the agent uses feedback from the model to generate alternative queries that adhere to best practices and answer the original question. The start_debugger method orchestrates this process, returning a potentially refined query along with validity information and an audit trail."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from abc import ABC\n",
        "\n",
        "from vertexai.language_models import CodeChatModel\n",
        "from vertexai.generative_models import GenerativeModel\n",
        "\n",
        "import pandas as pd\n",
        "import json  \n",
        "\n",
        "\n",
        "\n",
        "class DebugSQLAgent(Agent, ABC): \n",
        "    \"\"\" \n",
        "    This Chat Agent runs the debugging loop.\n",
        "    \"\"\" \n",
        "\n",
        "    agentType: str = \"DebugSQLAgent\"\n",
        "\n",
        "    def __init__(self, chat_model_id = 'gemini-1.0-pro'): \n",
        "        self.chat_model_id = chat_model_id\n",
        "        # self.model = CodeChatModel.from_pretrained(\"codechat-bison-32k\")\n",
        "\n",
        "\n",
        "    def init_chat(self, tables_schema,tables_detailed_schema,sql_example=\"-No examples provided..-\"):\n",
        "        context_prompt = f\"\"\"\n",
        "        You are an BigQuery SQL guru. This session is trying to troubleshoot an BigQuery SQL query.  As the user provides versions of the query and the errors returned by BigQuery,\n",
        "        return a new alternative SQL query that fixes the errors. It is important that the query still answer the original question.\n",
        "\n",
        "\n",
        "        Guidelines:\n",
        "        - Join as minimal tables as possible.\n",
        "        - When joining tables ensure all join columns are the same data_type.\n",
        "        - Analyze the database and the table schema provided as parameters and undestand the relations (column and table relations).\n",
        "        - Use always SAFE_CAST. If performing a SAFE_CAST, use only Bigquery supported datatypes.\n",
        "        - Always SAFE_CAST and then use aggregate functions\n",
        "        - Don't include any comments in code.\n",
        "        - Remove ```sql and ``` from the output and generate the SQL in single line.\n",
        "        - Tables should be refered to using a fully qualified name with enclosed in ticks (`) e.g. `project_id.owner.table_name`.\n",
        "        - Use all the non-aggregated columns from the \"SELECT\" statement while framing \"GROUP BY\" block.\n",
        "        - Return syntactically and symantically correct SQL for BigQuery with proper relation mapping i.e project_id, owner, table and column relation.\n",
        "        - Use ONLY the column names (column_name) mentioned in Table Schema. DO NOT USE any other column names outside of this.\n",
        "        - Associate column_name mentioned in Table Schema only to the table_name specified under Table Schema.\n",
        "        - Use SQL 'AS' statement to assign a new name temporarily to a table column or even a table wherever needed.\n",
        "        - Table names are case sensitive. DO NOT uppercase or lowercase the table names.\n",
        "        - Always enclose subqueries and union queries in brackets.\n",
        "        - Refer to the examples provided i.e. {sql_example}\n",
        "\n",
        "        Parameters:\n",
        "        - table metadata: {tables_schema}\n",
        "        - column metadata: {tables_detailed_schema}\n",
        "        - SQL example: {sql_example}\n",
        "\n",
        "        \"\"\"\n",
        "\n",
        "        \n",
        "        if self.chat_model_id == 'codechat-bison-32k':\n",
        "            chat_model = CodeChatModel.from_pretrained(\"codechat-bison-32k\")\n",
        "            chat_session = chat_model.start_chat(context=context_prompt)\n",
        "        elif self.chat_model_id == 'gemini-1.0-pro':\n",
        "            chat_model = GenerativeModel(\"gemini-1.0-pro-001\")\n",
        "            chat_session = chat_model.start_chat(response_validation=False)\n",
        "            chat_session.send_message(context_prompt)\n",
        "        elif self.chat_model_id == 'gemini-ultra':\n",
        "            chat_model = GenerativeModel(\"gemini-1.0-ultra-001\")\n",
        "            chat_session = chat_model.start_chat(response_validation=False)\n",
        "            chat_session.send_message(context_prompt)\n",
        "        else:\n",
        "            raise ValueError('Invalid chat_model_id')\n",
        "        \n",
        "        return chat_session\n",
        "\n",
        "\n",
        "    def rewrite_sql_chat(self, chat_session, question, error_df):\n",
        "\n",
        "\n",
        "        context_prompt = f\"\"\"\n",
        "            What is an alternative SQL statement to address the error mentioned below?\n",
        "            Present a different SQL from previous ones. It is important that the query still answer the original question.\n",
        "            All columns selected must be present on tables mentioned on the join section.\n",
        "            Avoid repeating suggestions.\n",
        "\n",
        "            Original SQL:\n",
        "            {question}\n",
        "\n",
        "            Error:\n",
        "            {error_df}\n",
        "\n",
        "            \"\"\"\n",
        "\n",
        "        if self.chat_model_id =='codechat-bison-32k':\n",
        "            response = chat_session.send_message(context_prompt)\n",
        "            resp_return = (str(response.candidates[0])).replace(\"```sql\", \"\").replace(\"```\", \"\")\n",
        "        elif self.chat_model_id =='gemini-1.0-pro':\n",
        "            response = chat_session.send_message(context_prompt, stream=False)\n",
        "            resp_return = (str(response.text)).replace(\"```sql\", \"\").replace(\"```\", \"\")\n",
        "        elif self.chat_model_id == 'gemini-ultra':\n",
        "            response = chat_session.send_message(context_prompt, stream=False)\n",
        "            resp_return = (str(response.text)).replace(\"```sql\", \"\").replace(\"```\", \"\")\n",
        "        else:\n",
        "            raise ValueError('Invalid chat_model_id')\n",
        "\n",
        "        return resp_return\n",
        "\n",
        "\n",
        "    def start_debugger  (self,\n",
        "                        source_type,\n",
        "                        query,\n",
        "                        user_question, \n",
        "                        SQLChecker,\n",
        "                        tables_schema, \n",
        "                        tables_detailed_schema,\n",
        "                        AUDIT_TEXT, \n",
        "                        similar_sql=\"-No examples provided..-\", \n",
        "                        DEBUGGING_ROUNDS = 2,\n",
        "                        LLM_VALIDATION=True):\n",
        "        i = 0  \n",
        "        STOP = False \n",
        "        invalid_response = False \n",
        "        chat_session = self.init_chat(tables_schema,tables_detailed_schema,similar_sql)\n",
        "        sql = query.replace(\"```sql\",\"\").replace(\"```\",\"\").replace(\"EXPLAIN ANALYZE \",\"\")\n",
        "\n",
        "        AUDIT_TEXT=AUDIT_TEXT+\"\\n\\nEntering the debugging steps!\"\n",
        "        while (not STOP):\n",
        "\n",
        "            # Check if LLM Validation is enabled \n",
        "            if LLM_VALIDATION: \n",
        "                # sql = query.replace(\"```sql\",\"\").replace(\"```\",\"\").replace(\"EXPLAIN ANALYZE \",\"\")\n",
        "                json_syntax_result = SQLChecker.check(user_question,tables_schema,tables_detailed_schema, sql) \n",
        "\n",
        "            else: \n",
        "                json_syntax_result['valid'] = True \n",
        "\n",
        "            if json_syntax_result['valid'] is True:\n",
        "                # Testing SQL Execution\n",
        "                if LLM_VALIDATION: \n",
        "                    AUDIT_TEXT=AUDIT_TEXT+\"\\nGenerated SQL is syntactically correct as per LLM Validation!\"\n",
        "                \n",
        "                else: \n",
        "                    AUDIT_TEXT=AUDIT_TEXT+\"\\nLLM Validation is deactivated. Jumping directly to dry run execution.\"\n",
        "\n",
        "                    \n",
        "                correct_sql, exec_result_df = connector.test_sql_plan_execution(sql)\n",
        "                print(\"exec_result_df:\" + exec_result_df)\n",
        "                if not correct_sql:\n",
        "                        AUDIT_TEXT=AUDIT_TEXT+\"\\nGenerated SQL failed on execution! Here is the feedback from bigquery dryrun/ explain plan:  \\n\" + str(exec_result_df)\n",
        "                        rewrite_result = self.rewrite_sql_chat(chat_session, sql, exec_result_df)\n",
        "                        print('\\n Rewritten and Cleaned SQL: ' + str(rewrite_result))\n",
        "                        AUDIT_TEXT=AUDIT_TEXT+\"\\nRewritten and Cleaned SQL: \\n' + str({rewrite_result})\"\n",
        "                        sql = str(rewrite_result).replace(\"```sql\",\"\").replace(\"```\",\"\").replace(\"EXPLAIN ANALYZE \",\"\")\n",
        "\n",
        "                else: STOP = True\n",
        "            else:\n",
        "                print(f'\\nGenerated qeury failed on syntax check as per LLM Validation!\\nError Message from LLM:  {json_syntax_result} \\nRewriting the query...')\n",
        "                AUDIT_TEXT=AUDIT_TEXT+'\\nGenerated qeury failed on syntax check as per LLM Validation! \\nError Message from LLM:  '+ str(json_syntax_result) + '\\nRewriting the query...'\n",
        "                \n",
        "                syntax_err_df = pd.read_json(json.dumps(json_syntax_result))\n",
        "                rewrite_result=self.rewrite_sql_chat(chat_session, sql, syntax_err_df)\n",
        "                print(rewrite_result)\n",
        "                AUDIT_TEXT=AUDIT_TEXT+'\\n Rewritten SQL: ' + str(rewrite_result)\n",
        "                sql=str(rewrite_result).replace(\"```sql\",\"\").replace(\"```\",\"\").replace(\"EXPLAIN ANALYZE \",\"\")\n",
        "            i+=1\n",
        "            if i > DEBUGGING_ROUNDS:\n",
        "                AUDIT_TEXT=AUDIT_TEXT+ \"Exceeded the number of iterations for correction!\"\n",
        "                AUDIT_TEXT=AUDIT_TEXT+ \"The generated SQL can be invalid!\"\n",
        "                STOP = True\n",
        "                invalid_response=True\n",
        "            # After the while is completed\n",
        "        if i > DEBUGGING_ROUNDS:\n",
        "            invalid_response=True\n",
        "        # print(AUDIT_TEXT)\n",
        "        return sql, invalid_response, AUDIT_TEXT"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Agent class for validation\n",
        "\n",
        " The `ValidateSQLAgent` class validates SQL queries using a language model. It assesses the query's validity against predefined guidelines, including column existence, join conditions, table relationships, and formatting conventions. Given a user question, table schema, and column descriptions, it generates a JSON response indicating the query's validity and any errors found. This aids in ensuring the accuracy and correctness of SQL queries before execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import json \n",
        "from abc import ABC\n",
        "\n",
        "\n",
        "class ValidateSQLAgent(Agent, ABC): \n",
        "    \"\"\" \n",
        "    This Chat Agent checks the SQL for vailidity\n",
        "    \"\"\" \n",
        "\n",
        "    agentType: str = \"ValidateSQLAgent\"\n",
        "\n",
        "\n",
        "    def check(self, user_question, tables_schema, columns_schema, generated_sql):\n",
        "\n",
        "        context_prompt = f\"\"\"\n",
        "\n",
        "            Classify the SQL query: {generated_sql} as valid or invalid?\n",
        "\n",
        "            Guidelines to be valid:\n",
        "            - all column_name in the query must exist in the table_name.\n",
        "            - If a join includes d.country_id and table_alias d is equal to table_name DEPT, then country_id column_name must exist with table_name DEPT in the table column metadata. If not, the sql is invalid\n",
        "            - all join columns must be the same data_type.\n",
        "            - table relationships must be correct.\n",
        "            - Tables should be refered to using a fully qualified name including owner and table name.\n",
        "            - Use table_alias.column_name when referring to columns. Example: dept_id=hr.dept_id\n",
        "            - Capitalize the table names on SQL \"where\" condition.\n",
        "            - Use the columns from the \"SELECT\" statement while framing \"GROUP BY\" block.\n",
        "            - Always refer the column name with rightly mapped table-name as seen in the table schema.\n",
        "            - Must be syntactically and symantically correct SQL with proper relation mapping i.e owner, table and column relation.\n",
        "            - Always the table should be refered as schema.table_name.\n",
        "\n",
        "\n",
        "        Parameters:\n",
        "        - SQL query: {generated_sql}\n",
        "        - table schema: {tables_schema}\n",
        "        - column description: {columns_schema}\n",
        "\n",
        "        Respond using a valid JSON format with two elements valid and errors. Remove ```json and ``` from the output:\n",
        "        {{ \"valid\": true or false, \"errors\":errors }}\n",
        "\n",
        "        Initial user question:\n",
        "        {user_question}\n",
        "\n",
        "\n",
        "        \"\"\"\n",
        "\n",
        "        \n",
        "        if self.model_id =='gemini-1.0-pro':\n",
        "            context_query = self.model.generate_content(context_prompt, stream=False)\n",
        "            generated_sql = str(context_query.candidates[0].text)\n",
        "\n",
        "        else:\n",
        "            context_query = self.model.predict(context_prompt, max_output_tokens = 8000, temperature=0)\n",
        "            generated_sql = str(context_query.candidates[0])\n",
        "\n",
        "\n",
        "        json_syntax_result = json.loads(str(generated_sql).replace(\"```json\",\"\").replace(\"```\",\"\"))\n",
        "\n",
        "        # print('\\n SQL Syntax Validity:' + str(json_syntax_result['valid']))\n",
        "        # print('\\n SQL Syntax Error Description:' +str(json_syntax_result['errors']) + '\\n')\n",
        "        \n",
        "        return json_syntax_result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Agent for generating table and column descriptions\n",
        "\n",
        "The `DescriptionAgent` class automatically generates concise descriptions for BigQuery tables and columns. It uses an LLM to analyze column metadata and table details, producing descriptions that can aid in understanding the data and improving SQL query generation. The agent iterates through dataframes containing table and column information, filling in missing descriptions with LLM-generated content."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from abc import ABC\n",
        "\n",
        "\n",
        "class DescriptionAgent(Agent, ABC): \n",
        "    \"\"\" \n",
        "    Generates table and column descriptions. \n",
        "    \"\"\" \n",
        "\n",
        "    agentType: str = \"DescriptionAgent\"\n",
        "\n",
        "    def generate_llm_response(self,prompt):\n",
        "        context_query = self.model.generate_content(prompt,safety_settings=self.safety_settings,stream=False)\n",
        "        return str(context_query.candidates[0].text).replace(\"```sql\", \"\").replace(\"```\", \"\")\n",
        "\n",
        "\n",
        "    def generate_missing_descriptions(self,source,table_desc_df, column_name_df):\n",
        "        llm_generated=0\n",
        "        for index, row in table_desc_df.iterrows():\n",
        "            if row['table_description'] is None or row['table_description']=='NA':\n",
        "                q=f\"table_name == '{row['table_name']}' and table_schema == '{row['table_schema']}'\"\n",
        "                if source=='bigquery':\n",
        "                    context_prompt = f\"\"\"\n",
        "                        Generate table description short and crisp for the table {row['project_id']}.{row['table_schema']}.{row['table_name']}\n",
        "                        Remember that these desciprtion should help LLMs to help build better SQL for any quries related to this table.\n",
        "                        Parameters:\n",
        "                        - column metadata: {column_name_df.query(q).to_markdown(index = False)}\n",
        "                        - table metadata: {table_desc_df.query(q).to_markdown(index = False)}\n",
        "                        \n",
        "                        DO NOT generate description more than two lines\n",
        "                    \"\"\"\n",
        "\n",
        "\n",
        "                table_desc_df.at[index,'table_description']=self.generate_llm_response(context_prompt)\n",
        "                # print(row['table_description'])\n",
        "                llm_generated=llm_generated+1\n",
        "        print(\"\\nLLM generated \"+ str(llm_generated) + \" Table Descriptions\")\n",
        "        llm_generated = 0\n",
        "\n",
        "        \n",
        "        for index, row in column_name_df.iterrows():\n",
        "            # print(row['column_description'])\n",
        "            if row['column_description'] is None or row['column_description']=='':\n",
        "                q=f\"table_name == '{row['table_name']}' and table_schema == '{row['table_schema']}'\"\n",
        "                if source=='bigquery':\n",
        "                    context_prompt = f\"\"\"\n",
        "                    Generate short and crisp description for the column {row['project_id']}.{row['table_schema']}.{row['table_name']}.{row['column_name']}\n",
        "\n",
        "                    Remember that this description should help LLMs to help generate better SQL for any queries related to these columns.\n",
        "\n",
        "                    Consider the below information to generate a good comment\n",
        "\n",
        "                    Name of the column : {row['column_name']}\n",
        "                    Data type of the column is : {row['data_type']}\n",
        "                    Details of the table of this column are below:\n",
        "                    {table_desc_df.query(q).to_markdown(index=False)}\n",
        "                    Column Contrainst of this column are : {row['column_constraints']}\n",
        "\n",
        "                    DO NOT generate description more than two lines\n",
        "                \"\"\"\n",
        "                \n",
        "\n",
        "                column_name_df.at[index,'column_description']=self.generate_llm_response(prompt=context_prompt)\n",
        "                # print(row['column_description'])\n",
        "                llm_generated=llm_generated+1\n",
        "        print(\"\\nLLM generated \"+ str(llm_generated) + \" Column Descriptions\")\n",
        "        return table_desc_df,column_name_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Response Agent\n",
        "The 'ResponseAgent' is a specialized chat agent that translates structured data from SQL queries into natural language responses. It bridges the gap between technical database results and user-friendly communication by leveraging language models like Gemini Pro. By analyzing both the user's question and the SQL output, it crafts informative and relevant answers, enhancing the user's interaction with complex data. This agent is adaptable, supporting multiple language models with minor adjustments, and is designed to make data more accessible and understandable to users."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import json \n",
        "from abc import ABC\n",
        "\n",
        "class ResponseAgent(Agent, ABC): \n",
        "    \"\"\"\n",
        "    A specialized Chat Agent designed to provide natural language responses to user questions based on SQL query results.\n",
        "\n",
        "    This agent acts as a bridge between structured data returned from SQL queries and the user's natural language input. It leverages a language model (e.g., Gemini Pro or others) to interpret the query results and craft informative, human-readable answers.\n",
        "\n",
        "    Key Features:\n",
        "\n",
        "    * **Natural Language Generation:**  Transforms SQL results into user-friendly responses.\n",
        "    * **Model Flexibility:** Supports multiple language models (currently handles Gemini Pro and others with slight adjustments).\n",
        "    * **Contextual Understanding:**  Incorporates the user's original question and the SQL results to provide accurate and relevant answers. \n",
        "\n",
        "    Attributes:\n",
        "        agentType (str): Identifies this agent as a \"ResponseAgent\".\n",
        "        model_id (str): Indicates the specific language model being used.\n",
        "\n",
        "    Methods:\n",
        "        run(user_question, sql_result):\n",
        "            Generates a natural language response based on the user's question and the SQL results.\n",
        "            \n",
        "    Example:\n",
        "\n",
        "        response_agent = ResponseAgent(model_id='gemini-1.0-pro')\n",
        "        response = response_agent.run(\"How many customers are in California?\", sql_result) \n",
        "        # response might be: \"There are 153 customers in California based on the data.\"\n",
        "    \"\"\"\n",
        "\n",
        "    agentType: str = \"ResponseAgent\"\n",
        "\n",
        "    # TODO: Make the LLM Validator optional\n",
        "    def run(self, user_question, sql_result):\n",
        "\n",
        "        context_prompt = f\"\"\"\n",
        "\n",
        "            You are a Data Assistant that helps to answer users' questions on their data within their databases.\n",
        "            The user has provided the following question in natural language: \"{str(user_question)}\"\n",
        "\n",
        "            The system has returned the following result after running the SQL query: \"{str(sql_result)}\".\n",
        "\n",
        "            Provide a natural sounding response to the user to answer the question with the SQL result provided to you. \n",
        "        \"\"\"\n",
        "\n",
        "        \n",
        "        if self.model_id =='gemini-1.0-pro':\n",
        "            context_query = self.model.generate_content(context_prompt, stream=False)\n",
        "            generated_sql = str(context_query.candidates[0].text)\n",
        "\n",
        "        else:\n",
        "            context_query = self.model.predict(context_prompt, max_output_tokens = 8000, temperature=0)\n",
        "            generated_sql = str(context_query.candidates[0])\n",
        "        \n",
        "        return generated_sql"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Helper function: chunking \n",
        "\n",
        "The `get_embedding_chunked` function efficiently generates embeddings for large amounts of text. It divides the input text (textinput) into smaller batches, processes them in parallel using the EmbedderAgent, and stores the resulting embeddings for each chunk. The function returns a pandas DataFrame containing the original text chunks along with their corresponding embedding vectors, facilitating further analysis or storage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def get_embedding_chunked(textinput, batch_size): \n",
        "    embedder = EmbedderAgent('vertex')\n",
        "\n",
        "    for i in range(0, len(textinput), batch_size):\n",
        "        request = [x[\"content\"] for x in textinput[i : i + batch_size]]\n",
        "        response = embedder.create(request) # Vertex Textmodel Embedder \n",
        "\n",
        "        # Store the retrieved vector embeddings for each chunk back.\n",
        "        for x, e in zip(textinput[i : i + batch_size], response):\n",
        "            x[\"embedding\"] = e\n",
        "\n",
        "    # Store the generated embeddings in a pandas dataframe.\n",
        "    out_df = pd.DataFrame(textinput)\n",
        "    return out_df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# **1. Vector Store Setup** (Run once)\n",
        "---\n",
        "\n",
        "This section walks through the Vector Store Setup needed for running the Open Data QnA application. \n",
        "\n",
        "It covers the following steps: \n",
        "> 1. Configuration: Environment and Databases setup including logging on Bigquery for analytics\n",
        "\n",
        "> 2. Creation of Table, Column and Known Good Query Embeddings in the Vector Store  for Retreival Augmented Generation(RAG)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 📈 **1.1 Set Up your Data Source and Vector Store**\n",
        "\n",
        "This section assumes that a datasource is already set up in your GCP project. \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### ⚙️  **Database Setup for Vector Store: BigQuery-vector**\n",
        "\n",
        "Create dataset on Big Query to store the embeddings tables.\n",
        "If Bigquery is the vector store, the same database is used for logging. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Create a new data set on Bigquery to use as Vector store; the same will be used for logging as well\n",
        "if VECTOR_STORE == 'bigquery-vector':\n",
        "  BQ_OPENDATAQNA_DATASET_NAME = \"opendataqna\" #@param {type:\"string\"} - name of the dataset in Vector Store\n",
        "\n",
        "  from google.cloud import bigquery\n",
        "  import google.api_core \n",
        "  client=bigquery.Client(project=PROJECT_ID)\n",
        "  dataset_ref = f\"{PROJECT_ID}.{BQ_OPENDATAQNA_DATASET_NAME}\"\n",
        "\n",
        "\n",
        "  # Create the dataset if it does not exist already\n",
        "  try:\n",
        "      client.get_dataset(dataset_ref)\n",
        "      print(\"Destination Dataset exists\")\n",
        "  except google.api_core.exceptions.NotFound:\n",
        "      print(\"Cannot find the dataset hence creating.......\")\n",
        "      dataset=bigquery.Dataset(dataset_ref)\n",
        "      dataset.location=DATASET_REGION\n",
        "      client.create_dataset(dataset)\n",
        "      print(str(dataset_ref)+\" is created\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##  **1.2. Create Embeddings in Vector Store for RAG** "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 🖋️ **Create Table and Column Embeddings**\n",
        "\n",
        "In this step, table and column metadata is retreived from the data source and embeddings are generated for both.\n",
        "For this, we first specify helper functions for retrieving table and column schemas. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Helper function to return table schema details: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def return_table_schema_sql(project_id, dataset, table_names=None):\n",
        "    \"\"\"\n",
        "    Returns the SQL query to get table schema info, optionally filtering by specific tables.\n",
        "    \"\"\"\n",
        "    user_dataset = f\"{project_id}.{dataset}\"\n",
        "\n",
        "    table_filter_clause = \"\"\n",
        "    if table_names:\n",
        "        # Extract individual table names from the input string\n",
        "        #table_names = [name.strip() for name in table_names[1:-1].split(\",\")]  # Handle the string as a list\n",
        "        formatted_table_names = [f\"'{name}'\" for name in table_names]\n",
        "        table_filter_clause = f\"\"\"AND TABLE_NAME IN ({', '.join(formatted_table_names)})\"\"\"\n",
        "\n",
        "\n",
        "    table_schema_sql = f\"\"\"\n",
        "    (SELECT\n",
        "        TABLE_CATALOG as project_id, TABLE_SCHEMA as table_schema , TABLE_NAME as table_name,  OPTION_VALUE as table_description,\n",
        "        (SELECT STRING_AGG(column_name, ', ') from `{user_dataset}.INFORMATION_SCHEMA.COLUMNS` where TABLE_NAME= t.TABLE_NAME and TABLE_SCHEMA=t.TABLE_SCHEMA) as table_columns\n",
        "    FROM\n",
        "        `{user_dataset}.INFORMATION_SCHEMA.TABLE_OPTIONS` as t\n",
        "    WHERE\n",
        "        OPTION_NAME = \"description\"\n",
        "        {table_filter_clause}\n",
        "    ORDER BY\n",
        "        project_id, table_schema, table_name)\n",
        "\n",
        "    UNION ALL\n",
        "\n",
        "    (SELECT\n",
        "        TABLE_CATALOG as project_id, TABLE_SCHEMA as table_schema , TABLE_NAME as table_name,  \"NA\" as table_description,\n",
        "        (SELECT STRING_AGG(column_name, ', ') from `{user_dataset}.INFORMATION_SCHEMA.COLUMNS` where TABLE_NAME= t.TABLE_NAME and TABLE_SCHEMA=t.TABLE_SCHEMA) as table_columns\n",
        "    FROM\n",
        "        `{user_dataset}.INFORMATION_SCHEMA.TABLES` as t \n",
        "    WHERE \n",
        "        NOT EXISTS (SELECT 1   FROM\n",
        "        `{user_dataset}.INFORMATION_SCHEMA.TABLE_OPTIONS`  \n",
        "    WHERE\n",
        "        OPTION_NAME = \"description\" AND  TABLE_NAME= t.TABLE_NAME and TABLE_SCHEMA=t.TABLE_SCHEMA)\n",
        "        {table_filter_clause}\n",
        "    ORDER BY\n",
        "        project_id, table_schema, table_name)\n",
        "    \"\"\"\n",
        "    return table_schema_sql\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Helper function to return column schema details: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "def return_column_schema_sql(project_id, dataset, table_names=None):\n",
        "    \"\"\"\n",
        "    Returns the SQL query to get column schema info, optionally filtering by specific tables.\n",
        "    \"\"\"\n",
        "    user_dataset = f\"{project_id}.{dataset}\"\n",
        "    \n",
        "    table_filter_clause = \"\"\n",
        "    if table_names:\n",
        "        # table_names = [name.strip() for name in table_names[1:-1].split(\",\")]  # Handle the string as a list\n",
        "        formatted_table_names = [f\"'{name}'\" for name in table_names]\n",
        "        table_filter_clause = f\"\"\"AND C.TABLE_NAME IN ({', '.join(formatted_table_names)})\"\"\"\n",
        "\n",
        "    column_schema_sql = f\"\"\"\n",
        "    SELECT\n",
        "        C.TABLE_CATALOG as project_id, C.TABLE_SCHEMA as table_schema, C.TABLE_NAME as table_name, C.COLUMN_NAME as column_name,\n",
        "        C.DATA_TYPE as data_type, C.DESCRIPTION as column_description, CASE WHEN T.CONSTRAINT_TYPE=\"PRIMARY KEY\" THEN \"This Column is a Primary Key for this table\" WHEN \n",
        "        T.CONSTRAINT_TYPE = \"FOREIGN_KEY\" THEN \"This column is Foreign Key\" ELSE NULL END as column_constraints\n",
        "    FROM\n",
        "        `{user_dataset}.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS` C \n",
        "    LEFT JOIN \n",
        "        `{user_dataset}.INFORMATION_SCHEMA.TABLE_CONSTRAINTS` T \n",
        "        ON C.TABLE_CATALOG = T.TABLE_CATALOG AND\n",
        "           C.TABLE_SCHEMA = T.TABLE_SCHEMA AND \n",
        "           C.TABLE_NAME = T.TABLE_NAME AND  \n",
        "           T.ENFORCED ='YES'\n",
        "    LEFT JOIN \n",
        "        `{user_dataset}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE` K\n",
        "        ON K.CONSTRAINT_NAME=T.CONSTRAINT_NAME AND C.COLUMN_NAME = K.COLUMN_NAME \n",
        "    WHERE\n",
        "        1=1\n",
        "        {table_filter_clause} \n",
        "    ORDER BY\n",
        "        project_id, table_schema, table_name, column_name;\n",
        "\"\"\"\n",
        "\n",
        "    return column_schema_sql\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Retrieve table and column dataframes with schema details and descriptions: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "client = bigquery.Client(project=project_id)\n",
        "\n",
        "table_schema_sql = return_table_schema_sql(project_id, BQ_DATASET_NAME, BQ_TABLE_LIST)\n",
        "table_desc_df = client.query_and_wait(table_schema_sql).to_dataframe()\n",
        "\n",
        "column_schema_sql = return_column_schema_sql(project_id, BQ_DATASET_NAME, BQ_TABLE_LIST) \n",
        "column_name_df = client.query_and_wait(column_schema_sql).to_dataframe()\n",
        "\n",
        "\n",
        "descriptor = DescriptionAgent('gemini-1.0-pro')\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "#GENERATE MISSING DESCRIPTIONS\n",
        "table_desc_df,column_name_df= descriptor.generate_missing_descriptions('bigquery',table_desc_df,column_name_df)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "table_desc_df['table_description']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "table_desc_df.head(10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "column_name_df.head(10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Function to generate embeddings: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import pandas as pd \n",
        "\n",
        "def retrieve_embeddings(): \n",
        "    \"\"\" Augment all the DB schema blocks to create document for embedding \"\"\"\n",
        "\n",
        "    #TABLE EMBEDDINGS\n",
        "    table_details_chunked = []\n",
        "\n",
        "    for _, row_aug in table_desc_df.iterrows():\n",
        "        cur_project_name =str(row_aug['project_id'])\n",
        "        cur_table_name = str(row_aug['table_name'])\n",
        "        cur_table_schema = str(row_aug['table_schema'])\n",
        "        curr_col_names = str(row_aug['table_columns'])\n",
        "        curr_tbl_desc = str(row_aug['table_description'])\n",
        "\n",
        "\n",
        "        table_detailed_description=f\"\"\"\n",
        "        Full Table Name : {cur_project_name}.{cur_table_schema}.{cur_table_name} |\n",
        "        Table Columns List: [{curr_col_names}] |\n",
        "        Table Description: {curr_tbl_desc} \"\"\"\n",
        "\n",
        "        r = {\"table_schema\": cur_table_schema,\"table_name\": cur_table_name,\"content\": table_detailed_description}\n",
        "        table_details_chunked.append(r)\n",
        "\n",
        "    table_details_embeddings = get_embedding_chunked(table_details_chunked, 10)\n",
        "\n",
        "\n",
        "    ### COLUMN EMBEDDING ###\n",
        "    \"\"\"\n",
        "    This SQL returns a df containing the cols table_schema, table_name, column_name, data_type, column_description, table_description, primary_key, column_constraints\n",
        "    for the schema specified above, e.g. 'retail'\n",
        "    \"\"\"\n",
        "\n",
        "    column_details_chunked = []\n",
        "\n",
        "    for _, row_aug in column_name_df.iterrows():\n",
        "        cur_project_name =str(row_aug['project_id'])\n",
        "        cur_table_name = str(row_aug['table_name'])\n",
        "        cur_table_owner = str(row_aug['table_schema'])\n",
        "        curr_col_name = str(row_aug['table_schema'])+'.'+str(row_aug['table_name'])+'.'+str(row_aug['column_name'])\n",
        "        curr_col_datatype = str(row_aug['data_type'])\n",
        "        curr_col_description = str(row_aug['column_description'])\n",
        "        curr_col_constraints = str(row_aug['column_constraints'])\n",
        "        curr_column_name = str(row_aug['column_name'])\n",
        "\n",
        "\n",
        "        column_detailed_description=f\"\"\"\n",
        "        Column Name: {curr_col_name}|\n",
        "        Full Table Name : {cur_project_name}.{cur_table_schema}.{cur_table_name} |\n",
        "        Data type: {curr_col_datatype}|\n",
        "        Column description: {curr_col_description}|\n",
        "        Column Constraints: {curr_col_constraints} \"\"\"\n",
        "\n",
        "        r = {\"table_schema\": cur_table_owner,\"table_name\": cur_table_name,\"column_name\":curr_column_name, \"content\": column_detailed_description}\n",
        "        column_details_chunked.append(r)\n",
        "\n",
        "    column_details_embeddings = get_embedding_chunked(column_details_chunked, 10)\n",
        "\n",
        "\n",
        "    return table_details_embeddings, column_details_embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Generate embeddings: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Create Table and Column Embeddings\n",
        "table_schema_embeddings, col_schema_embeddings = retrieve_embeddings()\n",
        "\n",
        "\n",
        "print(\"Table and Column embeddings are created\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "table_schema_embeddings.head(10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 💾 **Save the Table and Column Embeddings in the Vector Store**\n",
        "The table and column embeddings created in the above step are save to the Vector Store chosen"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from google.cloud import bigquery\n",
        "\n",
        "async def store_schema_embeddings(table_details_embeddings, \n",
        "                            tablecolumn_details_embeddings, \n",
        "                            project_id,\n",
        "                            schema):\n",
        "    \"\"\" \n",
        "    Store the vectorised table and column details in the DB table.\n",
        "    This code may run for a few minutes.  \n",
        "    \"\"\"\n",
        "         \n",
        "    client=bigquery.Client(project=project_id)\n",
        "\n",
        "    #Store table embeddings\n",
        "    client.query_and_wait(f'''CREATE TABLE IF NOT EXISTS `{project_id}.{schema}.table_details_embeddings` (\n",
        "        source_type string NOT NULL, table_schema string NOT NULL, table_name string NOT NULL, content string, embedding ARRAY<FLOAT64>)''')\n",
        "    #job_config = bigquery.LoadJobConfig(write_disposition=\"WRITE_TRUNCATE\")\n",
        "    table_details_embeddings['source_type']='BigQuery'\n",
        "    for _, row in table_details_embeddings.iterrows():\n",
        "        client.query_and_wait(f'''DELETE FROM `{project_id}.{schema}.table_details_embeddings`\n",
        "                WHERE table_schema= '{row[\"table_schema\"]}' and table_name= '{row[\"table_name\"]}' '''\n",
        "                    )\n",
        "    client.load_table_from_dataframe(table_details_embeddings,f'{project_id}.{schema}.table_details_embeddings')\n",
        "\n",
        "\n",
        "    #Store column embeddings\n",
        "    client.query_and_wait(f'''CREATE TABLE IF NOT EXISTS `{project_id}.{schema}.tablecolumn_details_embeddings` (\n",
        "        source_type string NOT NULL, table_schema string NOT NULL, table_name string NOT NULL, column_name string NOT NULL,\n",
        "        content string, embedding ARRAY<FLOAT64>)''')\n",
        "    #job_config = bigquery.LoadJobConfig(write_disposition=\"WRITE_TRUNCATE\")\n",
        "    tablecolumn_details_embeddings['source_type']='BigQuery'\n",
        "    for _, row in tablecolumn_details_embeddings.iterrows():\n",
        "        client.query_and_wait(f'''DELETE FROM `{project_id}.{schema}.tablecolumn_details_embeddings`\n",
        "                WHERE table_schema= '{row[\"table_schema\"]}' and table_name= '{row[\"table_name\"]}' and column_name= '{row[\"column_name\"]}' '''\n",
        "                    )\n",
        "    client.load_table_from_dataframe(tablecolumn_details_embeddings,f'{project_id}.{schema}.tablecolumn_details_embeddings')\n",
        "\n",
        "    return \"Embeddings are stored successfully\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next cell may take a while depending on the size of your data source. It stores the embeddings back to the vector db."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "await(store_schema_embeddings(table_details_embeddings=table_schema_embeddings, \n",
        "                                tablecolumn_details_embeddings=col_schema_embeddings, \n",
        "                                project_id=project_id,\n",
        "                                schema=BQ_OPENDATAQNA_DATASET_NAME                               \n",
        "                                ))\n",
        "\n",
        "\n",
        "print(\"Table and Column embeddings are saved to vector store\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 🗄️ **Load Known Good SQL into Vector Store**\n",
        "Known Good Queries are used to create query cache for Few shot examples. Creating a query cache is highly recommended for best outcomes! \n",
        "\n",
        "The following cell will load the Natural Language Question and Known Good SQL pairs into our Vector Store. There pairs are loaded from `known_good_sql.csv` file inside scripts folder. If you have your own Question-SQL examples, curate them in .csv file before running the cell below. \n",
        "\n",
        "If no Known Good Queries are available at this time to create query cache, you can use [3_LoadKnownGoodSQL.ipynb](3_LoadKnownGoodSQL.ipynb) to load them later!!\" Empty table for KGQ embedding will be created!\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Format of the Known Good SQL File (known_good_sql.csv)\n",
        "\n",
        "prompt | sql | database_name [3 columns]\n",
        "\n",
        "prompt ==> User Question \n",
        "\n",
        "sql ==> SQL for the user question (Note that the sql should enclosed in quotes and only in single line. Please remove the line  break)\n",
        "\n",
        "database_name ==>This name should exactly  match the SCHEMA   NAME for Postgres Source or BQ_DATASET_NAME"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import pandas as pd\n",
        "from google.cloud import bigquery\n",
        "\n",
        "\n",
        "embedder = EmbedderAgent('vertex')\n",
        "\n",
        "\n",
        "async def setup_kgq_table( project_id,\n",
        "                            schema):\n",
        "    \"\"\" \n",
        "    This function sets up or refreshes the Vector Store for Known Good Queries (KGQ)\n",
        "    \"\"\"\n",
        "\n",
        "    # Create BQ Client\n",
        "    client=bigquery.Client(project=project_id)\n",
        "\n",
        "    # Delete an old table\n",
        "    client.query_and_wait(f'''DROP TABLE IF EXISTS `{project_id}.{schema}.example_prompt_sql_embeddings`''')\n",
        "    # Create a new emptry table\n",
        "    client.query_and_wait(f'''CREATE TABLE IF NOT EXISTS `{project_id}.{schema}.example_prompt_sql_embeddings` (\n",
        "                            table_schema string NOT NULL, example_user_question string NOT NULL, example_generated_sql string NOT NULL,\n",
        "                            embedding ARRAY<FLOAT64>)''')\n",
        "        \n",
        "\n",
        "\n",
        "async def store_kgq_embeddings(df_kgq, \n",
        "                            project_id,\n",
        "                            schema\n",
        "                            ):\n",
        "    \"\"\" \n",
        "    Create and save the Known Good Query Embeddings to Vector Store  \n",
        "    \"\"\"\n",
        "\n",
        "    client=bigquery.Client(project=project_id)\n",
        "    \n",
        "    example_sql_details_chunked = []\n",
        "\n",
        "    for _, row_aug in df_kgq.iterrows():\n",
        "\n",
        "        example_user_question = str(row_aug['prompt'])\n",
        "        example_generated_sql = str(row_aug['sql'])\n",
        "        example_database_name = str(row_aug['database_name'])\n",
        "        emb =  embedder.create(example_user_question)\n",
        "        \n",
        "\n",
        "        r = {\"example_database_name\":example_database_name,\"example_user_question\": example_user_question,\"example_generated_sql\": example_generated_sql,\"embedding\": emb}\n",
        "        example_sql_details_chunked.append(r)\n",
        "\n",
        "    example_prompt_sql_embeddings = pd.DataFrame(example_sql_details_chunked)\n",
        "\n",
        "    for _, row in example_prompt_sql_embeddings.iterrows():\n",
        "            client.query_and_wait(f'''DELETE FROM `{project_id}.{schema}.example_prompt_sql_embeddings`\n",
        "                        WHERE table_schema= '{row[\"example_database_name\"]}' and example_user_question= '{row[\"example_user_question\"]}' '''\n",
        "                            )\n",
        "                # embedding=np.array(row[\"embedding\"])\n",
        "            cleaned_sql = row[\"example_generated_sql\"].replace(\"\\n\", \" \")\n",
        "            client.query_and_wait(f'''INSERT INTO `{project_id}.{schema}.example_prompt_sql_embeddings` \n",
        "                VALUES (\"{row[\"example_database_name\"]}\",\"{row[\"example_user_question\"]}\" , \n",
        "                \"{cleaned_sql}\",{row[\"embedding\"]} )''')\n",
        "                    \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next cell stores the kgq to the vector db:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "if EXAMPLES:\n",
        "    print(\"Examples are provided, creating KGQ embeddings and saving them to Vector store.....\")\n",
        "    \n",
        "    current_dir = os.getcwd()\n",
        "    root_dir = os.path.expanduser('~')  # Start at the user's home directory\n",
        "\n",
        "    while current_dir != root_dir:\n",
        "        for dirpath, dirnames, filenames in os.walk(current_dir):\n",
        "            config_path = os.path.join(dirpath, 'known_good_sql.csv')\n",
        "            if os.path.exists(config_path):\n",
        "                file_path = config_path  # Update root_dir to the found directory\n",
        "                break  # Stop outer loop once found\n",
        "\n",
        "        current_dir = os.path.dirname(current_dir)\n",
        "\n",
        "    print(\"Known Good SQL Found at Path :: \"+file_path)\n",
        "\n",
        "    # Load the file\n",
        "    df_kgq = pd.read_csv(file_path)\n",
        "    df_kgq = df_kgq.loc[:, [\"prompt\", \"sql\", \"database_name\"]]\n",
        "    df_kgq = df_kgq.dropna()\n",
        "\n",
        "    # Add KGQ to the vector store\n",
        "    await(store_kgq_embeddings(df_kgq,\n",
        "                                project_id=project_id,\n",
        "                                schema=BQ_OPENDATAQNA_DATASET_NAME\n",
        "                                ))\n",
        "\n",
        "    print('Done!!')\n",
        "\n",
        "else:\n",
        "    print(\"⚠️ WARNING: No Known Good Queries are provided to create query cache for Few shot examples!\")\n",
        "    print(\"Creating a query cache is highly recommended for best outcomes\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 🥁 If all the above steps are executed suucessfully, the following should be set up:\n",
        "\n",
        "* GCP project and all the required IAM permissions\n",
        "\n",
        "* Environment to run the solution\n",
        "\n",
        "* Data source and Vector store for the solution"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "__________________________________________________________________________________________________________________"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# **2. Run the Open Data QnA Pipeline**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 🔗 **3A. Connect to Datasource, Vector Source and Vertex AI**\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "from google.cloud import aiplatform\n",
        "import vertexai\n",
        "\n",
        "found_in_vector = 'N'\n",
        "final_sql='Not Generated Yet'\n",
        "\n",
        "vertexai.init(project=project_id, location=BQ_REGION)\n",
        "aiplatform.init(project=project_id, location=BQ_REGION)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "BQ_DATASET_NAME = project_id+'.'+BQ_DATASET_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "###  ❓ **Ask your Natural Language Question**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "print(\"\\033[1mData Source:- \"+ DATA_SOURCE)\n",
        "print(\"Vector Store:- \"+ VECTOR_STORE)\n",
        "    \n",
        "# Suggested question for 'fda_food' dataset: \"What are the top 5 cities with highest recalls?\"\n",
        "# Suggested question for 'google_dei' dataset: \"How many asian men were part of the leadership workforce in 2021?\"\n",
        "\n",
        "# user_question = input(prompt_for_question) #Uncomment if you want to ask question yourself\n",
        "user_question = 'How many movies have a rating higher than four?' # Or Enter Question here\n",
        "\n",
        "print(\"User Question:- \"+user_question)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 🏃 **Run the Pipeline**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Fetch the USER_DATABASE based on data source\n",
        "\n",
        "call_await = False\n",
        "\n",
        "num_table_matches = 5\n",
        "num_column_matches = 10\n",
        "table_similarity_threshold = 0.3\n",
        "column_similarity_threshold = 0.3 \n",
        "example_similarity_threshold = 0.3 \n",
        "num_sql_matches=3\n",
        "\n",
        "DEBUGGING_ROUNDS = 2\n",
        "RUN_DEBUGGER = True \n",
        "LLM_VALIDATION=True\n",
        "EXECUTE_FINAL_SQL = True \n",
        "\n",
        "\n",
        "# Fetch the embedding of the user's input question \n",
        "embedded_question = embedder.create(user_question)\n",
        "\n",
        "USER_DATABASE = BQ_OPENDATAQNA_DATASET_NAME\n",
        "\n",
        "SQLBuilder = BuildSQLAgent('gemini-1.0-pro')\n",
        "SQLDebugger = DebugSQLAgent('gemini-1.0-pro')\n",
        "SQLChecker = ValidateSQLAgent('gemini-1.0-pro')\n",
        "\n",
        "# Reset AUDIT_TEXT\n",
        "AUDIT_TEXT = ''\n",
        "\n",
        "AUDIT_TEXT = AUDIT_TEXT + \"\\nUser Question : \" + str(user_question) + \"\\nUser Database : \" + str(USER_DATABASE)\n",
        "process_step = \"\\n\\nGet Exact Match: \"\n",
        "\n",
        "# Look for exact matches in known questions IF kgq is enabled \n",
        "if EXAMPLES: \n",
        "    exact_sql_history = connector.getExactMatches(user_question) \n",
        "\n",
        "else: exact_sql_history = None \n",
        "\n",
        "# If exact user query has been found, retrieve the SQL and skip Generation Pipeline \n",
        "if exact_sql_history is not None:\n",
        "    found_in_vector = 'Y' \n",
        "    final_sql = exact_sql_history\n",
        "    invalid_response = False\n",
        "    AUDIT_TEXT = AUDIT_TEXT + \"\\nExact match has been found! Going to retreive the SQL query from cache and serve!\"\n",
        "\n",
        "\n",
        "else:\n",
        "    # No exact match found. Proceed looking for similar entries in db IF kgq is enabled \n",
        "    if EXAMPLES: \n",
        "        AUDIT_TEXT = AUDIT_TEXT +  process_step + \"\\nNo exact match found in query cache, retreiving revelant schema and known good queries for few shot examples using similarity search....\"\n",
        "        process_step = \"\\n\\nGet Similar Match: \"\n",
        "        if call_await:\n",
        "            similar_sql = await connector.getSimilarMatches('example', USER_DATABASE, embedded_question, num_sql_matches, example_similarity_threshold)\n",
        "        else:\n",
        "            similar_sql = connector.getSimilarMatches('example', USER_DATABASE, embedded_question, num_sql_matches, example_similarity_threshold)\n",
        "\n",
        "    else: similar_sql = \"No similar SQLs provided...\"\n",
        "\n",
        "    process_step = \"\\n\\nGet Table and Column Schema: \"\n",
        "    # Retrieve matching tables and columns\n",
        "    if call_await: \n",
        "        table_matches =  await connector.getSimilarMatches('table', USER_DATABASE, embedded_question, num_table_matches, table_similarity_threshold)\n",
        "        column_matches =  await connector.getSimilarMatches('column', USER_DATABASE, embedded_question, num_column_matches, column_similarity_threshold)\n",
        "    else:\n",
        "        table_matches =  connector.getSimilarMatches('table', USER_DATABASE, embedded_question, num_table_matches, table_similarity_threshold)\n",
        "        column_matches =  connector.getSimilarMatches('column', USER_DATABASE, embedded_question, num_column_matches, column_similarity_threshold)\n",
        "\n",
        "    AUDIT_TEXT = AUDIT_TEXT +  process_step + \"\\nRetrieved Similar Known Good Queries, Table Schema and Column Schema: \\n\" + '\\nRetrieved Tables: \\n' + str(table_matches) + '\\n\\nRetrieved Columns: \\n' + str(column_matches) + '\\n\\nRetrieved Known Good Queries: \\n' + str(similar_sql)\n",
        "    \n",
        "    \n",
        "    # If similar table and column schemas found: \n",
        "    if len(table_matches.replace('Schema(values):','').replace(' ','')) > 0 or len(column_matches.replace('Column name(type):','').replace(' ','')) > 0 :\n",
        "\n",
        "        # GENERATE SQL\n",
        "        process_step = \"\\n\\nBuild SQL: \"\n",
        "        generated_sql = SQLBuilder.build_sql(DATA_SOURCE,user_question,table_matches,column_matches,similar_sql)\n",
        "        final_sql=generated_sql\n",
        "        AUDIT_TEXT = AUDIT_TEXT + process_step +  \"\\nGenerated SQL : \" + str(generated_sql)\n",
        "        \n",
        "        if 'unrelated_answer' in generated_sql :\n",
        "            invalid_response=True\n",
        "\n",
        "        # If agent assessment is valid, proceed with checks  \n",
        "        else:\n",
        "            invalid_response=False\n",
        "\n",
        "            if RUN_DEBUGGER: \n",
        "                generated_sql, invalid_response, AUDIT_TEXT = SQLDebugger.start_debugger(DATA_SOURCE, generated_sql, user_question, SQLChecker, table_matches, column_matches, AUDIT_TEXT, similar_sql, DEBUGGING_ROUNDS, LLM_VALIDATION) \n",
        "                # AUDIT_TEXT = AUDIT_TEXT + '\\n Feedback from Debugger: \\n' + feedback_text\n",
        "\n",
        "            final_sql=generated_sql\n",
        "            AUDIT_TEXT = AUDIT_TEXT + \"\\nFinal SQL after Debugger: \\n\" +str(final_sql)\n",
        "\n",
        "\n",
        "    # No matching table found \n",
        "    else:\n",
        "        invalid_response=True\n",
        "        print('No tables found in Vector ...')\n",
        "        AUDIT_TEXT = AUDIT_TEXT + \"\\nNo tables have been found in the Vector DB. The question cannot be answered with the provide data source!\"\n",
        "\n",
        "print(f'\\n\\n AUDIT_TEXT: \\n {AUDIT_TEXT}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "final_sql"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "invalid_response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Run against db "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "Responder = ResponseAgent('gemini-1.0-pro')\n",
        "\n",
        "\n",
        "if not invalid_response:\n",
        "    try: \n",
        "        if EXECUTE_FINAL_SQL is True:\n",
        "                final_exec_result_df=connector.retrieve_df(final_sql.replace(\"```sql\",\"\").replace(\"```\",\"\").replace(\"EXPLAIN ANALYZE \",\"\"))\n",
        "                print('\\nQuestion: ' + user_question + '\\n')\n",
        "                # print('\\n Final SQL Execution Result: \\n')\n",
        "                # print(final_exec_result_df)\n",
        "                response = final_exec_result_df\n",
        "                _resp=Responder.run(user_question, response)\n",
        "                AUDIT_TEXT = AUDIT_TEXT + \"\\nModel says \" + str(_resp) \n",
        "\n",
        "\n",
        "        else:  # Do not execute final SQL\n",
        "                print(\"Not executing final SQL since EXECUTE_FINAL_SQL variable is False\\n \")\n",
        "                response = \"Please enable the Execution of the final SQL so I can provide an answer\"\n",
        "                _resp=Responder.run(user_question, response)\n",
        "                AUDIT_TEXT = AUDIT_TEXT + \"\\nModel says \" + str(_resp) \n",
        "\n",
        "    except ValueError: \n",
        "          print('')\n",
        "    # except Exception as e: \n",
        "    #     print(f\"An error occured. Aborting... Error Message: {e}\")\n",
        "        \n",
        "else:  # Do not execute final SQL\n",
        "    print(\"Not executing final SQL as it is invalid, please debug!\")\n",
        "    response = \"I am sorry, I could not come up with a valid SQL.\"\n",
        "    _resp=Responder.run(user_question, response)\n",
        "    AUDIT_TEXT = AUDIT_TEXT + \"\\nModel says \" + str(_resp)\n",
        "\n",
        "print(\"Final Answer:\" + str(_resp))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 🗑 **Clean Up Notebook Resources**\n",
        "Make sure to delete your Cloud SQL instance and BigQuery Datasets when your are finished with this notebook to avoid further costs. 💸 💰\n",
        "\n",
        "Uncomment and run the cell below to delete "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# #delete BigQuery Dataset using bq utility\n",
        "# !bq rm -r -f -d {BQ_DATASET_NAME}\n",
        "\n",
        "# #delete BigQuery 'Open Data QnA' Vector Store Dataset using bq utility\n",
        "# !bq rm -r -f -d {BQ_OPENDATAQNA_DATASET_NAME}\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}