here we go

README.md

@@ -5,7 +5,7 @@ colorFrom: indigo
 colorTo: gray
 sdk: streamlit
 sdk_version: 1.32.2
-app_file: minimal_agent.py
+app_file: app.py
 pinned: false
 license: bsd-2-clause
 ---

app.py

@@ -1,121 +1,68 @@
 import streamlit as st
-from openai import OpenAI
-import duckdb
-
-EXAMPLE_PROMPTS = [
-    "What are the unique values of Mang Type?",
-    "What is the mean species richness of each GAP_Sts?",
-    "STAT, LOC, FED, and DIST Mang types are all 'Public' land types. PVT and NGO are 'Private' land types. What is the average richness in public, private, and other land types?",
-]
-
-TITLE = "Protected Lands SQL Assistant"
-# DESCRIPTION = """[DBRX Instruct](https://huggingface.co/databricks/dbrx-instruct) is a mixture-of-experts (MoE) large language model trained by the Mosaic Research team at Databricks. Users can interact with this model in the [DBRX Playground](https://huggingface.co/spaces/databricks/dbrx-instruct), subject to the terms and conditions below. 
-# This demo is powered by [Databricks Foundation Model APIs](https://docs.databricks.com/en/machine-learning/foundation-models/index.html).
-DESCRIPTION="""
-This is a test
-"""
-
-
-
-duckdb.install_extension("spatial")
-duckdb.load_extension("spatial")
-duckdb.install_extension("httpfs")
-duckdb.load_extension("httpfs")
-
-duckdb.sql("create or replace view pad as select * from read_parquet('https://data.source.coop/cboettig/pad-us-3/pad-mobi.parquet')")
-
-st.title(TITLE)
-
-# Set OpenAI API key from Streamlit secrets
-client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])
-
-# Set a default model
-if "openai_model" not in st.session_state:
-    st.session_state["openai_model"] = "gpt-3.5-turbo"
-# "gpt-4"
-
-# Initialize chat history
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-
-setup = '''
-You are a database administrator, and expert in SQL. You will be helping me write complex SQL queries. I will explain you my needs, you will generate SQL queries against my database.
-
-My application does: Conservation prioritization of protected areas to help meet US 30x30 conservation goals. 
-
-Please reply only with the SQL code that I will need to execute. Do not include an explanation of the code.
-Please reply with raw SQL, do not use markdown formatting.
-
-The database is a POSTGIS Postgres database, please take it into consideration when generating PLSQL/SQL. Please avoid ST_Within queries if possible, because they are so slow.
-
-I will provide you with a description of the structure of my tables. You must remember them and use them for generating SQL queries. Once you read them all, just remember them for future and say nothing else.
-
-Here are the tables :
-
-Table "pad"
-┌─────────────┬─────────────┬─────────┬─────────┬─────────┬─────────┐
-│ column_name │ column_type │  null   │   key   │ default │  extra  │
-│   varchar   │   varchar   │ varchar │ varchar │ varchar │ varchar │
-├─────────────┼─────────────┼─────────┼─────────┼─────────┼─────────┤
-│ FID         │ INTEGER     │ YES     │ NULL    │ NULL    │ NULL    │
-│ time        │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ rsr         │ DOUBLE      │ YES     │ NULL    │ NULL    │ NULL    │
-│ richness    │ DOUBLE      │ YES     │ NULL    │ NULL    │ NULL    │
-│ bucket      │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ FeatClass   │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ Mang_Name   │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ Mang_Type   │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ Des_Tp      │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ Pub_Access  │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ GAP_Sts     │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ IUCN_Cat    │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ Unit_Nm     │ VARCHAR     │ YES     │ NULL    │ NULL    │ NULL    │
-│ area        │ DOUBLE      │ YES     │ NULL    │ NULL    │ NULL    │
-│ geometry    │ BLOB        │ YES     │ NULL    │ NULL    │ NULL    │
-├─────────────┴─────────────┴─────────┴─────────┴─────────┴─────────┤
-└───────────────────────────────────────────────────────────────────┘
-'''
-
-# Display chat messages from history on app rerun
-#for message in st.session_state.messages:
-#    with st.chat_message(message["role"]):
-#        st.markdown(message["content"])
-
-# Accept user input
-if prompt := st.chat_input("What is the total area in each GAP_Sts?"):
-    # Add user message to chat history
-    st.session_state.messages.append({"role": "system", "content": setup})
-
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    # Display user message in chat message container
-    # with st.chat_message("user"):
-        #st.markdown(prompt)
-
-  # Display assistant response in chat message container
+from pathlib import Path
+from langchain.llms.openai import OpenAI
+from langchain.agents import create_sql_agent
+from langchain.sql_database import SQLDatabase
+from langchain.agents.agent_types import AgentType
+from langchain.callbacks import StreamlitCallbackHandler
+from langchain.agents.agent_toolkits import SQLDatabaseToolkit
+from sqlalchemy import create_engine
+import sqlite3
+import os
+from langchain_openai import ChatOpenAI
+
+
+os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
+openai_api_key = st.secrets["OPENAI_API_KEY"]
+
+st.set_page_config(page_title="Chat with Protected Areas Database", page_icon="🦜")
+st.title("🦜 LangChain: Chat with Protected Areas Database")
+
+INJECTION_WARNING = """
+                    Experimental!
+                    """
+LOCALDB = "duckdb:///pad.duckdb"
+
+# User inputs
+radio_opt = ["US Protected Areas v3"]
+selected_opt = st.sidebar.radio(label="Choose suitable option", options=radio_opt)
+if radio_opt.index(selected_opt) == 1:
+    st.sidebar.warning(INJECTION_WARNING, icon="⚠️")
+    db_uri = st.sidebar.text_input(
+        label="Database URI", placeholder="duckdb:///pad.duckdb"
+    )
+else:
+    db_uri = LOCALDB
+
+#st.sidebar.text_input(label="OpenAI API Key", type="password")
+
+
+# Setup agent
+llm = OpenAI(openai_api_key=openai_api_key, model="gpt-3.5-turbo", temperature=0, streaming=True)
+
+@st.cache_resource(ttl="2h")
+def configure_db(db_uri):        
+    return SQLDatabase.from_uri(database_uri=db_uri, view_support=True)
+
+db = configure_db(db_uri)
+
+llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+agent = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)
+
+if "messages" not in st.session_state or st.sidebar.button("Clear message history"):
+    st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
+
+for msg in st.session_state.messages:
+    st.chat_message(msg["role"]).write(msg["content"])
+
+user_query = st.chat_input(placeholder="Ask me anything!")
+
+if user_query:
+    st.session_state.messages.append({"role": "user", "content": user_query})
+    st.chat_message("user").write(user_query)
+
     with st.chat_message("assistant"):
-        stream = client.chat.completions.create(
-            model=st.session_state["openai_model"],
-            messages=[
-                {"role": m["role"], "content": m["content"]}
-                for m in st.session_state.messages
-            ],
-            stream=True,
-        )
-        response = st.write_stream(stream)
-        st.divider()
-        # st.write(response)
-        df = duckdb.sql(response).df()
-        st.table(df)
-    st.session_state.messages.append({"role": "assistant", "content": response})
-
-
-
-with st.sidebar:
-    with st.container():
-        st.title("Examples")
-        for prompt in EXAMPLE_PROMPTS:
-            st.markdown(prompt)
-
-
-            # st.button(prompt, args=(prompt,)), on_click=handle_user_input)
+        st_cb = StreamlitCallbackHandler(st.container())
+        response = agent.run(user_query, callbacks=[st_cb])
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        st.write(response)