Added Colors

app.py

@@ -121,9 +121,11 @@ Consider these types of questions when recommending a visualization:
 5. Proportions (e.g., "What is the market share of the products?" - Pie Chart)
 6. Correlations (e.g., "Is there a correlation between marketing spend and revenue?" - Scatter Plot)
 
-Provide your response in the following format:
+Generate a JSON object. The JSON object should have the following structure:
+
 Recommended Visualization: [Chart type or "None"]. ONLY use the following names: bar, horizontal_bar, line, pie, scatter, none
 Reason: [Brief explanation for your recommendation]
+
 '''
     human_message = '''
 User question: {question}
@@ -139,19 +141,17 @@ Recommend a visualization:
                 ])
 
     final_prompt = prompt.format_prompt(question=text_query,
-                                        sql_query=sql_query, results=sql_result[:20])
+                                        sql_query=sql_query, results=sql_result)
     response = run_llm(final_prompt)
     response = response.replace('```', '')
-    lines = response.strip().split('\n')
-    print(lines)
-    visualization = lines[0].split(': ')[1]
-    reason = lines[1].split(': ')[1]
-
+    response = response.replace('json', '')
+    json_data = json.loads(response)
+    visualization = json_data['Recommended Visualization']
+    reason = json_data['Reason']
+    print(visualization, reason)
     return visualization, reason
 
 
-
-
 def format_data(text_query, sql_query, sql_result, visualization_type):
     instruction = graph_instructions[visualization_type]
 

plot_utils.py

@@ -1,13 +1,20 @@
 import matplotlib.pyplot as plt
+import numpy as np
 
 def plot_bar_chart(data):
-
     fig, ax = plt.subplots()
     labels = data['labels']
     values = data['values']
     
-    for value in values:
-        ax.bar(labels, value['data'], label=value['label'])
+    colors = ['darkviolet', 'indigo', 'blueviolet']  # Define the colors for each group
+    width = 0.2  # Width of the bars
+    x = np.arange(len(labels))  # Label locations
+    
+    for i, value in enumerate(values):
+        ax.bar(x + i * width, value['data'], width, label=value['label'], color=colors[i])
+    
+    ax.set_xticks(x + width / 2 * (len(values) - 1))  
+    ax.set_xticklabels(labels)
     
     ax.set_title('Bar Chart')
     ax.set_xlabel('Labels')
@@ -21,7 +28,7 @@ def plot_horizontal_bar_chart(data):
     values = data['values']
     
     for value in values:
-        ax.barh(labels, value['data'], label=value['label'])
+        ax.barh(labels, value['data'], label=value['label'], color=['darkviolet', 'indigo', 'blueviolet'])
     
     ax.set_title('Horizontal Bar Chart')
     ax.set_xlabel('Values')