Added Dataframe for Result

app.py

@@ -34,10 +34,7 @@ print('Loading Model...')
 tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
 
 quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type= "nf4")
+    load_in_4bit=True)
 
 model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", quantization_config=quantization_config,
                                              device_map="auto", torch_dtype=torch.bfloat16)
@@ -49,13 +46,14 @@ print('Model Loaded...')
 print(f'Model Device: {model.device}')
 
 
-def generate_output(schema, prompt, generated_sql_query, result_output):
-    return {
-        table_schema: schema,
-        input_prompt: prompt,
-        generated_query: generated_sql_query,
-        result_output: result_output,
-    }
+# def generate_output(schema, prompt, generated_sql_query, result_output, result_visulaized):
+#     return {
+#         table_schema: schema,
+#         input_prompt: prompt,
+#         generated_query: generated_sql_query,
+#         result_output: result_output,
+#         result_visulaized: result_visulaized
+#     }
 
 # Get Databases
 def get_schemas():
@@ -149,17 +147,12 @@ Recommend a visualization:
 
     final_prompt = prompt.format_prompt(question=text_query,
                                         sql_query=sql_query, results=sql_result)
-    print(final_prompt)
     response = run_llm(final_prompt)
 
-    print(f"Visualization Type: {response}")
-
     lines = response.strip().split('\n')
     visualization = lines[0].split(': ')[1]
     reason = lines[1].split(': ')[1]
 
-
-    print(visualization, reason)
     return visualization, reason
 
 def format_data(text_query, sql_query, sql_result, visualization_type):
@@ -173,20 +166,17 @@ def format_data(text_query, sql_query, sql_result, visualization_type):
 
     prompt = template.format_prompt(question=text_query, sql_query=sql_query,
                                     results=sql_result, instructions=instruction)
-    print(f'New Template for Data fromat {prompt}')
 
     formatted_data = run_llm(prompt)
     print(f'Formatted Data {formatted_data}')
-
     return json.loads(formatted_data.replace('.', '').strip())
 
-def visualize_result(text_query, visualization_type, reason, sql_query,
+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)
-        print(data)
         return plot_bar_chart(data)
 
     elif visualization_type == 'horizontal_bar':
@@ -227,36 +217,39 @@ def main(table, text_query):
         generated_sql_query = run_llm(prompt)
         print(f'Generated SQL Query: {generated_sql_query}')
     except Exception as e:
-        return generate_output(schema, prompt, '', fig)
+        return generate_output(schema, prompt, '', fig, pd.DataFrame([{"error": f"❌ Unable to generate the SQL query. {e}"}]))
 
     try:
-        sql_query_result = conn.sql(generated_sql_query).fetchall()
+        sql_query_result_raw = conn.sql(generated_sql_query)
+        sql_query_result,  = sql_query_result_raw.fetchall()
+        sql_query_df = sql_query_result_raw.df()
+        
         print(f"SQL Query Result: {sql_query_result}")
     except Exception as e:
-        return generate_output(schema, prompt, generated_sql_query, fig)
+        return generate_output(schema, prompt, generated_sql_query, fig, pd.DataFrame([{"error": f"❌ Unable to execute the SQL query. {e}"}]))
 
     try:
-      print('Generating Visualization Type')
       visualization_type, reason = get_visualization_type(text_query=text_query,
                                                   sql_query=generated_sql_query, sql_result=sql_query_result)
 
       print(f"Visualization Type: {visualization_type}")
+      print(f"Reason: {reason}")
 
     except Exception as e:
         return generate_output(schema, prompt, generated_sql_query,
-                               fig)
+                               fig, sql_query_df)
     try:
 
         plot = visualize_result(text_query=text_query, sql_query=generated_sql_query,
-                      sql_result=sql_query_result, visualization_type=visualization_type, reason=reason)
-        print(f"Plot: {plot}")
+                      sql_result=sql_query_result, visualization_type=visualization_type)
+
     except Exception as e:
-        return generate_output(schema, prompt, generated_sql_query, fig)
+        return generate_output(schema, prompt, generated_sql_query, fig, sql_query_df)
 
-    return generate_output(schema, prompt, generated_sql_query, plot)
+    return generate_output(schema, prompt, generated_sql_query, plot, sql_query_df)
 
-def generate_output(schema, prompt, generated_sql_query, result_output):
-    return [schema, prompt, generated_sql_query, result_output]
+def generate_output(schema, prompt, generated_sql_query, fig, sql_query_df):
+    return [schema, prompt, generated_sql_query, fig, sql_query_df]
 
 custom_css = """
 .gradio-container {
@@ -301,9 +294,10 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"
                     generate_query_button = gr.Button("Run Query", variant="primary")
 
             with gr.Tabs():
+                with gr.Tab("Result"):
+                    query_result_output = gr.DataFrame(label="Query Results", value=[], interactive=False)
                 with gr.Tab("Plot"):
-                    result_output = gr.Plot()
-                    # result_output = gr.Textbox(lines=TAB_LINES, label="Generated SQL Query", value="", interactive=False)
+                    result_plot = gr.Plot()
                 with gr.Tab("SQL Query"):
                     generated_query = gr.Textbox(lines=TAB_LINES, label="Generated SQL Query", value="", interactive=False)
                 with gr.Tab("Prompt"):
@@ -312,7 +306,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"
                     table_schema = gr.Textbox(lines=TAB_LINES, label="Table Schema", value="", interactive=False)
 
         schema_dropdown.change(update_tables, inputs=schema_dropdown, outputs=tables_dropdown)
-        generate_query_button.click(main, inputs=[tables_dropdown, query_input], outputs=[table_schema, input_prompt, generated_query, result_output])
+        generate_query_button.click(main, inputs=[tables_dropdown, query_input], outputs=[table_schema, input_prompt, generated_query, result_plot, query_result_output])
 
 if __name__ == "__main__":
     demo.launch(debug=True)