add regression function

app.py

@@ -36,7 +36,7 @@ with gr.Blocks(css=css, delete_cache=(3600,3600)) as demo:
     title = gr.HTML("<h1 style='text-align:center;'>Virtual Data Analyst</h1>")
     description = gr.HTML("""<p style='text-align:center;'>Upload a data file and chat with our virtual data analyst 
                           to get insights on your data set. Currently accepts CSV, TSV, TXT, XLS, XLSX, XML, and JSON files. 
-                          Can now generate charts and graphs!
+                          Can now generate charts and graphs! Can run linear regressions!
                           Try a sample file to get started!</p>
                           <p style='text-align:center;'>This tool is under active development. If you experience bugs with use, 
                           open a discussion in the community tab and I will respond.</p>""")
@@ -63,7 +63,8 @@ with gr.Blocks(css=css, delete_cache=(3600,3600)) as demo:
                                         ["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 table of levels of education versus average balance, percent married, percent with a loan, and percent in default?"]
+                                        ["Can you generate a table of levels of education versus average balance, percent married, percent with a loan, and percent in default?"],
+                                        ["Can we predict the relationship between the number of contacts performed before this campaign and the average balance?"],
                                     ]
             elif "online_retail_data" in filename:
                 example_questions = [
@@ -71,7 +72,8 @@ with gr.Blocks(css=css, delete_cache=(3600,3600)) as demo:
                                         ["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 table of revenue per month?"]
+                                        ["Can you generate a table of revenue per month?"],
+                                        ["Can we predict how time of day affects revenue in this data set?"],
                                     ]
             else:
                 try:

functions/__init__.py

@@ -1,5 +1,6 @@
 from .sqlite_functions import SQLiteQuery, sqlite_query_func
 from .chart_functions import chart_generation_func, table_generation_func
 from .chat_functions import example_question_generator, chatbot_with_fc
+from .stat_functions import regression_func
 
-__all__ = ["SQLiteQuery","sqlite_query_func","chart_generation_func","table_generation_func","example_question_generator","chatbot_with_fc"]
+__all__ = ["SQLiteQuery","sqlite_query_func","chart_generation_func","table_generation_func","regression_func","example_question_generator","chatbot_with_fc"]

functions/chart_functions.py

@@ -1,6 +1,7 @@
 from typing import List
 from typing import Dict
 import plotly.io as pio
+import plotly.express as px
 import pandas as pd
 from utils import TEMP_DIR
 import os
@@ -11,31 +12,59 @@ load_dotenv()
 
 root_url = os.getenv("ROOT_URL")
 
-def chart_generation_func(data: List[dict], session_hash: str, layout: Dict[str,str]={}):
+def chart_generation_func(data: List[str], x_column: str, y_column: str, graph_type: str, session_hash: str, layout: Dict[str,str]={}, category: str=""):
     print("CHART GENERATION")
     print(data)
+    print(x_column)
+    print(y_column)
+    print(category)
     print(layout)
     try:
         dir_path = TEMP_DIR / str(session_hash)
         chart_path = f'{dir_path}/chart.html'
+        csv_query_path = f'{dir_path}/query.csv'
 
         #Processing data to account for variation from LLM
         data_list = []
-        layout_dict = {}       
+        layout_dict = {}
+
+        df = pd.read_csv(csv_query_path)
+
+        if graph_type == "bar":
+         if category in df.columns:
+            initial_graph = px.bar(df, x=x_column, y=y_column, color=category, barmode="group")
+         else:
+            initial_graph = px.bar(df, x=x_column, y=y_column, barmode="group")
+        elif graph_type == "scatter":
+         if category in df.columns:
+            initial_graph = px.scatter(df, x=x_column, y=y_column, color=category)  
+         else:
+            initial_graph = px.scatter(df, x=x_column, y=y_column)
+        elif graph_type == "line":
+         if category in df.columns:
+            initial_graph = px.line(df, x=x_column, y=y_column, color=category)
+         else:
+            initial_graph = px.line(df, x=x_column, y=y_column)
+        elif graph_type == "pie":
+         if category in df.columns:
+            initial_graph = px.pie(df, x=x_column, y=y_column, color=category)
+         else:
+            initial_graph = px.pie(df, x=x_column, y=y_column)
 
         if isinstance(data, list):
            data_list = data
         else:
-           data_list.append(data)
-
-        data_dict_list = [] 
-        for data_obj in data_list:   
+           data_list.append(data)  
+           
+        for index, data_obj in enumerate(data_list):   
          if isinstance(data_obj, str):
+            data_obj = data_obj.replace("\n", "")
+            if not data_obj.startswith('{') and not data_obj.endswith('}'):
+               data_obj = "{" + data_obj + "}"
             data_dict = ast.literal_eval(data_obj)
          else:
             data_dict = data_obj
-         data_dict_list.append(data_dict)
-
+    
         if isinstance(layout, list):
            layout_obj = layout[0]
         else:
@@ -46,9 +75,15 @@ def chart_generation_func(data: List[dict], session_hash: str, layout: Dict[str,
         else:
            layout_dict = layout_obj
            
+        fig = initial_graph.to_dict()
+
+        fig["layout"] = layout_dict
 
-        fig = dict({"data": data_dict_list,
-                    "layout": layout_dict})
+        for key, value in data_dict.items():
+           if key not in ["x","y"]:
+            for data_item in fig["data"]:
+               data_item[key] = value
+     
         pio.write_html(fig, chart_path, full_html=False)
 
         chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/chart.html'
@@ -60,7 +95,7 @@ def chart_generation_func(data: List[dict], session_hash: str, layout: Dict[str,
     except Exception as e:
       print("CHART ERROR")
       print(e)
-      reply = f"""There was an error generating the Plotly Chart from {data} and {layout}
+      reply = f"""There was an error generating the Plotly Chart from {x_column}, {y_column}, {graph_type}, and {layout}
               The error is {e},
               You should probably try again.
               """

functions/chat_functions.py

@@ -36,10 +36,10 @@ 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, table_generation_func
+    from functions import sqlite_query_func, chart_generation_func, table_generation_func, regression_func
     import tools
 
-    available_functions = {"sql_query_func": sqlite_query_func, "chart_generation_func": chart_generation_func, "table_generation_func":table_generation_func }
+    available_functions = {"sql_query_func": sqlite_query_func, "chart_generation_func": chart_generation_func, "table_generation_func":table_generation_func, "regression_func":regression_func }
 
     if message_dict[session_hash] != None:
         message_dict[session_hash].append(ChatMessage.from_user(message))
@@ -47,8 +47,9 @@ 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 chart function that uses plotly dictionaries to generate charts and graphs and returns an iframe that we can display in our chat window.
-                You also have access to a function, called table_generation_func, that builds table formatted html and generates a link to download as CSV."""
+                You also have access to a chart function that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate charts and graphs and returns an iframe that we can display in our chat window.
+                You also have access to a function, called table_generation_func, that builds table formatted html and generates a link to download as CSV.
+                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."""
             )
         ]
         messages.append(ChatMessage.from_user(message))

functions/sqlite_functions.py

@@ -15,11 +15,13 @@ class SQLiteQuery:
       self.connection = sqlite3.connect(sql_database, check_same_thread=False)
 
     @component.output_types(results=List[str], queries=List[str])
-    def run(self, queries: List[str]):
+    def run(self, queries: List[str], session_hash):
         print("ATTEMPTING TO RUN 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)
           results.append(f"{result}")
         self.connection.close()
         return {"results": results, "queries": queries}
@@ -30,8 +32,12 @@ def sqlite_query_func(queries: List[str], session_hash):
     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]}
+      result = sql_query.run(queries, session_hash)
+      if len(result["results"][0]) > 300:
+        print("QUERY TOO LARGE")
+        return {"reply": "query result too large to be processed by llm, the query results are in our query.csv file"}
+      else:   
+        return {"reply": result["results"][0]}
 
     except Exception as e:
       reply = f"""There was an error running the SQL Query = {queries}

functions/stat_functions.py

@@ -0,0 +1,56 @@
+
+import pandas as pd
+from typing import List
+from utils import TEMP_DIR
+import plotly.express as px
+import plotly.io as pio
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+root_url = os.getenv("ROOT_URL")
+
+def basic_stats_function(data_list: List[str]):
+   return
+
+def regression_func(independent_variables: List[str], dependent_variable: str, session_hash, category: str='', regression_type: str="ols"):
+    print("LINEAR REGRESSION CALCULATION")
+    print(independent_variables)
+    print(dependent_variable)
+    try:
+        dir_path = TEMP_DIR / str(session_hash)
+        chart_path = f'{dir_path}/chart.html'
+        csv_query_path = f'{dir_path}/query.csv'
+
+        df = pd.read_csv(csv_query_path)
+
+        if category in df.columns:
+           fig = px.scatter(df, x=independent_variables, y=dependent_variable, color=category, trendline="ols")
+        else:
+           fig = px.scatter(df, x=independent_variables, y=dependent_variable, trendline="ols")
+
+        pio.write_html(fig, chart_path, full_html=False)
+
+        chart_url = f'{root_url}/gradio_api/file/temp/{session_hash}/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>'
+
+        results_frame = px.get_trendline_results(fig)
+
+        print("RESULTS")
+        print(results_frame)
+        print(results_frame.at[0, 'px_fit_results'])
+        results = results_frame.at[0, 'px_fit_results']
+        print(results.summary())
+
+        return {"reply": '{"regression_summary": %s, "regression_chart": %s' % (str(results.summary()), str(iframe))} 
+    
+    except Exception as e:
+      print("LINEAR REGRESSION ERROR")
+      print(e)
+      reply = f"""There was an error generating the linear regression calculation from {independent_variables} and {dependent_variable}
+              The error is {e},
+              You should probably try again.
+              """
+      return {"reply": reply}

requirements.txt

@@ -4,3 +4,4 @@ gradio
 pandas
 plotly
 openpyxl
+statsmodels

tools.py

@@ -17,7 +17,9 @@ def tools_call(session_hash):
             "type": "function",
             "function": {
                 "name": "sql_query_func",
-                "description": f"This a tool useful to query a SQLite table called 'data_source' with the following Columns: {columns}",
+                "description": f"""This a tool useful to query a SQLite table called 'data_source' with the following Columns: {columns}.
+                This function also saves the results of the query to csv file called query.csv. This is useful when query results are too large to process
+                or need to be used in an another function.""",
                 "parameters": {
                     "type": "object",
                     "properties": {
@@ -37,8 +39,9 @@ def tools_call(session_hash):
         "type": "function",
             "function": {
                 "name": "chart_generation_func",
-                "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 
+                "description": f"""This a chart generation tool useful to generate charts and graphs from queried data from our SQL table called 'data_source.
+                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 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": {
@@ -46,7 +49,40 @@ def tools_call(session_hash):
                     "properties": {
                         "data": {
                             "type": "array",
-                            "description": """The list containing a dictionary that contains the 'data' portion of the plotly chart generation. Infer this from the user's message.""",
+                            "description": """The list containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user. 
+                            Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query. 
+                            Infer this from the user's message.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        },
+                        "x_column": {
+                            "type": "string",
+                            "description": f"""The column in our query.csv file that contain the x values of the graph.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        },
+                        "y_column": {
+                            "type": "string",
+                            "description": f"""The column in our query.csv file that contain the y values of the graph.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        },
+                        "category": {
+                            "type": "string",
+                            "description": f"""An optional column in our query.csv file that contain a parameter that will define the category for the data.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        },
+                        "graph_type": {
+                            "type": "string",
+                            "description": f"""The type of plotly graph we wish to generate.
+                             This graph_type value can be one of ['bar','scatter','line','pie']. 
+                             Do not send any values outside of this list as the function will fail. 
+                             Infer this from the user's message.""",
                             "items": {
                                 "type": "string",
                             }
@@ -59,7 +95,7 @@ def tools_call(session_hash):
                             }
                         }
                     },
-                    "required": ["data"],
+                    "required": ["graph_type","x_column","y_column","layout"],
                 },
             }, 
         },
@@ -67,11 +103,10 @@ def tools_call(session_hash):
         "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() 
+                "description": f"""This an table generation tool useful to format data as a table from queried data from our SQL table called 'data_source. 
+                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. 
-                This link should open in a new window. Do not edit the string returned by the function in any way when displaying to the user, as the user needs 
-                all of the information returned by the function in it's exact state.""",
+                Do not edit the string returned by the function in any way when displaying to the user.""",
                 "parameters": {
                     "type": "object",
                     "properties": {
@@ -90,5 +125,54 @@ def tools_call(session_hash):
                     "required": ["data"],
                 },
             }, 
+        },
+        {
+        "type": "function",
+            "function": {
+                "name": "regression_func",
+                "description": f"""This a tool to calculate regressions on our SQLite table called 'data_source'.
+                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).""",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "independent_variables": {
+                            "type": "array",
+                            "description": f"""A list of strings that states the independent variables in our data set which should be column names in our query.csv file that is generated
+                            in the 'sql_query_func' function. This will allow us to identify the data to use for our independent variables.
+                            Infer this from the user's message.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        },
+                        "dependent_variable": {
+                            "type": "string",
+                            "description": f"""A string that states the dependent variables in our data set which should be a column name in our query.csv file that is generated
+                            in the 'sql_query_func' function. This will allow us to identify the data to use for our dependent variables.
+                            Infer this from the user's message.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        },
+                        "category": {
+                            "type": "string",
+                            "description": f"""An optional column in our query.csv file that contain a parameter that will define the category for the data. 
+                            Do not send value if no category is needed or specified. This category must be present in our query.csv file to be valid.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        },
+                        "regression_type": {
+                            "type": "string",
+                            "description": f"""A parameter that specifies the type of regression being used from the trendline options that plotly offers. Defaults to 'ols'.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        },
+                    },
+                    "required": ["independent_variables","dependent_variable"],
+                },
+            }, 
         }
     ]