wehub/run-llama--llama_index
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Build Package / build (ubuntu-latest) (push) Failing after 1s
Details
CodeQL / Analyze (python) (push) Failing after 1s
Details
Core Typecheck / core-typecheck (push) Failing after 1s
Details
Linting / lint (push) Failing after 1s
Details
llama-dev tests / test-llama-dev (push) Failing after 1s
Details
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Details
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Details
Build Package / build (windows-latest) (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:26:52 +08:00
commit a0c8464e58
10008 changed files with 4625987 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/examples/query_engine/SQLRouterQueryEngine.ipynb
+554
View File
@@ -0,0 +1,554 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/query_engine/SQLRouterQueryEngine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"# SQL Router Query Engine\n",
"\n",
"In this tutorial, we define a custom router query engine that can route to either a SQL database or a vector database.\n",
"\n",
"**NOTE:** Any Text-to-SQL application should be aware that executing \n",
"arbitrary SQL queries can be a security risk. It is recommended to\n",
"take precautions as needed, such as using restricted roles, read-only\n",
"databases, sandboxing, etc."
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Setup"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%pip install llama-index-readers-wikipedia"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pip install llama-index"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# NOTE: This is ONLY necessary in jupyter notebook.\n",
"# Details: Jupyter runs an event-loop behind the scenes.\n",
"# This results in nested event-loops when we start an event-loop to make async queries.\n",
"# This is normally not allowed, we use nest_asyncio to allow it for convenience.\n",
"import nest_asyncio\n",
"\n",
"nest_asyncio.apply()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO:numexpr.utils:Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
"Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
"INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n",
"NumExpr defaulting to 8 threads.\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n"
]
}
],
"source": [
"import logging\n",
"import sys\n",
"\n",
"logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
"logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
"\n",
"from llama_index.core import VectorStoreIndex, SQLDatabase\n",
"from llama_index.readers.wikipedia import WikipediaReader"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Create Database Schema + Test Data\n",
"\n",
"Here we introduce a toy scenario where there are 100 tables (too big to fit into the prompt)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from sqlalchemy import (\n",
" create_engine,\n",
" MetaData,\n",
" Table,\n",
" Column,\n",
" String,\n",
" Integer,\n",
" select,\n",
" column,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"engine = create_engine(\"sqlite:///:memory:\", future=True)\n",
"metadata_obj = MetaData()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# create city SQL table\n",
"table_name = \"city_stats\"\n",
"city_stats_table = Table(\n",
" table_name,\n",
" metadata_obj,\n",
" Column(\"city_name\", String(16), primary_key=True),\n",
" Column(\"population\", Integer),\n",
" Column(\"country\", String(16), nullable=False),\n",
")\n",
"\n",
"metadata_obj.create_all(engine)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"dict_keys(['city_stats'])"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# print tables\n",
"metadata_obj.tables.keys()"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"We introduce some test data into the `city_stats` table"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from sqlalchemy import insert\n",
"\n",
"rows = [\n",
" {\"city_name\": \"Toronto\", \"population\": 2930000, \"country\": \"Canada\"},\n",
" {\"city_name\": \"Tokyo\", \"population\": 13960000, \"country\": \"Japan\"},\n",
" {\"city_name\": \"Berlin\", \"population\": 3645000, \"country\": \"Germany\"},\n",
"]\n",
"for row in rows:\n",
" stmt = insert(city_stats_table).values(**row)\n",
" with engine.begin() as connection:\n",
" cursor = connection.execute(stmt)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[('Toronto', 2930000, 'Canada'), ('Tokyo', 13960000, 'Japan'), ('Berlin', 3645000, 'Germany')]\n"
]
}
],
"source": [
"with engine.connect() as connection:\n",
" cursor = connection.exec_driver_sql(\"SELECT * FROM city_stats\")\n",
" print(cursor.fetchall())"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Load Data\n",
"\n",
"We first show how to convert a Document into a set of Nodes, and insert into a DocumentStore."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Requirement already satisfied: wikipedia in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (1.4.0)\n",
"Requirement already satisfied: requests<3.0.0,>=2.0.0 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from wikipedia) (2.28.2)\n",
"Requirement already satisfied: beautifulsoup4 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from wikipedia) (4.12.2)\n",
"Requirement already satisfied: idna<4,>=2.5 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (3.4)\n",
"Requirement already satisfied: charset-normalizer<4,>=2 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (3.1.0)\n",
"Requirement already satisfied: certifi>=2017.4.17 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (2022.12.7)\n",
"Requirement already satisfied: urllib3<1.27,>=1.21.1 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (1.26.15)\n",
"Requirement already satisfied: soupsieve>1.2 in /Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages (from beautifulsoup4->wikipedia) (2.4.1)\n",
"\n",
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.2\u001b[0m\n",
"\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
]
}
],
"source": [
"# install wikipedia python package\n",
"!pip install wikipedia"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"cities = [\"Toronto\", \"Berlin\", \"Tokyo\"]\n",
"wiki_docs = WikipediaReader().load_data(pages=cities)"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Build SQL Index"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"sql_database = SQLDatabase(engine, include_tables=[\"city_stats\"])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core.query_engine import NLSQLTableQueryEngine"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
"> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 0 tokens\n",
"> [build_index_from_nodes] Total embedding token usage: 0 tokens\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages/langchain/sql_database.py:227: UserWarning: This method is deprecated - please use `get_usable_table_names`.\n",
" warnings.warn(\n"
]
}
],
"source": [
"sql_query_engine = NLSQLTableQueryEngine(\n",
" sql_database=sql_database,\n",
" tables=[\"city_stats\"],\n",
")"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Build Vector Index"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
"> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 20744 tokens\n",
"> [build_index_from_nodes] Total embedding token usage: 20744 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
"> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 21947 tokens\n",
"> [build_index_from_nodes] Total embedding token usage: 21947 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
"> [build_index_from_nodes] Total LLM token usage: 0 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 12786 tokens\n",
"> [build_index_from_nodes] Total embedding token usage: 12786 tokens\n"
]
}
],
"source": [
"# build a separate vector index per city\n",
"# You could also choose to define a single vector index across all docs, and annotate each chunk by metadata\n",
"vector_indices = []\n",
"for wiki_doc in wiki_docs:\n",
" vector_index = VectorStoreIndex.from_documents([wiki_doc])\n",
" vector_indices.append(vector_index)"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Define Query Engines, Set as Tools"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"vector_query_engines = [index.as_query_engine() for index in vector_indices]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core.tools import QueryEngineTool\n",
"\n",
"\n",
"sql_tool = QueryEngineTool.from_defaults(\n",
" query_engine=sql_query_engine,\n",
" description=(\n",
" \"Useful for translating a natural language query into a SQL query over\"\n",
" \" a table containing: city_stats, containing the population/country of\"\n",
" \" each city\"\n",
" ),\n",
")\n",
"vector_tools = []\n",
"for city, query_engine in zip(cities, vector_query_engines):\n",
" vector_tool = QueryEngineTool.from_defaults(\n",
" query_engine=query_engine,\n",
" description=f\"Useful for answering semantic questions about {city}\",\n",
" )\n",
" vector_tools.append(vector_tool)"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Define Router Query Engine"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core.query_engine import RouterQueryEngine\n",
"from llama_index.core.selectors import LLMSingleSelector\n",
"\n",
"query_engine = RouterQueryEngine(\n",
" selector=LLMSingleSelector.from_defaults(),\n",
" query_engine_tools=([sql_tool] + vector_tools),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO:llama_index.query_engine.router_query_engine:Selecting query engine 0: Useful for translating a natural language query into a SQL query over a table containing: city_stats, containing the population/country of each city.\n",
"Selecting query engine 0: Useful for translating a natural language query into a SQL query over a table containing: city_stats, containing the population/country of each city.\n",
"INFO:llama_index.indices.struct_store.sql_query:> Table desc str: Schema of table city_stats:\n",
"Table 'city_stats' has columns: city_name (VARCHAR(16)), population (INTEGER), country (VARCHAR(16)) and foreign keys: .\n",
"\n",
"> Table desc str: Schema of table city_stats:\n",
"Table 'city_stats' has columns: city_name (VARCHAR(16)), population (INTEGER), country (VARCHAR(16)) and foreign keys: .\n",
"\n",
"INFO:llama_index.token_counter.token_counter:> [query] Total LLM token usage: 347 tokens\n",
"> [query] Total LLM token usage: 347 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [query] Total embedding token usage: 0 tokens\n",
"> [query] Total embedding token usage: 0 tokens\n",
" Tokyo has the highest population, with 13,960,000 people.\n"
]
}
],
"source": [
"response = query_engine.query(\"Which city has the highest population?\")\n",
"print(str(response))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO:llama_index.query_engine.router_query_engine:Selecting query engine 2: Useful for answering semantic questions about Berlin.\n",
"Selecting query engine 2: Useful for answering semantic questions about Berlin.\n",
"INFO:llama_index.token_counter.token_counter:> [retrieve] Total LLM token usage: 0 tokens\n",
"> [retrieve] Total LLM token usage: 0 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [retrieve] Total embedding token usage: 8 tokens\n",
"> [retrieve] Total embedding token usage: 8 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [get_response] Total LLM token usage: 2031 tokens\n",
"> [get_response] Total LLM token usage: 2031 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [get_response] Total embedding token usage: 0 tokens\n",
"> [get_response] Total embedding token usage: 0 tokens\n",
"\n",
"Berlin is home to many historical museums, including the Altes Museum, Neues Museum, Alte Nationalgalerie, Pergamon Museum, and Bode Museum, which are all located on Museum Island. The Gemäldegalerie (Painting Gallery) focuses on the paintings of the \"old masters\" from the 13th to the 18th centuries, while the Neue Nationalgalerie (New National Gallery, built by Ludwig Mies van der Rohe) specializes in 20th-century European painting. The Hamburger Bahnhof, in Moabit, exhibits a major collection of modern and contemporary art. The expanded Deutsches Historisches Museum reopened in the Zeughaus with an overview of German history spanning more than a millennium. The Bauhaus Archive is a museum of 20th-century design from the famous Bauhaus school. Museum Berggruen houses the collection of noted 20th century collector Heinz Berggruen, and features an extensive assortment of works by Picasso, Matisse, Cézanne, and Giacometti, among others. The Kupferstichkabinett Berlin (Museum of Prints and Drawings) is part of the Staatlichen Museen z\n"
]
}
],
"source": [
"response = query_engine.query(\"Tell me about the historical museums in Berlin\")\n",
"print(str(response))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"INFO:llama_index.query_engine.router_query_engine:Selecting query engine 0: Useful for translating a natural language query into a SQL query over a table containing: city_stats, containing the population/country of each city.\n",
"Selecting query engine 0: Useful for translating a natural language query into a SQL query over a table containing: city_stats, containing the population/country of each city.\n",
"INFO:llama_index.indices.struct_store.sql_query:> Table desc str: Schema of table city_stats:\n",
"Table 'city_stats' has columns: city_name (VARCHAR(16)), population (INTEGER), country (VARCHAR(16)) and foreign keys: .\n",
"\n",
"> Table desc str: Schema of table city_stats:\n",
"Table 'city_stats' has columns: city_name (VARCHAR(16)), population (INTEGER), country (VARCHAR(16)) and foreign keys: .\n",
"\n",
"INFO:llama_index.token_counter.token_counter:> [query] Total LLM token usage: 334 tokens\n",
"> [query] Total LLM token usage: 334 tokens\n",
"INFO:llama_index.token_counter.token_counter:> [query] Total embedding token usage: 0 tokens\n",
"> [query] Total embedding token usage: 0 tokens\n",
" Toronto is from Canada, Tokyo is from Japan, and Berlin is from Germany.\n"
]
}
],
"source": [
"response = query_engine.query(\"Which countries are each city from?\")\n",
"print(str(response))"
]
}
],
"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": 4
}