sango07/Chat-with-Database

1import streamlit as st2from langchain import OpenAI, SQLDatabase3from langchain_experimental.sql import SQLDatabaseChain4from langchain_openai import ChatOpenAI5import os6import time7from langchain_core.messages import AIMessage, HumanMessage8from langchain_core.prompts import PromptTemplate9from langchain_community.callbacks import get_openai_callback10 11# Set page configuration at the very beginning12st.set_page_config(13    page_title="PostgreSQL Query Assistant", 14    page_icon="🤖", 15    layout="wide"16)17 18# Custom CSS for enhanced styling19def local_css():20    st.markdown("""21    <style>22    .main-container {23        background-color: #f0f2f6;24        padding: 2rem;25        border-radius: 15px;26    }27    .stApp {28        background-color: #ffffff;29    }30    .stChatInput {31        border-radius: 15px !important;32        border: 2px solid #3366cc !important;33    }34    .chat-header {35        background-color: #3366cc;36        color: white;37        padding: 15px;38        border-radius: 10px;39        margin-bottom: 20px;40    }41    .chat-message {42        margin-bottom: 10px;43        padding: 10px;44        border-radius: 10px;45    }46    .human-message {47        background-color: #e6f2ff;48        border-left: 4px solid #3366cc;49    }50    .ai-message {51        background-color: #f0f0f0;52        border-left: 4px solid #666666;53    }54    .sidebar .stTextInput > div > div > input {55        border-radius: 10px !important;56    }57    .sidebar {58        background-color: #f8f9fa;59        border-radius: 15px;60        padding: 15px;61    }62    </style>63    """, unsafe_allow_html=True)64 65def init_database(user: str, password: str, host: str, port: str, database: str, sslmode: str = None):66    """Initialize a connection to the PostgreSQL database."""67    try:68        db_uri = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"69        if sslmode:70            db_uri += f"?sslmode={sslmode}"71        72        db = SQLDatabase.from_uri(db_uri)73        return db74    except Exception as e:75        st.error(f"Unable to connect to the database: {e}")76        return None77 78def answer_sql(question: str, db, chat_history: list, llm):79    """Generate SQL answer based on the user's question and database content."""80    try:81        prompt = PromptTemplate(82            input_variables=['input', 'table_info', 'top_k'],83            template="""You are a PostgreSQL expert. Given an input question,84                        first create a syntactically correct PostgreSQL query to run,85                        then look at the results of the query and return the answer to the input question.86                        Unless the user specifies in the question a specific number of records to obtain, query for at most {top_k} results using the LIMIT clause as per PostgreSQL.87                        Wrap each column name in double quotes (") to denote them as delimited identifiers.88                        Only use the following tables:\n{table_info}\n\nQuestion: {input}')"""89        )90 91        db_chain = SQLDatabaseChain(92            llm=llm, 93            database=db, 94            top_k=100, 95            verbose=True, 96            use_query_checker=True, 97            prompt=prompt, 98            return_intermediate_steps=True99        )100 101        with get_openai_callback() as cb:102            response = db_chain.invoke({103                "query": question,104                "chat_history": chat_history,105            })["result"]106 107            # Optional: Log token usage (you can remove this or add logging as needed)108            print(f"Total Tokens: {cb.total_tokens}")109            print(f"Total Cost (USD): ${cb.total_cost}")110 111        return response112    except Exception as e:113        st.error(f"An error occurred: {e}")114        return "Sorry, I couldn't process your request."115 116def main():117    # Apply custom CSS118    local_css()119 120    # Main container121    with st.container():122        # Header123        st.markdown("<div class='chat-header'><h1 style='text-align: center;'>🤖 PostgreSQL Query Assistant</h1></div>", unsafe_allow_html=True)124 125        # Sidebar for connection126        with st.sidebar:127            st.image("https://www.postgresql.org/media/img/about/press/elephant.png", use_container_width=True)128            st.header("Database Connection")129            130            # Connection details131            with st.expander("Database Credentials", expanded=True):132                openai_api_key = st.text_input("OpenAI API Key", type="password", help="Required for natural language to SQL conversion")133                134                db_type = st.radio("Database Type", ("Local", "Cloud"))135                136                if db_type == "Local":137                    host = st.text_input("Host", value="localhost")138                    port = st.text_input("Port", value="5432")139                    user = st.text_input("Username", value="postgres")140                    password = st.text_input("Password", type="password")141                    database = st.text_input("Database Name", value="testing_3")142                    sslmode = None143                else:144                    host = st.text_input("Host (e.g., your-db-host.aws.com)")145                    port = st.text_input("Port", value="5432")146                    user = st.text_input("Username")147                    password = st.text_input("Password", type="password")148                    database = st.text_input("Database Name")149                    sslmode = st.selectbox("SSL Mode", ["require", "verify-ca", "verify-full", "disable"])150 151                connect_btn = st.button("🔌 Connect to Database")152 153        # Main chat area154        chat_container = st.container()155 156        # Initialize or load session state157        if 'chat_history' not in st.session_state:158            st.session_state.chat_history = [159                AIMessage(content="👋 Hi there! I'm your PostgreSQL Query Assistant. Connect to your database and ask me anything!")160            ]161        162        if 'db_connected' not in st.session_state:163            st.session_state.db_connected = False164 165        # Connection handling166        if connect_btn:167            if not openai_api_key:168                st.error("Please provide an OpenAI API Key")169            else:170                os.environ["OPENAI_API_KEY"] = openai_api_key171                llm = ChatOpenAI(temperature=0.7, model="gpt-3.5-turbo")172                173                db = init_database(user, password, host, port, database, sslmode)174                175                if db:176                    st.session_state.db = db177                    st.session_state.llm = llm178                    st.session_state.db_connected = True179                    st.success("🎉 Successfully connected to the database!")180 181        # Display chat history182        with chat_container:183            for message in st.session_state.chat_history:184                if isinstance(message, AIMessage):185                    with st.chat_message("assistant", avatar="🤖"):186                        st.markdown(f"<div class='chat-message ai-message'>{message.content}</div>", unsafe_allow_html=True)187                elif isinstance(message, HumanMessage):188                    with st.chat_message("user", avatar="👤"):189                        st.markdown(f"<div class='chat-message human-message'>{message.content}</div>", unsafe_allow_html=True)190 191        # Chat input and processing192        if st.session_state.db_connected:193            user_query = st.chat_input("Ask a question about your database...")194            195            if user_query:196                # Add user message to chat history197                st.session_state.chat_history.append(HumanMessage(content=user_query))198                199                # Display user message200                with st.chat_message("user", avatar="👤"):201                    st.markdown(f"<div class='chat-message human-message'>{user_query}</div>", unsafe_allow_html=True)202                203                # Generate and display AI response204                with st.chat_message("assistant", avatar="🤖"):205                    with st.spinner("Generating response..."):206                        response = answer_sql(207                            user_query, 208                            st.session_state.db, 209                            st.session_state.chat_history, 210                            st.session_state.llm211                        )212                    st.markdown(f"<div class='chat-message ai-message'>{response}</div>", unsafe_allow_html=True)213                214                # Add AI response to chat history215                st.session_state.chat_history.append(AIMessage(content=response))216        else:217            st.warning("Please connect to a database to start querying.")218 219if __name__ == "__main__":220    main()