update to llamaindex v0.6.13

app.py

@@ -15,13 +15,15 @@ from constants import (
     DEFAULT_BUSINESS_TABLE_DESCRP,
     DEFAULT_VIOLATIONS_TABLE_DESCRP,
     DEFAULT_INSPECTIONS_TABLE_DESCRP,
-    DEFAULT_LC_TOOL_DESCRP
+    DEFAULT_LC_TOOL_DESCRP,
 )
 from utils import get_sql_index_tool, get_llm
 
 
 @st.cache_resource
-def initialize_index(llm_name, model_temperature, table_context_dict, api_key, sql_path=DEFAULT_SQL_PATH):
+def initialize_index(
+    llm_name, model_temperature, table_context_dict, api_key, sql_path=DEFAULT_SQL_PATH
+):
     """Create the GPTSQLStructStoreIndex object."""
     llm = get_llm(llm_name, model_temperature, api_key)
 
@@ -30,14 +32,18 @@ def initialize_index(llm_name, model_temperature, table_context_dict, api_key, s
 
     context_container = None
     if table_context_dict is not None:
-        context_builder = SQLContextContainerBuilder(sql_database, context_dict=table_context_dict)
+        context_builder = SQLContextContainerBuilder(
+            sql_database, context_dict=table_context_dict
+        )
         context_container = context_builder.build_context_container()
-    
+
     service_context = ServiceContext.from_defaults(llm_predictor=LLMPredictor(llm=llm))
-    index = GPTSQLStructStoreIndex([], 
-                                   sql_database=sql_database, 
-                                   sql_context_container=context_container, 
-                                   service_context=service_context)
+    index = GPTSQLStructStoreIndex(
+        [],
+        sql_database=sql_database,
+        sql_context_container=context_container,
+        service_context=service_context,
+    )
 
     return index
 
@@ -45,63 +51,97 @@ def initialize_index(llm_name, model_temperature, table_context_dict, api_key, s
 @st.cache_resource
 def initialize_chain(llm_name, model_temperature, lc_descrp, api_key, _sql_index):
     """Create a (rather hacky) custom agent and sql_index tool."""
-    sql_tool = Tool(name="SQL Index", 
-                    func=get_sql_index_tool(_sql_index, _sql_index.sql_context_container.context_dict), 
-                    description=lc_descrp)
+    sql_tool = Tool(
+        name="SQL Index",
+        func=get_sql_index_tool(
+            _sql_index, _sql_index.sql_context_container.context_dict
+        ),
+        description=lc_descrp,
+    )
 
     llm = get_llm(llm_name, model_temperature, api_key=api_key)
 
     memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 
-    agent_chain = initialize_agent([sql_tool], llm, agent="chat-conversational-react-description", verbose=True, memory=memory)
+    agent_chain = initialize_agent(
+        [sql_tool],
+        llm,
+        agent="chat-conversational-react-description",
+        verbose=True,
+        memory=memory,
+    )
 
     return agent_chain
 
 
 st.title("🦙 Llama Index SQL Sandbox 🦙")
-st.markdown((
-    "This sandbox uses a sqlite database by default, powered by [Llama Index](https://gpt-index.readthedocs.io/en/latest/index.html) ChatGPT, and LangChain.\n\n" 
-    "The database contains information on health violations and inspections at restaurants in San Francisco."
-    "This data is spread across three tables - businesses, inspections, and violations.\n\n"
-    "Using the setup page, you can adjust LLM settings, change the context for the SQL tables, and change the tool description for Langchain."
-    "The other tabs will perform chatbot and text2sql operations.\n\n"
-    "Read more about LlamaIndexes structured data support [here!](https://gpt-index.readthedocs.io/en/latest/guides/tutorials/sql_guide.html)"
-))
+st.markdown(
+    (
+        "This sandbox uses a sqlite database by default, powered by [Llama Index](https://gpt-index.readthedocs.io/en/latest/index.html) ChatGPT, and LangChain.\n\n"
+        "The database contains information on health violations and inspections at restaurants in San Francisco."
+        "This data is spread across three tables - businesses, inspections, and violations.\n\n"
+        "Using the setup page, you can adjust LLM settings, change the context for the SQL tables, and change the tool description for Langchain."
+        "The other tabs will perform chatbot and text2sql operations.\n\n"
+        "Read more about LlamaIndexes structured data support [here!](https://gpt-index.readthedocs.io/en/latest/guides/tutorials/sql_guide.html)"
+    )
+)
 
 
-setup_tab, llama_tab, lc_tab = st.tabs(["Setup", "Llama Index", "Langchain+Llama Index"])
+setup_tab, llama_tab, lc_tab = st.tabs(
+    ["Setup", "Llama Index", "Langchain+Llama Index"]
+)
 
 with setup_tab:
     st.subheader("LLM Setup")
     api_key = st.text_input("Enter your OpenAI API key here", type="password")
-    llm_name = st.selectbox('Which LLM?', ["text-davinci-003", "gpt-3.5-turbo", "gpt-4"])
-    model_temperature = st.slider("LLM Temperature", min_value=0.0, max_value=1.0, step=0.1)
+    llm_name = st.selectbox(
+        "Which LLM?", ["text-davinci-003", "gpt-3.5-turbo", "gpt-4"]
+    )
+    model_temperature = st.slider(
+        "LLM Temperature", min_value=0.0, max_value=1.0, step=0.1
+    )
 
     st.subheader("Table Setup")
-    business_table_descrp = st.text_area("Business table description", value=DEFAULT_BUSINESS_TABLE_DESCRP)
-    violations_table_descrp = st.text_area("Business table description", value=DEFAULT_VIOLATIONS_TABLE_DESCRP)
-    inspections_table_descrp = st.text_area("Business table description", value=DEFAULT_INSPECTIONS_TABLE_DESCRP)
-
-    table_context_dict = {"businesses": business_table_descrp, 
-                          "inspections": inspections_table_descrp, 
-                          "violations": violations_table_descrp}
-    
+    business_table_descrp = st.text_area(
+        "Business table description", value=DEFAULT_BUSINESS_TABLE_DESCRP
+    )
+    violations_table_descrp = st.text_area(
+        "Business table description", value=DEFAULT_VIOLATIONS_TABLE_DESCRP
+    )
+    inspections_table_descrp = st.text_area(
+        "Business table description", value=DEFAULT_INSPECTIONS_TABLE_DESCRP
+    )
+
+    table_context_dict = {
+        "businesses": business_table_descrp,
+        "inspections": inspections_table_descrp,
+        "violations": violations_table_descrp,
+    }
+
     use_table_descrp = st.checkbox("Use table descriptions?", value=True)
     lc_descrp = st.text_area("LangChain Tool Description", value=DEFAULT_LC_TOOL_DESCRP)
 
 with llama_tab:
     st.subheader("Text2SQL with Llama Index")
     if st.button("Initialize Index", key="init_index_1"):
-        st.session_state['llama_index'] = initialize_index(llm_name, model_temperature, table_context_dict if use_table_descrp else None, api_key)
-    
+        st.session_state["llama_index"] = initialize_index(
+            llm_name,
+            model_temperature,
+            table_context_dict if use_table_descrp else None,
+            api_key,
+        )
+
     if "llama_index" in st.session_state:
-        query_text = st.text_input("Query:", value="Which restaurant has the most violations?")
+        query_text = st.text_input(
+            "Query:", value="Which restaurant has the most violations?"
+        )
+        use_nl = st.checkbox("Return natural language response?")
         if st.button("Run Query") and query_text:
             with st.spinner("Getting response..."):
                 try:
-                    response =  st.session_state['llama_index'].query(query_text)
+                    response = st.session_state["llama_index"].as_query_engine(synthesize_response=use_nl).query(query_text)
                     response_text = str(response)
-                    response_sql = response.extra_info['sql_query']
+                    response_sql = response.extra_info["sql_query"]
                 except Exception as e:
                     response_text = "Error running SQL Query."
                     response_sql = str(e)
@@ -119,19 +159,31 @@ with lc_tab:
     st.subheader("Langchain + Llama Index SQL Demo")
 
     if st.button("Initialize Agent"):
-        st.session_state['llama_index'] = initialize_index(llm_name, model_temperature, table_context_dict if use_table_descrp else None, api_key)
-        st.session_state['lc_agent'] = initialize_chain(llm_name, model_temperature, lc_descrp, api_key, st.session_state['llama_index'])
-        st.session_state['chat_history'] = []
-
-    model_input = st.text_input("Message:", value="Which restaurant has the most violations?")
-    if 'lc_agent' in st.session_state and st.button("Send"):
+        st.session_state["llama_index"] = initialize_index(
+            llm_name,
+            model_temperature,
+            table_context_dict if use_table_descrp else None,
+            api_key,
+        )
+        st.session_state["lc_agent"] = initialize_chain(
+            llm_name,
+            model_temperature,
+            lc_descrp,
+            api_key,
+            st.session_state["llama_index"],
+        )
+        st.session_state["chat_history"] = []
+
+    model_input = st.text_input(
+        "Message:", value="Which restaurant has the most violations?"
+    )
+    if "lc_agent" in st.session_state and st.button("Send"):
         model_input = "User: " + model_input
-        st.session_state['chat_history'].append(model_input)
+        st.session_state["chat_history"].append(model_input)
         with st.spinner("Getting response..."):
-            response =  st.session_state['lc_agent'].run(input=model_input)
-        st.session_state['chat_history'].append(response)
+            response = st.session_state["lc_agent"].run(input=model_input)
+        st.session_state["chat_history"].append(response)
 
-    if 'chat_history' in st.session_state:
-        for msg in st.session_state['chat_history']:
+    if "chat_history" in st.session_state:
+        for msg in st.session_state["chat_history"]:
             st_message(msg.split("User: ")[-1], is_user="User: " in msg)
-        

constants.py

@@ -21,4 +21,4 @@ DEFAULT_INGEST_DOCUMENT = (
     "The restaurant KING-KONG had an routine unscheduled inspection on 2023/12/31. "
     "The business achieved a score of 50. We two violations, a high risk "
     "vermin infestation as well as a high risk food holding temperatures."
-)
+)

requirements.txt

@@ -1,4 +1,5 @@
-langchain==0.0.123
-llama-index==0.5.1
+altair==4.2.2
+langchain==0.0.154
+llama-index==0.6.13
 streamlit==1.19.0
 streamlit-chat==0.0.2.2

utils.py

@@ -5,21 +5,22 @@ from langchain.chat_models import ChatOpenAI
 
 def get_sql_index_tool(sql_index, table_context_dict):
     table_context_str = "\n".join(table_context_dict.values())
+
     def run_sql_index_query(query_text):
         try:
-            response = sql_index.query(query_text)
+            response = sql_index.as_query_engine(synthesize_response=False).query(query_text)
         except Exception as e:
             return f"Error running SQL {e}.\nNot able to retrieve answer."
         text = str(response)
-        sql = response.extra_info['sql_query']
+        sql = response.extra_info["sql_query"]
         return f"Here are the details on the SQL table: {table_context_str}\nSQL Query Used: {sql}\nSQL Result: {text}\n"
-        #return f"SQL Query Used: {sql}\nSQL Result: {text}\n"
-    return run_sql_index_query
+        # return f"SQL Query Used: {sql}\nSQL Result: {text}\n"
 
+    return run_sql_index_query
 
 
 def get_llm(llm_name, model_temperature, api_key):
-    os.environ['OPENAI_API_KEY'] = api_key
+    os.environ["OPENAI_API_KEY"] = api_key
     if llm_name == "text-davinci-003":
         return OpenAI(temperature=model_temperature, model_name=llm_name)
     else: