Update app.py

app.py

@@ -1,26 +1,61 @@
+import streamlit as st
 from langchain import OpenAI, SQLDatabase
 from langchain_experimental.sql import SQLDatabaseChain
-from langchain_openai import AzureChatOpenAI, ChatOpenAI
-import pandas as pd
+from langchain_openai import ChatOpenAI
+import os
 import time
-from langchain_core.prompts.prompt import PromptTemplate
-import re
-from sqlalchemy import create_engine, text
-import psycopg2
-from psycopg2 import sql
-import streamlit as st
 from langchain_core.messages import AIMessage, HumanMessage
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import RunnablePassthrough
-from langchain_core.output_parsers import StrOutputParser
-from langchain_groq import ChatGroq
+from langchain_core.prompts import PromptTemplate
 from langchain_community.callbacks import get_openai_callback
-import os
-from langchain_openai import ChatOpenAI
-# llm = ChatOpenAI(temperature=0.7, model="gpt-3.5-turbo")
-
 
-def init_database(user: str, password: str, host: str, port: str, database: str, sslmode: str = None) -> SQLDatabase:
+# Custom CSS for enhanced styling
+def local_css():
+    st.markdown("""
+    <style>
+    .main-container {
+        background-color: #f0f2f6;
+        padding: 2rem;
+        border-radius: 15px;
+    }
+    .stApp {
+        background-color: #ffffff;
+    }
+    .stChatInput {
+        border-radius: 15px !important;
+        border: 2px solid #3366cc !important;
+    }
+    .chat-header {
+        background-color: #3366cc;
+        color: white;
+        padding: 15px;
+        border-radius: 10px;
+        margin-bottom: 20px;
+    }
+    .chat-message {
+        margin-bottom: 10px;
+        padding: 10px;
+        border-radius: 10px;
+    }
+    .human-message {
+        background-color: #e6f2ff;
+        border-left: 4px solid #3366cc;
+    }
+    .ai-message {
+        background-color: #f0f0f0;
+        border-left: 4px solid #666666;
+    }
+    .sidebar .stTextInput > div > div > input {
+        border-radius: 10px !important;
+    }
+    .sidebar {
+        background-color: #f8f9fa;
+        border-radius: 15px;
+        padding: 15px;
+    }
+    </style>
+    """, unsafe_allow_html=True)
+
+def init_database(user: str, password: str, host: str, port: str, database: str, sslmode: str = None):
     """Initialize a connection to the PostgreSQL database."""
     try:
         db_uri = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"
@@ -30,10 +65,10 @@ def init_database(user: str, password: str, host: str, port: str, database: str,
         db = SQLDatabase.from_uri(db_uri)
         return db
     except Exception as e:
-        st.error("Unable to connect to the database. Please check your credentials and try again.")
-        st.stop()  # Stop further execution if an error occurs
+        st.error(f"Unable to connect to the database: {e}")
+        return None
 
-def answer_sql(question: str, db: SQLDatabase, chat_history: list, llm) -> str:
+def answer_sql(question: str, db, chat_history: list, llm):
     """Generate SQL answer based on the user's question and database content."""
     try:
         prompt = PromptTemplate(
@@ -46,11 +81,6 @@ def answer_sql(question: str, db: SQLDatabase, chat_history: list, llm) -> str:
                         Only use the following tables:\n{table_info}\n\nQuestion: {input}')"""
         )
 
-        QUERY = f"""
-        Given an input question, look at the results of the query and return the answer in natural language to the user's question with all the records of SQLResult.
-        {question}
-        """
-
         db_chain = SQLDatabaseChain(
             llm=llm, 
             database=db, 
@@ -63,87 +93,128 @@ def answer_sql(question: str, db: SQLDatabase, chat_history: list, llm) -> str:
 
         with get_openai_callback() as cb:
             response = db_chain.invoke({
-                "query": QUERY.format(question=question),
+                "query": question,
                 "chat_history": chat_history,
             })["result"]
 
-            print("*" * 55)
-            print(f"Total Tokens : {cb.total_tokens}")
-            print(f"Prompt Tokens : {cb.prompt_tokens}")
-            print(f"Completion Tokens : {cb.completion_tokens}")
-            print(f"Total Cost (USD) : ${cb.total_cost}")
-            print("*" * 55)
+            # Optional: Log token usage (you can remove this or add logging as needed)
+            print(f"Total Tokens: {cb.total_tokens}")
+            print(f"Total Cost (USD): ${cb.total_cost}")
 
         return response
     except Exception as e:
-        st.error("A technical error occurred. Please try again later.")
-        st.stop()
-
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = [
-      AIMessage(content="Hello! I'm your SQL assistant. Ask me anything about your database."),
-    ]
-
-st.set_page_config(page_title="Chat with Postgres", page_icon=":speech_balloon:")
-st.title("Chat with Postgres DB")
-st.sidebar.image("https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcSfbBOY1t6ZMwLejpwbGVQ9p3LKplwt45yxEzeDsEEPibRm4JqIYF3xav53PNRLJwWkdw&usqp=CAU", use_container_width=True)
-
-# Get API key from user
-with st.sidebar:
-    st.subheader("API Key and Database Credentials")
-    
-    # Take OpenAI API key from the user
-    openai_api_key = st.text_input("Enter your OpenAI API Key:", type="password")
-    
-    # Database connection fields
-    db_type = st.radio("Is your PostgreSQL database on a local server or in the cloud?", ("Local", "Cloud"))
-    if db_type == "Local":
-        st.write("Enter your local database credentials.")
-        host = st.text_input("Host", value="localhost")
-        port = st.text_input("Port", value="5432")
-        user = st.text_input("User", value="postgres")
-        password = st.text_input("Password", type="password")
-        database = st.text_input("Database", value="testing_3")
-    elif db_type == "Cloud":
-        st.write("Enter your cloud database credentials.")
-        host = st.text_input("Host (e.g., your-db-host.aws.com)")
-        port = st.text_input("Port", value="5432")
-        user = st.text_input("User")
-        password = st.text_input("Password", type="password")
-        database = st.text_input("Database")
-        sslmode = st.selectbox("SSL Mode", ["require", "verify-ca", "verify-full", "disable"])
-
-    if st.button("Connect"):
-        if openai_api_key:
-            os.environ["OPENAI_API_KEY"] = openai_api_key  # Set the OpenAI API key in the environment
-            llm = ChatOpenAI(temperature=0.7, model="gpt-3.5-turbo") # Initialize model with user's API key
-            try:
-                db = init_database(user, password, host, port, database, sslmode if db_type == "Cloud" else None)
-                st.session_state.db = db
-                st.session_state.llm = llm
-                st.success("Connected to the database!")
-            except Exception as e:
-                st.error("Failed to connect to the database. Please check your details and try again.")
+        st.error(f"An error occurred: {e}")
+        return "Sorry, I couldn't process your request."
+
+def main():
+    # Apply custom CSS
+    local_css()
+
+    # Set page configuration
+    st.set_page_config(
+        page_title="PostgreSQL Query Assistant", 
+        page_icon="🤖", 
+        layout="wide"
+    )
+
+    # Main container
+    with st.container():
+        # Header
+        st.markdown("<div class='chat-header'><h1 style='text-align: center;'>🤖 PostgreSQL Query Assistant</h1></div>", unsafe_allow_html=True)
+
+        # Sidebar for connection
+        with st.sidebar:
+            st.image("https://www.postgresql.org/media/img/about/press/elephant.png", use_container_width=True)
+            st.header("Database Connection")
+            
+            # Connection details
+            with st.expander("Database Credentials", expanded=True):
+                openai_api_key = st.text_input("OpenAI API Key", type="password", help="Required for natural language to SQL conversion")
+                
+                db_type = st.radio("Database Type", ("Local", "Cloud"))
+                
+                if db_type == "Local":
+                    host = st.text_input("Host", value="localhost")
+                    port = st.text_input("Port", value="5432")
+                    user = st.text_input("Username", value="postgres")
+                    password = st.text_input("Password", type="password")
+                    database = st.text_input("Database Name", value="testing_3")
+                    sslmode = None
+                else:
+                    host = st.text_input("Host (e.g., your-db-host.aws.com)")
+                    port = st.text_input("Port", value="5432")
+                    user = st.text_input("Username")
+                    password = st.text_input("Password", type="password")
+                    database = st.text_input("Database Name")
+                    sslmode = st.selectbox("SSL Mode", ["require", "verify-ca", "verify-full", "disable"])
+
+                connect_btn = st.button("🔌 Connect to Database")
+
+        # Main chat area
+        chat_container = st.container()
+
+        # Initialize or load session state
+        if 'chat_history' not in st.session_state:
+            st.session_state.chat_history = [
+                AIMessage(content="👋 Hi there! I'm your PostgreSQL Query Assistant. Connect to your database and ask me anything!")
+            ]
+        
+        if 'db_connected' not in st.session_state:
+            st.session_state.db_connected = False
+
+        # Connection handling
+        if connect_btn:
+            if not openai_api_key:
+                st.error("Please provide an OpenAI API Key")
+            else:
+                os.environ["OPENAI_API_KEY"] = openai_api_key
+                llm = ChatOpenAI(temperature=0.7, model="gpt-3.5-turbo")
+                
+                db = init_database(user, password, host, port, database, sslmode)
+                
+                if db:
+                    st.session_state.db = db
+                    st.session_state.llm = llm
+                    st.session_state.db_connected = True
+                    st.success("🎉 Successfully connected to the database!")
+
+        # Display chat history
+        with chat_container:
+            for message in st.session_state.chat_history:
+                if isinstance(message, AIMessage):
+                    with st.chat_message("assistant", avatar="🤖"):
+                        st.markdown(f"<div class='chat-message ai-message'>{message.content}</div>", unsafe_allow_html=True)
+                elif isinstance(message, HumanMessage):
+                    with st.chat_message("user", avatar="👤"):
+                        st.markdown(f"<div class='chat-message human-message'>{message.content}</div>", unsafe_allow_html=True)
+
+        # Chat input and processing
+        if st.session_state.db_connected:
+            user_query = st.chat_input("Ask a question about your database...")
+            
+            if user_query:
+                # Add user message to chat history
+                st.session_state.chat_history.append(HumanMessage(content=user_query))
+                
+                # Display user message
+                with st.chat_message("user", avatar="👤"):
+                    st.markdown(f"<div class='chat-message human-message'>{user_query}</div>", unsafe_allow_html=True)
+                
+                # Generate and display AI response
+                with st.chat_message("assistant", avatar="🤖"):
+                    with st.spinner("Generating response..."):
+                        response = answer_sql(
+                            user_query, 
+                            st.session_state.db, 
+                            st.session_state.chat_history, 
+                            st.session_state.llm
+                        )
+                    st.markdown(f"<div class='chat-message ai-message'>{response}</div>", unsafe_allow_html=True)
+                
+                # Add AI response to chat history
+                st.session_state.chat_history.append(AIMessage(content=response))
         else:
-            st.error("Please enter your OpenAI API key.")
-
-# Main chat interface
-for message in st.session_state.chat_history:
-    if isinstance(message, AIMessage):
-        with st.chat_message("AI"):
-            st.markdown(message.content)
-    elif isinstance(message, HumanMessage):
-        with st.chat_message("Human"):
-            st.markdown(message.content)
-
-user_query = st.chat_input("Type a message...")
-if user_query:
-    st.session_state.chat_history.append(HumanMessage(content=user_query))
-    with st.chat_message("Human"):
-        st.markdown(user_query)
-
-    with st.chat_message("AI"):
-        response = answer_sql(user_query, st.session_state.db, st.session_state.chat_history, st.session_state.llm)
-        st.markdown(response)
-
-    st.session_state.chat_history.append(AIMessage(content=response))
+            st.warning("Please connect to a database to start querying.")
+
+if __name__ == "__main__":
+    main()