{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to better prompt when doing SQL question-answering\n",
    "\n",
    "In this guide we'll go over prompting strategies to improve SQL query generation using [create_sql_query_chain](https://api.python.langchain.com/en/latest/chains/langchain.chains.sql_database.query.create_sql_query_chain.html). We'll largely focus on methods for getting relevant database-specific information in your prompt.\n",
    "\n",
    "We will cover: \n",
    "\n",
    "- How the dialect of the LangChain [SQLDatabase](https://api.python.langchain.com/en/latest/utilities/langchain_community.utilities.sql_database.SQLDatabase.html) impacts the prompt of the chain;\n",
    "- How to format schema information into the prompt using `SQLDatabase.get_context`;\n",
    "- How to build and select few-shot examples to assist the model.\n",
    "\n",
    "## Setup\n",
    "\n",
    "First, get required packages and set environment variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet  langchain langchain-community langchain-experimental langchain-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment the below to use LangSmith. Not required.\n",
    "# import os\n",
    "# os.environ[\"LANGCHAIN_API_KEY\"] = getpass.getpass()\n",
    "# os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The below example will use a SQLite connection with Chinook database. Follow [these installation steps](https://database.guide/2-sample-databases-sqlite/) to create `Chinook.db` in the same directory as this notebook:\n",
    "\n",
    "* Save [this file](https://raw.githubusercontent.com/lerocha/chinook-database/master/ChinookDatabase/DataSources/Chinook_Sqlite.sql) as `Chinook_Sqlite.sql`\n",
    "* Run `sqlite3 Chinook.db`\n",
    "* Run `.read Chinook_Sqlite.sql`\n",
    "* Test `SELECT * FROM Artist LIMIT 10;`\n",
    "\n",
    "Now, `Chinhook.db` is in our directory and we can interface with it using the SQLAlchemy-driven `SQLDatabase` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sqlite\n",
      "['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']\n",
      "[(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham')]\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.utilities import SQLDatabase\n",
    "\n",
    "db = SQLDatabase.from_uri(\"sqlite:///Chinook.db\", sample_rows_in_table_info=3)\n",
    "print(db.dialect)\n",
    "print(db.get_usable_table_names())\n",
    "print(db.run(\"SELECT * FROM Artist LIMIT 10;\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dialect-specific prompting\n",
    "\n",
    "One of the simplest things we can do is make our prompt specific to the SQL dialect we're using. When using the built-in [create_sql_query_chain](https://api.python.langchain.com/en/latest/chains/langchain.chains.sql_database.query.create_sql_query_chain.html) and [SQLDatabase](https://api.python.langchain.com/en/latest/utilities/langchain_community.utilities.sql_database.SQLDatabase.html), this is handled for you for any of the following dialects:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['crate',\n",
       " 'duckdb',\n",
       " 'googlesql',\n",
       " 'mssql',\n",
       " 'mysql',\n",
       " 'mariadb',\n",
       " 'oracle',\n",
       " 'postgresql',\n",
       " 'sqlite',\n",
       " 'clickhouse',\n",
       " 'prestodb']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.chains.sql_database.prompt import SQL_PROMPTS\n",
    "\n",
    "list(SQL_PROMPTS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, using our current DB we can see that we'll get a SQLite-specific prompt.\n",
    "\n",
    "```{=mdx}\n",
    "import ChatModelTabs from \"@theme/ChatModelTabs\";\n",
    "\n",
    "<ChatModelTabs customVarName=\"llm\" />\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | output: false\n",
    "# | echo: false\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are a SQLite expert. Given an input question, first create a syntactically correct SQLite query to run, then look at the results of the query 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 LIMIT clause as per SQLite. 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. Wrap each column name in double quotes (\") to denote them as delimited identifiers.\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 date('now') function to get the current date, if the question involves \"today\".\n",
      "\n",
      "Use the following format:\n",
      "\n",
      "Question: Question here\n",
      "SQLQuery: SQL Query to run\n",
      "SQLResult: Result of the SQLQuery\n",
      "Answer: Final answer here\n",
      "\n",
      "Only use the following tables:\n",
      "\u001b[33;1m\u001b[1;3m{table_info}\u001b[0m\n",
      "\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)\n",
    "chain.get_prompts()[0].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table definitions and example rows\n",
    "\n",
    "In most SQL chains, we'll need to feed the model at least part of the database schema. Without this it won't be able to write valid queries. Our database comes with some convenience methods to give us the relevant context. Specifically, we can get the table names, their schemas, and a sample of rows from each table.\n",
    "\n",
    "Here we will use `SQLDatabase.get_context`, which provides available tables and their schemas:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['table_info', 'table_names']\n",
      "\n",
      "CREATE TABLE \"Album\" (\n",
      "\t\"AlbumId\" INTEGER NOT NULL, \n",
      "\t\"Title\" NVARCHAR(160) NOT NULL, \n",
      "\t\"ArtistId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"AlbumId\"), \n",
      "\tFOREIGN KEY(\"ArtistId\") REFERENCES \"Artist\" (\"ArtistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Album table:\n",
      "AlbumId\tTitle\tArtistId\n",
      "1\tFor Those About To Rock We Salute You\t1\n",
      "2\tBalls to the Wall\t2\n",
      "3\tRestless and Wild\t2\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Artist\" (\n",
      "\t\"ArtistId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"ArtistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Artist table:\n",
      "ArtistId\tName\n",
      "1\tAC/DC\n",
      "2\tAccept\n",
      "3\tAerosmith\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Customer\" (\n",
      "\t\"CustomerId\" INTEGER NOT NULL, \n",
      "\t\"FirstName\" NVARCHAR(40) NOT NULL, \n",
      "\t\"LastName\" NVARCHAR(20) NOT NULL, \n",
      "\t\"Company\" NVARCHAR(80), \n",
      "\t\"Address\" NVARCHAR(70), \n",
      "\t\"City\" NVARCHAR(40), \n",
      "\t\"State\" NVARCHAR(40), \n",
      "\t\"Country\" NVARCHAR(40), \n",
      "\t\"PostalCode\" NVARCHAR(10), \n",
      "\t\"Phone\" NVARCHAR(24), \n",
      "\t\"Fax\" NVARCHAR(24), \n",
      "\t\"Email\" NVARCHAR(60) NOT NULL, \n",
      "\t\"SupportRepId\" INTEGER, \n",
      "\tPRIMARY KEY (\"CustomerId\"), \n",
      "\tFOREIGN KEY(\"SupportRepId\") REFERENCES \"Employee\" (\"EmployeeId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Customer table:\n",
      "CustomerId\tFirstName\tLastName\tCompany\tAddress\tCity\tState\tCountry\tPostalCode\tPhone\tFax\tEmail\tSupportRepId\n",
      "1\tLuís\tGonçalves\tEmbraer - Empresa Brasileira de Aeronáutica S.A.\tAv. Brigadeiro Faria Lima, 2170\tSão José dos Campos\tSP\tBrazil\t12227-000\t+55 (12) 3923-5555\t+55 (12) 3923-5566\tluisg@embraer.com.br\t3\n",
      "2\tLeonie\tKöhler\tNone\tTheodor-Heuss-Straße 34\tStuttgart\tNone\tGermany\t70174\t+49 0711 2842222\tNone\tleonekohler@surfeu.de\t5\n",
      "3\tFrançois\tTremblay\tNone\t1498 rue Bélanger\tMontréal\tQC\tCanada\tH2G 1A7\t+1 (514) 721-4711\tNone\tftremblay@gmail.com\t3\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Employee\" (\n",
      "\t\"EmployeeId\" INTEGER NOT NULL, \n",
      "\t\"LastName\" NVARCHAR(20) NOT NULL, \n",
      "\t\"FirstName\" NVARCHAR(20) NOT NULL, \n",
      "\t\"Title\" NVARCHAR(30), \n",
      "\t\"ReportsTo\" INTEGER, \n",
      "\t\"BirthDate\" DATETIME, \n",
      "\t\"HireDate\" DATETIME, \n",
      "\t\"Address\" NVARCHAR(70), \n",
      "\t\"City\" NVARCHAR(40), \n",
      "\t\"State\" NVARCHAR(40), \n",
      "\t\"Country\" NVARCHAR(40), \n",
      "\t\"PostalCode\" NVARCHAR(10), \n",
      "\t\"Phone\" NVARCHAR(24), \n",
      "\t\"Fax\" NVARCHAR(24), \n",
      "\t\"Email\" NVARCHAR(60), \n",
      "\tPRIMARY KEY (\"EmployeeId\"), \n",
      "\tFOREIGN KEY(\"ReportsTo\") REFERENCES \"Employee\" (\"EmployeeId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Employee table:\n",
      "EmployeeId\tLastName\tFirstName\tTitle\tReportsTo\tBirthDate\tHireDate\tAddress\tCity\tState\tCountry\tPostalCode\tPhone\tFax\tEmail\n",
      "1\tAdams\tAndrew\tGeneral Manager\tNone\t1962-02-18 00:00:00\t2002-08-14 00:00:00\t11120 Jasper Ave NW\tEdmonton\tAB\tCanada\tT5K 2N1\t+1 (780) 428-9482\t+1 (780) 428-3457\tandrew@chinookcorp.com\n",
      "2\tEdwards\tNancy\tSales Manager\t1\t1958-12-08 00:00:00\t2002-05-01 00:00:00\t825 8 Ave SW\tCalgary\tAB\tCanada\tT2P 2T3\t+1 (403) 262-3443\t+1 (403) 262-3322\tnancy@chinookcorp.com\n",
      "3\tPeacock\tJane\tSales Support Agent\t2\t1973-08-29 00:00:00\t2002-04-01 00:00:00\t1111 6 Ave SW\tCalgary\tAB\tCanada\tT2P 5M5\t+1 (403) 262-3443\t+1 (403) 262-6712\tjane@chinookcorp.com\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Genre\" (\n",
      "\t\"GenreId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"GenreId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Genre table:\n",
      "GenreId\tName\n",
      "1\tRock\n",
      "2\tJazz\n",
      "3\tMetal\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Invoice\" (\n",
      "\t\"InvoiceId\" INTEGER NOT NULL, \n",
      "\t\"CustomerId\" INTEGER NOT NULL, \n",
      "\t\"InvoiceDate\" DATETIME NOT NULL, \n",
      "\t\"BillingAddress\" NVARCHAR(70), \n",
      "\t\"BillingCity\" NVARCHAR(40), \n",
      "\t\"BillingState\" NVARCHAR(40), \n",
      "\t\"BillingCountry\" NVARCHAR(40), \n",
      "\t\"BillingPostalCode\" NVARCHAR(10), \n",
      "\t\"Total\" NUMERIC(10, 2) NOT NULL, \n",
      "\tPRIMARY KEY (\"InvoiceId\"), \n",
      "\tFOREIGN KEY(\"CustomerId\") REFERENCES \"Customer\" (\"CustomerId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Invoice table:\n",
      "InvoiceId\tCustomerId\tInvoiceDate\tBillingAddress\tBillingCity\tBillingState\tBillingCountry\tBillingPostalCode\tTotal\n",
      "1\t2\t2021-01-01 00:00:00\tTheodor-Heuss-Straße 34\tStuttgart\tNone\tGermany\t70174\t1.98\n",
      "2\t4\t2021-01-02 00:00:00\tUllevålsveien 14\tOslo\tNone\tNorway\t0171\t3.96\n",
      "3\t8\t2021-01-03 00:00:00\tGrétrystraat 63\tBrussels\tNone\tBelgium\t1000\t5.94\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"InvoiceLine\" (\n",
      "\t\"InvoiceLineId\" INTEGER NOT NULL, \n",
      "\t\"InvoiceId\" INTEGER NOT NULL, \n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\t\"UnitPrice\" NUMERIC(10, 2) NOT NULL, \n",
      "\t\"Quantity\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"InvoiceLineId\"), \n",
      "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"InvoiceId\") REFERENCES \"Invoice\" (\"InvoiceId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from InvoiceLine table:\n",
      "InvoiceLineId\tInvoiceId\tTrackId\tUnitPrice\tQuantity\n",
      "1\t1\t2\t0.99\t1\n",
      "2\t1\t4\t0.99\t1\n",
      "3\t2\t6\t0.99\t1\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"MediaType\" (\n",
      "\t\"MediaTypeId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"MediaTypeId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from MediaType table:\n",
      "MediaTypeId\tName\n",
      "1\tMPEG audio file\n",
      "2\tProtected AAC audio file\n",
      "3\tProtected MPEG-4 video file\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Playlist\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120), \n",
      "\tPRIMARY KEY (\"PlaylistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Playlist table:\n",
      "PlaylistId\tName\n",
      "1\tMusic\n",
      "2\tMovies\n",
      "3\tTV Shows\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"PlaylistTrack\" (\n",
      "\t\"PlaylistId\" INTEGER NOT NULL, \n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"PlaylistId\", \"TrackId\"), \n",
      "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"PlaylistId\") REFERENCES \"Playlist\" (\"PlaylistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from PlaylistTrack table:\n",
      "PlaylistId\tTrackId\n",
      "1\t3402\n",
      "1\t3389\n",
      "1\t3390\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Track\" (\n",
      "\t\"TrackId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(200) NOT NULL, \n",
      "\t\"AlbumId\" INTEGER, \n",
      "\t\"MediaTypeId\" INTEGER NOT NULL, \n",
      "\t\"GenreId\" INTEGER, \n",
      "\t\"Composer\" NVARCHAR(220), \n",
      "\t\"Milliseconds\" INTEGER NOT NULL, \n",
      "\t\"Bytes\" INTEGER, \n",
      "\t\"UnitPrice\" NUMERIC(10, 2) NOT NULL, \n",
      "\tPRIMARY KEY (\"TrackId\"), \n",
      "\tFOREIGN KEY(\"MediaTypeId\") REFERENCES \"MediaType\" (\"MediaTypeId\"), \n",
      "\tFOREIGN KEY(\"GenreId\") REFERENCES \"Genre\" (\"GenreId\"), \n",
      "\tFOREIGN KEY(\"AlbumId\") REFERENCES \"Album\" (\"AlbumId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Track table:\n",
      "TrackId\tName\tAlbumId\tMediaTypeId\tGenreId\tComposer\tMilliseconds\tBytes\tUnitPrice\n",
      "1\tFor Those About To Rock (We Salute You)\t1\t1\t1\tAngus Young, Malcolm Young, Brian Johnson\t343719\t11170334\t0.99\n",
      "2\tBalls to the Wall\t2\t2\t1\tU. Dirkschneider, W. Hoffmann, H. Frank, P. Baltes, S. Kaufmann, G. Hoffmann\t342562\t5510424\t0.99\n",
      "3\tFast As a Shark\t3\t2\t1\tF. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman\t230619\t3990994\t0.99\n",
      "*/\n"
     ]
    }
   ],
   "source": [
    "context = db.get_context()\n",
    "print(list(context))\n",
    "print(context[\"table_info\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we don't have too many, or too wide of, tables, we can just insert the entirety of this information in our prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are a SQLite expert. Given an input question, first create a syntactically correct SQLite query to run, then look at the results of the query 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 LIMIT clause as per SQLite. 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. Wrap each column name in double quotes (\") to denote them as delimited identifiers.\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 date('now') function to get the current date, if the question involves \"today\".\n",
      "\n",
      "Use the following format:\n",
      "\n",
      "Question: Question here\n",
      "SQLQuery: SQL Query to run\n",
      "SQLResult: Result of the SQLQuery\n",
      "Answer: Final answer here\n",
      "\n",
      "Only use the following tables:\n",
      "\n",
      "CREATE TABLE \"Album\" (\n",
      "\t\"AlbumId\" INTEGER NOT NULL, \n",
      "\t\"Title\" NVARCHAR(160) NOT NULL, \n",
      "\t\"ArtistId\" INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (\"AlbumId\"), \n",
      "\tFOREIGN KEY(\"ArtistId\") REFERENCES \"Artist\" (\"ArtistId\")\n",
      ")\n",
      "\n",
      "/*\n",
      "3 rows from Album table:\n",
      "AlbumId\tTitle\tArtistId\n",
      "1\tFor Those About To Rock We Salute You\t1\n",
      "2\tBalls to the Wall\t2\n",
      "3\tRestless and Wild\t2\n",
      "*/\n",
      "\n",
      "\n",
      "CREATE TABLE \"Artist\" (\n",
      "\t\"ArtistId\" INTEGER NOT NULL, \n",
      "\t\"Name\" NVARCHAR(120)\n"
     ]
    }
   ],
   "source": [
    "prompt_with_context = chain.get_prompts()[0].partial(table_info=context[\"table_info\"])\n",
    "print(prompt_with_context.pretty_repr()[:1500])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we do have database schemas that are too large to fit into our model's context window, we'll need to come up with ways of inserting only the relevant table definitions into the prompt based on the user input. For more on this head to the [Many tables, wide tables, high-cardinality feature](/docs/how_to/sql_large_db) guide."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Few-shot examples\n",
    "\n",
    "Including examples of natural language questions being converted to valid SQL queries against our database in the prompt will often improve model performance, especially for complex queries.\n",
    "\n",
    "Let's say we have the following examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "examples = [\n",
    "    {\"input\": \"List all artists.\", \"query\": \"SELECT * FROM Artist;\"},\n",
    "    {\n",
    "        \"input\": \"Find all albums for the artist 'AC/DC'.\",\n",
    "        \"query\": \"SELECT * FROM Album WHERE ArtistId = (SELECT ArtistId FROM Artist WHERE Name = 'AC/DC');\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"List all tracks in the 'Rock' genre.\",\n",
    "        \"query\": \"SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Find the total duration of all tracks.\",\n",
    "        \"query\": \"SELECT SUM(Milliseconds) FROM Track;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"List all customers from Canada.\",\n",
    "        \"query\": \"SELECT * FROM Customer WHERE Country = 'Canada';\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"How many tracks are there in the album with ID 5?\",\n",
    "        \"query\": \"SELECT COUNT(*) FROM Track WHERE AlbumId = 5;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Find the total number of invoices.\",\n",
    "        \"query\": \"SELECT COUNT(*) FROM Invoice;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"List all tracks that are longer than 5 minutes.\",\n",
    "        \"query\": \"SELECT * FROM Track WHERE Milliseconds > 300000;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Who are the top 5 customers by total purchase?\",\n",
    "        \"query\": \"SELECT CustomerId, SUM(Total) AS TotalPurchase FROM Invoice GROUP BY CustomerId ORDER BY TotalPurchase DESC LIMIT 5;\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"Which albums are from the year 2000?\",\n",
    "        \"query\": \"SELECT * FROM Album WHERE strftime('%Y', ReleaseDate) = '2000';\",\n",
    "    },\n",
    "    {\n",
    "        \"input\": \"How many employees are there\",\n",
    "        \"query\": 'SELECT COUNT(*) FROM \"Employee\"',\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can create a few-shot prompt with them like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate\n",
    "\n",
    "example_prompt = PromptTemplate.from_template(\"User input: {input}\\nSQL query: {query}\")\n",
    "prompt = FewShotPromptTemplate(\n",
    "    examples=examples[:5],\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=\"You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run. Unless otherwise specificed, do not return more than {top_k} rows.\\n\\nHere is the relevant table info: {table_info}\\n\\nBelow are a number of examples of questions and their corresponding SQL queries.\",\n",
    "    suffix=\"User input: {input}\\nSQL query: \",\n",
    "    input_variables=[\"input\", \"top_k\", \"table_info\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run. Unless otherwise specificed, do not return more than 3 rows.\n",
      "\n",
      "Here is the relevant table info: foo\n",
      "\n",
      "Below are a number of examples of questions and their corresponding SQL queries.\n",
      "\n",
      "User input: List all artists.\n",
      "SQL query: SELECT * FROM Artist;\n",
      "\n",
      "User input: Find all albums for the artist 'AC/DC'.\n",
      "SQL query: SELECT * FROM Album WHERE ArtistId = (SELECT ArtistId FROM Artist WHERE Name = 'AC/DC');\n",
      "\n",
      "User input: List all tracks in the 'Rock' genre.\n",
      "SQL query: SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');\n",
      "\n",
      "User input: Find the total duration of all tracks.\n",
      "SQL query: SELECT SUM(Milliseconds) FROM Track;\n",
      "\n",
      "User input: List all customers from Canada.\n",
      "SQL query: SELECT * FROM Customer WHERE Country = 'Canada';\n",
      "\n",
      "User input: How many artists are there?\n",
      "SQL query: \n"
     ]
    }
   ],
   "source": [
    "print(prompt.format(input=\"How many artists are there?\", top_k=3, table_info=\"foo\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamic few-shot examples\n",
    "\n",
    "If we have enough examples, we may want to only include the most relevant ones in the prompt, either because they don't fit in the model's context window or because the long tail of examples distracts the model. And specifically, given any input we want to include the examples most relevant to that input.\n",
    "\n",
    "We can do just this using an ExampleSelector. In this case we'll use a [SemanticSimilarityExampleSelector](https://api.python.langchain.com/en/latest/example_selectors/langchain_core.example_selectors.semantic_similarity.SemanticSimilarityExampleSelector.html), which will store the examples in the vector database of our choosing. At runtime it will perform a similarity search between the input and our examples, and return the most semantically similar ones.\n",
    "\n",
    "We default to OpenAI embeddings here, but you can swap them out for the model provider of your choice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.vectorstores import FAISS\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",
    "    FAISS,\n",
    "    k=5,\n",
    "    input_keys=[\"input\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'input': 'List all artists.', 'query': 'SELECT * FROM Artist;'},\n",
       " {'input': 'How many employees are there',\n",
       "  'query': 'SELECT COUNT(*) FROM \"Employee\"'},\n",
       " {'input': 'How many tracks are there in the album with ID 5?',\n",
       "  'query': 'SELECT COUNT(*) FROM Track WHERE AlbumId = 5;'},\n",
       " {'input': 'Which albums are from the year 2000?',\n",
       "  'query': \"SELECT * FROM Album WHERE strftime('%Y', ReleaseDate) = '2000';\"},\n",
       " {'input': \"List all tracks in the 'Rock' genre.\",\n",
       "  'query': \"SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');\"}]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_selector.select_examples({\"input\": \"how many artists are there?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use it, we can pass the ExampleSelector directly in to our FewShotPromptTemplate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = FewShotPromptTemplate(\n",
    "    example_selector=example_selector,\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=\"You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run. Unless otherwise specificed, do not return more than {top_k} rows.\\n\\nHere is the relevant table info: {table_info}\\n\\nBelow are a number of examples of questions and their corresponding SQL queries.\",\n",
    "    suffix=\"User input: {input}\\nSQL query: \",\n",
    "    input_variables=[\"input\", \"top_k\", \"table_info\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are a SQLite expert. Given an input question, create a syntactically correct SQLite query to run. Unless otherwise specificed, do not return more than 3 rows.\n",
      "\n",
      "Here is the relevant table info: foo\n",
      "\n",
      "Below are a number of examples of questions and their corresponding SQL queries.\n",
      "\n",
      "User input: List all artists.\n",
      "SQL query: SELECT * FROM Artist;\n",
      "\n",
      "User input: How many employees are there\n",
      "SQL query: SELECT COUNT(*) FROM \"Employee\"\n",
      "\n",
      "User input: How many tracks are there in the album with ID 5?\n",
      "SQL query: SELECT COUNT(*) FROM Track WHERE AlbumId = 5;\n",
      "\n",
      "User input: Which albums are from the year 2000?\n",
      "SQL query: SELECT * FROM Album WHERE strftime('%Y', ReleaseDate) = '2000';\n",
      "\n",
      "User input: List all tracks in the 'Rock' genre.\n",
      "SQL query: SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');\n",
      "\n",
      "User input: how many artists are there?\n",
      "SQL query: \n"
     ]
    }
   ],
   "source": [
    "print(prompt.format(input=\"how many artists are there?\", top_k=3, table_info=\"foo\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Trying it out, we see that the model identifies the relevant table:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'SELECT COUNT(*) FROM Artist;'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = create_sql_query_chain(llm, db, prompt)\n",
    "chain.invoke({\"question\": \"how many artists are there?\"})"
   ]
  }
 ],
 "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}