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
1123 lines
41 KiB
Plaintext
1123 lines
41 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "5fb96106-e662-4a0f-9673-843c685c111e",
|
||
"metadata": {},
|
||
"source": [
|
||
"# Workflows for Advanced Text-to-SQL\n",
|
||
"\n",
|
||
"<a href=\"https://colab.research.google.com/github/jerryjliu/llama_index/blob/main/docs/examples/workflow/advanced_text_to_sql.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
|
||
"\n",
|
||
"In this guide we show you how to setup a text-to-SQL workflow over your data with our [workflows](https://docs.llamaindex.ai/en/stable/module_guides/workflow/#workflows) syntax.\n",
|
||
"\n",
|
||
"This gives you flexibility to enhance text-to-SQL with additional techniques. We show these in the below sections: \n",
|
||
"1. **Query-Time Table Retrieval**: Dynamically retrieve relevant tables in the text-to-SQL prompt.\n",
|
||
"2. **Query-Time Sample Row retrieval**: Embed/Index each row, and dynamically retrieve example rows for each table in the text-to-SQL prompt.\n",
|
||
"\n",
|
||
"Our out-of-the box workflows include our `NLSQLTableQueryEngine` and `SQLTableRetrieverQueryEngine`. (if you want to check out our text-to-SQL guide using these modules, take a look [here](https://developers.llamaindex.ai/python/examples/index_structs/struct_indices/sqlindexdemo/)). This guide implements an advanced version of those modules, giving you the utmost flexibility to apply this to your own setting.\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."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "d5c0cbbb-0d20-4233-b2f6-9e8ff3f6d04c",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Load and Ingest Data\n",
|
||
"\n",
|
||
"\n",
|
||
"### Load Data\n",
|
||
"We use the [WikiTableQuestions dataset](https://ppasupat.github.io/WikiTableQuestions/) (Pasupat and Liang 2015) as our test dataset.\n",
|
||
"\n",
|
||
"We go through all the csv's in one folder, store each in a sqlite database (we will then build an object index over each table schema)."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "dae8fb42-d65a-456c-b5bf-613efa43e9f1",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"%pip install llama-index-llms-openai"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "468a4900-a1ba-471d-a839-8cf56dfb5cf1",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"!wget \"https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip\" -O data.zip\n",
|
||
"!unzip data.zip"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "cba67be2-343b-4acd-b128-b26a79bf4504",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import pandas as pd\n",
|
||
"from pathlib import Path\n",
|
||
"\n",
|
||
"data_dir = Path(\"./WikiTableQuestions/csv/200-csv\")\n",
|
||
"csv_files = sorted([f for f in data_dir.glob(\"*.csv\")])\n",
|
||
"dfs = []\n",
|
||
"for csv_file in csv_files:\n",
|
||
" print(f\"processing file: {csv_file}\")\n",
|
||
" try:\n",
|
||
" df = pd.read_csv(csv_file)\n",
|
||
" dfs.append(df)\n",
|
||
" except Exception as e:\n",
|
||
" print(f\"Error parsing {csv_file}: {str(e)}\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "2bf32f7b-dad2-4aca-88b1-1ec0bc221ebb",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Extract Table Name and Summary from each Table\n",
|
||
"\n",
|
||
"Here we use gpt-4o-mini to extract a table name (with underscores) and summary from each table with our Pydantic program."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "0c3040ec-9523-492b-870b-754feffca0c7",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"mkdir: WikiTableQuestions_TableInfo: File exists\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"tableinfo_dir = \"WikiTableQuestions_TableInfo\"\n",
|
||
"!mkdir {tableinfo_dir}"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e921b5d0-662b-4991-b904-090f6cf3eb8c",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.prompts import ChatPromptTemplate\n",
|
||
"from llama_index.core.bridge.pydantic import BaseModel, Field\n",
|
||
"from llama_index.llms.openai import OpenAI\n",
|
||
"from llama_index.core.llms import ChatMessage\n",
|
||
"\n",
|
||
"\n",
|
||
"class TableInfo(BaseModel):\n",
|
||
" \"\"\"Information regarding a structured table.\"\"\"\n",
|
||
"\n",
|
||
" table_name: str = Field(\n",
|
||
" ..., description=\"table name (must be underscores and NO spaces)\"\n",
|
||
" )\n",
|
||
" table_summary: str = Field(\n",
|
||
" ..., description=\"short, concise summary/caption of the table\"\n",
|
||
" )\n",
|
||
"\n",
|
||
"\n",
|
||
"prompt_str = \"\"\"\\\n",
|
||
"Give me a summary of the table with the following JSON format.\n",
|
||
"\n",
|
||
"- The table name must be unique to the table and describe it while being concise. \n",
|
||
"- Do NOT output a generic table name (e.g. table, my_table).\n",
|
||
"\n",
|
||
"Do NOT make the table name one of the following: {exclude_table_name_list}\n",
|
||
"\n",
|
||
"Table:\n",
|
||
"{table_str}\n",
|
||
"\n",
|
||
"Summary: \"\"\"\n",
|
||
"prompt_tmpl = ChatPromptTemplate(\n",
|
||
" message_templates=[ChatMessage.from_str(prompt_str, role=\"user\")]\n",
|
||
")\n",
|
||
"\n",
|
||
"llm = OpenAI(model=\"gpt-4o-mini\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "15277c5d-9a12-4ae9-b7ef-de58139ea66b",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import json\n",
|
||
"\n",
|
||
"\n",
|
||
"def _get_tableinfo_with_index(idx: int) -> str:\n",
|
||
" results_gen = Path(tableinfo_dir).glob(f\"{idx}_*\")\n",
|
||
" results_list = list(results_gen)\n",
|
||
" if len(results_list) == 0:\n",
|
||
" return None\n",
|
||
" elif len(results_list) == 1:\n",
|
||
" path = results_list[0]\n",
|
||
" return TableInfo.parse_file(path)\n",
|
||
" else:\n",
|
||
" raise ValueError(\n",
|
||
" f\"More than one file matching index: {list(results_gen)}\"\n",
|
||
" )\n",
|
||
"\n",
|
||
"\n",
|
||
"table_names = set()\n",
|
||
"table_infos = []\n",
|
||
"for idx, df in enumerate(dfs):\n",
|
||
" table_info = _get_tableinfo_with_index(idx)\n",
|
||
" if table_info:\n",
|
||
" table_infos.append(table_info)\n",
|
||
" else:\n",
|
||
" while True:\n",
|
||
" df_str = df.head(10).to_csv()\n",
|
||
" table_info = llm.structured_predict(\n",
|
||
" TableInfo,\n",
|
||
" prompt_tmpl,\n",
|
||
" table_str=df_str,\n",
|
||
" exclude_table_name_list=str(list(table_names)),\n",
|
||
" )\n",
|
||
" table_name = table_info.table_name\n",
|
||
" print(f\"Processed table: {table_name}\")\n",
|
||
" if table_name not in table_names:\n",
|
||
" table_names.add(table_name)\n",
|
||
" break\n",
|
||
" else:\n",
|
||
" # try again\n",
|
||
" print(f\"Table name {table_name} already exists, trying again.\")\n",
|
||
" pass\n",
|
||
"\n",
|
||
" out_file = f\"{tableinfo_dir}/{idx}_{table_name}.json\"\n",
|
||
" json.dump(table_info.dict(), open(out_file, \"w\"))\n",
|
||
" table_infos.append(table_info)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "3f67987b-6906-42e0-b80c-3ba096b71b42",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Put Data in SQL Database\n",
|
||
"\n",
|
||
"We use `sqlalchemy`, a popular SQL database toolkit, to load all the tables."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "8d315556-e83a-4181-935a-1f067b12052a",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# put data into sqlite db\n",
|
||
"from sqlalchemy import (\n",
|
||
" create_engine,\n",
|
||
" MetaData,\n",
|
||
" Table,\n",
|
||
" Column,\n",
|
||
" String,\n",
|
||
" Integer,\n",
|
||
")\n",
|
||
"import re\n",
|
||
"\n",
|
||
"\n",
|
||
"# Function to create a sanitized column name\n",
|
||
"def sanitize_column_name(col_name):\n",
|
||
" # Remove special characters and replace spaces with underscores\n",
|
||
" return re.sub(r\"\\W+\", \"_\", col_name)\n",
|
||
"\n",
|
||
"\n",
|
||
"# Function to create a table from a DataFrame using SQLAlchemy\n",
|
||
"def create_table_from_dataframe(\n",
|
||
" df: pd.DataFrame, table_name: str, engine, metadata_obj\n",
|
||
"):\n",
|
||
" # Sanitize column names\n",
|
||
" sanitized_columns = {col: sanitize_column_name(col) for col in df.columns}\n",
|
||
" df = df.rename(columns=sanitized_columns)\n",
|
||
"\n",
|
||
" # Dynamically create columns based on DataFrame columns and data types\n",
|
||
" columns = [\n",
|
||
" Column(col, String if dtype == \"object\" else Integer)\n",
|
||
" for col, dtype in zip(df.columns, df.dtypes)\n",
|
||
" ]\n",
|
||
"\n",
|
||
" # Create a table with the defined columns\n",
|
||
" table = Table(table_name, metadata_obj, *columns)\n",
|
||
"\n",
|
||
" # Create the table in the database\n",
|
||
" metadata_obj.create_all(engine)\n",
|
||
"\n",
|
||
" # Insert data from DataFrame into the table\n",
|
||
" with engine.connect() as conn:\n",
|
||
" for _, row in df.iterrows():\n",
|
||
" insert_stmt = table.insert().values(**row.to_dict())\n",
|
||
" conn.execute(insert_stmt)\n",
|
||
" conn.commit()\n",
|
||
"\n",
|
||
"\n",
|
||
"# engine = create_engine(\"sqlite:///:memory:\")\n",
|
||
"engine = create_engine(\"sqlite:///wiki_table_questions.db\")\n",
|
||
"metadata_obj = MetaData()\n",
|
||
"for idx, df in enumerate(dfs):\n",
|
||
" tableinfo = _get_tableinfo_with_index(idx)\n",
|
||
" print(f\"Creating table: {tableinfo.table_name}\")\n",
|
||
" create_table_from_dataframe(df, tableinfo.table_name, engine, metadata_obj)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "982aa74e-323f-4c85-b0e2-1acf91d6b841",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# # setup Arize Phoenix for logging/observability\n",
|
||
"# import phoenix as px\n",
|
||
"# import llama_index.core\n",
|
||
"\n",
|
||
"# px.launch_app()\n",
|
||
"# llama_index.core.set_global_handler(\"arize_phoenix\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "f5492e07-234e-4b08-bf37-e4fac422927b",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Advanced Capability 1: Text-to-SQL with Query-Time Table Retrieval.\n",
|
||
"\n",
|
||
"We now show you how to setup an e2e text-to-SQL with table retrieval.\n",
|
||
"\n",
|
||
"### Define Modules\n",
|
||
"\n",
|
||
"Here we define the core modules.\n",
|
||
"1. Object index + retriever to store table schemas\n",
|
||
"2. SQLDatabase object to connect to the above tables + SQLRetriever.\n",
|
||
"3. Text-to-SQL Prompt\n",
|
||
"4. Response synthesis Prompt\n",
|
||
"5. LLM"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "926c03f4-0934-4e47-a11f-a7722e66feec",
|
||
"metadata": {},
|
||
"source": [
|
||
"Object index, retriever, SQLDatabase"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "709f0692-ddd1-4ff6-975a-1df5abfea1ab",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.objects import (\n",
|
||
" SQLTableNodeMapping,\n",
|
||
" ObjectIndex,\n",
|
||
" SQLTableSchema,\n",
|
||
")\n",
|
||
"from llama_index.core import SQLDatabase, VectorStoreIndex\n",
|
||
"\n",
|
||
"sql_database = SQLDatabase(engine)\n",
|
||
"\n",
|
||
"table_node_mapping = SQLTableNodeMapping(sql_database)\n",
|
||
"table_schema_objs = [\n",
|
||
" SQLTableSchema(table_name=t.table_name, context_str=t.table_summary)\n",
|
||
" for t in table_infos\n",
|
||
"] # add a SQLTableSchema for each table\n",
|
||
"\n",
|
||
"obj_index = ObjectIndex.from_objects(\n",
|
||
" table_schema_objs,\n",
|
||
" table_node_mapping,\n",
|
||
" VectorStoreIndex,\n",
|
||
")\n",
|
||
"obj_retriever = obj_index.as_retriever(similarity_top_k=3)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "cf8ce73a-5f21-4977-87ee-7752da4fed63",
|
||
"metadata": {},
|
||
"source": [
|
||
"SQLRetriever + Table Parser"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "3047347e-362e-46c2-9ac8-ee3b8730fe3c",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.retrievers import SQLRetriever\n",
|
||
"from typing import List\n",
|
||
"\n",
|
||
"sql_retriever = SQLRetriever(sql_database)\n",
|
||
"\n",
|
||
"\n",
|
||
"def get_table_context_str(table_schema_objs: List[SQLTableSchema]):\n",
|
||
" \"\"\"Get table context string.\"\"\"\n",
|
||
" context_strs = []\n",
|
||
" for table_schema_obj in table_schema_objs:\n",
|
||
" table_info = sql_database.get_single_table_info(\n",
|
||
" table_schema_obj.table_name\n",
|
||
" )\n",
|
||
" if table_schema_obj.context_str:\n",
|
||
" table_opt_context = \" The table description is: \"\n",
|
||
" table_opt_context += table_schema_obj.context_str\n",
|
||
" table_info += table_opt_context\n",
|
||
"\n",
|
||
" context_strs.append(table_info)\n",
|
||
" return \"\\n\\n\".join(context_strs)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "d5e693fd-18b4-4d3e-9877-7ca527d83f57",
|
||
"metadata": {},
|
||
"source": [
|
||
"Text-to-SQL Prompt + Output Parser"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "5c708fba-0551-4e14-938d-5c32eae52595",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer. You can order the results by a relevant column to return the most interesting examples in the database.\n",
|
||
"\n",
|
||
"Never query for all the columns from a specific table, only ask for a few relevant columns given the question.\n",
|
||
"\n",
|
||
"Pay attention to use only the column names that you can see in the schema description. Be careful to not query for columns that do not exist. Pay attention to which column is in which table. Also, qualify column names with the table name when needed. You are required to use the following format, each taking one line:\n",
|
||
"\n",
|
||
"Question: Question here\n",
|
||
"SQLQuery: SQL Query to run\n",
|
||
"SQLResult: Result of the SQLQuery\n",
|
||
"Answer: Final answer here\n",
|
||
"\n",
|
||
"Only use tables listed below.\n",
|
||
"{schema}\n",
|
||
"\n",
|
||
"Question: {query_str}\n",
|
||
"SQLQuery: \n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from llama_index.core.prompts.default_prompts import DEFAULT_TEXT_TO_SQL_PROMPT\n",
|
||
"from llama_index.core import PromptTemplate\n",
|
||
"from llama_index.core.llms import ChatResponse\n",
|
||
"\n",
|
||
"\n",
|
||
"def parse_response_to_sql(chat_response: ChatResponse) -> str:\n",
|
||
" \"\"\"Parse response to SQL.\"\"\"\n",
|
||
" response = chat_response.message.content\n",
|
||
" sql_query_start = response.find(\"SQLQuery:\")\n",
|
||
" if sql_query_start != -1:\n",
|
||
" response = response[sql_query_start:]\n",
|
||
" # TODO: move to removeprefix after Python 3.9+\n",
|
||
" if response.startswith(\"SQLQuery:\"):\n",
|
||
" response = response[len(\"SQLQuery:\") :]\n",
|
||
" sql_result_start = response.find(\"SQLResult:\")\n",
|
||
" if sql_result_start != -1:\n",
|
||
" response = response[:sql_result_start]\n",
|
||
" return response.strip().strip(\"```\").strip()\n",
|
||
"\n",
|
||
"\n",
|
||
"text2sql_prompt = DEFAULT_TEXT_TO_SQL_PROMPT.partial_format(\n",
|
||
" dialect=engine.dialect.name\n",
|
||
")\n",
|
||
"print(text2sql_prompt.template)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "bf67d6a9-f94c-4c27-b09b-2dbe1952727d",
|
||
"metadata": {},
|
||
"source": [
|
||
"Response Synthesis Prompt"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e1757f33-f1ca-4a31-8cce-36be9bb9a989",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"response_synthesis_prompt_str = (\n",
|
||
" \"Given an input question, synthesize a response from the query results.\\n\"\n",
|
||
" \"Query: {query_str}\\n\"\n",
|
||
" \"SQL: {sql_query}\\n\"\n",
|
||
" \"SQL Response: {context_str}\\n\"\n",
|
||
" \"Response: \"\n",
|
||
")\n",
|
||
"response_synthesis_prompt = PromptTemplate(\n",
|
||
" response_synthesis_prompt_str,\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "95b197ba-9c26-4e6f-9300-0e90a474911c",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
|
||
"llm = OpenAI(model=\"gpt-4o-mini\")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "655a23f3-06e3-45cf-9988-bf1cf5c4a5ad",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Define Workflow\n",
|
||
"\n",
|
||
"Now that the components are in place, let's define the full workflow! "
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "3f3b808c-df2a-4bc1-89a3-ba75f6377f68",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.workflow import (\n",
|
||
" Workflow,\n",
|
||
" StartEvent,\n",
|
||
" StopEvent,\n",
|
||
" step,\n",
|
||
" Context,\n",
|
||
" Event,\n",
|
||
")\n",
|
||
"\n",
|
||
"\n",
|
||
"class TableRetrieveEvent(Event):\n",
|
||
" \"\"\"Result of running table retrieval.\"\"\"\n",
|
||
"\n",
|
||
" table_context_str: str\n",
|
||
" query: str\n",
|
||
"\n",
|
||
"\n",
|
||
"class TextToSQLEvent(Event):\n",
|
||
" \"\"\"Text-to-SQL event.\"\"\"\n",
|
||
"\n",
|
||
" sql: str\n",
|
||
" query: str\n",
|
||
"\n",
|
||
"\n",
|
||
"class TextToSQLWorkflow1(Workflow):\n",
|
||
" \"\"\"Text-to-SQL Workflow that does query-time table retrieval.\"\"\"\n",
|
||
"\n",
|
||
" def __init__(\n",
|
||
" self,\n",
|
||
" obj_retriever,\n",
|
||
" text2sql_prompt,\n",
|
||
" sql_retriever,\n",
|
||
" response_synthesis_prompt,\n",
|
||
" llm,\n",
|
||
" *args,\n",
|
||
" **kwargs,\n",
|
||
" ) -> None:\n",
|
||
" \"\"\"Init params.\"\"\"\n",
|
||
" super().__init__(*args, **kwargs)\n",
|
||
" self.obj_retriever = obj_retriever\n",
|
||
" self.text2sql_prompt = text2sql_prompt\n",
|
||
" self.sql_retriever = sql_retriever\n",
|
||
" self.response_synthesis_prompt = response_synthesis_prompt\n",
|
||
" self.llm = llm\n",
|
||
"\n",
|
||
" @step\n",
|
||
" def retrieve_tables(\n",
|
||
" self, ctx: Context, ev: StartEvent\n",
|
||
" ) -> TableRetrieveEvent:\n",
|
||
" \"\"\"Retrieve tables.\"\"\"\n",
|
||
" table_schema_objs = self.obj_retriever.retrieve(ev.query)\n",
|
||
" table_context_str = get_table_context_str(table_schema_objs)\n",
|
||
" return TableRetrieveEvent(\n",
|
||
" table_context_str=table_context_str, query=ev.query\n",
|
||
" )\n",
|
||
"\n",
|
||
" @step\n",
|
||
" def generate_sql(\n",
|
||
" self, ctx: Context, ev: TableRetrieveEvent\n",
|
||
" ) -> TextToSQLEvent:\n",
|
||
" \"\"\"Generate SQL statement.\"\"\"\n",
|
||
" fmt_messages = self.text2sql_prompt.format_messages(\n",
|
||
" query_str=ev.query, schema=ev.table_context_str\n",
|
||
" )\n",
|
||
" chat_response = self.llm.chat(fmt_messages)\n",
|
||
" sql = parse_response_to_sql(chat_response)\n",
|
||
" return TextToSQLEvent(sql=sql, query=ev.query)\n",
|
||
"\n",
|
||
" @step\n",
|
||
" def generate_response(self, ctx: Context, ev: TextToSQLEvent) -> StopEvent:\n",
|
||
" \"\"\"Run SQL retrieval and generate response.\"\"\"\n",
|
||
" retrieved_rows = self.sql_retriever.retrieve(ev.sql)\n",
|
||
" fmt_messages = self.response_synthesis_prompt.format_messages(\n",
|
||
" sql_query=ev.sql,\n",
|
||
" context_str=str(retrieved_rows),\n",
|
||
" query_str=ev.query,\n",
|
||
" )\n",
|
||
" chat_response = llm.chat(fmt_messages)\n",
|
||
" return StopEvent(result=chat_response)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "2922dc47-1f63-4f6a-a795-c3b3e60fd354",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Visualize Workflow\n",
|
||
"\n",
|
||
"A really nice property of workflows is that you can both visualize the execution graph as well as a trace of the most recent execution."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "5ac1a1bc-7cc6-49e7-a7cd-31ca9e347e5b",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"text_to_sql_table_retrieval.html\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from llama_index.utils.workflow import draw_all_possible_flows\n",
|
||
"\n",
|
||
"draw_all_possible_flows(\n",
|
||
" TextToSQLWorkflow1, filename=\"text_to_sql_table_retrieval.html\"\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "74f3aaca-59b0-478e-9652-7ccddd1cb4d4",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from IPython.display import display, HTML\n",
|
||
"\n",
|
||
"# Read the contents of the HTML file\n",
|
||
"with open(\"text_to_sql_table_retrieval.html\", \"r\") as file:\n",
|
||
" html_content = file.read()\n",
|
||
"\n",
|
||
"# Display the HTML content\n",
|
||
"display(HTML(html_content))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "1e3743fc-1c75-4f6b-a97c-d95182384f58",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Run Some Queries! \n",
|
||
"\n",
|
||
"Now we're ready to run some queries across this entire workflow."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "87f17e5c-c0d3-4a82-b451-f532bbeee4c9",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"workflow = TextToSQLWorkflow1(\n",
|
||
" obj_retriever,\n",
|
||
" text2sql_prompt,\n",
|
||
" sql_retriever,\n",
|
||
" response_synthesis_prompt,\n",
|
||
" llm,\n",
|
||
" verbose=True,\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "61b4c13a-16c7-41f7-a3b9-8dd77831647a",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Running step retrieve_tables\n",
|
||
"Step retrieve_tables produced event TableRetrieveEvent\n",
|
||
"Running step generate_sql\n",
|
||
"Step generate_sql produced event TextToSQLEvent\n",
|
||
"Running step generate_response\n",
|
||
"Step generate_response produced event StopEvent\n",
|
||
"assistant: The Notorious B.I.G was signed to Bad Boy Records in 1993.\n",
|
||
"VERBOSE: True\n",
|
||
"> Table Info: Table 'bad_boy_artists_album_release_summary' has columns: Act (VARCHAR), Year_signed (INTEGER), _Albums_released_under_Bad_Boy (VARCHAR), . The table description is: A summary of artists signed to Bad Boy Records along with the year they were signed and the number of albums they released.\n",
|
||
"Here are some relevant example rows (values in the same order as columns above)\n",
|
||
"('The Notorious B.I.G', 1993, '5')\n",
|
||
"\n",
|
||
"> Table Info: Table 'filmography_of_diane_drummond' has columns: Year (INTEGER), Title (VARCHAR), Role (VARCHAR), Notes (VARCHAR), . The table description is: A list of film and television roles played by Diane Drummond from 1995 to 2001.\n",
|
||
"Here are some relevant example rows (values in the same order as columns above)\n",
|
||
"(2013, 'L.A. Slasher', 'The Actress', None)\n",
|
||
"\n",
|
||
"> Table Info: Table 'progressive_rock_album_chart_positions' has columns: Year (INTEGER), Title (VARCHAR), Chart_Positions_UK (VARCHAR), Chart_Positions_US (VARCHAR), Chart_Positions_NL (VARCHAR), Comments (VARCHAR), . The table description is: Chart positions of progressive rock albums in the UK, US, and NL from 1969 to 1981.\n",
|
||
"Here are some relevant example rows (values in the same order as columns above)\n",
|
||
"(1977, 'Novella', '–', '46', '–', '1977 (January in US, August in UK, as the band moved to the Warner Bros Music Group)')\n",
|
||
"\n",
|
||
"VERBOSE: True\n",
|
||
"> Table Info: Table 'bad_boy_artists_album_release_summary' has columns: Act (VARCHAR), Year_signed (INTEGER), _Albums_released_under_Bad_Boy (VARCHAR), . The table description is: A summary of artists signed to Bad Boy Records along with the year they were signed and the number of albums they released.\n",
|
||
"Here are some relevant example rows (values in the same order as columns above)\n",
|
||
"('The Notorious B.I.G', 1993, '5')\n",
|
||
"\n",
|
||
"> Table Info: Table 'filmography_of_diane_drummond' has columns: Year (INTEGER), Title (VARCHAR), Role (VARCHAR), Notes (VARCHAR), . The table description is: A list of film and television roles played by Diane Drummond from 1995 to 2001.\n",
|
||
"Here are some relevant example rows (values in the same order as columns above)\n",
|
||
"(2013, 'L.A. Slasher', 'The Actress', None)\n",
|
||
"\n",
|
||
"> Table Info: Table 'progressive_rock_album_chart_positions' has columns: Year (INTEGER), Title (VARCHAR), Chart_Positions_UK (VARCHAR), Chart_Positions_US (VARCHAR), Chart_Positions_NL (VARCHAR), Comments (VARCHAR), . The table description is: Chart positions of progressive rock albums in the UK, US, and NL from 1969 to 1981.\n",
|
||
"Here are some relevant example rows (values in the same order as columns above)\n",
|
||
"(1977, 'Novella', '–', '46', '–', '1977 (January in US, August in UK, as the band moved to the Warner Bros Music Group)')\n",
|
||
"\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"response = await workflow.run(\n",
|
||
" query=\"What was the year that The Notorious B.I.G was signed to Bad Boy?\"\n",
|
||
")\n",
|
||
"print(str(response))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "8d546d22-1326-406d-9050-7a8ec11d42d4",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Running step retrieve_tables\n",
|
||
"Step retrieve_tables produced event TableRetrieveEvent\n",
|
||
"Running step generate_sql\n",
|
||
"Step generate_sql produced event TextToSQLEvent\n",
|
||
"Running step generate_response\n",
|
||
"Step generate_response produced event StopEvent\n",
|
||
"assistant: William Friedkin won the Best Director award at the 1972 Academy Awards.\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"response = await workflow.run(\n",
|
||
" query=\"Who won best director in the 1972 academy awards\"\n",
|
||
")\n",
|
||
"print(str(response))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e9ed0f49-a2ed-4a3f-9ec7-bd2562323534",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Running step retrieve_tables\n",
|
||
"Step retrieve_tables produced event TableRetrieveEvent\n",
|
||
"Running step generate_sql\n",
|
||
"Step generate_sql produced event TextToSQLEvent\n",
|
||
"Running step generate_response\n",
|
||
"Step generate_response produced event StopEvent\n",
|
||
"assistant: Pasquale Preziosa has been serving since 25 February 2013 and is currently in office as the incumbent.\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"response = await workflow.run(query=\"What was the term of Pasquale Preziosa?\")\n",
|
||
"print(str(response))"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "02ce9fb2-3799-47a7-b170-eee83f0037d5",
|
||
"metadata": {},
|
||
"source": [
|
||
"## 2. Advanced Capability 2: Text-to-SQL with Query-Time Row Retrieval (along with Table Retrieval)\n",
|
||
"\n",
|
||
"One problem in the previous example is that if the user asks a query that asks for \"The Notorious BIG\" but the artist is stored as \"The Notorious B.I.G\", then the generated SELECT statement will likely not return any matches.\n",
|
||
"\n",
|
||
"We can alleviate this problem by fetching a small number of example rows per table. A naive option would be to just take the first k rows. Instead, we embed, index, and retrieve k relevant rows given the user query to give the text-to-SQL LLM the most contextually relevant information for SQL generation.\n",
|
||
"\n",
|
||
"We now extend our workflow."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "cbb5818c-22be-4ae2-882d-7aa0457f0f20",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Index Each Table\n",
|
||
"\n",
|
||
"We embed/index the rows of each table, resulting in one index per table."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "e5e22ed2-60c2-4b5b-b3b4-71cb9a364ba5",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Indexing rows in table: academy_awards_and_nominations_1972\n",
|
||
"Indexing rows in table: annual_traffic_accident_deaths\n",
|
||
"Indexing rows in table: bad_boy_artists_album_release_summary\n",
|
||
"Indexing rows in table: bbc_radio_services_cost_comparison_2012_2013\n",
|
||
"Indexing rows in table: binary_encoding_probabilities\n",
|
||
"Indexing rows in table: boxing_match_results_summary\n",
|
||
"Indexing rows in table: cancer_related_genes_and_functions\n",
|
||
"Indexing rows in table: diane_drummond_awards_nominations\n",
|
||
"Indexing rows in table: diane_drummond_oscar_nominations_and_wins\n",
|
||
"Indexing rows in table: diane_drummond_single_chart_performance\n",
|
||
"Indexing rows in table: euro_2020_group_stage_results\n",
|
||
"Indexing rows in table: experiment_drop_events_timeline\n",
|
||
"Indexing rows in table: filmography_of_diane_drummond\n",
|
||
"Indexing rows in table: grammy_awards_summary_for_wilco\n",
|
||
"Indexing rows in table: historical_college_football_records\n",
|
||
"Indexing rows in table: italian_ministers_term_dates\n",
|
||
"Indexing rows in table: kodachrome_film_types_and_dates\n",
|
||
"Indexing rows in table: missing_persons_case_summary\n",
|
||
"Indexing rows in table: monthly_climate_statistics\n",
|
||
"Indexing rows in table: monthly_climate_statistics_summary\n",
|
||
"Indexing rows in table: monthly_weather_statistics\n",
|
||
"Indexing rows in table: multilingual_greetings_and_phrases\n",
|
||
"Indexing rows in table: municipalities_merger_summary\n",
|
||
"Indexing rows in table: new_mexico_government_officials\n",
|
||
"Indexing rows in table: norwegian_club_performance_summary\n",
|
||
"Indexing rows in table: ohio_private_schools_summary\n",
|
||
"Indexing rows in table: progressive_rock_album_chart_positions\n",
|
||
"Indexing rows in table: regional_airports_usage_summary\n",
|
||
"Indexing rows in table: south_dakota_radio_stations\n",
|
||
"Indexing rows in table: triple_crown_winners_summary\n",
|
||
"Indexing rows in table: uk_ministers_and_titles_history\n",
|
||
"Indexing rows in table: voter_registration_status_by_party\n",
|
||
"Indexing rows in table: voter_registration_summary_by_party\n",
|
||
"Indexing rows in table: yamato_district_population_density\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from llama_index.core import VectorStoreIndex, load_index_from_storage\n",
|
||
"from sqlalchemy import text\n",
|
||
"from llama_index.core.schema import TextNode\n",
|
||
"from llama_index.core import StorageContext\n",
|
||
"import os\n",
|
||
"from pathlib import Path\n",
|
||
"from typing import Dict\n",
|
||
"\n",
|
||
"\n",
|
||
"def index_all_tables(\n",
|
||
" sql_database: SQLDatabase, table_index_dir: str = \"table_index_dir\"\n",
|
||
") -> Dict[str, VectorStoreIndex]:\n",
|
||
" \"\"\"Index all tables.\"\"\"\n",
|
||
" if not Path(table_index_dir).exists():\n",
|
||
" os.makedirs(table_index_dir)\n",
|
||
"\n",
|
||
" vector_index_dict = {}\n",
|
||
" engine = sql_database.engine\n",
|
||
" for table_name in sql_database.get_usable_table_names():\n",
|
||
" print(f\"Indexing rows in table: {table_name}\")\n",
|
||
" if not os.path.exists(f\"{table_index_dir}/{table_name}\"):\n",
|
||
" # get all rows from table\n",
|
||
" with engine.connect() as conn:\n",
|
||
" cursor = conn.execute(text(f'SELECT * FROM \"{table_name}\"'))\n",
|
||
" result = cursor.fetchall()\n",
|
||
" row_tups = []\n",
|
||
" for row in result:\n",
|
||
" row_tups.append(tuple(row))\n",
|
||
"\n",
|
||
" # index each row, put into vector store index\n",
|
||
" nodes = [TextNode(text=str(t)) for t in row_tups]\n",
|
||
"\n",
|
||
" # put into vector store index (use OpenAIEmbeddings by default)\n",
|
||
" index = VectorStoreIndex(nodes)\n",
|
||
"\n",
|
||
" # save index\n",
|
||
" index.set_index_id(\"vector_index\")\n",
|
||
" index.storage_context.persist(f\"{table_index_dir}/{table_name}\")\n",
|
||
" else:\n",
|
||
" # rebuild storage context\n",
|
||
" storage_context = StorageContext.from_defaults(\n",
|
||
" persist_dir=f\"{table_index_dir}/{table_name}\"\n",
|
||
" )\n",
|
||
" # load index\n",
|
||
" index = load_index_from_storage(\n",
|
||
" storage_context, index_id=\"vector_index\"\n",
|
||
" )\n",
|
||
" vector_index_dict[table_name] = index\n",
|
||
"\n",
|
||
" return vector_index_dict\n",
|
||
"\n",
|
||
"\n",
|
||
"vector_index_dict = index_all_tables(sql_database)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "379ab116-08b3-4ad6-83ad-32062893f3c8",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Define Expanded Table Parsing\n",
|
||
"\n",
|
||
"We expand the capability of our table parsing to not only return the relevant table schemas, but also return relevant rows per table schema.\n",
|
||
"\n",
|
||
"It now takes in both `table_schema_objs` (output of table retriever), but also the original `query_str` which will then be used for vector retrieval of relevant rows."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "ee4d51a3-e88f-49dc-b75b-c7ce0ca770f9",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.retrievers import SQLRetriever\n",
|
||
"from typing import List\n",
|
||
"\n",
|
||
"sql_retriever = SQLRetriever(sql_database)\n",
|
||
"\n",
|
||
"\n",
|
||
"def get_table_context_and_rows_str(\n",
|
||
" query_str: str,\n",
|
||
" table_schema_objs: List[SQLTableSchema],\n",
|
||
" verbose: bool = False,\n",
|
||
"):\n",
|
||
" \"\"\"Get table context string.\"\"\"\n",
|
||
" context_strs = []\n",
|
||
" for table_schema_obj in table_schema_objs:\n",
|
||
" # first append table info + additional context\n",
|
||
" table_info = sql_database.get_single_table_info(\n",
|
||
" table_schema_obj.table_name\n",
|
||
" )\n",
|
||
" if table_schema_obj.context_str:\n",
|
||
" table_opt_context = \" The table description is: \"\n",
|
||
" table_opt_context += table_schema_obj.context_str\n",
|
||
" table_info += table_opt_context\n",
|
||
"\n",
|
||
" # also lookup vector index to return relevant table rows\n",
|
||
" vector_retriever = vector_index_dict[\n",
|
||
" table_schema_obj.table_name\n",
|
||
" ].as_retriever(similarity_top_k=2)\n",
|
||
" relevant_nodes = vector_retriever.retrieve(query_str)\n",
|
||
" if len(relevant_nodes) > 0:\n",
|
||
" table_row_context = \"\\nHere are some relevant example rows (values in the same order as columns above)\\n\"\n",
|
||
" for node in relevant_nodes:\n",
|
||
" table_row_context += str(node.get_content()) + \"\\n\"\n",
|
||
" table_info += table_row_context\n",
|
||
"\n",
|
||
" if verbose:\n",
|
||
" print(f\"> Table Info: {table_info}\")\n",
|
||
"\n",
|
||
" context_strs.append(table_info)\n",
|
||
" return \"\\n\\n\".join(context_strs)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "b2c4b1d5-c2c5-439c-a1d4-28a9f31f81bc",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Define Expanded Workflow\n",
|
||
"\n",
|
||
"We re-use the workflow in section 1, but with an upgraded SQL parsing step after text-to-SQL generation.\n",
|
||
"\n",
|
||
"It is very easy to subclass and extend an existing workflow, and customizing existing steps to be more advanced. Here we define a new worfklow that overrides the existing `retrieve_tables` step in order to return the relevant rows."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "1ad3191c-6f64-4798-b4f9-bde13230ab8a",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.workflow import (\n",
|
||
" Workflow,\n",
|
||
" StartEvent,\n",
|
||
" StopEvent,\n",
|
||
" step,\n",
|
||
" Context,\n",
|
||
" Event,\n",
|
||
")\n",
|
||
"\n",
|
||
"\n",
|
||
"class TextToSQLWorkflow2(TextToSQLWorkflow1):\n",
|
||
" \"\"\"Text-to-SQL Workflow that does query-time row AND table retrieval.\"\"\"\n",
|
||
"\n",
|
||
" @step\n",
|
||
" def retrieve_tables(\n",
|
||
" self, ctx: Context, ev: StartEvent\n",
|
||
" ) -> TableRetrieveEvent:\n",
|
||
" \"\"\"Retrieve tables.\"\"\"\n",
|
||
" table_schema_objs = self.obj_retriever.retrieve(ev.query)\n",
|
||
" table_context_str = get_table_context_and_rows_str(\n",
|
||
" ev.query, table_schema_objs, verbose=self._verbose\n",
|
||
" )\n",
|
||
" return TableRetrieveEvent(\n",
|
||
" table_context_str=table_context_str, query=ev.query\n",
|
||
" )"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8bf37b20-013c-47ea-9570-e4f98444557b",
|
||
"metadata": {},
|
||
"source": [
|
||
"Since the overall sequence of steps is the same, the graph should look the same."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "bc4451a1-b10f-4d30-9adb-333f623f56c1",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"text_to_sql_table_retrieval.html\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from llama_index.utils.workflow import draw_all_possible_flows\n",
|
||
"\n",
|
||
"draw_all_possible_flows(\n",
|
||
" TextToSQLWorkflow2, filename=\"text_to_sql_table_retrieval.html\"\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "892f1581-2307-47c8-90fc-e15d709c0e36",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Run Some Queries\n",
|
||
"\n",
|
||
"We can now ask about relevant entries even if it doesn't exactly match the entry in the database."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "325aa74c-8a0d-4ec8-9bc7-89554f59e4d6",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"workflow2 = TextToSQLWorkflow2(\n",
|
||
" obj_retriever,\n",
|
||
" text2sql_prompt,\n",
|
||
" sql_retriever,\n",
|
||
" response_synthesis_prompt,\n",
|
||
" llm,\n",
|
||
" verbose=True,\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "48679514-22b1-4a9b-93bf-84979361c45c",
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Running step retrieve_tables\n",
|
||
"VERBOSE: True\n",
|
||
"> Table Info: Table 'bad_boy_artists_album_release_summary' has columns: Act (VARCHAR), Year_signed (INTEGER), _Albums_released_under_Bad_Boy (VARCHAR), . The table description is: A summary of artists signed to Bad Boy Records along with the year they were signed and the number of albums they released.\n",
|
||
"Here are some relevant example rows (values in the same order as columns above)\n",
|
||
"('The Notorious B.I.G', 1993, '5')\n",
|
||
"\n",
|
||
"> Table Info: Table 'filmography_of_diane_drummond' has columns: Year (INTEGER), Title (VARCHAR), Role (VARCHAR), Notes (VARCHAR), . The table description is: A list of film and television roles played by Diane Drummond from 1995 to 2001.\n",
|
||
"Here are some relevant example rows (values in the same order as columns above)\n",
|
||
"(2013, 'L.A. Slasher', 'The Actress', None)\n",
|
||
"\n",
|
||
"> Table Info: Table 'progressive_rock_album_chart_positions' has columns: Year (INTEGER), Title (VARCHAR), Chart_Positions_UK (VARCHAR), Chart_Positions_US (VARCHAR), Chart_Positions_NL (VARCHAR), Comments (VARCHAR), . The table description is: Chart positions of progressive rock albums in the UK, US, and NL from 1969 to 1981.\n",
|
||
"Here are some relevant example rows (values in the same order as columns above)\n",
|
||
"(1977, 'Novella', '–', '46', '–', '1977 (January in US, August in UK, as the band moved to the Warner Bros Music Group)')\n",
|
||
"\n",
|
||
"Step retrieve_tables produced event TableRetrieveEvent\n",
|
||
"Running step generate_sql\n",
|
||
"Step generate_sql produced event TextToSQLEvent\n",
|
||
"Running step generate_response\n",
|
||
"Step generate_response produced event StopEvent\n",
|
||
"assistant: The Notorious B.I.G. was signed to Bad Boy Records in 1993.\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"response = await workflow2.run(\n",
|
||
" query=\"What was the year that The Notorious BIG was signed to Bad Boy?\"\n",
|
||
")\n",
|
||
"print(str(response))"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "llama_index_v3",
|
||
"language": "python",
|
||
"name": "llama_index_v3"
|
||
},
|
||
"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
|
||
}
|