wip: langchain text to SQL agent

notebooks/langchain_text2SQL.ipynb

@@ -0,0 +1,330 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import getpass\n",
+    "import os\n",
+    "\n",
+    "if not os.environ.get(\"OPENAI_API_KEY\"):\n",
+    "    os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()\n",
+    "\n",
+    "# Comment out the below to opt-out of using LangSmith in this notebook. Not required.\n",
+    "# if not os.environ.get(\"LANGCHAIN_API_KEY\"):\n",
+    "#     os.environ[\"LANGCHAIN_API_KEY\"] = getpass.getpass()\n",
+    "#     os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "sqlite\n",
+      "['Drivers', 'Event', 'Laps', 'Sessions', 'Telemetry', 'Tracks', 'Weather']\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "\"[(1, 1271, 'VER', 80.0, 7087, 2, 26.0, 0, 8, -89.39, 4453.24, 0.0, 0, '2023-03-04 15:04:00.840000'), (2, 1271, 'VER', 82.0, 7112, 2, 24.0, 0, 8, -87.0, 4498.0, 0.0, 0, '2023-03-04 15:04:01.039000'), (3, 1271, 'VER', 85.0, 7401, 2, 25.0, 0, 8, -85.0, 4564.0, 0.0, 0, '2023-03-04 15:04:01.318000'), (4, 1271, 'VER', 90.0, 7908, 2, 31.0, 0, 8, -83.0, 4608.0, 0.0, 0, '2023-03-04 15:04:01.498000'), (5, 1271, 'VER', 92.0, 8132, 2, 35.0, 0, 8, -81.84, 4636.77, 0.0, 0, '2023-03-04 15:04:01.604000'), (6, 1271, 'VER', 95.0, 8378, 2, 36.0, 0, 8, -80.0, 4681.0, 0.0, 0, '2023-03-04 15:04:01.759000'), (7, 1271, 'VER', 99.0, 8625, 2, 37.0, 0, 8, -78.7, 4704.75, 0.0, 0, '2023-03-04 15:04:01.844000'), (8, 1271, 'VER', 101.0, 8925, 2, 40.0, 0, 8, -77.0, 4736.0, 0.0, 0, '2023-03-04 15:04:01.959000'), (9, 1271, 'VER', 103.0, 9225, 2, 43.0, 0, 8, -75.0, 4793.0, 0.0, 0, '2023-03-04 15:04:02.159000'), (10, 1271, 'VER', 105.0, 9526, 2, 46.0, 0, 8, -74.3, 4810.45, 0.0, 0, '2023-03-04 15:04:02.204000')]\""
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langchain_community.utilities import SQLDatabase\n",
+    "\n",
+    "db = SQLDatabase.from_uri(\"sqlite:///Bahrain_2023_Q.db\")\n",
+    "print(db.dialect)\n",
+    "print(db.get_usable_table_names())\n",
+    "db.run(\"SELECT * FROM Telemetry LIMIT 10;\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain_openai import ChatOpenAI\n",
+    "\n",
+    "llm = ChatOpenAI(model=\"gpt-4o-mini\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'SQLQuery: \\n```sql\\nSELECT \"driver_name\", MIN(\"lap_time_in_seconds\") AS \"fastest_lap_time\"\\nFROM \"Laps\"\\nGROUP BY \"driver_name\"\\nORDER BY \"fastest_lap_time\"\\nLIMIT 1;\\n```'"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langchain.chains import create_sql_query_chain\n",
+    "\n",
+    "chain = create_sql_query_chain(llm, db, )\n",
+    "response = chain.invoke({\"question\": \"Who is the fastest driver?\"})\n",
+    "response"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain_core.output_parsers import StrOutputParser\n",
+    "from langchain_core.prompts import ChatPromptTemplate\n",
+    "\n",
+    "system = \"\"\"Double check the user's {dialect} query for common mistakes, including:\n",
+    "- Only return SQL Query not anything else like ```sql ... ```\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",
+    "- Using the correct number of arguments for functions\n",
+    "- Casting to the correct data type\n",
+    "- Using the proper columns for joins\n",
+    "\n",
+    "If there are any of the above mistakes, rewrite the query.\n",
+    "If there are no mistakes, just reproduce the original query with no further commentary.\n",
+    "\n",
+    "Output the final SQL query only.\"\"\"\n",
+    "prompt = ChatPromptTemplate.from_messages(\n",
+    "    [(\"system\", system), (\"human\", \"{query}\")]\n",
+    ").partial(dialect=db.dialect)\n",
+    "validation_chain = prompt | llm | StrOutputParser()\n",
+    "\n",
+    "full_chain = {\"query\": chain} | validation_chain"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "SELECT \"driver_name\", MIN(\"lap_time_in_seconds\") AS \"fastest_lap_time\"\n",
+      "FROM \"Laps\"\n",
+      "GROUP BY \"driver_name\"\n",
+      "ORDER BY \"fastest_lap_time\" ASC\n",
+      "LIMIT 1;\n"
+     ]
+    }
+   ],
+   "source": [
+    "query = full_chain.invoke(\n",
+    "    {\n",
+    "        \"question\": \"WWho is the fastest driver?\"\n",
+    "    }\n",
+    ")\n",
+    "print(query)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\"[('VER', 89.708)]\""
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "db.run(query)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "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": [
+    "chain.get_prompts()[0].pretty_print()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "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": [
+    "full_chain.get_prompts()[0].pretty_print()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\"[('VER', 89.708)]\""
+      ]
+     },
+     "execution_count": 15,
+     "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)\n",
+    "chain = write_query | validation_chain | execute_query\n",
+    "chain.invoke({\"question\": \"Who is the fastest driver?\"})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Answer the question\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'The fastest driver is VER, with a fastest lap time of 89.708 seconds.'"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from operator import itemgetter\n",
+    "\n",
+    "from langchain_core.output_parsers import StrOutputParser\n",
+    "from langchain_core.prompts import PromptTemplate\n",
+    "from langchain_core.runnables import RunnablePassthrough\n",
+    "\n",
+    "answer_prompt = PromptTemplate.from_template(\n",
+    "    \"\"\"Given the following user question, corresponding SQL query, and SQL result, answer the user question.\n",
+    "\n",
+    "Question: {question}\n",
+    "SQL Query: {query}\n",
+    "SQL Result: {result}\n",
+    "Answer: \"\"\"\n",
+    ")\n",
+    "\n",
+    "chain = (\n",
+    "    RunnablePassthrough.assign(query=write_query).assign(\n",
+    "        result=itemgetter(\"query\") | validation_chain | execute_query\n",
+    "    )\n",
+    "    | answer_prompt\n",
+    "    | llm\n",
+    "    | StrOutputParser()\n",
+    ")\n",
+    "\n",
+    "chain.invoke({\"question\": \"Who is the fastest driver?\"})"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "llama",
+   "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.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}