This commit is contained in:
Ankush Gola
2024-06-12 14:29:14 -07:00
parent 8ca641cda6
commit 3ae9ca9184
+163 -162
View File
@@ -2,6 +2,10 @@
"cells": [
{
"cell_type": "markdown",
"id": "9f853e403eabd4f8",
"metadata": {
"collapsed": false
},
"source": [
"# An agent for interacting with a SQL database\n",
"\n",
@@ -20,27 +24,31 @@
"The end-to-end workflow will look something like below:\n",
"\n",
"![](sql-agent-diagram.png)"
],
"metadata": {
"collapsed": false
},
"id": "9f853e403eabd4f8"
]
},
{
"cell_type": "markdown",
"id": "b5a87813ffe7e4d2",
"metadata": {
"collapsed": false
},
"source": [
"## Set up environment\n",
"\n",
"We'll set up our environment variables for OpenAI, and optionally, to enable tracing with [LangSmith](https://smith.langchain.com)."
],
"metadata": {
"collapsed": false
},
"id": "b5a87813ffe7e4d2"
]
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": 1,
"id": "6c05a600f1afb5b6",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:00.532147Z",
"start_time": "2024-06-12T21:24:00.526043Z"
}
},
"outputs": [],
"source": [
"import os\n",
@@ -48,18 +56,14 @@
"os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
"os.environ[\"LANGSMITH_API_KEY\"] = \"lsv2_pt_...\"\n",
"os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\""
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:18:08.635369Z",
"start_time": "2024-06-12T20:18:08.630616Z"
}
},
"id": "6c05a600f1afb5b6"
]
},
{
"cell_type": "markdown",
"id": "877d8c85825089d8",
"metadata": {
"collapsed": false
},
"source": [
"## Configure the database\n",
"\n",
@@ -67,15 +71,19 @@
"Find more information about the database [here](https://www.sqlitetutorial.net/sqlite-sample-database/).\n",
"\n",
"For convenience, we have hosted the database (`Chinook.db`) on a public GCS bucket."
],
"metadata": {
"collapsed": false
},
"id": "877d8c85825089d8"
]
},
{
"cell_type": "code",
"execution_count": 19,
"execution_count": 2,
"id": "64b0bf1b14c2e902",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:09.918436Z",
"start_time": "2024-06-12T21:24:09.608563Z"
}
},
"outputs": [
{
"name": "stdout",
@@ -100,46 +108,46 @@
" print(\"File downloaded and saved as Chinook.db\")\n",
"else:\n",
" print(f\"Failed to download the file. Status code: {response.status_code}\")"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:26:06.460618Z",
"start_time": "2024-06-12T20:26:06.025590Z"
}
},
"id": "64b0bf1b14c2e902"
]
},
{
"cell_type": "markdown",
"source": [
"We will use a handy SQL database wrapper available in the `langchain_community` package to interact with the database. The wrapper provides a simple interface to execute SQL queries and fetch results. We will also use the `langchain_openai` package to interact with the OpenAI API for language models later in the tutorial."
],
"id": "61c8304aa5ceb6a5",
"metadata": {
"collapsed": false
},
"id": "61c8304aa5ceb6a5"
"source": [
"We will use a handy SQL database wrapper available in the `langchain_community` package to interact with the database. The wrapper provides a simple interface to execute SQL queries and fetch results. We will also use the `langchain_openai` package to interact with the OpenAI API for language models later in the tutorial."
]
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 3,
"id": "a60191bd3489f278",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:14.663745Z",
"start_time": "2024-06-12T21:24:13.527958Z"
}
},
"outputs": [],
"source": [
"%%capture --no-stderr --no-display\n",
"!pip install langchain_community langchain_openai"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:18:12.265718Z",
"start_time": "2024-06-12T20:18:11.265328Z"
}
},
"id": "a60191bd3489f278"
]
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": 4,
"id": "1f1e1f4f86ed54",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:15.891582Z",
"start_time": "2024-06-12T21:24:15.289782Z"
}
},
"outputs": [
{
"name": "stdout",
@@ -153,7 +161,7 @@
"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": 5,
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
@@ -165,39 +173,38 @@
"print(db.dialect)\n",
"print(db.get_usable_table_names())\n",
"db.run(\"SELECT * FROM Artist LIMIT 10;\")"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:18:12.902435Z",
"start_time": "2024-06-12T20:18:12.796315Z"
}
},
"id": "1f1e1f4f86ed54"
]
},
{
"cell_type": "markdown",
"id": "6959e93141d8099c",
"metadata": {
"collapsed": false
},
"source": [
"## Utility functions\n",
"\n",
"We will define a few utility functions to help us with the agent implementation. Specifically, we will wrap a `ToolNode` with a fallback to handle errors and surface them to the agent."
],
"metadata": {
"collapsed": false
},
"id": "6959e93141d8099c"
]
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": 5,
"id": "deae8460e4cf72b1",
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T21:24:17.557848Z",
"start_time": "2024-06-12T21:24:17.508550Z"
}
},
"outputs": [],
"source": [
"from typing import Any, Dict, List\n",
"from typing import Any\n",
"\n",
"from langchain_core.messages import ToolMessage\n",
"from langchain_core.runnables import RunnableLambda, RunnableWithFallbacks\n",
"from langchain_core.messages import ToolMessage, AIMessage\n",
"\n",
"from langgraph.prebuilt import ToolNode\n",
"\n",
"def create_tool_node_with_fallback(tools: list) -> RunnableWithFallbacks[Any, dict]:\n",
" \"\"\"\n",
@@ -219,18 +226,14 @@
" for tc in tool_calls\n",
" ]\n",
" }"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:18:14.197218Z",
"start_time": "2024-06-12T20:18:14.194149Z"
}
},
"id": "deae8460e4cf72b1"
]
},
{
"cell_type": "markdown",
"id": "d0196604f8cbb07b",
"metadata": {
"collapsed": false
},
"source": [
"## Define tools for the agent\n",
"\n",
@@ -241,15 +244,19 @@
"3. `db_query_tool`: Execute the query and fetch the results OR return an error message if the query fails\n",
"\n",
"For the first two tools, we will grab them from the `SQLDatabaseToolkit`, also available in the `langchain_community` package."
],
"metadata": {
"collapsed": false
},
"id": "d0196604f8cbb07b"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "452d049a3d2a4406",
"metadata": {
"ExecuteTime": {
"end_time": "2024-06-12T20:18:15.838940Z",
"start_time": "2024-06-12T20:18:15.734199Z"
},
"collapsed": false
},
"outputs": [
{
"name": "stdout",
@@ -286,29 +293,29 @@
"print(list_tables_tool.invoke(\"\"))\n",
"\n",
"print(get_schema_tool.invoke(\"Artist\"))"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:18:15.838940Z",
"start_time": "2024-06-12T20:18:15.734199Z"
}
},
"id": "452d049a3d2a4406"
]
},
{
"cell_type": "markdown",
"source": [
"The third will be defined manually. For the `db_query_tool`, we will execute the query against the database and return the results."
],
"id": "c16359edada327fa",
"metadata": {
"collapsed": false
},
"id": "c16359edada327fa"
"source": [
"The third will be defined manually. For the `db_query_tool`, we will execute the query against the database and return the results."
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "f7eb708ecb4c7cfc",
"metadata": {
"ExecuteTime": {
"end_time": "2024-06-12T20:39:35.759834Z",
"start_time": "2024-06-12T20:39:35.740255Z"
},
"collapsed": false
},
"outputs": [
{
"name": "stdout",
@@ -321,6 +328,7 @@
"source": [
"from langchain.agents import tool\n",
"\n",
"\n",
"@tool\n",
"def db_query_tool(query: str) -> str:\n",
" \"\"\"\n",
@@ -334,29 +342,29 @@
" return result\n",
"\n",
"print(db_query_tool.invoke(\"SELECT * FROM Artist LIMIT 10;\"))"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:39:35.759834Z",
"start_time": "2024-06-12T20:39:35.740255Z"
}
},
"id": "f7eb708ecb4c7cfc"
]
},
{
"cell_type": "markdown",
"source": [
"While not strictly a tool, we will prompt an LLM to check for common mistakes in the query and later add this as a node in the workflow."
],
"id": "f1d66db8b8621639",
"metadata": {
"collapsed": false
},
"id": "f1d66db8b8621639"
"source": [
"While not strictly a tool, we will prompt an LLM to check for common mistakes in the query and later add this as a node in the workflow."
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "293017e8f05ac2b3",
"metadata": {
"ExecuteTime": {
"end_time": "2024-06-12T20:18:19.658322Z",
"start_time": "2024-06-12T20:18:18.756256Z"
},
"collapsed": false
},
"outputs": [
{
"data": {
@@ -389,43 +397,44 @@
"query_check = query_check_prompt | ChatOpenAI(model=\"gpt-4o\", temperature=0).bind_tools([db_query_tool], tool_choice=\"required\")\n",
"\n",
"query_check.invoke({\"messages\": [(\"user\", \"SELET * FROM Artist LIMIT 10;\")]})"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:18:19.658322Z",
"start_time": "2024-06-12T20:18:18.756256Z"
}
},
"id": "293017e8f05ac2b3"
]
},
{
"cell_type": "markdown",
"id": "66f88452151e8188",
"metadata": {
"collapsed": false
},
"source": [
"## Define the workflow\n",
"\n",
"We will then define the workflow for the agent. The agent will first force-call the `list_tables_tool` to fetch the available tables from the database, then follow the steps mentioned at the beginning of the tutorial."
],
"metadata": {
"collapsed": false
},
"id": "66f88452151e8188"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "90d04ceea7b6b010",
"metadata": {
"ExecuteTime": {
"end_time": "2024-06-12T20:21:09.799829Z",
"start_time": "2024-06-12T20:21:09.765928Z"
},
"collapsed": false
},
"outputs": [],
"source": [
"from typing import Annotated, Literal\n",
"from typing_extensions import TypedDict\n",
"\n",
"from langchain_core.messages import AIMessage\n",
"from langchain_core.pydantic_v1 import BaseModel, Field\n",
"from langchain_openai import ChatOpenAI\n",
"from typing_extensions import TypedDict\n",
"\n",
"from langgraph.graph import END, StateGraph\n",
"from langgraph.graph.message import AnyMessage, add_messages\n",
"from langgraph.prebuilt.tool_node import ToolNode\n",
"from langchain_core.messages import AIMessage\n",
"from langchain_core.pydantic_v1 import BaseModel, Field\n",
"\n",
"\n",
"# Define the state for the agent\n",
"class State(TypedDict):\n",
@@ -555,29 +564,29 @@
"\n",
"# Compile the workflow into a runnable\n",
"app = workflow.compile()"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:21:09.799829Z",
"start_time": "2024-06-12T20:21:09.765928Z"
}
},
"id": "90d04ceea7b6b010"
]
},
{
"cell_type": "markdown",
"source": [
"## Visualize the graph"
],
"id": "6c344ae086ba8d22",
"metadata": {
"collapsed": false
},
"id": "6c344ae086ba8d22"
"source": [
"## Visualize the graph"
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "4f200d1813897000",
"metadata": {
"ExecuteTime": {
"end_time": "2024-06-12T20:21:11.813905Z",
"start_time": "2024-06-12T20:21:11.712945Z"
},
"collapsed": false
},
"outputs": [
{
"data": {
@@ -589,8 +598,8 @@
}
],
"source": [
"from IPython.display import Image, display\n",
"from langchain_core.runnables.graph import MermaidDrawMethod\n",
"from IPython.display import display, Image\n",
"\n",
"display(\n",
" Image(\n",
@@ -599,29 +608,29 @@
" )\n",
" )\n",
")"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:21:11.813905Z",
"start_time": "2024-06-12T20:21:11.712945Z"
}
},
"id": "4f200d1813897000"
]
},
{
"cell_type": "markdown",
"source": [
"## Run the agent"
],
"id": "bdf78dc68548522c",
"metadata": {
"collapsed": false
},
"id": "bdf78dc68548522c"
"source": [
"## Run the agent"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "956883cced0b8ec",
"metadata": {
"ExecuteTime": {
"end_time": "2024-06-12T20:21:21.878352Z",
"start_time": "2024-06-12T20:21:12.854570Z"
},
"collapsed": false
},
"outputs": [
{
"name": "stdout",
@@ -642,15 +651,7 @@
"source": [
"for event in app.stream({\"messages\": [(\"user\", \"Which sales agent made the most in sales in 2009?\")]}):\n",
" print(event)"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-06-12T20:21:21.878352Z",
"start_time": "2024-06-12T20:21:12.854570Z"
}
},
"id": "956883cced0b8ec"
]
}
],
"metadata": {