Error Fix

app.py

@@ -9,6 +9,7 @@ from langchain_core.prompts import ChatPromptTemplate
 from langchain_groq import ChatGroq
 import plotly.express as px
 import time
+import plotly.io as pio
 
 # Load environment variables
 load_dotenv()
@@ -43,6 +44,9 @@ current_context = {
     "table_name": None
 }
 
+# Add a global variable to store the current plot
+current_plot = None
+
 # Define the prompt with examples for SQL query generation
 query_prompt = ChatPromptTemplate.from_messages(
     [
@@ -122,8 +126,11 @@ interpret_prompt = ChatPromptTemplate.from_messages(
 
 def process_text_query(query, history):
     """Process a text query and update chat history"""
+    global current_plot
+    current_plot = None  # Reset the plot
+    
     if not query:
-        return "", history
+        return "", history, None
     
     # Add the user's query to history
     history.append({"role": "user", "content": query})
@@ -194,6 +201,7 @@ def process_text_query(query, history):
                 # Add visualization if requested
                 if is_visualization and not result_df.empty:
                     try:
+                        print("Visualization requested, attempting to create plot...")
                         # Determine the type of visualization based on the data
                         if len(result_df.columns) >= 2:
                             # Find numeric columns for y-axis
@@ -202,19 +210,25 @@ def process_text_query(query, history):
                             if len(numeric_cols) >= 1 and len(result_df) > 1:
                                 # Use the first column as x and first numeric column as y
                                 x_col = result_df.columns[0]
-                                y_col = numeric_cols[0]
+                                y_cols = numeric_cols[:3]  # Use up to 3 numeric columns
                                 
                                 # Create appropriate plot based on data characteristics
                                 if 'month' in result_df.columns or 'date' in result_df.columns or 'year' in result_df.columns:
                                     # Time series data - use line chart
-                                    fig = px.line(result_df, x=x_col, y=numeric_cols, title="Time Series Analysis")
+                                    fig = px.line(result_df, x=x_col, y=y_cols, title="Time Series Analysis")
                                 else:
                                     # Regular data - use bar chart
-                                    fig = px.bar(result_df, x=x_col, y=y_col, title="Data Visualization")
+                                    fig = px.bar(result_df, x=x_col, y=y_cols[0], title="Data Visualization")
                                 
-                                # Convert to HTML and add to response
-                                plot_html = fig.to_html(full_html=False, include_plotlyjs='cdn')
-                                response += f"\n\n**Visualization:**\n<div>{plot_html}</div>"
+                                # Store the figure for display
+                                current_plot = fig
+                                
+                                # Add note about visualization
+                                response += "\n\n**A visualization has been generated and is displayed below.**"
+                        
+                        # After creating the plot
+                        print(f"Plot created: {current_plot is not None}")
+                        
                     except Exception as viz_error:
                         print(f"Visualization error: {str(viz_error)}")
                         # Continue without visualization if there's an error
@@ -241,7 +255,7 @@ def process_text_query(query, history):
     # Add the response to history
     history.append({"role": "assistant", "content": response})
     
-    return "", history
+    return "", history, current_plot
 
 def process_file_upload(files):
     """Process uploaded files and index them"""
@@ -411,6 +425,9 @@ with gr.Blocks(title="AI Document Analysis & Voice Assistant") as demo:
     with gr.Tab("Chat"):
         chatbot = gr.Chatbot(height=400, type="messages")
         
+        # Add a plot component
+        plot_output = gr.Plot(label="Data Visualization", visible=False)
+        
         with gr.Row():
             with gr.Column(scale=8):
                 msg = gr.Textbox(
@@ -438,16 +455,18 @@ with gr.Blocks(title="AI Document Analysis & Voice Assistant") as demo:
         submit_btn.click(
             process_text_query, 
             inputs=[msg, chatbot], 
-            outputs=[msg, chatbot]
+            outputs=[msg, chatbot, plot_output],
+            postprocess=lambda _, __, fig: gr.update(value=fig, visible=fig is not None)
         )
         
         msg.submit(
             process_text_query, 
             inputs=[msg, chatbot], 
-            outputs=[msg, chatbot]
+            outputs=[msg, chatbot, plot_output],
+            postprocess=lambda _, __, fig: gr.update(value=fig, visible=fig is not None)
         )
         
-        clear_btn.click(lambda: None, None, chatbot, queue=False)
+        clear_btn.click(lambda: [None, None, gr.update(visible=False)], None, [chatbot, plot_output, plot_output], queue=False)
         clear_context_btn.click(clear_context, inputs=[], outputs=[chatbot])
         
         voice_btn.click(