feat/postgresql-integration (#25)

- sql integration (5acf39cc95da88ec7b5ae96fd388b6b40f5a1606)

app.py

@@ -1,6 +1,6 @@
 from utils import TEMP_DIR, message_dict
 import gradio as gr
-import data_file
+import templates.data_file as data_file, templates.sql_db as sql_db
 
 import os
 from getpass import getpass
@@ -18,7 +18,7 @@ def delete_db(req: gr.Request):
 if "OPENAI_API_KEY" not in os.environ:
     os.environ["OPENAI_API_KEY"] = getpass("Enter OpenAI API key:")
 
-css= ".file_marker .large{min-height:50px !important;} .example_btn{max-width:300px;} .padding{padding:0;} .description_component{overflow:visible !important;}"
+css= ".file_marker .large{min-height:50px !important;} .padding{padding:0;} .description_component{overflow:visible !important;}"
 head = """<meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <title>Virtual Data Analyst</title>
@@ -72,9 +72,8 @@ with gr.Blocks(theme=theme, css=css, head=head, delete_cache=(3600,3600)) as dem
                         </main>""")
     with gr.Tab("Data File"):    
         data_file.demo.render()
-    with gr.Tab("SQL Database - Coming Soon", interactive=False):
-        gr.Text("COMING SOON")
-    #    sql_db.demo.render()
+    with gr.Tab("SQL Database"):
+        sql_db.demo.render()
 
     footer = gr.HTML("""<!-- Footer -->
         <footer class="max-w-4xl mx-auto mt-12 text-center text-gray-500 text-sm">

data_sources/__init__.py

@@ -1,3 +1,4 @@
 from .upload_file import process_data_upload
+from .connect_sql_db import connect_sql_db
 
-__all__ = ["process_data_upload"]
+__all__ = ["process_data_upload","connect_sql_db"]

data_sources/connect_sql_db.py

@@ -0,0 +1,42 @@
+import psycopg2
+import os
+from utils import TEMP_DIR
+
+def connect_sql_db(url, sql_user, sql_port, sql_pass, sql_db_name, session_hash):
+    try:
+        conn = psycopg2.connect(
+            database=sql_db_name,
+            user=sql_user,
+            password=sql_pass,
+            host=url,  # e.g., "localhost" or an IP address
+            port=sql_port  # default is 5432
+        )
+        print("Connected to PostgreSQL")
+    
+        # Create a cursor object to execute SQL queries
+        cur = conn.cursor()
+        # Example: Execute a query
+        cur.execute("""SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'""")
+        table_tuples = cur.fetchall()
+        table_names = []
+        for table in table_tuples:
+            table_names.append(table[0])
+
+        print(table_names)
+    
+        # Close the cursor and connection
+        cur.close()
+        conn.close()
+        print("Connection closed.")
+
+        session_path = 'sql'
+
+        dir_path = TEMP_DIR / str(session_hash) / str(session_path)
+        os.makedirs(dir_path, exist_ok=True)
+
+        return ["success","<p style='color:green;text-align:center;font-size:18px;'>SQL database connected successful</p>", table_names]
+    except Exception as e:
+        print("UPLOAD ERROR")
+        print(e)
+        return ["error",f"<p style='color:red;text-align:center;font-size:18px;font-weight:bold;'>ERROR: {e}</p>"]
+        

data_sources/upload_file.py

@@ -74,7 +74,9 @@ def process_data_upload(data_file, session_hash):
             if df[column].dtype == 'object' and isinstance(df[column].iloc[0], list):
                 df[column] = df[column].explode()
 
-        dir_path = TEMP_DIR / str(session_hash)
+        session_path = 'file_upload'
+
+        dir_path = TEMP_DIR / str(session_hash) / str(session_path)
         os.makedirs(dir_path, exist_ok=True)
 
         connection = sqlite3.connect(f'{dir_path}/data_source.db')

functions/__init__.py

@@ -1,9 +1,9 @@
-from .sqlite_functions import SQLiteQuery, sqlite_query_func
+from .query_functions import SQLiteQuery, sqlite_query_func, PostgreSQLQuery, sql_query_func
 from .chart_functions import table_generation_func, scatter_chart_generation_func, \
 line_chart_generation_func, bar_chart_generation_func, pie_chart_generation_func, histogram_generation_func, scatter_chart_fig
-from .chat_functions import example_question_generator, chatbot_with_fc
+from .chat_functions import sql_example_question_generator, example_question_generator, chatbot_with_fc, sql_chatbot_with_fc
 from .stat_functions import regression_func
 
-__all__ = ["SQLiteQuery","sqlite_query_func","table_generation_func","scatter_chart_generation_func",
+__all__ = ["SQLiteQuery","sqlite_query_func","sql_query_func","table_generation_func","scatter_chart_generation_func",
            "line_chart_generation_func","bar_chart_generation_func","regression_func", "pie_chart_generation_func", "histogram_generation_func",
-           "scatter_chart_fig","example_question_generator","chatbot_with_fc"]
+           "scatter_chart_fig","sql_example_question_generator","example_question_generator","chatbot_with_fc","sql_chatbot_with_fc"]

functions/chart_functions.py

@@ -92,11 +92,11 @@ def scatter_chart_fig(df, x_column: List[str], y_column: str, category: str="",
 
    return fig
 
-def scatter_chart_generation_func(x_column: List[str], y_column: str, session_hash, data: List[dict]=[{}], layout: List[dict]=[{}], 
+def scatter_chart_generation_func(x_column: List[str], y_column: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], 
                                   category: str="", trendline: str="", trendline_options: List[dict]=[{}], marginal_x: str="", marginal_y: str="",
-                                  size: str=""):
+                                  size: str="", **kwargs):
    try:
-      dir_path = TEMP_DIR / str(session_hash)
+      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
       chart_path = f'{dir_path}/chart.html'
       csv_query_path = f'{dir_path}/query.csv'
 
@@ -129,7 +129,7 @@ def scatter_chart_generation_func(x_column: List[str], y_column: str, session_ha
    
       pio.write_html(fig, chart_path, full_html=False)
 
-      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/chart.html'
+      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
 
       iframe = '<div style=overflow:auto;><iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="' + chart_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n</div>'
 
@@ -144,10 +144,10 @@ def scatter_chart_generation_func(x_column: List[str], y_column: str, session_ha
             """
       return {"reply": reply}
    
-def line_chart_generation_func(x_column: str, y_column: str, session_hash, data: List[dict]=[{}], layout: List[dict]=[{}], 
-                                  category: str=""):
+def line_chart_generation_func(x_column: str, y_column: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], 
+                                  category: str="", **kwargs):
    try:
-      dir_path = TEMP_DIR / str(session_hash)
+      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
       chart_path = f'{dir_path}/chart.html'
       csv_query_path = f'{dir_path}/query.csv'
 
@@ -180,7 +180,7 @@ def line_chart_generation_func(x_column: str, y_column: str, session_hash, data:
    
       pio.write_html(fig, chart_path, full_html=False)
 
-      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/chart.html'
+      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
 
       iframe = '<div style=overflow:auto;><iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="' + chart_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n</div>'
 
@@ -195,10 +195,10 @@ def line_chart_generation_func(x_column: str, y_column: str, session_hash, data:
             """
       return {"reply": reply}
 
-def bar_chart_generation_func(x_column: str, y_column: str, session_hash, data: List[dict]=[{}], layout: List[dict]=[{}], 
-                                  category: str="", facet_row: str="", facet_col: str=""):
+def bar_chart_generation_func(x_column: str, y_column: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], 
+                                  category: str="", facet_row: str="", facet_col: str="", **kwargs):
    try:
-      dir_path = TEMP_DIR / str(session_hash)
+      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
       chart_path = f'{dir_path}/chart.html'
       csv_query_path = f'{dir_path}/query.csv'
 
@@ -235,7 +235,7 @@ def bar_chart_generation_func(x_column: str, y_column: str, session_hash, data:
    
       pio.write_html(fig, chart_path, full_html=False)
 
-      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/chart.html'
+      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
 
       iframe = '<div style=overflow:auto;><iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="' + chart_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n</div>'
 
@@ -250,9 +250,9 @@ def bar_chart_generation_func(x_column: str, y_column: str, session_hash, data:
             """
       return {"reply": reply}
 
-def pie_chart_generation_func(values: str, names: str, session_hash, data: List[dict]=[{}], layout: List[dict]=[{}]):
+def pie_chart_generation_func(values: str, names: str, session_hash, session_folder, data: List[dict]=[{}], layout: List[dict]=[{}], **kwargs):
    try:
-      dir_path = TEMP_DIR / str(session_hash)
+      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
       chart_path = f'{dir_path}/chart.html'
       csv_query_path = f'{dir_path}/query.csv'
 
@@ -282,7 +282,7 @@ def pie_chart_generation_func(values: str, names: str, session_hash, data: List[
    
       pio.write_html(fig, chart_path, full_html=False)
 
-      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/chart.html'
+      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
 
       iframe = '<div style=overflow:auto;><iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="' + chart_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n</div>'
 
@@ -297,16 +297,15 @@ def pie_chart_generation_func(values: str, names: str, session_hash, data: List[
             """
       return {"reply": reply}
    
-def histogram_generation_func(x_column: str, session_hash, y_column: str="", data: List[dict]=[{}], layout: List[dict]=[{}], histnorm: str="", category: str="",
-                              histfunc: str=""):   
+def histogram_generation_func(x_column: str, session_hash, session_folder, y_column: str="", data: List[dict]=[{}], layout: List[dict]=[{}], histnorm: str="", category: str="",
+                              histfunc: str="", **kwargs):   
    try:
-      dir_path = TEMP_DIR / str(session_hash)
+      dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
       chart_path = f'{dir_path}/chart.html'
       csv_query_path = f'{dir_path}/query.csv'
 
       df = pd.read_csv(csv_query_path)
 
-      print(df)
       print(x_column)
 
       function_args = {"data_frame":df, "x":x_column}
@@ -342,7 +341,7 @@ def histogram_generation_func(x_column: str, session_hash, y_column: str="", dat
    
       pio.write_html(fig, chart_path, full_html=False)
 
-      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/chart.html'
+      chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
 
       iframe = '<div style=overflow:auto;><iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="' + chart_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n</div>'
 
@@ -357,15 +356,14 @@ def histogram_generation_func(x_column: str, session_hash, y_column: str="", dat
             """
       return {"reply": reply}
 
-def table_generation_func(session_hash):
+def table_generation_func(session_hash, session_folder, **kwargs):
     print("TABLE GENERATION")
     try: 
-        dir_path = TEMP_DIR / str(session_hash)
+        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
         csv_query_path = f'{dir_path}/query.csv'
         table_path = f'{dir_path}/table.html'
 
         df = pd.read_csv(csv_query_path)
-        print(df)
 
         html_table = df.to_html()
         print(html_table)
@@ -373,7 +371,7 @@ def table_generation_func(session_hash):
         with open(table_path, "w") as file:
          file.write(html_table)
 
-        table_url = f'{root_url}/gradio_api/file/temp/{session_hash}/table.html'
+        table_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/table.html'
 
         iframe = '<div style=overflow:auto;><iframe\n scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="' + table_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n</div>'
         print(iframe)

functions/chat_functions.py

@@ -35,6 +35,25 @@ def example_question_generator(session_hash):
 
     return example_response["replies"][0].text
 
+def sql_example_question_generator(session_hash, db_tables, db_name):
+    example_response = None
+    example_messages = [
+        ChatMessage.from_system(
+            f"You are a helpful and knowledgeable agent who has access to an PostgreSQL database called {db_name}."
+        )
+    ]
+
+    example_messages.append(ChatMessage.from_user(text=f"""We have a PostgreSQL database with the following tables: {db_tables}. 
+                                                  We also have an AI agent with access to the same database that will be performing data analysis.
+                                                  Please return an array of seven strings, each one being a question for our data analysis agent
+                                                  that we can suggest that you believe will be insightful or helpful to a data analysis looking for
+                                                  data insights. Return nothing more than the array of questions because I need that specific data structure
+                                                  to process your response. No other response type or data structure will work."""))
+
+    example_response = chat_generator.run(messages=example_messages)
+
+    return example_response["replies"][0].text
+
 def chatbot_with_fc(message, history, session_hash):
     from functions import sqlite_query_func, table_generation_func, regression_func, scatter_chart_generation_func, \
         line_chart_generation_func,bar_chart_generation_func,pie_chart_generation_func,histogram_generation_func
@@ -46,8 +65,8 @@ def chatbot_with_fc(message, history, session_hash):
                            "histogram_generation_func":histogram_generation_func,
                            "regression_func":regression_func }
 
-    if message_dict[session_hash] != None:
-        message_dict[session_hash].append(ChatMessage.from_user(message))
+    if message_dict[session_hash]['file_upload'] != None:
+        message_dict[session_hash]['file_upload'].append(ChatMessage.from_user(message))
     else:
         messages = [
             ChatMessage.from_system(
@@ -58,35 +77,94 @@ def chatbot_with_fc(message, history, session_hash):
                 You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a bar graph and returns an iframe that we can display in our chat window.
                 You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a pie chart and returns an iframe that we can display in our chat window.
                 You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a histogram and returns an iframe that we can display in our chat window.
-                You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our sql query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe."""
+                You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our sql query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
+                Charts, tables, and visualizations are a very important part of your output. If you generate a chart, table, or visualization as part of your answer, please display it always."""
+            )
+        ]
+        messages.append(ChatMessage.from_user(message))
+        message_dict[session_hash]['file_upload'] = messages
+
+    response = chat_generator.run(messages=message_dict[session_hash]['file_upload'], generation_kwargs={"tools": tools.data_file_tools_call(session_hash)})
+
+    while True:
+        # if OpenAI response is a tool call
+        if response and response["replies"][0].meta["finish_reason"] == "tool_calls" or response["replies"][0].tool_calls:
+            function_calls = response["replies"][0].tool_calls
+            for function_call in function_calls:
+                message_dict[session_hash]['file_upload'].append(ChatMessage.from_assistant(tool_calls=[function_call]))
+                ## Parse function calling information
+                function_name = function_call.tool_name
+                function_args = function_call.arguments
+
+                ## Find the corresponding function and call it with the given arguments
+                function_to_call = available_functions[function_name]
+                function_response = function_to_call(**function_args, session_hash=session_hash, session_folder='file_upload')
+                print(function_name)
+                ## Append function response to the messages list using `ChatMessage.from_tool`
+                message_dict[session_hash]['file_upload'].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
+                response = chat_generator.run(messages=message_dict[session_hash]['file_upload'], generation_kwargs={"tools": tools.data_file_tools_call(session_hash)})
+
+        # Regular Conversation
+        else:
+            message_dict[session_hash]['file_upload'].append(response["replies"][0])
+            break
+
+    return response["replies"][0].text
+
+def sql_chatbot_with_fc(message, history, session_hash, db_url, db_port, db_user, db_pass, db_name, db_tables):
+    from functions import sql_query_func, table_generation_func, regression_func, scatter_chart_generation_func, \
+        line_chart_generation_func,bar_chart_generation_func,pie_chart_generation_func,histogram_generation_func
+    import tools.tools as tools
+
+    available_functions = {"sql_query_func": sql_query_func,"table_generation_func":table_generation_func,
+                           "line_chart_generation_func":line_chart_generation_func,"bar_chart_generation_func":bar_chart_generation_func,
+                           "scatter_chart_generation_func":scatter_chart_generation_func, "pie_chart_generation_func":pie_chart_generation_func,
+                           "histogram_generation_func":histogram_generation_func,
+                           "regression_func":regression_func }
+
+    if message_dict[session_hash]['sql'] != None:
+        message_dict[session_hash]['sql'].append(ChatMessage.from_user(message))
+    else:
+        messages = [
+            ChatMessage.from_system(
+                f"""You are a helpful and knowledgeable agent who has access to an PostgreSQL database which has a series of tables called {db_tables}. 
+                You also have access to a function, called table_generation_func, that can take a query.csv file generated from our sql query and returns an iframe that we can display in our chat window.
+                You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we can display in our chat window.
+                You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a line chart and returns an iframe that we can display in our chat window.
+                You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a bar graph and returns an iframe that we can display in our chat window.
+                You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a pie chart and returns an iframe that we can display in our chat window.
+                You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a histogram and returns an iframe that we can display in our chat window.
+                You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our sql query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
+                Charts, tables, and visualizations are a very important part of your output. If you generate a chart, table, or visualization as part of your answer, please display it always."""
             )
         ]
         messages.append(ChatMessage.from_user(message))
-        message_dict[session_hash] = messages
+        message_dict[session_hash]['sql'] = messages
 
-    response = chat_generator.run(messages=message_dict[session_hash], generation_kwargs={"tools": tools.data_file_tools_call(session_hash)})
+    response = chat_generator.run(messages=message_dict[session_hash]['sql'], generation_kwargs={"tools": tools.sql_tools_call(db_tables)})
 
     while True:
         # if OpenAI response is a tool call
         if response and response["replies"][0].meta["finish_reason"] == "tool_calls" or response["replies"][0].tool_calls:
             function_calls = response["replies"][0].tool_calls
             for function_call in function_calls:
-                message_dict[session_hash].append(ChatMessage.from_assistant(tool_calls=[function_call]))
+                message_dict[session_hash]['sql'].append(ChatMessage.from_assistant(tool_calls=[function_call]))
                 ## Parse function calling information
                 function_name = function_call.tool_name
                 function_args = function_call.arguments
 
                 ## Find the corresponding function and call it with the given arguments
                 function_to_call = available_functions[function_name]
-                function_response = function_to_call(**function_args, session_hash=session_hash)
+                function_response = function_to_call(**function_args, session_hash=session_hash, db_url=db_url, 
+                                                     db_port=db_port, db_user=db_user, db_pass=db_pass, db_name=db_name, session_folder='sql')
                 print(function_name)
                 ## Append function response to the messages list using `ChatMessage.from_tool`
-                message_dict[session_hash].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
-                response = chat_generator.run(messages=message_dict[session_hash], generation_kwargs={"tools": tools.data_file_tools_call(session_hash)})
+                message_dict[session_hash]['sql'].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
+                response = chat_generator.run(messages=message_dict[session_hash]['sql'], generation_kwargs={"tools": tools.sql_tools_call(db_tables)})
 
         # Regular Conversation
         else:
-            message_dict[session_hash].append(response["replies"][0])
+            message_dict[session_hash]['sql'].append(response["replies"][0])
             break
         
     return response["replies"][0].text

functions/{sqlite_functions.py → query_functions.py}

@@ -6,6 +6,7 @@ pd.set_option('display.max_columns', None)
 pd.set_option('display.width', None)
 pd.set_option('display.max_colwidth', None)
 import sqlite3
+import psycopg2
 from utils import TEMP_DIR
 
 @component
@@ -16,21 +17,21 @@ class SQLiteQuery:
 
     @component.output_types(results=List[str], queries=List[str])
     def run(self, queries: List[str], session_hash):
-        print("ATTEMPTING TO RUN QUERY")
+        print("ATTEMPTING TO RUN SQLITE QUERY")
         dir_path = TEMP_DIR / str(session_hash)
         results = []
         for query in queries:
           result = pd.read_sql(query, self.connection)
-          result.to_csv(f'{dir_path}/query.csv', index=False)
+          result.to_csv(f'{dir_path}/file_upload/query.csv', index=False)
           results.append(f"{result}")
         self.connection.close()
         return {"results": results, "queries": queries}
     
 
 
-def sqlite_query_func(queries: List[str], session_hash):
+def sqlite_query_func(queries: List[str], session_hash, **kwargs):
     dir_path = TEMP_DIR / str(session_hash)
-    sql_query = SQLiteQuery(f'{dir_path}/data_source.db')
+    sql_query = SQLiteQuery(f'{dir_path}/file_upload/data_source.db')
     try:
       result = sql_query.run(queries, session_hash)
       if len(result["results"][0]) > 1000:
@@ -45,3 +46,50 @@ def sqlite_query_func(queries: List[str], session_hash):
               You should probably try again.
               """
       return {"reply": reply}
+    
+@component
+class PostgreSQLQuery:
+
+    def __init__(self, url: str, sql_port: int, sql_user: str, sql_pass: str, sql_db_name: str):
+      self.connection = psycopg2.connect(
+            database=sql_db_name,
+            user=sql_user,
+            password=sql_pass,
+            host=url,  # e.g., "localhost" or an IP address
+            port=sql_port  # default is 5432
+        )
+
+    @component.output_types(results=List[str], queries=List[str])
+    def run(self, queries: List[str], session_hash):
+        print("ATTEMPTING TO RUN POSTGRESQL QUERY")
+        dir_path = TEMP_DIR / str(session_hash)
+        results = []
+        for query in queries:
+          print(query)
+          result = pd.read_sql_query(query, self.connection)
+          result.to_csv(f'{dir_path}/sql/query.csv', index=False)
+          results.append(f"{result}")
+        self.connection.close()
+        return {"results": results, "queries": queries}
+    
+
+
+def sql_query_func(queries: List[str], session_hash, db_url, db_port, db_user, db_pass, db_name, **kwargs):
+    sql_query = PostgreSQLQuery(db_url, db_port, db_user, db_pass, db_name)
+    try:
+      result = sql_query.run(queries, session_hash)
+      print("RESULT")
+      print(result)
+      if len(result["results"][0]) > 1000:
+        print("QUERY TOO LARGE")
+        return {"reply": "query result too large to be processed by llm, the query results are in our query.csv file. If you need to display the results directly, perhaps use the table_generation_func function."}
+      else:   
+        return {"reply": result["results"][0]}
+
+    except Exception as e:
+      reply = f"""There was an error running the SQL Query = {queries}
+              The error is {e},
+              You should probably try again.
+              """
+      print(reply)
+      return {"reply": reply}

functions/stat_functions.py

@@ -12,12 +12,12 @@ load_dotenv()
 
 root_url = os.getenv("ROOT_URL")
 
-def regression_func(independent_variables: List[str], dependent_variable: str, session_hash, category: str=''):
+def regression_func(independent_variables: List[str], dependent_variable: str, session_hash, session_folder, category: str='', **kwargs):
     print("LINEAR REGRESSION CALCULATION")
     print(independent_variables)
     print(dependent_variable)
     try:
-        dir_path = TEMP_DIR / str(session_hash)
+        dir_path = TEMP_DIR / str(session_hash) / str(session_folder)
         chart_path = f'{dir_path}/chart.html'
         csv_query_path = f'{dir_path}/query.csv'
 
@@ -32,7 +32,7 @@ def regression_func(independent_variables: List[str], dependent_variable: str, s
 
         pio.write_html(fig, chart_path, full_html=False)
 
-        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/chart.html'
+        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/{session_folder}/chart.html'
 
         iframe = '<div style=overflow:auto;><iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="' + chart_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n</div>'
 

requirements.txt

@@ -6,3 +6,4 @@ plotly
 openpyxl
 statsmodels
 xlrd
+psycopg2-binary

data_file.py → templates/data_file.py

@@ -68,7 +68,8 @@ with gr.Blocks() as demo:
     @gr.render(inputs=file_output)
     def data_options(filename, request: gr.Request):
         print(filename)
-        message_dict[request.session_hash] = None
+        message_dict[request.session_hash] = {}
+        message_dict[request.session_hash]['file_upload'] = None
         if filename:
             process_message = process_upload(filename, request.session_hash)
             gr.HTML(value=process_message[1], padding=False)
@@ -101,7 +102,9 @@ with gr.Blocks() as demo:
                                             ]
                         for example in generated_examples:
                             example_questions.append([example])
-                    except:
+                    except Exception as e:
+                        print("DATA FILE QUESTION GENERATION ERROR")
+                        print(e)
                         example_questions = [
                                             ["Describe the dataset"],
                                             ["List the columns in the dataset"],

templates/sql_db.py

@@ -0,0 +1,95 @@
+import ast
+import gradio as gr
+from functions import sql_example_question_generator, sql_chatbot_with_fc
+from data_sources import connect_sql_db
+from utils import message_dict
+
+def hide_info():
+    return gr.update(visible=False)
+
+with gr.Blocks() as demo:
+    description = gr.HTML("""
+                    <!-- Header -->
+                    <div class="max-w-4xl mx-auto mb-12 text-center">
+                        <div class="bg-blue-50 border border-blue-200 rounded-lg max-w-2xl mx-auto">
+                         <p>This tool allows users to communicate with and query real time data from a SQL DB (postgres for now, others can be added if requested) using natural
+                          language and the above features.</p>
+                         <p style="font-weight:bold;">Notice: the way this system is designed, no login information is retained and credentials are passed as session variables until the user leaves or 
+                          refreshes the page in which they disappear. They are never saved to any files. I also make use of the Pandas read_sql_query function to apply SQL 
+                          queries, which can't delete, drop, or add database lines to avoid unhappy accidents or glitches.
+                          That being said, it's probably not a good idea to connect a production database to a strange AI tool with an unfamiliar author. 
+                          This should be for demonstration purposes.</p>
+                          <p>Contact me if this is something you would like built in your organization, on your infrastructure, and with the requisite privacy and control a production
+                          database analytics tool requires.</p>
+                        </div>
+                    </div>
+                        """, elem_classes="description_component")
+    sql_url = gr.Textbox(label="URL", value="virtual-data-analyst-pg.cyetm2yjzppu.us-west-1.rds.amazonaws.com")
+    with gr.Row():
+        sql_port = gr.Textbox(label="Port", value="5432")
+        sql_user = gr.Textbox(label="Username", value="postgres")
+        sql_pass = gr.Textbox(label="Password", value="Vda-1988", type="password")
+        sql_db_name = gr.Textbox(label="Database Name", value="dvdrental")
+
+    submit = gr.Button(value="Submit")
+    submit.click(fn=hide_info, outputs=description)
+
+    @gr.render(inputs=[sql_url,sql_port,sql_user,sql_pass,sql_db_name], triggers=[submit.click])
+    def sql_chat(request: gr.Request, url=sql_url.value, sql_port=sql_port.value, sql_user=sql_user.value, sql_pass=sql_pass.value, sql_db_name=sql_db_name.value):
+        message_dict[request.session_hash]['sql'] = None
+        if url:
+            print("SQL APP")
+            print(request)
+            process_message = process_sql_db(url, sql_user, sql_port, sql_pass, sql_db_name, request.session_hash)
+            gr.HTML(value=process_message[1], padding=False)
+            if process_message[0] == "success":
+                if "virtual-data-analyst-pg.cyetm2yjzppu.us-west-1.rds.amazonaws.com" in url:
+                    example_questions = [
+                                            ["Describe the dataset"],
+                                            ["What is the total revenue generated by each store?"],
+                                            ["Can you generate and display a bar chart of film category to number of films in that category?"],
+                                            ["Can you generate a pie chart showing the top 10 most rented films by revenue vs all other films?"],
+                                            ["Can you generate a line chart of rental revenue over time?"],
+                                            ["What is the relationship between film length and rental frequency?"]
+                                        ]
+                else:
+                    try:
+                        generated_examples = ast.literal_eval(sql_example_question_generator(request.session_hash, process_message[2], sql_db_name))
+                        example_questions = [
+                                                ["Describe the dataset"]
+                                            ]
+                        for example in generated_examples:
+                            example_questions.append([example])
+                    except Exception as e:
+                        print("SQL QUESTION GENERATION ERROR")
+                        print(e)
+                        example_questions = [
+                                            ["Describe the dataset"],
+                                            ["List the columns in the dataset"],
+                                            ["What could this data be used for?"],
+                                        ]
+                session_hash = gr.Textbox(visible=False, value=request.session_hash)
+                db_url = gr.Textbox(visible=False, value=url)
+                db_port = gr.Textbox(visible=False, value=sql_port)
+                db_user = gr.Textbox(visible=False, value=sql_user)
+                db_pass = gr.Textbox(visible=False, value=sql_pass)
+                db_name = gr.Textbox(visible=False, value=sql_db_name)
+                db_tables = gr.Textbox(value=process_message[2], interactive=False, label="SQL Tables")
+                bot = gr.Chatbot(type='messages', label="CSV Chat Window", render_markdown=True, sanitize_html=False, show_label=True, render=False, visible=True, elem_classes="chatbot")
+                chat = gr.ChatInterface(
+                                    fn=sql_chatbot_with_fc,
+                                    type='messages',
+                                    chatbot=bot,
+                                    title="Chat with your Database",
+                                    examples=example_questions,
+                                    concurrency_limit=None,
+                                    additional_inputs=[session_hash, db_url, db_port, db_user, db_pass, db_name, db_tables]
+                                    )
+            
+    def process_sql_db(url, sql_user, sql_port, sql_pass, sql_db_name, session_hash):
+        if url:
+            process_message = connect_sql_db(url, sql_user, sql_port, sql_pass, sql_db_name, session_hash)
+        return process_message            
+
+if __name__ == "__main__":
+    demo.launch()

tools/chart_tools.py

@@ -3,7 +3,7 @@ chart_tools = [
         "type": "function",
             "function": {
                 "name": "scatter_chart_generation_func",
-                "description": f"""This is a scatter plot generation tool useful to generate scatter plots from queried data from our SQL table called 'data_source'.
+                "description": f"""This is a scatter plot generation tool useful to generate scatter plots from queried data from our data source that we are querying.
                 The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
                 Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
                 from the scatter_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
@@ -108,7 +108,7 @@ chart_tools = [
         "type": "function",
             "function": {
                 "name": "line_chart_generation_func",
-                "description": f"""This is a line chart generation tool useful to generate line charts from queried data from our SQL table called 'data_source'.
+                "description": f"""This is a line chart generation tool useful to generate line charts from queried data from our data source that we are querying.
                 The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
                 Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
                 from the line_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
@@ -164,7 +164,7 @@ chart_tools = [
         "type": "function",
             "function": {
                 "name": "bar_chart_generation_func",
-                "description": f"""This is a bar chart generation tool useful to generate line charts from queried data from our SQL table called 'data_source'.
+                "description": f"""This is a bar chart generation tool useful to generate line charts from queried data from our data source that we are querying.
                 The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
                 Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
                 from the bar_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
@@ -236,7 +236,7 @@ chart_tools = [
         "type": "function",
             "function": {
                 "name": "pie_chart_generation_func",
-                "description": f"""This is a pie chart generation tool useful to generate pie charts from queried data from our SQL table called 'data_source'.
+                "description": f"""This is a pie chart generation tool useful to generate pie charts from queried data from our data source that we are querying.
                 The data values will come from the columns of our query.csv (the 'values' and 'names' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
                 Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
                 from the pie_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
@@ -285,7 +285,7 @@ chart_tools = [
         "type": "function",
             "function": {
                 "name": "histogram_generation_func",
-                "description": f"""This is a histogram generation tool useful to generate histograms from queried data from our SQL table called 'data_source'.
+                "description": f"""This is a histogram generation tool useful to generate histograms from queried data from our data source that we are querying.
                 The data values will come from the columns of our query.csv (the 'values' and 'names' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
                 Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
                 from the histogram_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
@@ -360,7 +360,7 @@ chart_tools = [
         "type": "function",
             "function": {
                 "name": "table_generation_func",
-                "description": f"""This an table generation tool useful to format data as a table from queried data from our SQL table called 'data_source'.
+                "description": f"""This an table generation tool useful to format data as a table from queried data from our data source that we are querying.
                 Takes no parameters as it uses data queried in our query.csv file to build the table.
                 Call this function after running our SQLite query and generating query.csv.
                 Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 

tools/stats_tools.py

@@ -3,7 +3,7 @@ stats_tools = [
         "type": "function",
             "function": {
                 "name": "regression_func",
-                "description": f"""This a tool to calculate regressions on our SQLite table called 'data_source'.
+                "description": f"""This a tool to calculate regressions on our data source that we are querying.
                 We can run queries with our 'sql_query_func' function and they will be available to use in this function via the query.csv file that is generated.
                 Returns a dictionary of values that includes a regression_summary and a regression chart (which is an iframe displaying the
                 linear regression in chart form and should be shown to the user).""",

tools/tools.py

@@ -1,11 +1,12 @@
 import sqlite3
+import psycopg2
 from .stats_tools import stats_tools
 from .chart_tools import chart_tools
 from utils import TEMP_DIR
 
 def data_file_tools_call(session_hash):
     dir_path = TEMP_DIR / str(session_hash)
-    connection = sqlite3.connect(f'{dir_path}/data_source.db')
+    connection = sqlite3.connect(f'{dir_path}/file_upload/data_source.db')
     print("Querying Database in Tools.py");
     cur=connection.execute('select * from data_source')
     columns = [i[0] for i in cur.description]
@@ -46,22 +47,24 @@ def data_file_tools_call(session_hash):
 
     return tools_calls
 
-def graphql_tools_call(sessions_hash):
+def sql_tools_call(db_tables):
+    
+    table_string = (db_tables[:625] + '..') if len(db_tables) > 625 else db_tables
         
     tools_calls = [
         {
             "type": "function",
             "function": {
-                "name": "graphql_query_func",
-                "description": f"""This is a tool useful to query a GraphQL endpoint with the following Columns: {column_string}.
-                There may also be more columns in the table if the number of columns is too large to process. 
+                "name": "sql_query_func",
+                "description": f"""This is a tool useful to query a PostgreSQL database with the following tables, {table_string}.
+                There may also be more tables in the database if the number of columns is too large to process.
                 This function also saves the results of the query to csv file called query.csv.""",
                 "parameters": {
                     "type": "object",
                     "properties": {
                         "queries": {
                             "type": "array",
-                            "description": "The graphQL query to use in the search. Infer this from the user's message. It should be a question or a statement",
+                            "description": "The PostgreSQL query to use in the search. Infer this from the user's message. It should be a question or a statement",
                             "items": {
                                 "type": "string",
                             }
@@ -73,7 +76,7 @@ def graphql_tools_call(sessions_hash):
         },
     ]
 
-    tools_calls.append(chart_tools)
-    tools_calls.append(stats_tools)
+    tools_calls.extend(chart_tools)
+    tools_calls.extend(stats_tools)
 
-    return
+    return tools_calls