textToSQL.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2d62a9fa-bb82-4ef3-ab45-e2d0353bdb53",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: langchain-experimental in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (0.0.59)\n",
      "Requirement already satisfied: langchain-community<0.3,>=0.2 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-experimental) (0.2.1)\n",
      "Requirement already satisfied: langchain-core<0.3,>=0.2 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-experimental) (0.2.3)\n",
      "Requirement already satisfied: PyYAML>=5.3 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-community<0.3,>=0.2->langchain-experimental) (6.0.1)\n",
      "Requirement already satisfied: SQLAlchemy<3,>=1.4 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-community<0.3,>=0.2->langchain-experimental) (2.0.30)\n",
      "Requirement already satisfied: aiohttp<4.0.0,>=3.8.3 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-community<0.3,>=0.2->langchain-experimental) (3.9.5)\n",
      "Requirement already satisfied: dataclasses-json<0.7,>=0.5.7 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-community<0.3,>=0.2->langchain-experimental) (0.6.6)\n",
      "Requirement already satisfied: langchain<0.3.0,>=0.2.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-community<0.3,>=0.2->langchain-experimental) (0.2.1)\n",
      "Requirement already satisfied: langsmith<0.2.0,>=0.1.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-community<0.3,>=0.2->langchain-experimental) (0.1.67)\n",
      "Requirement already satisfied: numpy<2,>=1 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-community<0.3,>=0.2->langchain-experimental) (1.26.4)\n",
      "Requirement already satisfied: requests<3,>=2 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-community<0.3,>=0.2->langchain-experimental) (2.32.3)\n",
      "Requirement already satisfied: tenacity<9.0.0,>=8.1.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-community<0.3,>=0.2->langchain-experimental) (8.3.0)\n",
      "Requirement already satisfied: jsonpatch<2.0,>=1.33 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-core<0.3,>=0.2->langchain-experimental) (1.33)\n",
      "Requirement already satisfied: packaging<24.0,>=23.2 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-core<0.3,>=0.2->langchain-experimental) (23.2)\n",
      "Requirement already satisfied: pydantic<3,>=1 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain-core<0.3,>=0.2->langchain-experimental) (2.7.2)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain-community<0.3,>=0.2->langchain-experimental) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain-community<0.3,>=0.2->langchain-experimental) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain-community<0.3,>=0.2->langchain-experimental) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain-community<0.3,>=0.2->langchain-experimental) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain-community<0.3,>=0.2->langchain-experimental) (1.9.4)\n",
      "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from dataclasses-json<0.7,>=0.5.7->langchain-community<0.3,>=0.2->langchain-experimental) (3.21.2)\n",
      "Requirement already satisfied: typing-inspect<1,>=0.4.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from dataclasses-json<0.7,>=0.5.7->langchain-community<0.3,>=0.2->langchain-experimental) (0.9.0)\n",
      "Requirement already satisfied: jsonpointer>=1.9 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from jsonpatch<2.0,>=1.33->langchain-core<0.3,>=0.2->langchain-experimental) (2.4)\n",
      "Requirement already satisfied: langchain-text-splitters<0.3.0,>=0.2.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langchain<0.3.0,>=0.2.0->langchain-community<0.3,>=0.2->langchain-experimental) (0.2.0)\n",
      "Requirement already satisfied: orjson<4.0.0,>=3.9.14 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from langsmith<0.2.0,>=0.1.0->langchain-community<0.3,>=0.2->langchain-experimental) (3.10.3)\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from pydantic<3,>=1->langchain-core<0.3,>=0.2->langchain-experimental) (0.7.0)\n",
      "Requirement already satisfied: pydantic-core==2.18.3 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from pydantic<3,>=1->langchain-core<0.3,>=0.2->langchain-experimental) (2.18.3)\n",
      "Requirement already satisfied: typing-extensions>=4.6.1 in /Users/abhijit/Library/Python/3.11/lib/python/site-packages (from pydantic<3,>=1->langchain-core<0.3,>=0.2->langchain-experimental) (4.12.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from requests<3,>=2->langchain-community<0.3,>=0.2->langchain-experimental) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from requests<3,>=2->langchain-community<0.3,>=0.2->langchain-experimental) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from requests<3,>=2->langchain-community<0.3,>=0.2->langchain-experimental) (2.2.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from requests<3,>=2->langchain-community<0.3,>=0.2->langchain-experimental) (2024.2.2)\n",
      "Requirement already satisfied: mypy-extensions>=0.3.0 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from typing-inspect<1,>=0.4.0->dataclasses-json<0.7,>=0.5.7->langchain-community<0.3,>=0.2->langchain-experimental) (1.0.0)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.1.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "pip install langchain-experimental"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "53e63cec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: pymysql in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (1.1.1)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.1.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "pip install pymysql"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "04851b74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: cryptography in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (42.0.7)\n",
      "Requirement already satisfied: cffi>=1.12 in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from cryptography) (1.16.0)\n",
      "Requirement already satisfied: pycparser in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from cffi>=1.12->cryptography) (2.22)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.1.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "pip install cryptography"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8a7c5811",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: faiss-cpu in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (1.8.0)\n",
      "Requirement already satisfied: numpy in /Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages (from faiss-cpu) (1.26.4)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.1.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "pip install -U faiss-cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c5319ac2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: psutil in /Users/abhijit/Library/Python/3.11/lib/python/site-packages (5.9.8)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.1.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "pip install psutil"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "809a31c3",
   "metadata": {},
   "source": [
    "# Import statements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "e6841b54-66c4-47ec-be6b-3d1f0948c206",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.llms import Ollama\n",
    "import pymysql\n",
    "from langchain import SQLDatabase\n",
    "from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate\n",
    "from langchain.chains import create_sql_query_chain\n",
    "from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool\n",
    "import langchain\n",
    "from langchain_core.output_parsers import StrOutputParser, JsonOutputParser\n",
    "from langchain_community.embeddings import OllamaEmbeddings\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_core.example_selectors import SemanticSimilarityExampleSelector\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_core.messages import AIMessage\n",
    "from langchain_core.runnables import Runnable\n",
    "from pprint import pprint\n",
    "import psutil\n",
    "from enum import IntEnum\n",
    "\n",
    "# Debug flags\n",
    "langchain.verbose = False\n",
    "langchain.debug = False\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4223917e",
   "metadata": {},
   "source": [
    "# Connect to locally running LLM by using Ollama"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "c8078607-842f-4471-b2e8-6bcbca058484",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(' Hello! How can I help you today? Is there something specific you would like '\n",
      " \"to know or discuss? I'm here to answer any questions you might have.\\n\"\n",
      " '\\n'\n",
      " \"If you don't have a particular question in mind, feel free to ask about \"\n",
      " 'anything that interests you, such as technology, science, history, books, '\n",
      " \"movies, music, or even just chat about your day. I'm looking forward to our \"\n",
      " 'conversation!')\n"
     ]
    }
   ],
   "source": [
    "# This assumes that ollama is running on localhost at port 11434 and has \"mistral:instruct\" downloaded\n",
    "llm = Ollama(model=\"mistral:instruct\", num_thread = (psutil.cpu_count() - 1), keep_alive = -1, temperature=0, num_ctx=4098)\n",
    "pprint(llm.invoke(\"hi\")) # test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "120dc08b",
   "metadata": {},
   "source": [
    "# Connect to MySQL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "249752d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(\"[(1, 'Toyota', 'Camry', 2018, 'Red'), (2, 'Honda', 'Civic', 2019, 'Blue'), \"\n",
      " \"(3, 'Ford', 'Mustang', 2020, 'Black'), (4, 'BMW', '3 Series', 2017, 'Gray'), \"\n",
      " \"(5, 'Nissan', 'Altima', 2021, 'White'), (6, 'Toyota', 'Corolla', 2021, \"\n",
      " \"'Silver'), (7, 'Honda', 'CR-V', 2020, 'Blue'), (8, 'Ford', 'Escape', 2019, \"\n",
      " \"'Red'), (9, 'BMW', 'X5', 2018, 'Black'), (10, 'Nissan', 'Murano', 2021, \"\n",
      " \"'Gray')]\")\n",
      "(\"[(11, 1, 4, Decimal('2000.10')), (12, 2, 10, Decimal('4000.12')), (13, 3, 2, \"\n",
      " \"Decimal('6025.56')), (14, 4, 50, Decimal('3000.33')), (15, 5, 28, \"\n",
      " \"Decimal('6000.89')), (16, 6, 12, Decimal('98273.43')), (17, 7, 67, \"\n",
      " \"Decimal('2324.76')), (18, 8, 3, Decimal('4563.78')), (19, 9, 45, \"\n",
      " \"Decimal('87293.00')), (20, 10, 9, Decimal('13457.00'))]\")\n"
     ]
    }
   ],
   "source": [
    "# Connect to the database\n",
    "db = SQLDatabase.from_uri(\"mysql+pymysql://reader:reader@localhost/testdata\")\n",
    "pprint(db.run(\"SELECT * FROM cars;\")) # test\n",
    "pprint(db.run(\"SELECT * FROM sales;\")) # test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5327013e",
   "metadata": {},
   "source": [
    "# Chain to generate examples to use with few-shot prompting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "22b012c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'input': 'Which car brands are available in the cars table?',\n",
      "  'query': 'SELECT DISTINCT make FROM cars;'},\n",
      " {'input': 'How many red cars are there in the cars table?',\n",
      "  'query': \"SELECT COUNT(*) FROM cars WHERE color = 'Red';\"},\n",
      " {'input': 'What is the model of the car with ID 1 in the cars table?',\n",
      "  'query': 'SELECT model FROM cars WHERE carid = 1;'},\n",
      " {'input': 'Which car has the highest price in the sales table?',\n",
      "  'query': 'SELECT carid, MAX(price) as max_price FROM sales GROUP BY carid;'},\n",
      " {'input': 'What is the total number of cars sold for each brand in the sales '\n",
      "           'table?',\n",
      "  'query': 'SELECT make, SUM(quantity) as total_sold FROM sales JOIN cars ON '\n",
      "           'sales.carid = cars.carid GROUP BY make;'},\n",
      " {'input': 'Which car models were sold more than once in the sales table?',\n",
      "  'query': 'SELECT model FROM sales GROUP BY model HAVING COUNT(*) > 1;'},\n",
      " {'input': 'What is the average price of cars sold by brand in the sales '\n",
      "           'table?',\n",
      "  'query': 'SELECT make, AVG(price) as avg_price FROM sales JOIN cars ON '\n",
      "           'sales.carid = cars.carid GROUP BY make;'},\n",
      " {'input': 'Find the car with the lowest price that was sold more than once in '\n",
      "           'the sales table.',\n",
      "  'query': 'SELECT carid, MIN(price) as min_price FROM sales WHERE quantity > '\n",
      "           '1 GROUP BY carid HAVING min_price = (SELECT MIN(price) FROM '\n",
      "           'sales);'},\n",
      " {'input': 'Which cars were sold in the year 2019 according to the sales '\n",
      "           'table?',\n",
      "  'query': 'SELECT make, model FROM cars JOIN sales ON cars.carid = '\n",
      "           'sales.carid WHERE YEAR(sales.year) = 2019;'},\n",
      " {'input': 'What is the total revenue generated by each car brand in the sales '\n",
      "           'table?',\n",
      "  'query': 'SELECT make, SUM(quantity * price) as total_revenue FROM sales '\n",
      "           'JOIN cars ON sales.carid = cars.carid GROUP BY make;'}]\n"
     ]
    }
   ],
   "source": [
    "# Define schema for the generated examples\n",
    "class ExamplesGeneration(BaseModel):\n",
    "    input: str = Field(description=\"question based on the information in the tables\")\n",
    "    query: str = Field(description=f\"{db.dialect} query to answer the question\")\n",
    "\n",
    "# Create a parser to parse the generated response\n",
    "parser = JsonOutputParser(pydantic_object=ExamplesGeneration)\n",
    "\n",
    "examplesGenerationPromptStr= \"\"\"You are a {dialect} expert. \n",
    "Use the given tables information to create {top_k} sets of questions and their corresponding syntactically correct {dialect} queries. \n",
    "When creating questions and their corresponding {dialect} queries pay attention to use only the column names you can see in the tables below.\n",
    "Make sure each set of question and its corresponding {dialect} query is significantly different from the previously generated set.\n",
    "The generated sets should cover most type of questions and their corresponding {dialect} queries that can be aasked about the given tables.\n",
    "Strictly do not include columns that do not exist. \n",
    "Strictly include table using join or from when referencing its columns in the queries.\n",
    "when using join strictly reference columns with their corresponding table name\n",
    "Only provide {top_k} sets of questions and their corresponding {dialect} query as the response. Do not provide any additional explaination.\n",
    "\n",
    "Strictly avoid the following mistakes when generating {dialect} query:\n",
    "- Using NOT IN with NULL values\n",
    "- Using UNION when UNION ALL should have been used\n",
    "- Using BETWEEN for exclusive ranges\n",
    "- Data type mismatch in predicates\n",
    "- Properly quoting identifiers\n",
    "- Using the correct number of arguments for functions\n",
    "- Casting to the correct data type\n",
    "- Using the proper columns for joins\n",
    "\n",
    "Only use the following tables: {table_info}\n",
    "\n",
    "{formatInstructions}\n",
    "\n",
    "{input}\n",
    "\"\"\"\n",
    "# Create a prompt\n",
    "examplesGenerationPrompt = PromptTemplate(\n",
    "    template=examplesGenerationPromptStr,\n",
    "    input_variables=[\"input\"],\n",
    "    partial_variables={\"formatInstructions\": parser.get_format_instructions(),\n",
    "                       \"table_info\": db.get_context()[\"table_info\"],\n",
    "                        \"top_k\": 10,\n",
    "                        \"dialect\": db.dialect},\n",
    ")\n",
    "\n",
    "# Complete chain for generating examples\n",
    "examplesGenerationChain = examplesGenerationPrompt | llm | parser\n",
    "\n",
    "# Run the chain to generate examples\n",
    "generatedExamples = examplesGenerationChain.invoke({\"input\": \"\"})\n",
    "pprint(generatedExamples) # test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bf16180",
   "metadata": {},
   "source": [
    "# Create dynamic examples selector used to select examples relevant to user question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "5579d117",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ollama embeddings to convert generated examples and user questions into embeddings\n",
    "ollamaEmb = OllamaEmbeddings(\n",
    "    model=\"nomic-embed-text\"\n",
    ")\n",
    "\n",
    "# Selector uses Ollama embeddings to generate embeddings\n",
    "# It uses FAISS vector store to store embeddings and run similarity search\n",
    "exampleSelector = SemanticSimilarityExampleSelector.from_examples(\n",
    "    generatedExamples,\n",
    "    ollamaEmb,\n",
    "    FAISS,\n",
    "    k=5,\n",
    "    input_keys=['input'],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87f63947",
   "metadata": {},
   "source": [
    "# Chain to generate SQL from text by using Dynamic few-shot prompting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "0a2ba8a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\"SELECT COUNT(*) FROM cars WHERE color = 'Red';\"\n"
     ]
    }
   ],
   "source": [
    "# System prompt to generate SQL query from text\n",
    "systemPrompt = \"\"\"You are a {dialect} expert. \n",
    "Given an input question, create a syntactically correct {dialect} query to run. \n",
    "Pay attention to use only the column names you can see in the tables below.\n",
    "Be careful to not query for columns that do not exist. \n",
    "Also, pay attention to which column is in which table. \n",
    "Unless otherwise specificed, do not return more than {top_k} rows.\n",
    "If the input question is not related to the tables then say 'Not relevant question'.\n",
    "If the input question cannot be answered using the tables then say 'Cannot be answered'.\n",
    "If you don't know the answer to the input question then say 'I don't know'.\n",
    "Only provide {dialect} as the response. Do not provide any additional explaination.\n",
    "\n",
    "Avoid the following mistakes when generating {dialect} query:\n",
    "- Using NOT IN with NULL values\n",
    "- Using UNION when UNION ALL should have been used\n",
    "- Using BETWEEN for exclusive ranges\n",
    "- Data type mismatch in predicates\n",
    "- Properly quoting identifiers\n",
    "- Using the correct number of arguments for functions\n",
    "- Casting to the correct data type\n",
    "- Using the proper columns for joins\n",
    "\n",
    "Only use the following tables: {table_info}\n",
    "\n",
    "Below are a number of examples of questions and their corresponding SQL queries.\"\"\"\n",
    "\n",
    "# Prompt to format list of selected few-shot examples\n",
    "examplePrompt = PromptTemplate.from_template(\"User: {input}\\nAssistant: {query}\")\n",
    "\n",
    "# Prompt object to convert user question to SQL\n",
    "prompt = FewShotPromptTemplate(\n",
    "    example_selector=exampleSelector, \n",
    "    example_prompt=examplePrompt,\n",
    "    prefix=systemPrompt,\n",
    "    suffix=\"User: {input}\\nAssistant: \",\n",
    "    input_variables=[\"input\", \"top_k\", \"table_info\"],\n",
    "    partial_variables = {\"dialect\" : db.dialect},\n",
    ")\n",
    "# pprint(prompt.format(input = \"\", top_k = 10, table_info = db.get_context()[\"table_info\"])) # test\n",
    "# Chain to convert the text from the user into a SQL query\n",
    "writeQueryChain = create_sql_query_chain(llm, db, prompt) | StrOutputParser()\n",
    "\n",
    "# Filter out cases where llm could not answer the question\n",
    "def llmApproval(msg: AIMessage) -> Runnable:\n",
    "    errMsgs = ('not relevant question', 'cannot be answered', 'i don\\'t know')\n",
    "    for errMsg in errMsgs:\n",
    "        if errMsg in msg.lower():\n",
    "            raise ValueError(f\"Input question not approved by llm:\\n\\n{msg}\")\n",
    "    return msg\n",
    "\n",
    "# Chain that generated a SQL query after validating that the question is relevant\n",
    "queryChain = writeQueryChain | llmApproval\n",
    "\n",
    "pprint(queryChain.invoke({\"question\" : \"How many cars are of color red?\"})) # test\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9869e057",
   "metadata": {},
   "source": [
    "# Chain to generate explaination for a query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "7dc0aa6e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(' This SQL query counts the number of rows in the \"cars\" table where the '\n",
      " '\"color\" column is equal to \\'Red\\'. It uses the COUNT function to count all '\n",
      " 'rows and the WHERE clause to filter only red cars.')\n"
     ]
    }
   ],
   "source": [
    "explainPromptStr = \"\"\" You are an assistant who is excellent at explaing {dialect} queries. \n",
    "Please explain the given {dialect} query. \n",
    "Be consice and lets do it step by step. \n",
    "The response should not be more than 50 words.\n",
    "\n",
    "User: {query}\n",
    "Assistant: \"\"\"\n",
    "explainPromptChain = PromptTemplate.from_template(explainPromptStr)\n",
    "# Chain that takes the generated query and generates an explaination for the query\n",
    "explainChain = explainPromptChain | llm | StrOutputParser()\n",
    "pprint(explainChain.invoke({\"query\" : \"SELECT COUNT(*) FROM cars WHERE color = 'Red';\", \"dialect\" :  db.dialect})) # test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6576c5b",
   "metadata": {},
   "source": [
    "# Chain to execute SQL query against database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "bcc28ed0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#chain to take a SQL query and execute against a database\n",
    "executeQueryChain = QuerySQLDataBaseTool(db=db)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65e18130",
   "metadata": {},
   "source": [
    "# Helper methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "c03d907a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Actions\n",
    "class Action(IntEnum):\n",
    "    NONE = 0\n",
    "    EXECUTE = 1\n",
    "    EXPLAIN = 2\n",
    "    REGENERATE = 3    \n",
    "\n",
    "# Get approval from user before executing the query\n",
    "def humanApproval(msg, explainMsg):\n",
    "    inputMsg = (\n",
    "        f\"Here is the query generated by llm : \\n\\n{msg}\\n\\n\"\n",
    "        \"Enter 'Execute' to run the query.\"\n",
    "        \"\\nEnter 'Explain' to get explaination for the query.\"\n",
    "        \"\\nEnter 'Regenerate' to regenerate the query.\"\n",
    "        \"\\nEnter anything except above values will be treated as abort.\"\n",
    "        \"\\nNote : All inputs are case insensetive.\"\n",
    "    )\n",
    "    \n",
    "    if explainMsg != \"\":\n",
    "        inputMsg = explainMsg + \"\\n\" + inputMsg\n",
    "\n",
    "    resp = input(inputMsg)\n",
    "    if resp.lower() == \"execute\":\n",
    "        return Action.EXECUTE\n",
    "    elif resp.lower() == \"explain\":\n",
    "        return Action.EXPLAIN\n",
    "    elif resp.lower() == \"regenerate\":\n",
    "        return Action.REGENERATE\n",
    "    else :\n",
    "        raise ValueError(f\"Tool invocations not approved by user:\\n\\n{msg}\")\n",
    "\n",
    "# Get user question\n",
    "def getQuery():\n",
    "    inputMsg = \"Please enter the question (enter 'bye' to exit) : \"\n",
    "    inputStr = input(inputMsg)\n",
    "    if inputStr.lower() == \"bye\":\n",
    "        raise KeyboardInterrupt\n",
    "    else:\n",
    "        return inputStr \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d894786",
   "metadata": {},
   "source": [
    "# Main method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "7a7d8409",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "――――――――――\n",
      "Input question not approved by llm:\n",
      "\n",
      "Not relevant question.\n",
      "――――――――――\n",
      "User Question : total quantity of cars sold?\n",
      "SQL Query : SELECT SUM(quantity) as total_quantity FROM sales;\n",
      "Explaination :  This SQL query calculates the total sum of the 'quantity' column from the 'sales' table and labels it as 'total_quantity'. It helps to quickly find out the total quantity of items sold.\n",
      "Result : [(Decimal('230'),)]\n",
      "――――――――――\n",
      "Exiting...\n"
     ]
    }
   ],
   "source": [
    "while(True):\n",
    "    print('―' * 10)\n",
    "    query = \"\"\n",
    "    explaination = \"\"\n",
    "    userQuestion = \"\"\n",
    "    try:\n",
    "        userQuestion = getQuery()\n",
    "        query = queryChain.invoke({\"question\" : userQuestion})\n",
    "    except ValueError as err:\n",
    "        print(err)\n",
    "        continue\n",
    "    except KeyboardInterrupt:\n",
    "        print(\"Exiting...\")\n",
    "        break\n",
    "    while(True):\n",
    "        # Once we generate a query, we ask user if they wish to :\n",
    "        #  - Execute the query\n",
    "        #  - Get explaination\n",
    "        #  - Abort\n",
    "        action = Action.NONE\n",
    "        try: \n",
    "            action = humanApproval(query, explaination)\n",
    "        except ValueError as err:\n",
    "            print(err)\n",
    "            break\n",
    "        else:\n",
    "            if action == Action.EXECUTE:\n",
    "                try:\n",
    "                    result = executeQueryChain.invoke(query)\n",
    "                except Exception as err:\n",
    "                    print(f\"Failed to run generated query. Error : {err}\")\n",
    "                    continue\n",
    "                print(f\"User Question : {userQuestion}\")\n",
    "                print(f\"SQL Query : {query}\")\n",
    "                if explaination != \"\":\n",
    "                    print(f\"Explaination : {explaination}\")\n",
    "                print(f\"Result : {result}\")\n",
    "                break\n",
    "            if action == Action.EXPLAIN:\n",
    "                explaination = explainChain.invoke({\"query\" : query, \"dialect\" :  db.dialect})\n",
    "                continue\n",
    "            if action == Action.REGENERATE:\n",
    "                try:\n",
    "                    query = queryChain.invoke({\"question\" : userQuestion})\n",
    "                    explaination = \"\"\n",
    "                    continue\n",
    "                except ValueError as err:\n",
    "                    print(err)\n",
    "                    break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8e261e6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}