SQL.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "93066545-02c2-4b52-b92e-b52bf313adc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"COHERE_API_KEY\"] = 'fOKNiUz8cq3c7LfrrzrX7E70vgGpZStUEgdt3Ohj'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "37852370-4efe-48f2-865f-36e9d73c754e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_cohere import ChatCohere\n",
    "from langchain.chains import create_sql_query_chain\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool\n",
    "from langchain_community.embeddings import HuggingFaceInstructEmbeddings\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain.prompts import SemanticSimilarityExampleSelector\n",
    "from langchain.prompts.prompt import PromptTemplate\n",
    "from langchain.prompts import FewShotPromptTemplate\n",
    "from langchain.chains.sql_database.prompt import PROMPT_SUFFIX, _mysql_prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "438841f1-43af-4939-b8ff-44a7bcb639d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['discounts', 't_shirts']\n"
     ]
    }
   ],
   "source": [
    "db_user = \"root\"\n",
    "db_password = \"Donsurya23!\"\n",
    "db_host = \"localhost\"\n",
    "db_name = \"atliq_tshirts\"\n",
    "\n",
    "db = SQLDatabase.from_uri(f\"mysql+pymysql://{db_user}:{db_password}@{db_host}/{db_name}\", sample_rows_in_table_info=3)\n",
    "print(db.get_usable_table_names())\n",
    "# db.run(\"SELECT * FROM discounts\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a0a2f154-9254-4af1-b8e7-1ef0dd4ae7b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatCohere(model=\"command-r-plus\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cf96a451-b9de-4a60-bec9-6e3b4fa318d8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'```sql\\nSELECT COUNT(*) AS `num_discounts` FROM `discounts`;\\n```'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = create_sql_query_chain(llm, db)\n",
    "response = chain.invoke({\"question\": \"How many discounts are there\"})\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0efc8d25-c67b-4747-a25d-ed915bf977d4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'[(9,)]'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clean_response = response[7:].split(\";\")[0]\n",
    "clean_response\n",
    "db.run(clean_response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "76b77cd2-3a89-4fb3-97d5-4da49449b75a",
   "metadata": {},
   "outputs": [],
   "source": [
    "execute_query = QuerySQLDataBaseTool(db=db)\n",
    "write_query = create_sql_query_chain(llm, db)\n",
    "chain = write_query | (lambda x: (print(x), x[7:].split(\";\")[0])[1]) | execute_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "56eb67e6-c6f3-465f-a37f-3da92304c673",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```sql\n",
      "SELECT `stock_quantity` \n",
      "FROM `t_shirts` \n",
      "WHERE `brand` = 'Nike' AND `color` = 'White' AND `size` = 'XS';\n",
      "```\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'[(20,)]'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qns1 = chain.invoke({\"question\": \"How many tshirts do we have left for nike in extra small size and white color?\"})\n",
    "qns1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "91db9399-4572-4cae-a2cf-503271a95078",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```sql\n",
      "SELECT SUM(`price` * `stock_quantity`) AS `total_value`\n",
      "FROM `t_shirts`\n",
      "WHERE `size` = 'S';\n",
      "```\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"[(Decimal('24886'),)]\""
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# qns2 = chain.invoke({\"question\": \"How much is the price of the inventory for all small size tshirts?\"})\n",
    "qns2 = chain.invoke({\"question\": \"SELECT SUM(price*stock_quantity) FROM atliq_tshirts.t_shirts WHERE size = 'S';\"})\n",
    "qns2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ac692f9e-6c49-4d61-b42e-1b3d4d555888",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```sql\n",
      "SELECT SUM(a.total_amount * ((100 - COALESCE(d.pct_discount, 0)) / 100)) AS total_revenue\n",
      "FROM (\n",
      "    SELECT SUM(price * stock_quantity) AS total_amount, t_shirt_id\n",
      "    FROM t_shirts\n",
      "    WHERE brand = 'Levi'\n",
      "    GROUP BY t_shirt_id\n",
      ") a\n",
      "LEFT JOIN discounts d ON a.t_shirt_id = d.t_shirt_id;\n",
      "```\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"[(Decimal('22286.300000'),)]\""
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# qns3 = chain.invoke({\"question\": \"If we have to sell all the Levi’s T-shirts today with discounts applied. How much revenue our store will generate (post discounts)?\"})\n",
    "qns3 = chain.invoke({\"question\": \"\"\"\n",
    "select sum(a.total_amount * ((100-COALESCE(discounts.pct_discount,0))/100)) as total_revenue from\n",
    "(select sum(price*stock_quantity) as total_amount, t_shirt_id from t_shirts where brand = 'Levi'\n",
    "group by t_shirt_id) a left join discounts on a.t_shirt_id = discounts.t_shirt_id\n",
    " \"\"\"})\n",
    "qns3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "8840f33d-dc7a-4186-9946-bc28bcb1ac6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"[(Decimal('23146'),)]\""
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# qns4 = chain.invoke({\"question\": \"select SUM(price * stock_quantity) FROM t_shirts WHERE brand = 'Levi';\"})\n",
    "qns4 = db.run(\"select SUM(price * stock_quantity) FROM t_shirts WHERE brand = 'Levi';\")\n",
    "qns4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "daa1f4de-cb99-49a4-8451-88901c7af561",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"[(Decimal('170'),)]\""
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# qns5 = chain.invoke({\"question\": \"How many white color Levi's t shirts we have available?\"})\n",
    "# qns5 = chain.invoke({\"question\": \"SELECT sum(stock_quantity) FROM t_shirts WHERE brand = 'Levi' AND color = 'White'\"})\n",
    "qns5 = db.run(\"SELECT sum(stock_quantity) FROM t_shirts WHERE brand = 'Levi' AND color = 'White'\")\n",
    "qns5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "a646c7f2-1759-4fe1-b7c5-225d41e2b665",
   "metadata": {},
   "outputs": [],
   "source": [
    "few_shots = [\n",
    "    {'Question' : \"How many t-shirts do we have left for Nike in XS size and white color?\",\n",
    "     'SQLQuery' : \"SELECT sum(stock_quantity) FROM t_shirts WHERE brand = 'Nike' AND color = 'White' AND size = 'XS'\",\n",
    "     'SQLResult': \"Result of the SQL query\",\n",
    "     'Answer' : qns1},\n",
    "    {'Question': \"How much is the total price of the inventory for all S-size t-shirts?\",\n",
    "     'SQLQuery':\"SELECT SUM(price*stock_quantity) FROM t_shirts WHERE size = 'S'\",\n",
    "     'SQLResult': \"Result of the SQL query\",\n",
    "     'Answer': qns2},\n",
    "    {'Question': \"If we have to sell all the Levi’s T-shirts today with discounts applied. How much revenue  our store will generate (post discounts)?\" ,\n",
    "     'SQLQuery' : \"\"\"SELECT sum(a.total_amount * ((100-COALESCE(discounts.pct_discount,0))/100)) as total_revenue from\n",
    "(select sum(price*stock_quantity) as total_amount, t_shirt_id from t_shirts where brand = 'Levi'\n",
    "group by t_shirt_id) a left join discounts on a.t_shirt_id = discounts.t_shirt_id\n",
    " \"\"\",\n",
    "     'SQLResult': \"Result of the SQL query\",\n",
    "     'Answer': qns3} ,\n",
    "     {'Question' : \"If we have to sell all the Levi’s T-shirts today. How much revenue our store will generate without discount?\" ,\n",
    "      'SQLQuery': \"SELECT SUM(price * stock_quantity) FROM t_shirts WHERE brand = 'Levi'\",\n",
    "      'SQLResult': \"Result of the SQL query\",\n",
    "      'Answer' : qns4},\n",
    "    {'Question': \"How many white color Levi's shirt I have?\",\n",
    "     'SQLQuery' : \"SELECT sum(stock_quantity) FROM t_shirts WHERE brand = 'Levi' AND color = 'White'\",\n",
    "     'SQLResult': \"Result of the SQL query\",\n",
    "     'Answer' : qns5\n",
    "     }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "a3ee8d98-18cb-46e6-8f71-bba00b11f2f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load INSTRUCTOR_Transformer\n",
      "max_seq_length  512\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/donc/anaconda3/lib/python3.12/site-packages/transformers/tokenization_utils_base.py:1617: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be deprecated in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "embeddings = HuggingFaceInstructEmbeddings(\n",
    "    model_name=\"sentence-transformers/all-mpnet-base-v2\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "0714e4cc-6f27-4cb5-807a-b34ed7ef0f91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[\"How many t-shirts do we have left for Nike in XS size and white color? SELECT sum(stock_quantity) FROM t_shirts WHERE brand = 'Nike' AND color = 'White' AND size = 'XS' Result of the SQL query [(20,)]\",\n",
       " \"How much is the total price of the inventory for all S-size t-shirts? SELECT SUM(price*stock_quantity) FROM t_shirts WHERE size = 'S' Result of the SQL query [(Decimal('24886'),)]\",\n",
       " \"If we have to sell all the Levi’s T-shirts today with discounts applied. How much revenue  our store will generate (post discounts)? SELECT sum(a.total_amount * ((100-COALESCE(discounts.pct_discount,0))/100)) as total_revenue from\\n(select sum(price*stock_quantity) as total_amount, t_shirt_id from t_shirts where brand = 'Levi'\\ngroup by t_shirt_id) a left join discounts on a.t_shirt_id = discounts.t_shirt_id\\n  Result of the SQL query [(Decimal('22286.300000'),)]\",\n",
       " \"If we have to sell all the Levi’s T-shirts today. How much revenue our store will generate without discount? SELECT SUM(price * stock_quantity) FROM t_shirts WHERE brand = 'Levi' Result of the SQL query [(Decimal('23146'),)]\",\n",
       " \"How many white color Levi's shirt I have? SELECT sum(stock_quantity) FROM t_shirts WHERE brand = 'Levi' AND color = 'White' Result of the SQL query [(Decimal('170'),)]\"]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_vectorize = [\" \".join(example.values()) for example in few_shots]\n",
    "to_vectorize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "154ff6be-7bd5-4cd6-ad64-6144deb44443",
   "metadata": {},
   "outputs": [],
   "source": [
    "vectorstore = Chroma.from_texts(to_vectorize, embeddings, metadatas=few_shots)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "80c89cdf-d59f-4939-af22-15df63fd5216",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'Answer': '[(20,)]',\n",
       "  'Question': 'How many t-shirts do we have left for Nike in XS size and white color?',\n",
       "  'SQLQuery': \"SELECT sum(stock_quantity) FROM t_shirts WHERE brand = 'Nike' AND color = 'White' AND size = 'XS'\",\n",
       "  'SQLResult': 'Result of the SQL query'},\n",
       " {'Answer': \"[(Decimal('170'),)]\",\n",
       "  'Question': \"How many white color Levi's shirt I have?\",\n",
       "  'SQLQuery': \"SELECT sum(stock_quantity) FROM t_shirts WHERE brand = 'Levi' AND color = 'White'\",\n",
       "  'SQLResult': 'Result of the SQL query'}]"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_selector = SemanticSimilarityExampleSelector(\n",
    "    vectorstore=vectorstore,\n",
    "    k=2,\n",
    ")\n",
    "example_selector.select_examples({\"Question\": \"How many Adidas T shirts I have left in my store?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "ab485594-b7b6-42ee-bc6e-8a98c6d5f194",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_prompt = PromptTemplate(\n",
    "    input_variables=[\"Question\", \"SQLQuery\", \"SQLResult\",\"Answer\",],\n",
    "    template=\"\\nQuestion: {Question}\\nSQLQuery: {SQLQuery}\\nSQLResult: {SQLResult}\\nAnswer: {Answer}\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "52f80880-76dd-4caa-92bb-07d92ccebf16",
   "metadata": {},
   "outputs": [],
   "source": [
    "few_shot_prompt = FewShotPromptTemplate(\n",
    "    example_selector=example_selector,\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=_mysql_prompt,\n",
    "    suffix=PROMPT_SUFFIX,\n",
    "    input_variables=[\"input\", \"table_info\", \"top_k\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "def3849f-ae3f-42a5-b655-45f58f8922d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_write_query = create_sql_query_chain(llm, db, prompt=few_shot_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "62de5709-22c7-4e2a-94ac-e13d72a81ed7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'SQLQuery: SELECT count(DISTINCT `t_shirt_id`) AS total_tshirts FROM t_shirts;\\nSQLResult:'"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_write_query.invoke({\"question\": \"how many tshirts are there?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "501fef09-dec9-4f28-9ff1-31d2f34be092",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_chain = new_write_query | (lambda x: (print(x), x[10:].split(\";\")[0])[1]) | execute_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "236a19fa-f60f-434a-97b5-3bb4ae8ed8a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SQLQuery: SELECT sum(`stock_quantity`) AS total_white_levi_shirts FROM `t_shirts` WHERE `brand` = 'Levi' AND `color` = 'White';\n",
      "SQLResult:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"[(Decimal('170'),)]\""
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_chain.invoke({\"question\": \"How many white color Levi's shirt I have?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "389548db-0183-46ec-aa91-70e0a56d6727",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}