db select option

Author: vikram-floweraura

pages/Ask Your Database.py

@@ -8,7 +8,7 @@
 from langchain.llms import OpenAI
 from langchain.callbacks import get_openai_callback
 from langchain import OpenAI, SQLDatabase, SQLDatabaseChain
-
+import psycopg2
 import openai
 import os
 
@@ -22,7 +22,8 @@
 
 with st.form("my_form"):
    st.write("Inside the form")
-   db_type    = st.text_input("select db type")
+   all_db_types = ["Mysql","Postgres","MongoDB"]
+   db_type    = st.selectbox("Select DB Type",all_db_types)
    dbusername = st.text_input("select db username")
    dbpassword = st.text_input("select db password")
    dbname     = st.text_input("select db name")
@@ -31,46 +32,41 @@
    # Every form must have a submit button.
    submitted = st.form_submit_button("Submit")
    if submitted:
-        if db_type=='mysql':
+        if db_type=='Mysql':
             conn = "mysql+pymysql://"+dbusername+":"+dbpassword+"@"+dbhost+"/"+dbname
 
-st.write("Outside the form")
-
+        if db_type=='Postgres':
+            conn = "postgresql+psycopg2://"+dbusername+":"+dbpassword+"@"+dbhost+"/"+dbname
+            
+        st.write(conn)
+        
+        if conn!='':
+            db =  SQLDatabase.from_uri(
+                conn,
+                )
+            # setup llm
+            llm = OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)
 
 
+            # Create db chain
 
+            QUERY = """
+            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.
+            Use the following format:
 
-st.write(conn)
-# db =  SQLDatabase.from_uri(
-#     "mysql+pymysql://bakewish:EmtTSRwDeKOk@admin.bakewish.in/bakewish_030123",
-#     )
-if conn!='':
-    db =  SQLDatabase.from_uri(
-        conn,
-        )
-    
-st.write("db connection established")
-#st.write(db)
-
-# setup llm
-llm = OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)
+            Question: "Question here"
+            SQLQuery: "SQL Query to run"
+            SQLResult: "Result of the SQLQuery"
+            Answer: "Final answer here"
 
-# Create db chain
+            """
+            # Setup the database chain
+            db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)
 
-QUERY = """
-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.
-Use the following format:
+            user_question = st.text_input("write a query:")
 
-Question: "Question here"
-SQLQuery: "SQL Query to run"
-SQLResult: "Result of the SQLQuery"
-Answer: "Final answer here"
+            st.write(db_chain.run(user_question))
 
-"""
-
-# Setup the database chain
-db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)
-
-user_question = st.text_input("write a query:")
+st.write("Outside the form")
 
-st.write(db_chain.run(user_question))
+#st.write(db)