code: get_ipython().run_line_magic('pip', 'install llama-index-llms-openai')
  
  
  get_ipython().system('wget "https://github.com/ppasupat/WikiTableQuestions/releases/download/v1.0.2/WikiTableQuestions-1.0.2-compact.zip" -O data.zip')
  get_ipython().system('unzip data.zip')
  
  
  import pandas as pd
  from pathlib import Path
  
  data_dir = Path("./WikiTableQuestions/csv/200-csv")
  csv_files = sorted([f for f in data_dir.glob("*.csv")])
  dfs = []
  for csv_file in csv_files:
      print(f"processing file: {csv_file}")
      try:
          df = pd.read_csv(csv_file)
          dfs.append(df)
      except Exception as e:
          print(f"Error parsing {csv_file}: {str(e)}")
  
  
  
  tableinfo_dir = "WikiTableQuestions_TableInfo"
  get_ipython().system('mkdir {tableinfo_dir}')
  
  
  from llama_index.core.program import LLMTextCompletionProgram
  from llama_index.core.bridge.pydantic import BaseModel, Field
  from llama_index.llms.openai import OpenAI
  
  
  class TableInfo(BaseModel):
      """Information regarding a structured table."""
  
      table_name: str = Field(
          ..., description="table name (must be underscores and NO spaces)"
      )
      table_summary: str = Field(
          ..., description="short, concise summary/caption of the table"
      )
  
  
  prompt_str = """\
  Give me a summary of the table with the following JSON format.
  
  - The table name must be unique to the table and describe it while being concise. 
  - Do NOT output a generic table name (e.g. table, my_table).
  
  Do NOT make the table name one of the following: {exclude_table_name_list}
  
  Table:
  {table_str}
  
  Summary: """
  
  program = LLMTextCompletionProgram.from_defaults(
      output_cls=TableInfo,
      llm=OpenAI(model="gpt-3.5-turbo"),
      prompt_template_str=prompt_str,
  )
  
  
  import json
  
  
  def _get_tableinfo_with_index(idx: int) -> str:
      results_gen = Path(tableinfo_dir).glob(f"{idx}_*")
      results_list = list(results_gen)
      if len(results_list) == 0:
          return None
      elif len(results_list) == 1:
          path = results_list[0]
          return TableInfo.parse_file(path)
      else:
          raise ValueError(
              f"More than one file matching index: {list(results_gen)}"
          )
  
  
  table_names = set()
  table_infos = []
  for idx, df in enumerate(dfs):
      table_info = _get_tableinfo_with_index(idx)
      if table_info:
          table_infos.append(table_info)
      else:
          while True:
              df_str = df.head(10).to_csv()
              table_info = program(
                  table_str=df_str,
                  exclude_table_name_list=str(list(table_names)),
              )
              table_name = table_info.table_name
              print(f"Processed table: {table_name}")
              if table_name not in table_names:
                  table_names.add(table_name)
                  break
              else:
                  print(f"Table name {table_name} already exists, trying again.")
                  pass
  
          out_file = f"{tableinfo_dir}/{idx}_{table_name}.json"
          json.dump(table_info.dict(), open(out_file, "w"))
      table_infos.append(table_info)
  
  
  
  from sqlalchemy import (
      create_engine,
      MetaData,
      Table,
      Column,
      String,
      Integer,
  )
  import re
  
  
  def sanitize_column_name(col_name):
      return re.sub(r"\W+", "_", col_name)
  
  
  def create_table_from_dataframe(
      df: pd.DataFrame, table_name: str, engine, metadata_obj
  ):
      sanitized_columns = {col: sanitize_column_name(col) for col in df.columns}
      df = df.rename(columns=sanitized_columns)
  
      columns = [
          Column(col, String if dtype == "object" else Integer)
          for col, dtype in zip(df.columns, df.dtypes)
      ]
  
      table = Table(table_name, metadata_obj, *columns)
  
      metadata_obj.create_all(engine)
  
      with engine.connect() as conn:
          for _, row in df.iterrows():
              insert_stmt = table.insert().values(**row.to_dict())
              conn.execute(insert_stmt)
          conn.commit()
  
  
  engine = create_engine("sqlite:///:memory:")
  metadata_obj = MetaData()
  for idx, df in enumerate(dfs):
      tableinfo = _get_tableinfo_with_index(idx)
      print(f"Creating table: {tableinfo.table_name}")
      create_table_from_dataframe(df, tableinfo.table_name, engine, metadata_obj)
  
  
  import phoenix as px
  import llama_index.core
  
  px.launch_app()
  llama_index.core.set_global_handler("arize_phoenix")
  
  
  
  
  from llama_index.core.objects import (
      SQLTableNodeMapping,
      ObjectIndex,
      SQLTableSchema,
  )
  from llama_index.core import SQLDatabase, VectorStoreIndex
  
  sql_database = SQLDatabase(engine)
  
  table_node_mapping = SQLTableNodeMapping(sql_database)
  table_schema_objs = [
      SQLTableSchema(table_name=t.table_name, context_str=t.table_summary)
      for t in table_infos
  ]  # add a SQLTableSchema for each table
  
  obj_index = ObjectIndex.from_objects(
      table_schema_objs,
      table_node_mapping,
      VectorStoreIndex,
  )
  obj_retriever = obj_index.as_retriever(similarity_top_k=3)
  
  
  
  from llama_index.core.retrievers import SQLRetriever
  from typing import List
  from llama_index.core.query_pipeline import FnComponent
  
  sql_retriever = SQLRetriever(sql_database)
  
  
  def get_table_context_str(table_schema_objs: List[SQLTableSchema]):
      """Get table context string."""
      context_strs = []
      for table_schema_obj in table_schema_objs:
          table_info = sql_database.get_single_table_info(
              table_schema_obj.table_name
          )
          if table_schema_obj.context_str:
              table_opt_context = " The table description is: "
              table_opt_context += table_schema_obj.context_str
              table_info += table_opt_context
  
          context_strs.append(table_info)
      return "\n\n".join(context_strs)
  
  
  table_parser_component = FnComponent(fn=get_table_context_str)
  
  
  
  from llama_index.core.prompts.default_prompts import DEFAULT_TEXT_TO_SQL_PROMPT
  from llama_index.core import PromptTemplate
  from llama_index.core.query_pipeline import FnComponent
  from llama_index.core.llms import ChatResponse
  
  
  def parse_response_to_sql(response: ChatResponse) -> str:
      """Parse response to SQL."""
      response = response.message.content
      sql_query_start = response.find("SQLQuery:")
      if sql_query_start != -1:
          response = response[sql_query_start:]
          if response.startswith("SQLQuery:"):
              response = response[len("SQLQuery:") :]
      sql_result_start = response.find("SQLResult:")
      if sql_result_start != -1:
          response = response[:sql_result_start]
      return response.strip().strip("```").strip()
  
  
  sql_parser_component = FnComponent(fn=parse_response_to_sql)
  
  text2sql_prompt = DEFAULT_TEXT_TO_SQL_PROMPT.partial_format(
      dialect=engine.dialect.name
  )
  print(text2sql_prompt.template)
  
  
  
  response_synthesis_prompt_str = (
      "Given an input question, synthesize a response from the query results.\n"
      "Query: {query_str}\n"
      "SQL: {sql_query}\n"
      "SQL Response: {context_str}\n"
      "Response: "
  )
  response_synthesis_prompt = PromptTemplate(
      response_synthesis_prompt_str,
  )
  
  
  llm = OpenAI(model="gpt-3.5-turbo")
  
  
  
  from llama_index.core.query_pipeline import (
      QueryPipeline as QP,
      Link,
      InputComponent,
      CustomQueryComponent,
  )
  
  qp = QP(
      modules={
          "input": InputComponent(),
          "table_retriever": obj_retriever,
          "table_output_parser": table_parser_component,
          "text2sql_prompt": text2sql_prompt,
          "text2sql_llm": llm,
          "sql_output_parser": sql_parser_component,
          "sql_retriever": sql_retriever,
          "response_synthesis_prompt": response_synthesis_prompt,
          "response_synthesis_llm": llm,
      },
      verbose=True,
  )
  
  
  qp.add_chain(["input", "table_retriever", "table_output_parser"])
  qp.add_link("input", "text2sql_prompt", dest_key="query_str")
  qp.add_link("table_output_parser", "text2sql_prompt", dest_key="schema")
  qp.add_chain(
      ["text2sql_prompt", "text2sql_llm", "sql_output_parser", "sql_retriever"]
  )
  qp.add_link(
      "sql_output_parser", "response_synthesis_prompt", dest_key="sql_query"
  )
  qp.add_link(
      "sql_retriever", "response_synthesis_prompt", dest_key="context_str"
  )
  qp.add_link("input", "response_synthesis_prompt", dest_key="query_str")
  qp.add_link("response_synthesis_prompt", "response_synthesis_llm")
  
  
  
  from pyvis.network import Network
  
  net = Network(notebook=True, cdn_resources="in_line", directed=True)
  net.from_nx(qp.dag)
  net.show("text2sql_dag.html")
  
  
  
  response = qp.run(
      query="What was the year that The Notorious B.I.G was signed to Bad Boy?"
  )
  print(str(response))
  
  
  response = qp.run(query="Who won best director in the 1972 academy awards")
  print(str(response))
  
  
  response = qp.run(query="What was the term of Pasquale Preziosa?")
  print(str(response))
  
  
  
  
  from llama_index.core import VectorStoreIndex, load_index_from_storage
  from sqlalchemy import text
  from llama_index.core.schema import TextNode
  from llama_index.core import StorageContext
  import os
  from pathlib import Path
  from typing import Dict
  
  
  def index_all_tables(
      sql_database: SQLDatabase, table_index_dir: str = "table_index_dir"
  ) -> Dict[str, VectorStoreIndex]:
      """Index all tables."""
      if not Path(table_index_dir).exists():
          os.makedirs(table_index_dir)
  
      vector_index_dict = {}
      engine = sql_database.engine
      for table_name in sql_database.get_usable_table_names():
          print(f"Indexing rows in table: {table_name}")
          if not os.path.exists(f"{table_index_dir}/{table_name}"):
              with engine.connect() as conn:
                  cursor = conn.execute(text(f'SELECT * FROM "{table_name}"'))
                  result = cursor.fetchall()
                  row_tups = []
                  for row in result:
                      row_tups.append(tuple(row))
  
              nodes = [TextNode(text=str(t)) for t in row_tups]
  
              index = VectorStoreIndex(nodes)
  
              index.set_index_id("vector_index")
              index.storage_context.persist(f"{table_index_dir}/{table_name}")
          else:
              storage_context = StorageContext.from_defaults(
                  persist_dir=f"{table_index_dir}/{table_name}"
              )
              index = load_index_from_storage(
                  storage_context, index_id="vector_index"
              )
          vector_index_dict[table_name] = index
  
      return vector_index_dict
  
  
  vector_index_dict = index_all_tables(sql_database)
  
  
  test_retriever = vector_index_dict["Bad_Boy_Artists"].as_retriever(
      similarity_top_k=1
  )
  nodes = test_retriever.retrieve("P. Diddy")
  print(nodes[0].get_content())
  
  
  
  from llama_index.core.retrievers import SQLRetriever
  from typing import List
  from llama_index.core.query_pipeline import FnComponent
  
  sql_retriever = SQLRetriever(sql_database)
  
  
  def get_table_context_and_rows_str(
      query_str: str, table_schema_objs: List[SQLTableSchema]
  ):
      """Get table context string."""
      context_strs = []
      for table_schema_obj in table_schema_objs:
          table_info = sql_database.get_single_table_info(
              table_schema_obj.table_name
          )
          if table_schema_obj.context_str:
              table_opt_context = " The table description is: "
              table_opt_context += table_schema_obj.context_str
              table_info += table_opt_context
  
          vector_retriever = vector_index_dict[
              table_schema_obj.table_name
          ].as_retriever(similarity_top_k=2)
          relevant_nodes = vector_retriever.retrieve(query_str)
          if len(relevant_nodes) > 0:
              table_row_context = "\nHere are some relevant example rows (values in the same order as columns above)\n"
              for node in relevant_nodes:
                  table_row_context += str(node.get_content()) + "\n"
              table_info += table_row_context
  
          context_strs.append(table_info)
      return "\n\n".join(context_strs)
  
  
  table_parser_component = FnComponent(fn=get_table_context_and_rows_str)
  
  
  
  from llama_index.core.query_pipeline import (
      QueryPipeline as QP,
      Link,
      InputComponent,
      CustomQueryComponent,
  )
  
  qp = QP(
      modules={
          "input": InputComponent(),
          "table_retriever": obj_retriever,
          "table_output_parser": table_parser_component,
          "text2sql_prompt": text2sql_prompt,
          "text2sql_llm": llm,
          "sql_output_parser": sql_parser_component,
          "sql_retriever": sql_retriever,
          "response_synthesis_prompt": response_synthesis_prompt,
          "response_synthesis_llm": llm,
      },
      verbose=True,
  )
  
  
  qp.add_link("input", "table_retriever")
  qp.add_link("input", "table_output_parser", dest_key="query_str")
  qp.add_link(
      "table_retriever", "table_output_parser", dest_key="table_schema_objs"
  )
  qp.add_link("input", "text2sql_prompt", dest_key="query_str")
  qp.add_link("table_output_parser", "text2sql_prompt", dest_key="schema")
  qp.add_chain(
      ["text2sql_prompt", "text2sql_llm", "sql_output_parser", "sql_retriever"]
  )
  qp.add_link(
      "sql_output_parser", "response_synthesis_prompt", dest_key="sql_query"
  )
  qp.add_link(
      "sql_retriever", "response_synthesis_prompt", dest_key="context_str"
  )
  qp.add_link("input", "response_synthesis_prompt", dest_key="query_str")
  qp.add_link("response_synthesis_prompt", "response_synthesis_llm")
  
  
  from pyvis.network import Network
  
  net = Network(notebook=True, cdn_resources="in_line", directed=True)
  net.from_nx(qp.dag)
  net.show("text2sql_dag.html")
  
  
  
  response = qp.run(
      query="What was the year that The Notorious BIG was signed to Bad Boy?"
  )
  print(str(response))
apis: llama_index.core.StorageContext.from_defaults llama_index.core.objects.ObjectIndex.from_objects llama_index.core.load_index_from_storage llama_index.core.objects.SQLTableSchema llama_index.core.query_pipeline.FnComponent llama_index.core.prompts.default_prompts.DEFAULT_TEXT_TO_SQL_PROMPT.partial_format llama_index.core.PromptTemplate llama_index.core.query_pipeline.InputComponent llama_index.core.VectorStoreIndex llama_index.core.bridge.pydantic.Field llama_index.core.objects.SQLTableNodeMapping llama_index.llms.openai.OpenAI llama_index.core.retrievers.SQLRetriever llama_index.core.SQLDatabase
extract_api: [((323, 363), 'pathlib.Path', 'Path', (['"""./WikiTableQuestions/csv/200-csv"""'], {}), "('./WikiTableQuestions/csv/200-csv')\n", (327, 363), False, 'from pathlib import Path\n'), ((3874, 3909), 'sqlalchemy.create_engine', 'create_engine', (['"""sqlite:///:memory:"""'], {}), "('sqlite:///:memory:')\n", (3887, 3909), False, 'from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer\n'), ((3925, 3935), 'sqlalchemy.MetaData', 'MetaData', ([], {}), '()\n', (3933, 3935), False, 'from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer\n'), ((4195, 4210), 'phoenix.launch_app', 'px.launch_app', ([], {}), '()\n', (4208, 4210), True, 'import phoenix as px\n'), ((4446, 4465), 'llama_index.core.SQLDatabase', 'SQLDatabase', (['engine'], {}), '(engine)\n', (4457, 4465), False, 'from llama_index.core import SQLDatabase, VectorStoreIndex\n'), ((4488, 4521), 'llama_index.core.objects.SQLTableNodeMapping', 'SQLTableNodeMapping', (['sql_database'], {}), '(sql_database)\n', (4507, 4521), False, 'from llama_index.core.objects import SQLTableNodeMapping, ObjectIndex, SQLTableSchema\n'), ((4696, 4781), 'llama_index.core.objects.ObjectIndex.from_objects', 'ObjectIndex.from_objects', (['table_schema_objs', 'table_node_mapping', 'VectorStoreIndex'], {}), '(table_schema_objs, table_node_mapping,\n    VectorStoreIndex)\n', (4720, 4781), False, 'from llama_index.core.objects import SQLTableNodeMapping, ObjectIndex, SQLTableSchema\n'), ((5005, 5031), 'llama_index.core.retrievers.SQLRetriever', 'SQLRetriever', (['sql_database'], {}), '(sql_database)\n', (5017, 5031), False, 'from llama_index.core.retrievers import SQLRetriever\n'), ((5628, 5665), 'llama_index.core.query_pipeline.FnComponent', 'FnComponent', ([], {'fn': 'get_table_context_str'}), '(fn=get_table_context_str)\n', (5639, 5665), False, 'from llama_index.core.query_pipeline import FnComponent\n'), ((6454, 6491), 'llama_index.core.query_pipeline.FnComponent', 'FnComponent', ([], {'fn': 'parse_response_to_sql'}), '(fn=parse_response_to_sql)\n', (6465, 6491), False, 'from llama_index.core.query_pipeline import FnComponent\n'), ((6511, 6581), 'llama_index.core.prompts.default_prompts.DEFAULT_TEXT_TO_SQL_PROMPT.partial_format', 'DEFAULT_TEXT_TO_SQL_PROMPT.partial_format', ([], {'dialect': 'engine.dialect.name'}), '(dialect=engine.dialect.name)\n', (6552, 6581), False, 'from llama_index.core.prompts.default_prompts import DEFAULT_TEXT_TO_SQL_PROMPT\n'), ((6871, 6916), 'llama_index.core.PromptTemplate', 'PromptTemplate', (['response_synthesis_prompt_str'], {}), '(response_synthesis_prompt_str)\n', (6885, 6916), False, 'from llama_index.core import PromptTemplate\n'), ((6932, 6961), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (6938, 6961), False, 'from llama_index.llms.openai import OpenAI\n'), ((8202, 8264), 'pyvis.network.Network', 'Network', ([], {'notebook': '(True)', 'cdn_resources': '"""in_line"""', 'directed': '(True)'}), "(notebook=True, cdn_resources='in_line', directed=True)\n", (8209, 8264), False, 'from pyvis.network import Network\n'), ((10585, 10611), 'llama_index.core.retrievers.SQLRetriever', 'SQLRetriever', (['sql_database'], {}), '(sql_database)\n', (10597, 10611), False, 'from llama_index.core.retrievers import SQLRetriever\n'), ((11737, 11783), 'llama_index.core.query_pipeline.FnComponent', 'FnComponent', ([], {'fn': 'get_table_context_and_rows_str'}), '(fn=get_table_context_and_rows_str)\n', (11748, 11783), False, 'from llama_index.core.query_pipeline import FnComponent\n'), ((13153, 13215), 'pyvis.network.Network', 'Network', ([], {'notebook': '(True)', 'cdn_resources': '"""in_line"""', 'directed': '(True)'}), "(notebook=True, cdn_resources='in_line', directed=True)\n", (13160, 13215), False, 'from pyvis.network import Network\n'), ((1015, 1087), 'llama_index.core.bridge.pydantic.Field', 'Field', (['...'], {'description': '"""table name (must be underscores and NO spaces)"""'}), "(..., description='table name (must be underscores and NO spaces)')\n", (1020, 1087), False, 'from llama_index.core.bridge.pydantic import BaseModel, Field\n'), ((1127, 1196), 'llama_index.core.bridge.pydantic.Field', 'Field', (['...'], {'description': '"""short, concise summary/caption of the table"""'}), "(..., description='short, concise summary/caption of the table')\n", (1132, 1196), False, 'from llama_index.core.bridge.pydantic import BaseModel, Field\n'), ((3182, 3211), 're.sub', 're.sub', (['"""\\\\W+"""', '"""_"""', 'col_name'], {}), "('\\\\W+', '_', col_name)\n", (3188, 3211), False, 'import re\n'), ((3586, 3627), 'sqlalchemy.Table', 'Table', (['table_name', 'metadata_obj', '*columns'], {}), '(table_name, metadata_obj, *columns)\n', (3591, 3627), False, 'from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer\n'), ((4548, 4616), 'llama_index.core.objects.SQLTableSchema', 'SQLTableSchema', ([], {'table_name': 't.table_name', 'context_str': 't.table_summary'}), '(table_name=t.table_name, context_str=t.table_summary)\n', (4562, 4616), False, 'from llama_index.core.objects import SQLTableNodeMapping, ObjectIndex, SQLTableSchema\n'), ((520, 541), 'pandas.read_csv', 'pd.read_csv', (['csv_file'], {}), '(csv_file)\n', (531, 541), True, 'import pandas as pd\n'), ((1634, 1663), 'llama_index.llms.openai.OpenAI', 'OpenAI', ([], {'model': '"""gpt-3.5-turbo"""'}), "(model='gpt-3.5-turbo')\n", (1640, 1663), False, 'from llama_index.llms.openai import OpenAI\n'), ((3460, 3513), 'sqlalchemy.Column', 'Column', (['col', "(String if dtype == 'object' else Integer)"], {}), "(col, String if dtype == 'object' else Integer)\n", (3466, 3513), False, 'from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer\n'), ((9087, 9115), 'os.makedirs', 'os.makedirs', (['table_index_dir'], {}), '(table_index_dir)\n', (9098, 9115), False, 'import os\n'), ((1785, 1804), 'pathlib.Path', 'Path', (['tableinfo_dir'], {}), '(tableinfo_dir)\n', (1789, 1804), False, 'from pathlib import Path\n'), ((7135, 7151), 'llama_index.core.query_pipeline.InputComponent', 'InputComponent', ([], {}), '()\n', (7149, 7151), False, 'from llama_index.core.query_pipeline import QueryPipeline as QP, Link, InputComponent, CustomQueryComponent\n'), ((9308, 9357), 'os.path.exists', 'os.path.exists', (['f"""{table_index_dir}/{table_name}"""'], {}), "(f'{table_index_dir}/{table_name}')\n", (9322, 9357), False, 'import os\n'), ((9719, 9742), 'llama_index.core.VectorStoreIndex', 'VectorStoreIndex', (['nodes'], {}), '(nodes)\n', (9735, 9742), False, 'from llama_index.core import VectorStoreIndex, load_index_from_storage\n'), ((9912, 9987), 'llama_index.core.StorageContext.from_defaults', 'StorageContext.from_defaults', ([], {'persist_dir': 'f"""{table_index_dir}/{table_name}"""'}), "(persist_dir=f'{table_index_dir}/{table_name}')\n", (9940, 9987), False, 'from llama_index.core import StorageContext\n'), ((10038, 10103), 'llama_index.core.load_index_from_storage', 'load_index_from_storage', (['storage_context'], {'index_id': '"""vector_index"""'}), "(storage_context, index_id='vector_index')\n", (10061, 10103), False, 'from llama_index.core import VectorStoreIndex, load_index_from_storage\n'), ((11957, 11973), 'llama_index.core.query_pipeline.InputComponent', 'InputComponent', ([], {}), '()\n', (11971, 11973), False, 'from llama_index.core.query_pipeline import QueryPipeline as QP, Link, InputComponent, CustomQueryComponent\n'), ((9047, 9068), 'pathlib.Path', 'Path', (['table_index_dir'], {}), '(table_index_dir)\n', (9051, 9068), False, 'from pathlib import Path\n'), ((9440, 9477), 'sqlalchemy.text', 'text', (['f"""SELECT * FROM "{table_name}\\""""'], {}), '(f\'SELECT * FROM "{table_name}"\')\n', (9444, 9477), False, 'from sqlalchemy import text\n')]