Update app.py

app.py

@@ -1,247 +1,180 @@
 import gradio as gr
 import pandas as pd
-import matplotlib.pyplot as plt
-import matplotlib
-matplotlib.use('Agg')  # Use non-interactive backend
 import os
 import tempfile
-import base64
-from io import BytesIO
+import matplotlib.pyplot as plt
 from pandasai import SmartDataframe
 from langchain_groq.chat_models import ChatGroq
+from dotenv import load_dotenv
+import io
+import base64
 
-# HARDCODED API KEY - REPLACE WITH YOUR ACTUAL KEY
-API_KEY = "gsk_YOUR_ACTUAL_API_KEY_HERE"  # Replace with your real API key
-
-# Global variables to store data
-current_df = None
-llm = None
+# Load environment variables
+load_dotenv()
 
-def initialize_llm():
-    """Initialize the Groq LLM"""
-    global llm
+def process_data(file, query, api_key):
+    """
+    Process the uploaded CSV file with the given query using PandasAI
+    """
     try:
-        if API_KEY == "gsk_YOUR_ACTUAL_API_KEY_HERE":
-            return "❌ Please replace 'gsk_YOUR_ACTUAL_API_KEY_HERE' with your actual Groq API key", None
+        # Validate inputs
+        if file is None:
+            return "Please upload a CSV file.", None
+        
+        if not query.strip():
+            return "Please enter a query.", None
+            
+        if not api_key.strip():
+            return "Please enter your Groq API key.", None
+        
+        # Read the CSV file
+        df_data = pd.read_csv(file.name)
         
+        # Initialize Groq LLM
         llm = ChatGroq(
-            model_name="mixtral-8x7b-32768",
-            api_key=API_KEY,
+            model_name="mixtral-8x7b-32768",  # Using a more stable model
+            api_key=api_key.strip(),
             temperature=0
         )
-        return "✅ Groq LLM initialized successfully", llm
-    except Exception as e:
-        return f"❌ Failed to initialize Groq LLM: {str(e)}", None
-
-def process_csv(file):
-    """Process uploaded CSV file"""
-    global current_df
-    
-    if file is None:
-        return "No file uploaded", None, None
-    
-    try:
-        # Read the CSV file
-        current_df = pd.read_csv(file.name)
         
-        # Create preview
-        preview = current_df.head().to_html(classes='table table-striped', table_id='data-preview')
+        # Create SmartDataframe
+        df = SmartDataframe(df_data, config={
+            "llm": llm,
+            "save_charts": True,
+            "save_charts_path": tempfile.gettempdir(),
+            "open_charts": False,
+            "enable_cache": False
+        })
         
-        # Create info
-        info = f"""
-        **File Info:**
-        - Shape: {current_df.shape[0]} rows × {current_df.shape[1]} columns
-        - Columns: {', '.join(current_df.columns.tolist())}
-        """
+        # Process the query
+        result = df.chat(query)
         
-        return "✅ CSV file loaded successfully", preview, info
+        # Handle different types of results
+        if result is None:
+            return "No result returned. Please try a different query.", None
         
-    except Exception as e:
-        return f"❌ Error reading CSV: {str(e)}", None, None
-
-def chat_with_data(query):
-    """Process user query and return response"""
-    global current_df, llm
-    
-    if current_df is None:
-        return "❌ Please upload a CSV file first", None
-    
-    if llm is None:
-        status, _ = initialize_llm()
-        if llm is None:
-            return status, None
-    
-    if not query.strip():
-        return "❌ Please enter a query", None
-    
-    try:
-        # Create temporary directory for charts
-        temp_dir = tempfile.mkdtemp()
+        # Check if result is a plot/chart
+        chart_path = None
+        chart_files = [f for f in os.listdir(tempfile.gettempdir()) if f.endswith(('.png', '.jpg', '.jpeg'))]
         
-        # Create SmartDataframe
-        sdf = SmartDataframe(
-            current_df, 
-            config={
-                "llm": llm,
-                "verbose": True,
-                "save_charts": True,
-                "save_charts_path": temp_dir,
-                "custom_whitelisted_dependencies": ["matplotlib", "seaborn", "plotly"]
-            }
-        )
+        if chart_files:
+            # Get the most recent chart file
+            chart_files.sort(key=lambda x: os.path.getmtime(os.path.join(tempfile.gettempdir(), x)), reverse=True)
+            chart_path = os.path.join(tempfile.gettempdir(), chart_files[0])
         
-        # Process the query
-        result = sdf.chat(query)
+        # Format the text result
+        if isinstance(result, pd.DataFrame):
+            result_text = f"Query Result:\n\n{result.to_string()}"
+        elif isinstance(result, (int, float)):
+            result_text = f"Query Result: {result}"
+        elif isinstance(result, str):
+            result_text = f"Query Result:\n{result}"
+        else:
+            result_text = f"Query Result:\n{str(result)}"
         
-        # Handle different types of results
-        if isinstance(result, str):
-            # Text response
-            return f"📢 **Response:**\n{result}", None
+        return result_text, chart_path
+        
+    except Exception as e:
+        error_msg = f"Error processing query: {str(e)}"
+        return error_msg, None
+
+def create_interface():
+    """
+    Create the Gradio interface
+    """
+    with gr.Blocks(title="PandasAI with Groq", theme=gr.themes.Soft()) as demo:
+        gr.Markdown(
+            """
+            # 📊 PandasAI Data Analysis with Groq
+            
+            Upload a CSV file and ask questions about your data. The AI will analyze and visualize your data accordingly.
             
-        elif hasattr(result, 'savefig'):
-            # Matplotlib figure
-            try:
-                # Save figure to bytes
-                img_buffer = BytesIO()
-                result.savefig(img_buffer, format='png', dpi=150, bbox_inches='tight')
-                img_buffer.seek(0)
+            **Instructions:**
+            1. Get your Groq API key from [https://console.groq.com/keys](https://console.groq.com/keys)
+            2. Upload your CSV file
+            3. Enter your query (e.g., "Show top 5 countries by population", "Create a bar plot of sales by region")
+            4. Click Submit to get results
+            """
+        )
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Input components
+                api_key_input = gr.Textbox(
+                    label="Groq API Key",
+                    placeholder="Enter your Groq API key here...",
+                    type="password",
+                    info="Your API key is not stored and only used for this session"
+                )
                 
-                # Save to temporary file for Gradio
-                temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
-                with open(temp_file.name, 'wb') as f:
-                    f.write(img_buffer.getvalue())
+                file_input = gr.File(
+                    label="Upload CSV File",
+                    file_types=[".csv"],
+                    info="Upload your CSV data file"
+                )
                 
-                plt.close(result)  # Close the figure to free memory
+                query_input = gr.Textbox(
+                    label="Your Query",
+                    placeholder="e.g., 'Which are the top 5 countries by population?' or 'Create a bar plot of the top 5 countries'",
+                    lines=3,
+                    info="Ask questions about your data or request visualizations"
+                )
                 
-                return "📈 **Chart Generated:**", temp_file.name
+                submit_btn = gr.Button("🚀 Submit Query", variant="primary")
                 
-            except Exception as chart_error:
-                return f"❌ Error saving chart: {str(chart_error)}", None
+            with gr.Column(scale=2):
+                # Output components
+                result_output = gr.Textbox(
+                    label="Analysis Result",
+                    lines=10,
+                    interactive=False,
+                    show_copy_button=True
+                )
                 
-        elif isinstance(result, pd.DataFrame):
-            # DataFrame result
-            html_table = result.to_html(classes='table table-striped', max_rows=100)
-            return f"📊 **Data Result:**\n{html_table}", None
-            
-        else:
-            # Other types of results
-            return f"📊 **Result:**\n{str(result)}", None
-            
-    except Exception as e:
-        error_msg = f"❌ Error: {str(e)}"
+                chart_output = gr.Image(
+                    label="Generated Visualization",
+                    show_label=True
+                )
         
-        # Provide specific error guidance
-        if "403" in str(e):
-            error_msg += "\n\n🔐 **403 Forbidden Error** - This usually means:\n"
-            error_msg += "- Invalid API key\n"
-            error_msg += "- API key doesn't have permission for this model\n"
-            error_msg += "- Rate limit exceeded\n"
-            error_msg += "- Model name is incorrect"
-        elif "rate limit" in str(e).lower():
-            error_msg += "\n\n⏰ **Rate Limit** - Please wait a moment before trying again"
-        elif "timeout" in str(e).lower():
-            error_msg += "\n\n⏱️ **Timeout** - The query took too long. Try a simpler request"
-            
-        return error_msg, None
-
-def get_debug_info():
-    """Get debug information"""
-    if API_KEY and API_KEY != "gsk_YOUR_ACTUAL_API_KEY_HERE":
-        return f"✅ API Key loaded successfully\nKey starts with: {API_KEY[:10]}..."
-    else:
-        return "❌ Replace 'gsk_YOUR_ACTUAL_API_KEY_HERE' with your actual API key"
-
-# Initialize LLM on startup
-init_status, _ = initialize_llm()
-
-# Create Gradio interface
-with gr.Blocks(title="📊 CSV Chat with Groq + PandasAI", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 📊 Chat with Your CSV using PandasAI + Groq")
-    
-    with gr.Row():
-        with gr.Column(scale=2):
-            # File upload section
-            gr.Markdown("## 📁 Upload CSV File")
-            file_input = gr.File(
-                label="Upload your CSV file",
-                file_types=[".csv"],
-                type="filepath"
-            )
-            
-            upload_status = gr.Textbox(
-                label="Upload Status",
-                interactive=False,
-                value=init_status
-            )
-            
-            # Data preview section
-            gr.Markdown("## 📋 Data Preview")
-            data_preview = gr.HTML(label="Data Preview")
-            data_info = gr.Markdown()
-            
-        with gr.Column(scale=1):
-            # Debug and help section
-            gr.Markdown("## 🔧 Debug Info")
-            debug_btn = gr.Button("Show Debug Info")
-            debug_info = gr.Textbox(label="Debug Information", interactive=False)
-            
-            gr.Markdown("## 📝 Example Queries")
-            gr.Markdown("""
-            - "Show me the first 10 rows"
-            - "What are the column names?"
-            - "Create a histogram of [column_name]"
-            - "Show me the summary statistics"
-            - "Plot the top 5 values in [column_name]"
-            - "Create a bar chart showing [column1] vs [column2]"
-            """)
-    
-    # Chat section
-    gr.Markdown("## 💬 Chat with Your Data")
-    
-    with gr.Row():
-        query_input = gr.Textbox(
-            label="Ask a question or request a chart",
-            placeholder="What would you like to know about your data?",
-            lines=3,
-            scale=4
+        # Example queries
+        gr.Markdown(
+            """
+            ### 💡 Example Queries:
+            - "Which are the top 5 countries by population?"
+            - "Create a bar plot of the top 10 countries by population"
+            - "Show me a pie chart of the top 5 countries"
+            - "Calculate the total population of the top 3 countries"
+            - "What is the average population across all countries?"
+            - "Create a scatter plot showing the relationship between two columns"
+            """
+        )
+        
+        # Event handlers
+        submit_btn.click(
+            fn=process_data,
+            inputs=[file_input, query_input, api_key_input],
+            outputs=[result_output, chart_output],
+            show_progress=True
+        )
+        
+        # Allow Enter key to submit
+        query_input.submit(
+            fn=process_data,
+            inputs=[file_input, query_input, api_key_input],
+            outputs=[result_output, chart_output],
+            show_progress=True
         )
-        submit_btn = gr.Button("Submit Query", variant="primary", scale=1)
-    
-    # Results section
-    with gr.Row():
-        with gr.Column():
-            response_output = gr.Markdown(label="Response")
-        with gr.Column():
-            chart_output = gr.Image(label="Generated Chart", type="filepath")
-    
-    # Event handlers
-    file_input.change(
-        fn=process_csv,
-        inputs=[file_input],
-        outputs=[upload_status, data_preview, data_info]
-    )
-    
-    debug_btn.click(
-        fn=get_debug_info,
-        outputs=[debug_info]
-    )
-    
-    submit_btn.click(
-        fn=chat_with_data,
-        inputs=[query_input],
-        outputs=[response_output, chart_output]
-    )
     
-    query_input.submit(
-        fn=chat_with_data,
-        inputs=[query_input],
-        outputs=[response_output, chart_output]
-    )
+    return demo
 
-# Launch the app
 if __name__ == "__main__":
+    # Create and launch the interface
+    demo = create_interface()
     demo.launch(
-        share=False,  # Set to True if you want a public link
-        debug=True,
-        show_error=True
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True,
+        show_tips=True,
+        enable_queue=True
     )