wehub-resource-sync
a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
990 lines
30 KiB
Plaintext
990 lines
30 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "a0d0efc8-3e03-4f66-8f6d-a907b6b6d7c1",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Recursive Retriever + Query Engine Demo \n",
|
|
"\n",
|
|
"In this demo, we walk through a use case of showcasing our \"RecursiveRetriever\" module over hierarchical data.\n",
|
|
"\n",
|
|
"The concept of recursive retrieval is that we not only explore the directly most relevant nodes, but also explore\n",
|
|
"node relationships to additional retrievers/query engines and execute them. For instance, a node may represent a concise summary of a structured table,\n",
|
|
"and link to a SQL/Pandas query engine over that structured table. Then if the node is retrieved, we want to also query the underlying query engine for the answer.\n",
|
|
"\n",
|
|
"This can be especially useful for documents with hierarchical relationships. In this example, we walk through a Wikipedia article about billionaires (in PDF form), which contains both text and a variety of embedded structured tables. We first create a Pandas query engine over each table, but also represent each table by an `IndexNode` (stores a link to the query engine); this Node is stored along with other Nodes in a vector store. \n",
|
|
"\n",
|
|
"During query-time, if an `IndexNode` is fetched, then the underlying query engine/retriever will be queried. \n",
|
|
"\n",
|
|
"**Notes about Setup**\n",
|
|
"\n",
|
|
"We use `camelot` to extract text-based tables from PDFs."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1a141a3b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index-embeddings-openai\n",
|
|
"%pip install llama-index-readers-file pymupdf\n",
|
|
"%pip install llama-index-llms-openai\n",
|
|
"%pip install llama-index-experimental"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "199ec60c-ea45-46d3-ba21-66db6e16726f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import camelot\n",
|
|
"\n",
|
|
"# https://en.wikipedia.org/wiki/The_World%27s_Billionaires\n",
|
|
"from llama_index.core import VectorStoreIndex\n",
|
|
"from llama_index.experimental.query_engine import PandasQueryEngine\n",
|
|
"from llama_index.core.schema import IndexNode\n",
|
|
"from llama_index.llms.openai import OpenAI\n",
|
|
"\n",
|
|
"from llama_index.readers.file import PyMuPDFReader\n",
|
|
"from typing import List"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "f3980b8a",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Default Settings"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "043e65f3",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"\n",
|
|
"os.environ[\"OPENAI_API_KEY\"] = \"YOUR_API_KEY\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b1d1dcc5",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.embeddings.openai import OpenAIEmbedding\n",
|
|
"from llama_index.llms.openai import OpenAI\n",
|
|
"from llama_index.core import Settings\n",
|
|
"\n",
|
|
"Settings.llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
|
|
"Settings.embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\")"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "c41f1a9c-939c-4a89-866e-557f43fc330b",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Load in Document (and Tables)\n",
|
|
"\n",
|
|
"We use our `PyMuPDFReader` to read in the main text of the document.\n",
|
|
"\n",
|
|
"We also use `camelot` to extract some structured tables from the document"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4435fcb9-bf67-4d76-9d38-f5cd19086fae",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"file_path = \"billionaires_page.pdf\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "e5e78b24-42e0-4d65-980d-13f8c738012f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# initialize PDF reader\n",
|
|
"reader = PyMuPDFReader()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "09976ddc-fe47-4eb6-b577-d5c1c9185974",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"docs = reader.load(file_path)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "5f086ec9-17e3-40cb-b92d-4e121322ff51",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# use camelot to parse tables\n",
|
|
"def get_tables(path: str, pages: List[int]):\n",
|
|
" table_dfs = []\n",
|
|
" for page in pages:\n",
|
|
" table_list = camelot.read_pdf(path, pages=str(page))\n",
|
|
" table_df = table_list[0].df\n",
|
|
" table_df = (\n",
|
|
" table_df.rename(columns=table_df.iloc[0])\n",
|
|
" .drop(table_df.index[0])\n",
|
|
" .reset_index(drop=True)\n",
|
|
" )\n",
|
|
" table_dfs.append(table_df)\n",
|
|
" return table_dfs"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f6653cb1-d0f1-4408-94d3-31827e1a3115",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"table_dfs = get_tables(file_path, pages=[3, 25])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "24e90368-b982-4d05-91b3-0d9dae39eebc",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"<div>\n",
|
|
"<style scoped>\n",
|
|
" .dataframe tbody tr th:only-of-type {\n",
|
|
" vertical-align: middle;\n",
|
|
" }\n",
|
|
"\n",
|
|
" .dataframe tbody tr th {\n",
|
|
" vertical-align: top;\n",
|
|
" }\n",
|
|
"\n",
|
|
" .dataframe thead th {\n",
|
|
" text-align: right;\n",
|
|
" }\n",
|
|
"</style>\n",
|
|
"<table border=\"1\" class=\"dataframe\">\n",
|
|
" <thead>\n",
|
|
" <tr style=\"text-align: right;\">\n",
|
|
" <th></th>\n",
|
|
" <th>No.</th>\n",
|
|
" <th>Name</th>\n",
|
|
" <th>Net worth\\n(USD)</th>\n",
|
|
" <th>Age</th>\n",
|
|
" <th>Nationality</th>\n",
|
|
" <th>Primary source(s) of wealth</th>\n",
|
|
" </tr>\n",
|
|
" </thead>\n",
|
|
" <tbody>\n",
|
|
" <tr>\n",
|
|
" <th>0</th>\n",
|
|
" <td>1</td>\n",
|
|
" <td>Bernard Arnault &\\nfamily</td>\n",
|
|
" <td>$211 billion</td>\n",
|
|
" <td>74</td>\n",
|
|
" <td>France</td>\n",
|
|
" <td>LVMH</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>1</th>\n",
|
|
" <td>2</td>\n",
|
|
" <td>Elon Musk</td>\n",
|
|
" <td>$180 billion</td>\n",
|
|
" <td>51</td>\n",
|
|
" <td>United\\nStates</td>\n",
|
|
" <td>Tesla, SpaceX, X Corp.</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>2</th>\n",
|
|
" <td>3</td>\n",
|
|
" <td>Jeff Bezos</td>\n",
|
|
" <td>$114 billion</td>\n",
|
|
" <td>59</td>\n",
|
|
" <td>United\\nStates</td>\n",
|
|
" <td>Amazon</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>3</th>\n",
|
|
" <td>4</td>\n",
|
|
" <td>Larry Ellison</td>\n",
|
|
" <td>$107 billion</td>\n",
|
|
" <td>78</td>\n",
|
|
" <td>United\\nStates</td>\n",
|
|
" <td>Oracle Corporation</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>4</th>\n",
|
|
" <td>5</td>\n",
|
|
" <td>Warren Buffett</td>\n",
|
|
" <td>$106 billion</td>\n",
|
|
" <td>92</td>\n",
|
|
" <td>United\\nStates</td>\n",
|
|
" <td>Berkshire Hathaway</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>5</th>\n",
|
|
" <td>6</td>\n",
|
|
" <td>Bill Gates</td>\n",
|
|
" <td>$104 billion</td>\n",
|
|
" <td>67</td>\n",
|
|
" <td>United\\nStates</td>\n",
|
|
" <td>Microsoft</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>6</th>\n",
|
|
" <td>7</td>\n",
|
|
" <td>Michael Bloomberg</td>\n",
|
|
" <td>$94.5 billion</td>\n",
|
|
" <td>81</td>\n",
|
|
" <td>United\\nStates</td>\n",
|
|
" <td>Bloomberg L.P.</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>7</th>\n",
|
|
" <td>8</td>\n",
|
|
" <td>Carlos Slim & family</td>\n",
|
|
" <td>$93 billion</td>\n",
|
|
" <td>83</td>\n",
|
|
" <td>Mexico</td>\n",
|
|
" <td>Telmex, América Móvil, Grupo\\nCarso</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>8</th>\n",
|
|
" <td>9</td>\n",
|
|
" <td>Mukesh Ambani</td>\n",
|
|
" <td>$83.4 billion</td>\n",
|
|
" <td>65</td>\n",
|
|
" <td>India</td>\n",
|
|
" <td>Reliance Industries</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>9</th>\n",
|
|
" <td>10</td>\n",
|
|
" <td>Steve Ballmer</td>\n",
|
|
" <td>$80.7 billion</td>\n",
|
|
" <td>67</td>\n",
|
|
" <td>United\\nStates</td>\n",
|
|
" <td>Microsoft</td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table>\n",
|
|
"</div>"
|
|
],
|
|
"text/plain": [
|
|
" No. Name Net worth\\n(USD) Age Nationality \\\n",
|
|
"0 1 Bernard Arnault &\\nfamily $211 billion 74 France \n",
|
|
"1 2 Elon Musk $180 billion 51 United\\nStates \n",
|
|
"2 3 Jeff Bezos $114 billion 59 United\\nStates \n",
|
|
"3 4 Larry Ellison $107 billion 78 United\\nStates \n",
|
|
"4 5 Warren Buffett $106 billion 92 United\\nStates \n",
|
|
"5 6 Bill Gates $104 billion 67 United\\nStates \n",
|
|
"6 7 Michael Bloomberg $94.5 billion 81 United\\nStates \n",
|
|
"7 8 Carlos Slim & family $93 billion 83 Mexico \n",
|
|
"8 9 Mukesh Ambani $83.4 billion 65 India \n",
|
|
"9 10 Steve Ballmer $80.7 billion 67 United\\nStates \n",
|
|
"\n",
|
|
" Primary source(s) of wealth \n",
|
|
"0 LVMH \n",
|
|
"1 Tesla, SpaceX, X Corp. \n",
|
|
"2 Amazon \n",
|
|
"3 Oracle Corporation \n",
|
|
"4 Berkshire Hathaway \n",
|
|
"5 Microsoft \n",
|
|
"6 Bloomberg L.P. \n",
|
|
"7 Telmex, América Móvil, Grupo\\nCarso \n",
|
|
"8 Reliance Industries \n",
|
|
"9 Microsoft "
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"# shows list of top billionaires in 2023\n",
|
|
"table_dfs[0]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0ba98e98-bfb6-4c54-8fa6-b0d7abf381c8",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/html": [
|
|
"<div>\n",
|
|
"<style scoped>\n",
|
|
" .dataframe tbody tr th:only-of-type {\n",
|
|
" vertical-align: middle;\n",
|
|
" }\n",
|
|
"\n",
|
|
" .dataframe tbody tr th {\n",
|
|
" vertical-align: top;\n",
|
|
" }\n",
|
|
"\n",
|
|
" .dataframe thead th {\n",
|
|
" text-align: right;\n",
|
|
" }\n",
|
|
"</style>\n",
|
|
"<table border=\"1\" class=\"dataframe\">\n",
|
|
" <thead>\n",
|
|
" <tr style=\"text-align: right;\">\n",
|
|
" <th></th>\n",
|
|
" <th>Year</th>\n",
|
|
" <th>Number of billionaires</th>\n",
|
|
" <th>Group's combined net worth</th>\n",
|
|
" </tr>\n",
|
|
" </thead>\n",
|
|
" <tbody>\n",
|
|
" <tr>\n",
|
|
" <th>0</th>\n",
|
|
" <td>2023[2]</td>\n",
|
|
" <td>2,640</td>\n",
|
|
" <td>$12.2 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>1</th>\n",
|
|
" <td>2022[6]</td>\n",
|
|
" <td>2,668</td>\n",
|
|
" <td>$12.7 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>2</th>\n",
|
|
" <td>2021[11]</td>\n",
|
|
" <td>2,755</td>\n",
|
|
" <td>$13.1 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>3</th>\n",
|
|
" <td>2020</td>\n",
|
|
" <td>2,095</td>\n",
|
|
" <td>$8.0 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>4</th>\n",
|
|
" <td>2019</td>\n",
|
|
" <td>2,153</td>\n",
|
|
" <td>$8.7 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>5</th>\n",
|
|
" <td>2018</td>\n",
|
|
" <td>2,208</td>\n",
|
|
" <td>$9.1 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>6</th>\n",
|
|
" <td>2017</td>\n",
|
|
" <td>2,043</td>\n",
|
|
" <td>$7.7 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>7</th>\n",
|
|
" <td>2016</td>\n",
|
|
" <td>1,810</td>\n",
|
|
" <td>$6.5 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>8</th>\n",
|
|
" <td>2015[18]</td>\n",
|
|
" <td>1,826</td>\n",
|
|
" <td>$7.1 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>9</th>\n",
|
|
" <td>2014[67]</td>\n",
|
|
" <td>1,645</td>\n",
|
|
" <td>$6.4 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>10</th>\n",
|
|
" <td>2013[68]</td>\n",
|
|
" <td>1,426</td>\n",
|
|
" <td>$5.4 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>11</th>\n",
|
|
" <td>2012</td>\n",
|
|
" <td>1,226</td>\n",
|
|
" <td>$4.6 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>12</th>\n",
|
|
" <td>2011</td>\n",
|
|
" <td>1,210</td>\n",
|
|
" <td>$4.5 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>13</th>\n",
|
|
" <td>2010</td>\n",
|
|
" <td>1,011</td>\n",
|
|
" <td>$3.6 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>14</th>\n",
|
|
" <td>2009</td>\n",
|
|
" <td>793</td>\n",
|
|
" <td>$2.4 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>15</th>\n",
|
|
" <td>2008</td>\n",
|
|
" <td>1,125</td>\n",
|
|
" <td>$4.4 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>16</th>\n",
|
|
" <td>2007</td>\n",
|
|
" <td>946</td>\n",
|
|
" <td>$3.5 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>17</th>\n",
|
|
" <td>2006</td>\n",
|
|
" <td>793</td>\n",
|
|
" <td>$2.6 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>18</th>\n",
|
|
" <td>2005</td>\n",
|
|
" <td>691</td>\n",
|
|
" <td>$2.2 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>19</th>\n",
|
|
" <td>2004</td>\n",
|
|
" <td>587</td>\n",
|
|
" <td>$1.9 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>20</th>\n",
|
|
" <td>2003</td>\n",
|
|
" <td>476</td>\n",
|
|
" <td>$1.4 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>21</th>\n",
|
|
" <td>2002</td>\n",
|
|
" <td>497</td>\n",
|
|
" <td>$1.5 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>22</th>\n",
|
|
" <td>2001</td>\n",
|
|
" <td>538</td>\n",
|
|
" <td>$1.8 trillion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>23</th>\n",
|
|
" <td>2000</td>\n",
|
|
" <td>470</td>\n",
|
|
" <td>$898 billion</td>\n",
|
|
" </tr>\n",
|
|
" <tr>\n",
|
|
" <th>24</th>\n",
|
|
" <td>Sources: Forbes.[18][67][66][68]</td>\n",
|
|
" <td></td>\n",
|
|
" <td></td>\n",
|
|
" </tr>\n",
|
|
" </tbody>\n",
|
|
"</table>\n",
|
|
"</div>"
|
|
],
|
|
"text/plain": [
|
|
" Year Number of billionaires \\\n",
|
|
"0 2023[2] 2,640 \n",
|
|
"1 2022[6] 2,668 \n",
|
|
"2 2021[11] 2,755 \n",
|
|
"3 2020 2,095 \n",
|
|
"4 2019 2,153 \n",
|
|
"5 2018 2,208 \n",
|
|
"6 2017 2,043 \n",
|
|
"7 2016 1,810 \n",
|
|
"8 2015[18] 1,826 \n",
|
|
"9 2014[67] 1,645 \n",
|
|
"10 2013[68] 1,426 \n",
|
|
"11 2012 1,226 \n",
|
|
"12 2011 1,210 \n",
|
|
"13 2010 1,011 \n",
|
|
"14 2009 793 \n",
|
|
"15 2008 1,125 \n",
|
|
"16 2007 946 \n",
|
|
"17 2006 793 \n",
|
|
"18 2005 691 \n",
|
|
"19 2004 587 \n",
|
|
"20 2003 476 \n",
|
|
"21 2002 497 \n",
|
|
"22 2001 538 \n",
|
|
"23 2000 470 \n",
|
|
"24 Sources: Forbes.[18][67][66][68] \n",
|
|
"\n",
|
|
" Group's combined net worth \n",
|
|
"0 $12.2 trillion \n",
|
|
"1 $12.7 trillion \n",
|
|
"2 $13.1 trillion \n",
|
|
"3 $8.0 trillion \n",
|
|
"4 $8.7 trillion \n",
|
|
"5 $9.1 trillion \n",
|
|
"6 $7.7 trillion \n",
|
|
"7 $6.5 trillion \n",
|
|
"8 $7.1 trillion \n",
|
|
"9 $6.4 trillion \n",
|
|
"10 $5.4 trillion \n",
|
|
"11 $4.6 trillion \n",
|
|
"12 $4.5 trillion \n",
|
|
"13 $3.6 trillion \n",
|
|
"14 $2.4 trillion \n",
|
|
"15 $4.4 trillion \n",
|
|
"16 $3.5 trillion \n",
|
|
"17 $2.6 trillion \n",
|
|
"18 $2.2 trillion \n",
|
|
"19 $1.9 trillion \n",
|
|
"20 $1.4 trillion \n",
|
|
"21 $1.5 trillion \n",
|
|
"22 $1.8 trillion \n",
|
|
"23 $898 billion \n",
|
|
"24 "
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"# shows list of top billionaires\n",
|
|
"table_dfs[1]"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "02967be4-be85-4046-b4e2-e5dd8e65628a",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Create Pandas Query Engines\n",
|
|
"\n",
|
|
"We create a pandas query engine over each structured table.\n",
|
|
"\n",
|
|
"These can be executed on their own to answer queries about each table.\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."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "2f38bbf0-502a-44f7-b33b-443fcd90583a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# define query engines over these tables\n",
|
|
"llm = OpenAI(model=\"gpt-4\")\n",
|
|
"\n",
|
|
"df_query_engines = [\n",
|
|
" PandasQueryEngine(table_df, llm=llm) for table_df in table_dfs\n",
|
|
"]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9fbc43a0-5655-4c8a-80d4-71c7fe9d7275",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"$180 billion\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"response = df_query_engines[0].query(\n",
|
|
" \"What's the net worth of the second richest billionaire in 2023?\"\n",
|
|
")\n",
|
|
"print(str(response))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "78dbe39d-eea0-4c7f-bea7-2c0f6f6591cd",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"793\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"response = df_query_engines[1].query(\n",
|
|
" \"How many billionaires were there in 2009?\"\n",
|
|
")\n",
|
|
"print(str(response))"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "e9bca53c-8766-42e4-96c7-145e7f14be34",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Build Vector Index\n",
|
|
"\n",
|
|
"Build vector index over the chunked document as well as over the additional `IndexNode` objects linked to the tables."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6a8b6605-823a-4607-be0e-99c67d5a90ff",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core import Settings\n",
|
|
"\n",
|
|
"doc_nodes = Settings.node_parser.get_nodes_from_documents(docs)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6a3d249c-2242-4158-88ea-d16b67815107",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# define index nodes\n",
|
|
"summaries = [\n",
|
|
" (\n",
|
|
" \"This node provides information about the world's richest billionaires\"\n",
|
|
" \" in 2023\"\n",
|
|
" ),\n",
|
|
" (\n",
|
|
" \"This node provides information on the number of billionaires and\"\n",
|
|
" \" their combined net worth from 2000 to 2023.\"\n",
|
|
" ),\n",
|
|
"]\n",
|
|
"\n",
|
|
"df_nodes = [\n",
|
|
" IndexNode(text=summary, index_id=f\"pandas{idx}\")\n",
|
|
" for idx, summary in enumerate(summaries)\n",
|
|
"]\n",
|
|
"\n",
|
|
"df_id_query_engine_mapping = {\n",
|
|
" f\"pandas{idx}\": df_query_engine\n",
|
|
" for idx, df_query_engine in enumerate(df_query_engines)\n",
|
|
"}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "abf87458-cd3d-4ac4-b934-20ee8a9a820a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# construct top-level vector index + query engine\n",
|
|
"vector_index = VectorStoreIndex(doc_nodes + df_nodes)\n",
|
|
"vector_retriever = vector_index.as_retriever(similarity_top_k=1)"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "6ed917f6-e407-4ebb-9c15-36aedd207c6f",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Use `RecursiveRetriever` in our `RetrieverQueryEngine`\n",
|
|
"\n",
|
|
"We define a `RecursiveRetriever` object to recursively retrieve/query nodes. We then put this in our `RetrieverQueryEngine` along with a `ResponseSynthesizer` to synthesize a response.\n",
|
|
"\n",
|
|
"We pass in mappings from id to retriever and id to query engine. We then pass in a root id representing the retriever we query first."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9a7d6031-f7a0-45c2-9a84-813b4e3fcf28",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# baseline vector index (that doesn't include the extra df nodes).\n",
|
|
"# used to benchmark\n",
|
|
"vector_index0 = VectorStoreIndex(doc_nodes)\n",
|
|
"vector_query_engine0 = vector_index0.as_query_engine()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "e38d1e90-7d83-46bb-99d7-0892aef4d3ca",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.retrievers import RecursiveRetriever\n",
|
|
"from llama_index.core.query_engine import RetrieverQueryEngine\n",
|
|
"from llama_index.core import get_response_synthesizer\n",
|
|
"\n",
|
|
"recursive_retriever = RecursiveRetriever(\n",
|
|
" \"vector\",\n",
|
|
" retriever_dict={\"vector\": vector_retriever},\n",
|
|
" query_engine_dict=df_id_query_engine_mapping,\n",
|
|
" verbose=True,\n",
|
|
")\n",
|
|
"\n",
|
|
"response_synthesizer = get_response_synthesizer(response_mode=\"compact\")\n",
|
|
"\n",
|
|
"query_engine = RetrieverQueryEngine.from_args(\n",
|
|
" recursive_retriever, response_synthesizer=response_synthesizer\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6af6823d-bd07-4088-8422-bd9aa3224b08",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\u001b[36;1m\u001b[1;3mRetrieving with query id None: What's the net worth of the second richest billionaire in 2023?\n",
|
|
"\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: pandas0\n",
|
|
"\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id pandas0: What's the net worth of the second richest billionaire in 2023?\n",
|
|
"\u001b[0m\u001b[32;1m\u001b[1;3mGot response: $180 billion\n",
|
|
"\u001b[0m"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"response = query_engine.query(\n",
|
|
" \"What's the net worth of the second richest billionaire in 2023?\"\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6122ff8e-f84d-47cc-b363-44621e0623ab",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"\"Query: What's the net worth of the second richest billionaire in 2023?\\nResponse: $180\\xa0billion\""
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"response.source_nodes[0].node.get_content()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4096f30c-c6a3-4f74-b496-1240fdc08fd4",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"'$180 billion.'"
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"str(response)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "5cf133a7-2532-4179-af21-d495eb547083",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\u001b[36;1m\u001b[1;3mRetrieving with query id None: How many billionaires were there in 2009?\n",
|
|
"\u001b[0m\u001b[38;5;200m\u001b[1;3mRetrieved node with id, entering: pandas1\n",
|
|
"\u001b[0m\u001b[36;1m\u001b[1;3mRetrieving with query id pandas1: How many billionaires were there in 2009?\n",
|
|
"\u001b[0m\u001b[32;1m\u001b[1;3mGot response: 793\n",
|
|
"\u001b[0m"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"response = query_engine.query(\"How many billionaires were there in 2009?\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "489e4c59-a40c-47b8-a788-e80558ce7e3a",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"'793'"
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"str(response)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "551781f3-8761-4385-966d-a0a6b0526ed6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"response = vector_query_engine0.query(\n",
|
|
" \"How many billionaires were there in 2009?\"\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3656db61-32bc-49c6-9774-8439323358a5",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"print(response.source_nodes[0].node.get_content())"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "89b47692-95f0-4077-9439-d8af237e5c16",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Based on the context information, it is not possible to determine the exact number of billionaires in 2009. The provided information only mentions the number of billionaires in 2013 and 2014.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(str(response))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1fbe211d-d59b-4f87-897c-0fa03e42641e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"response.source_nodes[0].node.get_content()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "960c81f6-b7a6-43ac-aa4b-9ef20a5400b1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"response = query_engine.query(\n",
|
|
" \"Which billionaires are excluded from this list?\"\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "cec5f5a8-b9dc-464c-86a4-682a23f90f2b",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Royal families and dictators whose wealth is contingent on a position are excluded from this list.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(str(response))"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": ".venv",
|
|
"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"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|