Spaces:
Sleeping
Sleeping
File size: 3,886 Bytes
d8a070a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | {
"cells": [
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"True"
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import os\n",
"from dotenv import load_dotenv\n",
"from pyprojroot import here\n",
"from langchain.chains import create_sql_query_chain\n",
"from langchain_community.agent_toolkits import create_sql_agent\n",
"from langchain_openai import ChatOpenAI\n",
"from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit\n",
"from langchain_community.utilities import SQLDatabase\n",
"\n",
"load_dotenv()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Set the environment variable and load the LLM**"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"os.environ['OPENAI_API_KEY'] = os.getenv(\"OPEN_AI_API_KEY\")\n",
"\n",
"\n",
"llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
"# llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
"# llm = ChatOpenAI(model=\"gpt-4o\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Load and test the sqlite db**"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"sqlite\n",
"['aircrafts_data', 'airports_data', 'boarding_passes', 'bookings', 'car_rentals', 'flights', 'hotels', 'seats', 'ticket_flights', 'tickets', 'trip_recommendations']\n"
]
},
{
"data": {
"text/plain": [
"\"[('773', 'Boeing 777-300', 11100), ('763', 'Boeing 767-300', 7900), ('SU9', 'Sukhoi Superjet-100', 3000), ('320', 'Airbus A320-200', 5700), ('321', 'Airbus A321-200', 5600), ('319', 'Airbus A319-100', 6700), ('733', 'Boeing 737-300', 4200), ('CN1', 'Cessna 208 Caravan', 1200), ('CR2', 'Bombardier CRJ-200', 2700)]\""
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sqldb_directory = here(\"data/travel.sqlite\")\n",
"db = SQLDatabase.from_uri(f\"sqlite:///{sqldb_directory}\")\n",
"print(db.dialect)\n",
"print(db.get_usable_table_names())\n",
"db.run(\"SELECT * FROM aircrafts_data LIMIT 10;\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Create the SQL agent and run a test query**"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'SELECT COUNT(*) AS total_rows FROM aircrafts_data;'"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"chain = create_sql_query_chain(llm, db)\n",
"response = chain.invoke({\"question\": \"How many rows are there in the aircrafts_data table?\"})\n",
"response"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'[(9,)]'"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"db.run(response)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "rag-sqlagent",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
|