File size: 4,059 Bytes
a0a4b07 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import streamlit as st
from pathlib import Path
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_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from sqlalchemy import create_engine
import sqlite3
from langchain_groq import ChatGroq
st.set_page_config(page_title="LangChain: Chat with SQL DB", page_icon="🦜")
st.title("🦜 LangChain: Chat with SQL DB")
LOCALDB = "USE_LOCALDB"
MYSQL = "USE_MYSQL"
radio_opt = ["Use SQLLite 3 Database- STUDENT.db", "Connect to you MySQL Database"]
selected_opt = st.sidebar.radio(label="Choose the DB which you want to chat", options=radio_opt)
if radio_opt.index(selected_opt) == 1:
db_uri = MYSQL
mysql_host = st.sidebar.text_input("Provide MySQL Host")
mysql_user = st.sidebar.text_input("MYSQL User")
mysql_password = st.sidebar.text_input("MYSQL password (leave blank if no password)", type="password")
mysql_db = st.sidebar.text_input("MySQL database")
else:
db_uri = LOCALDB
# Get API key from user input
api_key = st.sidebar.text_input(label="GROQ API Key", type="password")
if not api_key:
st.info("Please add the Groq API key to continue")
st.stop()
if db_uri == MYSQL and not all([mysql_host, mysql_user, mysql_db]):
st.info("Please enter all required MySQL database information")
st.stop()
# LLM model - using user-provided API key
llm = ChatGroq(groq_api_key=api_key, model_name="llama-3.1-8b-instant", streaming=True)
@st.cache_resource(ttl="2h")
def configure_db(db_uri, mysql_host=None, mysql_user=None, mysql_password=None, mysql_db=None):
if db_uri == LOCALDB:
dbfilepath = (Path(__file__).parent/"STUDENT.db").absolute()
print(dbfilepath)
creator = lambda: sqlite3.connect(f"file:{dbfilepath}?mode=ro", uri=True)
return SQLDatabase(create_engine("sqlite:///", creator=creator))
elif db_uri == MYSQL:
if not (mysql_host and mysql_user and mysql_db):
st.error("Please provide MySQL host, user, and database name.")
st.stop()
# Handle optional password
if mysql_password:
connection_string = f"mysql+mysqlconnector://{mysql_user}:{mysql_password}@{mysql_host}/{mysql_db}"
else:
connection_string = f"mysql+mysqlconnector://{mysql_user}@{mysql_host}/{mysql_db}"
return SQLDatabase(create_engine(connection_string))
# Configure database
try:
if db_uri == MYSQL:
db = configure_db(db_uri, mysql_host, mysql_user, mysql_password, mysql_db)
else:
db = configure_db(db_uri)
except Exception as e:
st.error(f"Error connecting to database: {str(e)}")
st.stop()
# Toolkit
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent = create_sql_agent(
llm=llm,
toolkit=toolkit,
verbose=True,
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION
)
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 anything from the database")
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"):
try:
streamlit_callback = StreamlitCallbackHandler(st.container())
response = agent.run(user_query, callbacks=[streamlit_callback])
st.session_state.messages.append({"role": "assistant", "content": response})
st.write(response)
except Exception as e:
error_message = f"An error occurred: {str(e)}"
st.error(error_message)
st.session_state.messages.append({"role": "assistant", "content": error_message}) |