Update app.py

app.py

@@ -1,12 +1,11 @@
 from langchain import OpenAI, SQLDatabase
 from langchain_experimental.sql import SQLDatabaseChain
-# from langchain_openai import AzureChatOpenAI,ChatOpenAI
+from langchain_openai import AzureChatOpenAI, ChatOpenAI
 import pandas as pd
 import time
 from langchain_core.prompts.prompt import PromptTemplate
 import re
 from sqlalchemy import create_engine, text
-import pandas as pd
 import psycopg2
 from psycopg2 import sql
 import streamlit as st
@@ -15,138 +14,116 @@ from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.runnables import RunnablePassthrough
 from langchain_core.output_parsers import StrOutputParser
 from langchain_groq import ChatGroq
-import os
 from langchain_community.callbacks import get_openai_callback
-
 import os
-from langchain_groq import ChatGroq
-os.environ["GROQ_API_KEY"]="gsk_......................"
+
+os.environ["GROQ_API_KEY"] = "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"
 llm = ChatGroq(model="llama-3.1-70b-versatile", temperature=0.25)
 
-def init_database(user: str, password: str, host: str, port: str, database: str) -> SQLDatabase:
-  db_uri = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"
-  return SQLDatabase.from_uri(db_uri)
 
+def init_database(user: str, password: str, host: str, port: str, database: str, sslmode: str = None) -> SQLDatabase:
+    try:
+        db_uri = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"
+        if sslmode:
+            db_uri += f"?sslmode={sslmode}"
+        
+        # Attempt to create a database connection
+        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
 
-def answer_sql(question: str, db: SQLDatabase, chat_history: list):    
 
+def answer_sql(question: str, db: SQLDatabase, chat_history: list):
     try:
-
         # setup llm
         llm = ChatGroq(model="llama-3.1-70b-versatile", temperature=0.25)
 
-
-        #There is a table named "data_description" in the database, this table give details about all other tables & columns it contains. Use this information to write a query.
-    
-
-        prompt=PromptTemplate(input_variables=['input', 'table_info', 'top_k'], 
-                        template="""You are a PostgreSQL expert. Given an input question, 
-                        first create a syntactically correct PostgreSQL query to run, 
-                        then look at the results of the query and return the answer to the input question. 
-                        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. 
-                        You can order the results to return the most informative data in the database.\n
-                        Never query for all columns from a table. You must query only the columns that are needed to answer the question.
+        prompt = PromptTemplate(input_variables=['input', 'table_info', 'top_k'],
+                        template="""You are a PostgreSQL expert. Given an input question,
+                        first create a syntactically correct PostgreSQL query to run,
+                        then look at the results of the query and return the answer to the input question.
+                        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.
                         Wrap each column name in double quotes (") to denote them as delimited identifiers.
-                        Pay attention to use only the column names you can see in the tables below.
-                        Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.
-                        Pay attention to use CURRENT_DATE function to get the current date, if the question involves "today".
-                        Use the following format:\
-                            Question: Question here
-                            SQLQuery: SQL Query to run
-                            SQLResult: Result of the SQLQuery
-                            Answer: Final answer here
-                            Only use the following tables:\n{table_info}\n\nQuestion: {input}')""")
-
- 
-        QUERY = """
-
-        Given an input question, look at the results of the query and return the answer in natural language to the users question with all the records of SQLResult. Be careful not to truncate the records in output while returning answer. Pay attention to return answer in tabular format only.
-
-        Use the following format:
-
-        Question: Question here
-        SQLQuery: SQL Query to run
-        SQLResult: Result of the SQLQuery
-        Answer: Final answer here
+                        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_time_start = time.time() #start time of db
-
-        # Setup the database chain
-        db_chain = SQLDatabaseChain(llm=llm, database=db,top_k=100,verbose=True,use_query_checker=True,prompt=prompt,return_intermediate_steps=True) # verbose=True
-        
-        db_chain_time_end = time.time() #end time of db
-
-        question = QUERY.format(question=question)
-
+        db_chain = SQLDatabaseChain(
+            llm=llm, database=db, top_k=100, verbose=True, use_query_checker=True, prompt=prompt, return_intermediate_steps=True
+        )
 
         with get_openai_callback() as cb:
-
-            response_time_start = time.time()
-
             response = db_chain.invoke({
-                "query": question,
+                "query": QUERY.format(question=question),
                 "chat_history": chat_history,
             })["result"]
 
-            response_time_end = time.time()
-
-            
-
-            token_info = cb
-            print("*"*55)
-            print()
-            print(f"Overall_response_execution_time : {response_time_end-response_time_start}")
+            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()
-            print("*"*55)
+            print("*" * 55)
 
         return response
-    
+
     except Exception as e:
-        st.error("Some technical error occured. Please try again after some time!")
-        st.stop()  # Stop further execution if another error occurs
+        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 a your SQL assistant. Ask me anything about your database."),
+      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)
 
+# Step 1: Prompt user to select database type (local or cloud)
 with st.sidebar:
-    st.subheader("Postgres Credentials")
-    st.write("Enter your Credentials & Connect")
-    
-    st.text_input("Host", value="localhost", key="Host")
-    st.text_input("Port", value="5432", key="Port")
-    st.text_input("User", value="postgres", key="User")
-    st.text_input("Password", type="password", value="QKadmin", key="Password")
-    st.text_input("Database", value="testing_3", key="Database")
-    
-    if st.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!")
-    
+    st.subheader("Database Setup")
+    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")
+
+        # Connect Button
+        if st.button("Connect"):
+            with st.spinner("Connecting to the local database..."):
+                db = init_database(user, password, host, port, database)
+                st.session_state.db = db
+                st.success("Connected to local database!")
+
+    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 (default: 5432)", 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"])
+
+        # Connect Button
+        if st.button("Connect"):
+            with st.spinner("Connecting to the cloud database..."):
+                db = init_database(user, password, host, port, database, sslmode)
+                st.session_state.db = db
+                st.success("Connected to cloud database!")
+
+# Main chat interface
 for message in st.session_state.chat_history:
     if isinstance(message, AIMessage):
         with st.chat_message("AI"):
@@ -156,14 +133,13 @@ for message in st.session_state.chat_history:
             st.markdown(message.content)
 
 user_query = st.chat_input("Type a message...")
-if user_query is not None and user_query.strip() != "":
+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.markdown(response)
-        
+
     st.session_state.chat_history.append(AIMessage(content=response))