run-llama--llama_index/docs/examples/query_engine/pandas_query_engine.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8329aae0",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/query_engine/pandas_query_engine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e45f9b60-cd6b-4c15-958f-1feca5438128",
   "metadata": {},
   "source": [
    "# Pandas Query Engine\n",
    "\n",
    "This guide shows you how to use our `PandasQueryEngine`: convert natural language to Pandas python code using LLMs.\n",
    "\n",
    "The input to the `PandasQueryEngine` is a Pandas dataframe, and the output is a response. The LLM infers dataframe operations to perform in order to retrieve the result.\n",
    "\n",
    "**WARNING:** This tool provides the LLM access to the `eval` function.\n",
    "Arbitrary code execution is possible on the machine running this tool.\n",
    "While some level of filtering is done on code, this tool is not recommended \n",
    "to be used in a production setting without heavy sandboxing or virtual machines.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a6a0f96",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82fa3d16",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index llama-index-experimental"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "119eb42b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "from IPython.display import Markdown, display\n",
    "\n",
    "import pandas as pd\n",
    "from llama_index.experimental.query_engine import PandasQueryEngine\n",
    "\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5ece7d73-0f67-4ff5-95e5-249a25bd118c",
   "metadata": {},
   "source": [
    "## Let's start on a Toy DataFrame\n",
    "\n",
    "Here let's load a very simple dataframe containing city and population pairs, and run the `PandasQueryEngine` on it.\n",
    "\n",
    "By setting `verbose=True` we can see the intermediate generated instructions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1484fe58-4853-4a76-bffc-435a9cce3e2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test on some sample data\n",
    "df = pd.DataFrame(\n",
    "    {\n",
    "        \"city\": [\"Toronto\", \"Tokyo\", \"Berlin\"],\n",
    "        \"population\": [2930000, 13960000, 3645000],\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fea2edb-b3d4-4313-a656-d6edb00d93c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = PandasQueryEngine(df=df, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "451836bc-b073-4838-8ab8-3def7d2c4d9d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "> Pandas Instructions:\n",
      "```\n",
      "df['city'][df['population'].idxmax()]\n",
      "```\n",
      "> Pandas Output: Tokyo\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"What is the city with the highest population?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4253d4c3-f3e5-4779-bcd1-2e6e2818305f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<b>Tokyo</b>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e10b7da-b355-49b2-9f80-f17541d4f850",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "df['city'][df['population'].idxmax()]\n"
     ]
    }
   ],
   "source": [
    "# get pandas python instructions\n",
    "print(response.metadata[\"pandas_instruction_str\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c86a122-b65f-40af-92fe-b3054d526eb4",
   "metadata": {},
   "source": [
    "We can also take the step of using an LLM to synthesize a response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79829afa-5d7e-4e7c-a5c8-433f8e5fceb7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "> Pandas Instructions:\n",
      "```\n",
      "df.loc[df['population'].idxmax()]\n",
      "```\n",
      "> Pandas Output: city             Tokyo\n",
      "population    13960000\n",
      "Name: 1, dtype: object\n",
      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "The city with the highest population is Tokyo, with a population of 13,960,000.\n"
     ]
    }
   ],
   "source": [
    "query_engine = PandasQueryEngine(df=df, verbose=True, synthesize_response=True)\n",
    "response = query_engine.query(\n",
    "    \"What is the city with the highest population? Give both the city and population\",\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1de5eaf3-6129-47b1-b630-faf9138a04c5",
   "metadata": {},
   "source": [
    "## Analyzing the Titanic Dataset\n",
    "\n",
    "The Titanic dataset is one of the most popular tabular datasets in introductory machine learning\n",
    "Source: https://www.kaggle.com/c/titanic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e84f82e",
   "metadata": {},
   "source": [
    "#### Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc2e1896",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-01-13 17:45:15--  https://raw.githubusercontent.com/jerryjliu/llama_index/main/docs/examples/data/csv/titanic_train.csv\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8002::154, 2606:50c0:8001::154, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 57726 (56K) [text/plain]\n",
      "Saving to: ‘titanic_train.csv’\n",
      "\n",
      "titanic_train.csv   100%[===================>]  56.37K  --.-KB/s    in 0.009s  \n",
      "\n",
      "2024-01-13 17:45:15 (6.45 MB/s) - ‘titanic_train.csv’ saved [57726/57726]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget 'https://raw.githubusercontent.com/jerryjliu/llama_index/main/docs/examples/data/csv/titanic_train.csv' -O 'titanic_train.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "809f18c8-e38b-449e-b5ee-c2ea700f8698",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"./titanic_train.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb1758de-6310-4ed5-ae02-2dbf50d2c55f",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = PandasQueryEngine(df=df, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9dd658d-b62c-4e3b-aee9-0a06f57de032",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "> Pandas Instructions:\n",
      "```\n",
      "df['survived'].corr(df['age'])\n",
      "```\n",
      "> Pandas Output: -0.07722109457217755\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"What is the correlation between survival and age?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60474389-341b-4187-87b2-83811546dcea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<b>-0.07722109457217755</b>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Markdown(f\"<b>{response}</b>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af999a1f-fea6-4734-82e6-4450f1a06a3b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "df['survived'].corr(df['age'])\n"
     ]
    }
   ],
   "source": [
    "# get pandas python instructions\n",
    "print(response.metadata[\"pandas_instruction_str\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1896e6d8-9324-4ea8-a903-bcbd94c5fb4b",
   "metadata": {},
   "source": [
    "## Additional Steps\n",
    "\n",
    "### Analyzing / Modifying prompts\n",
    "\n",
    "Let's look at the prompts! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7007101-20f0-4081-bee3-ef8e2e7f3796",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import PromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a4a924d-315a-48e5-a404-f029c8962fc5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You are working with a pandas dataframe in Python.\n",
      "The name of the dataframe is `df`.\n",
      "This is the result of `print(df.head())`:\n",
      "{df_str}\n",
      "\n",
      "Follow these instructions:\n",
      "{instruction_str}\n",
      "Query: {query_str}\n",
      "\n",
      "Expression:\n"
     ]
    }
   ],
   "source": [
    "query_engine = PandasQueryEngine(df=df, verbose=True)\n",
    "prompts = query_engine.get_prompts()\n",
    "print(prompts[\"pandas_prompt\"].template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c561301-3cfe-4d61-86e3-9d43847534f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Given an input question, synthesize a response from the query results.\n",
      "Query: {query_str}\n",
      "\n",
      "Pandas Instructions (optional):\n",
      "{pandas_instructions}\n",
      "\n",
      "Pandas Output: {pandas_output}\n",
      "\n",
      "Response: \n"
     ]
    }
   ],
   "source": [
    "print(prompts[\"response_synthesis_prompt\"].template)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f30b563-ba0d-445f-83e9-46a8f480a83a",
   "metadata": {},
   "source": [
    "You can update prompts as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0118dfb-80e2-40d2-88dc-b8f6c62d786b",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_prompt = PromptTemplate(\n",
    "    \"\"\"\\\n",
    "You are working with a pandas dataframe in Python.\n",
    "The name of the dataframe is `df`.\n",
    "This is the result of `print(df.head())`:\n",
    "{df_str}\n",
    "\n",
    "Follow these instructions:\n",
    "{instruction_str}\n",
    "Query: {query_str}\n",
    "\n",
    "Expression: \"\"\"\n",
    ")\n",
    "\n",
    "query_engine.update_prompts({\"pandas_prompt\": new_prompt})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93f95796-2044-4ff4-829c-9c4ff21f924f",
   "metadata": {},
   "source": [
    "This is the instruction string (that you can customize by passing in `instruction_str` on initialization)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62d0fb52-685e-438b-890a-a88c2cb83292",
   "metadata": {},
   "outputs": [],
   "source": [
    "instruction_str = \"\"\"\\\n",
    "1. Convert the query to executable Python code using Pandas.\n",
    "2. The final line of code should be a Python expression that can be called with the `eval()` function.\n",
    "3. The code should represent a solution to the query.\n",
    "4. PRINT ONLY THE EXPRESSION.\n",
    "5. Do not quote the expression.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29c78a8e-d3ff-4c0e-b8c3-7fc79de83839",
   "metadata": {},
   "source": [
    "### Implementing Query Engine using Query Pipeline Syntax\n",
    "\n",
    "If you want to learn to construct your own Pandas Query Engine using our Query Pipeline syntax and the prompt components above, check out our below tutorial.\n",
    "\n",
    "[Setting up a Pandas DataFrame query engine with Query Pipelines](https://docs.llamaindex.ai/en/stable/examples/pipeline/query_pipeline_pandas.html)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}