Add files

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*sqlite filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,12 @@
 ---
 title: Llama Index Sql Sandbox
-emoji: ⚡
+emoji: 🦙
 colorFrom: blue
-colorTo: green
+colorTo: pink
 sdk: streamlit
-sdk_version: 1.17.0
+sdk_version: 1.19.0
 app_file: app.py
-pinned: false
+pinned: true
 license: mit
 ---
 

app.py

@@ -0,0 +1,137 @@
+import os
+import streamlit as st
+from streamlit_chat import message as st_message
+from sqlalchemy import create_engine
+
+from langchain.agents import Tool, initialize_agent
+from langchain.chains.conversation.memory import ConversationBufferMemory
+
+from llama_index import GPTSQLStructStoreIndex, LLMPredictor, ServiceContext
+from llama_index import SQLDatabase as llama_SQLDatabase
+from llama_index.indices.struct_store import SQLContextContainerBuilder
+
+from constants import (
+    DEFAULT_SQL_PATH,
+    DEFAULT_BUSINESS_TABLE_DESCRP,
+    DEFAULT_VIOLATIONS_TABLE_DESCRP,
+    DEFAULT_INSPECTIONS_TABLE_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):
+    """Create the GPTSQLStructStoreIndex object."""
+    llm = get_llm(llm_name, model_temperature, api_key)
+
+    engine = create_engine(sql_path)
+    sql_database = llama_SQLDatabase(engine)
+
+    context_container = None
+    if table_context_dict is not None:
+        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)
+
+    return index
+
+
+@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)
+
+    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)
+
+    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)"
+))
+
+
+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)
+
+    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}
+    
+    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)
+    
+    if "llama_index" in st.session_state:
+        query_text = st.text_input("Query:", value="Which restaurant has the most violations?")
+        if st.button("Run Query") and query_text:
+            with st.spinner("Getting response..."):
+                try:
+                    response =  st.session_state['llama_index'].query(query_text)
+                    response_text = str(response)
+                    response_sql = response.extra_info['sql_query']
+                except Exception as e:
+                    response_text = "Error running SQL Query."
+                    response_sql = str(e)
+
+            col1, col2 = st.columns(2)
+            with col1:
+                st.text("SQL Result:")
+                st.markdown(response_text)
+
+            with col2:
+                st.text("SQL Query:")
+                st.markdown(response_sql)
+
+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"):
+        model_input = "User: " + 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)
+
+    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

@@ -0,0 +1,24 @@
+DEFAULT_SQL_PATH = "sqlite:///sfscores.sqlite"
+DEFAULT_BUSINESS_TABLE_DESCRP = (
+    "This table gives information on the IDs, addresses, and other location "
+    "information for several restaurants in San Francisco. This table will "
+    "need to be referenced when users ask about specific businesses."
+)
+DEFAULT_VIOLATIONS_TABLE_DESCRP = (
+    "This table gives information on which business IDs have recorded health violations, "
+    "including the date, risk, and description of each violation. The user may query "
+    "about specific businesses, whose names can be found by mapping the business_id "
+    "to the 'businesses' table."
+)
+DEFAULT_INSPECTIONS_TABLE_DESCRP = (
+    "This table gives information on when each business ID was inspected, including "
+    "the score, date, and type of inspection. The user may query about specific "
+    "businesses, whose names can be found by mapping the business_id to the 'businesses' table."
+)
+DEFAULT_LC_TOOL_DESCRP = "Useful for when you want to answer queries about violations and inspections of businesses."
+
+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

@@ -0,0 +1,4 @@
+langchain==0.0.123
+llama-index==0.5.1
+streamlit==1.19.0
+streamlit-chat==0.0.2.2

sfscores.sqlite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:240deebf58f54606266cdd4a5dca14c48d58f8b530941c0249a9b23f00589afa
+size 9639936

sql_index.json

@@ -0,0 +1 @@
+{"index_struct_id": "b52fad59-0c00-4392-b775-f9cd3fdb6deb", "docstore": {"docs": {"b52fad59-0c00-4392-b775-f9cd3fdb6deb": {"text": null, "doc_id": "b52fad59-0c00-4392-b775-f9cd3fdb6deb", "embedding": null, "doc_hash": "08a14830cef184731c6b6a0bdd67fa351d923556941aa99027b276bd839a07a4", "extra_info": null, "context_dict": {}, "__type__": "sql"}}, "ref_doc_info": {"b52fad59-0c00-4392-b775-f9cd3fdb6deb": {"doc_hash": "08a14830cef184731c6b6a0bdd67fa351d923556941aa99027b276bd839a07a4"}}}, "sql_context_container": {"context_dict": {"violations": "Schema of table violations:\nTable 'violations' has columns: business_id (TEXT), date (TEXT), ViolationTypeID (TEXT), risk_category (TEXT), description (TEXT) and foreign keys: .\nContext of table violations:\nThis table gives information on which business IDs have recorded health violations, including the date, risk, and description of each violation. The user may query about specific businesses, whose names can be found by mapping the business_id to the 'businesses' table.", "businesses": "Schema of table businesses:\nTable 'businesses' has columns: business_id (INTEGER), name (VARCHAR(64)), address (VARCHAR(50)), city (VARCHAR(23)), postal_code (VARCHAR(9)), latitude (FLOAT), longitude (FLOAT), phone_number (BIGINT), TaxCode (VARCHAR(4)), business_certificate (INTEGER), application_date (DATE), owner_name (VARCHAR(99)), owner_address (VARCHAR(74)), owner_city (VARCHAR(22)), owner_state (VARCHAR(14)), owner_zip (VARCHAR(15)) and foreign keys: .\nContext of table businesses:\nThis table gives information on the IDs, addresses, and other location information for several restaruants in San Fransisco. This table will need to be referenced when users ask about specific bussinesses.", "inspections": "Schema of table inspections:\nTable 'inspections' has columns: business_id (TEXT), Score (INTEGER), date (TEXT), type (VARCHAR(33)) and foreign keys: .\nContext of table inspections:\nThis table gives information on when each bussiness ID was inspected, including the score, date, and type of inspection. The user may query about specific businesses, whose names can be found by mapping the business_id to the 'businesses' table."}, "context_str": null}}

utils.py

@@ -0,0 +1,26 @@
+import os
+from langchain import OpenAI
+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)
+        except Exception as e:
+            return f"Error running SQL {e}.\nNot able to retrieve answer."
+        text = str(response)
+        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
+
+
+
+def get_llm(llm_name, model_temperature, 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:
+        return ChatOpenAI(temperature=model_temperature, model_name=llm_name)