Fixes for chart generation (#20)

- Fixes for chart generation (970eb42c0a8e47c90808b79383cc8a15af9261bd)

functions/chart_functions.py

@@ -6,6 +6,7 @@ import pandas as pd
 from utils import TEMP_DIR
 import os
 import ast
+import json
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -22,12 +23,20 @@ def llm_chart_data_scrub(data, layout):
    else:
       data_list.append(data) 
 
+   false_replace = [':false', ': false']
+   false_value = ':False'   
+   true_replace = [':true', ': true']
+   true_value = ':True' 
+
    data_dict = {}   
    for data_obj in data_list:   
       if isinstance(data_obj, str):
          data_obj = data_obj.replace("\n", "")
-         if not data_obj.startswith('{') or not data_obj.endswith('}'):
-            data_obj = "{" + data_obj + "}"
+         for replace in false_replace:
+            data_obj = data_obj.replace(replace, false_value)
+         for replace in true_replace:
+            data_obj = data_obj.replace(replace, true_value)
+         print(data_obj)
          data_dict = ast.literal_eval(data_obj)
       else:
          data_dict = data_obj
@@ -38,6 +47,11 @@ def llm_chart_data_scrub(data, layout):
       layout_obj = layout
 
    if layout_obj and isinstance(layout_obj, str):
+      for replace in false_replace:
+         layout_obj = layout_obj.replace(replace, false_value)
+      for replace in true_replace:
+         layout_obj = layout_obj.replace(replace, true_value)
+      print(layout_obj)
       layout_dict = ast.literal_eval(layout_obj)
    else:
       layout_dict = layout_obj
@@ -348,6 +362,7 @@ def table_generation_func(session_hash):
     try: 
         dir_path = TEMP_DIR / str(session_hash)
         csv_query_path = f'{dir_path}/query.csv'
+        table_path = f'{dir_path}/table.html'
 
         df = pd.read_csv(csv_query_path)
         print(df)
@@ -355,12 +370,19 @@ def table_generation_func(session_hash):
         html_table = df.to_html()
         print(html_table)
 
-        return {"reply": html_table}
+        with open(table_path, "w") as file:
+         file.write(html_table)
+
+        table_url = f'{root_url}/gradio_api/file/temp/{session_hash}/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)
+        return {"reply": iframe}
     
     except Exception as e:
       print("TABLE ERROR")
       print(e)
-      reply = f"""There was an error generating the Pandas DataFrame table from {data}
+      reply = f"""There was an error generating the Pandas DataFrame table results.
               The error is {e},
               You should probably try again.
               """

functions/chat_functions.py

@@ -52,7 +52,7 @@ def chatbot_with_fc(message, history, session_hash):
         messages = [
             ChatMessage.from_system(
                 """You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source'. 
-                You also have access to a function, called table_generation_func, that builds table formatted html.
+                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.

functions/sqlite_functions.py

@@ -35,7 +35,7 @@ def sqlite_query_func(queries: List[str], session_hash):
       result = sql_query.run(queries, session_hash)
       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"}
+        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]}
 

tools.py

@@ -401,9 +401,8 @@ def tools_call(session_hash):
                 "description": f"""This an table generation tool useful to format data as a table from queried data from our SQL table called 'data_source'.
                 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 html string generated from the pandas library and pandas.to_html() 
-                function which will be displayed inline in our chat window.
-                Do not edit the string returned by the function in any way when displaying to the user.""",
+                Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
+                from the table_generation_func function in any way and always display the iframe fully to the user in the chat window.""",
                 "parameters": {},
             }, 
         },