Update src/streamlit_app.py

src/streamlit_app.py

@@ -0,0 +1,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})