jupyter/QA-SQL-LLM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "95617592-a31a-42d7-8977-561235bc3668",
   "metadata": {},
   "source": [
    "https://python.langchain.com/docs/tutorials/sql_qa/\n",
    "\n",
    "https://python.langchain.com/docs/how_to/sql_prompting/\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7f8fbe10-29ee-4d35-b3a3-99864b245d37",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain_community.utilities import SQLDatabase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4447f22a-aab1-4bc3-8d1e-f903d29345f8",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "db = SQLDatabase.from_uri(\"iris://superuser:SYS@iris:1972/LLMRAG\", sample_rows_in_table_info=3, schema='Holefoods')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7d20679a-1095-4de3-9b4f-252016807c74",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iris\n",
      "['Country', 'Outlet', 'Product', 'Region', 'SalesTransaction']\n"
     ]
    }
   ],
   "source": [
    "print(db.dialect)\n",
    "print(db.get_usable_table_names())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "65c2349a-4a90-41bb-aa0d-0563c790b7b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(1, 1, Decimal('2.66'), '2', None, datetime.date(2024, 7, 17), Decimal('0.1'), None, None, 33, 'SKU-101', None, 1, None), (2, 1, Decimal('5.36'), '2', None, datetime.date(2024, 7, 6), Decimal('0.1'), 39.673679, -89.711212, 26, 'SKU-222', None, 1, '62629'), (3, 1, Decimal('1.56'), '2', None, datetime.date(2021, 9, 10), Decimal('0.2'), None, None, 2, 'SKU-296', None, 1, None), (4, 1, Decimal('1.56'), None, None, datetime.date(2020, 3, 25), Decimal('0.2'), 42.291074, -86.274294, 26, 'SKU-296', None, 1, '49043'), (5, 1, Decimal('0.98'), '2', 'Shopper ranted about high prices', datetime.date(2023, 6, 11), Decimal('0.5'), None, None, 12, 'SKU-287', None, 1, None), (6, 1, Decimal('12.95'), '2', None, datetime.date(2020, 9, 24), 0, None, None, 1, 'SKU-195', None, 1, None), (7, 1, Decimal('44.75'), '2', None, datetime.date(2021, 10, 2), 0, 39.193793, -87.713073, 26, 'SKU-900', None, 5, '62478'), (8, 1, Decimal('0.92'), '2', None, datetime.date(2021, 6, 6), Decimal('0.2'), 30.529819, -96.714292, 27, 'SKU-451', None, 1, '77836'), (9, 1, Decimal('5.95'), '2', None, datetime.date(2021, 6, 27), 0, None, None, 33, 'SKU-222', None, 1, None), (10, 1, Decimal('3.83'), '2', None, datetime.date(2023, 6, 24), Decimal('0.1'), None, None, 19, 'SKU-204', None, 1, None)]\n"
     ]
    }
   ],
   "source": [
    "print(db.run(\"SELECT TOP 10 * FROM Holefoods.SalesTransaction\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3b94bdcc-a0c1-4b1f-baf3-399d8884012a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['table_info', 'table_names']\n",
      "\n",
      "CREATE TABLE \"Holefoods\".\"Country\" (\n",
      "\t\"ID\" IDENTITY DEFAULT $i(^HoleFoods.CountryD) NOT NULL, \n",
      "\t\"Name\" VARCHAR(90), \n",
      "\t\"Region\" BIGINT, \n",
      "\tCONSTRAINT \"RowIDField_As_PKey\" PRIMARY KEY (\"ID\")\n",
      ") WITH %CLASSPARAMETER ALLOWIDENTITYINSERT = 1\n",
      "\n",
      "/*\n",
      "3 rows from Country table:\n",
      "ID\tName\tRegion\n",
      "1\tChina\t1\n",
      "2\tIndia\t1\n",
      "3\tJapan\t1\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Holefoods\".\"Outlet\" (\n",
      "\t\"ID\" IDENTITY DEFAULT $i(^HoleFoods.OutletD) NOT NULL, \n",
      "\t\"City\" VARCHAR(100), \n",
      "\t\"Country\" BIGINT, \n",
      "\t\"Latitude\" DOUBLE, \n",
      "\t\"Longitude\" DOUBLE, \n",
      "\t\"Population\" INTEGER, \n",
      "\t\"Type\" VARCHAR(50), \n",
      "\tCONSTRAINT \"RowIDField_As_PKey\" PRIMARY KEY (\"ID\")\n",
      ") WITH %CLASSPARAMETER ALLOWIDENTITYINSERT = 1\n",
      "\n",
      "/*\n",
      "3 rows from Outlet table:\n",
      "ID\tCity\tCountry\tLatitude\tLongitude\tPopulation\tType\n",
      "1\tBeijing\t1\t39.86\t116.412\t17400000\tNone\n",
      "2\tShanghai\t1\t31.224\t121.466\t16738000\tNone\n",
      "3\tBangalore\t2\t12.963\t77.587\t6200000\tNone\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Holefoods\".\"Product\" (\n",
      "\t\"ID\" VARCHAR(22) NOT NULL, \n",
      "\t\"Category\" VARCHAR(100), \n",
      "\t\"Name\" VARCHAR(120), \n",
      "\t\"Price\" NUMERIC(15, 2), \n",
      "\t\"SKU\" VARCHAR(22) NOT NULL, \n",
      "\tCONSTRAINT \"IDKEYField_As_PKey\" PRIMARY KEY (\"SKU\")\n",
      ") WITH %CLASSPARAMETER ALLOWIDENTITYINSERT = 1\n",
      "\n",
      "/*\n",
      "3 rows from Product table:\n",
      "ID\tCategory\tName\tPrice\tSKU\n",
      "SKU-101\tSnack\tBagels (dozen)\t2.95\tSKU-101\n",
      "SKU-192\tSnack\tDonuts (dozen)\t2.95\tSKU-192\n",
      "SKU-195\tSnack\tFree-range Donuts (dozen)\t12.95\tSKU-195\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Holefoods\".\"Region\" (\n",
      "\t\"ID\" IDENTITY DEFAULT $i(^HoleFoods.RegionD) NOT NULL, \n",
      "\t\"Name\" VARCHAR(80), \n",
      "\tCONSTRAINT \"RowIDField_As_PKey\" PRIMARY KEY (\"ID\")\n",
      ") WITH %CLASSPARAMETER ALLOWIDENTITYINSERT = 1\n",
      "\n",
      "/*\n",
      "3 rows from Region table:\n",
      "ID\tName\n",
      "1\tAsia\n",
      "2\tEurope\n",
      "3\tN. America\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Holefoods\".\"SalesTransaction\" (\n",
      "\t\"ID\" IDENTITY DEFAULT $i(^HoleFoods.TransactionD) NOT NULL, \n",
      "\t\"Actual\" BIT, \n",
      "\t\"AmountOfSale\" NUMERIC(15, 2), \n",
      "\t\"Channel\" VARCHAR(50), \n",
      "\t\"Comment\" VARCHAR(500), \n",
      "\t\"DateOfSale\" DATE, \n",
      "\t\"Discount\" NUMERIC(15, 2), \n",
      "\t\"Latitude\" DOUBLE, \n",
      "\t\"Longitude\" DOUBLE, \n",
      "\t\"Outlet\" BIGINT, \n",
      "\t\"Product\" VARCHAR(22), \n",
      "\t\"TargetAmount\" NUMERIC(15, 2), \n",
      "\t\"UnitsSold\" INTEGER, \n",
      "\t\"ZipCode\" VARCHAR(25), \n",
      "\tCONSTRAINT \"RowIDField_As_PKey\" PRIMARY KEY (\"ID\")\n",
      ") WITH %CLASSPARAMETER ALLOWIDENTITYINSERT = 1\n",
      "\n",
      "/*\n",
      "3 rows from SalesTransaction table:\n",
      "ID\tActual\tAmountOfSale\tChannel\tComment\tDateOfSale\tDiscount\tLatitude\tLongitude\tOutlet\tProduct\tTargetAmount\tUnitsSold\tZipCode\n",
      "1\t1\t2.66\t2\tNone\t2024-07-17\t0.1\tNone\tNone\t33\tSKU-101\tNone\t1\tNone\n",
      "2\t1\t5.36\t2\tNone\t2024-07-06\t0.1\t39.673679\t-89.711212\t26\tSKU-222\tNone\t1\t62629\n",
      "3\t1\t1.56\t2\tNone\t2021-09-10\t0.2\tNone\tNone\t2\tSKU-296\tNone\t1\tNone\n",
      "*/\n"
     ]
    }
   ],
   "source": [
    "context = db.get_context()\n",
    "print(list(context))\n",
    "print(context[\"table_info\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9109ea9c-e187-45ac-9cdf-28e0de00f30b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      " ········\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"MISTRAL_API_KEY\"] = getpass.getpass()\n",
    "\n",
    "from langchain_mistralai import ChatMistralAI\n",
    "\n",
    "llm = ChatMistralAI(model=\"mistral-large-latest\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "29c1eb15-98f4-418a-a997-63370ec0ba03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      " ········\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1229f3af-b5b4-404e-9546-0799705d7b9e",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "# Define the custom prompt template\n",
    "template = '''\n",
    "You are an InterSystems IRIS SQL expert. \n",
    "Given an input question, first create a syntactically correct InterSystems IRIS SQL query to run and return the answer to the input question.\n",
    "Unless the user specifies in the question a specific number of examples to obtain, query for at most {top_k} results using the TOP as defined in InterSystems IRIS syntax: ```SELECT [DISTINCT] TOP int select-item, select-item,...```\n",
    "Always specify table names using schema as prefix.\n",
    "Do not use LIMIT clause as it is not correct in IRIS dialect.\n",
    "Do not end SQL sentences with an ;\n",
    "Do not enclose fields in quotes or double quotes.\n",
    "Do not enclose table names in quotes or double quotes.\n",
    "You can order the results to return the most informative data in the database.\n",
    "Never query for all columns from a table. You must query only the columns that are needed to answer the question.\n",
    "Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.\n",
    "Pay attention to use CAST(CURRENT_DATE as date) function to get the current date, if the question involves \"today\".\n",
    "\n",
    "Return only plain SQL without any formatting.\n",
    "\n",
    "Only use the following tables:\n",
    "\n",
    "{table_info}.\n",
    "Question: {input}'''\n",
    "\n",
    "# Create the PromptTemplate object\n",
    "custom_prompt = PromptTemplate(\n",
    "    input_variables=[\"input\", \"table_info\", \"top_k\", \"dialect\"],\n",
    "    template=template\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "0081f33b-17b1-42b9-9d2a-d102b54d908a",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "You are an InterSystems IRIS SQL expert. \n",
      "Given an input question, first create a syntactically correct InterSystems IRIS SQL query to run and return the answer to the input question.\n",
      "Unless the user specifies in the question a specific number of examples to obtain, query for at most 5 results using the TOP as defined in InterSystems IRIS syntax: ```SELECT [DISTINCT] TOP int select-item, select-item,...```\n",
      "Always specify table names using schema as prefix.\n",
      "Do not use LIMIT clause as it is not correct in IRIS dialect.\n",
      "Do not end SQL sentences with an ;\n",
      "Do not enclose fields in quotes or double quotes.\n",
      "Do not enclose table names in quotes or double quotes.\n",
      "You can order the results to return the most informative data in the database.\n",
      "Never query for all columns from a table. You must query only the columns that are needed to answer the question.\n",
      "Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.\n",
      "Pay attention to use CAST(CURRENT_DATE as date) function to get the current date, if the question involves \"today\".\n",
      "\n",
      "Return only plain SQL without any formatting.\n",
      "\n",
      "Only use the following tables:\n",
      "\n",
      "\u001b[33;1m\u001b[1;3m{table_info}\u001b[0m.\n",
      "Question: \u001b[33;1m\u001b[1;3m{input}\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from langchain.chains import create_sql_query_chain\n",
    "\n",
    "chain = create_sql_query_chain(llm, db, prompt=custom_prompt)\n",
    "chain.get_prompts()[0].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "c8885df3-f8b1-40c0-81b9-2f88856e8f02",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.callbacks.tracers import ConsoleCallbackHandler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "cac2d092-3d5b-4afd-add1-4fc0d37d6e35",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'SELECT COUNT(DISTINCT SKU) FROM Holefoods.Product'"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke({\"question\": \"how many products are there?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "2c00c29a-f139-4ae7-8064-81eb58edc2ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'SELECT COUNT(*) FROM Holefoods.SalesTransaction WHERE YEAR(DateOfSale) = 2023'"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke({\"question\": \"how many sales in 2023?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "96097a7b-d9a5-43da-b16a-5ccd12227709",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'SELECT DISTINCT TOP 5 Category FROM Holefoods.Product'"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke({\"question\": \"what are the product categories?\"})\n",
    "# chain.invoke({\"question\": \"what are the product categories?\"}, config={'callbacks': [ConsoleCallbackHandler()]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "8088217f-8758-4fb2-9243-ea2f1bee60da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"SELECT TOP 5 p.Category, SUM(st.UnitsSold) AS TotalUnitsSold \\nFROM Holefoods.SalesTransaction st \\nJOIN Holefoods.Product p ON st.Product = p.SKU \\nWHERE st.DateOfSale >= CAST('2023-01-01' AS date) AND st.DateOfSale < CAST('2024-01-01' AS date) \\nGROUP BY p.Category \\nORDER BY TotalUnitsSold DESC\""
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke({\"question\": \"what are the most sold product categories during 2023?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06fb484c-3f89-457c-8c77-242bf404b0f3",
   "metadata": {},
   "source": [
    "# Execute query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "0ce39360-3c24-4528-8337-0cae7968759c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"[('PASTA', 104), ('VEGETABLE', 90), ('SNACK', 79), ('CANDY', 43), ('CEREAL', 37)]\""
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool\n",
    "\n",
    "execute_query = QuerySQLDataBaseTool(db=db)\n",
    "write_query = create_sql_query_chain(llm, db, prompt=custom_prompt)\n",
    "chain = write_query | execute_query\n",
    "chain.invoke({\"question\": \"what are the most sold product categories during 2023?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4129d621-01b0-4de9-ad70-039b4981f6b3",
   "metadata": {},
   "source": [
    "# Dynamic Fewshot examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "fe98b814-c781-4f94-a770-95e50f295a29",
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = [\n",
    "    { \n",
    "        \"input\": \"List all regions.\", \n",
    "        \"query\": \"SELECT ID, Name FROM HoleFoods.Region;\"\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"List all countries.\",\n",
    "        \"query\": \"SELECT c.ID, c.Name, r.Name Region FROM HoleFoods.Country c JOIN HoleFoods.Region r on c.Region=r.ID\"\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"What are the different product categories ?\",\n",
    "        \"query\": \"SELECT DISTINCT(Category) Categories FROM HoleFoods.Product\"\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"How many pasta products where sold online in 2023 ?\",\n",
    "        \"query\": \"SELECT SUM(UnitsSold) FROM HoleFoods.SalesTransaction st JOIN HoleFoods.Product p ON st.Product=p.ID WHERE st.Channel='Online' AND YEAR(st.DateOfSale) = 2023 AND p.Category = 'Pasta'\"\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Find all snack products\",\n",
    "        \"query\": \"SELECT SKU, Name, Price FROM HoleFoods.Product p WHERE p.Category='Snack'\"\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Find all candy products\",\n",
    "        \"query\": \"SELECT SKU, Name, Price FROM HoleFoods.Product p WHERE p.Category='Candy'\"\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"How many products were sold in Europe in 2022 ?\",\n",
    "        \"query\": \"SELECT SUM(UnitsSold) FROM HoleFoods.SalesTransaction st JOIN HoleFoods.Outlet o ON st.Outlet=o.ID JOIN HoleFoods.Country c ON o.Country=c.ID JOIN HoleFoods.Region r ON c.Region=r.ID WHERE r.Name='Europe' AND YEAR(st.DateOfSale) = 2022\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "de9efe81-d946-4191-91a0-2249b2c32ab1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_iris import IRISVector\n",
    "from langchain_core.example_selectors import SemanticSimilarityExampleSelector\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "example_selector = SemanticSimilarityExampleSelector.from_examples(\n",
    "    examples,\n",
    "    OpenAIEmbeddings(),\n",
    "    IRISVector,\n",
    "    k=3,\n",
    "    input_keys=[\"input\"],\n",
    "    connection_string='iris://superuser:SYS@iris:1972/LLMRAG',\n",
    "    collection_name=\"sql_samples\",\n",
    "    pre_delete_collection=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "7ecd2d43-ffb3-42e1-b433-cbfd2935eac5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'input': 'How many products were sold in Europe in 2022 ?',\n",
       "  'query': \"SELECT SUM(UnitsSold) FROM HoleFoods.SalesTransaction st JOIN HoleFoods.Outlet o ON st.Outlet=o.ID JOIN HoleFoods.Country c ON o.Country=c.ID JOIN HoleFoods.Region r ON c.Region=r.ID WHERE r.Name='Europe' AND YEAR(st.DateOfSale) = 2022\"},\n",
       " {'input': 'How many pasta products where sold online in 2023 ?',\n",
       "  'query': \"SELECT SUM(UnitsSold) FROM HoleFoods.SalesTransaction st JOIN HoleFoods.Product p ON st.Product=p.ID WHERE st.Channel='Online' AND YEAR(st.DateOfSale) = 2023 AND p.Category = 'Pasta'\"},\n",
       " {'input': 'What are the different product categories ?',\n",
       "  'query': 'SELECT DISTINCT(Category) Categories FROM HoleFoods.Product'}]"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_selector.select_examples({\"input\": \"how many products were sold in America?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "2403d14d-b637-412f-8692-45ef752e399c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import FewShotPromptTemplate\n",
    "\n",
    "example_prompt = PromptTemplate.from_template(\"User input: {input}\\nSQL query: {query}\")\n",
    "\n",
    "prompt = FewShotPromptTemplate(\n",
    "    example_selector=example_selector,\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=template,\n",
    "    suffix=\"User input: {input}\\nSQL query: \",\n",
    "    input_variables=[\"input\", \"top_k\", \"table_info\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "1608f9c8-4f62-473d-8374-1360c9b4d809",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "You are an InterSystems IRIS SQL expert. \n",
      "Given an input question, first create a syntactically correct InterSystems IRIS SQL query to run and return the answer to the input question.\n",
      "Unless the user specifies in the question a specific number of examples to obtain, query for at most 3 results using the TOP as defined in InterSystems IRIS syntax: ```SELECT [DISTINCT] TOP int select-item, select-item,...```\n",
      "Always specify table names using schema as prefix.\n",
      "Do not use LIMIT clause as it is not correct in IRIS dialect.\n",
      "Do not end SQL sentences with an ;\n",
      "Do not enclose fields in quotes or double quotes.\n",
      "Do not enclose table names in quotes or double quotes.\n",
      "You can order the results to return the most informative data in the database.\n",
      "Never query for all columns from a table. You must query only the columns that are needed to answer the question.\n",
      "Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.\n",
      "Pay attention to use CAST(CURRENT_DATE as date) function to get the current date, if the question involves \"today\".\n",
      "\n",
      "Return only plain SQL without any formatting.\n",
      "\n",
      "Only use the following tables:\n",
      "\n",
      "foo.\n",
      "Question: how many products were sold in Europe?\n",
      "\n",
      "User input: How many products were sold in Europe in 2022 ?\n",
      "SQL query: SELECT SUM(UnitsSold) FROM HoleFoods.SalesTransaction st JOIN HoleFoods.Outlet o ON st.Outlet=o.ID JOIN HoleFoods.Country c ON o.Country=c.ID JOIN HoleFoods.Region r ON c.Region=r.ID WHERE r.Name='Europe' AND YEAR(st.DateOfSale) = 2022\n",
      "\n",
      "User input: How many pasta products where sold online in 2023 ?\n",
      "SQL query: SELECT SUM(UnitsSold) FROM HoleFoods.SalesTransaction st JOIN HoleFoods.Product p ON st.Product=p.ID WHERE st.Channel='Online' AND YEAR(st.DateOfSale) = 2023 AND p.Category = 'Pasta'\n",
      "\n",
      "User input: What are the different product categories ?\n",
      "SQL query: SELECT DISTINCT(Category) Categories FROM HoleFoods.Product\n",
      "\n",
      "User input: how many products were sold in Europe?\n",
      "SQL query: \n"
     ]
    }
   ],
   "source": [
    "print(prompt.format(input=\"how many products were sold in Europe?\", top_k=3, table_info=\"foo\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "56b92a33-d644-41f6-a5cf-73f2bc74b44f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"SELECT SUM(UnitsSold) FROM HoleFoods.SalesTransaction st JOIN HoleFoods.Outlet o ON st.Outlet=o.ID JOIN HoleFoods.Country c ON o.Country=c.ID JOIN HoleFoods.Region r ON c.Region=r.ID WHERE r.Name='N. America'\""
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = create_sql_query_chain(llm, db, prompt)\n",
    "query = chain.invoke({\"question\": \"how many products were sold in America?\"})\n",
    "query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "4c10bffb-f763-460c-bb7f-0e5fd0098a86",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(531,)]\n"
     ]
    }
   ],
   "source": [
    "print(db.run(query))"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}