14_text2sql_agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "85032306-c629-4c28-9395-79500591dbb6",
   "metadata": {},
   "source": [
    "# Text2SQL Agent to Interact with CSV Data\n",
    "\n",
    "## System Architecture\n",
    "\n",
    "COnsider it as an agent with a set of tools such as `search_cache()` and `generate_SQL_query()`, and `run_sql_query()`.\n",
    "\n",
    "![](images/agentic_text2sql.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c4832ae-65b3-466f-8ecf-80331153d9a4",
   "metadata": {},
   "source": [
    "## Data Ingestion Pipeline\n",
    "\n",
    "1. Read CSV\n",
    "2. Create Database schema\n",
    "3. Create a table\n",
    "4. Load table with CSV data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1c4991b8-da07-45b9-9972-cb0b27b6617f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Table 'movies' created with schema: \"Movie\" TEXT, \"LeadStudio\" TEXT, \"RottenTomatoes\" REAL, \"AudienceScore\" REAL, \"Story\" TEXT, \"Genre\" TEXT, \"TheatersOpenWeek\" REAL, \"OpeningWeekend\" REAL, \"BOAvgOpenWeekend\" REAL, \"DomesticGross\" REAL, \"ForeignGross\" REAL, \"WorldGross\" REAL, \"Budget\" REAL, \"Profitability\" REAL, \"OpenProfit\" REAL, \"Year\" INTEGER\n",
      "Data loaded into 'movies' table in 'movies_db.db' SQLite database.\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import sqlite3\n",
    "\n",
    "def csv_to_sqlite(csv_file, db_name, table_name):\n",
    "    # Read the CSV file into a pandas DataFrame\n",
    "    df = pd.read_csv(csv_file)\n",
    "    \n",
    "    # Connect to the SQLite database (it will create the database file if it doesn't exist)\n",
    "    conn = sqlite3.connect(db_name)\n",
    "    cursor = conn.cursor()\n",
    "\n",
    "    # Infer the schema based on the DataFrame columns and data types\n",
    "    def create_table_from_df(df, table_name):\n",
    "        # Get column names and types\n",
    "        col_types = []\n",
    "        for col in df.columns:\n",
    "            dtype = df[col].dtype\n",
    "            if dtype == 'int64':\n",
    "                col_type = 'INTEGER'\n",
    "            elif dtype == 'float64':\n",
    "                col_type = 'REAL'\n",
    "            else:\n",
    "                col_type = 'TEXT'\n",
    "            col_types.append(f'\"{col}\" {col_type}')\n",
    "        \n",
    "        # Create the table schema\n",
    "        col_definitions = \", \".join(col_types)\n",
    "        create_table_query = f'CREATE TABLE IF NOT EXISTS {table_name} ({col_definitions});'\n",
    "        # print(create_table_query)\n",
    "        \n",
    "        # Execute the table creation query\n",
    "        cursor.execute(create_table_query)\n",
    "        print(f\"Table '{table_name}' created with schema: {col_definitions}\")\n",
    "    \n",
    "    # Create table schema\n",
    "    create_table_from_df(df, table_name)\n",
    "    \n",
    "    # Insert CSV data into the SQLite table\n",
    "    df.to_sql(table_name, conn, if_exists='replace', index=False)\n",
    "    \n",
    "    # Commit and close the connection\n",
    "    conn.commit()\n",
    "    conn.close()\n",
    "    print(f\"Data loaded into '{table_name}' table in '{db_name}' SQLite database.\")\n",
    "\n",
    "csv_file = \"movies.csv\"\n",
    "db_name = \"movies_db.db\"\n",
    "table_name = \"movies\"\n",
    "csv_to_sqlite(csv_file, db_name, table_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2f5cf440-8871-420b-8bf3-ab56e3d35d48",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_sql_query(db_name, query):\n",
    "    \"\"\"\n",
    "    Executes a SQL query on a SQLite database and returns the results.\n",
    "\n",
    "    Args:\n",
    "        db_name (str): The name of the SQLite database file.\n",
    "        query (str): The SQL query to run.\n",
    "\n",
    "    Returns:\n",
    "        list: Query result as a list of tuples, or an empty list if no results or error occurred.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # Connect to the SQLite database\n",
    "        conn = sqlite3.connect(db_name)\n",
    "        cursor = conn.cursor()\n",
    "\n",
    "        # Execute the SQL query\n",
    "        cursor.execute(query)\n",
    "\n",
    "        # Fetch all results\n",
    "        results = cursor.fetchall()\n",
    "\n",
    "        # Close the connection\n",
    "        conn.close()\n",
    "\n",
    "        # Return results or an empty list if no results were found\n",
    "        return results if results else []\n",
    "    \n",
    "    except sqlite3.Error as e:\n",
    "        print(f\"An error occurred while executing the query: {e}\")\n",
    "        return []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "018d0654-889c-4530-806e-e33a441e4b4e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(970,)\n"
     ]
    }
   ],
   "source": [
    "query = f\"SELECT count(*) FROM {table_name};\"\n",
    "results = run_sql_query(db_name, query)\n",
    "\n",
    "if results:\n",
    "    for row in results:\n",
    "        print(row)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da4ac953-862d-4ab4-bd48-f1a6161dfe43",
   "metadata": {},
   "source": [
    "## Ask Natural Language Questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "5f8d7cae-86b6-4e8b-937c-8cd0d51f0471",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import faiss\n",
    "import numpy as np\n",
    "import os\n",
    "from openai import OpenAI\n",
    "from litellm import completion\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5211a0e2-f105-4c65-b841-d69668963367",
   "metadata": {},
   "outputs": [],
   "source": [
    "OPENAI_API_KEY = os.environ.get(\"OPENAI_API_KEY\")\n",
    "client = OpenAI(api_key=OPENAI_API_KEY)\n",
    "\n",
    "# Initialize the FAISS index\n",
    "dimension = 1536  # Dimension size for OpenAI embeddings (may vary by model)\n",
    "index = faiss.IndexFlatL2(dimension)  # L2 distance index\n",
    "\n",
    "# Cache will hold (user_question, sql_query, response)\n",
    "cache = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8181215d-9aab-4b23-bd4c-11e090300735",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function to get embeddings from OpenAI or any embedding model\n",
    "def get_embeddings(text):\n",
    "    \"\"\"\n",
    "    Converts a text string into a vector embedding using OpenAI embeddings.\n",
    "    \n",
    "    Args:\n",
    "        text (str): The text string to convert.\n",
    "    \n",
    "    Returns:\n",
    "        np.array: A vector representation of the text.\n",
    "    \"\"\"\n",
    "    response = client.embeddings.create(input=text, model=\"text-embedding-3-small\")\n",
    "    embedding = np.array(response.data[0].embedding)\n",
    "    return embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "d21c6fee-584c-40c1-a521-ebe0596cf032",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_cache(question_embedding, threshold=0.1):\n",
    "    \"\"\"\n",
    "    Searches the FAISS index for a similar question.\n",
    "    \n",
    "    Args:\n",
    "        question_embedding (np.array): The embedding of the user's question.\n",
    "        threshold (float): The similarity threshold for considering a hit.\n",
    "    \n",
    "    Returns:\n",
    "        tuple: (sql_query, response) if a hit is found, otherwise None.\n",
    "    \"\"\"\n",
    "    if index.ntotal > 0:\n",
    "        distances, indices = index.search(np.array([question_embedding]), k=1)\n",
    "        # print(distances)\n",
    "        # print(indices)\n",
    "        # Check if the closest distance is below the threshold\n",
    "        if distances[0][0] < threshold:\n",
    "            cache_index = indices[0][0]\n",
    "            return cache[cache_index][1], cache[cache_index][2]\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "14671fe2-40b7-48bd-b3f4-7e004c241211",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Schema for movies:\n",
      "(0, 'Movie', 'TEXT', 0, None, 0)\n",
      "(1, 'LeadStudio', 'TEXT', 0, None, 0)\n",
      "(2, 'RottenTomatoes', 'REAL', 0, None, 0)\n",
      "(3, 'AudienceScore', 'REAL', 0, None, 0)\n",
      "(4, 'Story', 'TEXT', 0, None, 0)\n",
      "(5, 'Genre', 'TEXT', 0, None, 0)\n",
      "(6, 'TheatersOpenWeek', 'REAL', 0, None, 0)\n",
      "(7, 'OpeningWeekend', 'REAL', 0, None, 0)\n",
      "(8, 'BOAvgOpenWeekend', 'REAL', 0, None, 0)\n",
      "(9, 'DomesticGross', 'REAL', 0, None, 0)\n",
      "(10, 'ForeignGross', 'REAL', 0, None, 0)\n",
      "(11, 'WorldGross', 'REAL', 0, None, 0)\n",
      "(12, 'Budget', 'REAL', 0, None, 0)\n",
      "(13, 'Profitability', 'REAL', 0, None, 0)\n",
      "(14, 'OpenProfit', 'REAL', 0, None, 0)\n",
      "(15, 'Year', 'INTEGER', 0, None, 0)\n"
     ]
    }
   ],
   "source": [
    "def get_table_schema(db_name, table_name):\n",
    "    \"\"\"\n",
    "    Retrieves the schema (columns and data types) for a given table in the SQLite database.\n",
    "    \n",
    "    Args:\n",
    "        db_name (str): The name of the SQLite database file.\n",
    "        table_name (str): The name of the table.\n",
    "    \n",
    "    Returns:\n",
    "        list: A list of tuples with column name, data type, and other info.\n",
    "    \"\"\"\n",
    "    conn = sqlite3.connect(db_name)\n",
    "    cursor = conn.cursor()\n",
    "\n",
    "    # Use PRAGMA to get the table schema\n",
    "    cursor.execute(f\"PRAGMA table_info({table_name});\")\n",
    "    schema = cursor.fetchall()\n",
    "\n",
    "    conn.close()\n",
    "    return schema\n",
    "\n",
    "table_name = 'movies'\n",
    "schema = get_table_schema(db_name, table_name)\n",
    "print(f\"Schema for {table_name}:\")\n",
    "for col in schema:\n",
    "    print(col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ef07dfd8-2109-40de-a849-5b9917eb497d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are an expert in writing SQL queries for relational databases. \n",
      "    You will be provided with a database schema and a natural \n",
      "    language question, and your task is to generate an accurate SQL query.\n",
      "    \n",
      "    The database has a table named 'movies' with the following schema:\n",
      "\n",
      "Columns:\n",
      "- Movie (TEXT)\n",
      "- LeadStudio (TEXT)\n",
      "- RottenTomatoes (REAL)\n",
      "- AudienceScore (REAL)\n",
      "- Story (TEXT)\n",
      "- Genre (TEXT)\n",
      "- TheatersOpenWeek (REAL)\n",
      "- OpeningWeekend (REAL)\n",
      "- BOAvgOpenWeekend (REAL)\n",
      "- DomesticGross (REAL)\n",
      "- ForeignGross (REAL)\n",
      "- WorldGross (REAL)\n",
      "- Budget (REAL)\n",
      "- Profitability (REAL)\n",
      "- OpenProfit (REAL)\n",
      "- Year (INTEGER)\n",
      "\n",
      "Please generate a SQL query based on the following natural language question. ONLY return the SQL query.\n"
     ]
    }
   ],
   "source": [
    "def generate_llm_prompt(table_name, table_schema):\n",
    "    \"\"\"\n",
    "    Generates a prompt to provide context about a table's schema for LLM to convert natural language to SQL.\n",
    "    \n",
    "    Args:\n",
    "        table_name (str): The name of the table.\n",
    "        table_schema (list): A list of tuples where each tuple contains information about the columns in the table.\n",
    "    \n",
    "    Returns:\n",
    "        str: The generated prompt to be used by the LLM.\n",
    "    \"\"\"\n",
    "    prompt = f\"\"\"You are an expert in writing SQL queries for relational databases. \n",
    "    You will be provided with a database schema and a natural \n",
    "    language question, and your task is to generate an accurate SQL query.\n",
    "    \n",
    "    The database has a table named '{table_name}' with the following schema:\\n\\n\"\"\"\n",
    "    \n",
    "    prompt += \"Columns:\\n\"\n",
    "\n",
    "    for col in table_schema:\n",
    "        column_name = col[1]\n",
    "        column_type = col[2]\n",
    "        prompt += f\"- {column_name} ({column_type})\\n\"\n",
    "    \n",
    "    prompt += \"\\nPlease generate a SQL query based on the following natural language question. ONLY return the SQL query.\"\n",
    "\n",
    "    return prompt\n",
    "\n",
    "table_name = \"movies\"\n",
    "schema = get_table_schema(db_name, table_name)\n",
    "# Generate the prompt\n",
    "llm_prompt = generate_llm_prompt(table_name, schema)\n",
    "print(llm_prompt)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "ff37a2e8-165b-4016-b6e1-2965ad3fe370",
   "metadata": {},
   "outputs": [],
   "source": [
    "def handle_user_question(user_question):\n",
    "    \"\"\"\n",
    "    Handles the user's question by first searching the cache, and if there's no hit, generating a SQL query and response.\n",
    "    \n",
    "    Args:\n",
    "        user_question (str): The user's natural language question.\n",
    "    \n",
    "    Returns:\n",
    "        list: The response to the user's question.\n",
    "    \"\"\"\n",
    "    # Convert the user's question to an embedding\n",
    "    question_embedding = get_embeddings(user_question)\n",
    "    \n",
    "    # Step 1: Search cache for similar questions\n",
    "    cache_hit = search_cache(question_embedding)\n",
    "    if cache_hit:\n",
    "        sql_query, response = cache_hit\n",
    "        print(f\"Cache hit! SQL Query: {sql_query}\")\n",
    "        return response\n",
    "    \n",
    "    # Step 2: No hit, go to LLM for SQL generation\n",
    "    print(\"Cache miss! Generating SQL from LLM...\")\n",
    "    sql_query = generate_sql_query(user_question)\n",
    "    \n",
    "    # Step 3: Run the SQL query on the database\n",
    "    response = run_sql_query(db_name, sql_query)\n",
    "    \n",
    "    # Step 4: Store question, SQL, and response in cache\n",
    "    cache.append((user_question, sql_query, response))\n",
    "    index.add(np.array([question_embedding]))  # Add question embedding to FAISS index\n",
    "    \n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b3158629-6950-4a7a-bbcc-17acd0b00e98",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_sql_query(question):\n",
    "    table_name = 'movies'\n",
    "    db_name = 'movies_db.db'\n",
    "    table_schema = get_table_schema(db_name, table_name)\n",
    "    llm_prompt = generate_llm_prompt(table_name, table_schema)\n",
    "    user_prompt = \"\"\"Question: {question}\"\"\"\n",
    "    response = completion(\n",
    "        api_key=OPENAI_API_KEY,\n",
    "        model=\"gpt-4o-mini\",\n",
    "        messages=[\n",
    "            {\"content\": llm_prompt.format(table_name=table_name),\"role\": \"system\"}, \n",
    "            {\"content\": user_prompt.format(question=question),\"role\": \"user\"}],\n",
    "        max_tokens=1000    \n",
    "    )\n",
    "    answer = response.choices[0].message.content\n",
    "    display(Markdown(answer))\n",
    "    query = answer.replace(\"```sql\", \"\").replace(\"```\", \"\")\n",
    "    query = query.strip()\n",
    "    return query\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "cd4ea65f-d9ef-4874-a16e-36f62d447c0a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cache miss! Generating SQL from LLM...\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "```sql\n",
       "SELECT COUNT(*) AS ActionMovieCount\n",
       "FROM movies\n",
       "WHERE Genre = 'Action';\n",
       "```"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[(166,)]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# question = \"total number of movies are made by Warner Bros company in year 2008?\"\n",
    "# question = \"how many movies have RottenTomatoes scores lower than 85?\"\n",
    "question = \"how many movies with action genre are in the database\"\n",
    "handle_user_question(question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "9b089cbb-b6e8-4309-be04-4eee20b20619",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('total number of movies are made by Warner Bros company in year 2008?',\n",
       "  \"SELECT COUNT(*) \\nFROM movies \\nWHERE LeadStudio = 'Warner Bros' AND Year = 2008;\",\n",
       "  [(21,)]),\n",
       " ('how many movies have RottenTomatoes scores greater than 85?',\n",
       "  'SELECT COUNT(*) \\nFROM movies \\nWHERE RottenTomatoes > 85;',\n",
       "  [(120,)]),\n",
       " ('how many movies have RottenTomatoes scores lower than 85?',\n",
       "  'SELECT COUNT(*) \\nFROM movies \\nWHERE RottenTomatoes < 85;',\n",
       "  [(782,)]),\n",
       " ('how many movies with action genre are in the database',\n",
       "  \"SELECT COUNT(*) AS ActionMovieCount\\nFROM movies\\nWHERE Genre = 'Action';\",\n",
       "  [(166,)])]"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cache"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}