Changed SQL Prompt

app.py

@@ -85,14 +85,16 @@ def get_table_schema(table):
 def get_prompt(schema, query_input):
     text = f"""
     ### Instruction:
-    Your task is to generate valid duckdb SQL query to answer the following question. Only Respond with SQL query do not inlcude anything else.
+    Your task is to generate a valid DuckDB SQL query to answer the following question. Select only the columns that are necessary for visualization or that will make visualization easier (e.g., numeric or categorical data). Respond only with the SQL query and do not include anything else.
+    
     ### Input:
     Here is the database schema that the SQL query will run on:
     {schema}
-
+    
     ### Question:
     {query_input}
-    ### Response (use duckdb shorthand if possible):
+
+    ### Response (use DuckDB shorthand if possible):
     """
     return text
 
@@ -151,19 +153,52 @@ Recommend a visualization:
     print(visualization, reason)
     return visualization, reason
 
+def format_bar_data(results, question, sql_query):
+    if isinstance(results, str):
+        results = eval(results)
+
+    if len(results[0]) == 2:
+        labels = [str(row[0]) for row in results]
+        data = [float(row[1]) if row[1] is not None else 0 for row in results]
+
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", "You are a data labeling expert. Given a question, SQL Query used, and some data, provide a concise and relevant label for the data series."),
+            ("human", "Question: {question}\n SQL Query: {sql_query} \nData (first few rows): {data}\n\n Provide a concise label name for this y axis. Just give me the name."),
+        ])
+        prompt = prompt.format_prompt(question=question, data=str(results[:2]),
+                                     sql_query=sql_query)
+        label = run_llm(prompt)
+        label = label.replace('**Answer:**', '')
+        values = [{"data": data, "label": label}]
+    elif len(results[0]) == 3:
+        categories = set(row[1] for row in results)
+        labels = list(categories)
+        entities = set(row[0] for row in results)
+        values = []
+        for entity in entities:
+            entity_data = [float(row[2]) for row in results if row[0] == entity]
+            values.append({"data": entity_data, "label": str(entity)})
+    else:
+        raise ValueError("Unexpected data format in results")
+
+    formatted_data = {
+        "labels": labels,
+        "values": values
+    }
+
+    return formatted_data
 
 def format_data(text_query, sql_query, sql_result, visualization_type):
     instruction = graph_instructions[visualization_type]
 
     template = ChatPromptTemplate.from_messages([
-    ("system", "You are a Data expert who formats data according to the required needs. You are given the question asked by the user, it's sql query, the result of the query and the format you need to format it in."),
-    ("human", "For the given question: {question}\n\nSQL query: {sql_query}\n\Result: {results}\n\nUse the following example to structure the data: {instructions}. If there is None in Result please change it to '0'. Just give the json string. Do not format it."),
+    ("system", "You are a Data expert who formats data according to the required needs. You are given the question asked by the user,Given the question, SQL query, and results, your job is to Understand the question and SQL query. Ensure the result matches the query and can be easily visualized."),
+    ("human", "For the given question: {question}\n\nSQL query: {sql_query}\n\Result: {results}\n\nUse the following example to structure the data: {instructions}. If there is None in Result please change it to '0'. Just give the json string. Do not format it. Do not give backticks."),
     ])
 
 
     prompt = template.format_prompt(question=text_query, sql_query=sql_query,
                                     results=sql_result, instructions=instruction)
-    print(prompt)
     formatted_data = run_llm(prompt)
     formatted_data = formatted_data.replace('```', '')
     formatted_data = formatted_data.replace('json', '')
@@ -174,8 +209,11 @@ def visualize_result(text_query, visualization_type, sql_query,
                       sql_result):
 
     if visualization_type == 'bar':
-        data = format_data(text_query=text_query, sql_query=sql_query,
-                           sql_result=sql_result, visualization_type=visualization_type)
+        try:
+            data =  format_bar_data(sql_result, text_query, sql_query)
+        except Exception as e:
+            data =  format_data(text_query=text_query, sql_query=sql_query,
+                            sql_result=sql_result, visualization_type=visualization_type)
         return plot_bar_chart(data)
 
     elif visualization_type == 'horizontal_bar':

plot_utils.py

@@ -6,9 +6,9 @@ def plot_bar_chart(data):
     labels = data['labels']
     values = data['values']
     
-    colors = ['darkviolet', 'indigo', 'blueviolet']  # Define the colors for each group
-    width = 0.2  # Width of the bars
-    x = np.arange(len(labels))  # Label locations
+    colors = ['darkviolet', 'indigo', 'blueviolet']
+    width = 0.2  
+    x = np.arange(len(labels)) 
     
     for i, value in enumerate(values):
         ax.bar(x + i * width, value['data'], width, label=value['label'], color=colors[i])