docs/docs/tutorials/sql_qa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a Question/Answering system over SQL data\n",
    "\n",
    ":::info Prerequisites\n",
    "\n",
    "This guide assumes familiarity with the following concepts:\n",
    "\n",
    "- [Chat models](/docs/concepts/chat_models)\n",
    "- [Tools](/docs/concepts/tools)\n",
    "- [Agents](/docs/concepts/agents)\n",
    "- [LangGraph](/docs/concepts/architecture/#langgraph)\n",
    "\n",
    ":::\n",
    "\n",
    "Enabling a LLM system to query structured data can be qualitatively different from unstructured text data. Whereas in the latter it is common to generate text that can be searched against a vector database, the approach for structured data is often for the LLM to write and execute queries in a DSL, such as SQL. In this guide we'll go over the basic ways to create a Q&A system over tabular data in databases. We will cover implementations using both [chains](/docs/tutorials/sql_qa#chains) and [agents](/docs/tutorials/sql_qa#agents). These systems will allow us to ask a question about the data in a database and get back a natural language answer. The main difference between the two is that our agent can query the database in a loop as many times as it needs to answer the question.\n",
    "\n",
    "## ⚠️ Security note ⚠️\n",
    "\n",
    "Building Q&A systems of SQL databases requires executing model-generated SQL queries. There are inherent risks in doing this. Make sure that your database connection permissions are always scoped as narrowly as possible for your chain/agent's needs. This will mitigate though not eliminate the risks of building a model-driven system. For more on general security best practices, [see here](/docs/security).\n",
    "\n",
    "\n",
    "## Architecture\n",
    "\n",
    "At a high-level, the steps of these systems are:\n",
    "\n",
    "1. **Convert question to SQL query**: Model converts user input to a SQL query.\n",
    "2. **Execute SQL query**: Execute the query.\n",
    "3. **Answer the question**: Model responds to user input using the query results.\n",
    "\n",
    "Note that querying data in CSVs can follow a similar approach. See our [how-to guide](/docs/how_to/sql_csv) on question-answering over CSV data for more detail.\n",
    "\n",
    "![sql_usecase.png](../../static/img/sql_usecase.png)\n",
    "\n",
    "## Setup\n",
    "\n",
    "First, get required packages and set environment variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install --upgrade --quiet langchain-community langgraph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "# Comment out the below to opt-out of using LangSmith in this notebook. Not required.\n",
    "if not os.environ.get(\"LANGSMITH_API_KEY\"):\n",
    "    os.environ[\"LANGSMITH_API_KEY\"] = getpass.getpass()\n",
    "    os.environ[\"LANGSMITH_TRACING\"] = \"true\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample data\n",
    "\n",
    "The below example will use a SQLite connection with the Chinook database, which is a sample database that represents a digital media store. Follow [these installation steps](https://database.guide/2-sample-databases-sqlite/) to create `Chinook.db` in the same directory as this notebook. You can also download and build the database via the command line:\n",
    "```bash\n",
    "curl -s https://raw.githubusercontent.com/lerocha/chinook-database/master/ChinookDatabase/DataSources/Chinook_Sqlite.sql | sqlite3 Chinook.db\n",
    "```\n",
    "\n",
    "Now, `Chinook.db` is in our directory and we can interface with it using the SQLAlchemy-driven `SQLDatabase` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sqlite\n",
      "['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"[(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')]\""
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_community.utilities import SQLDatabase\n",
    "\n",
    "db = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")\n",
    "print(db.dialect)\n",
    "print(db.get_usable_table_names())\n",
    "db.run(\"SELECT * FROM Artist LIMIT 10;\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! We've got a SQL database that we can query. Now let's try hooking it up to an LLM.\n",
    "\n",
    "## Chains {#chains}\n",
    "\n",
    "Chains are compositions of predictable steps. In [LangGraph](/docs/concepts/architecture/#langgraph), we can represent a chain via simple sequence of nodes. Let's create a sequence of steps that, given a question, does the following:\n",
    "- converts the question into a SQL query;\n",
    "- executes the query;\n",
    "- uses the result to answer the original question.\n",
    "\n",
    "There are scenarios not supported by this arrangement. For example, this system will execute a SQL query for any user input-- even \"hello\". Importantly, as we'll see below, some questions require more than one query to answer. We will address these scenarios in the Agents section.\n",
    "\n",
    "### Application state\n",
    "\n",
    "The LangGraph [state](https://langchain-ai.github.io/langgraph/concepts/low_level/#state) of our application controls what data is input to the application, transferred between steps, and output by the application. It is typically a `TypedDict`, but can also be a [Pydantic BaseModel](https://langchain-ai.github.io/langgraph/how-tos/state-model/).\n",
    "\n",
    "For this application, we can just keep track of the input question, generated query, query result, and generated answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing_extensions import TypedDict\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    question: str\n",
    "    query: str\n",
    "    result: str\n",
    "    answer: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we just need functions that operate on this state and populate its contents.\n",
    "\n",
    "### Convert question to SQL query\n",
    "\n",
    "The first step is to take the user input and convert it to a SQL query. To reliably obtain SQL queries (absent markdown formatting and explanations or clarifications), we will make use of LangChain's [structured output](/docs/concepts/structured_outputs/) abstraction.\n",
    "\n",
    "Let's select a chat model for our application:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "import ChatModelTabs from \"@theme/ChatModelTabs\";\n",
    "\n",
    "<ChatModelTabs customVarName=\"llm\" />\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | output: false\n",
    "# | echo: false\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o\", temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's provide some instructions for our model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m System Message \u001b[0m================================\n",
      "\n",
      "\n",
      "Given an input question, create a syntactically correct \u001b[33;1m\u001b[1;3m{dialect}\u001b[0m query to\n",
      "run to help find the answer. Unless the user specifies in his question a\n",
      "specific number of examples they wish to obtain, always limit your query to\n",
      "at most \u001b[33;1m\u001b[1;3m{top_k}\u001b[0m results. You can order the results by a relevant column to\n",
      "return the most interesting examples in the database.\n",
      "\n",
      "Never query for all the columns from a specific table, only ask for a the\n",
      "few relevant columns given the question.\n",
      "\n",
      "Pay attention to use only the column names that you can see in the schema\n",
      "description. Be careful to not query for columns that do not exist. Also,\n",
      "pay attention to which column is in which table.\n",
      "\n",
      "Only use the following tables:\n",
      "\u001b[33;1m\u001b[1;3m{table_info}\u001b[0m\n",
      "\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Question: \u001b[33;1m\u001b[1;3m{input}\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "system_message = \"\"\"\n",
    "Given an input question, create a syntactically correct {dialect} query to\n",
    "run to help find the answer. Unless the user specifies in his question a\n",
    "specific number of examples they wish to obtain, always limit your query to\n",
    "at most {top_k} results. You can order the results by a relevant column to\n",
    "return the most interesting examples in the database.\n",
    "\n",
    "Never query for all the columns from a specific table, only ask for a the\n",
    "few relevant columns given the question.\n",
    "\n",
    "Pay attention to use only the column names that you can see in the schema\n",
    "description. Be careful to not query for columns that do not exist. Also,\n",
    "pay attention to which column is in which table.\n",
    "\n",
    "Only use the following tables:\n",
    "{table_info}\n",
    "\"\"\"\n",
    "\n",
    "user_prompt = \"Question: {input}\"\n",
    "\n",
    "query_prompt_template = ChatPromptTemplate(\n",
    "    [(\"system\", system_message), (\"user\", user_prompt)]\n",
    ")\n",
    "\n",
    "for message in query_prompt_template.messages:\n",
    "    message.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The prompt includes several parameters we will need to populate, such as the SQL dialect and table schemas. LangChain's [SQLDatabase](https://python.langchain.com/api_reference/community/utilities/langchain_community.utilities.sql_database.SQLDatabase.html) object includes methods to help with this. Our `write_query` step will just populate these parameters and prompt a model to generate the SQL query:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing_extensions import Annotated\n",
    "\n",
    "\n",
    "class QueryOutput(TypedDict):\n",
    "    \"\"\"Generated SQL query.\"\"\"\n",
    "\n",
    "    query: Annotated[str, ..., \"Syntactically valid SQL query.\"]\n",
    "\n",
    "\n",
    "def write_query(state: State):\n",
    "    \"\"\"Generate SQL query to fetch information.\"\"\"\n",
    "    prompt = query_prompt_template.invoke(\n",
    "        {\n",
    "            \"dialect\": db.dialect,\n",
    "            \"top_k\": 10,\n",
    "            \"table_info\": db.get_table_info(),\n",
    "            \"input\": state[\"question\"],\n",
    "        }\n",
    "    )\n",
    "    structured_llm = llm.with_structured_output(QueryOutput)\n",
    "    result = structured_llm.invoke(prompt)\n",
    "    return {\"query\": result[\"query\"]}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test it out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'query': 'SELECT COUNT(*) as employee_count FROM Employee;'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "write_query({\"question\": \"How many Employees are there?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Execute query\n",
    "\n",
    "**This is the most dangerous part of creating a SQL chain.** Consider carefully if it is OK to run automated queries over your data. Minimize the database connection permissions as much as possible. Consider adding a human approval step to you chains before query execution (see below).\n",
    "\n",
    "To execute the query, we will load a tool from [langchain-community](/docs/concepts/architecture/#langchain-community). Our `execute_query` node will just wrap this tool:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.tools.sql_database.tool import QuerySQLDatabaseTool\n",
    "\n",
    "\n",
    "def execute_query(state: State):\n",
    "    \"\"\"Execute SQL query.\"\"\"\n",
    "    execute_query_tool = QuerySQLDatabaseTool(db=db)\n",
    "    return {\"result\": execute_query_tool.invoke(state[\"query\"])}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Testing this step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'result': '[(8,)]'}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "execute_query({\"query\": \"SELECT COUNT(EmployeeId) AS EmployeeCount FROM Employee;\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate answer\n",
    "\n",
    "Finally, our last step generates an answer to the question given the information pulled from the database:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_answer(state: State):\n",
    "    \"\"\"Answer question using retrieved information as context.\"\"\"\n",
    "    prompt = (\n",
    "        \"Given the following user question, corresponding SQL query, \"\n",
    "        \"and SQL result, answer the user question.\\n\\n\"\n",
    "        f'Question: {state[\"question\"]}\\n'\n",
    "        f'SQL Query: {state[\"query\"]}\\n'\n",
    "        f'SQL Result: {state[\"result\"]}'\n",
    "    )\n",
    "    response = llm.invoke(prompt)\n",
    "    return {\"answer\": response.content}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Orchestrating with LangGraph\n",
    "\n",
    "Finally, we compile our application into a single `graph` object. In this case, we are just connecting the three steps into a single sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import START, StateGraph\n",
    "\n",
    "graph_builder = StateGraph(State).add_sequence(\n",
    "    [write_query, execute_query, generate_answer]\n",
    ")\n",
    "graph_builder.add_edge(START, \"write_query\")\n",
    "graph = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LangGraph also comes with built-in utilities for visualizing the control flow of your application:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAKMAAAFNCAIAAADjN0iRAAAAAXNSR0IArs4c6QAAIABJREFUeJztnXdcU9f7x092SEhYYc+wQRQkiCgoLkDcGxdWrXUVV90djg7t0G9bq7aO1m2tHdY6ce9VFURkQ5hhhoTsdZPfH7dFfpURNOQC575fvHjd3HvOc57cT85zz733DILBYAA4EEDE2gEcM4ErDQu40rCAKw0LuNKwgCsNC2SsHWieqhKlUqqXS3WI1qBW6rF2p20odCKZRGCwSQwWycGDTiIRsPbovxA61f107hMpP1POz5R7BjH0egOTRbZxpGpUXUBpqgVRXKtRSBC1Aqnkq9z8Gd4hzIA+LAq1s0TNzqJ05v2Ge38JvYIZ3BAmN4RJpnSWE/R6lGTLizLlgkKlX29WZIIt1u6ATqG0UKBOPVzlzLXoP8aOZkHC1hmT8/CCMO26OD7Z0bunJbaeYKx07hPpk6uiUfOc2bYUDN3oULQa/c3fatl2FGwrN5ZKF2fJ855I45OdsHLAnDy8ICSSCH3iMRMbM6XTrouqS9XD34JCZpQH5+tkYmTYdEdMSsem4VOaoyjNVUAlMwAgagSHziCl3RBhUjoGSssadM9ui8cudDV/0ZgTM44jrtGW5ynMXzQGSt89XRfAY5m/3E5CrwFWt07Vmb9ccytdW6EWVWv8w+FV2s6ZxnGl5j6WmrlccyudebchZjzHzIV2NmLGcPLTu7XSOo0+57HUzZdhzkI7IQw2WSFBqktV5izUrErzX8i5PZjmLBEAcPLkyU2bNr1GxrVr1545c6YDPAIAAG4Ik58p7yDjzWJWpSv5Kt8wcz8UzM7ONnNGY/DpZVknUHec/Vcxq9JVxSqWTUe9J01LS5s3b96gQYMGDBjw9ttvP336FAAwf/78M2fOnD17NiIiIjc3FwBw8eLFGTNmDBgwYOjQoStWrCgvL0eznzx5Mi4u7ubNm3Fxcd98801ERIRAINi8efOgQYM6wlu2LbksT9kRllvEYEZ+2lgkFWk7wrJCoRg4cOBnn31WVFRUWFi4devW6OjohoYGqVQ6Y8aM9evXi0QinU6XmZnJ4/F27drF5/MzMzMXLFgwdepU1MIff/wRHR29cOHCO3fulJeXV1dX83i8EydOiMXijnDYYDDsWVeoUug6yPirmLUnglyCMNkd8raqqqpKLpePGDGCy+UCAFatWhUXF0elUul0OplMplKp1tbWAABPT88jR474+fmRyWQAwPTp09977736+npbW1sCgaBSqaZPnx4dHQ0AUKvVAAAGg2FlZdURDgMAmGySXIKY7fWd+ZQ26A0WDCKB2CGdMTw8PDw9PT/88MNJkyZFRUUFBATweLxXk1laWlZUVOzcubOsrEylUmm1WgCARCKxtf3nxUPPnj07wr1moTNJesR8Lx3Md50mEAmAQFBIdR1hnEQi7d+/f9iwYadOnZo5c+bo0aPPnTv3arJLly6tW7cuJCRkx44dx48f/+CDD/6TwNLSfA1GUY2GyTZfTTNriwyNVx1k3MbGZvny5adPnz558mRkZOTGjRtfbTyfOnUqIiJi0aJFXl5eHA5HpTLrHW1T9IhBrdRbWJqv54VZlXb2pis7pk5XVFTcuHED3fb29n7//feJRGJhYSG6p/HNrEajQS/YKBcvXmx69FU67pWurEHnFWzWRwtmVZrjSitI75DHBVVVVWvWrDl69GhxcXFJScn+/fuJRCJ60WWxWLm5ubm5uWKxOCQk5MGDB5mZmZWVlVu3buVwOACArKysVys3jUaj0WhPnz7Nzc3V6Uz/6+Q/l7Nszdsx12ytfIPBIJdo939Y1EHGz549O3Xq1Ojo6IEDB86ePfv27dvo/jt37gwZMiQ6OvrevXtisfi9996LiYlJSEjYs2cPgiCLFy+Oioq6cOHCqVOneDyeVvvyJnDv3r3R0dFDhgyRSCQm9/bUrvLSXLnJzbaCufucXDpS1Xuwtb0b3ZyFdjYQneGvHyrGp7iZs1Bzv8sKiGDdP1dv5kI7G/fPCb3M/vzf3GM4PIOYT6+KKwqUrr4WzSZISUnJzMxs9hCCICRS843VzZs3x8bGmtTTl7T0QBRBEPQGr9mjV65cQZ/P/AelDMn5WzLvU29Tu9kGGPQYrC5VZdxpiGuh45xCoUDP4KvodLpmzx0AwMLCoqVDb45U2vy7ZLSl1lK5LFbzvS0enBfaOFLN3+sGm76hz+80CKvUgyY5mL9obHl+t0EoUA+ajMEXx6ZvaM8YK4MePLooxKR0rCh6Lst9LMVEZox79j+5KkJ0hk4ybKmjyU+TFqTLEuc4Y+UAlgPdeENtdFp96uEqDH0wD0+uiLCVGftxWQCAvKfSm7/X9B1u12uAtRHJuxgF6bK7Z+pC+rF5wzAOXdgrDQDQqpF7Z+uLnst6xVhzQ5i2TlSsPXpTpCIt/4W8JFtBphCiR3PYdtiPL+wUSqPIxLqMO2J+phzRGXxCmSQSkckms23JSBcYKA9IJIJUrFVIEKUMqeQrVXI9twczMJLl6NFZngZ2IqUbaajTCoqUMrFOLtERSQRpvYlfMDx79iw4OJhCMWU9s7Qm63UGBpvEtCY7utPt3WgmNG4SOqPSHU1CQsKxY8fQF1nw0LUnmcAxHlxpWIBRaX9/fwKh080i1dHAqHReXh6ErRMYlWaz2XidhgK0txDWXpgbGJV2coJrfhUUGJWuqur+71ReBUalg4KC8Os0FGRnZ+PXaZxuC4xKN46shAoYla6vh7HDOYxKczgcvEUGBXV1dXiLDKfbAqPSXC4Xj95QwOfz8eiN022BUenAwECsXcAAGJXOycnB2gUMgFFpOIFR6eDgYLztDQVZWVl42xun2wKj0ngvYFjAewHjdGdgVBrv7w0LeH9vWPD29sbrNBQUFRXhdRqn2wKj0g4ODnj0hoKamho8ekMBPloHFvDROrCA12lYwOs0LLi6umLtAgZANPPc8OHDqVSqwWAQCoU2NjYkEglBEDs7u8OHD2Ptmjkw9zocGEIkEgUCAbpdXV2NLlC6fPlyrP0yExBFbx6P958AxuVy4+LisPPIrECk9MyZM5vOZcNgMGbMmIGpR2YFIqUDAgLCwsIaP3p7e8fHx2PqkVmBSGkAQHJysqOjI1qhp0+fjrU7ZgUupQMDA3v37m0wGLhcLlQVuqPa3ojOIK7VSOt1+s53B5cwYFZpjnps3LiizA5ZNfdNIBgAw4pk60il0ExfA01/P515vyH7oVSj1Dt40JWyjlpXvFtCJBFkYq1GifiFs/qNtDOtcRMrnXG7obxAGTPeEcIHyyYk7YYQUSODp5hyDTVTRomsB5KyPMWACU64zG9I70F2FDr51qlaE9o0mdJ6xJB5vyF6XPPrkuK0l9BY2/oqjbhWYyqDJlNaKtIpZQiJDFdjvkMhEon1VZ1SaXvXzrI2VPfAxpEmazDZClKmq4IGoJLjLW1TolHr9aY7o3iwhQVcaVjAlYYFXGlYwJWGBVxpWMCVhgVcaVjAlYYFXGlYwJWGhS6m9MZNa1auWoS1F12SLjaGY9SoCTqtFt3etHltVFTM8ITRGPvURehiSveJiGrczsvLjoqKwdSdrgRm0fvQ4X3vrVzY+HHW7InjJ74cOPPxJ+vXvb+Mzy8cPDTi3r1bs+dOXrR4VtPoPXhoRGWV4IsvN48eOwjNcvVa6sJFyYkjYyZMit+5a7tKpWrTh9ramrXrlyYk9p80ZfjBQ3v37d+Z/NYEAEBObtbgoRE5uVmNKWcmj/v+h2/QbbFYtOXzDUnTRg4fEb04ZXZa+mN0/6k/T46fGHf37s3xE+O+/+GbpcvnrV7zbtPiPtqwanHK7Dc+c68JZkoH+Adl52TqdDoAQH29sKamymAwlJWVoEcznqdF8PpSKBQAwKHDe5OmJK9etaFp9pMnzgMAlqSsPnrkNADgzp0bn372AY/Xd9/en9es3njr9tXtX3/Wpg9bP9/A5xds3fLt9q++F4vrUy+dJZPbCHJ6vX7tuiUvXmSsXbNpz/dHAwOC161fWlRUAACgUCgqlfKPUyfWrtk0duzkkYnjnjx9VFf3T18wpVL59+P7GF5rMFPa3z9IpVIVFOYBANKfPfHx8Q8ICM54ngYAKK8oEwrreOF9AYEAAAgLi0gcPsbb27dpdjbbCh2KYcW2AgAcP3EwNDT8nXkpbq7uUX2j35m35MqVCzU11a04UFtbk5b+ePq0OeG9+3h6cpctXUuntd1n5vGTh3n5OatWfojmSnl3laOj8x+nTgAACASCSqWaNHF6VN9oF2fX2NhhTCbz6rWLaMb7D24bDIYhgxPe+My9JpgpbWtr5+ri9iLzGQAgI+Npz5CwHsG9nmemox/t7Dhcrg+aMji4Z+um9Hp9Xl52BO/lJTwslAcAKCrKbyVXSSkfAODr449+JBAIgUEhbbqdnZ1JoVBQ+2hXr149excU5DYmaPSWTqcPGZxw6fI59OOtW1cHxAy2tLRss4gOAssWWXh45PPM9IkTp6U/e7LgnaU0Oj019Qwaunm8vo3JmMw2zo5KpUIQ5OChPYeP7Gu6X1hf10oupVIBAGAwmC8LarLdEgqFXKvVJiT2b9yDIIit7cte+E29HTFi3F9nfi8oyHNz83j46O7Hm7e1ab/jwFjpnbu2icWi0tLiHiGhVAq1pra6rq4249nTObMXGmHgH+h0OplMnjB+6sgR45rut7ZpbfVhOt0CAKBWv2y4SaUSdOPV/uqqf5MxmZZUKnXfnuNNjxKJzYfGAP8gP9+AGzcv+/kFstlWvPBI47+UycFS6d5hEUJh3cXUM1yuD5vFRmPpteuplVWCcONOCjoAhUgk+vkFVldXenh4ofu1Wm1NbTVqsyXc3TwBAHn5OUFBIWjVfJGVgVZxtHLLZFI0pUhULxT+Ex4CA3toNBoEQRovLlVVldbWNi2Vkpg49rffj1dUlMXHjWzpB2EesCzbysrazzfg1J+/9OrZG90TEhL2x6kT3t6+dnac1vPSaDQajfYs42l+Qa5Op5uaNOvW7WvHfz5YVlaSX5C7ZetHS5e9LZe3NsbOycm5R49eR4/9+PDRvbz8nM+/2Nh4yMHBycrK+tLlczqdTiqT7vjuS7QBCADghUf6+QZs2fpRevqTyirBlasX5y+YfvqvX1sqZdiwRKGw9s7dGwlYP+HB+GloeHhkTU11r17h6MeePcOqq6vCextVoadNnX3z5pVVqxcrVcqBA4a8v/6Tq9cuzp2XtHrNu1qd9uvte5jMNq67H7z/qYe710cbVq5dt8TFxa1/v4HofiqVum7t5uzszNFjB6UsmTNkSIKbm4derwcAkEikLz7/juvtu3HzmtlzJh05uj85eV7SlOSWimBZssLCIoKCQtxc3dtzYkyPyUbglecpH6XWx83qwhNAfbvji/RnTw78eNKENsVi0fSZY9as3jgodlh78z66WGfnRA6LtTaJJ13saWgXokHSIKgo27l7u6en98ABQ7B2p1sr/fx5+vsftjgJ1dEjp63+vfp2BKmpZ/bt3xnaK3z1qg3YtsVQunP0VqvV9SJhS0cdHZw6gwCtgEdvY6HRaM5OLlh70Vno1D9qHBOCKw0LuNKwgCsNC7jSsIArDQu40rCAKw0LuNKwYDKlCWTAsO7OT9zMD5VOpNFNJpDJDNm70EoyZaayhgMAEBQqrB0pprJmMqWpdKJHELNOoDSVQcjRavQkEnDyMNlkfqa8Tg+eYn/z12qtWm9Cm9By+UhF9Bg7AtFkc+2aeNZnpQw5/EkxL4HDsqFYcaig883k3pkhEIBUrG2o1Ty5LJywxJXjQjOl8Y5YGe1RqrCiQKXXA2m91uTG3xy1Wk2lUjvh1NQUKoHGIDlz6RFxNjQLkmmNQ7QGXiMJCQnHjh3jcNroftrNwO+nYQFXGhZgVBpffxoW8PWnYcHb2xuv01BQVFSE12koCAgIwOs0FOTm5uJ1Ggq4XC5ep6GAz+fjdRqn2wKj0r6+vnj0hoKCggI8euN0W2BUmk6n49EbClQqFR69oYDFYuF1GgqkUilep3G6LTAq7eIC4+QnMCotEAiwdgEDYFQaTmBUGu9zAgt4nxOc7gyMSuO9gGEB7wWM052BUWm87Q0LeNsbFmxsbPA6DQUikQiv0zjdFhiV9vf3x6M3FOTl5eHRGwoCAwOxdgEDYFQ6JycHaxcwAEalAwICsHYBA2BUOjc314hU3Q0YlYbzOg3RzHOTJ0+m0WhEIrGgoMDV1RXdptPpe/fuxdo1cwDRjNxFRUWNt9F8Ph9dYnjp0qVY+2UmIIreffr0+c8ed3f3KVOmYOSOuYFI6dmzZ7PZ7MaPRCJx/PjxFIrJpkrv5ECkdFRUlL+/f2O7xM3NberUqVg7ZT4gUhoA8NZbb1lZWaFX6MmTJ5NIJp5ZuTMDl9L9+vULDAw0GAwuLi5JSUlYu2NWXrPtbdAbZGId6IJvhKZOmsPPr5o0bqa8QQ9A11tIgsYgUmmvUz/bfT/Nz5Q/uyUuL1DaOdPUCuQ1isR5EwwGQKaA0FjrXjHtW2u+fUpnPZLk/i3rk8ixsqO230kc0yCt1764J7KwJMaMbceyA+1Q+sV9SdFz2aAkGIekdkKeXqkDBEPsBHsj0xsb8TVqfV6aFJe58xA+jKOU6atLVEamN1ZpoUCtVcHyhLyrQCIRasvVRiY2VmlJvc7Jy+INvMIxPfbudLlEZ2RiY5VGtAalHG9pdy60aoNKYeyNIlxPTmAGVxoWcKVhAVcaFnClYQFXGhZwpWEBVxoWcKVhAVcaFnClYQFXGha6s9KbNq+9mHoGay86C91Z6by8bKxd6ER0oNI6ne7goT2zZk9MSOw/c9b403/9hu6/cvXi0LjI/IJ/hrZmZj4bPDTi5q2rrWQBAGi12n37d05OSkwcGbNk2duZmc/Q/YkjY345eaQx2VfbPlmwcCYAYPDQiMoqwRdfbh49dhB66Oq11IWLkhNHxkyYFL9z13aVqu3eGrW1NWvXL01I7D9pyvCDh/bu278z+a0JrZcLABCLRVs+35A0beTwEdGLU2anpT9G9/P5hYOHRty7d2v23MmLFs9aunze6jXvNi3uow2rFqfMfq2T3TYdqPQPe7795eSRGdPm/Lj/l8mTZuzcte3c+T8BAMOGDo+Kivl2xxcGgwFBkB3ffTkodljswKGtZAEAfP/D1+fO/7l40XvffL3P1dV9zboUQWVFK6WfPHEeALAkZfXRI6cBAHfu3Pj0sw94vL779v68ZvXGW7evbv/6sza/wtbPN/D5BVu3fLv9q+/F4vrUS2fJ5Db6Tev1+rXrlrx4kbF2zaY93x8NDAhet35pUVEBAAAdGXTo8N6kKcmrV20YmTjuydNHdXW1aEalUvn34/vDE0a35xy3g45SWqFQnP7r16QpyQkJo9xc3ceOmZQQP+r4zwfRoyuWrS8pLrqYeuavM7/X1FYvXbIGACCTyVrKIpfLz53/c1byO4MHxQX4B61c8UGfiH4VFWWtOMBmWwEAGAyGFdsKAHD8xMHQ0PB35qW4ubpH9Y1+Z96SK1cu1NRUt2KhtrYmLf3x9Glzwnv38fTkLlu6lk6jt/nFHz95mJefs2rlh2iulHdXOTo6/3HqBAAA7R4fFhaROHyMt7dvbOwwJpN59dpFNOP9B7cNBsOQwQntO9FG01FKl5QU6XS6CF5U457QUJ5AUK5QKAAAHI79woXL9+zdceDA90tSVtvY2AIACgvzWspSXFyo0WiCAnug+ykUyuZNX/aJiGqu5GbQ6/V5edlNLYeF8gAARUX5rX2FUj4AwNfHH/1IIBACg0LaLCs7O5NCoaD20XF+vXr2Lih4OQtDcHBPdINOpw8ZnHDp8jn0461bVwfEDLa0tDTyS7WXjho/LVfIAQArVi5oHLKMdjeuFwkZDAYAYOiQ4bu//x+JRB4QMxhNoGg5i1QqAQDQjKhSzaJSqRAEOXhoz+Ej+5ruF9bXtZJLqVQAABgMZuMeZpPtllAo5FqtNiGxf+MeBEFsbe1eGmG+1HLEiHF/nfm9oCDPzc3j4aO7H2/e1p6v1T46Smn0BH3w/qfeXN+m+x3sHdGNAwd/4HAcdFrtocN735mX0ngKms2CKo3+FP7Df+aQ02ia6StJp9PJZPKE8VNHjhjXdL+1jW0rX4FOtwAAqNUvG26oG62Xy2RaUqnUfXuONz1KJDYfOwP8g/x8A27cvOznF8hmW/HCI1vx5w3pKKW9PL0pFIpIVO8R64XuEYtFBAKBSqUCAHJys37/4+evvtyl0Wg++HDFwIFDA/yDvL39Wsri7uZJp9OfZTwNCQlFo/GKlQtGDB+bkDCKwWDKZNLGcguL8inkl0Oi0ahAJBL9/AKrqys9PP6xrNVqa2qr2Sw2aBl3N08AQF5+TlBQCFo1X2RlNFbxlsoNDOyh0WgQBOFyfdBDVVWV1tY2LZWSmDj2t9+PV1SUxceNbOkHYRI6yjSDwRg1asLBQ3uuXb8kqKxIS3+8as3iz7/chN5KfbXt46FDh/cOi+gb2X9AzOAvv9qs0+ksLS1bymJpaZk4fMyx4z9dunQuNy/7f19vycvLDukZBgDw9w+6c/dGQ4NYq9UeO35AImlAHaDRaDQa7VnG0/yCXJ1ONzVp1q3b147/fLCsrCS/IHfL1o+WLntbLm8mSDTi5OTco0evo8d+fPjoXl5+zudfbGx6tKVyeeGRfr4BW7Z+lJ7+pLJKcOXqxfkLpp/+69eWShk2LFEorL1z90ZCh7W6UYwdrZP1QFKWr+o/xsF40zqd7sjR/amXzgqFdba2dv37DXx77ruWlpaHj+z/7ffjhw/+jv7S6+pqZ8+dNGnijNlvzW8pCwBArVbv3f/d9euXlEoFl+s7f96SsDAeAKC8ouzLrzbn5+ewWOwRieO0Ws3ff9/fu+cYAODQ4X0nfjlEpdKOHvmTZcm6cvXizycOlpYWM5mWISGh8+ctaaziLVFZJdi27ZPnmelMpuWY0RMlkob0Z08O/Hiy9XJFovrv93zz8OFdlUrp5OQyauT4yZNmoFmSZ43/6stdEby+TUtZ9/4yhUK+45v9xp9blJxHDQqJJnaiUQN2OlDp7se3O75oVNpUiMWi6TPHrFm9cVDssPbmbZfSEM1d1NlokDQIKsp27t7u6ek9cMCQji4OdqUbn5W+yro1m6OjYzuu6NTUM/v27wztFb561YYObYuhwB69K6taXM3UxtqWTn/NO3jzgEfvduDsBMs44e781hKnKbjSsIArDQu40rCAKw0LuNKwgCsNC7jSsIArDQvGKk0mAwtLiOZI7hJQKEQ6w1gFjU1nZU8VFCrewCsc01NdprS0MfZ5trFK27tRqXQ81Hcu9Ije0cPYdzBG130SsWc0+/KR1jrT45iTB2drbBwoHBeakenbN+tzSZb8wYX6iOEca3va680njvOG6PUGYaU6677ImUvnDWmxI+KrtHsm98piZdo1cVmegmFJ7qLzSyJ6hEgkdb3lBgAAAJDIBCsOJXSglV9vVrsyvv4aeCoF0kXX6544ceLevXvt7OyMSNvpoNGJ4LXO+uv3RKAzuupNlxZRUOkEmgVcVx+4vi3MwKg0l8vtotedNwFGpfl8Pjwr9DYCo9JBQUF4nYaC7OxsvE5DAV6nYQGv07DAYrHwOg0FUqkUr9M43RYYlQ4ODsbaBQyAUemsrCysXcAAGJWGExiV9vDwwNveUFBaWoq3vXG6LTAqzWaz8egNBRKJBI/eUEAkEvE6DQV6vR6v0zjdFhiVtrGxwaM3FIhEIjx643RbYFQa7wUMC3gvYJzuDIxK431DYQHvG4rTnYFRabwXMCzgvYBhAW+RwQLeIoMFFxdY1t5oCoxKCwQtrqfTjYFRaWdnZ6xdwAAYla6srMTaBQyAUenAwEC87Q0FOTk5ELa9X3+OwS4Hj8czGAxEIlGv16P/SSRScnLy0qVLsXbNHEBUp0NDQ9ENdLlQIpHo5uY2ffp0rP0yExApPXXqVFtb26Z74uPjORwOdh6ZFYiUjo+P9/T0bPzo7u6elJSEqUdmBSKlAQBJSUnW1tbodkJCwn+qePcGLqXj4+O5XC5aoadMmYK1O2YFLqUBAFOmTGEymXFxcVBV6De9yxIUKosyFTXlaqUMUckQAgFoNHqTutch6LRaEpncJR6eWDBJRBLBwpJk7073DKB7BTNf29TrKK2UIX9fEmc9bKBbUtiOTDKNTKaRyVQSmUKE5d7cXBgQg1at02kQRItIquWSWqU/j80bYmXnbOxCK420T2mDwXD9V2HeU4mTvx2LY0GidNVp+7soBoNBJlTWFNQ7uNMGTbJjWVOMz9sOpcsLNNd/rbGwZnC8rF7XVRzTIBbI5EJZrwFWPftZGpnFWKWzH0nunRN593XtEpc3SCjLqPbpQe8/yqimpVFt7/IC1aPLEp8oN1zmToV7L8fiPE3aTYkxiduu0yXZ8pt/ijzCYHx73yWozq3zCaFGDGtjlbQ26rRCqrt4qBqXuTPjGMDJeiQvyZG3nqwNpc/9WO3JczKpYzimxz3M6dqJWr2+tfDcmtJ5T6UaLZFu2e5bNxwzQyAQWI6s+2eFraRpTenbfwrtfeB6ZNh14XhZZ9xu0KhbfEbZotKFGVILazrV4vWXQ8QxMxyudfoNcUtHW1Q6L01hYWXsKtadjY1bE+pF0HXqtrSzyEuTtXS0RaVLsuRs+9d/no4hInGVXNHiT7sbY8GmKaWITKxr9mjz99M1paobp0QO/g5tWi8qSf/z7PbqWj7H1m308GVXbh5wcfKdMHoNAEAmF5258G1h8VO5Quzs6DcibrGvNw8AcO/R76lX986duf30+f/V1BYzGFZDY+f05Y1BDZYLcs5f3l0uyEF0Wj+fPmMSV9jaOAMADp9YDwDB0d7rxt1jyVM+Cw6MKavIOn95d0VlnlardnLwThy2yN83sqDoyQ8HFqOmegQOnDPjKwTRXbl5IP35ZZG40trKcWD/af0jJ7b5vZo1DgCoruGyPdh/AAAKOUlEQVR/9d3UhXN2375/gl/6jEgghoYMG5O4gkQiAQAePj596/6JelEFhUL39uo9bsR7ao3yqx1Ji9/+wdurNwAgLePSsV8/mjB6DepDTW3xlzuSli74ycOtR1rGpZt3j1fX8mk0Ru+e8YnDFlGpdDQ+DY2dnVfwML/o8cfrL9FojFbcruPX9+xLD+A1szR183VaLkW0Rrx/1GrVB4+vodOYS+f/OH706vOXd9eLKgAgoDM27ju0vLjsedKEDcsXHnJ3Ddp/ZHllVQEAgEQkq1SyKzd/mjV16ycfXOWFjfjjzBfihhq0Ov7w02Iigbho7u6Fc3cpFJI9B1O0Og0AgESiVNUUlgty5iV/7eEeotWq9x1eTiZRF7z13bKFBzzdex44vlrcUMP1DJ055TMAwPJFh6ZN3AQAOJv63c07R4cMfGtVyvGB/aedPve/h49Pt/m9mjUOACCRyACA0xe+Hjwg+eP1l2ZM/uTuw1+fZ10HABQVp/16esuAfkkrU46/PfN/CnnDkV8+cLT3srZyLC7NQC0XFadZWznyi9PRj4XFaRZ0lptLUGbWzWO/fuTvG7ny3aNJ4z/KeHHtt7+2omlIJPKDv/90cvRdNHc3mdzGfZAeAVKRttlDzSutkOqI5LbfU2Xl3lEoGiaMWePqEuDL5Y0ftVIirUMP5Rc+qqjMmTz2fT/vCEcH7tgR79lYO995cBI9iuh1gwfMsrZyJBAIkeGjEUQnqMoHANz/+w9AIMyY/Imzo6+7a/C0SZvqRRXPX1xDcwnry6dO3OjDDbdkWhOJpEVzdydN2ODqEuDk4D186AKtVlVcmkEikek0JgCAYcGm05lKlezew99iY2b26T2SY+feP3JiRO+R124fbv17tWS8MUFojyFeHr0AAH4+fexsXMsrsgEAVTVFFAqtT+9RHFs3T/eQmUmfjUlcDgDw5UbwS5+hGQuLn/bljS0q+UfpouI0P58+RCLx2u3D3l7hI+IWc+zcg/z7j4x/9+mzi+KGagAAAAQKhT4qIcXLoxcaOVqBRCXJxEizh5pvWmtVegqD2rpRNPjQ6ZZODt7oR65nGJPxTy+tkvJMEoniww3/99wRvT3DKirzGvO6OPqhGwwLNgBApZICAErLMj1cgy0s/gk+NtZOtjauFZV54aHDAQD2dp5Mxj+v0Ugksk6n/fPcNkFVvlIpNQADAEChbPiPh4LKPESv8/eJbNzjww1/+OS0Wq1oJQy2adzZya9xm05nKVVSAIAPl0cAhF37F0TyRvv7RNrauLBZduiv4c9z2w0Gg0wuqhOW9Y+cePXWwXqRwNbGpbj02ZCBs/V6fbkgO37IO402vb3CAQCVVQXWVo4AAC+Pnm1qgUKxoCBIe5QmkglahaZNuwqlBK1AjTD+VUKtViCIdt3mAY2H9HqEZWn30ifK/wtEaHNBqZILqnLXbopp3I8g2sY4Qae/fENXW1e658C7vt4R0yZusmLb6/X6T7eNftVDtVoBAPjhp8Xg5bsZAwBAKhO2onSbxin/P4qiPwVHe6+U+fuv3z5y/tKu35RbPdxCxo5Y4eke4ufdR6mSVtUU1dQWOzv6MZnW7q7BRcXp6NXK3ydSq1Xp9cila/suX/+xqdlmv3jr6NQ6Pan5J2XNK81gkfU6RZt2KRSaRqtqukehaPjXOSaZTH1v8ZGmRwmENh6+0ulMrkfYpLHrmu6kUpuRJP35Zb0emTH5E/QXIxJXtWQQADB98sfOjj5N91tZObbihpHGX8XFyW/G5I8RBOGXpl+8suenoys/XH2GzeY42nOLSzMEVfneXmEAAK5HKL/0mQEY7Gzd7Gxd9Xo9iUSOiUpqbJaiWDLb/dhKp0FYrs1H+OZPPYNFQrTNB4Gm2Nm6KRQNdfXl6MeikvTG2xsP1x46nQbRIw72XugfmUyzYrfRmPd0D6mrL7OzdWvMBQCBzWqm870O0VIo9MbA8OTZhf8kQIOEs5MfiUSRyeobDTIYVgyGNYXc2rWpTePNUlKWiV7LSSSSL5c3fOgCuUIslQoBAH4+kcWlGUXFaWgL3MszlF+SXlzyDL2sEIlEV+dAkbiy0UlbG1cikcxgsI0ptymIRse0bo/Sjh50mVDdpt0g/2gKhYbeLBWVpJ+9+F2jKr7efVydA37+bVMB/0m9SPD0WerXu5PvPfqtdYNREePVasWJPz6uEOTW1pVevv7jtp3TyipevJrSw62HXCF+9PSMRFp39+FvZeVZlkwbQWW+UiVDL/PZeXeraoos6Jb9+oxPvb4v/fllYX1FQdGTPQeX/HLq49bdaMV4K7ly8+8fOLY6I/NanbC8QpB758EvNtbONtZOAABf74iCosc1tXyuZygAgOvRq7auNK/wod+/DYhBMTOfZ12/dutQTW1JhSD3+G8bd+2fr1K18XrqVdRSjaN788+7mo/eJDLBiWshrVOwOK3dvbFZdslTtvx18Zvtu2Y6O/qOHbHi19Nb0KpAIpHmzfrm7MUdh0+s12iUttYuwwbNjY1uYxCUrY3zwrm7z13auWv/fCKR5OTgM2fGNk/3ZtojPQIHDIqeeS51518Xvgny6z914sZb936+fvswgUgcN2JloF+/Mxe+5XqGLpy7e/TwZRZ01rlLOyXSOpalXXDAgMS4Ra270Yrx2P7TWso1NHaODtGdSd0hkdTS6ZZeHr3mJX+N9t3w4YZLZUJ7O09Lpg0AwMKC5eTgXVVT6MvloXl79Rg8beLm67cPp17di+ZdNHc3eukxHq1Kp9MgDi0o3WJPhIw74qzHaqeANoYtyRUN1H8DnVan2bglbmR8SnTU5Ha5iGMShGUSa7Z26NTmL5EtvsAI7MN+cq2sddNKlWzr1xP8vPvEDX6bAAg37h4jEIg9gwe/sc84r4NKrOiRYNfS0dZ6F907K6woNdhzW+u2UlKWef7yrvKKHAKR6OLkPzJ+cbPBtrPBL0n/8ejKlo6uX/FH4417V6GhWk5Qy8csaLF3UBv9yHavKgwc5EEkdbdBPVqtWipr8b29tZUTOsa6C1Fwr2zyMlcrTos9wNtQOi9NmnZL7uhv3zHu4ZiG+lKxswehX2Jr999t/HL9e7NcuRRhscjUvuGYDEmNjGTQtC6zUf29Y8bYcRwJNQW42J0RSa1cK5WPXdh2512jrkaxE+yYTF1tYb0pfMMxGWKBRF7VMDHFqMkx2zEu61FqfUm+lu3EpjHbfs2F06EgWkRUIWGzkPiZrT3Ab0r7xlqW5Mivn6yjMmkOPjZkGt6ZEAMMBkNtoai+XDpwAie4bzsejL/O+Omsh5IXD2RyCcK0Y7AdmVSLrjHqvEujVSHSWrlMqCCRDH6hzMiENsbmvMrrz4lQyVfmp8urStQ1JUoqnUSxIFEsSAYdPlbelBAIBKVUo1YiDp4Wtg4UvzCmZ9BrduM0zRyDCqlO3oBoVF1g6ouuBYVGYLDIDDaJSHzTqAnRbJKQ08We+eG8NrjSsIArDQu40rCAKw0LuNKw8H9pcEeNZi1IqgAAAABJRU5ErkJggg==",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test our application! Note that we can stream the results of individual steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'write_query': {'query': 'SELECT COUNT(*) as employee_count FROM Employee;'}}\n",
      "{'execute_query': {'result': '[(8,)]'}}\n",
      "{'generate_answer': {'answer': 'There are 8 employees in total.'}}\n"
     ]
    }
   ],
   "source": [
    "for step in graph.stream(\n",
    "    {\"question\": \"How many employees are there?\"}, stream_mode=\"updates\"\n",
    "):\n",
    "    print(step)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check out the [LangSmith trace](https://smith.langchain.com/public/30a79380-6ba6-46af-8bd9-5d1df0b9ccca/r)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Human-in-the-loop\n",
    "\n",
    "LangGraph supports a number of features that can be useful for this workflow. One of them is [human-in-the-loop](https://langchain-ai.github.io/langgraph/concepts/human_in_the_loop/): we can interrupt our application before sensitive steps (such as the execution of a SQL query) for human review. This is enabled by LangGraph's [persistence](https://langchain-ai.github.io/langgraph/concepts/persistence/) layer, which saves run progress to your storage of choice. Below, we specify storage in-memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "memory = MemorySaver()\n",
    "graph = graph_builder.compile(checkpointer=memory, interrupt_before=[\"execute_query\"])\n",
    "\n",
    "# Now that we're using persistence, we need to specify a thread ID\n",
    "# so that we can continue the run after review.\n",
    "config = {\"configurable\": {\"thread_id\": \"1\"}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAKMAAAF3CAIAAABG1mxaAAAAAXNSR0IArs4c6QAAIABJREFUeJztnXdcE/f/xz/ZIYGwCRsSNqIgIFLBLQLujQvrqqvubYerrba1frXWUVfd1KrVusW9Vx2IyAiBsHdIyN75/XE28tMQUDOQzz0fPHhc7u7z/rzvXvd53+fuPgOj1WoBCgRgLe0AiplAlYYFVGlYQJWGBVRpWECVhgW8pR3QT1WxVCrUiIUqtVIrl2os7U7zEMhYPA5DoeEoNjgXbzIOh7G0R2+DaVXP03lPhZwsMSdL7BNC0Wi0VBu8PZ2okH0CShOtsPxahUSglkvUlRyZZyCFGUYN6mRDILaWqNlalM560HD/DNc3lMIIozLCqHhCazlBH0ZxjrgwS1xRIA3oaBOT6GBpd0CrUJpbIU8/WOXGsOoyyJFkhbOsM0bn0UXu8xv8vql0Zntry3piYaXzngqfXuMNmOpGcyBY0A2TolRobp2opTkSLFu4Lal0UbaY9VTYN9XVUg6Yk0cXuVgcplNfi4ltMaWf3+BVl8iTPodCZoSHF+pEfHWfsXSL5G6Zik9JrqQkTwKVzACA2H5OZAru+U2eRXK3gNKiBtWLO/zBMzzMn7XFiR/ixK9RlrEk5s/aAkrfO10XFGVj/nxbCR262t4+VWf+fM2tdG25nFetCIyEV2lHN5KTBzHvidDM+Zpb6ax7DfFDncycaWsjfpBTfkabVlql0OQ+EXr6U8yZaSuEQsNLBOrqEpk5MzWr0pxXYkY7qjlzBAAcO3Zs9erVH5Bw2bJlZ8+eNYFHAADACKNyssQmMq4XsypdyZH5R5j7pWBOTo6ZE7YEvw7WdRVy09l/F7MqXVUks7E31XfS58+fT506tUePHl27dp0yZcqzZ88AANOmTTt79uy5c+eio6Pz8vIAAJcuXRo3blzXrl179+69YMGCsrIyJPmxY8cSEhJu3bqVkJCwefPm6OjoioqKNWvW9OjRwxTe0hzwpSypKSw3idaM/LGqUMhTmsKyRCLp1q3bDz/8UFhYWFBQsH79+ri4uIaGBqFQOG7cuBUrVvB4PJVKlZWVFRUVtW3bNg6Hk5WVNX369NGjRyMWTp48GRcXN2PGjLt375aVlVVXV0dFRR09epTP55vCYa1Wu3N5gUyiMpHxdzFrSwSxQE2lmeRrVVVVlVgs7tevH4PBAAAsXrw4ISGBSCSSyWQ8Hk8kEu3s7AAAPj4+hw4dCggIwOPxAICxY8cuXLiwvr7ewcEBg8HIZLKxY8fGxcUBAORyOQCAQqHY2tqawmEAAJWGEwvUZvt8Zz6ltRqtFQWLwZqkMYa3t7ePj88333wzYsSI2NjYoKCgqKiod3eztrYuLy/funVraWmpTCZTKpUAAIFA4ODw+sND+/btTeGeXshUnEZtvo8O5rtPY7AYgMFIhCpTGMfhcHv27OnTp8+pU6fGjx8/cODA8+fPv7vb5cuXly9fHhYWtmXLlrS0tK+//vqtHaytzVdh5NUoqDTzlTSz1siQeGUi4/b29vPnzz99+vSxY8diYmJWrVr1buX51KlT0dHRM2fO9PX1dXJyksnM+kTbGI1aK5dqrKzN1/LCrEq7MclS05Tp8vLymzdvIstMJvOrr77CYrEFBQXIGt2XWYVCgdywES5dutR467uY7pOuqEHlG2rWVwtmVdrJg8TOMMnrgqqqqqVLlx4+fLioqKi4uHjPnj1YLBa56drY2OTl5eXl5fH5/LCwsIcPH2ZlZVVWVq5fv97JyQkAkJ2d/W7hJpFIJBLp2bNneXl5KpXxr07OS7GNg3kb5pqtlq/VasUC5Z5vCk1k/Ny5c6NHj46Li+vWrdvEiRPv3LmDrL97926vXr3i4uLu37/P5/MXLlwYHx+fmJi4c+dOtVo9a9as2NjYixcvnjp1KioqSql88xC4a9euuLi4Xr16CQQCo3t7altZSZ7Y6GYNYO42J5cPVXXsaefsSTZnpq0NtUp75vfyobM9zZmpub9lBUXbPDhfb+ZMWxsPznN9zf7+39x9OHxCqM+u8cvZUg9/K707zJ49OysrS+8mtVqNw+mvrK5Zs6Z79+5G9fQNTb0QVavVyAOe3q1Xr15F3s+8hVSkzv1XMPV7prHdbAYLtBisLpFl3m1IaKLhnEQiQc7gu6hUKr3nDgBgZWXV1KaPRyjU/y0Zqak1la+Njf7WFg8vcO3pRPO3urFM29CXdxu4VfIeI1zMn7VleXmvgVsh7zHSAgdumbah7eNttRrw+BLXIrlbisKXorwnQovIbOGW/U+v8dQqbSvptmRq8p8L2Rmi5ElulnLAkh3donrbq5Sa9INVFvTBPDy9yrOszJbvlwUAYD0T3vq7pnOSY4eudi3Y/RODnSG6d7Yu7DNaVB8Lhy7LKw0AUMrV98/VF74UdYi3Y4RRHVyJlvboYxHylJxX4uIcCZ6AiRvoRHO0fP/CVqE0goivyrzL52SJ1SqtXzgVh8NSaXiaA179CXSUBzgcRshXSgRqqUhdyZHKxBpGO2pwjA3du7W8DWxFSutoqFNWFEpFfJVYoMLiMMJ6I39gePHiRWhoKIFgzHJmbYfXqLQUGo5qh6d7kZ09SUY0bhRao9KmJjEx8ciRI8iHLHj4tAeZQGk5qNKwAKPSgYGBGEyrG0XK1MCoNIvFgrB2AqPSNBoNLdNQgLQWsrQX5gZGpV1d4RpfBQFGpauq2v43lXeBUemQkBD0Pg0FOTk56H0apc0Co9K6npVQAaPS9fUwNjiHUWknJye0RgYFdXV1aI0Mpc0Co9IMBgON3lDA4XDQ6I3SZoFR6eDgYEu7YAFgVDo3N9fSLlgAGJWGExiVDg0NReveUJCdnY3WvVHaLDAqjbYChgW0FTBKWwZGpdH23rCAtveGBSaTiZZpKCgsLETLNEqbBUalXVxc0OgNBTU1NWj0hgK0tw4soL11YAEt07CAlmlY8PDwsLQLFgCikeeSkpKIRKJWq+Vyufb29jgcTq1WOzo6Hjx40NKumQNzz8NhQbBYbEVFBbJcXV2NTFA6f/58S/tlJiCK3lFRUW8FMAaDkZCQYDmPzApESo8fP77xWDYUCmXcuHEW9cisQKR0UFBQRESE7ieTyezbt69FPTIrECkNAEhNTaXT6UiBHjt2rKXdMStwKR0cHNyxY0etVstgMKAq0Kaqe6tVWn6tQliv0rS+J7jErhNKcuWDE4YUZplk1tyPAaMFFFucA51IIBm/BBr/eTrrQUPOI6FCqnHxJktFpppXvE2CxWFEfKVCqg6ItPmsv6NxjRtZ6cw7DWVsafxQOoQvlo3I85tctVzdc5Qx51AzZpTIfigoZUm6DnNFZf5IOvZwJJDxt0/VGtGm0ZTWqLVZDxrihuiflxTlfQnv7lBfpeDXKoxl0GhKC3kqqUiNw8NVmTcpWCy2vqpVKu3s0Vrmhmob2NNJogajzSBlvCKoBTIxWtM2Jgq5RmO8M4oGW1hAlYYFVGlYQJWGBVRpWECVhgVUaVhAlYYFVGlYQJWGBVRpWPjElF61eumixTMt7cUnySfWh2PAgGEqpRJZXr1mWWxsfFLiQAv79InwiSndKTpWt8xi5cTGxlvUnU8Ji0XvAwd3L1w0Q/dzwsThQ4e/6Tiz9rsVy7+ax+EU9Owdff/+7YmTR86cNaFx9O7ZO7qyquKnn9cMHNwDSXLtevqMmanJ/eOHjei7ddtGmUzWrA+1tTXLVsxNTO4yYlTS/gO7du/Zmvr5MABAbl52z97RuXnZuj3Hpw7Z8ftmZJnP5637cWXKmP5J/eJmzZ74POMJsv7UP8eGDk+4d+/W0OEJO37fPHf+1CVLv2yc3bcrF8+aPfGjz9wHYjGlgwJDcnKzVCoVAKC+nltTU6XVaktLi5GtmS+fR0d1JhAIAIADB3eljEpdsnhl4+THjl4AAMyZveTwodMAgLt3b37/w9dRUZ137/pz6ZJVt+9c27jph2Z9WP/jSg6HvX7drxs37ODz69Mvn8PjmwlyGo1m2fI5r15lLlu6eueOw8FBoctXzC0sZAMACASCTCY9eerosqWrBw8e2T95yNNnj+vqXrcFk0ql/z55YMF7jcWUDgwMkclk7AIWACDjxVM/v8CgoNDMl88BAGXlpVxuXVRkZ4DBAAAiIqKTkwYxmf6Nk9NotkhXDFuaLQAg7ej+8PDIL6bO9vTwiu0c98XUOVevXqypqTbgQG1tzfOMJ2PHTIrs2MnHhzFv7jIyqfk2M0+ePmLl5y5e9A2SavaXi+l0t5OnjgIAMBiMTCYbMXxsbOc4dzeP7t37UKnUa9cvIQkfPLyj1Wp79Uz86DP3gVhMaQcHRw93z1dZLwAAmZnP2odFtAvt8DIrA/np6OjEYPghe4aGtjdsSqPRsFg50VFvbuER4VEAgMLCfAOpiks4AAB/v0DkJwaDCQ4Ja9btnJwsAoGA2EeaenVo35HNztPtoPOWTCb36pl4+cp55Oft29e6xve0trZuNgsTYckaWWRkzMusjOHDx2S8eDr9i7kkMjk9/SwSuqOiOut2o1KbOTsymUytVu8/sPPgod2N13Pr6wykkkolAAAKhfomo0bLTSGRiJVKZWJyF90atVrt4PCmFX5jb/v1G3Lm7N9sNsvT0/vR43tr1/zSrH3TYWGlt277hc/nlZQUtQsLJxKINbXVdXW1mS+eTZo4owUGXkMmk/F4/LCho/v3G9J4vZ29odmHyWQrAIBc/qbiJhQKkIV326vL/tuNSrUmEom7d6Y13orF6g+NQYEhAf5BN29dCQgIptFsoyJjWn5QRseSSneMiOZy6y6ln2Uw/Gg2NCSWXr+RXllVEdmyk4J0QMFisQEBwdXVld7evsh6pVJZU1uN2GwKL08fAAArPzckJAwpmq+yM5EijhRukUiI7Mnj1XO5r8NDcHA7hUKhVqt1N5eqqko7O/umcklOHnzi77Ty8tK+Cf2buiDMgyXztrW1C/APOvXPXx3ad0TWhIVFnDx1lMn0d3R0MpyWRCKRSKQXmc/y2XkqlWp0yoTbd66n/bm/tLQ4n523bv23c+dNEYsN9bFzdXVr167D4SN7Hz2+z8rP/fGnVbpNLi6utrZ2l6+cV6lUQpFwy28/IxVAAEBUZEyAf9C69d9mZDytrKq4eu3StOljT5853lQuffokc7m1d+/dTLT0Gx4Lvw2NjIypqanu0CES+dm+fUR1dVVkxxYV6DGjJ966dXXxkllSmbRb115frfju2vVLk6emLFn6pVKl3LRxJ5XazH3366++9/by/XblomXL57i7e3b5rBuynkgkLl+2Jicna+DgHrPnTOrVK9HT01uj0QAAcDjcTz/+xmD6r1qzdOKkEYcO70lNnZoyKrWpLGysbSIiokNCwjw9vN7nxBgfo/XAK2NJH6fXJ0z4hAeA+nXLTxkvnu7be8yINvl83tjxg5YuWdWje5/3Tfv4Up2jKz6iu51RPPnE3oZ+QjQIGirKS7du3+jjw+zWtZel3WnTSr98mfHVN00OQnX40Gnb/+6+piA9/ezuPVvDO0QuWbzSsnUxhLYcveVyeT2P29RWuotraxDAAGj0bikkEsnN1d3SXrQWWvVFjWJEUKVhAVUaFlClYQFVGhZQpWEBVRoWUKVhAVUaFoymNAYPKHZt+Y2b+SGSsSSy0QQymiFnd1JxlshY1lAAABUFEjs6wVjWjKY0kYz1DqHWVUiNZRBylAoNDgdcvY02mJ8x79M9RznfOl6tlGuMaBNarhwqjxvkiMEabaxdI4/6LBWpD35XFJXoZGNPsHUigtY3kntrBoMBQr6yoVbx9Ap32BwPJ3eSMY2bYma0x+nccrZMowHCeqXRjX88crmcSCS2wqGpCUQMiYJzY5CjE+xJVjjjGodoDjwdiYmJR44ccXJqpvlpGwN9noYFVGlYgFFpdP5pWEDnn4YFJpOJlmkoKCwsRMs0FAQFBaFlGgry8vLQMg0FDAYDLdNQwOFw0DKN0maBUWl/f380ekMBm81GozdKmwVGpclkMhq9oUAmk6HRGwpsbGzQMg0FQqEQLdMobRYYlXZ3h3HwExiVrqiosLQLFgBGpeEERqXRNiewgLY5QWnLwKg02goYFtBWwChtGRiVRuvesIDWvWHB3t4eLdNQwOPx0DKN0maBUenAwEA0ekMBi8VCozcUBAcHW9oFCwCj0rm5uZZ2wQLAqHRQUJClXbAAMCqdl5fXgr3aGjAqDed9GqKR50aOHEkikbBYLJvN9vDwQJbJZPKuXbss7Zo5gGhE7sLCQt1jNIfDQaYYnjt3rqX9MhMQRe9OnTq9tcbLy2vUqFEWcsfcQKT0xIkTaTSa7icWix06dCiBYLSh0ls5ECkdGxsbGBioq5d4enqOHj3a0k6ZD4iUBgB8/vnntra2yB165MiROJyRR1ZuzcCl9GeffRYcHKzVat3d3VNSUiztjllpad1bwFUacaYACzJ6xCROftWIIePFDRoAPvmJJLRaLdUWj8M1L00zz9PVxbInV3lFr8RuTCsBtzUOwA85BAKGz1W6+pDDu9v5h1sb2NOQ0mX50junauOH0WmORGybKNBtFUG94snlOmYYtX2cbVP7NKl0KUty7yy3/1QvU3qIYkxu/13l4UeO6G6nd2uTNbJn1/m9x8HYz/jTpdtw15JciVig0rtVv9LiBhW3Qk429vQuKKZGpdRyKxR6N+lXml+r8AykmNgrFOND97ES8PRXnPUrrdVgRDz9QQClNSOTaFQK/RUvuN6cwAyqNCygSsMCqjQsoErDAqo0LKBKwwKqNCygSsMCqjQsoErDAqo0LLRlpVevWXYp/aylvWgttGWlWawcS7vQijCh0iqVav+BnRMmDk9M7jJ+wtDTZ04g669eu9Q7ISaf/bpra1bWi569o2/dvmYgCQBAqVTu3rN1ZEpycv/4OfOmZGW9QNYn94//69gh3W4bfvlu+ozxAICevaMrqyp++nnNwME9kE3XrqfPmJma3D9+2Ii+W7dtlMlkzR5CbW3NshVzE5O7jBiVtP/Art17tqZ+PsxwvgAAPp+37seVKWP6J/WLmzV74vOMJ8h6DqegZ+/o+/dvT5w8cuasCXPnT12y9MvG2X27cvGs2RM/6GQ3jwmV/n3nr38dOzRuzKS9e/4aOWLc1m2/nL/wDwCgT++k2Nj4X7f8pNVq1Wr1lt9+7tG9T/duvQ0kAQDs+H3T+Qv/zJq5cPOm3R4eXkuXz66oLDeQ+7GjFwAAc2YvOXzoNADg7t2b3//wdVRU5927/ly6ZNXtO9c2bvqh2UNY/+NKDoe9ft2vGzfs4PPr0y+fw+ObaTet0WiWLZ/z6lXmsqWrd+44HBwUunzF3MJCNgAA6Rl04OCulFGpSxav7J885Omzx3V1tUhCqVT675MHSYkD3+ccvwemUlokEp0+czxlVGpi4gBPD6/Bg0Yk9h2Q9ud+ZOuCeSuKiwovpZ89c/bvmtrquXOWGk4iFovPX/hnQuoXPXskBAWGLFrwdafoz8rLSw04QKPZAgAoFIotzRYAkHZ0f3h45BdTZ3t6eMV2jvti6pyrVy/W1FQbsFBbW/M848nYMZMiO3by8WHMm7uMTCI3e+BPnj5i5ecuXvQNkmr2l4vpdLeTp44CAAAGAwCIiIhOThrEZPp3796HSqVeu34JSfjg4R2tVturZ+L7negWYyqlCwpYKpUqOipWtyY8PKqiokwikQAAnJycZ8yYv3PXln37dsyZvcTe3sFwkqKiAoVCERLcDllPIBDWrP65U3Ssvpz1oNFoWKycxpYjwqMAAIWF+QZSFZdwAAD+foHITwwGExwS1mxeOTlZBAIBsY/08+vQviOb/WYUhtDQ9sgCmUzu1TPx8pXzyM/bt691je9pbW2ozfbHYKr+0xKJGACwYNF0XZdlpLlxPY9LoVAAAL17JW3f8T8cDt81vmezSYRCAQCA1IIipReZTKZWq/cf2Hnw0O7G67n1dQZSSaUSAACFQtWtoTZabgqJRKxUKhOTu+jWqNVqBwfHN0aob7Ts12/ImbN/s9ksT0/vR4/vrV3zy/sc1vthKqWR4/n6q++ZDP/G612c6cjCvv2/Ozm5qJTKAwd3fTF1tuEkiNLIpfAWb40hp1DI392HTCbj8fhhQ0f37zek8Xo7ewcDh0AmWwEA5PI3FTfEDcP5UqnWRCJx9860xluxWP2xMygwJMA/6OatKwEBwTSabVRkjAF/PhJTKc1kBhAIBB6v3ru7L7KGz+dhMBgikQgAyM3L/vvknxt+3qZQKL7+ZkG3br2DAkMMJPHy9CGTyS8yn4WFhSPReMGi6f2SBicmDqBQqCKRUJdvQWE+Af+mSzQSFbBYbEBAcHV1pbf3a8tKpbKmtppmQwNN4+XpAwBg5eeGhIQhRfNVdqauiDeVb3BwO4VCoVarGQw/ZFNVVaWdnX1TuSQnDz7xd1p5eWnfhP5NXRBGwVSmra2tBwwYtv/Azus3LldUlj/PeLJ46awff16NPEpt+GVt795JHSOiO8d06Rrf8+cNa1QqlYEk1tbWyUmDjqT9cfny+TxWzv82rWOxcsLaRwAAAgND7t672dDAVyqVR9L2CQQNiAMkEolEIr3IfJbPzlOpVKNTJty+cz3tz/2lpcX57Lx167+dO2+KWKwnSOhwdXVr167D4SN7Hz2+z8rP/fGnVY23NpVvVGRMgH/QuvXfZmQ8rayquHrt0rTpY0+fOd5ULn36JHO5tXfv3Uw0Wa0bQX9vnTKW9HF6fcIEj48xrVKpDh3ek375HJdb5+Dg2OWzblMmf2ltbX3w0J4Tf6cd3P83cqXX1dVOnDxixPBxEz+f1lQSAIBcLt+157cbNy5LpRIGw3/a1DkREVEAgLLy0p83rMnPz7WxofVLHqJUKv7998GunUcAAAcO7j761wEikXT40D821jZXr1368+j+kpIiKtU6LCx82tQ5uiLeFJVVFb/88t3LrAwq1XrQwOECQUPGi6f79h4znC+PV79j5+ZHj+7JZFJXV/cB/YeOHDEOSZI6YeiGn7dFR3VunMvyr+ZJJOItm/d8zNlGeHypztEVr7fDjgmVbnv8uuUnndLGgs/njR0/aOmSVT269/l4awaUhmjsotZGg6Chorx06/aNPj7Mbl17mTo72JXWvSt9l+VL18TFdTdd1unpZ3fv2RreIXLJ4pUmrYshwB69K6uanM3U3s6BTP7AJ3hLgUbvJnFzhaXncFv+aonSGFRpWECVhgVUaVhAlYYFVGlYQJWGBeM8T2e8eCQSC/E42J/OTQEWi42KjPv4wWyNo41MLmEyfVvSJAPlfSGS8EaZsM84SoeHdyKTiEYxhfIOagwGBz56thTjKG1FMlU7NxQAwMfLjNbIIAJVGhZQpWEBVRoWUKVhAVUaFlClYQFVGhZQpWEBVRoWUKVhAVUaFlClYQFVGhYso7RCoUhIjC0u5jS1Q01N9d27N83gycuXGWw26wMSLlg4/Z/TTfaKfou0P/ePGJX07crFH5CRsbCM0ng8/mjaOQPdl/86fqi0rNgMnvz6208Kpf6pxAyg0WhY+Tlh7cJbsrNIJPpj344f1235bq0JhzFpFsu0/Ppj346qqopvvv5h956tVVUVVlaU6urKsrKSeXOXxcbGb9q8/szZv93dPVUqVer4KSf+Tjt95gQGg6HRbGfNXBgaEiaXy5P7x8+cMf/CxdPz5y5//O/9mtrqBj6PSrWeN3fZ0OEJx/+66OTkDABY/9MqezuHGdPnTZs+LiIiuryitKGBr9FoVn6znk53nTRlVGlp8YZf1s6aubDlIyEBAIqLOVgs9vada6vWLBWLRSNHjBs3dhIA4N8nD/fs2SoSi0gkUsrI1MTEATk5WavWLMXhcD+s/2bh/K9sbe127NxcWVmuUqkiO3b6ctYiEom0e89Wnf+rVv74rhGjnHPLKJ2fnxsZGQMAYBewqqoqNm3c6eDgePjIH2lH98fGxo8dM+nM2b937jhsbW198uTRc+dPbdq408nJ+crViytXLT6adg4ZPwqPJyCd1v86fqi6unLDT9scHBwfPb7v4OCIyIxkNHbMJJVKVVRcGBAQvHb1BhwO9/0PXx8+snfxom/GpHx+8tTRnb8fbuzbxUtntu/431sOHz50GhnXDCE375VcLg8MDJk8aWZOTtas2RN790oSCBvWfrd83feb27ePKCsrmTZjXEBAcEhI2KCBI7JzXq77flN9PXfajHFfTJmdmDhAJpMtWjLz+Ikj48dN5hQV6Pxn5ee+a4TJ9AcfjYWUZuelpExAlJgzewkyiBMGgyERSQCAfHauu7untbW1TCbbf3DXV8vXIsp169pr3fpvq2uq8vNz3VzdBw8a8dpafu7UyV8iRvLzcwMCgpH1crm8pKQoMCC4pKQIADBzxgKkhaWbmwcypCiL/WZnHclJg5KTBhn2Pzf3VXLSoPi4HgCAoKBQDAZTW1t9+Mje/v2Gtm8fAQDw9PT29fXLyc1iMv3z83MDA4IBABcunmb4+iFllEwmd4qOzc55+Zb/e/du02vk48+5BZSuq6vl8eoDAoKRhcj/hmYqLMz38wtEjjwoMAQZi04oFGze8iPY8jqttbU1lULNz8+NiemC9C7n8err6mo7d45DdtClRcYTIhKJnp7eV69e9PML1A3qVlVd6exMR3bu3SvpAw4hN/dVaupU3eFotVpnZ/rzjCf57Lybt64g66VSKTLOYX5+bt+E/gCAp08fder0mc6IQNBApVq/5X9TRj4eCyidn5/r7uZhY22T9TKDTnfVRUVWfm5cXA9kIbxDJABArpC7uNCPpp17y0IeK2fQwOG6ZTrdVTcMVEEBq2fPvsjykycP/f2DsFgsu4Cl20Gr1WZmPpsyaZZWqy0szJ85Y8FbxpuN3nK5vJDDdrB/PZjcixdPnZ1dnJycVSrVb1v+8HD3bJxQKBJWVlUgkUOlVjUePO/ps8ejRo5v7L9KpdJrxChYoO7N+i/AsvJzA/xfB0+JRFLWiK93AAARQUlEQVReXhoYGAIAKCkpcnJyAQAwfP1EIiEyanBDA3/tdys4nAKVSlVYmK9LyGLl6JYBAEqVUqVSIUZOnjqKZMRm53E4bB6vHgBw6p9jBDyhW7fedXW1YrHY2cnlLfeSkwadPX3zrb/GN+mCAhYA4NHjewAAgVCQdnT/6FET8Hh8gH/Qw4d3dcMWX712CbmsbW3tXFzoAID2YRH37t1UKpVarfbPoweQgRYb+9+UEaNggTLNZucFB7dDFnS3yYICFpVKRa7l8A6Rmzavk8mkA/oPXbFs7br13yoVChweP3DAMAbDj81mabVaX18mkrDxjRkAMCH1iz17t168eNrfP8jXl4kMV1hQmD992tzlK+aKJWIHB8fvv/sfmUzGYrG+vswvpo/9+aetAf5BLfc/O+dlbOd4mUyW+vkwjUbTp3fSkCGjAAArlq/d9Ov6kyf/xGAw0dGxcV26I8cY+J97qeOn/rZ1w6Qpo3A4HJPh//OPW8lk8lv+6zViFNr+OCfV1VVjxw+6cO4OiUSytC8mx+TjnKSnn6uuqWq8RqVS6R0KO65Ldz+/AKNk2kIKClient4wyGwY4yhtrKd7U8AuYL015CyctP3ekRP+exyCHPRbFiygSsMCqjQsoErDAqo0LKBKwwKqNCygSsMCqjQsoErDAqo0LKBKwwKqNCw0oTRGa+NI0L8JpRVDpuIIRP1jT+pX2sGVWJxtaCZAlNZJVaHE1kl/EdWvNMUG7+pDlgiUJnYMxcjg8MDZU3/rmibv0zFJ9lcONzmXFEor5FpaRVC0DclK//jQ+lsMItRVys/trogf6mrrRCRTPnZ8aRQToVRo+DXyZ1e5HXva+XVocqheQ0oDABq4yn/T64teiW2dCbzqNhLM1Ro1FoszxpjZlodAwsqlas8Aq4497D38rQzs2YzSOmRiDaatPJENHz58165djo6OlnbEGGi1pJaF25a2GCRT24rOACjVEiIZQ7JqO0fUEuA6WpiBUWkGg4ExytQWnxQwKs3hcFpYO2lLwKh0SEgIWqahICcnBy3TUICWaVhAyzQs2NjYoGUaCoRCIVqmUdosMCodGhpqaRcsAIxKZ2dnW9oFCwCj0nACo9Le3t5o3RsKSkpK0Lo3SpsFRqVpNBoavaFAIBCg0RsKsFgsWqahQKPRoGUapc0Co9L29vZo9IYCHo+HRm+UNguMSqOtgGEBbQWM0paBUWm0bSgsoG1DUdoyMCqNtgKGBbQVMCygNTJYQGtksODu7m5pFywAjEpXVMA4zhqMSru5uVnaBQsAo9KVlZWWdsECwKh0cHAwWveGgtzcXAjr3i0dY7ANEBUVpdVqsVisRqNB/uNwuNTU1Llz51raNXMAUZkODw9HFrBYLPLf09Nz7NixlvbLTECk9OjRox0cHBqv6du3r5OTk+U8MisQKd23b18fHx/dTy8vr5SUFIt6ZFYgUhoAkJKSYmdnhywnJia+VcTbNnAp3bdvXwaDgRToUaNGWdodswKX0gCAUaNGUanUhIQEqAr0xz5lVRRIC7MkNWVyqUgtE6kxGKBQaIzqnklQKZU4PP6TeHliRcVhcRgra5yzF9kniOwbSv1gUx+itFSk/vcyP/tRA9maQKNT8SQ8noTHE3F4AhaWZ3NzoVVrlXKVSqFWK9WCarGgVhoYRYvqZevopn8CHQO8n9JarfbGcS7rmcA10NHGyQpHQKdhMStarVbEldaw6128SD1GONrYvcfsde+hdBlbceN4jZUdxcnX9kNdRTEO/AqRmCvq0NW2/WdNTqbzFi1VOuex4P55HrOzxydxe4OE0sxqv3bkLgNaVLVsUd27jC17fEXgF+uJytyq8OpAL2Ipnt8StGTn5st0cY741j887wgYv95/ElTn1fmFEaP72BverZkyLRGqLh2oRmVuzdCDnLIfi4tzm5lxthmlz++t9olyNapjKMbHK8L1+tFajcZQeDakNOuZUKHEkq3f+9ENxcxgMBgbus2Dc1wD+xhS+s4/XGc/uF4Zfro4+dpl3mlQyJt8R9mk0gWZQis7MtGqpdMholgcJ4Zdxk1+U1ubVJr1XGJlSzaZV6Zl1frEeh50jbqtHa1Yz0VNbW1S6eJsMc35w9+nWxAev0osafLSbsNY0UhSoVrEV+ndqv95uqZEdvMUzyXQpVnrhcUZ/5zbWF3LcXLwHJg07+qtfe6u/sMGLgUAiMS8sxd/LSh6Jpbw3egB/RJm+TOjAAD3H/+dfm3X5PEbT1/4X01tEYVi27v7pM5RgxCDZRW5F65sL6vIVauUAX6dBiUvcLB3AwAcPLoCAAzd2ffmvSOpo34IDY4vLc++cGV7eSVLqZS7ujCT+8wM9I9hFz79fd8sxFS74G6Txm1Qq1VXb+3LeHmFx6+0s6V36zKmS8zwZo9Lr3EAQHUNZ8Nvo2dM2n7nwVFOyQssBhse1mdQ8gIcDgcAePTk9O0HR+t55QQCmenbcUi/hXKFdMOWlFlTfmf6dgQAPM+8fOT4t8MGLkV8qKkt+nlLytzpf3h7tnueefnWvbTqWg6JROnYvm9yn5lEIhmJT727T2SxH+UXPlm74jKJRDHgdh2nvn1nclCUzbub9JdpsVCtbMH3R6VSvj9tKZlEnTtt79CBSy5c2V7PKwcAg4zYuPvA/KLSlynDVs6fccDLI2TPofmVVWwAAA6Ll8lEV2/9MWH0+u++vhYV0e/k2Z/4DTVIcfz9j1lYDHbm5O0zJm+TSAQ7989WqhQAAByOUFVTUFaROzV1k7dXmFIp331wPh5HnP75b/Nm7PPxar8vbQm/oYbhEz5+1A8AgPkzD4wZvhoAcC79t1t3D/fq9vni2Wnduow5ff5/j56cbva49BoHAOBweADA6YubenZNXbvi8riR3917dPxl9g0AQGHR8+On13X9LGXR7LQp4/8nETcc+utrurOvnS29qCQTsVxY9NzOls4pykB+FhQ9tyLbeLqHZGXfOnL820D/mEVfHk4Z+m3mq+snzqxH9sHh8A///ceV7j9z8nY8vpnnII0aCHn6J4TXr7REqMLim/9OlZ13VyJpGDZoqYd7kD8jauiARQJhHbIpv+BxeWXuyMFfBTCj6S6Mwf0W2tu53X14DNmq1qh6dp1gZ0vHYDAxkQPValVFVT4A4MG/JwEGM27kd250fy+P0DEjVtfzyl++uo6k4taXjR6+yo8RaU21w2JxMydvTxm20sM9yNWFmdR7ulIpKyrJxOHwZBIVAECxopHJVKlMdP/Rie7x4zt17O/k6NUlZnh0x/7X7xw0fFxNGdftEN6ul693BwBAgF8nR3uPsvIcAEBVTSGBQOrUcYCTg6ePV9j4lB8GJc8HAPgzojklL5CEBUXPOkcNLix+rXRh0fMAv05YLPb6nYNM38h+CbOcHL1CArv07/vlsxeX+A3VAAAAMAQCeUDibF/vDkjkMACOiBPx1Xo36a9aK2UaAoVo2CgSfMhka1cXJvKT4RNBpbxupVVcloXDEfwYkf+dOyzTJ6K8kqVL604PQBYoVjQAgEwmBACUlGZ5e4RaWb0OPvZ2rg72HuWVrMjwJACAs6MPlfL6MxoOh1eplP+c/6WiKl8qFWqBFgAgkTa85WFFJUutUQX6xejW+DEiHz09LZdLDITBZo27uQbolslkG6lMCADwY0RhAGbbnukxUQMD/WIc7N1pNo7I1fDP+Y1arVYk5tVxS7vEDL92e389r8LB3r2o5EWvbhM1Gk1ZRU7fXl/obDJ9IwEAlVVsO1s6AMDXu32zWiAQrAhq9fsojcVjlBJFs3YlUgFSgHRQ/lNCLpeo1crla7rqNmk0ahtrxzc+Ef5fIEKqC1KZuKIqb9nqeN16tVqpixNk8psvdLV1JTv3fenPjB4zfLUtzVmj0Xz/y8B3PZTLJQCA3/+YBd58m9ECAIQirgGlmzVO+P9RFLkU6M6+s6ftuXHn0IXL205I13t7hg3ut8DHKyyA2UkqE1bVFNbUFrnRA6hUOy+P0MKiDORuFegXo1TKNBr15eu7r9zY29is3gM3jEqu0uD0vynTrzTFBq9RSZq1SyCQFEpZ4zUSScN/zlHxeOLCWYcab8Vgmnn5SiZTGd4RIwYvb7ySSNQjScbLKxqNetzI75ArhsevasogAGDsyLVudL/G621t6QbcaKHxd3F3DRg3cq1areaUZFy6uvOPw4u+WXKWRnOiOzOKSjIrqvKZvhEAAIZ3OKfkhRZoHR08HR08NBoNDoePj03RVUsRrKnv/dpKpVDbeOiP8PpPPcUGp1bqDwKNcXTwlEga6urLkJ+FxRm6xxtvj3YqlUKtUbs4+yJ/eDzJltZMZd7HK6yuvtTRwVOXCgAMzUZP43uVWkkgkHWB4emLi2/tgAQJN9cAHI4gEtXrDFIothSKHQFv6N7UrHG9FJdmIfdyHA7nz4hK6j1dLOELhVwAQIBfTFFJZmHRc6QG7usTzinOKCp+gdxWsFish1swj1+pc9LB3gOLxVMotJbk2xi1QkW1ex+l6d5kEVferN2QwDgCgYQ8LBUWZ5y79JtOFX9mJw+3oD9PrGZzntbzKp69SN+0PfX+4xOGDcZGD5XLJUdPri2vyKutK7lyY+8vW8eUlr96d09vz3ZiCf/xs7MCYd29RydKy7KtqfYVlflSmQi5zeew7lXVFFqRrT/rNDT9xu6Ml1e49eXswqc798/569Raw24YMG4gVV7+g31HlmRmXa/jlpVX5N19+Je9nZu9nSsAwJ8ZzS58UlPLYfiEAwAY3h1q60pYBY8C/qtA9Igf/zL7xvXbB2pqi8sr8tJOrNq2Z5pM1sznqXeRCxV0L/3vu/RHbxwe48qwEtZJbJwMPb3RbBxTR607c2nzxm3j3ej+g/stOH56HVIUcDjc1Ambz13acvDoCoVC6mDn3qfH5O5xzXSCcrB3mzF5+/nLW7ftmYbF4lxd/CaN+8XHS099pF1w1x5x48+nbz1zcXNIQJfRw1fdvv/njTsHMVjskH6LggM+O3vxV4ZP+IzJ2wcmzbMi25y/vFUgrLOxdgwN6pqcMNOwGwaMd+8ypqlUvbtPUqlVZ9O3CAS1ZLK1r3eHqambkLYbfoxIoYjr7OhjTbUHAFhZ2bi6MKtqCvwZUUjaDu16jhm+5sadg+nXdiFpZ07ejtx6Wo5SplIp1C5NKN1kS4TMu/zsJ3LXoGa6LYklDcT/Ap1SpVi1LqF/39lxsSPfy0UUo8AtFdjRlL1H679FNvkBI7gT7en1UsOmpTLR+k3DApidEnpOwQDMzXtHMBhs+9CeH+0zyocg40vaJTo2tdVQ66L757jlJVpnhqFmK8WlWReubCsrz8Vgse6ugf37ztIbbFsbnOKMvYcXNbV1xYKTugf3T4WGajFGLh40vcnWQc20I9u+uCC4hzcW19Y69SiVcqGoye/2drauSB/rTwj2/dKR8zxsnZpsAd6M0qznwue3xfRAZ9O4h2Ic6kv4bt6Yz5INPX83c+UGdrTxYBC4RTxj+4ZiNAQ1IpxWYVjmFrX3jh/k6ETH1LBRsVsjglqxUigePKP5xrstuht1H+ZIpapqC+qN4RuK0eBXCMRVDcNnt2hwzPfol/U4vb44X0lzpZGozX/mQjEpaqWaVy6g2aj7jjf0Ar8x79fXsjhXfONYHZFKcvGzx5PQxoQWQKvV1hbw6suE3YY5hXZ+jxfjH9J/OvuR4NVDkVigpjpSaHQq0erT6HX+SaOUqYW1YhFXgsNpA8KpMYnN9M15lw8fE6GSI83PEFcVy2uKpUQyjmCFI1jhtCq0r7wxwWAwUqFCLlW7+Fg5uBACIqg+IR/YjNM4YwxKhCpxg1oh+wSGvvi0IJAwFBs8hYbDYj82akI0miTkfGLv/FA+GFRpWECVhgVUaVhAlYYFVGlY+D+LJJD1VR4a4QAAAABJRU5ErkJggg==",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's repeat the same run, adding in a simple yes/no approval step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'write_query': {'query': 'SELECT COUNT(EmployeeId) AS EmployeeCount FROM Employee;'}}\n",
      "{'__interrupt__': ()}\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Do you want to go to execute query? (yes/no):  yes\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'execute_query': {'result': '[(8,)]'}}\n",
      "{'generate_answer': {'answer': 'There are 8 employees.'}}\n"
     ]
    }
   ],
   "source": [
    "for step in graph.stream(\n",
    "    {\"question\": \"How many employees are there?\"},\n",
    "    config,\n",
    "    stream_mode=\"updates\",\n",
    "):\n",
    "    print(step)\n",
    "\n",
    "try:\n",
    "    user_approval = input(\"Do you want to go to execute query? (yes/no): \")\n",
    "except Exception:\n",
    "    user_approval = \"no\"\n",
    "\n",
    "if user_approval.lower() == \"yes\":\n",
    "    # If approved, continue the graph execution\n",
    "    for step in graph.stream(None, config, stream_mode=\"updates\"):\n",
    "        print(step)\n",
    "else:\n",
    "    print(\"Operation cancelled by user.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See [this](https://langchain-ai.github.io/langgraph/concepts/human_in_the_loop/) LangGraph guide for more detail and examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next steps\n",
    "\n",
    "For more complex query-generation, we may want to create few-shot prompts or add query-checking steps. For advanced techniques like this and more check out:\n",
    "\n",
    "* [Prompting strategies](/docs/how_to/sql_prompting): Advanced prompt engineering techniques.\n",
    "* [Query checking](/docs/how_to/sql_query_checking): Add query validation and error handling.\n",
    "* [Large databases](/docs/how_to/sql_large_db): Techniques for working with large databases."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agents {#agents}\n",
    "\n",
    "[Agents](/docs/concepts/agents) leverage the reasoning capabilities of LLMs to make decisions during execution. Using agents allows you to offload additional discretion over the query generation and execution process. Although their behavior is less predictable than the above \"chain\", they feature some advantages:\n",
    "\n",
    "- They can query the database as many times as needed to answer the user question.\n",
    "- They can recover from errors by running a generated query, catching the traceback and regenerating it correctly.\n",
    "- They can answer questions based on the databases' schema as well as on the databases' content (like describing a specific table).\n",
    "\n",
    "\n",
    "Below we assemble a minimal SQL agent. We will equip it with a set of tools using LangChain's [SQLDatabaseToolkit](https://python.langchain.com/api_reference/community/agent_toolkits/langchain_community.agent_toolkits.sql.toolkit.SQLDatabaseToolkit.html). Using LangGraph's [pre-built ReAct agent constructor](https://langchain-ai.github.io/langgraph/how-tos/#langgraph.prebuilt.chat_agent_executor.create_react_agent), we can do this in one line.\n",
    "\n",
    ":::tip\n",
    "\n",
    "Check out LangGraph's [SQL Agent Tutorial](https://langchain-ai.github.io/langgraph/tutorials/sql-agent/) for a more advanced formulation of a SQL agent.\n",
    "\n",
    ":::\n",
    "\n",
    "The `SQLDatabaseToolkit` includes tools that can:\n",
    "\n",
    "* Create and execute queries\n",
    "* Check query syntax\n",
    "* Retrieve table descriptions\n",
    "* ... and more"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[QuerySQLDatabaseTool(description=\"Input to this tool is a detailed and correct SQL query, output is a result from the database. If the query is not correct, an error message will be returned. If an error is returned, rewrite the query, check the query, and try again. If you encounter an issue with Unknown column 'xxxx' in 'field list', use sql_db_schema to query the correct table fields.\", db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x10d5f9120>),\n",
       " InfoSQLDatabaseTool(description='Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables. Be sure that the tables actually exist by calling sql_db_list_tables first! Example Input: table1, table2, table3', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x10d5f9120>),\n",
       " ListSQLDatabaseTool(db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x10d5f9120>),\n",
       " QuerySQLCheckerTool(description='Use this tool to double check if your query is correct before executing it. Always use this tool before executing a query with sql_db_query!', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x10d5f9120>, llm=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x119315480>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x119317550>, root_client=<openai.OpenAI object at 0x10d5f8df0>, root_async_client=<openai.AsyncOpenAI object at 0x1193154e0>, model_name='gpt-4o', temperature=0.0, model_kwargs={}, openai_api_key=SecretStr('**********')), llm_chain=LLMChain(verbose=False, prompt=PromptTemplate(input_variables=['dialect', 'query'], input_types={}, partial_variables={}, template='\\n{query}\\nDouble check the {dialect} query above for common mistakes, including:\\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\\nIf there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.\\n\\nOutput the final SQL query only.\\n\\nSQL Query: '), llm=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x119315480>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x119317550>, root_client=<openai.OpenAI object at 0x10d5f8df0>, root_async_client=<openai.AsyncOpenAI object at 0x1193154e0>, model_name='gpt-4o', temperature=0.0, model_kwargs={}, openai_api_key=SecretStr('**********')), output_parser=StrOutputParser(), llm_kwargs={}))]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_community.agent_toolkits import SQLDatabaseToolkit\n",
    "\n",
    "toolkit = SQLDatabaseToolkit(db=db, llm=llm)\n",
    "\n",
    "tools = toolkit.get_tools()\n",
    "\n",
    "tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### System Prompt\n",
    "\n",
    "We will also want to load a system prompt for our agent. This will consist of instructions for how to behave. Note that the prompt below has several parameters, which we assign below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_message = \"\"\"\n",
    "You are an agent designed to interact with a SQL database.\n",
    "Given an input question, create a syntactically correct {dialect} query to run,\n",
    "then look at the results of the query and return the answer. Unless the user\n",
    "specifies a specific number of examples they wish to obtain, always limit your\n",
    "query to at most {top_k} results.\n",
    "\n",
    "You can order the results by a relevant column to return the most interesting\n",
    "examples in the database. Never query for all the columns from a specific table,\n",
    "only ask for the relevant columns given the question.\n",
    "\n",
    "You MUST double check your query before executing it. If you get an error while\n",
    "executing a query, rewrite the query and try again.\n",
    "\n",
    "DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the\n",
    "database.\n",
    "\n",
    "To start you should ALWAYS look at the tables in the database to see what you\n",
    "can query. Do NOT skip this step.\n",
    "\n",
    "Then you should query the schema of the most relevant tables.\n",
    "\"\"\".format(\n",
    "    dialect=\"SQLite\",\n",
    "    top_k=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initializing agent\n",
    "\n",
    "We will use a prebuilt [LangGraph](/docs/concepts/architecture/#langgraph) agent to build our agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "agent_executor = create_react_agent(llm, tools, prompt=system_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider how the agent responds to the below question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Which country's customers spent the most?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  sql_db_list_tables (call_tFp7HYD6sAAmCShgeqkVZH6Q)\n",
      " Call ID: call_tFp7HYD6sAAmCShgeqkVZH6Q\n",
      "  Args:\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: sql_db_list_tables\n",
      "\n",
      "Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  sql_db_schema (call_KJZ1Jx6JazyDdJa0uH1UeiOz)\n",
      " Call ID: call_KJZ1Jx6JazyDdJa0uH1UeiOz\n",
      "  Args:\n",
      "    table_names: Customer, Invoice\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: sql_db_schema\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 \"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",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  sql_db_query_checker (call_AQuTGbgH63u4gPgyV723yrjX)\n",
      " Call ID: call_AQuTGbgH63u4gPgyV723yrjX\n",
      "  Args:\n",
      "    query: SELECT c.Country, SUM(i.Total) as TotalSpent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.Country ORDER BY TotalSpent DESC LIMIT 1;\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: sql_db_query_checker\n",
      "\n",
      "```sql\n",
      "SELECT c.Country, SUM(i.Total) as TotalSpent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.Country ORDER BY TotalSpent DESC LIMIT 1;\n",
      "```\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  sql_db_query (call_B88EwU44nwwpQL5M9nlcemSU)\n",
      " Call ID: call_B88EwU44nwwpQL5M9nlcemSU\n",
      "  Args:\n",
      "    query: SELECT c.Country, SUM(i.Total) as TotalSpent FROM Customer c JOIN Invoice i ON c.CustomerId = i.CustomerId GROUP BY c.Country ORDER BY TotalSpent DESC LIMIT 1;\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: sql_db_query\n",
      "\n",
      "[('USA', 523.06)]\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The country whose customers spent the most is the USA, with a total spending of 523.06.\n"
     ]
    }
   ],
   "source": [
    "question = \"Which country's customers spent the most?\"\n",
    "\n",
    "for step in agent_executor.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}]},\n",
    "    stream_mode=\"values\",\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also use the [LangSmith trace](https://smith.langchain.com/public/8af422aa-b651-4bfe-8683-e2a7f4ccd82c/r) to visualize these steps and associated metadata.\n",
    "\n",
    "Note that the agent executes multiple queries until it has the information it needs:\n",
    "1. List available tables;\n",
    "2. Retrieves the schema for three tables;\n",
    "3. Queries multiple of the tables via a join operation.\n",
    "\n",
    "The agent is then able to use the result of the final query to generate an answer to the original question.\n",
    "\n",
    "The agent can similarly handle qualitative questions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Describe the playlisttrack table\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  sql_db_list_tables (call_fMF8eTmX5TJDJjc3Mhdg52TI)\n",
      " Call ID: call_fMF8eTmX5TJDJjc3Mhdg52TI\n",
      "  Args:\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: sql_db_list_tables\n",
      "\n",
      "Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  sql_db_schema (call_W8Vkk4NEodkAAIg8nexAszUH)\n",
      " Call ID: call_W8Vkk4NEodkAAIg8nexAszUH\n",
      "  Args:\n",
      "    table_names: PlaylistTrack\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: sql_db_schema\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",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The `PlaylistTrack` table is designed to associate tracks with playlists. It has the following structure:\n",
      "\n",
      "- **PlaylistId**: An integer that serves as a foreign key referencing the `Playlist` table. It is part of the composite primary key.\n",
      "- **TrackId**: An integer that serves as a foreign key referencing the `Track` table. It is also part of the composite primary key.\n",
      "\n",
      "The primary key for this table is a composite key consisting of both `PlaylistId` and `TrackId`, ensuring that each track can be uniquely associated with a playlist. The table enforces referential integrity by linking to the `Track` and `Playlist` tables through foreign keys.\n"
     ]
    }
   ],
   "source": [
    "question = \"Describe the playlisttrack table\"\n",
    "\n",
    "for step in agent_executor.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}]},\n",
    "    stream_mode=\"values\",\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dealing with high-cardinality columns\n",
    "\n",
    "In order to filter columns that contain proper nouns such as addresses, song names or artists, we first need to double-check the spelling in order to filter the data correctly. \n",
    "\n",
    "We can achieve this by creating a vector store with all the distinct proper nouns that exist in the database. We can then have the agent query that vector store each time the user includes a proper noun in their question, to find the correct spelling for that word. In this way, the agent can make sure it understands which entity the user is referring to before building the target query.\n",
    "\n",
    "First we need the unique values for each entity we want, for which we define a function that parses the result into a list of elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['In Through The Out Door',\n",
       " 'Transmission',\n",
       " 'Battlestar Galactica (Classic), Season',\n",
       " 'A Copland Celebration, Vol. I',\n",
       " 'Quiet Songs']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ast\n",
    "import re\n",
    "\n",
    "\n",
    "def query_as_list(db, query):\n",
    "    res = db.run(query)\n",
    "    res = [el for sub in ast.literal_eval(res) for el in sub if el]\n",
    "    res = [re.sub(r\"\\b\\d+\\b\", \"\", string).strip() for string in res]\n",
    "    return list(set(res))\n",
    "\n",
    "\n",
    "artists = query_as_list(db, \"SELECT Name FROM Artist\")\n",
    "albums = query_as_list(db, \"SELECT Title FROM Album\")\n",
    "albums[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using this function, we can create a **retriever tool** that the agent can execute at its discretion.\n",
    "\n",
    "Let's select an [embeddings model](/docs/integrations/text_embedding/) and [vector store](/docs/integrations/vectorstores/) for this step:\n",
    "\n",
    "**Select an embedding model**:\n",
    "\n",
    "import EmbeddingTabs from \"@theme/EmbeddingTabs\";\n",
    "\n",
    "<EmbeddingTabs/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | output: false\n",
    "# | echo: false\n",
    "\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "embeddings = OpenAIEmbeddings()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Select a vector store**:\n",
    "\n",
    "import VectorStoreTabs from \"@theme/VectorStoreTabs\";\n",
    "\n",
    "<VectorStoreTabs/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# | output: false\n",
    "# | echo: false\n",
    "\n",
    "from langchain_core.vectorstores import InMemoryVectorStore\n",
    "\n",
    "vector_store = InMemoryVectorStore(embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now construct a retrieval tool that can search over relevant proper nouns in the database:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.agent_toolkits import create_retriever_tool\n",
    "\n",
    "_ = vector_store.add_texts(artists + albums)\n",
    "retriever = vector_store.as_retriever(search_kwargs={\"k\": 5})\n",
    "description = (\n",
    "    \"Use to look up values to filter on. Input is an approximate spelling \"\n",
    "    \"of the proper noun, output is valid proper nouns. Use the noun most \"\n",
    "    \"similar to the search.\"\n",
    ")\n",
    "retriever_tool = create_retriever_tool(\n",
    "    retriever,\n",
    "    name=\"search_proper_nouns\",\n",
    "    description=description,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try it out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Alice In Chains\n",
      "\n",
      "Alanis Morissette\n",
      "\n",
      "Pearl Jam\n",
      "\n",
      "Pearl Jam\n",
      "\n",
      "Audioslave\n"
     ]
    }
   ],
   "source": [
    "print(retriever_tool.invoke(\"Alice Chains\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This way, if the agent determines it needs to write a filter based on an artist along the lines of \"Alice Chains\", it can first use the retriever tool to observe relevant values of a column.\n",
    "\n",
    "Putting this together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add to system message\n",
    "suffix = (\n",
    "    \"If you need to filter on a proper noun like a Name, you must ALWAYS first look up \"\n",
    "    \"the filter value using the 'search_proper_nouns' tool! Do not try to \"\n",
    "    \"guess at the proper name - use this function to find similar ones.\"\n",
    ")\n",
    "\n",
    "system = f\"{system_message}\\n\\n{suffix}\"\n",
    "\n",
    "tools.append(retriever_tool)\n",
    "\n",
    "agent = create_react_agent(llm, tools, prompt=system)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "How many albums does alis in chain have?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  search_proper_nouns (call_8ryjsRPLAr79mM3Qvnq6gTOH)\n",
      " Call ID: call_8ryjsRPLAr79mM3Qvnq6gTOH\n",
      "  Args:\n",
      "    query: alis in chain\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: search_proper_nouns\n",
      "\n",
      "Alice In Chains\n",
      "\n",
      "Aisha Duo\n",
      "\n",
      "Xis\n",
      "\n",
      "Da Lama Ao Caos\n",
      "\n",
      "A-Sides\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  sql_db_list_tables (call_NJjtCpU89MBMplssjn1z0xzq)\n",
      " Call ID: call_NJjtCpU89MBMplssjn1z0xzq\n",
      "  Args:\n",
      "  search_proper_nouns (call_1BfrueC9koSIyi4OfMu2Ao8q)\n",
      " Call ID: call_1BfrueC9koSIyi4OfMu2Ao8q\n",
      "  Args:\n",
      "    query: Alice In Chains\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: search_proper_nouns\n",
      "\n",
      "Alice In Chains\n",
      "\n",
      "Pearl Jam\n",
      "\n",
      "Pearl Jam\n",
      "\n",
      "Foo Fighters\n",
      "\n",
      "Soundgarden\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  sql_db_schema (call_Kn09w9jd9swcNzIZ1b5MlKID)\n",
      " Call ID: call_Kn09w9jd9swcNzIZ1b5MlKID\n",
      "  Args:\n",
      "    table_names: Album, Artist\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: sql_db_schema\n",
      "\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",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  sql_db_query (call_WkHRiPcBoGN9bc58MIupRHKP)\n",
      " Call ID: call_WkHRiPcBoGN9bc58MIupRHKP\n",
      "  Args:\n",
      "    query: SELECT COUNT(*) FROM Album WHERE ArtistId = (SELECT ArtistId FROM Artist WHERE Name = 'Alice In Chains')\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: sql_db_query\n",
      "\n",
      "[(1,)]\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Alice In Chains has released 1 album in the database.\n"
     ]
    }
   ],
   "source": [
    "question = \"How many albums does alis in chain have?\"\n",
    "\n",
    "for step in agent.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}]},\n",
    "    stream_mode=\"values\",\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, both in the streamed steps and in the [LangSmith trace](https://smith.langchain.com/public/1d757ed2-5688-4458-9400-023594e2c5a7/r), the agent used the `search_proper_nouns` tool in order to check how to correctly query the database for this specific artist."
   ]
  }
 ],
 "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
}