ZeeshanML/Chat-With-SQL-Database

1import streamlit as st 2from pathlib import Path3from langchain_community.agent_toolkits.sql.base import create_sql_agent4from langchain_community.utilities.sql_database import SQLDatabase5from langchain.agents.agent_types import AgentType6from langchain_community.callbacks.streamlit.streamlit_callback_handler import StreamlitCallbackHandler7from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit8from sqlalchemy import create_engine9import sqlite310from langchain_groq import ChatGroq11import os12 13st.set_page_config(page_title = "LangChain: Chat with SQL DB", page_icon = "🦜")14st.title("🦜🔗 Chat with SQL DB")15 16LOCALDB = "USE_LOCALDB"17MYSQL = "USE_MYSQL"18 19radio_option = ["Use SQLite - student.db", "Connect Your MySQL Database"]20 21selected_option = st.sidebar.radio("Select the database", options=radio_option)22 23if radio_option.index(selected_option) == 1:24    db_uri = MYSQL25    mysql_host = st.sidebar.text_input("MySQL Host")26    mysql_user = st.sidebar.text_input("MySQL User")27    mysql_password = st.sidebar.text_input("MySQL Password", type="password")28    mysql_db = st.sidebar.text_input("MySQL Database")29else:30    db_uri = LOCALDB31 32 33groq_api_key = st.sidebar.text_input("Groq API Key", type="password")34 35# Check for the API key and stop execution until it's provided36if not groq_api_key:37    st.warning("Please enter your Groq API Key.")38    st.stop()39 40os.environ["GROQ_API_KEY"] = groq_api_key41 42# Initialize the LLM after the API key is set43llm = ChatGroq(groq_api_key=groq_api_key, model_name = "Llama3-8b-8192", streaming = True)44 45def configure_db(db_uri, mysql_host=None, mysql_user=None, mysql_password=None, mysql_db=None):46    if db_uri == LOCALDB:47        db_file_path = (Path(__file__).parent/"student.db").absolute()48        creator = lambda: sqlite3.connect(f"file:{db_file_path}?mode=ro", uri=True)49        return SQLDatabase(create_engine("sqlite://", creator=creator))50    elif db_uri == MYSQL:51        if not (mysql_host and mysql_user and mysql_password and mysql_db):52            st.error("Please fill all the fields.")53            st.stop()54        return SQLDatabase(create_engine(f"mysql+mysqlconnector://{mysql_user}:{mysql_password}@{mysql_host}/{mysql_db}"))55 56# Configure the database connection57if db_uri == MYSQL:58    db = configure_db(db_uri, mysql_host, mysql_user, mysql_password, mysql_db)59else:60    db = configure_db(db_uri)61 62toolkit = SQLDatabaseToolkit(db=db, llm=llm)63 64agent = create_sql_agent(65    llm=llm,66    toolkit=toolkit,67    verbose=True,68    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,69    handle_parsing_errors=True70)71 72# Initialize session state for chat messages73if "messages" not in st.session_state or st.sidebar.button("Clear Chat"):74    st.session_state["messages"] = [75        {"role": "assistant", "content": "Hi, I am your assistant who can search in the database. How can I help you?"}76    ]77 78# Display chat messages79for message in st.session_state["messages"]:80    st.chat_message(message["role"]).write(message["content"])81 82# Get user input for query83user_query = st.chat_input(placeholder="Enter your query here...")84 85if user_query:86    st.session_state["messages"].append({"role": "user", "content": user_query})87    st.chat_message("user").write(user_query)88 89    # Assistant response90    with st.chat_message("assistant"):91        streamlit_callback = StreamlitCallbackHandler(st.container())92        response = agent.run(user_query, callbacks=[streamlit_callback])93        st.session_state["messages"].append({"role": "assistant", "content": response})94        st.write(response)95