app.py

import streamlit as st
import os
from dotenv import load_dotenv
from langchain_core.messages import AIMessage, HumanMessage
from langchain.prompts import PromptTemplate
from utilities.text_processor import process_texts, chat_with_text_files
from utilities.pdf_processor import process_pdfs
from utilities.csv_processor import process_csv, generate_insights
from utilities.excel_processor import process_excel
from utilities.sql import init_database, SQLChain

# Load environment variables
load_dotenv()

# Page config
st.set_page_config(page_title="Chat with QueryBot", page_icon=":speech_balloon:", layout="wide")

# Initialize session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = [
        AIMessage(content="Hello! I'm your assistant named QueryBot. Click on the settings icon then connect to the database and start chatting."),
    ]
if "show_settings" not in st.session_state:
    st.session_state.show_settings = False

# Function to handle specific words responses
specific_words_responses = {
    "hello querybot": "Hello! How can I assist you further?",
    "hi": "Hello! How can I help you today?",
    "hello": "Hello! How can I help you today?",
    "thank you": "You're welcome! If you have any other questions, feel free to ask.",
    "thanks": "You're welcome! If you have any other questions, feel free to ask.",
    "bye": "Goodbye! Have a great day!",
    "exit": "Goodbye! Have a great day!"
}

# Toggle sidebar visibility
if st.button("⚙️ Settings"):
    st.session_state.show_settings = not st.session_state.show_settings

# Sidebar settings
if st.session_state.show_settings:
    st.sidebar.subheader("Settings")
    st.sidebar.write("This is a chat application. Select a database type, connect, and start chatting.")

    # Dropdown for selecting database type
    db_type = st.sidebar.selectbox("Database Type", ["Select a database", "MySQL", "CSV", "Excel", "PDF", "TXT"])

    # Database connection fields based on selected type
    if db_type == "MySQL":
        st.sidebar.text_input("Host", value="localhost", key="Host")
        st.sidebar.text_input("Port", value="3306", key="Port")
        st.sidebar.text_input("User", value="root", key="User")
        st.sidebar.text_input("Password", type="password", value="passcode", key="Password")
        st.sidebar.text_input("Database", value="RestaurantMenu", key="Database")
        if st.sidebar.button("Connect"):
            with st.spinner("Connecting to database..."):
                db = init_database(
                    st.session_state["User"],
                    st.session_state["Password"],
                    st.session_state["Host"],
                    st.session_state["Port"],
                    st.session_state["Database"]
                )
                st.session_state.db = db
                st.success("Connected to database!")

    elif db_type == "CSV":
        csv_files = st.sidebar.file_uploader("Upload CSV files", type=["csv"], accept_multiple_files=True)
        if csv_files:
            with st.spinner("Processing CSV files..."):
                combined_df, dfs = process_csv(csv_files)
                st.session_state.dfs = dfs
                st.session_state.df = combined_df
                st.success("CSV files processed!")

    elif db_type == "Excel":
        excel_files = st.sidebar.file_uploader("Upload Excel files", type=["xls", "xlsx"], accept_multiple_files=True)
        if excel_files:
            with st.spinner("Processing Excel files..."):
                combined_df, dfs = process_excel(excel_files)
                st.session_state.dfs = dfs
                st.session_state.df = combined_df
                st.success("Excel files processed!")

    elif db_type == "PDF":
        pdf_files = st.sidebar.file_uploader("Upload PDF files", type=["pdf"], accept_multiple_files=True)
        if pdf_files:
            with st.spinner("Processing PDF files..."):
                combined_text, texts = process_pdfs(pdf_files)
                st.session_state.texts = texts
                st.success("PDF files processed!")

    elif db_type == "TXT":
        text_files = st.sidebar.file_uploader("Upload TXT files", type=["txt"], accept_multiple_files=True)
        if text_files:
            with st.spinner("Processing text files..."):
                combined_text, texts = process_texts(text_files)
                st.session_state.texts = texts
                st.success("Text files processed!")

# Show chat history
for chat in st.session_state.chat_history:
    if isinstance(chat, HumanMessage):
        with st.chat_message("Human"):
            st.markdown(f"✍️ {chat.content}")
    elif isinstance(chat, AIMessage):
        with st.chat_message("AI"):
            st.markdown(f"🟢 {chat.content}")

# Handle user input and interaction
if user_input := st.chat_input(placeholder="Ask your query..."):
    st.session_state.chat_history.append(HumanMessage(content=user_input))
    with st.chat_message("Human"):
        st.markdown(f"✍️ {user_input}")

    response = specific_words_responses.get(user_input.lower())
    if response is None:
        with st.spinner("Processing..."):
            if db_type in ["TXT", "PDF"]:
                if "texts" in st.session_state:
                    texts = st.session_state.texts
                    response = chat_with_text_files(user_input, texts)
                else:
                    response = "No texts available for processing. Please upload text or PDF files."
            elif db_type in ["CSV", "Excel"]:
                if "df" in st.session_state:
                    df = st.session_state.df
                    response = generate_insights(user_input, df)
                else:
                    response = "No CSV or Excel file processed. Please upload CSV or Excel files."
            elif db_type == "MySQL":
                if "db" in st.session_state:
                    chain = SQLChain(st.session_state.db)
                    response = chain.invoke_chain(user_input, st.session_state.chat_history)
                else:
                    response = "Database not connected. Please check your settings."
            else:
                response = "No valid database or files loaded. Please check your settings."

    st.session_state.chat_history.append(AIMessage(content=response))
    with st.chat_message("AI"):
        st.markdown(f"🟢 {response}")