import gradio as gr import pandas as pd import os import tempfile 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 import re # Load environment variables # load_dotenv() # # Hardcoded API key - Replace with your actual Groq API key # GROQ_API_KEY = "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" GROQ_API_KEY = os.getenv("GROQ_API_KEY") # Global variables to store data current_dataframe = None current_smart_df = None last_query_result = None def analyze_chart_feasibility(query, df_data): """ Analyze if the query can generate a meaningful chart """ query_lower = query.lower() # Chart-related keywords chart_keywords = [ 'plot', 'chart', 'graph', 'visualize', 'visualization', 'bar', 'line', 'pie', 'scatter', 'histogram', 'heatmap', 'boxplot', 'distribution' ] # Statistical keywords that might benefit from visualization stat_keywords = [ 'top', 'bottom', 'highest', 'lowest', 'compare', 'comparison', 'trend', 'relationship', 'correlation', 'by category', 'group by' ] # Check if query explicitly asks for a chart explicit_chart = any(keyword in query_lower for keyword in chart_keywords) # Check if query has statistical nature that could be visualized statistical_nature = any(keyword in query_lower for keyword in stat_keywords) # Check data characteristics numeric_columns = df_data.select_dtypes(include=['number']).columns.tolist() categorical_columns = df_data.select_dtypes(include=['object', 'category']).columns.tolist() can_create_chart = False chart_recommendation = "" reasoning = "" if explicit_chart: can_create_chart = True reasoning = "Query explicitly requests a chart/visualization." chart_recommendation = "Chart will be generated as requested." elif statistical_nature and len(numeric_columns) > 0: can_create_chart = True reasoning = f"Query involves statistical analysis with {len(numeric_columns)} numeric columns available for visualization." # Suggest appropriate chart types if 'top' in query_lower or 'bottom' in query_lower: chart_recommendation = "Recommended: Bar chart to show rankings/comparisons." elif 'relationship' in query_lower or 'correlation' in query_lower: chart_recommendation = "Recommended: Scatter plot to show relationships." elif 'distribution' in query_lower: chart_recommendation = "Recommended: Histogram or box plot for distribution analysis." else: chart_recommendation = "Recommended: Bar chart or line chart based on data nature." else: reasoning = "Query appears to be asking for specific values, calculations, or text-based information that doesn't require visualization." chart_recommendation = "Chart generation not recommended for this type of query." return can_create_chart, reasoning, chart_recommendation def process_query_only(file, query): """ Process the query without generating charts """ global current_dataframe, current_smart_df, last_query_result try: # Validate inputs if file is None: return "Please upload a CSV file.", "", "" if not query.strip(): return "Please enter a query.", "", "" # Read the CSV file if not already loaded or if file changed if current_dataframe is None: current_dataframe = pd.read_csv(file.name) # Initialize Groq LLM llm = ChatGroq( model_name="llama-3.3-70b-versatile", api_key=GROQ_API_KEY, temperature=0 ) # Create SmartDataframe current_smart_df = SmartDataframe(current_dataframe, config={ "llm": llm, "save_charts": False, # Disabled for query-only mode "enable_cache": False }) # Analyze chart feasibility can_chart, reasoning, recommendation = analyze_chart_feasibility(query, current_dataframe) # Process the query result = current_smart_df.chat(query) last_query_result = result # Handle different types of results if result is None: return "No result returned. Please try a different query.", reasoning, recommendation # 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)}" return result_text, reasoning, recommendation except Exception as e: error_msg = f"Error processing query: {str(e)}" return error_msg, "", "" def generate_chart(query): """ Generate chart based on the query and last result """ global current_dataframe, current_smart_df, last_query_result try: if current_smart_df is None: return "Please run a query first before generating charts.", None if not query.strip(): return "Please enter a query for chart generation.", None # Clean up old chart files chart_files = [f for f in os.listdir(tempfile.gettempdir()) if f.endswith(('.png', '.jpg', '.jpeg'))] for file in chart_files: try: os.remove(os.path.join(tempfile.gettempdir(), file)) except: pass # Create a chart-focused version of the query chart_query = query if not any(keyword in query.lower() for keyword in ['plot', 'chart', 'graph', 'visualize']): # Add visualization instruction to the query chart_query = f"Create a chart or visualization for: {query}" # Reconfigure SmartDataframe for chart generation llm = ChatGroq( model_name="llama-3.3-70b-versatile", api_key=GROQ_API_KEY, temperature=0 ) chart_smart_df = SmartDataframe(current_dataframe, config={ "llm": llm, "save_charts": True, "save_charts_path": tempfile.gettempdir(), "open_charts": False, "enable_cache": False }) # Generate chart result = chart_smart_df.chat(chart_query) # Look for generated chart chart_path = None chart_files = [f for f in os.listdir(tempfile.gettempdir()) if f.endswith(('.png', '.jpg', '.jpeg'))] 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]) return "Chart generated successfully!", chart_path else: return "Chart could not be generated. The query might not be suitable for visualization or there might be an issue with the data.", None except Exception as e: error_msg = f"Error generating chart: {str(e)}" return error_msg, None def reset_data(): """ Reset the loaded data to allow new file upload """ global current_dataframe, current_smart_df, last_query_result current_dataframe = None current_smart_df = None last_query_result = None return "Data reset. Please upload a new file.", "", "", None, None def create_interface(): """ Create the Gradio interface """ with gr.Blocks(title="Enhanced PandasAI with Groq", theme=gr.themes.Soft()) as demo: gr.Markdown( """ # 📊 Enhanced PandasAI Data Analysis with Groq Upload a CSV file and analyze your data with separate query and chart generation capabilities. **Instructions:** 1. Upload your CSV file 2. Enter your query and click "Analyze Query" to get text results and chart feasibility analysis 3. If chart is recommended, click "Generate Chart" to create visualizations 4. Use "Reset Data" to load a new file """ ) with gr.Row(): with gr.Column(scale=1): # Input components file_input = gr.File( label="Upload CSV File", file_types=[".csv"] ) query_input = gr.Textbox( label="Your Query", placeholder="e.g., 'Which are the top 5 countries by population?' or 'Show relationship between two columns'", lines=3 ) with gr.Row(): analyze_btn = gr.Button("🔍 Analyze Query", variant="primary") chart_btn = gr.Button("📊 Generate Chart", variant="secondary") reset_btn = gr.Button("🔄 Reset Data", variant="stop") with gr.Column(scale=2): # Output components result_output = gr.Textbox( label="Analysis Result", lines=8, interactive=False ) with gr.Row(): with gr.Column(): feasibility_output = gr.Textbox( label="Chart Feasibility Analysis", lines=3, interactive=False ) with gr.Column(): recommendation_output = gr.Textbox( label="Chart Recommendation", lines=3, interactive=False ) chart_status = gr.Textbox( label="Chart Generation Status", lines=2, interactive=False ) chart_output = gr.Image( label="Generated Visualization" ) # Example section gr.Markdown( """ ### 💡 Example Workflow: **Step 1 - Data Analysis Queries:** - "What are the top 10 countries by population?" - "Calculate the average population of all countries" - "Which country has the highest GDP?" **Step 2 - Chart Generation:** - After running a query, click "Generate Chart" to visualize the results - The system will analyze if your query can be effectively visualized - Charts work best with comparative, ranking, or relationship-based queries **Query Types that work well for charts:** - Ranking queries (top/bottom N items) - Comparisons between categories - Relationships between variables - Distribution analysis """ ) # Event handlers analyze_btn.click( fn=process_query_only, inputs=[file_input, query_input], outputs=[result_output, feasibility_output, recommendation_output] ) chart_btn.click( fn=generate_chart, inputs=[query_input], outputs=[chart_status, chart_output] ) reset_btn.click( fn=reset_data, outputs=[chart_status, feasibility_output, recommendation_output, chart_output, result_output] ) # Allow Enter key to analyze query query_input.submit( fn=process_query_only, inputs=[file_input, query_input], outputs=[result_output, feasibility_output, recommendation_output] ) return demo if __name__ == "__main__": # Create and launch the interface demo = create_interface() demo.launch( server_name="0.0.0.0", server_port=7860, share=False )