Table generation and download (#12)

- Table generation and download (e7b4bfb12b525c5770dd25408a1a433d15cc5200)

app.py

@@ -10,4 +10,4 @@ if "OPENAI_API_KEY" not in os.environ:
     os.environ["OPENAI_API_KEY"] = getpass("Enter OpenAI API key:")
 
 ## Uncomment the line below to launch the chat app with UI
-demo.launch(debug=True)
+demo.launch(debug=True, allowed_paths=["temp/"])

data_sources/upload_file.py

@@ -3,6 +3,8 @@ import sqlite3
 import csv
 import json
 import time
+import os
+from utils import TEMP_DIR
 
 def is_file_done_saving(file_path):
     try:
@@ -65,8 +67,10 @@ 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()
 
-    
-    connection = sqlite3.connect(f'data_source_{session_hash}.db')
+    dir_path = TEMP_DIR / str(session_hash)
+    os.makedirs(dir_path, exist_ok=True)
+
+    connection = sqlite3.connect(f'{dir_path}/data_source.db')
     print("Opened database successfully");
     print(df.columns)
 

functions/__init__.py

@@ -1,5 +1,5 @@
 from .sqlite_functions import SQLiteQuery, sqlite_query_func
-from .chart_functions import chart_generation_func
+from .chart_functions import chart_generation_func, table_generation_func
 from .chat_functions import demo
 
-__all__ = ["SQLiteQuery","sqlite_query_func","chart_generation_func","demo"]
+__all__ = ["SQLiteQuery","sqlite_query_func","chart_generation_func","table_generation_func","demo"]

functions/chart_functions.py

@@ -1,5 +1,7 @@
 from typing import List
 from quickchart import QuickChart
+import pandas as pd
+from utils import TEMP_DIR
 
 def chart_generation_func(queries: List[str], session_hash):
     print("CHART GENERATION")
@@ -22,3 +24,16 @@ def chart_generation_func(queries: List[str], session_hash):
     iframe = '<div style=overflow:auto;><iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="' + interactive_url + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n <p>Edit, share, and download this graph <a target="_blank" href="' + edit_url + '">here</a></p></div>'
 
     return {"reply": iframe}
+
+def table_generation_func(data: List[str], session_hash):
+    dir_path = TEMP_DIR / str(session_hash)
+    print("TABLE GENERATION")
+    print(data)
+    df = pd.DataFrame(data)
+    csv_path = f'{dir_path}/data.csv'
+    df.to_csv(csv_path)
+    download_path = f'gradio_api/file/temp/{session_hash}/data.csv'
+    html_table = df.to_html() + f'<p>Download as a <a href="{download_path}">CSV</a></p>'
+    print(html_table)
+
+    return {"reply": html_table}

functions/chat_functions.py

@@ -1,4 +1,5 @@
 from data_sources import process_data_upload
+from utils import TEMP_DIR
 
 import gradio as gr
 
@@ -27,7 +28,8 @@ def example_question_generator(session_hash):
             "You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source'."
         )
     ]
-    connection = sqlite3.connect(f'data_source_{session_hash}.db')
+    dir_path = TEMP_DIR / str(session_hash)
+    connection = sqlite3.connect(f'{dir_path}/data_source.db')
     print("Querying questions");
     cur=connection.execute('select * from data_source')
     columns = [i[0] for i in cur.description]
@@ -48,11 +50,11 @@ def example_question_generator(session_hash):
     return example_response["replies"][0].text
 
 def chatbot_with_fc(message, history, session_hash):
-    from functions import sqlite_query_func, chart_generation_func
+    from functions import sqlite_query_func, chart_generation_func, table_generation_func
     from pipelines import rag_pipeline_func
     import tools
 
-    available_functions = {"sql_query_func": sqlite_query_func, "rag_pipeline_func": rag_pipeline_func, "chart_generation_func": chart_generation_func}
+    available_functions = {"sql_query_func": sqlite_query_func, "rag_pipeline_func": rag_pipeline_func, "chart_generation_func": chart_generation_func, "table_generation_func":table_generation_func }
 
     if message_dict[session_hash] != None:
         message_dict[session_hash].append(ChatMessage.from_user(message))
@@ -92,9 +94,10 @@ def chatbot_with_fc(message, history, session_hash):
     return response["replies"][0].text
 
 def delete_db(req: gr.Request):
-    db_file_path = f'data_source_{req.session_hash}.db'
-    if os.path.exists(db_file_path):
-        os.remove(db_file_path)
+    import shutil
+    dir_path = TEMP_DIR / str(req.session_hash)
+    if os.path.exists(dir_path):
+        shutil.rmtree(dir_path)
         message_dict[req.session_hash] = None
 
 def run_example(input):
@@ -139,14 +142,16 @@ with gr.Blocks(css=css, delete_cache=(3600,3600)) as demo:
                                         ["Describe the dataset"],
                                         ["What levels of education have the highest and lowest average balance?"],
                                         ["What job is most and least common for a yes response from the individuals, not counting 'unknown'?"],
-                                        ["Can you generate a bar chart of education vs. average balance?"]
+                                        ["Can you generate a bar chart of education vs. average balance?"],
+                                        ["Can you generate a table of levels of education versus average balance, percent married, percent with a loan, and percent in default?"]
                                     ]
             elif "online_retail_data" in filename:
                 example_questions = [
                                         ["Describe the dataset"],
                                         ["What month had the highest revenue?"],
                                         ["Is revenue higher in the morning or afternoon?"],
-                                        ["Can you generate a line graph of revenue per month?"]
+                                        ["Can you generate a line graph of revenue per month?"],
+                                        ["Can you generate a table of revenue per month?"]
                                     ]
             else:
                 try:

functions/sqlite_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
+from utils import TEMP_DIR
 
 @component
 class SQLiteQuery:
@@ -26,7 +27,8 @@ class SQLiteQuery:
 
 
 def sqlite_query_func(queries: List[str], session_hash):
-    sql_query = SQLiteQuery(f'data_source_{session_hash}.db')
+    dir_path = TEMP_DIR / str(session_hash)
+    sql_query = SQLiteQuery(f'{dir_path}/data_source.db')
     try:
       result = sql_query.run(queries)
       return {"reply": result["results"][0]}

tools.py

@@ -1,7 +1,9 @@
 import sqlite3
+from utils import TEMP_DIR
 
 def tools_call(session_hash):
-    connection = sqlite3.connect(f'data_source_{session_hash}.db')
+    dir_path = TEMP_DIR / str(session_hash)
+    connection = sqlite3.connect(f'{dir_path}/data_source.db')
     print("Querying Database in Tools.py");
     cur=connection.execute('select * from data_source')
     columns = [i[0] for i in cur.description]
@@ -35,7 +37,7 @@ def tools_call(session_hash):
         "type": "function",
             "function": {
                 "name": "chart_generation_func",
-                "description": f"This an chart generation tool useful to generate charts and graphs from queried data from our SQL table called 'data_source with the following Columns: {columns}. Returns an iframe string which will be displayed inline in our chat window. Do not edit the string returned from the chart_generation_func function in any way and display it fully to the user. You can add your own text supplementary to it for context if desired.",
+                "description": f"This a chart generation tool useful to generate charts and graphs from queried data from our SQL table called 'data_source with the following Columns: {columns}. Returns an iframe string which will be displayed inline in our chat window. Do not edit the string returned from the chart_generation_func function in any way and display it fully to the user in the chat window. You can add your own text supplementary to it for context if desired.",
                 "parameters": {
                     "type": "object",
                     "properties": {
@@ -43,7 +45,36 @@ def tools_call(session_hash):
                             "type": "array",
                             "description": """The data points to use in the chart generation. Infer this from the user's message. 
                             Send a chart.js dictionary with options that correspond to the users request. But also format this dictionary as a string as this will allow javascript to be interpreted by the API we are using.
-                            Return nothing else.""",
+                            Send nothing else.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        }
+                    },
+                    "required": ["question"],
+                },
+            }, 
+        },
+        {
+        "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 with the following Columns: {columns}. Returns an html string generated from the pandas library and pandas.to_html() 
+                function which will be displayed inline in our chat window. There will also be a link to download the CSV included in the HTML string. 
+                Do not change or edit this link in any way. Do not edit the HTML returned by the function in any way.
+                Do not edit the string returned from the table_generation_func function in any way and display it fully to the user in the chat window. 
+                You can add your own text next to the returned string for context if desired.""",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "data": {
+                            "type": "array",
+                            "description": """The data points to use in the table generation. Infer this from the user's message. 
+                            Send a python dictionary object with query data that correspond to data that will be converted into a pandas DataFrame and that correspond to the users request.
+                            The keys of this python dictionary object will be the names of the columns and values will be a list of values for each object.
+                            Make sure this is a dictionary object and not a string or an array.
+                            Send nothing else.""",
                             "items": {
                                 "type": "string",
                             }

utils.py

@@ -0,0 +1,5 @@
+from pathlib import Path
+
+current_dir = Path(__file__).parent
+
+TEMP_DIR = current_dir / 'temp'