Adds charts and graphs functionality from quickchart

functions/__init__.py

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

functions/chart_functions.py

@@ -0,0 +1,24 @@
+from typing import List
+from haystack import component
+import pandas as pd
+from quickchart import QuickChart
+import json
+import ast
+
+def chart_generation_func(queries: List[str], session_hash):
+    print("CHART GENERATION")
+    query_dict = queries[0]
+    print(query_dict)
+
+    qc = QuickChart()
+    qc.width = 1000
+    qc.height = 500
+
+    # Config can be set as a string or as a nested dict
+    qc.config = query_dict
+
+    # You can get the chart URL...
+
+    iframe = '<iframe\n    scrolling="yes"\n    width="1000px"\n    height="500px"\n    src="' + qc.get_url() + '"\n    frameborder="0"\n    allowfullscreen\n></iframe>\n'
+
+    return {"reply": iframe}

functions/chat_functions.py

@@ -23,11 +23,12 @@ messages = [
 ]
 
 def chatbot_with_fc(message, history, session_hash):
-    from functions import sqlite_query_func
+    from functions import sqlite_query_func, chart_generation_func
     from pipelines import rag_pipeline_func
     import tools
 
-    available_functions = {"sql_query_func": sqlite_query_func, "rag_pipeline_func": rag_pipeline_func}
+    available_functions = {"sql_query_func": sqlite_query_func, "rag_pipeline_func": rag_pipeline_func, "chart_generation_func": chart_generation_func}
+
     messages.append(ChatMessage.from_user(message))
     response = chat_generator.run(messages=messages, generation_kwargs={"tools": tools.tools_call(session_hash)})
 
@@ -41,9 +42,10 @@ def chatbot_with_fc(message, history, session_hash):
                 function_name = function_call.tool_name
                 function_args = function_call.arguments
 
-                ## Find the correspoding function and call it with the given 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)
+                print(function_name)
                 ## Append function response to the messages list using `ChatMessage.from_tool`
                 messages.append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
                 response = chat_generator.run(messages=messages, generation_kwargs={"tools": tools.tools_call(session_hash)})
@@ -73,7 +75,12 @@ css= ".file_marker .large{min-height:50px !important;} .example_btn{max-width:30
 
 with gr.Blocks(css=css) 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. Try a sample file to get started!</p>")
+    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!
+                          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>""")
     example_file_1 = gr.File(visible=False, value="samples/bank_marketing_campaign.csv")
     example_file_2 = gr.File(visible=False, value="samples/online_retail_data.csv")
     with gr.Row():
@@ -94,12 +101,14 @@ with gr.Blocks(css=css) 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?"]
                                     ]
             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?"]
                                     ]
             else:
                 example_questions = [
@@ -108,7 +117,7 @@ with gr.Blocks(css=css) as demo:
                                         ["What could this data be used for?"],
                                     ]
             parameters = gr.Textbox(visible=False, value=request.session_hash)
-            bot = gr.Chatbot(type='messages', label="CSV Chat Window", show_label=True, render=False, visible=True, elem_classes="chatbot")
+            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=chatbot_with_fc,
                                 type='messages',

requirements.txt

@@ -2,4 +2,5 @@ haystack-ai
 python-dotenv
 gradio
 pandas
-openpyxl
+quickchart.io
+openpyxl

tools.py

@@ -15,7 +15,7 @@ def tools_call(session_hash):
             "type": "function",
             "function": {
                 "name": "sql_query_func",
-                "description": f"This a tool useful to query a SQL 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}",
                 "parameters": {
                     "type": "object",
                     "properties": {
@@ -35,7 +35,7 @@ def tools_call(session_hash):
         "type": "function",
             "function": {
                 "name": "rag_pipeline_func",
-                "description": f"This a tool useful to query a SQL 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}",
                 "parameters": {
                     "type": "object",
                     "properties": {
@@ -50,5 +50,27 @@ def tools_call(session_hash):
                     "required": ["question"],
                 },
             }, 
+        },
+        {
+        "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.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "queries": {
+                            "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.""",
+                            "items": {
+                                "type": "string",
+                            }
+                        }
+                    },
+                    "required": ["question"],
+                },
+            }, 
         }
     ]