GenAI_Capstone_NL_to_SQL.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 131
    },
    "id": "z2XwfWrcyB-8",
    "outputId": "f2976008-ee10-4c2e-c844-a31442f8b2fd"
   },
   "outputs": [],
   "source": [
    "### Install Google Cloud Vertex AI\n",
    "# !pip install google-cloud-aiplatform\n",
    "# # Install langchain\n",
    "# ! pip install langchain\n",
    "# # Install pandas\n",
    "# ! pip install pandas\n",
    "# # Install Python client for Google Search API\n",
    "# ! pip install google-api-python-client\n",
    "# # Install google-cloud-language\n",
    "# ! pip install google-cloud-language\n",
    "# # PDF loader\n",
    "# ! pip install pypdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "bWYjF2ANyMqw"
   },
   "outputs": [],
   "source": [
    "# from google.colab import auth as google_auth\n",
    "# google_auth.authenticate_user()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "mFHx3EjAyNVr"
   },
   "outputs": [],
   "source": [
    "import vertexai\n",
    "from vertexai.preview.language_models import CodeGenerationModel\n",
    "\n",
    "WIDTH=600 # @param\n",
    "PROJECT_ID=\"my-project-0004-346516\" # @param\n",
    "TEXT_MODEL=\"code-bison@001\"#'text-bison@001' # @param "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "ktlt9C44yX9V"
   },
   "outputs": [],
   "source": [
    "def nl_to_sql(\n",
    "    prefix_input: str,\n",
    "    ) :\n",
    "    \"\"\"Codey Model Prediction.\"\"\"\n",
    "    vertexai.init(project=PROJECT_ID, location=\"us-central1\")\n",
    "    parameters = {\n",
    "        \"temperature\": 0.5,\n",
    "        \"max_output_tokens\": 256\n",
    "    }\n",
    "    model = CodeGenerationModel.from_pretrained(TEXT_MODEL)\n",
    "    response = model.predict(\n",
    "        prefix = prefix_input,\n",
    "        **parameters\n",
    "    )\n",
    "    print(f\"Question: {QUESTION}\")\n",
    "    print(f\"Response from Model: {response.text}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "BGPl5g1HyXx4"
   },
   "outputs": [],
   "source": [
    "def build_prefix(\n",
    "    question_input: str,\n",
    "    ) :\n",
    "\n",
    "    \"\"\"Open prefix file and append question input from user.\"\"\"\n",
    "    file1 = open(\"prefix_file.txt\")\n",
    "    mystuff = file1.read()\n",
    "    mystuff += \"Question: {}\\nSQL result: \".format(question_input)\n",
    "    nl_to_sql(mystuff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "QeofS7XdyXqN"
   },
   "outputs": [],
   "source": [
    "_prefix_file = 'prefix_file.txt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "gUyzrUwryeqJ"
   },
   "outputs": [],
   "source": [
    "QUESTION = \"Top 3 user id with highest number of orders from order_line_items\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "htQ8V5v6ycvl",
    "outputId": "a2657992-25e2-4ebf-aaa1-741be907f8b9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting prefix_file.txt\n"
     ]
    }
   ],
   "source": [
    "%%writefile {_prefix_file}\n",
    "You are a database engineer that has profound SQL skills. You are using the following schema as reference to the database. Combine tables using JOIN if needed.\n",
    "\n",
    "Schema =\n",
    "[\n",
    "    {\n",
    "        \"table_name\": \"surveys\",\n",
    "        \"column_names\": [\n",
    "            primary_order_id,\n",
    "            survey_type,\n",
    "            survey_responded\n",
    "         ]\n",
    "     }, {\n",
    "        \"table_name\": \"order_line_items\",\n",
    "        \"column_names\": [\n",
    "            source_system,\n",
    "            primary_order_id,\n",
    "            order_id,\n",
    "            user_id,\n",
    "            order_line_item_quantity,\n",
    "            sale_price\n",
    "         ]\n",
    "     }, {\n",
    "        \"table_name\": \"orders\",\n",
    "        \"column_names\": [\n",
    "            sent_to_warehouse,\n",
    "            received_by_warehouse,\n",
    "            split_shipment_order_type\n",
    "         ]\n",
    "     }, {\n",
    "        \"table_name\": \"monthly_subscription_order_retention\",\n",
    "        \"column_names\": [\n",
    "            next_subscription_order_churn_type,\n",
    "            sub_count\n",
    "         ]\n",
    "     },  {\n",
    "         \"table_name\":\"user_clusters\",\n",
    "         \"column_names\": [\n",
    "            user_id,\n",
    "            cluster_version,\n",
    "            user_added_to_cluster_date,\n",
    "            bd_id\n",
    "         ]\n",
    "      }, {\n",
    "         \"table_name\":\"order_delivery_performance\",\n",
    "         \"column_names\": [\n",
    "            order_id,\n",
    "            shipment_id,\n",
    "            tracking_number,\n",
    "            ship2d_busines\n",
    "         ]\n",
    "      }\n",
    "]\n",
    "\n",
    "Here's a simple training example for you to learn,\n",
    "The question is \"Tell me the most used coupon code for every coupon type.\"\n",
    "And the sql query would be something like \"SELECT coupon_type, coupon_code, SUM(coupon_redemptions) AS total_redemptions\n",
    "FROM bark-dataengineering-poc.data_mart.promos\n",
    "GROUP BY coupon_type, coupon_code\n",
    "ORDER BY total_redemptions DESC;\"\n",
    "\n",
    "Give me the sql query result of the following question.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JZhdlRxZyizu",
    "outputId": "79e6d900-ee4f-411f-f74f-482a91330605"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: Top 3 user id with highest number of orders from order_line_items\n",
      "Response from Model: SELECT user_id, COUNT(1) AS order_count\n",
      "FROM order_line_items\n",
      "GROUP BY user_id\n",
      "ORDER BY order_count DESC\n",
      "LIMIT 3;\n"
     ]
    }
   ],
   "source": [
    "build_prefix(QUESTION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "environment": {
   "kernel": "py311",
   "name": "pytorch-gpu.1-12.m98",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/pytorch-gpu.1-12:m98"
  },
  "kernelspec": {
   "display_name": "py311",
   "language": "python",
   "name": "py311"
  },
  "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}