app.py

import streamlit as st
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"):
        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)