backup/langchain_text2SQL.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "if not os.environ.get(\"OPENAI_API_KEY\"):\n",
    "    raise EnvironmentError(\n",
    "        \"OPENAI_API_KEY not found in environment variables.\")\n",
    "\n",
    "# Comment out the below to opt-out of using LangSmith in this notebook. Not required.\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"briefing_generator\"\n",
    "if not os.environ.get(\"LANGCHAIN_API_KEY\"):\n",
    "    raise EnvironmentError(\n",
    "        \"LANGCHAIN_API_KEY not found in environment variables.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_community.utilities import SQLDatabase\n",
    "\n",
    "db = SQLDatabase.from_uri(\"sqlite:///../db/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": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'```sql\\nSQLQuery: SELECT \"driver_name\", MIN(\"lap_time_in_seconds\") AS \"fastest_lap_time\" FROM \"Laps\" GROUP BY \"driver_name\" ORDER BY \"fastest_lap_time\" LIMIT 1\\n```'"
      ]
     },
     "execution_count": 6,
     "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": "markdown",
   "metadata": {},
   "source": [
    "## Fixing SQL errors\n",
    "\n",
    "We will use a validation chain to fix the SQL errors like this ``sql ... `\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": 8,
   "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\": \"Who is the fastest driver?\"\n",
    "    }\n",
    ")\n",
    "print(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"[('VER', 89.708)]\""
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "db.run(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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": 11,
   "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": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"[('VER', 89.708)]\""
      ]
     },
     "execution_count": 12,
     "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": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The fastest driver is VER, with a lap time of 89.708 seconds.'"
      ]
     },
     "execution_count": 13,
     "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?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agents\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "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 0x106a960d0>),\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 0x106a960d0>),\n",
       " ListSQLDatabaseTool(db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x106a960d0>),\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 0x106a960d0>, llm=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x11c49ec90>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x11c4cb410>, root_client=<openai.OpenAI object at 0x11c4cb8d0>, root_async_client=<openai.AsyncOpenAI object at 0x11c47fa90>, model_name='gpt-4o-mini', 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 0x11c49ec90>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x11c4cb410>, root_client=<openai.OpenAI object at 0x11c4cb8d0>, root_async_client=<openai.AsyncOpenAI object at 0x11c47fa90>, model_name='gpt-4o-mini', model_kwargs={}, openai_api_key=SecretStr('**********')), output_parser=StrOutputParser(), llm_kwargs={}))]"
      ]
     },
     "execution_count": 14,
     "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SystemMessage(content='You are an agent designed to interact with a SQL database.\\nGiven an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.\\nUnless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.\\nYou can order the results by a relevant column to return the most interesting examples in the database.\\nNever query for all the columns from a specific table, only ask for the relevant columns given the question.\\nYou have access to tools for interacting with the database.\\nOnly use the below tools. Only use the information returned by the below tools to construct your final answer.\\nYou MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.\\n\\nDO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.\\n\\nTo start you should ALWAYS look at the tables in the database to see what you can query.\\nDo NOT skip this step.\\nThen you should query the schema of the most relevant tables.', additional_kwargs={}, response_metadata={})"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import SystemMessage\n",
    "\n",
    "SQL_PREFIX = \"\"\"You are an agent designed to interact with a SQL database.\n",
    "Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.\n",
    "Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.\n",
    "You can order the results by a relevant column to return the most interesting examples in the database.\n",
    "Never query for all the columns from a specific table, only ask for the relevant columns given the question.\n",
    "You have access to tools for interacting with the database.\n",
    "Only use the below tools. Only use the information returned by the below tools to construct your final answer.\n",
    "You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.\n",
    "\n",
    "DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.\n",
    "\n",
    "To start you should ALWAYS look at the tables in the database to see what you can query.\n",
    "Do NOT skip this step.\n",
    "Then you should query the schema of the most relevant tables.\"\"\"\n",
    "\n",
    "system_message = SystemMessage(content=SQL_PREFIX)\n",
    "system_message"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use a prebuilt LangGraph agent to build our agent\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/09/6cwl3kq50bv50t8lrvqxmmfc0000gn/T/ipykernel_86930/2283420487.py:4: LangGraphDeprecationWarning: Parameter 'messages_modifier' in function 'create_react_agent' is deprecated as of version 0.1.9 and will be removed in version 0.3.0. Use 'state_modifier' parameter instead.\n",
      "  agent_executor = create_react_agent(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_oy9eFT0bKTRTlyByDdybkvDw', 'function': {'arguments': '{}', 'name': 'sql_db_list_tables'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 12, 'prompt_tokens': 547, 'total_tokens': 559, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-85d5fbe0-5a67-49e2-85a2-f6be2dfe5b17-0', tool_calls=[{'name': 'sql_db_list_tables', 'args': {}, 'id': 'call_oy9eFT0bKTRTlyByDdybkvDw', 'type': 'tool_call'}], usage_metadata={'input_tokens': 547, 'output_tokens': 12, 'total_tokens': 559, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content='Drivers, Event, Laps, Sessions, Telemetry, Tracks, Weather', name='sql_db_list_tables', id='6595a9fb-d4d6-4ba3-8b1c-e153d5b0daad', tool_call_id='call_oy9eFT0bKTRTlyByDdybkvDw')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_u4GgVmf1WbLtt3Suy4nYTZG3', 'function': {'arguments': '{\"table_names\": \"Drivers\"}', 'name': 'sql_db_schema'}, 'type': 'function'}, {'id': 'call_VPcwJvUKTHLgPcxlKVCgsIsZ', 'function': {'arguments': '{\"table_names\": \"Laps\"}', 'name': 'sql_db_schema'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 49, 'prompt_tokens': 584, 'total_tokens': 633, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-dc918370-84c0-476d-9b59-dc13a4de2439-0', tool_calls=[{'name': 'sql_db_schema', 'args': {'table_names': 'Drivers'}, 'id': 'call_u4GgVmf1WbLtt3Suy4nYTZG3', 'type': 'tool_call'}, {'name': 'sql_db_schema', 'args': {'table_names': 'Laps'}, 'id': 'call_VPcwJvUKTHLgPcxlKVCgsIsZ', 'type': 'tool_call'}], usage_metadata={'input_tokens': 584, 'output_tokens': 49, 'total_tokens': 633, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content='\\nCREATE TABLE \"Drivers\" (\\n\\tdriver_id INTEGER, \\n\\tdriver_name TEXT NOT NULL, \\n\\tteam TEXT NOT NULL, \\n\\tPRIMARY KEY (driver_id)\\n)\\n\\n/*\\n3 rows from Drivers table:\\ndriver_id\\tdriver_name\\tteam\\n1\\tMax Verstappen\\tRed Bull Racing\\n2\\tSergio Perez\\tRed Bull Racing\\n3\\tCharles Leclerc\\tFerrari\\n*/', name='sql_db_schema', id='77207876-0f98-4dd8-8ecb-ffd10b51217b', tool_call_id='call_u4GgVmf1WbLtt3Suy4nYTZG3'), ToolMessage(content='\\nCREATE TABLE \"Laps\" (\\n\\tlap_id INTEGER, \\n\\tsession_id INTEGER, \\n\\tdriver_name TEXT NOT NULL, \\n\\tlap_number INTEGER NOT NULL, \\n\\tstint INTEGER, \\n\\tsector_1_speed_trap_in_km REAL, \\n\\tsector_2_speed_trap_in_km REAL, \\n\\tfinish_line_speed_trap_in_km REAL, \\n\\tlongest_strait_speed_trap_in_km REAL, \\n\\tis_personal_best BOOLEAN, \\n\\ttyre_compound TEXT, \\n\\ttyre_life_in_laps INTEGER, \\n\\tis_fresh_tyre BOOLEAN, \\n\\tposition INTEGER, \\n\\tlap_time_in_seconds REAL, \\n\\tsector_1_time_in_seconds REAL, \\n\\tsector_2_time_in_seconds REAL, \\n\\tsector_3_time_in_seconds REAL, \\n\\tlap_start_time_in_datetime DATETIME, \\n\\tpin_in_time_in_datetime DATETIME, \\n\\tpin_out_time_in_datetime DATETIME, \\n\\tPRIMARY KEY (lap_id), \\n\\tFOREIGN KEY(session_id) REFERENCES \"Sessions\" (session_id), \\n\\tUNIQUE (session_id, driver_name, lap_number)\\n)\\n\\n/*\\n3 rows from Laps table:\\nlap_id\\tsession_id\\tdriver_name\\tlap_number\\tstint\\tsector_1_speed_trap_in_km\\tsector_2_speed_trap_in_km\\tfinish_line_speed_trap_in_km\\tlongest_strait_speed_trap_in_km\\tis_personal_best\\ttyre_compound\\ttyre_life_in_laps\\tis_fresh_tyre\\tposition\\tlap_time_in_seconds\\tsector_1_time_in_seconds\\tsector_2_time_in_seconds\\tsector_3_time_in_seconds\\tlap_start_time_in_datetime\\tpin_in_time_in_datetime\\tpin_out_time_in_datetime\\n1\\t1\\tVER\\t1\\tNone\\tNone\\tNone\\tNone\\t178.0\\tFalse\\tSOFT\\t1\\tTrue\\tNone\\tNone\\tNone\\tNone\\tNone\\t2023-03-04 15:04:00.840000\\t2023-03-04 15:06:07.931000\\t2023-03-04 15:04:00.840000\\n2\\t1\\tVER\\t2\\tNone\\tNone\\tNone\\t288.0\\t141.0\\tFalse\\tSOFT\\t2\\tFalse\\tNone\\tNone\\tNone\\t53.666\\t38.509\\t2023-03-04 15:12:21.456000\\tNone\\t2023-03-04 15:12:21.456000\\n3\\t1\\tVER\\t3\\tNone\\tNone\\tNone\\t288.0\\t322.0\\tTrue\\tSOFT\\t3\\tFalse\\tNone\\t91.295\\t29.152\\t39.195\\t22.948\\t2023-03-04 15:14:33.391000\\tNone\\tNone\\n*/', name='sql_db_schema', id='b19bb039-85b0-4893-a826-d2d0288134f6', tool_call_id='call_VPcwJvUKTHLgPcxlKVCgsIsZ')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_5MdS7VtbFs4MrJ5ksZMlF2wN', 'function': {'arguments': '{\"query\":\"SELECT driver_name, MIN(lap_time_in_seconds) AS fastest_lap FROM Laps GROUP BY driver_name ORDER BY fastest_lap LIMIT 5;\"}', 'name': 'sql_db_query_checker'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 46, 'prompt_tokens': 1387, 'total_tokens': 1433, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-09cb8288-2acb-48b2-9ed3-879719cd3f6e-0', tool_calls=[{'name': 'sql_db_query_checker', 'args': {'query': 'SELECT driver_name, MIN(lap_time_in_seconds) AS fastest_lap FROM Laps GROUP BY driver_name ORDER BY fastest_lap LIMIT 5;'}, 'id': 'call_5MdS7VtbFs4MrJ5ksZMlF2wN', 'type': 'tool_call'}], usage_metadata={'input_tokens': 1387, 'output_tokens': 46, 'total_tokens': 1433, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content='```sql\\nSELECT driver_name, MIN(lap_time_in_seconds) AS fastest_lap FROM Laps GROUP BY driver_name ORDER BY fastest_lap LIMIT 5;\\n```', name='sql_db_query_checker', id='da986a31-e492-4dee-9dfe-528f5063b97b', tool_call_id='call_5MdS7VtbFs4MrJ5ksZMlF2wN')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_rgS5p53gmkoIkAzHzrV1Hv3o', 'function': {'arguments': '{\"query\":\"SELECT driver_name, MIN(lap_time_in_seconds) AS fastest_lap FROM Laps GROUP BY driver_name ORDER BY fastest_lap LIMIT 5;\"}', 'name': 'sql_db_query'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 45, 'prompt_tokens': 1478, 'total_tokens': 1523, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 1280}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-04a34824-f5a6-4623-8078-7954ea60a6f6-0', tool_calls=[{'name': 'sql_db_query', 'args': {'query': 'SELECT driver_name, MIN(lap_time_in_seconds) AS fastest_lap FROM Laps GROUP BY driver_name ORDER BY fastest_lap LIMIT 5;'}, 'id': 'call_rgS5p53gmkoIkAzHzrV1Hv3o', 'type': 'tool_call'}], usage_metadata={'input_tokens': 1478, 'output_tokens': 45, 'total_tokens': 1523, 'input_token_details': {'cache_read': 1280}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content=\"[('VER', 89.708), ('PER', 89.846), ('LEC', 90.0), ('SAI', 90.154), ('ALO', 90.336)]\", name='sql_db_query', id='f1e96f9f-540a-49db-b63d-7fdd1b0d3998', tool_call_id='call_rgS5p53gmkoIkAzHzrV1Hv3o')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='The fastest drivers based on their best lap times are:\\n\\n1. Max Verstappen (VER) - 89.708 seconds\\n2. Sergio Perez (PER) - 89.846 seconds\\n3. Charles Leclerc (LEC) - 90.0 seconds\\n4. Carlos Sainz (SAI) - 90.154 seconds\\n5. Fernando Alonso (ALO) - 90.336 seconds', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 87, 'prompt_tokens': 1575, 'total_tokens': 1662, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 1408}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'stop', 'logprobs': None}, id='run-7bff3793-8bc0-4eae-9c76-efc6559efcfc-0', usage_metadata={'input_tokens': 1575, 'output_tokens': 87, 'total_tokens': 1662, 'input_token_details': {'cache_read': 1408}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "agent_executor = create_react_agent(\n",
    "    llm, tools, messages_modifier=system_message)\n",
    "\n",
    "for s in agent_executor.stream(\n",
    "    {\"messages\": [HumanMessage(content=\"Who is the fastest driver?\")]}\n",
    "):\n",
    "    print(s)\n",
    "    print(\"----\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a retriever tool\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Lewis Hamilton',\n",
       " 'Charles Leclerc',\n",
       " 'Lando Norris',\n",
       " 'Nico Hulkenberg',\n",
       " 'Yuki Tsunoda']"
      ]
     },
     "execution_count": 17,
     "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",
    "drivers = query_as_list(db, \"SELECT driver_name FROM Drivers\")\n",
    "drivers[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents.agent_toolkits import create_retriever_tool\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "vector_db = FAISS.from_texts(drivers, OpenAIEmbeddings())\n",
    "retriever = vector_db.as_retriever(search_kwargs={\"k\": 5})\n",
    "description = \"\"\"Use to look up values to filter on. Input is an approximate spelling of the proper noun, output is \\\n",
    "valid proper nouns. Use the noun most similar to the search.\"\"\"\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\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Oscar Piastri\n",
      "\n",
      "Fernando Alonso\n",
      "\n",
      "Pierre Gasly\n",
      "\n",
      "Sergio Perez\n",
      "\n",
      "George Russell\n"
     ]
    }
   ],
   "source": [
    "print(retriever_tool.invoke(\"PIA\"))"
   ]
  },
  {
   "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 \"PIA\", it can first use the retriever tool to observe relevant values of a column.\n",
    "\n",
    "Putting this together:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/09/6cwl3kq50bv50t8lrvqxmmfc0000gn/T/ipykernel_86930/3211409952.py:23: LangGraphDeprecationWarning: Parameter 'messages_modifier' in function 'create_react_agent' is deprecated as of version 0.1.9 and will be removed in version 0.3.0. Use 'state_modifier' parameter instead.\n",
      "  agent = create_react_agent(llm, tools, messages_modifier=system_message)\n"
     ]
    }
   ],
   "source": [
    "system = \"\"\"You are an agent designed to interact with a SQL database.\n",
    "Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.\n",
    "Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.\n",
    "You can order the results by a relevant column to return the most interesting examples in the database.\n",
    "Never query for all the columns from a specific table, only ask for the relevant columns given the question.\n",
    "You have access to tools for interacting with the database.\n",
    "Only use the given tools. Only use the information returned by the tools to construct your final answer.\n",
    "You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.\n",
    "\n",
    "DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.\n",
    "\n",
    "You have access to the following tables: {table_names}\n",
    "\n",
    "If you need to filter on a proper noun, you must ALWAYS first look up the filter value using the \"search_proper_nouns\" tool!\n",
    "Do not try to guess at the proper name - use this function to find similar ones.\"\"\".format(\n",
    "    table_names=db.get_usable_table_names()\n",
    ")\n",
    "\n",
    "system_message = SystemMessage(content=system)\n",
    "\n",
    "tools.append(retriever_tool)\n",
    "\n",
    "agent = create_react_agent(llm, tools, messages_modifier=system_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_V2mYVxQ9Tbe8nklNbMCEpyPX', 'function': {'arguments': '{\"query\":\"Lewis Hamilton\"}', 'name': 'search_proper_nouns'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 18, 'prompt_tokens': 649, 'total_tokens': 667, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-d435d5b9-fd52-46ea-bc9a-e25cee3cc02a-0', tool_calls=[{'name': 'search_proper_nouns', 'args': {'query': 'Lewis Hamilton'}, 'id': 'call_V2mYVxQ9Tbe8nklNbMCEpyPX', 'type': 'tool_call'}], usage_metadata={'input_tokens': 649, 'output_tokens': 18, 'total_tokens': 667, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content='Lewis Hamilton\\n\\nLando Norris\\n\\nNico Hulkenberg\\n\\nValtteri Bottas\\n\\nFernando Alonso', name='search_proper_nouns', id='e442b4ac-e0b1-49de-aa79-6b920abf7ee7', tool_call_id='call_V2mYVxQ9Tbe8nklNbMCEpyPX')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_vQBxLBhzI0H8uGWMadgizqt2', 'function': {'arguments': '{\"query\":\"SELECT COUNT(*) AS FastestLaps FROM Laps WHERE Driver = \\'Lewis Hamilton\\' AND FastestLap = 1\"}', 'name': 'sql_db_query_checker'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 40, 'prompt_tokens': 699, 'total_tokens': 739, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-1708a779-f9d3-4dde-b6bf-824cc41dd6f7-0', tool_calls=[{'name': 'sql_db_query_checker', 'args': {'query': \"SELECT COUNT(*) AS FastestLaps FROM Laps WHERE Driver = 'Lewis Hamilton' AND FastestLap = 1\"}, 'id': 'call_vQBxLBhzI0H8uGWMadgizqt2', 'type': 'tool_call'}], usage_metadata={'input_tokens': 699, 'output_tokens': 40, 'total_tokens': 739, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content=\"```sql\\nSELECT COUNT(*) AS FastestLaps FROM Laps WHERE Driver = 'Lewis Hamilton' AND FastestLap = 1\\n```\", name='sql_db_query_checker', id='c026375c-7c11-423d-95cf-6ff56b53e53a', tool_call_id='call_vQBxLBhzI0H8uGWMadgizqt2')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_WpFTMmqGVLpwHzxnyGgr2b11', 'function': {'arguments': '{\"query\":\"SELECT COUNT(*) AS FastestLaps FROM Laps WHERE Driver = \\'Lewis Hamilton\\' AND FastestLap = 1\"}', 'name': 'sql_db_query'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 39, 'prompt_tokens': 779, 'total_tokens': 818, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-74afbb9c-9cb9-4661-a63e-d97cb44e09c5-0', tool_calls=[{'name': 'sql_db_query', 'args': {'query': \"SELECT COUNT(*) AS FastestLaps FROM Laps WHERE Driver = 'Lewis Hamilton' AND FastestLap = 1\"}, 'id': 'call_WpFTMmqGVLpwHzxnyGgr2b11', 'type': 'tool_call'}], usage_metadata={'input_tokens': 779, 'output_tokens': 39, 'total_tokens': 818, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content=\"Error: (sqlite3.OperationalError) no such column: Driver\\n[SQL: SELECT COUNT(*) AS FastestLaps FROM Laps WHERE Driver = 'Lewis Hamilton' AND FastestLap = 1]\\n(Background on this error at: https://sqlalche.me/e/20/e3q8)\", name='sql_db_query', id='9f2ad90b-32de-4f30-ab5d-bbdadfb6e7ac', tool_call_id='call_WpFTMmqGVLpwHzxnyGgr2b11')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_2M65yPLFFCjHsTc2rLUobtt4', 'function': {'arguments': '{\"table_names\":\"Laps\"}', 'name': 'sql_db_schema'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 17, 'prompt_tokens': 892, 'total_tokens': 909, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-ab8e7251-4de5-4539-a5ca-e66af17016dc-0', tool_calls=[{'name': 'sql_db_schema', 'args': {'table_names': 'Laps'}, 'id': 'call_2M65yPLFFCjHsTc2rLUobtt4', 'type': 'tool_call'}], usage_metadata={'input_tokens': 892, 'output_tokens': 17, 'total_tokens': 909, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content='\\nCREATE TABLE \"Laps\" (\\n\\tlap_id INTEGER, \\n\\tsession_id INTEGER, \\n\\tdriver_name TEXT NOT NULL, \\n\\tlap_number INTEGER NOT NULL, \\n\\tstint INTEGER, \\n\\tsector_1_speed_trap_in_km REAL, \\n\\tsector_2_speed_trap_in_km REAL, \\n\\tfinish_line_speed_trap_in_km REAL, \\n\\tlongest_strait_speed_trap_in_km REAL, \\n\\tis_personal_best BOOLEAN, \\n\\ttyre_compound TEXT, \\n\\ttyre_life_in_laps INTEGER, \\n\\tis_fresh_tyre BOOLEAN, \\n\\tposition INTEGER, \\n\\tlap_time_in_seconds REAL, \\n\\tsector_1_time_in_seconds REAL, \\n\\tsector_2_time_in_seconds REAL, \\n\\tsector_3_time_in_seconds REAL, \\n\\tlap_start_time_in_datetime DATETIME, \\n\\tpin_in_time_in_datetime DATETIME, \\n\\tpin_out_time_in_datetime DATETIME, \\n\\tPRIMARY KEY (lap_id), \\n\\tFOREIGN KEY(session_id) REFERENCES \"Sessions\" (session_id), \\n\\tUNIQUE (session_id, driver_name, lap_number)\\n)\\n\\n/*\\n3 rows from Laps table:\\nlap_id\\tsession_id\\tdriver_name\\tlap_number\\tstint\\tsector_1_speed_trap_in_km\\tsector_2_speed_trap_in_km\\tfinish_line_speed_trap_in_km\\tlongest_strait_speed_trap_in_km\\tis_personal_best\\ttyre_compound\\ttyre_life_in_laps\\tis_fresh_tyre\\tposition\\tlap_time_in_seconds\\tsector_1_time_in_seconds\\tsector_2_time_in_seconds\\tsector_3_time_in_seconds\\tlap_start_time_in_datetime\\tpin_in_time_in_datetime\\tpin_out_time_in_datetime\\n1\\t1\\tVER\\t1\\tNone\\tNone\\tNone\\tNone\\t178.0\\tFalse\\tSOFT\\t1\\tTrue\\tNone\\tNone\\tNone\\tNone\\tNone\\t2023-03-04 15:04:00.840000\\t2023-03-04 15:06:07.931000\\t2023-03-04 15:04:00.840000\\n2\\t1\\tVER\\t2\\tNone\\tNone\\tNone\\t288.0\\t141.0\\tFalse\\tSOFT\\t2\\tFalse\\tNone\\tNone\\tNone\\t53.666\\t38.509\\t2023-03-04 15:12:21.456000\\tNone\\t2023-03-04 15:12:21.456000\\n3\\t1\\tVER\\t3\\tNone\\tNone\\tNone\\t288.0\\t322.0\\tTrue\\tSOFT\\t3\\tFalse\\tNone\\t91.295\\t29.152\\t39.195\\t22.948\\t2023-03-04 15:14:33.391000\\tNone\\tNone\\n*/', name='sql_db_schema', id='58fa77cf-9ca2-44ce-a6b4-7c64c258de10', tool_call_id='call_2M65yPLFFCjHsTc2rLUobtt4')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_aYYhIdyMvFlatT5NmwZFz9a4', 'function': {'arguments': '{\"query\":\"SELECT COUNT(*) AS FastestLaps FROM Laps WHERE driver_name = \\'Lewis Hamilton\\' AND is_personal_best = 1\"}', 'name': 'sql_db_query_checker'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 42, 'prompt_tokens': 1514, 'total_tokens': 1556, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-7be6718b-adb2-443f-9913-cba293c2a94c-0', tool_calls=[{'name': 'sql_db_query_checker', 'args': {'query': \"SELECT COUNT(*) AS FastestLaps FROM Laps WHERE driver_name = 'Lewis Hamilton' AND is_personal_best = 1\"}, 'id': 'call_aYYhIdyMvFlatT5NmwZFz9a4', 'type': 'tool_call'}], usage_metadata={'input_tokens': 1514, 'output_tokens': 42, 'total_tokens': 1556, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content=\"```sql\\nSELECT COUNT(*) AS FastestLaps FROM Laps WHERE driver_name = 'Lewis Hamilton' AND is_personal_best = 1\\n```\", name='sql_db_query_checker', id='ee92b264-a20a-4276-b93d-0eb88da8d647', tool_call_id='call_aYYhIdyMvFlatT5NmwZFz9a4')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_Tp5C2y8jZLjLAZoH2qTmwrTH', 'function': {'arguments': '{\"query\":\"SELECT COUNT(*) AS FastestLaps FROM Laps WHERE driver_name = \\'Lewis Hamilton\\' AND is_personal_best = 1\"}', 'name': 'sql_db_query'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 41, 'prompt_tokens': 1598, 'total_tokens': 1639, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 1408}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-e7a82348-e88d-466e-9e6d-7e2ef17297de-0', tool_calls=[{'name': 'sql_db_query', 'args': {'query': \"SELECT COUNT(*) AS FastestLaps FROM Laps WHERE driver_name = 'Lewis Hamilton' AND is_personal_best = 1\"}, 'id': 'call_Tp5C2y8jZLjLAZoH2qTmwrTH', 'type': 'tool_call'}], usage_metadata={'input_tokens': 1598, 'output_tokens': 41, 'total_tokens': 1639, 'input_token_details': {'cache_read': 1408}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n",
      "{'tools': {'messages': [ToolMessage(content='[(0,)]', name='sql_db_query', id='3748a9fc-1225-43ab-8c7d-a2ff46a37c98', tool_call_id='call_Tp5C2y8jZLjLAZoH2qTmwrTH')]}}\n",
      "----\n",
      "{'agent': {'messages': [AIMessage(content='Lewis Hamilton has achieved 0 fastest laps.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 10, 'prompt_tokens': 1652, 'total_tokens': 1662, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 1536}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'stop', 'logprobs': None}, id='run-91be8556-79ff-4d3f-8128-d912462b21f0-0', usage_metadata={'input_tokens': 1652, 'output_tokens': 10, 'total_tokens': 1662, 'input_token_details': {'cache_read': 1536}, 'output_token_details': {'reasoning': 0}})]}}\n",
      "----\n"
     ]
    }
   ],
   "source": [
    "for s in agent.stream(\n",
    "    {\"messages\": [HumanMessage(\n",
    "        content=\"How many fastest laps did Lewis Hamilton?\")]}\n",
    "):\n",
    "    print(s)\n",
    "    print(\"----\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to multipart ingest runs: langsmith.utils.LangSmithError: Failed to POST https://api.smith.langchain.com/runs/multipart in LangSmith API. HTTPError('400 Client Error: Bad Request for url: https://api.smith.langchain.com/runs/multipart', '{\"detail\":\"Empty request\"}')\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='How many fastest laps did Verstappen achieve?', additional_kwargs={}, response_metadata={}, id='8ccf219e-5657-4b54-8ef1-a0652669dc5a'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_O6knt1bzCUHBEOi3kBB819FI', 'function': {'arguments': '{\"query\":\"Verstappen\"}', 'name': 'search_proper_nouns'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 19, 'prompt_tokens': 650, 'total_tokens': 669, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-2d3a3fb6-a33d-4eb0-82a6-fbb702a61bd4-0', tool_calls=[{'name': 'search_proper_nouns', 'args': {'query': 'Verstappen'}, 'id': 'call_O6knt1bzCUHBEOi3kBB819FI', 'type': 'tool_call'}], usage_metadata={'input_tokens': 650, 'output_tokens': 19, 'total_tokens': 669, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}}),\n",
       "  ToolMessage(content='Max Verstappen\\n\\nNico Hulkenberg\\n\\nLewis Hamilton\\n\\nLance Stroll\\n\\nValtteri Bottas', name='search_proper_nouns', id='94669590-6e71-48d6-b2e1-674593107eac', tool_call_id='call_O6knt1bzCUHBEOi3kBB819FI'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_xs6oJjmQm0nrkEfsHUGlq0l9', 'function': {'arguments': '{\"query\":\"SELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = \\'Max Verstappen\\' AND is_fastest_lap = 1;\"}', 'name': 'sql_db_query_checker'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 45, 'prompt_tokens': 703, 'total_tokens': 748, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-4ef76281-738f-478a-a41a-af9710e961be-0', tool_calls=[{'name': 'sql_db_query_checker', 'args': {'query': \"SELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = 'Max Verstappen' AND is_fastest_lap = 1;\"}, 'id': 'call_xs6oJjmQm0nrkEfsHUGlq0l9', 'type': 'tool_call'}], usage_metadata={'input_tokens': 703, 'output_tokens': 45, 'total_tokens': 748, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}}),\n",
       "  ToolMessage(content=\"```sql\\nSELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = 'Max Verstappen' AND is_fastest_lap = 1;\\n```\", name='sql_db_query_checker', id='4b1ac6b7-c831-4e13-9459-a5a8d6218f04', tool_call_id='call_xs6oJjmQm0nrkEfsHUGlq0l9'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_958H9JRomjsCZ3UWWM80VBQR', 'function': {'arguments': '{\"query\":\"SELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = \\'Max Verstappen\\' AND is_fastest_lap = 1;\"}', 'name': 'sql_db_query'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 44, 'prompt_tokens': 792, 'total_tokens': 836, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-0b3546dc-01be-4062-b8b1-086d9e3e26bd-0', tool_calls=[{'name': 'sql_db_query', 'args': {'query': \"SELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = 'Max Verstappen' AND is_fastest_lap = 1;\"}, 'id': 'call_958H9JRomjsCZ3UWWM80VBQR', 'type': 'tool_call'}], usage_metadata={'input_tokens': 792, 'output_tokens': 44, 'total_tokens': 836, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}}),\n",
       "  ToolMessage(content=\"Error: (sqlite3.OperationalError) no such column: is_fastest_lap\\n[SQL: SELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = 'Max Verstappen' AND is_fastest_lap = 1;]\\n(Background on this error at: https://sqlalche.me/e/20/e3q8)\", name='sql_db_query', id='54cb39eb-bac0-4a43-b991-b2b422817e7a', tool_call_id='call_958H9JRomjsCZ3UWWM80VBQR'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_L2G2iHs2GkXm9xpbRjsXFt4N', 'function': {'arguments': '{\"table_names\":\"Laps\"}', 'name': 'sql_db_schema'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 17, 'prompt_tokens': 919, 'total_tokens': 936, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-c802e2d9-a552-4521-8859-44364fa9ef95-0', tool_calls=[{'name': 'sql_db_schema', 'args': {'table_names': 'Laps'}, 'id': 'call_L2G2iHs2GkXm9xpbRjsXFt4N', 'type': 'tool_call'}], usage_metadata={'input_tokens': 919, 'output_tokens': 17, 'total_tokens': 936, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}}),\n",
       "  ToolMessage(content='\\nCREATE TABLE \"Laps\" (\\n\\tlap_id INTEGER, \\n\\tsession_id INTEGER, \\n\\tdriver_name TEXT NOT NULL, \\n\\tlap_number INTEGER NOT NULL, \\n\\tstint INTEGER, \\n\\tsector_1_speed_trap_in_km REAL, \\n\\tsector_2_speed_trap_in_km REAL, \\n\\tfinish_line_speed_trap_in_km REAL, \\n\\tlongest_strait_speed_trap_in_km REAL, \\n\\tis_personal_best BOOLEAN, \\n\\ttyre_compound TEXT, \\n\\ttyre_life_in_laps INTEGER, \\n\\tis_fresh_tyre BOOLEAN, \\n\\tposition INTEGER, \\n\\tlap_time_in_seconds REAL, \\n\\tsector_1_time_in_seconds REAL, \\n\\tsector_2_time_in_seconds REAL, \\n\\tsector_3_time_in_seconds REAL, \\n\\tlap_start_time_in_datetime DATETIME, \\n\\tpin_in_time_in_datetime DATETIME, \\n\\tpin_out_time_in_datetime DATETIME, \\n\\tPRIMARY KEY (lap_id), \\n\\tFOREIGN KEY(session_id) REFERENCES \"Sessions\" (session_id), \\n\\tUNIQUE (session_id, driver_name, lap_number)\\n)\\n\\n/*\\n3 rows from Laps table:\\nlap_id\\tsession_id\\tdriver_name\\tlap_number\\tstint\\tsector_1_speed_trap_in_km\\tsector_2_speed_trap_in_km\\tfinish_line_speed_trap_in_km\\tlongest_strait_speed_trap_in_km\\tis_personal_best\\ttyre_compound\\ttyre_life_in_laps\\tis_fresh_tyre\\tposition\\tlap_time_in_seconds\\tsector_1_time_in_seconds\\tsector_2_time_in_seconds\\tsector_3_time_in_seconds\\tlap_start_time_in_datetime\\tpin_in_time_in_datetime\\tpin_out_time_in_datetime\\n1\\t1\\tVER\\t1\\tNone\\tNone\\tNone\\tNone\\t178.0\\tFalse\\tSOFT\\t1\\tTrue\\tNone\\tNone\\tNone\\tNone\\tNone\\t2023-03-04 15:04:00.840000\\t2023-03-04 15:06:07.931000\\t2023-03-04 15:04:00.840000\\n2\\t1\\tVER\\t2\\tNone\\tNone\\tNone\\t288.0\\t141.0\\tFalse\\tSOFT\\t2\\tFalse\\tNone\\tNone\\tNone\\t53.666\\t38.509\\t2023-03-04 15:12:21.456000\\tNone\\t2023-03-04 15:12:21.456000\\n3\\t1\\tVER\\t3\\tNone\\tNone\\tNone\\t288.0\\t322.0\\tTrue\\tSOFT\\t3\\tFalse\\tNone\\t91.295\\t29.152\\t39.195\\t22.948\\t2023-03-04 15:14:33.391000\\tNone\\tNone\\n*/', name='sql_db_schema', id='3db63be6-4e5e-4994-817e-414d087edf65', tool_call_id='call_L2G2iHs2GkXm9xpbRjsXFt4N'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_EIQeVdEYOGNiGCOMEnjzCfAn', 'function': {'arguments': '{\"query\":\"SELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = \\'VER\\' AND is_personal_best = 1;\"}', 'name': 'sql_db_query_checker'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 42, 'prompt_tokens': 1541, 'total_tokens': 1583, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-c912c195-8503-4a0d-9879-4d970aec1170-0', tool_calls=[{'name': 'sql_db_query_checker', 'args': {'query': \"SELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = 'VER' AND is_personal_best = 1;\"}, 'id': 'call_EIQeVdEYOGNiGCOMEnjzCfAn', 'type': 'tool_call'}], usage_metadata={'input_tokens': 1541, 'output_tokens': 42, 'total_tokens': 1583, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}}),\n",
       "  ToolMessage(content=\"```sql\\nSELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = 'VER' AND is_personal_best = 1;\\n```\", name='sql_db_query_checker', id='cde57c67-a90c-43ef-bcc8-d4f6fbc2d96b', tool_call_id='call_EIQeVdEYOGNiGCOMEnjzCfAn'),\n",
       "  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_e52OzWNQWFn5ZcVz28Ii2TSG', 'function': {'arguments': '{\"query\":\"SELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = \\'VER\\' AND is_personal_best = 1;\"}', 'name': 'sql_db_query'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 41, 'prompt_tokens': 1624, 'total_tokens': 1665, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 1536}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-f9e102db-f36a-448d-a0e8-188c7d38f969-0', tool_calls=[{'name': 'sql_db_query', 'args': {'query': \"SELECT COUNT(*) as fastest_laps_count FROM Laps WHERE driver_name = 'VER' AND is_personal_best = 1;\"}, 'id': 'call_e52OzWNQWFn5ZcVz28Ii2TSG', 'type': 'tool_call'}], usage_metadata={'input_tokens': 1624, 'output_tokens': 41, 'total_tokens': 1665, 'input_token_details': {'cache_read': 1536}, 'output_token_details': {'reasoning': 0}}),\n",
       "  ToolMessage(content='[(24,)]', name='sql_db_query', id='470aa971-0d39-4026-9e8e-e4dc5d219051', tool_call_id='call_e52OzWNQWFn5ZcVz28Ii2TSG'),\n",
       "  AIMessage(content='Max Verstappen achieved a total of 24 fastest laps.', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 13, 'prompt_tokens': 1678, 'total_tokens': 1691, 'completion_tokens_details': {'audio_tokens': None, 'reasoning_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 1536}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_482c22a7bc', 'finish_reason': 'stop', 'logprobs': None}, id='run-99c62bec-1d38-4ea2-ba44-67ffdc425335-0', usage_metadata={'input_tokens': 1678, 'output_tokens': 13, 'total_tokens': 1691, 'input_token_details': {'cache_read': 1536}, 'output_token_details': {'reasoning': 0}})]}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.invoke({\"messages\": [HumanMessage(\n",
    "    content=\"How many fastest laps did Verstappen achieve?\")]})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}