Youtube_NL2SQL_llama_Index.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPlxEAJuBIziMBf/yvMhjFO",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/PradipNichite/Youtube-Tutorials/blob/main/Youtube_NL2SQL_llama_Index.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "###Realted Videos"
      ],
      "metadata": {
        "id": "H9IAFY6Gohar"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "LangChain, SQL Agents & OpenAI LLMs: Query Database Using Natural Language https://youtu.be/VG9KYCS0-8E?si=EAqu-9DS6QX9eitR\n",
        "\n",
        "Mastering LlamaIndex : Create, Save & Load Indexes, Customize LLMs, Prompts & Embeddings: https://youtu.be/XGBQ_f-Yy48?si=ZWs0meSlQpLPcCE0"
      ],
      "metadata": {
        "id": "tbhVKp3YoUPY"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "xcNg74LCPwUd"
      },
      "outputs": [],
      "source": [
        "!pip install llama-index pymysql -q"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "WyxrjYaooYvb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import logging\n",
        "import sys\n",
        "\n",
        "logging.basicConfig(stream=sys.stdout, level=logging.INFO, force=True)\n",
        "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
        "from IPython.display import Markdown, display"
      ],
      "metadata": {
        "id": "97SFvNIQTbML"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "https://www.mysqltutorial.org/mysql-sample-database.aspx"
      ],
      "metadata": {
        "id": "68GQHrv7f6pI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "db_user = \"admin\"\n",
        "db_password = \"llama_nl2sql\"\n",
        "db_host = \"langchainsql.cl0j8hicdoox.us-east-1.rds.amazonaws.com\"\n",
        "db_name = \"classicmodels\""
      ],
      "metadata": {
        "id": "ZVQ5c1AgP3Zs"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from sqlalchemy import create_engine, text\n",
        "\n",
        "# Construct the connection string\n",
        "connection_string = f\"mysql+pymysql://{db_user}:{db_password}@{db_host}/{db_name}\"\n",
        "\n",
        "# Create an engine instance\n",
        "engine = create_engine(connection_string)\n",
        "\n",
        "# Test the connection using raw SQL\n",
        "with engine.connect() as connection:\n",
        "    result = connection.execute(text(\"select * from customers limit 3\"))\n",
        "    for row in result:\n",
        "        print(row)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XZm6fC_7P366",
        "outputId": "7da302bf-2b9c-4682-b343-9fbb8408bb80"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(103, 'Atelier graphique', 'Schmitt', 'Carine ', '40.32.2555', '54, rue Royale', None, 'Nantes', None, '44000', 'France', 1370, Decimal('21000.00'))\n",
            "(112, 'Signal Gift Stores', 'King', 'Jean', '7025551838', '8489 Strong St.', None, 'Las Vegas', 'NV', '83030', 'USA', 1166, Decimal('71800.00'))\n",
            "(114, 'Australian Collectors, Co.', 'Ferguson', 'Peter', '03 9520 4555', '636 St Kilda Road', 'Level 3', 'Melbourne', 'Victoria', '3004', 'Australia', 1611, Decimal('117300.00'))\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "table_details = {\n",
        "    \"customers\": \"stores customer’s data.\",\n",
        "    \"products\": \"stores a list of scale model cars.\",\n",
        "    \"productlines\": \"stores a list of product line categories.\",\n",
        "    \"orders\": \"stores sales orders placed by customers.\",\n",
        "    \"orderdetails\": \"stores sales order line items for each sales order.\",\n",
        "    \"payments\": \"stores payments made by customers based on their accounts.\",\n",
        "    \"employees\": \"stores all employee information as well as the organization structure such as who reports to whom.\",\n",
        "    \"offices\": \"stores sales office data.\"\n",
        "}\n"
      ],
      "metadata": {
        "id": "IuLV4gB3ZmDH"
      },
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "https://gpt-index.readthedocs.io/en/latest/api_reference/struct_store.html#llama_index.utilities.sql_wrapper.SQLDatabase"
      ],
      "metadata": {
        "id": "b7zMnMc-ZHLn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import SQLDatabase\n",
        "tables = [\"customers\",\"orders\"]\n",
        "# sql_database = SQLDatabase(engine, include_tables=tables,sample_rows_in_table_info=5)\n",
        "sql_database = SQLDatabase(engine, sample_rows_in_table_info=2)\n",
        "sql_database"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IklEeNaTP6Y4",
        "outputId": "2928475f-f685-40c2-a6ce-29df81d1c1bf"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:numexpr.utils:NumExpr defaulting to 2 threads.\n",
            "NumExpr defaulting to 2 threads.\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<llama_index.langchain_helpers.sql_wrapper.SQLDatabase at 0x7adf90cebb80>"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "list(sql_database._all_tables)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VaqVhL1HlzEf",
        "outputId": "78426bba-45ef-4bd6-abcf-e0901d221a74"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['payments',\n",
              " 'employees',\n",
              " 'customers',\n",
              " 'orders',\n",
              " 'productlines',\n",
              " 'products',\n",
              " 'orderdetails',\n",
              " 'offices']"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(sql_database.table_info)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QUPPAQP9QM38",
        "outputId": "ae976f72-707e-4dda-bbe0-18363508d02f"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "CREATE TABLE customers (\n",
            "\t`customerNumber` INTEGER NOT NULL, \n",
            "\t`customerName` VARCHAR(50) NOT NULL, \n",
            "\t`contactLastName` VARCHAR(50) NOT NULL, \n",
            "\t`contactFirstName` VARCHAR(50) NOT NULL, \n",
            "\tphone VARCHAR(50) NOT NULL, \n",
            "\t`addressLine1` VARCHAR(50) NOT NULL, \n",
            "\t`addressLine2` VARCHAR(50), \n",
            "\tcity VARCHAR(50) NOT NULL, \n",
            "\tstate VARCHAR(50), \n",
            "\t`postalCode` VARCHAR(15), \n",
            "\tcountry VARCHAR(50) NOT NULL, \n",
            "\t`salesRepEmployeeNumber` INTEGER, \n",
            "\t`creditLimit` DECIMAL(10, 2), \n",
            "\tPRIMARY KEY (`customerNumber`), \n",
            "\tCONSTRAINT customers_ibfk_1 FOREIGN KEY(`salesRepEmployeeNumber`) REFERENCES employees (`employeeNumber`)\n",
            ")DEFAULT CHARSET=latin1 ENGINE=InnoDB\n",
            "\n",
            "/*\n",
            "2 rows from customers table:\n",
            "customerNumber\tcustomerName\tcontactLastName\tcontactFirstName\tphone\taddressLine1\taddressLine2\tcity\tstate\tpostalCode\tcountry\tsalesRepEmployeeNumber\tcreditLimit\n",
            "103\tAtelier graphique\tSchmitt\tCarine \t40.32.2555\t54, rue Royale\tNone\tNantes\tNone\t44000\tFrance\t1370\t21000.00\n",
            "112\tSignal Gift Stores\tKing\tJean\t7025551838\t8489 Strong St.\tNone\tLas Vegas\tNV\t83030\tUSA\t1166\t71800.00\n",
            "*/\n",
            "\n",
            "\n",
            "CREATE TABLE employees (\n",
            "\t`employeeNumber` INTEGER NOT NULL, \n",
            "\t`lastName` VARCHAR(50) NOT NULL, \n",
            "\t`firstName` VARCHAR(50) NOT NULL, \n",
            "\textension VARCHAR(10) NOT NULL, \n",
            "\temail VARCHAR(100) NOT NULL, \n",
            "\t`officeCode` VARCHAR(10) NOT NULL, \n",
            "\t`reportsTo` INTEGER, \n",
            "\t`jobTitle` VARCHAR(50) NOT NULL, \n",
            "\tPRIMARY KEY (`employeeNumber`), \n",
            "\tCONSTRAINT employees_ibfk_1 FOREIGN KEY(`reportsTo`) REFERENCES employees (`employeeNumber`), \n",
            "\tCONSTRAINT employees_ibfk_2 FOREIGN KEY(`officeCode`) REFERENCES offices (`officeCode`)\n",
            ")DEFAULT CHARSET=latin1 ENGINE=InnoDB\n",
            "\n",
            "/*\n",
            "2 rows from employees table:\n",
            "employeeNumber\tlastName\tfirstName\textension\temail\tofficeCode\treportsTo\tjobTitle\n",
            "1002\tMurphy\tDiane\tx5800\tdmurphy@classicmodelcars.com\t1\tNone\tPresident\n",
            "1056\tPatterson\tMary\tx4611\tmpatterso@classicmodelcars.com\t1\t1002\tVP Sales\n",
            "*/\n",
            "\n",
            "\n",
            "CREATE TABLE offices (\n",
            "\t`officeCode` VARCHAR(10) NOT NULL, \n",
            "\tcity VARCHAR(50) NOT NULL, \n",
            "\tphone VARCHAR(50) NOT NULL, \n",
            "\t`addressLine1` VARCHAR(50) NOT NULL, \n",
            "\t`addressLine2` VARCHAR(50), \n",
            "\tstate VARCHAR(50), \n",
            "\tcountry VARCHAR(50) NOT NULL, \n",
            "\t`postalCode` VARCHAR(15) NOT NULL, \n",
            "\tterritory VARCHAR(10) NOT NULL, \n",
            "\tPRIMARY KEY (`officeCode`)\n",
            ")DEFAULT CHARSET=latin1 ENGINE=InnoDB\n",
            "\n",
            "/*\n",
            "2 rows from offices table:\n",
            "officeCode\tcity\tphone\taddressLine1\taddressLine2\tstate\tcountry\tpostalCode\tterritory\n",
            "1\tSan Francisco\t+1 650 219 4782\t100 Market Street\tSuite 300\tCA\tUSA\t94080\tNA\n",
            "2\tBoston\t+1 215 837 0825\t1550 Court Place\tSuite 102\tMA\tUSA\t02107\tNA\n",
            "*/\n",
            "\n",
            "\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",
            ")DEFAULT CHARSET=latin1 ENGINE=InnoDB\n",
            "\n",
            "/*\n",
            "2 rows from orderdetails table:\n",
            "orderNumber\tproductCode\tquantityOrdered\tpriceEach\torderLineNumber\n",
            "10100\tS18_1749\t30\t136.00\t3\n",
            "10100\tS18_2248\t50\t55.09\t2\n",
            "*/\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",
            ")DEFAULT CHARSET=latin1 ENGINE=InnoDB\n",
            "\n",
            "/*\n",
            "2 rows from orders table:\n",
            "orderNumber\torderDate\trequiredDate\tshippedDate\tstatus\tcomments\tcustomerNumber\n",
            "10100\t2003-01-06\t2003-01-13\t2003-01-10\tShipped\tNone\t363\n",
            "10101\t2003-01-09\t2003-01-18\t2003-01-11\tShipped\tCheck on availability.\t128\n",
            "*/\n",
            "\n",
            "\n",
            "CREATE TABLE payments (\n",
            "\t`customerNumber` INTEGER NOT NULL, \n",
            "\t`checkNumber` VARCHAR(50) NOT NULL, \n",
            "\t`paymentDate` DATE NOT NULL, \n",
            "\tamount DECIMAL(10, 2) NOT NULL, \n",
            "\tPRIMARY KEY (`customerNumber`, `checkNumber`), \n",
            "\tCONSTRAINT payments_ibfk_1 FOREIGN KEY(`customerNumber`) REFERENCES customers (`customerNumber`)\n",
            ")DEFAULT CHARSET=latin1 ENGINE=InnoDB\n",
            "\n",
            "/*\n",
            "2 rows from payments table:\n",
            "customerNumber\tcheckNumber\tpaymentDate\tamount\n",
            "103\tHQ336336\t2004-10-19\t6066.78\n",
            "103\tJM555205\t2003-06-05\t14571.44\n",
            "*/\n",
            "\n",
            "\n",
            "CREATE TABLE productlines (\n",
            "\t`productLine` VARCHAR(50) NOT NULL, \n",
            "\t`textDescription` VARCHAR(4000), \n",
            "\t`htmlDescription` MEDIUMTEXT, \n",
            "\timage MEDIUMBLOB, \n",
            "\tPRIMARY KEY (`productLine`)\n",
            ")DEFAULT CHARSET=latin1 ENGINE=InnoDB\n",
            "\n",
            "/*\n",
            "2 rows from productlines table:\n",
            "productLine\ttextDescription\thtmlDescription\timage\n",
            "Classic Cars\tAttention car enthusiasts: Make your wildest car ownership dreams come true. Whether you are looking\tNone\tNone\n",
            "Motorcycles\tOur motorcycles are state of the art replicas of classic as well as contemporary motorcycle legends \tNone\tNone\n",
            "*/\n",
            "\n",
            "\n",
            "CREATE TABLE products (\n",
            "\t`productCode` VARCHAR(15) NOT NULL, \n",
            "\t`productName` VARCHAR(70) NOT NULL, \n",
            "\t`productLine` VARCHAR(50) NOT NULL, \n",
            "\t`productScale` VARCHAR(10) NOT NULL, \n",
            "\t`productVendor` VARCHAR(50) NOT NULL, \n",
            "\t`productDescription` TEXT NOT NULL, \n",
            "\t`quantityInStock` SMALLINT NOT NULL, \n",
            "\t`buyPrice` DECIMAL(10, 2) NOT NULL, \n",
            "\t`MSRP` DECIMAL(10, 2) NOT NULL, \n",
            "\tPRIMARY KEY (`productCode`), \n",
            "\tCONSTRAINT products_ibfk_1 FOREIGN KEY(`productLine`) REFERENCES productlines (`productLine`)\n",
            ")DEFAULT CHARSET=latin1 ENGINE=InnoDB\n",
            "\n",
            "/*\n",
            "2 rows from products table:\n",
            "productCode\tproductName\tproductLine\tproductScale\tproductVendor\tproductDescription\tquantityInStock\tbuyPrice\tMSRP\n",
            "S10_1678\t1969 Harley Davidson Ultimate Chopper\tMotorcycles\t1:10\tMin Lin Diecast\tThis replica features working kickstand, front suspension, gear-shift lever, footbrake lever, drive \t7933\t48.81\t95.70\n",
            "S10_1949\t1952 Alpine Renault 1300\tClassic Cars\t1:10\tClassic Metal Creations\tTurnable front wheels; steering function; detailed interior; detailed engine; opening hood; opening \t7305\t98.58\t214.30\n",
            "*/\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import openai\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"sk-bE\"\n",
        "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
      ],
      "metadata": {
        "id": "FFxFOyxoQRKj"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "https://gpt-index.readthedocs.io/en/latest/examples/callbacks/TokenCountingHandler.html"
      ],
      "metadata": {
        "id": "JttyItMycui0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import tiktoken\n",
        "from llama_index.callbacks import CallbackManager, TokenCountingHandler\n",
        "token_counter = TokenCountingHandler(\n",
        "    tokenizer=tiktoken.encoding_for_model(\"gpt-3.5-turbo\").encode\n",
        ")\n",
        "\n",
        "callback_manager = CallbackManager([token_counter])"
      ],
      "metadata": {
        "id": "kVIxgmsszwfi"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index import ServiceContext, LLMPredictor, OpenAIEmbedding, PromptHelper\n",
        "from llama_index.llms import OpenAI\n",
        "from llama_index import ServiceContext\n",
        "llm = OpenAI(temperature=0.1, model=\"gpt-3.5-turbo\")\n",
        "\n",
        "service_context = ServiceContext.from_defaults(\n",
        "  llm=llm,callback_manager=callback_manager\n",
        ")"
      ],
      "metadata": {
        "id": "l20LYCv2SF7N"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.indices.struct_store.sql_query import NLSQLTableQueryEngine\n",
        "\n",
        "query_engine = NLSQLTableQueryEngine(\n",
        "    sql_database=sql_database,\n",
        "    service_context=service_context\n",
        ")"
      ],
      "metadata": {
        "id": "q1Rzr7uMQnOA"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# query_str = \"Which customer has the highest orders?\"\n",
        "query_str = \"when was order number 10100 shiiped?\"\n",
        "# query_str =\"What are the most frequently mentioned keywords or phrases in the comments made by sales representatives\"\n",
        "response = query_engine.query(query_str)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nW7aMjT4aaSc",
        "outputId": "c21eff9b-de95-4d58-f6e1-b0ecd0ebf1ed"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:llama_index.indices.struct_store.sql_query:> Table desc str: Table 'customers' has columns: customerNumber (INTEGER), customerName (VARCHAR(50)), contactLastName (VARCHAR(50)), contactFirstName (VARCHAR(50)), phone (VARCHAR(50)), addressLine1 (VARCHAR(50)), addressLine2 (VARCHAR(50)), city (VARCHAR(50)), state (VARCHAR(50)), postalCode (VARCHAR(15)), country (VARCHAR(50)), salesRepEmployeeNumber (INTEGER), creditLimit (DECIMAL(10, 2)), and foreign keys: ['salesRepEmployeeNumber'] -> employees.['employeeNumber'].\n",
            "\n",
            "Table 'employees' has columns: employeeNumber (INTEGER), lastName (VARCHAR(50)), firstName (VARCHAR(50)), extension (VARCHAR(10)), email (VARCHAR(100)), officeCode (VARCHAR(10)), reportsTo (INTEGER), jobTitle (VARCHAR(50)), and foreign keys: ['reportsTo'] -> employees.['employeeNumber'], ['officeCode'] -> offices.['officeCode'].\n",
            "\n",
            "Table 'offices' has columns: officeCode (VARCHAR(10)), city (VARCHAR(50)), phone (VARCHAR(50)), addressLine1 (VARCHAR(50)), addressLine2 (VARCHAR(50)), state (VARCHAR(50)), country (VARCHAR(50)), postalCode (VARCHAR(15)), territory (VARCHAR(10)), and foreign keys: .\n",
            "\n",
            "Table 'orderdetails' has columns: orderNumber (INTEGER), productCode (VARCHAR(15)), quantityOrdered (INTEGER), priceEach (DECIMAL(10, 2)), orderLineNumber (SMALLINT), and foreign keys: ['orderNumber'] -> orders.['orderNumber'], ['productCode'] -> products.['productCode'].\n",
            "\n",
            "Table 'orders' has columns: orderNumber (INTEGER), orderDate (DATE), requiredDate (DATE), shippedDate (DATE), status (VARCHAR(15)), comments (TEXT), customerNumber (INTEGER), and foreign keys: ['customerNumber'] -> customers.['customerNumber'].\n",
            "\n",
            "Table 'payments' has columns: customerNumber (INTEGER), checkNumber (VARCHAR(50)), paymentDate (DATE), amount (DECIMAL(10, 2)), and foreign keys: ['customerNumber'] -> customers.['customerNumber'].\n",
            "\n",
            "Table 'productlines' has columns: productLine (VARCHAR(50)), textDescription (VARCHAR(4000)), htmlDescription (MEDIUMTEXT), image (MEDIUMBLOB), and foreign keys: .\n",
            "\n",
            "Table 'products' has columns: productCode (VARCHAR(15)), productName (VARCHAR(70)), productLine (VARCHAR(50)), productScale (VARCHAR(10)), productVendor (VARCHAR(50)), productDescription (TEXT), quantityInStock (SMALLINT), buyPrice (DECIMAL(10, 2)), MSRP (DECIMAL(10, 2)), and foreign keys: ['productLine'] -> productlines.['productLine'].\n",
            "> Table desc str: Table 'customers' has columns: customerNumber (INTEGER), customerName (VARCHAR(50)), contactLastName (VARCHAR(50)), contactFirstName (VARCHAR(50)), phone (VARCHAR(50)), addressLine1 (VARCHAR(50)), addressLine2 (VARCHAR(50)), city (VARCHAR(50)), state (VARCHAR(50)), postalCode (VARCHAR(15)), country (VARCHAR(50)), salesRepEmployeeNumber (INTEGER), creditLimit (DECIMAL(10, 2)), and foreign keys: ['salesRepEmployeeNumber'] -> employees.['employeeNumber'].\n",
            "\n",
            "Table 'employees' has columns: employeeNumber (INTEGER), lastName (VARCHAR(50)), firstName (VARCHAR(50)), extension (VARCHAR(10)), email (VARCHAR(100)), officeCode (VARCHAR(10)), reportsTo (INTEGER), jobTitle (VARCHAR(50)), and foreign keys: ['reportsTo'] -> employees.['employeeNumber'], ['officeCode'] -> offices.['officeCode'].\n",
            "\n",
            "Table 'offices' has columns: officeCode (VARCHAR(10)), city (VARCHAR(50)), phone (VARCHAR(50)), addressLine1 (VARCHAR(50)), addressLine2 (VARCHAR(50)), state (VARCHAR(50)), country (VARCHAR(50)), postalCode (VARCHAR(15)), territory (VARCHAR(10)), and foreign keys: .\n",
            "\n",
            "Table 'orderdetails' has columns: orderNumber (INTEGER), productCode (VARCHAR(15)), quantityOrdered (INTEGER), priceEach (DECIMAL(10, 2)), orderLineNumber (SMALLINT), and foreign keys: ['orderNumber'] -> orders.['orderNumber'], ['productCode'] -> products.['productCode'].\n",
            "\n",
            "Table 'orders' has columns: orderNumber (INTEGER), orderDate (DATE), requiredDate (DATE), shippedDate (DATE), status (VARCHAR(15)), comments (TEXT), customerNumber (INTEGER), and foreign keys: ['customerNumber'] -> customers.['customerNumber'].\n",
            "\n",
            "Table 'payments' has columns: customerNumber (INTEGER), checkNumber (VARCHAR(50)), paymentDate (DATE), amount (DECIMAL(10, 2)), and foreign keys: ['customerNumber'] -> customers.['customerNumber'].\n",
            "\n",
            "Table 'productlines' has columns: productLine (VARCHAR(50)), textDescription (VARCHAR(4000)), htmlDescription (MEDIUMTEXT), image (MEDIUMBLOB), and foreign keys: .\n",
            "\n",
            "Table 'products' has columns: productCode (VARCHAR(15)), productName (VARCHAR(70)), productLine (VARCHAR(50)), productScale (VARCHAR(10)), productVendor (VARCHAR(50)), productDescription (TEXT), quantityInStock (SMALLINT), buyPrice (DECIMAL(10, 2)), MSRP (DECIMAL(10, 2)), and foreign keys: ['productLine'] -> productlines.['productLine'].\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(response.response)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EcrUO1D6af7E",
        "outputId": "c15aba3b-afc7-49a5-9045-2ea7e48f3dd7"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Order number 10100 was shipped on January 10, 2003.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(response.metadata['sql_query'])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "172e64M5QpaP",
        "outputId": "b57e198e-80a7-4e63-9a7c-7a4fae19c6c8"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "SELECT shippedDate\n",
            "FROM orders\n",
            "WHERE orderNumber = 10100\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(response.metadata['result'])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FWMZcq4UQuIG",
        "outputId": "2da05c8a-6038-443b-8fe9-57670ae39f9c"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[(datetime.date(2003, 1, 10),)]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(token_counter.total_llm_token_count)"
      ],
      "metadata": {
        "id": "Ovl2gD15079v",
        "outputId": "3566ab7b-4442-49d5-bced-3843ab134598",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "904\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "token_counter.reset_counts()"
      ],
      "metadata": {
        "id": "f_kEFjRbaQ42"
      },
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.objects import ObjectIndex\n",
        "from llama_index.objects import SQLTableNodeMapping, SQLTableSchema\n",
        "import pandas as pd\n",
        "\n",
        "tables = list(sql_database._all_tables)\n",
        "table_node_mapping = SQLTableNodeMapping(sql_database)\n",
        "table_schema_objs = []\n",
        "for table in tables:\n",
        "    table_schema_objs.append((SQLTableSchema(table_name = table, context_str = table_details[table])))"
      ],
      "metadata": {
        "id": "kDZ4ZTpJQzn5"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "table_schema_objs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iHfmYPhNSBEM",
        "outputId": "6bf053f8-b094-4be9-b9d8-17ee6ff87885"
      },
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[SQLTableSchema(table_name='payments', context_str='stores payments made by customers based on their accounts.'),\n",
              " SQLTableSchema(table_name='employees', context_str='stores all employee information as well as the organization structure such as who reports to whom.'),\n",
              " SQLTableSchema(table_name='customers', context_str='stores customer’s data.'),\n",
              " SQLTableSchema(table_name='orders', context_str='stores sales orders placed by customers.'),\n",
              " SQLTableSchema(table_name='productlines', context_str='stores a list of product line categories.'),\n",
              " SQLTableSchema(table_name='products', context_str='stores a list of scale model cars.'),\n",
              " SQLTableSchema(table_name='orderdetails', context_str='stores sales order line items for each sales order.'),\n",
              " SQLTableSchema(table_name='offices', context_str='stores sales office data.')]"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from llama_index.indices.struct_store import SQLTableRetrieverQueryEngine\n",
        "from llama_index import VectorStoreIndex\n",
        "\n",
        "\n",
        "obj_index = ObjectIndex.from_objects(\n",
        "    table_schema_objs,\n",
        "    table_node_mapping,\n",
        "    VectorStoreIndex,\n",
        "    service_context=service_context\n",
        ")"
      ],
      "metadata": {
        "id": "I6IGxVIkSdz_"
      },
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query_engine = SQLTableRetrieverQueryEngine(\n",
        "    sql_database, obj_index.as_retriever(similarity_top_k=3), service_context=service_context\n",
        ")\n",
        "response = query_engine.query(\"How many customers we have?\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2In0lOM3jZmk",
        "outputId": "bd11b1b7-ffd4-4852-9195-06be87d00b28"
      },
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:llama_index.indices.struct_store.sql_query:> Table desc str: Table 'customers' has columns: customerNumber (INTEGER), customerName (VARCHAR(50)), contactLastName (VARCHAR(50)), contactFirstName (VARCHAR(50)), phone (VARCHAR(50)), addressLine1 (VARCHAR(50)), addressLine2 (VARCHAR(50)), city (VARCHAR(50)), state (VARCHAR(50)), postalCode (VARCHAR(15)), country (VARCHAR(50)), salesRepEmployeeNumber (INTEGER), creditLimit (DECIMAL(10, 2)), and foreign keys: ['salesRepEmployeeNumber'] -> employees.['employeeNumber']. The table description is: stores customer’s data.\n",
            "\n",
            "Table 'orders' has columns: orderNumber (INTEGER), orderDate (DATE), requiredDate (DATE), shippedDate (DATE), status (VARCHAR(15)), comments (TEXT), customerNumber (INTEGER), and foreign keys: ['customerNumber'] -> customers.['customerNumber']. The table description is: stores sales orders placed by customers.\n",
            "\n",
            "Table 'payments' has columns: customerNumber (INTEGER), checkNumber (VARCHAR(50)), paymentDate (DATE), amount (DECIMAL(10, 2)), and foreign keys: ['customerNumber'] -> customers.['customerNumber']. The table description is: stores payments made by customers based on their accounts.\n",
            "> Table desc str: Table 'customers' has columns: customerNumber (INTEGER), customerName (VARCHAR(50)), contactLastName (VARCHAR(50)), contactFirstName (VARCHAR(50)), phone (VARCHAR(50)), addressLine1 (VARCHAR(50)), addressLine2 (VARCHAR(50)), city (VARCHAR(50)), state (VARCHAR(50)), postalCode (VARCHAR(15)), country (VARCHAR(50)), salesRepEmployeeNumber (INTEGER), creditLimit (DECIMAL(10, 2)), and foreign keys: ['salesRepEmployeeNumber'] -> employees.['employeeNumber']. The table description is: stores customer’s data.\n",
            "\n",
            "Table 'orders' has columns: orderNumber (INTEGER), orderDate (DATE), requiredDate (DATE), shippedDate (DATE), status (VARCHAR(15)), comments (TEXT), customerNumber (INTEGER), and foreign keys: ['customerNumber'] -> customers.['customerNumber']. The table description is: stores sales orders placed by customers.\n",
            "\n",
            "Table 'payments' has columns: customerNumber (INTEGER), checkNumber (VARCHAR(50)), paymentDate (DATE), amount (DECIMAL(10, 2)), and foreign keys: ['customerNumber'] -> customers.['customerNumber']. The table description is: stores payments made by customers based on their accounts.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(response)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lh8c1D-MUrGN",
        "outputId": "cc82b2a6-ff82-4b8b-e144-d56dff78e7d7"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "We have 122 customers.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(token_counter.total_llm_token_count)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bVAv5lCyhCcU",
        "outputId": "7553df59-b4d3-46db-a884-47f2d1d3d961"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "514\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "yEdRxks6kG4z"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}