search/relational-data-querying-llm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e80822d5-60ed-46ca-a940-65fb8322bb68",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Relational Data Querying Using LLMs\n",
    "\n",
    "This notebook demonstrates how relational data can be queried in natural language using LLMs. \n",
    "\n",
    "### Use Case\n",
    "We have a relational database with multiple tables. We want to ask questions in natural language, get them translated into SQL, and then executed. The final result is a relational table.  \n",
    "\n",
    "### Prototype: Approach and Data\n",
    "The solution approach is to fetch the dataschema first, and then provide it as a contex for generating SQL based on a natural language query. We implement this using a standard SQL querying chain provided by Langchain. \n",
    "\n",
    "We use the `classicmodel` database [1]. The database contains a set of tables which is typical for order management systems. We use only two tables: orders and order details (line items). We assume an external MySQL instance with `classicmodel` database loaded into it.\n",
    "\n",
    "### Usage and Productization\n",
    "This is a basic prototype that demonstrates the concept. Productization would typically require to pre-process the schema to select only the relevant tables, handle SQL execution errors, and post-process the results to elect the optimal representation such as a line plot, bar chart, or textual summary.\n",
    "\n",
    "### References:\n",
    "1. https://www.mysqltutorial.org/mysql-sample-database.aspx (copy available in `tensor-house-data/search/orders-database`)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c435a99e-0b2c-4c9d-9fef-530b9fd77055",
   "metadata": {},
   "source": [
    "## Environment Setup and Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c81a3815-c2bb-45e9-8b20-218ef9ec41a9",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2024-01-14T14:36:51.714154Z",
     "start_time": "2024-01-14T14:36:51.664573Z"
    }
   },
   "outputs": [],
   "source": [
    "#\n",
    "# Imports\n",
    "#\n",
    "from langchain_community.llms.vertexai import VertexAI\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "from langchain_experimental.sql import SQLDatabaseChain\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "#\n",
    "# Initialize LLM provider\n",
    "# (google-cloud-aiplatform must be installed)\n",
    "#\n",
    "from google.cloud import aiplatform\n",
    "aiplatform.init(\n",
    "    project='<< specify your project name here >>',\n",
    "    location='us-central1'\n",
    ")\n",
    "\n",
    "#\n",
    "# Initialize the DB connection\n",
    "#\n",
    "db = SQLDatabase.from_uri(\"mysql+pymysql://root:@localhost/classicmodels\", include_tables=['orders', 'orderdetails'], sample_rows_in_table_info=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8af0cfa5-5976-4b1e-a669-c455845014a1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-14T14:36:53.484620Z",
     "start_time": "2024-01-14T14:36:52.559446Z"
    }
   },
   "outputs": [],
   "source": [
    "_DEFAULT_TEMPLATE = \"\"\"Given an input question, create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.\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",
    "\n",
    "{table_info}\n",
    "\n",
    "Determine which tables need to be joined before generating the SQL query. Limit the number of rows in the SQL result by {top_k}. Do NOT add any formatting symbols or keywords.\n",
    "\n",
    "Question: {input}\"\"\"\n",
    "custom_prompt = PromptTemplate(\n",
    "    input_variables=[\"input\", \"table_info\", \"dialect\", \"top_k\"], template=_DEFAULT_TEMPLATE\n",
    ")\n",
    "\n",
    "llm = VertexAI(temperature=0, verbose=True)\n",
    "db_chain = SQLDatabaseChain.from_llm(llm, \n",
    "                                     db, \n",
    "                                     top_k=3,\n",
    "                                     verbose=True, \n",
    "                                     prompt=custom_prompt,\n",
    "                                     return_intermediate_steps=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "503605aa-47af-40d4-a594-f82428de9294",
   "metadata": {},
   "source": [
    "## Example Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e4f2432c-ccd0-42a1-9281-02aa8cf19c85",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-14T14:36:58.812892Z",
     "start_time": "2024-01-14T14:36:56.149939Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new SQLDatabaseChain chain...\u001B[0m\n",
      "How many orders are there?\n",
      "SQLQuery:\u001B[32;1m\u001B[1;3mQuestion: How many orders are there?\n",
      "SQLQuery: SELECT COUNT(*) AS order_count FROM orders LIMIT 3;\u001B[0m\n",
      "SQLResult: \u001B[33;1m\u001B[1;3m[(326,)]\u001B[0m\n",
      "Answer:\u001B[32;1m\u001B[1;3mQuestion: How many orders are there?\n",
      "SQLQuery: SELECT COUNT(*) AS order_count FROM orders;\u001B[0m\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "#\n",
    "# Convert the input query into SQL and execute it - basic example\n",
    "#\n",
    "result = db_chain(\"How many orders are there?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "95195869-f133-4aef-8c10-b40b2ea8e1da",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-14T14:35:06.999011Z",
     "start_time": "2024-01-14T14:35:02.968935Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001B[1m> Entering new SQLDatabaseChain chain...\u001B[0m\n",
      "What was the revenue for each date?\n",
      "SQLQuery:\u001B[32;1m\u001B[1;3mQuestion: What was the revenue for each date?\n",
      "SQLQuery: SELECT orderDate, SUM(quantityOrdered * priceEach) AS revenue\n",
      "FROM orderdetails\n",
      "JOIN orders ON orderdetails.orderNumber = orders.orderNumber\n",
      "GROUP BY orderDate\n",
      "LIMIT 3;\u001B[0m\n",
      "SQLResult: \u001B[33;1m\u001B[1;3m[(datetime.date(2003, 1, 6), Decimal('10223.83')), (datetime.date(2003, 1, 9), Decimal('10549.01')), (datetime.date(2003, 1, 10), Decimal('5494.78'))]\u001B[0m\n",
      "Answer:\u001B[32;1m\u001B[1;3mQuestion: What was the revenue for each date?\n",
      "SQLQuery:SELECT orderDate, SUM(quantityOrdered * priceEach) AS revenue\n",
      "FROM orderdetails\n",
      "JOIN orders ON orderdetails.orderNumber = orders.orderNumber\n",
      "GROUP BY orderDate\n",
      "LIMIT 3;\u001B[0m\n",
      "\u001B[1m> Finished chain.\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "#\n",
    "# Convert the input query into SQL and execute it - more complex examples\n",
    "#\n",
    "result = db_chain(\"What was the revenue for each date?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "344ab602-75d4-42c7-b3c8-cbfdd53ec7e2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-14T14:35:18.259748Z",
     "start_time": "2024-01-14T14:35:18.246183Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "CREATE TABLE orderdetails (\n",
      "\t`orderNumber` INTEGER NOT NULL, \n",
      "\t`productCode` VARCHAR(15) NOT NULL, \n",
      "\t`quantityOrdered` INTEGER NOT NULL, \n",
      "\t`priceEach` DECIMAL(10, 2) NOT NULL, \n",
      "\t`orderLineNumber` SMALLINT NOT NULL, \n",
      "\tPRIMARY KEY (`orderNumber`, `productCode`), \n",
      "\tCONSTRAINT orderdetails_ibfk_1 FOREIGN KEY(`orderNumber`) REFERENCES orders (`orderNumber`), \n",
      "\tCONSTRAINT orderdetails_ibfk_2 FOREIGN KEY(`productCode`) REFERENCES products (`productCode`)\n",
      ")ENGINE=InnoDB DEFAULT CHARSET=latin1\n",
      "\n",
      "\n",
      "CREATE TABLE orders (\n",
      "\t`orderNumber` INTEGER NOT NULL, \n",
      "\t`orderDate` DATE NOT NULL, \n",
      "\t`requiredDate` DATE NOT NULL, \n",
      "\t`shippedDate` DATE, \n",
      "\tstatus VARCHAR(15) NOT NULL, \n",
      "\tcomments TEXT, \n",
      "\t`customerNumber` INTEGER NOT NULL, \n",
      "\tPRIMARY KEY (`orderNumber`), \n",
      "\tCONSTRAINT orders_ibfk_1 FOREIGN KEY(`customerNumber`) REFERENCES customers (`customerNumber`)\n",
      ")ENGINE=InnoDB DEFAULT CHARSET=latin1\n"
     ]
    }
   ],
   "source": [
    "#\n",
    "# Print the table info used in the query\n",
    "#\n",
    "print(result[\"intermediate_steps\"][0]['table_info'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "74a54b8d9bf3673e"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "py310"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}